Fix: eliminate unnecessary repartitioning for small datasets

[shashidhar-bm]

Which issue does this PR close?

• Closes Eliminate Repartitioning for Small Datasets #18595

Rationale for this change

Aggregate planning was repartitioning tiny datasets, leading to unnecessary cost for lightweight queries

The planner now uses statistics to keep small inputs in single-partition mode when there aren't enough rows to benefit from the overhead of repartitioning, thus avoiding this performance drain.

What changes are included in this PR?

• Added the function has_sufficient_rows_for_repartition(...) in datafusion/core/src/physical_planner.rs.
• Wired this check into the aggregate planning branch to only build AggregateMode::FinalPartitioned when $num_rows \geq batch_size$.
• Extended planner coverage with tests (e.g., hash_agg_small_dataset_single_mode) and dataframe insta snapshots to assert the new single-mode plan and its explain output.
• Regenerated all affected SQL logic suites (e.g., aggregate, group_by, tpch) to reflect the planning change.

Are these changes tested?

Yes.

Are there any user-facing changes?

Yes.

[shashidhar-bm]

Sorry for opening a new PR—my goal was to keep the diff clean and reviewable, so I rebuilt the changes clearly from scratch in this pr.

[Copilot]

Pull request overview

This PR optimizes aggregate query planning by eliminating unnecessary repartitioning for small datasets that don't have enough rows to benefit from parallel processing overhead.

Key Changes:

• Added has_sufficient_rows_for_repartition() function to check if input has enough rows (≥ batch_size) to warrant repartitioning
• Modified aggregate planning to skip FinalPartitioned mode for small datasets, using Single or Final mode instead
• Updated existing test with larger dataset (10,000 rows) to ensure repartitioning still occurs for large data
• Added new test case to verify single-mode execution for small datasets (6 rows)

Reviewed changes

Copilot reviewed 19 out of 19 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
datafusion/core/src/physical_planner.rs	Core implementation: adds row count check before repartitioning aggregates; updates test dataset size and adds new small dataset test
datafusion/core/tests/dataframe/mod.rs	Updates dataframe test snapshots to reflect new single-mode aggregation plans
datafusion/sqllogictest/test_files/*.slt	Regenerates expected query plans across multiple test files to show Single/Final modes instead of FinalPartitioned for small datasets
datafusion/sqllogictest/test_files/encrypted_parquet.slt	Changes test from expecting error to expecting results with rowsort
datafusion/sqllogictest/test_files/clickbench_extended.slt	Updates result ordering for non-deterministic query results

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 18 out of 18 changed files in this pull request and generated 2 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[alamb]

Thanks @ShashidharM0118 -- this is looking good. The basic idea certainly looks right to me.

I have a suggestion for exactly what setting to use. Let me know if that makes sense

[alamb]

I think there is already a config setting that would be appropriate for this setting (rather than using the batch size): https://datafusion.apache.org/user-guide/configs.html

datafusion.optimizer.repartition_file_min_size

What do you think?

[shashidhar-bm]

hi @alamb, thanks for the suggestion to use datafusion.optimizer.repartition_file_min_size — it makes more sense to me, so I’ve switched the heuristic over to be size-based.

I have two concerns:

1. I’ve added a follow-up num_rows check that is used only when total_byte_size is missing. In other words, the code first applies the size-based threshold, and only if no size statistics are available does it fall back to num_rows >= batch_size. Does this fallback approach look reasonable?

2. The hash_agg_group_by_partitioned_on_dicts test in physical_planner.rs (lines 3275–3304) was originally written to assert that a partitioned aggregate is produced on dictionary keys, but with the new size-based heuristic the small in-memory dict dataset no longer triggers repartitioning and the test now asserts mode: Single, which no longer matches the test name or original intent.

[gene-bordegaray]

Throwing my two cents in here. I think this configuration would be great as letting users "turn knobs" is a great way for extensibility in datafusion and have experiemented it with myself.

I see a use for this configuration in my work and I think this fallback behavior should not exist with the min_size configuration. As a user I prefer if I turn a knob to say declare a min_size that it sticks to this behavior without this fallback behavior.

Let me know your thoughts on this.

[shashidhar-bm]

Throwing my two cents in here. I think this configuration would be great as letting users "turn knobs" is a great way for extensibility in datafusion and have experiemented it with myself.

I see a use for this configuration in my work and I think this fallback behavior should not exist with the min_size configuration. As a user I prefer if I turn a knob to say declare a min_size that it sticks to this behavior without this fallback behavior.

Let me know your thoughts on this.

understood,
I removed fallback,