Weijun-H commented on code in PR #25910: URL: https://github.com/apache/datafusion/pull/25910#discussion_r4400364576
########## 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: Could this record an upper bound on decoded size, including IPC padding? With 4,096 UInt8 columns × 17 rows and a 64 KiB bound, the recorded maximum is 36,864 B but decoding needs 98,304 B. Two runs reserve 147,456 B yet hold 196,608 B. The sizing gap predates this PR, but splitting can exhaust the existing 2× margin. ########## datafusion/physical-plan/src/sorts/sort.rs: ########## @@ -998,6 +1008,38 @@ impl ExternalSorter { } } +/// The bound on the in-memory size of each batch that [`ExternalSorter`] spills. +/// +/// The merge of the spill files starts with `sort_spill_reservation_bytes` of headroom. +/// It reads at least [`MIN_SPILL_MERGE_FAN_IN`] files, and reserves +/// [`get_reserved_bytes_for_record_batch_size`] of the largest batch for each, with +/// read-ahead reduced to one batch under pressure. The bound is the largest batch size for +/// which that minimum merge fits in the headroom, found from the same formula. +/// +/// The bound is at least [`MIN_SPILL_MAX_BATCH_BYTES`], so that a small or zero headroom +/// does not write, read and merge tiny batches. Below that headroom the merge cannot fit +/// in it anyway, and takes the rest from the memory pool, as it does without a bound. +fn spill_max_batch_bytes(sort_spill_reservation_bytes: usize) -> usize { + let fits = |batch: usize| { + MIN_SPILL_MERGE_FAN_IN * get_reserved_bytes_for_record_batch_size(batch, batch) + <= sort_spill_reservation_bytes + }; + // Binary search for the largest batch size that fits: `fits` is true up to it. + let (mut fitting, mut too_large) = (0, sort_spill_reservation_bytes + 1); Review Comment: Large reservations overflow here (e.g. 9223372036854775807 on 64-bit targets). Could this use `(reservation / (2 * MIN_SPILL_MERGE_FAN_IN)).max(MIN_SPILL_MAX_BATCH_BYTES)` instead? Division avoids both the search and the overflow; please add a boundary 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]
