handle single wildcard as well

Author: gruuya

datafusion/sql/src/set_expr.rs

@@ -203,6 +203,21 @@ fn alias_set_expr(set_expr: &mut SetExpr, schema: &DFSchemaRef) {
 // 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) {
+    // Figure out how many (qualified) wildcards we got. We only handle
+    // the case of a single unqualified wildcard; for multiple or qualified
+    // wildcards we can't reliably determine column counts, so bail out.
+    let (wildcard_count, qualified_wildcard_count) =
+        items.iter().fold((0, 0), |(wc, qwc), item| match item {
+            SelectItem::Wildcard(_) => (wc + 1, qwc),
+            SelectItem::QualifiedWildcard(_, _) => (wc, qwc + 1),
+            _ => (wc, qwc),
+        });
+    if qualified_wildcard_count > 0 || wildcard_count > 1 {
+        return;
+    }
+
+    let wildcard_expansion = schema.fields().len().saturating_sub(items.len() - 1);
+
     let mut col_idx = 0;
     for item in items.iter_mut() {
         match item {
@@ -222,8 +237,11 @@ fn alias_select_items(items: &mut [SelectItem], schema: &DFSchemaRef) {
             SelectItem::ExprWithAlias { .. } => {
                 col_idx += 1;
             }
-            SelectItem::Wildcard(_) | SelectItem::QualifiedWildcard(_, _) => {
-                // Wildcards expand to multiple columns - skip position tracking
+            SelectItem::Wildcard(_) => {
+                col_idx += wildcard_expansion;
+            }
+            SelectItem::QualifiedWildcard(_, _) => {
+                unreachable!("qualified wildcards are handled above")
             }
         }
     }

datafusion/sql/tests/sql_integration.rs

@@ -2678,6 +2678,25 @@ fn union_all_with_duplicate_expressions() {
     );
 }
 
+#[test]
+fn union_with_qualified_and_duplicate_expressions() {
+    let sql = "\
+        SELECT 0 a, id b, price c, 0 d FROM test_decimal \
+        UNION SELECT 1, *, 1 FROM test_decimal";
+    let plan = logical_plan(sql).unwrap();
+    assert_snapshot!(
+        plan,
+        @"
+    Distinct:
+      Union
+        Projection: Int64(0) AS a, test_decimal.id AS b, test_decimal.price AS c, Int64(0) AS d
+          TableScan: test_decimal
+        Projection: Int64(1) AS a, test_decimal.id, test_decimal.price, Int64(1) AS d
+          TableScan: test_decimal
+    "
+    );
+}
+
 #[test]
 fn intersect_with_duplicate_expressions() {
     let sql = "\