fix: handle case of non-literal expressions with same name

Author: gruuya

datafusion/sql/src/set_expr.rs

@@ -199,23 +199,27 @@ fn alias_set_expr(set_expr: &mut SetExpr, schema: &DFSchemaRef) {
     }
 }
 
-// Adds aliases to literal value expressions where missing, based on the input schema, as these are
-// the ones that can cause duplicate name issues (e.g. `SELECT 0, 0` has two columns named `Int64(0)`).
-// Other expressions typically have unique names.
+// Aliases unnamed expressions in the provided select items using the provided schema.
+// This helps with set expression queries where the right side has duplicate expressions,
+// but the left side has unique column names, which control the output schema anyway.
 fn alias_select_items(items: &mut [SelectItem], schema: &DFSchemaRef) {
     let mut col_idx = 0;
     for item in items.iter_mut() {
         match item {
-            SelectItem::UnnamedExpr(expr) if is_literal_value(expr) => {
-                if let Some(field) = schema.fields().get(col_idx) {
+            SelectItem::UnnamedExpr(expr) => {
+                if !matches!(
+                    expr,
+                    SQLExpr::Identifier(_) | SQLExpr::CompoundIdentifier(_)
+                ) && let Some(field) = schema.fields().get(col_idx)
+                {
                     *item = SelectItem::ExprWithAlias {
                         expr: expr.clone(),
                         alias: Ident::new(field.name()),
                     };
                 }
                 col_idx += 1;
             }
-            SelectItem::UnnamedExpr(_) | SelectItem::ExprWithAlias { .. } => {
+            SelectItem::ExprWithAlias { .. } => {
                 col_idx += 1;
             }
             SelectItem::Wildcard(_) | SelectItem::QualifiedWildcard(_, _) => {
@@ -224,14 +228,3 @@ fn alias_select_items(items: &mut [SelectItem], schema: &DFSchemaRef) {
         }
     }
 }
-
-/// Returns true if the expression is a literal value that could cause duplicate names.
-fn is_literal_value(expr: &SQLExpr) -> bool {
-    matches!(
-        expr,
-        SQLExpr::Value(_)
-            | SQLExpr::UnaryOp { .. }
-            | SQLExpr::TypedString { .. }
-            | SQLExpr::Interval(_)
-    )
-}

datafusion/sql/tests/sql_integration.rs

@@ -2656,55 +2656,82 @@ fn union_all_by_name_same_column_names() {
 }
 
 #[test]
-fn union_all_with_duplicate_literals() {
-    let sql = "SELECT 0 a, 0 b UNION ALL SELECT 1, 1";
+fn union_all_with_duplicate_expressions() {
+    let sql = "\
+        SELECT 0 a, 0 b \
+        UNION ALL SELECT 1, 1 \
+        UNION ALL SELECT count(*), count(*) FROM orders";
     let plan = logical_plan(sql).unwrap();
     assert_snapshot!(
         plan,
         @r"
     Union
-      Projection: Int64(0) AS a, Int64(0) AS b
-        EmptyRelation: rows=1
-      Projection: Int64(1) AS a, Int64(1) AS b
-        EmptyRelation: rows=1
+      Union
+        Projection: Int64(0) AS a, Int64(0) AS b
+          EmptyRelation: rows=1
+        Projection: Int64(1) AS a, Int64(1) AS b
+          EmptyRelation: rows=1
+      Projection: count(*) AS a, count(*) AS b
+        Aggregate: groupBy=[[]], aggr=[[count(*)]]
+          TableScan: orders
     "
     );
 }
 
 #[test]
-fn intersect_with_duplicate_literals() {
-    let sql = "SELECT 1 as a, 0 as b, 0 as c INTERSECT SELECT 1, 0, 0";
+fn intersect_with_duplicate_expressions() {
+    let sql = "\
+        SELECT 0 a, 0 b \
+        INTERSECT SELECT 1, 1 \
+        INTERSECT SELECT count(*), count(*) FROM orders";
     let plan = logical_plan(sql).unwrap();
     assert_snapshot!(
         plan,
         @r"
-    LeftSemi Join: left.a = right.a, left.b = right.b, left.c = right.c
+    LeftSemi Join: left.a = right.a, left.b = right.b
       Distinct:
         SubqueryAlias: left
-          Projection: Int64(1) AS a, Int64(0) AS b, Int64(0) AS c
-            EmptyRelation: rows=1
+          LeftSemi Join: left.a = right.a, left.b = right.b
+            Distinct:
+              SubqueryAlias: left
+                Projection: Int64(0) AS a, Int64(0) AS b
+                  EmptyRelation: rows=1
+            SubqueryAlias: right
+              Projection: Int64(1) AS a, Int64(1) AS b
+                EmptyRelation: rows=1
       SubqueryAlias: right
-        Projection: Int64(1) AS a, Int64(0) AS b, Int64(0) AS c
-          EmptyRelation: rows=1
+        Projection: count(*) AS a, count(*) AS b
+          Aggregate: groupBy=[[]], aggr=[[count(*)]]
+            TableScan: orders
     "
     );
 }
 
 #[test]
-fn except_with_duplicate_literals() {
-    let sql = "SELECT 1 as a, 0 as b, 0 as c EXCEPT SELECT 2, 0, 0";
+fn except_with_duplicate_expressions() {
+    let sql = "\
+        SELECT 0 a, 0 b \
+        EXCEPT SELECT 1, 1 \
+        EXCEPT SELECT count(*), count(*) FROM orders";
     let plan = logical_plan(sql).unwrap();
     assert_snapshot!(
         plan,
         @r"
-    LeftAnti Join: left.a = right.a, left.b = right.b, left.c = right.c
+    LeftAnti Join: left.a = right.a, left.b = right.b
       Distinct:
         SubqueryAlias: left
-          Projection: Int64(1) AS a, Int64(0) AS b, Int64(0) AS c
-            EmptyRelation: rows=1
+          LeftAnti Join: left.a = right.a, left.b = right.b
+            Distinct:
+              SubqueryAlias: left
+                Projection: Int64(0) AS a, Int64(0) AS b
+                  EmptyRelation: rows=1
+            SubqueryAlias: right
+              Projection: Int64(1) AS a, Int64(1) AS b
+                EmptyRelation: rows=1
       SubqueryAlias: right
-        Projection: Int64(2) AS a, Int64(0) AS b, Int64(0) AS c
-          EmptyRelation: rows=1
+        Projection: count(*) AS a, count(*) AS b
+          Aggregate: groupBy=[[]], aggr=[[count(*)]]
+            TableScan: orders
     "
     );
 }