From e39304102c3946a4e583e1a5a420a11ad03fff35 Mon Sep 17 00:00:00 2001 From: ryux1 Date: Thu, 13 Aug 2026 05:07:43 +0200 Subject: [PATCH 1/4] perf: reduce record batch memory accounting overhead --- datafusion/common/Cargo.toml | 4 + .../common/benches/record_batch_memory.rs | 60 +++ datafusion/common/src/utils/memory.rs | 351 ++++++++++++++++-- 3 files changed, 392 insertions(+), 23 deletions(-) create mode 100644 datafusion/common/benches/record_batch_memory.rs 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..39c4e4117d73f --- /dev/null +++ b/datafusion/common/benches/record_batch_memory.rs @@ -0,0 +1,60 @@ +// 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}; +use arrow::datatypes::{DataType, Field, 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(num_columns: usize) -> RecordBatch { + let fields = (0..num_columns) + .map(|index| Field::new(format!("col_{index}"), DataType::Int64, false)) + .collect::>(); + let columns = (0..num_columns) + .map(|index| { + Arc::new(Int64Array::from_iter_values( + (0..8192).map(|value| value + index as i64), + )) as ArrayRef + }) + .collect::>(); + + RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap() +} + +fn benchmark_record_batch_memory_size(c: &mut Criterion) { + let mut group = c.benchmark_group("record_batch_memory_size"); + + for num_columns in [1, 4, 16, 64] { + let batch = make_batch(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(); +} + +criterion_group!(benches, benchmark_record_batch_memory_size); +criterion_main!(benches); diff --git a/datafusion/common/src/utils/memory.rs b/datafusion/common/src/utils/memory.rs index 21c084119e120..c945a19bb3a6d 100644 --- a/datafusion/common/src/utils/memory.rs +++ b/datafusion/common/src/utils/memory.rs @@ -19,11 +19,17 @@ use crate::error::_exec_datafusion_err; use crate::{HashSet, Result}; -use arrow::array::ArrayData; +use arrow::array::types::ByteArrayType; +use arrow::array::{Array, AsArray, downcast_run_array}; +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; +const INLINE_BUFFER_IDS: usize = 16; + /// Estimates the memory size required for a hash table prior to allocation. /// /// # Parameters @@ -151,7 +157,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, } @@ -167,9 +173,8 @@ impl RecordBatchMemoryCounter { let mut total_size = 0; for array in batch.columns() { - let array_data = array.to_data(); - count_array_data_memory_size( - &array_data, + count_array_memory_size( + array.as_ref(), &mut self.counted_buffers, &mut total_size, ); @@ -185,31 +190,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>, +/// 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 + } +} + +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::(), + counted_buffers, + total_size, + ), + DataType::LargeBinary => count_byte_array_memory_size( + array.as_binary::(), + counted_buffers, + total_size, + ), + DataType::Utf8 => count_byte_array_memory_size( + array.as_string::(), + counted_buffers, + total_size, + ), + DataType::LargeUtf8 => count_byte_array_memory_size( + array.as_string::(), + 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::(), + counted_buffers, + total_size, + ), + DataType::LargeList(_) => count_list_array_memory_size( + array.as_list::(), + counted_buffers, + total_size, + ), + DataType::ListView(_) => { + let array = array.as_list_view::(); + 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::(); + 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( + array: &arrow::array::GenericByteArray, + 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( + array: &arrow::array::GenericListArray, + counted_buffers: &mut BufferIdSet, + total_size: &mut usize, +) { + count_buffer_memory_size( + array.offsets().inner().inner(), + counted_buffers, + total_size, + ); + count_array_memory_size(array.values().as_ref(), counted_buffers, total_size); +} + #[cfg(test)] mod tests { use std::{collections::HashSet, mem::size_of}; @@ -247,10 +446,37 @@ 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, Float64Array, Int32Array, ListArray, new_null_array}; + use arrow::datatypes::{DataType, Field, Int32Type, 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 + } + #[test] fn test_get_record_batch_memory_size() { let schema = Arc::new(Schema::new(vec![ @@ -359,6 +585,85 @@ 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 run_ends = Arc::new(Field::new("run_ends", DataType::Int32, false)); + let run_values = Arc::new(Field::new("values", DataType::Utf8, true)); + 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), + DataType::RunEndEncoded(run_ends, run_values), + ]; + + for data_type in data_types { + let array = new_null_array(&data_type, 3); + let mut total_size = 0; + count_array_memory_size( + array.as_ref(), + &mut BufferIdSet::default(), + &mut total_size, + ); + assert_eq!( + total_size, + array_data_memory_size(array.as_ref()), + "{data_type}" + ); + } + } + #[test] fn test_get_record_batch_memory_size_nested_array() { let schema = Arc::new(Schema::new(vec![ From 0d288d040e7c9f6800cfe713b9415ce9fc0ebbd4 Mon Sep 17 00:00:00 2001 From: ryux1 Date: Thu, 13 Aug 2026 19:38:57 +0200 Subject: [PATCH 2/4] bench: add row-count and nested memory accounting cases Extend the microbenchmark across row counts and List/Struct layouts. Exercise all legal run-end index types in the ArrayData parity test. --- .../common/benches/record_batch_memory.rs | 108 ++++++++++++++++-- datafusion/common/src/utils/memory.rs | 47 +++++++- 2 files changed, 139 insertions(+), 16 deletions(-) diff --git a/datafusion/common/benches/record_batch_memory.rs b/datafusion/common/benches/record_batch_memory.rs index 39c4e4117d73f..fc047892ee8fc 100644 --- a/datafusion/common/benches/record_batch_memory.rs +++ b/datafusion/common/benches/record_batch_memory.rs @@ -18,32 +18,76 @@ use std::hint::black_box; use std::sync::Arc; -use arrow::array::{ArrayRef, Int64Array}; -use arrow::datatypes::{DataType, Field, Schema}; +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(num_columns: usize) -> RecordBatch { - let fields = (0..num_columns) - .map(|index| Field::new(format!("col_{index}"), DataType::Int64, false)) +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..8192).map(|value| value + index as i64), + (0..num_rows).map(|value| value as i64 + index as i64), )) as ArrayRef }) .collect::>(); - RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap() + 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 benchmark_record_batch_memory_size(c: &mut Criterion) { - let mut group = c.benchmark_group("record_batch_memory_size"); +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"); for num_columns in [1, 4, 16, 64] { - let batch = make_batch(num_columns); + let batch = make_primitive_batch(8192, num_columns); group.bench_with_input( BenchmarkId::from_parameter(num_columns), &batch, @@ -56,5 +100,47 @@ fn benchmark_record_batch_memory_size(c: &mut Criterion) { group.finish(); } -criterion_group!(benches, benchmark_record_batch_memory_size); +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!( + benches, + benchmark_column_count, + benchmark_row_count, + benchmark_array_layout +); criterion_main!(benches); diff --git a/datafusion/common/src/utils/memory.rs b/datafusion/common/src/utils/memory.rs index c945a19bb3a6d..bc8b380b0c3be 100644 --- a/datafusion/common/src/utils/memory.rs +++ b/datafusion/common/src/utils/memory.rs @@ -446,8 +446,13 @@ mod tests { #[cfg(test)] mod record_batch_tests { use super::*; - use arrow::array::{ArrayData, Float64Array, Int32Array, ListArray, new_null_array}; - use arrow::datatypes::{DataType, Field, Int32Type, Schema, UnionFields, UnionMode}; + use arrow::array::{ + ArrayData, ArrayRef, Float64Array, Int16Array, Int32Array, Int64Array, ListArray, + RunArray, StringArray, 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 { @@ -624,8 +629,6 @@ mod record_batch_tests { ), false, )); - let run_ends = Arc::new(Field::new("run_ends", DataType::Int32, false)); - let run_values = Arc::new(Field::new("values", DataType::Utf8, true)); let data_types = vec![ DataType::Boolean, DataType::Int32, @@ -645,7 +648,6 @@ mod record_batch_tests { DataType::Union(union_fields, UnionMode::Dense), DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), DataType::Map(map_entries, false), - DataType::RunEndEncoded(run_ends, run_values), ]; for data_type in data_types { @@ -662,6 +664,41 @@ mod record_batch_tests { "{data_type}" ); } + + 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 { + let mut total_size = 0; + count_array_memory_size( + array.as_ref(), + &mut BufferIdSet::default(), + &mut total_size, + ); + assert_eq!(total_size, array_data_memory_size(array.as_ref())); + } } #[test] From e022e54660bf27bfd9aaa6929176a8fb8f9152f9 Mon Sep 17 00:00:00 2001 From: ryux1 Date: Fri, 14 Aug 2026 00:01:24 +0200 Subject: [PATCH 3/4] document(common): explain inline buffer capacity Clarify that the 16-entry threshold avoids allocations for small buffer sets while bounding inline storage and linear lookup. The threshold is a performance heuristic rather than a semantic limit. --- datafusion/common/src/utils/memory.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/datafusion/common/src/utils/memory.rs b/datafusion/common/src/utils/memory.rs index bc8b380b0c3be..3cdfa9f1b4ccd 100644 --- a/datafusion/common/src/utils/memory.rs +++ b/datafusion/common/src/utils/memory.rs @@ -28,6 +28,10 @@ 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. From b6414cd886e2901abc90c1f64c545da5759f1b9b Mon Sep 17 00:00:00 2001 From: ryux1 Date: Sat, 15 Aug 2026 15:50:34 +0200 Subject: [PATCH 4/4] perf(common): cover shared memory accounting paths Move buffer traversal helpers onto RecordBatchMemoryCounter and replace panic-based array dispatch with safe generic accounting fallbacks. Add concrete byte/list view parity coverage and benchmark a counter reused across zero-copy slices. --- .../common/benches/record_batch_memory.rs | 48 +- datafusion/common/src/utils/memory.rs | 417 +++++++++--------- 2 files changed, 257 insertions(+), 208 deletions(-) diff --git a/datafusion/common/benches/record_batch_memory.rs b/datafusion/common/benches/record_batch_memory.rs index fc047892ee8fc..2479d6ac987cb 100644 --- a/datafusion/common/benches/record_batch_memory.rs +++ b/datafusion/common/benches/record_batch_memory.rs @@ -15,6 +15,10 @@ // 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; @@ -22,7 +26,9 @@ 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; +use datafusion_common::utils::memory::{ + RecordBatchMemoryCounter, get_record_batch_memory_size, +}; fn make_batch(columns: Vec) -> RecordBatch { let fields = columns @@ -86,6 +92,8 @@ fn make_struct_batch(num_rows: usize, num_columns: usize) -> RecordBatch { 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( @@ -103,6 +111,8 @@ fn benchmark_column_count(c: &mut Criterion) { 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( @@ -120,6 +130,8 @@ fn benchmark_row_count(c: &mut Criterion) { 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)), @@ -137,10 +149,42 @@ fn benchmark_array_layout(c: &mut Criterion) { 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_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 3cdfa9f1b4ccd..fd405e06a262e 100644 --- a/datafusion/common/src/utils/memory.rs +++ b/datafusion/common/src/utils/memory.rs @@ -19,8 +19,11 @@ use crate::error::_exec_datafusion_err; use crate::{HashSet, Result}; -use arrow::array::types::ByteArrayType; -use arrow::array::{Array, AsArray, downcast_run_array}; +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; @@ -174,24 +177,187 @@ 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() { - count_array_memory_size( - array.as_ref(), - &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 } + + fn count_buffer_memory_size(&mut self, buffer: &Buffer) { + if self.counted_buffers.insert(buffer.data_ptr().addr()) { + self.memory_usage += buffer.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()), + } + } + + 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 @@ -237,182 +403,6 @@ impl BufferIdSet { } } -fn count_buffer_memory_size( - buffer: &Buffer, - counted_buffers: &mut BufferIdSet, - total_size: &mut usize, -) { - if counted_buffers.insert(buffer.data_ptr().addr()) { - *total_size += buffer.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); - } - - 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::(), - counted_buffers, - total_size, - ), - DataType::LargeBinary => count_byte_array_memory_size( - array.as_binary::(), - counted_buffers, - total_size, - ), - DataType::Utf8 => count_byte_array_memory_size( - array.as_string::(), - counted_buffers, - total_size, - ), - DataType::LargeUtf8 => count_byte_array_memory_size( - array.as_string::(), - 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::(), - counted_buffers, - total_size, - ), - DataType::LargeList(_) => count_list_array_memory_size( - array.as_list::(), - counted_buffers, - total_size, - ), - DataType::ListView(_) => { - let array = array.as_list_view::(); - 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::(); - 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( - array: &arrow::array::GenericByteArray, - 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( - array: &arrow::array::GenericListArray, - counted_buffers: &mut BufferIdSet, - total_size: &mut usize, -) { - count_buffer_memory_size( - array.offsets().inner().inner(), - counted_buffers, - total_size, - ); - count_array_memory_size(array.values().as_ref(), counted_buffers, total_size); -} - #[cfg(test)] mod tests { use std::{collections::HashSet, mem::size_of}; @@ -451,8 +441,9 @@ mod tests { mod record_batch_tests { use super::*; use arrow::array::{ - ArrayData, ArrayRef, Float64Array, Int16Array, Int32Array, Int64Array, ListArray, - RunArray, StringArray, new_null_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, @@ -486,6 +477,12 @@ mod record_batch_tests { 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![ @@ -656,17 +653,31 @@ mod record_batch_tests { for data_type in data_types { let array = new_null_array(&data_type, 3); - let mut total_size = 0; - count_array_memory_size( - array.as_ref(), - &mut BufferIdSet::default(), - &mut total_size, - ); - assert_eq!( - total_size, - array_data_memory_size(array.as_ref()), - "{data_type}" - ); + 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"]); @@ -695,13 +706,7 @@ mod record_batch_tests { ]; for array in run_arrays { - let mut total_size = 0; - count_array_memory_size( - array.as_ref(), - &mut BufferIdSet::default(), - &mut total_size, - ); - assert_eq!(total_size, array_data_memory_size(array.as_ref())); + assert_array_memory_size_matches(array.as_ref()); } }