You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/09/06 22:37:56 UTC
(datafusion) branch main updated: fix: stop pre-allocating and undercounting a hash table per COUNT(DISTINCT) group (#24857)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new d25ffaab49 fix: stop pre-allocating and undercounting a hash table per COUNT(DISTINCT) group (#24857)
d25ffaab49 is described below
commit d25ffaab493c18690028410656ed12ea6c6fb0b6
Author: Adrian Garcia Badaracco <17...@users.noreply.github.com>
AuthorDate: Sun Sep 6 22:27:02 2026 +0000
fix: stop pre-allocating and undercounting a hash table per COUNT(DISTINCT) group (#24857)
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.
```sql
-- 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 #24888 and #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 #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.
---------
Co-authored-by: Claude Opus 5 <no...@anthropic.com>
---
datafusion/core/tests/memory_limit/mod.rs | 138 +++++++-
.../src/aggregate/count_distinct/bytes.rs | 155 +++++++++
.../benches/arrow_bytes_map.rs | 12 +-
datafusion/physical-expr-common/src/binary_map.rs | 370 +++++++++++++++++++--
.../physical-expr-common/src/binary_view_map.rs | 184 ++++++++--
.../group_values/single_group_by/bytes.rs | 73 +++-
.../group_values/single_group_by/bytes_view.rs | 72 +++-
7 files changed, 944 insertions(+), 60 deletions(-)
diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs
index 9172e34a87..f873f95601 100644
--- a/datafusion/core/tests/memory_limit/mod.rs
+++ b/datafusion/core/tests/memory_limit/mod.rs
@@ -26,7 +26,10 @@ mod nlj_spill_unmatched;
mod repartition_mem_limit;
mod union_nullable_spill;
mod view_spill_compaction;
-use arrow::array::{ArrayRef, DictionaryArray, Int32Array, RecordBatch, StringViewArray};
+use arrow::array::{
+ ArrayRef, DictionaryArray, Int32Array, Int64Array, RecordBatch, StringArray,
+ StringViewArray,
+};
use arrow::compute::SortOptions;
use arrow::datatypes::{Int32Type, SchemaRef};
use arrow_schema::{DataType, Field, Schema};
@@ -220,6 +223,68 @@ mod count_distinct_spill {
}
}
+/// 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 used to be created pre-allocated, which costs far more than the
+/// handful of values a group typically holds, so the query's memory use
+/// tracked the number of groups rather than the amount of data.
+///
+/// With 4,000 groups holding 2 distinct values each, this query needed about
+/// 35.5 MB of budget before the per group pre-allocation was removed and
+/// about 1.9 MB after, so an 8 MB limit is a failure before the change and a
+/// success after it. Spilling is disabled, so completing means the query
+/// genuinely fit in the budget.
+///
+/// The `avg(payload)` 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. 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: averaging per group averages of
+/// different sizes gives the wrong answer, so the rule can never accept it.
+/// That is why ClickBench Q9 keeps its distinct aggregate. Do not replace
+/// this with `count(*)`: `count` is only incidentally rejected today, and
+/// <https://github.com/apache/datafusion/pull/24859> proposes accepting it,
+/// which would rewrite the query and leave this test passing by construction.
+#[tokio::test]
+async fn group_by_count_distinct_utf8() {
+ TestCase::new()
+ .with_query(
+ "select group_key, count(distinct value), avg(payload) from t group by group_key",
+ )
+ .with_scenario(Scenario::GroupedDistinctStrings {
+ groups: 4_000,
+ string_view: false,
+ })
+ .with_config(SessionConfig::new().with_target_partitions(1))
+ .with_memory_limit(8_000_000)
+ .with_expected_success()
+ .run()
+ .await
+}
+
+/// The `Utf8View` counterpart of [`group_by_count_distinct_utf8`], covering
+/// the separate view flavoured hash set. The same query over a `Utf8View`
+/// column needed about 123 MB of budget before the change and about 2.5 MB
+/// after, so 16 MB separates the two.
+#[tokio::test]
+async fn group_by_count_distinct_utf8_view() {
+ TestCase::new()
+ .with_query(
+ "select group_key, count(distinct value), avg(payload) from t group by group_key",
+ )
+ .with_scenario(Scenario::GroupedDistinctStrings {
+ groups: 4_000,
+ string_view: true,
+ })
+ .with_config(SessionConfig::new().with_target_partitions(1))
+ .with_memory_limit(16_000_000)
+ .with_expected_success()
+ .run()
+ .await
+}
+
#[tokio::test]
async fn join_by_key_multiple_partitions() {
let config = SessionConfig::new().with_target_partitions(2);
@@ -1077,6 +1142,14 @@ enum Scenario {
/// If true, splits all input batches into 1 row each
single_row_batches: bool,
},
+
+ /// `groups` distinct integer keys paired with a short string value, for
+ /// grouped aggregates that build one accumulator per group.
+ GroupedDistinctStrings {
+ groups: usize,
+ /// If true, the value column is `Utf8View` rather than `Utf8`
+ string_view: bool,
+ },
}
impl Scenario {
@@ -1151,6 +1224,15 @@ impl Scenario {
let table = SortedTableProvider::new(batches, sort_information);
Arc::new(table)
}
+ Self::GroupedDistinctStrings {
+ groups,
+ string_view,
+ } => {
+ let batches = grouped_distinct_string_batches(*groups, *string_view);
+ let table =
+ MemTable::try_new(batches[0].schema(), vec![batches]).unwrap();
+ Arc::new(table)
+ }
}
}
@@ -1174,8 +1256,62 @@ impl Scenario {
// Use default rules
None
}
+ Self::GroupedDistinctStrings { .. } => {
+ // Disable the rules that would add a repartition, so the test
+ // measures the aggregate's budget rather than a repartition's
+ Some(vec![Arc::new(JoinSelection::new())])
+ }
+ }
+ }
+}
+
+/// Number of distinct string values held by every group produced by
+/// [`grouped_distinct_string_batches`]
+const DISTINCT_VALUES_PER_GROUP: usize = 2;
+
+/// Returns batches of 1024 rows with `groups` distinct keys in `group_key`,
+/// each key paired with [`DISTINCT_VALUES_PER_GROUP`] distinct short strings
+/// in `value` and an `Int64` `payload` to aggregate over. The values are
+/// `Utf8View` if `string_view` is set, `Utf8` otherwise.
+fn grouped_distinct_string_batches(groups: usize, string_view: bool) -> Vec<RecordBatch> {
+ let value_type = if string_view {
+ DataType::Utf8View
+ } else {
+ DataType::Utf8
+ };
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("group_key", DataType::Int32, false),
+ Field::new("value", value_type, false),
+ Field::new("payload", DataType::Int64, false),
+ ]));
+
+ const ROWS_PER_BATCH: usize = 1024;
+
+ let rows = groups * DISTINCT_VALUES_PER_GROUP;
+ let mut keys = Vec::with_capacity(rows);
+ let mut values = Vec::with_capacity(rows);
+ for value in 0..DISTINCT_VALUES_PER_GROUP {
+ for group in 0..groups {
+ keys.push(group as i32);
+ values.push(format!("value-{value}"));
}
}
+
+ keys.chunks(ROWS_PER_BATCH)
+ .zip(values.chunks(ROWS_PER_BATCH))
+ .map(|(keys, values)| {
+ let payloads: ArrayRef =
+ Arc::new(Int64Array::from_iter_values(keys.iter().map(|k| *k as i64)));
+ let keys: ArrayRef = Arc::new(Int32Array::from(keys.to_vec()));
+ let values: ArrayRef = if string_view {
+ Arc::new(StringViewArray::from_iter_values(values))
+ } else {
+ Arc::new(StringArray::from_iter_values(values))
+ };
+ RecordBatch::try_new(Arc::clone(&schema), vec![keys, values, payloads])
+ .unwrap()
+ })
+ .collect()
}
fn access_log_batches() -> Vec<RecordBatch> {
diff --git a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/bytes.rs b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/bytes.rs
index f6df4182a8..d955d343ad 100644
--- a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/bytes.rs
+++ b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/bytes.rs
@@ -39,6 +39,18 @@ use std::mem::size_of_val;
pub struct BytesDistinctCountAccumulator<O: OffsetSizeTrait>(ArrowBytesSet<O>);
impl<O: OffsetSizeTrait> BytesDistinctCountAccumulator<O> {
+ /// The set deliberately does not pre-allocate. `GroupsAccumulatorAdapter`
+ /// creates one accumulator per group, so a grouped `COUNT(DISTINCT)` over a
+ /// high cardinality key holds hundreds of thousands of these at once and
+ /// most of them see only a handful of values.
+ ///
+ /// The ungrouped path builds one of these and grows it to hold every
+ /// distinct value in the input, so it is the caller that had a use for the
+ /// warm up. It loses nothing here: the set grows into exactly the
+ /// capacities a pre-allocated one reaches, which
+ /// `ungrouped_utf8_accumulator_is_never_worse_than_a_pre_allocated_set`
+ /// pins. That is why this constructor needs no signal distinguishing the
+ /// two callers.
pub fn new(output_type: OutputType) -> Self {
Self(ArrowBytesSet::new(output_type))
}
@@ -100,6 +112,8 @@ impl<O: OffsetSizeTrait> Accumulator for BytesDistinctCountAccumulator<O> {
pub struct BytesViewDistinctCountAccumulator(ArrowBytesViewSet);
impl BytesViewDistinctCountAccumulator {
+ /// See [`BytesDistinctCountAccumulator::new`] for why the set does not
+ /// pre-allocate.
pub fn new(output_type: OutputType) -> Self {
Self(ArrowBytesViewSet::new(output_type))
}
@@ -151,3 +165,144 @@ impl Accumulator for BytesViewDistinctCountAccumulator {
size_of_val(self) + self.0.size()
}
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow::array::{StringArray, StringViewArray};
+ use datafusion_physical_expr_common::binary_map::INITIAL_MAP_CAPACITY;
+ use datafusion_physical_expr_common::binary_view_map::INITIAL_MAP_CAPACITY as INITIAL_VIEW_MAP_CAPACITY;
+ use std::sync::Arc;
+
+ /// The batch size the aggregate stream drives an ungrouped accumulator with.
+ const BATCH_SIZE: usize = 8192;
+
+ /// Distinct value counts spanning both sides of the warm up capacity, up to
+ /// ones where the two constructors have converged.
+ const CARDINALITIES: [usize; 7] = [0, 1, 100, 1_000, 10_000, 100_000, 500_000];
+
+ /// Cardinalities small enough that the warm up dominates what the set
+ /// holds. This is the per group population, where `GroupsAccumulatorAdapter`
+ /// holds one accumulator per group and most see a handful of values.
+ const PER_GROUP_SCALE: usize = 100;
+
+ /// Cardinalities at which a lazily built set has grown into exactly the
+ /// capacities a pre-allocated one reaches. This is the ungrouped
+ /// population, one set holding every distinct value in the input.
+ const UNGROUPED_SCALE: usize = 10_000;
+
+ /// Longer than the map's inline value length, so the value lands in the
+ /// value buffer rather than inside the hash table entry.
+ fn distinct_value(i: usize) -> String {
+ format!("distinct value number {i}")
+ }
+
+ fn batches(distinct_values: usize, view: bool) -> Vec<ArrayRef> {
+ (0..distinct_values)
+ .step_by(BATCH_SIZE)
+ .map(|start| {
+ let values = (start..(start + BATCH_SIZE).min(distinct_values))
+ .map(distinct_value);
+ if view {
+ Arc::new(StringViewArray::from_iter_values(values)) as ArrayRef
+ } else {
+ Arc::new(StringArray::from_iter_values(values)) as ArrayRef
+ }
+ })
+ .collect()
+ }
+
+ /// The property that decides whether the ungrouped path can afford to share
+ /// the lazy constructor with the per group path.
+ fn assert_lazy_is_not_worse(
+ distinct_values: usize,
+ lazy_size: usize,
+ pre_allocated_size: usize,
+ ) {
+ // The guarantee that lets both paths share one lazy constructor.
+ assert!(
+ lazy_size <= pre_allocated_size,
+ "at {distinct_values} distinct values the lazy set reported \
+ {lazy_size} bytes against the {pre_allocated_size} bytes the \
+ pre-allocated one reported"
+ );
+
+ if distinct_values <= PER_GROUP_SCALE {
+ // The warm up is pure overhead here, and removing it is the whole
+ // point of the change.
+ assert!(
+ lazy_size < pre_allocated_size,
+ "at {distinct_values} distinct values the lazy set should be \
+ strictly cheaper, but reported {lazy_size} bytes against \
+ {pre_allocated_size}"
+ );
+ }
+
+ if distinct_values >= UNGROUPED_SCALE {
+ // The hash table's bucket count is a power of two fixed by the
+ // number of entries, and the value buffer grows on a power of two
+ // ladder, so a set that starts empty lands on exactly the
+ // capacities a pre-allocated one reaches. The ungrouped path gives
+ // up nothing by starting empty, which is why these accumulators
+ // need no signal telling them apart from the per group ones.
+ assert_eq!(
+ lazy_size, pre_allocated_size,
+ "at {distinct_values} distinct values the lazy and pre-allocated \
+ sets should have converged"
+ );
+ }
+ }
+
+ /// An ungrouped `COUNT(DISTINCT <string>)` builds a single accumulator that
+ /// grows to hold every distinct value in its input, which is the population
+ /// the warm up existed for. It must be no more expensive without one.
+ #[test]
+ fn ungrouped_utf8_accumulator_is_never_worse_than_a_pre_allocated_set() {
+ for distinct_values in CARDINALITIES {
+ let mut accumulator =
+ BytesDistinctCountAccumulator::<i32>::new(OutputType::Utf8);
+ let mut pre_allocated = ArrowBytesSet::<i32>::with_capacity(
+ OutputType::Utf8,
+ INITIAL_MAP_CAPACITY,
+ );
+
+ for batch in batches(distinct_values, false) {
+ accumulator.update_batch(&[Arc::clone(&batch)]).unwrap();
+ pre_allocated.insert(&batch);
+ }
+
+ assert_eq!(accumulator.0.non_null_len(), distinct_values);
+ assert_lazy_is_not_worse(
+ distinct_values,
+ accumulator.0.size(),
+ pre_allocated.size(),
+ );
+ }
+ }
+
+ /// The `Utf8View` counterpart of
+ /// [`ungrouped_utf8_accumulator_is_never_worse_than_a_pre_allocated_set`].
+ #[test]
+ fn ungrouped_utf8_view_accumulator_is_never_worse_than_a_pre_allocated_set() {
+ for distinct_values in CARDINALITIES {
+ let mut accumulator =
+ BytesViewDistinctCountAccumulator::new(OutputType::Utf8View);
+ let mut pre_allocated = ArrowBytesViewSet::with_capacity(
+ OutputType::Utf8View,
+ INITIAL_VIEW_MAP_CAPACITY,
+ );
+
+ for batch in batches(distinct_values, true) {
+ accumulator.update_batch(&[Arc::clone(&batch)]).unwrap();
+ pre_allocated.insert(&batch);
+ }
+
+ assert_eq!(accumulator.0.non_null_len(), distinct_values);
+ assert_lazy_is_not_worse(
+ distinct_values,
+ accumulator.0.size(),
+ pre_allocated.size(),
+ );
+ }
+ }
+}
diff --git a/datafusion/physical-expr-common/benches/arrow_bytes_map.rs b/datafusion/physical-expr-common/benches/arrow_bytes_map.rs
index 7c8cdc3b4c..68351a8395 100644
--- a/datafusion/physical-expr-common/benches/arrow_bytes_map.rs
+++ b/datafusion/physical-expr-common/benches/arrow_bytes_map.rs
@@ -17,7 +17,9 @@
use arrow::array::{ArrayRef, StringArray};
use criterion::{Criterion, Throughput, criterion_group, criterion_main};
-use datafusion_physical_expr_common::binary_map::{ArrowBytesMap, OutputType};
+use datafusion_physical_expr_common::binary_map::{
+ ArrowBytesMap, INITIAL_MAP_CAPACITY, OutputType,
+};
use std::hint::black_box;
use std::sync::Arc;
@@ -57,7 +59,13 @@ fn bench_arrow_bytes_map(c: &mut Criterion) {
for (name, values) in cases {
group.bench_function(name, |b| {
b.iter(|| {
- let mut map = ArrowBytesMap::<i32, usize>::new(OutputType::Utf8);
+ // The `long_low_cardinality` case is defined by the distinct
+ // values fitting in the pre-allocated buffer, so this benchmark
+ // measures the pre-allocating constructor.
+ let mut map = ArrowBytesMap::<i32, usize>::with_capacity(
+ OutputType::Utf8,
+ INITIAL_MAP_CAPACITY,
+ );
let mut next_payload = 0;
map.insert_if_new(
&values,
diff --git a/datafusion/physical-expr-common/src/binary_map.rs b/datafusion/physical-expr-common/src/binary_map.rs
index 7543e6b297..16e00b6549 100644
--- a/datafusion/physical-expr-common/src/binary_map.rs
+++ b/datafusion/physical-expr-common/src/binary_map.rs
@@ -28,7 +28,7 @@ use arrow::buffer::{Buffer, NullBuffer, OffsetBuffer, ScalarBuffer};
use arrow::datatypes::DataType;
use datafusion_common::hash_utils::RandomState;
use datafusion_common::hash_utils::create_hashes;
-use datafusion_common::utils::proxy::{HashTableAllocExt, VecAllocExt};
+use datafusion_common::utils::proxy::VecAllocExt;
use datafusion_common::{Result, exec_err};
use std::any::type_name;
use std::fmt::Debug;
@@ -55,10 +55,21 @@ pub enum OutputType {
pub struct ArrowBytesSet<O: OffsetSizeTrait>(ArrowBytesMap<O, ()>);
impl<O: OffsetSizeTrait> ArrowBytesSet<O> {
+ /// Creates a set that does not pre-allocate its hash table or value buffer.
+ ///
+ /// See [`ArrowBytesMap::new`] for when to prefer this over
+ /// [`Self::with_capacity`].
pub fn new(output_type: OutputType) -> Self {
Self(ArrowBytesMap::new(output_type))
}
+ /// Creates a set with room for `map_capacity` entries.
+ ///
+ /// See [`ArrowBytesMap::with_capacity`].
+ pub fn with_capacity(output_type: OutputType, map_capacity: usize) -> Self {
+ Self(ArrowBytesMap::with_capacity(output_type, map_capacity))
+ }
+
/// Return the contents of this set and replace it with a new empty
/// set with the same output type
pub fn take(&mut self) -> Self {
@@ -217,8 +228,13 @@ where
output_type: OutputType,
/// Underlying hash set for each distinct value
map: hashbrown::hash_table::HashTable<Entry<O, V>>,
- /// Total size of the map in bytes
- map_size: usize,
+ /// Hash table capacity to re-create the map with in [`Self::take`], so a
+ /// map built with [`Self::with_capacity`] keeps its pre-allocation when it
+ /// is emptied and reused
+ initial_map_capacity: usize,
+ /// Value buffer capacity to re-create the buffer with in [`Self::take`],
+ /// for the same reason as `initial_map_capacity`
+ initial_buffer_capacity: usize,
/// In progress buffer containing all values
buffer: Vec<u8>,
/// Offsets into `buffer` for each distinct value. These offsets as used
@@ -236,20 +252,71 @@ where
null: Option<(V, usize)>,
}
-/// The size, in number of entries, of the initial hash table
-const INITIAL_MAP_CAPACITY: usize = 128;
-/// The initial size, in bytes, of the string data
+/// The size, in number of entries, of the hash table pre-allocated by
+/// [`ArrowBytesMap::with_capacity`]. It is a warm up size for maps that go on
+/// to hold many values, not a bound on what the map can hold.
+pub const INITIAL_MAP_CAPACITY: usize = 128;
+/// The size, in bytes, of the string data buffer pre-allocated by
+/// [`ArrowBytesMap::with_capacity`]
pub const INITIAL_BUFFER_CAPACITY: usize = 8 * 1024;
+
+/// Appends `value` to a map's value buffer, growing the buffer on a power of
+/// two ladder.
+///
+/// `Vec` doubles from wherever its first allocation landed, so a buffer started
+/// empty by [`ArrowBytesMap::new`] and one started at
+/// [`INITIAL_BUFFER_CAPACITY`] by [`ArrowBytesMap::with_capacity`] sit on
+/// different ladders, and can hold the same values at capacities differing by
+/// up to 2x in either direction depending on the value lengths. Rounding every
+/// growth up to a power of two puts both on one ladder, which is what makes a
+/// lazily allocated map never larger than a pre-allocated one holding the same
+/// values. Growth stays geometric, so appending is still amortized constant
+/// time.
+fn push_value_bytes(buffer: &mut Vec<u8>, value: &[u8]) {
+ let required = buffer.len() + value.len();
+ if required > buffer.capacity() {
+ let target = required.checked_next_power_of_two().unwrap_or(required);
+ buffer.reserve_exact(target - buffer.len());
+ }
+ buffer.extend_from_slice(value);
+}
+
impl<O: OffsetSizeTrait, V> ArrowBytesMap<O, V>
where
V: Debug + PartialEq + Eq + Clone + Copy + Default,
{
+ /// Creates a map that does not pre-allocate its hash table or value buffer.
+ ///
+ /// Use this when maps are created in large numbers and most of them stay
+ /// small, such as the per group `COUNT(DISTINCT)` accumulators that
+ /// `GroupsAccumulatorAdapter` creates one of per group. There the
+ /// pre-allocation dwarfs the values the map actually holds.
pub fn new(output_type: OutputType) -> Self {
+ Self::new_inner(output_type, 0, 0)
+ }
+
+ /// Creates a map whose hash table is pre-allocated for `map_capacity`
+ /// entries and whose value buffer is pre-allocated with
+ /// [`INITIAL_BUFFER_CAPACITY`] bytes.
+ ///
+ /// Use this for the few long lived maps that are each expected to hold many
+ /// values, such as the single map backing a `GROUP BY` on one string
+ /// column. The capacities are preserved across [`Self::take`].
+ pub fn with_capacity(output_type: OutputType, map_capacity: usize) -> Self {
+ Self::new_inner(output_type, map_capacity, INITIAL_BUFFER_CAPACITY)
+ }
+
+ fn new_inner(
+ output_type: OutputType,
+ map_capacity: usize,
+ buffer_capacity: usize,
+ ) -> Self {
Self {
output_type,
- map: hashbrown::hash_table::HashTable::with_capacity(INITIAL_MAP_CAPACITY),
- map_size: 0,
- buffer: Vec::with_capacity(INITIAL_BUFFER_CAPACITY),
+ map: hashbrown::hash_table::HashTable::with_capacity(map_capacity),
+ initial_map_capacity: map_capacity,
+ initial_buffer_capacity: buffer_capacity,
+ buffer: Vec::with_capacity(buffer_capacity),
offsets: vec![O::default()], // first offset is always 0
random_state: RandomState::default(),
hashes_buffer: vec![],
@@ -260,11 +327,31 @@ where
/// Return the contents of this map and replace it with a new empty map with
/// the same output type
pub fn take(&mut self) -> Self {
- let mut new_self = Self::new(self.output_type);
+ let mut new_self = Self::new_inner(
+ self.output_type,
+ self.initial_map_capacity,
+ self.initial_buffer_capacity,
+ );
swap(self, &mut new_self);
new_self
}
+ /// Empties this map and releases every allocation it holds, so
+ /// [`Self::size`] drops to approximately zero.
+ ///
+ /// This is the difference from [`Self::take`]: `take` restores the
+ /// capacities the map was configured with so the emptied map stays warm for
+ /// continued use, which is what emitting wants, whereas here the point is
+ /// to hand the memory back, as before spilling or before a downstream sort.
+ /// The configured capacities are remembered, so the next [`Self::take`]
+ /// warms the map up again.
+ pub fn clear_and_release(&mut self) {
+ let mut released = Self::new_inner(self.output_type, 0, 0);
+ released.initial_map_capacity = self.initial_map_capacity;
+ released.initial_buffer_capacity = self.initial_buffer_capacity;
+ *self = released;
+ }
+
/// Inserts each value from `values` into the map, invoking `payload_fn` for
/// each value if *not* already present, deferring the allocation of the
/// payload until it is needed.
@@ -406,7 +493,7 @@ where
// Put the small values into buffer and offsets so it appears
// the output array, but store the actual bytes inline for
// comparison
- self.buffer.extend_from_slice(value);
+ push_value_bytes(&mut self.buffer, value);
self.offsets.push(O::usize_as(self.buffer.len()));
let payload = make_payload_fn(Some(value));
let new_header = Entry {
@@ -415,11 +502,8 @@ where
offset_or_inline: inline,
payload,
};
- self.map.insert_accounted(
- new_header,
- |header| header.hash,
- &mut self.map_size,
- );
+ self.map
+ .insert_unique(hash, new_header, |header| header.hash);
payload
}
}
@@ -447,7 +531,7 @@ where
// appears the output array, and store that offset
// so the bytes can be compared if needed
let offset = self.buffer.len(); // offset of start for data
- self.buffer.extend_from_slice(value);
+ push_value_bytes(&mut self.buffer, value);
self.offsets.push(O::usize_as(self.buffer.len()));
let payload = make_payload_fn(Some(value));
@@ -457,11 +541,8 @@ where
offset_or_inline: offset,
payload,
};
- self.map.insert_accounted(
- new_header,
- |header| header.hash,
- &mut self.map_size,
- );
+ self.map
+ .insert_unique(hash, new_header, |header| header.hash);
payload
}
};
@@ -486,7 +567,8 @@ where
let Self {
output_type,
map: _,
- map_size: _,
+ initial_map_capacity: _,
+ initial_buffer_capacity: _,
offsets,
buffer,
random_state: _,
@@ -591,7 +673,11 @@ where
/// Return the total size, in bytes, of memory used to store the data in
/// this set, not including `self`
pub fn size(&self) -> usize {
- self.map_size
+ // `HashTable::allocation_size` reports the whole hashbrown allocation,
+ // which is the entry array plus the control bytes plus the trailing
+ // group, so it is larger than `capacity() * size_of::<Entry<O, V>>()`.
+ // It is a constant time layout calculation, not a walk of the table.
+ self.map.allocation_size()
+ self.buffer.capacity() * size_of::<u8>()
+ self.offsets.allocated_size()
+ self.hashes_buffer.allocated_size()
@@ -615,7 +701,7 @@ where
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ArrowBytesMap")
.field("map", &"<map>")
- .field("map_size", &self.map_size)
+ .field("map_allocation_size", &self.map.allocation_size())
.field("buffer", &self.buffer)
.field("random_state", &self.random_state)
.field("hashes_buffer", &self.hashes_buffer)
@@ -663,6 +749,240 @@ mod tests {
use arrow::array::{BinaryArray, LargeBinaryArray, StringArray};
use std::collections::HashMap;
+ /// A lower bound on the bytes a hashbrown table holding `entries` entries
+ /// of type `T` must allocate: one entry slot and one control byte each.
+ /// Derived independently of the production accounting so it can bracket it.
+ fn min_table_bytes<T>(entries: usize) -> usize {
+ entries * (size_of::<T>() + 1)
+ }
+
+ #[test]
+ fn map_new_does_not_allocate() {
+ let map = ArrowBytesMap::<i32, ()>::new(OutputType::Utf8);
+
+ assert_eq!(map.map.capacity(), 0);
+ assert_eq!(map.map.allocation_size(), 0);
+ assert_eq!(map.buffer.capacity(), 0);
+ // Only the single leading zero offset is allocated.
+ assert!(map.size() < 128, "expected {} to be tiny", map.size());
+ }
+
+ #[test]
+ fn map_with_capacity_reports_the_real_hash_table_allocation() {
+ let map = ArrowBytesMap::<i32, ()>::with_capacity(
+ OutputType::Utf8,
+ INITIAL_MAP_CAPACITY,
+ );
+
+ assert!(map.map.capacity() >= INITIAL_MAP_CAPACITY);
+ assert_eq!(map.buffer.capacity(), INITIAL_BUFFER_CAPACITY);
+
+ // Before this accounting was corrected the map reported its hash table
+ // as costing zero bytes until the table grew past its pre-allocation.
+ let table_bytes = map.map.allocation_size();
+ let lower_bound = min_table_bytes::<Entry<i32, ()>>(map.map.capacity());
+ assert!(
+ table_bytes >= lower_bound,
+ "expected {table_bytes} to be at least {lower_bound}"
+ );
+ assert!(
+ table_bytes <= 2 * lower_bound + 64,
+ "expected {table_bytes} to be within a small factor of {lower_bound}"
+ );
+ assert!(map.size() >= table_bytes + INITIAL_BUFFER_CAPACITY);
+ }
+
+ #[test]
+ fn take_preserves_the_capacity_the_map_was_built_with() {
+ let mut preallocated = ArrowBytesMap::<i32, ()>::with_capacity(
+ OutputType::Utf8,
+ INITIAL_MAP_CAPACITY,
+ );
+ let capacity = preallocated.map.capacity();
+ preallocated.take();
+ assert_eq!(preallocated.map.capacity(), capacity);
+ assert_eq!(preallocated.buffer.capacity(), INITIAL_BUFFER_CAPACITY);
+
+ let mut lazy = ArrowBytesMap::<i32, ()>::new(OutputType::Utf8);
+ lazy.take();
+ assert_eq!(lazy.map.capacity(), 0);
+ assert_eq!(lazy.buffer.capacity(), 0);
+ }
+
+ #[test]
+ fn clear_and_release_frees_the_preallocation_that_take_keeps() {
+ let mut map = ArrowBytesMap::<i32, ()>::with_capacity(
+ OutputType::Utf8,
+ INITIAL_MAP_CAPACITY,
+ );
+ let values: ArrayRef = Arc::new(StringArray::from_iter_values(
+ (0..1_000).map(|i| format!("distinct value number {i}")),
+ ));
+ map.insert_if_new(&values, |_| (), |_| ());
+
+ let populated_size = map.size();
+ assert!(populated_size > INITIAL_BUFFER_CAPACITY);
+
+ // `take` deliberately keeps the map warm, so it does not release the
+ // configured capacities.
+ map.take();
+ let taken_size = map.size();
+ assert!(
+ taken_size > INITIAL_BUFFER_CAPACITY,
+ "expected take to retain the warm up allocations, got {taken_size}"
+ );
+
+ map.clear_and_release();
+ let released_size = map.size();
+ assert_eq!(map.map.allocation_size(), 0);
+ assert_eq!(map.buffer.capacity(), 0);
+ assert!(
+ released_size < 128,
+ "expected the released map to report approximately zero bytes, got {released_size}"
+ );
+
+ // The configured capacities survive, so the map warms back up when it
+ // is emitted from again.
+ map.take();
+ assert!(map.map.capacity() >= INITIAL_MAP_CAPACITY);
+ assert_eq!(map.buffer.capacity(), INITIAL_BUFFER_CAPACITY);
+ }
+
+ #[test]
+ fn push_value_bytes_ignores_an_empty_value() {
+ let mut buffer = Vec::new();
+ push_value_bytes(&mut buffer, b"");
+ assert_eq!(buffer.capacity(), 0);
+ assert!(buffer.is_empty());
+
+ push_value_bytes(&mut buffer, b"abcd");
+ let capacity = buffer.capacity();
+ push_value_bytes(&mut buffer, b"");
+ assert_eq!(buffer, b"abcd");
+ assert_eq!(buffer.capacity(), capacity);
+ }
+
+ #[test]
+ fn push_value_bytes_grows_only_past_a_power_of_two() {
+ let mut buffer = Vec::new();
+
+ // A first append that is not itself a power of two still lands on the
+ // ladder: a plain `Vec` would have stopped at a capacity of 9.
+ push_value_bytes(&mut buffer, &[1u8; 9]);
+ assert_eq!(buffer.len(), 9);
+ assert_eq!(buffer.capacity(), 16);
+
+ // Filling exactly to the capacity does not grow it.
+ push_value_bytes(&mut buffer, &[2u8; 7]);
+ assert_eq!(buffer.len(), 16);
+ assert_eq!(buffer.capacity(), 16);
+
+ // The next byte moves up one rung.
+ push_value_bytes(&mut buffer, &[3u8; 1]);
+ assert_eq!(buffer.len(), 17);
+ assert_eq!(buffer.capacity(), 32);
+
+ // So does an append that overshoots several rungs at once.
+ push_value_bytes(&mut buffer, &[4u8; 100]);
+ assert_eq!(buffer.len(), 117);
+ assert_eq!(buffer.capacity(), 128);
+
+ let mut expected = vec![1u8; 9];
+ expected.extend_from_slice(&[2u8; 7]);
+ expected.extend_from_slice(&[3u8; 1]);
+ expected.extend_from_slice(&[4u8; 100]);
+ assert_eq!(buffer, expected);
+ }
+
+ #[test]
+ fn push_value_bytes_does_not_grow_a_buffer_with_room() {
+ let mut buffer = Vec::with_capacity(INITIAL_BUFFER_CAPACITY);
+ let start = buffer.as_ptr();
+
+ let value = vec![7u8; 512];
+ for _ in 0..16 {
+ push_value_bytes(&mut buffer, &value);
+ }
+
+ assert_eq!(buffer.len(), INITIAL_BUFFER_CAPACITY);
+ assert_eq!(buffer.capacity(), INITIAL_BUFFER_CAPACITY);
+ assert!(
+ std::ptr::eq(buffer.as_ptr(), start),
+ "buffer was reallocated"
+ );
+ assert!(buffer.iter().all(|&byte| byte == 7));
+ }
+
+ #[test]
+ fn push_value_bytes_rounds_a_single_oversized_append_up_to_a_power_of_two() {
+ let value = vec![9u8; INITIAL_BUFFER_CAPACITY + 1];
+ let rounded = (INITIAL_BUFFER_CAPACITY + 1).next_power_of_two();
+
+ // A value larger than the pre-allocation lands on the same rung
+ // whether the buffer started empty or pre-allocated, which is what
+ // keeps a lazily allocated map from outgrowing a pre-allocated one.
+ let mut lazy = Vec::new();
+ push_value_bytes(&mut lazy, &value);
+ assert_eq!(lazy.len(), value.len());
+ assert_eq!(lazy.capacity(), rounded);
+ assert_eq!(lazy, value);
+
+ let mut pre_allocated = Vec::with_capacity(INITIAL_BUFFER_CAPACITY);
+ push_value_bytes(&mut pre_allocated, &value);
+ assert_eq!(pre_allocated.capacity(), rounded);
+ assert_eq!(pre_allocated, value);
+ }
+
+ #[test]
+ fn lazy_and_pre_allocated_buffers_grow_on_the_same_ladder() {
+ // Value lengths chosen so the buffer requirement lands between powers
+ // of two, which is where the two ladders used to diverge.
+ for value_len in [9usize, 13, 24, 37] {
+ let mut lazy = ArrowBytesMap::<i32, ()>::new(OutputType::Utf8);
+ let mut pre_allocated = ArrowBytesMap::<i32, ()>::with_capacity(
+ OutputType::Utf8,
+ INITIAL_MAP_CAPACITY,
+ );
+
+ for batch in 0..8 {
+ let values: ArrayRef =
+ Arc::new(StringArray::from_iter_values((0..1_000).map(|i| {
+ let value = format!("{}:{i}", batch * 1_000 + i);
+ format!("{value:value_len$}")
+ })));
+ lazy.insert_if_new(&values, |_| (), |_| ());
+ pre_allocated.insert_if_new(&values, |_| (), |_| ());
+
+ assert_eq!(lazy.buffer.len(), pre_allocated.buffer.len());
+ assert_eq!(
+ lazy.buffer.capacity(),
+ pre_allocated.buffer.capacity(),
+ "value length {value_len}, batch {batch}: a buffer that \
+ started empty reached {} bytes of capacity against {} for \
+ one that started at INITIAL_BUFFER_CAPACITY",
+ lazy.buffer.capacity(),
+ pre_allocated.buffer.capacity(),
+ );
+ }
+ }
+ }
+
+ #[test]
+ fn a_lazy_buffer_stays_below_the_pre_allocated_floor_while_it_is_small() {
+ let mut lazy = ArrowBytesMap::<i32, ()>::new(OutputType::Utf8);
+ let values: ArrayRef = Arc::new(StringArray::from_iter_values(
+ (0..10).map(|i| format!("distinct value number {i}")),
+ ));
+ lazy.insert_if_new(&values, |_| (), |_| ());
+
+ assert!(
+ lazy.buffer.capacity() < INITIAL_BUFFER_CAPACITY,
+ "expected a small lazy buffer to stay under the {INITIAL_BUFFER_CAPACITY} \
+ byte pre-allocation, got {}",
+ lazy.buffer.capacity()
+ );
+ }
+
#[test]
fn string_set_empty() {
let mut set = ArrowBytesSet::<i32>::new(OutputType::Utf8);
diff --git a/datafusion/physical-expr-common/src/binary_view_map.rs b/datafusion/physical-expr-common/src/binary_view_map.rs
index 0457825dec..29f4014c5f 100644
--- a/datafusion/physical-expr-common/src/binary_view_map.rs
+++ b/datafusion/physical-expr-common/src/binary_view_map.rs
@@ -28,10 +28,9 @@ use arrow::buffer::{Buffer, ScalarBuffer};
use arrow::datatypes::{BinaryViewType, ByteViewType, DataType, StringViewType};
use datafusion_common::hash_utils::RandomState;
use datafusion_common::hash_utils::create_hashes;
-use datafusion_common::utils::proxy::{HashTableAllocExt, VecAllocExt};
+use datafusion_common::utils::proxy::VecAllocExt;
use datafusion_common::{Result, exec_err};
use std::fmt::Debug;
-use std::mem::size_of;
use std::sync::Arc;
/// HashSet optimized for storing string or binary values that can produce that
@@ -40,10 +39,21 @@ use std::sync::Arc;
pub struct ArrowBytesViewSet(ArrowBytesViewMap<()>);
impl ArrowBytesViewSet {
+ /// Creates a set that does not pre-allocate its hash table.
+ ///
+ /// See [`ArrowBytesViewMap::new`] for when to prefer this over
+ /// [`Self::with_capacity`].
pub fn new(output_type: OutputType) -> Self {
Self(ArrowBytesViewMap::new(output_type))
}
+ /// Creates a set with room for `map_capacity` entries.
+ ///
+ /// See [`ArrowBytesViewMap::with_capacity`].
+ pub fn with_capacity(output_type: OutputType, map_capacity: usize) -> Self {
+ Self(ArrowBytesViewMap::with_capacity(output_type, map_capacity))
+ }
+
/// Inserts each value from `values` into the set
pub fn insert(&mut self, values: &ArrayRef) {
fn make_payload_fn(_value: Option<&[u8]>) {}
@@ -55,9 +65,7 @@ impl ArrowBytesViewSet {
/// Return the contents of this map and replace it with a new empty map with
/// the same output type
pub fn take(&mut self) -> Self {
- let mut new_self = Self::new(self.0.output_type);
- std::mem::swap(self, &mut new_self);
- new_self
+ Self(self.0.take())
}
/// Converts this set into a `StringViewArray` or `BinaryViewArray`
@@ -129,8 +137,10 @@ where
output_type: OutputType,
/// Underlying hash set for each distinct value
map: hashbrown::hash_table::HashTable<Entry<V>>,
- /// Total size of the map in bytes
- map_size: usize,
+ /// Hash table capacity to re-create the map with in [`Self::take`], so a
+ /// map built with [`Self::with_capacity`] keeps its pre-allocation when it
+ /// is emptied and reused
+ initial_map_capacity: usize,
/// Views for all stored values (in insertion order)
views: Vec<u128>,
@@ -151,21 +161,36 @@ where
null: Option<(V, usize)>,
}
-/// The size, in number of entries, of the initial hash table
-const INITIAL_MAP_CAPACITY: usize = 512;
+/// The size, in number of entries, of the hash table pre-allocated by
+/// [`ArrowBytesViewMap::with_capacity`]. It is a warm up size for maps that go
+/// on to hold many values, not a bound on what the map can hold.
+pub const INITIAL_MAP_CAPACITY: usize = 512;
impl<V> ArrowBytesViewMap<V>
where
V: Debug + PartialEq + Eq + Clone + Copy + Default,
{
+ /// Creates a map that does not pre-allocate its hash table.
+ ///
+ /// Use this when maps are created in large numbers and most of them stay
+ /// small, such as the per group `COUNT(DISTINCT)` accumulators that
+ /// `GroupsAccumulatorAdapter` creates one of per group. There the
+ /// pre-allocation dwarfs the values the map actually holds.
pub fn new(output_type: OutputType) -> Self {
- let map = hashbrown::hash_table::HashTable::with_capacity(INITIAL_MAP_CAPACITY);
- let map_size = map.capacity() * size_of::<Entry<V>>();
+ Self::with_capacity(output_type, 0)
+ }
+ /// Creates a map whose hash table is pre-allocated for `map_capacity`
+ /// entries.
+ ///
+ /// Use this for the few long lived maps that are each expected to hold many
+ /// values, such as the single map backing a `GROUP BY` on one string
+ /// column. The capacity is preserved across [`Self::take`].
+ pub fn with_capacity(output_type: OutputType, map_capacity: usize) -> Self {
Self {
output_type,
- map,
- map_size,
+ map: hashbrown::hash_table::HashTable::with_capacity(map_capacity),
+ initial_map_capacity: map_capacity,
views: Vec::new(),
in_progress: Vec::new(),
completed: Vec::new(),
@@ -179,11 +204,27 @@ where
/// Return the contents of this map and replace it with a new empty map with
/// the same output type
pub fn take(&mut self) -> Self {
- let mut new_self = Self::new(self.output_type);
+ let mut new_self =
+ Self::with_capacity(self.output_type, self.initial_map_capacity);
std::mem::swap(self, &mut new_self);
new_self
}
+ /// Empties this map and releases every allocation it holds, so
+ /// [`Self::size`] drops to approximately zero.
+ ///
+ /// This is the difference from [`Self::take`]: `take` restores the capacity
+ /// the map was configured with so the emptied map stays warm for continued
+ /// use, which is what emitting wants, whereas here the point is to hand the
+ /// memory back, as before spilling or before a downstream sort. The
+ /// configured capacity is remembered, so the next [`Self::take`] warms the
+ /// map up again.
+ pub fn clear_and_release(&mut self) {
+ let mut released = Self::with_capacity(self.output_type, 0);
+ released.initial_map_capacity = self.initial_map_capacity;
+ *self = released;
+ }
+
/// Inserts each value from `values` into the map, invoking `payload_fn` for
/// each value if *not* already present, deferring the allocation of the
/// payload until it is needed.
@@ -374,8 +415,7 @@ where
payload,
};
- self.map
- .insert_accounted(new_header, |h| h.hash, &mut self.map_size);
+ self.map.insert_unique(hash, new_header, |h| h.hash);
payload
};
observe_payload_fn(payload);
@@ -540,13 +580,19 @@ where
pub fn size(&self) -> usize {
// All fields below own their allocations. Count retained capacity rather
// than used length because this value drives memory accounting.
+ //
+ // `HashTable::allocation_size` reports the whole hashbrown allocation,
+ // which is the entry array plus the control bytes plus the trailing
+ // group, so it is larger than `capacity() * size_of::<Entry<V>>()`. It
+ // is a constant time layout calculation, not a walk of the table.
+ let map_size = self.map.allocation_size();
let views_size = self.views.allocated_size();
let in_progress_size = self.in_progress.allocated_size();
let completed_size = self.completed.allocated_size()
+ self.completed.iter().map(Buffer::capacity).sum::<usize>();
let nulls_size = self.nulls.allocated_size();
- self.map_size
+ map_size
+ views_size
+ in_progress_size
+ completed_size
@@ -562,7 +608,7 @@ where
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ArrowBytesMap")
.field("map", &"<map>")
- .field("map_size", &self.map_size)
+ .field("map_allocation_size", &self.map.allocation_size())
.field("views_len", &self.views.len())
.field("completed_buffers", &self.completed.len())
.field("random_state", &self.random_state)
@@ -597,6 +643,7 @@ where
mod tests {
use arrow::array::{GenericByteViewArray, StringViewArray};
use datafusion_common::HashMap;
+ use std::mem::size_of;
use super::*;
@@ -785,11 +832,104 @@ mod tests {
assert_eq!(set.len(), 10);
}
+ /// A lower bound on the bytes a hashbrown table holding `entries` entries
+ /// of type `T` must allocate: one entry slot and one control byte each.
+ /// Derived independently of the production accounting so it can bracket it.
+ fn min_table_bytes<T>(entries: usize) -> usize {
+ entries * (size_of::<T>() + 1)
+ }
+
#[test]
- fn test_size_counts_initial_hash_table_capacity() {
+ fn map_new_does_not_allocate() {
let map = ArrowBytesViewMap::<()>::new(OutputType::Utf8View);
- assert_eq!(map.size(), map.map.capacity() * size_of::<Entry<()>>());
+ assert_eq!(map.map.capacity(), 0);
+ assert_eq!(map.map.allocation_size(), 0);
+ assert_eq!(map.size(), 0);
+ }
+
+ #[test]
+ fn test_size_counts_initial_hash_table_capacity() {
+ let map = ArrowBytesViewMap::<()>::with_capacity(
+ OutputType::Utf8View,
+ INITIAL_MAP_CAPACITY,
+ );
+
+ assert!(map.map.capacity() >= INITIAL_MAP_CAPACITY);
+ assert_eq!(map.size(), map.map.allocation_size());
+
+ // Before this accounting was corrected the map reported exactly
+ // `capacity() * size_of::<Entry<V>>()`, which leaves out the control
+ // bytes and undercounts the real allocation by roughly half.
+ let lower_bound = min_table_bytes::<Entry<()>>(map.map.capacity());
+ assert!(
+ map.size() >= lower_bound,
+ "expected {} to be at least {lower_bound}",
+ map.size()
+ );
+ assert!(
+ map.size() <= 2 * lower_bound + 64,
+ "expected {} to be within a small factor of {lower_bound}",
+ map.size()
+ );
+ assert!(
+ map.size() > map.map.capacity() * size_of::<Entry<()>>(),
+ "expected {} to exceed {}",
+ map.size(),
+ map.map.capacity() * size_of::<Entry<()>>()
+ );
+ }
+
+ #[test]
+ fn take_preserves_the_capacity_the_map_was_built_with() {
+ let mut preallocated = ArrowBytesViewMap::<()>::with_capacity(
+ OutputType::Utf8View,
+ INITIAL_MAP_CAPACITY,
+ );
+ let capacity = preallocated.map.capacity();
+ preallocated.take();
+ assert_eq!(preallocated.map.capacity(), capacity);
+
+ let mut lazy = ArrowBytesViewMap::<()>::new(OutputType::Utf8View);
+ lazy.take();
+ assert_eq!(lazy.map.capacity(), 0);
+ }
+
+ #[test]
+ fn clear_and_release_frees_the_preallocation_that_take_keeps() {
+ let mut map = ArrowBytesViewMap::<()>::with_capacity(
+ OutputType::Utf8View,
+ INITIAL_MAP_CAPACITY,
+ );
+ let values: ArrayRef = Arc::new(StringViewArray::from_iter_values(
+ (0..1_000).map(|i| format!("distinct value number {i}")),
+ ));
+ map.insert_if_new(&values, |_| (), |_| ());
+
+ let warm_size = map.map.allocation_size();
+ assert!(warm_size > 0);
+
+ // `take` deliberately keeps the map warm, so it does not release the
+ // configured capacity.
+ map.take();
+ let taken_size = map.size();
+ assert!(
+ taken_size >= min_table_bytes::<Entry<()>>(INITIAL_MAP_CAPACITY),
+ "expected take to retain the warm up allocation, got {taken_size}"
+ );
+
+ map.clear_and_release();
+ assert_eq!(map.map.allocation_size(), 0);
+ let released_size = map.size();
+ assert!(
+ released_size < 128,
+ "expected the released map to report approximately zero bytes, got {released_size}"
+ );
+
+ // The configured capacity survives, so the map warms back up when it is
+ // emitted from again.
+ map.take();
+ assert!(map.map.capacity() >= INITIAL_MAP_CAPACITY);
}
#[test]
@@ -822,7 +962,7 @@ mod tests {
.any(|buffer| buffer.capacity() > buffer.len())
);
- let expected_size = map.map_size
+ let expected_size = map.map.allocation_size()
+ map.views.allocated_size()
+ map.in_progress.allocated_size()
+ map.completed.allocated_size()
@@ -832,7 +972,7 @@ mod tests {
assert_eq!(map.size(), expected_size);
// Verify the retained-capacity delta independently from the production formula.
- let legacy_size = map.map_size
+ let legacy_size = map.map.allocation_size()
+ map.views.len() * size_of::<u128>()
+ map.in_progress.capacity()
+ map.completed.iter().map(Buffer::len).sum::<usize>()
diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs
index 34ec36be31..9f7b4b4e91 100644
--- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs
+++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs
@@ -22,7 +22,9 @@ use crate::aggregates::group_values::GroupValues;
use arrow::array::{Array, ArrayRef, OffsetSizeTrait};
use datafusion_common::Result;
use datafusion_expr::{EmitTo, GroupSelection};
-use datafusion_physical_expr_common::binary_map::{ArrowBytesMap, OutputType};
+use datafusion_physical_expr_common::binary_map::{
+ ArrowBytesMap, INITIAL_MAP_CAPACITY, OutputType,
+};
/// A [`GroupValues`] storing single column of Utf8/LargeUtf8/Binary/LargeBinary values
///
@@ -38,7 +40,9 @@ pub struct GroupValuesBytes<O: OffsetSizeTrait> {
impl<O: OffsetSizeTrait> GroupValuesBytes<O> {
pub fn new(output_type: OutputType) -> Self {
Self {
- map: ArrowBytesMap::new(output_type),
+ // One map holds every group value for the whole query, so it is
+ // worth pre-allocating the hash table and the value buffer.
+ map: ArrowBytesMap::with_capacity(output_type, INITIAL_MAP_CAPACITY),
num_groups: 0,
}
}
@@ -133,8 +137,67 @@ impl<O: OffsetSizeTrait> GroupValues for GroupValuesBytes<O> {
}
fn clear_shrink(&mut self, _num_rows: usize) {
- // in theory we could potentially avoid this reallocation and clear the
- // contents of the maps, but for now we just reset the map from the beginning
- self.map.take();
+ // Callers use this to hand memory back before spilling or sorting, so
+ // release the map's allocations rather than restoring the warm up
+ // capacities that `take` keeps for the emit path.
+ self.map.clear_and_release();
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ use std::sync::Arc;
+
+ use arrow::array::StringArray;
+ use datafusion_physical_expr_common::binary_map::INITIAL_BUFFER_CAPACITY;
+
+ /// `clear_shrink` is how the aggregate stream hands memory back before it
+ /// spills and before the spilled batch is sorted, so the memory it releases
+ /// has to actually show up in the size it reports afterwards.
+ #[test]
+ fn clear_shrink_releases_the_reported_memory() {
+ let mut group_values = GroupValuesBytes::<i32>::new(OutputType::Utf8);
+ let empty = size_of::<GroupValuesBytes<i32>>();
+
+ // The map is pre-allocated at construction, so it is already well above
+ // its own struct size before a single row is interned.
+ let warm_size = group_values.size();
+ assert!(
+ warm_size > empty + INITIAL_BUFFER_CAPACITY,
+ "expected the pre-allocated map to report more than {} bytes, got {warm_size}",
+ empty + INITIAL_BUFFER_CAPACITY
+ );
+
+ let values: ArrayRef = Arc::new(StringArray::from_iter_values(
+ (0..1_000).map(|i| format!("group value number {i}")),
+ ));
+ let mut groups = vec![];
+ group_values
+ .intern(&[Arc::clone(&values)], &mut groups)
+ .unwrap();
+ let populated_size = group_values.size();
+ assert!(populated_size > warm_size);
+
+ group_values.clear_shrink(0);
+
+ // Everything the map held is gone: what remains is the struct itself
+ // plus the single leading zero offset.
+ let released_size = group_values.size();
+ assert!(
+ released_size < empty + 128,
+ "expected clear_shrink to release the map, got {released_size} with a struct size of {empty}"
+ );
+ assert!(
+ released_size * 10 < populated_size,
+ "expected {released_size} to be far below {populated_size}"
+ );
+
+ // The map still works, and warms back up on the next emit.
+ group_values.intern(&[values], &mut groups).unwrap();
+ assert!(group_values.size() > released_size);
+ group_values.emit(EmitTo::All).unwrap();
+ assert!(group_values.size() > empty + INITIAL_BUFFER_CAPACITY);
}
}
diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs
index 997a7ce166..23ea4e7ed3 100644
--- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs
+++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs
@@ -19,7 +19,9 @@ use crate::aggregates::group_values::GroupValues;
use arrow::array::{Array, ArrayRef};
use datafusion_expr::{EmitTo, GroupSelection};
use datafusion_physical_expr::binary_map::OutputType;
-use datafusion_physical_expr_common::binary_view_map::ArrowBytesViewMap;
+use datafusion_physical_expr_common::binary_view_map::{
+ ArrowBytesViewMap, INITIAL_MAP_CAPACITY,
+};
use std::mem::size_of;
/// A [`GroupValues`] storing single column of Utf8View/BinaryView values
@@ -36,7 +38,9 @@ pub struct GroupValuesBytesView {
impl GroupValuesBytesView {
pub fn new(output_type: OutputType) -> Self {
Self {
- map: ArrowBytesViewMap::new(output_type),
+ // One map holds every group value for the whole query, so it is
+ // worth pre-allocating the hash table.
+ map: ArrowBytesViewMap::with_capacity(output_type, INITIAL_MAP_CAPACITY),
num_groups: 0,
}
}
@@ -135,8 +139,66 @@ impl GroupValues for GroupValuesBytesView {
}
fn clear_shrink(&mut self, _num_rows: usize) {
- // in theory we could potentially avoid this reallocation and clear the
- // contents of the maps, but for now we just reset the map from the beginning
- self.map.take();
+ // Callers use this to hand memory back before spilling or sorting, so
+ // release the map's allocations rather than restoring the warm up
+ // capacity that `take` keeps for the emit path.
+ self.map.clear_and_release();
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ use std::sync::Arc;
+
+ use arrow::array::StringViewArray;
+
+ /// `clear_shrink` is how the aggregate stream hands memory back before it
+ /// spills and before the spilled batch is sorted, so the memory it releases
+ /// has to actually show up in the size it reports afterwards.
+ #[test]
+ fn clear_shrink_releases_the_reported_memory() {
+ let mut group_values = GroupValuesBytesView::new(OutputType::Utf8View);
+ let empty = size_of::<GroupValuesBytesView>();
+
+ // The hash table is pre-allocated at construction, so the map is
+ // already well above its own struct size before a single row is
+ // interned.
+ let warm_size = group_values.size();
+ assert!(
+ warm_size > empty + INITIAL_MAP_CAPACITY,
+ "expected the pre-allocated map to report more than {} bytes, got {warm_size}",
+ empty + INITIAL_MAP_CAPACITY
+ );
+
+ let values: ArrayRef = Arc::new(StringViewArray::from_iter_values(
+ (0..1_000).map(|i| format!("group value number {i}")),
+ ));
+ let mut groups = vec![];
+ group_values
+ .intern(&[Arc::clone(&values)], &mut groups)
+ .unwrap();
+ let populated_size = group_values.size();
+ assert!(populated_size > warm_size);
+
+ group_values.clear_shrink(0);
+
+ // Everything the map held is gone: what remains is the struct itself.
+ let released_size = group_values.size();
+ assert!(
+ released_size < empty + 128,
+ "expected clear_shrink to release the map, got {released_size} with a struct size of {empty}"
+ );
+ assert!(
+ released_size * 10 < populated_size,
+ "expected {released_size} to be far below {populated_size}"
+ );
+
+ // The map still works, and warms back up on the next emit.
+ group_values.intern(&[values], &mut groups).unwrap();
+ assert!(group_values.size() > released_size);
+ group_values.emit(EmitTo::All).unwrap();
+ assert!(group_values.size() > empty + INITIAL_MAP_CAPACITY);
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org