adriangb commented on code in PR #25910: URL: https://github.com/apache/datafusion/pull/25910#discussion_r4446016407
########## datafusion/physical-plan/src/spill/spill_pieces.rs: ########## @@ -0,0 +1,209 @@ +// 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. + +//! 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() { Review Comment: Fixed in 1182cb94a: the recorded size now counts per-buffer IPC padding (schema alignment, or 64 B when compressed). Added your example as a test, plus a check across 6 types × {none, LZ4, Zstd}. -- 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]
