Plan V2 random access from known row selections

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;
 
@@ -128,7 +130,8 @@ impl RandomAccessor for ArrowIpcRandomAccessor {
 
 /// Random accessor for Vortex format files.
 ///
-/// The file handle is opened at construction time and reused across `take()` calls.
+/// The file handle is opened at construction time and reused across `take()` calls. Rows are
+/// read with the V2 scan, which requires files without list layouts.
 pub struct VortexRandomAccessor {
     name: String,
     format: Format,
@@ -142,11 +145,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 +162,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,16 +183,16 @@ 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?;
 
-        // We canonicalize / decompress for equivalence to Arrow's `RecordBatch`es.
+        // Keep the historical benchmark contract: canonicalize the outer array, leaving nested
+        // values lazy.
         let mut ctx = SESSION.create_execution_ctx();
         let canonical = array.execute::<Canonical>(&mut ctx)?.into_array();
         Ok(RandomAccessorRet::ArrayRef(canonical))

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

@@ -19,6 +19,7 @@ mod optimize;
 pub mod optimizer;
 pub mod pipeline;
 mod plans;
+pub(crate) mod selection;
 mod typed;
 mod vtable;
 

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

@@ -43,9 +43,11 @@ use crate::layouts::chunked::ChunkedLayout;
 use crate::layouts::dict::DictLayout;
 use crate::layouts::flat::FlatLayout;
 use crate::layouts::struct_::StructLayout;
+use crate::plan::physical::ConcatPlan;
 use crate::plan::physical::EvalPlan;
 use crate::plan::physical::Filter;
 use crate::plan::physical::PackPlan;
+use crate::plan::physical::PlanChildren;
 use crate::plan::physical::PlanRef;
 use crate::plan::physical::SegmentScan;
 use crate::plan::physical::Take;
@@ -60,6 +62,7 @@ use crate::plan::physical::operator::SourceReady;
 use crate::plan::physical::operator::Step;
 use crate::plan::physical::operators::PackSource;
 use crate::plan::physical::optimize;
+use crate::plan::physical::selection::selected;
 use crate::segments::SegmentId;
 use crate::test::SESSION;
 
@@ -834,3 +837,88 @@ fn pack_readiness_tracks_changes_and_consumption(
     }
     Ok(())
 }
+
+#[rstest]
+#[case::sparse(2..19, vec![0, 5, 14, 16])]
+#[case::contiguous(2..19, (5..10).collect())]
+#[case::empty(2..19, vec![])]
+fn bound_selection_matches_projection(
+    #[case] rows: Range<u64>,
+    #[case] indices: Vec<usize>,
+) -> VortexResult<()> {
+    let mut store = Store::default();
+    let (plan, expected) = fixture(&mut store)?;
+    let mask = Mask::from_indices((rows.end - rows.start) as usize, indices);
+    let plan = selected(plan, rows.clone(), mask.clone());
+    let result = run(
+        &store,
+        &plan,
+        rows.clone(),
+        Mask::new_true(mask.len()),
+        delivery(Delivery::Lifo),
+    )?;
+    assert_view(&expected, &rows, &mask, result.arrays)
+}
+
+#[test]
+fn bound_selection_never_initializes_skipped_chunks() -> VortexResult<()> {
+    let mut store = Store::default();
+    let array = PrimitiveArray::from_iter(0..20i32).into_array();
+    let first = lower(&store.flat(&array.slice(0..5)?)?)?;
+    let last = lower(&store.flat(&array.slice(15..20)?)?)?;
+    let children = PlanChildren::lazy(4, move |index| match index {
+        0 => Ok(first.clone()),
+        3 => Ok(last.clone()),
+        _ => Err(vortex_err!("unselected chunk {index} was initialized")),
+    });
+    // SAFETY: All four logical chunks have five rows of the same primitive dtype. The two
+    // unselected initializers deliberately fail to detect discovery as well as compilation.
+    let plan = unsafe {
+        ConcatPlan::from_children_unchecked(
+            array.dtype().clone(),
+            20,
+            vec![0, 5, 10, 15].into(),
+            children,
+        )
+    }
+    .into_plan();
+    let mask = Mask::from_indices(20, [2, 17]);
+    let plan = selected(plan, 0..20, mask.clone());
+    let result = run(
+        &store,
+        &plan,
+        0..20,
+        Mask::new_true(20),
+        delivery(Delivery::Fifo),
+    )?;
+    assert_eq!(reads(&result.events), 2);
+    assert_view(&array, &(0..20), &mask, result.arrays)
+}
+
+#[test]
+fn bound_selection_keeps_whole_dictionary_values() -> VortexResult<()> {
+    let mut store = Store::default();
+    let values = PrimitiveArray::from_iter(100..120i32).into_array();
+    let codes = PrimitiveArray::from_iter((0..20u32).rev()).into_array();
+    let layout = DictLayout::new(
+        store.chunked(&values, &[5, 5, 5, 5])?,
+        store.chunked(&codes, &[5, 5, 5, 5])?,
+    )
+    .into_layout();
+    let mask = Mask::from_indices(20, [1, 18]);
+    let plan = selected(lower(&layout)?, 0..20, mask);
+    let result = run(
+        &store,
+        &plan,
+        0..20,
+        Mask::new_true(20),
+        delivery(Delivery::Lifo),
+    )?;
+    assert_eq!(reads(&result.events), 6, "all values and two code chunks");
+    assert_arrays_eq!(
+        join(values.dtype(), result.arrays)?,
+        PrimitiveArray::from_iter([118, 101]).into_array(),
+        &mut SESSION.create_execution_ctx()
+    );
+    Ok(())
+}

vortex-layout/src/plan/physical/plans/concat.rs

@@ -26,6 +26,7 @@ use crate::plan::physical::pipeline::Chain;
 use crate::plan::physical::pipeline::Compiler;
 use crate::plan::physical::pipeline::Reach;
 use crate::plan::physical::pipeline::overlapping;
+use crate::plan::physical::selection::selected_slice;
 use crate::segments::SegmentId;
 
 /// Concatenates its children row-wise.
@@ -157,20 +158,26 @@ impl PlanVTable for Concat {
         mask: &Mask,
         compiler: &mut Compiler<'_>,
     ) -> VortexResult<Option<Chain>> {
-        let mut chains = Vec::with_capacity(overlapping(plan, &rows).size_hint().1.unwrap_or(0));
+        let capacity = if mask.all_true() {
+            overlapping(plan, &rows).size_hint().1.unwrap_or(0)
+        } else {
+            0
+        };
+        let mut chains = Vec::with_capacity(capacity);
         for (index, start, local) in overlapping(plan, &rows) {
-            let local_mask = mask.slice(
+            let Some((span, local_mask)) = selected_slice(
+                mask,
                 usize::try_from(local.start - rows.start)?
                     ..usize::try_from(local.end - rows.start)?,
-            );
-            // A chunk the mask selects nothing of is never built, so never read.
-            if local_mask.all_false() {
+            ) else {
                 continue;
-            }
+            };
             let chunk = plan.child_required(index)?;
-            if let Some(chain) =
-                compiler.compile(&chunk, local.start - start..local.end - start, &local_mask)?
-            {
+            if let Some(chain) = compiler.compile(
+                &chunk,
+                rows.start + span.start as u64 - start..rows.start + span.end as u64 - start,
+                &local_mask,
+            )? {
                 chains.push(chain);
             }
         }

vortex-layout/src/plan/physical/plans/filter.rs

@@ -103,13 +103,13 @@ impl PlanVTable for Filter {
         mask: &Mask,
         compiler: &mut Compiler<'_>,
     ) -> VortexResult<Option<Chain>> {
+        if mask.all_false() {
+            return Ok(None);
+        }
         let child = plan.child_plan()?;
         if let Some(scan) = child.as_opt::<SegmentScan>() {
             return compiler.scan(scan, rows, Some(mask.clone()));
         }
-        if mask.all_false() {
-            return Ok(None);
-        }
         let len = usize::try_from(rows.end - rows.start)?;
         let predicate = plan.dtype().is_boolean();
         let chain = compiler.compile(&child, rows, &Mask::new_true(len))?;