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..fc047892ee8fc --- /dev/null +++ b/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) -> 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"); + + 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!( + 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 21c084119e120..3cdfa9f1b4ccd 100644 --- a/datafusion/common/src/utils/memory.rs +++ b/datafusion/common/src/utils/memory.rs @@ -19,11 +19,21 @@ 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; +/// 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 +161,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 +177,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 +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>, +/// 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 +450,42 @@ 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, 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 { + 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 +594,117 @@ 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); + 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}" + ); + } + + 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] fn test_get_record_batch_memory_size_nested_array() { let schema = Arc::new(Schema::new(vec![