feat(coalesce): add size function

arrow-select/src/coalesce.rs

@@ -583,6 +583,22 @@ impl BatchCoalescer {
     pub fn next_completed_batch(&mut self) -> Option<RecordBatch> {
         self.completed.pop_front()
     }
+
+    /// Returns the number of bytes used by this data structure.
+    pub fn size(&self) -> usize {
+        self.in_progress_arrays.capacity() * size_of::<Box<dyn InProgressArray>>()
+            + self
+                .in_progress_arrays
+                .iter()
+                .map(|array| array.size())
+                .sum::<usize>()
+            + self.completed.capacity() * size_of::<RecordBatch>()
+            + self
+                .completed
+                .iter()
+                .map(|batch| batch.get_array_memory_size())
+                .sum::<usize>()
+    }
 }
 
 impl BatchCoalescer {
@@ -742,6 +758,9 @@ trait InProgressArray: std::fmt::Debug + Send + Sync {
 
     /// Finish the currently in-progress array and return it as an `ArrayRef`
     fn finish(&mut self) -> Result<ArrayRef, ArrowError>;
+
+    /// Get the number of bytes this array is using
+    fn size(&self) -> usize;
 }
 
 #[cfg(test)]
@@ -2683,4 +2702,112 @@ mod tests {
             "unexpected error: {err}"
         );
     }
+
+    #[test]
+    fn test_size_grows_with_buffering_and_shrinks_when_draining() {
+        let batch = uint32_batch(0..8);
+        let mut coalescer = BatchCoalescer::new(batch.schema(), 21);
+        let baseline = coalescer.size();
+        assert!(baseline > 0, "size includes container capacities");
+
+        // Buffer rows without completing a batch
+        coalescer.push_batch(batch.clone()).unwrap();
+        assert!(coalescer.next_completed_batch().is_none());
+        let buffered = coalescer.size();
+        assert!(
+            buffered > baseline,
+            "buffering rows should grow size ({buffered} > {baseline})"
+        );
+
+        // Push enough to complete several batches
+        for _ in 0..10 {
+            coalescer.push_batch(batch.clone()).unwrap();
+        }
+        let peak = coalescer.size();
+        assert!(peak > buffered);
+
+        // Draining completed batches must never grow the reported size
+        let mut prev = peak;
+        let mut drained_any = false;
+        while coalescer.next_completed_batch().is_some() {
+            drained_any = true;
+            let now = coalescer.size();
+            assert!(now <= prev, "size grew while draining: {now} > {prev}");
+            prev = now;
+        }
+        assert!(drained_any);
+        assert!(
+            prev < peak,
+            "draining completed batches should release memory"
+        );
+    }
+
+    #[test]
+    fn test_size_string_view_buffers_released_after_drain() {
+        // Long strings spill into external data buffers, exercising the byte-view
+        // size accounting through the real coalescer path (including compaction).
+        let batch = stringview_batch_repeated(
+            1000,
+            [Some("this string is definitely longer than 12 bytes")],
+        );
+        let mut coalescer = BatchCoalescer::new(batch.schema(), 4096);
+        let baseline = coalescer.size();
+
+        for _ in 0..20 {
+            coalescer.push_batch(batch.clone()).unwrap();
+        }
+        let peak = coalescer.size();
+        assert!(
+            peak > baseline,
+            "buffered string view data should grow size ({peak} > {baseline})"
+        );
+
+        coalescer.finish_buffered_batch().unwrap();
+        while coalescer.next_completed_batch().is_some() {}
+
+        // Once fully drained the in-progress byte-view buffers are released.
+        let drained = coalescer.size();
+        assert!(
+            drained < peak,
+            "draining should release buffered data ({drained} < {peak})"
+        );
+    }
+
+    /// Every byte added to the accounting must eventually be removed: running the
+    /// exact same push/finish/drain sequence twice must report identical sizes.
+    /// This catches accounting leaks and drift without hard-coding magic numbers.
+    #[test]
+    fn test_size_accounting_conserved_across_cycles() {
+        // Primitive column: internal capacities stabilize after the first cycle
+        // (unlike byte-view, whose buffer sizer keeps growing), so the readings
+        // are deterministic across cycles.
+        let batch = uint32_batch(0..8);
+        let mut coalescer = BatchCoalescer::new(batch.schema(), 4096);
+
+        let run_cycle = |coalescer: &mut BatchCoalescer| {
+            for _ in 0..20 {
+                coalescer.push_batch(batch.clone()).unwrap();
+            }
+            coalescer.finish_buffered_batch().unwrap();
+            let peak = coalescer.size();
+            while coalescer.next_completed_batch().is_some() {}
+            (peak, coalescer.size())
+        };
+
+        let (peak1, drained1) = run_cycle(&mut coalescer);
+        let (peak2, drained2) = run_cycle(&mut coalescer);
+
+        assert_eq!(
+            peak1, peak2,
+            "identical work must report identical peak size"
+        );
+        assert_eq!(
+            drained1, drained2,
+            "fully-drained size must be stable across cycles (no accounting leak)"
+        );
+        assert!(
+            drained1 < peak1,
+            "draining must release the accounted memory"
+        );
+    }
 }

