From c935c02bf3f2b3078265324bdab909f836db6b0b Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 26 Oct 2026 14:12:43 +0000 Subject: [PATCH] Emit projection pieces as separate batches ProjectionMorsel wrapped the pieces of a split in one ChunkedArray, which the Arrow export then concatenated, re-merging nested dictionaries. Emit pieces as separate batches in row order, grouping only pieces under 4,096 rows. Split out of vortex-data/vortex#10761. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01SQocaYaDvJHfmorKXnyvJC --- .../src/scan/execution/projection.rs | 75 +++++++++++++------ vortex-layout/src/scan/v2/tests.rs | 46 ++++++++++++ 2 files changed, 100 insertions(+), 21 deletions(-) diff --git a/vortex-layout/src/scan/execution/projection.rs b/vortex-layout/src/scan/execution/projection.rs index 141b7377c26..50eb62d773c 100644 --- a/vortex-layout/src/scan/execution/projection.rs +++ b/vortex-layout/src/scan/execution/projection.rs @@ -1,9 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::collections::VecDeque; use std::sync::Arc; use std::vec::IntoIter; +use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::arrays::ChunkedArray; use vortex_array::dtype::DType; @@ -30,6 +32,8 @@ use crate::scan::execution::adaptor::Piece; use crate::scan::execution::adaptor::ScanAdaptor; const MAX_OUTPUT_ROWS: usize = 65_536; +/// Pieces smaller than this are grouped with their neighbours into one batch. +const MIN_OUTPUT_ROWS: usize = 4_096; /// Keeps sparse selections together and subdivides dense selections to avoid large copies when /// converting chunked projection results to Arrow. @@ -97,15 +101,19 @@ impl Planner for ProjectionPlanner { } } -/// Runs the projection plan over a split's selected rows and returns them as one batch, in row -/// order. A selection that projects to no rows finishes without a batch. +/// Runs the projection plan over a split's selected rows and returns them in row order, one +/// batch per piece the scan produced. Pieces under [`MIN_OUTPUT_ROWS`] are grouped with their +/// neighbours. Emitting pieces separately avoids concatenating them on conversion to Arrow, +/// which rebuilds nested dictionaries. A selection that projects to no rows finishes without a +/// batch. struct ProjectionMorsel { scan: Arc, dtype: DType, selected: SelectedRows, adaptor: Option, pieces: Vec, - done: bool, + /// Batches ready to emit once the scan has finished, in row order. + batches: Option>, } impl ProjectionMorsel { @@ -117,24 +125,39 @@ impl ProjectionMorsel { selected, adaptor: None, pieces: Vec::new(), - done: false, + batches: None, } } - fn finish(&mut self) -> VortexResult { - self.done = true; + /// Groups the finished pieces into batches of at least [`MIN_OUTPUT_ROWS`] rows where + /// possible. + fn group_pieces(&mut self) -> VortexResult> { self.pieces.sort_by_key(|piece| piece.rows.start); - let mut arrays: Vec<_> = std::mem::take(&mut self.pieces) - .into_iter() - .map(|piece| piece.array) - .filter(|array| !array.is_empty()) - .collect(); - let array = match arrays.len() { - 0 => return Ok(MorselOutput::Done), - 1 => arrays.remove(0), - _ => ChunkedArray::try_new(arrays, self.dtype.clone())?.into_array(), - }; - Ok(MorselOutput::Batch(array)) + let mut batches = VecDeque::new(); + let mut group = Vec::new(); + let mut group_rows = 0; + for piece in std::mem::take(&mut self.pieces) { + if piece.array.is_empty() { + continue; + } + group_rows += piece.array.len(); + group.push(piece.array); + if group_rows >= MIN_OUTPUT_ROWS { + batches.push_back(self.concat(std::mem::take(&mut group))?); + group_rows = 0; + } + } + if !group.is_empty() { + batches.push_back(self.concat(group)?); + } + Ok(batches) + } + + fn concat(&self, mut arrays: Vec) -> VortexResult { + if arrays.len() == 1 { + return Ok(arrays.remove(0)); + } + Ok(ChunkedArray::try_new(arrays, self.dtype.clone())?.into_array()) } } @@ -148,8 +171,14 @@ impl IoConsumer for ProjectionMorsel { impl Morsel for ProjectionMorsel { fn state(&self) -> State { + if let Some(batches) = &self.batches { + return if batches.is_empty() { + State::Done + } else { + State::NeedsCompute + }; + } match &self.adaptor { - _ if self.done => State::Done, None => State::NeedsCompute, Some(adaptor) => match adaptor.state() { State::Done => State::NeedsCompute, @@ -159,8 +188,11 @@ impl Morsel for ProjectionMorsel { } fn compute(&mut self) -> VortexResult { - if self.done { - vortex_bail!("ProjectionMorsel: compute called after Done"); + if let Some(batches) = &mut self.batches { + return match batches.pop_front() { + Some(batch) => Ok(MorselOutput::Batch(batch)), + None => vortex_bail!("ProjectionMorsel: compute called after Done"), + }; } let Some(adaptor) = self.adaptor.as_mut() else { let scan = Scan::try_new( @@ -181,7 +213,8 @@ impl Morsel for ProjectionMorsel { return Ok(MorselOutput::Continue); }; if adaptor.state() == State::Done { - return self.finish(); + self.batches = Some(self.group_pieces()?); + return Ok(MorselOutput::Continue); } Ok(match adaptor.compute()? { AdaptorStep::Yield => MorselOutput::Continue, diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index b69f5390704..401bdc480a8 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -21,6 +21,7 @@ use vortex_array::ArrayContext; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::Chunked; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; @@ -698,6 +699,51 @@ async fn primitive_column() -> VortexResult<()> { Ok(()) } +/// Stored chunks within one projection morsel are emitted as separate batches rather than one +/// chunked array, which the Arrow export would concatenate. +#[tokio::test] +async fn large_projection_pieces_are_separate_batches() -> VortexResult<()> { + const PIECE_ROWS: usize = MAX_ROWS_PER_CHUNK / 4; + let session = new_session().with_tokio(); + let segments = Arc::new(TestSegments::default()); + let (mut sequence, eof) = SequenceId::root().split(); + let arrays = (0..4) + .map(|index| { + let start = i32::try_from(index * PIECE_ROWS)?; + Ok(( + sequence.advance(), + PrimitiveArray::from_iter(start..start + i32::try_from(PIECE_ROWS)?).into_array(), + )) + }) + .collect::>>()?; + let dtype = arrays[0].1.dtype().clone(); + let layout = chunked_flat() + .write_stream( + ArrayContext::empty().into(), + Arc::::clone(&segments), + SequentialStreamAdapter::new(dtype, stream::iter(arrays.into_iter().map(Ok))) + .sendable(), + eof, + &session, + ) + .await?; + let segments: Arc = segments; + let scan = ScanRun::try_new(builder(&session, &segments, &layout)?)?; + let batches: Vec<_> = scan.execute_stream()?.try_collect().await?; + assert_eq!( + batches.iter().map(|batch| batch.len()).collect::>(), + vec![PIECE_ROWS; 4] + ); + assert!(batches.iter().all(|batch| !batch.is::())); + let total = i32::try_from(4 * PIECE_ROWS)?; + assert_arrays_eq!( + ChunkedArray::try_new(batches, layout.dtype().clone())?, + PrimitiveArray::from_iter(0..total), + &mut session.create_execution_ctx() + ); + Ok(()) +} + /// Stored chunk boundaries do not subdivide a fixed work chunk. #[rstest] #[tokio::test]