Plan V2 random access by the segments a selection reads

AGENTS.md

@@ -223,7 +223,7 @@ These CI checks are the ones most often missed when adding files rather than edi
   default to Snappy.
   - parquet-rs: `Compression::ZSTD(ZstdLevel::try_new(3)?)`
   - DuckDB: `(FORMAT parquet, COMPRESSION zstd, COMPRESSION_LEVEL 3)`
-  - tpchgen-cli: `--parquet-compression='ZSTD(3)'`
+  - tpchgen-cli: `tpchgen-cli parquet --compression='ZSTD(3)'`
 - Benchmarks that deliberately compare Parquet codecs, such as the GPU compress benchmark's
   `--codec`, are exempt.
 

benchmarks/random-access-bench/Cargo.toml

@@ -26,12 +26,20 @@ tokio = { workspace = true, features = ["full"] }
 url = { workspace = true }
 vortex-bench = { workspace = true }
 
+[dev-dependencies]
+divan = { workspace = true }
+vortex = { workspace = true }
+
 [features]
 lance = ["dep:lance-bench"]
 
 [[bin]]
 name = "random-access-bench"
 test = false
 
+[[bench]]
+name = "v2_take"
+harness = false
+
 [lints]
 workspace = true

benchmarks/random-access-bench/README.md

@@ -22,12 +22,25 @@ untimed warm-up, then reuses the open file handle. Reopen mode includes file ope
 work in each timed iteration. CI drives the full matrix via
 [`scripts/random-access-split.py`](../../scripts/random-access-split.py).
 
+Vortex measurements use the planner-driven V2 scan builder.
+
 ## Running locally
 
 ```bash
 cargo run -p random-access-bench --profile release_debug --features lance
 ```
 
+## Micro benchmarks
+
+The `v2_take` divan benchmark runs each Vortex query with the default and V2 scans over the
+files in the page cache, plus the per-scan cost of a V2 scan that selects no rows. Write the
+files first:
+
+```bash
+cargo run -p random-access-bench --profile release_debug -- --prepare-data --formats vortex
+cargo bench -p random-access-bench --bench v2_take
+```
+
 ## Running against S3
 
 The same benchmark can read its data from an object store instead of local disk. The remote

benchmarks/random-access-bench/benches/v2_take.rs

