ryux1 commented on code in PR #24319: URL: https://github.com/apache/datafusion/pull/24319#discussion_r3789485421
########## datafusion/common/benches/record_batch_memory.rs: ########## @@ -0,0 +1,146 @@ +// 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. + +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::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"); + + 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"); + + 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"); + + 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(); +} + +criterion_group!( Review Comment: Added benchmark-level and group-specific documentation in b6414cd88. It now states that batch construction is outside the timed region and explains how the column-count, row-count, layout, and shared-slice groups isolate buffer traversal and identity-deduplication overhead. ########## datafusion/common/src/utils/memory.rs: ########## @@ -185,31 +194,225 @@ impl RecordBatchMemoryCounter { } } -/// 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>>, +/// 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 + } +} + +fn count_buffer_memory_size( Review Comment: Applied in b6414cd88: buffer counting and recursive array traversal are now private methods on RecordBatchMemoryCounter and update the counter state directly. ########## datafusion/common/src/utils/memory.rs: ########## @@ -185,31 +194,225 @@ impl RecordBatchMemoryCounter { } } -/// 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>>, +/// 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 + } +} + +fn count_buffer_memory_size( + buffer: &Buffer, + counted_buffers: &mut BufferIdSet, 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 + if counted_buffers.insert(buffer.data_ptr().addr()) { + *total_size += 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( + array: &dyn Array, + counted_buffers: &mut BufferIdSet, + total_size: &mut usize, +) { + if let Some(nulls) = array.nulls() { + count_buffer_memory_size(nulls.buffer(), counted_buffers, total_size); } - // Count all children `ArrayData` recursively - for child in array_data.child_data() { - count_array_data_memory_size(child, counted_buffers, total_size); + downcast_primitive_array! { + array => count_buffer_memory_size( + array.values().inner(), + counted_buffers, + total_size, + ), + DataType::Null => {} + DataType::Boolean => count_buffer_memory_size( + array.as_boolean().values().inner(), + counted_buffers, + total_size, + ), + DataType::Binary => count_byte_array_memory_size( + array.as_binary::<i32>(), + counted_buffers, + total_size, + ), + DataType::LargeBinary => count_byte_array_memory_size( + array.as_binary::<i64>(), + counted_buffers, + total_size, + ), + DataType::Utf8 => count_byte_array_memory_size( + array.as_string::<i32>(), + counted_buffers, + total_size, + ), + DataType::LargeUtf8 => count_byte_array_memory_size( + array.as_string::<i64>(), + counted_buffers, + total_size, + ), + DataType::BinaryView => { + let array = array.as_binary_view(); + count_buffer_memory_size(array.views().inner(), counted_buffers, total_size); + for buffer in array.data_buffers() { + count_buffer_memory_size(buffer, counted_buffers, total_size); + } + } + DataType::Utf8View => { + let array = array.as_string_view(); + count_buffer_memory_size(array.views().inner(), counted_buffers, total_size); + for buffer in array.data_buffers() { + count_buffer_memory_size(buffer, counted_buffers, total_size); + } + } + DataType::FixedSizeBinary(_) => count_buffer_memory_size( + array.as_fixed_size_binary().values(), + counted_buffers, + total_size, + ), + DataType::List(_) => count_list_array_memory_size( + array.as_list::<i32>(), + counted_buffers, + total_size, + ), + DataType::LargeList(_) => count_list_array_memory_size( + array.as_list::<i64>(), + counted_buffers, + total_size, + ), + DataType::ListView(_) => { + let array = array.as_list_view::<i32>(); + count_buffer_memory_size(array.offsets().inner(), counted_buffers, total_size); + count_buffer_memory_size(array.sizes().inner(), counted_buffers, total_size); + count_array_memory_size(array.values().as_ref(), counted_buffers, total_size); + } + DataType::LargeListView(_) => { + let array = array.as_list_view::<i64>(); + count_buffer_memory_size(array.offsets().inner(), counted_buffers, total_size); + count_buffer_memory_size(array.sizes().inner(), counted_buffers, total_size); + count_array_memory_size(array.values().as_ref(), counted_buffers, total_size); + } + DataType::FixedSizeList(_, _) => count_array_memory_size( + array.as_fixed_size_list().values().as_ref(), + counted_buffers, + total_size, + ), + DataType::Struct(_) => { + for child in array.as_struct().columns() { + count_array_memory_size(child.as_ref(), counted_buffers, total_size); + } + } + DataType::Union(_, _) => { + let array = array.as_union(); + count_buffer_memory_size(array.type_ids().inner(), counted_buffers, total_size); + if let Some(offsets) = array.offsets() { + count_buffer_memory_size(offsets.inner(), counted_buffers, total_size); + } + for (type_id, _) in array.fields().iter() { + count_array_memory_size( + array.child(type_id).as_ref(), + counted_buffers, + total_size, + ); + } + } + DataType::Dictionary(_, _) => { + let array = array.as_any_dictionary(); + count_array_memory_size(array.keys(), counted_buffers, total_size); + count_array_memory_size(array.values().as_ref(), counted_buffers, total_size); + } + DataType::Map(_, _) => { + let array = array.as_map(); + count_buffer_memory_size( + array.offsets().inner().inner(), + counted_buffers, + total_size, + ); + count_array_memory_size(array.entries(), counted_buffers, total_size); + } + DataType::RunEndEncoded(_, _) => downcast_run_array! { + array => { + count_buffer_memory_size( + array.run_ends().inner().inner(), + counted_buffers, + total_size, + ); + count_array_memory_size( + array.values().as_ref(), + counted_buffers, + total_size, + ); + }, + _ => unreachable!(), Review Comment: Removed the runtime panic paths in b6414cd88. All current Arrow layouts are handled explicitly (with primitive types dispatched by downcast_primitive_array); legal REE index widths remain specialized, while malformed/custom or future Array implementations fall back to generic ArrayData traversal instead of panicking. ########## datafusion/common/src/utils/memory.rs: ########## @@ -185,31 +194,225 @@ impl RecordBatchMemoryCounter { } } -/// 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>>, +/// 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 + } +} + +fn count_buffer_memory_size( + buffer: &Buffer, + counted_buffers: &mut BufferIdSet, 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 + if counted_buffers.insert(buffer.data_ptr().addr()) { + *total_size += 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( + array: &dyn Array, + counted_buffers: &mut BufferIdSet, + total_size: &mut usize, +) { + if let Some(nulls) = array.nulls() { + count_buffer_memory_size(nulls.buffer(), counted_buffers, total_size); } - // Count all children `ArrayData` recursively - for child in array_data.child_data() { - count_array_data_memory_size(child, counted_buffers, total_size); + downcast_primitive_array! { + array => count_buffer_memory_size( + array.values().inner(), + counted_buffers, + total_size, + ), + DataType::Null => {} + DataType::Boolean => count_buffer_memory_size( + array.as_boolean().values().inner(), + counted_buffers, + total_size, + ), + DataType::Binary => count_byte_array_memory_size( + array.as_binary::<i32>(), + counted_buffers, + total_size, + ), + DataType::LargeBinary => count_byte_array_memory_size( + array.as_binary::<i64>(), + counted_buffers, + total_size, + ), + DataType::Utf8 => count_byte_array_memory_size( + array.as_string::<i32>(), + counted_buffers, + total_size, + ), + DataType::LargeUtf8 => count_byte_array_memory_size( + array.as_string::<i64>(), + counted_buffers, + total_size, + ), + DataType::BinaryView => { + let array = array.as_binary_view(); + count_buffer_memory_size(array.views().inner(), counted_buffers, total_size); + for buffer in array.data_buffers() { + count_buffer_memory_size(buffer, counted_buffers, total_size); + } + } + DataType::Utf8View => { + let array = array.as_string_view(); + count_buffer_memory_size(array.views().inner(), counted_buffers, total_size); + for buffer in array.data_buffers() { + count_buffer_memory_size(buffer, counted_buffers, total_size); + } + } + DataType::FixedSizeBinary(_) => count_buffer_memory_size( + array.as_fixed_size_binary().values(), + counted_buffers, + total_size, + ), + DataType::List(_) => count_list_array_memory_size( + array.as_list::<i32>(), + counted_buffers, + total_size, + ), + DataType::LargeList(_) => count_list_array_memory_size( + array.as_list::<i64>(), + counted_buffers, + total_size, + ), + DataType::ListView(_) => { + let array = array.as_list_view::<i32>(); + count_buffer_memory_size(array.offsets().inner(), counted_buffers, total_size); + count_buffer_memory_size(array.sizes().inner(), counted_buffers, total_size); + count_array_memory_size(array.values().as_ref(), counted_buffers, total_size); + } + DataType::LargeListView(_) => { + let array = array.as_list_view::<i64>(); + count_buffer_memory_size(array.offsets().inner(), counted_buffers, total_size); + count_buffer_memory_size(array.sizes().inner(), counted_buffers, total_size); + count_array_memory_size(array.values().as_ref(), counted_buffers, total_size); + } + DataType::FixedSizeList(_, _) => count_array_memory_size( + array.as_fixed_size_list().values().as_ref(), + counted_buffers, + total_size, + ), + DataType::Struct(_) => { + for child in array.as_struct().columns() { + count_array_memory_size(child.as_ref(), counted_buffers, total_size); + } + } + DataType::Union(_, _) => { + let array = array.as_union(); + count_buffer_memory_size(array.type_ids().inner(), counted_buffers, total_size); + if let Some(offsets) = array.offsets() { + count_buffer_memory_size(offsets.inner(), counted_buffers, total_size); + } + for (type_id, _) in array.fields().iter() { + count_array_memory_size( + array.child(type_id).as_ref(), + counted_buffers, + total_size, + ); + } + } + DataType::Dictionary(_, _) => { + let array = array.as_any_dictionary(); + count_array_memory_size(array.keys(), counted_buffers, total_size); + count_array_memory_size(array.values().as_ref(), counted_buffers, total_size); + } + DataType::Map(_, _) => { + let array = array.as_map(); + count_buffer_memory_size( + array.offsets().inner().inner(), + counted_buffers, + total_size, + ); + count_array_memory_size(array.entries(), counted_buffers, total_size); + } + DataType::RunEndEncoded(_, _) => downcast_run_array! { + array => { + count_buffer_memory_size( + array.run_ends().inner().inner(), + counted_buffers, + total_size, + ); + count_array_memory_size( + array.values().as_ref(), + counted_buffers, + total_size, + ); + }, + _ => unreachable!(), + } + _ => unreachable!("unsupported array type: {}", array.data_type()), } } +fn count_byte_array_memory_size<T: ByteArrayType>( + array: &arrow::array::GenericByteArray<T>, + counted_buffers: &mut BufferIdSet, Review Comment: Applied in b6414cd88 as part of the counter-method refactor: byte, byte-view, list, list-view, run-array, and generic ArrayData traversal helpers are now methods on RecordBatchMemoryCounter. ########## datafusion/common/src/utils/memory.rs: ########## @@ -185,31 +194,225 @@ impl RecordBatchMemoryCounter { } } -/// 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>>, +/// 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 + } +} + +fn count_buffer_memory_size( + buffer: &Buffer, + counted_buffers: &mut BufferIdSet, 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 + if counted_buffers.insert(buffer.data_ptr().addr()) { + *total_size += 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( + array: &dyn Array, + counted_buffers: &mut BufferIdSet, + total_size: &mut usize, +) { + if let Some(nulls) = array.nulls() { + count_buffer_memory_size(nulls.buffer(), counted_buffers, total_size); } - // Count all children `ArrayData` recursively - for child in array_data.child_data() { - count_array_data_memory_size(child, counted_buffers, total_size); + downcast_primitive_array! { + array => count_buffer_memory_size( + array.values().inner(), + counted_buffers, + total_size, + ), + DataType::Null => {} + DataType::Boolean => count_buffer_memory_size( + array.as_boolean().values().inner(), + counted_buffers, + total_size, + ), + DataType::Binary => count_byte_array_memory_size( + array.as_binary::<i32>(), + counted_buffers, + total_size, + ), + DataType::LargeBinary => count_byte_array_memory_size( + array.as_binary::<i64>(), + counted_buffers, + total_size, + ), + DataType::Utf8 => count_byte_array_memory_size( + array.as_string::<i32>(), + counted_buffers, + total_size, + ), + DataType::LargeUtf8 => count_byte_array_memory_size( + array.as_string::<i64>(), + counted_buffers, + total_size, + ), + DataType::BinaryView => { + let array = array.as_binary_view(); + count_buffer_memory_size(array.views().inner(), counted_buffers, total_size); + for buffer in array.data_buffers() { + count_buffer_memory_size(buffer, counted_buffers, total_size); + } + } + DataType::Utf8View => { + let array = array.as_string_view(); + count_buffer_memory_size(array.views().inner(), counted_buffers, total_size); + for buffer in array.data_buffers() { + count_buffer_memory_size(buffer, counted_buffers, total_size); + } + } + DataType::FixedSizeBinary(_) => count_buffer_memory_size( + array.as_fixed_size_binary().values(), + counted_buffers, + total_size, + ), + DataType::List(_) => count_list_array_memory_size( + array.as_list::<i32>(), + counted_buffers, + total_size, + ), + DataType::LargeList(_) => count_list_array_memory_size( + array.as_list::<i64>(), + counted_buffers, + total_size, + ), + DataType::ListView(_) => { + let array = array.as_list_view::<i32>(); + count_buffer_memory_size(array.offsets().inner(), counted_buffers, total_size); + count_buffer_memory_size(array.sizes().inner(), counted_buffers, total_size); + count_array_memory_size(array.values().as_ref(), counted_buffers, total_size); + } + DataType::LargeListView(_) => { + let array = array.as_list_view::<i64>(); + count_buffer_memory_size(array.offsets().inner(), counted_buffers, total_size); + count_buffer_memory_size(array.sizes().inner(), counted_buffers, total_size); + count_array_memory_size(array.values().as_ref(), counted_buffers, total_size); + } + DataType::FixedSizeList(_, _) => count_array_memory_size( + array.as_fixed_size_list().values().as_ref(), + counted_buffers, + total_size, + ), + DataType::Struct(_) => { + for child in array.as_struct().columns() { + count_array_memory_size(child.as_ref(), counted_buffers, total_size); + } + } + DataType::Union(_, _) => { + let array = array.as_union(); + count_buffer_memory_size(array.type_ids().inner(), counted_buffers, total_size); + if let Some(offsets) = array.offsets() { + count_buffer_memory_size(offsets.inner(), counted_buffers, total_size); + } + for (type_id, _) in array.fields().iter() { + count_array_memory_size( + array.child(type_id).as_ref(), + counted_buffers, + total_size, + ); + } + } + DataType::Dictionary(_, _) => { + let array = array.as_any_dictionary(); + count_array_memory_size(array.keys(), counted_buffers, total_size); + count_array_memory_size(array.values().as_ref(), counted_buffers, total_size); + } + DataType::Map(_, _) => { + let array = array.as_map(); + count_buffer_memory_size( + array.offsets().inner().inner(), + counted_buffers, + total_size, + ); + count_array_memory_size(array.entries(), counted_buffers, total_size); + } + DataType::RunEndEncoded(_, _) => downcast_run_array! { + array => { + count_buffer_memory_size( + array.run_ends().inner().inner(), + counted_buffers, + total_size, + ); + count_array_memory_size( + array.values().as_ref(), + counted_buffers, + total_size, + ); + }, + _ => unreachable!(), + } + _ => unreachable!("unsupported array type: {}", array.data_type()), } } +fn count_byte_array_memory_size<T: ByteArrayType>( + array: &arrow::array::GenericByteArray<T>, + counted_buffers: &mut BufferIdSet, + total_size: &mut usize, +) { + count_buffer_memory_size( + array.offsets().inner().inner(), + counted_buffers, + total_size, + ); + count_buffer_memory_size(array.values(), counted_buffers, total_size); +} + +fn count_list_array_memory_size<O: arrow::array::OffsetSizeTrait>( Review Comment: Added dedicated byte-view and list-view accounting methods in b6414cd88, including views/data buffers and list-view offsets/sizes/child values. The parity test now uses concrete non-empty BinaryView, Utf8View, ListView, and LargeListView arrays against the previous ArrayData traversal. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
