feat: switch vortex compress to use v2 scan api

benchmarks/compress-bench/src/vortex.rs

@@ -19,6 +19,7 @@ use vortex::expr::root;
 use vortex::expr::select;
 use vortex::file::OpenOptionsSessionExt;
 use vortex::file::WriteOptionsSessionExt;
+use vortex::layout::scan::v2::ScanBuilder;
 use vortex_arrow::ArrowSessionExt;
 use vortex_bench::Format;
 use vortex_bench::SESSION;
@@ -81,8 +82,9 @@ impl Compressor for VortexCompressor {
 
 /// Scan every record batch out of an in-memory Vortex file.
 async fn vortex_decompress_read(data: Bytes) -> Result<()> {
-    let mut scan = SESSION.open_options().open_buffer(data)?.scan()?;
-    let source_dtype = scan.dtype()?;
+    let file = SESSION.open_options().open_buffer(data)?;
+    let mut scan = ScanBuilder::new(SESSION.clone(), file.scan_file());
+    let source_dtype = scan.dtype().clone();
     let root_columns = source_dtype
         .as_struct_fields_opt()
         .map_or(0, |fields| fields.nfields());
@@ -94,7 +96,7 @@ async fn vortex_decompress_read(data: Bytes) -> Result<()> {
             .optimize_recursive()?;
         scan = scan.with_projection(projection);
     }
-    let schema = Arc::new(SESSION.arrow().to_arrow_schema(&scan.dtype()?)?);
+    let schema = Arc::new(SESSION.arrow().to_arrow_schema(scan.dtype())?);
 
     let stream = scan.into_record_batch_stream(schema)?;
     pin_mut!(stream);

vortex-file/src/file.rs

@@ -21,6 +21,8 @@ use vortex_layout::LayoutReader;
 use vortex_layout::scan::layout::LayoutReaderDataSource;
 use vortex_layout::scan::scan_builder::ScanBuilder;
 use vortex_layout::scan::split_by::SplitBy;
+use vortex_layout::scan::v2::ScanFile;
+use vortex_layout::scan::v2::SegmentLocation;
 use vortex_layout::segments::SegmentSource;
 use vortex_scan::DataSourceRef;
 use vortex_session::VortexSession;
@@ -221,6 +223,28 @@ impl VortexFile {
         ))
     }
 
+    /// The parts of this file the planner-driven scan executor reads through: its root layout,
+    /// where each of its segments lives, and its segment source.
+    ///
+    /// Pass this to [`vortex_layout::scan::v2::ScanBuilder::new`] to build a V2 scan explicitly.
+    pub fn scan_file(&self) -> ScanFile {
+        let locations = self
+            .footer
+            .segment_specs_with_metadata()
+            .iter()
+            .map(|spec| SegmentLocation {
+                offset: spec.offset,
+                length: spec.length,
+                alignment: spec.alignment,
+            })
+            .collect();
+        ScanFile {
+            layout: Arc::clone(self.footer.layout()),
+            locations,
+            segments: Arc::clone(&self.segment_source),
+        }
+    }
+
     /// Returns `true` if file-level statistics prove the expression cannot
     /// match any rows in this file.
     ///

vortex-layout/Cargo.toml

@@ -92,6 +92,10 @@ harness = false
 name = "aggregate_partial_merge"
 harness = false
 
+[[bench]]
+name = "pipeline_scan"
+harness = false
+
 [lints]
 workspace = true
 

vortex-layout/benches/pipeline_scan.rs

