feat: add size-aware join selection for pipelined jobs by andygrove · Pull Request #2685 · apache/datafusion-ballista

andygrove · 2026-10-28T19:55:30Z

Which issue does this PR close?

Part of #2652. This is split out of #2657 to make that PR smaller to review.

Rationale for this change

A pipelined job is planned with the static planner, so AQE never picks its join strategies from measured sizes. Every hash table in a pipelined job is resident at the same time, a hash join build cannot spill, and a failed pipelined task fails the whole job. Planning a large join as a partitioned hash join is therefore much riskier than in blocking mode, and the planner needs to fall back to a join that can spill when the build side may not fit.

What changes are included in this PR?

state/pipelined_join_selection.rs adds:
- PipelinedJoinSelection, a physical optimizer rule. It turns every Partitioned hash join whose estimated build size per partition exceeds ballista.optimizer.hash_join_max_build_partition_bytes into a SortMergeJoinExec. The per-partition size is the build side's estimate divided by target_partitions. A build of unknown size is treated as too large. This is stricter than AQE, which keeps the hash join when the size is unknown. CollectLeft joins and small builds are left alone, and a limit of 0 disables the check.
- create_pipelined_physical_plan, which plans a logical plan with prefer_hash_join = true, as AQE does. It inserts the rule right after DataFusion's join_selection, so EnsureRequirements adds the sorts and repartitions the sort-merge joins need. Rewriting the finished plan instead would lose orderings: a hash join keeps its probe side's order, so the planner can drop the sort for an ORDER BY above it.
The size estimate reuses AQE's estimate_output_byte_size, which is now pub(crate). It uses total_byte_size when known and estimates from the row count otherwise.

Nothing calls this yet, so the module carries #[allow(dead_code)]. #2657 calls create_pipelined_physical_plan from TaskManager::submit_job for pipelined jobs and removes the attribute.

This PR adds tests:

Unit tests for the rule cover an oversized, small and unknown-size build, estimation from the row count, keeping a projection over the join, the zero limit and CollectLeft joins.
Planner tests check that small joins are planned as hash joins whatever prefer_hash_join says, that joins over the build limit are planned as sort-merge joins, and that a demoted join still returns rows in ORDER BY order when the hash join plan relied on the probe side's order.

Are there any user-facing changes?

No. The rule is not used by any query yet.

A pipelined job is planned without AQE and holds every hash table of the job at once, none of which can spill. Plan its joins as hash joins, then demote a partitioned hash join to a sort-merge join when its estimated build size per partition exceeds ballista.optimizer.hash_join_max_build_partition_bytes or is unknown. Nothing plans pipelined jobs yet; the scheduler wires this in later.

Rewriting the finished plan could lose an ordering that a parent relied on: a hash join keeps its probe side's order, so the planner may drop the sort for an ORDER BY above it, and the sort-merge join that replaced the join returned rows ordered by the join keys instead. Run the rule right after JoinSelection, so EnsureRequirements sorts and partitions the sort-merge join inputs, and estimate the per-partition build size from target_partitions. Warn when the session has no JoinSelection rule, and make the CollectLeft test use a build side that would otherwise be demoted.

avantgardnerio · 2026-10-28T23:45:24Z

