diff --git a/datafusion/core/tests/sql/joins.rs b/datafusion/core/tests/sql/joins.rs index dd32830dd5eb8..38f03a6b9f145 100644 --- a/datafusion/core/tests/sql/joins.rs +++ b/datafusion/core/tests/sql/joins.rs @@ -27,6 +27,71 @@ use datafusion_sql::unparser::plan_to_sql; use super::*; +#[tokio::test] +async fn volatile_join_filter_preserves_evaluations_below_min() -> Result<()> { + use arrow::array::record_batch; + use datafusion::logical_expr::{ColumnarValue, Volatility, create_udf}; + use std::sync::atomic::{AtomicUsize, Ordering}; + + // Request sort-merge joins; these small inputs remain in a single partition. + let mut config = SessionConfig::new().with_target_partitions(2); + config.options_mut().optimizer.prefer_hash_join = false; + let ctx = SessionContext::new_with_config(config); + + // Return true, false, true, ... across successive input rows, including + // across batches. Both arguments keep the predicate at the join. + let evaluations = AtomicUsize::new(0); + ctx.register_udf(create_udf( + "alternating", + vec![DataType::Int32, DataType::Int32], + DataType::Boolean, + Volatility::Volatile, + Arc::new(move |args| { + let args = ColumnarValue::values_to_arrays(args)?; + let len = args[0].len(); + let first = evaluations.fetch_add(len, Ordering::Relaxed); + let values = BooleanArray::from_iter( + (first..first + len).map(|i| Some(i.is_multiple_of(2))), + ); + Ok(ColumnarValue::Array(Arc::new(values))) + }), + )); + + ctx.register_batch( + "l", + record_batch!(("id", Int32, vec![1, 2]), ("x", Int32, vec![20, 10]))?, + )?; + ctx.register_batch( + "r", + record_batch!(("id", Int32, vec![1, 1, 2]), ("y", Int32, vec![0, 0, 0]))?, + )?; + + // The inner join evaluates both pairs for id=1, so id=2 receives the + // third (true) result. A semi join would stop after the first match for + // id=1, give id=2 the second (false) result, and incorrectly return 20. + let df = ctx + .sql( + "SELECT MIN(l.x) AS minimum FROM l JOIN r \ + ON l.id = r.id AND alternating(l.x, r.y)", + ) + .await?; + let plan = df.create_physical_plan().await?; + let formatted = displayable(plan.as_ref()).indent(true).to_string(); + assert_contains!(formatted, "SortMergeJoinExec: join_type=Inner"); + let batches = collect(plan, ctx.task_ctx()).await?; + assert_batches_eq!( + [ + "+---------+", + "| minimum |", + "+---------+", + "| 10 |", + "+---------+", + ], + &batches + ); + Ok(()) +} + #[tokio::test] async fn join_change_in_planner() -> Result<()> { let config = SessionConfig::new().with_target_partitions(8); diff --git a/datafusion/optimizer/src/eliminate_join.rs b/datafusion/optimizer/src/eliminate_join.rs index 56aa8887065be..c1626314caa23 100644 --- a/datafusion/optimizer/src/eliminate_join.rs +++ b/datafusion/optimizer/src/eliminate_join.rs @@ -28,9 +28,10 @@ //! //! 1. None of R's columns are referenced above the join. //! 2. R does not observably multiply L's rows. This holds when either the -//! join's ancestors are duplicate-insensitive (e.g., DISTINCT) or we can use -//! functional dependencies to prove that each L row matches at most one R -//! row (R is provably unique on the join keys). +//! join's ancestors are duplicate-insensitive (e.g., DISTINCT) and its +//! conditions are repeatable, or we can use functional dependencies to +//! prove that each L row matches at most one R row (R is provably unique +//! on the join keys). //! //! * A left outer join `L ⟕ R` can be removed entirely, i.e. replaced by `L`, //! under the same two conditions. Unlike an inner join, a left join @@ -38,11 +39,11 @@ //! columns are unused and R cannot multiply L's rows the join has no //! observable effect at all. Such joins commonly appear in generated SQL //! and in queries over views that join in lookup tables the query does not -//! read. A join filter does not prevent this rewrite: for a left join it -//! only decides whether a left row is matched or null-padded, and either -//! way the row is emitted. Symmetrically, a right outer join `L ⟖ R` can be -//! replaced by `R` when L's columns are unused and L cannot multiply R's -//! rows. +//! read. A repeatable join filter does not prevent this rewrite: for a left +//! join it only decides whether a left row is matched or null-padded, and +//! either way the row is emitted. Symmetrically, a right outer join `L ⟖ R` +//! can be replaced by `R` when L's columns are unused and L cannot multiply +//! R's rows. //! //! # Overview //! @@ -56,11 +57,12 @@ //! set across its two inputs. //! * `duplicate_insensitive` — whether emitting each row once instead of many //! times will not change the output. A duplicate-collapsing node (e.g., -//! DISTINCT, GROUP BY with no aggregate functions, or the existence side of a -//! semi/anti/mark join) sets it `true` for its subtree, and it propagates -//! downward until a node that makes the row count observable again (a `LIMIT`, -//! a top-N sort, ...) clears it. It is therefore fixed by the nearest such -//! node, not by the whole ancestor chain: a collapsing node shields its subtree, +//! DISTINCT, an `Aggregate` plan node whose aggregate expressions all ignore +//! duplicate input rows, or the existence side of a semi/anti/mark join) sets +//! it `true` for its subtree, and it propagates downward until a node that +//! makes the row count observable again (a `LIMIT`, a top-N sort, a volatile +//! expression, ...) clears it. It is therefore fixed by the nearest such node, +//! not by the whole ancestor chain: a collapsing node shields its subtree, //! so a duplicate-sensitive node further above does not matter. //! //! At each join, `rewritten_join_type` combines this context with the side's @@ -69,7 +71,9 @@ //! node types just forward the context to their single child via //! `rewrite_single_input`; nodes that alter column requirements or //! duplicate-sensitivity (projection, aggregate, sort, ...) adjust it first. -use crate::utils::for_each_referenced_index; +use crate::utils::{ + for_each_referenced_index, is_duplicate_insensitive_aggregate, is_repeatable, +}; use crate::{OptimizerConfig, OptimizerRule}; use datafusion_common::tree_node::{Transformed, TreeNode}; use datafusion_common::{ @@ -219,6 +223,8 @@ fn rewrite_node( }) => { // Narrows `live` to the columns the projection's expressions reference. let child_live = LiveColumns::try_new(&expr, input.schema())?; + let duplicate_insensitive = + duplicate_insensitive && expr.iter().all(is_repeatable); rewrite_single_input(input, child_live, duplicate_insensitive, |input| { Ok(LogicalPlan::Projection(Projection::try_new_with_schema( expr, input, schema, @@ -231,6 +237,8 @@ fn rewrite_node( // Adds the predicate's columns to `live` (a side used only by the filter stays live). let mut child_live = live; child_live.extend_from([&predicate], input.schema())?; + let duplicate_insensitive = + duplicate_insensitive && is_repeatable(&predicate); rewrite_single_input(input, child_live, duplicate_insensitive, |input| { Ok(LogicalPlan::Filter(Filter::new(predicate, input))) }) @@ -248,12 +256,14 @@ fn rewrite_node( input.schema(), )?; - // A grouping aggregate with no aggregate functions (`GROUP BY` with - // an empty `aggr_expr`) only observes which group-key values exist, - // not how many rows produced them, so its input is duplicate- - // insensitive. - let child_duplicate_insensitive = - !group_expr.is_empty() && aggr_expr.is_empty(); + // The input can ignore repeated rows when grouping expressions are + // repeatable and every aggregate ignores duplicates, either by + // nature (`min`) or because it deduplicates its own input + // (`count(DISTINCT x)`). This covers grouping-only and global + // aggregates. One sensitive aggregate makes input multiplicity + // observable, even beneath an insensitive ancestor. + let child_duplicate_insensitive = group_expr.iter().all(is_repeatable) + && aggr_expr.iter().all(is_duplicate_insensitive_aggregate); rewrite_single_input( input, @@ -293,7 +303,13 @@ fn rewrite_node( .extend_from(sort_expr.iter().map(|s| &s.expr), input.schema())?; } - rewrite_single_input(input, child_live, true, |input| { + let duplicate_insensitive = + on_expr.iter().chain(&select_expr).all(is_repeatable) + && sort_expr + .iter() + .flatten() + .all(|sort| is_repeatable(&sort.expr)); + rewrite_single_input(input, child_live, duplicate_insensitive, |input| { Ok(LogicalPlan::Distinct(Distinct::On(DistinctOn { on_expr, select_expr, @@ -310,7 +326,9 @@ fn rewrite_node( // A `fetch` (top-N) makes the row count observable, so duplicate- // insensitivity does not survive past it. - let child_duplicate_insensitive = duplicate_insensitive && fetch.is_none(); + let child_duplicate_insensitive = duplicate_insensitive + && fetch.is_none() + && expr.iter().all(|sort| is_repeatable(&sort.expr)); rewrite_single_input( input, child_live, @@ -409,6 +427,17 @@ fn rewrite_join( let (visible_left, visible_right) = split_join_output_columns(&join, live); + // A semi join may stop evaluating conditions after finding a match. + // If the conditions are not repeatable, skipping evaluations can change + // later matches. Removing duplicate input rows can have the same effect. + // Require repeatable conditions for this join and rewrites of its inputs. + let repeatable = join + .on + .iter() + .all(|(left, right)| is_repeatable(left) && is_repeatable(right)) + && join.filter.iter().all(is_repeatable); + let duplicate_insensitive = duplicate_insensitive && repeatable; + let rewritten_join_type = match rewritten_join_type( &join, &visible_left, @@ -452,12 +481,12 @@ fn rewrite_join( let left = rewrite_subtree( Arc::unwrap_or_clone(join.left), left_live, - left_dup_insensitive, + left_dup_insensitive && repeatable, )?; let right = rewrite_subtree( Arc::unwrap_or_clone(join.right), right_live, - right_dup_insensitive, + right_dup_insensitive && repeatable, )?; let changed = @@ -546,9 +575,10 @@ fn rewritten_join_type( )); // A LEFT JOIN preserves every left row, so with a redundant right side the - // join has no observable effect and can be replaced by its left input. A - // join filter cannot prevent this: it only decides whether a left row is - // matched or null-padded, and either way the row is emitted. + // join has no observable effect and can be replaced by its left input. + // A filter only decides whether a row is matched or null-padded. When + // relying on duplicate-insensitivity, the caller has already required + // repeatable join conditions. if join.join_type == JoinType::Left && can_remove_right { return JoinRewrite::ReplaceWithLeft; } @@ -641,21 +671,29 @@ fn side_unique_on_join<'a>( #[cfg(test)] mod tests { use crate::OptimizerContext; + use crate::OptimizerRule; use crate::assert_optimized_plan_eq_snapshot; use crate::eliminate_join::EliminateJoin; + use crate::test::udfs::PlacementTestUDF; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::{ Constraint, Constraints, NullEquality, Result, ScalarValue, SplitPoint, }; use datafusion_expr::JoinType::Inner; + use datafusion_expr::function::AccumulatorArgs; use datafusion_expr::{ - Expr, JoinType, Partitioning, RangePartitioning, col, exists, lit, + Accumulator, AggregateUDF, AggregateUDFImpl, DistinctHandling, Expr, + ExprFunctionExt, JoinType, LogicalPlan, Partitioning, RangePartitioning, + ScalarUDF, Signature, Volatility, col, exists, lit, logical_plan::builder::{ LogicalPlanBuilder, table_scan, table_source_with_constraints, }, - out_ref_col, + out_ref_col, scalar_subquery, }; - use datafusion_functions_aggregate::expr_fn::count; + use datafusion_functions_aggregate::expr_fn::{ + corr, count, count_distinct, max, min, regr_count, stddev, + }; + use std::hash::{Hash, Hasher}; use std::sync::Arc; macro_rules! assert_optimized_plan_equal { @@ -752,9 +790,8 @@ mod tests { } #[test] - fn aggregate_with_aggregates_is_not_duplicate_insensitive() -> Result<()> { - // A `GROUP BY` *with* aggregate functions observes how many rows fall in - // each group, so its input is not duplicate-insensitive. With a non-unique + fn count_is_not_duplicate_insensitive() -> Result<()> { + // COUNT observes how many rows fall in each group. With a non-unique // right side the join must stay an inner join: collapsing it to a semi // join would drop matching duplicates and undercount `count(l.id)`. let plan = left_join_right()? @@ -769,6 +806,438 @@ mod tests { ") } + #[test] + fn insensitive_aggregates_enable_semi_joins() -> Result<()> { + for column in ["l.x", "r.x"] { + let aggr_expr = vec![ + min(col(column)).alias("minimum"), + max(col(column)).distinct().build()?, + ]; + // Both global and grouped aggregates ignore duplicate input rows. + for group_expr in [vec![], vec![col(column)]] { + let plan = left_join_right()? + .aggregate(group_expr, aggr_expr.clone())? + .build()?; + let result = + EliminateJoin::new().rewrite(plan, &OptimizerContext::new())?; + assert!(result.transformed); + let LogicalPlan::Aggregate(aggregate) = result.data else { + panic!("expected aggregate"); + }; + assert_eq!(aggregate.aggr_expr, aggr_expr); + let LogicalPlan::Join(join) = aggregate.input.as_ref() else { + panic!("expected join"); + }; + assert_eq!( + join.join_type, + if column == "l.x" { + JoinType::LeftSemi + } else { + JoinType::RightSemi + } + ); + } + } + Ok(()) + } + + #[test] + fn global_min_removes_unused_outer_join() -> Result<()> { + for (join_type, column, table) in + [(JoinType::Left, "l.x", "l"), (JoinType::Right, "r.x", "r")] + { + let left = scan("l", &test_schema(), Constraints::default())?; + let right = scan("r", &test_schema(), Constraints::default())?; + let plan = LogicalPlanBuilder::from(left) + .join(right, join_type, (vec!["l.id"], vec!["r.id"]), None)? + .aggregate(Vec::::new(), vec![min(col(column))])? + .build()?; + let optimized = EliminateJoin::new() + .rewrite(plan, &OptimizerContext::new())? + .data; + let expected = LogicalPlanBuilder::from(scan( + table, + &test_schema(), + Constraints::default(), + )?) + .aggregate(Vec::::new(), vec![min(col(column))])? + .build()?; + assert_eq!(optimized, expected); + } + Ok(()) + } + + #[test] + fn distinct_sensitive_aggregates_enable_semi_joins() -> Result<()> { + // A `Sensitive` function called with DISTINCT deduplicates its own + // input, so it cannot observe rows repeated by the join. + for aggr_expr in [ + vec![count_distinct(col("l.x"))], + vec![count_distinct(col("l.x")), count_distinct(col("l.y"))], + vec![ + min(col("l.x")), + count(col("l.x")) + .distinct() + .filter(col("l.y").gt(lit(0))) + .build()?, + ], + ] { + let plan = left_join_right()? + .aggregate(vec![col("l.id")], aggr_expr)? + .build()?; + let result = EliminateJoin::new().rewrite(plan, &OptimizerContext::new())?; + assert!(result.transformed); + let LogicalPlan::Aggregate(aggregate) = result.data else { + panic!("expected aggregate"); + }; + let LogicalPlan::Join(join) = aggregate.input.as_ref() else { + panic!("expected join"); + }; + assert_eq!(join.join_type, JoinType::LeftSemi); + } + Ok(()) + } + + #[test] + fn duplicate_sensitive_aggregates_block_rewrite() -> Result<()> { + // One aggregate that observes repeated rows keeps the join, even + // beside aggregates that do not. DISTINCT does not qualify an + // `Unsupported` function: its accumulator does not deduplicate, and + // may silently compute the non-distinct answer. + for sensitive in [ + count(col("l.x")), + stddev(col("l.x")).distinct().build()?, + corr(col("l.x"), col("l.y")).distinct().build()?, + regr_count(col("l.x"), col("l.y")).distinct().build()?, + ] { + let plan = left_join_right()? + .aggregate( + Vec::::new(), + vec![min(col("l.x")), count_distinct(col("l.x")), sensitive], + )? + .build()?; + assert!( + !EliminateJoin::new() + .rewrite(plan, &OptimizerContext::new())? + .transformed + ); + } + Ok(()) + } + + #[test] + fn sensitive_aggregate_blocks_insensitive_ancestor() -> Result<()> { + let plan = left_join_right()? + .aggregate(vec![col("l.x")], vec![count(col("l.id")).alias("n")])? + .aggregate(Vec::::new(), vec![min(col("n"))])? + .build()?; + assert!( + !EliminateJoin::new() + .rewrite(plan, &OptimizerContext::new())? + .transformed + ); + Ok(()) + } + + #[test] + fn subquery_aggregate_argument_blocks_rewrite() -> Result<()> { + // Expr's usual volatility check does not descend into a subquery plan. + let volatile = ScalarUDF::from( + PlacementTestUDF::new().with_volatility(Volatility::Volatile), + ) + .call(vec![lit(1)]); + let subquery = LogicalPlanBuilder::empty(true) + .project(vec![volatile])? + .build()?; + let plan = left_join_right()? + .aggregate( + vec![col("l.x")], + vec![min(scalar_subquery(Arc::new(subquery)))], + )? + .build()?; + assert!( + !EliminateJoin::new() + .rewrite(plan, &OptimizerContext::new())? + .transformed + ); + Ok(()) + } + + #[test] + fn aggregate_filter_and_ordering_keep_columns_live() -> Result<()> { + for aggr in [ + min(col("l.x")).filter(col("r.y").gt(lit(0))).build()?, + min(col("l.x")) + .order_by(vec![col("r.y").sort(true, false)]) + .build()?, + ] { + let plan = left_join_right()? + .aggregate(Vec::::new(), vec![aggr])? + .build()?; + assert!( + !EliminateJoin::new() + .rewrite(plan, &OptimizerContext::new())? + .transformed + ); + } + + let plan = left_join_right()? + .aggregate( + Vec::::new(), + vec![min(col("l.x")).filter(col("l.y").gt(lit(0))).build()?], + )? + .build()?; + assert_optimized_plan_equal!(plan, @r" + Aggregate: groupBy=[[]], aggr=[[min(l.x) FILTER (WHERE l.y > Int32(0))]] + LeftSemi Join: l.id = r.id + TableScan: l + TableScan: r + ") + } + + fn volatile_expr() -> Expr { + ScalarUDF::from(PlacementTestUDF::new().with_volatility(Volatility::Volatile)) + .call(vec![col("l.x")]) + } + + #[test] + fn volatile_aggregate_expressions_block_rewrite() -> Result<()> { + for (group_expr, aggr) in [ + (vec![], min(volatile_expr())), + (vec![volatile_expr()], min(col("l.x"))), + ( + vec![], + min(col("l.x")) + .filter(volatile_expr().gt(lit(0_u32))) + .build()?, + ), + ( + vec![], + min(col("l.x")) + .order_by(vec![volatile_expr().sort(true, false)]) + .build()?, + ), + ] { + let plan = left_join_right()? + .aggregate(group_expr, vec![aggr])? + .build()?; + assert!( + !EliminateJoin::new() + .rewrite(plan, &OptimizerContext::new())? + .transformed + ); + } + Ok(()) + } + + #[test] + fn volatile_intervening_expressions_block_rewrite() -> Result<()> { + for input in [ + left_join_right()?.project(vec![col("l.x"), volatile_expr().alias("v")])?, + left_join_right()?.filter(volatile_expr().gt(lit(0_u32)))?, + left_join_right()?.sort(vec![volatile_expr().sort(true, false)])?, + ] { + let plan = input + .aggregate(Vec::::new(), vec![min(col("l.x"))])? + .build()?; + assert!( + !EliminateJoin::new() + .rewrite(plan, &OptimizerContext::new())? + .transformed + ); + } + Ok(()) + } + + #[test] + fn join_conditions_must_be_repeatable() -> Result<()> { + for volatility in [Volatility::Stable, Volatility::Volatile] { + let udf = + ScalarUDF::from(PlacementTestUDF::new().with_volatility(volatility)); + for (left_key, right_key, filter) in [ + (udf.call(vec![col("l.id")]), col("r.id"), None), + (col("l.id"), udf.call(vec![col("r.id")]), None), + ( + col("l.id"), + col("r.id"), + Some(udf.call(vec![col("l.x")]).gt(lit(0_u32))), + ), + ] { + let plan = LogicalPlanBuilder::from(scan( + "l", + &test_schema(), + Constraints::default(), + )?) + .join_with_expr_keys( + scan("r", &test_schema(), Constraints::default())?, + Inner, + (vec![left_key], vec![right_key]), + filter, + )? + .aggregate(Vec::::new(), vec![min(col("l.x"))])? + .build()?; + let result = EliminateJoin::new() + .rewrite(plan.clone(), &OptimizerContext::new())?; + assert_eq!( + result.transformed, + volatility != Volatility::Volatile, + "{volatility:?}: {}", + plan.display_indent(), + ); + } + } + Ok(()) + } + + #[test] + fn existence_side_rewrites_require_repeatable_join_conditions() -> Result<()> { + // There is no duplicate-insensitive ancestor. Only the semi join's + // existence side can make the nested inner join eligible for rewriting. + for join_type in [JoinType::LeftSemi, JoinType::RightSemi] { + for volatility in [Volatility::Stable, Volatility::Volatile] { + let inner = left_join_right()?.build()?; + let other = scan("s", &test_schema(), Constraints::default())?; + let (left, right, keys) = if join_type == JoinType::LeftSemi { + (other, inner, (vec!["s.id"], vec!["l.id"])) + } else { + (inner, other, (vec!["l.id"], vec!["s.id"])) + }; + let predicate = + ScalarUDF::from(PlacementTestUDF::new().with_volatility(volatility)) + .call(vec![col("l.x")]) + .gt(lit(0_u32)); + let plan = LogicalPlanBuilder::from(left) + .join(right, join_type, keys, Some(predicate))? + .build()?; + let result = + EliminateJoin::new().rewrite(plan, &OptimizerContext::new())?; + let repeatable = volatility != Volatility::Volatile; + assert_eq!( + result.transformed, repeatable, + "{join_type:?}, {volatility:?}" + ); + let LogicalPlan::Join(join) = result.data else { + panic!("expected semi join"); + }; + assert_eq!(join.join_type, join_type); + let existence_side = if join_type == JoinType::LeftSemi { + join.right + } else { + join.left + }; + let LogicalPlan::Join(nested) = existence_side.as_ref() else { + panic!("expected nested join"); + }; + assert_eq!( + nested.join_type, + if repeatable { + JoinType::LeftSemi + } else { + Inner + }, + ); + } + } + Ok(()) + } + + /// An aggregate that declares the given [`DistinctHandling`]. + #[derive(Debug, PartialEq, Eq)] + struct CustomUdaf { + signature: Signature, + distinct_handling: DistinctHandling, + } + + impl CustomUdaf { + fn new(volatility: Volatility, distinct_handling: DistinctHandling) -> Self { + Self { + signature: Signature::any(1, volatility), + distinct_handling, + } + } + } + + impl Hash for CustomUdaf { + fn hash(&self, state: &mut H) { + self.signature.hash(state); + } + } + + impl AggregateUDFImpl for CustomUdaf { + fn name(&self) -> &str { + "custom_udaf" + } + fn signature(&self) -> &Signature { + &self.signature + } + fn return_type(&self, _: &[DataType]) -> Result { + Ok(DataType::Int32) + } + fn accumulator(&self, _: AccumulatorArgs) -> Result> { + unimplemented!("logical optimizer test") + } + fn distinct_handling(&self) -> DistinctHandling { + self.distinct_handling + } + } + + #[test] + fn distinct_only_qualifies_sensitive_udaf() -> Result<()> { + for (distinct_handling, distinct, expect_rewrite) in [ + (DistinctHandling::Insensitive, false, true), + (DistinctHandling::Insensitive, true, true), + (DistinctHandling::Sensitive, false, false), + // The accumulator deduplicates its own input. + (DistinctHandling::Sensitive, true, true), + (DistinctHandling::Unsupported, false, false), + // The accumulator does not implement DISTINCT, so the input's + // repeated rows stay observable. + (DistinctHandling::Unsupported, true, false), + ] { + let udf = AggregateUDF::from(CustomUdaf::new( + Volatility::Immutable, + distinct_handling, + )); + let mut aggr = udf.call(vec![col("l.x")]); + if distinct { + aggr = aggr.distinct().build()?; + } + let plan = left_join_right()? + .aggregate(Vec::::new(), vec![aggr])? + .build()?; + let result = EliminateJoin::new().rewrite(plan, &OptimizerContext::new())?; + assert_eq!( + result.transformed, expect_rewrite, + "{distinct_handling:?}, distinct={distinct}" + ); + } + Ok(()) + } + + #[test] + fn aliased_udaf_uses_declared_handling() -> Result<()> { + for volatility in [ + Volatility::Immutable, + Volatility::Stable, + Volatility::Volatile, + ] { + let udf = AggregateUDF::from(CustomUdaf::new( + volatility, + DistinctHandling::Insensitive, + )) + .with_aliases(["custom_alias"]); + let plan = left_join_right()? + .aggregate( + Vec::::new(), + vec![udf.call(vec![col("l.x")]).alias("result")], + )? + .build()?; + let result = EliminateJoin::new().rewrite(plan, &OptimizerContext::new())?; + assert_eq!(result.transformed, volatility != Volatility::Volatile); + } + Ok(()) + } + #[test] fn duplicate_insensitive_context_propagates_through_join_tree() -> Result<()> { let left = scan("l", &test_schema(), Constraints::default())?; @@ -867,18 +1336,16 @@ mod tests { #[test] fn correlated_subquery_outer_ref_prevents_rewrite() -> Result<()> { - // The aggregate makes the parent duplicate-insensitive, so absent any - // other use of the right side the join would collapse to a semi join. - // But the `EXISTS` subquery correlates on `r.y`, so the right side is - // still needed and the join must stay an inner join. Otherwise the - // semi join would drop `r`, orphaning the correlated `r.y` reference. + // The right side is unique, so the subquery's repeatability barrier + // alone cannot prevent a semi-join rewrite. Tracking the correlated + // `r.y` reference must keep the right side live and the join inner. let subquery = LogicalPlanBuilder::from(scan("s", &test_schema(), Constraints::default())?) .filter(col("s.id").eq(out_ref_col(DataType::Int32, "r.y")))? .project(vec![lit(1)])? .build()?; - let plan = left_join_right()? + let plan = left_join_right_with_constraints(primary_key_on_id())? .filter(exists(Arc::new(subquery)))? .aggregate(vec![col("l.x")], Vec::::new())? .build()?; @@ -1176,6 +1643,37 @@ mod tests { ") } + #[test] + fn distinct_on_expressions_must_be_repeatable() -> Result<()> { + for volatility in [Volatility::Stable, Volatility::Volatile] { + let expr = + ScalarUDF::from(PlacementTestUDF::new().with_volatility(volatility)) + .call(vec![col("l.x")]); + for (on_expr, select_expr, sort_expr) in [ + (vec![expr.clone()], vec![col("l.x")], None), + (vec![col("l.x")], vec![expr.clone()], None), + ( + vec![col("l.x")], + vec![col("l.x")], + Some(vec![col("l.x").sort(true, false), expr.sort(true, false)]), + ), + ] { + let plan = left_join_right()? + .distinct_on(on_expr, select_expr, sort_expr)? + .build()?; + let result = EliminateJoin::new() + .rewrite(plan.clone(), &OptimizerContext::new())?; + assert_eq!( + result.transformed, + volatility != Volatility::Volatile, + "{volatility:?}: {}", + plan.display_indent(), + ); + } + } + Ok(()) + } + #[test] fn existing_semi_join_passes_through_unchanged() -> Result<()> { // A join that is already a semi join is threaded through unchanged: the rule @@ -1237,7 +1735,7 @@ mod tests { name: &str, schema: &Schema, constraints: Constraints, - ) -> Result { + ) -> Result { if constraints.is_empty() { table_scan(Some(name), schema, None)?.build() } else { diff --git a/datafusion/optimizer/src/unions_to_filter.rs b/datafusion/optimizer/src/unions_to_filter.rs index 2eda0daba1308..7bc39a084f758 100644 --- a/datafusion/optimizer/src/unions_to_filter.rs +++ b/datafusion/optimizer/src/unions_to_filter.rs @@ -18,6 +18,7 @@ //! Rewrites `UNION DISTINCT` branches that differ only by filter predicates //! into a single filtered branch plus `DISTINCT`. +use crate::utils::is_repeatable; use crate::{OptimizerConfig, OptimizerRule}; use datafusion_common::Result; use datafusion_common::tree_node::{ @@ -180,7 +181,7 @@ fn extract_branch(plan: LogicalPlan) -> Result> { LogicalPlan::Filter(Filter { predicate, input, .. }) => { - if !is_mergeable_predicate(&predicate) { + if !is_repeatable(&predicate) { debug!( "unions_to_filter skipped: branch predicate contains volatility or a subquery" ); @@ -334,10 +335,6 @@ fn align_plan_to_schema( )?)) } -fn is_mergeable_predicate(expr: &Expr) -> bool { - !expr.is_volatile() && !expr_contains_subquery(expr) -} - /// Check every expression in the retained source, including its descendants. /// Merging branches also merges their source evaluations, so the same /// restrictions as for projection wrappers apply throughout the source. @@ -345,7 +342,7 @@ fn source_is_safe(source: &LogicalPlan) -> Result { let mut safe = true; source.apply(|node| { node.apply_expressions(|expr| { - if is_mergeable_predicate(expr) { + if is_repeatable(expr) { Ok(TreeNodeRecursion::Continue) } else { safe = false; @@ -370,30 +367,17 @@ fn collect_table_sources(source: &LogicalPlan) -> Result bool { wrappers.iter().all(|w| match w { - Wrapper::Projection { expr, .. } => expr - .iter() - .all(|e| !e.is_volatile() && !expr_contains_subquery(e)), + Wrapper::Projection { expr, .. } => expr.iter().all(is_repeatable), Wrapper::SubqueryAlias { .. } => true, }) } -fn expr_contains_subquery(expr: &Expr) -> bool { - expr.exists(|e| match e { - Expr::ScalarSubquery(_) - | Expr::Exists(_) - | Expr::InSubquery(_) - | Expr::SetComparison(_) => Ok(true), - _ => Ok(false), - }) - .expect("boolean expression walk is infallible") -} - #[cfg(test)] mod tests { use super::*; @@ -401,11 +385,10 @@ mod tests { use crate::assert_optimized_plan_eq_snapshot; use crate::test::test_table_scan_with_name; use arrow::datatypes::DataType; - use datafusion_common::{Result, Spans}; - use datafusion_expr::expr::{SetComparison, SetQuantifier}; + use datafusion_common::Result; use datafusion_expr::{ - ColumnarValue, Expr, Operator, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, - Signature, Subquery, Volatility, col, lit, + ColumnarValue, Expr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, + Volatility, col, lit, }; macro_rules! assert_optimized_plan_equal { @@ -455,27 +438,6 @@ mod tests { ScalarUDF::new_from_impl(VolatileTestUdf).call(vec![]) } - #[test] - fn set_comparison_is_detected_as_subquery() { - let subquery = LogicalPlanBuilder::from(test_table_scan_with_name("t2").unwrap()) - .project(vec![col("b")]) - .unwrap() - .build() - .unwrap(); - let expr = Expr::SetComparison(SetComparison::new( - Box::new(col("a")), - Subquery { - subquery: Arc::new(subquery), - outer_ref_columns: vec![], - spans: Spans::new(), - }, - Operator::Gt, - SetQuantifier::Any, - )); - - assert!(expr_contains_subquery(&expr)); - } - fn assert_not_rewritten(plan: LogicalPlan) { let mut options = datafusion_common::config::ConfigOptions::default(); options.optimizer.enable_unions_to_filter = true; diff --git a/datafusion/optimizer/src/utils.rs b/datafusion/optimizer/src/utils.rs index d4ac31e8a517c..034797c8b1b2f 100644 --- a/datafusion/optimizer/src/utils.rs +++ b/datafusion/optimizer/src/utils.rs @@ -30,7 +30,9 @@ use datafusion_expr::execution_props::ExecutionProps; use datafusion_expr::expr::{Exists, InSubquery, SetComparison}; use datafusion_expr::expr_rewriter::replace_col; use datafusion_expr::physical_planning_context::PhysicalPlanningContext; -use datafusion_expr::{ColumnarValue, Expr, logical_plan::LogicalPlan}; +use datafusion_expr::{ + ColumnarValue, DistinctHandling, Expr, Volatility, logical_plan::LogicalPlan, +}; use datafusion_physical_expr::create_physical_expr; use log::{debug, trace}; use std::sync::Arc; @@ -39,6 +41,51 @@ use std::sync::Arc; /// as it was initially placed here and then moved elsewhere. pub use datafusion_expr::expr_rewriter::NamePreserver; +/// Whether an expression is free of volatile scalar functions and subqueries. +/// Subqueries are conservative barriers because their plans may contain +/// volatile expressions that [`Expr::is_volatile`] does not visit. +pub(crate) fn is_repeatable(expr: &Expr) -> bool { + !expr + .exists(|expr| { + Ok(expr.is_volatile_node() + || matches!( + expr, + Expr::Exists(_) + | Expr::InSubquery(_) + | Expr::SetComparison(_) + | Expr::ScalarSubquery(_) + )) + }) + .expect("expression traversal is infallible") +} + +/// Whether an aggregate can safely ignore repeated input rows. This holds when +/// the function declares [`DistinctHandling::Insensitive`], or when it is +/// called with `DISTINCT` and declares [`DistinctHandling::Sensitive`], so that +/// its accumulator removes the repeated rows itself. Arguments, FILTER, and +/// ORDER BY must also be repeatable: MIN(random()) still observes repetitions. +pub(crate) fn is_duplicate_insensitive_aggregate(mut expr: &Expr) -> bool { + while let Expr::Alias(alias) = expr { + expr = &alias.expr; + } + let Expr::AggregateFunction(aggregate) = expr else { + return false; + }; + let ignores_duplicates = match aggregate.func.distinct_handling() { + DistinctHandling::Insensitive => true, + DistinctHandling::Sensitive => aggregate.params.distinct, + // The accumulator does not implement `DISTINCT` and may silently + // compute the non-distinct answer, so the flag proves nothing. + // Variants added in the future are treated the same way. + _ => false, + }; + // Expr::is_volatile checks scalar functions only; check the aggregate + // function's own volatility separately. + ignores_duplicates + && aggregate.func.signature().volatility != Volatility::Volatile + && is_repeatable(expr) +} + /// Invokes `f` with the index, within `schema`, of every column referenced by /// `expr` — including columns reached through a correlated subquery's outer /// references. Columns absent from `schema` are skipped. @@ -251,7 +298,34 @@ fn coerce(expr: Expr, schema: &DFSchema) -> Result { #[cfg(test)] mod tests { use super::*; - use datafusion_expr::{Operator, binary_expr, case, col, in_list, is_null, lit}; + use crate::test::test_table_scan_with_name; + use datafusion_common::Spans; + use datafusion_expr::expr::SetQuantifier; + use datafusion_expr::logical_plan::builder::LogicalPlanBuilder; + use datafusion_expr::{ + Operator, Subquery, binary_expr, case, col, in_list, is_null, lit, + }; + + #[test] + fn set_comparison_is_not_repeatable() { + let subquery = LogicalPlanBuilder::from(test_table_scan_with_name("t2").unwrap()) + .project(vec![col("b")]) + .unwrap() + .build() + .unwrap(); + let expr = Expr::SetComparison(SetComparison::new( + Box::new(col("a")), + Subquery { + subquery: Arc::new(subquery), + outer_ref_columns: vec![], + spans: Spans::new(), + }, + Operator::Gt, + SetQuantifier::Any, + )); + + assert!(!is_repeatable(&expr)); + } #[test] fn expr_is_restrict_null_predicate() -> Result<()> { diff --git a/datafusion/sqllogictest/test_files/eliminate_join_distinct.slt b/datafusion/sqllogictest/test_files/eliminate_join_distinct.slt new file mode 100644 index 0000000000000..64116e0a629a7 --- /dev/null +++ b/datafusion/sqllogictest/test_files/eliminate_join_distinct.slt @@ -0,0 +1,219 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Duplicate-insensitive aggregates may ignore join fanout, including NULLs. +statement ok +CREATE TABLE l(id INT, g INT, x INT); + +statement ok +CREATE TABLE r(id INT, y INT); + +statement ok +INSERT INTO l VALUES (1, 0, 10), (1, 0, 20), (2, 0, NULL), (3, 1, 30), (4, 2, NULL), (NULL, 3, 99); + +statement ok +INSERT INTO r VALUES (1, 100), (1, 200), (2, 300), (4, 400), (4, 500), (NULL, 600); + +statement ok +SET datafusion.explain.logical_plan_only = true; + +query TT +EXPLAIN SELECT l.g, MIN(l.x), MAX(l.x) FROM l JOIN r ON l.id = r.id GROUP BY l.g; +---- +logical_plan +01)Aggregate: groupBy=[[l.g]], aggr=[[min(l.x), max(l.x)]] +02)--Projection: l.g, l.x +03)----LeftSemi Join: l.id = r.id +04)------TableScan: l projection=[id, g, x] +05)------TableScan: r projection=[id] + +query III rowsort +SELECT l.g, MIN(l.x), MAX(l.x) FROM l JOIN r ON l.id = r.id GROUP BY l.g; +---- +0 10 20 +2 NULL NULL + +query II rowsort +SELECT l.g, APPROX_DISTINCT(l.x) FROM l JOIN r ON l.id = r.id GROUP BY l.g; +---- +0 2 +2 0 + +query IBBII rowsort +SELECT l.g, BOOL_AND(l.x > 15), BOOL_OR(l.x > 15), BIT_AND(l.x), BIT_OR(l.x) FROM l JOIN r ON l.id = r.id GROUP BY l.g; +---- +0 false true 0 30 +2 NULL NULL NULL NULL + +# An insensitive aggregate with a sensitive companion must retain join fanout. +query IIII rowsort +SELECT l.g, MIN(l.x), COUNT(*), SUM(l.x) FROM l JOIN r ON l.id = r.id GROUP BY l.g; +---- +0 10 5 60 +2 NULL 2 NULL + +# The right side remains live through an aggregate FILTER. +# The filter removes all non-NULL values, so ignoring it would change the result. +query TT +EXPLAIN SELECT MIN(l.x) FILTER (WHERE r.y > 250) FROM l JOIN r ON l.id = r.id; +---- +logical_plan +01)Aggregate: groupBy=[[]], aggr=[[min(l.x) FILTER (WHERE r.y > Int32(250)) AS min(l.x) FILTER (WHERE r.y > Int64(250))]] +02)--Projection: l.x, r.y +03)----Inner Join: l.id = r.id +04)------TableScan: l projection=[id, x] +05)------TableScan: r projection=[id, y] + +query I +SELECT MIN(l.x) FILTER (WHERE r.y > 250) FROM l JOIN r ON l.id = r.id; +---- +NULL + +# The unused non-preserved side of an outer join disappears entirely. +query TT +EXPLAIN SELECT MIN(l.x) FROM l LEFT JOIN r ON l.id = r.id; +---- +logical_plan +01)Aggregate: groupBy=[[]], aggr=[[min(l.x)]] +02)--TableScan: l projection=[x] + +query III rowsort +SELECT l.g, MIN(l.x), MAX(l.x) FROM l LEFT JOIN r ON l.id = r.id GROUP BY l.g; +---- +0 10 20 +1 30 30 +2 NULL NULL +3 99 99 + +# The surviving join side can also be the right side. +query TT +EXPLAIN SELECT MAX(r.y) FROM l JOIN r ON l.id = r.id; +---- +logical_plan +01)Aggregate: groupBy=[[]], aggr=[[max(r.y)]] +02)--Projection: r.y +03)----RightSemi Join: l.id = r.id +04)------TableScan: l projection=[id] +05)------TableScan: r projection=[id, y] + +query I +SELECT MAX(r.y) FROM l JOIN r ON l.id = r.id; +---- +500 + +# Empty join inputs retain global-aggregate semantics. +statement ok +CREATE TABLE empty_r(id INT); + +query II +SELECT MIN(l.x), APPROX_DISTINCT(l.x) FROM l JOIN empty_r ON l.id = empty_r.id; +---- +NULL 0 + +query I +SELECT MIN(l.x) FROM l LEFT JOIN empty_r ON l.id = empty_r.id; +---- +10 + +query I +SELECT MIN(empty_r.id) FROM empty_r LEFT JOIN l ON l.id = empty_r.id; +---- +NULL + +# Aggregates that deduplicate their own input ignore join fanout as well. +query TT +EXPLAIN SELECT l.g, COUNT(DISTINCT l.x), SUM(DISTINCT l.x) FROM l JOIN r ON l.id = r.id GROUP BY l.g; +---- +logical_plan +01)Aggregate: groupBy=[[l.g]], aggr=[[count(DISTINCT l.x), sum(DISTINCT CAST(l.x AS Int64))]] +02)--Projection: l.g, l.x +03)----LeftSemi Join: l.id = r.id +04)------TableScan: l projection=[id, g, x] +05)------TableScan: r projection=[id] + +query III rowsort +SELECT l.g, COUNT(DISTINCT l.x), SUM(DISTINCT l.x) FROM l JOIN r ON l.id = r.id GROUP BY l.g; +---- +0 2 30 +2 0 NULL + +query TT +EXPLAIN SELECT COUNT(DISTINCT l.x) FILTER (WHERE l.g = 0), MIN(l.x) FROM l JOIN r ON l.id = r.id; +---- +logical_plan +01)Aggregate: groupBy=[[]], aggr=[[count(DISTINCT l.x) FILTER (WHERE l.g = Int32(0)) AS count(DISTINCT l.x) FILTER (WHERE l.g = Int64(0)), min(l.x)]] +02)--Projection: l.g, l.x +03)----LeftSemi Join: l.id = r.id +04)------TableScan: l projection=[id, g, x] +05)------TableScan: r projection=[id] + +query II +SELECT COUNT(DISTINCT l.x) FILTER (WHERE l.g = 0), MIN(l.x) FROM l JOIN r ON l.id = r.id; +---- +2 10 + +query ? +SELECT ARRAY_AGG(DISTINCT l.x ORDER BY l.x) FROM l JOIN r ON l.id = r.id; +---- +[10, 20, NULL] + +# A DISTINCT aggregate with a non-DISTINCT sensitive companion must retain join fanout. +query TT +EXPLAIN SELECT COUNT(DISTINCT l.x), COUNT(l.g) FROM l JOIN r ON l.id = r.id; +---- +logical_plan +01)Aggregate: groupBy=[[]], aggr=[[count(DISTINCT l.x), count(l.g)]] +02)--Projection: l.g, l.x +03)----Inner Join: l.id = r.id +04)------TableScan: l projection=[id, g, x] +05)------TableScan: r projection=[id] + +query II +SELECT COUNT(DISTINCT l.x), COUNT(l.g) FROM l JOIN r ON l.id = r.id; +---- +2 7 + +# REGR_COUNT does not implement DISTINCT and counts every joined row, so +# DISTINCT does not hide the join fanout and the join must stay an inner join. +query TT +EXPLAIN SELECT REGR_COUNT(DISTINCT l.x, l.g) FROM l JOIN r ON l.id = r.id; +---- +logical_plan +01)Aggregate: groupBy=[[]], aggr=[[regr_count(DISTINCT CAST(l.x AS Float64), CAST(l.g AS Float64))]] +02)--Projection: l.g, l.x +03)----Inner Join: l.id = r.id +04)------TableScan: l projection=[id, g, x] +05)------TableScan: r projection=[id] + +query I +SELECT REGR_COUNT(DISTINCT l.x, l.g) FROM l JOIN r ON l.id = r.id; +---- +4 + +query I +SELECT REGR_COUNT(DISTINCT l.x, l.g) FROM l LEFT JOIN r ON l.id = r.id; +---- +6 + +# A limit below the aggregate still observes repeated join rows. +query I +SELECT MAX(x) FROM (SELECT l.x FROM l JOIN r ON l.id = r.id ORDER BY l.x NULLS LAST LIMIT 2); +---- +10 + +statement ok +RESET datafusion.explain.logical_plan_only; diff --git a/datafusion/sqllogictest/test_files/joins.slt b/datafusion/sqllogictest/test_files/joins.slt index d2453af739c07..d42de4e97d82f 100644 --- a/datafusion/sqllogictest/test_files/joins.slt +++ b/datafusion/sqllogictest/test_files/joins.slt @@ -1375,10 +1375,9 @@ group by join_t1.t1_id ---- logical_plan 01)Aggregate: groupBy=[[join_t1.t1_id]], aggr=[[count(DISTINCT join_t1.t1_int), count(DISTINCT join_t1.t1_name)]] -02)--Projection: join_t1.t1_id, join_t1.t1_name, join_t1.t1_int -03)----Inner Join: join_t1.t1_id = join_t2.t2_id -04)------TableScan: join_t1 projection=[t1_id, t1_name, t1_int] -05)------TableScan: join_t2 projection=[t2_id] +02)--LeftSemi Join: join_t1.t1_id = join_t2.t2_id +03)----TableScan: join_t1 projection=[t1_id, t1_name, t1_int] +04)----TableScan: join_t2 projection=[t2_id] statement ok set datafusion.explain.logical_plan_only = false;