Split SplitBy::Layout into Layout and LayoutSubSplitting variants

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<Arc<[Range<u64>]>> {
     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()

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<Vec<Range<u64>>> {
         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()

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<ArrayRef> = 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(

vortex-layout/src/scan/scan_builder.rs

@@ -95,7 +95,7 @@ impl ScanBuilder<ArrayRef> {
             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<A: 'static + Send> ScanBuilder<A> {
         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

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<Vec<u64>> {
         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<str>,
-            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<str>,
+        dtype: DType,
+        row_count: u64,
+        interior_splits: Vec<u64>,
+    }
 
-        impl LayoutReader for DupReader {
-            fn name(&self) -> &Arc<str> {
-                &self.name
+    impl StubReader {
+        fn new(row_count: u64, interior_splits: Vec<u64>) -> 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<str> {
+            &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<u64>,
-                _: &Expression,
-                _: Mask,
-            ) -> VortexResult<MaskFuture> {
-                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<u64>,
-                _: &Expression,
-                _: MaskFuture,
-            ) -> VortexResult<MaskFuture> {
-                unimplemented!()
-            }
+        fn pruning_evaluation(
+            &self,
+            _: &Range<u64>,
+            _: &Expression,
+            _: Mask,
+        ) -> VortexResult<MaskFuture> {
+            unimplemented!()
+        }
 
-            fn projection_evaluation(
-                &self,
-                _: &Range<u64>,
-                _: &Expression,
-                _: MaskFuture,
-            ) -> VortexResult<BoxFuture<'static, VortexResult<ArrayRef>>> {
-                unimplemented!()
-            }
+        fn filter_evaluation(
+            &self,
+            _: &Range<u64>,
+            _: &Expression,
+            _: MaskFuture,
+        ) -> VortexResult<MaskFuture> {
+            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<u64>,
+            _: &Expression,
+            _: MaskFuture,
+        ) -> VortexResult<BoxFuture<'static, VortexResult<ArrayRef>>> {
+            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]