feat: Support EXISTS subquery

Author: MazterQyou

datafusion/core/src/logical_plan/plan.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,
@@ -282,13 +282,16 @@ pub struct Subquery {
 pub enum SubqueryType {
     /// Scalar (SELECT, WHERE) evaluating to one value
     Scalar,
-    // This will be extended with `Exists` and `AnyAll` types.
+    /// EXISTS(...) evaluating to true if at least one row was produced
+    Exists,
+    // This will be extended with `AnyAll` type.
 }
 
 impl Display for SubqueryType {
     fn fmt(&self, f: &mut Formatter) -> fmt::Result {
         let subquery_type = match self {
             SubqueryType::Scalar => "Scalar",
+            SubqueryType::Exists => "Exists",
         };
         write!(f, "{}", subquery_type)
     }
@@ -315,15 +318,31 @@ impl Subquery {
     pub fn transform_dfschema(schema: &DFSchema, typ: SubqueryType) -> DFSchema {
         match typ {
             SubqueryType::Scalar => schema.clone(),
-            // Schema will be transformed for `Exists` and `AnyAll`
+            SubqueryType::Exists => {
+                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()
+            } // Schema will be transformed for `AnyAll` as well
         }
     }
 
     /// Transform Arrow field according to subquery type
     pub fn transform_field(field: &Field, typ: SubqueryType) -> Field {
         match typ {
             SubqueryType::Scalar => field.clone(),
-            // Field will be transformed for `Exists` and `AnyAll`
+            SubqueryType::Exists => Field::new(field.name(), DataType::Boolean, false),
+            // Field will be transformed for `AnyAll` as well
         }
     }
 }

datafusion/core/src/physical_plan/subquery.rs

@@ -30,7 +30,7 @@ use std::task::{Context, Poll};
 use crate::error::{DataFusionError, Result};
 use crate::logical_plan::{Subquery, SubqueryType};
 use crate::physical_plan::{DisplayFormatType, ExecutionPlan, Partitioning};
-use arrow::array::new_null_array;
+use arrow::array::{new_null_array, BooleanArray};
 use arrow::datatypes::{Schema, SchemaRef};
 use arrow::error::{ArrowError, Result as ArrowResult};
 use arrow::record_batch::RecordBatch;
@@ -215,12 +215,23 @@ impl ExecutionPlan for SubqueryExec {
                                             .to_string(),
                                     )),
                                 },
+                                SubqueryType::Exists => match subquery_batch
+                                    .column(0)
+                                    .len()
+                                {
+                                    0 => subquery_arrays[subquery_i]
+                                        .push(Arc::new(BooleanArray::from(vec![false]))),
+                                    _ => subquery_arrays[subquery_i]
+                                        .push(Arc::new(BooleanArray::from(vec![true]))),
+                                },
                             };
                         } else {
                             match subquery_type {
                                 SubqueryType::Scalar => {
                                     subquery_arrays[subquery_i].push(null_array())
                                 }
+                                SubqueryType::Exists => subquery_arrays[subquery_i]
+                                    .push(Arc::new(BooleanArray::from(vec![false]))),
                             };
                         }
                     }

datafusion/core/src/sql/planner.rs

@@ -2300,11 +2300,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
                 })
             }
 
-            // FIXME: Exists is unsupported but all the queries we need return false
-            SQLExpr::Exists(_) => {
-                warn!("EXISTS(...) is not supported yet. Replacing with scalar `false` value.");
-                Ok(Expr::Literal(ScalarValue::Boolean(Some(false))))
-            }
+            SQLExpr::Exists(q) => self.subquery_to_plan(q, SubqueryType::Exists, schema),
 
             // FIXME: ArraySubquery is unsupported but all the queries we need return empty array
             SQLExpr::ArraySubquery(_) => {
@@ -5074,6 +5070,22 @@ mod tests {
         quick_test(sql, expected);
     }
 
+    #[test]
+    fn subquery_exists() {
+        let sql = "select person.id, exists(select person.id) from person where exists(select 1 where false)";
+        let expected = "Projection: #person.id, #__subquery-1.person.id\
+                        \n  Subquery: types=[Exists]\
+                        \n    Filter: #__subquery-0.Int64(1)\
+                        \n      Subquery: types=[Exists]\
+                        \n        TableScan: person projection=None\
+                        \n        Projection: Int64(1), alias=__subquery-0\
+                        \n          Filter: Boolean(false)\
+                        \n            EmptyRelation\
+                        \n    Projection: ^#person.id, alias=__subquery-1\
+                        \n      EmptyRelation";
+        quick_test(sql, expected);
+    }
+
     #[test]
     fn join_on_disjunction_condition() {
         let sql = "SELECT id, order_id \

datafusion/core/tests/sql/subquery.rs

@@ -148,6 +148,27 @@ async fn subquery_select_and_where_with_from() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn subquery_exists() -> Result<()> {
+    let ctx = SessionContext::new();
+    register_aggregate_simple_csv(&ctx).await?;
+
+    let sql = "SELECT DISTINCT c1 FROM aggregate_simple o WHERE EXISTS(SELECT 1 FROM aggregate_simple p WHERE o.c1 * 2 = p.c1) ORDER BY c1";
+    let actual = execute_to_batches(&ctx, sql).await;
+
+    let expected = vec![
+        "+---------+",
+        "| c1      |",
+        "+---------+",
+        "| 0.00001 |",
+        "| 0.00002 |",
+        "+---------+",
+    ];
+    assert_batches_eq!(expected, &actual);
+
+    Ok(())
+}
+
 #[tokio::test]
 async fn subquery_projection_pushdown() -> Result<()> {
     let ctx = SessionContext::new();