1717
1818use crate :: planner:: { ContextProvider , PlannerContext , SqlToRel } ;
1919use datafusion_common:: {
20- DataFusionError , Diagnostic , Result , Span , not_impl_err, plan_err,
20+ DFSchemaRef , DataFusionError , Diagnostic , Result , Span , not_impl_err, plan_err,
2121} ;
2222use datafusion_expr:: { LogicalPlan , LogicalPlanBuilder } ;
23- use sqlparser:: ast:: { SetExpr , SetOperator , SetQuantifier , Spanned } ;
23+ use sqlparser:: ast:: {
24+ Expr as SQLExpr , Ident , SelectItem , SetExpr , SetOperator , SetQuantifier , Spanned ,
25+ } ;
2426
2527impl < S : ContextProvider > SqlToRel < ' _ , S > {
2628 #[ cfg_attr( feature = "recursive_protection" , recursive:: recursive) ]
@@ -42,7 +44,28 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
4244 let left_span = Span :: try_from_sqlparser_span ( left. span ( ) ) ;
4345 let right_span = Span :: try_from_sqlparser_span ( right. span ( ) ) ;
4446 let left_plan = self . set_expr_to_plan ( * left, planner_context) ;
45- let right_plan = self . set_expr_to_plan ( * right, planner_context) ;
47+
48+ // For non-*ByName operations, add missing aliases to right side using left schema's
49+ // column names. This allows queries like
50+ // `SELECT 1 c1, 0 c2, 0 c3 UNION ALL SELECT 2, 0, 0`
51+ // where the right side has duplicate literal values.
52+ // We only do this if the left side succeeded.
53+ let right = if let Ok ( plan) = & left_plan
54+ && plan. schema ( ) . fields ( ) . len ( ) > 1
55+ && matches ! (
56+ set_quantifier,
57+ SetQuantifier :: All
58+ | SetQuantifier :: Distinct
59+ | SetQuantifier :: None
60+ ) {
61+ alias_set_expr ( * right, plan. schema ( ) )
62+ } else {
63+ * right
64+ } ;
65+
66+ let right_plan = self . set_expr_to_plan ( right, planner_context) ;
67+
68+ // Handle errors from both sides, collecting them if both failed
4669 let ( left_plan, right_plan) = match ( left_plan, right_plan) {
4770 ( Ok ( left_plan) , Ok ( right_plan) ) => ( left_plan, right_plan) ,
4871 ( Err ( left_err) , Err ( right_err) ) => {
@@ -160,3 +183,75 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
160183 }
161184 }
162185}
186+
187+ // Adds aliases to SELECT items in a SetExpr using the provided schema.
188+ // This ensures that unnamed expressions on the right side of a UNION/INTERSECT/EXCEPT
189+ // get aliased with the column names from the left side, allowing queries like
190+ // `SELECT 1 AS a, 0 AS b, 0 AS c UNION ALL SELECT 2, 0, 0` to work correctly.
191+ fn alias_set_expr ( set_expr : SetExpr , schema : & DFSchemaRef ) -> SetExpr {
192+ match set_expr {
193+ SetExpr :: Select ( mut select) => {
194+ alias_select_items ( & mut select. projection , schema) ;
195+ SetExpr :: Select ( select)
196+ }
197+ SetExpr :: SetOperation {
198+ op,
199+ left,
200+ right,
201+ set_quantifier,
202+ } => {
203+ // For nested set operations, only alias the leftmost branch
204+ // since that's what determines the output column names
205+ SetExpr :: SetOperation {
206+ op,
207+ left : Box :: new ( alias_set_expr ( * left, schema) ) ,
208+ right,
209+ set_quantifier,
210+ }
211+ }
212+ SetExpr :: Query ( mut query) => {
213+ // Handle parenthesized queries like (SELECT ... UNION ALL SELECT ...)
214+ query. body = Box :: new ( alias_set_expr ( * query. body , schema) ) ;
215+ SetExpr :: Query ( query)
216+ }
217+ // For other cases (Values, etc.), return as-is
218+ other => other,
219+ }
220+ }
221+
222+ // Adds aliases to literal value expressions where missing, based on the input schema, as these are
223+ // the ones that can cause duplicate name issues (e.g. `SELECT 0, 0` has two columns named `Int64(0)`).
224+ // Other expressions typically have unique names.
225+ fn alias_select_items ( items : & mut [ SelectItem ] , schema : & DFSchemaRef ) {
226+ let mut col_idx = 0 ;
227+ for item in items. iter_mut ( ) {
228+ match item {
229+ SelectItem :: UnnamedExpr ( expr) if is_literal_value ( expr) => {
230+ if let Some ( field) = schema. fields ( ) . get ( col_idx) {
231+ * item = SelectItem :: ExprWithAlias {
232+ expr : expr. clone ( ) ,
233+ alias : Ident :: new ( field. name ( ) ) ,
234+ } ;
235+ }
236+ col_idx += 1 ;
237+ }
238+ SelectItem :: UnnamedExpr ( _) | SelectItem :: ExprWithAlias { .. } => {
239+ col_idx += 1 ;
240+ }
241+ SelectItem :: Wildcard ( _) | SelectItem :: QualifiedWildcard ( _, _) => {
242+ // Wildcards expand to multiple columns - skip position tracking
243+ }
244+ }
245+ }
246+ }
247+
248+ /// Returns true if the expression is a literal value that could cause duplicate names.
249+ fn is_literal_value ( expr : & SQLExpr ) -> bool {
250+ matches ! (
251+ expr,
252+ SQLExpr :: Value ( _)
253+ | SQLExpr :: UnaryOp { .. }
254+ | SQLExpr :: TypedString { .. }
255+ | SQLExpr :: Interval ( _)
256+ )
257+ }
0 commit comments