Skip to content

Commit 4d167a1

Browse files
mohitgurav20kosiew
andauthored
fix(physical-plan): replace approximate map_size accounting with exact HashTable::allocation_size() in GroupValuesRows (#25750)
## Which issue does this PR close? Closes #25735 Related to #25188 ## Rationale for this change **Background** `GroupValuesRows` is the general-purpose fallback implementation used for grouping complex nested types (like `Struct`, `List`, `Map`) and multi-column schemas. Its `size()` method is highly critical because it tells DataFusion's memory pool exactly how much memory is being consumed, which dictates when the system should safely spill to disk to prevent Out of Memory (OOM) errors. **The Problem** Previously, the memory size of the hash table was poorly estimated. It naively calculated the size as `capacity * size_of::<(u64, usize)>()`. This completely ignored the actual internal memory overhead of the hash table (like control bytes and entry slots), meaning DataFusion was under-reporting its memory usage. ## What changes are included in this PR? This PR completely removes the inaccurate manual `self.map_size` calculation. Instead, it directly calls `self.map.allocation_size()` inside the `size()` function. This ensures that DataFusion's memory pool accurately accounts for every single byte of overhead used by the underlying hash table. ## Are these changes tested? Yes! I have added two comprehensive tests to verify the exact memory accounting: - `test_exact_hash_table_allocation_accounting_multi_column`: Validates precise memory tracking during inserts, partial emits, full emits, and table shrink operations for multi-column groupings. - `test_exact_hash_table_allocation_accounting_nested_keys`: Validates the exact tracking for nested types (Lists) to ensure the row converter and map sizes are perfectly aligned. ## Are there any user-facing changes? No API changes. This is strictly an internal reliability fix to improve the accuracy of query execution memory tracking and prevent potential OOM crashes. Co-authored-by: kosiew <[email protected]>
1 parent e6dfef5 commit 4d167a1

1 file changed

Lines changed: 180 additions & 11 deletions

File tree

  • datafusion/physical-plan/src/aggregates/group_values

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

Lines changed: 180 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -28,11 +28,10 @@ use datafusion_common::Result;
2828
use datafusion_common::hash_utils::RandomState;
2929
use datafusion_common::hash_utils::create_hashes;
3030
use datafusion_common::utils::normalize_float_zero;
31-
use datafusion_execution::memory_pool::proxy::{HashTableAllocExt, VecAllocExt};
31+
use datafusion_execution::memory_pool::proxy::VecAllocExt;
3232
use datafusion_expr::{EmitTo, GroupSelection};
3333
use hashbrown::hash_table::HashTable;
3434
use log::debug;
35-
use std::mem::size_of;
3635
use std::sync::Arc;
3736

3837
/// A [`GroupValues`] making use of [`Rows`]
@@ -60,9 +59,6 @@ pub struct GroupValuesRows {
6059
/// values: (hash, group_index)
6160
map: HashTable<(u64, usize)>,
6261

63-
/// The size of `map` in bytes
64-
map_size: usize,
65-
6662
/// The actual group by values, stored in arrow [`Row`] format.
6763
/// `group_values[i]` holds the group value for group_index `i`.
6864
///
@@ -107,7 +103,6 @@ impl GroupValuesRows {
107103
schema,
108104
row_converter,
109105
map,
110-
map_size: 0,
111106
group_values: None,
112107
hashes_buffer: Default::default(),
113108
rows_buffer,
@@ -168,10 +163,10 @@ impl GroupValues for GroupValuesRows {
168163
group_values.push(group_rows.row(row));
169164

170165
// for hasher function, use precomputed hash value
171-
self.map.insert_accounted(
166+
self.map.insert_unique(
167+
target_hash,
172168
(target_hash, group_idx),
173169
|(hash, _group_index)| *hash,
174-
&mut self.map_size,
175170
);
176171
group_idx
177172
}
@@ -186,9 +181,11 @@ impl GroupValues for GroupValuesRows {
186181

187182
fn size(&self) -> usize {
188183
let group_values_size = self.group_values.as_ref().map(|v| v.size()).unwrap_or(0);
184+
// `HashTable::allocation_size` reports the complete hashbrown allocation,
185+
// which includes the entry array, control bytes, and trailing group layout.
189186
self.row_converter.size()
190187
+ group_values_size
191-
+ self.map_size
188+
+ self.map.allocation_size()
192189
+ self.rows_buffer.size()
193190
+ self.hashes_buffer.allocated_size()
194191
}
@@ -289,7 +286,6 @@ impl GroupValues for GroupValuesRows {
289286
});
290287
self.map.clear();
291288
self.map.shrink_to(num_rows, |_| 0); // hasher does not matter since the map is cleared
292-
self.map_size = self.map.capacity() * size_of::<(u64, usize)>();
293289
self.hashes_buffer.clear();
294290
self.hashes_buffer.shrink_to(num_rows);
295291
}
@@ -443,8 +439,9 @@ pub(crate) fn encode_array_if_necessary(
443439
#[cfg(test)]
444440
mod tests {
445441
use super::*;
446-
use arrow::array::{AsArray, ListArray};
442+
use arrow::array::{AsArray, Int32Array, ListArray, StringArray};
447443
use arrow::datatypes::{Field, Int32Type, Schema};
444+
use std::mem::size_of;
448445

449446
#[test]
450447
fn preserving_nested_row_values() -> Result<()> {
@@ -484,4 +481,176 @@ mod tests {
484481
assert_eq!(groups, vec![3]);
485482
Ok(())
486483
}
484+
485+
#[test]
486+
fn test_exact_hash_table_allocation_accounting_multi_column() -> Result<()> {
487+
let schema = Arc::new(Schema::new(vec![
488+
Field::new("a", DataType::Int32, true),
489+
Field::new("b", DataType::Utf8, true),
490+
]));
491+
let mut group_values = GroupValuesRows::try_new(schema)?;
492+
493+
// 1. Initial state: map is unallocated
494+
assert_eq!(group_values.map.capacity(), 0);
495+
assert_eq!(group_values.map.allocation_size(), 0);
496+
assert_eq!(
497+
group_values.size(),
498+
group_values.row_converter.size()
499+
+ group_values.rows_buffer.size()
500+
+ group_values.hashes_buffer.allocated_size()
501+
);
502+
503+
// 2. Insert multi-column keys and observe table allocation growth
504+
let num_distinct = 500;
505+
let col_a = Arc::new(Int32Array::from_iter_values(0..num_distinct)) as ArrayRef;
506+
let col_b = Arc::new(StringArray::from_iter_values(
507+
(0..num_distinct).map(|i| format!("key_string_{i}")),
508+
)) as ArrayRef;
509+
510+
let mut groups = vec![];
511+
group_values.intern(&[col_a, col_b], &mut groups)?;
512+
assert_eq!(groups.len(), num_distinct as usize);
513+
514+
let table_bytes = group_values.map.allocation_size();
515+
assert!(group_values.map.capacity() >= num_distinct as usize);
516+
assert!(table_bytes > 0);
517+
518+
// Hashbrown's layout contains entry slots, control bytes, and trailing group bytes.
519+
// It must strictly exceed the naive entry-only size: capacity * size_of::<(u64, usize)>().
520+
let entry_only_bytes = group_values.map.capacity() * size_of::<(u64, usize)>();
521+
assert!(
522+
table_bytes > entry_only_bytes,
523+
"allocation_size ({table_bytes}) must exceed entry-only capacity ({entry_only_bytes})"
524+
);
525+
526+
// GroupValuesRows::size() should accurately account for the map's allocation_size
527+
let expected_size = group_values.row_converter.size()
528+
+ group_values
529+
.group_values
530+
.as_ref()
531+
.map(|v| v.size())
532+
.unwrap_or(0)
533+
+ table_bytes
534+
+ group_values.rows_buffer.size()
535+
+ group_values.hashes_buffer.allocated_size();
536+
assert_eq!(group_values.size(), expected_size);
537+
538+
// 3. Partial emit (EmitTo::First) - table retains entries and capacity
539+
let emit_count = 200;
540+
let emitted = group_values.emit(EmitTo::First(emit_count))?;
541+
assert_eq!(emitted[0].len(), emit_count);
542+
assert_eq!(group_values.len(), (num_distinct as usize) - emit_count);
543+
544+
// Retained capacity is still charged until the table releases it
545+
assert_eq!(group_values.map.allocation_size(), table_bytes);
546+
assert_eq!(
547+
group_values.size(),
548+
group_values.row_converter.size()
549+
+ group_values
550+
.group_values
551+
.as_ref()
552+
.map(|v| v.size())
553+
.unwrap_or(0)
554+
+ table_bytes
555+
+ group_values.rows_buffer.size()
556+
+ group_values.hashes_buffer.allocated_size()
557+
);
558+
559+
// 4. Emit all (EmitTo::All) - map is cleared but capacity remains allocated for reuse
560+
let remaining_count = group_values.len();
561+
let emitted_all = group_values.emit(EmitTo::All)?;
562+
assert_eq!(emitted_all[0].len(), remaining_count);
563+
assert_eq!(group_values.len(), 0);
564+
565+
// Clear preserves table capacity
566+
assert_eq!(group_values.map.allocation_size(), table_bytes);
567+
568+
// 5. clear_shrink: releases/shrinks map capacity
569+
group_values.clear_shrink(16);
570+
let shrunken_table_bytes = group_values.map.allocation_size();
571+
assert!(
572+
shrunken_table_bytes < table_bytes,
573+
"clear_shrink should shrink table from {table_bytes} to {shrunken_table_bytes}"
574+
);
575+
assert_eq!(
576+
group_values.size(),
577+
group_values.row_converter.size()
578+
+ group_values
579+
.group_values
580+
.as_ref()
581+
.map(|v| v.size())
582+
.unwrap_or(0)
583+
+ shrunken_table_bytes
584+
+ group_values.rows_buffer.size()
585+
+ group_values.hashes_buffer.allocated_size()
586+
);
587+
588+
// 6. Reinsert new distinct values into the shrunken table
589+
let new_col_a =
590+
Arc::new(Int32Array::from_iter_values(num_distinct..num_distinct * 2))
591+
as ArrayRef;
592+
let new_col_b = Arc::new(StringArray::from_iter_values(
593+
(num_distinct..num_distinct * 2).map(|i| format!("key_string_{i}")),
594+
)) as ArrayRef;
595+
group_values.intern(&[new_col_a, new_col_b], &mut groups)?;
596+
assert_eq!(groups.len(), num_distinct as usize);
597+
assert!(group_values.map.allocation_size() > shrunken_table_bytes);
598+
assert_eq!(
599+
group_values.size(),
600+
group_values.row_converter.size()
601+
+ group_values
602+
.group_values
603+
.as_ref()
604+
.map(|v| v.size())
605+
.unwrap_or(0)
606+
+ group_values.map.allocation_size()
607+
+ group_values.rows_buffer.size()
608+
+ group_values.hashes_buffer.allocated_size()
609+
);
610+
611+
Ok(())
612+
}
613+
614+
#[test]
615+
fn test_exact_hash_table_allocation_accounting_nested_keys() -> Result<()> {
616+
let field = Arc::new(Field::new_list_field(DataType::Int32, true));
617+
let schema = Arc::new(Schema::new(vec![Field::new(
618+
"nested",
619+
DataType::List(field),
620+
true,
621+
)]));
622+
let mut group_values = GroupValuesRows::try_new(schema)?;
623+
624+
assert_eq!(group_values.map.capacity(), 0);
625+
assert_eq!(group_values.map.allocation_size(), 0);
626+
627+
let input = Arc::new(ListArray::from_iter_primitive::<Int32Type, _, _>(
628+
(0..256).map(|i| Some(vec![Some(i), Some(i * 2)])),
629+
)) as ArrayRef;
630+
let mut groups = vec![];
631+
group_values.intern(&[input], &mut groups)?;
632+
assert_eq!(groups.len(), 256);
633+
634+
let table_bytes = group_values.map.allocation_size();
635+
assert!(table_bytes > 0);
636+
assert_eq!(
637+
group_values.size(),
638+
group_values.row_converter.size()
639+
+ group_values
640+
.group_values
641+
.as_ref()
642+
.map(|v| v.size())
643+
.unwrap_or(0)
644+
+ table_bytes
645+
+ group_values.rows_buffer.size()
646+
+ group_values.hashes_buffer.allocated_size()
647+
);
648+
649+
// Shrink to 0: table should completely deallocate
650+
group_values.clear_shrink(0);
651+
assert_eq!(group_values.map.capacity(), 0);
652+
assert_eq!(group_values.map.allocation_size(), 0);
653+
654+
Ok(())
655+
}
487656
}

0 commit comments

Comments
 (0)