perf(scan): bound V2 row chunks by serialized size

Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk>

Author: joseph-isaacs

vortex-layout/src/scan/v2/mod.rs

@@ -11,7 +11,8 @@
 //! Import this module's [`ScanBuilder`] to select V2 explicitly. It accepts a [`ScanFile`]
 //! directly and supports projection, row ranges, and selections. Unsupported layouts return an
 //! error; filters and limits are not part of this builder's API. Work is partitioned into fixed
-//! 65,536-row chunks, independently of the stored layout's chunk boundaries.
+//! chunks of up to 65,536 rows, independently of the stored layout's chunk boundaries.
+//! Full-file projections use smaller 8,192-row multiples when serialized bytes per row are high.
 
 mod io;
 mod split;
@@ -63,19 +64,34 @@ use crate::scan::v2::io::segment_ranges;
 use crate::scan::v2::split::SplitPlan;
 use crate::segments::SegmentSource;
 
-/// Eight of the default writer's 8,192-row blocks per scan task.
-const ROWS_PER_CHUNK: usize = 8 * 8192;
+const ROW_BLOCK_SIZE: usize = 8192;
+const MAX_ROWS_PER_CHUNK: usize = 8 * ROW_BLOCK_SIZE;
+const TARGET_BYTES_PER_CHUNK: u64 = 8 * 1024 * 1024;
+
+/// Limit estimated bytes per task without changing the writer-aligned row granularity.
+fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> usize {
+    if row_count <= MAX_ROWS_PER_CHUNK as u64 {
+        return MAX_ROWS_PER_CHUNK;
+    }
+    let bytes: u64 = locations
+        .iter()
+        .map(|location| u64::from(location.length))
+        .sum();
+    let bytes_per_row = bytes.div_ceil(row_count).max(1);
+    let blocks = TARGET_BYTES_PER_CHUNK / bytes_per_row / ROW_BLOCK_SIZE as u64;
+    blocks.clamp(1, 8) as usize * ROW_BLOCK_SIZE
+}
 
 /// Fixed boundaries in the file's row domain, clipped to the requested range.
-fn row_chunks(rows: Range<u64>) -> Vec<Range<u64>> {
+fn row_chunks(rows: Range<u64>, rows_per_chunk: usize) -> Vec<Range<u64>> {
     if rows.is_empty() {
         return Vec::new();
     }
-    let first = rows.start / ROWS_PER_CHUNK as u64 * ROWS_PER_CHUNK as u64;
+    let first = rows.start / rows_per_chunk as u64 * rows_per_chunk as u64;
     (first..rows.end)
-        .step_by(ROWS_PER_CHUNK)
+        .step_by(rows_per_chunk)
         .map(|start| {
-            start.max(rows.start)..start.saturating_add(ROWS_PER_CHUNK as u64).min(rows.end)
+            start.max(rows.start)..start.saturating_add(rows_per_chunk as u64).min(rows.end)
         })
         .collect()
 }
@@ -282,7 +298,12 @@ impl<A: 'static + Send> PreparedScan<A> {
         } else {
             optimize(plan_row_idx_expression(builder.projection, root)?)?
         };
-        let splits = row_chunks(builder.row_range.unwrap_or(0..row_count));
+        let chunk_rows = if identity {
+            rows_per_chunk(row_count, &file.locations)
+        } else {
+            MAX_ROWS_PER_CHUNK
+        };
+        let splits = row_chunks(builder.row_range.unwrap_or(0..row_count), chunk_rows);
         let io = Arc::new(SegmentScanIo::new(
             file.segments,
             segment_ranges(&file.locations),

vortex-layout/src/scan/v2/tests.rs

@@ -38,10 +38,11 @@ use vortex_scan::strict_sorted_buffer::StrictSortedBuffer;
 use vortex_session::VortexSession;
 
 use super::PreparedScan;
-use super::ROWS_PER_CHUNK;
+use super::MAX_ROWS_PER_CHUNK;
 use super::ScanBuilder;
 use super::ScanFile;
 use super::row_chunks;
+use super::rows_per_chunk;
 use crate::LayoutRef;
 use crate::LayoutStrategy;
 use crate::layouts::chunked::writer::ChunkedLayoutStrategy;
@@ -309,7 +310,31 @@ async fn arrow_stream_projects_across_chunks() -> VortexResult<()> {
 #[case::empty(5..5, vec![])]
 #[case::near_max(u64::MAX - 3..u64::MAX, vec![u64::MAX - 3..u64::MAX])]
 fn fixed_row_chunks(#[case] rows: Range<u64>, #[case] expected: Vec<Range<u64>>) {
-    assert_eq!(row_chunks(rows), expected);
+    assert_eq!(row_chunks(rows, 65_536), expected);
+}
+
+#[rstest]
+#[case::empty(0, 0, 65_536)]
+#[case::single_chunk(1000, 100_000_000, 65_536)]
+#[case::no_bytes(100_000, 0, 65_536)]
+#[case::small_rows(1_000_000, 64_000_000, 65_536)]
+#[case::large_rows(1_000_000, 384_000_000, 16_384)]
+#[case::minimum(1_000_000, 2_000_000_000, 8192)]
+fn chunk_byte_budget(#[case] rows: u64, #[case] bytes: u32, #[case] expected: usize) {
+    let locations = [SegmentLocation {
+        offset: 0,
+        length: bytes,
+        alignment: Alignment::none(),
+    }];
+    assert_eq!(rows_per_chunk(rows, &locations), expected);
+}
+
+#[test]
+fn smaller_chunks_keep_file_alignment() {
+    assert_eq!(
+        row_chunks(7500..33_000, 16_384),
+        [7500..16_384, 16_384..32_768, 32_768..33_000]
+    );
 }
 
 struct CountingSegments {
@@ -362,7 +387,7 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> {
     let splits = scan.split_plans();
     assert_eq!(
         splits.len(),
-        (STORED_CHUNKS * STORED_CHUNK_ROWS).div_ceil(ROWS_PER_CHUNK)
+        (STORED_CHUNKS * STORED_CHUNK_ROWS).div_ceil(MAX_ROWS_PER_CHUNK)
     );
     for (index, split) in splits.into_iter().enumerate() {
         let rows = split.scope.rows.clone();