From bc20b0c3fd4e47c919575eab86c79c89d2e0dc68 Mon Sep 17 00:00:00 2001 From: Luke Kim <80174+lukekim@users.noreply.github.com> Date: Fri, 12 Jun 2026 16:24:10 -0700 Subject: [PATCH 1/2] perf(scan): memoize FlatReader decoded array (filter+projection decode once) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A column referenced by BOTH the filter (WHERE) and the projection (SELECT) — the common filtered-scan-returning-the-filter-column case — went through array_future() twice, decoding the chunk twice. Memoize the shared decode in a OnceLock (mirroring DictReader::values_array) so it decodes once. The segment is still re-requested on every call so the segment reader keeps prioritization visibility; only the decode is cached. Bounded by the reader's lifetime (readers are held via Weak in the opener), so it does not reintroduce the unbounded retention that prompted removing the prior cache. --- vortex-layout/src/layouts/flat/reader.rs | 68 +++++++++++++++--------- 1 file changed, 43 insertions(+), 25 deletions(-) diff --git a/vortex-layout/src/layouts/flat/reader.rs b/vortex-layout/src/layouts/flat/reader.rs index 954bbe5ab5c..273c5a0f846 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() } } From b84bf733be54b39d1df0fdfa3d1985bafa36f9fc Mon Sep 17 00:00:00 2001 From: Luke Kim <80174+lukekim@users.noreply.github.com> Date: Fri, 12 Jun 2026 16:39:02 -0700 Subject: [PATCH 2/2] test(scan): assert FlatReader memoizes the decode (filter+projection decode once) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit array_future_memoizes_the_decode: the decode cache (OnceLock) starts empty, is populated after the first array_future() (the WHERE-filter touch), and the second touch (the SELECT projection of the same column) reuses it — so a column feeding both filter and projection decodes once. Mirrors the existing flat-reader test harness (TestSegments + FlatLayoutStrategy). --- vortex-layout/src/layouts/flat/reader.rs | 45 ++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/vortex-layout/src/layouts/flat/reader.rs b/vortex-layout/src/layouts/flat/reader.rs index 273c5a0f846..f991e83a123 100644 --- a/vortex-layout/src/layouts/flat/reader.rs +++ b/vortex-layout/src/layouts/flat/reader.rs @@ -379,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); + }) + } }