perf(scan): overlap remote split reads within decode concurrency

SPICE_PATCHES.md

@@ -62,6 +62,8 @@ exists: a grep proves presence, a test proves behaviour.
 | 14 | Arrow `Map` alias | `1a6dc54f1` (`lukim/map`) | `Map` columns unwritable: `Array encoding not implemented for Arrow data type Map(...)` on every flush (spiceai/spiceai#13524) | `cargo test -p vortex-arrow --lib executor::map` | Yes — 0.86.1 native `Map` (`#9111`) |
 | 17 | Byte-capped data blocks for string/binary columns, with encodings chosen once per row block | this change | A string or binary column is written in blocks of 8,192 rows whatever its row width, because `effective_block_len` sizes blocks only for fixed-width types; a ~3 KB/row JSON column gets ~12.5 MiB blocks, and a one-row lookup reads and copies a whole block (spiceai/spiceai#15323). `WriteStrategyBuilder::with_data_block_max_bytes` splits such a block into power-of-two row counts under the cap. Searching encodings on each small block costs several full compressions per block (a block of at most 1,024 rows is its own sample), so `CompressingStrategy::with_group_rows` and `CascadingCompressor::plan_for`/`compress_with_plan` choose encodings once per row block, on the sample the compressor would draw from the unsplit block, and apply them to its split blocks (a choice made on the first split block alone fits only it: incompressible rows ahead of JSON picked FSST for the JSON, +15% size and 2.4x `LIKE` CPU) | `cargo test -p vortex-layout --lib repartition::tests::a_block_size_maximum_splits_wide_strings_into_power_of_two_blocks` **and** `cargo test -p vortex-btrblocks --lib canonical_compressor::tests::a_replayed_plan_roundtrips_other_data_of_the_same_dtype` **and** `cargo test -p vortex-layout --lib compressed::tests::a_group_is_encoded_as_its_unsplit_data_would_be` **and** `cargo test -p vortex-file --test test_write_table test_data_block_max_bytes_roundtrip` | Proposed — additive; off unless the cap is set |
 | 18 | Chunked reads skip chunks a known mask excludes | this change | A read of a row range given an already-resolved mask (`MaskFuture::ready`, as a point read passes) requested every chunk in the range, even those the mask selects no row of, because the flat reader fetches its segment before it looks at the mask. Once string blocks are split under a byte cap (row 17), a point read of one row of an 8,192-row range fetched all of its split blocks, so a lookup got slower rather than faster (spiceai/spiceai#15323). `MaskFuture::known` exposes a mask known up front, and `ChunkedReader` skips chunks it excludes | `cargo test -p vortex-layout --lib chunked::reader::test::a_known_mask_reads_only_the_chunks_it_selects` | Proposed — upstream defect, not Spice behaviour |
+| 19 | Overlap remote split I/O without increasing concurrent decodes | spiceai/vortex#107 | Cold object-store scans hold their limited split slots while awaiting segment bytes, leaving read capacity idle; shared pending arrays must hand off a permit for every waiting split | `cargo test -p vortex-layout --lib scan::split_gate::tests` and `grep -q 'mark_pending_wait(array)' vortex-layout/src/layouts/flat/reader.rs` and `grep -q 'mark_pending_wait(values_eval)' vortex-layout/src/layouts/dict/reader.rs` | No |
+| 20 | Callback for object-store I/O waits | spiceai/vortex#107 | The embedding query executor must be able to mark metadata, GET, and response-stream waits without classifying Vortex decode work as remote I/O | `cargo test -p vortex-io --features object_store,tokio --lib io_wait_marker_covers_head_and_both_range_read_paths` | No |
 
 **Confidence:** rows 1–4, 13 and 14 are verified against the branches. Rows 5–12 are seeded from an
 audit of non-merge commits on `spiceai-54` authored by Spice engineers; their descriptions

vortex-file/src/open.rs

@@ -22,6 +22,7 @@ use vortex_layout::segments::SegmentCacheSourceAdapter;
 use vortex_layout::segments::SegmentId;
 use vortex_layout::segments::SegmentSource;
 use vortex_layout::segments::SharedSegmentSource;
+use vortex_layout::segments::WaitMarkedSegmentSource;
 use vortex_layout::session::LayoutSessionExt;
 use vortex_metrics::DefaultMetricsRegistry;
 use vortex_metrics::Label;
@@ -326,9 +327,10 @@ impl VortexOpenOptions {
         )));
 
         // Wrap up the segment source to first resolve segments from the initial read cache.
-        let segment_source: Arc<dyn SegmentSource> = Arc::new(SegmentCacheSourceAdapter::new(
-            segment_cache,
-            segment_source,
+        // The outermost source, so each request pending on a cache fill or a shared read counts
+        // as a wait for the scan split that made it.
+        let segment_source: Arc<dyn SegmentSource> = Arc::new(WaitMarkedSegmentSource::new(
+            SegmentCacheSourceAdapter::new(segment_cache, segment_source),
         ));
 
         let metadata = if self.include_metadata {

vortex-io/src/object_store/read_at.rs

@@ -4,22 +4,25 @@
 use std::io;
 use std::sync::Arc;
 
+use bytes::Bytes;
 use futures::FutureExt;
 use futures::SinkExt;
 use futures::StreamExt;
 use futures::channel::mpsc;
 use futures::future::BoxFuture;
 use futures::stream;
+use futures::stream::BoxStream;
 use object_store::GetOptions;
 use object_store::GetRange;
+use object_store::GetResult;
 use object_store::GetResultPayload;
+use object_store::ObjectMeta;
 use object_store::ObjectStore;
 use object_store::ObjectStoreExt;
 use object_store::path::Path as ObjectPath;
 use vortex_array::buffer::BufferHandle;
 use vortex_array::memory::BufferAllocatorRef;
 use vortex_buffer::Alignment;
-use vortex_error::VortexError;
 use vortex_error::VortexResult;
 use vortex_error::vortex_ensure;
 
@@ -34,6 +37,32 @@ use crate::std_file::read_exact_at;
 /// Default number of concurrent requests to allow.
 pub const DEFAULT_CONCURRENCY: usize = 192;
 
+/// Marks waits on a remote object store without coupling Vortex to a query engine.
+///
+/// The callback receives the network request and its response byte stream separately. It must
+/// preserve their outputs. In particular, it should not mark work that decodes a Vortex array or
+/// copies bytes after a stream item is ready. Applications can leave the callback unset for local
+/// and in-memory stores.
+pub trait ObjectStoreIoWait: Send + Sync + 'static {
+    /// Wrap a remote object metadata request.
+    fn head(
+        &self,
+        future: BoxFuture<'static, object_store::Result<ObjectMeta>>,
+    ) -> BoxFuture<'static, object_store::Result<ObjectMeta>>;
+
+    /// Wrap a remote object range request.
+    fn get(
+        &self,
+        future: BoxFuture<'static, object_store::Result<GetResult>>,
+    ) -> BoxFuture<'static, object_store::Result<GetResult>>;
+
+    /// Wrap each pending wait for the next bytes from a remote response.
+    fn stream(
+        &self,
+        stream: BoxStream<'static, object_store::Result<Bytes>>,
+    ) -> BoxStream<'static, object_store::Result<Bytes>>;
+}
+
 /// An object store backed I/O source.
 pub struct ObjectStoreReadAt {
     store: Arc<dyn ObjectStore>,
@@ -43,6 +72,7 @@ pub struct ObjectStoreReadAt {
     allocator: BufferAllocatorRef,
     concurrency: usize,
     coalesce_config: Option<CoalesceConfig>,
+    io_wait: Option<Arc<dyn ObjectStoreIoWait>>,
 }
 
 impl ObjectStoreReadAt {
@@ -72,9 +102,17 @@ impl ObjectStoreReadAt {
             allocator,
             concurrency: DEFAULT_CONCURRENCY,
             coalesce_config: Some(CoalesceConfig::object_storage()),
+            io_wait: None,
         }
     }
 
+    /// Mark remote object-store waits using the application's I/O-wait contract.
+    /// Leave this unset for local and in-memory object stores.
+    pub fn with_io_wait(mut self, io_wait: Arc<dyn ObjectStoreIoWait>) -> Self {
+        self.io_wait = Some(io_wait);
+        self
+    }
+
     /// Set the concurrency for this source.
     pub fn with_concurrency(mut self, concurrency: usize) -> Self {
         self.concurrency = concurrency;
@@ -94,6 +132,7 @@ async fn read_object_store_range(
     io_handle: Handle,
     allocator: BufferAllocatorRef,
     request: ReadAtRequest,
+    io_wait: Option<Arc<dyn ObjectStoreIoWait>>,
 ) -> VortexResult<BufferHandle> {
     let ReadAtRequest {
         offset,
@@ -105,15 +144,22 @@ async fn read_object_store_range(
     // SAFETY: each return path checks that every byte was initialized.
     unsafe { buffer.set_len(length) };
 
-    let response = store
-        .get_opts(
-            &path,
-            GetOptions {
-                range: Some(GetRange::Bounded(range.clone())),
-                ..Default::default()
-            },
-        )
-        .await?;
+    let get_range = range.clone();
+    let get = async move {
+        store
+            .get_opts(
+                &path,
+                GetOptions {
+                    range: Some(GetRange::Bounded(get_range)),
+                    ..Default::default()
+                },
+            )
+            .await
+    };
+    let response = match &io_wait {
+        Some(marker) => marker.get(get.boxed()).await?,
+        None => get.await?,
+    };
 
     let buffer = match response.payload {
         #[cfg(not(target_arch = "wasm32"))]
@@ -128,7 +174,11 @@ async fn read_object_store_range(
         GetResultPayload::File(..) => {
             unreachable!("File payload not supported on wasm32")
         }
-        GetResultPayload::Stream(mut byte_stream) => {
+        GetResultPayload::Stream(byte_stream) => {
+            let mut byte_stream = match &io_wait {
+                Some(marker) => marker.stream(byte_stream),
+                None => byte_stream,
+            };
             let mut written = 0usize;
             while let Some(bytes) = byte_stream.next().await {
                 let bytes = bytes?;
@@ -175,12 +225,14 @@ impl VortexReadAt for ObjectStoreReadAt {
     fn size(&self) -> BoxFuture<'static, VortexResult<u64>> {
         let store = Arc::clone(&self.store);
         let path = self.path.clone();
+        let io_wait = self.io_wait.clone();
         async move {
-            store
-                .head(&path)
-                .await
-                .map(|h| h.size)
-                .map_err(VortexError::from)
+            let head = async move { store.head(&path).await };
+            let meta = match io_wait {
+                Some(marker) => marker.head(head.boxed()).await?,
+                None => head.await?,
+            };
+            Ok(meta.size)
         }
         .boxed()
     }
@@ -196,13 +248,15 @@ impl VortexReadAt for ObjectStoreReadAt {
         let handle = self.handle.clone();
         let allocator = self.allocator.clone();
         let io_handle = handle.clone();
+        let io_wait = self.io_wait.clone();
         handle
             .spawn_io(read_object_store_range(
                 store,
                 path,
                 io_handle,
                 allocator,
                 ReadAtRequest::new(offset, length, alignment),
+                io_wait,
             ))
             .boxed()
     }
@@ -216,6 +270,7 @@ impl VortexReadAt for ObjectStoreReadAt {
         let path = self.path.clone();
         let handle = self.handle.clone();
         let allocator = self.allocator.clone();
+        let io_wait = self.io_wait.clone();
         let concurrency = self.concurrency.max(1);
         let (mut send, recv) = mpsc::channel(concurrency);
         let io_handle = handle.clone();
@@ -229,9 +284,12 @@ impl VortexReadAt for ObjectStoreReadAt {
                 let path = path.clone();
                 let io_handle = io_handle.clone();
                 let allocator = allocator.clone();
+                let io_wait = io_wait.clone();
                 async move {
-                    let result =
-                        read_object_store_range(store, path, io_handle, allocator, request).await;
+                    let result = read_object_store_range(
+                        store, path, io_handle, allocator, request, io_wait,
+                    )
+                    .await;
                     (request, result)
                 }
             });
@@ -277,6 +335,39 @@ mod tests {
         spawn_io_count: AtomicUsize,
     }
 
+    #[derive(Default)]
+    struct CountingIoWait {
+        heads: AtomicUsize,
+        gets: AtomicUsize,
+        streams: AtomicUsize,
+    }
+
+    impl ObjectStoreIoWait for CountingIoWait {
+        fn head(
+            &self,
+            future: BoxFuture<'static, object_store::Result<ObjectMeta>>,
+        ) -> BoxFuture<'static, object_store::Result<ObjectMeta>> {
+            self.heads.fetch_add(1, Ordering::SeqCst);
+            future
+        }
+
+        fn get(
+            &self,
+            future: BoxFuture<'static, object_store::Result<GetResult>>,
+        ) -> BoxFuture<'static, object_store::Result<GetResult>> {
+            self.gets.fetch_add(1, Ordering::SeqCst);
+            future
+        }
+
+        fn stream(
+            &self,
+            stream: BoxStream<'static, object_store::Result<Bytes>>,
+        ) -> BoxStream<'static, object_store::Result<Bytes>> {
+            self.streams.fetch_add(1, Ordering::SeqCst);
+            stream
+        }
+    }
+
     impl Executor for CountingExecutor {
         fn spawn(&self, fut: BoxFuture<'static, ()>) -> AbortHandleRef {
             self.spawn_count.fetch_add(1, Ordering::SeqCst);
@@ -364,4 +455,34 @@ mod tests {
 
         Ok(())
     }
+
+    #[tokio::test]
+    async fn io_wait_marker_covers_head_and_both_range_read_paths() -> anyhow::Result<()> {
+        let executor = Arc::new(CountingExecutor::default());
+        let runtime = Arc::clone(&executor) as Arc<dyn Executor>;
+        let handle = Handle::new(Arc::downgrade(&runtime));
+        let store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
+        let path = ObjectPath::from("test.bin");
+        store.put(&path, PutPayload::from_static(TEST_DATA)).await?;
+
+        let marker = Arc::new(CountingIoWait::default());
+        let reader = ObjectStoreReadAt::new(store, path, handle)
+            .with_io_wait(Arc::clone(&marker) as Arc<dyn ObjectStoreIoWait>);
+        assert_eq!(reader.size().await?, TEST_DATA.len() as u64);
+        assert_eq!(reader.read_at(7, 5, Alignment::new(1)).await?.len(), 5);
+        let requests: Arc<[ReadAtRequest]> = Arc::from([
+            ReadAtRequest::new(0, 6, Alignment::new(1)),
+            ReadAtRequest::new(7, 5, Alignment::new(1)),
+        ]);
+        let results = reader.read_ranges(requests).collect::<Vec<_>>().await;
+        assert_eq!(results.len(), 2);
+        for (_, result) in results {
+            result?;
+        }
+
+        assert_eq!(marker.heads.load(Ordering::SeqCst), 1);
+        assert_eq!(marker.gets.load(Ordering::SeqCst), 3);
+        assert_eq!(marker.streams.load(Ordering::SeqCst), 3);
+        Ok(())
+    }
 }

vortex-layout/Cargo.toml

@@ -62,7 +62,12 @@ insta = { workspace = true }
 rand = { workspace = true }
 rstest = { workspace = true }
 temp-env = { workspace = true }
-tokio = { workspace = true, features = ["rt", "macros"] }
+tokio = { workspace = true, features = [
+    "rt",
+    "macros",
+    "rt-multi-thread",
+    "time",
+] }
 vortex-array = { path = "../vortex-array", features = ["_test-harness"] }
 vortex-io = { path = "../vortex-io", features = ["tokio"] }
 

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

@@ -42,6 +42,7 @@ use crate::RowSplits;
 use crate::SplitRange;
 use crate::layouts::SharedArrayFuture;
 use crate::segments::SegmentSource;
+use crate::segments::mark_pending_wait;
 
 pub struct DictReader {
     layout: DictLayout,
@@ -273,7 +274,10 @@ impl LayoutReader for DictReader {
 
         Ok(MaskFuture::new(mask.len(), async move {
             // Join on the I/O futures first, before the mask.
-            let (codes, values) = try_join!(codes_eval, values_eval.map_err(VortexError::from))?;
+            let (codes, values) = try_join!(
+                codes_eval,
+                mark_pending_wait(values_eval).map_err(VortexError::from)
+            )?;
             let mask = mask.await?;
 
             let mut ctx = session.create_execution_ctx();
@@ -325,7 +329,10 @@ impl LayoutReader for DictReader {
         };
         let all_values_referenced = self.layout.has_all_values_referenced();
         Ok(async move {
-            let (values, codes) = try_join!(values_eval.map_err(VortexError::from), codes_eval)?;
+            let (values, codes) = try_join!(
+                mark_pending_wait(values_eval).map_err(VortexError::from),
+                codes_eval
+            )?;
 
             // SAFETY: Layout was validated at write time.
             //  * The codes dtype is guaranteed to be an integer type from the layout

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

@@ -27,6 +27,7 @@ use crate::reader::LayoutReader;
 use crate::reader::RowSplits;
 use crate::reader::SplitRange;
 use crate::segments::SegmentSource;
+use crate::segments::mark_pending_wait;
 
 /// The threshold of mask density below which we will evaluate the expression only over the
 /// selected rows, and above which we evaluate the expression over all rows and then select
@@ -144,7 +145,7 @@ impl LayoutReader for FlatReader {
             // TODO(ngates): if the mask density is low enough, or if the mask is dense within a range
             //  (as often happens with zone map pruning), then we could slice/filter the array prior
             //  to evaluating the expression.
-            let mut array = array.clone().await?;
+            let mut array = mark_pending_wait(array).await?;
             let mask = mask.await?;
 
             // Slice the array based on the row mask.
@@ -201,7 +202,7 @@ impl LayoutReader for FlatReader {
         Ok(async move {
             trace!("Flat array evaluation {} - {}", name, expr);
 
-            let mut array = array.clone().await?;
+            let mut array = mark_pending_wait(array).await?;
             let mask = mask.await?;
 
             // Slice the array based on the row mask.

vortex-layout/src/scan/mod.rs

@@ -8,6 +8,7 @@ pub mod multi;
 pub mod repeated_scan;
 pub mod scan_builder;
 pub mod split_by;
+pub(crate) mod split_gate;
 mod splits;
 mod tasks;
 #[cfg(test)]