Skip to content

Commit 1f0615a

Browse files
authored
Fix ArrowBytesViewMap retained capacity accounting (#24257)
## Which issue does this PR close? * Part of #23393 ## Rationale for this change `ArrowBytesViewMap::size()` can underreport memory retained by the map. It previously omitted the initial allocation of the hash table, counted `views` by length rather than capacity, counted completed buffers by used length rather than retained capacity, and did not include the backing allocation of the `completed` vector. Because this value is used for memory accounting, it should reflect heap allocations owned by the map while continuing to exclude `self` and external input-array buffers. ## What changes are included in this PR? * Initialize `map_size` from the hash table's initial allocated capacity. * Account for `views` and the in-progress buffer using their allocated sizes. * Include the allocation backing the `completed` vector. * Account for each completed Arrow `Buffer` by retained capacity rather than used length. * Clarify that `size()` excludes both `self` and input-array buffers. ## Are these changes tested? Yes. This PR adds: * `test_size_counts_initial_hash_table_capacity`, which verifies that a newly created map reports the initial hash table allocation. * `test_size_counts_retained_buffer_capacities`, which verifies unused `views` capacity, `completed` vector storage, retained completed-buffer capacity, and that re-inserting duplicate values does not increase the reported size. ## Are there any user-facing changes? No public API or query-result behavior changes. This fixes internal memory accounting so `ArrowBytesViewMap::size()` more accurately reports allocations retained by the map. ## LLM-generated code disclosure This PR includes LLM-generated code and comments. All LLM-generated content has been manually reviewed.
1 parent 84a8e6d commit 1f0615a

1 file changed

Lines changed: 80 additions & 6 deletions

File tree

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

Lines changed: 80 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -155,10 +155,13 @@ where
155155
V: Debug + PartialEq + Eq + Clone + Copy + Default,
156156
{
157157
pub fn new(output_type: OutputType) -> Self {
158+
let map = hashbrown::hash_table::HashTable::with_capacity(INITIAL_MAP_CAPACITY);
159+
let map_size = map.capacity() * size_of::<Entry<V>>();
160+
158161
Self {
159162
output_type,
160-
map: hashbrown::hash_table::HashTable::with_capacity(INITIAL_MAP_CAPACITY),
161-
map_size: 0,
163+
map,
164+
map_size,
162165
views: Vec::new(),
163166
in_progress: Vec::new(),
164167
completed: Vec::new(),
@@ -469,11 +472,14 @@ where
469472
}
470473

471474
/// Return the total size, in bytes, of memory used to store the data in
472-
/// this set, not including `self`
475+
/// this set, not including `self` or input-array buffers.
473476
pub fn size(&self) -> usize {
474-
let views_size = self.views.len() * size_of::<u128>();
475-
let in_progress_size = self.in_progress.capacity();
476-
let completed_size: usize = self.completed.iter().map(|b| b.len()).sum();
477+
// All fields below own their allocations. Count retained capacity rather
478+
// than used length because this value drives memory accounting.
479+
let views_size = self.views.allocated_size();
480+
let in_progress_size = self.in_progress.allocated_size();
481+
let completed_size = self.completed.allocated_size()
482+
+ self.completed.iter().map(Buffer::capacity).sum::<usize>();
477483
let nulls_size = self.nulls.allocated_size();
478484

479485
self.map_size
@@ -715,6 +721,74 @@ mod tests {
715721
assert_eq!(set.len(), 10);
716722
}
717723

724+
#[test]
725+
fn test_size_counts_initial_hash_table_capacity() {
726+
let map = ArrowBytesViewMap::<()>::new(OutputType::Utf8View);
727+
728+
assert_eq!(map.size(), map.map.capacity() * size_of::<Entry<()>>());
729+
}
730+
731+
#[test]
732+
fn test_size_counts_retained_buffer_capacities() {
733+
let first = "a".repeat(BYTE_VIEW_MAX_BLOCK_SIZE / 2 + 1);
734+
let second = "b".repeat(BYTE_VIEW_MAX_BLOCK_SIZE / 2 + 1);
735+
let third = "c".repeat(BYTE_VIEW_MAX_BLOCK_SIZE / 2 + 1);
736+
let values: ArrayRef = Arc::new(StringViewArray::from(vec![
737+
first.as_str(),
738+
second.as_str(),
739+
third.as_str(),
740+
]));
741+
742+
let mut map = ArrowBytesViewMap::new(OutputType::Utf8View);
743+
map.insert_if_new(&values, |_| (), |_| {});
744+
745+
// Make unused vector capacity explicit; the completed buffers were created
746+
// by the map's flush path.
747+
map.views.shrink_to_fit();
748+
map.views.reserve_exact(1);
749+
map.completed.shrink_to_fit();
750+
map.completed.reserve_exact(1);
751+
752+
// The map owns these allocations; `values` and its Arrow buffers remain external.
753+
assert!(map.views.capacity() > map.views.len());
754+
assert!(map.completed.capacity() > map.completed.len());
755+
assert!(
756+
map.completed
757+
.iter()
758+
.any(|buffer| buffer.capacity() > buffer.len())
759+
);
760+
761+
let expected_size = map.map_size
762+
+ map.views.allocated_size()
763+
+ map.in_progress.allocated_size()
764+
+ map.completed.allocated_size()
765+
+ map.completed.iter().map(Buffer::capacity).sum::<usize>()
766+
+ map.nulls.allocated_size()
767+
+ map.hashes_buffer.allocated_size();
768+
assert_eq!(map.size(), expected_size);
769+
770+
// Verify the retained-capacity delta independently from the production formula.
771+
let legacy_size = map.map_size
772+
+ map.views.len() * size_of::<u128>()
773+
+ map.in_progress.capacity()
774+
+ map.completed.iter().map(Buffer::len).sum::<usize>()
775+
+ map.nulls.allocated_size()
776+
+ map.hashes_buffer.allocated_size();
777+
let retained_capacity_delta = (map.views.capacity() - map.views.len())
778+
* size_of::<u128>()
779+
+ map.completed.capacity() * size_of::<Buffer>()
780+
+ map
781+
.completed
782+
.iter()
783+
.map(|buffer| buffer.capacity() - buffer.len())
784+
.sum::<usize>();
785+
assert_eq!(map.size() - legacy_size, retained_capacity_delta);
786+
787+
let size_after_insert = map.size();
788+
map.insert_if_new(&values, |_| (), |_| {});
789+
assert_eq!(map.size(), size_after_insert);
790+
}
791+
718792
#[derive(Debug, PartialEq, Eq, Default, Clone, Copy)]
719793
struct TestPayload {
720794
// store the string value to check against input

0 commit comments

Comments
 (0)