adriangb commented on code in PR #25385:
URL: https://github.com/apache/datafusion/pull/25385#discussion_r4039635681
##########
datafusion/optimizer/src/eliminate_join.rs:
##########
@@ -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,
Review Comment:
Subqueries also clear the flag (see `is_repeatable`). Please name them here.
```suggestion
//! 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, an
//! expression that is not repeatable such as `random()` or a subquery,
...)
//! clears it. It is therefore fixed by the nearest such node, not by the
//! whole ancestor chain: a collapsing node shields its subtree,
```
##########
datafusion/optimizer/src/eliminate_join.rs:
##########
@@ -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::<Expr>::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::<Expr>::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::<Expr>::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::<Expr>::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::<Expr>::new(), vec![aggr])?
+ .build()?;
+ assert!(
+ !EliminateJoin::new()
+ .rewrite(plan, &OptimizerContext::new())?
+ .transformed
+ );
+ }
+
+ let plan = left_join_right()?
+ .aggregate(
+ Vec::<Expr>::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::<Expr>::new(), vec![min(col("l.x"))])?
+ .build()?;
+ assert!(
+ !EliminateJoin::new()
+ .rewrite(plan, &OptimizerContext::new())?
+ .transformed
+ );
+ }
+ Ok(())
+ }
Review Comment:
These tests only assert `!transformed`. If a different rule or a build error
stops the rewrite, the tests still pass. Please add a `Stable` control, as
`join_conditions_must_be_repeatable` does.
```suggestion
fn udf_expr(volatility: Volatility) -> Expr {
ScalarUDF::from(PlacementTestUDF::new().with_volatility(volatility))
.call(vec![col("l.x")])
}
#[test]
fn volatile_aggregate_expressions_block_rewrite() -> Result<()> {
// `Stable` is the control: the same shape with a repeatable
// expression is rewritten.
for volatility in [Volatility::Stable, Volatility::Volatile] {
let expr = udf_expr(volatility);
for (group_expr, aggr) in [
(vec![], min(expr.clone())),
(vec![expr.clone()], min(col("l.x"))),
(
vec![],
min(col("l.x"))
.filter(expr.clone().gt(lit(0_u32)))
.build()?,
),
(
vec![],
min(col("l.x"))
.order_by(vec![expr.clone().sort(true, false)])
.build()?,
),
] {
let plan = left_join_right()?
.aggregate(group_expr, vec![aggr])?
.build()?;
let result = EliminateJoin::new()
.rewrite(plan.clone(), &OptimizerContext::new())?;
assert_eq!(
result.transformed,
volatility != Volatility::Volatile,
"{volatility:?}: {}",
plan.display_indent(),
);
}
}
Ok(())
}
#[test]
fn volatile_intervening_expressions_block_rewrite() -> Result<()> {
for volatility in [Volatility::Stable, Volatility::Volatile] {
let expr = udf_expr(volatility);
for input in [
left_join_right()?.project(vec![col("l.x"),
expr.clone().alias("v")])?,
left_join_right()?.filter(expr.clone().gt(lit(0_u32)))?,
left_join_right()?.sort(vec![expr.clone().sort(true,
false)])?,
] {
let plan = input
.aggregate(Vec::<Expr>::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(())
}
```
##########
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.
Review Comment:
`regr_count` declares `Unsupported`, and its accumulator ignores
`is_distinct`. Thus `4` is the count without DISTINCT. The distinct count is
`2`. When plan-time checks for `Unsupported` are added, this result will
change. Please write this in the comment.
```suggestion
# REGR_COUNT declares `DistinctHandling::Unsupported`: its accumulator
ignores
# `is_distinct` and counts every joined row (4 below, not 2), so DISTINCT
does
# not hide the join fanout and the join must stay an inner join. Plan-time
# enforcement of `Unsupported` is a follow-up; update the results below when
# it lands.
```
##########
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]
Review Comment:
The plan is now a semi join, and #22644 is closed. Please update the comment
above this query (lines 1367–1368). GitHub cannot attach a suggestion there,
because those lines are not in the diff.
```diff
-# A similar query with two DISTINCT aggregates is currently not rewritten
-# TODO: https://github.com/apache/datafusion/issues/22644
+# A similar query with two DISTINCT aggregates is also rewritten: each
+# `count(DISTINCT ...)` removes its own duplicates, so the join's duplicates
+# are not observable (see
https://github.com/apache/datafusion/issues/22644).
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
