Support predicate subqueries in projections

datafusion/optimizer/src/decorrelate_predicate_subquery.rs

@@ -28,16 +28,16 @@ use crate::{OptimizerConfig, OptimizerRule};
 use datafusion_common::alias::AliasGenerator;
 use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
 use datafusion_common::{
-    Column, DFSchemaRef, ExprSchema, NullEquality, Result, assert_or_internal_err,
-    plan_err,
+    Column, DFSchemaRef, ExprSchema, NullEquality, Result, ScalarValue,
+    assert_or_internal_err, plan_err,
 };
 use datafusion_expr::expr::{Exists, InSubquery};
 use datafusion_expr::expr_rewriter::create_col_from_scalar_expr;
 use datafusion_expr::logical_plan::{JoinType, Subquery};
 use datafusion_expr::utils::{conjunction, expr_to_columns, split_conjunction_owned};
 use datafusion_expr::{
     BinaryExpr, Expr, Filter, LogicalPlan, LogicalPlanBuilder, Operator, exists,
-    in_subquery, lit, not, not_exists, not_in_subquery,
+    in_subquery, lit, not, not_exists, not_in_subquery, when,
 };
 
 use log::debug;
@@ -69,6 +69,37 @@ impl OptimizerRule for DecorrelatePredicateSubquery {
             })?
             .data;
 
+        if let LogicalPlan::Projection(projection) = plan {
+            if !projection.expr.iter().any(has_subquery) {
+                return Ok(Transformed::no(LogicalPlan::Projection(projection)));
+            }
+
+            let original_projection = projection.clone();
+            let mut cur_input = Arc::unwrap_or_clone(projection.input);
+            let mut rewritten_exprs = Vec::with_capacity(projection.expr.len());
+            for expr in projection.expr {
+                let original_name = expr.schema_name().to_string();
+                let (new_input, mut rewritten_expr) =
+                    rewrite_inner_subqueries(cur_input, expr, config, true)?;
+                if has_subquery(&rewritten_expr) {
+                    return Ok(Transformed::no(LogicalPlan::Projection(
+                        original_projection,
+                    )));
+                }
+                cur_input = new_input;
+
+                if rewritten_expr.schema_name().to_string() != original_name {
+                    rewritten_expr = rewritten_expr.alias(original_name);
+                }
+                rewritten_exprs.push(rewritten_expr);
+            }
+
+            let new_plan = LogicalPlanBuilder::from(cur_input)
+                .project(rewritten_exprs)?
+                .build()?;
+            return Ok(Transformed::yes(new_plan));
+        }
+
         let LogicalPlan::Filter(filter) = plan else {
             return Ok(Transformed::no(plan));
         };
