fix UNION ALL bug: thread 'main' panicked at 'index out of bounds: th…

…e len is 1 but the index is 1', ./src/datatypes/schema.rs:165:10

datafusion/src/optimizer/projection_push_down.rs

@@ -395,7 +395,6 @@ fn optimize_plan(
                 .filter(|f| new_required_columns.contains(&f.qualified_column()))
                 .map(|f| f.field())
                 .collect::<HashSet<&Field>>();
-
             let new_inputs = inputs
                 .iter()
                 .map(|input_plan| {
@@ -416,10 +415,17 @@ fn optimize_plan(
                     )
                 })
                 .collect::<Result<Vec<_>>>()?;
-
+            let new_schema = DFSchema::new(
+                schema
+                    .fields()
+                    .iter()
+                    .filter(|f| union_required_fields.contains(f.field()))
+                    .cloned()
+                    .collect(),
+            )?;
             Ok(LogicalPlan::Union {
                 inputs: new_inputs,
-                schema: schema.clone(),
+                schema: Arc::new(new_schema),
                 alias: alias.clone(),
             })
         }

datafusion/tests/sql.rs

@@ -5076,3 +5076,14 @@ async fn union_distinct() -> Result<()> {
     assert_eq!(expected, actual);
     Ok(())
 }
+
+#[tokio::test]
+async fn union_all_with_aggregate() -> Result<()> {
+    let mut ctx = ExecutionContext::new();
+    let sql =
+        "SELECT SUM(d) FROM (SELECT 1 as c, 2 as d UNION ALL SELECT 1 as c, 3 AS d) as a";
+    let actual = execute(&mut ctx, sql).await;
+    let expected = vec![vec!["5"]];
+    assert_eq!(expected, actual);
+    Ok(())
+}