Repository navigation
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #25988 +/- ##
==========================================
+ Coverage 82.59% 82.73% +0.14%
==========================================
Files 1145 1147 +2
Lines 444096 449377 +5281
Branches 444096 449377 +5281
==========================================
+ Hits 366806 371798 +4992
+ Misses 55021 54920 -101
- Partials 22269 22659 +390 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
run benchmarks |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing codex/25936-numeric-distinct-count (3ae5962) to 416002a (merge-base) diff Run configurationrun benchmark clickbench_partitionedResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing codex/25936-numeric-distinct-count (3ae5962) to 416002a (merge-base) diff Run configurationrun benchmark tpchResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing codex/25936-numeric-distinct-count (3ae5962) to 416002a (merge-base) diff Run configurationrun benchmark tpcdsResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing codex/25936-numeric-distinct-count (3ae5962) to 416002a (merge-base) diff Run configurationrun benchmark tpchCPU Details (lscpu)Details
Resource Usagetpch — base (merge-base)
tpch — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing codex/25936-numeric-distinct-count (3ae5962) to 416002a (merge-base) diff Run configurationrun benchmark tpcdsCPU Details (lscpu)Details
Resource Usagetpcds — base (merge-base)
tpcds — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing codex/25936-numeric-distinct-count (3ae5962) to 416002a (merge-base) diff Run configurationrun benchmark clickbench_partitionedCPU Details (lscpu)Details
Resource Usageclickbench_partitioned — base (merge-base)
clickbench_partitioned — branch
File an issue against this benchmark runner |
|
Synced with main and fixed the spill-test conflict. The optimization still only applies with @Dandandan, could you rerun the benchmarks when you have a chance? The earlier run was on the version that also changed multi-partition plans. I also tried to rerun Rust CI on the fork, but GitHub rejects it with "Actions has been disabled for this repository", even though Actions is enabled in the repo settings. |
jayzhan211
left a comment
There was a problem hiding this comment.
Optional cleanup: in single_distinct_to_groupby.rs around line 70, the rule's doc comment still reads as if the rewrite always applies. A sentence on when grouped DISTINCT stays native would help the next reader.
| // keeping DISTINCT instead repartitions lists of values by group and merges | ||
| // them in the final accumulator. That can concentrate work on a few groups. | ||
| // Native accumulator support alone does not establish that this is cheaper. | ||
| if use_native_grouped_distinct |
There was a problem hiding this comment.
@QinXi-ai Thanks for working on this. With spilling on, this turns queries that finish under a memory limit into ResourcesExhausted: the rewrite dedupes through sorted (g, x) spill runs, but the native accumulator has to hold every distinct value of the largest group in memory. ClickBench Q8 at one partition finishes at -m 16M on main and fails at -m 128M here (it passes at 256M). The equal-memory table in #25936 used -d 0, where neither plan can spill.
SET datafusion.execution.target_partitions = 1;
SET datafusion.execution.batch_size = 128;
SET datafusion.runtime.memory_limit = '1M';
-- main: 2, 100000. This PR: Resources exhausted
SELECT count(*), sum(d) FROM (
SELECT v % 2 AS k, count(DISTINCT v) AS d
FROM generate_series(1, 100000) AS t(v)
GROUP BY v % 2
);The rule can't see the memory pool, so it can't choose per query. Could you make the native plan opt-in so the rewrite stays the default? #26000 adds a very similar flag to this rule, so it's worth settling on one.
!group_expr.is_empty()
&& config.options().execution.target_partitions == 1
&& config.options().optimizer.single_distinct_prefer_nativePlease add the query above to aggregate_memory_spill.slt as the flag-off test, plus an EXPLAIN with the flag on showing the native plan. If a maintainer would rather keep it on by default, it needs an upgrade guide note, and the "Query results ... are unchanged" line in the description needs correcting.
There was a problem hiding this comment.
Thanks, I missed this because the spill test only had one distinct value per group. I've made the native path opt-in with single_distinct_prefer_native, defaulting to false. Your 1 MiB reproducer now passes with the default; the test also checks the native EXPLAIN and memory error with the flag on. I used your prefer_native name since native can use more memory too. Agreed on sharing one setting with #26000.
| CAST((value * 48271) % 999983 AS VARCHAR) AS string_x, | ||
| CAST((value * 48271) % 999983 AS DOUBLE) AS float_x, | ||
| CASE WHEN value % 100 < 90 THEN 0 ELSE value % 8 END AS skew_g, | ||
| CAST(CASE WHEN value % 100 < 90 THEN 0 ELSE value % 8 END AS VARCHAR) AS string_g |
There was a problem hiding this comment.
The suite doesn't include the shape where the native plan loses: many groups with duplicate rows. With 1M groups and four rows per group, I measured it about 1.4x slower than main (about 130 ms vs 93 ms median, on a busy laptop), because every row probes both the group table and seen, while the rewrite's outer aggregate only sees each pair once. The rule has no statistics to avoid this, so could you add a q10 over these columns and state this worst case in the description?
| CAST(CASE WHEN value % 100 < 90 THEN 0 ELSE value % 8 END AS VARCHAR) AS string_g | |
| CAST(CASE WHEN value % 100 < 90 THEN 0 ELSE value % 8 END AS VARCHAR) AS string_g, | |
| value % 1000000 AS many_g, | |
| value % 2 AS dup_x |
There was a problem hiding this comment.
Added Q10 with those columns and reproduced the slowdown: at 4M rows and one target partition, I got 398.12 ms with native enabled vs 284.87 ms with it off, about 40% slower (15 retained samples each). I've put this alongside the original workload's improvement in the description and added off/on commands to the suite README.
| input.schema(), | ||
| count_rollup.as_ref(), | ||
| !group_expr.is_empty() | ||
| && config.options().execution.target_partitions == 1, |
There was a problem hiding this comment.
target_partitions == 1 doesn't guarantee one input partition: a UNION ALL or a table created at higher parallelism still plans Partial, Coalesce, Final, which is the list-merge path this PR avoids for parallel plans. A 4-partition table with a hot group ran about 1.3x slower than main in my runs, since Partial dedupes first and Final rebuilds the set. Could you add this test before the switch to four partitions, so that path is covered, and mention the case in the comment above use_native_grouped_distinct?
query II rowsort
SELECT g, count(DISTINCT x)
FROM (SELECT g, x FROM numeric_counts UNION ALL SELECT g, x FROM numeric_counts)
GROUP BY g;
----
1 2
2 0
3 1There was a problem hiding this comment.
You're right, one target partition can still merge several input partitions. Added your UNION ALL query and a table built at four partitions then queried at one. Both have result checks and physical-plan snapshots showing Partial/Coalesce/Final. The comment now calls out these cases.
| Ok(true) | ||
| } | ||
|
|
||
| fn all_have_groups_accumulators( |
There was a problem hiding this comment.
Optional cleanup: this repeats the loop in rewrite_pays_for_count, and the two only differ in which answer they look for. One helper covers both:
-fn all_have_groups_accumulators(
+/// Each distinct aggregate's answer to whether a specialized `GroupsAccumulator`
+/// exists for its argument types.
+fn distinct_groups_support(
distinct_aggs: &[(&Arc<AggregateUDF>, &[Expr])],
input_schema: &DFSchema,
-) -> Result<bool> {
- for (func, args) in distinct_aggs {
- let arg_types = args
- .iter()
- .map(|arg| arg.get_type(input_schema))
- .collect::<Result<Vec<_>>>()?;
- // An unknown answer must preserve the existing rewrite.
- if func.groups_accumulator_supported_for_types(&arg_types, true) != Some(true) {
- return Ok(false);
- }
- }
- Ok(true)
+) -> Result<Vec<Option<bool>>> {
+ distinct_aggs
+ .iter()
+ .map(|(func, args)| {
+ let arg_types = args
+ .iter()
+ .map(|arg| arg.get_type(input_schema))
+ .collect::<Result<Vec<_>>>()?;
+ Ok(func.groups_accumulator_supported_for_types(&arg_types, true))
+ })
+ .collect()
}The call here becomes distinct_groups_support(&distinct_aggs, input_schema)?.iter().all(|support| *support == Some(true)), and rewrite_pays_for_count becomes Ok(distinct_groups_support(distinct_aggs, input_schema)?.contains(&Some(false))).
There was a problem hiding this comment.
Done, both checks now share the type lookup through distinct_groups_support. I also added the opt-in exception to the rule's doc comment.
Preserve the spill-capable rewrite by default and document the native plan's memory and duplicate-heavy workload tradeoffs. Cover low-memory execution and multiple input partitions, and add the million-group duplicate benchmark.
Regenerate configuration documentation after merging main 4bf32cb. Generated-by: OpenAI Codex
Which issue does this PR close?
Closes #25936.
Rationale for this change
Grouped integer
count(DISTINCT x)can use its native GroupsAccumulator to avoid an extra aggregate over distinct(group, value)pairs. This helps the original single-target-partition workload, but the choice also affects spilling and performance on other data shapes.The native plan is now opt-in. The default keeps the existing two-phase rewrite. Under a 1 MiB memory limit, the reviewer's query with two groups and 100,000 distinct values completes with the rewrite and returns
2, 100000; the native plan runs out of memory. Native DISTINCT must fit each group's distinct values in memory, including when merging spilled or partial states.What changes are included in this PR?
Add
datafusion.optimizer.single_distinct_prefer_native, defaultfalse. When enabled, grouped all-DISTINCT aggregates retain the native plan only if:execution.target_partitions == 1;Unknown support preserves the rewrite. Integer DISTINCT counts currently qualify; string and floating-point arguments, global aggregates, mixed aggregates and multiple target partitions keep their existing behavior.
The earlier GKE ClickBench comparison regressed with native parallel plans, so multiple target partitions continue to distribute deduplication by both the group key and distinct value.
One target partition can still have multiple input partitions.
UNION ALLand a table created with four partitions can use Partial → Coalesce → Final and rebuild the DISTINCT set in Final. This is covered by execution tests and physical-plan snapshots.Add a ten-query SQL benchmark suite covering argument types, grouping-key types, cardinality, skew and duplicates. Q10 has one million groups with four copies of the same value per group at the default four million rows. Native execution probes both the group table and the distinct set on every row; the rewrite's outer aggregate only sees deduplicated pairs. This is an unfavorable case for the native path, and the logical rule has no statistics to cost it.
This overlaps with the option proposed in #26000.
prefer_nativedescribes the plan choice without promising lower memory use; the two PRs should converge on one setting.Local comparison, November 7
Compared main
636590bwith review-fix commit4a0bacd4, flag off and on, on Linux (Xeon Gold 6238T, CPU affinity 0–7), Rust 1.99.1,release-nonlto, four million rows, greedy / 1 GiB. All ten queries ran at 1, 4 and 12 target partitions in three rounds, reversing version order in the second round. Each round had six iterations per query, with the first discarded: 15 retained samples per version and configuration. All result assertions passed. The main executable was reused from the November 2 build; the review fix was freshly rebuilt, and all versions used the same updated ten-query suite. No builds ran during timing.Single-partition medians:
Q10 reproduces the reviewer's regression. Its memory-pool peak also rises from 69.88 MiB with the flag off to 95.88 MiB with it on. Q01 goes the other way, from 191.45 to 168.17 MiB. Both the runtime and memory tradeoffs are reasons to keep this opt-in.
The primary comparison is off/on in the same executable. Main and flag-off use the same plans but have timing differences on this shared host. At four and twelve target partitions both flag values keep the rewrite; observed off/on median differences range from −8.0% to +11.2%. Separate physical-plan checks for Q01 and Q10 passed for all three versions at all three partition settings (18 checks).
What is the testing strategy for this PR?
SQL tests cover the default and both flag values in one session, integer types and NULLs, the low-memory spill reproducer, native plan selection,
UNION ALL, a pre-partitioned table, and the rewrite at four target partitions. The configuration documentation and information-schema snapshots include the new default.On Linux at
4a0bacd4, formatting, generated configuration docs, Prettier, strict all-targets/all-features Clippy and the fulldev/rust_lint.shpassed. The extended workspace run passed 13,395 Rust tests (8 ignored, none failed) and all 547 SQL files. The CLI run passed 112 tests, and the focused SQL run passed all seven files.Are there any user-facing changes?
The new setting lets users opt into native grouped integer DISTINCT counts with one target partition. Default planning stays unchanged.