arrow-select/src/coalesce/byte_view.rs

@@ -54,6 +54,10 @@ pub(crate) struct InProgressByteViewArray<B: ByteViewType> {
     /// Phantom so we can use the same struct for both StringViewArray and
     /// BinaryViewArray
     _phantom: PhantomData<B>,
+    /// The size in bytes the [`Buffer`]s in [`Self::completed`] is taking
+    completed_buffers_size: usize,
+    /// The size in bytes from [`Self::source`] that it is being used in [`Self::completed`]
+    size_of_completed_buffers_from_current_source: usize,
 }
 
 struct Source {
@@ -89,6 +93,8 @@ impl<B: ByteViewType> InProgressByteViewArray<B> {
             nulls: NullBufferBuilder::new(batch_size), // no allocation
             current: None,
             completed: vec![],
+            completed_buffers_size: 0,
+            size_of_completed_buffers_from_current_source: 0,
             buffer_source,
             _phantom: PhantomData,
         }
@@ -109,7 +115,10 @@ impl<B: ByteViewType> InProgressByteViewArray<B> {
         let Some(next_buffer) = self.current.take() else {
             return;
         };
-        self.completed.push(next_buffer.into());
+        let buffer: Buffer = next_buffer.into();
+
+        self.completed_buffers_size += buffer.capacity();
+        self.completed.push(buffer);
     }
 
     fn append_views_by_filter(&mut self, views: &[u128], filter: &FilterPredicate) {
@@ -169,10 +178,26 @@ impl<B: ByteViewType> InProgressByteViewArray<B> {
 
     /// Append views to self.views, updating the buffer index if necessary
     #[inline(never)]
-    fn append_views_and_update_buffer_index(&mut self, views: &[u128], buffers: &[Buffer]) {
+    fn append_views_and_update_buffer_index(
+        &mut self,
+        views: &[u128],
+        buffers: &[Buffer],
+        is_reused: bool,
+    ) {
         if let Some(buffer) = self.current.take() {
-            self.completed.push(buffer.into());
+            let buffer: Buffer = buffer.into();
+            self.completed_buffers_size += buffer.capacity();
+            self.completed.push(buffer);
         }
+
+        let buffers_size = buffers.iter().map(|b| b.capacity()).sum::<usize>();
+        if !is_reused {
+            self.completed_buffers_size += buffers_size;
+        } else if self.size_of_completed_buffers_from_current_source == 0 {
+            // Don't double count buffers size if already counted that
+            self.size_of_completed_buffers_from_current_source += buffers_size;
+        }
+
         let starting_buffer: u32 = self.completed.len().try_into().expect("too many buffers");
         self.completed.extend_from_slice(buffers);
 
@@ -251,8 +276,9 @@ impl<B: ByteViewType> InProgressByteViewArray<B> {
         let remaining_view_buffer_size = view_buffer_size - string_bytes_to_copy;
 
         self.append_views_and_copy_strings_inner(first_views, current, buffers);
-        let completed = self.current.take().expect("completed");
-        self.completed.push(completed.into());
+        let completed: Buffer = self.current.take().expect("completed").into();
+        self.completed_buffers_size += completed.capacity();
+        self.completed.push(completed);
 
         // Copy any remaining views into a new buffer
         let remaining_views = &views[num_view_to_current..];
@@ -328,13 +354,18 @@ impl<B: ByteViewType> InProgressByteViewArray<B> {
             }
             b.as_u128()
         });
+
         self.views.extend(new_views);
         self.current = Some(dst_buffer);
     }
 }
 
 impl<B: ByteViewType> InProgressArray for InProgressByteViewArray<B> {
     fn set_source(&mut self, source: Option<ArrayRef>) {
+        // If used values from source, add only the size that was used
+        self.completed_buffers_size += self.size_of_completed_buffers_from_current_source;
+        self.size_of_completed_buffers_from_current_source = 0;
+
         self.source = source.map(|array| {
             let s = array.as_byte_view::<B>();
 
@@ -358,7 +389,7 @@ impl<B: ByteViewType> InProgressArray for InProgressByteViewArray<B> {
                 need_gc,
                 ideal_buffer_size,
             }
-        })
+        });
     }
 
     fn copy_rows(&mut self, offset: usize, len: usize) -> Result<(), ArrowError> {
@@ -397,7 +428,7 @@ impl<B: ByteViewType> InProgressArray for InProgressByteViewArray<B> {
         if source.need_gc {
             self.append_views_and_copy_strings(views, source.ideal_buffer_size, buffers);
         } else {
-            self.append_views_and_update_buffer_index(views, buffers);
+            self.append_views_and_update_buffer_index(views, buffers, true);
         }
         self.source = Some(source);
         Ok(())
@@ -452,7 +483,7 @@ impl<B: ByteViewType> InProgressArray for InProgressByteViewArray<B> {
         } else {
             self.nulls.append_n_non_nulls(filter.count());
         }
-        self.append_views_and_update_buffer_index(filtered.views(), filtered.data_buffers());
+        self.append_views_and_update_buffer_index(filtered.views(), filtered.data_buffers(), false);
         Ok(())
     }
 
@@ -464,12 +495,27 @@ impl<B: ByteViewType> InProgressArray for InProgressByteViewArray<B> {
         let nulls = self.nulls.finish();
         self.nulls = NullBufferBuilder::new(self.batch_size);
 
+        // Not reusing anything since we took all complete
+        self.size_of_completed_buffers_from_current_source = 0;
+        self.completed_buffers_size = 0;
+
         // Safety: we created valid views and buffers above and the
         // input arrays had value data and nulls
         let new_array =
             unsafe { GenericByteViewArray::<B>::new_unchecked(views.into(), buffers, nulls) };
         Ok(Arc::new(new_array))
     }
+
+    fn size(&self) -> usize {
+        self.completed_buffers_size
+            + self.current.as_ref().map_or(0, |c| c.capacity())
+            + self.nulls.allocated_size()
+            + self.views.capacity() * size_of::<u128>()
+            + self
+                .source
+                .as_ref()
+                .map_or(0, |s| s.array.get_array_memory_size())
+    }
 }
 
 const STARTING_BLOCK_SIZE: usize = 4 * 1024; // (note the first size used is actually 8KiB)
@@ -607,4 +653,87 @@ mod tests {
             "expected filtered output to reuse the source data buffer"
         );
     }
