andygrove commented on code in PR #25853: URL: https://github.com/apache/datafusion/pull/25853#discussion_r4436292566
########## datafusion/physical-plan/src/spill/spilled_row_sizes.rs: ########## @@ -0,0 +1,398 @@ +// 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. Each view counts the bytes of its value, so views that share bytes +/// are counted as if compaction copied them, although it can write shared bytes +/// only once. +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()) + } Review Comment: Thanks, that reproduces here, and so does a single-stage version (`target_partitions = 1`), which fails with `Failed to allocate additional 233.6 MB`. Fixed in c5cc3af12, but differently from the diff above: - `normalized_keys()` asserts that the dictionary has values, so it panics on the empty dictionary that a batch of null lists holds. - Compacting the dictionaries of every chunk makes repeated values expensive to spill. The `array_agg` state dictionary holds a copy of each input batch's dictionary, and each chunk references a different subset of those copies, so each chunk writes its own. On `array_agg` over 100 distinct 200-byte values (2M rows, 100K groups, 24 MB limit), the diff above spills 371 MB, against 14.7 MB on main, where every chunk shares the dictionary and the spill file holds it once. So `SpilledRowSizes` now decides per dictionary. When fewer than two keys reference each value on average, a row counts the values its keys reference, and `IncrementalSortIterator` compacts each chunk to those values. Other dictionaries count only their keys and stay shared, as on main. Replay also compacts the dictionaries of nested columns, as it does their views, since a merged batch can hold the dictionaries of all the batches its rows come from. The compaction only runs in aggregate spilling, so the spill files of other operators don't change. Your query and the 1000-byte variant now pass at 128 MB, as do single-stage and `max_spill_merge_fan_in = 2` versions, and the low-cardinality case above spills 10.5 MB and 14.7 MB (20- and 200-byte values), as on main. Case I now has a version of your query at the file's 24 MB limit, with `target_partitions` 2 and 1, and a query whose values repeat within each group, which needs the replay compaction. This still leaves one case as on main: a large dictionary whose values repeat within a group but not across groups. For example, `array_agg` over 5M rows whose values each repeat 10 times in one group still fails at 24 and 64 MB. Compacting would be cheap there, but telling it apart from values that groups share needs counting distinct values per row, which I'd leave for a follow-up. ########## datafusion/sqllogictest/test_files/aggregate_memory_spill.slt: ########## @@ -239,6 +239,146 @@ 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 --- Review Comment: Added in c5cc3af12, as a single-stage query. At the file's 24 MB limit with 500K rows, the two-partition version already passes on main, while the single-stage one fails there with `Failed to allocate additional 48.3 MB`. ########## datafusion/sqllogictest/test_files/aggregate_memory_spill.slt: ########## @@ -239,6 +239,146 @@ 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. Review Comment: It comes from the Case I queries, but not from the spilling changes. Run on its own, the first Case I query (two partitions) peaks at 103.6 MB allocated with 10.9 KB reserved, sampled when a partial aggregate changes its reservation. With a 2 GB limit, where nothing spills and none of this PR's code runs, it's the same (107 MB). With `target_partitions = 1` and 2 GB, the drift is 103.5 MB with 48 B reserved, sampled when the aggregate releases its reservation after emitting. So it's memory that the aggregation holds outside any reservation, such as the batches an aggregate has emitted and the next operator hasn't reserved, which is what #25650 tracks. ########## datafusion/physical-plan/src/spill/spilled_row_sizes.rs: ########## @@ -0,0 +1,398 @@ +// 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. Each view counts the bytes of its value, so views that share bytes +/// are counted as if compaction copied them, although it can write shared bytes +/// only once. +pub(crate) struct SpilledRowSizes { Review Comment: Agreed. Nothing in it is specific to aggregation, since it sizes rows by how they're written to a spill file, but only aggregate spilling uses it so far. A sort could use it the same way to bound its spill batches. Happy to consolidate it in the #26711 follow-up. ########## datafusion/physical-plan/src/sorts/merge.rs: ########## @@ -146,11 +156,22 @@ impl<C: CursorValues> SortPreservingMergeStream<C> { poll_reset_epochs: vec![0; stream_count], loser_tree: vec![], batch_size, + max_batch_bytes: None, + avg_row_bytes: vec![0; stream_count], + in_progress_bytes: 0, fetch, produced: 0, } } + /// Also emits a batch once its rows take about `max_batch_bytes`, as Review Comment: Only aggregate spilling needs it so far. `batch_size` bounds rows, and a row of aggregate state can hold a whole group, such as an `array_agg` list, so the batches that spilling writes and replay merges can be far larger than their row count suggests. A global `max_batch_bytes` like `batch_size` makes sense to me, but every operator that builds batches would need to estimate row sizes, as this PR does for the two that aggregate spilling uses: `IncrementalSortIterator` sizes each row as it is spilled, and the merge uses the average row size of each input batch. I kept the option crate-private so it can become a config later without an API change. It could be part of #26711. -- 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]
