This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-24319-bb038a6739091567fe32c196d8c568a98e4f42e2 in repository https://gitbox.apache.org/repos/asf/datafusion.git
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]