@@ -0,0 +1,126 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Each random-access query with the default and V2 scans, over files in the page cache.
+//!
+//! Run `random-access-bench --prepare-data --formats vortex` first to write the files.
+
+#![expect(clippy::unwrap_used)]
+
+use std::hint::black_box;
+use std::sync::LazyLock;
+
+use divan::Bencher;
+use mimalloc::MiMalloc;
+use random_access_bench::AccessPattern;
+use random_access_bench::generate_indices;
+use tokio::runtime::Runtime;
+use vortex::array::stream::ArrayStreamAdapter;
+use vortex::array::stream::ArrayStreamExt;
+use vortex::buffer::Buffer;
+use vortex::file::OpenOptionsSessionExt;
+use vortex::file::VortexFile;
+use vortex::layout::scan::v2::ScanBuilder;
+use vortex::scan::strict_sorted_buffer::StrictSortedBuffer;
+use vortex_bench::Format;
+use vortex_bench::SESSION;
+use vortex_bench::datasets::feature_vectors::FeatureVectorsData;
+use vortex_bench::datasets::nested_structs::NestedStructsData;
+use vortex_bench::datasets::taxi_data::TaxiData;
+use vortex_bench::random_access::BenchDataset;
+
+#[global_allocator]
+static GLOBAL: MiMalloc = MiMalloc;
+
+static RUNTIME: LazyLock<Runtime> = LazyLock::new(|| Runtime::new().unwrap());
+
+/// The queries V2 can scan: nested lists use list layouts, which V2 does not support.
+const QUERIES: &[&str] = &[
+    "taxi/correlated",
+    "taxi/uniform",
+    "feature-vectors/correlated",
+    "feature-vectors/uniform",
+    "nested-structs/correlated",
+    "nested-structs/uniform",
+];
+
+fn main() {
+    divan::main();
+}
+
+struct Query {
+    file: VortexFile,
+    indices: Vec<u64>,
+}
+
+impl Query {
+    fn open(query: &str) -> Self {
+        let (dataset, pattern) = query.split_once('/').unwrap();
+        let dataset: Box<dyn BenchDataset> = match dataset {
+            "taxi" => Box::new(TaxiData),
+            "feature-vectors" => Box::new(FeatureVectorsData),
+            _ => Box::new(NestedStructsData),
+        };
+        let pattern = match pattern {
+            "correlated" => AccessPattern::Correlated,
+            _ => AccessPattern::Uniform,
+        };
+        let indices = generate_indices(dataset.as_ref(), pattern);
+        let file = RUNTIME.block_on(async {
+            let path = dataset.path(Format::OnDiskVortex).await.unwrap();
+            SESSION
+                .open_options()
+                .with_layout_reader_cache()
+                .open_path(path)
+                .await
+                .unwrap()
+        });
+        Self { file, indices }
+    }
+
+    fn indices(&self) -> StrictSortedBuffer<u64> {
+        StrictSortedBuffer::try_new(Buffer::from(self.indices.clone())).unwrap()
+    }
+
+    fn take_v1(&self) {
+        let array = RUNTIME.block_on(
+            self.file
+                .scan()
+                .unwrap()
+                .with_row_indices(self.indices())
+                .into_array_stream()
+                .unwrap()
+                .read_all(),
+        );
+        black_box(array.unwrap());
+    }
+
+    fn scan_v2(&self, builder: impl FnOnce(ScanBuilder) -> ScanBuilder) {
+        let stream = builder(ScanBuilder::new(SESSION.clone(), self.file.scan_file()))
+            .into_stream()
+            .unwrap();
+        let array =
+            RUNTIME.block_on(ArrayStreamAdapter::new(self.file.dtype().clone(), stream).read_all());
+        black_box(array.unwrap());
+    }
+}
+
+/// The default scan, the target V2 should beat.
+#[divan::bench(args = QUERIES)]
+fn take_v1(bencher: Bencher, query: &str) {
+    let query = Query::open(query);
+    bencher.bench_local(|| query.take_v1());
+}
+
+#[divan::bench(args = QUERIES)]
+fn take_v2(bencher: Bencher, query: &str) {
+    let query = Query::open(query);
+    bencher.bench_local(|| query.scan_v2(|scan| scan.with_row_indices(query.indices())));
+}
+
+/// A V2 scan that selects no rows: its per-scan cost before any split is planned.
+#[divan::bench(args = QUERIES)]
+fn setup_v2(bencher: Bencher, query: &str) {
+    let query = Query::open(query);
+    bencher.bench_local(|| query.scan_v2(|scan| scan.with_row_range(0..0)));
+}

benchmarks/random-access-bench/src/lib.rs

@@ -72,7 +72,7 @@ const POISSON_EXPECTED_COUNT: usize = 100;
 const CACHED_WARMUP_DURATION: Duration = Duration::from_secs(1);
 
 /// Generate indices for the given dataset and access pattern.
