Support predicate subqueries in projections

datafusion/optimizer/src/decorrelate_predicate_subquery.rs

@@ -29,16 +29,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;
@@ -70,6 +70,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));
         };
@@ -105,7 +136,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);
                 }
@@ -140,6 +171,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();
@@ -160,12 +192,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))
@@ -179,6 +222,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
@@ -428,16 +513,17 @@ fn build_join(
 
         // Keep only columns that actually belong to the RIGHT child, and sort by their
         // position in the right schema for deterministic order.
-        let mut right_cols_idx_and_col: Vec<(usize, Column)> = needed
+        let mut right_col_indices: Vec<usize> = needed
             .into_iter()
-            .filter_map(|c| right_schema.index_of_column(&c).ok().map(|idx| (idx, c)))
+            .filter_map(|column| right_schema.index_of_column(&column).ok())
             .collect();
 
-        right_cols_idx_and_col.sort_by_key(|(idx, _)| *idx);
+        right_col_indices.sort_unstable();
+        right_col_indices.dedup();
 
-        let right_proj_exprs: Vec<Expr> = right_cols_idx_and_col
+        let right_proj_exprs: Vec<Expr> = right_col_indices
             .into_iter()
-            .map(|(_, c)| Expr::Column(c))
+            .map(|index| Expr::Column(Column::from(right_schema.qualified_field(index))))
             .collect();
 
         let right_projected = if !right_proj_exprs.is_empty() {
@@ -1203,6 +1289,63 @@ 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;N, mark:Boolean;N, mark:Boolean;N]
+            LeftMark Join:  Filter: Boolean(true) [a:UInt32, b:UInt32, c:UInt32, mark:Boolean;N, mark:Boolean;N]
+              LeftMark Join:  Filter: test.c = __correlated_sq_1.c [a:UInt32, b:UInt32, c:UInt32, mark:Boolean;N]
+                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]
+    fn unsupported_correlated_in_projection_is_left_unchanged() -> Result<()> {
+        let subquery = Arc::new(
+            LogicalPlanBuilder::from(scan_tpch_table("orders"))
+                .filter(
+                    out_ref_col(DataType::Int64, "customer.c_custkey")
+                        .eq(col("orders.o_custkey")),
+                )?
+                .limit(0, Some(1))?
+                .project(vec![col("orders.o_custkey")])?
+                .build()?,
+        );
+        let plan = LogicalPlanBuilder::from(scan_tpch_table("customer"))
+            .project(vec![
+                in_subquery(col("customer.c_custkey"), subquery).alias("is_present"),
+            ])?
+            .build()?;
+
+        let result = DecorrelatePredicateSubquery::new()
+            .rewrite(plan.clone(), &crate::OptimizerContext::new())?;
+
+        assert!(!result.transformed);
+        assert_eq!(result.data, plan);
+        Ok(())
+    }
+
     /// 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/projection_pushdown.slt

@@ -2295,11 +2295,11 @@ SET datafusion.execution.target_partitions = 4;
 # Without the fix, extraction reuses `__datafusion_extracted_1` and planning
 # aborts with: Optimizer rule 'push_down_leaf_projections' failed Schema error:
 # Schema contains duplicate unqualified field name __datafusion_extracted_1.
-# With the fix, the generator starts at 2, as the plan below shows.
+# With the fix, generated aliases remain distinct, as the plan below shows.
 #
-# This is `EXPLAIN` under `logical_plan_only` rather than an executed query
-# because a surviving `InSubquery` expression has no physical plan; the logical
-# plan is both the observable result and exactly what regressed.
+# Keep this as `EXPLAIN` under `logical_plan_only`: the logical plan exposes both
+# the collision-free extracted aliases and the mark joins used to preserve the
+# three-valued semantics of `IN` in a projection.
 #####################
 
 statement ok
@@ -2321,13 +2321,22 @@ SELECT
 FROM simple_struct;
 ----
 logical_plan
-01)Projection: simple_struct.id, simple_struct.id IN (<subquery>) AS has_matching_label
-02)--Subquery:
-03)----Projection: simple_struct.id
-04)------Filter: __datafusion_extracted_2 > Int64(120) AND __datafusion_extracted_1 != Utf8("delta")
-05)--------Projection: get_field(simple_struct.s, Utf8("value")) AS __datafusion_extracted_2, simple_struct.id, get_field(simple_struct.s, Utf8("label")) AS __datafusion_extracted_1
-06)----------TableScan: simple_struct projection=[id, s], partial_filters=[get_field(simple_struct.s, Utf8("value")) > Int64(120)]
-07)--TableScan: simple_struct projection=[id]
+01)Projection: simple_struct.id, __correlated_sq_1.mark IS NOT DISTINCT FROM Boolean(true) OR __correlated_sq_2.mark IS NOT DISTINCT FROM Boolean(true) AND __correlated_sq_1.mark IS DISTINCT FROM Boolean(true) AND Boolean(NULL) AS has_matching_label
+02)--LeftMark Join:
+03)----LeftMark Join:
+04)------LeftMark Join: simple_struct.id = __correlated_sq_1.id
+05)--------TableScan: simple_struct projection=[id]
+06)--------SubqueryAlias: __correlated_sq_1
+07)----------Projection: simple_struct.id
+08)------------Filter: __datafusion_extracted_4 > Int64(120) AND __datafusion_extracted_1 != Utf8("delta")
+09)--------------Projection: get_field(simple_struct.s, Utf8("value")) AS __datafusion_extracted_4, simple_struct.id, get_field(simple_struct.s, Utf8("label")) AS __datafusion_extracted_1
+10)----------------TableScan: simple_struct projection=[id, s], partial_filters=[get_field(simple_struct.s, Utf8("value")) > Int64(120)]
+11)------EmptyRelation: rows=0
+12)----SubqueryAlias: __correlated_sq_3
+13)------Projection: simple_struct.id
+14)--------Filter: __datafusion_extracted_6 > Int64(120) AND __datafusion_extracted_1 != Utf8("delta")
+15)----------Projection: get_field(simple_struct.s, Utf8("value")) AS __datafusion_extracted_6, simple_struct.id, get_field(simple_struct.s, Utf8("label")) AS __datafusion_extracted_1
+16)------------TableScan: simple_struct projection=[id, s], partial_filters=[Boolean(true), get_field(simple_struct.s, Utf8("value")) > Int64(120)]
 
 statement ok
 set datafusion.explain.logical_plan_only = false;

datafusion/sqllogictest/test_files/subquery_projection.slt

@@ -0,0 +1,57 @@
+# 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 B
+SELECT 'a' IN (SELECT column1 FROM VALUES ('b'), ('c'), ('d'));
+----
+false
+
+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