Size filtered V2 scan splits by projected bytes and threads

vortex-layout/src/scan/execution/chunk.rs

@@ -7,6 +7,7 @@ use std::ops::Range;
 use std::sync::Arc;
 use std::vec::IntoIter;
 
+use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_io::request::IoConsumer;
 use vortex_io::request::IoRequestId;
@@ -35,13 +36,41 @@ pub(crate) fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> u
     if row_count <= MAX_ROWS_PER_CHUNK as u64 {
         return MAX_ROWS_PER_CHUNK;
     }
-    let bytes: u64 = locations
+    let blocks = blocks_per_target_bytes(total_bytes(locations), row_count);
+    blocks.clamp(1, 8) as usize * ROW_BLOCK_SIZE
+}
+
+/// Rows per chunk of a filtered scan over `rows` rows of the file.
+///
+/// Large chunks amortize filter evaluation over cheap projections. A projection that reads many
+/// bytes per row, or a row range that would otherwise give fewer chunks than threads, gets smaller
+/// chunks so the surviving rows still spread across the runtime's threads.
+pub(crate) fn filtered_rows_per_chunk(
+    rows: u64,
+    row_count: u64,
+    projected_bytes: u64,
+    parallelism: usize,
+) -> usize {
+    let max_blocks = (FILTERED_ROWS_PER_CHUNK / ROW_BLOCK_SIZE) as u64;
+    let min_blocks = (MAX_ROWS_PER_CHUNK / ROW_BLOCK_SIZE) as u64;
+    let by_bytes = blocks_per_target_bytes(projected_bytes, row_count).max(1);
+    let by_threads = (rows / parallelism.max(1) as u64 / ROW_BLOCK_SIZE as u64).max(min_blocks);
+    let blocks = by_bytes.min(by_threads).min(max_blocks);
+    usize::try_from(blocks).vortex_expect("at most FILTERED_ROWS_PER_CHUNK blocks") * ROW_BLOCK_SIZE
+}
+
+/// The row blocks that hold [`TARGET_BYTES_PER_CHUNK`] when `bytes` spread over `row_count` rows.
+fn blocks_per_target_bytes(bytes: u64, row_count: u64) -> u64 {
+    let bytes_per_row = bytes.div_ceil(row_count.max(1)).max(1);
+    TARGET_BYTES_PER_CHUNK / bytes_per_row / ROW_BLOCK_SIZE as u64
+}
+
+/// The serialized bytes of every segment in `locations`.
+pub(crate) fn total_bytes(locations: &[SegmentLocation]) -> 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
+        .sum()
 }
 
 /// Fixed boundaries in the file's row domain, clipped to the requested range.

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

@@ -31,12 +31,15 @@ use vortex_scan::selection::Selection;
 use vortex_session::VortexSession;
 
 pub use self::chunk::ChunkPlanner;
+#[cfg(test)]
 pub(crate) use self::chunk::FILTERED_ROWS_PER_CHUNK;
 pub(crate) use self::chunk::MAX_ROWS_PER_CHUNK;
 pub(crate) use self::chunk::ROW_BLOCK_SIZE;
+pub(crate) use self::chunk::filtered_rows_per_chunk;
 #[cfg(test)]
 pub(crate) use self::chunk::row_chunks;
 pub(crate) use self::chunk::rows_per_chunk;
+pub(crate) use self::chunk::total_bytes;
 use self::filter::FilterPlanner;
 pub use self::filter::FilterPlans;
 use self::projection::ProjectionPlanner;

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

@@ -10,6 +10,7 @@ use parking_lot::Mutex;
 use rustc_hash::FxHashMap;
 use vortex_array::expr::BoundExpression;
 use vortex_array::expr::ExactBoundExpr;
+use vortex_error::VortexResult;
 
 use super::io::SegmentRanges;
 use crate::plan::physical::PlanRef;
@@ -43,6 +44,8 @@ struct Inner {
     /// Plans of filters, keyed by the exact filter expression. A pruning plan remembers the zones it
     /// excluded, so sharing the plans prunes a file's zones once for every scan of that filter.
     filters: Mutex<FxHashMap<ExactBoundExpr, Arc<FilterPlans>>>,
+    /// Serialized bytes the whole-file projection reads, by projection.
+    projection_bytes: Mutex<FxHashMap<ExactBoundExpr, u64>>,
     values: ValueCache,
 }
 
@@ -112,6 +115,22 @@ impl ScanCache {
         )
     }
 
+    /// The serialized bytes a whole-file scan of `projection` reads, computing them with
+    /// `compute` on first use.
+    pub(super) fn projection_bytes(
+        &self,
+        projection: &BoundExpression,
+        compute: impl FnOnce() -> VortexResult<u64>,
+    ) -> VortexResult<u64> {
+        let key = ExactBoundExpr(projection.clone());
+        if let Some(bytes) = self.0.projection_bytes.lock().get(&key) {
+            return Ok(*bytes);
+        }
+        let bytes = compute()?;
+        self.0.projection_bytes.lock().insert(key, bytes);
+        Ok(bytes)
+    }
+
     pub(super) fn values(&self) -> ValueCache {
         self.0.values.clone()
     }

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

@@ -13,9 +13,10 @@
 //! Import this module's [`ScanBuilder`] to prepare a V2 scan. It accepts a [`ScanFile`] directly
 //! and supports projection, filters, row ranges, and selections. Unsupported layouts return an
 //! error when the plan is prepared; limits are not part of this builder's API. Work is cut into
