adriangb opened a new pull request, #24857: URL: https://github.com/apache/datafusion/pull/24857
## 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 `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 backing a `GROUP BY` on one string column. 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. Both maps also misreported the table's footprint. `ArrowBytesViewMap` seeded its `map_size` with `capacity() * size_of::<Entry<V>>()`, which leaves out the control bytes. `ArrowBytesMap` seeded it with 0 despite pre-allocating, and `HashTableAllocExt::insert_accounted` only charges on growth, so any map staying under its pre-allocated capacity reported its table as free forever. `clear_shrink` is the other half of the release path. `GroupValuesBytes::clear_shrink` and `GroupValuesBytesView::clear_shrink` went through `take`, 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<()>`) | | --- | --- | --- | | `Entry` size | 24 bytes | 32 bytes | | pre-allocated capacity requested | 128 | 512 | | hashbrown `capacity()` | 224 | 896 | | real table allocation | 6,408 bytes | 33,800 bytes | | old `capacity() * size_of::<Entry>()` | 5,376 bytes | 28,672 bytes | | undercount | 1.19x | 1.18x | One per-group accumulator holding a single 24-byte distinct value, `size()` in bytes: | | before, actual | before, reported | after, actual and reported | | --- | --- | --- | --- | | `BytesDistinctCountAccumulator` | 14,648 | 8,240 | 180 | | `BytesViewDistinctCountAccumulator` | 33,920 | 28,792 | 260 | The `ArrowBytesMap` row 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 through `GroupsAccumulatorAdapter`. 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? - Split the constructors of both maps: `new` allocates nothing, `with_capacity` keeps the previous pre-allocating behavior. The capacity is remembered so `take` re-creates the map the way it was built. - `GroupValuesBytes` and `GroupValuesBytesView` move to `with_capacity`. The two distinct-count accumulators stay on `new`. - Drop `map_size` in favour of `HashTable::allocation_size`, which is exact, covers the control bytes, and is a constant time layout calculation, so `size()` stays cheap. - Add `clear_and_release`, which drops every allocation the map holds and remembers the configured capacity so the map warms back up on the next `take`. `GroupValuesBytes::clear_shrink` and `GroupValuesBytesView::clear_shrink` now call it. - `datafusion/physical-expr-common/benches/arrow_bytes_map.rs` moves to `with_capacity` so it keeps measuring the pre-allocating constructor. Its `long_low_cardinality` case 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? This is a sizing and accounting change, so no query results change. It is covered by the existing suites for every crate it 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), and `datafusion-physical-plan -- group_values` (96). `cargo clippy --all-targets` is 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.rs` above. ## Are there any user-facing changes? Yes, in `datafusion-physical-expr-common`. `ArrowBytesMap::new` and `ArrowBytesViewMap::new` no longer pre-allocate; callers wanting the previous behavior should use the new `with_capacity`. Both types also gain `clear_and_release`. The change to `new` is a behavior change to an existing public constructor rather than an addition, so please let me know if you would like the `api change` label. Grouped `COUNT(DISTINCT)` on string and binary columns uses substantially less memory and reports its usage to the `MemoryPool` accurately, so a query that previously hit a memory limit may now succeed. No query results change. ## Follow-ups, not in this PR - The same undercount class remains at five other production `insert_accounted` call 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. - The `count_distinct_groups` benchmarks in `datafusion/functions-aggregate/benches/count_distinct.rs` cover `Int64`, `Int32` and `UInt32` only, so the headline win has per-accumulator byte measurements but no criterion evidence. - `GroupsAccumulatorAdapter` has no way to tell an accumulator it is one of many, so the ungrouped `COUNT(DISTINCT)` also loses its warm-up here. A capacity hint would let the two paths differ. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