+
+    /// Build a compacted BinaryViewArray whose values all spill into external
+    /// data buffers. `gc()` makes the buffers dense so the coalescer reuses them
+    /// (`need_gc == false`) rather than copying/compacting them.
+    fn non_inline_array(n: usize) -> (BinaryViewArray, usize) {
+        let values = (0..n)
+            .map(|i| format!("This value is longer than 12 bytes: {i}").into_bytes())
+            .collect::<Vec<_>>();
+        let array = BinaryViewArray::from_iter(values.iter().map(|v| Some(v.as_slice()))).gc();
+        assert!(!array.data_buffers().is_empty());
+        let buffer_capacity = array.data_buffers().iter().map(|b| b.capacity()).sum();
+        (array, buffer_capacity)
+    }
+
+    #[test]
+    fn test_size_empty() {
+        let in_progress = InProgressByteViewArray::<BinaryViewType>::new(64);
+        assert_eq!(in_progress.size(), 0);
+    }
+
+    #[test]
+    fn test_size_reused_buffers_not_double_counted() {
+        let (array, buffer_capacity) = non_inline_array(64);
+        let source: ArrayRef = Arc::new(array);
+
+        let mut in_progress = InProgressByteViewArray::<BinaryViewType>::new(64);
+        in_progress.set_source(Some(Arc::clone(&source)));
+
+        in_progress.copy_rows(0, 60).unwrap();
+        in_progress.copy_rows(60, 4).unwrap();
+
+        // The reused buffers now live in `completed`, but while the source is
+        // still set they are counted via the source, not `completed_buffers_size`,
+        // to avoid double counting them.
+        assert_eq!(in_progress.completed_buffers_size, 0);
+        assert_eq!(
+            in_progress.size_of_completed_buffers_from_current_source,
+            buffer_capacity,
+        );
+
+        // Setting a new source commits the pending reused-buffer bytes, since the
+        // old source (and its double count) is dropped.
+        let other: ArrayRef = Arc::new(BinaryViewArray::from_iter(std::iter::once(Some(
+            b"short".as_slice(),
+        ))));
+        in_progress.set_source(Some(Arc::clone(&other)));
+        assert_eq!(in_progress.completed_buffers_size, buffer_capacity);
+        assert_eq!(in_progress.size_of_completed_buffers_from_current_source, 0);
+
+        // finish() releases everything.
+        in_progress.finish().unwrap();
+        assert_eq!(in_progress.completed_buffers_size, 0);
+        assert_eq!(in_progress.size_of_completed_buffers_from_current_source, 0);
+    }
+
+    #[test]
+    fn size_should_be_the_same_if_copying_multiple_time_from_same_source_or_once() {
+        let (array, _buffer_capacity) = non_inline_array(64);
+        let source: ArrayRef = Arc::new(array);
+
+        let in_progress_size_with_split = {
+            let mut in_progress = InProgressByteViewArray::<BinaryViewType>::new(64);
+            in_progress.set_source(Some(Arc::clone(&source)));
+
+            in_progress.copy_rows(0, 60).unwrap();
+            in_progress.copy_rows(60, 4).unwrap();
+            in_progress.set_source(None);
+
+            in_progress.size()
+        };
+
+        let in_progress_size_without_split = {
+            let mut in_progress = InProgressByteViewArray::<BinaryViewType>::new(64);
+            in_progress.set_source(Some(Arc::clone(&source)));
+
+            in_progress.copy_rows(0, 64).unwrap();
+            in_progress.set_source(None);
+
+            in_progress.size()
+        };
+
+        assert_eq!(in_progress_size_with_split, in_progress_size_without_split);
+    }
 }