[arrow-datafusion] branch master updated: fix schema handling bug in projection push down for UNION ALL plan (#1088)

[houqp]

This is an automated email from the ASF dual-hosted git repository.

houqp pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git


The following commit(s) were added to refs/heads/master by this push:
     new 5a7daa3  fix schema handling bug in projection push down for UNION ALL 
plan (#1088)
5a7daa3 is described below

commit 5a7daa3899b889dc5f50ef65a952991a4000aba4
Author: carlos <wxd963996...@gmail.com>
AuthorDate: Sat Oct 9 13:37:55 2021 +0800

    fix schema handling bug in projection push down for UNION ALL plan (#1088)
---
 datafusion/src/optimizer/projection_push_down.rs | 12 +++++++++---
 datafusion/tests/sql.rs                          | 11 +++++++++++
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/datafusion/src/optimizer/projection_push_down.rs 
b/datafusion/src/optimizer/projection_push_down.rs
index 0a84c9d..aba248c 100644
--- a/datafusion/src/optimizer/projection_push_down.rs
+++ b/datafusion/src/optimizer/projection_push_down.rs
@@ -397,7 +397,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| {
@@ -418,10 +417,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(),
             })
         }
diff --git a/datafusion/tests/sql.rs b/datafusion/tests/sql.rs
index 5ae3585..32da908 100644
--- a/datafusion/tests/sql.rs
+++ b/datafusion/tests/sql.rs
@@ -5080,3 +5080,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(())
+}