adriangb commented on code in PR #25910:
URL: https://github.com/apache/datafusion/pull/25910#discussion_r4446054858
##########
datafusion/physical-plan/src/spill/in_progress_spill_file.rs:
##########
@@ -100,6 +100,15 @@ impl InProgressSpillFile {
/// Appends a `RecordBatch` using the asynchronous spill writer.
pub async fn append_batch_async(&mut self, batch: &RecordBatch) ->
Result<usize> {
+ Ok(self.append_batch_async_with_stats(batch).await?.0)
+ }
+
+ /// Like [`Self::append_batch_async`], but returns both the memory size
and the row
+ /// count of the batch written.
+ pub(crate) async fn append_batch_async_with_stats(
+ &mut self,
+ batch: &RecordBatch,
+ ) -> Result<(usize, usize)> {
Review Comment:
Let's use a named newtype struct here to avoid mixing up our usizes
##########
datafusion/physical-plan/src/spill/spill_pieces.rs:
##########
@@ -0,0 +1,196 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for
additional information regarding copyright ownership
+// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF)
under one or more contributor license agreements
+
+//! Splits a batch into byte-bounded pieces before it is written to a spill
file.
+
+use std::ops::Range;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, AsArray, RecordBatch, RecordBatchOptions,
make_array};
+use arrow::datatypes::{DataType, Schema};
+use arrow_data::transform::MutableArrayData;
+use arrow_select::dictionary::garbage_collect_any_dictionary;
+use datafusion_common::Result;
+
+use super::gc_view_arrays;
+use super::spill_manager::GetSlicedSize;
+
+/// The pieces that one appended batch is written as, in row order.
+///
+/// Without a bound, or when the batch is already within it, there is one
piece: the batch
+/// after [`gc_view_arrays`], as it is written without a bound. Its
dictionaries are not
+/// changed, so the batches of a spill file can keep sharing one dictionary.
+///
+/// When the batch is larger than `max_batch_bytes`, it is cut into equal row
ranges by its
+/// average row size. Each range is copied out with [`compact`], so that it
holds only its
+/// own rows, and measured with [`GetSlicedSize`]. A range that is still too
large (the rows
+/// differ in size) is halved until it fits or has one row. Pieces are
produced one at a
+/// time, so that the caller can write each one before the next is copied.
+///
+/// Schemas that [`compact`] cannot narrow (see [`can_split`]) are written
whole.
+pub(crate) struct SpillPieces {
+ batch: RecordBatch,
+ max_batch_bytes: usize,
+ /// Row ranges still to produce, the next one last.
+ pending: Vec<Range<usize>>,
+ /// The whole batch, when it is written as one piece.
+ whole: Option<(RecordBatch, usize)>,
+}
+
+impl SpillPieces {
+ pub(crate) fn new(
+ batch: &RecordBatch,
+ max_batch_bytes: Option<usize>,
+ ) -> Result<Self> {
+ let batch = gc_view_arrays(batch)?;
+ let size = batch.get_sliced_size()?;
+ let rows = batch.num_rows();
+ let whole = |batch, size| Self {
+ whole: Some((batch, size)),
+ batch: RecordBatch::new_empty(Arc::new(Schema::empty())),
+ max_batch_bytes: 0,
+ pending: vec![],
+ };
+ let Some(max_batch_bytes) = max_batch_bytes
+ .filter(|max| size > *max && rows > 1 &&
can_split(batch.schema_ref()))
+ else {
+ return Ok(whole(batch, size));
+ };
+ // A shared dictionary can hold many values that no row of this batch
uses. Drop
+ // them once before choosing the row ranges, so that the ranges are
sized by what
+ // the batch uses and the pieces do not each prune the whole
dictionary again.
+ let batch = prune_dictionaries(&batch)?;
+ let size = batch.get_sliced_size()?;
+ if size <= max_batch_bytes {
+ return Ok(whole(batch, size));
+ }
+ // The sliced size over-counts nested columns, whose children are not
narrowed by
+ // a slice, so this can give more pieces than needed but not fewer.
+ let pieces = size.div_ceil(max_batch_bytes.max(1)).min(rows);
+ let rows_per_piece = rows.div_ceil(pieces);
+ let pending = (0..rows)
+ .step_by(rows_per_piece)
+ .map(|start| start..(start + rows_per_piece).min(rows))
+ .rev()
+ .collect();
+ Ok(Self {
+ batch,
+ max_batch_bytes,
+ pending,
+ whole: None,
+ })
+ }
+}
+
+impl Iterator for SpillPieces {
+ /// A piece and its size, as [`GetSlicedSize`] measures it.
+ type Item = Result<(RecordBatch, usize)>;
Review Comment:
A named newtype struct here would be nice
##########
datafusion/physical-plan/src/spill/spill_pieces.rs:
##########
@@ -0,0 +1,196 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for
additional information regarding copyright ownership
+// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF)
under one or more contributor license agreements
+
+//! Splits a batch into byte-bounded pieces before it is written to a spill
file.
+
+use std::ops::Range;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, AsArray, RecordBatch, RecordBatchOptions,
make_array};
+use arrow::datatypes::{DataType, Schema};
+use arrow_data::transform::MutableArrayData;
+use arrow_select::dictionary::garbage_collect_any_dictionary;
+use datafusion_common::Result;
+
+use super::gc_view_arrays;
+use super::spill_manager::GetSlicedSize;
+
+/// The pieces that one appended batch is written as, in row order.
+///
+/// Without a bound, or when the batch is already within it, there is one
piece: the batch
+/// after [`gc_view_arrays`], as it is written without a bound. Its
dictionaries are not
+/// changed, so the batches of a spill file can keep sharing one dictionary.
+///
+/// When the batch is larger than `max_batch_bytes`, it is cut into equal row
ranges by its
+/// average row size. Each range is copied out with [`compact`], so that it
holds only its
+/// own rows, and measured with [`GetSlicedSize`]. A range that is still too
large (the rows
+/// differ in size) is halved until it fits or has one row. Pieces are
produced one at a
+/// time, so that the caller can write each one before the next is copied.
+///
+/// Schemas that [`compact`] cannot narrow (see [`can_split`]) are written
whole.
+pub(crate) struct SpillPieces {
+ batch: RecordBatch,
+ max_batch_bytes: usize,
+ /// Row ranges still to produce, the next one last.
+ pending: Vec<Range<usize>>,
+ /// The whole batch, when it is written as one piece.
+ whole: Option<(RecordBatch, usize)>,
+}
+
+impl SpillPieces {
+ pub(crate) fn new(
+ batch: &RecordBatch,
+ max_batch_bytes: Option<usize>,
+ ) -> Result<Self> {
+ let batch = gc_view_arrays(batch)?;
+ let size = batch.get_sliced_size()?;
+ let rows = batch.num_rows();
+ let whole = |batch, size| Self {
+ whole: Some((batch, size)),
+ batch: RecordBatch::new_empty(Arc::new(Schema::empty())),
+ max_batch_bytes: 0,
+ pending: vec![],
+ };
+ let Some(max_batch_bytes) = max_batch_bytes
+ .filter(|max| size > *max && rows > 1 &&
can_split(batch.schema_ref()))
+ else {
+ return Ok(whole(batch, size));
+ };
+ // A shared dictionary can hold many values that no row of this batch
uses. Drop
+ // them once before choosing the row ranges, so that the ranges are
sized by what
+ // the batch uses and the pieces do not each prune the whole
dictionary again.
+ let batch = prune_dictionaries(&batch)?;
+ let size = batch.get_sliced_size()?;
+ if size <= max_batch_bytes {
+ return Ok(whole(batch, size));
+ }
+ // The sliced size over-counts nested columns, whose children are not
narrowed by
+ // a slice, so this can give more pieces than needed but not fewer.
+ let pieces = size.div_ceil(max_batch_bytes.max(1)).min(rows);
+ let rows_per_piece = rows.div_ceil(pieces);
+ let pending = (0..rows)
+ .step_by(rows_per_piece)
+ .map(|start| start..(start + rows_per_piece).min(rows))
+ .rev()
+ .collect();
+ Ok(Self {
+ batch,
+ max_batch_bytes,
+ pending,
+ whole: None,
+ })
+ }
+}
+
+impl Iterator for SpillPieces {
+ /// A piece and its size, as [`GetSlicedSize`] measures it.
+ type Item = Result<(RecordBatch, usize)>;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ if let Some(whole) = self.whole.take() {
+ return Some(Ok(whole));
+ }
+ while let Some(range) = self.pending.pop() {
+ let piece = match compact(&self.batch.slice(range.start,
range.len())) {
+ Ok(piece) => piece,
+ Err(e) => return Some(Err(e)),
+ };
+ let size = match piece.get_sliced_size() {
+ Ok(size) => size,
+ Err(e) => return Some(Err(e)),
+ };
+ if size > self.max_batch_bytes && range.len() > 1 {
+ let mid = range.start + range.len() / 2;
+ self.pending.push(mid..range.end);
+ self.pending.push(range.start..mid);
+ continue;
+ }
+ return Some(Ok((piece, size)));
+ }
+ None
+ }
+}
+
+/// Copies a slice so that it holds only its own rows: list children are
narrowed to the
+/// range the slice uses, view arrays are garbage collected, and top-level
dictionaries drop
+/// the values that no key uses.
+fn compact(slice: &RecordBatch) -> Result<RecordBatch> {
Review Comment:
Surely there must be other copies / implementations of this functionality or
similar in datafusion and/or arrow-rs.
##########
datafusion/physical-plan/src/spill/spill_pieces.rs:
##########
@@ -0,0 +1,196 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for
additional information regarding copyright ownership
+// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF)
under one or more contributor license agreements
+
+//! Splits a batch into byte-bounded pieces before it is written to a spill
file.
+
+use std::ops::Range;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, AsArray, RecordBatch, RecordBatchOptions,
make_array};
+use arrow::datatypes::{DataType, Schema};
+use arrow_data::transform::MutableArrayData;
+use arrow_select::dictionary::garbage_collect_any_dictionary;
+use datafusion_common::Result;
+
+use super::gc_view_arrays;
+use super::spill_manager::GetSlicedSize;
+
+/// The pieces that one appended batch is written as, in row order.
+///
+/// Without a bound, or when the batch is already within it, there is one
piece: the batch
+/// after [`gc_view_arrays`], as it is written without a bound. Its
dictionaries are not
+/// changed, so the batches of a spill file can keep sharing one dictionary.
+///
+/// When the batch is larger than `max_batch_bytes`, it is cut into equal row
ranges by its
+/// average row size. Each range is copied out with [`compact`], so that it
holds only its
+/// own rows, and measured with [`GetSlicedSize`]. A range that is still too
large (the rows
+/// differ in size) is halved until it fits or has one row. Pieces are
produced one at a
+/// time, so that the caller can write each one before the next is copied.
+///
+/// Schemas that [`compact`] cannot narrow (see [`can_split`]) are written
whole.
+pub(crate) struct SpillPieces {
+ batch: RecordBatch,
+ max_batch_bytes: usize,
+ /// Row ranges still to produce, the next one last.
+ pending: Vec<Range<usize>>,
+ /// The whole batch, when it is written as one piece.
+ whole: Option<(RecordBatch, usize)>,
+}
+
+impl SpillPieces {
+ pub(crate) fn new(
+ batch: &RecordBatch,
+ max_batch_bytes: Option<usize>,
+ ) -> Result<Self> {
+ let batch = gc_view_arrays(batch)?;
+ let size = batch.get_sliced_size()?;
+ let rows = batch.num_rows();
+ let whole = |batch, size| Self {
+ whole: Some((batch, size)),
+ batch: RecordBatch::new_empty(Arc::new(Schema::empty())),
+ max_batch_bytes: 0,
+ pending: vec![],
+ };
+ let Some(max_batch_bytes) = max_batch_bytes
+ .filter(|max| size > *max && rows > 1 &&
can_split(batch.schema_ref()))
+ else {
+ return Ok(whole(batch, size));
+ };
+ // A shared dictionary can hold many values that no row of this batch
uses. Drop
+ // them once before choosing the row ranges, so that the ranges are
sized by what
+ // the batch uses and the pieces do not each prune the whole
dictionary again.
+ let batch = prune_dictionaries(&batch)?;
+ let size = batch.get_sliced_size()?;
+ if size <= max_batch_bytes {
+ return Ok(whole(batch, size));
+ }
+ // The sliced size over-counts nested columns, whose children are not
narrowed by
+ // a slice, so this can give more pieces than needed but not fewer.
+ let pieces = size.div_ceil(max_batch_bytes.max(1)).min(rows);
+ let rows_per_piece = rows.div_ceil(pieces);
+ let pending = (0..rows)
+ .step_by(rows_per_piece)
+ .map(|start| start..(start + rows_per_piece).min(rows))
+ .rev()
+ .collect();
+ Ok(Self {
+ batch,
+ max_batch_bytes,
+ pending,
+ whole: None,
+ })
+ }
+}
+
+impl Iterator for SpillPieces {
+ /// A piece and its size, as [`GetSlicedSize`] measures it.
+ type Item = Result<(RecordBatch, usize)>;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ if let Some(whole) = self.whole.take() {
+ return Some(Ok(whole));
+ }
+ while let Some(range) = self.pending.pop() {
+ let piece = match compact(&self.batch.slice(range.start,
range.len())) {
+ Ok(piece) => piece,
+ Err(e) => return Some(Err(e)),
+ };
+ let size = match piece.get_sliced_size() {
+ Ok(size) => size,
+ Err(e) => return Some(Err(e)),
+ };
+ if size > self.max_batch_bytes && range.len() > 1 {
+ let mid = range.start + range.len() / 2;
+ self.pending.push(mid..range.end);
+ self.pending.push(range.start..mid);
+ continue;
+ }
+ return Some(Ok((piece, size)));
+ }
+ None
+ }
+}
+
+/// Copies a slice so that it holds only its own rows: list children are
narrowed to the
+/// range the slice uses, view arrays are garbage collected, and top-level
dictionaries drop
+/// the values that no key uses.
+fn compact(slice: &RecordBatch) -> Result<RecordBatch> {
+ let columns = slice
+ .columns()
+ .iter()
+ .map(|column| {
+ let data = column.to_data();
+ let mut copy = MutableArrayData::new(vec![&data], false,
data.len());
+ copy.try_extend(0, 0, data.len())?;
+ let copied = make_array(copy.freeze());
+ Ok(match copied.data_type() {
+ DataType::Dictionary(_, _) => {
+ garbage_collect_any_dictionary(copied.as_any_dictionary())?
+ }
+ _ => copied,
+ })
+ })
+ .collect::<Result<Vec<ArrayRef>>>()?;
+ let options =
RecordBatchOptions::new().with_row_count(Some(slice.num_rows()));
+ let piece = RecordBatch::try_new_with_options(slice.schema(), columns,
&options)?;
+ gc_view_arrays(&piece)
+}
+
+/// Drops the values of top-level dictionaries that no key uses.
+fn prune_dictionaries(batch: &RecordBatch) -> Result<RecordBatch> {
+ if !batch
+ .schema_ref()
+ .fields()
+ .iter()
+ .any(|field| matches!(field.data_type(), DataType::Dictionary(_, _)))
+ {
+ return Ok(batch.clone());
+ }
+ let columns = batch
+ .columns()
+ .iter()
+ .map(|column| {
+ Ok(match column.data_type() {
+ DataType::Dictionary(_, _) => {
+ garbage_collect_any_dictionary(column.as_any_dictionary())?
+ }
+ _ => Arc::clone(column),
+ })
+ })
+ .collect::<Result<Vec<ArrayRef>>>()?;
+ let options =
RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));
+ Ok(RecordBatch::try_new_with_options(
+ batch.schema(),
+ columns,
+ &options,
+ )?)
+}
+
+/// Whether [`compact`] narrows every column of `schema` to its rows. It does
not for list
+/// views and unions, whose children a row range does not bound, run-end
encoded arrays, or
+/// dictionaries below the top level, whose values it does not prune. A batch
with such a
+/// column is written whole.
+fn can_split(schema: &Schema) -> bool {
Review Comment:
Are we sure these are not handled? Again maybe some existing shared
functionality would handle them. Even if not I'm not convinced this list is
accurate.
##########
datafusion/physical-plan/src/spill/spill_pieces.rs:
##########
@@ -112,6 +119,43 @@ impl Iterator for SpillPieces {
}
}
+/// The capacity granularity of arrow's `MutableBuffer`, which a compressed
buffer is
+/// decompressed into.
+const DECOMPRESSED_ALIGNMENT: usize = 64;
Review Comment:
I'm also surprised there is no existing functionality in arrow-rs or
datafusion to calculate this sort of padding / alignment of IPC files. We
should double check, and consider if this is the right place to put it. It
might be a better fit for arrow-rs.
##########
datafusion/physical-plan/src/spill/spill_pieces.rs:
##########
@@ -0,0 +1,196 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for
additional information regarding copyright ownership
+// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF)
under one or more contributor license agreements
+
+//! Splits a batch into byte-bounded pieces before it is written to a spill
file.
+
+use std::ops::Range;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, AsArray, RecordBatch, RecordBatchOptions,
make_array};
+use arrow::datatypes::{DataType, Schema};
+use arrow_data::transform::MutableArrayData;
+use arrow_select::dictionary::garbage_collect_any_dictionary;
+use datafusion_common::Result;
+
+use super::gc_view_arrays;
+use super::spill_manager::GetSlicedSize;
+
+/// The pieces that one appended batch is written as, in row order.
+///
+/// Without a bound, or when the batch is already within it, there is one
piece: the batch
+/// after [`gc_view_arrays`], as it is written without a bound. Its
dictionaries are not
+/// changed, so the batches of a spill file can keep sharing one dictionary.
+///
+/// When the batch is larger than `max_batch_bytes`, it is cut into equal row
ranges by its
+/// average row size. Each range is copied out with [`compact`], so that it
holds only its
+/// own rows, and measured with [`GetSlicedSize`]. A range that is still too
large (the rows
+/// differ in size) is halved until it fits or has one row. Pieces are
produced one at a
+/// time, so that the caller can write each one before the next is copied.
+///
+/// Schemas that [`compact`] cannot narrow (see [`can_split`]) are written
whole.
+pub(crate) struct SpillPieces {
+ batch: RecordBatch,
+ max_batch_bytes: usize,
+ /// Row ranges still to produce, the next one last.
+ pending: Vec<Range<usize>>,
+ /// The whole batch, when it is written as one piece.
+ whole: Option<(RecordBatch, usize)>,
+}
+
+impl SpillPieces {
+ pub(crate) fn new(
+ batch: &RecordBatch,
+ max_batch_bytes: Option<usize>,
+ ) -> Result<Self> {
+ let batch = gc_view_arrays(batch)?;
+ let size = batch.get_sliced_size()?;
+ let rows = batch.num_rows();
+ let whole = |batch, size| Self {
+ whole: Some((batch, size)),
+ batch: RecordBatch::new_empty(Arc::new(Schema::empty())),
+ max_batch_bytes: 0,
+ pending: vec![],
+ };
+ let Some(max_batch_bytes) = max_batch_bytes
+ .filter(|max| size > *max && rows > 1 &&
can_split(batch.schema_ref()))
+ else {
+ return Ok(whole(batch, size));
+ };
+ // A shared dictionary can hold many values that no row of this batch
uses. Drop
+ // them once before choosing the row ranges, so that the ranges are
sized by what
+ // the batch uses and the pieces do not each prune the whole
dictionary again.
+ let batch = prune_dictionaries(&batch)?;
+ let size = batch.get_sliced_size()?;
+ if size <= max_batch_bytes {
+ return Ok(whole(batch, size));
+ }
+ // The sliced size over-counts nested columns, whose children are not
narrowed by
+ // a slice, so this can give more pieces than needed but not fewer.
+ let pieces = size.div_ceil(max_batch_bytes.max(1)).min(rows);
+ let rows_per_piece = rows.div_ceil(pieces);
+ let pending = (0..rows)
+ .step_by(rows_per_piece)
+ .map(|start| start..(start + rows_per_piece).min(rows))
+ .rev()
+ .collect();
+ Ok(Self {
+ batch,
+ max_batch_bytes,
+ pending,
+ whole: None,
+ })
+ }
+}
+
+impl Iterator for SpillPieces {
+ /// A piece and its size, as [`GetSlicedSize`] measures it.
+ type Item = Result<(RecordBatch, usize)>;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ if let Some(whole) = self.whole.take() {
+ return Some(Ok(whole));
+ }
+ while let Some(range) = self.pending.pop() {
+ let piece = match compact(&self.batch.slice(range.start,
range.len())) {
+ Ok(piece) => piece,
+ Err(e) => return Some(Err(e)),
+ };
+ let size = match piece.get_sliced_size() {
+ Ok(size) => size,
+ Err(e) => return Some(Err(e)),
+ };
+ if size > self.max_batch_bytes && range.len() > 1 {
+ let mid = range.start + range.len() / 2;
+ self.pending.push(mid..range.end);
+ self.pending.push(range.start..mid);
+ continue;
+ }
+ return Some(Ok((piece, size)));
+ }
+ None
+ }
+}
+
+/// Copies a slice so that it holds only its own rows: list children are
narrowed to the
+/// range the slice uses, view arrays are garbage collected, and top-level
dictionaries drop
+/// the values that no key uses.
+fn compact(slice: &RecordBatch) -> Result<RecordBatch> {
+ let columns = slice
+ .columns()
+ .iter()
+ .map(|column| {
+ let data = column.to_data();
+ let mut copy = MutableArrayData::new(vec![&data], false,
data.len());
+ copy.try_extend(0, 0, data.len())?;
+ let copied = make_array(copy.freeze());
+ Ok(match copied.data_type() {
+ DataType::Dictionary(_, _) => {
+ garbage_collect_any_dictionary(copied.as_any_dictionary())?
+ }
+ _ => copied,
+ })
+ })
+ .collect::<Result<Vec<ArrayRef>>>()?;
+ let options =
RecordBatchOptions::new().with_row_count(Some(slice.num_rows()));
+ let piece = RecordBatch::try_new_with_options(slice.schema(), columns,
&options)?;
+ gc_view_arrays(&piece)
+}
+
+/// Drops the values of top-level dictionaries that no key uses.
+fn prune_dictionaries(batch: &RecordBatch) -> Result<RecordBatch> {
Review Comment:
Similarly here, I'm surprised this functionality doesn't exist. IIRC it
does, and I also remember us putting in place bounds like "only gc / compact if
more than X% of the space is wasted". Maybe that matters less here since
spilling by nature involves "slow down drastically to be able to actually
finish", but still if we can avoid copying all of the data for a marginal gain
the tradeoff of e.g. inflating the reservation on the other end might be worth
it.
--
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]