@@ -0,0 +1,132 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Isolates scan preparation, compilation, and segment delivery from Arrow conversion and IO.
+
+#![expect(clippy::unwrap_used)]
+
+use std::hint::black_box;
+
+use divan::Bencher;
+use mimalloc::MiMalloc;
+use vortex_array::ArrayContext;
+use vortex_array::IntoArray;
+use vortex_array::arrays::PrimitiveArray;
+use vortex_array::buffer::BufferHandle;
+use vortex_array::dtype::Nullability;
+use vortex_array::dtype::StructFields;
+use vortex_array::serde::SerializeOptions;
+use vortex_buffer::Alignment;
+use vortex_buffer::ByteBufferMut;
+use vortex_layout::plan::physical::PackPlan;
+use vortex_layout::plan::physical::PlanRef;
+use vortex_layout::plan::physical::SegmentScanPlan;
+use vortex_layout::plan::physical::pipeline::Scan;
+use vortex_layout::plan::physical::pipeline::Split;
+use vortex_layout::plan::physical::pipeline::Turn;
+use vortex_layout::segments::SegmentId;
+use vortex_layout::session::LayoutSession;
+use vortex_session::VortexSession;
+use vortex_session::registry::ReadContext;
+
+#[global_allocator]
+static GLOBAL: MiMalloc = MiMalloc;
+
+const ROWS: u64 = 1024;
+
+fn main() {
+    divan::main();
+}
+
+struct Fixture {
+    session: VortexSession,
+    plan: PlanRef,
+    segment: BufferHandle,
+}
+
+impl Fixture {
+    fn new(columns: u32) -> Self {
+        let session = vortex_array::array_session().with::<LayoutSession>();
+        let array = PrimitiveArray::from_iter(0_i32..1024).into_array();
+        let context = ArrayContext::empty();
+        let buffers = array
+            .serialize(
+                &context,
+                &session,
+                &SerializeOptions {
+                    offset: 0,
+                    include_padding: true,
+                },
+            )
+            .unwrap();
+        let mut bytes = ByteBufferMut::empty_aligned(Alignment::new(64));
+        for buffer in buffers {
+            bytes.extend_from_slice(buffer.as_ref());
+        }
+        let fields = StructFields::from_iter(
+            (0..columns).map(|index| (format!("c{index}"), array.dtype().clone())),
+        );
+        let children = (0..columns)
+            .map(|index| {
+                SegmentScanPlan::new(
+                    array.dtype().clone(),
+                    ROWS,
+                    SegmentId::from(index),
+                    ReadContext::new(context.to_ids()),
+                    None,
+                )
+                .into_plan()
+            })
+            .collect();
+        Self {
+            session,
+            plan: PackPlan::try_new(fields, Nullability::NonNullable, ROWS, children, None)
+                .unwrap()
+                .into_plan(),
+            segment: BufferHandle::new_host(bytes.freeze()),
+        }
+    }
+
+    fn scan(&self) -> Scan {
+        Scan::try_new(
+            self.session.clone(),
+            self.plan.clone(),
+            vec![Split::all(0..ROWS)],
+        )
+        .unwrap()
+    }
+}
+
+#[divan::bench(args = [1, 100, 1000])]
+fn prepare(bencher: Bencher, columns: u32) {
+    let fixture = Fixture::new(columns);
+    bencher.bench_local(|| black_box(fixture.scan()));
+}
+
+#[divan::bench(args = [1, 100, 1000])]
+fn prepare_and_compile(bencher: Bencher, columns: u32) {
+    let fixture = Fixture::new(columns);
+    bencher.bench_local(|| {
+        let mut scan = fixture.scan();
+        assert!(matches!(scan.step().unwrap(), Turn::Read(_)));
+        black_box(scan);
+    });
+}
+
+#[divan::bench(args = [1, 100, 1000])]
+fn scan(bencher: Bencher, columns: u32) {
+    let fixture = Fixture::new(columns);
+    bencher.bench_local(|| {
+        let mut scan = fixture.scan();
+        loop {
+            match scan.step().unwrap() {
+                Turn::Read(read) => scan.deliver(read.id, fixture.segment.clone()).unwrap(),
+                Turn::Output(_, array) => {
+                    black_box(array);
+                }
+                Turn::Waiting => unreachable!("every read is delivered immediately"),
+                Turn::Done => break,
+            }
+        }
+    });
+}

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

@@ -192,7 +192,7 @@ impl Compiler<'_> {
             }
             return port;
         }
