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;

datafusion/core/src/logical_plan/expr_visitor.rs

@@ -26,7 +26,7 @@ use crate::error::DataFusionError;
 use crate::logical_plan::dfschema::DFSchemaRef;
 use crate::sql::parser::FileType;
 use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
-use datafusion_common::DFSchema;
+use datafusion_common::{DFField, DFSchema};
 use std::fmt::Formatter;
 use std::{
     collections::HashSet,
@@ -268,21 +268,79 @@ pub struct Limit {
 #[derive(Clone)]
 pub struct Subquery {
     /// The list of sub queries
-    pub subqueries: Vec<LogicalPlan>,
+    pub subqueries: Vec<(LogicalPlan, SubqueryType)>,
     /// The incoming logical plan
     pub input: Arc<LogicalPlan>,
     /// The schema description of the output
     pub schema: DFSchemaRef,
 }
 
+/// Subquery type
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub enum SubqueryType {
+    /// Scalar (SELECT, WHERE) evaluating to one value
+    Scalar,
+    /// EXISTS(...) evaluating to true if at least one row was produced
+    Exists,
+    /// ANY(...)/ALL(...)
+    AnyAll,
+}
+
 impl Subquery {
     /// Merge schema of main input and correlated subquery columns
-    pub fn merged_schema(input: &LogicalPlan, subqueries: &[LogicalPlan]) -> DFSchema {
-        subqueries.iter().fold((**input.schema()).clone(), |a, b| {
-            let mut res = a;
-            res.merge(b.schema());
-            res
-        })
+    pub fn merged_schema(
+        input: &LogicalPlan,
+        subqueries: &[(LogicalPlan, SubqueryType)],
+    ) -> DFSchema {
+        subqueries
+            .iter()
+            .fold((**input.schema()).clone(), |input_schema, (plan, typ)| {
+                let mut res = input_schema;
+                let subquery_schema = Self::transform_dfschema(plan.schema(), *typ);
+                res.merge(&subquery_schema);
+                res
+            })
+    }
+
+    /// Transform DataFusion schema according to subquery type
+    pub fn transform_dfschema(schema: &DFSchema, typ: SubqueryType) -> DFSchema {
+        match typ {
+            SubqueryType::Scalar => schema.clone(),
+            SubqueryType::Exists | SubqueryType::AnyAll => {
+                let new_fields = schema
+                    .fields()
+                    .iter()
+                    .map(|field| {
+                        let new_field = Subquery::transform_field(field.field(), typ);
+                        if let Some(qualifier) = field.qualifier() {
+                            DFField::from_qualified(qualifier, new_field)
+                        } else {
+                            DFField::from(new_field)
+                        }
+                    })
+                    .collect();
+                DFSchema::new_with_metadata(new_fields, schema.metadata().clone())
+                    .unwrap()
+            }
+        }
+    }
+
+    /// Transform Arrow field according to subquery type
+    pub fn transform_field(field: &Field, typ: SubqueryType) -> Field {
+        match typ {
+            SubqueryType::Scalar => field.clone(),
+            SubqueryType::Exists => Field::new(field.name(), DataType::Boolean, false),
+            SubqueryType::AnyAll => {
+                let item = Field::new_dict(
+                    "item",
+                    field.data_type().clone(),
+                    true,
+                    field.dict_id().unwrap_or(0),
+                    field.dict_is_ordered().unwrap_or(false),
+                );
+                Field::new(field.name(), DataType::List(Box::new(item)), false)
+            }
+        }
     }
 }
 
@@ -585,7 +643,7 @@ impl LogicalPlan {
                 input, subqueries, ..
             }) => vec![input.as_ref()]
                 .into_iter()
-                .chain(subqueries.iter())
+                .chain(subqueries.iter().map(|(q, _)| q))
                 .collect(),
             LogicalPlan::Filter(Filter { input, .. }) => vec![input],
             LogicalPlan::Repartition(Repartition { input, .. }) => vec![input],
@@ -728,7 +786,7 @@ impl LogicalPlan {
                 input, subqueries, ..
             }) => {
                 input.accept(visitor)?;
-                for input in subqueries {
+                for (input, _) in subqueries {
                     if !input.accept(visitor)? {
                         return Ok(false);
                     }

datafusion/core/src/logical_plan/mod.rs

@@ -508,6 +508,10 @@ impl ExprIdentifierVisitor<'_> {
                 desc.push_str("InList-");
                 desc.push_str(&negated.to_string());
             }
+            Expr::InSubquery { negated, .. } => {
+                desc.push_str("InSubquery-");
+                desc.push_str(&negated.to_string());
+            }
             Expr::Wildcard => {
                 desc.push_str("Wildcard-");
             }

datafusion/core/src/logical_plan/plan.rs

@@ -456,7 +456,7 @@ fn optimize_plan(
             input, subqueries, ..
         }) => {
             let mut subquery_required_columns = HashSet::new();
-            for subquery in subqueries.iter() {
+            for subquery in subqueries.iter().map(|(q, _)| q) {
                 let mut inputs = vec![subquery];
                 while !inputs.is_empty() {
                     let mut next_inputs = Vec::new();

datafusion/core/src/optimizer/common_subexpr_eliminate.rs

@@ -392,6 +392,7 @@ impl<'a> ConstEvaluator<'a> {
             | Expr::OuterColumn(_, _)
             | Expr::WindowFunction { .. }
             | Expr::Sort { .. }
+            | Expr::InSubquery { .. }
             | Expr::Wildcard
             | Expr::QualifiedWildcard { .. } => false,
             Expr::ScalarFunction { fun, .. } => Self::volatility_ok(fun.volatility()),

datafusion/core/src/optimizer/projection_push_down.rs

@@ -93,6 +93,7 @@ impl ExpressionVisitor for ColumnNameVisitor<'_> {
             | Expr::AggregateFunction { .. }
             | Expr::AggregateUDF { .. }
             | Expr::InList { .. }
+            | Expr::InSubquery { .. }
             | Expr::Wildcard
             | Expr::QualifiedWildcard { .. }
             | Expr::GetIndexedField { .. } => {}
@@ -161,13 +162,17 @@ pub fn from_plan(
                 alias: alias.clone(),
             }))
         }
-        LogicalPlan::Subquery(Subquery { schema, .. }) => {
-            Ok(LogicalPlan::Subquery(Subquery {
-                subqueries: inputs[1..inputs.len()].to_vec(),
-                input: Arc::new(inputs[0].clone()),
-                schema: schema.clone(),
-            }))
-        }
+        LogicalPlan::Subquery(Subquery {
+            schema, subqueries, ..
+        }) => Ok(LogicalPlan::Subquery(Subquery {
+            subqueries: inputs[1..inputs.len()]
+                .iter()
+                .zip(subqueries.iter())
+                .map(|(input, (_, t))| (input.clone(), *t))
+                .collect(),
+            input: Arc::new(inputs[0].clone()),
+            schema: schema.clone(),
+        })),
         LogicalPlan::TableUDFs(TableUDFs { .. }) => {
             Ok(LogicalPlan::TableUDFs(TableUDFs {
                 expr: expr.to_vec(),
@@ -390,6 +395,9 @@ pub fn expr_sub_expressions(expr: &Expr) -> Result<Vec<Expr>> {
             }
             Ok(expr_list)
         }
+        Expr::InSubquery { expr, subquery, .. } => {
+            Ok(vec![expr.as_ref().to_owned(), subquery.as_ref().to_owned()])
+        }
         Expr::Wildcard { .. } => Err(DataFusionError::Internal(
             "Wildcard expressions are not valid in a logical query plan".to_owned(),
         )),
@@ -410,10 +418,11 @@ pub fn rewrite_expression(expr: &Expr, expressions: &[Expr]) -> Result<Expr> {
             op: *op,
             right: Box::new(expressions[1].clone()),
         }),
-        Expr::AnyExpr { op, .. } => Ok(Expr::AnyExpr {
+        Expr::AnyExpr { op, all, .. } => Ok(Expr::AnyExpr {
             left: Box::new(expressions[0].clone()),
             op: *op,
             right: Box::new(expressions[1].clone()),
+            all: *all,
         }),
         Expr::Like(Like {
             negated,
@@ -593,6 +602,11 @@ pub fn rewrite_expression(expr: &Expr, expressions: &[Expr]) -> Result<Expr> {
                 Ok(expr)
             }
         }
+        Expr::InSubquery { negated, .. } => Ok(Expr::InSubquery {
+            expr: Box::new(expressions[0].clone()),
+            subquery: Box::new(expressions[1].clone()),
+            negated: *negated,
+        }),
         Expr::Wildcard => Err(DataFusionError::Internal(
             "Wildcard expressions are not valid in a logical query plan".to_owned(),
         )),