feat: implement exists subquery

datafusion/Cargo.toml

@@ -79,6 +79,7 @@ num-traits = { version = "0.2", optional = true }
 pyo3 = { version = "0.15", optional = true }
 tempfile = "3"
 parking_lot = "0.12"
+uuid = { version = "0.8", features = ["v4"] }
 
 [dev-dependencies]
 criterion = "0.3"

datafusion/src/datasource/listing/helpers.rs

@@ -107,7 +107,8 @@ impl ExpressionVisitor for ApplicabilityVisitor<'_> {
             | Expr::AggregateFunction { .. }
             | Expr::Sort { .. }
             | Expr::WindowFunction { .. }
-            | Expr::Wildcard => {
+            | Expr::Wildcard
+            | Expr::Exists(_) => {
                 *self.is_applicable = false;
                 Recursion::Stop(self)
             }

datafusion/src/logical_plan/dfschema.rs

@@ -32,7 +32,7 @@ use std::fmt::{Display, Formatter};
 pub type DFSchemaRef = Arc<DFSchema>;
 
 /// DFSchema wraps an Arrow schema and adds relation names
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct DFSchema {
     /// Fields
     fields: Vec<DFField>,
@@ -403,7 +403,7 @@ impl Display for DFSchema {
 }
 
 /// DFField wraps an Arrow field and adds an optional qualifier
-#[derive(Debug, Clone, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct DFField {
     /// Optional qualifier (usually a table or relation name)
     qualifier: Option<String>,

datafusion/src/logical_plan/expr.rs

@@ -373,6 +373,8 @@ pub enum Expr {
     },
     /// Represents a reference to all fields in a schema.
     Wildcard,
+    /// Exists subquery
+    Exists(Box<LogicalPlan>),
 }
 
 /// Fixed seed for the hashing so that Ords are consistent across runs
@@ -508,9 +510,163 @@ impl Expr {
 
                 get_indexed_field(&data_type, key).map(|x| x.data_type().clone())
             }
+            Expr::Exists(logical_plan) => {
+                let df_fields = logical_plan.schema().fields();
+                let mut fields = Vec::with_capacity(df_fields.len());
+                let _ = df_fields
+                    .iter()
+                    .map(|df_field| fields.push(df_field.field().clone()));
+                match fields.len() {
+                    1 => Ok(fields[0].data_type().clone()),
+                    _ => Ok(DataType::Struct(fields)),
+                }
+            }
         }
     }
 
+    /// rewrite the subquery to apply to the outer logical plan
+    pub fn rewrite_to(&self, outer: &mut LogicalPlan) -> Result<Expr> {
+        Ok(match self {
+            Expr::Alias(expr, name) => {
+                Expr::Alias(Box::new(expr.rewrite_to(outer).unwrap()), name.clone())
+            }
+            Expr::Column(col) => Expr::Column(col.clone()),
+            Expr::ScalarVariable(vars) => Expr::ScalarVariable(vars.clone()),
+            Expr::Literal(val) => Expr::Literal(val.clone()),
+            Expr::BinaryExpr { left, op, right } => Expr::BinaryExpr {
+                left: Box::new(left.rewrite_to(outer).unwrap()),
+                op: op.clone(),
+                right: Box::new(right.rewrite_to(outer).unwrap()),
+            },
+            Expr::Not(expr) => Expr::Not(Box::new(expr.rewrite_to(outer).unwrap())),
+            Expr::IsNotNull(expr) => {
+                Expr::IsNotNull(Box::new(expr.rewrite_to(outer).unwrap()))
+            }
+            Expr::IsNull(expr) => Expr::IsNull(Box::new(expr.rewrite_to(outer).unwrap())),
+            Expr::Negative(expr) => {
+                Expr::Negative(Box::new(expr.rewrite_to(outer).unwrap()))
+            }
+            Expr::GetIndexedField { expr, key } => Expr::GetIndexedField {
+                expr: Box::new(expr.rewrite_to(outer).unwrap()),
+                key: key.clone(),
+            },
+            Expr::Between {
+                expr,
+                negated,
+                low,
+                high,
+            } => Expr::Between {
+                expr: Box::new(expr.rewrite_to(outer).unwrap()),
+                negated: *negated,
+                low: Box::new(low.rewrite_to(outer).unwrap()),
+                high: Box::new(high.rewrite_to(outer).unwrap()),
+            },
+            Expr::Case {
+                expr,
+                when_then_expr,
+                else_expr,
+            } => Expr::Case {
+                expr: if let Some(e) = expr {
+                    Option::from(Box::new(e.rewrite_to(outer).unwrap()))
+                } else {
+                    None
+                },
+                when_then_expr: when_then_expr
+                    .into_iter()
+                    .map(|(when, then)| {
+                        (
+                            Box::new(when.rewrite_to(outer).unwrap()),
+                            Box::new(then.rewrite_to(outer).unwrap()),
+                        )
+                    })
+                    .collect(),
+                else_expr: if let Some(e) = else_expr {
+                    Option::from(Box::new(e.rewrite_to(outer).unwrap()))
+                } else {
+                    None
+                },
+            },
+            Expr::Cast { expr, data_type } => Expr::Cast {
+                expr: Box::new(expr.rewrite_to(outer).unwrap()),
+                data_type: data_type.clone(),
+            },
+            Expr::TryCast { expr, data_type } => Expr::TryCast {
+                expr: Box::new(expr.rewrite_to(outer).unwrap()),
+                data_type: data_type.clone(),
+            },
+            Expr::Sort {
+                expr,
+                asc,
+                nulls_first,
+            } => Expr::Sort {
+                expr: Box::new(expr.rewrite_to(outer).unwrap()),
+                asc: *asc,
+                nulls_first: *nulls_first,
+            },
+            Expr::ScalarFunction { fun, args } => Expr::ScalarFunction {
+                fun: fun.clone(),
+                args: args.iter().map(|a| a.rewrite_to(outer).unwrap()).collect(),
+            },
+            Expr::ScalarUDF { fun, args } => Expr::ScalarUDF {
+                fun: fun.clone(),
+                args: args.iter().map(|a| a.rewrite_to(outer).unwrap()).collect(),
+            },
+            Expr::AggregateFunction {
+                fun,
+                args,
+                distinct,
+            } => Expr::AggregateFunction {
+                fun: fun.clone(),
+                args: args.iter().map(|a| a.rewrite_to(outer).unwrap()).collect(),
+                distinct: *distinct,
+            },
+            Expr::WindowFunction {
+                fun,
+                args,
+                partition_by,
+                order_by,
+                window_frame,
+            } => Expr::WindowFunction {
+                fun: fun.clone(),
+                args: args.iter().map(|a| a.rewrite_to(outer).unwrap()).collect(),
+                partition_by: partition_by
+                    .iter()
+                    .map(|e| e.rewrite_to(outer).unwrap())
+                    .collect(),
+                order_by: order_by
+                    .iter()
+                    .map(|e| e.rewrite_to(outer).unwrap())
+                    .collect(),
+                window_frame: window_frame.clone(),
+            },
+            Expr::AggregateUDF { fun, args } => Expr::AggregateUDF {
+                fun: fun.clone(),
+                args: args.iter().map(|a| a.rewrite_to(outer).unwrap()).collect(),
+            },
+            Expr::InList {
+                expr,
+                list,
+                negated,
+            } => Expr::InList {
+                expr: Box::new(expr.rewrite_to(outer).unwrap()),
+                list: list.iter().map(|e| e.rewrite_to(outer).unwrap()).collect(),
+                negated: *negated,
+            },
+            Expr::Wildcard => Expr::Wildcard,
+            Expr::Exists(logical_plan) => {
+                *outer = outer
+                    .clone()
+                    .process_subquery(logical_plan.as_ref())
+                    .unwrap();
+                // finally, return the column contains `bool`
+                Expr::Column(Column {
+                    relation: None,
+                    name: "exists_subquery_col".to_string(),
+                })
+            }
+        })
+    }
+
     /// Returns the nullability of the expression based on [ExprSchema].
     ///
     /// Note: [DFSchema] implements [ExprSchema].
