From f091f2d6c59d21721ad655f879191a5cb47761db Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 30 Oct 2026 18:56:31 +0000 Subject: [PATCH] Isolate: export each DuckDB V2 part as one batch One of the three changes of #11060, on its own, to benchmark it in isolation. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MZ3d6HFyGGeUGy8brEr6BM Signed-off-by: Joe Isaacs --- vortex-duckdb/src/file_reader.rs | 3 +++ vortex-layout/src/scan/execution/mod.rs | 3 +++ vortex-layout/src/scan/execution/projection.rs | 16 ++++++++++------ vortex-layout/src/scan/v2/mod.rs | 15 +++++++++++++++ 4 files changed, 31 insertions(+), 6 deletions(-) diff --git a/vortex-duckdb/src/file_reader.rs b/vortex-duckdb/src/file_reader.rs index 7c24823c9e1..4593a505251 100644 --- a/vortex-duckdb/src/file_reader.rs +++ b/vortex-duckdb/src/file_reader.rs @@ -272,6 +272,9 @@ fn v2_builder(state: &FileScanState, global: &GlobalState) -> V2ScanBuilder { .with_some_filter(filter.filter.clone()) .with_selection(filter.row_selection.clone()) .with_row_range(rows) + // DuckDB reads a batch in vectors of its own, so a part is exported as one batch, as a + // V1 split is. + .with_max_batch_rows(usize::MAX) } /// What a thread's split yields next. diff --git a/vortex-layout/src/scan/execution/mod.rs b/vortex-layout/src/scan/execution/mod.rs index 1411b96f01e..cd7f2c9c900 100644 --- a/vortex-layout/src/scan/execution/mod.rs +++ b/vortex-layout/src/scan/execution/mod.rs @@ -41,6 +41,7 @@ pub(crate) use self::chunk::rows_per_chunk; pub(crate) use self::chunk::scan_chunks; use self::filter::FilterPlanner; pub use self::filter::FilterPlans; +pub(crate) use self::projection::MAX_BATCH_ROWS; use self::projection::ProjectionPlanner; pub use self::prune::PruneNext; pub use self::prune::PrunePlanner; @@ -102,6 +103,8 @@ pub struct PreparedScan { pub(super) values: ValueCache, /// The rows of one chunk, the unit of work a chunk planner emits. pub(super) chunk_rows: usize, + /// The most rows a projection batch of densely selected rows holds. + pub(super) max_batch_rows: usize, } impl PreparedScan { diff --git a/vortex-layout/src/scan/execution/projection.rs b/vortex-layout/src/scan/execution/projection.rs index 141b7377c26..3d2d931c5a8 100644 --- a/vortex-layout/src/scan/execution/projection.rs +++ b/vortex-layout/src/scan/execution/projection.rs @@ -29,10 +29,13 @@ use crate::scan::execution::adaptor::AdaptorStep; use crate::scan::execution::adaptor::Piece; use crate::scan::execution::adaptor::ScanAdaptor; -const MAX_OUTPUT_ROWS: usize = 65_536; +/// The default cap on the rows of a projection batch; see +/// [`ScanBuilder::with_max_batch_rows`](crate::scan::v2::ScanBuilder::with_max_batch_rows). +pub(crate) const MAX_BATCH_ROWS: usize = 65_536; -/// Keeps sparse selections together and subdivides dense selections to avoid large copies when -/// converting chunked projection results to Arrow. +/// Keeps sparse selections together and subdivides dense selections into batches of at most the +/// scan's `max_batch_rows`, to avoid large copies when converting chunked projection results to +/// Arrow. pub(super) struct ProjectionPlanner { scan: Arc, selected: IntoIter, @@ -41,13 +44,14 @@ pub(super) struct ProjectionPlanner { impl ProjectionPlanner { /// Creates a planner for `selected`. pub(super) fn new(scan: Arc, selected: SelectedRows) -> Self { - let selected = if selected.mask.true_count() <= MAX_OUTPUT_ROWS { + let max_rows = scan.max_batch_rows; + let selected = if selected.mask.true_count() <= max_rows { vec![selected] } else { (0..selected.mask.len()) - .step_by(MAX_OUTPUT_ROWS) + .step_by(max_rows) .filter_map(|start| { - let end = (start + MAX_OUTPUT_ROWS).min(selected.mask.len()); + let end = (start + max_rows).min(selected.mask.len()); let mask = selected.mask.slice(start..end); if mask.all_false() { return None; diff --git a/vortex-layout/src/scan/v2/mod.rs b/vortex-layout/src/scan/v2/mod.rs index 2c401bd46d6..4c2bc713d76 100644 --- a/vortex-layout/src/scan/v2/mod.rs +++ b/vortex-layout/src/scan/v2/mod.rs @@ -49,6 +49,7 @@ 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_BATCH_ROWS; use crate::scan::execution::MAX_ROWS_PER_CHUNK; pub use crate::scan::execution::PreparedScan; pub use crate::scan::execution::PruneNext; @@ -100,6 +101,7 @@ pub struct ScanBuilder { row_range: Option>, selection: Selection, row_offset: u64, + max_batch_rows: usize, } impl ScanBuilder { @@ -113,6 +115,7 @@ impl ScanBuilder { row_range: None, selection: Selection::default(), row_offset: 0, + max_batch_rows: MAX_BATCH_ROWS, } } @@ -161,6 +164,17 @@ impl ScanBuilder { self } + /// Caps the rows of a batch the projection of densely selected rows emits, 65,536 by default. + /// Converting a batch to Arrow copies it whole, so a consumer that does gains from small + /// batches; one that reads a batch in vectors of its own, as DuckDB does, gains from batches + /// as large as its units of work, each of which slices and decodes the segments it overlaps + /// once. + pub fn with_max_batch_rows(mut self, max_batch_rows: usize) -> Self { + assert!(max_batch_rows > 0, "a batch holds at least one row"); + self.max_batch_rows = max_batch_rows; + self + } + /// Lowers the file's layout and plans the filter and the projection, from the file's cache /// where an earlier scan planned them already. Touches no segment. Lowering and optimizing a /// wide file is real CPU, so a driver runs this where it runs compute. @@ -228,6 +242,7 @@ impl ScanBuilder { row_offset: self.row_offset, values, chunk_rows, + max_batch_rows: self.max_batch_rows, }) } }