-//! chunks of up to 65,536 rows, or 1,048,576 rows for filtered scans, independently of the stored
-//! layout's chunk boundaries. Full-file projections use smaller 8,192-row multiples when
-//! serialized bytes per row are high.
+//! chunks of up to 65,536 rows, or up to 1,048,576 rows for filtered scans, independently of the
+//! stored layout's chunk boundaries. Filtered scans use smaller chunks when the projection reads
+//! many bytes per row or the row range would give fewer chunks than threads. Full-file
+//! projections use smaller 8,192-row multiples when serialized bytes per row are high.
 //!
 //! A [`ScanFile`] carrying a [`ScanCache`] shares its lowered plan, filter plans and pruned
 //! zones, segment index, and dictionary values with every other scan of that file, like a cached
@@ -33,26 +34,30 @@ use vortex_array::dtype::DType;
 use vortex_array::expr::BoundExpression;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
+use vortex_mask::Mask;
 use vortex_scan::selection::Selection;
 use vortex_scan::strict_sorted_buffer::StrictSortedBuffer;
 use vortex_session::VortexSession;
+use vortex_utils::parallelism::get_available_parallelism;
 
 pub use self::cache::ScanCache;
 pub use self::cache::WeakScanCache;
 use crate::LayoutRef;
 use crate::plan::physical::lower;
 use crate::plan::physical::optimize;
+use crate::plan::physical::pipeline::Scan;
 use crate::plan::physical::plan_row_idx_expression;
 pub use crate::scan::execution::ChunkPlanner;
-use crate::scan::execution::FILTERED_ROWS_PER_CHUNK;
 pub use crate::scan::execution::FilterPlans;
 use crate::scan::execution::MAX_ROWS_PER_CHUNK;
 pub use crate::scan::execution::PreparedScan;
 pub use crate::scan::execution::PruneNext;
 pub use crate::scan::execution::PrunePlanner;
 use crate::scan::execution::ROW_BLOCK_SIZE;
 pub use crate::scan::execution::SegmentLocation;
+use crate::scan::execution::filtered_rows_per_chunk;
 use crate::scan::execution::rows_per_chunk;
+use crate::scan::execution::total_bytes;
 use crate::scan::v2::io::segment_ranges;
 use crate::segments::SegmentSource;
 
@@ -187,11 +192,33 @@ impl ScanBuilder {
                     plan_row_idx_expression(self.projection.clone(), root)?,
                     &self.session,
                 )?;
-                cache.insert_projection(self.projection, plan)
+                cache.insert_projection(self.projection.clone(), plan)
             }
         };
+        let rows = self.row_range.unwrap_or(0..row_count);
         let chunk_rows = if filter.is_some() {
-            FILTERED_ROWS_PER_CHUNK
+            let projected_bytes = if identity {
+                total_bytes(&file.locations)
+            } else {
+                cache.projection_bytes(&self.projection, || {
+                    let segments = Scan::segments(
+                        &self.session,
+                        &projection,
+                        0..row_count,
+                        &Mask::new_true(usize::try_from(row_count)?),
+                    )?;
+                    Ok(segments
+                        .iter()
+                        .map(|segment| u64::from(file.locations[**segment as usize].length))
+                        .sum())
+                })?
+            };
+            filtered_rows_per_chunk(
+                rows.end - rows.start,
+                row_count,
+                projected_bytes,
+                get_available_parallelism().unwrap_or(1),
+            )
         } else if identity {
             rows_per_chunk(row_count, &file.locations)
         } else {
@@ -204,7 +231,7 @@ impl ScanBuilder {
         let values = cache.values();
         Ok(PreparedScan {
             session: self.session,
-            rows: self.row_range.unwrap_or(0..row_count),
+            rows,
             selection: self.selection,
             segments: file.segments,
             ranges,

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

@@ -71,7 +71,9 @@ use crate::layouts::table::TableStrategy;
 use crate::layouts::zoned::writer::ZonedLayoutOptions;
 use crate::layouts::zoned::writer::ZonedStrategy;
 use crate::scan::arrow::to_record_batch;
+use crate::scan::execution::FILTERED_ROWS_PER_CHUNK;
 use crate::scan::execution::SegmentLocation;
+use crate::scan::execution::filtered_rows_per_chunk;
 use crate::scan::execution::row_chunks;
 use crate::scan::scan_builder::ScanBuilder as V1ScanBuilder;
 use crate::segments::SegmentFuture;
@@ -871,3 +873,22 @@ async fn shared_cache_prunes_zones_once() -> VortexResult<()> {
     assert_eq!(reads, vec![2, 1]);
     Ok(())
 }
+
+#[rstest]
+#[case::cheap_projection_large_range(16 << 20, 16 << 20, 16 << 20, 4, FILTERED_ROWS_PER_CHUNK)]
+#[case::wide_projection(1 << 20, 1 << 20, 100 << 20, 1, 10 * ROW_BLOCK_SIZE)]
+#[case::very_wide_projection(1 << 20, 1 << 20, 10 << 30, 1, ROW_BLOCK_SIZE)]
+#[case::partition_spread_over_threads(1 << 20, 16 << 20, 16 << 20, 16, MAX_ROWS_PER_CHUNK)]
+#[case::partition_split_per_thread(1 << 20, 16 << 20, 16 << 20, 4, 32 * ROW_BLOCK_SIZE)]
+fn filtered_chunks_shrink_for_wide_projections_and_few_rows(
+    #[case] rows: u64,
+    #[case] row_count: u64,
+    #[case] projected_bytes: u64,
+    #[case] parallelism: usize,
+    #[case] expected: usize,
+) {
+    assert_eq!(
+        filtered_rows_per_chunk(rows, row_count, projected_bytes, parallelism),
+        expected
+    );
+}