Reuse deserialized segments across partial V2 scans

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,31 +67,43 @@ pub(crate) struct ScanSource {
     unread: Option<SegmentId>,
     /// The segment has been emitted.
     done: bool,
+    cache: Option<DecodedSegments>,
+    decoded: Option<ArrayRef>,
 }
 
 impl ScanSource {
     pub(crate) fn new(
         plan: SegmentScanPlan,
         slice: Option<Range<usize>>,
         filter: Option<Mask>,
+        cache: Option<DecodedSegments>,
     ) -> Self {
         Self {
             unread: Some(plan.segment_id()),
             plan,
             slice,
             filter,
             done: false,
+            cache,
+            decoded: None,
         }
     }
 }
 
 impl Source for ScanSource {
     fn ready(&mut self, cx: &mut Cx<'_>) -> VortexResult<SourceReady> {
         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(),

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,
         ))))
     }
 }

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

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::<SegmentScan>());
+        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<()> {

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<Mutex<FxHashMap<PlanByPtr, ArrayRef>>>);
+pub struct ValueCache {
+    values: Arc<Mutex<FxHashMap<PlanByPtr, ArrayRef>>>,
+    decoded: Option<DecodedSegments>,
+}
 
 impl ValueCache {
     /// The values `plan` produced earlier.
     pub(crate) fn get(&self, plan: &PlanRef) -> Option<ArrayRef> {
-        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<DecodedSegments> {
+        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<DecodedCache>);
+
+const DECODED_SHARDS: usize = 4;
+
+struct DecodedCache {
+    shards: [Mutex<DecodedShard>; DECODED_SHARDS],
+    capacity: u64,
+}
+
+#[derive(Default)]
+struct DecodedShard {
+    arrays: FxHashMap<PlanByPtr, (ArrayRef, u64)>,
+    order: VecDeque<PlanByPtr>,
+    weight: u64,
+}
+
+impl DecodedSegments {
+    fn shard(&self, plan: &PlanByPtr) -> &Mutex<DecodedShard> {
+        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<ArrayRef> {
+        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);
+    }
 }