+        .position(|rule| rule.name() == "join_selection")
+    {
+        Some(index) => index + 1,
+        None => {


If join_selection is missing, the rule lands at position 0 where every join is still Auto, so it never fires. But prefer_hash_join = true is still forced below, so the result is all hash joins with no size check, which is less safe than plain static planning. Would plan_err! fit better here than warn!?

Agreed that forcing hash joins with no size check is the wrong outcome. I went a slightly different way from plan_err!: without join_selection the job now falls back to plain static planning and keeps the session's own prefer_hash_join. That way a custom optimizer setup doesn't fail the job, and it's no less safe than static planning. Happy to switch it to an error if you'd rather be strict.

avantgardnerio · 2026-10-28T23:45:24Z

+//! its probe side's order, so the planner may have dropped a sort above it.
+//!
+//! This is stricter than AQE, which keeps a hash join when the size is
+//! unknown: a pipelined job holds every hash table of the job at once, a hash


holds every hash table of the job at once

The check is per join, per partition, so N joins each just under the limit all stay hash joins and are all resident together. Is a job-wide or per-executor budget planned for #2657, or is per-join the intended bound?

Per join is the intended bound for now. #2657 doesn't add a job-wide or per-executor budget. I added a paragraph to the module docs saying the limit doesn't cap the total of all the hash tables a job holds at once. A job-level budget seems worth a follow-up issue.

Filed the job-wide budget follow-up as #2690.

avantgardnerio · 2026-10-28T23:45:24Z

+            let Some(join) = plan.downcast_ref::<HashJoinExec>() else {
+                return Ok(Transformed::no(plan));
+            };
+            // `CollectLeft` joins are left to the distributed planner's


The distributed planner can turn a CollectLeft join back into Partitioned after this rule has run (planner.rs, unsafe join type or over broadcast_join_threshold_bytes), so those joins never get size-checked. In practice they are under DataFusion's hash_join_single_partition_threshold, so it only matters if that threshold is raised a lot. Maybe worth a sentence in the module docs?

Good catch, confirmed in planner.rs. Added a paragraph to the module docs about it.

avantgardnerio · 2026-10-28T23:45:24Z

+        .compute(plan, &StatisticsArgs::new())
+        .ok()?;
+    if let Some(&bytes) = stats.total_byte_size.get_value()
+        && bytes != 0


The reasoning for treating zero as unknown fits Inexact(0), but an Exact(0) build side is known empty and gets demoted to a sort-merge join too. Only costs a needless sort, so minor.

Fixed. An exact zero is now treated as a known empty build and keeps its hash join. An inexact zero is still treated as unknown. Added a test for both.

avantgardnerio · 2026-10-28T23:45:24Z

+                    );
+                    Ok(Transformed::yes(sort_merge_join))
+                }
+                // A join type the sort-merge join does not support keeps its


A join type the sort-merge join does not support

In DataFusion 55.2 SortMergeJoinExec::try_new accepts every JoinType, Mark included; it only errors on key or sort-option validation. So this arm is really "construction failed", and it keeps a possibly oversized hash join. Should it propagate the error instead?

You're right, it only fails on key or sort-option validation. The error is now propagated instead of quietly keeping the hash join.

avantgardnerio · 2026-10-28T23:45:24Z

+            // The inputs are not repartitioned yet. A `Partitioned` join's
+            // will be, into `target_partitions` partitions.
+            let build_partition_bytes = estimated_bytes(join.left().as_ref())
+                .map(|bytes| bytes.div_ceil(config.execution.target_partitions.max(1)));


Nit: with repartition_joins off, or when EnforceDistribution reuses an existing hash partitioning with a different count, the join runs with fewer input partitions than target_partitions, so each one is larger than this estimate.

With repartition_joins off, DataFusion's planner makes the join CollectLeft, so the rule never sees it. The existing hash partitioning case is real, though. The rule now divides by the smaller of target_partitions and any input's existing hash partition count, and there's a test for it.

avantgardnerio · 2026-10-28T23:45:24Z

+                probe,
+                on,
+                None,
+                &JoinType::Inner,


Nit: only Inner is tested. A Left or Full join with a filter would also cover the sort-merge join's filter path after demotion.

Added a test that runs LEFT JOIN and FULL JOIN with a filter as a hash join and as a demoted sort-merge join, and checks they return the same rows.

avantgardnerio

I'm sorry. "I had my AI review your AI" is the modern version of "if I had more time, I'd have written you a shorter letter".

Please take these auto generated comments as a grain of salt, and I trust you with provisional approval.

andygrove · 2026-10-29T03:50:16Z

I'm sorry. "I had my AI review your AI" is the modern version of "if I had more time, I'd have written you a shorter letter".

Please take these auto generated comments as a grain of salt, and I trust you with provisional approval.

Plan without the rule when the session has no join_selection, propagate sort-merge join construction errors, keep hash joins over known empty build sides, use the partition count of already hash partitioned inputs, and document the per-join scope and late CollectLeft demotions.

andygrove · 2026-10-29T03:57:28Z

