diff --git a/vortex-layout/src/layouts/flat/reader.rs b/vortex-layout/src/layouts/flat/reader.rs index 954bbe5ab5c..f991e83a123 100644 --- a/vortex-layout/src/layouts/flat/reader.rs +++ b/vortex-layout/src/layouts/flat/reader.rs @@ -4,6 +4,7 @@ use std::ops::BitAnd; use std::ops::Range; use std::sync::Arc; +use std::sync::OnceLock; use futures::FutureExt; use futures::future::BoxFuture; @@ -39,6 +40,14 @@ pub struct FlatReader { name: Arc, segment_source: Arc, session: VortexSession, + /// Memoized decoded array for this layout. A column referenced by both the + /// filter and the projection (the common filtered-scan-returning-the-filter- + /// column case) would otherwise decode twice; this caches the shared decode + /// so it happens once. The segment is still re-requested on every + /// `array_future()` call so the segment reader keeps prioritization + /// visibility (see there). Bounded by the reader's lifetime — readers are + /// held via `Weak` in the opener, mirroring `DictReader`'s `values_array`. + array: OnceLock, } impl FlatReader { @@ -53,38 +62,47 @@ impl FlatReader { name, segment_source, session, + array: OnceLock::new(), } } /// Register the segment request and return a future that would resolve into the deserialised array. fn array_future(&self) -> SharedArrayFuture { - let row_count = - usize::try_from(self.layout.row_count()).vortex_expect("row count must fit in usize"); - - // We create the segment_fut here to ensure we give the segment reader visibility into - // how to prioritize this segment, even if the `array` future has already been initialized. - // This is gross... see the function's TODO for a maybe better solution? + // We create the segment_fut here to ensure we give the segment reader visibility into how + // to prioritize this segment, even if the `array` future has already been initialized (it + // is, on the second of a filter+projection pair over the same column). This is gross... see + // the function's TODO for a maybe better solution? let segment_fut = self.segment_source.request(self.layout.segment_id()); - let ctx = self.layout.array_ctx().clone(); - let session = self.session.clone(); - let dtype = self.layout.dtype().clone(); - let array_tree = self.layout.array_tree().cloned(); - async move { - let segment = segment_fut.await?; - let parts = if let Some(array_tree) = array_tree { - // Use the pre-stored flatbuffer from layout metadata combined with segment buffers. - SerializedArray::from_flatbuffer_and_segment(array_tree, segment)? - } else { - // Parse the flatbuffer from the segment itself. - SerializedArray::try_from(segment)? - }; - parts - .decode(&dtype, row_count, &ctx, &session) - .map_err(Arc::new) - } - .boxed() - .shared() + // Memoize the decoded array so a column used by BOTH the filter and the projection decodes + // once instead of twice. On a cache hit the init closure is not run, so the just-requested + // `segment_fut` is dropped after registering its prioritization hint and the cached decode + // (which owns the first call's `segment_fut`) is reused. + self.array + .get_or_init(move || { + let row_count = usize::try_from(self.layout.row_count()) + .vortex_expect("row count must fit in usize"); + let ctx = self.layout.array_ctx().clone(); + let session = self.session.clone(); + let dtype = self.layout.dtype().clone(); + let array_tree = self.layout.array_tree().cloned(); + async move { + let segment = segment_fut.await?; + let parts = if let Some(array_tree) = array_tree { + // Use the pre-stored flatbuffer from layout metadata combined with segment buffers. + SerializedArray::from_flatbuffer_and_segment(array_tree, segment)? + } else { + // Parse the flatbuffer from the segment itself. + SerializedArray::try_from(segment)? + }; + parts + .decode(&dtype, row_count, &ctx, &session) + .map_err(Arc::new) + } + .boxed() + .shared() + }) + .clone() } } @@ -361,4 +379,49 @@ mod test { assert_arrays_eq!(result, expected); }) } + + #[test] + fn array_future_memoizes_the_decode() { + block_on(|handle| async { + let session = SESSION.clone().with_handle(handle); + let ctx = ArrayContext::empty(); + let segments = Arc::new(TestSegments::default()); + let (ptr, eof) = SequenceId::root().split(); + let array = + PrimitiveArray::new(buffer![1, 2, 3, 4, 5], Validity::AllValid).into_array(); + let layout = FlatLayoutStrategy::default() + .write_stream( + ctx, + Arc::::clone(&segments), + array.to_array_stream().sequenced(ptr), + eof, + &session, + ) + .await + .unwrap(); + + let reader = layout.new_reader("".into(), segments, &SESSION).unwrap(); + let flat = reader + .as_any() + .downcast_ref::() + .expect("flat layout must yield a FlatReader"); + + // The decoded array is not cached until the column is first touched. + assert!(flat.array.get().is_none(), "decode cache must start empty"); + + // First touch (e.g. the WHERE filter) decodes the chunk and memoizes it. + let a1 = flat.array_future().await.unwrap(); + assert!( + flat.array.get().is_some(), + "decode must be memoized after the first array_future()" + ); + + // Second touch (e.g. the SELECT projection of the same column) reuses + // the cached decode instead of decoding the chunk again. + let a2 = flat.array_future().await.unwrap(); + + assert_arrays_eq!(a1, array.clone()); + assert_arrays_eq!(a2, array); + }) + } }