feat: Improve subquery support

Author: MazterQyou

Cargo.lock

@@ -44,4 +44,4 @@ cranelift-module = { version = "0.82.0", optional = true }
 ordered-float = "2.10"
 parquet = { git = 'https://github.com/cube-js/arrow-rs.git', rev = "096ef28dde6b1ae43ce89ba2c3a9d98295f2972e", features = ["arrow"], optional = true }
 pyo3 = { version = "0.16", optional = true }
-sqlparser = { git = 'https://github.com/cube-js/sqlparser-rs.git', rev = "b3b40586d4c32a218ffdfcb0462e7e216cf3d6eb" }
+sqlparser = { git = 'https://github.com/cube-js/sqlparser-rs.git', rev = "2229652dc8fae8f45cbec344b4a1e40cf1bb69d9" }

datafusion-cli/Cargo.lock

@@ -79,7 +79,7 @@ pin-project-lite= "^0.2.7"
 pyo3 = { version = "0.16", optional = true }
 rand = "0.8"
 smallvec = { version = "1.6", features = ["union"] }
-sqlparser = { git = 'https://github.com/cube-js/sqlparser-rs.git', rev = "b3b40586d4c32a218ffdfcb0462e7e216cf3d6eb" }
+sqlparser = { git = 'https://github.com/cube-js/sqlparser-rs.git', rev = "2229652dc8fae8f45cbec344b4a1e40cf1bb69d9" }
 tempfile = "3"
 tokio = { version = "1.0", features = ["macros", "rt", "rt-multi-thread", "sync", "fs", "parking_lot"] }
 tokio-stream = "0.1"

datafusion/common/Cargo.toml

@@ -97,6 +97,7 @@ impl ExpressionVisitor for ApplicabilityVisitor<'_> {
             | Expr::ILike { .. }
             | Expr::SimilarTo { .. }
             | Expr::InList { .. }
+            | Expr::InSubquery { .. }
             | Expr::GetIndexedField { .. }
             | Expr::Case { .. } => Recursion::Continue(self),
 

datafusion/core/Cargo.toml

@@ -26,7 +26,7 @@ use crate::error::{DataFusionError, Result};
 use crate::logical_plan::expr_schema::ExprSchemable;
 use crate::logical_plan::plan::{
     Aggregate, Analyze, EmptyRelation, Explain, Filter, Join, Projection, Sort, Subquery,
-    TableScan, TableUDFs, ToStringifiedPlan, Union, Window,
+    SubqueryType, TableScan, TableUDFs, ToStringifiedPlan, Union, Window,
 };
 use crate::optimizer::utils;
 use crate::prelude::*;
@@ -528,12 +528,15 @@ impl LogicalPlanBuilder {
     pub fn subquery(
         &self,
         subqueries: impl IntoIterator<Item = impl Into<LogicalPlan>>,
+        types: impl IntoIterator<Item = SubqueryType>,
     ) -> Result<Self> {
         let subqueries = subqueries.into_iter().map(|l| l.into()).collect::<Vec<_>>();
-        let schema = Arc::new(Subquery::merged_schema(&self.plan, &subqueries));
+        let types = types.into_iter().collect::<Vec<_>>();
+        let schema = Arc::new(Subquery::merged_schema(&self.plan, &subqueries, &types));
         Ok(Self::from(LogicalPlan::Subquery(Subquery {
             input: Arc::new(self.plan.clone()),
             subqueries,
+            types,
             schema,
         })))
     }

datafusion/core/src/datasource/listing/helpers.rs

@@ -122,10 +122,16 @@ impl ExprRewritable for Expr {
                 op,
                 right: rewrite_boxed(right, rewriter)?,
             },
-            Expr::AnyExpr { left, op, right } => Expr::AnyExpr {
+            Expr::AnyExpr {
+                left,
+                op,
+                right,
+                all,
+            } => Expr::AnyExpr {
                 left: rewrite_boxed(left, rewriter)?,
                 op,
                 right: rewrite_boxed(right, rewriter)?,
+                all,
             },
             Expr::Like(Like {
                 negated,
@@ -263,6 +269,15 @@ impl ExprRewritable for Expr {
                 list: rewrite_vec(list, rewriter)?,
                 negated,
             },
+            Expr::InSubquery {
+                expr,
+                subquery,
+                negated,
+            } => Expr::InSubquery {
+                expr: rewrite_boxed(expr, rewriter)?,
+                subquery: rewrite_boxed(subquery, rewriter)?,
+                negated,
+            },
             Expr::Wildcard => Expr::Wildcard,
             Expr::QualifiedWildcard { qualifier } => {
                 Expr::QualifiedWildcard { qualifier }

datafusion/core/src/logical_plan/builder.rs

@@ -111,6 +111,7 @@ impl ExprSchemable for Expr {
             | Expr::IsNull(_)
             | Expr::Between { .. }
             | Expr::InList { .. }
+            | Expr::InSubquery { .. }
             | Expr::AnyExpr { .. }
             | Expr::IsNotNull(_) => Ok(DataType::Boolean),
             Expr::BinaryExpr {
@@ -158,7 +159,7 @@ impl ExprSchemable for Expr {
             | Expr::Between { expr, .. }
             | Expr::InList { expr, .. } => expr.nullable(input_schema),
             Expr::Column(c) => input_schema.nullable(c),
-            Expr::OuterColumn(_, _) => Ok(true),
+            Expr::OuterColumn(_, _) | Expr::InSubquery { .. } => Ok(true),
             Expr::Literal(value) => Ok(value.is_null()),
             Expr::Case {
                 when_then_expr,

datafusion/core/src/logical_plan/expr_rewriter.rs

@@ -191,6 +191,10 @@ impl ExprVisitable for Expr {
                 list.iter()
                     .try_fold(visitor, |visitor, arg| arg.accept(visitor))
             }
+            Expr::InSubquery { expr, subquery, .. } => {
+                let visitor = expr.accept(visitor)?;
+                subquery.accept(visitor)
+            }
         }?;
 
         visitor.post_visit(self)

datafusion/core/src/logical_plan/expr_schema.rs

@@ -68,6 +68,6 @@ pub use plan::{
     CreateCatalogSchema, CreateExternalTable, CreateMemoryTable, CrossJoin, Distinct,
     DropTable, EmptyRelation, Filter, JoinConstraint, JoinType, Limit, LogicalPlan,
     Partitioning, PlanType, PlanVisitor, Repartition, StringifiedPlan, Subquery,
-    TableScan, ToStringifiedPlan, Union, Values,
+    SubqueryType, TableScan, ToStringifiedPlan, Union, Values,
 };
 pub use registry::FunctionRegistry;