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

[adriangb]

A grouped COUNT(DISTINCT <string>) over 4,000 groups holding 2 short strings each needs a 36 MB memory budget. It needs 2.0 MB after this change.

Every group gets its own hash table, and each table is allocated at warm-up size before the group holds anything, so the memory the query needs tracks the number of groups rather than the amount of data. The query also reports less memory than it holds, so a memory limit does not stop it at the right point.

Reproduction

This needs only datafusion-cli. There is no patch, no custom allocator and no data file.

-- repro.sql
SET datafusion.execution.target_partitions = 1;

-- 4,000 groups with 2 distinct short strings in each.
-- avg() stops SingleDistinctToGroupBy from rewriting the distinct aggregate away.
SELECT g, count(DISTINCT s) AS d, avg(p) AS a
FROM (
  SELECT v % 4000 AS g, 'v' || CAST(v AS VARCHAR) AS s, v AS p
  FROM generate_series(0, 7999) AS t(v)
)
GROUP BY g
ORDER BY g
LIMIT 3;

datafusion-cli -m 8M -f repro.sql

s is a Utf8View column, so this exercises ArrowBytesViewMap. The 8,000 rows arrive in one batch, so the aggregate builds all 4,000 accumulators before it can emit or spill.

Current main at 20d1c56761 fails:

Resources exhausted: Additional allocation failed for SingleHashAggregateStream[0] with top memory
consumers (across reservations) as:
  DataFusion-Cli#1(can spill: false) consumed 0.0 B, peak 0.0 B,
  SingleHashAggregateStream[0]#2(can spill: true) consumed 0.0 B, peak 48.0 B,
  TopK[0]#3(can spill: false) consumed 0.0 B, peak 0.0 B.
Error: Failed to allocate additional 111.1 MB for SingleHashAggregateStream[0] with 0.0 B already
allocated for this reservation - 8.0 MB remain available for the total memory pool:
greedy(used: 0.0 B, pool_size: 8.0 MB)

This branch returns the rows:

+---+---+--------+
| g | d | a      |
+---+---+--------+
| 0 | 2 | 2000.0 |
| 1 | 2 | 2001.0 |
| 2 | 2 | 2002.0 |
+---+---+--------+
3 row(s) fetched.
Elapsed 0.006 seconds.

Those 4,000 accumulators hold 8,000 short strings, which is about 100 KB of data. The base asks the pool for 111.1 MB to hold it. This branch runs the same query inside -m 3M. Both builds return the same rows, and the base does so at -m 200M. Each run takes well under a second, and the outcome repeats exactly over three runs on each side.

How much it improves

The minimum memory limit at which that query completes, bisected on each side:

value column	before	after
Utf8	fails 34 MB, passes 36 MB	fails 1.8 MB, passes 2.0 MB
Utf8View	fails 120 MB, passes 124 MB	fails 2.4 MB, passes 2.6 MB

clickbench_extended at DATAFUSION_RUNTIME_MEMORY_LIMIT: 4G, pool peak over six runs:

query	base	this branch	change
Q2, grouped, 4 string distincts	98.1 to 98.6 MiB	11.7 MiB in all six runs	-88.1%
Q1, ungrouped string distincts	3.4 MiB	2.4 MiB	-29.4%
Q0, ungrouped, high cardinality	796.8 to 834.8 MiB	846.5 to 885.5 MiB	+6.3%

Q2 is the only query in any benchmark suite that puts a grouped COUNT(DISTINCT) on a non-integer column.

Q0 costs more, and it is the one disclosed cost of this PR. Those extra bytes are memory Q0 always held and the pool could not see, not new allocation; appendix A has the three-build decomposition that separates the two. Latency does not move anywhere, which is what an allocation-sizing change should do.

Which issue does this PR close?

No existing issue. I found this when I investigated a production out of memory. I can file an issue if you want it in the changelog.

Rationale for this change

Pre-allocating is right for the one long-lived map behind a GROUP BY on a string column. It is wrong for the distinct-count accumulators, because GroupsAccumulatorAdapter creates one accumulator per group and most groups hold a handful of values. There the warm-up dwarfs the data.

Both maps also under-report the table they hold. ArrowBytesViewMap left the control bytes out. ArrowBytesMap charged the table only when it grew, so a map that stayed inside its pre-allocation reported its table as free forever. A memory limit acted on a number that was too small.

