[arrow-datafusion] branch master updated: nested query fix (#2402)

[agrove]

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

agrove 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 946194682 nested query fix (#2402)
946194682 is described below

commit 946194682248abdeadfe358eaf9d7824b72262ed
Author: comphead <comph...@users.noreply.github.com>
AuthorDate: Mon May 2 12:47:18 2022 -0700

    nested query fix (#2402)
---
 datafusion/core/src/sql/planner.rs |  1 +
 datafusion/core/tests/sql/expr.rs  | 24 ++++++++++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/datafusion/core/src/sql/planner.rs 
b/datafusion/core/src/sql/planner.rs
index cf974c0d1..ba1a551c9 100644
--- a/datafusion/core/src/sql/planner.rs
+++ b/datafusion/core/src/sql/planner.rs
@@ -324,6 +324,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
                     }
                 }
             }
+            SetExpr::Query(q) => self.query_to_plan(*q, ctes),
             _ => Err(DataFusionError::NotImplemented(format!(
                 "Query {} not implemented yet",
                 set_expr
diff --git a/datafusion/core/tests/sql/expr.rs 
b/datafusion/core/tests/sql/expr.rs
index 5f0da3458..5b2d50213 100644
--- a/datafusion/core/tests/sql/expr.rs
+++ b/datafusion/core/tests/sql/expr.rs
@@ -15,6 +15,8 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use datafusion::datasource::empty::EmptyTable;
+
 use super::*;
 
 #[tokio::test]
@@ -1127,3 +1129,25 @@ async fn csv_query_sqrt_sqrt() -> Result<()> {
     assert_float_eq(&expected, &actual);
     Ok(())
 }
+
+#[tokio::test]
+async fn nested_subquery() -> Result<()> {
+    let ctx = SessionContext::new();
+    let schema = Schema::new(vec![
+        Field::new("id", DataType::Int16, false),
+        Field::new("a", DataType::Int16, false),
+    ]);
+    let empty_table = Arc::new(EmptyTable::new(Arc::new(schema)));
+    ctx.register_table("t1", empty_table.clone())?;
+    ctx.register_table("t2", empty_table)?;
+    let sql = "SELECT COUNT(*) as cnt \
+    FROM (\
+        (SELECT id FROM t1) EXCEPT \
+        (SELECT id FROM t2)\
+        ) foo";
+    let actual = execute_to_batches(&ctx, sql).await;
+    // the purpose of this test is just to make sure the query produces a 
valid plan
+    let expected = vec!["+-----+", "| cnt |", "+-----+", "| 0   |", "+-----+"];
+    assert_batches_eq!(expected, &actual);
+    Ok(())
+}