From 25ac855c00a6b1cabb2cfc99041688befe8a27c6 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Sun, 6 Sep 2026 12:26:13 +0300 Subject: [PATCH 1/2] Support predicate subqueries in projections --- .../src/decorrelate_predicate_subquery.rs | 147 ++++++++++++++++-- .../sqllogictest/test_files/predicates.slt | 4 +- .../test_files/projection_pushdown.slt | 31 ++-- .../test_files/subquery_projection.slt | 57 +++++++ 4 files changed, 212 insertions(+), 27 deletions(-) create mode 100644 datafusion/sqllogictest/test_files/subquery_projection.slt diff --git a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs index 5f623f1bef6f6..7165215d80691 100644 --- a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs +++ b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs @@ -29,8 +29,8 @@ 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; @@ -38,7 +38,7 @@ 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, +) -> Result> { + 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 = 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 = right_cols_idx_and_col + let right_proj_exprs: Vec = 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,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;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 for single NOT IN subquery filter #[test] fn not_in_subquery_simple() -> Result<()> { diff --git a/datafusion/sqllogictest/test_files/predicates.slt b/datafusion/sqllogictest/test_files/predicates.slt index 9e28d1fb08044..9530ed2d92475 100644 --- a/datafusion/sqllogictest/test_files/predicates.slt +++ b/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: , 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) diff --git a/datafusion/sqllogictest/test_files/projection_pushdown.slt b/datafusion/sqllogictest/test_files/projection_pushdown.slt index 1f9176f7137bd..c92c95fdfbc59 100644 --- a/datafusion/sqllogictest/test_files/projection_pushdown.slt +++ b/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 () 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; diff --git a/datafusion/sqllogictest/test_files/subquery_projection.slt b/datafusion/sqllogictest/test_files/subquery_projection.slt new file mode 100644 index 0000000000000..20de82531561a --- /dev/null +++ b/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 From 707d9df8732f7b2c09d3aa1b1797fdc99a11119a Mon Sep 17 00:00:00 2001 From: osipovartem Date: Sun, 6 Sep 2026 17:47:20 +0300 Subject: [PATCH 2/2] test: cover unsupported projection subquery fallback --- .../src/decorrelate_predicate_subquery.rs | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs index 7165215d80691..4a12b4ab7b17a 100644 --- a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs +++ b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs @@ -1320,6 +1320,32 @@ mod tests { ) } + #[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<()> {