From 1ea9321fcb24318bbfa983a90d2d0f597af1a1a9 Mon Sep 17 00:00:00 2001 From: Joseph Isaacs Date: Wed, 12 Aug 2026 15:10:42 +0000 Subject: [PATCH 1/8] Add batched positional reads Signed-off-by: Joseph Isaacs --- vortex-file/src/segments/source.rs | 134 ++++++++++++++++++++++---- vortex-io/src/compat/read_at.rs | 8 ++ vortex-io/src/object_store/read_at.rs | 84 ++++++++++++++++ vortex-io/src/read_at.rs | 96 ++++++++++++++++++ 4 files changed, 303 insertions(+), 19 deletions(-) diff --git a/vortex-file/src/segments/source.rs b/vortex-file/src/segments/source.rs index 3af33362b05..bbeafa37d2d 100644 --- a/vortex-file/src/segments/source.rs +++ b/vortex-file/src/segments/source.rs @@ -20,10 +20,11 @@ use parking_lot::Mutex; use vortex_array::buffer::BufferHandle; use vortex_buffer::Alignment; use vortex_buffer::ByteBuffer; +use vortex_error::VortexError; use vortex_error::VortexResult; -use vortex_error::vortex_bail; use vortex_error::vortex_err; use vortex_error::vortex_panic; +use vortex_io::ReadAtRequest; use vortex_io::VortexReadAt; use vortex_io::runtime::Handle; use vortex_io::runtime::JoinOutcome; @@ -140,29 +141,52 @@ impl FileSegmentSource { let drive_fut = async move { stream - .map(move |req| { + .ready_chunks(concurrency) + .for_each(move |reqs| { let reader = reader.clone(); async move { - let result = reader - .read_at(req.offset(), req.len(), req.alignment()) - .await; - let result = result.and_then(|buffer| { - if req.len() != buffer.len() { - vortex_bail!( - "FileSegmentSource: expected buffer of length {} but received {}. {:?}", - req.len(), - buffer.len(), - req - ) + let requests = reqs + .iter() + .map(|req| { + ReadAtRequest::new(req.offset(), req.len(), req.alignment()) + }) + .collect::>() + .into(); + match reader.read_ranges(requests).await { + Ok(buffers) if buffers.len() == reqs.len() => { + for (req, buffer) in reqs.into_iter().zip(buffers) { + let result = if req.len() == buffer.len() { + Ok(buffer) + } else { + Err(vortex_err!( + "FileSegmentSource: expected buffer of length {} but received {}. {:?}", + req.len(), + buffer.len(), + req + )) + }; + req.resolve(result); + } } - Ok(buffer) - }); - - req.resolve(result); + Ok(buffers) => { + let error = Arc::new(vortex_err!( + "FileSegmentSource: expected {} buffers but received {}", + reqs.len(), + buffers.len() + )); + for req in reqs { + req.resolve(Err(VortexError::from(Arc::clone(&error)))); + } + } + Err(error) => { + let error = Arc::new(error); + for req in reqs { + req.resolve(Err(VortexError::from(Arc::clone(&error)))); + } + } + } } }) - .buffer_unordered(concurrency) - .collect::<()>() .await }; @@ -383,6 +407,7 @@ mod tests { use std::panic::AssertUnwindSafe; use futures::future::BoxFuture; + use vortex_error::vortex_bail; use vortex_io::runtime::tokio::TokioRuntime; use vortex_layout::segments::SegmentSource; use vortex_metrics::DefaultMetricsRegistry; @@ -510,6 +535,77 @@ mod tests { ); } + #[derive(Clone)] + struct ReadRangesOnly { + calls: Arc, + } + + impl VortexReadAt for ReadRangesOnly { + fn concurrency(&self) -> usize { + 4 + } + + fn size(&self) -> BoxFuture<'static, VortexResult> { + async { Ok(16) }.boxed() + } + + fn read_at( + &self, + _offset: u64, + _length: usize, + _alignment: Alignment, + ) -> BoxFuture<'static, VortexResult> { + async { panic!("read_at should not be called") }.boxed() + } + + fn read_ranges( + &self, + requests: Arc<[ReadAtRequest]>, + ) -> BoxFuture<'static, VortexResult>> { + self.calls.fetch_add(1, Ordering::Relaxed); + async move { + Ok(requests + .iter() + .map(|request| { + BufferHandle::new_host( + ByteBuffer::from(vec![0; request.length]).aligned(request.alignment), + ) + }) + .collect()) + } + .boxed() + } + } + + #[tokio::test] + async fn read_driver_batches_ready_requests() -> VortexResult<()> { + let calls = Arc::new(AtomicUsize::new(0)); + let segments: Arc<[SegmentSpec]> = (0..4) + .map(|i| SegmentSpec { + offset: i * 4, + length: 4, + alignment: Alignment::none(), + }) + .collect(); + let metrics = DefaultMetricsRegistry::default(); + let source = FileSegmentSource::open( + segments, + ReadRangesOnly { + calls: Arc::clone(&calls), + }, + TokioRuntime::current(), + RequestMetrics::new(&metrics, vec![]), + ); + + let results = future::join_all((0..4).map(|i| source.request(SegmentId::from(i)))).await; + + for result in results { + assert_eq!(result?.len(), 4); + } + assert_eq!(calls.load(Ordering::Relaxed), 1); + Ok(()) + } + #[derive(Clone)] struct SlowErrReadAt; diff --git a/vortex-io/src/compat/read_at.rs b/vortex-io/src/compat/read_at.rs index 4fc49785d28..353947167cd 100644 --- a/vortex-io/src/compat/read_at.rs +++ b/vortex-io/src/compat/read_at.rs @@ -10,6 +10,7 @@ use vortex_buffer::Alignment; use vortex_error::VortexResult; use crate::CoalesceConfig; +use crate::ReadAtRequest; use crate::VortexReadAt; use crate::compat::Compat; @@ -40,4 +41,11 @@ impl VortexReadAt for Compat { ) -> BoxFuture<'static, VortexResult> { Compat::new(self.inner().read_at(offset, length, alignment)).boxed() } + + fn read_ranges( + &self, + requests: Arc<[ReadAtRequest]>, + ) -> BoxFuture<'static, VortexResult>> { + Compat::new(self.inner().read_ranges(requests)).boxed() + } } diff --git a/vortex-io/src/object_store/read_at.rs b/vortex-io/src/object_store/read_at.rs index 086d70c1bcf..d1f8c427185 100644 --- a/vortex-io/src/object_store/read_at.rs +++ b/vortex-io/src/object_store/read_at.rs @@ -20,8 +20,10 @@ use vortex_buffer::Alignment; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use crate::CoalesceConfig; +use crate::ReadAtRequest; use crate::VortexReadAt; use crate::runtime::Handle; #[cfg(not(target_arch = "wasm32"))] @@ -181,6 +183,62 @@ impl VortexReadAt for ObjectStoreReadAt { }) .boxed() } + + fn read_ranges( + &self, + requests: Arc<[ReadAtRequest]>, + ) -> BoxFuture<'static, VortexResult>> { + let store = Arc::clone(&self.store); + let path = self.path.clone(); + let handle = self.handle.clone(); + let allocator = Arc::clone(&self.allocator); + + handle + .spawn_io(async move { + let ranges = requests + .iter() + .map(|request| { + let end = request + .offset + .checked_add(request.length as u64) + .ok_or_else(|| { + vortex_err!( + "Read range overflows u64: offset={}, length={}", + request.offset, + request.length + ) + })?; + Ok(request.offset..end) + }) + .collect::>>()?; + let bytes = store.get_ranges(&path, &ranges).await?; + vortex_ensure!( + bytes.len() == requests.len(), + "Object store returned {} ranges but expected {}", + bytes.len(), + requests.len() + ); + + requests + .iter() + .zip(bytes) + .map(|(request, bytes)| { + vortex_ensure!( + bytes.len() == request.length, + "Object store returned {} bytes but expected {} for range {}..{}", + bytes.len(), + request.length, + request.offset, + request.offset + request.length as u64 + ); + let mut buffer = allocator.allocate(request.length, request.alignment)?; + buffer.as_mut_slice().copy_from_slice(&bytes); + Ok(BufferHandle::new_host(buffer.freeze())) + }) + .collect() + }) + .boxed() + } } #[cfg(test)] @@ -258,4 +316,30 @@ mod tests { Ok(()) } + + #[tokio::test] + async fn read_ranges_uses_spawn_io() -> anyhow::Result<()> { + let executor = Arc::new(CountingExecutor::default()); + let runtime = Arc::clone(&executor) as Arc; + let handle = Handle::new(Arc::downgrade(&runtime)); + + let store = Arc::new(InMemory::new()) as Arc; + let path = ObjectPath::from("test.bin"); + store.put(&path, PutPayload::from_static(TEST_DATA)).await?; + + let reader = ObjectStoreReadAt::new(store, path, handle); + let buffers = reader + .read_ranges(Arc::from([ + ReadAtRequest::new(7, 5, Alignment::new(1)), + ReadAtRequest::new(0, 6, Alignment::new(1)), + ])) + .await?; + + assert_eq!(buffers[0].to_host().await.as_slice(), b"store"); + assert_eq!(buffers[1].to_host().await.as_slice(), b"object"); + assert_eq!(executor.spawn_io_count.load(Ordering::SeqCst), 1); + assert_eq!(executor.spawn_count.load(Ordering::SeqCst), 0); + + Ok(()) + } } diff --git a/vortex-io/src/read_at.rs b/vortex-io/src/read_at.rs index aa9a8a03abf..ee0d185c3e1 100644 --- a/vortex-io/src/read_at.rs +++ b/vortex-io/src/read_at.rs @@ -4,7 +4,10 @@ use std::sync::Arc; use futures::FutureExt; +use futures::StreamExt; +use futures::TryStreamExt; use futures::future::BoxFuture; +use futures::stream; use vortex_array::buffer::BufferHandle; use vortex_buffer::Alignment; use vortex_buffer::ByteBuffer; @@ -27,6 +30,28 @@ pub struct CoalesceConfig { pub max_size: u64, } +/// A positional read request used by [`VortexReadAt::read_ranges`]. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ReadAtRequest { + /// The byte offset at which to start reading. + pub offset: u64, + /// The exact number of bytes to read. + pub length: usize, + /// The required alignment of the returned buffer. + pub alignment: Alignment, +} + +impl ReadAtRequest { + /// Creates a positional read request. + pub const fn new(offset: u64, length: usize, alignment: Alignment) -> Self { + Self { + offset, + length, + alignment, + } + } +} + impl CoalesceConfig { /// Creates a new coalesce configuration. pub const fn new(distance: u64, max_size: u64) -> Self { @@ -89,6 +114,26 @@ pub trait VortexReadAt: Send + Sync + 'static { length: usize, alignment: Alignment, ) -> BoxFuture<'static, VortexResult>; + + /// Request multiple asynchronous positional reads. + /// + /// Results are returned in request order. The default implementation executes + /// [`VortexReadAt::read_at`] calls concurrently, bounded by [`VortexReadAt::concurrency`]. + /// If any request fails, the entire operation fails. Implementations can override this to use + /// a native multi-range operation. + fn read_ranges( + &self, + requests: Arc<[ReadAtRequest]>, + ) -> BoxFuture<'static, VortexResult>> { + let reads = requests + .iter() + .map(|request| self.read_at(request.offset, request.length, request.alignment)) + .collect::>(); + stream::iter(reads) + .buffered(self.concurrency().max(1)) + .try_collect() + .boxed() + } } impl VortexReadAt for Arc { @@ -116,6 +161,13 @@ impl VortexReadAt for Arc { ) -> BoxFuture<'static, VortexResult> { self.as_ref().read_at(offset, length, alignment) } + + fn read_ranges( + &self, + requests: Arc<[ReadAtRequest]>, + ) -> BoxFuture<'static, VortexResult>> { + self.as_ref().read_ranges(requests) + } } impl VortexReadAt for Arc { @@ -143,6 +195,13 @@ impl VortexReadAt for Arc { ) -> BoxFuture<'static, VortexResult> { self.as_ref().read_at(offset, length, alignment) } + + fn read_ranges( + &self, + requests: Arc<[ReadAtRequest]>, + ) -> BoxFuture<'static, VortexResult>> { + self.as_ref().read_ranges(requests) + } } impl VortexReadAt for ByteBuffer { @@ -316,6 +375,26 @@ impl VortexReadAt for InstrumentedReadAt { } .boxed() } + + fn read_ranges( + &self, + requests: Arc<[ReadAtRequest]>, + ) -> BoxFuture<'static, VortexResult>> { + let durations = self.metrics.durations.clone(); + let sizes = self.metrics.sizes.clone(); + let total_size = self.metrics.total_size.clone(); + let read_fut = self.read.read_ranges(Arc::clone(&requests)); + async move { + let _timer = durations.time(); + let buffers = read_fut.await; + for request in requests.iter() { + sizes.update(request.length as f64); + total_size.add(request.length as u64); + } + buffers + } + .boxed() + } } #[cfg(test)] @@ -356,6 +435,23 @@ mod tests { assert_eq!(result.to_host().await.as_ref(), &[2, 3, 4]); } + #[tokio::test] + async fn test_byte_buffer_read_ranges() -> VortexResult<()> { + let data = ByteBuffer::from(vec![1, 2, 3, 4, 5, 6]); + let requests = Arc::from([ + ReadAtRequest::new(4, 2, Alignment::none()), + ReadAtRequest::new(0, 1, Alignment::none()), + ReadAtRequest::new(2, 3, Alignment::none()), + ]); + + let results = data.read_ranges(requests).await?; + let expected: [&[u8]; 3] = [&[5, 6], &[1], &[3, 4, 5]]; + for (result, expected) in results.into_iter().zip(expected) { + assert_eq!(result.to_host().await.as_ref(), expected); + } + Ok(()) + } + #[tokio::test] async fn test_byte_buffer_read_out_of_bounds() { let data = ByteBuffer::from(vec![1, 2, 3]); From e9bf1dd76f81394500fc564a6a07dd4ca06b74c5 Mon Sep 17 00:00:00 2001 From: Joseph Isaacs Date: Wed, 12 Aug 2026 16:12:46 +0000 Subject: [PATCH 2/8] Stream batched positional read results Signed-off-by: Joseph Isaacs --- vortex-file/src/read/driver.rs | 63 ++++++++++-- vortex-file/src/segments/source.rs | 96 +++++++++--------- vortex-io/src/compat/read_at.rs | 7 +- vortex-io/src/object_store/read_at.rs | 84 --------------- vortex-io/src/read_at.rs | 141 +++++++++++++++++++------- 5 files changed, 208 insertions(+), 183 deletions(-) diff --git a/vortex-file/src/read/driver.rs b/vortex-file/src/read/driver.rs index 7f6dc3b2f7c..c393901b724 100644 --- a/vortex-file/src/read/driver.rs +++ b/vortex-file/src/read/driver.rs @@ -30,13 +30,14 @@ pin_project! { /// an ordering of `(has_been_polled, insertion_order)`, skipping any canceled requests, and /// then coalescing with other nearby requests within the configured `window`. /// - /// The output of this stream is expected to be buffered by the desired I/O concurrency, and - /// driven to completion. + /// The output contains up to `batch_size` immediately eligible physical requests. A poll never + /// waits to fill a batch. pub(crate) struct IoRequestStream { #[pin] events: S, inner_done: bool, coalesce_window: Option, + batch_size: usize, state: State, } } @@ -48,15 +49,18 @@ impl IoRequestStream { events: S, coalesce_window: Option, coalesced_buffer_alignment: Alignment, + batch_size: usize, metrics: RequestMetrics, ) -> Self where S: Stream + Unpin + Send + 'static, { + assert!(batch_size > 0, "I/O request batch size must be non-zero"); IoRequestStream { events, inner_done: false, coalesce_window, + batch_size, state: State::new(metrics, coalesced_buffer_alignment), } } @@ -66,7 +70,7 @@ impl Stream for IoRequestStream where S: Stream + Unpin + Send + 'static, { - type Item = IoRequest; + type Item = Vec; fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let mut this = self.project(); @@ -87,9 +91,16 @@ where } } - // Try to get a coalesced request - if let Some(coalesced) = this.state.next(this.coalesce_window.as_ref()) { - return Poll::Ready(Some(coalesced)); + // Return up to batch_size requests that are eligible now. Do not wait to fill the batch. + let mut batch = Vec::with_capacity(*this.batch_size); + while batch.len() < *this.batch_size { + let Some(request) = this.state.next(this.coalesce_window.as_ref()) else { + break; + }; + batch.push(request); + } + if !batch.is_empty() { + return Poll::Ready(Some(batch)); } // If the inner stream is done, and we have no more _polled_ requests, we're done @@ -374,9 +385,10 @@ mod tests { event_stream, coalesce_window, coalesced_buffer_alignment, + 1024, metrics, ); - io_stream.collect().await + io_stream.concat().await } #[tokio::test] @@ -415,6 +427,36 @@ mod tests { assert_eq!(offsets, vec![0, 100, 200]); // req1, req2, req3 } + #[tokio::test] + async fn test_bounded_request_batches() { + let mut events = Vec::new(); + let mut receivers = Vec::new(); + for id in 0..5 { + let (request, recv) = create_request(id, id as u64 * 10, 10); + events.push(ReadEvent::Request(request)); + events.push(ReadEvent::Polled(id)); + receivers.push(recv); + } + + let metrics_registry = DefaultMetricsRegistry::default(); + let metrics = RequestMetrics::new(&metrics_registry, vec![]); + let batches = + IoRequestStream::new(stream::iter(events), None, Alignment::none(), 2, metrics) + .collect::>() + .await; + + assert_eq!(receivers.len(), 5); + assert_eq!(batches.iter().map(Vec::len).collect::>(), [2, 2, 1]); + assert_eq!( + batches + .into_iter() + .flatten() + .map(|request| request.offset()) + .collect::>(), + [0, 10, 20, 30, 40] + ); + } + #[tokio::test] async fn test_coalesce_adjacent() { let (req1, _rx1) = create_request(1, 0, 10); @@ -734,10 +776,11 @@ mod tests { max_size: 1024, }), Alignment::none(), + 1024, metrics, ); - let outputs: Vec = io_stream.collect().await; + let outputs: Vec = io_stream.concat().await; assert_eq!(outputs.len(), 2); let snapshot = metrics_registry.snapshot(); @@ -788,9 +831,9 @@ mod tests { let metrics_registry = DefaultMetricsRegistry::default(); let metrics = RequestMetrics::new(&metrics_registry, vec![]); // No coalescing window - should be individual requests - let io_stream = IoRequestStream::new(event_stream, None, Alignment::none(), metrics); + let io_stream = IoRequestStream::new(event_stream, None, Alignment::none(), 1024, metrics); - let outputs: Vec = io_stream.collect().await; + let outputs: Vec = io_stream.concat().await; assert_eq!(outputs.len(), 2); // Check metrics diff --git a/vortex-file/src/segments/source.rs b/vortex-file/src/segments/source.rs index bbeafa37d2d..44cd37dbad2 100644 --- a/vortex-file/src/segments/source.rs +++ b/vortex-file/src/segments/source.rs @@ -20,7 +20,7 @@ use parking_lot::Mutex; use vortex_array::buffer::BufferHandle; use vortex_buffer::Alignment; use vortex_buffer::ByteBuffer; -use vortex_error::VortexError; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_error::vortex_panic; @@ -135,13 +135,13 @@ impl FileSegmentSource { StreamExt::boxed(recv), coalesce_config, max_alignment, + concurrency, metrics, ) .boxed(); let drive_fut = async move { stream - .ready_chunks(concurrency) .for_each(move |reqs| { let reader = reader.clone(); async move { @@ -152,38 +152,41 @@ impl FileSegmentSource { }) .collect::>() .into(); - match reader.read_ranges(requests).await { - Ok(buffers) if buffers.len() == reqs.len() => { - for (req, buffer) in reqs.into_iter().zip(buffers) { - let result = if req.len() == buffer.len() { - Ok(buffer) - } else { - Err(vortex_err!( - "FileSegmentSource: expected buffer of length {} but received {}. {:?}", - req.len(), - buffer.len(), - req - )) - }; - req.resolve(result); + let mut remaining = reqs.into_iter().map(Some).collect::>(); + let mut results = reader.read_ranges(requests); + while let Some((request, result)) = results.next().await { + let Some(position) = remaining.iter().position(|req| { + req.as_ref().is_some_and(|req| { + req.offset() == request.offset + && req.len() == request.length + && req.alignment() == request.alignment + }) + }) else { + tracing::warn!(?request, "reader returned an unknown range"); + continue; + }; + let req = remaining[position] + .take() + .vortex_expect("matched request is present"); + let result = result.and_then(|buffer| { + if req.len() != buffer.len() { + return Err(vortex_err!( + "FileSegmentSource: expected buffer of length {} but received {}. {:?}", + req.len(), + buffer.len(), + req + )); } - } - Ok(buffers) => { - let error = Arc::new(vortex_err!( - "FileSegmentSource: expected {} buffers but received {}", - reqs.len(), - buffers.len() - )); - for req in reqs { - req.resolve(Err(VortexError::from(Arc::clone(&error)))); - } - } - Err(error) => { - let error = Arc::new(error); - for req in reqs { - req.resolve(Err(VortexError::from(Arc::clone(&error)))); - } - } + Ok(buffer) + }); + req.resolve(result); + } + for req in remaining.into_iter().flatten() { + let error = vortex_err!( + "FileSegmentSource: read_ranges ended before resolving request. {:?}", + req + ); + req.resolve(Err(error)); } } }) @@ -558,22 +561,19 @@ mod tests { async { panic!("read_at should not be called") }.boxed() } - fn read_ranges( - &self, - requests: Arc<[ReadAtRequest]>, - ) -> BoxFuture<'static, VortexResult>> { + fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> vortex_io::ReadAtStream { self.calls.fetch_add(1, Ordering::Relaxed); - async move { - Ok(requests - .iter() - .map(|request| { - BufferHandle::new_host( - ByteBuffer::from(vec![0; request.length]).aligned(request.alignment), - ) - }) - .collect()) - } - .boxed() + let results = requests + .iter() + .copied() + .map(|request| { + let buffer = BufferHandle::new_host( + ByteBuffer::from(vec![0; request.length]).aligned(request.alignment), + ); + (request, Ok(buffer)) + }) + .collect::>(); + futures::stream::iter(results).boxed() } } diff --git a/vortex-io/src/compat/read_at.rs b/vortex-io/src/compat/read_at.rs index 353947167cd..3d9cc93b1a6 100644 --- a/vortex-io/src/compat/read_at.rs +++ b/vortex-io/src/compat/read_at.rs @@ -4,6 +4,7 @@ use std::sync::Arc; use futures::FutureExt; +use futures::StreamExt; use futures::future::BoxFuture; use vortex_array::buffer::BufferHandle; use vortex_buffer::Alignment; @@ -11,6 +12,7 @@ use vortex_error::VortexResult; use crate::CoalesceConfig; use crate::ReadAtRequest; +use crate::ReadAtStream; use crate::VortexReadAt; use crate::compat::Compat; @@ -42,10 +44,7 @@ impl VortexReadAt for Compat { Compat::new(self.inner().read_at(offset, length, alignment)).boxed() } - fn read_ranges( - &self, - requests: Arc<[ReadAtRequest]>, - ) -> BoxFuture<'static, VortexResult>> { + fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> ReadAtStream { Compat::new(self.inner().read_ranges(requests)).boxed() } } diff --git a/vortex-io/src/object_store/read_at.rs b/vortex-io/src/object_store/read_at.rs index d1f8c427185..086d70c1bcf 100644 --- a/vortex-io/src/object_store/read_at.rs +++ b/vortex-io/src/object_store/read_at.rs @@ -20,10 +20,8 @@ use vortex_buffer::Alignment; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_ensure; -use vortex_error::vortex_err; use crate::CoalesceConfig; -use crate::ReadAtRequest; use crate::VortexReadAt; use crate::runtime::Handle; #[cfg(not(target_arch = "wasm32"))] @@ -183,62 +181,6 @@ impl VortexReadAt for ObjectStoreReadAt { }) .boxed() } - - fn read_ranges( - &self, - requests: Arc<[ReadAtRequest]>, - ) -> BoxFuture<'static, VortexResult>> { - let store = Arc::clone(&self.store); - let path = self.path.clone(); - let handle = self.handle.clone(); - let allocator = Arc::clone(&self.allocator); - - handle - .spawn_io(async move { - let ranges = requests - .iter() - .map(|request| { - let end = request - .offset - .checked_add(request.length as u64) - .ok_or_else(|| { - vortex_err!( - "Read range overflows u64: offset={}, length={}", - request.offset, - request.length - ) - })?; - Ok(request.offset..end) - }) - .collect::>>()?; - let bytes = store.get_ranges(&path, &ranges).await?; - vortex_ensure!( - bytes.len() == requests.len(), - "Object store returned {} ranges but expected {}", - bytes.len(), - requests.len() - ); - - requests - .iter() - .zip(bytes) - .map(|(request, bytes)| { - vortex_ensure!( - bytes.len() == request.length, - "Object store returned {} bytes but expected {} for range {}..{}", - bytes.len(), - request.length, - request.offset, - request.offset + request.length as u64 - ); - let mut buffer = allocator.allocate(request.length, request.alignment)?; - buffer.as_mut_slice().copy_from_slice(&bytes); - Ok(BufferHandle::new_host(buffer.freeze())) - }) - .collect() - }) - .boxed() - } } #[cfg(test)] @@ -316,30 +258,4 @@ mod tests { Ok(()) } - - #[tokio::test] - async fn read_ranges_uses_spawn_io() -> anyhow::Result<()> { - let executor = Arc::new(CountingExecutor::default()); - let runtime = Arc::clone(&executor) as Arc; - let handle = Handle::new(Arc::downgrade(&runtime)); - - let store = Arc::new(InMemory::new()) as Arc; - let path = ObjectPath::from("test.bin"); - store.put(&path, PutPayload::from_static(TEST_DATA)).await?; - - let reader = ObjectStoreReadAt::new(store, path, handle); - let buffers = reader - .read_ranges(Arc::from([ - ReadAtRequest::new(7, 5, Alignment::new(1)), - ReadAtRequest::new(0, 6, Alignment::new(1)), - ])) - .await?; - - assert_eq!(buffers[0].to_host().await.as_slice(), b"store"); - assert_eq!(buffers[1].to_host().await.as_slice(), b"object"); - assert_eq!(executor.spawn_io_count.load(Ordering::SeqCst), 1); - assert_eq!(executor.spawn_count.load(Ordering::SeqCst), 0); - - Ok(()) - } } diff --git a/vortex-io/src/read_at.rs b/vortex-io/src/read_at.rs index ee0d185c3e1..82190bdd699 100644 --- a/vortex-io/src/read_at.rs +++ b/vortex-io/src/read_at.rs @@ -2,12 +2,13 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::sync::Arc; +use std::time::Instant; use futures::FutureExt; use futures::StreamExt; -use futures::TryStreamExt; use futures::future::BoxFuture; use futures::stream; +use futures::stream::BoxStream; use vortex_array::buffer::BufferHandle; use vortex_buffer::Alignment; use vortex_buffer::ByteBuffer; @@ -52,6 +53,9 @@ impl ReadAtRequest { } } +/// A stream of positional read results, yielded as each request completes. +pub type ReadAtStream = BoxStream<'static, (ReadAtRequest, VortexResult)>; + impl CoalesceConfig { /// Creates a new coalesce configuration. pub const fn new(distance: u64, max_size: u64) -> Self { @@ -117,21 +121,20 @@ pub trait VortexReadAt: Send + Sync + 'static { /// Request multiple asynchronous positional reads. /// - /// Results are returned in request order. The default implementation executes - /// [`VortexReadAt::read_at`] calls concurrently, bounded by [`VortexReadAt::concurrency`]. - /// If any request fails, the entire operation fails. Implementations can override this to use - /// a native multi-range operation. - fn read_ranges( - &self, - requests: Arc<[ReadAtRequest]>, - ) -> BoxFuture<'static, VortexResult>> { + /// Each item includes its request and result, and is yielded as soon as that read completes. + /// A failed request does not prevent other results from being yielded. Callers should submit + /// batches no larger than [`VortexReadAt::concurrency`]. + fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> ReadAtStream { let reads = requests .iter() - .map(|request| self.read_at(request.offset, request.length, request.alignment)) + .copied() + .map(|request| { + let read = self.read_at(request.offset, request.length, request.alignment); + async move { (request, read.await) } + }) .collect::>(); stream::iter(reads) - .buffered(self.concurrency().max(1)) - .try_collect() + .buffer_unordered(self.concurrency().max(1)) .boxed() } } @@ -162,10 +165,7 @@ impl VortexReadAt for Arc { self.as_ref().read_at(offset, length, alignment) } - fn read_ranges( - &self, - requests: Arc<[ReadAtRequest]>, - ) -> BoxFuture<'static, VortexResult>> { + fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> ReadAtStream { self.as_ref().read_ranges(requests) } } @@ -196,10 +196,7 @@ impl VortexReadAt for Arc { self.as_ref().read_at(offset, length, alignment) } - fn read_ranges( - &self, - requests: Arc<[ReadAtRequest]>, - ) -> BoxFuture<'static, VortexResult>> { + fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> ReadAtStream { self.as_ref().read_ranges(requests) } } @@ -376,36 +373,62 @@ impl VortexReadAt for InstrumentedReadAt { .boxed() } - fn read_ranges( - &self, - requests: Arc<[ReadAtRequest]>, - ) -> BoxFuture<'static, VortexResult>> { + fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> ReadAtStream { let durations = self.metrics.durations.clone(); let sizes = self.metrics.sizes.clone(); let total_size = self.metrics.total_size.clone(); - let read_fut = self.read.read_ranges(Arc::clone(&requests)); - async move { - let _timer = durations.time(); - let buffers = read_fut.await; - for request in requests.iter() { + let start = Instant::now(); + self.read + .read_ranges(requests) + .map(move |(request, result)| { + durations.update(start.elapsed()); sizes.update(request.length as f64); total_size.add(request.length as u64); - } - buffers - } - .boxed() + (request, result) + }) + .boxed() } } #[cfg(test)] mod tests { use std::sync::Arc; + use std::time::Duration; use vortex_buffer::Alignment; use vortex_buffer::ByteBuffer; use super::*; + struct DelayedReadAt; + + impl VortexReadAt for DelayedReadAt { + fn concurrency(&self) -> usize { + 2 + } + + fn size(&self) -> BoxFuture<'static, VortexResult> { + async { Ok(2) }.boxed() + } + + fn read_at( + &self, + offset: u64, + _length: usize, + _alignment: Alignment, + ) -> BoxFuture<'static, VortexResult> { + async move { + if offset == 0 { + tokio::time::sleep(Duration::from_millis(50)).await; + } + Ok(BufferHandle::new_host(ByteBuffer::from(vec![ + u8::try_from(offset).vortex_expect("test offset fits in u8"), + ]))) + } + .boxed() + } + } + #[test] fn test_coalesce_config_in_memory() { let config = CoalesceConfig::in_memory(); @@ -444,14 +467,58 @@ mod tests { ReadAtRequest::new(2, 3, Alignment::none()), ]); - let results = data.read_ranges(requests).await?; - let expected: [&[u8]; 3] = [&[5, 6], &[1], &[3, 4, 5]]; - for (result, expected) in results.into_iter().zip(expected) { - assert_eq!(result.to_host().await.as_ref(), expected); + let results = data.read_ranges(requests).collect::>().await; + for (request, result) in results { + let expected: &[u8] = match request.offset { + 0 => &[1], + 2 => &[3, 4, 5], + 4 => &[5, 6], + offset => panic!("unexpected offset: {offset}"), + }; + assert_eq!(result?.to_host().await.as_ref(), expected); } Ok(()) } + #[tokio::test] + async fn test_read_ranges_keeps_streaming_after_an_error() -> VortexResult<()> { + let data = ByteBuffer::from(vec![1, 2, 3]); + let requests = Arc::from([ + ReadAtRequest::new(100, 1, Alignment::none()), + ReadAtRequest::new(1, 2, Alignment::none()), + ]); + + let results = data.read_ranges(requests).collect::>().await; + + assert_eq!(results.len(), 2); + assert!(results.iter().any(|(_, result)| result.is_err())); + let (_, valid) = results + .into_iter() + .find(|(request, _)| request.offset == 1) + .vortex_expect("valid request result is present"); + assert_eq!(valid?.to_host().await.as_ref(), &[2, 3]); + Ok(()) + } + + #[tokio::test] + async fn test_read_ranges_yields_in_completion_order() -> VortexResult<()> { + let requests = Arc::from([ + ReadAtRequest::new(0, 1, Alignment::none()), + ReadAtRequest::new(1, 1, Alignment::none()), + ]); + let mut results = DelayedReadAt.read_ranges(requests); + + let (first_request, first_result) = results + .next() + .await + .vortex_expect("first result is present"); + + assert_eq!(first_request.offset, 1); + assert_eq!(first_result?.to_host().await.as_ref(), &[1]); + assert_eq!(results.count().await, 1); + Ok(()) + } + #[tokio::test] async fn test_byte_buffer_read_out_of_bounds() { let data = ByteBuffer::from(vec![1, 2, 3]); From c0204c41d846d4128d5e2d1de3f60a2f5332aa69 Mon Sep 17 00:00:00 2001 From: Joseph Isaacs Date: Wed, 12 Aug 2026 17:39:12 +0000 Subject: [PATCH 3/8] Monitor positional read batch sizes Signed-off-by: Joseph Isaacs --- vortex-file/src/read/driver.rs | 31 ++++++++++++++++++++++++ vortex-file/src/segments/source.rs | 38 +++++++++++++++++++++++++++--- 2 files changed, 66 insertions(+), 3 deletions(-) diff --git a/vortex-file/src/read/driver.rs b/vortex-file/src/read/driver.rs index c393901b724..74c8e614164 100644 --- a/vortex-file/src/read/driver.rs +++ b/vortex-file/src/read/driver.rs @@ -337,10 +337,12 @@ impl State { #[cfg(test)] mod tests { use futures::StreamExt; + use futures::channel::mpsc; use futures::stream; use vortex_array::buffer::BufferHandle; use vortex_buffer::Alignment; use vortex_error::VortexResult; + use vortex_error::vortex_panic; use vortex_metrics::DefaultMetricsRegistry; use vortex_metrics::MetricValue; use vortex_metrics::MetricsRegistry; @@ -457,6 +459,35 @@ mod tests { ); } + #[test] + fn test_partial_batch_emits_without_waiting_for_more_events() { + let (sender, receiver) = mpsc::unbounded(); + let (request, _recv) = create_request(1, 0, 10); + assert!(sender.unbounded_send(ReadEvent::Request(request)).is_ok()); + assert!(sender.unbounded_send(ReadEvent::Polled(1)).is_ok()); + + let metrics_registry = DefaultMetricsRegistry::default(); + let metrics = RequestMetrics::new(&metrics_registry, vec![]); + let mut batches = Box::pin(IoRequestStream::new( + receiver, + None, + Alignment::none(), + 32, + metrics, + )); + + // Keep `sender` alive: the input is pending, not finished, and the partial batch must still + // be returned by the current poll rather than waiting for 31 more requests. + let waker = futures::task::noop_waker(); + let mut context = Context::from_waker(&waker); + let Poll::Ready(Some(batch)) = batches.as_mut().poll_next(&mut context) else { + vortex_panic!("partial batch was not emitted by the current poll"); + }; + assert_eq!(batch.len(), 1); + assert_eq!(batch[0].offset(), 0); + drop(sender); + } + #[tokio::test] async fn test_coalesce_adjacent() { let (req1, _rx1) = create_request(1, 0, 10); diff --git a/vortex-file/src/segments/source.rs b/vortex-file/src/segments/source.rs index 44cd37dbad2..9819dde046a 100644 --- a/vortex-file/src/segments/source.rs +++ b/vortex-file/src/segments/source.rs @@ -136,7 +136,7 @@ impl FileSegmentSource { coalesce_config, max_alignment, concurrency, - metrics, + metrics.clone(), ) .boxed(); @@ -144,7 +144,18 @@ impl FileSegmentSource { stream .for_each(move |reqs| { let reader = reader.clone(); + let metrics = metrics.clone(); async move { + metrics.read_ranges_calls.add(1); + metrics.read_ranges_num_ranges.update(reqs.len() as f64); + if reqs.len() > 1 { + metrics.read_ranges_multi.add(1); + } + tracing::trace!( + target: "vortex_file::read_ranges", + num_ranges = reqs.len(), + "submitting positional read batch" + ); let requests = reqs .iter() .map(|req| { @@ -336,6 +347,7 @@ impl Drop for ReadFuture { } /// Metrics emitted by the file segment request driver. +#[derive(Clone)] pub struct RequestMetrics { /// Number of individual segment requests observed by the driver. pub individual_requests: Counter, @@ -343,6 +355,12 @@ pub struct RequestMetrics { pub coalesced_requests: Counter, /// Distribution of how many segment requests were merged into each physical read. pub num_requests_coalesced: Histogram, + /// Number of calls made to [`VortexReadAt::read_ranges`](vortex_io::VortexReadAt::read_ranges). + pub read_ranges_calls: Counter, + /// Number of `read_ranges` calls containing more than one physical range. + pub read_ranges_multi: Counter, + /// Distribution of physical range counts submitted per `read_ranges` call. + pub read_ranges_num_ranges: Histogram, } impl RequestMetrics { @@ -356,8 +374,17 @@ impl RequestMetrics { .add_labels(labels.clone()) .counter("io.requests.coalesced"), num_requests_coalesced: MetricBuilder::new(metrics_registry) - .add_labels(labels) + .add_labels(labels.clone()) .histogram("io.requests.coalesced.num_coalesced"), + read_ranges_calls: MetricBuilder::new(metrics_registry) + .add_labels(labels.clone()) + .counter("io.read_ranges.calls"), + read_ranges_multi: MetricBuilder::new(metrics_registry) + .add_labels(labels.clone()) + .counter("io.read_ranges.multi_range_calls"), + read_ranges_num_ranges: MetricBuilder::new(metrics_registry) + .add_labels(labels) + .histogram("io.read_ranges.num_ranges"), } } } @@ -588,13 +615,14 @@ mod tests { }) .collect(); let metrics = DefaultMetricsRegistry::default(); + let request_metrics = RequestMetrics::new(&metrics, vec![]); let source = FileSegmentSource::open( segments, ReadRangesOnly { calls: Arc::clone(&calls), }, TokioRuntime::current(), - RequestMetrics::new(&metrics, vec![]), + request_metrics.clone(), ); let results = future::join_all((0..4).map(|i| source.request(SegmentId::from(i)))).await; @@ -603,6 +631,10 @@ mod tests { assert_eq!(result?.len(), 4); } assert_eq!(calls.load(Ordering::Relaxed), 1); + assert_eq!(request_metrics.read_ranges_calls.value(), 1); + assert_eq!(request_metrics.read_ranges_multi.value(), 1); + assert_eq!(request_metrics.read_ranges_num_ranges.count(), 1); + assert_eq!(request_metrics.read_ranges_num_ranges.total(), 4.0); Ok(()) } From 2bd880e4fa4f37484b53cde9ce68b417078a991e Mon Sep 17 00:00:00 2001 From: Joseph Isaacs Date: Wed, 12 Aug 2026 21:28:52 +0000 Subject: [PATCH 4/8] Keep positional read concurrency saturated Signed-off-by: Joseph Isaacs --- Cargo.lock | 1 + vortex-file/Cargo.toml | 1 + vortex-file/src/read/mod.rs | 1 + vortex-file/src/segments/source.rs | 292 +++++++++++++++++++++++------ 4 files changed, 238 insertions(+), 57 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index cd2662d2372..4c0c45a8b37 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10117,6 +10117,7 @@ dependencies = [ name = "vortex-file" version = "0.1.0" dependencies = [ + "async-stream", "async-trait", "bytes", "codspeed-divan-compat", diff --git a/vortex-file/Cargo.toml b/vortex-file/Cargo.toml index fc406f7133e..72a47d88f50 100644 --- a/vortex-file/Cargo.toml +++ b/vortex-file/Cargo.toml @@ -17,6 +17,7 @@ version = { workspace = true } all-features = true [dependencies] +async-stream = { workspace = true } async-trait = { workspace = true } bytes = { workspace = true } flatbuffers = { workspace = true } diff --git a/vortex-file/src/read/mod.rs b/vortex-file/src/read/mod.rs index a812b81f63b..f1b18e9a5b1 100644 --- a/vortex-file/src/read/mod.rs +++ b/vortex-file/src/read/mod.rs @@ -5,5 +5,6 @@ mod driver; mod request; pub(crate) use driver::IoRequestStream; +pub(crate) use request::IoRequest; pub(crate) use request::ReadRequest; pub(crate) use request::RequestId; diff --git a/vortex-file/src/segments/source.rs b/vortex-file/src/segments/source.rs index 9819dde046a..93793025e6b 100644 --- a/vortex-file/src/segments/source.rs +++ b/vortex-file/src/segments/source.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use std::any::Any; +use std::collections::VecDeque; use std::future::Future; use std::pin::Pin; use std::sync::Arc; @@ -16,6 +17,7 @@ use futures::channel::mpsc; use futures::future; use futures::future::BoxFuture; use futures::future::Shared; +use futures::stream::SelectAll; use parking_lot::Mutex; use vortex_array::buffer::BufferHandle; use vortex_buffer::Alignment; @@ -38,6 +40,7 @@ use vortex_metrics::MetricBuilder; use vortex_metrics::MetricsRegistry; use crate::SegmentSpec; +use crate::read::IoRequest; use crate::read::IoRequestStream; use crate::read::ReadRequest; use crate::read::RequestId; @@ -86,6 +89,23 @@ type SharedDriver = Shared>; /// observe completion takes the payload and re-raises it; later readers report a graceful error. type DriverPanic = Arc>>>; +fn validate_read_result( + request: &IoRequest, + result: VortexResult, +) -> VortexResult { + result.and_then(|buffer| { + if request.len() != buffer.len() { + return Err(vortex_err!( + "FileSegmentSource: expected buffer of length {} but received {}. {:?}", + request.len(), + buffer.len(), + request + )); + } + Ok(buffer) + }) +} + pub struct FileSegmentSource { segments: Arc<[SegmentSpec]>, /// A queue for sending read request events to the I/O stream. @@ -141,67 +161,111 @@ impl FileSegmentSource { .boxed(); let drive_fut = async move { - stream - .for_each(move |reqs| { - let reader = reader.clone(); - let metrics = metrics.clone(); - async move { - metrics.read_ranges_calls.add(1); - metrics.read_ranges_num_ranges.update(reqs.len() as f64); - if reqs.len() > 1 { - metrics.read_ranges_multi.add(1); + let mut batches = stream.fuse(); + let mut pending = VecDeque::::new(); + let mut reads = SelectAll::new(); + let mut num_active = 0usize; + let mut batches_done = false; + + loop { + if !batches_done { + loop { + match batches.next().now_or_never() { + Some(Some(batch)) => pending.extend(batch), + Some(None) => { + batches_done = true; + break; + } + None => break, } - tracing::trace!( - target: "vortex_file::read_ranges", - num_ranges = reqs.len(), - "submitting positional read batch" - ); - let requests = reqs - .iter() - .map(|req| { - ReadAtRequest::new(req.offset(), req.len(), req.alignment()) - }) - .collect::>() - .into(); - let mut remaining = reqs.into_iter().map(Some).collect::>(); - let mut results = reader.read_ranges(requests); - while let Some((request, result)) = results.next().await { - let Some(position) = remaining.iter().position(|req| { - req.as_ref().is_some_and(|req| { - req.offset() == request.offset - && req.len() == request.length - && req.alignment() == request.alignment - }) - }) else { - tracing::warn!(?request, "reader returned an unknown range"); - continue; - }; - let req = remaining[position] - .take() - .vortex_expect("matched request is present"); - let result = result.and_then(|buffer| { - if req.len() != buffer.len() { - return Err(vortex_err!( - "FileSegmentSource: expected buffer of length {} but received {}. {:?}", - req.len(), - buffer.len(), - req - )); - } - Ok(buffer) - }); - req.resolve(result); + } + } + + while num_active < concurrency && !pending.is_empty() { + let batch_len = (concurrency - num_active).min(pending.len()); + let reqs = pending.drain(..batch_len).collect::>(); + num_active += batch_len; + + metrics.read_ranges_calls.add(1); + metrics.read_ranges_num_ranges.update(batch_len as f64); + if batch_len > 1 { + metrics.read_ranges_multi.add(1); + } + tracing::trace!( + target: "vortex_file::read_ranges", + num_ranges = batch_len, + num_active, + "submitting positional read batch" + ); + + let requests = reqs + .iter() + .map(|req| ReadAtRequest::new(req.offset(), req.len(), req.alignment())) + .collect::>() + .into(); + let mut remaining = reqs.into_iter().map(Some).collect::>(); + let mut results = reader.read_ranges(requests); + reads.push( + async_stream::stream! { + while let Some((request, result)) = results.next().await { + let Some(position) = remaining.iter().position(|req| { + req.as_ref().is_some_and(|req| { + req.offset() == request.offset + && req.len() == request.length + && req.alignment() == request.alignment + }) + }) else { + tracing::warn!(?request, "reader returned an unknown range"); + continue; + }; + let req = remaining[position] + .take() + .vortex_expect("matched request is present"); + yield (req, result); + } + for req in remaining.into_iter().flatten() { + let error = vortex_err!( + "FileSegmentSource: read_ranges ended before resolving request. {:?}", + req + ); + yield (req, Err(error)); + } } - for req in remaining.into_iter().flatten() { - let error = vortex_err!( - "FileSegmentSource: read_ranges ended before resolving request. {:?}", - req - ); - req.resolve(Err(error)); + .boxed(), + ); + } + + if batches_done && num_active == 0 { + break; + } + if num_active == 0 { + match batches.next().await { + Some(batch) => pending.extend(batch), + None => batches_done = true, + } + continue; + } + + let next_read = reads.next(); + let next = if batches_done { + future::Either::Left((next_read.await, batches.next())) + } else { + future::select(next_read, batches.next()).await + }; + match next { + future::Either::Left((result, _)) => { + if let Some((req, result)) = result { + num_active -= 1; + let result = validate_read_result(&req, result); + req.resolve(result); } } - }) - .await + future::Either::Right((batch, _)) => match batch { + Some(batch) => pending.extend(batch), + None => batches_done = true, + }, + } + } }; // Spawn the driver so the runtime makes I/O progress independently of any reader. Readers @@ -638,6 +702,120 @@ mod tests { Ok(()) } + #[derive(Clone)] + struct ControlledReadRanges { + active: Arc, + max_active: Arc, + batch_sizes: Arc>>, + permits: Arc, + } + + impl VortexReadAt for ControlledReadRanges { + fn concurrency(&self) -> usize { + 4 + } + + fn size(&self) -> BoxFuture<'static, VortexResult> { + async { Ok(24) }.boxed() + } + + fn read_at( + &self, + _offset: u64, + _length: usize, + _alignment: Alignment, + ) -> BoxFuture<'static, VortexResult> { + async { panic!("read_at should not be called") }.boxed() + } + + fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> vortex_io::ReadAtStream { + self.batch_sizes.lock().push(requests.len()); + let active = self.active.fetch_add(requests.len(), Ordering::SeqCst) + requests.len(); + self.max_active.fetch_max(active, Ordering::SeqCst); + + let reads = requests + .iter() + .copied() + .map(|request| { + let active = Arc::clone(&self.active); + let permits = Arc::clone(&self.permits); + async move { + let Ok(permit) = permits.acquire_owned().await else { + vortex_panic!("test semaphore unexpectedly closed"); + }; + permit.forget(); + active.fetch_sub(1, Ordering::SeqCst); + let buffer = BufferHandle::new_host( + ByteBuffer::from(vec![0; request.length]).aligned(request.alignment), + ); + (request, Ok(buffer)) + } + }) + .collect::>(); + futures::stream::iter(reads).buffer_unordered(4).boxed() + } + } + + #[tokio::test] + async fn read_driver_refills_global_concurrency_across_batches() -> VortexResult<()> { + let active = Arc::new(AtomicUsize::new(0)); + let max_active = Arc::new(AtomicUsize::new(0)); + let batch_sizes = Arc::new(Mutex::new(Vec::new())); + let permits = Arc::new(tokio::sync::Semaphore::new(0)); + let segments: Arc<[SegmentSpec]> = (0..6) + .map(|i| SegmentSpec { + offset: i * 4, + length: 4, + alignment: Alignment::none(), + }) + .collect(); + let metrics = DefaultMetricsRegistry::default(); + let source = FileSegmentSource::open( + segments, + ControlledReadRanges { + active: Arc::clone(&active), + max_active: Arc::clone(&max_active), + batch_sizes: Arc::clone(&batch_sizes), + permits: Arc::clone(&permits), + }, + TokioRuntime::current(), + RequestMetrics::new(&metrics, vec![]), + ); + let reads = TokioRuntime::current().spawn(async move { + future::join_all((0..6).map(|i| source.request(SegmentId::from(i)))).await + }); + + assert!( + tokio::time::timeout(std::time::Duration::from_secs(1), async { + while active.load(Ordering::SeqCst) != 4 { + tokio::task::yield_now().await; + } + }) + .await + .is_ok() + ); + + permits.add_permits(1); + assert!( + tokio::time::timeout(std::time::Duration::from_secs(1), async { + while batch_sizes.lock().len() < 2 || active.load(Ordering::SeqCst) != 4 { + tokio::task::yield_now().await; + } + }) + .await + .is_ok() + ); + assert_eq!(batch_sizes.lock().as_slice(), [4, 1]); + assert_eq!(max_active.load(Ordering::SeqCst), 4); + + permits.add_permits(5); + for result in reads.await { + assert_eq!(result?.len(), 4); + } + assert_eq!(max_active.load(Ordering::SeqCst), 4); + Ok(()) + } + #[derive(Clone)] struct SlowErrReadAt; From aa3e14e4134fec52624375e2e9881b5f306edc1d Mon Sep 17 00:00:00 2001 From: Joseph Isaacs Date: Thu, 13 Aug 2026 10:25:01 +0000 Subject: [PATCH 5/8] Optimize batched object store reads Signed-off-by: Joe Isaacs --- vortex-file/src/segments/source.rs | 72 +++++++++ vortex-io/src/object_store/read_at.rs | 222 +++++++++++++++++++------- 2 files changed, 233 insertions(+), 61 deletions(-) diff --git a/vortex-file/src/segments/source.rs b/vortex-file/src/segments/source.rs index 93793025e6b..c9de86892b2 100644 --- a/vortex-file/src/segments/source.rs +++ b/vortex-file/src/segments/source.rs @@ -816,6 +816,78 @@ mod tests { Ok(()) } + #[tokio::test] + async fn read_driver_keeps_slots_full_while_a_straggler_is_in_flight() -> VortexResult<()> { + let active = Arc::new(AtomicUsize::new(0)); + let max_active = Arc::new(AtomicUsize::new(0)); + let batch_sizes = Arc::new(Mutex::new(Vec::new())); + let permits = Arc::new(tokio::sync::Semaphore::new(0)); + let segments: Arc<[SegmentSpec]> = (0..8) + .map(|i| SegmentSpec { + offset: i * 4, + length: 4, + alignment: Alignment::none(), + }) + .collect(); + let metrics = DefaultMetricsRegistry::default(); + let source = FileSegmentSource::open( + segments, + ControlledReadRanges { + active: Arc::clone(&active), + max_active: Arc::clone(&max_active), + batch_sizes: Arc::clone(&batch_sizes), + permits: Arc::clone(&permits), + }, + TokioRuntime::current(), + RequestMetrics::new(&metrics, vec![]), + ); + let reads = TokioRuntime::current().spawn(async move { + future::join_all((0..8).map(|i| source.request(SegmentId::from(i)))).await + }); + + wait_for_active_reads(&active, 4).await; + + // Complete three reads while leaving one original read blocked as a straggler. Each freed + // slot must be refilled before the next completion; a batch-barrier implementation would + // instead fall from four active reads to one and submit no replacement work. + for expected_calls in 2..=4 { + permits.add_permits(1); + assert!( + tokio::time::timeout(std::time::Duration::from_secs(1), async { + while batch_sizes.lock().len() < expected_calls + || active.load(Ordering::SeqCst) != 4 + { + tokio::task::yield_now().await; + } + }) + .await + .is_ok() + ); + } + + assert_eq!(batch_sizes.lock().as_slice(), [4, 1, 1, 1]); + assert_eq!(active.load(Ordering::SeqCst), 4); + assert_eq!(max_active.load(Ordering::SeqCst), 4); + + permits.add_permits(5); + for result in reads.await { + assert_eq!(result?.len(), 4); + } + Ok(()) + } + + async fn wait_for_active_reads(active: &AtomicUsize, expected: usize) { + assert!( + tokio::time::timeout(std::time::Duration::from_secs(1), async { + while active.load(Ordering::SeqCst) != expected { + tokio::task::yield_now().await; + } + }) + .await + .is_ok() + ); + } + #[derive(Clone)] struct SlowErrReadAt; diff --git a/vortex-io/src/object_store/read_at.rs b/vortex-io/src/object_store/read_at.rs index 086d70c1bcf..462bc498f82 100644 --- a/vortex-io/src/object_store/read_at.rs +++ b/vortex-io/src/object_store/read_at.rs @@ -5,8 +5,11 @@ use std::io; use std::sync::Arc; use futures::FutureExt; +use futures::SinkExt; use futures::StreamExt; +use futures::channel::mpsc; use futures::future::BoxFuture; +use futures::stream; use object_store::GetOptions; use object_store::GetRange; use object_store::GetResultPayload; @@ -22,6 +25,8 @@ use vortex_error::VortexResult; use vortex_error::vortex_ensure; use crate::CoalesceConfig; +use crate::ReadAtRequest; +use crate::ReadAtStream; use crate::VortexReadAt; use crate::runtime::Handle; #[cfg(not(target_arch = "wasm32"))] @@ -79,6 +84,75 @@ impl ObjectStoreReadAt { } } +async fn read_object_store_range( + store: Arc, + path: ObjectPath, + io_handle: Handle, + allocator: HostAllocatorRef, + request: ReadAtRequest, +) -> VortexResult { + let ReadAtRequest { + offset, + length, + alignment, + } = request; + let range = offset..(offset + length as u64); + let mut buffer = allocator.allocate(length, alignment)?; + + let response = store + .get_opts( + &path, + GetOptions { + range: Some(GetRange::Bounded(range.clone())), + ..Default::default() + }, + ) + .await?; + + let buffer = match response.payload { + #[cfg(not(target_arch = "wasm32"))] + GetResultPayload::File(file, _) => io_handle + .spawn_blocking(move || { + read_exact_at(&file, buffer.as_mut_slice(), range.start)?; + Ok::<_, io::Error>(buffer) + }) + .await + .map_err(io::Error::other)?, + #[cfg(target_arch = "wasm32")] + GetResultPayload::File(..) => { + unreachable!("File payload not supported on wasm32") + } + GetResultPayload::Stream(mut byte_stream) => { + let mut written = 0usize; + while let Some(bytes) = byte_stream.next().await { + let bytes = bytes?; + let end = written + bytes.len(); + vortex_ensure!( + end <= length, + "Object store stream returned too many bytes: {} > expected {} (range: {:?})", + end, + length, + range + ); + buffer.as_mut_slice()[written..end].copy_from_slice(&bytes); + written = end; + } + + vortex_ensure!( + written == length, + "Object store stream returned {} bytes but expected {} bytes (range: {:?})", + written, + length, + range + ); + + buffer + } + }; + + Ok(BufferHandle::new_host(buffer.freeze())) +} + impl VortexReadAt for ObjectStoreReadAt { fn uri(&self) -> Option<&Arc> { Some(&self.uri) @@ -115,70 +189,62 @@ impl VortexReadAt for ObjectStoreReadAt { let path = self.path.clone(); let handle = self.handle.clone(); let allocator = Arc::clone(&self.allocator); - let range = offset..(offset + length as u64); + let io_handle = handle.clone(); + handle + .spawn_io(read_object_store_range( + store, + path, + io_handle, + allocator, + ReadAtRequest::new(offset, length, alignment), + )) + .boxed() + } - // Requires to deal with borrowed lifetimes + fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> ReadAtStream { + if requests.is_empty() { + return stream::empty().boxed(); + } + + let store = Arc::clone(&self.store); + let path = self.path.clone(); + let handle = self.handle.clone(); + let allocator = Arc::clone(&self.allocator); + let concurrency = self.concurrency.max(1); + let (mut send, recv) = mpsc::channel(concurrency); let io_handle = handle.clone(); - handle - .spawn_io(async move { - let mut buffer = allocator.allocate(length, alignment)?; - - let response = store - .get_opts( - &path, - GetOptions { - range: Some(GetRange::Bounded(range.clone())), - ..Default::default() - }, - ) - .await?; - - let buffer = match response.payload { - #[cfg(not(target_arch = "wasm32"))] - GetResultPayload::File(file, _) => { - io_handle - .spawn_blocking(move || { - read_exact_at(&file, buffer.as_mut_slice(), range.start)?; - Ok::<_, io::Error>(buffer) - }) - .await - .map_err(io::Error::other)? - } - #[cfg(target_arch = "wasm32")] - GetResultPayload::File(..) => { - unreachable!("File payload not supported on wasm32") - } - GetResultPayload::Stream(mut byte_stream) => { - let mut written = 0usize; - while let Some(bytes) = byte_stream.next().await { - let bytes = bytes?; - let end = written + bytes.len(); - vortex_ensure!( - end <= length, - "Object store stream returned too many bytes: {} > expected {} (range: {:?})", - end, - length, - range - ); - buffer.as_mut_slice()[written..end].copy_from_slice(&bytes); - written = end; - } - - vortex_ensure!( - written == length, - "Object store stream returned {} bytes but expected {} bytes (range: {:?})", - written, - length, - range - ); - - buffer - } - }; - - Ok(BufferHandle::new_host(buffer.freeze())) - }) + // A single runtime task drives all GETs, avoiding one spawn per range. Do not use + // ObjectStore::get_ranges here: it returns one Vec after every range completes, whereas + // VortexReadAt::read_ranges must expose each result as soon as it is ready. + let task = handle.spawn_io(async move { + let reads = requests.iter().copied().map(|request| { + let store = Arc::clone(&store); + let path = path.clone(); + let io_handle = io_handle.clone(); + let allocator = Arc::clone(&allocator); + async move { + let result = + read_object_store_range(store, path, io_handle, allocator, request).await; + (request, result) + } + }); + + let mut reads = stream::iter(reads).buffer_unordered(concurrency); + while let Some(result) = reads.next().await { + if send.send(result).await.is_err() { + break; + } + } + }); + + async_stream::stream! { + let mut recv = recv; + while let Some(result) = recv.next().await { + yield result; + } + task.await; + } .boxed() } } @@ -258,4 +324,38 @@ mod tests { Ok(()) } + + #[tokio::test] + async fn read_ranges_uses_one_io_task() -> anyhow::Result<()> { + let executor = Arc::new(CountingExecutor::default()); + let runtime = Arc::clone(&executor) as Arc; + let handle = Handle::new(Arc::downgrade(&runtime)); + + let store = Arc::new(InMemory::new()) as Arc; + let path = ObjectPath::from("test.bin"); + store.put(&path, PutPayload::from_static(TEST_DATA)).await?; + + let reader = ObjectStoreReadAt::new(store, path, handle); + let requests: Arc<[ReadAtRequest]> = Arc::from([ + ReadAtRequest::new(0, 6, Alignment::new(1)), + ReadAtRequest::new(7, 5, Alignment::new(1)), + ReadAtRequest::new(18, 4, Alignment::new(1)), + ]); + let results = reader.read_ranges(requests).collect::>().await; + + assert_eq!(results.len(), 3); + for (request, result) in results { + let buffer = result?; + let offset = usize::try_from(request.offset)?; + assert_eq!(buffer.len(), request.length); + assert_eq!( + buffer.to_host().await.as_slice(), + &TEST_DATA[offset..offset + request.length] + ); + } + assert_eq!(executor.spawn_io_count.load(Ordering::SeqCst), 1); + assert_eq!(executor.spawn_count.load(Ordering::SeqCst), 0); + + Ok(()) + } } From fc61ebb9bc566157b9604758325060372d5608c1 Mon Sep 17 00:00:00 2001 From: Joseph Isaacs Date: Thu, 13 Aug 2026 10:35:46 +0000 Subject: [PATCH 6/8] Fix redundant rustdoc link Signed-off-by: Joe Isaacs --- vortex-file/src/segments/source.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex-file/src/segments/source.rs b/vortex-file/src/segments/source.rs index c9de86892b2..d686a7570b5 100644 --- a/vortex-file/src/segments/source.rs +++ b/vortex-file/src/segments/source.rs @@ -419,7 +419,7 @@ pub struct RequestMetrics { pub coalesced_requests: Counter, /// Distribution of how many segment requests were merged into each physical read. pub num_requests_coalesced: Histogram, - /// Number of calls made to [`VortexReadAt::read_ranges`](vortex_io::VortexReadAt::read_ranges). + /// Number of calls made to [`VortexReadAt::read_ranges`]. pub read_ranges_calls: Counter, /// Number of `read_ranges` calls containing more than one physical range. pub read_ranges_multi: Counter, From 0a48a0c4f0b6c2f08a716237876606b4d6e8fa23 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 13 Aug 2026 12:24:45 +0000 Subject: [PATCH 7/8] Clarify I/O request stream behavior Signed-off-by: Joe Isaacs --- vortex-file/src/read/driver.rs | 23 +++++++++++------------ 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/vortex-file/src/read/driver.rs b/vortex-file/src/read/driver.rs index 74c8e614164..69d3851d079 100644 --- a/vortex-file/src/read/driver.rs +++ b/vortex-file/src/read/driver.rs @@ -23,20 +23,18 @@ use crate::segments::ReadEvent; use crate::segments::RequestMetrics; pin_project! { - /// A stream that performs coalescing and prioritization of I/O requests. + /// Converts request lifecycle events into batches of physical reads. /// - /// Takes an input stream of [`ReadRequest`]s and buffers all ready requests into local state. - /// When polled for the next request, this stream will choose the next best request based on - /// an ordering of `(has_been_polled, insertion_order)`, skipping any canceled requests, and - /// then coalescing with other nearby requests within the configured `window`. - /// - /// The output contains up to `batch_size` immediately eligible physical requests. A poll never - /// waits to fill a batch. + /// Polled requests become eligible in registration order. An eligible request may absorb nearby + /// registered requests according to `coalesce_window`. Each poll emits every physical read + /// currently available, up to `batch_size`; it never waits for a full batch. pub(crate) struct IoRequestStream { #[pin] events: S, + // True after the event source closes; buffered requests may still remain. inner_done: bool, coalesce_window: Option, + // Maximum physical reads returned by one stream item. batch_size: usize, state: State, } @@ -75,7 +73,8 @@ where fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let mut this = self.project(); - // First, try to drain all immediately available requests from the inner stream + // Apply all events already available before choosing work. This gives coalescing visibility + // into registered neighbors without delaying emission for future events. loop { match this.events.as_mut().poll_next(cx) { Poll::Ready(Some(event)) => { @@ -91,7 +90,7 @@ where } } - // Return up to batch_size requests that are eligible now. Do not wait to fill the batch. + // Emit a partial batch immediately so the downstream driver can fill free I/O slots. let mut batch = Vec::with_capacity(*this.batch_size); while batch.len() < *this.batch_size { let Some(request) = this.state.next(this.coalesce_window.as_ref()) else { @@ -103,12 +102,12 @@ where return Poll::Ready(Some(batch)); } - // If the inner stream is done, and we have no more _polled_ requests, we're done + // Unpolled requests cannot initiate I/O, so a closed source is done once none are eligible. if *this.inner_done && this.state.polled_requests.is_empty() { return Poll::Ready(None); } - // Otherwise, we need more data from the inner stream + // A new poll/drop/register event will wake us. Poll::Pending } } From e8759c77e4f93d0858693f14794755741b2b10c3 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 13 Aug 2026 14:50:34 +0100 Subject: [PATCH 8/8] Refactor read driver as an explicit stream (#9395) Signed-off-by: Joe Isaacs --- vortex-file/src/segments/source.rs | 312 +++++++++++++++++++---------- 1 file changed, 203 insertions(+), 109 deletions(-) diff --git a/vortex-file/src/segments/source.rs b/vortex-file/src/segments/source.rs index d686a7570b5..33aed2fa9d2 100644 --- a/vortex-file/src/segments/source.rs +++ b/vortex-file/src/segments/source.rs @@ -12,11 +12,14 @@ use std::task::Context; use std::task::Poll; use futures::FutureExt; +use futures::Stream; use futures::StreamExt; use futures::channel::mpsc; use futures::future; use futures::future::BoxFuture; use futures::future::Shared; +use futures::stream::BoxStream; +use futures::stream::Fuse; use futures::stream::SelectAll; use parking_lot::Mutex; use vortex_array::buffer::BufferHandle; @@ -27,6 +30,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_error::vortex_panic; use vortex_io::ReadAtRequest; +use vortex_io::ReadAtStream; use vortex_io::VortexReadAt; use vortex_io::runtime::Handle; use vortex_io::runtime::JoinOutcome; @@ -106,6 +110,171 @@ fn validate_read_result( }) } +type IoBatchStream = Fuse>>; + +enum ReadRangeResultsState { + Reading(ReadAtStream), + Missing, +} + +/// Matches streamed range results back to their logical requests. +struct ReadRangeResults { + state: ReadRangeResultsState, + remaining: Vec>, +} + +impl ReadRangeResults { + fn new(results: ReadAtStream, requests: Vec) -> Self { + Self { + state: ReadRangeResultsState::Reading(results), + remaining: requests.into_iter().map(Some).collect(), + } + } +} + +impl Stream for ReadRangeResults { + type Item = (IoRequest, VortexResult); + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + loop { + match &mut self.state { + ReadRangeResultsState::Reading(results) => match results.poll_next_unpin(cx) { + Poll::Ready(Some((request, result))) => { + let Some(position) = self.remaining.iter().position(|req| { + req.as_ref().is_some_and(|req| { + req.offset() == request.offset + && req.len() == request.length + && req.alignment() == request.alignment + }) + }) else { + tracing::warn!(?request, "reader returned an unknown range"); + continue; + }; + let req = self.remaining[position] + .take() + .vortex_expect("matched request is present"); + return Poll::Ready(Some((req, result))); + } + Poll::Ready(None) => self.state = ReadRangeResultsState::Missing, + Poll::Pending => return Poll::Pending, + }, + ReadRangeResultsState::Missing => { + let Some(req) = self.remaining.iter_mut().find_map(Option::take) else { + return Poll::Ready(None); + }; + let error = vortex_err!( + "FileSegmentSource: read_ranges ended before resolving request. {:?}", + req + ); + return Poll::Ready(Some((req, Err(error)))); + } + } + } + } +} + +/// Drives request batches while keeping the reader's concurrency slots occupied. +struct ReadDriver { + reader: Arc, + batches: IoBatchStream, + pending: VecDeque, + reads: SelectAll, + num_active: usize, + batches_done: bool, + concurrency: usize, + metrics: RequestMetrics, +} + +impl ReadDriver { + fn new( + reader: R, + batches: BoxStream<'static, Vec>, + concurrency: usize, + metrics: RequestMetrics, + ) -> Self { + Self { + reader: Arc::new(reader), + batches: batches.fuse(), + pending: VecDeque::new(), + reads: SelectAll::new(), + num_active: 0, + batches_done: false, + concurrency, + metrics, + } + } + + fn submit_pending(&mut self) { + while self.num_active < self.concurrency && !self.pending.is_empty() { + let batch_len = (self.concurrency - self.num_active).min(self.pending.len()); + let reqs = self.pending.drain(..batch_len).collect::>(); + self.num_active += batch_len; + + self.metrics.read_ranges_calls.add(1); + self.metrics.read_ranges_num_ranges.update(batch_len as f64); + if batch_len > 1 { + self.metrics.read_ranges_multi.add(1); + } + tracing::trace!( + target: "vortex_file::read_ranges", + num_ranges = batch_len, + num_active = self.num_active, + "submitting positional read batch" + ); + + let requests = reqs + .iter() + .map(|req| ReadAtRequest::new(req.offset(), req.len(), req.alignment())) + .collect::>() + .into(); + let results = self.reader.read_ranges(requests); + self.reads.push(ReadRangeResults::new(results, reqs)); + } + } +} + +impl Stream for ReadDriver { + type Item = (); + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.as_mut().get_mut(); + + // Observe every batch already available so submission can fill all free slots at once. + if !this.batches_done { + loop { + match this.batches.poll_next_unpin(cx) { + Poll::Ready(Some(batch)) => this.pending.extend(batch), + Poll::Ready(None) => { + this.batches_done = true; + break; + } + Poll::Pending => break, + } + } + } + + this.submit_pending(); + + if this.batches_done && this.num_active == 0 { + return Poll::Ready(None); + } + + match this.reads.poll_next_unpin(cx) { + Poll::Ready(Some((req, result))) => { + this.num_active -= 1; + let result = validate_read_result(&req, result); + req.resolve(result); + Poll::Ready(Some(())) + } + Poll::Ready(None) if this.num_active == 0 => Poll::Pending, + Poll::Ready(None) => { + vortex_panic!("read result streams ended with active requests") + } + Poll::Pending => Poll::Pending, + } + } +} + pub struct FileSegmentSource { segments: Arc<[SegmentSpec]>, /// A queue for sending read request events to the I/O stream. @@ -160,113 +329,7 @@ impl FileSegmentSource { ) .boxed(); - let drive_fut = async move { - let mut batches = stream.fuse(); - let mut pending = VecDeque::::new(); - let mut reads = SelectAll::new(); - let mut num_active = 0usize; - let mut batches_done = false; - - loop { - if !batches_done { - loop { - match batches.next().now_or_never() { - Some(Some(batch)) => pending.extend(batch), - Some(None) => { - batches_done = true; - break; - } - None => break, - } - } - } - - while num_active < concurrency && !pending.is_empty() { - let batch_len = (concurrency - num_active).min(pending.len()); - let reqs = pending.drain(..batch_len).collect::>(); - num_active += batch_len; - - metrics.read_ranges_calls.add(1); - metrics.read_ranges_num_ranges.update(batch_len as f64); - if batch_len > 1 { - metrics.read_ranges_multi.add(1); - } - tracing::trace!( - target: "vortex_file::read_ranges", - num_ranges = batch_len, - num_active, - "submitting positional read batch" - ); - - let requests = reqs - .iter() - .map(|req| ReadAtRequest::new(req.offset(), req.len(), req.alignment())) - .collect::>() - .into(); - let mut remaining = reqs.into_iter().map(Some).collect::>(); - let mut results = reader.read_ranges(requests); - reads.push( - async_stream::stream! { - while let Some((request, result)) = results.next().await { - let Some(position) = remaining.iter().position(|req| { - req.as_ref().is_some_and(|req| { - req.offset() == request.offset - && req.len() == request.length - && req.alignment() == request.alignment - }) - }) else { - tracing::warn!(?request, "reader returned an unknown range"); - continue; - }; - let req = remaining[position] - .take() - .vortex_expect("matched request is present"); - yield (req, result); - } - for req in remaining.into_iter().flatten() { - let error = vortex_err!( - "FileSegmentSource: read_ranges ended before resolving request. {:?}", - req - ); - yield (req, Err(error)); - } - } - .boxed(), - ); - } - - if batches_done && num_active == 0 { - break; - } - if num_active == 0 { - match batches.next().await { - Some(batch) => pending.extend(batch), - None => batches_done = true, - } - continue; - } - - let next_read = reads.next(); - let next = if batches_done { - future::Either::Left((next_read.await, batches.next())) - } else { - future::select(next_read, batches.next()).await - }; - match next { - future::Either::Left((result, _)) => { - if let Some((req, result)) = result { - num_active -= 1; - let result = validate_read_result(&req, result); - req.resolve(result); - } - } - future::Either::Right((batch, _)) => match batch { - Some(batch) => pending.extend(batch), - None => batches_done = true, - }, - } - } - }; + let drive_fut = ReadDriver::new(reader, stream, concurrency, metrics).collect::<()>(); // Spawn the driver so the runtime makes I/O progress independently of any reader. Readers // join it (below) only to surface a panic raised while driving reads. @@ -508,6 +571,37 @@ mod tests { use super::*; + fn io_request(id: RequestId, offset: u64, length: usize) -> IoRequest { + let (callback, _receiver) = oneshot::channel(); + IoRequest::new_single(ReadRequest { + id, + offset, + length, + alignment: Alignment::none(), + callback, + }) + } + + #[tokio::test] + async fn read_range_results_matches_results_and_reports_missing_requests() { + let requests = vec![io_request(0, 0, 4), io_request(1, 4, 4)]; + let unknown = ReadAtRequest::new(8, 4, Alignment::none()); + let returned = ReadAtRequest::new(4, 4, Alignment::none()); + let buffer = BufferHandle::new_host(ByteBuffer::from(vec![0; 4])); + let results = + futures::stream::iter([(unknown, Ok(buffer.clone())), (returned, Ok(buffer))]).boxed(); + + let resolved = ReadRangeResults::new(results, requests) + .collect::>() + .await; + + assert_eq!(resolved.len(), 2); + assert_eq!(resolved[0].0.offset(), 4); + assert!(resolved[0].1.is_ok()); + assert_eq!(resolved[1].0.offset(), 0); + assert!(resolved[1].1.is_err()); + } + #[derive(Clone)] struct PanickingReadAt; @@ -652,7 +746,7 @@ mod tests { async { panic!("read_at should not be called") }.boxed() } - fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> vortex_io::ReadAtStream { + fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> ReadAtStream { self.calls.fetch_add(1, Ordering::Relaxed); let results = requests .iter() @@ -728,7 +822,7 @@ mod tests { async { panic!("read_at should not be called") }.boxed() } - fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> vortex_io::ReadAtStream { + fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> ReadAtStream { self.batch_sizes.lock().push(requests.len()); let active = self.active.fetch_add(requests.len(), Ordering::SeqCst) + requests.len(); self.max_active.fetch_max(active, Ordering::SeqCst);