clear_shrink is the third part. The aggregate stream calls it to hand memory back before it spills and before a downstream sort. It restored the warm-up capacity instead of releasing it, so nothing came back.

What changes are included in this PR?

• new on both maps allocates nothing. A new with_capacity keeps the previous pre-allocating behavior. GroupValuesBytes and GroupValuesBytesView use with_capacity, and the two distinct-count accumulators use new. A map remembers how it was built, so take warms it back up the way it started.
• size() reports HashTable::allocation_size(), the real hashbrown allocation including the control bytes, in place of the old estimate.
• A new clear_and_release drops every allocation the map holds, and clear_shrink calls it.
• The value buffer rounds each growth up to a power of two. A lazily grown buffer and a pre-allocated one then sit on one ladder, so a lazy map is never the larger of the two for the same values. Growth stays geometric.
• benches/arrow_bytes_map.rs moves to with_capacity so it keeps measuring the pre-allocating constructor.

What is the testing strategy for this PR?

Two tests in datafusion/core/tests/memory_limit/mod.rs, group_by_count_distinct_utf8 and group_by_count_distinct_utf8_view, turn the headline claim into a pass or a fail rather than a number. They run the reproduction query over 4,000 groups with spilling disabled and target_partitions pinned to 1, so completing means the query fits the budget rather than spills out of it. The limits are 8 MB and 16 MB, at least 4x clear of both cliffs in the table above. Both tests fail on the merge base and pass here, over five consecutive runs. The avg(payload) in the query is load bearing; appendix C says why.

Unit tests cover what the memory-limit tests cannot see: that new allocates nothing, that with_capacity reports a table size bracketed by an independently derived lower bound, that take preserves the configured capacity, that clear_shrink drops the reported size to near zero, and that a lazily grown buffer never exceeds a pre-allocated one holding the same values.

