This is an automated email from the ASF dual-hosted git repository.
Jefffrey pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new f9bf62845c feat(coalesce): add size function (#10331)
f9bf62845c is described below
commit f9bf62845ca459c16938359e9378b34a4d8c51d9
Author: Raz Luvaton <[email protected]>
AuthorDate: Fri Jul 24 05:09:06 2026 +0300
feat(coalesce): add size function (#10331)
# Which issue does this PR close?
N/A
But this is prerequisite for:
- https://github.com/apache/datafusion/issues/23385
# Rationale for this change
We want to be able to reserve memory for the data that `BatchCoalescer`
is using in DataFusion and other places
# What changes are included in this PR?
added `size` functions and implement sizing + tests
# Are these changes tested?
Yes
# Are there any user-facing changes?
yes, new function
---------
Co-authored-by: Mikhail Zabaluev <[email protected]>
---
arrow-select/src/coalesce.rs | 127 +++++++++++++++++++++++++++++
arrow-select/src/coalesce/byte_view.rs | 145 +++++++++++++++++++++++++++++++--
arrow-select/src/coalesce/generic.rs | 139 ++++++++++++++++++++++++++++++-
arrow-select/src/coalesce/primitive.rs | 77 +++++++++++++++++
4 files changed, 478 insertions(+), 10 deletions(-)
diff --git a/arrow-select/src/coalesce.rs b/arrow-select/src/coalesce.rs
index 0f1086be82..ad1fe2ab92 100644
--- a/arrow-select/src/coalesce.rs
+++ b/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"
+ );
+ }
}
diff --git a/arrow-select/src/coalesce/byte_view.rs
b/arrow-select/src/coalesce/byte_view.rs
index 32f92d044e..f43572b37c 100644
--- a/arrow-select/src/coalesce/byte_view.rs
+++ b/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,6 +354,7 @@ impl<B: ByteViewType> InProgressByteViewArray<B> {
}
b.as_u128()
});
+
self.views.extend(new_views);
self.current = Some(dst_buffer);
}
@@ -335,6 +362,10 @@ impl<B: ByteViewType> InProgressByteViewArray<B> {
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);
+ }
}
diff --git a/arrow-select/src/coalesce/generic.rs
b/arrow-select/src/coalesce/generic.rs
index 4fa64273ec..5e99fb1045 100644
--- a/arrow-select/src/coalesce/generic.rs
+++ b/arrow-select/src/coalesce/generic.rs
@@ -18,7 +18,7 @@
use super::InProgressArray;
use crate::concat::concat;
use crate::filter::FilterPredicate;
-use arrow_array::ArrayRef;
+use arrow_array::{Array, ArrayRef};
use arrow_schema::ArrowError;
/// Generic implementation for [`InProgressArray`] that works with any type of
@@ -32,8 +32,14 @@ use arrow_schema::ArrowError;
pub(crate) struct GenericInProgressArray {
/// The current source
source: Option<ArrayRef>,
+
+ /// Is [`Self::source`] referenced in [`Self::buffered_arrays`]
+ source_data_referenced_in_buffers: bool,
/// The buffered array slices
buffered_arrays: Vec<ArrayRef>,
+
+ /// The number of bytes the arrays in [`Self::buffered_arrays`] takes
+ total_size_of_non_shared_buffers: usize,
}
impl GenericInProgressArray {
@@ -42,12 +48,22 @@ impl GenericInProgressArray {
Self {
source: None,
buffered_arrays: vec![],
+ total_size_of_non_shared_buffers: 0,
+ source_data_referenced_in_buffers: false,
}
}
}
impl InProgressArray for GenericInProgressArray {
fn set_source(&mut self, source: Option<ArrayRef>) {
- self.source = source
+ if let Some(old_source) = self.source.take() {
+ // If the source is still referenced in buffered_arrays,
+ // then count it now
+ if self.source_data_referenced_in_buffers {
+ self.total_size_of_non_shared_buffers +=
old_source.get_array_memory_size();
+ }
+ }
+ self.source_data_referenced_in_buffers = false;
+ self.source = source;
}
fn copy_rows(&mut self, offset: usize, len: usize) -> Result<(),
ArrowError> {
@@ -56,6 +72,9 @@ impl InProgressArray for GenericInProgressArray {
"Internal Error: GenericInProgressArray: source not
set".to_string(),
)
})?;
+ // No need to count the size of the array that was pushed to
buffered_arrays
+ // since the data is shared with the `source` memory that is already
counted
+ self.source_data_referenced_in_buffers = true;
let array = source.slice(offset, len);
self.buffered_arrays.push(array);
Ok(())
@@ -67,6 +86,7 @@ impl InProgressArray for GenericInProgressArray {
filter: &FilterPredicate,
) -> Result<(), ArrowError> {
let array = filter.filter(source.as_ref())?;
+ self.total_size_of_non_shared_buffers += array.get_array_memory_size();
self.buffered_arrays.push(array);
Ok(())
}
@@ -82,6 +102,121 @@ impl InProgressArray for GenericInProgressArray {
.collect::<Vec<_>>(),
)?;
self.buffered_arrays.clear();
+ self.total_size_of_non_shared_buffers = 0;
+ self.source_data_referenced_in_buffers = false;
Ok(array)
}
+
+ fn size(&self) -> usize {
+ self.total_size_of_non_shared_buffers
+ + self.buffered_arrays.capacity() * size_of::<ArrayRef>()
+ + self
+ .source
+ .as_ref()
+ .map_or(0, |a| a.get_array_memory_size())
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow_array::Int32Array;
+ use std::sync::Arc;
+
+ fn arr(range: std::ops::Range<i32>) -> ArrayRef {
+ Arc::new(Int32Array::from_iter_values(range))
+ }
+
+ #[test]
+ fn test_size_empty() {
+ let in_progress = GenericInProgressArray::new();
+ assert_eq!(in_progress.size(), 0);
+ }
+
+ #[test]
+ fn test_source_is_counted_in_memory_and_released_when_not_used() {
+ let mut in_progress = GenericInProgressArray::new();
+
+ // Starting with no used memory
+ assert_eq!(in_progress.size(), 0);
+ {
+ let source1 = arr(0..100);
+ in_progress.set_source(Some(Arc::clone(&source1)));
+
+ // We are holding on source so this is the size
+ assert_eq!(in_progress.size(), source1.get_array_memory_size());
+ }
+
+ {
+ // Replacing an unused source drops the old size and counts only
the new one
+ let source2 = arr(0..40);
+ in_progress.set_source(Some(Arc::clone(&source2)));
+ assert_eq!(in_progress.size(), source2.get_array_memory_size());
+ }
+
+ // Drop the used source
+ in_progress.set_source(None);
+
+ // The source is no longer being held, so it should reduce the size
+ assert_eq!(in_progress.size(), 0);
+ }
+
+ #[test]
+ fn test_double_copy_on_same_source_should_not_double_count() {
+ let mut in_progress = GenericInProgressArray::new();
+
+ // Starting with no used memory
+ assert_eq!(in_progress.size(), 0);
+
+ let source = arr(0..100);
+ in_progress.set_source(Some(Arc::clone(&source)));
+
+ // We are holding on source so this is the size
+ let size_before_copy = in_progress.size();
+ assert_eq!(size_before_copy, source.get_array_memory_size());
+
+ for _ in 0..2 {
+ // Only copy a subset
+ in_progress.copy_rows(0, 98).unwrap();
+
+ // The size should now account for the buffered but not the actual
array data since we are still holding on it
+ assert!(
+ in_progress.size() > size_before_copy,
+ "size after copy {} should be greater than before copy
{size_before_copy}",
+ in_progress.size()
+ );
+ {
+ let in_progress_size = in_progress.size() as f64;
+ let source_size = source.get_array_memory_size();
+ let size_if_source_and_sliced_would_be_counted = (source_size
as f64) * 1.8;
+ assert!(
+ in_progress_size <
size_if_source_and_sliced_would_be_counted,
+ "size after copy {in_progress_size} should not include the
source and sliced array (should be greater than
{size_if_source_and_sliced_would_be_counted}), source size is {source_size}"
+ );
+ }
+ }
+
+ let size_before_clear_source = in_progress.size();
+
+ // Drop the used source
+ in_progress.set_source(None);
+
+ // The source is still being held in the buffered
+ assert_eq!(in_progress.size(), size_before_clear_source);
+
+ {
+ let source2 = arr(0..40);
+ in_progress.set_source(Some(Arc::clone(&source2)));
+ assert_eq!(
+ in_progress.size(),
+ size_before_clear_source + source2.get_array_memory_size()
+ );
+ in_progress.set_source(None);
+ }
+
+ in_progress.finish().unwrap();
+
+ // There is still some memory being held by some leftover capacity but
not arrays
+ assert!(in_progress.size() < source.get_array_memory_size());
+ }
}
diff --git a/arrow-select/src/coalesce/primitive.rs
b/arrow-select/src/coalesce/primitive.rs
index ac831fe089..21ddd956cc 100644
--- a/arrow-select/src/coalesce/primitive.rs
+++ b/arrow-select/src/coalesce/primitive.rs
@@ -226,6 +226,14 @@ impl<T: ArrowPrimitiveType + Debug> InProgressArray for
InProgressPrimitiveArray
.with_data_type(self.data_type.clone());
Ok(Arc::new(array))
}
+
+ fn size(&self) -> usize {
+ self.source
+ .as_ref()
+ .map_or(0, |source| source.get_array_memory_size())
+ + self.current.capacity() * std::mem::size_of::<T::Native>()
+ + self.nulls.allocated_size()
+ }
}
#[cfg(test)]
@@ -305,4 +313,73 @@ mod tests {
]);
assert_eq!(result, &expected);
}
+
+ #[test]
+ fn test_size_empty() {
+ // A fresh in-progress array has allocated nothing yet
+ let in_progress = InProgressPrimitiveArray::<Int32Type>::new(64,
DataType::Int32);
+ assert_eq!(in_progress.size(), 0);
+ }
+
+ #[test]
+ fn test_size_counts_source() {
+ let mut in_progress = InProgressPrimitiveArray::<Int32Type>::new(64,
DataType::Int32);
+ let source: ArrayRef = Arc::new(Int32Array::from_iter_values(0..100));
+ in_progress.set_source(Some(Arc::clone(&source)));
+ // Nothing copied yet, so size is exactly the source's memory
+ assert_eq!(in_progress.size(), source.get_array_memory_size());
+ }
+
+ #[test]
+ fn test_size_counts_values_buffer_and_resets_on_finish() {
+ const BATCH_SIZE: usize = 64;
+ let mut in_progress =
+ InProgressPrimitiveArray::<Int32Type>::new(BATCH_SIZE,
DataType::Int32);
+ // Non-null source: the nulls builder stays empty (allocated_size ==
0),
+ // so the only growth is the values buffer.
+ let source: ArrayRef = Arc::new(Int32Array::from_iter_values(0..100));
+ let source_size = source.get_array_memory_size();
+ in_progress.set_source(Some(Arc::clone(&source)));
+
+ in_progress.copy_rows(0, 50).unwrap();
+ assert!(
+ in_progress.size() >= source_size + 50 * size_of::<i32>(),
+ "values buffer under-counted: {} < {} + {} * {}",
+ in_progress.size(),
+ source_size,
+ 50,
+ size_of::<i32>(),
+ );
+
+ // finish() takes the buffered values/nulls but keeps the source, so
the
+ // reported size drops back to exactly the source.
+ in_progress.finish().unwrap();
+ assert_eq!(in_progress.size(), source_size);
+ }
+
+ #[test]
+ fn test_size_counts_null_buffer() {
+ const BATCH_SIZE: usize = 64;
+
+ let in_progress_bytes = |source: ArrayRef| {
+ let mut in_progress =
+ InProgressPrimitiveArray::<Int32Type>::new(BATCH_SIZE,
DataType::Int32);
+ let source_len = source.len();
+ in_progress.set_source(Some(source));
+ in_progress.copy_rows(0, source_len / 2).unwrap();
+ in_progress.size()
+ };
+
+ // All values valid: the nulls builder never allocates.
+ let all_valid =
in_progress_bytes(Arc::new(Int32Array::from_iter_values(0..100)));
+ // Some values null: copying materializes a null buffer that must
count.
+ let with_nulls = in_progress_bytes(Arc::new(Int32Array::from_iter(
+ (0..100).map(|i| (i % 2 == 0).then_some(i)),
+ )));
+
+ assert!(
+ with_nulls > all_valid,
+ "null buffer must be included in size(): with_nulls={with_nulls}
all_valid={all_valid}"
+ );
+ }
}
