diff --git a/datafusion/expr/src/utils.rs b/datafusion/expr/src/utils.rs index 44e4fb4415dc5..03fada098fde7 100644 --- a/datafusion/expr/src/utils.rs +++ b/datafusion/expr/src/utils.rs @@ -1505,6 +1505,10 @@ pub fn add_filter(plan: LogicalPlan, predicates: &[&Expr]) -> Result) -> Result<(Vec, Vec)> { let mut joins = vec![]; let mut others = vec![]; diff --git a/datafusion/optimizer/src/decorrelate.rs b/datafusion/optimizer/src/decorrelate.rs index b81f5cb44708a..af5de0b29ac2a 100644 --- a/datafusion/optimizer/src/decorrelate.rs +++ b/datafusion/optimizer/src/decorrelate.rs @@ -22,6 +22,7 @@ use std::sync::Arc; use crate::push_down_filter::lr_is_preserved; use crate::simplify_expressions::ExprSimplifier; +use crate::utils::replace_qualified_name; use datafusion_common::nested_struct::is_injective_cast; use datafusion_common::tree_node::{ @@ -33,11 +34,10 @@ use datafusion_common::{ plan_err, }; use datafusion_expr::expr::{Alias, GroupingSet}; +use datafusion_expr::expr_rewriter::strip_outer_reference; use datafusion_expr::logical_plan::{Join, JoinType}; use datafusion_expr::simplify::SimplifyContext; -use datafusion_expr::utils::{ - collect_subquery_cols, conjunction, find_join_exprs, split_conjunction, -}; +use datafusion_expr::utils::{collect_subquery_cols, conjunction, split_conjunction}; use datafusion_expr::{ Aggregate, BinaryExpr, Cast, Distinct, EmptyRelation, Expr, ExprSchemable, Extension, FetchType, LogicalPlan, LogicalPlanBuilder, MaterializedCte, Operator, SkipType, @@ -408,7 +408,7 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr { } } let (mut join_filters, subquery_filters) = - find_join_exprs(subquery_filter_exprs)?; + find_join_exprs(subquery_filter_exprs); if let Some(in_predicate) = &self.in_predicate_opt { // in_predicate may be already included in the join filters, remove it from the join filters first. join_filters = remove_duplicated_filter(join_filters, in_predicate)?; @@ -1124,6 +1124,60 @@ fn collect_local_correlated_cols( } } +/// Extracts correlated predicates, such as a comparison between a column from the +/// subquery and a column from the outer scope. +/// +/// Preserves [`Expr::OuterReferenceColumn`] markers so callers can distinguish +/// inner and outer columns when qualifying join filters (see [`build_join_filter`]). +/// +/// # Arguments +/// +/// * `exprs` - Subquery filter predicates to classify. +/// +/// # Return value +/// +/// Tuple of (correlated join filters, remaining subquery filters). Correlated +/// self-equalities are discarded from both lists. +pub(crate) fn find_join_exprs(exprs: Vec<&Expr>) -> (Vec, Vec) { + let mut joins = vec![]; + let mut others = vec![]; + for filter in exprs { + // Predicates containing outer references become join filters. + if filter.contains_outer() { + // Compare with the outer-reference markers intact: inner and outer + // columns with the same qualified name belong to different scopes. + if !matches!(filter, Expr::BinaryExpr(BinaryExpr { left, op: Operator::Eq, right }) if left.eq(right)) + { + joins.push(filter.clone()); + } + } else { + others.push(filter.clone()); + } + } + (joins, others) +} + +/// Combines join filters and qualifies their inner columns with the subquery alias. +/// +/// Aliasing must precede stripping outer references so inner and outer columns +/// with the same qualified name remain distinguishable. +pub(crate) fn build_join_filter<'a>( + join_filters: Vec, + correlated_cols: impl IntoIterator, + subquery_alias: &str, +) -> Result> { + let correlated_cols = correlated_cols + .into_iter() + .cloned() + .collect::>(); + conjunction(join_filters) + .map(|filter| { + replace_qualified_name(filter, &correlated_cols, subquery_alias) + .map(strip_outer_reference) + }) + .transpose() +} + fn remove_duplicated_filter( filters: Vec, in_predicate: &Expr, @@ -1141,6 +1195,7 @@ fn remove_duplicated_filter( Ok(filters .into_iter() .filter(|filter| { + let filter = &strip_outer_reference(filter.clone()); if filter == in_predicate { return false; } @@ -1305,3 +1360,38 @@ fn filter_exprs_evaluation_result_on_empty_batch( }; Ok(pull_up_expr) } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::datatypes::DataType; + use datafusion_expr::{col, out_ref_col}; + + #[test] + fn test_find_join_exprs_preserves_outer_references() { + let outer = out_ref_col(DataType::Int32, "t.a"); + let correlated = col("t.a").eq(outer.clone()); + let self_equality = outer.clone().eq(outer); + let local = col("t.b").gt(lit(2)); + + assert_eq!( + find_join_exprs(vec![&correlated, &self_equality, &local]), + (vec![correlated], vec![local]) + ); + } + + #[test] + fn test_build_join_filter() -> Result<()> { + assert_eq!(build_join_filter(vec![], [], "sq")?, None); + + let column = Column::from("t.a"); + let outer = out_ref_col(DataType::Int32, "t.a"); + let filters = vec![col("t.a").eq(outer.clone()), outer.gt(lit(1))]; + + assert_eq!( + build_join_filter(filters, [&column, &column], "sq")?, + Some(col("sq.a").eq(col("t.a")).and(col("t.a").gt(lit(1)))) + ); + Ok(()) + } +} diff --git a/datafusion/optimizer/src/decorrelate_lateral_join.rs b/datafusion/optimizer/src/decorrelate_lateral_join.rs index fe484ff7373a6..aa81b1cc87f78 100644 --- a/datafusion/optimizer/src/decorrelate_lateral_join.rs +++ b/datafusion/optimizer/src/decorrelate_lateral_join.rs @@ -23,6 +23,7 @@ use crate::decorrelate::{PullUpCorrelatedExpr, UN_MATCHED_ROW_INDICATOR}; use crate::optimizer::ApplyOrder; use crate::utils::evaluates_to_null; use crate::{OptimizerConfig, OptimizerRule}; +use datafusion_expr::expr_rewriter::strip_outer_reference; use datafusion_expr::{Expr, Join, expr}; use datafusion_common::tree_node::{ @@ -149,6 +150,8 @@ fn rewrite_internal(join: Join) -> Result> { (rewritten_subquery, correlation_filter, original_join_filter) }; + let correlation_filter = correlation_filter.map(strip_outer_reference); + // For LEFT lateral joins, verify that all column references in the // correlation filter are resolvable within the join's left and right // schemas. If the lateral subquery references columns from an outer scope, diff --git a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs index 9c20108da261c..febbd3a6c84b5 100644 --- a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs +++ b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs @@ -21,10 +21,11 @@ use std::ops::Deref; use std::sync::Arc; use crate::analyzer::type_coercion::TypeCoercionRewriter; -use crate::decorrelate::{PullUpCorrelatedExpr, UN_MATCHED_ROW_INDICATOR}; +use crate::decorrelate::{ + PullUpCorrelatedExpr, UN_MATCHED_ROW_INDICATOR, build_join_filter, +}; use crate::extract_equijoin_predicate::split_eq_and_noneq_join_predicate; use crate::optimizer::ApplyOrder; -use crate::utils::replace_qualified_name; use crate::{OptimizerConfig, OptimizerRule}; use datafusion_common::alias::AliasGenerator; @@ -772,7 +773,11 @@ fn subquery_operands_are_bare_columns( filter .apply(|e| { let refs = e.column_refs(); - if !refs.is_empty() && refs.iter().all(|c| subquery_cols.contains(*c)) { + // `column_refs` skips outer references, which the join filters keep, so check them separately. + if !refs.is_empty() + && !e.contains_outer() + && refs.iter().all(|c| subquery_cols.contains(*c)) + { bare = matches!(e, Expr::Column(_)); return Ok(if bare { TreeNodeRecursion::Jump @@ -853,17 +858,11 @@ fn build_join( let sub_query_alias = LogicalPlanBuilder::from(new_plan) .alias(alias.to_string())? .build()?; - let mut all_correlated_cols = BTreeSet::new(); - pull_up - .correlated_subquery_cols_map - .values() - .for_each(|cols| all_correlated_cols.extend(cols.iter().cloned())); - - // alias the join filter - let join_filter_opt = conjunction(pull_up.join_filters) - .map_or(Ok(None), |filter| { - replace_qualified_name(filter, &all_correlated_cols, alias).map(Some) - })?; + let join_filter_opt = build_join_filter( + pull_up.join_filters, + pull_up.correlated_subquery_cols_map.values().flatten(), + alias, + )?; if let Some(expr_map) = count_bug_compensation && !expr_map.is_empty() @@ -2662,6 +2661,34 @@ mod tests { ) } + #[test] + fn in_subquery_correlated_with_same_table() -> Result<()> { + let outer_scan = test_table_scan()?; + let subquery_scan = test_table_scan()?; + let subquery = LogicalPlanBuilder::from(subquery_scan) + .filter(out_ref_col(DataType::UInt32, "test.a").gt(col("test.a")))? + .project(vec![col("c")])? + .build()?; + + let plan = LogicalPlanBuilder::from(outer_scan) + .filter(in_subquery(col("test.b"), Arc::new(subquery)))? + .project(vec![col("test.b")])? + .build()?; + + // The outer `test.a` and the subquery's `test.a` stay apart in the join filter. + assert_optimized_plan_equal!( + plan, + @r" + Projection: test.b [b:UInt32] + LeftSemi Join: Filter: test.b = __correlated_sq_1.c AND test.a > __correlated_sq_1.a [a:UInt32, b:UInt32, c:UInt32] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + SubqueryAlias: __correlated_sq_1 [c:UInt32, a:UInt32] + Projection: test.c, test.a [c:UInt32, a:UInt32] + TableScan: test [a:UInt32, b:UInt32, c:UInt32] + " + ) + } + /// Test for multiple exists subqueries in the same filter expression #[test] fn multiple_exists_subqueries() -> Result<()> { diff --git a/datafusion/optimizer/src/scalar_subquery_to_join.rs b/datafusion/optimizer/src/scalar_subquery_to_join.rs index 91383a24a6e80..0a1ea917ce987 100644 --- a/datafusion/optimizer/src/scalar_subquery_to_join.rs +++ b/datafusion/optimizer/src/scalar_subquery_to_join.rs @@ -17,12 +17,14 @@ //! [`ScalarSubqueryToJoin`] rewriting scalar subquery filters to `JOIN`s -use std::collections::{BTreeSet, HashMap}; +use std::collections::HashMap; use std::sync::Arc; -use crate::decorrelate::{PullUpCorrelatedExpr, UN_MATCHED_ROW_INDICATOR}; +use crate::decorrelate::{ + PullUpCorrelatedExpr, UN_MATCHED_ROW_INDICATOR, build_join_filter, +}; use crate::optimizer::ApplyOrder; -use crate::utils::{evaluates_to_null, replace_qualified_name}; +use crate::utils::evaluates_to_null; use crate::{OptimizerConfig, OptimizerRule}; use crate::analyzer::type_coercion::TypeCoercionRewriter; @@ -33,7 +35,6 @@ use datafusion_common::tree_node::{ use datafusion_common::{Column, Result, ScalarValue, assert_or_internal_err, plan_err}; use datafusion_expr::expr_rewriter::create_col_from_scalar_expr; use datafusion_expr::logical_plan::{JoinType, Subquery}; -use datafusion_expr::utils::conjunction; use datafusion_expr::{Expr, LogicalPlan, LogicalPlanBuilder, lit, not, when}; /// Optimizer rule that rewrites scalar subquery filters to joins and places an @@ -361,19 +362,11 @@ fn build_join( .alias(subquery_alias.to_string())? .build()?; - let all_correlated_cols: BTreeSet = pull_up - .correlated_subquery_cols_map - .values() - .flatten() - .cloned() - .collect(); - - // Correlated columns now live in the decorrelated subquery's output, - // so re-qualify them with the subquery alias. - let join_filter_opt = - conjunction(pull_up.join_filters).map_or(Ok(None), |filter| { - replace_qualified_name(filter, &all_correlated_cols, subquery_alias).map(Some) - })?; + let join_filter_opt = build_join_filter( + pull_up.join_filters, + pull_up.correlated_subquery_cols_map.values().flatten(), + subquery_alias, + )?; // When pull-up did not extract any usable join keys (a correlated subquery // whose predicate references only outer columns), fall back to `ON true`: diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index 45ec0345f75a1..d4a42cca0c801 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -4076,3 +4076,122 @@ DROP TABLE nested_agg_outer; statement ok DROP TABLE nested_agg_inner; + +# Regression tests for #25258 - Unaliased set comparison subqueries return wrong results + +statement ok +SET datafusion.explain.logical_plan_only = true; + +statement ok +CREATE TABLE emp (id INT, name TEXT, salary INT); + +statement ok +INSERT INTO emp VALUES +(1,'Alice',5000), +(2,'Bob',6000), +(3,'Carol',7000), +(4,'Dave',8000), +(5,'Eve',9000), +(6,'Frank',10000); + +## Un-aliased + +query I +SELECT id +FROM emp +WHERE salary > ANY (SELECT salary FROM emp WHERE id <= 3); +---- +2 +3 +4 +5 +6 + +query TT +EXPLAIN SELECT id +FROM emp +WHERE salary > ANY (SELECT salary FROM emp WHERE id <= 3); +---- +logical_plan +01)Projection: emp.id +02)--LeftMark Join: Filter: emp.salary > __correlated_sq_2.salary IS NULL +03)----Projection: emp.id, emp.salary +04)------Filter: __correlated_sq_1.mark IS NOT DISTINCT FROM Boolean(true) +05)--------LeftMark Join: Filter: emp.salary > __correlated_sq_1.salary IS TRUE +06)----------TableScan: emp projection=[id, salary] +07)----------SubqueryAlias: __correlated_sq_1 +08)------------Projection: emp.salary +09)--------------Filter: emp.id <= Int32(3) +10)----------------TableScan: emp projection=[id, salary] +11)----SubqueryAlias: __correlated_sq_2 +12)------Projection: emp.salary +13)--------Filter: emp.id <= Int32(3) +14)----------TableScan: emp projection=[id, salary] + +query I +SELECT id +FROM emp +WHERE salary > ALL (SELECT salary FROM emp WHERE id <= 2); +---- +3 +4 +5 +6 + +query TT +EXPLAIN SELECT id +FROM emp +WHERE salary > ALL (SELECT salary FROM emp WHERE id <= 2); +---- +logical_plan +01)Projection: emp.id +02)--Filter: __correlated_sq_2.mark IS DISTINCT FROM Boolean(true) +03)----Projection: emp.id, __correlated_sq_2.mark +04)------LeftMark Join: Filter: emp.salary > __correlated_sq_2.salary IS NULL +05)--------Projection: emp.id, emp.salary +06)----------Filter: __correlated_sq_1.mark IS DISTINCT FROM Boolean(true) +07)------------LeftMark Join: Filter: emp.salary > __correlated_sq_1.salary IS FALSE +08)--------------TableScan: emp projection=[id, salary] +09)--------------SubqueryAlias: __correlated_sq_1 +10)----------------Projection: emp.salary +11)------------------Filter: emp.id <= Int32(2) +12)--------------------TableScan: emp projection=[id, salary] +13)--------SubqueryAlias: __correlated_sq_2 +14)----------Projection: emp.salary +15)------------Filter: emp.id <= Int32(2) +16)--------------TableScan: emp projection=[id, salary] + +query I +SELECT id +FROM emp +WHERE salary = ANY (SELECT salary FROM emp WHERE id <= 2); +---- +1 +2 + +query TT +EXPLAIN SELECT id +FROM emp +WHERE salary = ANY (SELECT salary FROM emp WHERE id <= 2); +---- +logical_plan +01)Projection: emp.id +02)--LeftMark Join: Filter: emp.salary = __correlated_sq_2.salary IS NULL +03)----Projection: emp.id, emp.salary +04)------Filter: __correlated_sq_1.mark IS NOT DISTINCT FROM Boolean(true) +05)--------LeftMark Join: Filter: emp.salary = __correlated_sq_1.salary IS TRUE +06)----------TableScan: emp projection=[id, salary] +07)----------SubqueryAlias: __correlated_sq_1 +08)------------Projection: emp.salary +09)--------------Filter: emp.id <= Int32(2) +10)----------------TableScan: emp projection=[id, salary] +11)----SubqueryAlias: __correlated_sq_2 +12)------Projection: emp.salary +13)--------Filter: emp.id <= Int32(2) +14)----------TableScan: emp projection=[id, salary] + +statement ok +DROP TABLE emp; + +statement ok +RESET datafusion.explain.logical_plan_only;