diff --git a/vortex-array/src/executor.rs b/vortex-array/src/executor.rs index d6070ac1a4d..8fe78f4a0ad 100644 --- a/vortex-array/src/executor.rs +++ b/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, #[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::().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::(); + 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>, + kernels: Option<&KernelSnapshot>, ctx: &mut ExecutionCtx, ) -> VortexResult> { 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> { - let tmp_session = ctx.session().clone(); - let kernels = tmp_session.get_opt::(); + let kernels = ctx.kernels().cloned(); for (slot_idx, slot) in array.slots().iter().enumerate() { let Some(child) = slot else { continue }; diff --git a/vortex-array/src/optimizer/kernels.rs b/vortex-array/src/optimizer/kernels.rs index d38bc9402d1..5f5c9cbb220 100644 --- a/vortex-array/src/optimizer/kernels.rs +++ b/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>>, +} + +impl KernelSnapshot { + /// Look up the [`ExecuteParentFn`]s registered for `(parent, child)`. + pub(crate) fn find_execute_parent( + &self, + parent: Id, + child: Id, + ) -> Option> { + let id = hash_fn_id(parent, child); + self.execute_parent.get(&id).cloned() + } } fn hash_fn_id(parent: Id, child: Id) -> u64 { diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 2ba10d96684..982e692b0f2 100644 --- a/vortex-file/src/tests.rs +++ b/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<()> { diff --git a/vortex-layout/src/scan/split_by.rs b/vortex-layout/src/scan/split_by.rs index 34b0edbf4c8..e1deb42e815 100644 --- a/vortex-layout/src/scan/split_by.rs +++ b/vortex-layout/src/scan/split_by.rs @@ -5,11 +5,19 @@ 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. /// @@ -17,7 +25,9 @@ use crate::SplitRange; #[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, max_span: u64) -> Vec { + 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::::new(), 100), + Vec::::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![ + 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:?}" + ); + } + } + } }