-fn generate_indices(dataset: &dyn BenchDataset, pattern: AccessPattern) -> Vec<u64> {
+pub fn generate_indices(dataset: &dyn BenchDataset, pattern: AccessPattern) -> Vec<u64> {
     let row_count = dataset.row_count();
     let mut rng = StdRng::seed_from_u64(42);
 

vortex-bench/src/random_access/take.rs

@@ -29,10 +29,12 @@ use tokio::fs::File;
 use vortex::array::Canonical;
 use vortex::array::IntoArray;
 use vortex::array::VortexSessionExecute;
+use vortex::array::stream::ArrayStreamAdapter;
 use vortex::array::stream::ArrayStreamExt;
 use vortex::buffer::Buffer;
 use vortex::file::OpenOptionsSessionExt;
 use vortex::file::VortexFile;
+use vortex::layout::scan::v2::ScanBuilder;
 use vortex::scan::strict_sorted_buffer::StrictSortedBuffer;
 use vortex::utils::aliases::hash_map::HashMap;
 
@@ -126,7 +128,7 @@ impl RandomAccessor for ArrowIpcRandomAccessor {
     }
 }
 
-/// Random accessor for Vortex format files.
+/// Random accessor for Vortex format files, read with the V2 scan.
 ///
 /// The file handle is opened at construction time and reused across `take()` calls.
 pub struct VortexRandomAccessor {
@@ -142,11 +144,7 @@ impl VortexRandomAccessor {
         name: impl Into<String>,
         format: Format,
     ) -> anyhow::Result<Self> {
-        let file = SESSION
-            .open_options()
-            .with_layout_reader_cache()
-            .open_path(path.as_ref())
-            .await?;
+        let file = SESSION.open_options().open_path(path.as_ref()).await?;
         Ok(Self {
             name: name.into(),
             format,
@@ -163,7 +161,6 @@ impl VortexRandomAccessor {
     ) -> anyhow::Result<Self> {
         let file = SESSION
             .open_options()
-            .with_layout_reader_cache()
             .open_object_store(remote.store(), remote.key(path)?)
             .await?;
         Ok(Self {
@@ -185,12 +182,11 @@ impl RandomAccessor for VortexRandomAccessor {
     }
 
     async fn take(&self, indices: &[u64]) -> anyhow::Result<RandomAccessorRet> {
-        let indices_buf: Buffer<u64> = Buffer::from(indices.to_vec());
-        let array = self
-            .file
-            .scan()?
-            .with_row_indices(StrictSortedBuffer::try_new(indices_buf)?)
-            .into_array_stream()?
+        let indices = StrictSortedBuffer::try_new(Buffer::from(indices.to_vec()))?;
+        let stream = ScanBuilder::new(SESSION.clone(), self.file.scan_file())
+            .with_row_indices(indices)
+            .into_stream()?;
+        let array = ArrayStreamAdapter::new(self.file.dtype().clone(), stream)
             .read_all()
             .await?;
 

vortex-layout/src/layouts/chunked/mod.rs

@@ -117,6 +117,11 @@ impl Layout<Chunked> {
         .into_typed()
     }
 
+    /// The first row of each chunk, followed by the layout's row count.
+    pub fn chunk_offsets(&self) -> &[u64] {
+        &self.data().chunk_offsets
+    }
+
     /// Rebuild this layout with owned children.
     pub fn with_children(&self, children: Vec<LayoutRef>) -> Self {
         Self::new(

vortex-layout/src/plan/physical/children.rs

@@ -30,7 +30,7 @@ impl PlanChildren {
     ) -> Self {
         Self {
             initializer: Some(Arc::new(initializer)),
-            cache: (0..len).map(|_| OnceCell::new()).collect::<Vec<_>>().into(),
+            cache: (0..len).map(|_| OnceCell::new()).collect(),
         }
     }
 

vortex-layout/src/plan/physical/lower.rs

@@ -83,22 +83,16 @@ fn lower_flat(layout: &FlatLayout) -> SegmentScanPlan {
 }
 
 fn lower_chunked(layout: &ChunkedLayout) -> VortexResult<ConcatPlan> {
-    let mut row_offsets = Vec::with_capacity(layout.nchildren());
-    let mut row_count = 0u64;
-    for index in 0..layout.nchildren() {
-        row_offsets.push(row_count);
-        row_count = row_count
-            .checked_add(layout.child_row_count(index))
-            .ok_or_else(|| vortex_err!("Chunked row count overflow"))?;
-    }
+    let nchunks = layout.nchildren();
+    let chunks = layout.to_layout();
     // SAFETY: Chunked layout construction validates that every child has the parent dtype and
-    // that their row counts sum to the parent; offsets were computed with checked addition above.
+    // computes the chunk offsets with checked addition, ending at the parent's row count.
     Ok(unsafe {
         ConcatPlan::from_children_unchecked(
             layout.dtype().clone(),
             layout.row_count(),
-            row_offsets.into(),
-            lazy_children(layout.to_layout(), (0..layout.nchildren()).collect()),
+            layout.chunk_offsets()[..nchunks].into(),
+            PlanChildren::lazy(nchunks, move |index| lower_slot(&chunks, index)),
         )
     })
 }
@@ -159,13 +153,17 @@ fn lazy_children(layout: LayoutRef, slots: Vec<usize>) -> PlanChildren {
             .get(index)
             .copied()
             .ok_or_else(|| vortex_err!("Missing plan child slot {index}"))?;
-        let child = layout
-            .slot(slot)?
-            .ok_or_else(|| vortex_err!("Layout child slot {slot} is absent"))?;
-        lower(&child)
+        lower_slot(&layout, slot)
     })
 }
 
+fn lower_slot(layout: &LayoutRef, slot: usize) -> VortexResult<PlanRef> {
+    let child = layout
+        .slot(slot)?
+        .ok_or_else(|| vortex_err!("Layout child slot {slot} is absent"))?;
+    lower(&child)
+}
+
 fn lower_zoned(layout: &LayoutRef) -> VortexResult<PlanRef> {
     // Zoned and legacy stats layouts share a child shape: transparent data, auxiliary zones.
     let metadata = if let Some(layout) = layout.as_opt::<Zoned>() {

vortex-layout/src/plan/physical/pipeline/scan.rs

@@ -671,15 +671,14 @@ impl Scan {
             0,
             splits.iter().map(|split| split.rows.clone()).collect(),
         );
-        // Only the rows some split reads can share a segment between readers.
-        let start = splits
-            .iter()
-            .map(|split| split.rows.start)
-            .min()
-            .unwrap_or(0);
-        let end = splits.iter().map(|split| split.rows.end).max().unwrap_or(0);
+        // Only the rows some split reads can share a segment between readers, so the rows
+        // between splits are never visited.
         match &root {
-            Root::Plan(plan) => core.shares.add(plan, start..end)?,
+            Root::Plan(plan) => {
+                for split in &splits {
+                    core.shares.add(plan, split.rows.clone())?;
+                }
+            }
         }
         core.shares.retain_shared();
         core.dirty_flags = vec![false; 1];

vortex-layout/src/scan/execution/adaptor.rs

@@ -28,6 +28,8 @@ use crate::segments::SegmentId;
 /// Rows a running plan produced, in row order.
 #[derive(Clone)]
 pub(super) struct Piece {
+    /// The scan split that produced the piece. Splits are in row order.
+    pub(super) split: usize,
     /// Orders the piece among the outputs of one scan.
     pub(super) rows: Range<u64>,
     /// The selected rows the piece holds.
@@ -96,7 +98,7 @@ impl ScanAdaptor {
         }
         Ok(match self.step()? {
             Turn::Waiting | Turn::Done => AdaptorStep::Yield,
-            Turn::Output(_, array) => AdaptorStep::Piece(self.piece(array)),
+            Turn::Output(split, array) => AdaptorStep::Piece(self.piece(split, array)),
             Turn::Read(read) => {
                 // Publish every read the scan has ready together, so the owner can coalesce them.
                 let mut batch = vec![self.request(read.id, read.segment_id)?];
@@ -130,10 +132,11 @@ impl ScanAdaptor {
         Ok(turn)
     }
 
-    fn piece(&mut self, array: ArrayRef) -> Piece {
+    fn piece(&mut self, split: usize, array: ArrayRef) -> Piece {
         let start = self.next_piece;
         self.next_piece += 1;
         Piece {
+            split,
             rows: start..start + 1,
             array,
         }