Thanks @avantgardnerio. My AI thanks your AI.

comphead

Follow-up notes for #2657, since this is merged. The tests pass locally (14/14), and I checked the plan behavior against DataFusion 55.2.

Design question: demoting to SortMergeJoinExec adds a SortExec on both inputs, including the probe side, which is often the largest input. Sort is a pipeline breaker, so a demoted join stops being pipelined and pays for a full probe sort. Would a blocking stage boundary at that join, or planning the job with AQE, be cheaper?

Spark's canBuildLocalHashMapBySize supports treating unknown sizes as too large (spark.sql.defaultSizeInBytes) and using the per-partition average. It also requires muchSmaller (build size x shuffleHashJoinFactor <= probe size) before choosing a shuffled hash join, which this rule doesn't have. Unknown sizes are common: AQE's comment in supports_collect_by_thresholds notes that DataFusion drops total_byte_size at every join. I expect, but have not verified, that multi-join TPC-H plans will demote most joins. Pipelined TPC-H numbers before #2657 wires this in would settle it.

comphead · 2026-10-31T17:24:24Z

+/// repartition them into `target_partitions` partitions. An input that is
+/// already hash partitioned may be kept as it is, with its own partition
+/// count, so the smallest such count is used when it is lower.
+fn join_input_partitions(join: &HashJoinExec, target_partitions: usize) -> usize {


The per-partition estimate is too low when the join inputs are single-partition. In DataFusion 55.2, Partitioning::satisfaction treats partition_count() == 1 as satisfying HashPartitioned, so EnsureRequirements keeps such inputs as they are and the join runs as one task holding the whole build side. This function still divides by target_partitions because it only looks at Partitioning::Hash. It also misses Partitioning::Range, which satisfies key distribution in DataFusion 55.

Reproduced with the plan_pipelined tables: a 16,192-byte build with the limit set to 8,096 still plans as

HashJoinExec: mode=Partitioned, join_type=Inner, on=[(k@0, k@0)], projection=[k@0] DataSourceExec: partitions=1, partition_sizes=[1] DataSourceExec: partitions=1, partition_sizes=[1]

With 16 target partitions, a 1 GiB single-file build is estimated at exactly 64 MiB and stays a hash join that can't spill. Suggest treating count == 1 and Range(_, n) like Hash(_, n) in the filter_map.

comphead · 2026-10-31T17:24:24Z

+        );
+        return state.create_physical_plan(logical_plan).await;
+    };
+    rules.insert(index + 1, Arc::new(PipelinedJoinSelection::new()));


The module docs describe the CollectLeft gap, and AQE already closes part of it in code: aqe/planner.rs inserts DemoteUnsafeBroadcastJoinRule right after join_selection. Inserting [DemoteUnsafeBroadcastJoinRule, PipelinedJoinSelection] here would size-check broadcast-unsafe joins before EnsureRequirements runs. Joins over broadcast_join_threshold_bytes would still be missed, so the doc paragraph would shrink rather than disappear.

comphead · 2026-10-31T17:24:24Z

