From bdd36d1309e3e6888371798d4d4b08414503fc55 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Tue, 28 Jul 2026 00:31:33 +0000 Subject: [PATCH] feat(scan): make sub-splitting of wide chunk spans configurable MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Sub-division of chunk-boundary spans wider than the ideal split size (so a file with few, large chunks decodes across multiple cores) is now a scan config option rather than unconditional behaviour. `SplitBy::Layout` now splits only at the file's own chunk boundaries, and the new `SplitBy::LayoutSubSplitting` variant — the default — additionally sub-divides wide spans. `ScanBuilder::with_no_sub_splitting()` is a shorthand for selecting `SplitBy::Layout`. Callers that relied on the sub-dividing default (`VortexFile::splits`, the DataFusion repartitioner, Python datasets) now name `LayoutSubSplitting` explicitly, so their behaviour is unchanged. Signed-off-by: Robert Kruszewski Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01M9d52UfuZRn3cC9eixruzx --- vortex-datafusion/src/persistent/opener.rs | 2 +- vortex-file/src/file.rs | 4 +- vortex-file/src/tests.rs | 72 +++++---- vortex-layout/src/scan/scan_builder.rs | 12 +- vortex-layout/src/scan/split_by.rs | 168 +++++++++++++-------- vortex-python/src/dataset.rs | 12 +- 6 files changed, 172 insertions(+), 98 deletions(-) diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index d0ffb472ebe..63ecebabba0 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -505,7 +505,7 @@ fn natural_split_ranges_for_file( fn compute_natural_split_ranges(layout_reader: &dyn LayoutReader) -> DFResult]>> { let row_count = layout_reader.row_count(); let row_range = 0..row_count; - let split_points: Vec<_> = SplitBy::Layout + let split_points: Vec<_> = SplitBy::LayoutSubSplitting .splits(layout_reader, &row_range, &[FieldMask::All]) .map_err(|e| exec_datafusion_err!("Failed to compute Vortex natural splits: {e}"))? .into_iter() diff --git a/vortex-file/src/file.rs b/vortex-file/src/file.rs index 739d7c9fd8b..209432ab86c 100644 --- a/vortex-file/src/file.rs +++ b/vortex-file/src/file.rs @@ -247,10 +247,10 @@ impl VortexFile { /// Return the file's natural row splits as root-coordinate ranges. /// - /// These are the ranges that [`SplitBy::Layout`] would use for an all-fields scan. + /// These are the ranges that [`SplitBy::LayoutSubSplitting`] would use for an all-fields scan. pub fn splits(&self) -> VortexResult>> { let reader = self.layout_reader()?; - Ok(SplitBy::Layout + Ok(SplitBy::LayoutSubSplitting .splits(reader.as_ref(), &(0..reader.row_count()), &[FieldMask::All])? .into_iter() .tuple_windows() diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 44c0b76d2db..816710323b8 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -2076,19 +2076,19 @@ fn layout_has_dict(layout: &dyn DynLayout) -> bool { .any(|child| layout_has_dict(child.as_ref())) } -/// Mirrors the (private) `IDEAL_SPLIT_SIZE` that `SplitBy::Layout` uses to sub-divide wide -/// chunk-boundary spans: layout splits are never wider than this many rows. +/// Mirrors the (private) `IDEAL_SPLIT_SIZE` that `SplitBy::LayoutSubSplitting` uses to sub-divide +/// wide chunk-boundary spans: layout splits are never wider than this many rows. const MAX_SPLIT_ROWS: u64 = 100_000; -#[tokio::test] -#[cfg_attr(miri, ignore)] -async fn test_large_flat_chunk_scan_subdivides_splits() -> VortexResult<()> { - // A single flat (unchunked) 250k-row layout spans the 100k sub-split threshold, so the scan - // must decode it as multiple row-range splits. - let mut ctx = SESSION.create_execution_ctx(); - const N_ROWS: u64 = 250_000; +/// Rows in the [`large_flat_file`] fixture; spans the sub-split threshold. +const FLAT_N_ROWS: u64 = 250_000; + +/// A single flat (unchunked) [`FLAT_N_ROWS`]-row layout with alternating-sign values, so filters +/// select rows on both sides of any split boundary. Returns the opened file and original array. +async fn large_flat_file() -> VortexResult<(VortexFile, ArrayRef)> { let values = - Buffer::from_iter((0..N_ROWS as i32).map(|i| if i % 2 == 0 { i } else { -i })).into_array(); + Buffer::from_iter((0..FLAT_N_ROWS as i32).map(|i| if i % 2 == 0 { i } else { -i })) + .into_array(); let mut buf = ByteBufferMut::empty(); SESSION @@ -2097,14 +2097,23 @@ async fn test_large_flat_chunk_scan_subdivides_splits() -> VortexResult<()> { .write(&mut buf, values.to_array_stream()) .await?; - let file = SESSION.open_options().open_buffer(buf)?; + Ok((SESSION.open_options().open_buffer(buf)?, values)) +} + +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn test_large_flat_chunk_scan_subdivides_splits() -> VortexResult<()> { + // A single flat (unchunked) 250k-row layout spans the 100k sub-split threshold, so the scan + // must decode it as multiple row-range splits. + let mut ctx = SESSION.create_execution_ctx(); + let (file, values) = large_flat_file().await?; // Sub-division caps each split at MAX_SPLIT_ROWS while tiling the file exactly. let splits = file.splits()?; assert!(splits.len() > 1, "expected sub-divided splits: {splits:?}"); assert!(splits.iter().all(|r| r.end - r.start <= MAX_SPLIT_ROWS)); assert_eq!(splits.first().map(|r| r.start), Some(0)); - assert_eq!(splits.last().map(|r| r.end), Some(N_ROWS)); + assert_eq!(splits.last().map(|r| r.end), Some(FLAT_N_ROWS)); assert!(splits.windows(2).all(|w| w[0].end == w[1].start)); // A full scan across the sub-splits returns the original rows. @@ -2119,12 +2128,32 @@ async fn test_large_flat_chunk_scan_subdivides_splits() -> VortexResult<()> { .read_all() .await?; let expected = - Buffer::from_iter((0..N_ROWS as i32).filter(|i| i % 2 == 0 && *i > 0)).into_array(); + Buffer::from_iter((0..FLAT_N_ROWS as i32).filter(|i| i % 2 == 0 && *i > 0)).into_array(); assert_arrays_eq!(result, expected, &mut ctx); Ok(()) } +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn test_no_sub_splitting_keeps_large_chunk_whole() -> VortexResult<()> { + // The same over-wide single chunk as above, scanned with sub-splitting disabled: the scan + // follows the layout's chunk boundaries exactly, so the file decodes as one batch. + let mut ctx = SESSION.create_execution_ctx(); + let (file, values) = large_flat_file().await?; + + let mut chunks: Vec = file + .scan()? + .with_no_sub_splitting() + .into_array_stream()? + .try_collect() + .await?; + assert_eq!(chunks.len(), 1, "expected a single un-split chunk"); + assert_arrays_eq!(chunks.remove(0), values, &mut ctx); + + Ok(()) +} + #[rstest] #[case::unaligned(33_333)] #[case::exceeds_file(300_000)] @@ -2137,18 +2166,7 @@ async fn test_flat_chunk_scan_with_row_count_splits( // results whether the split size straddles the chunk arbitrarily or exceeds the file's // row count (a single split). let mut ctx = SESSION.create_execution_ctx(); - const N_ROWS: u64 = 250_000; - let values = - Buffer::from_iter((0..N_ROWS as i32).map(|i| if i % 2 == 0 { i } else { -i })).into_array(); - - let mut buf = ByteBufferMut::empty(); - SESSION - .write_options() - .with_strategy(Arc::new(FlatLayoutStrategy::default())) - .write(&mut buf, values.to_array_stream()) - .await?; - - let file = SESSION.open_options().open_buffer(buf)?; + let (file, values) = large_flat_file().await?; let result = file .scan()? @@ -2166,7 +2184,7 @@ async fn test_flat_chunk_scan_with_row_count_splits( .read_all() .await?; let expected = - Buffer::from_iter((0..N_ROWS as i32).filter(|i| i % 2 == 0 && *i > 0)).into_array(); + Buffer::from_iter((0..FLAT_N_ROWS as i32).filter(|i| i % 2 == 0 && *i > 0)).into_array(); assert_arrays_eq!(result, expected, &mut ctx); Ok(()) @@ -2177,7 +2195,7 @@ async fn test_flat_chunk_scan_with_row_count_splits( async fn test_string_chunks_stay_fine_grained_under_split_cap() -> VortexResult<()> { // Default writing targets ~1MiB uncompressed blocks, so ~120-byte strings chunk at a few // thousand rows (~8k with today's defaults). These natural boundaries sit far below the - // sub-split cap, and SplitBy::Layout must pass them through untouched. + // sub-split cap, and SplitBy::LayoutSubSplitting must pass them through untouched. let mut ctx = SESSION.create_execution_ctx(); const N_ROWS: usize = 40_000; let strings = VarBinArray::from_iter( diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index 11fd5c7b882..674a1143239 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -95,7 +95,7 @@ impl ScanBuilder { ordered: true, row_range: None, selection: Default::default(), - split_by: SplitBy::Layout, + split_by: SplitBy::LayoutSubSplitting, // We default to four tasks per worker thread, which allows for some I/O lookahead // without too much impact on work-stealing. concurrency: 4, @@ -190,6 +190,16 @@ impl ScanBuilder { self } + /// Split only at the layout's own chunk boundaries, without sub-dividing wide chunk spans. + /// + /// By default ([`SplitBy::LayoutSubSplitting`]) spans between adjacent chunk boundaries that + /// are wider than the ideal split size are sub-divided, so a file with few, large chunks + /// decodes across multiple cores. This shorthand for `with_split_by(SplitBy::Layout)` disables + /// that sub-division, yielding fewer, larger splits that follow the file's chunking exactly. + pub fn with_no_sub_splitting(self) -> Self { + self.with_split_by(SplitBy::Layout) + } + /// Returns the per-worker row-split concurrency. pub fn concurrency(&self) -> usize { self.concurrency diff --git a/vortex-layout/src/scan/split_by.rs b/vortex-layout/src/scan/split_by.rs index 6c2cde11b72..ff348776dfa 100644 --- a/vortex-layout/src/scan/split_by.rs +++ b/vortex-layout/src/scan/split_by.rs @@ -24,11 +24,15 @@ const MAX_SPLIT_ROWS: u64 = IDEAL_SPLIT_SIZE; /// Note that each split must fit into the platform's maximum usize. #[derive(Default, Copy, Clone, Debug)] pub enum SplitBy { - #[default] - /// Splits any time there is a chunk boundary in the file. Spans between adjacent boundaries - /// wider than `MAX_SPLIT_ROWS` are further sub-divided so that a file with few, large chunks - /// can still be decoded across multiple cores. + /// Splits any time there is a chunk boundary in the file, and nowhere else. This yields + /// splits that follow the file's chunking exactly, trading intra-file decode parallelism + /// for fewer, larger batches. Layout, + #[default] + /// Splits like [`SplitBy::Layout`], except that spans between adjacent chunk boundaries + /// wider than `MAX_SPLIT_ROWS` are further sub-divided so that a file with few, large + /// chunks can still be decoded across multiple cores. + LayoutSubSplitting, /// Splits every n rows. RowCount(usize), // UncompressedSize(u64), @@ -44,7 +48,7 @@ impl SplitBy { field_mask: &[FieldMask], ) -> VortexResult> { Ok(match *self { - SplitBy::Layout => { + SplitBy::Layout | SplitBy::LayoutSubSplitting => { // We usually have under 100 splits so reserving upfront saves // us some allocations let mut row_splits = RowSplits::new_capacity(128); @@ -54,7 +58,12 @@ impl SplitBy { &SplitRange::root(row_range.clone())?, &mut row_splits, )?; - subdivide_large_spans(row_splits.into_sorted_deduped(), MAX_SPLIT_ROWS) + let boundaries = row_splits.into_sorted_deduped(); + if matches!(self, SplitBy::LayoutSubSplitting) { + subdivide_large_spans(boundaries, MAX_SPLIT_ROWS) + } else { + boundaries + } } SplitBy::RowCount(n) => row_range .clone() @@ -203,78 +212,107 @@ mod test { assert_eq!(splits, vec![0u64, 3, 6, 9, 10]); } - #[test] - fn test_layout_splits_dedup() { - struct DupReader { - name: Arc, - dtype: DType, - } + /// A reader that registers a fixed set of interior split points (which may repeat) plus the + /// end of its row range. + struct StubReader { + name: Arc, + dtype: DType, + row_count: u64, + interior_splits: Vec, + } - impl LayoutReader for DupReader { - fn name(&self) -> &Arc { - &self.name + impl StubReader { + fn new(row_count: u64, interior_splits: Vec) -> Self { + Self { + name: Arc::from("stub"), + dtype: DType::Primitive(PType::U8, Nullability::NonNullable), + row_count, + interior_splits, } + } + } - fn as_any(&self) -> &dyn Any { - self - } + impl LayoutReader for StubReader { + fn name(&self) -> &Arc { + &self.name + } - fn dtype(&self) -> &DType { - &self.dtype - } + fn as_any(&self) -> &dyn Any { + self + } - fn row_count(&self) -> u64 { - 10 - } + fn dtype(&self) -> &DType { + &self.dtype + } - fn register_splits( - &self, - _field_mask: &[FieldMask], - split_range: &SplitRange, - splits: &mut RowSplits, - ) -> VortexResult<()> { - splits.push(split_range.row_offset() + 5); - splits.push(split_range.row_offset() + 5); - splits.push(split_range.root_row_range().end); - Ok(()) - } + fn row_count(&self) -> u64 { + self.row_count + } - fn pruning_evaluation( - &self, - _: &Range, - _: &Expression, - _: Mask, - ) -> VortexResult { - unimplemented!() + fn register_splits( + &self, + _field_mask: &[FieldMask], + split_range: &SplitRange, + splits: &mut RowSplits, + ) -> VortexResult<()> { + for split in &self.interior_splits { + splits.push(split_range.row_offset() + split); } + splits.push(split_range.root_row_range().end); + Ok(()) + } - fn filter_evaluation( - &self, - _: &Range, - _: &Expression, - _: MaskFuture, - ) -> VortexResult { - unimplemented!() - } + fn pruning_evaluation( + &self, + _: &Range, + _: &Expression, + _: Mask, + ) -> VortexResult { + unimplemented!() + } - fn projection_evaluation( - &self, - _: &Range, - _: &Expression, - _: MaskFuture, - ) -> VortexResult>> { - unimplemented!() - } + fn filter_evaluation( + &self, + _: &Range, + _: &Expression, + _: MaskFuture, + ) -> VortexResult { + unimplemented!() } - let reader = DupReader { - name: Arc::from("dup"), - dtype: DType::Primitive(PType::U8, Nullability::NonNullable), - }; - let splits = SplitBy::Layout - .splits(&reader, &(0..10), &[FieldMask::All]) - .unwrap(); + fn projection_evaluation( + &self, + _: &Range, + _: &Expression, + _: MaskFuture, + ) -> VortexResult>> { + unimplemented!() + } + } + + #[test] + fn test_layout_splits_dedup() -> VortexResult<()> { + let reader = StubReader::new(10, vec![5, 5]); + let splits = SplitBy::Layout.splits(&reader, &(0..10), &[FieldMask::All])?; assert_eq!(splits, vec![0u64, 5, 10]); + Ok(()) + } + + #[test] + fn test_layout_keeps_large_chunk_whole() -> VortexResult<()> { + // A single chunk wider than MAX_SPLIT_ROWS: sub-divided by SplitBy::LayoutSubSplitting, + // left whole by SplitBy::Layout. + let row_count = MAX_SPLIT_ROWS * 2 + 1; + let reader = StubReader::new(row_count, vec![]); + + let splits = + SplitBy::LayoutSubSplitting.splits(&reader, &(0..row_count), &[FieldMask::All])?; + assert!(splits.len() > 2, "expected sub-divided splits: {splits:?}"); + + let splits = SplitBy::Layout.splits(&reader, &(0..row_count), &[FieldMask::All])?; + assert_eq!(splits, vec![0, row_count]); + + Ok(()) } #[test] diff --git a/vortex-python/src/dataset.rs b/vortex-python/src/dataset.rs index 504ed05aeee..46c343e4341 100644 --- a/vortex-python/src/dataset.rs +++ b/vortex-python/src/dataset.rs @@ -189,7 +189,11 @@ impl PyVortexDataset { .scan()? .with_projection(projection) .with_some_filter(filter) - .with_split_by(split_by.map(SplitBy::RowCount).unwrap_or(SplitBy::Layout)); + .with_split_by( + split_by + .map(SplitBy::RowCount) + .unwrap_or(SplitBy::LayoutSubSplitting), + ); if let Some((l, r)) = row_range { scan = scan.with_row_range(l..r); } @@ -228,7 +232,11 @@ impl PyVortexDataset { .scan()? .with_projection(select(FieldNames::empty(), root())) .with_some_filter(filter) - .with_split_by(split_by.map(SplitBy::RowCount).unwrap_or(SplitBy::Layout)); + .with_split_by( + split_by + .map(SplitBy::RowCount) + .unwrap_or(SplitBy::LayoutSubSplitting), + ); if let Some((l, r)) = row_range { scan = scan.with_row_range(l..r); }