perf: Use agg DistinctHandling in join optimization

[neilconway]

Which issue does this PR close?

• N/A; see perf: Convert inner joins to semi joins when equivalent #22652 and feat: DISTINCT handling for Aggregate UDFs #25288 for context

Rationale for this change

If we can prove that parts of a query are insensitive to duplicates, the optimizer apply various simplifications, like replacing inner joins with semi-joins and removing unused outer-join inputs.

The join analysis was previously conservative and assumed that all aggregate expressions are duplicate sensitive. Since #25288 added a framework for classifying how an aggregate treats duplicate values, we can now apply that framework to optimize joins more effectively.

What changes are included in this PR?

• Extend EliminateJoin to recognize Aggregate plan nodes whose aggregate expressions are all either DistinctHandling::Insensitive or DistinctHandling::Sensitive and invoked with DISTINCT
• Guard duplicate-insensitivity propagation against volatile expressions and subqueries, to avoid changing query results
• Refactor code to share the existing volatility/subquery check with UnionsToFilter.
• Add tests

What is the testing strategy for this PR?

Existing tests pass; new tests added.

Are there any user-facing changes?

Some query plans might change (usually for the better).

[neilconway]

FYI @mkleen @adriangb @jayzhan211

[codecov-commenter]

Codecov Report

❌ Patch coverage is 75.74257% with 98 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.32%. Comparing base (b0b5471) to head (ac06738).
⚠️ Report is 27 commits behind head on main.

Files with missing lines	Patch %	Lines
datafusion/optimizer/src/eliminate_join.rs	73.05%	20 Missing and 77 partials ⚠️
datafusion/optimizer/src/utils.rs	97.56%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #25385      +/-   ##
==========================================
+ Coverage   81.93%   82.32%   +0.38%     
==========================================
  Files        1136     1137       +1     
  Lines      428779   432246    +3467     
  Branches   428779   432246    +3467     
==========================================
+ Hits       351319   355825    +4506     
+ Misses      56446    54851    -1595     
- Partials    21014    21570     +556     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[jayzhan211]

Thanks @neilconway , there is a suggestion

[jayzhan211]

Suggested change

	aggregate.func.distinct_handling() == DistinctHandling::Insensitive
	(aggregate.params.distinct || aggregate.func.distinct_handling() == DistinctHandling::Insensitive)

[jayzhan211]

Should we check distinct?

[neilconway]

Good idea! For Unsupported aggs, DISTINCT isn't quite strong enough, but I think this is sound if we check for either (a) Insensitive, or (b) Sensitive + DISTINCT. I made this change and added more tests.

[adriangb]

Thank you, Neil. The change looks correct to me, nice work!

Some minor update suggestions for the PR description:

-If we can prove that parts of a query are insensitive to duplicates, the optimizer apply various simplifications, like replacing inner joins with semi-joins and removing unused outer-join inputs.
+If we can prove that parts of a query are insensitive to duplicates, the optimizer can apply various simplifications, like replacing inner joins with semi-joins and removing unused outer-join inputs.

 Existing tests pass; new tests added.
+No TPC-H plans change.

Once this review is addressed I'm good to merge this 🚀

[adriangb]

Subqueries also clear the flag (see is_repeatable). Please name them here.

Suggested change

	//! 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,
	//! 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,

[adriangb]

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.

Suggested change

	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(())
	}
	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(())
	}

[adriangb]

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.

Suggested change

	# 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.
	# 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.

[adriangb]

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.

-# 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).