[Benchmark only] Export each DuckDB V2 part as one batch

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.

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 {

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<PreparedScan>,
     selected: IntoIter<SelectedRows>,
@@ -41,13 +44,14 @@ pub(super) struct ProjectionPlanner {
 impl ProjectionPlanner {
     /// Creates a planner for `selected`.
     pub(super) fn new(scan: Arc<PreparedScan>, 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;

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<Range<u64>>,
     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,
         })
     }
 }