From c888e4fa5bb45279a32eb055623c55bdbce6e151 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Mon, 26 Oct 2026 14:08:59 +0000 Subject: [PATCH] Size filtered V2 scan splits by projected bytes and threads Filtered V2 scans used fixed 1,048,576-row chunks. A DataFusion byte-range partition of a ~1M-row file then became a single chunk, so when the surviving rows sat in one partition, one thread did all the work. Size filtered chunks by the serialized bytes the projection reads per row, cached per projection in ScanCache, and by the partition's rows per thread. Cheap projections over large ranges keep 1M-row chunks. 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 --- vortex-layout/src/scan/execution/chunk.rs | 39 ++++++++++++++++++--- vortex-layout/src/scan/execution/mod.rs | 3 ++ vortex-layout/src/scan/v2/cache.rs | 19 +++++++++++ vortex-layout/src/scan/v2/mod.rs | 41 +++++++++++++++++++---- vortex-layout/src/scan/v2/tests.rs | 21 ++++++++++++ 5 files changed, 111 insertions(+), 12 deletions(-) diff --git a/vortex-layout/src/scan/execution/chunk.rs b/vortex-layout/src/scan/execution/chunk.rs index 382425d9bb0..dde62c8e89c 100644 --- a/vortex-layout/src/scan/execution/chunk.rs +++ b/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. diff --git a/vortex-layout/src/scan/execution/mod.rs b/vortex-layout/src/scan/execution/mod.rs index 0185f61d6a9..aea1ff6427e 100644 --- a/vortex-layout/src/scan/execution/mod.rs +++ b/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; diff --git a/vortex-layout/src/scan/v2/cache.rs b/vortex-layout/src/scan/v2/cache.rs index c621bf19289..9c53d7d6224 100644 --- a/vortex-layout/src/scan/v2/cache.rs +++ b/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>>, + /// Serialized bytes the whole-file projection reads, by projection. + projection_bytes: Mutex>, 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, + ) -> VortexResult { + 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() } diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index 1f1566ed175..2c27564f600 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/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,18 +34,20 @@ 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; @@ -52,7 +55,9 @@ 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, diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index eff68dfe996..e0e235a36f4 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/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 + ); +}