From 4b56bc11347ce3908b08c1994f1296eb7d69699f Mon Sep 17 00:00:00 2001 From: osipovartem Date: Thu, 3 Sep 2026 23:35:07 +0300 Subject: [PATCH 1/2] Support predicate subqueries in projections --- .../src/decorrelate_predicate_subquery.rs | 137 +++++++++++++++--- .../sqllogictest/test_files/predicates.slt | 4 +- .../test_files/subquery_projection.slt | 52 +++++++ 3 files changed, 175 insertions(+), 18 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 0609109ec6e58..6d03f313684c6 100644 --- a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs +++ b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs @@ -28,8 +28,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; @@ -37,7 +37,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; @@ -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)?; + 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)); }; @@ -158,26 +189,67 @@ fn rewrite_inner_subqueries( expr, 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)? { - Some((plan, exists_expr)) => { - cur_input = plan; - Ok(Transformed::yes(exists_expr)) - } - None if negated => Ok(Transformed::no(not_in_subquery(*expr, subquery))), - None => Ok(Transformed::no(in_subquery(*expr, subquery))), + }) => match in_subquery_mark_join( + &cur_input, + &subquery, + *expr.clone(), + negated, + alias, + )? { + Some((plan, exists_expr)) => { + cur_input = plan; + Ok(Transformed::yes(exists_expr)) } - } + None if negated => Ok(Transformed::no(not_in_subquery(*expr, subquery))), + None => Ok(Transformed::no(in_subquery(*expr, subquery))), + }, _ => Ok(Transformed::no(e)), })?; Ok((cur_input, expr_without_subqueries.data)) } +fn in_subquery_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 @@ -1134,6 +1206,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<()> { diff --git a/datafusion/sqllogictest/test_files/predicates.slt b/datafusion/sqllogictest/test_files/predicates.slt index b4482a3af1beb..dfdbfc9322e8e 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/subquery_projection.slt b/datafusion/sqllogictest/test_files/subquery_projection.slt new file mode 100644 index 0000000000000..3b361841faaaa --- /dev/null +++ b/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 From f235ff3cc5c1a3e6b245019b6e10bf0e2f58e328 Mon Sep 17 00:00:00 2001 From: osipovartem Date: Fri, 4 Sep 2026 00:38:40 +0300 Subject: [PATCH 2/2] Preserve filter subquery plans --- .../src/decorrelate_predicate_subquery.rs | 45 ++++++++++++------- 1 file changed, 29 insertions(+), 16 deletions(-) diff --git a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs index 6d03f313684c6..e68ff114543ad 100644 --- a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs +++ b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs @@ -80,7 +80,7 @@ impl OptimizerRule for DecorrelatePredicateSubquery { 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)?; + rewrite_inner_subqueries(cur_input, expr, config, true)?; if has_subquery(&rewritten_expr) { return Ok(Transformed::no(LogicalPlan::Projection( original_projection, @@ -135,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); } @@ -170,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(); @@ -189,26 +190,38 @@ fn rewrite_inner_subqueries( expr, subquery: Subquery { subquery, .. }, negated, - }) => match in_subquery_mark_join( - &cur_input, - &subquery, - *expr.clone(), - negated, - alias, - )? { - Some((plan, exists_expr)) => { - cur_input = plan; - Ok(Transformed::yes(exists_expr)) + }) => { + 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)) + } + None if negated => Ok(Transformed::no(not_in_subquery(*expr, subquery))), + None => Ok(Transformed::no(in_subquery(*expr, subquery))), } - None if negated => Ok(Transformed::no(not_in_subquery(*expr, subquery))), - None => Ok(Transformed::no(in_subquery(*expr, subquery))), - }, + } _ => Ok(Transformed::no(e)), })?; Ok((cur_input, expr_without_subqueries.data)) } -fn in_subquery_mark_join( +fn in_subquery_value_mark_join( left: &LogicalPlan, subquery: &LogicalPlan, expr: Expr,