diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 6f74e583b7c..725a7eb2ac4 100644 --- a/vortex-file/src/tests.rs +++ b/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( diff --git a/vortex-layout/src/layouts/flat/reader.rs b/vortex-layout/src/layouts/flat/reader.rs index aa7609f1659..70b1e8c81c2 100644 --- a/vortex-layout/src/layouts/flat/reader.rs +++ b/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, segment_source: Arc, 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>>>>, } 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, + requests: Arc, + } + + 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::::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(()) + }) + } }