diff --git a/datafusion/common/benches/record_batch_memory.rs b/datafusion/common/benches/record_batch_memory.rs index 2479d6ac987cb..014915713cd12 100644 --- a/datafusion/common/benches/record_batch_memory.rs +++ b/datafusion/common/benches/record_batch_memory.rs @@ -180,11 +180,44 @@ fn benchmark_shared_slices(c: &mut Criterion) { group.finish(); } +fn benchmark_count_uncount(c: &mut Criterion) { + let mut group = c.benchmark_group("record_batch_memory_size/count_uncount"); + + // Model the sort/window workload: batches are counted as they arrive and + // uncounted when they are dropped. Measure the full cycle. + for num_columns in [4, 16, 64] { + let batch = make_primitive_batch(8192, num_columns); + let slices = (0..32) + .map(|index| batch.slice(index * 256, 256)) + .collect::>(); + + group.bench_with_input( + BenchmarkId::from_parameter(num_columns), + &slices, + |bencher, slices| { + bencher.iter(|| { + let mut counter = RecordBatchMemoryCounter::new(); + for batch in black_box(slices) { + black_box(counter.count_batch(black_box(batch))); + } + for batch in black_box(slices) { + black_box(counter.uncount_batch(black_box(batch))); + } + assert_eq!(counter.memory_usage(), 0); + }); + }, + ); + } + + group.finish(); +} + criterion_group!( benches, benchmark_column_count, benchmark_row_count, benchmark_array_layout, - benchmark_shared_slices + benchmark_shared_slices, + benchmark_count_uncount ); criterion_main!(benches); diff --git a/datafusion/common/src/utils/memory.rs b/datafusion/common/src/utils/memory.rs index ac5e15301a9f8..e42fc40cf1007 100644 --- a/datafusion/common/src/utils/memory.rs +++ b/datafusion/common/src/utils/memory.rs @@ -36,8 +36,8 @@ use std::num::NonZero; use std::sync::Arc; /// Maximum number of distinct buffer IDs retained inline before promotion to -/// a [`HashSet`]. Sixteen keeps small buffer sets allocation-free while -/// limiting linear lookup and inline storage to 16 pointer-sized entries. +/// a heap-allocated map. Sixteen keeps small buffer sets allocation-free while +/// limiting linear lookup and inline storage to 16 entries. /// This is a performance heuristic, not a semantic limit. const INLINE_BUFFER_IDS: usize = 16; @@ -165,11 +165,36 @@ pub fn get_record_batch_memory_size(batch: &RecordBatch) -> usize { /// batch's buffers are kept alive by the batch even when only a sub-range is /// referenced, so counting unique buffers in full reflects the memory the /// batches actually retain. +/// +/// # Releasing batches +/// +/// Operators that retain batches incrementally and drop them later (sort, +/// window, sort-merge join, TopK) can use [`Self::uncount_batch`] to release +/// a batch's contribution. The counter tracks a reference count per buffer: +/// a buffer's capacity is added when its count goes from 0 to 1 and +/// subtracted when it returns to 0. +/// +/// **Contract**: a batch must be uncounted *before* its buffers are dropped. +/// The counter identifies buffers by their data-pointer address and does not +/// keep them alive. +/// +/// ```text +/// ───────────────────────┬──────────────────────────── +/// step │ memory_usage() +/// ───────────────────────┼──────────────────────────── +/// count_batch(slice1) │ 4 MB (parent buffer counted) +/// count_batch(slice2) │ 4 MB (already counted) +/// uncount_batch(slice1) │ 4 MB (slice2 still uses it) +/// uncount_batch(slice2) │ 0 +/// ───────────────────────┴──────────────────────────── +/// ``` #[derive(Debug, Default)] pub struct RecordBatchMemoryCounter { - /// Start addresses of `Buffer`s that have already been counted (instead of - /// actual used data region's pointer represented by current `Array`) - counted_buffers: BufferIdSet, + /// Reference-counted buffer tracker. Each buffer (identified by its + /// data-pointer address) maps to its current reference count. A buffer's + /// capacity is added to `memory_usage` when its count goes from 0→1 and + /// subtracted when it returns to 0. + counted_buffers: BufferIdMap, /// Array objects already counted by [`Self::count_batch_with_array_overhead`] counted_arrays: HashSet, /// Total memory of all counted allocations @@ -187,7 +212,7 @@ impl RecordBatchMemoryCounter { let previous_memory_usage = self.memory_usage; for array in batch.columns() { - self.count_array_memory_size(array.as_ref()); + self.visit_array_buffers(array.as_ref(), BufferOp::Count); } self.memory_usage - previous_memory_usage @@ -197,7 +222,7 @@ impl RecordBatchMemoryCounter { /// been counted before. pub fn count_array(&mut self, array: &dyn Array) -> usize { let previous_memory_usage = self.memory_usage; - self.count_array_memory_size(array); + self.visit_array_buffers(array, BufferOp::Count); self.memory_usage - previous_memory_usage } @@ -222,168 +247,238 @@ impl RecordBatchMemoryCounter { total_size } + /// Release the buffers of `batch`, returning the bytes freed. + /// + /// Each buffer's reference count is decremented; the buffer's capacity is + /// returned (and subtracted from [`Self::memory_usage`]) only when its + /// count reaches zero. Buffers that were never counted are silently + /// ignored (their count is already zero). + pub fn uncount_batch(&mut self, batch: &RecordBatch) -> usize { + let previous_memory_usage = self.memory_usage; + + for array in batch.columns() { + self.visit_array_buffers(array.as_ref(), BufferOp::Uncount); + } + + previous_memory_usage - self.memory_usage + } + + /// Release the buffers of `array`, returning the bytes freed. + /// + /// See [`Self::uncount_batch`] for semantics. + pub fn uncount_array(&mut self, array: &dyn Array) -> usize { + let previous_memory_usage = self.memory_usage; + self.visit_array_buffers(array, BufferOp::Uncount); + previous_memory_usage - self.memory_usage + } + + /// Inverse of [`Self::count_batch_with_array_overhead`]: releases buffer + /// memory and array-object overhead for `batch`. + /// + /// Array-object overhead is released when the array identity (pointer) is + /// removed from the tracked set. + pub fn uncount_batch_with_array_overhead(&mut self, batch: &RecordBatch) -> usize { + let mut total_released = self.uncount_batch(batch); + let mut array_overhead = 0; + + for array in batch.columns() { + array_overhead += + uncount_unique_array_object_memory_size(array, &mut self.counted_arrays); + } + + total_released += array_overhead; + self.memory_usage -= array_overhead; + total_released + } + /// Total memory of all counted allocations. pub fn memory_usage(&self) -> usize { self.memory_usage } - fn count_buffer_memory_size(&mut self, buffer: &Buffer) { - if self.counted_buffers.insert(buffer.data_ptr().addr()) { - self.memory_usage += buffer.capacity(); + /// Apply `op` to `buffer`: increment or decrement its reference count and + /// adjust `memory_usage` accordingly. + fn apply_buffer_op(&mut self, buffer: &Buffer, op: BufferOp) { + let addr = buffer.data_ptr().addr(); + match op { + BufferOp::Count => { + if self.counted_buffers.increment(addr) { + self.memory_usage += buffer.capacity(); + } + } + BufferOp::Uncount => { + if self.counted_buffers.decrement(addr) { + self.memory_usage = + self.memory_usage.saturating_sub(buffer.capacity()); + } + } } } - /// Count the memory usage of `array` and its children recursively. - fn count_array_memory_size(&mut self, array: &dyn Array) { + /// Walk `array`'s buffers recursively, applying `op` to each. + fn visit_array_buffers(&mut self, array: &dyn Array, op: BufferOp) { if let Some(nulls) = array.nulls() { - self.count_buffer_memory_size(nulls.buffer()); + self.apply_buffer_op(nulls.buffer(), op); } downcast_primitive_array! { - array => self.count_buffer_memory_size(array.values().inner()), + array => self.apply_buffer_op(array.values().inner(), op), DataType::Null => {} DataType::Boolean => { - self.count_buffer_memory_size(array.as_boolean().values().inner()); + self.apply_buffer_op(array.as_boolean().values().inner(), op); } DataType::Binary => { - self.count_byte_array_memory_size(array.as_binary::()); + self.visit_byte_array_buffers(array.as_binary::(), op); } DataType::LargeBinary => { - self.count_byte_array_memory_size(array.as_binary::()); + self.visit_byte_array_buffers(array.as_binary::(), op); } DataType::Utf8 => { - self.count_byte_array_memory_size(array.as_string::()); + self.visit_byte_array_buffers(array.as_string::(), op); } DataType::LargeUtf8 => { - self.count_byte_array_memory_size(array.as_string::()); + self.visit_byte_array_buffers(array.as_string::(), op); } DataType::BinaryView => { - self.count_byte_view_array_memory_size(array.as_binary_view()); + self.visit_byte_view_array_buffers(array.as_binary_view(), op); } DataType::Utf8View => { - self.count_byte_view_array_memory_size(array.as_string_view()); + self.visit_byte_view_array_buffers(array.as_string_view(), op); } DataType::FixedSizeBinary(_) => { - self.count_buffer_memory_size(array.as_fixed_size_binary().values()); + self.apply_buffer_op(array.as_fixed_size_binary().values(), op); } DataType::List(_) => { - self.count_list_array_memory_size(array.as_list::()); + self.visit_list_array_buffers(array.as_list::(), op); } DataType::LargeList(_) => { - self.count_list_array_memory_size(array.as_list::()); + self.visit_list_array_buffers(array.as_list::(), op); } DataType::ListView(_) => { - self.count_list_view_array_memory_size(array.as_list_view::()); + self.visit_list_view_array_buffers(array.as_list_view::(), op); } DataType::LargeListView(_) => { - self.count_list_view_array_memory_size(array.as_list_view::()); + self.visit_list_view_array_buffers(array.as_list_view::(), op); } DataType::FixedSizeList(_, _) => { - self.count_array_memory_size( + self.visit_array_buffers( array.as_fixed_size_list().values().as_ref(), + op, ); } DataType::Struct(_) => { for child in array.as_struct().columns() { - self.count_array_memory_size(child.as_ref()); + self.visit_array_buffers(child.as_ref(), op); } } DataType::Union(_, _) => { let array = array.as_union(); - self.count_buffer_memory_size(array.type_ids().inner()); + self.apply_buffer_op(array.type_ids().inner(), op); if let Some(offsets) = array.offsets() { - self.count_buffer_memory_size(offsets.inner()); + self.apply_buffer_op(offsets.inner(), op); } for (type_id, _) in array.fields().iter() { - self.count_array_memory_size(array.child(type_id).as_ref()); + self.visit_array_buffers(array.child(type_id).as_ref(), op); } } DataType::Dictionary(_, _) => { let array = array.as_any_dictionary(); - self.count_array_memory_size(array.keys()); - self.count_array_memory_size(array.values().as_ref()); + self.visit_array_buffers(array.keys(), op); + self.visit_array_buffers(array.values().as_ref(), op); } DataType::Map(_, _) => { let array = array.as_map(); - self.count_buffer_memory_size(array.offsets().inner().inner()); - self.count_array_memory_size(array.entries()); + self.apply_buffer_op(array.offsets().inner().inner(), op); + self.visit_array_buffers(array.entries(), op); } DataType::RunEndEncoded(run_ends, _) => match run_ends.data_type() { - DataType::Int16 => self.count_run_array_memory_size::(array), - DataType::Int32 => self.count_run_array_memory_size::(array), - DataType::Int64 => self.count_run_array_memory_size::(array), - // Arrow only permits Int16, Int32, and Int64 run-end indexes. A - // custom Array implementation may still expose malformed data; - // retain correct accounting for it without panicking. - _ => self.count_array_data_memory_size(&array.to_data()), + DataType::Int16 => self.visit_run_array_buffers::(array, op), + DataType::Int32 => self.visit_run_array_buffers::(array, op), + DataType::Int64 => self.visit_run_array_buffers::(array, op), + _ => self.visit_array_data_buffers(&array.to_data(), op), }, - // All currently supported non-primitive layouts are handled above. - // The Arrow macro requires a final arm for primitive variants that - // its nested dispatch has already consumed. Keep a safe generic - // fallback for custom or future Array implementations. - _ => self.count_array_data_memory_size(&array.to_data()), + _ => self.visit_array_data_buffers(&array.to_data(), op), } } - fn count_byte_array_memory_size( + fn visit_byte_array_buffers( &mut self, array: &GenericByteArray, + op: BufferOp, ) { - self.count_buffer_memory_size(array.offsets().inner().inner()); - self.count_buffer_memory_size(array.values()); + self.apply_buffer_op(array.offsets().inner().inner(), op); + self.apply_buffer_op(array.values(), op); } - fn count_byte_view_array_memory_size( + fn visit_byte_view_array_buffers( &mut self, array: &GenericByteViewArray, + op: BufferOp, ) { - self.count_buffer_memory_size(array.views().inner()); + self.apply_buffer_op(array.views().inner(), op); for buffer in array.data_buffers().iter() { - self.count_buffer_memory_size(buffer); + self.apply_buffer_op(buffer, op); } } - fn count_list_array_memory_size( + fn visit_list_array_buffers( &mut self, array: &GenericListArray, + op: BufferOp, ) { - self.count_buffer_memory_size(array.offsets().inner().inner()); - self.count_array_memory_size(array.values().as_ref()); + self.apply_buffer_op(array.offsets().inner().inner(), op); + self.visit_array_buffers(array.values().as_ref(), op); } - fn count_list_view_array_memory_size( + fn visit_list_view_array_buffers( &mut self, array: &GenericListViewArray, + op: BufferOp, ) { - self.count_buffer_memory_size(array.offsets().inner()); - self.count_buffer_memory_size(array.sizes().inner()); - self.count_array_memory_size(array.values().as_ref()); + self.apply_buffer_op(array.offsets().inner(), op); + self.apply_buffer_op(array.sizes().inner(), op); + self.visit_array_buffers(array.values().as_ref(), op); } - fn count_run_array_memory_size(&mut self, array: &dyn Array) { + fn visit_run_array_buffers( + &mut self, + array: &dyn Array, + op: BufferOp, + ) { if let Some(array) = array.as_any().downcast_ref::>() { - self.count_buffer_memory_size(array.run_ends().inner().inner()); - self.count_array_memory_size(array.values().as_ref()); + self.apply_buffer_op(array.run_ends().inner().inner(), op); + self.visit_array_buffers(array.values().as_ref(), op); } else { - // The DataType and concrete array implementation disagree. Use the - // generic representation rather than panic while accounting memory. - self.count_array_data_memory_size(&array.to_data()); + self.visit_array_data_buffers(&array.to_data(), op); } } - fn count_array_data_memory_size(&mut self, array_data: &arrow::array::ArrayData) { + fn visit_array_data_buffers( + &mut self, + array_data: &arrow::array::ArrayData, + op: BufferOp, + ) { for buffer in array_data.buffers() { - self.count_buffer_memory_size(buffer); + self.apply_buffer_op(buffer, op); } if let Some(nulls) = array_data.nulls() { - self.count_buffer_memory_size(nulls.buffer()); + self.apply_buffer_op(nulls.buffer(), op); } for child in array_data.child_data() { - self.count_array_data_memory_size(child); + self.visit_array_data_buffers(child, op); } } } +/// Direction of a buffer accounting operation. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum BufferOp { + /// Increment the reference count; add capacity on 0→1 transition. + Count, + /// Decrement the reference count; release capacity on 1→0 transition. + Uncount, +} + /// Counts the unique Array object memory retained by `array` and its children. fn count_unique_array_object_memory_size( array: &ArrayRef, @@ -410,6 +505,33 @@ fn count_unique_array_object_memory_size( .sum::() } +/// Inverse of [`count_unique_array_object_memory_size`]: removes array +/// identities from the tracked set and returns the released overhead. +fn uncount_unique_array_object_memory_size( + array: &ArrayRef, + counted_arrays: &mut HashSet, +) -> usize { + let array_ptr = Arc::as_ptr(array).cast::<()>() as usize; + if !counted_arrays.remove(&array_ptr) { + return 0; + } + + let children = array_children(array); + let children_overhead: usize = children + .iter() + .map(|child| child.get_array_memory_size() - child.get_buffer_memory_size()) + .sum(); + let own_overhead = array.get_array_memory_size() + - array.get_buffer_memory_size() + - children_overhead; + + own_overhead + + children + .into_iter() + .map(|child| uncount_unique_array_object_memory_size(child, counted_arrays)) + .sum::() +} + /// Returns the `ArrayRef` children whose object allocations may be shared. fn array_children(array: &ArrayRef) -> Vec<&ArrayRef> { match array.data_type() { @@ -464,46 +586,111 @@ fn array_children(array: &ArrayRef) -> Vec<&ArrayRef> { } } -/// Tracks a small number of buffers inline, avoiding a heap allocation for -/// typical batches, and promotes to a hash set when more buffers are seen. +/// Reference-counted buffer tracker with an inline fast path. +/// +/// Tracks buffer addresses with their reference counts. The first +/// [`INLINE_BUFFER_IDS`] distinct buffers are stored in a fixed-size inline +/// array to avoid heap allocation for typical batches. When more buffers are +/// encountered, the tracker promotes to a heap-allocated `HashMap`. #[derive(Debug)] -struct BufferIdSet { - inline: [Option>; INLINE_BUFFER_IDS], +struct BufferIdMap { + /// Inline storage: `(address, count)` pairs. + inline: [(NonZero, u32); INLINE_BUFFER_IDS], + /// Number of occupied inline slots. len: usize, - overflow: Option>>, + /// Overflow storage used when more than `INLINE_BUFFER_IDS` distinct + /// buffers are tracked. + overflow: Option, u32>>, } -impl Default for BufferIdSet { +impl Default for BufferIdMap { fn default() -> Self { + // SAFETY: NonZero has no validity invariant for zero-initialized + // memory because we gate access on `self.len`. We use a dummy NonZero + // value to satisfy the type system. Self { - inline: [None; INLINE_BUFFER_IDS], + inline: [(NonZero::::new(1).unwrap(), 0); INLINE_BUFFER_IDS], len: 0, overflow: None, } } } -impl BufferIdSet { - fn insert(&mut self, buffer_id: NonZero) -> bool { +impl BufferIdMap { + /// Increment the reference count for `buffer_id`. Returns `true` when the + /// count went from 0 to 1 (i.e., the buffer is newly tracked). + fn increment(&mut self, buffer_id: NonZero) -> bool { if let Some(overflow) = &mut self.overflow { - return overflow.insert(buffer_id); + let count = overflow.entry(buffer_id).or_insert(0); + *count += 1; + return *count == 1; } - if self.inline[..self.len].contains(&Some(buffer_id)) { - return false; + // Search inline slots. + for i in 0..self.len { + if self.inline[i].0 == buffer_id { + self.inline[i].1 += 1; + return false; // Already tracked, count > 1 now. + } } + // New buffer: try inline first. if self.len < INLINE_BUFFER_IDS { - self.inline[self.len] = Some(buffer_id); + self.inline[self.len] = (buffer_id, 1); self.len += 1; return true; } - let mut overflow = HashSet::with_capacity(INLINE_BUFFER_IDS + 1); - overflow.extend(self.inline.iter().flatten().copied()); - let inserted = overflow.insert(buffer_id); + // Promote to overflow. + let mut overflow = hashbrown::HashMap::with_capacity(INLINE_BUFFER_IDS + 1); + for i in 0..self.len { + overflow.insert(self.inline[i].0, self.inline[i].1); + } + overflow.insert(buffer_id, 1); self.overflow = Some(overflow); - inserted + true + } + + /// Decrement the reference count for `buffer_id`. Returns `true` when the + /// count reached zero (i.e., the buffer is no longer tracked). Returns + /// `false` if the buffer was not tracked or its count is still positive. + fn decrement(&mut self, buffer_id: NonZero) -> bool { + if let Some(overflow) = &mut self.overflow { + if let Some(count) = overflow.get_mut(&buffer_id) { + *count -= 1; + if *count == 0 { + overflow.remove(&buffer_id); + return true; + } + } + return false; + } + + // Search inline slots. + for i in 0..self.len { + if self.inline[i].0 == buffer_id { + self.inline[i].1 -= 1; + if self.inline[i].1 == 0 { + // Swap-remove: move last entry here. + self.len -= 1; + if i < self.len { + self.inline[i] = self.inline[self.len]; + } + return true; + } + return false; + } + } + + false + } + + /// Insert-only compatibility used by tests. Returns `true` if the buffer + /// was newly inserted. + #[cfg(test)] + #[expect(dead_code)] + fn insert(&mut self, buffer_id: NonZero) -> bool { + self.increment(buffer_id) } } @@ -585,7 +772,7 @@ mod record_batch_tests { fn assert_array_memory_size_matches(array: &dyn Array) { let mut counter = RecordBatchMemoryCounter::new(); - counter.count_array_memory_size(array); + counter.visit_array_buffers(array, BufferOp::Count); assert_eq!(counter.memory_usage(), array_data_memory_size(array)); } @@ -1103,4 +1290,920 @@ mod record_batch_tests { let size = get_record_batch_memory_size(&batch); assert_eq!(size, 8208); } + + // ---- uncount tests ---- + + #[test] + fn test_uncount_batch_round_trip() { + let batch = RecordBatch::try_from_iter(vec![( + "ints", + Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])) as ArrayRef, + )]) + .unwrap(); + + let mut counter = RecordBatchMemoryCounter::new(); + let counted = counter.count_batch(&batch); + assert!(counted > 0); + assert_eq!(counter.memory_usage(), counted); + + let released = counter.uncount_batch(&batch); + assert_eq!(released, counted); + assert_eq!(counter.memory_usage(), 0); + } + + #[test] + fn test_uncount_two_slices_example_from_issue() { + // The exact example from the issue description. + let array = Int32Array::from(vec![0; 1_000_000]); // 4 MB + let schema = + Arc::new(Schema::new(vec![Field::new("col", DataType::Int32, false)])); + let batch = RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap(); + let slice1 = batch.slice(0, 500_000); + let slice2 = batch.slice(500_000, 500_000); + + let buffer_size = get_record_batch_memory_size(&batch); + + let mut counter = RecordBatchMemoryCounter::new(); + + // count_batch(slice1) → buffer_size (parent buffer counted) + assert_eq!(counter.count_batch(&slice1), buffer_size); + assert_eq!(counter.memory_usage(), buffer_size); + + // count_batch(slice2) → 0 (already counted) + assert_eq!(counter.count_batch(&slice2), 0); + assert_eq!(counter.memory_usage(), buffer_size); + + // uncount_batch(slice1) → 0 released (slice2 still uses it) + assert_eq!(counter.uncount_batch(&slice1), 0); + assert_eq!(counter.memory_usage(), buffer_size); + + // uncount_batch(slice2) → buffer_size released + assert_eq!(counter.uncount_batch(&slice2), buffer_size); + assert_eq!(counter.memory_usage(), 0); + } + + #[test] + fn test_uncount_array_round_trip() { + let array = Int32Array::from(vec![1, 2, 3]); + let mut counter = RecordBatchMemoryCounter::new(); + + let counted = counter.count_array(&array); + assert!(counted > 0); + + let released = counter.uncount_array(&array); + assert_eq!(released, counted); + assert_eq!(counter.memory_usage(), 0); + } + + #[test] + fn test_uncount_batch_with_array_overhead_round_trip() { + let schema = Arc::new(Schema::new(vec![ + Field::new("ints", DataType::Int32, false), + Field::new("floats", DataType::Float64, false), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(Float64Array::from(vec![1., 2., 3.])), + ], + ) + .unwrap(); + + let mut counter = RecordBatchMemoryCounter::new(); + let counted = counter.count_batch_with_array_overhead(&batch); + assert!(counted > 0); + assert_eq!(counter.memory_usage(), counted); + + let released = counter.uncount_batch_with_array_overhead(&batch); + assert_eq!(released, counted); + assert_eq!(counter.memory_usage(), 0); + } + + #[test] + fn test_uncount_view_arrays_sharing_data_buffers() { + let array = StringViewArray::from_iter_values([ + "short", + "a payload longer than twelve bytes that goes to data buffers", + ]); + + let mut counter = RecordBatchMemoryCounter::new(); + let counted = counter.count_array(&array); + assert!(counted > 0); + + let released = counter.uncount_array(&array); + assert_eq!(released, counted); + assert_eq!(counter.memory_usage(), 0); + } + + #[test] + fn test_uncount_dictionaries_sharing_values() { + let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c"])); + let dict1 = DictionaryArray::::try_new( + Int32Array::from(vec![0, 1, 2]), + Arc::clone(&values), + ) + .unwrap(); + let dict2 = DictionaryArray::::try_new( + Int32Array::from(vec![2, 1, 0]), + Arc::clone(&values), + ) + .unwrap(); + + let mut counter = RecordBatchMemoryCounter::new(); + let counted1 = counter.count_array(&dict1); + let counted2 = counter.count_array(&dict2); + // dict2 shares the values buffer with dict1, so only keys are new. + assert!(counted2 < counted1); + + let total = counter.memory_usage(); + + // Uncount dict1: only its unique keys buffer should be freed. + let released1 = counter.uncount_array(&dict1); + assert_eq!(released1, counted2); // Only the keys that dict1 uniquely owns + // Shared values still counted. + assert_eq!(counter.memory_usage(), total - released1); + + // Uncount dict2: everything else freed. + let released2 = counter.uncount_array(&dict2); + assert_eq!(counter.memory_usage(), 0); + assert_eq!(released1 + released2, total); + } + + #[test] + fn test_uncount_nested_struct() { + let inner: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let fields = + Fields::from(vec![Arc::new(Field::new("v", DataType::Int32, false))]); + let outer = StructArray::new(fields, vec![inner], None); + + let mut counter = RecordBatchMemoryCounter::new(); + let counted = counter.count_array(&outer); + assert!(counted > 0); + + let released = counter.uncount_array(&outer); + assert_eq!(released, counted); + assert_eq!(counter.memory_usage(), 0); + } + + #[test] + fn test_uncount_with_overflow_promotion() { + // Create more than INLINE_BUFFER_IDS distinct buffers, then uncount them. + let n = INLINE_BUFFER_IDS + 4; + let arrays: Vec = (0..n) + .map(|i| Arc::new(Int32Array::from(vec![i as i32])) as ArrayRef) + .collect(); + + let mut counter = RecordBatchMemoryCounter::new(); + let mut total_counted = 0usize; + for array in &arrays { + total_counted += counter.count_array(array.as_ref()); + } + assert!(counter.counted_buffers.overflow.is_some()); + assert_eq!(counter.memory_usage(), total_counted); + + // Uncount all in reverse order. + let mut total_released = 0usize; + for array in arrays.iter().rev() { + total_released += counter.uncount_array(array.as_ref()); + } + assert_eq!(total_released, total_counted); + assert_eq!(counter.memory_usage(), 0); + } + + #[test] + fn test_uncount_never_counted_is_noop() { + let array = Int32Array::from(vec![1, 2, 3]); + let mut counter = RecordBatchMemoryCounter::new(); + + // Uncounting something never counted should return 0 and not panic. + assert_eq!(counter.uncount_array(&array), 0); + assert_eq!(counter.memory_usage(), 0); + } + + #[test] + fn test_randomized_count_uncount_sequence() { + // Reference model: track per-buffer refcount in a HashMap. + // Uses diverse array types and hundreds of operations. + use std::collections::HashMap; + + // Build a mix of array types to exercise different buffer layouts. + let arrays: Vec = vec![ + // Int32 (single values buffer) + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(Int64Array::from(vec![10, 20, 30])), + // String (offsets + values buffers) + Arc::new(StringArray::from(vec!["hello", "world", "foo"])), + // StringView (views + variadic data buffers) + Arc::new(StringViewArray::from_iter_values([ + "short", + "a string longer than twelve bytes that spills to data buffers", + "another long string that also spills to data buffers for sharing", + ])), + // Float64 + Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0])), + // Dictionary (keys + shared values) + { + let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c"])); + Arc::new( + DictionaryArray::::try_new( + Int32Array::from(vec![0, 1, 2]), + values, + ) + .unwrap(), + ) as ArrayRef + }, + // List (offsets + child values) + { + let child: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4])); + let field = Arc::new(Field::new_list_field(DataType::Int32, false)); + Arc::new(ListArray::new( + field, + OffsetBuffer::new(vec![0, 2, 4].into()), + child, + None, + )) as ArrayRef + }, + // Boolean (bit-packed values buffer) + Arc::new(arrow::array::BooleanArray::from(vec![true, false, true])), + ]; + + // Also add slices that share buffers with arrays[0] and arrays[2]. + let int_slice = arrays[0].slice(0, 2); + let str_slice = arrays[2].slice(1, 2); + + let all_arrays: Vec<&dyn Array> = arrays + .iter() + .map(|a| a.as_ref()) + .chain([int_slice.as_ref(), str_slice.as_ref()]) + .collect(); + + // Deterministic pseudo-random sequence using a simple LCG. + // 200 operations: mix of count and uncount. + let mut seed: u64 = 42; + let mut next_rand = || -> u64 { + seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1); + seed >> 33 + }; + + let mut counter = RecordBatchMemoryCounter::new(); + + // Reference model: map from buffer addr -> (refcount, capacity). + let mut ref_map: HashMap, (u32, usize)> = HashMap::new(); + let mut ref_memory: usize = 0; + + // Helper: walk an array's buffers and collect (addr, capacity) pairs. + fn collect_buffers(array: &dyn Array) -> Vec<(NonZero, usize)> { + let mut result = Vec::new(); + fn walk(array: &dyn Array, out: &mut Vec<(NonZero, usize)>) { + if let Some(nulls) = array.nulls() { + let buf = nulls.buffer(); + out.push((buf.data_ptr().addr(), buf.capacity())); + } + let data = array.to_data(); + for buf in data.buffers() { + out.push((buf.data_ptr().addr(), buf.capacity())); + } + for child in data.child_data() { + let child_array = arrow::array::make_array(child.clone()); + walk(child_array.as_ref(), out); + } + } + walk(array, &mut result); + result + } + + // Track per-array count so we don't uncount more than counted. + let mut array_live_counts = vec![0u32; all_arrays.len()]; + + for op_num in 0..200 { + let idx = (next_rand() as usize) % all_arrays.len(); + // Bias toward counting early, uncounting later. + let is_count = if array_live_counts[idx] == 0 { + true + } else if op_num < 80 { + next_rand() % 3 != 0 // 67% count early + } else { + next_rand() % 3 == 0 // 33% count later + }; + + let array = all_arrays[idx]; + let buffers = collect_buffers(array); + + if is_count { + counter.count_array(array); + array_live_counts[idx] += 1; + for (addr, cap) in &buffers { + let entry = ref_map.entry(*addr).or_insert((0, *cap)); + if entry.0 == 0 { + ref_memory += cap; + } + entry.0 += 1; + } + } else if array_live_counts[idx] > 0 { + counter.uncount_array(array); + array_live_counts[idx] -= 1; + for (addr, cap) in &buffers { + if let Some(entry) = ref_map.get_mut(addr) + && entry.0 > 0 + { + entry.0 -= 1; + if entry.0 == 0 { + ref_memory -= cap; + } + } + } + } + + assert_eq!( + counter.memory_usage(), + ref_memory, + "Mismatch at op #{op_num} (array {idx}, count={is_count})" + ); + } + + // Drain everything. + for (idx, array) in all_arrays.iter().enumerate() { + for _ in 0..array_live_counts[idx] { + counter.uncount_array(*array); + } + } + assert_eq!( + counter.memory_usage(), + 0, + "Memory not zero after full drain" + ); + } + + // --------------------------------------------------------------- + // Gap: view arrays sharing data buffers across slices (criterion 3) + // --------------------------------------------------------------- + + #[test] + fn test_uncount_view_array_slices_sharing_data_buffers() { + // Two slices of a StringViewArray share the same variadic data buffers. + let array = StringViewArray::from_iter_values([ + "short", + "a payload longer than twelve bytes that goes to data buffers", + "another long payload that also spills to a data buffer for sure", + "yet another long payload exceeding the twelve byte inline limit", + ]); + let slice1 = array.slice(0, 2); + let slice2 = array.slice(2, 2); + + let full_size = { + let mut c = RecordBatchMemoryCounter::new(); + c.count_array(&array) + }; + + let mut counter = RecordBatchMemoryCounter::new(); + let counted1 = counter.count_array(&slice1); + assert_eq!( + counted1, full_size, + "First slice should count all shared data buffers" + ); + + let counted2 = counter.count_array(&slice2); + assert_eq!(counted2, 0, "Second slice shares all buffers with first"); + assert_eq!(counter.memory_usage(), full_size); + + // Uncount first slice — buffers still held by slice2. + let released1 = counter.uncount_array(&slice1); + assert_eq!(released1, 0, "Shared buffers still referenced by slice2"); + assert_eq!(counter.memory_usage(), full_size); + + // Uncount second slice — everything freed. + let released2 = counter.uncount_array(&slice2); + assert_eq!(released2, full_size); + assert_eq!(counter.memory_usage(), 0); + } + + #[test] + fn test_uncount_binary_view_array_shared_data_buffers() { + // BinaryViewArray has the same data buffer sharing as StringViewArray. + let array = BinaryViewArray::from_iter_values([ + b"short".as_slice(), + b"a binary payload longer than twelve bytes for spillover", + b"another binary payload that also spills into the data buffer", + ]); + let slice1 = array.slice(0, 1); + let slice2 = array.slice(1, 2); + + let mut counter = RecordBatchMemoryCounter::new(); + let counted1 = counter.count_array(&slice1); + let counted2 = counter.count_array(&slice2); + let total = counted1 + counted2; + // Both slices together should equal counting the full array. + assert_eq!(counter.memory_usage(), { + let mut c = RecordBatchMemoryCounter::new(); + c.count_array(&array) + }); + + let released1 = counter.uncount_array(&slice1); + let released2 = counter.uncount_array(&slice2); + assert_eq!(released1 + released2, total); + assert_eq!(counter.memory_usage(), 0); + } + + // --------------------------------------------------------------- + // Gap: nested types beyond struct (criterion 5) + // --------------------------------------------------------------- + + #[test] + fn test_uncount_list_array_shared_child() { + let shared_child: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6])); + let field = Arc::new(Field::new_list_field(DataType::Int32, false)); + + let list1 = ListArray::new( + Arc::clone(&field), + OffsetBuffer::new(vec![0, 3].into()), + Arc::clone(&shared_child), + None, + ); + let list2 = ListArray::new( + field, + OffsetBuffer::new(vec![0, 3].into()), + Arc::clone(&shared_child), + None, + ); + + let mut counter = RecordBatchMemoryCounter::new(); + let counted1 = counter.count_array(&list1); + let counted2 = counter.count_array(&list2); + assert!( + counted2 < counted1, + "list2 shares child buffer, should add less" + ); + + let total = counter.memory_usage(); + let released1 = counter.uncount_array(&list1); + let released2 = counter.uncount_array(&list2); + assert_eq!(released1 + released2, total); + assert_eq!(counter.memory_usage(), 0); + } + + #[test] + fn test_uncount_map_array_shared_children() { + let shared_key: ArrayRef = Arc::new(Int32Array::from(vec![4, 5, 6])); + let shared_value: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let fields = Fields::from(vec![ + Arc::new(Field::new("key", DataType::Int32, false)), + Arc::new(Field::new("value", DataType::Int32, false)), + ]); + + let map1 = map_with_shared_children(&fields, &shared_key, &shared_value); + let map2 = map_with_shared_children(&fields, &shared_key, &shared_value); + + let mut counter = RecordBatchMemoryCounter::new(); + let counted1 = counter.count_array(map1.as_ref()); + let counted2 = counter.count_array(map2.as_ref()); + assert!(counted2 < counted1, "map2 shares key+value buffers"); + + let total = counter.memory_usage(); + let released1 = counter.uncount_array(map1.as_ref()); + let released2 = counter.uncount_array(map2.as_ref()); + assert_eq!(released1 + released2, total); + assert_eq!(counter.memory_usage(), 0); + } + + #[test] + fn test_uncount_union_array() { + let child: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let fields: UnionFields = + std::iter::once((0, Arc::new(Field::new("v", DataType::Int32, false)))) + .collect(); + let union_arr = UnionArray::try_new( + fields, + vec![0, 0, 0].into(), + Some(vec![0, 1, 2].into()), + vec![Arc::clone(&child)], + ) + .unwrap(); + + let mut counter = RecordBatchMemoryCounter::new(); + let counted = counter.count_array(&union_arr); + assert!(counted > 0); + + let released = counter.uncount_array(&union_arr); + assert_eq!(released, counted); + assert_eq!(counter.memory_usage(), 0); + } + + #[test] + fn test_uncount_union_slices_shared_child() { + let child: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let fields: UnionFields = + std::iter::once((0, Arc::new(Field::new("v", DataType::Int32, false)))) + .collect(); + let union_arr = Arc::new( + UnionArray::try_new( + fields, + vec![0, 0, 0].into(), + Some(vec![0, 1, 2].into()), + vec![Arc::clone(&child)], + ) + .unwrap(), + ) as ArrayRef; + let slice = union_arr.as_union().slice(1, 2); + + let mut counter = RecordBatchMemoryCounter::new(); + let counted1 = counter.count_array(union_arr.as_ref()); + let counted2 = counter.count_array(&slice); + // Slice shares the child buffer. + assert!(counted2 < counted1); + + let total = counter.memory_usage(); + let released1 = counter.uncount_array(union_arr.as_ref()); + let released2 = counter.uncount_array(&slice); + assert_eq!(released1 + released2, total); + assert_eq!(counter.memory_usage(), 0); + } + + #[test] + fn test_uncount_run_end_encoded_array() { + let values: ArrayRef = Arc::new(Int32Array::from(vec![10, 20, 30])); + let run_array = RunArray::::try_new( + &Int32Array::from(vec![1, 2, 3]), + values.as_ref(), + ) + .unwrap(); + + let mut counter = RecordBatchMemoryCounter::new(); + let counted = counter.count_array(&run_array); + assert!(counted > 0); + + let released = counter.uncount_array(&run_array); + assert_eq!(released, counted); + assert_eq!(counter.memory_usage(), 0); + } + + #[test] + fn test_uncount_run_end_encoded_slices_shared() { + let values: ArrayRef = Arc::new(Int32Array::from(vec![10, 20, 30])); + let run_array = Arc::new( + RunArray::::try_new( + &Int32Array::from(vec![1, 2, 3]), + values.as_ref(), + ) + .unwrap(), + ); + let slice = run_array + .as_any() + .downcast_ref::>() + .unwrap() + .slice(1, 2); + + let mut counter = RecordBatchMemoryCounter::new(); + let counted1 = counter.count_array(run_array.as_ref()); + let counted2 = counter.count_array(&slice); + assert!(counted2 < counted1, "Slice shares run-end + values buffers"); + + let total = counter.memory_usage(); + let released1 = counter.uncount_array(run_array.as_ref()); + let released2 = counter.uncount_array(&slice); + assert_eq!(released1 + released2, total); + assert_eq!(counter.memory_usage(), 0); + } + + // --------------------------------------------------------------- + // Gap: null bitmaps + // --------------------------------------------------------------- + + #[test] + fn test_uncount_array_with_nulls() { + // Arrays with null values have an extra validity bitmap buffer. + let array = Int32Array::from(vec![Some(1), None, Some(3), None, Some(5)]); + assert!(array.nulls().is_some(), "Array should have a null bitmap"); + + let mut counter = RecordBatchMemoryCounter::new(); + let counted = counter.count_array(&array); + assert!(counted > 0); + + let released = counter.uncount_array(&array); + assert_eq!(released, counted, "Should release values + null bitmap"); + assert_eq!(counter.memory_usage(), 0); + } + + #[test] + fn test_uncount_string_array_with_nulls() { + let array = StringArray::from(vec![Some("hello"), None, Some("world")]); + assert!(array.nulls().is_some()); + + let mut counter = RecordBatchMemoryCounter::new(); + let counted = counter.count_array(&array); + // String has: offsets + values + null bitmap = 3 buffers. + assert!(counted > 0); + + let released = counter.uncount_array(&array); + assert_eq!(released, counted); + assert_eq!(counter.memory_usage(), 0); + } + + // --------------------------------------------------------------- + // Gap: empty arrays + // --------------------------------------------------------------- + + #[test] + fn test_uncount_empty_array() { + let array = Int32Array::from(Vec::::new()); + assert_eq!(array.len(), 0); + + let mut counter = RecordBatchMemoryCounter::new(); + let counted = counter.count_array(&array); + // Even empty arrays may have allocated buffers with nonzero capacity. + let released = counter.uncount_array(&array); + assert_eq!(released, counted); + assert_eq!(counter.memory_usage(), 0); + } + + #[test] + fn test_uncount_empty_batch() { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Utf8, false), + ])); + let batch = RecordBatch::new_empty(schema); + assert_eq!(batch.num_rows(), 0); + + let mut counter = RecordBatchMemoryCounter::new(); + let counted = counter.count_batch(&batch); + let released = counter.uncount_batch(&batch); + assert_eq!(released, counted); + assert_eq!(counter.memory_usage(), 0); + } + + // --------------------------------------------------------------- + // Gap: Binary/Utf8 offset+data buffer pair + // --------------------------------------------------------------- + + #[test] + fn test_uncount_binary_array() { + let array = arrow::array::BinaryArray::from_iter_values([ + b"foo".as_slice(), + b"barbaz", + b"quxquuxcorge", + ]); + + let mut counter = RecordBatchMemoryCounter::new(); + let counted = counter.count_array(&array); + assert!(counted > 0); + + let released = counter.uncount_array(&array); + assert_eq!(released, counted); + assert_eq!(counter.memory_usage(), 0); + } + + #[test] + fn test_uncount_large_utf8_array() { + let array = + arrow::array::LargeStringArray::from(vec!["alpha", "bravo", "charlie"]); + + let mut counter = RecordBatchMemoryCounter::new(); + let counted = counter.count_array(&array); + assert!(counted > 0); + + let released = counter.uncount_array(&array); + assert_eq!(released, counted); + assert_eq!(counter.memory_usage(), 0); + } + + #[test] + fn test_uncount_string_slices_sharing_buffers() { + // Two slices of a StringArray share the offset and data buffers. + let array = StringArray::from(vec!["hello", "world", "foo", "bar"]); + let slice1 = array.slice(0, 2); + let slice2 = array.slice(2, 2); + + let full_size = { + let mut c = RecordBatchMemoryCounter::new(); + c.count_array(&array) + }; + + let mut counter = RecordBatchMemoryCounter::new(); + let counted1 = counter.count_array(&slice1); + assert_eq!(counted1, full_size, "First slice counts full buffers"); + + let counted2 = counter.count_array(&slice2); + assert_eq!(counted2, 0, "Second slice shares all buffers"); + + let released1 = counter.uncount_array(&slice1); + assert_eq!(released1, 0, "Buffers still referenced by slice2"); + + let released2 = counter.uncount_array(&slice2); + assert_eq!(released2, full_size); + assert_eq!(counter.memory_usage(), 0); + } + + // --------------------------------------------------------------- + // Gap: FixedSizeBinary + FixedSizeList + // --------------------------------------------------------------- + + #[test] + fn test_uncount_fixed_size_binary() { + let array = arrow::array::FixedSizeBinaryArray::try_from_iter( + vec![vec![1u8, 2, 3, 4], vec![5, 6, 7, 8], vec![9, 10, 11, 12]].into_iter(), + ) + .unwrap(); + + let mut counter = RecordBatchMemoryCounter::new(); + let counted = counter.count_array(&array); + assert!(counted > 0); + + let released = counter.uncount_array(&array); + assert_eq!(released, counted); + assert_eq!(counter.memory_usage(), 0); + } + + #[test] + fn test_uncount_fixed_size_list() { + let values = Int32Array::from(vec![1, 2, 3, 4, 5, 6]); + let field = Arc::new(Field::new("item", DataType::Int32, false)); + let array = + arrow::array::FixedSizeListArray::new(field, 3, Arc::new(values), None); + + let mut counter = RecordBatchMemoryCounter::new(); + let counted = counter.count_array(&array); + assert!(counted > 0); + + let released = counter.uncount_array(&array); + assert_eq!(released, counted); + assert_eq!(counter.memory_usage(), 0); + } + + // --------------------------------------------------------------- + // Gap: count_batch_with_array_overhead shared across batches + // --------------------------------------------------------------- + + #[test] + fn test_uncount_batch_with_array_overhead_shared_across_batches() { + // Two batches share the same ArrayRef -- array overhead should be + // counted once and released correctly. + let shared_col: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6])); + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int32, false)])); + let batch1 = + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::clone(&shared_col)]) + .unwrap() + .slice(0, 3); + let batch2 = + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::clone(&shared_col)]) + .unwrap() + .slice(3, 3); + + let mut counter = RecordBatchMemoryCounter::new(); + let counted1 = counter.count_batch_with_array_overhead(&batch1); + let _counted2 = counter.count_batch_with_array_overhead(&batch2); + assert!(counted1 > 0); + // batch2 shares the buffer, counted2 may only be array overhead for the slice. + + let total = counter.memory_usage(); + let released1 = counter.uncount_batch_with_array_overhead(&batch1); + let released2 = counter.uncount_batch_with_array_overhead(&batch2); + assert_eq!(released1 + released2, total); + assert_eq!(counter.memory_usage(), 0); + } + + // --------------------------------------------------------------- + // Gap: boolean array + // --------------------------------------------------------------- + + #[test] + fn test_uncount_boolean_array() { + let array = arrow::array::BooleanArray::from(vec![ + Some(true), + None, + Some(false), + Some(true), + ]); + + let mut counter = RecordBatchMemoryCounter::new(); + let counted = counter.count_array(&array); + assert!(counted > 0); + + let released = counter.uncount_array(&array); + assert_eq!(released, counted); + assert_eq!(counter.memory_usage(), 0); + } + + // --------------------------------------------------------------- + // Gap: deeply nested types (list of struct of dictionary) + // --------------------------------------------------------------- + + #[test] + fn test_uncount_deeply_nested() { + // List>> + let dict_values: ArrayRef = Arc::new(StringArray::from(vec!["x", "y", "z"])); + let dict_array = Arc::new( + DictionaryArray::::try_new( + Int32Array::from(vec![0, 1, 2]), + Arc::clone(&dict_values), + ) + .unwrap(), + ) as ArrayRef; + + let struct_fields = Fields::from(vec![Arc::new(Field::new( + "d", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + false, + ))]); + let struct_array = + StructArray::new(struct_fields.clone(), vec![dict_array], None); + + let list_field = + Arc::new(Field::new("item", DataType::Struct(struct_fields), false)); + let list_array = ListArray::new( + list_field, + OffsetBuffer::new(vec![0, 3].into()), + Arc::new(struct_array), + None, + ); + + let mut counter = RecordBatchMemoryCounter::new(); + let counted = counter.count_array(&list_array); + assert!(counted > 0); + + let released = counter.uncount_array(&list_array); + assert_eq!(released, counted); + assert_eq!(counter.memory_usage(), 0); + } + + // --------------------------------------------------------------- + // Gap: ListView and LargeListView + // --------------------------------------------------------------- + + #[test] + fn test_uncount_list_view_array() { + let values = Int32Array::from(vec![1, 2, 3, 4, 5]); + let list_view = ListViewArray::new( + Arc::new(Field::new("item", DataType::Int32, false)), + arrow::buffer::ScalarBuffer::from(vec![0i32, 2]), + arrow::buffer::ScalarBuffer::from(vec![3i32, 3]), + Arc::new(values), + None, + ); + + let mut counter = RecordBatchMemoryCounter::new(); + let counted = counter.count_array(&list_view); + assert!(counted > 0); + + let released = counter.uncount_array(&list_view); + assert_eq!(released, counted); + assert_eq!(counter.memory_usage(), 0); + } + + #[test] + fn test_uncount_large_list_view_array() { + let values = Int32Array::from(vec![10, 20, 30]); + let list_view = LargeListViewArray::new( + Arc::new(Field::new("item", DataType::Int32, false)), + arrow::buffer::ScalarBuffer::from(vec![0i64, 3]), + arrow::buffer::ScalarBuffer::from(vec![3i64, 0]), + Arc::new(values), + None, + ); + + let mut counter = RecordBatchMemoryCounter::new(); + let counted = counter.count_array(&list_view); + assert!(counted > 0); + + let released = counter.uncount_array(&list_view); + assert_eq!(released, counted); + assert_eq!(counter.memory_usage(), 0); + } + + // --------------------------------------------------------------- + // Gap: Null array type + // --------------------------------------------------------------- + + #[test] + fn test_uncount_null_array() { + let array = arrow::array::NullArray::new(10); + + let mut counter = RecordBatchMemoryCounter::new(); + let counted = counter.count_array(&array); + // Null arrays have no buffers -- should be 0 but not panic. + assert_eq!(counted, 0); + + let released = counter.uncount_array(&array); + assert_eq!(released, 0); + assert_eq!(counter.memory_usage(), 0); + } + + // --------------------------------------------------------------- + // Gap: multiple uncount of the same batch (idempotency / no underflow) + // --------------------------------------------------------------- + + #[test] + fn test_double_uncount_does_not_underflow() { + let array = Int32Array::from(vec![1, 2, 3]); + let mut counter = RecordBatchMemoryCounter::new(); + + counter.count_array(&array); + counter.uncount_array(&array); + assert_eq!(counter.memory_usage(), 0); + + // Second uncount of the same array should be a no-op, not underflow. + let released = counter.uncount_array(&array); + assert_eq!(released, 0); + assert_eq!(counter.memory_usage(), 0); + } }