fix: stop pre-allocating and undercounting a hash table per COUNT(DISTINCT) group

[adriangb]

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.

Split the constructors: new allocates nothing, with_capacity keeps the previous 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.

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. They now call a new 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.

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.

Motivation

A production query worker 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.

Tests

Two new tests in datafusion/core/tests/memory_limit/mod.rs, group_by_count_distinct_utf8 and group_by_count_distinct_utf8_view, make the headline claim a binary observable rather than a number: a grouped COUNT(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 and target_partitions pinned 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:

value column	before	after	limit the test uses
Utf8	~35.5 MB (fails at 35 MB, passes at 36 MB)	~1.9 MB (fails at 1.8 MB, passes at 2.0 MB)	8 MB
Utf8View	~123 MB (fails at 120 MB, passes at 124 MB)	~2.5 MB (fails at 2.4 MB, passes at 2.6 MB)	16 MB

Each 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, and avg specifically. Without a second aggregate, single_distinct_aggregation_to_group_by rewrites the distinct aggregate into a plain two stage GROUP BY that does not use these accumulators at all, and the tests would pass by construction. That rule tolerates a non-distinct sum, min or max beside the distinct aggregate, because it re-aggregates its own partial results over the deduplicated inner group by and those three compose with themselves. avg does not, so the rule can never accept it, which is why ClickBench Q9 keeps its distinct aggregate.

This matters because apache#24859 proposes adding count to that allow list. Checked by cherry-picking apache#24859 onto this branch and re-planning:

-- avg version, unchanged, still goes through GroupsAccumulatorAdapter
AggregateExec: mode=Final, gby=[group_key@0], aggr=[count(DISTINCT t.value), avg(t.payload)]
  AggregateExec: mode=Partial, gby=[group_key@0], aggr=[count(DISTINCT t.value), avg(t.payload)]
    DataSourceExec: partitions=1, partition_sizes=[1]

-- count(*) version, rewritten away, no per group accumulators left
ProjectionExec: expr=[group_key@0, count(alias1)@1 as count(DISTINCT t.value), ...]
  AggregateExec: mode=Final, gby=[group_key@0], aggr=[count(alias1), sum(alias2)]
    AggregateExec: mode=Partial, gby=[group_key@0], aggr=[count(alias1), sum(alias2)]
      AggregateExec: mode=Final, gby=[group_key@0, alias1@1], aggr=[count(1) as alias2]
        AggregateExec: mode=Partial, gby=[group_key@0, value@1 as alias1], aggr=[]
          DataSourceExec: partitions=1, partition_sizes=[1]

With apache#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. The avg form needs ~1.9 MB either way.

Existing suites, 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 full memory_limit module 34. cargo clippy --all-targets clean on all three crates. No query results change.

Benchmarks were not re-run for this revision. 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.

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.

🤖 Generated with Claude Code

[adriangb]

Rebased from the DataFusion 55 fork branch (friendlymatthew/pydantic-main-df55) onto main at da89c7c8, and retargeted the PR base to main. The four fork-only commits that used to sit under this branch are gone from the diff; the four commits here are unchanged in intent.

Conflicts were confined to imports: HashTableAllocExt is no longer used by either map file, while main has since added Result and exec_err there for the new keys method, and single_group_by/bytes.rs now also imports GroupSelection. Upstream had not touched the constructors or the size accounting, so nothing in this change became redundant.

All figures in the description were re-derived on the new base rather than carried over. The ArrowBytesViewMap undercount is 1.18x as before (28,672 reported against 33,800 real), and the per-accumulator numbers moved slightly: the view accumulator holding one 24-byte value now measures 33,920 actual and 28,792 reported before the change, 260 after. Benchmarks were not re-run.

Tests: 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, clippy --all-targets clean on all three crates.

[adriangb]

Note on the clear_shrink review suggestion, since the thread is resolved and the reasoning is not recorded anywhere.

The suggestion was to construct GroupValuesBytes with the non-pre-allocating ArrowBytesMap::new, so that clear_shrink could actually shrink. This branch does the opposite and keeps the warm-up at construction, adding clear_and_release so clear_shrink releases the allocation directly.

Two reasons:

1. GroupValuesBytes holds one map for the whole query, which is exactly the case the pre-allocation was designed for and the case arrow_bytes_map benchmarks. Removing the warm-up there would trade a real cost on the hot string GROUP BY path for a benefit only on the spill path.
2. The retention was not introduced by this branch. Before it, take() was let mut new_self = Self::new(self.output_type); swap(...), and new pre-allocated 128 entries plus an 8 KiB buffer, so clear_shrink retained exactly as much as it does after the constructor split. What this branch newly makes possible is releasing it at all, since before there was no non-pre-allocating constructor.

So the finding identified a real defect, and the fix here addresses it without giving up the warm-up where it earns its keep.

[macroscopeapp]

Macroscope has since reviewed this pull request. An earlier review was skipped by a cost limit; a review has now completed, so that notice no longer applies.