Skip to content

Commit 541031c

Browse files
joseph-isaacsclaude
andcommitted
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 #10761. Signed-off-by: Joe Isaacs <joe.isaacs@live.co.uk> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01SQocaYaDvJHfmorKXnyvJC
1 parent f3cbb32 commit 541031c

3 files changed

Lines changed: 102 additions & 10 deletions

File tree

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

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ use parking_lot::Mutex;
99
use rustc_hash::FxHashMap;
1010
use vortex_array::expr::BoundExpression;
1111
use vortex_array::expr::ExactBoundExpr;
12+
use vortex_error::VortexResult;
1213

1314
use super::io::SegmentRanges;
1415
use crate::plan::physical::PlanRef;
@@ -36,6 +37,8 @@ struct Inner {
3637
/// them, so a projection must reuse its plan for its values to be found again.
3738
projections: Mutex<FxHashMap<ExactBoundExpr, PlanRef>>,
3839
ranges: Mutex<Option<SegmentRanges>>,
40+
/// Serialized bytes the whole-file projection reads, by projection.
41+
projection_bytes: Mutex<FxHashMap<ExactBoundExpr, u64>>,
3942
values: ValueCache,
4043
}
4144

@@ -80,6 +83,22 @@ impl ScanCache {
8083
Arc::clone(self.0.ranges.lock().get_or_insert(ranges))
8184
}
8285

86+
/// The serialized bytes a whole-file scan of `projection` reads, computing them with
87+
/// `compute` on first use.
88+
pub(super) fn projection_bytes(
89+
&self,
90+
projection: &BoundExpression,
91+
compute: impl FnOnce() -> VortexResult<u64>,
92+
) -> VortexResult<u64> {
93+
let key = ExactBoundExpr(projection.clone());
94+
if let Some(bytes) = self.0.projection_bytes.lock().get(&key) {
95+
return Ok(*bytes);
96+
}
97+
let bytes = compute()?;
98+
self.0.projection_bytes.lock().insert(key, bytes);
99+
Ok(bytes)
100+
}
101+
83102
pub(super) fn values(&self) -> ValueCache {
84103
self.0.values.clone()
85104
}

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

Lines changed: 62 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,9 @@
1111
//! Import this module's [`ScanBuilder`] to select V2 explicitly. It accepts a [`ScanFile`]
1212
//! directly and supports projection, filters, row ranges, and selections. Unsupported layouts return an
1313
//! error; limits are not part of this builder's API. Work is partitioned into fixed
14-
//! chunks of up to 65,536 rows, or 1,048,576 rows for filtered scans, independently of the stored
15-
//! layout's chunk boundaries.
14+
//! chunks of up to 65,536 rows, or up to 1,048,576 rows for filtered scans, independently of the
15+
//! stored layout's chunk boundaries. Filtered scans use smaller chunks when the projection reads
16+
//! many bytes per row or the row range would give fewer chunks than threads.
1617
//! Full-file projections use smaller 8,192-row multiples when serialized bytes per row are high.
1718
//!
1819
//! A [`ScanFile`] carrying a [`ScanCache`] shares its lowered plan, segment index, and dictionary
@@ -45,6 +46,7 @@ use vortex_error::VortexExpect;
4546
use vortex_error::VortexResult;
4647
use vortex_io::request::IoService;
4748
use vortex_io::session::RuntimeSessionExt;
49+
use vortex_mask::Mask;
4850
use vortex_scan::planning::planner::WorkScope;
4951
use vortex_scan::selection::Selection;
5052
use vortex_scan::strict_sorted_buffer::StrictSortedBuffer;
@@ -55,6 +57,7 @@ pub use self::cache::ScanCache;
5557
use crate::LayoutRef;
5658
use crate::plan::physical::lower;
5759
use crate::plan::physical::optimize;
60+
use crate::plan::physical::pipeline::Scan;
5861
use crate::plan::physical::plan_row_idx_expression;
5962
use crate::scan::arrow::to_record_batch;
6063
use crate::scan::execution::FilterPlans;
@@ -76,13 +79,40 @@ fn rows_per_chunk(row_count: u64, locations: &[SegmentLocation]) -> usize {
7679
if row_count <= MAX_ROWS_PER_CHUNK as u64 {
7780
return MAX_ROWS_PER_CHUNK;
7881
}
79-
let bytes: u64 = locations
82+
let blocks = blocks_per_target_bytes(total_bytes(locations), row_count);
83+
blocks.clamp(1, 8) as usize * ROW_BLOCK_SIZE
84+
}
85+
86+
/// Rows per task of a filtered scan over `rows` rows of the file.
87+
///
88+
/// Large chunks amortize filter evaluation over cheap projections. A projection that reads many
89+
/// bytes per row, or a row range that would otherwise give fewer chunks than threads, gets smaller
90+
/// chunks so the surviving rows still spread across the runtime's threads.
91+
fn filtered_rows_per_chunk(
92+
rows: u64,
93+
row_count: u64,
94+
projected_bytes: u64,
95+
parallelism: usize,
96+
) -> usize {
97+
let max_blocks = (FILTERED_ROWS_PER_CHUNK / ROW_BLOCK_SIZE) as u64;
98+
let min_blocks = (MAX_ROWS_PER_CHUNK / ROW_BLOCK_SIZE) as u64;
99+
let by_bytes = blocks_per_target_bytes(projected_bytes, row_count).max(1);
100+
let by_threads = (rows / parallelism.max(1) as u64 / ROW_BLOCK_SIZE as u64).max(min_blocks);
101+
let blocks = by_bytes.min(by_threads).min(max_blocks);
102+
usize::try_from(blocks).vortex_expect("at most FILTERED_ROWS_PER_CHUNK blocks") * ROW_BLOCK_SIZE
103+
}
104+
105+
/// The row blocks that hold [`TARGET_BYTES_PER_CHUNK`] when `bytes` spread over `row_count` rows.
106+
fn blocks_per_target_bytes(bytes: u64, row_count: u64) -> u64 {
107+
let bytes_per_row = bytes.div_ceil(row_count.max(1)).max(1);
108+
TARGET_BYTES_PER_CHUNK / bytes_per_row / ROW_BLOCK_SIZE as u64
109+
}
110+
111+
fn total_bytes(locations: &[SegmentLocation]) -> u64 {
112+
locations
80113
.iter()
81114
.map(|location| u64::from(location.length))
82-
.sum();
83-
let bytes_per_row = bytes.div_ceil(row_count).max(1);
84-
let blocks = TARGET_BYTES_PER_CHUNK / bytes_per_row / ROW_BLOCK_SIZE as u64;
85-
blocks.clamp(1, 8) as usize * ROW_BLOCK_SIZE
115+
.sum()
86116
}
87117

