fix: account logical record batch slice memory (#8480)

* fix: account logical record batch slice memory Signed-off-by: jeremyhi <fengjiachun@gmail.com> * fix: account nested view payload memory Signed-off-by: jeremyhi <fengjiachun@gmail.com> * test: align scan memory expectation with logical size Signed-off-by: jeremyhi <fengjiachun@gmail.com> * refactor: name Arrow inline view limit Signed-off-by: jeremyhi <fengjiachun@gmail.com> * fix: account visible list child memory Signed-off-by: jeremyhi <fengjiachun@gmail.com> * fix: exclude null list child memory Signed-off-by: jeremyhi <fengjiachun@gmail.com> * fix: exclude null struct child memory Signed-off-by: jeremyhi <fengjiachun@gmail.com> * fix: account fixed-size list slices Signed-off-by: jeremyhi <fengjiachun@gmail.com> * fix: limit scan accounting to flat slices Signed-off-by: jeremyhi <fengjiachun@gmail.com> * docs: clarify nested view accounting scope Signed-off-by: jeremyhi <fengjiachun@gmail.com> * perf: reduce view memory accounting overhead Signed-off-by: jeremyhi <fengjiachun@gmail.com> * chore: fix benchmark license header Signed-off-by: jeremyhi <fengjiachun@gmail.com> * perf: optimize view slice accounting Signed-off-by: jeremyhi <fengjiachun@gmail.com> * test: cover sliced mixed-null view accounting Signed-off-by: jeremyhi <fengjiachun@gmail.com> * refactor: narrow slice accounting to arrow buffers Signed-off-by: jeremyhi <fengjiachun@gmail.com> * docs: note view accounting assumption Signed-off-by: jeremyhi <fengjiachun@gmail.com> --------- Signed-off-by: jeremyhi <fengjiachun@gmail.com>
2026-08-18 03:58:29 +00:00 · 2026-07-14 17:35:17 +08:00
parent 1e94c3adab
commit 9cde31a647
5 changed files with 270 additions and 6 deletions
@@ -34,3 +34,7 @@ tokio.workspace = true
 [[bench]]
 name = "iter_record_batch_rows"
 harness = false
+
+[[bench]]
+name = "memory_accounting"
+harness = false
@@ -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<VectorRef> = 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::<Vec<_>>();
+    batch(vec![
+        (
+            "i32",
+            ConcreteDataType::int32_datatype(),
+            false,
+            Arc::new(Int32Vector::from_slice(
+                (0..rows).map(|value| value as i32).collect::<Vec<_>>(),
+            )),
+        ),
+        (
+            "u64",
+            ConcreteDataType::uint64_datatype(),
+            false,
+            Arc::new(UInt64Vector::from_slice(
+                (0..rows).map(|value| value as u64).collect::<Vec<_>>(),
+            )),
+        ),
+        (
+            "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::<Vec<_>>();
+        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);
@@ -769,7 +769,7 @@ impl MemoryTrackedStream {
        batch: RecordBatch,
        cx: &mut Context<'_>,
    ) -> Poll<Option<Result<RecordBatch>>> {
-        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()])
@@ -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::<Vec<_>>();
+        let binary = (0_u32..1024)
+            .map(|value| (value % 3 != 0).then(|| value.to_le_bytes().to_vec()))
+            .collect::<Vec<_>>();
+        let columns: Vec<VectorRef> = 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::<usize>()
+        );
+    }
+
+    #[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::<Vec<_>>();
+        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::<Vec<_>>();
+        let logical_total = slices
+            .iter()
+            .map(RecordBatch::logical_slice_memory_size)
+            .sum::<usize>();
+        let buffer_total = slices
+            .iter()
+            .map(RecordBatch::buffer_memory_size)
+            .sum::<usize>();
+
+        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<VectorRef> =
+            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![
@@ -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())))