perf: Use agg `DistinctHandling` in join optimization by neilconway · Pull Request #25385 · apache/datafusion

neilconway · 2026-09-16T19:25:40Z

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 can 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. No TPC-H plans change.

Are there any user-facing changes?

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

neilconway · 2026-09-16T19:26:31Z

FYI @mkleen @adriangb @jayzhan211

codecov-commenter · 2026-09-16T19:48:16Z

Codecov Report

❌ Patch coverage is 77.06422% with 100 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.37%. Comparing base (de8803c) to head (6dc8e34).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
datafusion/optimizer/src/eliminate_join.rs	73.20%	18 Missing and 79 partials ⚠️
datafusion/optimizer/src/test/udfs.rs	93.10%	2 Missing ⚠️
datafusion/optimizer/src/utils.rs	97.56%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #25385      +/-   ##
==========================================
- Coverage   82.38%   82.37%   -0.02%     
==========================================
  Files        1138     1138              
  Lines      433776   434081     +305     
  Branches   433776   434081     +305     
==========================================
+ Hits       357372   357560     +188     
- Misses      54849    54885      +36     
- Partials    21555    21636      +81

☔ 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

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 · 2026-09-17T17:20:25Z

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


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,

Thanks! Fixed.

adriangb · 2026-09-17T17:20:25Z

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


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

}

Good catch, fixed.

adriangb · 2026-09-17T17:20:25Z

+# 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 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 · 2026-09-17T17:24:12Z

+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]


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

…distinct # Conflicts: # datafusion/optimizer/src/utils.rs

adriangb · 2026-09-19T16:26:09Z

LGTM!

adriangb · 2026-09-19T19:51:30Z

@neilconway i realized we should probably run some benchmarks, do you mind running relevant benchmarks vs the HEAD commit preceding the merge?

neilconway · 2026-09-20T02:45:55Z

@adriangb I ran a few benchmarks on optimizer/planner perf; is this what you had in mind?

   Workload                  Measurement           Before         After    Change
  ━━━━━━━━━━━━━━━━━━━━━━━━  ━━━━━━━━━━━━━━━━  ━━━━━━━━━━━━  ━━━━━━━━━━━━  ━━━━━━━━
   TPC-H, 21 statements      Optimizer only      7.737 ms      7.847 ms    +1.42%
  ────────────────────────  ────────────────  ────────────  ────────────  ────────
   TPC-DS, 103 statements    Optimizer only    145.898 ms    147.741 ms    +1.26%
  ────────────────────────  ────────────────  ────────────  ────────────  ────────
   TPC-H, 21 statements      Full planning      17.878 ms     18.056 ms    +1.00%
  ────────────────────────  ────────────────  ────────────  ────────────  ────────
   TPC-DS, 103 statements    Full planning     305.110 ms    305.990 ms    +0.29%
  ────────────────────────  ────────────────  ────────────  ────────────  ────────
   ClickBench, 58 queries    Full planning      62.894 ms     62.840 ms    −0.09%

Looks like some of the overhead would be relatively easy to address, I'll take a look.

adriangb · 2026-09-20T05:48:25Z

Looks like some of the overhead would be relatively easy to address, I'll take a look.

Thanks, good find.

I was thinking of the full query bench suites, i’ll trigger some.

.

76e27f0

github-actions Bot added the optimizer Optimizer rules label Sep 16, 2026

Add SLT tests

7d8b4e5

github-actions Bot added the sqllogictest SQL Logic Tests (.slt) label Sep 16, 2026

Improve comments, add test

ab9e061

github-actions Bot added the core Core DataFusion crate label Sep 16, 2026

jayzhan211 reviewed Sep 17, 2026

View reviewed changes

Comment thread datafusion/optimizer/src/utils.rs Outdated

neilconway added 2 commits September 17, 2026 11:22

Address review comment on DISTINCT aggs

ac06738

Cleanup and simplification

de3a71a

adriangb approved these changes Sep 17, 2026

View reviewed changes

neilconway added 2 commits September 19, 2026 10:45

Merge remote-tracking branch 'origin/main' into neilc/perf-semi-join-…

4d3869b

…distinct # Conflicts: # datafusion/optimizer/src/utils.rs

Address review comments

6dc8e34

neilconway enabled auto-merge September 19, 2026 16:05

neilconway added this pull request to the merge queue Sep 19, 2026

github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Sep 19, 2026

neilconway added this pull request to the merge queue Sep 19, 2026

Merged via the queue into apache:main with commit 42f3888 Sep 19, 2026
41 checks passed

neilconway deleted the neilc/perf-semi-join-distinct branch September 19, 2026 16:45

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

Conversation

neilconway commented Sep 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

What is the testing strategy for this PR?

Are there any user-facing changes?

Uh oh!

neilconway commented Sep 16, 2026

Uh oh!

codecov-commenter commented Sep 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

jayzhan211 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

adriangb left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

adriangb commented Sep 19, 2026

Uh oh!

Uh oh!

Uh oh!

adriangb commented Sep 19, 2026

Uh oh!

neilconway commented Sep 20, 2026

Uh oh!

adriangb commented Sep 20, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

neilconway commented Sep 16, 2026 •

edited

Loading

codecov-commenter commented Sep 16, 2026 •

edited

Loading

adriangb left a comment •

edited

Loading