perf(scan): intra-file decode parallelism + per-execution kernel cache + FlatReader decode memo

vortex-array/src/executor.rs

@@ -26,7 +26,6 @@ use vortex_error::VortexResult;
 use vortex_error::vortex_bail;
 use vortex_error::vortex_ensure;
 use vortex_error::vortex_panic;
-use vortex_session::Ref;
 use vortex_session::SessionExt;
 use vortex_session::VortexSession;
 
@@ -43,6 +42,7 @@ use crate::memory::HostAllocatorRef;
 use crate::memory::MemorySessionExt;
 use crate::optimizer::ArrayOptimizer;
 use crate::optimizer::kernels::ArrayKernels;
+use crate::optimizer::kernels::KernelSnapshot;
 use crate::stats::ArrayStats;
 use crate::stats::StatsSet;
 
@@ -305,6 +305,16 @@ struct StackFrame {
 #[derive(Debug, Clone)]
 pub struct ExecutionCtx {
     session: VortexSession,
+    /// Snapshot of the session's [`ArrayKernels`] execute-parent kernels, resolved once at
+    /// construction.
+    ///
+    /// The registry is session-scoped and mutable through its public `register_*` methods, so
+    /// this context sees the kernels as registered when it was created; later registrations are
+    /// picked up by the next context (contexts are created per evaluation). Caching the snapshot
+    /// avoids a per-array-node session clone plus a sharded `DashMap` `RwLock` probe in the hot
+    /// `execute_until` loop, and avoids holding the session-variable read guard across kernel
+    /// invocation. `None` means the session had no [`ArrayKernels`] when the context was created.
+    kernels: Option<KernelSnapshot>,
     #[cfg(debug_assertions)]
     id: usize,
     #[cfg(debug_assertions)]
@@ -314,8 +324,10 @@ pub struct ExecutionCtx {
 impl ExecutionCtx {
     /// Create a new execution context with the given session.
     pub fn new(session: VortexSession) -> Self {
+        let kernels = session.get_opt::<ArrayKernels>().map(|k| k.snapshot());
         Self {
             session,
+            kernels,
             #[cfg(debug_assertions)]
             id: {
                 static EXEC_CTX_ID: AtomicUsize = AtomicUsize::new(0);
@@ -331,6 +343,12 @@ impl ExecutionCtx {
         &self.session
     }
 
+    /// Get the [`KernelSnapshot`] resolved once for this execution context, if the session had an
+    /// [`ArrayKernels`] registry. Cheap to clone (one `Arc` clone).
+    pub(crate) fn kernels(&self) -> Option<&KernelSnapshot> {
+        self.kernels.as_ref()
+    }
+
     /// Get the session-scoped host allocator for this execution context.
     pub fn allocator(&self) -> HostAllocatorRef {
         self.session.allocator()
@@ -424,8 +442,7 @@ impl Executable for ArrayRef {
             }
         }
 
-        let tmp_session = ctx.session().clone();
-        let kernels = tmp_session.get_opt::<ArrayKernels>();
+        let kernels = ctx.kernels().cloned();
 
         for (slot_idx, slot) in array.slots().iter().enumerate() {
             let Some(child) = slot else { continue };
@@ -542,7 +559,7 @@ fn execute_parent_for_child(
     parent: &ArrayRef,
     child: &ArrayRef,
     slot_idx: usize,
-    kernels: Option<&Ref<ArrayKernels>>,
+    kernels: Option<&KernelSnapshot>,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<Option<ArrayRef>> {
     if let Some(kernels) = kernels
@@ -561,8 +578,7 @@ fn execute_parent_for_child(
 
 /// Try execute_parent on each occupied slot of the array.
 fn try_execute_parent(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Option<ArrayRef>> {
-    let tmp_session = ctx.session().clone();
-    let kernels = tmp_session.get_opt::<ArrayKernels>();
+    let kernels = ctx.kernels().cloned();
 
     for (slot_idx, slot) in array.slots().iter().enumerate() {
         let Some(child) = slot else { continue };

vortex-array/src/optimizer/kernels.rs

@@ -200,6 +200,40 @@ impl ArrayKernels {
         let id = hash_fn_id(parent, child);
         self.execute_parent.load().get(&id).cloned()
     }
+
+    /// Capture an owned, cheaply-cloneable [`KernelSnapshot`] of the currently-registered
+    /// execute-parent kernels.
+    ///
+    /// The [`ArcSwap`] is loaded once into an [`Arc`], so the snapshot is an `Arc` clone (no map
+    /// copy) that outlives the session-variable borrow. Registrations made after the snapshot is
+    /// taken are not visible through it.
+    pub(crate) fn snapshot(&self) -> KernelSnapshot {
+        KernelSnapshot {
+            execute_parent: self.execute_parent.load_full(),
+        }
+    }
+}
+
+/// An owned, point-in-time view of the execute-parent kernels registered on an [`ArrayKernels`]
+/// registry.
+///
+/// Holding the registry map directly (rather than re-probing the session per array node) lets the
+/// executor resolve [`ArrayKernels`] once per execution context. Cloning is one [`Arc`] clone.
+#[derive(Debug, Clone)]
+pub(crate) struct KernelSnapshot {
+    execute_parent: Arc<HashMap<ExecuteParentFnId, Arc<[ExecuteParentFn]>>>,
+}
+
+impl KernelSnapshot {
+    /// Look up the [`ExecuteParentFn`]s registered for `(parent, child)`.
+    pub(crate) fn find_execute_parent(
+        &self,
+        parent: Id,
+        child: Id,
+    ) -> Option<Arc<[ExecuteParentFn]>> {
+        let id = hash_fn_id(parent, child);
+        self.execute_parent.get(&id).cloned()
+    }
 }
 
 fn hash_fn_id(parent: Id, child: Id) -> u64 {

vortex-file/src/tests.rs

@@ -64,6 +64,7 @@ use vortex_buffer::buffer;
 use vortex_error::VortexResult;
 use vortex_io::session::RuntimeSession;
 use vortex_layout::Layout;
+use vortex_layout::layouts::flat::writer::FlatLayoutStrategy;
 use vortex_layout::scan::scan_builder::ScanBuilder;
 use vortex_layout::session::LayoutSession;
 use vortex_session::VortexSession;
@@ -1812,6 +1813,50 @@ fn assert_offsets_ordered(before: &[u64], after: &[u64], context: &str) {
     }
 }
 
+#[tokio::test]
+#[cfg_attr(miri, ignore)]
+async fn test_large_flat_chunk_scan_subdivides_splits() -> VortexResult<()> {
+    // A single flat (unchunked) 250k-row layout spans the 100k sub-split threshold, so the scan
+    // must decode it as multiple row-range splits.
+    const N_ROWS: u64 = 250_000;
+    let values =
+        Buffer::from_iter((0..N_ROWS as i32).map(|i| if i % 2 == 0 { i } else { -i })).into_array();
+
+    let mut buf = ByteBufferMut::empty();
+    SESSION
+        .write_options()
+        .with_strategy(Arc::new(FlatLayoutStrategy::default()))
+        .write(&mut buf, values.to_array_stream())
+        .await?;
+
+    let file = SESSION.open_options().open_buffer(buf)?;
+
+    // Sub-division caps each split at 100k rows while tiling the file exactly.
+    let splits = file.splits()?;
+    assert!(splits.len() > 1, "expected sub-divided splits: {splits:?}");
+    assert!(splits.iter().all(|r| r.end - r.start <= 100_000));
+    assert_eq!(splits.first().map(|r| r.start), Some(0));
+    assert_eq!(splits.last().map(|r| r.end), Some(N_ROWS));
+    assert!(splits.windows(2).all(|w| w[0].end == w[1].start));
+
+    // A full scan across the sub-splits returns the original rows.
+    let result = file.scan()?.into_array_stream()?.read_all().await?;
+    assert_arrays_eq!(result, values);
+
+    // A filtered scan crossing sub-split boundaries selects exactly the matching rows.
+    let result = file
+        .scan()?
+        .with_filter(gt(root(), lit(0i32)))
+        .into_array_stream()?
+        .read_all()
+        .await?;
+    let expected =
+        Buffer::from_iter((0..N_ROWS as i32).filter(|i| i % 2 == 0 && *i > 0)).into_array();
+    assert_arrays_eq!(result, expected);
+
+    Ok(())
+}
+
 #[tokio::test]
 #[cfg_attr(miri, ignore)]
 async fn test_segment_ordering_dict_codes_before_values() -> VortexResult<()> {

vortex-layout/src/scan/split_by.rs

@@ -5,19 +5,29 @@ use std::iter::once;
 use std::ops::Range;
 
 use vortex_array::dtype::FieldMask;
+use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 
 use crate::LayoutReader;
 use crate::RowSplits;
 use crate::SplitRange;
+use crate::scan::IDEAL_SPLIT_SIZE;
+
+/// Chunk-boundary spans wider than this are sub-divided into multiple row-range splits so that a
+/// file with few, large chunks can be decoded across multiple cores rather than one.
+///
+/// Reuses [`IDEAL_SPLIT_SIZE`] as the target span per split.
+const MAX_SPLIT_ROWS: u64 = IDEAL_SPLIT_SIZE;
 
 /// Defines how the Vortex file is split into batches for reading.
 ///
 /// Note that each split must fit into the platform's maximum usize.
 #[derive(Default, Copy, Clone, Debug)]
 pub enum SplitBy {
     #[default]
-    /// Splits any time there is a chunk boundary in the file.
+    /// Splits any time there is a chunk boundary in the file. Spans between adjacent boundaries
+    /// wider than `MAX_SPLIT_ROWS` are further sub-divided so that a file with few, large chunks
+    /// can still be decoded across multiple cores.
     Layout,
     /// Splits every n rows.
     RowCount(usize),
@@ -44,7 +54,7 @@ impl SplitBy {
                     &SplitRange::root(row_range.clone())?,
                     &mut row_splits,
                 )?;
-                row_splits.into_sorted_deduped()
+                subdivide_large_spans(row_splits.into_sorted_deduped(), MAX_SPLIT_ROWS)
             }
             SplitBy::RowCount(n) => row_range
                 .clone()
@@ -55,6 +65,60 @@ impl SplitBy {
     }
 }
 
+/// Sub-divide any gap between adjacent split boundaries that is wider than `max_span` into evenly
+/// sized row-range sub-splits.
+///
+/// `boundaries` is the sorted, deduplicated list of split points produced by the layout (chunk
+/// boundaries). Downstream consumers turn this list into half-open ranges by pairing adjacent
+/// entries (`tuple_windows().map(|(s, e)| s..e)`), so the row coverage is fully determined by the
+/// boundary set. This function only *inserts* points that lie strictly between two existing
+/// adjacent boundaries; it never moves or removes a boundary. Splitting `[lo, hi)` at an interior
+/// point `m` (with `lo < m < hi`) yields exactly `[lo, m) + [m, hi)`, so the union of ranges is
+/// unchanged: the rows are still partitioned contiguously, with no gaps and no overlaps, covering
+/// every row exactly once. The output remains sorted and deduplicated.
+fn subdivide_large_spans(boundaries: Vec<u64>, max_span: u64) -> Vec<u64> {
+    debug_assert!(boundaries.is_sorted(), "boundaries must be sorted");
+    debug_assert!(max_span > 0, "max_span must be non-zero");
+
+    // Fast path: nothing to split (also covers the empty / single-boundary cases).
+    if boundaries.len() < 2 || boundaries.windows(2).all(|w| w[1] - w[0] <= max_span) {
+        return boundaries;
+    }
+
+    let mut out = Vec::with_capacity(boundaries.len() * 2);
+    for window in boundaries.windows(2) {
+        let lo = window[0];
+        let hi = window[1];
+        // Always emit the lower boundary; the final `hi` is appended once after the loop.
+        out.push(lo);
+
+        let span = hi - lo;
+        if span > max_span {
+            // Number of sub-ranges so that each is <= max_span. `span > max_span` and
+            // `max_span >= 1` guarantee `sub_count >= 2`.
+            let sub_count = span.div_ceil(max_span);
+            // Even sub-range size (rounded up); the last sub-range absorbs any remainder and is
+            // bounded by `hi`. Inserted points `lo + k*sub_size` are strictly in `(lo, hi)`.
+            let sub_size = span.div_ceil(sub_count);
+            let mut point = lo + sub_size;
+            while point < hi {
+                out.push(point);
+                // Saturating: a sum past u64::MAX can never be < `hi`, so the loop exits.
+                point = point.saturating_add(sub_size);
+            }
+        }
+    }
+    // Append the final boundary (the `hi` of the last window).
+    out.push(*boundaries.last().vortex_expect("len >= 2 checked above"));
+
+    debug_assert!(out.is_sorted(), "subdivided boundaries must stay sorted");
+    debug_assert!(
+        out.windows(2).all(|w| w[0] < w[1]),
+        "subdivided boundaries must stay strictly increasing (deduped)"
+    );
+    out
+}
+
 #[cfg(test)]
 mod test {
     use std::any::Any;
@@ -212,4 +276,81 @@ mod test {
             .unwrap();
         assert_eq!(splits, vec![0u64, 5, 10]);
     }
+
+    #[test]
+    fn subdivide_below_threshold_is_noop() {
+        // Gaps all <= max_span: boundaries returned unchanged.
+        assert_eq!(subdivide_large_spans(vec![0, 5, 10], 100), vec![0, 5, 10]);
+        assert_eq!(subdivide_large_spans(vec![0, 100], 100), vec![0, 100]);
+        assert_eq!(
+            subdivide_large_spans(Vec::<u64>::new(), 100),
+            Vec::<u64>::new()
+        );
+        assert_eq!(subdivide_large_spans(vec![7], 100), vec![7]);
+    }
+
+    #[test]
+    fn subdivide_near_u64_max_does_not_overflow() {
+        // The increment past the last interior point would overflow without saturating math.
+        let hi = u64::MAX;
+        let out = subdivide_large_spans(vec![hi - 3, hi], 2);
+        assert_eq!(out, vec![hi - 3, hi - 1, hi]);
+    }
+
+    #[test]
+    fn subdivide_splits_large_single_chunk() {
+        // One large chunk [0, 1000) with max_span 100 -> 10 contiguous sub-splits.
+        let out = subdivide_large_spans(vec![0, 1000], 100);
+        assert_eq!(
+            out,
+            vec![0, 100, 200, 300, 400, 500, 600, 700, 800, 900, 1000]
+        );
+    }
+
+    #[test]
+    fn subdivide_only_large_gaps() {
+        // Mixed: [0,50) stays whole, [50, 350) splits into 100-row pieces, [350, 360) stays whole.
+        let out = subdivide_large_spans(vec![0, 50, 350, 360], 100);
+        assert_eq!(out, vec![0, 50, 150, 250, 350, 360]);
+    }
+
+    /// Property: for any sorted, deduped boundary set, subdivision (a) keeps the first and last
+    /// boundary, (b) stays strictly increasing, and (c) preserves exact row coverage — the union
+    /// of the half-open ranges the consumer derives is identical before and after.
+    #[test]
+    fn subdivide_preserves_exact_coverage() {
+        let cases: Vec<Vec<u64>> = vec![
+            vec![0, 1000],
+            vec![0, 7, 250_001],
+            vec![0, 5, 10, 15, 20, 25, 30],
+            vec![3, 1_000_003],
+            vec![0, 99_999, 100_000, 300_000],
+        ];
+        for boundaries in cases {
+            let out = subdivide_large_spans(boundaries.clone(), MAX_SPLIT_ROWS);
+            // (a) endpoints preserved
+            assert_eq!(out.first(), boundaries.first());
+            assert_eq!(out.last(), boundaries.last());
+            // (b) strictly increasing (sorted + deduped)
+            assert!(
+                out.windows(2).all(|w| w[0] < w[1]),
+                "not strictly increasing: {out:?}"
+            );
+            // (c) exact coverage: ranges from `out` tile the same span with no gap/overlap, and
+            // every original boundary is still present (so original ranges are sub-divided, never
+            // merged or shifted).
+            let total: u64 = out.windows(2).map(|w| w[1] - w[0]).sum();
+            let expected_total = boundaries.last().unwrap() - boundaries.first().unwrap();
+            assert_eq!(
+                total, expected_total,
+                "coverage span changed for {boundaries:?}"
+            );
+            for b in &boundaries {
+                assert!(
+                    out.contains(b),
+                    "original boundary {b} dropped from {out:?}"
+                );
+            }
+        }
+    }
 }