-        let ranges = std::mem::take(shares.reaches.get_mut(*segment as usize)?);
+        let ranges = shares.shared.remove(&segment)?;
         let mut claimants: SmallVec<[usize; 4]> = SmallVec::new();
         for rows in &ranges {
             claimants.extend(
@@ -240,10 +240,11 @@ pub(crate) struct Shares {
     ordered: bool,
     /// Splits that have finished, and so will claim nothing.
     finished: Vec<bool>,
-    /// By segment id, which a file numbers densely: for each segment more than one reader will
-    /// read and whose decoding pipeline is not built yet, the range of plan rows each of its
-    /// readers reads it for. Every split overlapping a range reads it once for that range.
-    reaches: Vec<SmallVec<[Range<u64>; 1]>>,
+    /// Segment readers collected before the scan starts. A split may read only a few of a
+    /// file's segments, so recording a reader must not allocate by the file-wide segment id.
+    reaches: Vec<(SegmentId, Range<u64>)>,
+    /// The row ranges of segments read more than once, until their decoding pipeline is built.
+    shared: FxHashMap<SegmentId, SmallVec<[Range<u64>; 1]>>,
     /// By split: the unclaimed ports of each shared segment the split reads.
     ports: Vec<FxHashMap<SegmentId, SmallVec<[PortId; 1]>>>,
 }
@@ -257,6 +258,7 @@ impl Shares {
             splits,
             ordered,
             reaches: Vec::new(),
+            shared: FxHashMap::default(),
         }
     }
 
@@ -290,31 +292,33 @@ impl Shares {
     pub(crate) fn add(&mut self, plan: &PlanRef, rows: Range<u64>) -> VortexResult<()> {
         let reaches = &mut self.reaches;
         plan.reach(rows, &Reach::Offset(0), &mut |segment, rows| {
-            record(reaches, segment, rows)
+            reaches.push((segment, rows))
         })
     }
 
     /// Keeps only the segments more than one reader reads. The rest are read by the one
     /// pipeline that needs them.
     pub(crate) fn retain_shared(&mut self) {
         let mut reaches = std::mem::take(&mut self.reaches);
-        for ranges in &mut reaches {
-            let shared = match ranges.as_slice() {
-                [] => false,
-                // One reader per split overlapping the range: shared when the range crosses a
-                // split boundary.
-                [rows] => self
+        reaches.sort_unstable_by_key(|(segment, _)| *segment);
+        for readers in reaches.chunk_by(|(left, _), (right, _)| left == right) {
+            let (segment, rows) = &readers[0];
+            // One reader per split overlapping the range: shared when the range crosses a
+            // split boundary, or when multiple plan nodes read the segment.
+            if readers.len() == 1
+                && self
                     .overlapping(rows)
                     .filter(|&split| self.overlaps(split, rows))
                     .nth(1)
-                    .is_some(),
-                _ => true,
-            };
-            if !shared {
-                *ranges = SmallVec::new();
+                    .is_none()
+            {
+                continue;
             }
+            self.shared.insert(
+                *segment,
+                readers.iter().map(|(_, rows)| rows.clone()).collect(),
+            );
         }
-        self.reaches = reaches;
     }
 
     /// Every port made for a reader that has not claimed it yet.
@@ -371,15 +375,6 @@ impl Reach {
     }
 }
 
-/// Adds a reader of `segment` reading it for `rows`.
-fn record(reaches: &mut Vec<SmallVec<[Range<u64>; 1]>>, segment: SegmentId, rows: Range<u64>) {
-    let index = *segment as usize;
-    if index >= reaches.len() {
-        reaches.resize_with(index + 1, SmallVec::new);
-    }
-    reaches[index].push(rows);
-}
-
 /// The chunks of `concat` overlapping `rows`: each chunk's index, where it starts, and the rows
 /// of `rows` it holds, in the concatenation's domain. Found by binary search, so a split pays
 /// for the chunks it overlaps, not for the chunks of the file.