Skip to content

Commit 104e098

Browse files
Dandandanclaude
Dandandan
and
claude
committed
Use NDV estimate to pre-allocate hash tables during aggregation
Use column distinct_count statistics (from Parquet metadata or other sources) to pre-size the hash table in GroupValues implementations, avoiding expensive rehashing during aggregation. The capacity hint is bounded by 128K entries to prevent over-allocation. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 5c653be commit 104e098

11 files changed

Lines changed: 84 additions & 31 deletions

File tree

‎datafusion/physical-expr-common/src/binary_map.rs‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -244,9 +244,14 @@ where
244244
V: Debug + PartialEq + Eq + Clone + Copy + Default,
245245
{
246246
pub fn new(output_type: OutputType) -> Self {
247+
Self::with_capacity(output_type, INITIAL_MAP_CAPACITY)
248+
}
249+
250+
pub fn with_capacity(output_type: OutputType, capacity: usize) -> Self {
251+
let capacity = capacity.max(INITIAL_MAP_CAPACITY);
247252
Self {
248253
output_type,
249-
map: hashbrown::hash_table::HashTable::with_capacity(INITIAL_MAP_CAPACITY),
254+
map: hashbrown::hash_table::HashTable::with_capacity(capacity),
250255
map_size: 0,
251256
buffer: BufferBuilder::new(INITIAL_BUFFER_CAPACITY),
252257
offsets: vec![O::default()], // first offset is always 0

‎datafusion/physical-expr-common/src/binary_view_map.rs‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -155,9 +155,14 @@ where
155155
V: Debug + PartialEq + Eq + Clone + Copy + Default,
156156
{
157157
pub fn new(output_type: OutputType) -> Self {
158+
Self::with_capacity(output_type, INITIAL_MAP_CAPACITY)
159+
}
160+
161+
pub fn with_capacity(output_type: OutputType, capacity: usize) -> Self {
162+
let capacity = capacity.max(INITIAL_MAP_CAPACITY);
158163
Self {
159164
output_type,
160-
map: hashbrown::hash_table::HashTable::with_capacity(INITIAL_MAP_CAPACITY),
165+
map: hashbrown::hash_table::HashTable::with_capacity(capacity),
161166
map_size: 0,
162167
views: Vec::new(),
163168
in_progress: Vec::new(),

‎datafusion/physical-plan/src/aggregates/group_values/mod.rs‎

Lines changed: 37 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -134,13 +134,18 @@ pub trait GroupValues: Send {
134134
pub fn new_group_values(
135135
schema: SchemaRef,
136136
group_ordering: &GroupOrdering,
137+
capacity_hint: Option<usize>,
137138
) -> Result<Box<dyn GroupValues>> {
139+
let capacity = capacity_hint.unwrap_or(0);
138140
if schema.fields.len() == 1 {
139141
let d = schema.fields[0].data_type();
140142

141143
macro_rules! downcast_helper {
142144
($t:ty, $d:ident) => {
143-
return Ok(Box::new(GroupValuesPrimitive::<$t>::new($d.clone())))
145+
return Ok(Box::new(GroupValuesPrimitive::<$t>::new(
146+
$d.clone(),
147+
capacity,
148+
)))
144149
};
145150
}
146151

@@ -176,22 +181,40 @@ pub fn new_group_values(
176181
downcast_helper!(Decimal128Type, d);
177182
}
178183
DataType::Utf8 => {
179-
return Ok(Box::new(GroupValuesBytes::<i32>::new(OutputType::Utf8)));
184+
return Ok(Box::new(GroupValuesBytes::<i32>::new(
185+
OutputType::Utf8,
186+
capacity,
187+
)));
180188
}
181189
DataType::LargeUtf8 => {
182-
return Ok(Box::new(GroupValuesBytes::<i64>::new(OutputType::Utf8)));
190+
return Ok(Box::new(GroupValuesBytes::<i64>::new(
191+
OutputType::Utf8,
192+
capacity,
193+
)));
183194
}
184195
DataType::Utf8View => {
185-
return Ok(Box::new(GroupValuesBytesView::new(OutputType::Utf8View)));
196+
return Ok(Box::new(GroupValuesBytesView::new(
197+
OutputType::Utf8View,
198+
capacity,
199+
)));
186200
}
187201
DataType::Binary => {
188-
return Ok(Box::new(GroupValuesBytes::<i32>::new(OutputType::Binary)));
202+
return Ok(Box::new(GroupValuesBytes::<i32>::new(
203+
OutputType::Binary,
204+
capacity,
205+
)));
189206
}
190207
DataType::LargeBinary => {
191-
return Ok(Box::new(GroupValuesBytes::<i64>::new(OutputType::Binary)));
208+
return Ok(Box::new(GroupValuesBytes::<i64>::new(
209+
OutputType::Binary,
210+
capacity,
211+
)));
192212
}
193213
DataType::BinaryView => {
194-
return Ok(Box::new(GroupValuesBytesView::new(OutputType::BinaryView)));
214+
return Ok(Box::new(GroupValuesBytesView::new(
215+
OutputType::BinaryView,
216+
capacity,
217+
)));
195218
}
196219
DataType::Boolean => {
197220
return Ok(Box::new(GroupValuesBoolean::new()));
@@ -202,11 +225,15 @@ pub fn new_group_values(
202225

203226
if multi_group_by::supported_schema(schema.as_ref()) {
204227
if matches!(group_ordering, GroupOrdering::None) {
205-
Ok(Box::new(GroupValuesColumn::<false>::try_new(schema)?))
228+
Ok(Box::new(GroupValuesColumn::<false>::try_new(
229+
schema, capacity,
230+
)?))
206231
} else {
207-
Ok(Box::new(GroupValuesColumn::<true>::try_new(schema)?))
232+
Ok(Box::new(GroupValuesColumn::<true>::try_new(
233+
schema, capacity,
234+
)?))
208235
}
209236
} else {
210-
Ok(Box::new(GroupValuesRows::try_new(schema)?))
237+
Ok(Box::new(GroupValuesRows::try_new(schema, capacity)?))
211238
}
212239
}

‎datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -260,8 +260,8 @@ impl<const STREAMING: bool> GroupValuesColumn<STREAMING> {
260260
// ========================================================================
261261

262262
/// Create a new instance of GroupValuesColumn if supported for the specified schema
263-
pub fn try_new(schema: SchemaRef) -> Result<Self> {
264-
let map = HashTable::with_capacity(0);
263+
pub fn try_new(schema: SchemaRef, capacity: usize) -> Result<Self> {
264+
let map = HashTable::with_capacity(capacity);
265265
Ok(Self {
266266
schema,
267267
map,
@@ -1268,7 +1268,7 @@ mod tests {
12681268
fn test_intern_for_vectorized_group_values() {
12691269
let data_set = VectorizedTestDataSet::new();
12701270
let mut group_values =
1271-
GroupValuesColumn::<false>::try_new(data_set.schema()).unwrap();
1271+
GroupValuesColumn::<false>::try_new(data_set.schema(), 0).unwrap();
12721272

12731273
data_set.load_to_group_values(&mut group_values);
12741274
let actual_batch = group_values.emit(EmitTo::All).unwrap();
@@ -1281,7 +1281,7 @@ mod tests {
12811281
fn test_emit_first_n_for_vectorized_group_values() {
12821282
let data_set = VectorizedTestDataSet::new();
12831283
let mut group_values =
1284-
GroupValuesColumn::<false>::try_new(data_set.schema()).unwrap();
1284+
GroupValuesColumn::<false>::try_new(data_set.schema(), 0).unwrap();
12851285

12861286
// 1~num_rows times to emit the groups
12871287
let num_rows = data_set.expected_batch.num_rows();
@@ -1332,7 +1332,7 @@ mod tests {
13321332

13331333
let field = Field::new_list_field(DataType::Int32, true);
13341334
let schema = Arc::new(Schema::new_with_metadata(vec![field], HashMap::new()));
1335-
let mut group_values = GroupValuesColumn::<false>::try_new(schema).unwrap();
1335+
let mut group_values = GroupValuesColumn::<false>::try_new(schema, 0).unwrap();
13361336

13371337
// Insert group index views and check if success to insert
13381338
insert_inline_group_index_view(&mut group_values, 0, 0);

‎datafusion/physical-plan/src/aggregates/group_values/row.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,7 @@ pub struct GroupValuesRows {
8282
}
8383

8484
impl GroupValuesRows {
85-
pub fn try_new(schema: SchemaRef) -> Result<Self> {
85+
pub fn try_new(schema: SchemaRef, capacity: usize) -> Result<Self> {
8686
// Print a debugging message, so it is clear when the (slower) fallback
8787
// GroupValuesRows is used.
8888
debug!("Creating GroupValuesRows for schema: {schema}");
@@ -94,9 +94,9 @@ impl GroupValuesRows {
9494
.collect(),
9595
)?;
9696

97-
let map = HashTable::with_capacity(0);
97+
let map = HashTable::with_capacity(capacity);
9898

99-
let starting_rows_capacity = 1000;
99+
let starting_rows_capacity = capacity.max(1000);
100100

101101
let starting_data_capacity = 64 * starting_rows_capacity;
102102
let rows_buffer =

‎datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,9 +36,9 @@ pub struct GroupValuesBytes<O: OffsetSizeTrait> {
3636
}
3737

3838
impl<O: OffsetSizeTrait> GroupValuesBytes<O> {
39-
pub fn new(output_type: OutputType) -> Self {
39+
pub fn new(output_type: OutputType, capacity: usize) -> Self {
4040
Self {
41-
map: ArrowBytesMap::new(output_type),
41+
map: ArrowBytesMap::with_capacity(output_type, capacity),
4242
num_groups: 0,
4343
}
4444
}

‎datafusion/physical-plan/src/aggregates/group_values/single_group_by/bytes_view.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,9 +34,9 @@ pub struct GroupValuesBytesView {
3434
}
3535

3636
impl GroupValuesBytesView {
37-
pub fn new(output_type: OutputType) -> Self {
37+
pub fn new(output_type: OutputType, capacity: usize) -> Self {
3838
Self {
39-
map: ArrowBytesViewMap::new(output_type),
39+
map: ArrowBytesViewMap::with_capacity(output_type, capacity),
4040
num_groups: 0,
4141
}
4242
}

‎datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -98,12 +98,13 @@ pub struct GroupValuesPrimitive<T: ArrowPrimitiveType> {
9898
}
9999

100100
impl<T: ArrowPrimitiveType> GroupValuesPrimitive<T> {
101-
pub fn new(data_type: DataType) -> Self {
101+
pub fn new(data_type: DataType, capacity: usize) -> Self {
102102
assert!(PrimitiveArray::<T>::is_compatible(&data_type));
103+
let capacity = capacity.max(128);
103104
Self {
104105
data_type,
105-
map: HashTable::with_capacity(128),
106-
values: Vec::with_capacity(128),
106+
map: HashTable::with_capacity(capacity),
107+
values: Vec::with_capacity(capacity),
107108
null_group: None,
108109
random_state: crate::aggregates::AGGREGATION_HASH_SEED,
109110
}

‎datafusion/physical-plan/src/aggregates/mod.rs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1202,7 +1202,10 @@ impl AggregateExec {
12021202
/// **Grouping sets:** `GROUPING SETS ((a), (b), (a, b))` with NDV(a) = 100, NDV(b) = 50
12031203
/// → set(a) = 100, set(b) = 50, set(a, b) = 100 × 50 = 5,000
12041204
/// → total = 100 + 50 + 5,000 = 5,150
1205-
fn compute_group_ndv(&self, child_statistics: &Statistics) -> Option<usize> {
1205+
pub(crate) fn compute_group_ndv(
1206+
&self,
1207+
child_statistics: &Statistics,
1208+
) -> Option<usize> {
12061209
let mut total: usize = 0;
12071210
for group_mask in &self.group_by.groups {
12081211
let mut set_product: usize = 1;

‎datafusion/physical-plan/src/aggregates/row_hash.rs‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -587,7 +587,18 @@ impl GroupedHashAggregateStream {
587587
_ => OutOfMemoryMode::ReportError,
588588
};
589589

590-
let group_values = new_group_values(group_schema, &group_ordering)?;
590+
// Use NDV estimate from child statistics to pre-allocate hash table,
591+
// bounded by 128K to avoid over-allocation.
592+
const MAX_NDV_CAPACITY: usize = 128 * 1024;
593+
let capacity_hint = agg
594+
.input
595+
.partition_statistics(None)
596+
.ok()
597+
.and_then(|stats| agg.compute_group_ndv(&stats))
598+
.map(|ndv: usize| ndv.min(MAX_NDV_CAPACITY));
599+
600+
let group_values =
601+
new_group_values(group_schema, &group_ordering, capacity_hint)?;
591602
let reservation = MemoryConsumer::new(name)
592603
// We interpret 'can spill' as 'can handle memory back pressure'.
593604
// This value needs to be set to true for the default memory pool implementations
@@ -1269,7 +1280,8 @@ impl GroupedHashAggregateStream {
12691280
.merging_group_by
12701281
.group_schema(&self.spill_state.spill_schema)?;
12711282
if group_schema.fields().len() > 1 {
1272-
self.group_values = new_group_values(group_schema, &self.group_ordering)?;
1283+
self.group_values =
1284+
new_group_values(group_schema, &self.group_ordering, None)?;
12731285
}
12741286

12751287
// Use `OutOfMemoryMode::ReportError` from this point on

0 commit comments

Comments
 (0)