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 = "10782e5d11fc0e2900c9359dddee0fbefbffd359" }
+sqlparser = { git = 'https://github.com/cube-js/sqlparser-rs.git', rev = "0bea7fa17907b96f32abf4bd6bb1cde43fe8d244" }

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 = "10782e5d11fc0e2900c9359dddee0fbefbffd359" }
+sqlparser = { git = 'https://github.com/cube-js/sqlparser-rs.git', rev = "0bea7fa17907b96f32abf4bd6bb1cde43fe8d244" }
 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

@@ -47,7 +47,8 @@ use super::{dfschema::ToDFSchema, expr_rewriter::coerce_plan_expr_for_schema, Di
 use super::{exprlist_to_fields, Expr, JoinConstraint, JoinType, LogicalPlan, PlanType};
 use crate::logical_plan::{
     columnize_expr, normalize_col, normalize_cols, rewrite_sort_cols_by_aggs, Column,
-    CrossJoin, DFField, DFSchema, DFSchemaRef, Limit, Partitioning, Repartition, Values,
+    CrossJoin, DFField, DFSchema, DFSchemaRef, Limit, Partitioning, Repartition,
+    SubqueryType, Values,
 };
 use crate::sql::utils::group_window_expr_by_sort_keys;
 
@@ -527,7 +528,7 @@ impl LogicalPlanBuilder {
     /// Apply correlated sub query
     pub fn subquery(
         &self,
-        subqueries: impl IntoIterator<Item = impl Into<LogicalPlan>>,
+        subqueries: impl IntoIterator<Item = impl Into<(LogicalPlan, SubqueryType)>>,
     ) -> Result<Self> {
         let subqueries = subqueries.into_iter().map(|l| l.into()).collect::<Vec<_>>();
         let schema = Arc::new(Subquery::merged_schema(&self.plan, &subqueries));

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;