perf(layout): share a flat layout's decode across the splits reading it

vortex-file/src/tests.rs

@@ -2528,6 +2528,8 @@ async fn test_sub_splitting_max_rows_caps_scan_batches() -> VortexResult<()> {
 #[rstest]
 #[case::unaligned(33_333)]
 #[case::exceeds_file(300_000)]
+// Hundreds of splits share one decode of the flat chunk while they are in flight.
+#[case::many_small(1_000)]
 #[tokio::test]
 #[cfg_attr(miri, ignore)]
 async fn test_flat_chunk_scan_with_row_count_splits(

vortex-layout/src/layouts/flat/reader.rs

@@ -7,6 +7,8 @@ use std::sync::Arc;
 
 use futures::FutureExt;
 use futures::future::BoxFuture;
+use futures::future::WeakShared;
+use parking_lot::Mutex;
 use tracing::trace;
 use vortex_array::ArrayRef;
 use vortex_array::MaskFuture;
@@ -15,6 +17,7 @@ use vortex_array::dtype::DType;
 use vortex_array::dtype::FieldMask;
 use vortex_array::expr::BoundExpression;
 use vortex_array::serde::SerializedArray;
+use vortex_error::SharedVortexResult;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_mask::Mask;
@@ -39,6 +42,13 @@ pub struct FlatReader {
     name: Arc<str>,
     segment_source: Arc<dyn SegmentSource>,
     session: VortexSession,
+    /// The decode of this layout's array while any split still holds it.
+    ///
+    /// A scan cuts a large flat layout into several splits, and each asks for the array. Sharing
+    /// the decode stops every split from deserializing (and validating) the whole array again,
+    /// while holding it weakly frees the array as soon as the last split using it is done, instead
+    /// of keeping every chunk of a file alive for the length of a scan.
+    array: Mutex<Option<WeakShared<BoxFuture<'static, SharedVortexResult<ArrayRef>>>>>,
 }
 
 impl FlatReader {
@@ -53,6 +63,7 @@ impl FlatReader {
             name,
             segment_source,
             session,
+            array: Mutex::new(None),
         }
     }
 
@@ -61,16 +72,21 @@ impl FlatReader {
         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?
+        let mut cached = self.array.lock();
+        if let Some(array) = cached.as_ref().and_then(WeakShared::upgrade) {
+            return array;
+        }
+
+        // We create the segment_fut here, before the returned future is polled, to give the
+        // segment reader visibility into how to prioritize this segment. A decode that is still
+        // shared above already holds a request for the same segment.
         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 array = 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.
@@ -84,7 +100,9 @@ impl FlatReader {
                 .map_err(Arc::new)
         }
         .boxed()
-        .shared()
+        .shared();
+        *cached = array.downgrade();
+        array
     }
 }
 
@@ -229,6 +247,8 @@ impl LayoutReader for FlatReader {
 #[cfg(test)]
 mod test {
     use std::sync::Arc;
+    use std::sync::atomic::AtomicUsize;
+    use std::sync::atomic::Ordering;
 
     use vortex_array::ArrayContext;
     use vortex_array::IntoArray;
@@ -248,6 +268,9 @@ mod test {
 
     use crate::LayoutStrategy;
     use crate::layouts::flat::writer::FlatLayoutStrategy;
+    use crate::segments::SegmentFuture;
+    use crate::segments::SegmentId;
+    use crate::segments::SegmentSource;
     use crate::segments::TestSegments;
     use crate::sequence::SequenceId;
     use crate::sequence::SequentialArrayStreamExt;
@@ -370,4 +393,59 @@ mod test {
             assert_arrays_eq!(result, expected, &mut ctx);
         })
     }
+
+    struct CountingSegmentSource {
+        inner: Arc<TestSegments>,
+        requests: Arc<AtomicUsize>,
+    }
+
+    impl SegmentSource for CountingSegmentSource {
+        fn request(&self, id: SegmentId) -> SegmentFuture {
+            self.requests.fetch_add(1, Ordering::Relaxed);
+            self.inner.request(id)
+        }
+    }
+
+    #[test]
+    fn splits_share_one_decode_until_dropped() -> VortexResult<()> {
+        block_on(|handle| async {
+            let session = new_session().with_handle(handle);
+            let mut ctx = session.create_execution_ctx();
+            let segments = Arc::new(TestSegments::default());
+            let (ptr, eof) = SequenceId::root().split();
+            let array = PrimitiveArray::from_iter(0..100i32).into_array();
+            let layout = FlatLayoutStrategy::default()
+                .write_stream(
+                    ArrayContext::empty().into(),
+                    Arc::<TestSegments>::clone(&segments),
+                    array.to_array_stream().sequenced(ptr),
+                    eof,
+                    &session,
+                )
+                .await?;
+
+            let requests = Arc::new(AtomicUsize::new(0));
+            let source = Arc::new(CountingSegmentSource {
+                inner: segments,
+                requests: Arc::clone(&requests),
+            });
+            let reader = layout.new_reader("".into(), source, &session, &Default::default())?;
+            let expr = root().bind(reader.dtype())?;
+
+            // Two splits of the same layout, both outstanding: one segment request, one decode.
+            let first = reader.projection_evaluation(&(0..50), &expr, MaskFuture::new_true(50))?;
+            let second =
+                reader.projection_evaluation(&(50..100), &expr, MaskFuture::new_true(50))?;
+            assert_eq!(requests.load(Ordering::Relaxed), 1);
+            assert_arrays_eq!(first.await?, array.slice(0..50)?, &mut ctx);
+            assert_arrays_eq!(second.await?, array.slice(50..100)?, &mut ctx);
+
+            // With every split done the decode is released, so a later split decodes again.
+            let third =
+                reader.projection_evaluation(&(0..100), &expr, MaskFuture::new_true(100))?;
+            assert_eq!(requests.load(Ordering::Relaxed), 2);
+            assert_arrays_eq!(third.await?, array, &mut ctx);
+            Ok(())
+        })
+    }
 }