fix: stop pre-allocating and undercounting a hash table per COUNT(DISTINCT) group - #24857
Draft
adriangb wants to merge 6 commits into
Draft
fix: stop pre-allocating and undercounting a hash table per COUNT(DISTINCT) group#24857adriangb wants to merge 6 commits into
adriangb wants to merge 6 commits into
Conversation
…wBytesViewMap Both maps tracked their hash table footprint in a `map_size` field that was only ever incremented by `HashTableAllocExt::insert_accounted`, which charges `capacity * size_of::<Entry>()` on growth and nothing else. That undercounts in two ways. `ArrowBytesViewMap::new` seeded `map_size` with `capacity() * size_of::<Entry<V>>()`, which ignores the control bytes and the trailing group that hashbrown allocates alongside the entry array, so the reported size was roughly half the real allocation. `ArrowBytesMap::new` seeded `map_size` with 0 despite pre-allocating a table for 128 entries. Since `insert_accounted` only charges when the table grows, any map holding fewer entries than the pre-allocated capacity reported its hash table as free forever. Drop the field and ask hashbrown for the exact figure with `HashTable::allocation_size`, which covers entries, control bytes and the trailing group. It is a constant time layout calculation, so `size()` stays cheap, and it cannot drift out of sync with the table the way an incrementally maintained counter can.
`ArrowBytesMap` and `ArrowBytesViewMap` always pre-allocated their hash table, and `ArrowBytesMap` also pre-allocated an 8 KiB value buffer. That is the right trade for the single map that backs a `GROUP BY` on one string column, which goes on to hold every group value in the query. It is the wrong trade for `BytesDistinctCountAccumulator` and `BytesViewDistinctCountAccumulator`, because `GroupsAccumulatorAdapter` creates one accumulator per group: a grouped `COUNT(DISTINCT)` over a high cardinality key holds hundreds of thousands of them at once, and most see only a handful of values, so the pre-allocation dwarfs the data. Split the constructors. `new` no longer allocates anything, and `with_capacity` keeps the previous behavior for the callers that want it. The capacity is stored so `take` re-creates the map the way it was built. The `GroupValuesBytes` and `GroupValuesBytesView` call sites move to `with_capacity`; the two distinct-count accumulators stay on `new`. The `arrow_bytes_map` benchmark also moves to `with_capacity`: its `long_low_cardinality` case is defined by the distinct values fitting inside the pre-allocated buffer.
Keep the comment about what `HashTable::allocation_size` covers next to the value it describes, and say what the test helper's lower bound is derived from.
`GroupValuesBytes::clear_shrink` and `GroupValuesBytesView::clear_shrink` reset their map with `take()`, which restores the capacity the map was configured with so the emptied map stays warm. That is what the emit path wants, but `clear_shrink` exists to hand memory back before spilling and before the spilled batch is sorted, so it left roughly 16 KiB (string and binary) and 34 KiB (view) reserved instead of releasing it. Add `clear_and_release` to `ArrowBytesMap` and `ArrowBytesViewMap`, which empties the map and drops its allocations while remembering the configured capacities so a later `take()` still warms the map up, and call it from the two `clear_shrink` implementations. The pre-allocation stays at construction, where the hot single column string `GROUP BY` path earns it.
github-actions
Bot
added
physical-expr
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #24857 +/- ##
==========================================
+ Coverage 81.60% 81.67% +0.07%
==========================================
Files 1123 1123
Lines 408898 410550 +1652
Branches 408898 410550 +1652
==========================================
+ Hits 333670 335320 +1650
- Misses 55625 55633 +8
+ Partials 19603 19597 -6 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Contributor
Author
|
run benchmark clickbench_partitioned external_aggr |
Contributor
Author
|
run benchmark clickbench_extended |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing claude/bytes-map-initial-capacity-accounting (84f07da) to da89c7c (merge-base) diff Run configurationrun benchmark external_aggr
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: "4G"Results will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing claude/bytes-map-initial-capacity-accounting (84f07da) to da89c7c (merge-base) diff Run configurationrun benchmark clickbench_partitioned
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: "4G"Results will be posted here when complete File an issue against this benchmark runner |
A grouped `COUNT(DISTINCT <string>)` gets one accumulator per group, and each of those owns a hash set of the distinct values it has seen. Those sets were created pre-allocated, so the query's memory use tracked the number of groups rather than the amount of data. Add two `memory_limit` tests that turn that into a binary observable, one for `Utf8` and one for `Utf8View`, over a new scenario of 4,000 groups holding 2 distinct values each. Measured against this branch's base commit with spilling disabled and `target_partitions` pinned to 1: | value column | budget needed before | budget needed after | | ------------ | -------------------- | ------------------- | | `Utf8` | ~35.5 MB | ~1.9 MB | | `Utf8View` | ~123 MB | ~2.7 MB | The tests run at 8 MB and 16 MB respectively, so each sits at least 4x above what the branch needs and at least 4x below what the base needs. Both fail on the base commit with `Resources exhausted` and pass here.
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing claude/bytes-map-initial-capacity-accounting (84f07da) to da89c7c (merge-base) diff Run configurationrun benchmark clickbench_partitioned
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: "4G"CPU Details (lscpu)Details
Memory Pool PeaksPeak Base:
Pool accounting vs. process RSS Max pool peak is the largest reservation any single query in the run reached; peak RSS covers the whole invocation, including data loading and allocator retention, and the two high-water marks need not coincide in time. The gap is therefore an upper bound on what the pool did not account for, not a measurement of it.
Resource Usageclickbench_partitioned — base (merge-base)
clickbench_partitioned — branch
File an issue against this benchmark runner |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing claude/bytes-map-initial-capacity-accounting (84f07da) to da89c7c (merge-base) diff Run configurationrun benchmark clickbench_extended
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: "4G"Results will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing claude/bytes-map-initial-capacity-accounting (84f07da) to da89c7c (merge-base) diff Run configurationrun benchmark external_aggr
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: "4G"CPU Details (lscpu)Details
Memory Pool PeaksPeak Base:
Pool accounting vs. process RSS Max pool peak is the largest reservation any single query in the run reached; peak RSS covers the whole invocation, including data loading and allocator retention, and the two high-water marks need not coincide in time. The gap is therefore an upper bound on what the pool did not account for, not a measurement of it.
Resource Usageexternal_aggr — base (merge-base)
external_aggr — branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing claude/bytes-map-initial-capacity-accounting (84f07da) to da89c7c (merge-base) diff Run configurationrun benchmark clickbench_extended
env:
DATAFUSION_RUNTIME_MEMORY_LIMIT: "4G"CPU Details (lscpu)Details
Memory Pool PeaksPeak Base:
Pool accounting vs. process RSS Max pool peak is the largest reservation any single query in the run reached; peak RSS covers the whole invocation, including data loading and allocator retention, and the two high-water marks need not coincide in time. The gap is therefore an upper bound on what the pool did not account for, not a measurement of it.
Resource Usageclickbench_extended — base (merge-base)
clickbench_extended — branch
File an issue against this benchmark runner |
The two grouped `COUNT(DISTINCT <string>)` memory limit tests only reach the per group accumulators while `single_distinct_aggregation_to_group_by` declines to rewrite the query. They leant on `count(*)` for that, which the rule rejects only because `count` is missing from the `sum`/`min`/`max` allow list. #24859 proposes adding `count` to that list, which would rewrite the query, remove the accumulators, and leave both tests passing at any memory limit while still looking like they test something. Aggregate `avg(payload)` over a new `Int64` column instead. `avg` cannot be added to that list: the rule re-aggregates its own partial results over the deduplicated inner group by, and averaging per group averages of different sizes gives the wrong answer. That is why ClickBench Q9 keeps its distinct aggregate under #24859. Verified from the physical plan with #24859 cherry-picked on top of this branch: the `avg` query still plans as `aggr=[count(DISTINCT t.value), avg(t.payload)]`, while the `count(*)` query becomes `aggr=[count(alias1), sum(alias2)]` over an inner `GROUP BY group_key, value`, and drops from needing ~1.9 MB to ~0.9 MB. Re-swept both thresholds against the base commit. `Utf8` needs ~35.5 MB before and ~1.9 MB after; `Utf8View` needs ~123 MB before and ~2.5 MB after, so the 8 MB and 16 MB limits keep at least 4x margin on each side and are unchanged.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Which issue does this PR close?
No existing issue. This was found while investigating a production out of memory. Happy to file one if you would like it tracked for the changelog.
Rationale for this change
ArrowBytesMapandArrowBytesViewMapalways pre-allocated their hash table, andArrowBytesMapalso pre-allocated an 8 KiB value buffer. That is the right trade for the single map backing aGROUP BYon one string column. It is the wrong trade forBytesDistinctCountAccumulatorandBytesViewDistinctCountAccumulator, becauseGroupsAccumulatorAdaptercreates one accumulator per group. A groupedCOUNT(DISTINCT)over a high cardinality key holds hundreds of thousands of them at once and most see only a handful of values, so the pre-allocation dwarfs the data.Both maps also misreported the table's footprint.
ArrowBytesViewMapseeded itsmap_sizewithcapacity() * size_of::<Entry<V>>(), which leaves out the control bytes.ArrowBytesMapseeded it with 0 despite pre-allocating, andHashTableAllocExt::insert_accountedonly charges on growth, so any map staying under its pre-allocated capacity reported its table as free forever.clear_shrinkis the other half of the release path.GroupValuesBytes::clear_shrinkandGroupValuesBytesView::clear_shrinkwent throughtake, which restores the configured warm-up capacity, so the memory the aggregate stream intends to hand back before spilling and before a downstream sort was never actually released.Measured on this base
Hash table sizing, one map, nothing inserted:
ArrowBytesMap(Utf8,Entry<i32, ()>)ArrowBytesViewMap(Utf8View,Entry<()>)Entrysizecapacity()capacity() * size_of::<Entry>()One per-group accumulator holding a single 24-byte distinct value,
size()in bytes:BytesDistinctCountAccumulatorBytesViewDistinctCountAccumulatorThe
ArrowBytesMaprow is the more extreme reporting error: the map really held 14,648 bytes and reported 8,240, because the whole 6,408-byte table was invisible to the old accounting.Production motivation
A production process running DataFusion died holding 10.98 GB, roughly 75% of it in per-group
COUNT(DISTINCT)accumulators allocated throughGroupsAccumulatorAdapter. At ~254,000 live view accumulators this change takes that from ~8.6 GB to ~66 MB, and makes the reported figure exact rather than ~1.3 GB short.What changes are included in this PR?
newallocates nothing,with_capacitykeeps the previous pre-allocating behavior. The capacity is remembered sotakere-creates the map the way it was built.GroupValuesBytesandGroupValuesBytesViewmove towith_capacity. The two distinct-count accumulators stay onnew.map_sizein favour ofHashTable::allocation_size, which is exact, covers the control bytes, and is a constant time layout calculation, sosize()stays cheap.clear_and_release, which drops every allocation the map holds and remembers the configured capacity so the map warms back up on the nexttake.GroupValuesBytes::clear_shrinkandGroupValuesBytesView::clear_shrinknow call it.datafusion/physical-expr-common/benches/arrow_bytes_map.rsmoves towith_capacityso it keeps measuring the pre-allocating constructor. Itslong_low_cardinalitycase is defined by the distinct values fitting inside the pre-allocated buffer, so switching it to the lazy constructor would change what the benchmark measures rather than how fast it runs.What is the testing strategy for this PR?
Two new tests in
datafusion/core/tests/memory_limit/mod.rs,group_by_count_distinct_utf8andgroup_by_count_distinct_utf8_view, make the headline claim a binary observable rather than a number: a groupedCOUNT(DISTINCT <string>)that could not run inside a memory limit before this change completes inside it now. They aggregate a new scenario of 4,000 groups holding 2 distinct values each, with spilling disabled andtarget_partitionspinned to 1, so completing means the query genuinely fit in the budget rather than spilled out of it.The minimum budget the same query needs, swept against this PR's base commit:
Utf8Utf8ViewEach limit sits at least 4x above what this branch needs and at least 4x below what the base needs, so neither test is on a cliff edge. Checked out onto the base commit both fail with
Resources exhausted: Additional allocation failed for FinalHashAggregateStream[0]; on this branch both pass, and they pass on 5 consecutive runs.The
avg(payload)in the test query is load bearing, andavgspecifically. Without a second aggregate,single_distinct_aggregation_to_group_byrewrites the distinct aggregate into a plain two stageGROUP BYthat does not use these accumulators at all, and the tests would pass by construction. That rule tolerates a non-distinctsum,minormaxbeside the distinct aggregate, because it re-aggregates its own partial results over the deduplicated inner group by and those three compose with themselves.avgdoes not, so the rule can never accept it, which is why ClickBench Q9 keeps its distinct aggregate.This matters because #24859 proposes adding
countto that allow list. Checked by cherry-picking #24859 onto this branch and re-planning:With #24859 applied the
count(*)form drops from needing ~1.9 MB to ~0.9 MB, so it would have passed the 8 MB test for the wrong reason and on the base commit too. Theavgform needs ~1.9 MB either way.No query results change. The rest is covered by the existing suites for every crate this touches, all run locally and passing:
datafusion-physical-expr-common(85 lib, 8 doc),datafusion-functions-aggregate-common(47),datafusion-functions-aggregate -- count_distinct(2),datafusion-physical-plan -- group_values(96), and the fullmemory_limitmodule (34).cargo clippy --all-targetsis clean on all three crates.The per-accumulator byte figures in the tables above were measured directly on this base rather than asserted in a test, since the exact numbers depend on the hashbrown layout.
Benchmarks were not re-run for this revision, for the reason given about
arrow_bytes_map.rsabove.Are there any user-facing changes?
Yes, in
datafusion-physical-expr-common.ArrowBytesMap::newandArrowBytesViewMap::newno longer pre-allocate; callers wanting the previous behavior should use the newwith_capacity. Both types also gainclear_and_release. The change tonewis a behavior change to an existing public constructor rather than an addition, so please let me know if you would like theapi changelabel. GroupedCOUNT(DISTINCT)on string and binary columns uses substantially less memory and reports its usage to theMemoryPoolaccurately, so a query that previously hit a memory limit may now succeed. No query results change.Follow-ups, not in this PR
insert_accountedcall sites (group_values/row.rs:171,multi_group_by/mod.rs:434,554,multi_group_by/dictionary.rs:197,584,array_agg.rs:989). All are one-map-per-query so the absolute error is bounded, and the fix is the same one-line swap.count_distinct_groupsbenchmarks indatafusion/functions-aggregate/benches/count_distinct.rscoverInt64,Int32andUInt32only, so the headline win has per-accumulator byte measurements but no criterion evidence.GroupsAccumulatorAdapterhas no way to tell an accumulator it is one of many, so the ungroupedCOUNT(DISTINCT)also loses its warm-up here. A capacity hint would let the two paths differ.