diff --git a/src/common/recordbatch/Cargo.toml b/src/common/recordbatch/Cargo.toml index efc6b6f60e..59929ba347 100644 --- a/src/common/recordbatch/Cargo.toml +++ b/src/common/recordbatch/Cargo.toml @@ -34,3 +34,7 @@ tokio.workspace = true [[bench]] name = "iter_record_batch_rows" harness = false + +[[bench]] +name = "memory_accounting" +harness = false diff --git a/src/common/recordbatch/benches/memory_accounting.rs b/src/common/recordbatch/benches/memory_accounting.rs new file mode 100644 index 0000000000..05fb63555a --- /dev/null +++ b/src/common/recordbatch/benches/memory_accounting.rs @@ -0,0 +1,117 @@ +// Copyright 2023 Greptime Team +// +// Licensed 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 common_recordbatch::RecordBatch; +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use datatypes::data_type::ConcreteDataType; +use datatypes::schema::{ColumnSchema, Schema}; +use datatypes::vectors::{Int32Vector, StringVector, UInt64Vector, VectorRef}; + +const LARGE_ROWS: usize = 8_192; + +fn batch(columns: Vec<(&str, ConcreteDataType, bool, VectorRef)>) -> RecordBatch { + let schema = Arc::new(Schema::new( + columns + .iter() + .map(|(name, data_type, nullable, _)| { + ColumnSchema::new(*name, data_type.clone(), *nullable) + }) + .collect(), + )); + let vectors: Vec = columns + .into_iter() + .map(|(_, _, _, vector)| vector) + .collect(); + RecordBatch::new(schema, vectors).unwrap() +} + +fn primitive_string_batch(rows: usize) -> RecordBatch { + let strings = (0..rows) + .map(|row| (row % 5 != 0).then(|| format!("value-{row:08}-with-payload"))) + .collect::>(); + batch(vec![ + ( + "i32", + ConcreteDataType::int32_datatype(), + false, + Arc::new(Int32Vector::from_slice( + (0..rows).map(|value| value as i32).collect::>(), + )), + ), + ( + "u64", + ConcreteDataType::uint64_datatype(), + false, + Arc::new(UInt64Vector::from_slice( + (0..rows).map(|value| value as u64).collect::>(), + )), + ), + ( + "string", + ConcreteDataType::string_datatype(), + true, + Arc::new(StringVector::from(strings)), + ), + ]) +} + +fn bench_one_call(c: &mut Criterion) { + let mut group = c.benchmark_group("memory_accounting/one_call"); + for rows in [128, LARGE_ROWS] { + let batch = primitive_string_batch(rows); + group.bench_function(BenchmarkId::new("buffer", rows), |b| { + b.iter(|| black_box(black_box(&batch).buffer_memory_size())) + }); + group.bench_function(BenchmarkId::new("logical", rows), |b| { + b.iter(|| black_box(black_box(&batch).logical_slice_memory_size())) + }); + } + group.finish(); +} + +fn bench_shared_buffer_slices(c: &mut Criterion) { + let backing = primitive_string_batch(LARGE_ROWS); + let mut group = c.benchmark_group("memory_accounting/shared_buffer_slices"); + for slice_len in [1, 16] { + let slices = (0..512) + .map(|index| backing.slice(index * slice_len, slice_len).unwrap()) + .collect::>(); + group.throughput(Throughput::Elements(slices.len() as u64)); + group.bench_function(BenchmarkId::new("buffer", slice_len), |b| { + b.iter(|| { + let total = black_box(&slices) + .iter() + .map(RecordBatch::buffer_memory_size) + .fold(0usize, usize::saturating_add); + black_box(total) + }) + }); + group.bench_function(BenchmarkId::new("logical", slice_len), |b| { + b.iter(|| { + let total = black_box(&slices) + .iter() + .map(RecordBatch::logical_slice_memory_size) + .fold(0usize, usize::saturating_add); + black_box(total) + }) + }); + } + group.finish(); +} + +criterion_group!(benches, bench_one_call, bench_shared_buffer_slices); +criterion_main!(benches); diff --git a/src/common/recordbatch/src/lib.rs b/src/common/recordbatch/src/lib.rs index 3db3d9c08b..46375d8fd7 100644 --- a/src/common/recordbatch/src/lib.rs +++ b/src/common/recordbatch/src/lib.rs @@ -769,7 +769,7 @@ impl MemoryTrackedStream { batch: RecordBatch, cx: &mut Context<'_>, ) -> Poll>> { - let additional = batch.buffer_memory_size(); + let additional = batch.logical_slice_memory_size(); let tracker = self.ready_tracker_mut(); if let Err(error) = tracker.try_track(additional) { @@ -860,6 +860,35 @@ mod tests { as usize } + #[tokio::test] + async fn test_memory_tracked_stream_charges_logical_slice_size() { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "payload", + ConcreteDataType::string_datatype(), + false, + )])); + let payloads: Vec<_> = (0..1024) + .map(|value| format!("payload-{value:04}")) + .collect(); + let batch = RecordBatch::new(schema, vec![Arc::new(StringVector::from(payloads)) as _]) + .unwrap() + .slice(512, 1) + .unwrap(); + let expected_bytes = aligned_tracked_bytes(batch.logical_slice_memory_size()); + assert!(expected_bytes < aligned_tracked_bytes(batch.buffer_memory_size())); + let tracker = QueryMemoryTracker::builder(MB, OnExhaustedPolicy::Fail).build(); + let mut stream = MemoryTrackedStream::new( + RecordBatches::try_new(batch.schema.clone(), vec![batch]) + .unwrap() + .as_stream(), + tracker.clone(), + ); + + stream.next().await.unwrap().unwrap(); + + assert_eq!(tracker.current(), expected_bytes); + } + #[test] fn test_recordbatches_try_from_columns() { let schema = Arc::new(Schema::new(vec![ColumnSchema::new( @@ -1052,7 +1081,7 @@ mod tests { }) .build(); let batch = large_string_batch(700 * 1024); - let expected_bytes = aligned_tracked_bytes(batch.buffer_memory_size()); + let expected_bytes = aligned_tracked_bytes(batch.logical_slice_memory_size()); let mut stream1 = MemoryTrackedStream::new( RecordBatches::try_new(batch.schema.clone(), vec![batch.clone()]) diff --git a/src/common/recordbatch/src/recordbatch.rs b/src/common/recordbatch/src/recordbatch.rs index 7280772d05..77a36836d6 100644 --- a/src/common/recordbatch/src/recordbatch.rs +++ b/src/common/recordbatch/src/recordbatch.rs @@ -269,6 +269,28 @@ impl RecordBatch { .sum() } + /// Returns the logical memory size of this batch's array slices. + /// + /// This sums Arrow's logical visible slice buffers rather than the capacity of their shared + /// backing buffers. View out-of-line payloads and nested custom payloads are not separately + /// traversed or accounted. It is not an exact measure of live physical memory. If Arrow cannot + /// calculate a slice's size, the full buffer size is used conservatively. + /// + /// Mito's current scan paths do not produce top-level View arrays. If they do in the future, + /// their out-of-line payload accounting must be reassessed here. + pub fn logical_slice_memory_size(&self) -> usize { + self.df_record_batch + .columns() + .iter() + .fold(0, |total, array| { + let array_size = array + .to_data() + .get_slice_memory_size() + .unwrap_or_else(|_| array.get_buffer_memory_size()); + total.saturating_add(array_size) + }) + } + /// Iterate the values as strings in the column at index `i`. /// /// Note that if the underlying array is not a valid GreptimeDB vector, an empty iterator is @@ -382,13 +404,14 @@ fn maybe_align_json_array_with_schema( mod tests { use std::sync::Arc; - use datatypes::arrow::array::{AsArray, BinaryArray, UInt32Array}; + use datatypes::arrow::array::{ + AsArray, BinaryArray, StringArray, StringViewArray, UInt32Array, + }; use datatypes::arrow::datatypes::{DataType, Field, Schema as ArrowSchema, UInt32Type}; - use datatypes::arrow_array::StringArray; use datatypes::data_type::ConcreteDataType; use datatypes::extension::json::{JsonExtensionType, JsonMetadata}; use datatypes::schema::{ColumnSchema, Schema}; - use datatypes::vectors::{StringVector, UInt32Vector}; + use datatypes::vectors::{BinaryVector, StringVector, UInt32Vector}; use super::*; @@ -475,6 +498,97 @@ mod tests { assert!(recordbatch.slice(1, 5).is_err()); } + #[test] + fn test_logical_slice_memory_size_for_visible_primitive_string_binary_slices() { + let schema = Arc::new(Schema::new(vec![ + ColumnSchema::new("numbers", ConcreteDataType::uint32_datatype(), false), + ColumnSchema::new("strings", ConcreteDataType::string_datatype(), true), + ColumnSchema::new("binary", ConcreteDataType::binary_datatype(), true), + ])); + let numbers: Vec<_> = (0..1024).collect(); + let strings = (0..1024) + .map(|value| (value % 3 != 0).then(|| format!("value-{value}"))) + .collect::>(); + let binary = (0_u32..1024) + .map(|value| (value % 3 != 0).then(|| value.to_le_bytes().to_vec())) + .collect::>(); + let columns: Vec = vec![ + Arc::new(UInt32Vector::from_slice(numbers)), + Arc::new(StringVector::from(strings)), + Arc::new(BinaryVector::from(binary)), + ]; + let batch = RecordBatch::new(schema, columns).unwrap(); + let slice = batch.slice(511, 3).unwrap(); + + assert!(slice.columns().iter().any(|column| column.null_count() > 0)); + assert!(slice.logical_slice_memory_size() < slice.buffer_memory_size()); + assert_eq!( + slice.logical_slice_memory_size(), + slice + .columns() + .iter() + .map(|column| column.to_data().get_slice_memory_size().unwrap()) + .sum::() + ); + } + + #[test] + fn test_logical_slice_memory_size_for_many_shared_buffer_slices() { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "strings", + ConcreteDataType::string_datatype(), + false, + )])); + let strings = (0..1024) + .map(|value| format!("shared-value-{value}")) + .collect::>(); + let backing = RecordBatch::new( + schema, + vec![Arc::new(StringVector::from(strings)) as VectorRef], + ) + .unwrap(); + let slices = (0..128) + .map(|index| backing.slice(index * 4, 2).unwrap()) + .collect::>(); + let logical_total = slices + .iter() + .map(RecordBatch::logical_slice_memory_size) + .sum::(); + let buffer_total = slices + .iter() + .map(RecordBatch::buffer_memory_size) + .sum::(); + + assert!(logical_total < buffer_total); + assert!(slices.iter().all(|slice| { + slice.logical_slice_memory_size() + == slice.column(0).to_data().get_slice_memory_size().unwrap() + })); + } + + #[test] + fn test_logical_slice_memory_size_uses_arrow_slice_scope_for_views() { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "strings", + ConcreteDataType::utf8_view_datatype(), + false, + )])); + let columns: Vec = + vec![Arc::new(StringVector::from(StringViewArray::from(vec![ + "unrelated backing payload", + "visible string view payload", + ])))]; + let batch = RecordBatch::new(schema, columns) + .unwrap() + .slice(1, 1) + .unwrap(); + + assert_eq!( + batch.column(0).to_data().get_slice_memory_size().unwrap(), + batch.logical_slice_memory_size() + ); + } + #[test] fn test_merge_record_batch() { let column_schemas = vec![ diff --git a/src/table/src/table/scan.rs b/src/table/src/table/scan.rs index 825d930a6f..9ce752553e 100644 --- a/src/table/src/table/scan.rs +++ b/src/table/src/table/scan.rs @@ -541,7 +541,7 @@ impl Stream for StreamWithMetricWrapper { Ok(record_batch) => { // we don't record elapsed time here // since it's calling storage api involving I/O ops - let batch_bytes = record_batch.buffer_memory_size(); + let batch_bytes = record_batch.logical_slice_memory_size(); this.metric.record_output_bytes(batch_bytes); this.metric.record_output(record_batch.num_rows()); Poll::Ready(Some(Ok(record_batch.into_df_record_batch())))