@@ -104,7 +135,7 @@ impl OptimizerRule for DecorrelatePredicateSubquery {
                 // The subquery expression is embedded within another expression
                 SubqueryPredicate::Embedded(expr) => {
                     let (plan, expr_without_subqueries) =
-                        rewrite_inner_subqueries(cur_input, expr, config)?;
+                        rewrite_inner_subqueries(cur_input, expr, config, false)?;
                     cur_input = plan;
                     other_exprs.push(expr_without_subqueries);
                 }
@@ -139,6 +170,7 @@ fn rewrite_inner_subqueries(
     outer: LogicalPlan,
     expr: Expr,
     config: &dyn OptimizerConfig,
+    materialize_in_value: bool,
 ) -> Result<(LogicalPlan, Expr)> {
     let mut cur_input = outer;
     let alias = config.alias_generator();
@@ -159,12 +191,23 @@ fn rewrite_inner_subqueries(
             subquery: Subquery { subquery, .. },
             negated,
         }) => {
-            let in_predicate = subquery
-                .head_output_expr()?
-                .map_or(plan_err!("single expression required."), |output_expr| {
-                    Ok(Expr::eq(*expr.clone(), output_expr))
-                })?;
-            match mark_join(&cur_input, &subquery, Some(&in_predicate), negated, alias)? {
+            let rewritten = if materialize_in_value {
+                in_subquery_value_mark_join(
+                    &cur_input,
+                    &subquery,
+                    *expr.clone(),
+                    negated,
+                    alias,
+                )?
+            } else {
+                let in_predicate = subquery
+                    .head_output_expr()?
+                    .map_or(plan_err!("single expression required."), |output_expr| {
+                        Ok(Expr::eq(*expr.clone(), output_expr))
+                    })?;
+                mark_join(&cur_input, &subquery, Some(&in_predicate), negated, alias)?
+            };
+            match rewritten {
                 Some((plan, exists_expr)) => {
                     cur_input = plan;
                     Ok(Transformed::yes(exists_expr))
@@ -178,6 +221,48 @@ fn rewrite_inner_subqueries(
     Ok((cur_input, expr_without_subqueries.data))
 }
 
+fn in_subquery_value_mark_join(
+    left: &LogicalPlan,
+    subquery: &LogicalPlan,
+    expr: Expr,
+    negated: bool,
+    alias: &Arc<AliasGenerator>,
+) -> Result<Option<(LogicalPlan, Expr)>> {
+    let output_expr = subquery
+        .head_output_expr()?
+        .map_or(plan_err!("single expression required."), Ok)?;
+    let in_predicate = Expr::eq(expr.clone(), output_expr.clone());
+    let Some((matched_plan, matched)) =
+        mark_join(left, subquery, Some(&in_predicate), false, alias)?
+    else {
+        return Ok(None);
+    };
+
+    // SQL IN needs three facts per outer row to distinguish FALSE from UNKNOWN.
+    let null_subquery = LogicalPlanBuilder::from(subquery.clone())
+        .filter(output_expr.is_null())?
+        .build()?;
+    let Some((null_plan, subquery_has_null)) =
+        mark_join(&matched_plan, &null_subquery, None, false, alias)?
+    else {
+        return Ok(None);
+    };
+    let Some((final_plan, subquery_non_empty)) =
+        mark_join(&null_plan, subquery, None, false, alias)?
+    else {
+        return Ok(None);
+    };
+
+    let unknown = subquery_has_null.or(expr.is_null().and(subquery_non_empty));
+    let result = when(matched, lit(true))
+        .when(unknown, lit(ScalarValue::Boolean(None)))
+        .otherwise(lit(false))?;
+    Ok(Some((
+        final_plan,
+        if negated { not(result) } else { result },
+    )))
+}
+
 enum SubqueryPredicate {
     // The subquery expression is at the top level of the filter and can be fully replaced by a
     // semi/anti join
@@ -1134,6 +1219,37 @@ mod tests {
         )
     }
 
+    #[test]
+    fn in_subquery_in_projection() -> Result<()> {
+        let plan = LogicalPlanBuilder::from(test_table_scan()?)
+            .project(vec![
+                in_subquery(col("c"), test_subquery_with_name("sq")?).alias("is_present"),
+            ])?
+            .build()?;
+
+        assert_optimized_plan_equal!(
+            plan,
+            @r"
+        Projection: CASE WHEN __correlated_sq_1.mark THEN Boolean(true) WHEN __correlated_sq_2.mark OR test.c IS NULL AND __correlated_sq_3.mark THEN Boolean(NULL) ELSE Boolean(false) END AS is_present [is_present:Boolean;N]
+          LeftMark Join:  Filter: Boolean(true) [a:UInt32, b:UInt32, c:UInt32, mark:Boolean, mark:Boolean, mark:Boolean]
+            LeftMark Join:  Filter: Boolean(true) [a:UInt32, b:UInt32, c:UInt32, mark:Boolean, mark:Boolean]
+              LeftMark Join:  Filter: test.c = __correlated_sq_1.c [a:UInt32, b:UInt32, c:UInt32, mark:Boolean]
+                TableScan: test [a:UInt32, b:UInt32, c:UInt32]
+                Projection: __correlated_sq_1.c [c:UInt32]
+                  SubqueryAlias: __correlated_sq_1 [c:UInt32]
+                    Projection: sq.c [c:UInt32]
+                      TableScan: sq [a:UInt32, b:UInt32, c:UInt32]
+              SubqueryAlias: __correlated_sq_2 [c:UInt32]
+                Filter: sq.c IS NULL [c:UInt32]
+                  Projection: sq.c [c:UInt32]
+                    TableScan: sq [a:UInt32, b:UInt32, c:UInt32]
+            SubqueryAlias: __correlated_sq_3 [c:UInt32]
+              Projection: sq.c [c:UInt32]
+                TableScan: sq [a:UInt32, b:UInt32, c:UInt32]
+        "
+        )
+    }
+
     /// Test for single NOT IN subquery filter
     #[test]
     fn not_in_subquery_simple() -> Result<()> {

datafusion/sqllogictest/test_files/predicates.slt

@@ -1001,12 +1001,14 @@ explain select x from t where x NOT IN (1,2,3,4,5) AND x IN (1,2,3);
 logical_plan EmptyRelation: rows=0
 physical_plan EmptyExec
 
-query error DataFusion error: This feature is not implemented: Physical plan does not support logical expression InSubquery\(InSubquery \{ expr: Literal\(Int64\(NULL\), None\), subquery: <subquery>, negated: false \}\)
+query BB
 WITH empty AS (SELECT 10 WHERE false)
 SELECT
     NULL IN (SELECT * FROM empty), -- should be false, as the right side is empty relation
     NULL NOT IN (SELECT * FROM empty) -- should be true, as the right side is empty relation
 FROM (SELECT 1) t;
+----
+false true
 
 query I
 WITH empty AS (SELECT 10 WHERE false)

datafusion/sqllogictest/test_files/subquery_projection.slt

@@ -0,0 +1,52 @@
+# IN subqueries in projection use SQL three-valued logic.
+
+query BB
+WITH empty AS (SELECT 10 WHERE false)
+SELECT
+    NULL IN (SELECT * FROM empty),
+    NULL NOT IN (SELECT * FROM empty)
+FROM (SELECT 1) t;
+----
+false true
+
+query BBBB
+WITH vals AS (SELECT * FROM (VALUES (1), (2), (NULL)) AS t(x))
+SELECT
+    1 IN (SELECT x FROM vals) AS found,
+    3 IN (SELECT x FROM vals) AS unknown,
+    3 NOT IN (SELECT x FROM vals) AS not_unknown,
+    NULL IN (SELECT x FROM vals) AS null_unknown;
+----
+true NULL NULL NULL
+
+query BB
+WITH vals AS (SELECT * FROM (VALUES (1), (2)) AS t(x))
+SELECT
+    EXISTS (SELECT x FROM vals WHERE x = 2),
+    NOT EXISTS (SELECT x FROM vals WHERE x = 3);
+----
+true true
+
+# Correlated IN distinguishes a match, a miss, a NULL-containing result,
+# and an empty result independently for each outer row.
+query IB rowsort
+WITH
+outer_values AS (
+    SELECT * FROM (VALUES (1), (2), (3), (4)) AS t(x)
+),
+inner_values AS (
+    SELECT * FROM (VALUES (1, 1), (2, 9), (3, NULL)) AS t(group_id, value)
+)
+SELECT
+    o.x,
+    o.x IN (
+        SELECT i.value
+        FROM inner_values i
+        WHERE i.group_id = o.x
+    )
+FROM outer_values o;
+----
+1 true
+2 false
+3 NULL
+4 false