diff --git a/datafusion/physical-expr-common/src/binary_map.rs b/datafusion/physical-expr-common/src/binary_map.rs index 44ca35c7f870..4f8fae1f4195 100644 --- a/datafusion/physical-expr-common/src/binary_map.rs +++ b/datafusion/physical-expr-common/src/binary_map.rs @@ -244,10 +244,12 @@ where V: Debug + PartialEq + Eq + Clone + Copy + Default, { pub fn new(output_type: OutputType) -> Self { + let map = hashbrown::hash_table::HashTable::with_capacity(INITIAL_MAP_CAPACITY); + let map_size = map.allocation_size(); Self { output_type, - map: hashbrown::hash_table::HashTable::with_capacity(INITIAL_MAP_CAPACITY), - map_size: 0, + map, + map_size, buffer: Vec::with_capacity(INITIAL_BUFFER_CAPACITY), offsets: vec![O::default()], // first offset is always 0 random_state: RandomState::default(), @@ -264,6 +266,21 @@ where new_self } + fn insert_entry( + map: &mut hashbrown::hash_table::HashTable>, + map_size: &mut usize, + entry: Entry, + ) { + let capacity = map.capacity(); + map.insert_accounted(entry, |entry| entry.hash, map_size); + + // `insert_accounted` estimates growth from capacity. Keep `map_size` + // consistent with the exact allocation recorded by `new` after a resize. + if map.capacity() != capacity { + *map_size = map.allocation_size(); + } + } + /// 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. @@ -414,11 +431,7 @@ where offset_or_inline: inline, payload, }; - self.map.insert_accounted( - new_header, - |header| header.hash, - &mut self.map_size, - ); + Self::insert_entry(&mut self.map, &mut self.map_size, new_header); payload } } @@ -456,11 +469,7 @@ where offset_or_inline: offset, payload, }; - self.map.insert_accounted( - new_header, - |header| header.hash, - &mut self.map_size, - ); + Self::insert_entry(&mut self.map, &mut self.map_size, new_header); payload } }; diff --git a/datafusion/physical-expr-common/src/binary_view_map.rs b/datafusion/physical-expr-common/src/binary_view_map.rs index 9d4b556393a2..ee04bc5eb8ad 100644 --- a/datafusion/physical-expr-common/src/binary_view_map.rs +++ b/datafusion/physical-expr-common/src/binary_view_map.rs @@ -155,10 +155,12 @@ where V: Debug + PartialEq + Eq + Clone + Copy + Default, { pub fn new(output_type: OutputType) -> Self { + let map = hashbrown::hash_table::HashTable::with_capacity(INITIAL_MAP_CAPACITY); + let map_size = map.allocation_size(); Self { output_type, - map: hashbrown::hash_table::HashTable::with_capacity(INITIAL_MAP_CAPACITY), - map_size: 0, + map, + map_size, views: Vec::new(), in_progress: Vec::new(), completed: Vec::new(), @@ -177,6 +179,21 @@ where new_self } + fn insert_entry( + map: &mut hashbrown::hash_table::HashTable>, + map_size: &mut usize, + entry: Entry, + ) { + let capacity = map.capacity(); + map.insert_accounted(entry, |entry| entry.hash, map_size); + + // `insert_accounted` estimates growth from capacity. Keep `map_size` + // consistent with the exact allocation recorded by `new` after a resize. + if map.capacity() != capacity { + *map_size = map.allocation_size(); + } + } + /// 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. @@ -367,8 +384,7 @@ where payload, }; - self.map - .insert_accounted(new_header, |h| h.hash, &mut self.map_size); + Self::insert_entry(&mut self.map, &mut self.map_size, new_header); payload }; observe_payload_fn(payload); diff --git a/datafusion/physical-expr-common/tests/memory_accounting.rs b/datafusion/physical-expr-common/tests/memory_accounting.rs new file mode 100644 index 000000000000..9f280dca09ba --- /dev/null +++ b/datafusion/physical-expr-common/tests/memory_accounting.rs @@ -0,0 +1,82 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use arrow::array::{ArrayRef, StringArray, StringViewArray}; +use datafusion_physical_expr_common::binary_map::{ArrowBytesMap, OutputType}; +use datafusion_physical_expr_common::binary_view_map::ArrowBytesViewMap; + +type LargePayload = [u8; 32]; + +#[test] +fn arrow_bytes_map_size_accounts_for_hash_table_allocations() { + let mut unit_map = ArrowBytesMap::::new(OutputType::Utf8); + let mut payload_map = ArrowBytesMap::::new(OutputType::Utf8); + + let initial_allocation_difference = payload_map.size() - unit_map.size(); + assert!(initial_allocation_difference > 0); + + let values: ArrayRef = Arc::new(StringArray::from_iter_values( + (0..1024).map(|index| format!("value-{index}")), + )); + unit_map.insert_if_new(&values, |_| (), |_| {}); + payload_map.insert_if_new(&values, |_| LargePayload::default(), |_| {}); + + let grown_allocation_difference = payload_map.size() - unit_map.size(); + assert!(grown_allocation_difference > initial_allocation_difference); + + let populated_unit_map = unit_map.take(); + let populated_payload_map = payload_map.take(); + assert_eq!( + populated_payload_map.size() - populated_unit_map.size(), + grown_allocation_difference + ); + assert_eq!( + payload_map.size() - unit_map.size(), + initial_allocation_difference + ); +} + +#[test] +fn arrow_bytes_view_map_size_accounts_for_hash_table_allocations() { + let mut unit_map = ArrowBytesViewMap::<()>::new(OutputType::Utf8View); + let mut payload_map = ArrowBytesViewMap::::new(OutputType::Utf8View); + + let initial_allocation_difference = payload_map.size() - unit_map.size(); + assert!(initial_allocation_difference > 0); + + let values: ArrayRef = Arc::new(StringViewArray::from_iter_values( + (0..1024).map(|index| format!("value-{index}")), + )); + unit_map.insert_if_new(&values, |_| (), |_| {}); + payload_map.insert_if_new(&values, |_| LargePayload::default(), |_| {}); + + let grown_allocation_difference = payload_map.size() - unit_map.size(); + assert!(grown_allocation_difference > initial_allocation_difference); + + let populated_unit_map = unit_map.take(); + let populated_payload_map = payload_map.take(); + assert_eq!( + populated_payload_map.size() - populated_unit_map.size(), + grown_allocation_difference + ); + assert_eq!( + payload_map.size() - unit_map.size(), + initial_allocation_difference + ); +}