comphead commented on code in PR #25853: URL: https://github.com/apache/datafusion/pull/25853#discussion_r4139564667
########## datafusion/sqllogictest/test_files/aggregate_memory_spill.slt: ########## @@ -233,6 +233,110 @@ FROM ( 08)--------------AggregateExec: mode=Partial, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[sum(t.v)], metrics=[<slt:ignore>early_emit_count=<slt:ignore>] <slt:ignore> +# --- Case H: a few groups with large state --- +# The array_agg states of all 32 groups do not fit in the memory limit, but a +# single group does. Spilling must write the groups in batches that can be +# read back within the limit, instead of one batch holding every group. +# See https://github.com/apache/datafusion/issues/25851 +statement ok +RESET datafusion.runtime.max_spill_merge_fan_in + +statement ok +SET datafusion.execution.batch_size = 8192 + +statement ok +SET datafusion.runtime.memory_limit = '24M' + +# Partial/final aggregation: the final aggregation spills the partial states of +# its partition. +statement ok +SET datafusion.execution.target_partitions = 2 + +query II +SELECT count(*), sum(n) +FROM ( + SELECT v % 32 AS k, cardinality(array_agg(concat(CAST(v AS VARCHAR), repeat('x', 100)))) AS n + FROM generate_series(1, 500000) AS t(v) + GROUP BY v % 32 +) +---- +32 500000 + +# Assert the final aggregation spilled, the `spilled_bytes` metric must be in MB +query TT +EXPLAIN ANALYZE +SELECT count(*), sum(n) +FROM ( + SELECT v % 32 AS k, cardinality(array_agg(concat(CAST(v AS VARCHAR), repeat('x', 100)))) AS n + FROM generate_series(1, 500000) AS t(v) + GROUP BY v % 32 +) +---- +<slt:ignore> +06)----------AggregateExec: mode=FinalPartitioned, gby=[t.v % Int64(32)@0 as t.v % Int64(32)], aggr=[<slt:ignore>], metrics=[<slt:ignore>spilled_bytes=<slt:ignore> MB, spilled_rows=<slt:ignore>] +07)------------RepartitionExec: <slt:ignore> + +# Single-stage aggregation: spills several runs as input arrives, which replay +# merges. +statement ok +SET datafusion.execution.target_partitions = 1 + +query II +SELECT count(*), sum(n) +FROM ( + SELECT v % 32 AS k, cardinality(array_agg(concat(CAST(v AS VARCHAR), repeat('x', 100)))) AS n + FROM generate_series(1, 500000) AS t(v) + GROUP BY v % 32 +) +---- +32 500000 + +# Assert spill happened, the `spilled_bytes` metric must be in MB +query TT +EXPLAIN ANALYZE +SELECT count(*), sum(n) +FROM ( + SELECT v % 32 AS k, cardinality(array_agg(concat(CAST(v AS VARCHAR), repeat('x', 100)))) AS n + FROM generate_series(1, 500000) AS t(v) + GROUP BY v % 32 +) +---- +<slt:ignore> +04)------AggregateExec: mode=Single, gby=[v@0 % 32 as t.v % Int64(32)], aggr=[<slt:ignore>], metrics=[<slt:ignore>spilled_bytes=<slt:ignore> MB, spilled_rows=<slt:ignore>] +05)--------ProjectionExec: <slt:ignore> + +# 14 large groups, adjacent in key order, spilled together with 62,500 small Review Comment: Every other case in this file follows its query with an `EXPLAIN ANALYZE` check that it spilled. This query and the next one have none, so if a later planner or memory change stopped them from spilling, they would keep passing without covering #25851. ########## datafusion/physical-plan/src/sorts/streaming_merge.rs: ########## @@ -326,6 +338,121 @@ mod tests { ExecutionPlanMetricsSet, SpillMetrics, }; + /// Merges five 1000 byte rows, one per batch, with 1000 small rows in + /// batches of 100, either from spill files or from in-memory streams. + async fn merge_wide_and_narrow_rows( + from_spill_files: bool, + max_batch_bytes: Option<usize>, + ) -> Result<Vec<RecordBatch>> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Utf8, false), + ])); + let batch = |keys: std::ops::Range<i32>, value_len: usize| { + let values = keys.clone().map(|_| "x".repeat(value_len)); + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from_iter_values(keys)), + Arc::new(StringArray::from_iter_values(values)), + ], + ) + .unwrap() + }; + let wide: Vec<_> = (0..5).map(|key| batch(key..key + 1, 1000)).collect(); + let narrow: Vec<_> = (0..10) + .map(|i| batch(10 + i * 100..10 + (i + 1) * 100, 1)) + .collect(); + + let sort: LexOrdering = + [PhysicalSortExpr::new_default(col("key", &schema)?)].into(); + let mut builder = StreamingMergeBuilder::new() + .with_schema(Arc::clone(&schema)) + .with_expressions(&sort) + .with_metrics(BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0)) + .with_batch_size(100) + .with_max_batch_bytes(max_batch_bytes) + .with_bypass_mempool(); + if from_spill_files { + let spill_manager = SpillManager::new( + Arc::new(TaskContext::default()).runtime_env(), + SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + Arc::clone(&schema), + ); + let spill = |batches: &[RecordBatch]| SortedSpillFile { + file: spill_manager + .spill_record_batch_and_finish(batches, "test") + .unwrap() + .unwrap(), + max_record_batch_memory: batches + .iter() + .map(RecordBatch::get_array_memory_size) + .max() + .unwrap(), + }; + builder = builder + .with_sorted_spill_files(vec![spill(&wide), spill(&narrow)]) + .with_spill_manager(spill_manager); + } else { + let stream = |batches: Vec<RecordBatch>| { + Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&schema), + futures::stream::iter(batches.into_iter().map(Ok)), + )) as SendableRecordBatchStream + }; + builder = builder.with_streams(vec![stream(wide), stream(narrow)]); + } + let batches = collect(builder.build()?).await?; + + let keys: Vec<i32> = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_primitive::<Int32Type>() + .values() + .to_vec() + }) + .collect(); + let expected: Vec<i32> = (0..5).chain(10..1010).collect(); + assert_eq!(keys, expected); + Ok(batches) + } + + #[tokio::test] + async fn test_max_batch_bytes_limits_merged_batches() -> Result<()> { + let num_rows = |batches: &[RecordBatch]| -> Vec<usize> { + batches.iter().map(RecordBatch::num_rows).collect() + }; + for from_spill_files in [true, false] { Review Comment: The byte cap is only set by aggregate replay, which always merges spill files. Each multi-level pass (`create_new_merge_sort`) then builds this same `SortPreservingMergeStream` over the spill-file streams, so the `false` run covers no code path that the `true` run misses. Keeping only the spill-file case would drop the `from_spill_files` flag and the in-memory branch of `merge_wide_and_narrow_rows`. ########## datafusion/physical-plan/src/spill/spilled_row_sizes.rs: ########## @@ -0,0 +1,426 @@ +// 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. + +//! Per-row sizes of a batch once it is written to a spill file. + +use arrow::array::{Array, AsArray, RecordBatch}; +use arrow::buffer::{OffsetBuffer, ScalarBuffer}; +use arrow::datatypes::DataType; +use arrow_data::MAX_INLINE_VIEW_LEN; + +use super::VIEW_SIZE_BYTES; + +/// Estimates how many bytes each row of a batch takes once written to a spill +/// file, where view arrays are compacted (see [`super::gc_view_arrays`]). +/// +/// A row counts its fixed-width values, the offsets and bytes of its +/// variable-width values, and the elements of its nested values. Buffers that +/// every row shares, such as validity bitmaps and dictionary values, are not +/// counted. +pub(crate) struct SpilledRowSizes { + /// Bytes that every row takes in the fixed-width columns + fixed: usize, + /// The other columns + variable: Vec<ColumnSize>, + /// See [`Self::max_row`] + max_row: usize, +} + +impl SpilledRowSizes { + pub(crate) fn new(batch: &RecordBatch) -> Self { + let mut fixed = 0; + let mut variable = vec![]; + for column in batch.columns() { + let column = ColumnSize::new(column.as_ref()); + match column.fixed() { + Some(width) => fixed += width, + None => variable.push(column), + } + } + let max_row = variable.iter().fold(fixed, |bytes, column| { + bytes.saturating_add(column.max_row()) + }); + Self { + fixed, + variable, + max_row, + } + } + + /// Returns the size in bytes of row `row`. + pub(crate) fn row(&self, row: usize) -> usize { + self.variable.iter().fold(self.fixed, |bytes, column| { + bytes + column.rows(row, row + 1) + }) + } + + /// Returns the size in bytes of every row, if they all have the same size. + pub(crate) fn fixed(&self) -> Option<usize> { + self.variable.is_empty().then_some(self.fixed) + } + + /// Returns an upper bound on the size in bytes of any row. + pub(crate) fn max_row(&self) -> usize { + self.max_row + } +} + +/// Size model of one column, see [`SpilledRowSizes`]. +enum ColumnSize { + /// Every row takes the same number of bytes. + Fixed(usize), + /// An offset per row, plus the bytes between consecutive offsets. + Bytes(Offsets), + /// A view per row, plus the bytes of values too long to be inlined. + Views(ScalarBuffer<u128>), + /// An offset per row, plus the elements between consecutive offsets. + List(Offsets, Box<ColumnSize>), + /// A fixed number of elements per row. + FixedSizeList(usize, Box<ColumnSize>), + /// A row of each field. + Struct(Vec<ColumnSize>), +} + +impl ColumnSize { + fn new(array: &dyn Array) -> Self { + match array.data_type() { + DataType::Utf8 => Self::Bytes(array.as_string::<i32>().offsets().into()), + DataType::LargeUtf8 => Self::Bytes(array.as_string::<i64>().offsets().into()), + DataType::Binary => Self::Bytes(array.as_binary::<i32>().offsets().into()), + DataType::LargeBinary => { + Self::Bytes(array.as_binary::<i64>().offsets().into()) + } + DataType::Utf8View => Self::Views(array.as_string_view().views().clone()), + DataType::BinaryView => Self::Views(array.as_binary_view().views().clone()), + DataType::List(_) => { + let list = array.as_list::<i32>(); + Self::List(list.offsets().into(), Box::new(Self::new(list.values()))) + } + DataType::LargeList(_) => { + let list = array.as_list::<i64>(); + Self::List(list.offsets().into(), Box::new(Self::new(list.values()))) + } + DataType::Map(_, _) => { + let map = array.as_map(); + Self::List(map.offsets().into(), Box::new(Self::new(map.entries()))) + } + DataType::FixedSizeList(_, size) => Self::FixedSizeList( + *size as usize, + Box::new(Self::new(array.as_fixed_size_list().values())), + ), + DataType::Struct(_) => Self::Struct( + array + .as_struct() + .columns() + .iter() + .map(|field| Self::new(field.as_ref())) + .collect(), + ), + DataType::Dictionary(key, _) => { + Self::Fixed(key.primitive_width().unwrap_or_default()) + } + DataType::FixedSizeBinary(size) => Self::Fixed(*size as usize), + DataType::Null | DataType::Boolean => Self::Fixed(0), + data_type => Self::Fixed(data_type.primitive_width().unwrap_or_else(|| { + // Other layouts are spread evenly over their rows + let bytes = array.to_data().get_slice_memory_size().unwrap_or_default(); + bytes / array.len().max(1) + })), + } + } + + /// Returns the size in bytes of every row, if they all have the same size. + fn fixed(&self) -> Option<usize> { + match self { + Self::Fixed(width) => Some(*width), + Self::FixedSizeList(size, values) => values.fixed().map(|width| size * width), + Self::Struct(fields) => fields.iter().map(Self::fixed).sum(), + Self::Bytes(_) | Self::Views(_) | Self::List(_, _) => None, + } + } + + /// Returns an upper bound on the size in bytes of a row. + fn max_row(&self) -> usize { + match self { + Self::Fixed(width) => *width, + Self::Bytes(offsets) => offsets.width() + offsets.max_len(), + Self::Views(views) => { + let max_len = views + .iter() + .map(|&view| view as u32) + .filter(|&len| len > MAX_INLINE_VIEW_LEN) + .max() + .unwrap_or_default(); + VIEW_SIZE_BYTES + max_len as usize + } + Self::List(offsets, values) => offsets + .width() + .saturating_add(offsets.max_len().saturating_mul(values.max_row())), + Self::FixedSizeList(size, values) => size.saturating_mul(values.max_row()), + Self::Struct(fields) => fields + .iter() + .fold(0, |bytes, field| bytes.saturating_add(field.max_row())), + } + } + + /// Returns the size in bytes of rows `start..end`. + fn rows(&self, start: usize, end: usize) -> usize { + let num_rows = end - start; + match self { + Self::Fixed(width) => num_rows * width, + Self::Bytes(offsets) => { + num_rows * offsets.width() + offsets.get(end) - offsets.get(start) + } + Self::Views(views) => views[start..end] + .iter() + .map(|&view| { + let len = view as u32; + if len > MAX_INLINE_VIEW_LEN { + VIEW_SIZE_BYTES + len as usize + } else { + VIEW_SIZE_BYTES + } + }) + .sum(), + Self::List(offsets, values) => { + num_rows * offsets.width() + + values.rows(offsets.get(start), offsets.get(end)) + } + Self::FixedSizeList(size, values) => values.rows(start * size, end * size), + Self::Struct(fields) => { + fields.iter().map(|field| field.rows(start, end)).sum() + } + } + } +} + +/// The offsets of a variable-width or list column. +enum Offsets { + Small(ScalarBuffer<i32>), + Large(ScalarBuffer<i64>), +} + +impl Offsets { + /// Returns the size in bytes of one offset. + fn width(&self) -> usize { + match self { + Self::Small(_) => size_of::<i32>(), + Self::Large(_) => size_of::<i64>(), + } + } + + fn get(&self, index: usize) -> usize { + match self { + Self::Small(offsets) => offsets[index] as usize, + Self::Large(offsets) => offsets[index] as usize, + } + } + + /// Returns the largest distance between consecutive offsets. + fn max_len(&self) -> usize { + match self { + Self::Small(offsets) => offsets + .windows(2) + .map(|pair| (pair[1] - pair[0]) as usize) + .max(), + Self::Large(offsets) => offsets + .windows(2) + .map(|pair| (pair[1] - pair[0]) as usize) + .max(), + } + .unwrap_or_default() + } +} + +impl From<&OffsetBuffer<i32>> for Offsets { + fn from(offsets: &OffsetBuffer<i32>) -> Self { + Self::Small(offsets.inner().clone()) + } +} + +impl From<&OffsetBuffer<i64>> for Offsets { + fn from(offsets: &OffsetBuffer<i64>) -> Self { + Self::Large(offsets.inner().clone()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ + ArrayRef, BooleanArray, DictionaryArray, FixedSizeBinaryArray, Int32Array, + Int64Array, ListArray, StringArray, StringViewArray, StructArray, + }; + use arrow::buffer::Buffer; + use arrow::datatypes::{Field, Int32Type, Int64Type}; + use std::sync::Arc; + + fn batch(columns: Vec<ArrayRef>) -> RecordBatch { + RecordBatch::try_from_iter( + columns + .into_iter() + .enumerate() + .map(|(i, column)| (format!("c{i}"), column)), + ) + .unwrap() + } + + fn row_sizes(columns: Vec<ArrayRef>) -> Vec<usize> { + let batch = batch(columns); + let sizes = SpilledRowSizes::new(&batch); + (0..batch.num_rows()).map(|row| sizes.row(row)).collect() + } + + #[test] + fn fixed_width_and_offsets() { + let ints: ArrayRef = Arc::new(Int64Array::from(vec![1, 2])); + let strings: ArrayRef = Arc::new(StringArray::from(vec!["a", "bbbb"])); + // 8 bytes of Int64, plus a 4 byte offset and the string bytes + assert_eq!(row_sizes(vec![ints, strings]), vec![13, 16]); + } + + #[test] + fn rows_of_fixed_width_columns_have_one_size() { + let ints: ArrayRef = Arc::new(Int64Array::from(vec![1, 2])); + let fixed: ArrayRef = Arc::new(FixedSizeBinaryArray::new( + 3, + Buffer::from(vec![0u8; 6]), + None, + )); + let strings: ArrayRef = Arc::new(StringArray::from(vec!["a", "bbbb"])); + let sizes = SpilledRowSizes::new(&batch(vec![Arc::clone(&ints), fixed])); + assert_eq!(sizes.fixed(), Some(11)); + assert_eq!( + SpilledRowSizes::new(&batch(vec![ints, strings])).fixed(), + None + ); + } + + #[test] + fn largest_row_bounds_every_row() { Review Comment: This test builds the same view and list arrays as `views_count_their_bytes_unless_inlined` and `lists_count_their_elements`, and nearly the same batch as `fixed_width_and_offsets`. Its `max_row` asserts could move into those three tests, next to the `row` asserts for the same arrays. `other_types` and `rows_of_fixed_width_columns_have_one_size` also build the same `FixedSizeBinary(3)` array, so they could be one test. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
