diff --git a/datafusion/expr/src/utils.rs b/datafusion/expr/src/utils.rs index 4ca12c5e0339e..f71f35be1b63f 100644 --- a/datafusion/expr/src/utils.rs +++ b/datafusion/expr/src/utils.rs @@ -1427,6 +1427,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 0c37f00b64355..08efaf4a12fd0 100644 --- a/datafusion/optimizer/src/decorrelate.rs +++ b/datafusion/optimizer/src/decorrelate.rs @@ -21,6 +21,7 @@ use std::collections::BTreeSet; use std::sync::Arc; use crate::simplify_expressions::ExprSimplifier; +use crate::utils::replace_qualified_name; use datafusion_common::tree_node::{ Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeRewriter, @@ -29,10 +30,9 @@ use datafusion_common::{ Column, DFSchemaRef, HashMap, Result, ScalarValue, assert_or_internal_err, plan_err, }; use datafusion_expr::expr::Alias; +use datafusion_expr::expr_rewriter::strip_outer_reference; 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::{ BinaryExpr, Cast, EmptyRelation, Expr, ExprSchemable, FetchType, LogicalPlan, LogicalPlanBuilder, Operator, expr, lit, @@ -183,7 +183,7 @@ impl TreeNodeRewriter for PullUpCorrelatedExpr { .filter(|e| e.contains_outer()) .all(|&e| can_pullup_over_aggregation(e)); 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)?; @@ -480,6 +480,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() { + // Check equality before stripping markers: inner and outer columns + // with the same qualified name still 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, @@ -496,6 +550,7 @@ fn remove_duplicated_filter( Ok(filters .into_iter() + .map(strip_outer_reference) .filter(|filter| { if filter == in_predicate { return false; @@ -652,3 +707,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 a8df5e69e3f33..f83e20a546aaf 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::{ @@ -153,6 +154,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 5f623f1bef6f6..85efb230851be 100644 --- a/datafusion/optimizer/src/decorrelate_predicate_subquery.rs +++ b/datafusion/optimizer/src/decorrelate_predicate_subquery.rs @@ -16,14 +16,12 @@ // under the License. //! [`DecorrelatePredicateSubquery`] converts `IN`/`EXISTS` subquery predicates to `SEMI`/`ANTI` joins -use std::collections::BTreeSet; use std::ops::Deref; use std::sync::Arc; -use crate::decorrelate::PullUpCorrelatedExpr; +use crate::decorrelate::{PullUpCorrelatedExpr, 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; @@ -374,17 +372,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.clone())); - - // 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, + )?; let join_filter = match (join_filter_opt, in_predicate_opt.cloned()) { ( @@ -765,7 +757,7 @@ mod tests { SubqueryAlias: __correlated_sq_2 [o_custkey:Int64] Projection: orders.o_custkey [o_custkey:Int64] TableScan: orders [o_orderkey:Int64, o_custkey:Int64, o_orderstatus:Utf8, o_totalprice:Float64;N] - " + " ) } diff --git a/datafusion/optimizer/src/scalar_subquery_to_join.rs b/datafusion/optimizer/src/scalar_subquery_to_join.rs index 44011a125ba96..e91a699f19129 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 @@ -365,19 +366,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 7aa24d4c7fe37..536c9ac2ce357 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -2737,3 +2737,138 @@ b 400 statement ok DROP TABLE metrics; + +# 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_3.salary IS TRUE +03)----Projection: emp.id, emp.salary +04)------LeftMark Join: Filter: emp.salary > __correlated_sq_2.salary IS NULL +05)--------Projection: emp.id, emp.salary +06)----------Filter: __correlated_sq_1.mark +07)------------LeftMark Join: Filter: emp.salary > __correlated_sq_1.salary IS TRUE +08)--------------TableScan: emp projection=[id, salary] +09)--------------SubqueryAlias: __correlated_sq_1 +10)----------------Projection: emp.salary +11)------------------Filter: emp.id <= Int32(3) +12)--------------------TableScan: emp projection=[id, salary] +13)--------SubqueryAlias: __correlated_sq_2 +14)----------Projection: emp.salary +15)------------Filter: emp.id <= Int32(3) +16)--------------TableScan: emp projection=[id, salary] +17)----SubqueryAlias: __correlated_sq_3 +18)------Projection: emp.salary +19)--------Filter: emp.id <= Int32(3) +20)----------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: NOT __correlated_sq_3.mark +03)----Projection: emp.id, __correlated_sq_3.mark +04)------LeftMark Join: Filter: emp.salary > __correlated_sq_3.salary IS NULL +05)--------Projection: emp.id, emp.salary +06)----------LeftMark Join: Filter: emp.salary > __correlated_sq_2.salary IS NULL +07)------------LeftAnti 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] +17)--------SubqueryAlias: __correlated_sq_3 +18)----------Projection: emp.salary +19)------------Filter: emp.id <= Int32(2) +20)--------------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_3.salary IS TRUE +03)----Projection: emp.id, emp.salary +04)------LeftMark Join: Filter: emp.salary = __correlated_sq_2.salary IS NULL +05)--------Projection: emp.id, emp.salary +06)----------Filter: __correlated_sq_1.mark +07)------------LeftMark Join: Filter: emp.salary = __correlated_sq_1.salary IS TRUE +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] +17)----SubqueryAlias: __correlated_sq_3 +18)------Projection: emp.salary +19)--------Filter: emp.id <= Int32(2) +20)----------TableScan: emp projection=[id, salary] + +statement ok +DROP TABLE emp; + +statement ok +RESET datafusion.explain.logical_plan_only;