diff --git a/datafusion/common/Cargo.toml b/datafusion/common/Cargo.toml index 1eb23089a4021..9ee199fe82f28 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 0000000000000..2479d6ac987cb --- /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) -> RecordBatch { + let fields = columns + .iter() + .enumerate() + .map(|(index, column)| { + Field::new(format!("col_{index}"), column.data_type().clone(), false) + }) + .collect::>(); + + 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::>(); + + 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::( + (0..num_rows).map(|row| { + let value = row as i64 + column as i64; + Some(vec![Some(value), Some(value + 1)]) + }), + )) as ArrayRef + }) + .collect::>(); + + 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::>(); + + 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::>(); + + 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 21c084119e120..fd405e06a262e 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>, + 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>, - 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::()); + } + DataType::LargeBinary => { + self.count_byte_array_memory_size(array.as_binary::()); + } + DataType::Utf8 => { + self.count_byte_array_memory_size(array.as_string::()); + } + DataType::LargeUtf8 => { + self.count_byte_array_memory_size(array.as_string::()); + } + 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::()); + } + DataType::LargeList(_) => { + self.count_list_array_memory_size(array.as_list::()); + } + DataType::ListView(_) => { + self.count_list_view_array_memory_size(array.as_list_view::()); + } + DataType::LargeListView(_) => { + self.count_list_view_array_memory_size(array.as_list_view::()); + } + 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::( + 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()), + }, + // 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( + &mut self, + array: &GenericByteArray, + ) { + self.count_buffer_memory_size(array.offsets().inner().inner()); + self.count_buffer_memory_size(array.values()); + } + + fn count_byte_view_array_memory_size( + &mut self, + array: &GenericByteViewArray, + ) { + 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( + &mut self, + array: &GenericListArray, + ) { + 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( + &mut self, + array: &GenericListViewArray, + ) { + 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(&mut self, array: &dyn Array) { + 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()); + } 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>; INLINE_BUFFER_IDS], + len: usize, + overflow: Option>>, +} + +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) -> 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>, + 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::>(); + let columns = (0..=INLINE_BUFFER_IDS) + .map(|value| Arc::new(Int32Array::from(vec![value as i32])) as _) + .collect::>(); + 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::() + ); + 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::([ + Some(vec![Some(1), Some(2)]), + None, + Some(vec![Some(3)]), + ])) as ArrayRef, + Arc::new(LargeListViewArray::from_iter_primitive::( + [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::::try_new( + &Int16Array::from(vec![2_i16, 5]), + &run_values, + ) + .unwrap(), + ) as ArrayRef, + Arc::new( + RunArray::::try_new( + &Int32Array::from(vec![2_i32, 5]), + &run_values, + ) + .unwrap(), + ) as ArrayRef, + Arc::new( + RunArray::::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![