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 + ); +}