Run locally on the rebased head, all passing: datafusion-physical-expr-common (87 lib, 8 doc), datafusion-functions-aggregate-common (49), datafusion-functions-aggregate -- count_distinct (2), datafusion-physical-plan -- group_values (96) and the memory_limit module (39, which includes the count_distinct_spill test that arrived on main in apache#24888 and apache#24918). cargo fmt --check and cargo clippy --all-targets -D warnings are clean on the changed crates. CI has not yet run this branch against the new base.

No query results change.

Are there any user-facing changes?

Yes, in datafusion-physical-expr-common. ArrowBytesMap::new and ArrowBytesViewMap::new no longer pre-allocate, and callers that want the previous behavior should use the new with_capacity. Both types also gain clear_and_release. This changes an existing public constructor rather than adding one, so tell me if you would like the api change label.

For users, a grouped COUNT(DISTINCT) on string and binary columns uses much less memory and reports its usage to the MemoryPool accurately. A query that previously hit a memory limit may now succeed.

---

Appendix A: Query 0 costs 6.3% more

Q0 is COUNT(DISTINCT) over three high-cardinality strings with no GROUP BY. It is a handful of maps that each grow to millions of entries, which is the opposite population from the one this PR targets. The pre-allocation was never the dominant cost there, so removing it buys nothing.

Over six runs the base spans 796.8 to 834.8 MiB and this branch spans 846.5 to 885.5 MiB. The ranges do not overlap, so the effect is real and not run-to-run noise.

Three local builds on a deterministic subset separate the two changes. The middle build differs from the base only in the accounting, because restoring the warm-up makes the constructors byte-identical to base:

build	Q0 pool peak
base	48,421,820
this branch with the warm-up restored	51,048,396
this branch	51,018,828

That decomposes the increase exactly. +2,626,576 is the accounting correction: allocation_size() charges the real hashbrown allocation, which is 4 * buckets + 5,384 more than the old formula, being the control bytes plus the 7/8 load-factor slack. -29,568 is the lazy constructor, which makes Q0 slightly better.

Reverting the accounting would restore an under-report of about 19% on this path. That under-report is the bug this PR exists to fix, and the memory-limit result above depends on fixing it.

Appendix B: what one accumulator costs

One per-group accumulator holding a single 24-byte value:

	before, actual	before, reported	after
BytesDistinctCountAccumulator	14,648 B	8,240 B	180 B
BytesViewDistinctCountAccumulator	33,920 B	28,792 B	260 B

The middle column is the reporting gap. The Utf8 map really held 14,648 bytes and reported 8,240, because the whole hash table was invisible to the old accounting.

These are measured directly rather than asserted in a test, since the exact numbers follow the hashbrown layout.

Appendix C: notes on the tests and the benchmarks

The memory-limit query uses avg(payload), not count(*). A non-distinct count lets SingleDistinctToGroupBy rewrite the distinct aggregate into a plain two-stage GROUP BY. The per-group accumulators would then never exist, and the tests would pass by construction on the base commit too. That rule accepts a non-distinct sum, min or max because each re-aggregates its own partial results correctly over the deduplicated inner group by. avg does not, so the rule can never admit it under any extension, including the one apache#24859 proposes.

The benchmark figures were measured against the previous merge base da89c7c85b. They have not been re-run against the current base 20d1c56761. The commits after 84f07da on this branch touch only datafusion/core/tests/memory_limit/mod.rs, so nothing on this branch since then can move a benchmark, but the base itself has moved.

Pool peak is the instrument here, not peak RSS. Pool peak reproduces to under 1% on a 98 MiB query and exactly on the 3.4 and 11.7 MiB ones. Peak RSS on this harness has a 4.1% standard deviation over 11 readings of identical code plus a 2.3% order bias, and shows no effect from this change once that null is accounted for.

Follow-ups, not in this PR

• The same undercount 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 and array_agg.rs:989. Each is 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 this path has no criterion coverage.

[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 skipped reviewing this pull request. Per-review cost limit exceeded (workspace setting).

This review would cost an estimated $16.75, which exceeds your per-review limit of $10.00.

The top 3 files driving up this estimate:

File	Diff Size	Estimate
datafusion/physical-plan/src/aggregates/hash_stream.rs	30.58KB	$1.53
datafusion/physical-plan/src/joins/nested_loop_join.rs	23.72KB	$1.19
datafusion/physical-expr/src/expressions/cast.rs	21.87KB	$1.09

Tip

To get this pull request reviewed, you can:

1. Comment @macroscope-app on this PR to request a manual review (monthly spend limits still apply).
2. Exclude the file(s) above from review by adding a pattern to your .macroscope/ignore.md — note that creating this file replaces Macroscope's built-in default ignores rather than extending them.
3. Raise your cost limit in your workspace billing settings.

^{Turn off this reminder going forward}

[github-actions]

Thank you for opening this pull request!

Reviewer note: cargo-semver-checks reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch).

Details

     Cloning apache/main
    Building datafusion v55.0.0 (current)
       Built [  86.535s] (current)
     Parsing datafusion v55.0.0 (current)
      Parsed [   0.026s] (current)
    Building datafusion v55.0.0 (baseline)
       Built [  44.416s] (baseline)
     Parsing datafusion v55.0.0 (baseline)
      Parsed [   0.027s] (baseline)
    Checking datafusion v55.0.0 -> v55.0.0 (no change; assume patch)
     Checked [   0.798s] 223 checks: 222 pass, 1 fail, 0 warn, 31 skip

--- failure method_requires_different_generic_type_params: method now requires a different number of generic type parameters ---

Description:
A method now requires a different number of generic type parameters than it used to. Uses of this method that supplied the previous number of generic types will be broken.
        ref: https://doc.rust-lang.org/reference/items/generics.html
       impl: https://github.com/obi1kenobi/cargo-semver-checks/tree/v0.50.0/src/lints/method_requires_different_generic_type_params.ron

Failed in:
  datafusion::prelude::dataframe::DataFrame::from_columns takes 1 generic types instead of 0, in /home/runner/work/datafusion/datafusion/datafusion/core/src/dataframe/mod.rs:2631
  datafusion::dataframe::DataFrame::from_columns takes 1 generic types instead of 0, in /home/runner/work/datafusion/datafusion/datafusion/core/src/dataframe/mod.rs:2631
  datafusion::prelude::DataFrame::from_columns takes 1 generic types instead of 0, in /home/runner/work/datafusion/datafusion/datafusion/core/src/dataframe/mod.rs:2631

     Summary semver requires new major version: 1 major and 0 minor checks failed
    Finished [ 134.209s] datafusion
    Building datafusion-expr v55.0.0 (current)
       Built [  21.617s] (current)
     Parsing datafusion-expr v55.0.0 (current)
      Parsed [   0.057s] (current)
    Building datafusion-expr v55.0.0 (baseline)
       Built [  21.745s] (baseline)
     Parsing datafusion-expr v55.0.0 (baseline)
      Parsed [   0.058s] (baseline)
    Checking datafusion-expr v55.0.0 -> v55.0.0 (no change; assume patch)
     Checked [   1.645s] 223 checks: 223 pass, 31 skip
     Summary no semver update required
    Finished [  46.621s] datafusion-expr
    Building datafusion-functions v55.0.0 (current)
       Built [  24.046s] (current)
     Parsing datafusion-functions v55.0.0 (current)
      Parsed [   0.064s] (current)
    Building datafusion-functions v55.0.0 (baseline)
       Built [  24.300s] (baseline)
     Parsing datafusion-functions v55.0.0 (baseline)
      Parsed [   0.065s] (baseline)
    Checking datafusion-functions v55.0.0 -> v55.0.0 (no change; assume patch)
     Checked [   0.530s] 223 checks: 223 pass, 31 skip
     Summary no semver update required
    Finished [  50.253s] datafusion-functions
    Building datafusion-functions-aggregate-common v55.0.0 (current)
       Built [  16.885s] (current)
     Parsing datafusion-functions-aggregate-common v55.0.0 (current)
      Parsed [   0.016s] (current)
    Building datafusion-functions-aggregate-common v55.0.0 (baseline)
       Built [  16.732s] (baseline)
     Parsing datafusion-functions-aggregate-common v55.0.0 (baseline)
      Parsed [   0.015s] (baseline)
    Checking datafusion-functions-aggregate-common v55.0.0 -> v55.0.0 (no change; assume patch)
     Checked [   0.168s] 223 checks: 223 pass, 31 skip
     Summary no semver update required
    Finished [  34.726s] datafusion-functions-aggregate-common
    Building datafusion-functions-nested v55.0.0 (current)
       Built [  27.608s] (current)
     Parsing datafusion-functions-nested v55.0.0 (current)
      Parsed [   0.029s] (current)
    Building datafusion-functions-nested v55.0.0 (baseline)
       Built [  27.081s] (baseline)
     Parsing datafusion-functions-nested v55.0.0 (baseline)
      Parsed [   0.030s] (baseline)
    Checking datafusion-functions-nested v55.0.0 -> v55.0.0 (no change; assume patch)
     Checked [   0.275s] 223 checks: 223 pass, 31 skip
     Summary no semver update required
    Finished [  56.423s] datafusion-functions-nested
    Building datafusion-optimizer v55.0.0 (current)
       Built [  20.650s] (current)
     Parsing datafusion-optimizer v55.0.0 (current)
      Parsed [   0.023s] (current)
    Building datafusion-optimizer v55.0.0 (baseline)
       Built [  21.332s] (baseline)
     Parsing datafusion-optimizer v55.0.0 (baseline)
      Parsed [   0.025s] (baseline)
    Checking datafusion-optimizer v55.0.0 -> v55.0.0 (no change; assume patch)
     Checked [   0.214s] 223 checks: 223 pass, 31 skip
     Summary no semver update required
    Finished [  43.369s] datafusion-optimizer
    Building datafusion-physical-expr v55.0.0 (current)
       Built [  23.480s] (current)
     Parsing datafusion-physical-expr v55.0.0 (current)
      Parsed [   0.037s] (current)
    Building datafusion-physical-expr v55.0.0 (baseline)
       Built [  22.796s] (baseline)
     Parsing datafusion-physical-expr v55.0.0 (baseline)
      Parsed [   0.043s] (baseline)
    Checking datafusion-physical-expr v55.0.0 -> v55.0.0 (no change; assume patch)
     Checked [   0.453s] 223 checks: 223 pass, 31 skip
     Summary no semver update required
    Finished [  47.956s] datafusion-physical-expr
    Building datafusion-physical-expr-adapter v55.0.0 (current)
       Built [  25.308s] (current)
     Parsing datafusion-physical-expr-adapter v55.0.0 (current)
      Parsed [   0.009s] (current)
    Building datafusion-physical-expr-adapter v55.0.0 (baseline)
       Built [  25.946s] (baseline)
     Parsing datafusion-physical-expr-adapter v55.0.0 (baseline)
      Parsed [   0.008s] (baseline)
    Checking datafusion-physical-expr-adapter v55.0.0 -> v55.0.0 (no change; assume patch)
     Checked [   0.107s] 223 checks: 223 pass, 31 skip
     Summary no semver update required
    Finished [  52.512s] datafusion-physical-expr-adapter
    Building datafusion-physical-expr-common v55.0.0 (current)
       Built [  19.048s] (current)
     Parsing datafusion-physical-expr-common v55.0.0 (current)
      Parsed [   0.018s] (current)
    Building datafusion-physical-expr-common v55.0.0 (baseline)
       Built [  18.975s] (baseline)
     Parsing datafusion-physical-expr-common v55.0.0 (baseline)
      Parsed [   0.016s] (baseline)
    Checking datafusion-physical-expr-common v55.0.0 -> v55.0.0 (no change; assume patch)
     Checked [   0.287s] 223 checks: 223 pass, 31 skip
     Summary no semver update required
    Finished [  39.669s] datafusion-physical-expr-common
    Building datafusion-physical-plan v55.0.0 (current)
       Built [  29.872s] (current)
     Parsing datafusion-physical-plan v55.0.0 (current)
      Parsed [   0.116s] (current)
    Building datafusion-physical-plan v55.0.0 (baseline)
       Built [  29.635s] (baseline)
     Parsing datafusion-physical-plan v55.0.0 (baseline)
      Parsed [   0.112s] (baseline)
    Checking datafusion-physical-plan v55.0.0 -> v55.0.0 (no change; assume patch)
     Checked [   0.924s] 223 checks: 223 pass, 31 skip
     Summary no semver update required
    Finished [  62.083s] datafusion-physical-plan
    Building datafusion-proto v55.0.0 (current)
       Built [  43.259s] (current)
     Parsing datafusion-proto v55.0.0 (current)
      Parsed [   0.014s] (current)
    Building datafusion-proto v55.0.0 (baseline)
       Built [  41.873s] (baseline)
     Parsing datafusion-proto v55.0.0 (baseline)
      Parsed [   0.020s] (baseline)
    Checking datafusion-proto v55.0.0 -> v55.0.0 (no change; assume patch)
     Checked [   0.221s] 223 checks: 223 pass, 31 skip
     Summary no semver update required
    Finished [  86.814s] datafusion-proto
    Building datafusion-sql v55.0.0 (current)
       Built [  33.339s] (current)
     Parsing datafusion-sql v55.0.0 (current)
      Parsed [   0.024s] (current)
    Building datafusion-sql v55.0.0 (baseline)
       Built [  32.775s] (baseline)
     Parsing datafusion-sql v55.0.0 (baseline)
      Parsed [   0.025s] (baseline)
    Checking datafusion-sql v55.0.0 -> v55.0.0 (no change; assume patch)
     Checked [   0.301s] 223 checks: 223 pass, 31 skip
     Summary no semver update required
    Finished [  68.187s] datafusion-sql
    Building datafusion-sqllogictest v55.0.0 (current)
       Built [  78.088s] (current)
     Parsing datafusion-sqllogictest v55.0.0 (current)
      Parsed [   0.017s] (current)
    Building datafusion-sqllogictest v55.0.0 (baseline)
       Built [  78.246s] (baseline)
     Parsing datafusion-sqllogictest v55.0.0 (baseline)
      Parsed [   0.018s] (baseline)
    Checking datafusion-sqllogictest v55.0.0 -> v55.0.0 (no change; assume patch)
     Checked [   0.118s] 223 checks: 223 pass, 31 skip
     Summary no semver update required
    Finished [ 159.436s] datafusion-sqllogictest
    Building datafusion-substrait v55.0.0 (current)
       Built [ 246.401s] (current)
     Parsing datafusion-substrait v55.0.0 (current)
      Parsed [   0.015s] (current)
    Building datafusion-substrait v55.0.0 (baseline)
       Built [ 254.615s] (baseline)
     Parsing datafusion-substrait v55.0.0 (baseline)
      Parsed [   0.015s] (baseline)
    Checking datafusion-substrait v55.0.0 -> v55.0.0 (no change; assume patch)
     Checked [   0.275s] 223 checks: 223 pass, 31 skip
     Summary no semver update required
    Finished [ 504.514s] datafusion-substrait

[adriangb]

Superseded by the upstream pull request: apache#24857

This copy existed only to run the change through review here before sending it upstream. That is done, so closing this one. The branch is unchanged and still backs the upstream pull request.