From a5e4a997da7f8bbc61c2db9a9e03918656abde37 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 27 Oct 2026 10:00:34 +0000 Subject: [PATCH] Reuse deserialized V2 segments across partial scans Keep full encoded arrays in a per-file cache with four FIFO shards and a 16 MiB total weighted capacity. Partial scans can reuse deserialization and skip read requests on a hit; whole-segment scans bypass the cache. Decode, weigh, and release evicted arrays outside shard locks. Cover independent row selections, plan identity, disabled caching, failed reads, and eviction. Preserve existing dictionary-value reuse and update read-count expectations for overlapping chunks. Signed-off-by: Joe Isaacs --- .../src/plan/physical/operators/scan.rs | 42 +++++- .../src/plan/physical/pipeline/compile.rs | 6 +- .../src/plan/physical/pipeline/mod.rs | 1 + .../src/plan/physical/pipeline/tests.rs | 129 ++++++++++++++++++ .../src/plan/physical/pipeline/values.rs | 98 ++++++++++++- vortex-layout/src/scan/v2/cache.rs | 15 +- vortex-layout/src/scan/v2/tests.rs | 16 ++- 7 files changed, 287 insertions(+), 20 deletions(-) diff --git a/vortex-layout/src/plan/physical/operators/scan.rs b/vortex-layout/src/plan/physical/operators/scan.rs index 842bd33ee87..ee42e806072 100644 --- a/vortex-layout/src/plan/physical/operators/scan.rs +++ b/vortex-layout/src/plan/physical/operators/scan.rs @@ -21,6 +21,7 @@ use crate::plan::physical::operator::Source; use crate::plan::physical::operator::SourceBlocked; use crate::plan::physical::operator::SourceReady; use crate::plan::physical::operator::Step; +use crate::plan::physical::pipeline::DecodedSegments; use crate::segments::SegmentId; /// Decodes the whole of a scan's segment. @@ -66,6 +67,8 @@ pub(crate) struct ScanSource { unread: Option, /// The segment has been emitted. done: bool, + cache: Option, + decoded: Option, } impl ScanSource { @@ -73,6 +76,7 @@ impl ScanSource { plan: SegmentScanPlan, slice: Option>, filter: Option, + cache: Option, ) -> Self { Self { unread: Some(plan.segment_id()), @@ -80,6 +84,8 @@ impl ScanSource { slice, filter, done: false, + cache, + decoded: None, } } } @@ -87,10 +93,17 @@ impl ScanSource { impl Source for ScanSource { fn ready(&mut self, cx: &mut Cx<'_>) -> VortexResult { if let Some(segment) = self.unread.take() { + self.decoded = self + .cache + .as_ref() + .and_then(|cache| cache.get(self.plan.to_plan())); + if self.decoded.is_some() { + return Ok(SourceReady::Ready); + } cx.request_read(segment); return Ok(SourceReady::Blocked(SourceBlocked::Io)); } - Ok(if self.done || cx.has_bytes() { + Ok(if self.done || self.decoded.is_some() || cx.has_bytes() { SourceReady::Ready } else { SourceReady::Blocked(SourceBlocked::Io) @@ -103,11 +116,30 @@ impl Operator for ScanSource { if self.done { return Ok(Step::End); } - let (_, bytes) = cx - .take_segment() - .ok_or_else(|| vortex_err!("Segment scan ran without its bytes"))?; + let array = match self.decoded.take() { + Some(array) => array, + None => { + let (_, bytes) = cx + .take_segment() + .ok_or_else(|| vortex_err!("Segment scan ran without its bytes"))?; + // Another chunk may have decoded the segment while this read was in flight. + match self + .cache + .as_ref() + .and_then(|cache| cache.get(self.plan.to_plan())) + { + Some(array) => array, + None => { + let array = decode(&self.plan, bytes, cx.session())?; + if let Some(cache) = &self.cache { + cache.insert(self.plan.to_plan(), array.clone()); + } + array + } + } + } + }; self.done = true; - let array = decode(&self.plan, bytes, cx.session())?; Ok(Step::Final(select( array, self.slice.as_ref(), diff --git a/vortex-layout/src/plan/physical/pipeline/compile.rs b/vortex-layout/src/plan/physical/pipeline/compile.rs index 0e18ebe4e4c..61a124d62e7 100644 --- a/vortex-layout/src/plan/physical/pipeline/compile.rs +++ b/vortex-layout/src/plan/physical/pipeline/compile.rs @@ -131,7 +131,7 @@ impl Compiler<'_> { } /// The rows `rows` of `scan`'s segment, filtered by `filter`, read by the split compiling. - /// A segment several readers read is decoded once and shared. + /// Readers of part of a segment can share its decoded array through the value cache. pub(crate) fn scan( &mut self, scan: &SegmentScanPlan, @@ -148,11 +148,15 @@ impl Compiler<'_> { } else { Some(usize::try_from(rows.start)?..usize::try_from(rows.end)?) }; + // Whole segments are usually consumed once. Retain only segments that span chunks, + // where neighboring readers would otherwise repeat their deserialization. + let decoded = slice.as_ref().and_then(|_| self.core.values.decoded()); self.core.segments.push(scan.segment_id()); Ok(Some(Chain::new(ScanSource::new( scan.clone(), slice, filter, + decoded, )))) } } diff --git a/vortex-layout/src/plan/physical/pipeline/mod.rs b/vortex-layout/src/plan/physical/pipeline/mod.rs index 869babeaf8d..60a4ca6dd78 100644 --- a/vortex-layout/src/plan/physical/pipeline/mod.rs +++ b/vortex-layout/src/plan/physical/pipeline/mod.rs @@ -55,6 +55,7 @@ pub use self::scan::ReadRequest; pub use self::scan::Scan; pub use self::scan::Split; pub use self::scan::Turn; +pub(crate) use self::values::DecodedSegments; pub use self::values::ValueCache; /// Batches an inlet holds before its sink is blocked, unless the plan joining it says diff --git a/vortex-layout/src/plan/physical/pipeline/tests.rs b/vortex-layout/src/plan/physical/pipeline/tests.rs index 804d56553d9..6ce7e40e06d 100644 --- a/vortex-layout/src/plan/physical/pipeline/tests.rs +++ b/vortex-layout/src/plan/physical/pipeline/tests.rs @@ -716,6 +716,135 @@ fn eval_applies_expression_to_selected_rows() -> VortexResult<()> { Ok(()) } +#[rstest] +#[case::cached(1024 * 1024, 0)] +#[case::disabled(0, 1)] +fn decoded_segments_keep_the_full_array( + #[case] capacity: u64, + #[case] second_reads: usize, +) -> VortexResult<()> { + let mut store = Store::default(); + let expected = PrimitiveArray::from_iter(0..ROWS as i32).into_array(); + let plan = lower(&store.flat(&expected)?)?; + let cache = ValueCache::with_decoded_segments(capacity); + for (rows, mask, expected_reads) in [ + (2..10, Sel::EveryOther.mask(8), 1), + (10..18, Mask::new_true(8), second_reads), + ] { + let run = run_with_values( + &store, + &plan, + rows.clone(), + mask.clone(), + cache.clone(), + delivery(Delivery::Fifo), + )?; + assert_eq!(reads(&run.events), expected_reads); + assert_view(&expected, &rows, &mask, run.arrays)?; + } + Ok(()) +} + +#[test] +fn decoded_segments_distinguish_plans_with_the_same_segment_id() -> VortexResult<()> { + let cache = ValueCache::with_decoded_segments(1024 * 1024); + for start in [0, 100] { + let mut store = Store::default(); + let expected = PrimitiveArray::from_iter(start..start + ROWS as i32).into_array(); + let plan = lower(&store.flat(&expected)?)?; + let run = run_with_values( + &store, + &plan, + 2..18, + Mask::new_true(16), + cache.clone(), + delivery(Delivery::Fifo), + )?; + assert_eq!(segments_read(&run.events), [0]); + assert_view(&expected, &(2..18), &Mask::new_true(16), run.arrays)?; + } + Ok(()) +} + +#[test] +fn decoded_segments_do_not_keep_failed_reads() -> VortexResult<()> { + let mut store = Store::default(); + let expected = PrimitiveArray::from_iter(0..ROWS as i32).into_array(); + let plan = lower(&store.flat(&expected)?)?; + let cache = ValueCache::with_decoded_segments(1024 * 1024); + let bytes = std::mem::replace( + &mut store.segments[0], + BufferHandle::new_host(ByteBufferMut::empty_aligned(Alignment::new(64)).freeze()), + ); + assert!( + run_with_values( + &store, + &plan, + 2..18, + Mask::new_true(16), + cache.clone(), + delivery(Delivery::Fifo), + ) + .is_err() + ); + store.segments[0] = bytes; + let run = run_with_values( + &store, + &plan, + 2..18, + Mask::new_true(16), + cache, + delivery(Delivery::Fifo), + )?; + assert_eq!(reads(&run.events), 1); + assert_view(&expected, &(2..18), &Mask::new_true(16), run.arrays) +} + +#[test] +fn decoded_segments_evict_and_reread() -> VortexResult<()> { + let mut store = Store::default(); + let expected = PrimitiveArray::from_iter(0..ROWS as i32).into_array(); + let capacity = 8 * 1024; + let cache = ValueCache::with_decoded_segments(capacity); + let decoded = cache.decoded().vortex_expect("decoded cache enabled"); + let mut plans = Vec::new(); + for _ in 0..16 { + let plan = lower(&store.flat(&expected)?)?.child_required(0)?; + assert!(plan.is::()); + let run = run_with_values( + &store, + &plan, + 2..18, + Mask::new_true(16), + cache.clone(), + delivery(Delivery::Fifo), + )?; + assert_eq!(reads(&run.events), 1); + assert_view(&expected, &(2..18), &Mask::new_true(16), run.arrays)?; + plans.push(plan); + } + let retained = plans + .iter() + .filter(|plan| decoded.get((*plan).clone()).is_some()) + .count(); + assert!(retained > 0); + assert!(retained as u64 * (expected.nbytes() + 1024) <= capacity); + let evicted = plans + .iter() + .find(|plan| decoded.get((*plan).clone()).is_none()) + .vortex_expect("the cache cannot retain all segments"); + let run = run_with_values( + &store, + evicted, + 2..18, + Mask::new_true(16), + cache, + delivery(Delivery::Fifo), + )?; + assert_eq!(reads(&run.events), 1); + assert_view(&expected, &(2..18), &Mask::new_true(16), run.arrays) +} + /// A column read by two fields of one struct is read by each field, and both get it. #[test] fn a_segment_two_readers_need_is_read_by_each() -> VortexResult<()> { diff --git a/vortex-layout/src/plan/physical/pipeline/values.rs b/vortex-layout/src/plan/physical/pipeline/values.rs index 18c2116b637..13f34122d6b 100644 --- a/vortex-layout/src/plan/physical/pipeline/values.rs +++ b/vortex-layout/src/plan/physical/pipeline/values.rs @@ -1,11 +1,14 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Whole values, such as dictionary values, kept between the scans that share them. +//! Whole values and decoded segments kept between the scans that share them. +use std::collections::VecDeque; +use std::hash::BuildHasher; use std::sync::Arc; use parking_lot::Mutex; +use rustc_hash::FxBuildHasher; use rustc_hash::FxHashMap; use vortex_array::ArrayRef; @@ -17,22 +20,107 @@ use crate::plan::physical::PlanRef; /// /// A dictionary's values are read whole by every reader of its codes, so readers sharing this /// cache can reuse them across filters, projections, and splits. The caller controls its lifetime; -/// only whole values that a reader requested are retained. +/// only whole values that a reader requested are retained. V2 scans also use a bounded cache of +/// full encoded arrays to share segment deserialization between chunks. #[derive(Clone, Default)] -pub struct ValueCache(Arc>>); +pub struct ValueCache { + values: Arc>>, + decoded: Option, +} impl ValueCache { /// The values `plan` produced earlier. pub(crate) fn get(&self, plan: &PlanRef) -> Option { - self.0.lock().get(&PlanByPtr(plan.clone())).cloned() + self.values.lock().get(&PlanByPtr(plan.clone())).cloned() } /// Records the values of `plan`, keeping the first when two race. pub(crate) fn insert(&self, plan: &PlanRef, values: ArrayRef) -> ArrayRef { - self.0 + self.values .lock() .entry(PlanByPtr(plan.clone())) .or_insert(values) .clone() } + + pub(crate) fn with_decoded_segments(capacity: u64) -> Self { + Self { + decoded: Some(DecodedSegments(Arc::new(DecodedCache { + shards: Default::default(), + capacity: capacity / DECODED_SHARDS as u64, + }))), + ..Self::default() + } + } + + pub(crate) fn decoded(&self) -> Option { + self.decoded.clone() + } +} + +/// Full encoded arrays shared by overlapping scans of one file. Entries are weighted by logical +/// buffer bytes plus an allowance for metadata. Slicing and filtering happen after lookup, +/// independently for each reader; no cache lock is held while decoding. +#[derive(Clone)] +pub(crate) struct DecodedSegments(Arc); + +const DECODED_SHARDS: usize = 4; + +struct DecodedCache { + shards: [Mutex; DECODED_SHARDS], + capacity: u64, +} + +#[derive(Default)] +struct DecodedShard { + arrays: FxHashMap, + order: VecDeque, + weight: u64, +} + +impl DecodedSegments { + fn shard(&self, plan: &PlanByPtr) -> &Mutex { + let hash = FxBuildHasher.hash_one(plan); + &self.0.shards[usize::from(hash.to_le_bytes()[0]) % DECODED_SHARDS] + } + + pub(crate) fn get(&self, plan: PlanRef) -> Option { + let key = PlanByPtr(plan); + self.shard(&key) + .lock() + .arrays + .get(&key) + .map(|(array, _)| array.clone()) + } + + pub(crate) fn insert(&self, plan: PlanRef, array: ArrayRef) { + // Include metadata even for arrays with no buffers, and weigh outside the shard lock. + let weight = array.nbytes().saturating_add(1024); + if weight > self.0.capacity { + return; + } + let key = PlanByPtr(plan); + let mut retired = Vec::new(); + { + let mut shard = self.shard(&key).lock(); + if shard.arrays.contains_key(&key) { + return; + } + // Readers progress through the file, so FIFO eviction needs no bookkeeping on hits. + while shard.weight > self.0.capacity - weight { + let Some(oldest) = shard.order.pop_front() else { + break; + }; + if let Some((array, weight)) = shard.arrays.remove(&oldest) { + shard.weight -= weight; + retired.push((oldest, array)); + } + } + shard.weight += weight; + shard.order.push_back(key.clone()); + shard.arrays.insert(key, (array, weight)); + } + // Releasing an encoded array can walk its children; never do that under the shard lock. + drop(retired); + } } diff --git a/vortex-layout/src/scan/v2/cache.rs b/vortex-layout/src/scan/v2/cache.rs index c621bf19289..08e76556d4f 100644 --- a/vortex-layout/src/scan/v2/cache.rs +++ b/vortex-layout/src/scan/v2/cache.rs @@ -25,14 +25,25 @@ use crate::scan::execution::FilterPlans; /// scans. The lowered plan keeps the children it has initialized, each distinct projection keeps /// its optimized plan, each distinct filter keeps its conjunct and pruning plans along with the /// zones those plans have already excluded, and dictionary values stay keyed by the cached plan -/// that produced them. +/// that produced them. Segments spanning scan chunks also share their deserialized arrays in a +/// cache with a 16 MiB weighted capacity. /// /// A cache describes one file's layout, segment locations, and segment source. Use a new cache /// when any of them changes. Cached plans and values live as long as the cache, one set per /// distinct projection. -#[derive(Clone, Default)] +#[derive(Clone)] pub struct ScanCache(Arc); +impl Default for ScanCache { + fn default() -> Self { + Self(Arc::new(Inner { + // Reuse neighboring chunks' deserialization without retaining the whole file. + values: ValueCache::with_decoded_segments(16 * 1024 * 1024), + ..Inner::default() + })) + } +} + #[derive(Default)] struct Inner { root: Mutex>, diff --git a/vortex-layout/src/scan/v2/tests.rs b/vortex-layout/src/scan/v2/tests.rs index eff68dfe996..1286ca373ca 100644 --- a/vortex-layout/src/scan/v2/tests.rs +++ b/vortex-layout/src/scan/v2/tests.rs @@ -556,22 +556,24 @@ async fn dictionary_values_are_reused_between_splits() -> VortexResult<()> { chunks.len(), (STORED_CHUNKS * STORED_CHUNK_ROWS).div_ceil(MAX_ROWS_PER_CHUNK) ); + let mut first_unread_codes_chunk = 0; for (index, rows) in chunks.into_iter().enumerate() { let before = counted.reads.load(Ordering::Relaxed); let arrays = run_rows(&scan, rows.clone()).await?; let reads = counted.reads.load(Ordering::Relaxed) - before; - let code_chunks = - (rows.end - 1) / STORED_CHUNK_ROWS as u64 - rows.start / STORED_CHUNK_ROWS as u64 + 1; + let end_codes_chunk = (rows.end - 1) / STORED_CHUNK_ROWS as u64 + 1; // The first chunk reads every values chunk once for the whole scan; its codes are one of - // them, and the session shares that read. Later chunks read only their codes. + // them, and the session shares that read. Later work chunks reuse the codes that overlap + // their predecessor and only read stored codes chunks they encounter for the first time. assert_eq!( reads, if index == 0 { STORED_CHUNKS } else { - usize::try_from(code_chunks)? + usize::try_from(end_codes_chunk - first_unread_codes_chunk)? } ); + first_unread_codes_chunk = end_codes_chunk; let expected = PrimitiveArray::from_iter(i32::try_from(rows.start)?..i32::try_from(rows.end)?); assert_arrays_eq!( @@ -686,7 +688,7 @@ async fn chunked_ints( /// A shared [`ScanCache`] keeps dictionary values between scans; without one, each scan reads /// them again. #[rstest] -#[case::cached(true, 1)] +#[case::cached(true, 0)] #[case::uncached(false, 4)] #[tokio::test] async fn scan_cache_reuses_values_between_scans( @@ -719,7 +721,7 @@ async fn scan_cache_reuses_values_between_scans( ); reads.push(counted.reads.load(Ordering::Relaxed) - before); } - // Every values chunk and the one selected codes chunk, then only the codes chunk if cached. + // Every values chunk and the selected codes chunk are retained by the shared cache. assert_eq!(reads, vec![4, second_scan_reads]); Ok(()) } @@ -779,7 +781,7 @@ async fn repeated_projections_reuse_their_values() -> VortexResult<()> { ); reads.push(counted.reads.load(Ordering::Relaxed) - before); } - assert_eq!(reads, vec![4, 1, 1]); + assert_eq!(reads, vec![4, 0, 0]); Ok(()) }