This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new c7be84a18c perf: Reduce record batch memory accounting overhead
(#24319)
c7be84a18c is described below
commit c7be84a18c77b974fafc26f2cf7265526af4d18d
Author: Ryu <[email protected]>
AuthorDate: Mon Aug 17 19:29:25 2026 +0000
perf: Reduce record batch memory accounting overhead (#24319)
## Which issue does this PR close?
- Closes #24310.
## Rationale for this change
Record batch memory accounting runs on hot execution paths. The current
implementation materializes ArrayData for every array and allocates a
hash set even for small batches, adding measurable overhead to queries
that frequently update memory reservations.
## What changes are included in this PR?
- Traverse Arrow arrays directly and recursively count their backing
buffers without materializing ArrayData.
- Track the first 16 buffer identities inline, then promote to a hash
set for wider batches or counters spanning many batches.
- Preserve shared-buffer deduplication and full buffer-capacity
accounting semantics.
- Add parity coverage against the previous ArrayData traversal for
primitive, binary/view, list/view, fixed-size, struct, union,
dictionary, map, and all legal run-end index layouts.
- Add a focused Criterion benchmark across column counts, row counts,
and primitive/list/struct layouts.
Criterion point estimates from `cargo bench -p datafusion-common --bench
record_batch_memory`, measured sequentially on upstream `main` and this
PR on the same machine:
### Column count (8,192 rows, Int64)
| Columns | main | this PR | Speedup |
|---:|---:|---:|---:|
| 1 | 81.999 ns | 10.645 ns | 7.70x |
| 4 | 308.11 ns | 38.669 ns | 7.97x |
| 16 | 1.2543 us | 173.55 ns | 7.23x |
| 64 | 4.7457 us | 1.3914 us | 3.41x |
### Row count (4 Int64 columns)
| Rows | main | this PR | Speedup |
|---:|---:|---:|---:|
| 1 | 299.20 ns | 39.939 ns | 7.49x |
| 128 | 302.99 ns | 38.914 ns | 7.79x |
| 8,192 | 307.64 ns | 39.740 ns | 7.74x |
| 65,536 | 310.61 ns | 38.994 ns | 7.97x |
### Array layout (4 columns, 8,192 rows)
| Layout | main | this PR | Speedup |
|---|---:|---:|---:|
| Primitive Int64 | 311.01 ns | 40.140 ns | 7.75x |
| List of Int64 | 778.29 ns | 82.865 ns | 9.39x |
| Struct of two Int64 fields | 1.0072 us | 113.70 ns | 8.86x |
Each list row contains two Int64 values. Each struct column contains two
Int64 child fields.
## Are these changes tested?
Yes.
- `cargo fmt --all -- --check`
- `cargo clippy --all-targets --all-features -- -D warnings`
- `cargo test -p datafusion-common utils::memory --lib` (10 passed)
- `RUST_BACKTRACE=1 cargo test --profile ci --exclude
datafusion-examples --exclude datafusion-benchmarks --exclude
datafusion-cli --workspace --lib --tests --bins --features
avro,json,backtrace,extended_tests,recursive_protection,parquet_encryption`
- Focused Criterion comparison shown above
## Are there any user-facing changes?
No API or behavior changes. This reduces CPU and allocation overhead in
record batch memory accounting.
AI assistance: OpenAI Codex assisted with implementation and test
execution. I reviewed the change and its behavior end to end.
---
datafusion/common/Cargo.toml | 4 +
datafusion/common/benches/record_batch_memory.rs | 190 +++++++++++
datafusion/common/src/utils/memory.rs | 415 +++++++++++++++++++++--
3 files changed, 577 insertions(+), 32 deletions(-)
diff --git a/datafusion/common/Cargo.toml b/datafusion/common/Cargo.toml
index 1eb23089a4..9ee199fe82 100644
--- a/datafusion/common/Cargo.toml
+++ b/datafusion/common/Cargo.toml
@@ -64,6 +64,10 @@ name = "scalar_to_array"
harness = false
name = "stats_merge"
+[[bench]]
+harness = false
+name = "record_batch_memory"
+
[dependencies]
arrow = { workspace = true }
arrow-ipc = { workspace = true }
diff --git a/datafusion/common/benches/record_batch_memory.rs
b/datafusion/common/benches/record_batch_memory.rs
new file mode 100644
index 0000000000..2479d6ac98
--- /dev/null
+++ b/datafusion/common/benches/record_batch_memory.rs
@@ -0,0 +1,190 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Measures the CPU overhead of accounting for the backing buffers retained by
+//! [`RecordBatch`]es. Batch construction is intentionally outside the timed
+//! region so the benchmarks isolate buffer traversal and identity
deduplication.
+
+use std::hint::black_box;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, Int64Array, ListArray, StructArray};
+use arrow::datatypes::{DataType, Field, Int64Type, Schema};
+use arrow::record_batch::RecordBatch;
+use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
+use datafusion_common::utils::memory::{
+ RecordBatchMemoryCounter, get_record_batch_memory_size,
+};
+
+fn make_batch(columns: Vec<ArrayRef>) -> RecordBatch {
+ let fields = columns
+ .iter()
+ .enumerate()
+ .map(|(index, column)| {
+ Field::new(format!("col_{index}"), column.data_type().clone(),
false)
+ })
+ .collect::<Vec<_>>();
+
+ RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap()
+}
+
+fn make_primitive_batch(num_rows: usize, num_columns: usize) -> RecordBatch {
+ let columns = (0..num_columns)
+ .map(|index| {
+ Arc::new(Int64Array::from_iter_values(
+ (0..num_rows).map(|value| value as i64 + index as i64),
+ )) as ArrayRef
+ })
+ .collect::<Vec<_>>();
+
+ make_batch(columns)
+}
+
+fn make_list_batch(num_rows: usize, num_columns: usize) -> RecordBatch {
+ let columns = (0..num_columns)
+ .map(|column| {
+ Arc::new(ListArray::from_iter_primitive::<Int64Type, _, _>(
+ (0..num_rows).map(|row| {
+ let value = row as i64 + column as i64;
+ Some(vec![Some(value), Some(value + 1)])
+ }),
+ )) as ArrayRef
+ })
+ .collect::<Vec<_>>();
+
+ make_batch(columns)
+}
+
+fn make_struct_batch(num_rows: usize, num_columns: usize) -> RecordBatch {
+ let columns = (0..num_columns)
+ .map(|column| {
+ let left = Arc::new(Int64Array::from_iter_values(
+ (0..num_rows).map(|row| row as i64 + column as i64),
+ )) as ArrayRef;
+ let right = Arc::new(Int64Array::from_iter_values(
+ (0..num_rows).map(|row| row as i64 - column as i64),
+ )) as ArrayRef;
+
+ Arc::new(StructArray::from(vec![
+ (Arc::new(Field::new("left", DataType::Int64, false)), left),
+ (Arc::new(Field::new("right", DataType::Int64, false)), right),
+ ])) as ArrayRef
+ })
+ .collect::<Vec<_>>();
+
+ make_batch(columns)
+}
+
+fn benchmark_column_count(c: &mut Criterion) {
+ let mut group = c.benchmark_group("record_batch_memory_size/column_count");
+
+ // Each primitive column contributes a distinct backing buffer, exercising
+ // both the inline buffer-ID path and hash-set promotion.
+ for num_columns in [1, 4, 16, 64] {
+ let batch = make_primitive_batch(8192, num_columns);
+ group.bench_with_input(
+ BenchmarkId::from_parameter(num_columns),
+ &batch,
+ |bencher, batch| {
+ bencher.iter(||
get_record_batch_memory_size(black_box(batch)));
+ },
+ );
+ }
+
+ group.finish();
+}
+
+fn benchmark_row_count(c: &mut Criterion) {
+ let mut group = c.benchmark_group("record_batch_memory_size/row_count");
+
+ // Buffer traversal should depend on the number of buffers, not the number
+ // of values stored in each buffer.
+ for num_rows in [1, 128, 8192, 65_536] {
+ let batch = make_primitive_batch(num_rows, 4);
+ group.bench_with_input(
+ BenchmarkId::from_parameter(num_rows),
+ &batch,
+ |bencher, batch| {
+ bencher.iter(||
get_record_batch_memory_size(black_box(batch)));
+ },
+ );
+ }
+
+ group.finish();
+}
+
+fn benchmark_array_layout(c: &mut Criterion) {
+ let mut group = c.benchmark_group("record_batch_memory_size/array_layout");
+
+ // Compare direct primitive-buffer accounting with recursive traversal of
+ // representative nested layouts.
+ for (name, batch) in [
+ ("primitive", make_primitive_batch(8192, 4)),
+ ("list", make_list_batch(8192, 4)),
+ ("struct", make_struct_batch(8192, 4)),
+ ] {
+ group.bench_with_input(
+ BenchmarkId::from_parameter(name),
+ &batch,
+ |bencher, batch| {
+ bencher.iter(||
get_record_batch_memory_size(black_box(batch)));
+ },
+ );
+ }
+
+ group.finish();
+}
+
+fn benchmark_shared_slices(c: &mut Criterion) {
+ let mut group =
c.benchmark_group("record_batch_memory_size/shared_slices");
+
+ // Model the hash-join build-side workload: one counter is reused across a
+ // sequence of zero-copy batch slices that retain the same backing buffers.
+ // Slicing happens outside the timed region; the benchmark measures
repeated
+ // identity lookups and the one-time accounting of each shared buffer.
+ 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::<Vec<_>>();
+
+ 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)));
+ }
+ black_box(counter.memory_usage())
+ });
+ },
+ );
+ }
+
+ group.finish();
+}
+
+criterion_group!(
+ benches,
+ benchmark_column_count,
+ benchmark_row_count,
+ benchmark_array_layout,
+ benchmark_shared_slices
+);
+criterion_main!(benches);
diff --git a/datafusion/common/src/utils/memory.rs
b/datafusion/common/src/utils/memory.rs
index 21c084119e..fd405e06a2 100644
--- a/datafusion/common/src/utils/memory.rs
+++ b/datafusion/common/src/utils/memory.rs
@@ -19,11 +19,24 @@
use crate::error::_exec_datafusion_err;
use crate::{HashSet, Result};
-use arrow::array::ArrayData;
+use arrow::array::types::{ByteArrayType, ByteViewType, RunEndIndexType};
+use arrow::array::{
+ Array, AsArray, GenericByteArray, GenericByteViewArray, GenericListArray,
+ GenericListViewArray, RunArray,
+};
+use arrow::buffer::Buffer;
+use arrow::datatypes::DataType;
+use arrow::downcast_primitive_array;
use arrow::record_batch::RecordBatch;
use std::mem::size_of;
use std::num::NonZero;
+/// 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.
+/// This is a performance heuristic, not a semantic limit.
+const INLINE_BUFFER_IDS: usize = 16;
+
/// Estimates the memory size required for a hash table prior to allocation.
///
/// # Parameters
@@ -151,7 +164,7 @@ pub fn get_record_batch_memory_size(batch: &RecordBatch) ->
usize {
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: HashSet<NonZero<usize>>,
+ counted_buffers: BufferIdSet,
/// Total memory of all unique buffers counted so far
memory_usage: usize,
}
@@ -164,49 +177,229 @@ impl RecordBatchMemoryCounter {
/// Count `batch`, returning the memory used by its buffers that have not
/// been counted before.
pub fn count_batch(&mut self, batch: &RecordBatch) -> usize {
- let mut total_size = 0;
+ let previous_memory_usage = self.memory_usage;
for array in batch.columns() {
- let array_data = array.to_data();
- count_array_data_memory_size(
- &array_data,
- &mut self.counted_buffers,
- &mut total_size,
- );
+ self.count_array_memory_size(array.as_ref());
}
- self.memory_usage += total_size;
- total_size
+ self.memory_usage - previous_memory_usage
}
/// Total memory of the unique buffers of all batches counted so far.
pub fn memory_usage(&self) -> usize {
self.memory_usage
}
-}
-/// Count the memory usage of `array_data` and its children recursively.
-fn count_array_data_memory_size(
- array_data: &ArrayData,
- counted_buffers: &mut HashSet<NonZero<usize>>,
- total_size: &mut usize,
-) {
- // Count memory usage for `array_data`
- for buffer in array_data.buffers() {
- if counted_buffers.insert(buffer.data_ptr().addr()) {
- *total_size += buffer.capacity();
- } // Otherwise the buffer's memory is already counted
+ fn count_buffer_memory_size(&mut self, buffer: &Buffer) {
+ if self.counted_buffers.insert(buffer.data_ptr().addr()) {
+ self.memory_usage += buffer.capacity();
+ }
}
- if let Some(null_buffer) = array_data.nulls()
- &&
counted_buffers.insert(null_buffer.inner().inner().data_ptr().addr())
- {
- *total_size += null_buffer.inner().inner().capacity();
+ /// Count the memory usage of `array` and its children recursively.
+ fn count_array_memory_size(&mut self, array: &dyn Array) {
+ if let Some(nulls) = array.nulls() {
+ self.count_buffer_memory_size(nulls.buffer());
+ }
+
+ downcast_primitive_array! {
+ array => self.count_buffer_memory_size(array.values().inner()),
+ DataType::Null => {}
+ DataType::Boolean => {
+
self.count_buffer_memory_size(array.as_boolean().values().inner());
+ }
+ DataType::Binary => {
+ self.count_byte_array_memory_size(array.as_binary::<i32>());
+ }
+ DataType::LargeBinary => {
+ self.count_byte_array_memory_size(array.as_binary::<i64>());
+ }
+ DataType::Utf8 => {
+ self.count_byte_array_memory_size(array.as_string::<i32>());
+ }
+ DataType::LargeUtf8 => {
+ self.count_byte_array_memory_size(array.as_string::<i64>());
+ }
+ DataType::BinaryView => {
+ self.count_byte_view_array_memory_size(array.as_binary_view());
+ }
+ DataType::Utf8View => {
+ self.count_byte_view_array_memory_size(array.as_string_view());
+ }
+ DataType::FixedSizeBinary(_) => {
+
self.count_buffer_memory_size(array.as_fixed_size_binary().values());
+ }
+ DataType::List(_) => {
+ self.count_list_array_memory_size(array.as_list::<i32>());
+ }
+ DataType::LargeList(_) => {
+ self.count_list_array_memory_size(array.as_list::<i64>());
+ }
+ DataType::ListView(_) => {
+
self.count_list_view_array_memory_size(array.as_list_view::<i32>());
+ }
+ DataType::LargeListView(_) => {
+
self.count_list_view_array_memory_size(array.as_list_view::<i64>());
+ }
+ DataType::FixedSizeList(_, _) => {
+ self.count_array_memory_size(
+ array.as_fixed_size_list().values().as_ref(),
+ );
+ }
+ DataType::Struct(_) => {
+ for child in array.as_struct().columns() {
+ self.count_array_memory_size(child.as_ref());
+ }
+ }
+ DataType::Union(_, _) => {
+ let array = array.as_union();
+ self.count_buffer_memory_size(array.type_ids().inner());
+ if let Some(offsets) = array.offsets() {
+ self.count_buffer_memory_size(offsets.inner());
+ }
+ for (type_id, _) in array.fields().iter() {
+
self.count_array_memory_size(array.child(type_id).as_ref());
+ }
+ }
+ DataType::Dictionary(_, _) => {
+ let array = array.as_any_dictionary();
+ self.count_array_memory_size(array.keys());
+ self.count_array_memory_size(array.values().as_ref());
+ }
+ DataType::Map(_, _) => {
+ let array = array.as_map();
+ self.count_buffer_memory_size(array.offsets().inner().inner());
+ self.count_array_memory_size(array.entries());
+ }
+ DataType::RunEndEncoded(run_ends, _) => match run_ends.data_type()
{
+ DataType::Int16 => {
+
self.count_run_array_memory_size::<arrow::datatypes::Int16Type>(
+ array,
+ );
+ }
+ DataType::Int32 => {
+
self.count_run_array_memory_size::<arrow::datatypes::Int32Type>(
+ array,
+ );
+ }
+ DataType::Int64 => {
+
self.count_run_array_memory_size::<arrow::datatypes::Int64Type>(
+ 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()),
+ },
+ // 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()),
+ }
}
- // Count all children `ArrayData` recursively
- for child in array_data.child_data() {
- count_array_data_memory_size(child, counted_buffers, total_size);
+ fn count_byte_array_memory_size<T: ByteArrayType>(
+ &mut self,
+ array: &GenericByteArray<T>,
+ ) {
+ self.count_buffer_memory_size(array.offsets().inner().inner());
+ self.count_buffer_memory_size(array.values());
+ }
+
+ fn count_byte_view_array_memory_size<T: ByteViewType>(
+ &mut self,
+ array: &GenericByteViewArray<T>,
+ ) {
+ self.count_buffer_memory_size(array.views().inner());
+ for buffer in array.data_buffers() {
+ self.count_buffer_memory_size(buffer);
+ }
+ }
+
+ fn count_list_array_memory_size<O: arrow::array::OffsetSizeTrait>(
+ &mut self,
+ array: &GenericListArray<O>,
+ ) {
+ self.count_buffer_memory_size(array.offsets().inner().inner());
+ self.count_array_memory_size(array.values().as_ref());
+ }
+
+ fn count_list_view_array_memory_size<O: arrow::array::OffsetSizeTrait>(
+ &mut self,
+ array: &GenericListViewArray<O>,
+ ) {
+ 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());
+ }
+
+ fn count_run_array_memory_size<R: RunEndIndexType>(&mut self, array: &dyn
Array) {
+ if let Some(array) = array.as_any().downcast_ref::<RunArray<R>>() {
+ self.count_buffer_memory_size(array.run_ends().inner().inner());
+ self.count_array_memory_size(array.values().as_ref());
+ } 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());
+ }
+ }
+
+ fn count_array_data_memory_size(&mut self, array_data:
&arrow::array::ArrayData) {
+ for buffer in array_data.buffers() {
+ self.count_buffer_memory_size(buffer);
+ }
+ if let Some(nulls) = array_data.nulls() {
+ self.count_buffer_memory_size(nulls.buffer());
+ }
+ for child in array_data.child_data() {
+ self.count_array_data_memory_size(child);
+ }
+ }
+}
+
+/// 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.
+#[derive(Debug)]
+struct BufferIdSet {
+ inline: [Option<NonZero<usize>>; INLINE_BUFFER_IDS],
+ len: usize,
+ overflow: Option<HashSet<NonZero<usize>>>,
+}
+
+impl Default for BufferIdSet {
+ fn default() -> Self {
+ Self {
+ inline: [None; INLINE_BUFFER_IDS],
+ len: 0,
+ overflow: None,
+ }
+ }
+}
+
+impl BufferIdSet {
+ fn insert(&mut self, buffer_id: NonZero<usize>) -> bool {
+ if let Some(overflow) = &mut self.overflow {
+ return overflow.insert(buffer_id);
+ }
+
+ if self.inline[..self.len].contains(&Some(buffer_id)) {
+ return false;
+ }
+
+ if self.len < INLINE_BUFFER_IDS {
+ self.inline[self.len] = Some(buffer_id);
+ 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);
+ self.overflow = Some(overflow);
+ inserted
}
}
@@ -247,10 +440,49 @@ mod tests {
#[cfg(test)]
mod record_batch_tests {
use super::*;
- use arrow::array::{Float64Array, Int32Array, ListArray};
- use arrow::datatypes::{DataType, Field, Int32Type, Schema};
+ use arrow::array::{
+ ArrayData, ArrayRef, BinaryViewArray, Float64Array, Int16Array,
Int32Array,
+ Int64Array, LargeListViewArray, ListArray, ListViewArray, RunArray,
StringArray,
+ StringViewArray, new_null_array,
+ };
+ use arrow::datatypes::{
+ DataType, Field, Int16Type, Int32Type, Int64Type, Schema, UnionFields,
UnionMode,
+ };
use std::sync::Arc;
+ fn array_data_memory_size(array: &dyn Array) -> usize {
+ fn count(
+ array_data: &ArrayData,
+ counted_buffers: &mut HashSet<NonZero<usize>>,
+ total_size: &mut usize,
+ ) {
+ for buffer in array_data.buffers() {
+ if counted_buffers.insert(buffer.data_ptr().addr()) {
+ *total_size += buffer.capacity();
+ }
+ }
+ if let Some(nulls) = array_data.nulls() {
+ let buffer = nulls.inner().inner();
+ if counted_buffers.insert(buffer.data_ptr().addr()) {
+ *total_size += buffer.capacity();
+ }
+ }
+ for child in array_data.child_data() {
+ count(child, counted_buffers, total_size);
+ }
+ }
+
+ let mut total_size = 0;
+ count(&array.to_data(), &mut HashSet::default(), &mut total_size);
+ total_size
+ }
+
+ fn assert_array_memory_size_matches(array: &dyn Array) {
+ let mut counter = RecordBatchMemoryCounter::new();
+ counter.count_array_memory_size(array);
+ assert_eq!(counter.memory_usage(), array_data_memory_size(array));
+ }
+
#[test]
fn test_get_record_batch_memory_size() {
let schema = Arc::new(Schema::new(vec![
@@ -359,6 +591,125 @@ mod record_batch_tests {
assert_eq!(counter.memory_usage(),
get_record_batch_memory_size(&batch));
}
+ #[test]
+ fn test_record_batch_memory_counter_promotes_buffer_set() {
+ let fields = (0..=INLINE_BUFFER_IDS)
+ .map(|index| Field::new(format!("col_{index}"), DataType::Int32,
false))
+ .collect::<Vec<_>>();
+ let columns = (0..=INLINE_BUFFER_IDS)
+ .map(|value| Arc::new(Int32Array::from(vec![value as i32])) as _)
+ .collect::<Vec<_>>();
+ let batch = RecordBatch::try_new(Arc::new(Schema::new(fields)),
columns).unwrap();
+
+ let mut counter = RecordBatchMemoryCounter::new();
+ assert_eq!(
+ counter.count_batch(&batch),
+ (INLINE_BUFFER_IDS + 1) * size_of::<i32>()
+ );
+ assert!(counter.counted_buffers.overflow.is_some());
+ assert_eq!(counter.count_batch(&batch), 0);
+ }
+
+ #[test]
+ fn test_array_memory_size_matches_array_data_layouts() {
+ let list_field = Arc::new(Field::new_list_field(DataType::Int32,
true));
+ let struct_fields = vec![Field::new("value", DataType::Int32,
true)].into();
+ let union_fields = UnionFields::try_new(
+ vec![0],
+ vec![Field::new("value", DataType::Int32, true)],
+ )
+ .unwrap();
+ let map_entries = Arc::new(Field::new(
+ "entries",
+ DataType::Struct(
+ vec![
+ Field::new("key", DataType::Utf8, false),
+ Field::new("value", DataType::Int32, true),
+ ]
+ .into(),
+ ),
+ false,
+ ));
+ let data_types = vec![
+ DataType::Boolean,
+ DataType::Int32,
+ DataType::Binary,
+ DataType::LargeBinary,
+ DataType::FixedSizeBinary(4),
+ DataType::BinaryView,
+ DataType::Utf8,
+ DataType::LargeUtf8,
+ DataType::Utf8View,
+ DataType::List(Arc::clone(&list_field)),
+ DataType::LargeList(Arc::clone(&list_field)),
+ DataType::ListView(Arc::clone(&list_field)),
+ DataType::LargeListView(Arc::clone(&list_field)),
+ DataType::FixedSizeList(Arc::clone(&list_field), 2),
+ DataType::Struct(struct_fields),
+ DataType::Union(union_fields, UnionMode::Dense),
+ DataType::Dictionary(Box::new(DataType::Int32),
Box::new(DataType::Utf8)),
+ DataType::Map(map_entries, false),
+ ];
+
+ for data_type in data_types {
+ let array = new_null_array(&data_type, 3);
+ assert_array_memory_size_matches(array.as_ref());
+ }
+
+ // Exercise the view-specific buffers with concrete, non-empty values.
+ let view_arrays = [
+ Arc::new(BinaryViewArray::from_iter_values([
+ b"short".as_slice(),
+ b"a payload longer than twelve bytes".as_slice(),
+ ])) as ArrayRef,
+ Arc::new(StringViewArray::from_iter_values([
+ "short",
+ "a payload longer than twelve bytes",
+ ])) as ArrayRef,
+ Arc::new(ListViewArray::from_iter_primitive::<Int32Type, _, _>([
+ Some(vec![Some(1), Some(2)]),
+ None,
+ Some(vec![Some(3)]),
+ ])) as ArrayRef,
+ Arc::new(LargeListViewArray::from_iter_primitive::<Int32Type, _,
_>(
+ [Some(vec![Some(1), Some(2)]), None, Some(vec![Some(3)])],
+ )) as ArrayRef,
+ ];
+
+ for array in view_arrays {
+ assert_array_memory_size_matches(array.as_ref());
+ }
+
+ let run_values = StringArray::from(vec!["alpha", "beta"]);
+ let run_arrays = [
+ Arc::new(
+ RunArray::<Int16Type>::try_new(
+ &Int16Array::from(vec![2_i16, 5]),
+ &run_values,
+ )
+ .unwrap(),
+ ) as ArrayRef,
+ Arc::new(
+ RunArray::<Int32Type>::try_new(
+ &Int32Array::from(vec![2_i32, 5]),
+ &run_values,
+ )
+ .unwrap(),
+ ) as ArrayRef,
+ Arc::new(
+ RunArray::<Int64Type>::try_new(
+ &Int64Array::from(vec![2_i64, 5]),
+ &run_values,
+ )
+ .unwrap(),
+ ) as ArrayRef,
+ ];
+
+ for array in run_arrays {
+ assert_array_memory_size_matches(array.as_ref());
+ }
+ }
+
#[test]
fn test_get_record_batch_memory_size_nested_array() {
let schema = Arc::new(Schema::new(vec![
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
