Skip to content

Commit db5cfd0

Browse files
author
Xiaowei Tan
committed
feat: support releasing batches in RecordBatchMemoryCounter
Add uncount_batch, uncount_array, and uncount_batch_with_array_overhead methods to RecordBatchMemoryCounter. The internal BufferIdSet is changed to a reference-counted BufferIdMap so that a buffer's capacity is subtracted from memory_usage only when its refcount drops to zero. The per-type buffer walk is shared between counting and uncounting via a visit_array_buffers visitor, avoiding code duplication. Comprehensive test suite covers all acceptance criteria from #26140: - Count/uncount round trip (batch, array, batch_with_array_overhead) - Two-slice shared-buffer example from the issue - View arrays sharing data buffers (StringView slices, BinaryView) - Dictionaries sharing values across arrays - All nested types: List, Map, Union, RunEndEncoded, FixedSizeList, FixedSizeBinary, deeply nested (List<Struct<Dict>>) - Slice-sharing for Union and RunEndEncoded - Null bitmaps, empty arrays, empty batches, NullArray - Binary/LargeBinary/Utf8/LargeUtf8 offset+data buffer pairs - String slices sharing offset+data buffers - ListView and LargeListView - Boolean arrays (bit-packed) - >16 distinct buffers with inline-to-hash promotion + removals - count_batch_with_array_overhead shared across batches - Double-uncount idempotency (no underflow) - Randomized 200-op sequence across 10 diverse array types (Int32, Int64, String, StringView, Float64, Dictionary, List, Boolean, plus slices) verified against a reference model - record_batch_memory benchmark with count_uncount case, no regression Closes #26140
1 parent a23b89f commit db5cfd0

2 files changed

Lines changed: 1229 additions & 89 deletions

File tree

‎datafusion/common/benches/record_batch_memory.rs‎

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -180,11 +180,44 @@ fn benchmark_shared_slices(c: &mut Criterion) {
180180
group.finish();
181181
}
182182

183+
fn benchmark_count_uncount(c: &mut Criterion) {
184+
let mut group = c.benchmark_group("record_batch_memory_size/count_uncount");
185+
186+
// Model the sort/window workload: batches are counted as they arrive and
187+
// uncounted when they are dropped. Measure the full cycle.
188+
for num_columns in [4, 16, 64] {
189+
let batch = make_primitive_batch(8192, num_columns);
190+
let slices = (0..32)
191+
.map(|index| batch.slice(index * 256, 256))
192+
.collect::<Vec<_>>();
193+
194+
group.bench_with_input(
195+
BenchmarkId::from_parameter(num_columns),
196+
&slices,
197+
|bencher, slices| {
198+
bencher.iter(|| {
199+
let mut counter = RecordBatchMemoryCounter::new();
200+
for batch in black_box(slices) {
201+
black_box(counter.count_batch(black_box(batch)));
202+
}
203+
for batch in black_box(slices) {
204+
black_box(counter.uncount_batch(black_box(batch)));
205+
}
206+
assert_eq!(counter.memory_usage(), 0);
207+
});
208+
},
209+
);
210+
}
211+
212+
group.finish();
213+
}
214+
183215
criterion_group!(
184216
benches,
185217
benchmark_column_count,
186218
benchmark_row_count,
187219
benchmark_array_layout,
188-
benchmark_shared_slices
220+
benchmark_shared_slices,
221+
benchmark_count_uncount
189222
);
190223
criterion_main!(benches);

0 commit comments

Comments
 (0)