From 554dff9b130714ea98673de2cc2555a967377f14 Mon Sep 17 00:00:00 2001 From: comphead Date: Mon, 2 May 2022 01:21:42 -0700 Subject: [PATCH 1/2] nested query fix --- datafusion/core/src/sql/planner.rs | 1 + datafusion/core/tests/sql/expr.rs | 20 ++++++++++++++++++++ 2 files changed, 21 insertions(+) diff --git a/datafusion/core/src/sql/planner.rs b/datafusion/core/src/sql/planner.rs index ca8cb23bc5e9e..e6b50355ca2a7 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 5f0da3458babf..6646221d339f8 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,21 @@ 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 id FROM t1"; + 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!["++", "++"]; + assert_batches_eq!(expected, &actual); + Ok(()) +} From 60f563736268343e07e971803d801ac7c2f6adfe Mon Sep 17 00:00:00 2001 From: comphead Date: Mon, 2 May 2022 01:36:30 -0700 Subject: [PATCH 2/2] nested query test fix --- datafusion/core/tests/sql/expr.rs | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/datafusion/core/tests/sql/expr.rs b/datafusion/core/tests/sql/expr.rs index 6646221d339f8..5b2d50213069c 100644 --- a/datafusion/core/tests/sql/expr.rs +++ b/datafusion/core/tests/sql/expr.rs @@ -1140,10 +1140,14 @@ async fn nested_subquery() -> Result<()> { 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 id FROM t1"; + 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!["++", "++"]; + let expected = vec!["+-----+", "| cnt |", "+-----+", "| 0 |", "+-----+"]; assert_batches_eq!(expected, &actual); Ok(()) }