@@ -569,6 +725,14 @@ impl Expr {
                 let data_type = expr.get_type(input_schema)?;
                 get_indexed_field(&data_type, key).map(|x| x.is_nullable())
             }
+            Expr::Exists(logical_plan) => {
+                for df_field in logical_plan.schema().fields() {
+                    if !df_field.is_nullable() {
+                        return Ok(false);
+                    }
+                }
+                Ok(true)
+            }
         }
     }
 
@@ -752,7 +916,7 @@ impl Expr {
     /// pre_visit(Column("foo"))
     /// pre_visit(Column("bar"))
     /// post_visit(Column("bar"))
-    /// post_visit(Column("bar"))
+    /// post_visit(Column("foo"))
     /// post_visit(BinaryExpr(GT))
     /// ```
     ///
@@ -783,7 +947,8 @@ impl Expr {
             Expr::Column(_)
             | Expr::ScalarVariable(_)
             | Expr::Literal(_)
-            | Expr::Wildcard => Ok(visitor),
+            | Expr::Wildcard
+            | Expr::Exists(_) => Ok(visitor),
             Expr::BinaryExpr { left, right, .. } => {
                 let visitor = left.accept(visitor)?;
                 right.accept(visitor)
@@ -872,7 +1037,7 @@ impl Expr {
     /// ```text
     /// pre_visit(BinaryExpr(GT))
     /// pre_visit(Column("foo"))
-    /// mutatate(Column("foo"))
+    /// mutate(Column("foo"))
     /// pre_visit(Column("bar"))
     /// mutate(Column("bar"))
     /// mutate(BinaryExpr(GT))
@@ -1010,6 +1175,7 @@ impl Expr {
                 expr: rewrite_boxed(expr, rewriter)?,
                 key,
             },
+            Expr::Exists(logical_plan) => Expr::Exists(logical_plan),
         };
 
         // now rewrite this expression itself
@@ -2063,6 +2229,9 @@ impl fmt::Debug for Expr {
             Expr::GetIndexedField { ref expr, key } => {
                 write!(f, "({:?})[{}]", expr, key)
             }
+            Expr::Exists(logical_plan) => {
+                write!(f, "exists subquery: {:?}", logical_plan)
+            }
         }
     }
 }
@@ -2224,6 +2393,9 @@ fn create_name(e: &Expr, input_schema: &DFSchema) -> Result<String> {
         Expr::Wildcard => Err(DataFusionError::Internal(
             "Create name does not support wildcard".to_string(),
         )),
+        Expr::Exists(logical_plan) => {
+            Ok(format!("Exists subquery {}", logical_plan.display()))
+        }
     }
 }