+/// unknown, since statistics can report zero for inputs they did not measure.
+/// An exact zero is a known empty input.
+fn estimated_bytes(plan: &dyn ExecutionPlan) -> Option<usize> {
+    let stats = StatisticsContext::new()


DataFusion 55.2 JoinSelection reads statistics through StatisticsRegistry when datafusion.optimizer.use_statistics_registry is on (via optimize_with_context and get_stats). This rule always uses StatisticsContext::new(), so with the flag on, a build over a join can look Absent here even though JoinSelection swapped or kept it based on registry estimates. The flag defaults to false, so this only affects users who opt in. Implementing optimize_with_context would keep the two rules consistent.

comphead · 2026-10-31T17:24:25Z

+fn estimated_bytes(plan: &dyn ExecutionPlan) -> Option<usize> {
+    let stats = StatisticsContext::new()
+        .compute(plan, &StatisticsArgs::new())
+        .ok()?;


.ok()? turns a statistics error into "unknown size", which silently demotes the join. DynamicJoinExec::to_actual_join propagates the same error with ?. Returning Result<Option<usize>> would match it and keep statistics bugs visible.

comphead · 2026-10-31T17:24:25Z

+        return Ok(sort_merge_join);
+    };
+    let schema = sort_merge_join.schema();
+    let exprs = projection.iter().map(|&index| {


This mapping re-implements ProjectionExprs::from_indices from DataFusion 55.2 (datafusion/physical-expr/src/projection.rs). Something like ProjectionExec::try_new(ProjectionExprs::from_indices(projection, &schema).iter().cloned(), sort_merge_join) should do the same (not compiled).

The SortMergeJoinExec::try_new call above also duplicates DynamicJoinExec::to_sort_merge_join, and estimated_bytes overlaps AQE's supports_collect_by_thresholds with a different zero policy. Keeping the estimator next to estimate_output_byte_size in dynamic_join.rs, with the zero handling stated once, would stop the two from drifting.

comphead · 2026-10-31T17:24:25Z

+
+impl PipelinedJoinSelection {
+    /// Creates a new `PipelinedJoinSelection` rule.
+    pub fn new() -> Self {


Nit: new() is redundant next to #[derive(Default)]. The sibling DemoteUnsafeBroadcastJoinRule uses only Default. Also, the prefer_hash_join reasoning is repeated in the module doc, the create_pipelined_physical_plan doc and the inline comment there. Once is enough.

comphead · 2026-10-31T17:24:25Z

+
+    /// Plans `sql` over tables `a` and `b`, each with columns `k` (0..1000) and
+    /// `v` (1000 down to 1), declared sorted on `v`.
+    async fn plan_pipelined(


These planner tests are all single-partition. Each MemTable has one partition, so with_target_partitions(4) has no effect and the plans contain no RepartitionExec. The realistic shape (hash repartition, per-partition sort, then merge for ORDER BY) is untested, and this is also why the join_input_partitions underestimate went unnoticed. Splitting each table into two or more partitions would cover both.

No test runs plan_query_stages or a BallistaCodec round-trip on a demoted plan yet. That can reasonably wait for #2657.

comphead · 2026-10-31T17:24:25Z

+    }
+
+    #[tokio::test]
+    async fn plans_joins_over_the_build_limit_as_sort_merge_joins() {


This test is covered by demoting_keeps_the_order_a_parent_relies_on and demoted_outer_joins_with_filters_match_hash_joins, which both assert ["SortMergeJoinExec"] at limit 1. It could be removed.

comphead · 2026-10-31T17:24:25Z

+    }
+
+    #[test]
+    fn demotes_hash_join_with_oversized_build() {


This is the None-projection case of keeps_projection_when_demoting. The size-case tests (oversized, small, unknown, empty, row-estimated) also form one table of (num_rows, bytes) -> operator and would fit #[rstest] cases. rstest is already a dev-dependency of the scheduler.

comphead · 2026-10-31T17:24:25Z

+    }
+
+    /// Runs `plan` and returns its rows as sorted, formatted lines.
+    async fn sorted_rows(


This re-implements assert_batches_sorted_eq! and also sorts the table border lines. Likewise, assert_batches_eq! on the collected batches would replace the manual Int64Array extraction in the ordering test.

andygrove requested a review from avantgardnerio October 28, 2026 20:32

avantgardnerio reviewed Oct 28, 2026

View reviewed changes

avantgardnerio approved these changes Oct 28, 2026

View reviewed changes

andygrove mentioned this pull request Oct 29, 2026

Bound the total hash join build memory of a pipelined job #2690

Open

andygrove added this pull request to the merge queue Oct 29, 2026

Merged via the queue into apache:main with commit 07de2ca Oct 29, 2026
25 checks passed

andygrove deleted the pipelined-join-selection branch October 29, 2026 04:20

comphead reviewed Oct 31, 2026

View reviewed changes

Conversation

andygrove commented Oct 28, 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?

Are there any user-facing changes?

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!

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!

avantgardnerio left a comment

Choose a reason for hiding this comment

Uh oh!

andygrove commented Oct 29, 2026

Uh oh!

andygrove commented Oct 29, 2026

Uh oh!

Uh oh!

comphead left a comment

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

andygrove commented Oct 28, 2026 •

edited

Loading