jayzhan211 commented on code in PR #25853:
URL: https://github.com/apache/datafusion/pull/25853#discussion_r4435947546
##########
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:
Optional: every Case H query uses `array_agg`, so only the list-of-views
path is covered end to end. A `string_agg` case exercises top-level bytes
state; it fails on main and passes with this PR at `-m 128M`:
```sql
SELECT count(*), sum(n) FROM (
SELECT value % 16 AS k, length(string_agg(concat(CAST(value AS VARCHAR),
repeat('x', 100)), ',')) AS n
FROM generate_series(1, 3000000)
GROUP BY value % 16);
```
##########
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:
Dictionary-typed state still hits #25851: a dictionary row is sized by its
key only, and `take` keeps every dictionary value, so a `List<Dictionary>`
state ends up in one chunk, and even a one-row chunk carries the whole
dictionary. Since Comet passes dictionaries, I think this belongs in this PR.
Fails on this PR with `Failed to allocate additional 160.4 MB`
(`datafusion-cli -m 128M`):
```sql
SET datafusion.execution.target_partitions = 2;
SELECT count(*), sum(n) FROM (
SELECT k, cardinality(array_agg(v)) AS n FROM (
SELECT value % 16 AS k,
arrow_cast(concat(CAST(value AS VARCHAR), repeat('x', 100)),
'Dictionary(Int32, Utf8)') AS v
FROM generate_series(1, 3000000))
GROUP BY k);
```
I tried the change below locally. The query above returns 16 groups and
3,000,000 values, and so does a variant with 1000-byte values over 300,000
rows, which still fails with only the GC part. `dictionaries_count_their_keys`
needs updating. The GC in `gc_array` applies to every spill; if you'd rather
keep sort unchanged, it could run only on the chunks `IncrementalSortIterator`
yields when `max_batch_bytes` is set (I haven't tested that variant).
```diff
--- a/datafusion/physical-plan/src/spill/spilled_row_sizes.rs
+++ b/datafusion/physical-plan/src/spill/spilled_row_sizes.rs
/// A row of each field.
Struct(Vec<ColumnSize>),
+ /// A key per row, plus the dictionary value it references.
+ Dictionary(usize, Vec<usize>, Box<ColumnSize>),
}
@@ fn new
DataType::Dictionary(key, _) => {
- Self::Fixed(key.primitive_width().unwrap_or_default())
+ let dictionary = array.as_any_dictionary();
+ Self::Dictionary(
+ key.primitive_width().unwrap_or_default(),
+ dictionary.normalized_keys(),
+ Box::new(Self::new(dictionary.values().as_ref())),
+ )
}
@@ fn fixed
- Self::Bytes(_) | Self::Views(_) | Self::List(_, _) => None,
+ Self::Bytes(_)
+ | Self::Views(_)
+ | Self::List(_, _)
+ | Self::Dictionary(_, _, _) => None,
@@ fn max_row
+ Self::Dictionary(width, _, values) =>
width.saturating_add(values.max_row()),
@@ fn rows
+ Self::Dictionary(width, keys, values) => keys[start..end]
+ .iter()
+ .map(|&key| width + values.rows(key, key + 1))
+ .sum(),
--- a/datafusion/physical-plan/src/spill/mod.rs
+++ b/datafusion/physical-plan/src/spill/mod.rs
@@ fn gc_array
+ DataType::Dictionary(_, _) => {
+ let dictionary = array.as_any_dictionary();
+ let gc =
+
arrow_select::dictionary::garbage_collect_any_dictionary(dictionary)?;
+ let shrunk =
+ gc.as_any_dictionary().values().len() <
dictionary.values().len();
+ let (values, values_mutated) =
gc_array(gc.as_any_dictionary().values())?;
+ if !shrunk && !values_mutated {
+ return Ok((Arc::clone(array), false));
+ }
+ Ok((gc.as_any_dictionary().with_values(values), true))
+ }
_ => gc_array_children(array),
```
(This also needs `arrow-select = { workspace = true }` in
`datafusion/physical-plan/Cargo.toml`.) Could you add the query above to Case H
as well?
--
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]
