fix: account for Arrow bytes map hash table allocations

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<Entry<O, V>>,
+        map_size: &mut usize,
+        entry: Entry<O, V>,
+    ) {
+        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
                 }
             };

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<Entry<V>>,
+        map_size: &mut usize,
+        entry: Entry<V>,
+    ) {
+        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);

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::<i32, ()>::new(OutputType::Utf8);
+    let mut payload_map = ArrowBytesMap::<i32, LargePayload>::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::<LargePayload>::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
+    );
+}