88118
/// Fixed boundaries in the file's row domain, clipped to the requested range.
@@ -309,18 +339,40 @@ impl<A: 'static + Send> PreparedScan<A> {
309339
plan_row_idx_expression(builder.projection.clone(), root)?,
310340
&builder.session,
311341
)?;
312-
cache.insert_projection(builder.projection, plan)
342+
cache.insert_projection(builder.projection.clone(), plan)
313343
}
314344
}
315345
};
346+
let rows = builder.row_range.unwrap_or(0..row_count);
316347
let chunk_rows = if filter.is_some() {
317-
FILTERED_ROWS_PER_CHUNK
348+
let projected_bytes = if identity {
349+
total_bytes(&file.locations)
350+
} else {
351+
cache.projection_bytes(&builder.projection, || {
352+
let segments = Scan::segments(
353+
&builder.session,
354+
&projection,
355+
0..row_count,
356+
&Mask::new_true(usize::try_from(row_count)?),
357+
)?;
358+
Ok(segments
359+
.iter()
360+
.map(|segment| u64::from(file.locations[**segment as usize].length))
361+
.sum())
362+
})?
363+
};
364+
filtered_rows_per_chunk(
365+
rows.end - rows.start,
366+
row_count,
367+
projected_bytes,
368+
get_available_parallelism().unwrap_or(1),
369+
)
318370
} else if identity {
319371
rows_per_chunk(row_count, &file.locations)
320372
} else {
321373
MAX_ROWS_PER_CHUNK
322374
};
323-
let splits = row_chunks(builder.row_range.unwrap_or(0..row_count), chunk_rows);
375+
let splits = row_chunks(rows, chunk_rows);
324376
let ranges = match cache.ranges() {
325377
Some(ranges) => ranges,
326378
None => cache.insert_ranges(segment_ranges(&file.locations)),

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

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,12 +48,14 @@ use vortex_io::session::RuntimeSessionExt;
4848
use vortex_scan::strict_sorted_buffer::StrictSortedBuffer;
4949
use vortex_session::VortexSession;
5050

51+
use super::FILTERED_ROWS_PER_CHUNK;
5152
use super::MAX_ROWS_PER_CHUNK;
5253
use super::PreparedScan;
5354
use super::ROW_BLOCK_SIZE;
5455
use super::ScanBuilder;
5556
use super::ScanCache;
5657
use super::ScanFile;
58+
use super::filtered_rows_per_chunk;
5759
use super::row_chunks;
5860
use super::rows_per_chunk;
5961
use crate::LayoutRef;
@@ -803,3 +805,22 @@ fn partition_rows_tile_the_file_on_row_blocks() {
803805
"a range inside one block owns no rows"
804806
);
805807
}
808+
809+
#[rstest]
810+
#[case::cheap_projection_large_range(16 << 20, 16 << 20, 16 << 20, 4, FILTERED_ROWS_PER_CHUNK)]
811+
#[case::wide_projection(1 << 20, 1 << 20, 100 << 20, 1, 10 * ROW_BLOCK_SIZE)]
812+
#[case::very_wide_projection(1 << 20, 1 << 20, 10 << 30, 1, ROW_BLOCK_SIZE)]
813+
#[case::partition_spread_over_threads(1 << 20, 16 << 20, 16 << 20, 16, MAX_ROWS_PER_CHUNK)]
814+
#[case::partition_split_per_thread(1 << 20, 16 << 20, 16 << 20, 4, 32 * ROW_BLOCK_SIZE)]
815+
fn filtered_chunks_shrink_for_wide_projections_and_few_rows(
816+
#[case] rows: u64,
817+
#[case] row_count: u64,
818+
#[case] projected_bytes: u64,
819+
#[case] parallelism: usize,
820+
#[case] expected: usize,
821+
) {
822+
assert_eq!(
823+
filtered_rows_per_chunk(rows, row_count, projected_bytes, parallelism),
824+
expected
825+
);
826+
}

0 commit comments

Comments
 (0)