Fix ArrowBytesViewMap retained capacity accounting

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

@@ -155,10 +155,13 @@ 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.capacity() * size_of::<Entry<V>>();
+
         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(),
@@ -469,11 +472,14 @@ where
     }
 
     /// Return the total size, in bytes, of memory used to store the data in
-    /// this set, not including `self`
+    /// this set, not including `self` or input-array buffers.
     pub fn size(&self) -> usize {
-        let views_size = self.views.len() * size_of::<u128>();
-        let in_progress_size = self.in_progress.capacity();
-        let completed_size: usize = self.completed.iter().map(|b| b.len()).sum();
+        // All fields below own their allocations. Count retained capacity rather
+        // than used length because this value drives memory accounting.
+        let views_size = self.views.allocated_size();
+        let in_progress_size = self.in_progress.allocated_size();
+        let completed_size = self.completed.allocated_size()
+            + self.completed.iter().map(Buffer::capacity).sum::<usize>();
         let nulls_size = self.nulls.allocated_size();
 
         self.map_size
@@ -715,6 +721,74 @@ mod tests {
         assert_eq!(set.len(), 10);
     }
 
+    #[test]
+    fn test_size_counts_initial_hash_table_capacity() {
+        let map = ArrowBytesViewMap::<()>::new(OutputType::Utf8View);
+
+        assert_eq!(map.size(), map.map.capacity() * size_of::<Entry<()>>());
+    }
+
+    #[test]
+    fn test_size_counts_retained_buffer_capacities() {
+        let first = "a".repeat(BYTE_VIEW_MAX_BLOCK_SIZE / 2 + 1);
+        let second = "b".repeat(BYTE_VIEW_MAX_BLOCK_SIZE / 2 + 1);
+        let third = "c".repeat(BYTE_VIEW_MAX_BLOCK_SIZE / 2 + 1);
+        let values: ArrayRef = Arc::new(StringViewArray::from(vec![
+            first.as_str(),
+            second.as_str(),
+            third.as_str(),
+        ]));
+
+        let mut map = ArrowBytesViewMap::new(OutputType::Utf8View);
+        map.insert_if_new(&values, |_| (), |_| {});
+
+        // Make unused vector capacity explicit; the completed buffers were created
+        // by the map's flush path.
+        map.views.shrink_to_fit();
+        map.views.reserve_exact(1);
+        map.completed.shrink_to_fit();
+        map.completed.reserve_exact(1);
+
+        // The map owns these allocations; `values` and its Arrow buffers remain external.
+        assert!(map.views.capacity() > map.views.len());
+        assert!(map.completed.capacity() > map.completed.len());
+        assert!(
+            map.completed
+                .iter()
+                .any(|buffer| buffer.capacity() > buffer.len())
+        );
+
+        let expected_size = map.map_size
+            + map.views.allocated_size()
+            + map.in_progress.allocated_size()
+            + map.completed.allocated_size()
+            + map.completed.iter().map(Buffer::capacity).sum::<usize>()
+            + map.nulls.allocated_size()
+            + map.hashes_buffer.allocated_size();
+        assert_eq!(map.size(), expected_size);
+
+        // Verify the retained-capacity delta independently from the production formula.
+        let legacy_size = map.map_size
+            + map.views.len() * size_of::<u128>()
+            + map.in_progress.capacity()
+            + map.completed.iter().map(Buffer::len).sum::<usize>()
+            + map.nulls.allocated_size()
+            + map.hashes_buffer.allocated_size();
+        let retained_capacity_delta = (map.views.capacity() - map.views.len())
+            * size_of::<u128>()
+            + map.completed.capacity() * size_of::<Buffer>()
+            + map
+                .completed
+                .iter()
+                .map(|buffer| buffer.capacity() - buffer.len())
+                .sum::<usize>();
+        assert_eq!(map.size() - legacy_size, retained_capacity_delta);
+
+        let size_after_insert = map.size();
+        map.insert_if_new(&values, |_| (), |_| {});
+        assert_eq!(map.size(), size_after_insert);
+    }
+
     #[derive(Debug, PartialEq, Eq, Default, Clone, Copy)]
     struct TestPayload {
         // store the string value to check against input