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

@@ -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/datasource/listing/helpers.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,
+    SubqueryNode, 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,
@@ -267,14 +267,37 @@ pub struct Limit {
 /// Evaluates correlated sub queries
 #[derive(Clone)]
 pub struct Subquery {
-    /// The list of sub queries
+    /// The list of sub queries (SubqueryNode)
     pub subqueries: Vec<LogicalPlan>,
     /// The incoming logical plan
     pub input: Arc<LogicalPlan>,
     /// The schema description of the output
     pub schema: DFSchemaRef,
 }
 
+/// Subquery node defines single subquery with its type
+#[derive(Clone)]
+pub struct SubqueryNode {
+    /// The logical plan of subquery
+    pub input: Arc<LogicalPlan>,
+    /// The subquery type
+    pub typ: SubqueryType,
+    /// The schema description of the output
+    pub schema: DFSchemaRef,
+}
+
+/// Subquery type
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd)]
+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,
+    // [NOT] IN(...) is not defined as it is implicitly evaluated as ANY = (...) / ALL <> (...)
+}
+
 impl Subquery {
     /// Merge schema of main input and correlated subquery columns
     pub fn merged_schema(input: &LogicalPlan, subqueries: &[LogicalPlan]) -> DFSchema {
@@ -284,6 +307,72 @@ impl Subquery {
             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),
+            // ANY/ALL subquery converts subquery result rows into a list
+            // and uses existing code evaluating ANY with a list to evaluate the result
+            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)
+            }
+        }
+    }
+}
+
+impl SubqueryNode {
+    /// Creates a new SubqueryNode evaluating the schema based on subquery type
+    pub fn new(input: LogicalPlan, typ: SubqueryType) -> Self {
+        let schema = Subquery::transform_dfschema(input.schema(), typ);
+        Self {
+            input: Arc::new(input),
+            typ,
+            schema: Arc::new(schema),
+        }
+    }
+}
+
+impl Display for SubqueryType {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        let name = match self {
+            Self::Scalar => "Scalar",
+            Self::Exists => "Exists",
+            Self::AnyAll => "AnyAll",
+        };
+        write!(f, "{}", name)
+    }
 }
 
 /// Values expression. See
@@ -402,6 +491,8 @@ pub enum LogicalPlan {
     Limit(Limit),
     /// Evaluates correlated sub queries
     Subquery(Subquery),
+    /// Single subquery node with subquery type
+    SubqueryNode(SubqueryNode),
     /// Creates an external table.
     CreateExternalTable(CreateExternalTable),
     /// Creates an in memory table.
@@ -439,6 +530,7 @@ impl LogicalPlan {
             }) => projected_schema,
             LogicalPlan::Projection(Projection { schema, .. }) => schema,
             LogicalPlan::Subquery(Subquery { schema, .. }) => schema,
+            LogicalPlan::SubqueryNode(SubqueryNode { schema, .. }) => schema,
             LogicalPlan::Filter(Filter { input, .. }) => input.schema(),
             LogicalPlan::Distinct(Distinct { input }) => input.schema(),
             LogicalPlan::Window(Window { schema, .. }) => schema,
@@ -498,7 +590,8 @@ impl LogicalPlan {
                 schemas.insert(0, schema);
                 schemas
             }
-            LogicalPlan::Union(Union { schema, .. }) => {
+            LogicalPlan::Union(Union { schema, .. })
+            | LogicalPlan::SubqueryNode(SubqueryNode { schema, .. }) => {
                 vec![schema]
             }
             LogicalPlan::Extension(extension) => vec![extension.node.schema()],
@@ -569,6 +662,7 @@ impl LogicalPlan {
             | LogicalPlan::Analyze { .. }
             | LogicalPlan::Explain { .. }
             | LogicalPlan::Subquery(_)
+            | LogicalPlan::SubqueryNode(_)
             | LogicalPlan::Union(_)
             | LogicalPlan::Distinct(_) => {
                 vec![]
@@ -587,6 +681,7 @@ impl LogicalPlan {
                 .into_iter()
                 .chain(subqueries.iter())
                 .collect(),
+            LogicalPlan::SubqueryNode(SubqueryNode { input, .. }) => vec![input],
             LogicalPlan::Filter(Filter { input, .. }) => vec![input],
             LogicalPlan::Repartition(Repartition { input, .. }) => vec![input],
             LogicalPlan::Window(Window { input, .. }) => vec![input],
@@ -735,6 +830,9 @@ impl LogicalPlan {
                 }
                 true
             }
+            LogicalPlan::SubqueryNode(SubqueryNode { input, .. }) => {
+                input.accept(visitor)?
+            }
             LogicalPlan::Filter(Filter { input, .. }) => input.accept(visitor)?,
             LogicalPlan::Repartition(Repartition { input, .. }) => {
                 input.accept(visitor)?
@@ -1064,6 +1162,9 @@ impl LogicalPlan {
                         Ok(())
                     }
                     LogicalPlan::Subquery(Subquery { .. }) => write!(f, "Subquery"),
+                    LogicalPlan::SubqueryNode(SubqueryNode { typ, .. }) => {
+                        write!(f, "SubqueryNode: type={:?}", typ)
+                    }
                     LogicalPlan::Filter(Filter {
                         predicate: ref expr,
                         ..