From 1ea9321fcb24318bbfa983a90d2d0f597af1a1a9 Mon Sep 17 00:00:00 2001 From: Joseph Isaacs Date: Wed, 12 Aug 2026 15:10:42 +0000 Subject: [PATCH 01/17] 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 02/17] 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 03/17] 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 04/17] 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 965f5027c41a907b5fdbd39af6091abaa7b877e7 Mon Sep 17 00:00:00 2001 From: Joseph Isaacs Date: Thu, 13 Aug 2026 10:25:01 +0000 Subject: [PATCH 05/17] Optimize batched object store reads --- 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 d99da1c4c8c34ae366cea7d69ea74d7448979b14 Mon Sep 17 00:00:00 2001 From: Joseph Isaacs Date: Thu, 13 Aug 2026 07:58:30 +0000 Subject: [PATCH 06/17] Tune local read coalescing --- vortex-io/src/read_at.rs | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/vortex-io/src/read_at.rs b/vortex-io/src/read_at.rs index 82190bdd699..e2788b4ee86 100644 --- a/vortex-io/src/read_at.rs +++ b/vortex-io/src/read_at.rs @@ -69,7 +69,10 @@ impl CoalesceConfig { /// Configuration appropriate for local filesystem access. pub const fn file() -> Self { - Self::new(1 << 20, 4 << 20) // 1MB distance, 4MB max + // Local random reads are cheap enough that reading gaps between segments costs more than + // issuing another operation. Adjacent and overlapping requests still coalesce, while the + // 4 MiB cap preserves useful batching for scans. + Self::new(0, 4 << 20) } /// Configuration appropriate for object storage (S3, GCS, etc.). @@ -439,7 +442,7 @@ mod tests { #[test] fn test_coalesce_config_file() { let config = CoalesceConfig::file(); - assert_eq!(config.distance, 1 << 20); // 1MB + assert_eq!(config.distance, 0); assert_eq!(config.max_size, 4 << 20); // 4MB } From 03ac2a8f5227fe75c9c4330a8e5469624b197f67 Mon Sep 17 00:00:00 2001 From: Joseph Isaacs Date: Thu, 13 Aug 2026 07:58:36 +0000 Subject: [PATCH 07/17] Add optional io_uring local reads --- Cargo.lock | 2 + Cargo.toml | 1 + vortex-io/Cargo.toml | 10 + vortex-io/benches/uring_read_at.rs | 671 +++++++++++++++++++++++++++++ vortex-io/src/std_file/mod.rs | 2 + vortex-io/src/std_file/read_at.rs | 13 + vortex-io/src/std_file/uring.rs | 357 +++++++++++++++ 7 files changed, 1056 insertions(+) create mode 100644 vortex-io/benches/uring_read_at.rs create mode 100644 vortex-io/src/std_file/uring.rs diff --git a/Cargo.lock b/Cargo.lock index 4c0c45a8b37..6fddca6f7db 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10229,6 +10229,7 @@ dependencies = [ "custom-labels", "futures", "glob", + "io-uring", "itertools 0.14.0", "kanal", "object_store", @@ -10236,6 +10237,7 @@ dependencies = [ "parking_lot", "pin-project-lite", "rstest", + "rustix", "smol", "tempfile", "tokio", diff --git a/Cargo.toml b/Cargo.toml index 44e154eadc3..0b91b2b10d4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -167,6 +167,7 @@ geoarrow = "0.8.0" geoarrow-cast = "0.8.0" get_dir = "0.5.0" glob = "0.3.2" +io-uring = "0.7.13" goldenfile = "1" half = { version = "2.7.1", features = ["std", "num-traits"] } hashbrown = "0.17.1" diff --git a/vortex-io/Cargo.toml b/vortex-io/Cargo.toml index b3f6448484a..96aa2a80098 100644 --- a/vortex-io/Cargo.toml +++ b/vortex-io/Cargo.toml @@ -44,6 +44,9 @@ vortex-utils = { workspace = true } [target.'cfg(unix)'.dependencies] custom-labels = { workspace = true } +[target.'cfg(target_os = "linux")'.dependencies] +io-uring = { workspace = true } + [target.'cfg(not(target_arch = "wasm32"))'.dependencies] # Smol is our default impl, so we don't want it to be optional, but it cannot be part of wasm smol = { workspace = true } @@ -60,6 +63,13 @@ rstest = { workspace = true } tempfile = { workspace = true } tokio = { workspace = true, features = ["full"] } +[target.'cfg(target_os = "linux")'.dev-dependencies] +rustix = { workspace = true } + +[[bench]] +name = "uring_read_at" +harness = false + [features] object_store = ["dep:object_store", "vortex-error/object_store"] tokio = ["tokio/fs", "tokio/rt-multi-thread"] diff --git a/vortex-io/benches/uring_read_at.rs b/vortex-io/benches/uring_read_at.rs new file mode 100644 index 00000000000..ea191c4e272 --- /dev/null +++ b/vortex-io/benches/uring_read_at.rs @@ -0,0 +1,671 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Fixed-workload positional-read comparison. Run with `--help` for usage. + +#[cfg(not(target_os = "linux"))] +fn main() { + eprintln!("this benchmark is Linux-only"); +} + +#[cfg(target_os = "linux")] +mod bench { + use std::env; + use std::fs::File; + use std::hint::black_box; + use std::io; + use std::os::fd::AsRawFd; + use std::os::unix::fs::FileExt; + use std::path::Path; + use std::path::PathBuf; + use std::sync::Arc; + use std::sync::Barrier; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use std::sync::mpsc::Receiver; + use std::sync::mpsc::SyncSender; + use std::sync::mpsc::TryRecvError; + use std::sync::mpsc::sync_channel; + use std::thread; + use std::time::Duration; + use std::time::Instant; + + use io_uring::IoUring; + use io_uring::opcode; + use io_uring::types; + use parking_lot::Mutex; + use rustix::fs::Advice; + use rustix::fs::fadvise; + use vortex_utils::aliases::hash_map::HashMap; + + pub fn main() -> io::Result<()> { + let config = Config::parse()?; + if config.help { + help(); + return Ok(()); + } + let file = Arc::new(File::open(&config.path)?); + let file_len = file.metadata()?.len(); + let max_len = config.sizes.iter().copied().max().unwrap_or(0); + if max_len == 0 || file_len < max_len as u64 { + return Err(invalid( + "input file is smaller than the largest non-zero read", + )); + } + fadvise(&*file, 0, None, Advice::Random)?; + prepare_cache(&file, file_len, config.cache)?; + + let device_before = config.device.as_deref().map(device_stats).transpose()?; + let started = Instant::now(); + let mut result = run(Arc::clone(&file), file_len, &config)?; + let elapsed = started.elapsed(); + let device_after = config.device.as_deref().map(device_stats).transpose()?; + result.latencies.sort_unstable(); + let seconds = elapsed.as_secs_f64(); + + println!( + "mode={} engine_threads={} clients={} requests={} sizes={} cpu_ns={} cache={}", + config.mode, + config.engine_threads, + config.clients, + config.requests, + config + .sizes + .iter() + .map(usize::to_string) + .collect::>() + .join(","), + config.cpu_ns, + config.cache, + ); + println!( + "elapsed_s={seconds:.6} logical_reads={} kernel_read_ops={} submission_calls={} ops_per_submit={:.2} bytes={} throughput_mib_s={:.2} reads_s={:.0}", + result.latencies.len(), + result.kernel_ops, + result.submissions, + result.kernel_ops as f64 / result.submissions.max(1) as f64, + result.bytes, + result.bytes as f64 / 1_048_576.0 / seconds, + result.latencies.len() as f64 / seconds, + ); + println!( + "latency_us_p50={:.1} latency_us_p95={:.1} latency_us_p99={:.1} latency_us_max={:.1} checksum={}", + percentile(&result.latencies, 50).as_secs_f64() * 1e6, + percentile(&result.latencies, 95).as_secs_f64() * 1e6, + percentile(&result.latencies, 99).as_secs_f64() * 1e6, + result + .latencies + .last() + .copied() + .unwrap_or_default() + .as_secs_f64() + * 1e6, + result.checksum, + ); + if let Some((before, after)) = device_before.zip(device_after) { + println!( + "device_read_ios={} device_read_mib={:.2} device_read_ms={} device_inflight_end={}", + after.read_ios.saturating_sub(before.read_ios), + after.sectors.saturating_sub(before.sectors) as f64 * 512.0 / 1_048_576.0, + after.read_ms.saturating_sub(before.read_ms), + after.inflight, + ); + } + Ok(()) + } + + fn run(file: Arc, file_len: u64, config: &Config) -> io::Result { + let engine = match config.mode { + Mode::Inline => None, + Mode::Pool => Some(Arc::new(Engine::new( + Arc::clone(&file), + EngineKind::Pread, + config.engine_threads, + config.queue_depth, + )?)), + Mode::Uring => Some(Arc::new(Engine::new( + Arc::clone(&file), + EngineKind::Uring, + config.engine_threads, + config.queue_depth, + )?)), + }; + let next = Arc::new(AtomicUsize::new(0)); + let barrier = Arc::new(Barrier::new(config.clients + 1)); + let mut joins = Vec::with_capacity(config.clients); + for client_id in 0..config.clients { + let file = Arc::clone(&file); + let engine = engine.as_ref().map(Arc::clone); + let next = Arc::clone(&next); + let barrier = Arc::clone(&barrier); + let sizes = Arc::clone(&config.sizes); + let request_count = config.requests; + let cpu_ns = config.cpu_ns; + joins.push(thread::spawn(move || -> io::Result { + let mut row = ClientRow::default(); + barrier.wait(); + loop { + let request_id = next.fetch_add(1, Ordering::Relaxed); + if request_id >= request_count { + break; + } + let len = sizes[request_id % sizes.len()]; + let offset = (random_at(request_id as u64) + % ((file_len - len as u64) / 4096 + 1)) + * 4096; + let started = Instant::now(); + let buffer = match &engine { + Some(engine) => engine.read(offset, len, client_id)?, + None => { + let mut buffer = vec![0; len]; + file.read_exact_at(&mut buffer, offset)?; + buffer + } + }; + row.latencies.push(started.elapsed()); + row.bytes += len as u64; + row.checksum = row.checksum.wrapping_add(sample(&buffer)); + busy_cpu(cpu_ns, row.checksum); + } + Ok(row) + })); + } + barrier.wait(); + let mut result = ResultRow::default(); + for join in joins { + let row = join + .join() + .map_err(|_| io::Error::other("client panicked"))??; + result.latencies.extend(row.latencies); + result.bytes += row.bytes; + result.checksum = result.checksum.wrapping_add(row.checksum); + } + result.kernel_ops = match engine { + Some(engine) => { + let stats = engine.shutdown()?; + result.submissions = stats.submissions; + stats.operations + } + None => { + result.submissions = config.requests as u64; + config.requests as u64 + } + }; + Ok(result) + } + + struct Engine { + senders: Vec>, + joins: Mutex>>>, + } + + #[derive(Clone, Copy)] + enum EngineKind { + Pread, + Uring, + } + + impl Engine { + fn new(file: Arc, kind: EngineKind, n: usize, depth: usize) -> io::Result { + if n == 0 { + return Err(invalid("engine-threads must be non-zero")); + } + let mut senders = Vec::with_capacity(n); + let mut joins = Vec::with_capacity(n); + for id in 0..n { + let (tx, rx) = sync_channel(depth); + let file = Arc::clone(&file); + joins.push( + thread::Builder::new() + .name(format!("read-engine-{id}")) + .spawn(move || match kind { + EngineKind::Pread => pread_worker(file, rx), + EngineKind::Uring => uring_worker(file, rx, depth), + })?, + ); + senders.push(tx); + } + Ok(Self { + senders, + joins: Mutex::new(joins), + }) + } + + fn read(&self, offset: u64, len: usize, shard: usize) -> io::Result> { + let (complete, receive) = sync_channel(1); + self.senders[shard % self.senders.len()] + .send(Message::Read(Request { + offset, + buffer: vec![0; len], + filled: 0, + complete, + })) + .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "engine stopped"))?; + receive + .recv() + .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "completion dropped"))? + } + + fn shutdown(&self) -> io::Result { + for sender in &self.senders { + sender + .send(Message::Stop) + .map_err(|_| io::Error::other("engine stopped"))?; + } + let mut stats = WorkerStats::default(); + for join in self.joins.lock().drain(..) { + let worker = join + .join() + .map_err(|_| io::Error::other("engine panicked"))??; + stats.operations += worker.operations; + stats.submissions += worker.submissions; + } + Ok(stats) + } + } + + enum Message { + Read(Request), + Stop, + } + + struct Request { + offset: u64, + buffer: Vec, + filled: usize, + complete: SyncSender>>, + } + + fn pread_worker(file: Arc, rx: Receiver) -> io::Result { + let mut operations = 0; + while let Ok(message) = rx.recv() { + match message { + Message::Read(mut request) => { + let result = file + .read_exact_at(&mut request.buffer, request.offset) + .map(|()| request.buffer); + operations += 1; + drop(request.complete.send(result)); + } + Message::Stop => break, + } + } + Ok(WorkerStats { + operations, + submissions: operations, + }) + } + + fn uring_worker( + file: Arc, + rx: Receiver, + depth: usize, + ) -> io::Result { + let entries = u32::try_from(depth.next_power_of_two()).map_err(io::Error::other)?; + let mut ring: IoUring = IoUring::builder() + .setup_single_issuer() + .setup_defer_taskrun() + .build(entries)?; + let mut pending: HashMap = HashMap::with_capacity(depth); + let mut next_id = 1_u64; + let mut operations = 0; + let mut submissions = 0; + let mut stopping = false; + loop { + let completions = ring + .completion() + .map(|cqe| (cqe.user_data(), cqe.result())) + .collect::>(); + for (user_data, completion_result) in completions { + let Some(mut request) = pending.remove(&user_data) else { + return Err(io::Error::other("unknown completion")); + }; + operations += 1; + match completion_result { + result if result < 0 => drop( + request + .complete + .send(Err(io::Error::from_raw_os_error(-result))), + ), + 0 => drop(request.complete.send(Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "io_uring read reached EOF", + )))), + result => { + request.filled += result as usize; + if request.filled == request.buffer.len() { + drop(request.complete.send(Ok(request.buffer))); + } else { + push(&mut ring, &file, request, &mut pending, &mut next_id)?; + } + } + } + } + if stopping && pending.is_empty() { + break; + } + let mut accepted = 0; + while pending.len() < depth { + let message = if pending.is_empty() && accepted == 0 { + rx.recv() + .map_err(|_| io::Error::other("request queue disconnected"))? + } else { + match rx.try_recv() { + Ok(message) => message, + Err(TryRecvError::Empty) => break, + Err(TryRecvError::Disconnected) => { + stopping = true; + break; + } + } + }; + match message { + Message::Read(request) => { + push(&mut ring, &file, request, &mut pending, &mut next_id)?; + accepted += 1; + } + Message::Stop => { + stopping = true; + break; + } + } + } + if !pending.is_empty() { + ring.submit_and_wait(1)?; + submissions += 1; + } + } + Ok(WorkerStats { + operations, + submissions, + }) + } + + fn push( + ring: &mut IoUring, + file: &File, + request: Request, + pending: &mut HashMap, + next_id: &mut u64, + ) -> io::Result<()> { + let id = *next_id; + *next_id = next_id.wrapping_add(1); + let remaining = request.buffer.len() - request.filled; + let len = u32::try_from(remaining.min(u32::MAX as usize)).map_err(io::Error::other)?; + let pointer = unsafe { request.buffer.as_ptr().add(request.filled).cast_mut() }; + let entry = opcode::Read::new(types::Fd(file.as_raw_fd()), pointer, len) + .offset(request.offset + request.filled as u64) + .build() + .user_data(id); + // SAFETY: `pending` owns the stable allocation until this operation's CQE is reaped. + unsafe { + ring.submission() + .push(&entry) + .map_err(|_| io::Error::new(io::ErrorKind::WouldBlock, "SQ full"))?; + } + pending.insert(id, request); + Ok(()) + } + + fn prepare_cache(file: &File, file_len: u64, mode: Cache) -> io::Result<()> { + match mode { + Cache::Keep => Ok(()), + Cache::Cold => { + file.sync_all()?; + fadvise(file, 0, None, Advice::DontNeed)?; + Ok(()) + } + Cache::Warm => { + let mut buffer = vec![0; 1024 * 1024]; + let mut offset = 0; + while offset < file_len { + let len = usize::try_from((file_len - offset).min(buffer.len() as u64)) + .map_err(io::Error::other)?; + file.read_exact_at(&mut buffer[..len], offset)?; + offset += len as u64; + } + black_box(sample(&buffer)); + Ok(()) + } + } + } + + fn busy_cpu(ns: u64, seed: u64) { + if ns == 0 { + return; + } + let start = Instant::now(); + let duration = Duration::from_nanos(ns); + let mut value = seed; + while start.elapsed() < duration { + for _ in 0..64 { + value = value.wrapping_mul(0x9e37_79b9_7f4a_7c15).rotate_left(17) + ^ 0xe703_7ed1_a0b4_28db; + } + } + black_box(value); + } + + fn sample(buffer: &[u8]) -> u64 { + u64::from(buffer[0]) + ^ (u64::from(buffer[buffer.len() / 2]) << 8) + ^ (u64::from(buffer[buffer.len() - 1]) << 16) + } + + fn percentile(values: &[Duration], p: usize) -> Duration { + values + .get((values.len().saturating_sub(1)) * p / 100) + .copied() + .unwrap_or_default() + } + + #[derive(Default)] + struct ResultRow { + latencies: Vec, + bytes: u64, + checksum: u64, + kernel_ops: u64, + submissions: u64, + } + + #[derive(Default)] + struct WorkerStats { + operations: u64, + submissions: u64, + } + + #[derive(Default)] + struct ClientRow { + latencies: Vec, + bytes: u64, + checksum: u64, + } + + fn random_at(index: u64) -> u64 { + let mut z = index.wrapping_add(0x9e37_79b9_7f4a_7c15); + z = (z ^ z >> 30).wrapping_mul(0xbf58_476d_1ce4_e5b9); + z = (z ^ z >> 27).wrapping_mul(0x94d0_49bb_1331_11eb); + z ^ z >> 31 + } + + #[derive(Clone, Copy)] + enum Mode { + Inline, + Pool, + Uring, + } + impl std::fmt::Display for Mode { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "{}", + match self { + Self::Inline => "pread-inline", + Self::Pool => "pread-pool", + Self::Uring => "uring", + } + ) + } + } + #[derive(Clone, Copy)] + enum Cache { + Keep, + Cold, + Warm, + } + impl std::fmt::Display for Cache { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "{}", + match self { + Self::Keep => "keep", + Self::Cold => "cold", + Self::Warm => "warm", + } + ) + } + } + + struct Config { + path: PathBuf, + mode: Mode, + clients: usize, + engine_threads: usize, + queue_depth: usize, + requests: usize, + sizes: Arc<[usize]>, + cpu_ns: u64, + cache: Cache, + device: Option, + help: bool, + } + + impl Config { + fn parse() -> io::Result { + let mut c = Self { + path: PathBuf::new(), + mode: Mode::Uring, + clients: 32, + engine_threads: 1, + queue_depth: 256, + requests: 10_000, + sizes: Arc::from([64 * 1024]), + cpu_ns: 0, + cache: Cache::Keep, + device: None, + help: false, + }; + let mut args = env::args().skip(1); + while let Some(arg) = args.next() { + let value = args.next(); + match arg.as_str() { + "--help" | "-h" => c.help = true, + "--path" => c.path = value.ok_or_else(|| invalid("missing path"))?.into(), + "--mode" => { + c.mode = match value.as_deref() { + Some("pread-inline") => Mode::Inline, + Some("pread-pool") => Mode::Pool, + Some("uring") => Mode::Uring, + _ => return Err(invalid("bad mode")), + } + } + "--clients" => c.clients = number(value, &arg)?, + "--engine-threads" => c.engine_threads = number(value, &arg)?, + "--queue-depth" => c.queue_depth = number(value, &arg)?, + "--requests" => c.requests = number(value, &arg)?, + "--cpu-ns" => c.cpu_ns = number(value, &arg)?, + "--sizes" => { + c.sizes = value + .ok_or_else(|| invalid("missing sizes"))? + .split(',') + .map(size) + .collect::>>()? + .into() + } + "--cache" => { + c.cache = match value.as_deref() { + Some("keep") => Cache::Keep, + Some("cold") => Cache::Cold, + Some("warm") => Cache::Warm, + _ => return Err(invalid("bad cache")), + } + } + "--device" => c.device = value, + _ => return Err(invalid(format!("unknown argument {arg}"))), + } + } + if !c.help && c.path.as_os_str().is_empty() { + return Err(invalid("--path is required")); + } + if c.clients == 0 + || c.engine_threads == 0 + || c.queue_depth == 0 + || c.requests == 0 + || c.sizes.is_empty() + { + return Err(invalid("counts and sizes must be non-zero")); + } + Ok(c) + } + } + + fn number(value: Option, name: &str) -> io::Result { + value + .ok_or_else(|| invalid(format!("missing {name}")))? + .parse() + .map_err(|_| invalid(format!("bad {name}"))) + } + fn size(value: &str) -> io::Result { + let (n, multiplier) = match value.as_bytes().last() { + Some(b'K' | b'k') => (&value[..value.len() - 1], 1024), + Some(b'M' | b'm') => (&value[..value.len() - 1], 1024 * 1024), + _ => (value, 1), + }; + n.parse::() + .ok() + .and_then(|n| n.checked_mul(multiplier)) + .filter(|n| *n > 0) + .ok_or_else(|| invalid(format!("bad size {value}"))) + } + fn invalid(message: impl Into) -> io::Error { + io::Error::new(io::ErrorKind::InvalidInput, message.into()) + } + fn help() { + println!( + "usage: uring_read_at --path FILE [--mode pread-inline|pread-pool|uring] [--clients N] [--engine-threads N] [--queue-depth N] [--requests N] [--sizes 4K,64K,1M] [--cpu-ns N] [--cache keep|cold|warm] [--device nvme0n1]" + ); + } + + #[derive(Default)] + struct DeviceStats { + read_ios: u64, + sectors: u64, + read_ms: u64, + inflight: u64, + } + fn device_stats(device: &str) -> io::Result { + let values = + std::fs::read_to_string(Path::new("/sys/class/block").join(device).join("stat"))? + .split_whitespace() + .map(|v| v.parse::().map_err(io::Error::other)) + .collect::>>()?; + if values.len() < 9 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "short block stat", + )); + } + Ok(DeviceStats { + read_ios: values[0], + sectors: values[2], + read_ms: values[3], + inflight: values[8], + }) + } +} + +#[cfg(target_os = "linux")] +fn main() -> std::io::Result<()> { + bench::main() +} diff --git a/vortex-io/src/std_file/mod.rs b/vortex-io/src/std_file/mod.rs index c2e4b12bf40..02818381991 100644 --- a/vortex-io/src/std_file/mod.rs +++ b/vortex-io/src/std_file/mod.rs @@ -2,5 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors mod read_at; +#[cfg(target_os = "linux")] +mod uring; pub use read_at::*; diff --git a/vortex-io/src/std_file/read_at.rs b/vortex-io/src/std_file/read_at.rs index 3d59a595f70..3fdc830015f 100644 --- a/vortex-io/src/std_file/read_at.rs +++ b/vortex-io/src/std_file/read_at.rs @@ -125,6 +125,19 @@ impl VortexReadAt for FileReadAt { let handle = self.handle.clone(); let allocator = Arc::clone(&self.allocator); async move { + #[cfg(target_os = "linux")] + if let Some(submission) = super::uring::try_admit(length) { + let buffer = allocator.allocate(length, alignment)?; + if buffer.is_empty() { + return Ok(BufferHandle::new_host(buffer.freeze())); + } + let receive = submission.read_at(Arc::clone(&file), offset, buffer); + let buffer = receive.into_future().await.map_err(|_| { + io::Error::new(io::ErrorKind::BrokenPipe, "io_uring completion dropped") + })??; + return Ok(BufferHandle::new_host(buffer.freeze())); + } + handle .spawn_blocking(move || { let mut buffer = allocator.allocate(length, alignment)?; diff --git a/vortex-io/src/std_file/uring.rs b/vortex-io/src/std_file/uring.rs new file mode 100644 index 00000000000..7673971c18d --- /dev/null +++ b/vortex-io/src/std_file/uring.rs @@ -0,0 +1,357 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Process-wide `io_uring` engine for local positional reads. + +use std::env; +use std::fs::File; +use std::io; +use std::os::fd::AsRawFd; +use std::sync::Arc; +use std::sync::OnceLock; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::sync::mpsc; +use std::sync::mpsc::Receiver; +use std::sync::mpsc::TryRecvError; +use std::thread; + +use io_uring::IoUring; +use io_uring::opcode; +use io_uring::types; +use vortex_array::memory::WritableHostBuffer; +use vortex_utils::aliases::hash_map::HashMap; +use vortex_utils::parallelism::get_available_parallelism; + +const DEFAULT_QUEUE_DEPTH: usize = 256; +const DEFAULT_MIN_READ_SIZE: usize = 1024 * 1024; +const MAX_RINGS: usize = 4; + +type Completion = oneshot::Sender>; + +static ENGINE: OnceLock>> = OnceLock::new(); + +/// Submit a positional read to the shared engine. +/// +/// `None` means that io_uring is disabled or unavailable and the caller should use its portable +/// blocking-I/O path. Setting `VORTEX_IO_URING=1` enables the engine. The ring and queue counts can +/// be overridden for benchmarking with `VORTEX_IO_URING_RINGS` and +/// `VORTEX_IO_URING_QUEUE_DEPTH`; `VORTEX_IO_URING_MAX_IN_FLIGHT` controls when excess requests +/// spill back to the blocking-I/O path and `VORTEX_IO_URING_MIN_READ_SIZE` controls the minimum +/// request size. +pub(super) fn try_admit(length: usize) -> Option { + let engine = ENGINE + .get_or_init(|| match UringEngine::from_environment() { + Ok(engine) => engine.map(Arc::new), + Err(error) => { + tracing::debug!(%error, "io_uring unavailable; using blocking positional reads"); + None + } + }) + .as_ref()?; + let admission = engine.try_admit(length)?; + + Some(Submission { + engine: Arc::clone(engine), + admission, + }) +} + +pub(super) struct Submission { + engine: Arc, + admission: Admission, +} + +impl Submission { + pub(super) fn read_at( + self, + file: Arc, + offset: u64, + buffer: WritableHostBuffer, + ) -> oneshot::Receiver> { + let (complete, receive) = oneshot::channel(); + let request = Request { + file, + offset, + buffer, + filled: 0, + complete, + _admission: self.admission, + }; + let sender_index = + self.engine.next.fetch_add(1, Ordering::Relaxed) % self.engine.senders.len(); + if let Err(error) = self.engine.senders[sender_index].send(request) { + drop(error.0.complete.send(Err(io::Error::new( + io::ErrorKind::BrokenPipe, + "io_uring worker stopped", + )))); + } + receive + } +} + +struct UringEngine { + senders: Vec>, + next: AtomicUsize, + in_flight: Arc, + max_in_flight: usize, + min_read_size: usize, +} + +impl UringEngine { + fn from_environment() -> io::Result> { + if !env::var("VORTEX_IO_URING") + .is_ok_and(|value| matches!(value.as_str(), "1" | "true" | "on" | "yes")) + { + return Ok(None); + } + + let available = get_available_parallelism().unwrap_or(1); + // One owner per four available CPUs retained the batching advantage without making the + // owner thread a page-cache bottleneck. Storage-bound workloads naturally need fewer. + let default_rings = available.div_ceil(4).clamp(1, MAX_RINGS); + let rings = read_env_usize("VORTEX_IO_URING_RINGS", default_rings)?.clamp(1, 64); + let depth = read_env_usize("VORTEX_IO_URING_QUEUE_DEPTH", DEFAULT_QUEUE_DEPTH)? + .clamp(8, 32_768) + .next_power_of_two(); + let max_in_flight = + read_env_usize("VORTEX_IO_URING_MAX_IN_FLIGHT", rings)?.clamp(1, rings * depth); + let min_read_size = read_env_usize("VORTEX_IO_URING_MIN_READ_SIZE", DEFAULT_MIN_READ_SIZE)?; + + let mut senders = Vec::with_capacity(rings); + for id in 0..rings { + let (send, receive) = mpsc::channel(); + let (ready_send, ready_receive) = mpsc::sync_channel(1); + thread::Builder::new() + .name(format!("vortex-io-uring-{id}")) + .spawn(move || match new_ring(depth) { + Ok(ring) => { + drop(ready_send.send(Ok(()))); + worker(ring, receive, depth); + } + Err(error) => { + let startup_error = io::Error::new(error.kind(), error.to_string()); + drop(ready_send.send(Err(startup_error))); + } + })?; + // SINGLE_ISSUER and DEFER_TASKRUN bind the ring to its owner task, so ring setup must + // happen inside the owner thread. This handshake still detects setup failure before + // publishing the engine. + ready_receive.recv().map_err(|_| { + io::Error::new(io::ErrorKind::BrokenPipe, "io_uring worker failed to start") + })??; + senders.push(send); + } + tracing::debug!( + rings, + depth, + max_in_flight, + min_read_size, + "started local-file io_uring engine" + ); + Ok(Some(Self { + senders, + next: AtomicUsize::new(0), + in_flight: Arc::new(AtomicUsize::new(0)), + max_in_flight, + min_read_size, + })) + } + + fn try_admit(&self, length: usize) -> Option { + if length < self.min_read_size { + return None; + } + self.in_flight + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| { + (current < self.max_in_flight).then_some(current + 1) + }) + .ok()?; + Some(Admission(Arc::clone(&self.in_flight))) + } +} + +struct Admission(Arc); + +impl Drop for Admission { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::Relaxed); + } +} + +fn read_env_usize(name: &str, default: usize) -> io::Result { + match env::var(name) { + Ok(value) => value.parse().map_err(|error| { + io::Error::new( + io::ErrorKind::InvalidInput, + format!("invalid {name}={value:?}: {error}"), + ) + }), + Err(env::VarError::NotPresent) => Ok(default), + Err(error) => Err(io::Error::new(io::ErrorKind::InvalidInput, error)), + } +} + +fn new_ring(depth: usize) -> io::Result { + let entries = u32::try_from(depth).map_err(io::Error::other)?; + IoUring::builder() + .setup_single_issuer() + .setup_defer_taskrun() + .build(entries) + .or_else(|_| IoUring::new(entries)) +} + +struct Request { + file: Arc, + offset: u64, + buffer: WritableHostBuffer, + filled: usize, + complete: Completion, + _admission: Admission, +} + +fn worker(ring: IoUring, receive: Receiver, depth: usize) { + let result = run_worker(ring, &receive, depth); + if let Err(error) = result { + tracing::warn!(%error, "local-file io_uring worker stopped"); + } +} + +fn run_worker(ring: IoUring, receive: &Receiver, depth: usize) -> io::Result<()> { + let mut pending: HashMap = HashMap::with_capacity(depth); + // Declared after `pending` so the ring is closed (and the kernel has released all requests) + // before any in-flight buffers are dropped on an error return. + let mut ring = ring; + let mut next_id = 1_u64; + + loop { + let completions = ring + .completion() + .map(|cqe| (cqe.user_data(), cqe.result())) + .collect::>(); + for (id, result) in completions { + let Some(mut request) = pending.remove(&id) else { + return Err(io::Error::other("io_uring returned an unknown completion")); + }; + match result { + result if result < 0 => { + drop( + request + .complete + .send(Err(io::Error::from_raw_os_error(-result))), + ); + } + 0 => { + drop(request.complete.send(Err(io::Error::new( + io::ErrorKind::UnexpectedEof, + "io_uring read reached EOF", + )))); + } + result => { + request.filled += result as usize; + if request.filled == request.buffer.len() { + drop(request.complete.send(Ok(request.buffer))); + } else { + push(&mut ring, request, &mut pending, &mut next_id)?; + } + } + } + } + + let mut accepted = 0; + while pending.len() < depth { + let request = if pending.is_empty() && accepted == 0 { + match receive.recv() { + Ok(request) => request, + Err(_) => return Ok(()), + } + } else { + match receive.try_recv() { + Ok(request) => request, + Err(TryRecvError::Empty) => break, + Err(TryRecvError::Disconnected) => { + if pending.is_empty() { + return Ok(()); + } + break; + } + } + }; + push(&mut ring, request, &mut pending, &mut next_id)?; + accepted += 1; + } + + if !pending.is_empty() { + ring.submit_and_wait(1)?; + } + } +} + +fn push( + ring: &mut IoUring, + mut request: Request, + pending: &mut HashMap, + next_id: &mut u64, +) -> io::Result<()> { + let id = *next_id; + *next_id = next_id.wrapping_add(1); + let remaining = request.buffer.len() - request.filled; + let length = u32::try_from(remaining.min(u32::MAX as usize)).map_err(io::Error::other)?; + let pointer = request.buffer.as_mut_slice()[request.filled..].as_mut_ptr(); + let entry = opcode::Read::new(types::Fd(request.file.as_raw_fd()), pointer, length) + .offset(request.offset + request.filled as u64) + .build() + .user_data(id); + // SAFETY: `pending` retains the file and stable buffer allocation until the CQE is reaped. + unsafe { + ring.submission() + .push(&entry) + .map_err(|_| io::Error::new(io::ErrorKind::WouldBlock, "io_uring SQ is full"))?; + } + pending.insert(id, request); + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::io::Write; + + use vortex_array::memory::DefaultHostAllocator; + use vortex_array::memory::HostAllocator; + use vortex_buffer::Alignment; + + use super::*; + + #[test] + fn reads_into_owned_host_buffer() -> anyhow::Result<()> { + let mut file = tempfile::tempfile()?; + file.write_all(b"abcdefgh")?; + let file = Arc::new(file); + let (send, receive) = mpsc::channel(); + let owner = thread::spawn(move || { + let ring = new_ring(8)?; + run_worker(ring, &receive, 8) + }); + + let buffer = DefaultHostAllocator.allocate(4, Alignment::none())?; + let (complete, completed) = oneshot::channel(); + send.send(Request { + file, + offset: 2, + buffer, + filled: 0, + complete, + _admission: Admission(Arc::new(AtomicUsize::new(1))), + }) + .map_err(|_| anyhow::anyhow!("io_uring request channel closed"))?; + + let buffer = futures::executor::block_on(completed.into_future()) + .map_err(|_| anyhow::anyhow!("io_uring completion channel closed"))??; + assert_eq!(buffer.freeze().as_slice(), b"cdef"); + drop(send); + owner + .join() + .map_err(|_| anyhow::anyhow!("io_uring owner panicked"))??; + Ok(()) + } +} From 40519350b018b021a7aaff07bb64e18a741f7305 Mon Sep 17 00:00:00 2001 From: Joseph Isaacs Date: Thu, 13 Aug 2026 07:58:40 +0000 Subject: [PATCH 08/17] Use direct reads for local scan files --- benchmarks/datafusion-bench/src/main.rs | 12 ++++--- vortex-io/src/object_store/filesystem.rs | 41 +++++++++++++++++++++--- 2 files changed, 45 insertions(+), 8 deletions(-) diff --git a/benchmarks/datafusion-bench/src/main.rs b/benchmarks/datafusion-bench/src/main.rs index 49a9bf92f58..fd1dc33f4a3 100644 --- a/benchmarks/datafusion-bench/src/main.rs +++ b/benchmarks/datafusion-bench/src/main.rs @@ -289,10 +289,14 @@ async fn register_v2_tables( .runtime_env() .object_store(table_url.object_store())?; - let fs: FileSystemRef = Arc::new(ObjectStoreFileSystem::new( - Arc::clone(&store), - SESSION.handle(), - )); + let fs: FileSystemRef = if benchmark_base.scheme() == "file" { + Arc::new(ObjectStoreFileSystem::local(SESSION.handle())) + } else { + Arc::new(ObjectStoreFileSystem::new( + Arc::clone(&store), + SESSION.handle(), + )) + }; let base_prefix = benchmark_base.path().trim_start_matches('/').to_string(); let fs = fs.with_prefix(base_prefix); diff --git a/vortex-io/src/object_store/filesystem.rs b/vortex-io/src/object_store/filesystem.rs index ca68f5f7efa..57b980975fd 100644 --- a/vortex-io/src/object_store/filesystem.rs +++ b/vortex-io/src/object_store/filesystem.rs @@ -5,6 +5,7 @@ use std::fmt::Debug; use std::fmt::Formatter; +use std::path::PathBuf; use std::sync::Arc; use async_trait::async_trait; @@ -21,6 +22,8 @@ use crate::filesystem::FileListing; use crate::filesystem::FileSystem; use crate::object_store::ObjectStoreReadAt; use crate::runtime::Handle; +#[cfg(not(target_arch = "wasm32"))] +use crate::std_file::FileReadAt; /// A [`FileSystem`] backed by an [`ObjectStore`]. // TODO(ngates): we could consider spawning a driver task inside this file system such that we can @@ -28,6 +31,7 @@ use crate::runtime::Handle; pub struct ObjectStoreFileSystem { store: Arc, handle: Handle, + local_root: Option, } impl Debug for ObjectStoreFileSystem { @@ -41,16 +45,21 @@ impl Debug for ObjectStoreFileSystem { impl ObjectStoreFileSystem { /// Create a new filesystem backed by the given object store and runtime handle. pub fn new(store: Arc, handle: Handle) -> Self { - Self { store, handle } + Self { + store, + handle, + local_root: None, + } } /// Create a new filesystem backed by a local file system object store and the given runtime /// handle. pub fn local(handle: Handle) -> Self { - Self::new( - Arc::new(object_store::local::LocalFileSystem::new()), + Self { + store: Arc::new(object_store::local::LocalFileSystem::new()), handle, - ) + local_root: Some(PathBuf::from("/")), + } } } @@ -105,6 +114,13 @@ impl FileSystem for ObjectStoreFileSystem { } async fn open_read(&self, path: &str) -> VortexResult> { + #[cfg(not(target_arch = "wasm32"))] + if let Some(root) = &self.local_root { + return Ok(Arc::new(FileReadAt::open( + root.join(path), + self.handle.clone(), + )?)); + } Ok(Arc::new(ObjectStoreReadAt::new( Arc::clone(&self.store), to_object_path(path), @@ -153,6 +169,23 @@ mod tests { Ok(ObjectStoreFileSystem::new(store, handle)) } + #[cfg(not(target_arch = "wasm32"))] + #[tokio::test] + async fn local_files_use_file_read_settings() -> VortexResult<()> { + let file = tempfile::NamedTempFile::new()?; + let handle = Handle::find().expect("tokio runtime available within #[tokio::test]"); + let reader = ObjectStoreFileSystem::local(handle) + .open_read(file.path().to_string_lossy().as_ref()) + .await?; + + assert_eq!( + reader.coalesce_config().expect("local coalescing").distance, + 0 + ); + assert_eq!(reader.concurrency(), crate::std_file::DEFAULT_CONCURRENCY); + Ok(()) + } + /// Regression test for #6599: globbing an exact path that exists must return that one file. /// `ObjectStore::list` never yields the prefix itself, so this would return nothing if the /// exact-path branch used `list`. From b78ca66fdafbffa1dd9b14acbb1d4ffe87f82864 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 13 Aug 2026 09:57:20 +0000 Subject: [PATCH 09/17] Add partial segment range requests Signed-off-by: Joe Isaacs --- vortex-file/src/read/driver.rs | 92 +++++++++++++++++--- vortex-file/src/read/request.rs | 3 + vortex-file/src/segments/source.rs | 119 ++++++++++++++++++++++---- vortex-io/src/compat/read_at.rs | 4 + vortex-io/src/object_store/read_at.rs | 13 +++ vortex-io/src/read_at.rs | 26 ++++++ vortex-io/src/std_file/read_at.rs | 5 ++ vortex-layout/src/segments/cache.rs | 34 ++++++++ vortex-layout/src/segments/shared.rs | 88 ++++++++++++++++++- vortex-layout/src/segments/source.rs | 53 ++++++++++++ vortex-layout/src/segments/test.rs | 7 ++ 11 files changed, 409 insertions(+), 35 deletions(-) diff --git a/vortex-file/src/read/driver.rs b/vortex-file/src/read/driver.rs index 74c8e614164..1d63b165bb5 100644 --- a/vortex-file/src/read/driver.rs +++ b/vortex-file/src/read/driver.rs @@ -147,14 +147,7 @@ impl State { fn on_event(&mut self, event: ReadEvent) { trace!(?event, "Received ReadEvent"); match event { - ReadEvent::Request(req) => { - if req.callback.is_closed() { - trace!(?req, "ReadRequest dropped before registration"); - return; - } - self.requests_by_offset.insert((req.offset, req.id)); - self.requests.insert(req.id, req); - } + ReadEvent::Request(req) => self.register(req), ReadEvent::Polled(req_id) => { if let Some(req) = self.requests.remove(&req_id) { if req.callback.is_closed() { @@ -178,6 +171,15 @@ impl State { } } + fn register(&mut self, request: ReadRequest) { + if request.callback.is_closed() { + trace!(?request, "ReadRequest dropped before registration"); + return; + } + self.requests_by_offset.insert((request.offset, request.id)); + self.requests.insert(request.id, request); + } + /// Get the next request, if any. fn next(&mut self, coalesce_window: Option<&CoalesceConfig>) -> Option { match coalesce_window { @@ -227,6 +229,10 @@ impl State { let first_req = self.next_uncoalesced()?; let mut requests = vec![first_req]; + let mut coalesce_distance = requests[0] + .coalesce_distance + .unwrap_or(window.distance) + .min(window.distance); let mut current_start = requests[0].offset; let mut current_end = requests[0].offset + requests[0].length as u64; let align = *self.coalesced_buffer_alignment as u64; @@ -242,8 +248,8 @@ impl State { found_new_requests = false; // Find the range we should scan for coalescing in this iteration - let scan_start = current_start.saturating_sub(window.distance); - let scan_end = current_end.saturating_add(window.distance); + let scan_start = current_start.saturating_sub(coalesce_distance); + let scan_end = current_end.saturating_add(coalesce_distance); // Look for requests that can be coalesced with our current range for &(req_offset, req_id) in self @@ -271,8 +277,12 @@ impl State { // Check if this request is within coalescing distance of our current range let req_end = req_offset + req.length as u64; - if (req_offset <= current_end + window.distance && req_end >= current_start) - || (req_end + window.distance >= current_start && req_offset <= current_end) + let request_distance = req + .coalesce_distance + .unwrap_or(window.distance) + .min(coalesce_distance); + if (req_offset <= current_end + request_distance && req_end >= current_start) + || (req_end + request_distance >= current_start && req_offset <= current_end) { // Calculate what the new range would be if we include this request let new_start = current_start.min(req_offset); @@ -287,6 +297,7 @@ impl State { current_start = new_start; current_end = new_end; + coalesce_distance = request_distance; let req = self .polled_requests .remove(&req_id) @@ -362,6 +373,7 @@ mod tests { offset, length, alignment: Alignment::none(), + coalesce_distance: None, callback: tx, }, rx, @@ -522,6 +534,56 @@ mod tests { } } + #[tokio::test] + async fn test_file_profile_coalesces_adjacent_pages() { + const PAGE_SIZE: usize = 64 * 1024; + let (mut req1, _rx1) = create_request(1, 0, PAGE_SIZE); + let (mut req2, _rx2) = create_request(2, PAGE_SIZE as u64, PAGE_SIZE); + req1.coalesce_distance = Some(16 * 1024); + req2.coalesce_distance = Some(16 * 1024); + + let outputs = collect_outputs( + vec![ + ReadEvent::Request(req1), + ReadEvent::Request(req2), + ReadEvent::Polled(1), + ReadEvent::Polled(2), + ], + Some(CoalesceConfig::file()), + ) + .await; + + assert_eq!(outputs.len(), 1); + assert_eq!(outputs[0].range(), 0..(2 * PAGE_SIZE) as u64); + } + + #[tokio::test] + async fn test_file_profile_does_not_cross_unrequested_page() { + const PAGE_SIZE: usize = 64 * 1024; + let (mut req1, _rx1) = create_request(1, 0, PAGE_SIZE); + let (mut req2, _rx2) = create_request(2, (2 * PAGE_SIZE) as u64, PAGE_SIZE); + req1.coalesce_distance = Some(16 * 1024); + req2.coalesce_distance = Some(16 * 1024); + + let outputs = collect_outputs( + vec![ + ReadEvent::Request(req1), + ReadEvent::Request(req2), + ReadEvent::Polled(1), + ReadEvent::Polled(2), + ], + Some(CoalesceConfig::file()), + ) + .await; + + assert_eq!(outputs.len(), 2); + assert_eq!(outputs[0].range(), 0..PAGE_SIZE as u64); + assert_eq!( + outputs[1].range(), + (2 * PAGE_SIZE) as u64..(3 * PAGE_SIZE) as u64 + ); + } + #[tokio::test] async fn test_coalesce_with_gap() { let (req1, _rx1) = create_request(1, 0, 10); @@ -559,6 +621,7 @@ mod tests { offset: 6, length: 5, alignment: Alignment::new(2), + coalesce_distance: None, callback: tx1, }; let req2 = ReadRequest { @@ -566,6 +629,7 @@ mod tests { offset: 12, length: 1, alignment: Alignment::new(4), + coalesce_distance: None, callback: tx2, }; @@ -634,6 +698,7 @@ mod tests { offset: 0, length: 10, alignment: Alignment::none(), + coalesce_distance: None, callback: tx1, }; let req2 = ReadRequest { @@ -641,6 +706,7 @@ mod tests { offset: 100, length: 10, alignment: Alignment::none(), + coalesce_distance: None, callback: tx2, }; @@ -670,6 +736,7 @@ mod tests { offset: 10, length: 4, alignment: Alignment::none(), + coalesce_distance: None, callback: tx1, }; state.on_event(ReadEvent::Request(req1)); @@ -684,6 +751,7 @@ mod tests { offset: 20, length: 8, alignment: Alignment::none(), + coalesce_distance: None, callback: tx2, }; state.on_event(ReadEvent::Request(req2)); diff --git a/vortex-file/src/read/request.rs b/vortex-file/src/read/request.rs index c4bb4bdc975..55aad0cf1b4 100644 --- a/vortex-file/src/read/request.rs +++ b/vortex-file/src/read/request.rs @@ -96,6 +96,8 @@ pub struct ReadRequest { pub(crate) offset: u64, pub(crate) length: usize, pub(crate) alignment: Alignment, + /// Optional per-request cap on the empty gap this request may coalesce across. + pub(crate) coalesce_distance: Option, pub(crate) callback: oneshot::Sender>, } @@ -106,6 +108,7 @@ impl Debug for ReadRequest { .field("offset", &self.offset) .field("length", &self.length) .field("alignment", &self.alignment) + .field("coalesce_distance", &self.coalesce_distance) .field("is_closed", &self.callback.is_closed()) .finish() } diff --git a/vortex-file/src/segments/source.rs b/vortex-file/src/segments/source.rs index c9de86892b2..1f2407757d6 100644 --- a/vortex-file/src/segments/source.rs +++ b/vortex-file/src/segments/source.rs @@ -4,6 +4,7 @@ use std::any::Any; use std::collections::VecDeque; use std::future::Future; +use std::ops::Range; use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::AtomicUsize; @@ -116,6 +117,8 @@ pub struct FileSegmentSource { driver_panic: DriverPanic, /// The next read request ID. next_id: Arc, + /// Preferred size of canonical byte ranges for the underlying source. + preferred_read_size: Option, } impl FileSegmentSource { @@ -130,6 +133,7 @@ impl FileSegmentSource { metrics: RequestMetrics, ) -> Self { let (send, recv) = mpsc::unbounded(); + let preferred_read_size = reader.preferred_read_size(); let max_alignment = segments .iter() @@ -293,43 +297,94 @@ impl FileSegmentSource { driver, driver_panic, next_id: Arc::new(AtomicUsize::new(0)), + preferred_read_size, } } } impl SegmentSource for FileSegmentSource { + fn preferred_read_size(&self) -> Option { + self.preferred_read_size + } + + fn segment_len(&self, id: SegmentId) -> Option { + self.segments + .get(*id as usize) + .map(|spec| u64::from(spec.length)) + } + fn request(&self, id: SegmentId) -> SegmentFuture { - // We eagerly register the read request here assuming the behaviour of [`FileSegmentSource`], where - // coalescing becomes effective prior to the future being polled. - let spec = *match self.segments.get(*id as usize) { + let Some(length) = self.segment_len(id) else { + return future::ready(Err(vortex_err!("Missing segment: {}", id))).boxed(); + }; + self.request_range_with_coalesce_distance(id, 0..length, None) + } + + fn request_range(&self, segment_id: SegmentId, range: Range) -> SegmentFuture { + self.request_range_with_coalesce_distance( + segment_id, + range, + self.preferred_read_size.map(|size| size / 4), + ) + } +} + +impl FileSegmentSource { + fn request_range_with_coalesce_distance( + &self, + segment_id: SegmentId, + range: Range, + coalesce_distance: Option, + ) -> SegmentFuture { + // We eagerly register the read request here assuming the behaviour of + // [`FileSegmentSource`], where coalescing becomes effective prior to polling. + let spec = *match self.segments.get(*segment_id as usize) { Some(spec) => spec, None => { - return future::ready(Err(vortex_err!("Missing segment: {}", id))).boxed(); + return future::ready(Err(vortex_err!("Missing segment: {}", segment_id))).boxed(); } }; + if range.start > range.end || range.end > u64::from(spec.length) { + return future::ready(Err(vortex_err!( + "Segment {} range {}..{} is out of bounds for a {}-byte segment", + segment_id, + range.start, + range.end, + spec.length + ))) + .boxed(); + } + let SegmentSpec { - offset, - length, - alignment, + offset, alignment, .. } = spec; + let Some(offset) = offset.checked_add(range.start) else { + return future::ready(Err(vortex_err!("Segment range offset overflow"))).boxed(); + }; + let Ok(length) = usize::try_from(range.end - range.start) else { + return future::ready(Err(vortex_err!("Segment range length does not fit usize"))) + .boxed(); + }; + let (send, recv) = oneshot::channel(); let id = self.next_id.fetch_add(1, Ordering::Relaxed); let event = ReadEvent::Request(ReadRequest { id, offset, - length: length as usize, + length, alignment, + coalesce_distance, callback: send, }); - // If we fail to submit the event, we create a future that has failed. - if let Err(e) = self.events.unbounded_send(event) { - return future::ready(Err(vortex_err!("Failed to submit read request: {e}"))).boxed(); + if let Err(error) = self.events.unbounded_send(event) { + return future::ready(Err(vortex_err!("Failed to submit read request: {error}"))) + .boxed(); } - let fut = ReadFuture { + ReadFuture { id, recv: recv.into_future(), polled: false, @@ -337,10 +392,8 @@ impl SegmentSource for FileSegmentSource { events: self.events.clone(), driver: self.driver.clone(), driver_panic: Arc::clone(&self.driver_panic), - }; - - // One allocation: we only box the returned SegmentFuture, not the inner ReadFuture. - fut.boxed() + } + .boxed() } } @@ -470,7 +523,20 @@ impl BufferSegmentSource { } impl SegmentSource for BufferSegmentSource { + fn segment_len(&self, id: SegmentId) -> Option { + self.segments + .get(*id as usize) + .map(|spec| u64::from(spec.length)) + } + fn request(&self, id: SegmentId) -> SegmentFuture { + let Some(length) = self.segment_len(id) else { + return future::ready(Err(vortex_err!("Missing segment: {}", id))).boxed(); + }; + self.request_range(id, 0..length) + } + + fn request_range(&self, id: SegmentId, range: Range) -> SegmentFuture { let spec = match self.segments.get(*id as usize) { Some(spec) => spec, None => { @@ -478,8 +544,19 @@ impl SegmentSource for BufferSegmentSource { } }; - let start = spec.offset as usize; - let end = start + spec.length as usize; + if range.start > range.end || range.end > u64::from(spec.length) { + return future::ready(Err(vortex_err!( + "Segment {} range {}..{} out of bounds for segment length {}", + *id, + range.start, + range.end, + spec.length + ))) + .boxed(); + } + + let start = spec.offset as usize + range.start as usize; + let end = spec.offset as usize + range.end as usize; if end > self.buffer.len() { return future::ready(Err(vortex_err!( "Segment {} range {}..{} out of bounds for buffer of length {}", @@ -491,7 +568,11 @@ impl SegmentSource for BufferSegmentSource { .boxed(); } - let slice = self.buffer.slice(start..end).aligned(spec.alignment); + let slice = if range.start == 0 { + self.buffer.slice(start..end).aligned(spec.alignment) + } else { + self.buffer.slice(start..end) + }; future::ready(Ok(BufferHandle::new_host(slice))).boxed() } } diff --git a/vortex-io/src/compat/read_at.rs b/vortex-io/src/compat/read_at.rs index 3d9cc93b1a6..8cc722d24aa 100644 --- a/vortex-io/src/compat/read_at.rs +++ b/vortex-io/src/compat/read_at.rs @@ -27,6 +27,10 @@ impl VortexReadAt for Compat { self.inner().coalesce_config() } + fn preferred_read_size(&self) -> Option { + self.inner().preferred_read_size() + } + fn concurrency(&self) -> usize { self.inner().concurrency() } diff --git a/vortex-io/src/object_store/read_at.rs b/vortex-io/src/object_store/read_at.rs index 462bc498f82..494bc650d44 100644 --- a/vortex-io/src/object_store/read_at.rs +++ b/vortex-io/src/object_store/read_at.rs @@ -25,6 +25,7 @@ use vortex_error::VortexResult; use vortex_error::vortex_ensure; use crate::CoalesceConfig; +use crate::OBJECT_STORAGE_PREFERRED_READ_SIZE; use crate::ReadAtRequest; use crate::ReadAtStream; use crate::VortexReadAt; @@ -44,6 +45,7 @@ pub struct ObjectStoreReadAt { allocator: HostAllocatorRef, concurrency: usize, coalesce_config: Option, + preferred_read_size: Option, } impl ObjectStoreReadAt { @@ -68,6 +70,7 @@ impl ObjectStoreReadAt { allocator, concurrency: DEFAULT_CONCURRENCY, coalesce_config: Some(CoalesceConfig::object_storage()), + preferred_read_size: Some(OBJECT_STORAGE_PREFERRED_READ_SIZE), } } @@ -82,6 +85,12 @@ impl ObjectStoreReadAt { self.coalesce_config = Some(config); self } + + /// Set the preferred size of independently requested byte ranges for this source. + pub fn with_preferred_read_size(mut self, preferred_read_size: u64) -> Self { + self.preferred_read_size = Some(preferred_read_size); + self + } } async fn read_object_store_range( @@ -162,6 +171,10 @@ impl VortexReadAt for ObjectStoreReadAt { self.coalesce_config } + fn preferred_read_size(&self) -> Option { + self.preferred_read_size + } + fn concurrency(&self) -> usize { self.concurrency } diff --git a/vortex-io/src/read_at.rs b/vortex-io/src/read_at.rs index e2788b4ee86..3ce0eb7d91d 100644 --- a/vortex-io/src/read_at.rs +++ b/vortex-io/src/read_at.rs @@ -22,6 +22,12 @@ use vortex_metrics::MetricBuilder; use vortex_metrics::MetricsRegistry; use vortex_metrics::Timer; +/// Preferred read size for local file sources, including SSDs. +pub const FILE_PREFERRED_READ_SIZE: u64 = 64 * 1024; + +/// Preferred read size for object storage sources. +pub const OBJECT_STORAGE_PREFERRED_READ_SIZE: u64 = 1 << 20; + /// Configuration for coalescing nearby I/O requests into single operations. #[derive(Clone, Copy, Debug)] pub struct CoalesceConfig { @@ -96,6 +102,14 @@ pub trait VortexReadAt: Send + Sync + 'static { None } + /// Preferred size of independently requested byte ranges for this source. + /// + /// Layout readers can use this hint when dividing large logical segments into canonical read + /// ranges. Returning `None` asks readers to preserve whole-segment reads. + fn preferred_read_size(&self) -> Option { + None + } + /// Maximum number of concurrent I/O requests for that should be pulled from this source. /// /// This value is used to control how many [`VortexReadAt::read_at`] calls can @@ -151,6 +165,10 @@ impl VortexReadAt for Arc { self.as_ref().coalesce_config() } + fn preferred_read_size(&self) -> Option { + self.as_ref().preferred_read_size() + } + fn concurrency(&self) -> usize { self.as_ref().concurrency() } @@ -182,6 +200,10 @@ impl VortexReadAt for Arc { self.as_ref().coalesce_config() } + fn preferred_read_size(&self) -> Option { + self.as_ref().preferred_read_size() + } + fn concurrency(&self) -> usize { self.as_ref().concurrency() } @@ -347,6 +369,10 @@ impl VortexReadAt for InstrumentedReadAt { self.read.coalesce_config() } + fn preferred_read_size(&self) -> Option { + self.read.preferred_read_size() + } + fn concurrency(&self) -> usize { self.read.concurrency() } diff --git a/vortex-io/src/std_file/read_at.rs b/vortex-io/src/std_file/read_at.rs index 3fdc830015f..28f1042ac17 100644 --- a/vortex-io/src/std_file/read_at.rs +++ b/vortex-io/src/std_file/read_at.rs @@ -23,6 +23,7 @@ use vortex_buffer::Alignment; use vortex_error::VortexResult; use crate::CoalesceConfig; +use crate::FILE_PREFERRED_READ_SIZE; use crate::VortexReadAt; use crate::runtime::Handle; @@ -102,6 +103,10 @@ impl VortexReadAt for FileReadAt { Some(CoalesceConfig::file()) } + fn preferred_read_size(&self) -> Option { + Some(FILE_PREFERRED_READ_SIZE) + } + fn concurrency(&self) -> usize { DEFAULT_CONCURRENCY } diff --git a/vortex-layout/src/segments/cache.rs b/vortex-layout/src/segments/cache.rs index 1f7f5e91f5c..ced1a27cabf 100644 --- a/vortex-layout/src/segments/cache.rs +++ b/vortex-layout/src/segments/cache.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::ops::Range; use std::sync::Arc; use async_trait::async_trait; @@ -146,6 +147,14 @@ impl SegmentCacheSourceAdapter { } impl SegmentSource for SegmentCacheSourceAdapter { + fn preferred_read_size(&self) -> Option { + self.source.preferred_read_size() + } + + fn segment_len(&self, id: SegmentId) -> Option { + self.source.segment_len(id) + } + fn request(&self, id: SegmentId) -> SegmentFuture { let cache = Arc::clone(&self.cache); let delegate = self.source.request(id); @@ -166,4 +175,29 @@ impl SegmentSource for SegmentCacheSourceAdapter { } .boxed() } + + fn request_range(&self, id: SegmentId, range: Range) -> SegmentFuture { + let cache = Arc::clone(&self.cache); + let delegate = self.source.request_range(id, range.clone()); + + async move { + if let Ok(Some(segment)) = cache.get(id).await { + let start = usize::try_from(range.start)?; + let end = usize::try_from(range.end)?; + if start > end || end > segment.len() { + return Err(vortex_error::vortex_err!( + "Segment {} range {}..{} is out of bounds for cached segment length {}", + id, + range.start, + range.end, + segment.len() + )); + } + tracing::debug!("Resolved segment {} range {:?} from cache", id, range); + return Ok(BufferHandle::new_host(segment.slice(start..end))); + } + delegate.await + } + .boxed() + } } diff --git a/vortex-layout/src/segments/shared.rs b/vortex-layout/src/segments/shared.rs index c794daf608e..2db73d2098b 100644 --- a/vortex-layout/src/segments/shared.rs +++ b/vortex-layout/src/segments/shared.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::ops::Range; use std::sync::Arc; use futures::FutureExt; @@ -22,25 +23,60 @@ use crate::segments::SegmentSource; /// request. pub struct SharedSegmentSource { inner: S, - in_flight: DashMap>, + in_flight: Arc>>, +} + +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +enum SegmentRequest { + Whole(SegmentId), + Range(SegmentId, Range), } type SharedSegmentFuture = BoxFuture<'static, SharedVortexResult>; +struct InFlightGuard { + in_flight: Arc>>, + request: SegmentRequest, +} + +impl Drop for InFlightGuard { + fn drop(&mut self) { + self.in_flight.remove(&self.request); + } +} + impl SharedSegmentSource { /// Create a new `SharedSegmentSource` wrapping the provided inner source. pub fn new(inner: S) -> Self { Self { inner, - in_flight: DashMap::default(), + in_flight: Arc::default(), } } } impl SegmentSource for SharedSegmentSource { + fn preferred_read_size(&self) -> Option { + self.inner.preferred_read_size() + } + + fn segment_len(&self, id: SegmentId) -> Option { + self.inner.segment_len(id) + } + fn request(&self, id: SegmentId) -> SegmentFuture { + self.request_shared(SegmentRequest::Whole(id)) + } + + fn request_range(&self, id: SegmentId, range: Range) -> SegmentFuture { + self.request_shared(SegmentRequest::Range(id, range)) + } +} + +impl SharedSegmentSource { + fn request_shared(&self, request: SegmentRequest) -> SegmentFuture { loop { - match self.in_flight.entry(id) { + match self.in_flight.entry(request.clone()) { Entry::Occupied(e) => { if let Some(shared_future) = e.get().upgrade() { return shared_future.map_err(VortexError::from).boxed(); @@ -50,7 +86,22 @@ impl SegmentSource for SharedSegmentSource { } } Entry::Vacant(e) => { - let future = self.inner.request(id).map_err(Arc::new).boxed().shared(); + let inner_future = match &request { + SegmentRequest::Whole(id) => self.inner.request(*id), + SegmentRequest::Range(id, range) => { + self.inner.request_range(*id, range.clone()) + } + }; + let guard = InFlightGuard { + in_flight: Arc::clone(&self.in_flight), + request: request.clone(), + }; + let future = async move { + let _guard = guard; + inner_future.await.map_err(Arc::new) + } + .boxed() + .shared(); e.insert( future .downgrade() @@ -69,6 +120,7 @@ mod tests { use std::sync::atomic::Ordering; use vortex_buffer::ByteBuffer; + use vortex_error::VortexResult; use super::*; use crate::segments::SegmentSink; @@ -80,6 +132,7 @@ mod tests { struct CountingSegmentSource { segments: TestSegments, request_count: Arc, + range_request_count: Arc, } impl SegmentSource for CountingSegmentSource { @@ -87,6 +140,11 @@ mod tests { self.request_count.fetch_add(1, Ordering::SeqCst); self.segments.request(id) } + + fn request_range(&self, id: SegmentId, range: Range) -> SegmentFuture { + self.range_request_count.fetch_add(1, Ordering::SeqCst); + self.segments.request_range(id, range) + } } #[tokio::test] @@ -116,6 +174,7 @@ mod tests { // The inner source should have been called only once assert_eq!(source.request_count.load(Ordering::Relaxed), 1); + assert!(shared_source.in_flight.is_empty()); } #[tokio::test] @@ -139,6 +198,7 @@ mod tests { let _future = shared_source.request(id); // Future is dropped here } + assert!(shared_source.in_flight.is_empty()); // A new request should still work correctly let result = shared_source.request(id).await; @@ -147,4 +207,24 @@ mod tests { // Should have made 2 requests since the first was dropped before completion assert_eq!(source.request_count.load(Ordering::Relaxed), 2); } + + #[tokio::test] + async fn test_shared_source_deduplicates_identical_ranges() -> VortexResult<()> { + let source = CountingSegmentSource::default(); + let data = ByteBuffer::from(vec![1, 2, 3, 4]); + let seq_id = SequenceId::root().downgrade(); + source.segments.write(seq_id, vec![data]).await?; + + let shared_source = SharedSegmentSource::new(source.clone()); + let id = SegmentId::from(0); + let (first, second) = futures::join!( + shared_source.request_range(id, 1..3), + shared_source.request_range(id, 1..3) + ); + assert_eq!(first?.unwrap_host(), ByteBuffer::from(vec![2, 3])); + assert_eq!(second?.unwrap_host(), ByteBuffer::from(vec![2, 3])); + assert_eq!(source.range_request_count.load(Ordering::Relaxed), 1); + assert!(shared_source.in_flight.is_empty()); + Ok(()) + } } diff --git a/vortex-layout/src/segments/source.rs b/vortex-layout/src/segments/source.rs index 5c709f5a7ad..45ab59b5c32 100644 --- a/vortex-layout/src/segments/source.rs +++ b/vortex-layout/src/segments/source.rs @@ -1,9 +1,13 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::ops::Range; + +use futures::FutureExt; use futures::future::BoxFuture; use vortex_array::buffer::BufferHandle; use vortex_error::VortexResult; +use vortex_error::vortex_bail; use crate::segments::SegmentId; /// Static future resolving to a segment byte buffer. @@ -14,6 +18,55 @@ pub type SegmentFuture = BoxFuture<'static, VortexResult>; /// Implementations may issue asynchronous file reads, object-store requests, cache lookups, or /// in-memory buffer slices. Returned futures must be independent and safe to poll concurrently. pub trait SegmentSource: 'static + Send + Sync { + /// Preferred size of independently requested byte ranges for this source. + /// + /// Layout readers can use this hint to divide a logical segment into canonical read ranges. + /// Returning `None` asks readers to preserve whole-segment reads. + fn preferred_read_size(&self) -> Option { + None + } + + /// Return the serialized length of `id`, when it is known without issuing I/O. + fn segment_len(&self, _id: SegmentId) -> Option { + None + } + /// Request a segment, returning a future that will eventually resolve to the segment data. fn request(&self, id: SegmentId) -> SegmentFuture; + + /// Request a byte range relative to the start of a segment. + /// + /// Sources backed by random-access storage should override this method. The default keeps + /// custom sources compatible by reading the segment and slicing it after bounds checking. + fn request_range(&self, id: SegmentId, range: Range) -> SegmentFuture { + let segment = self.request(id); + async move { + let segment = segment.await?; + let start = usize::try_from(range.start)?; + let end = usize::try_from(range.end)?; + if start > end || end > segment.len() { + vortex_bail!( + "Segment {} range {}..{} is out of bounds for a {}-byte segment", + id, + range.start, + range.end, + segment.len() + ); + } + Ok(segment.slice(start..end)) + } + .boxed() + } + + /// Register multiple ranges from one segment together. + /// + /// The returned futures correspond positionally to `ranges`. Sources can override this to + /// amortize registration while retaining independent canonical range futures for sharing and + /// coalescing. + fn request_ranges(&self, id: SegmentId, ranges: Vec>) -> Vec { + ranges + .into_iter() + .map(|range| self.request_range(id, range)) + .collect() + } } diff --git a/vortex-layout/src/segments/test.rs b/vortex-layout/src/segments/test.rs index d880d15cc1a..b6a3d3007e2 100644 --- a/vortex-layout/src/segments/test.rs +++ b/vortex-layout/src/segments/test.rs @@ -26,6 +26,13 @@ pub struct TestSegments { } impl SegmentSource for TestSegments { + fn segment_len(&self, id: SegmentId) -> Option { + self.segments + .lock() + .get(*id as usize) + .and_then(|buffer| u64::try_from(buffer.len()).ok()) + } + fn request(&self, id: SegmentId) -> SegmentFuture { let buffer = self.segments.lock().get(*id as usize).cloned(); async move { From e222ef29a3c00c53b3081d0090bc953764be11c0 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 13 Aug 2026 10:11:22 +0000 Subject: [PATCH 10/17] Reuse blocking workers across positional read batches Signed-off-by: Joe Isaacs --- vortex-io/src/std_file/read_at.rs | 52 +++++++++++++++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/vortex-io/src/std_file/read_at.rs b/vortex-io/src/std_file/read_at.rs index 28f1042ac17..f317618ea38 100644 --- a/vortex-io/src/std_file/read_at.rs +++ b/vortex-io/src/std_file/read_at.rs @@ -13,9 +13,14 @@ use std::os::unix::fs::FileExt; use std::os::windows::fs::FileExt; use std::path::Path; use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use futures::FutureExt; +use futures::StreamExt; +use futures::channel::mpsc; use futures::future::BoxFuture; +use futures::stream; use vortex_array::buffer::BufferHandle; use vortex_array::memory::DefaultHostAllocator; use vortex_array::memory::HostAllocatorRef; @@ -24,6 +29,8 @@ use vortex_error::VortexResult; use crate::CoalesceConfig; use crate::FILE_PREFERRED_READ_SIZE; +use crate::ReadAtRequest; +use crate::ReadAtStream; use crate::VortexReadAt; use crate::runtime::Handle; @@ -153,4 +160,49 @@ impl VortexReadAt for FileReadAt { } .boxed() } + + fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> ReadAtStream { + if requests.is_empty() { + return stream::empty().boxed(); + } + + let worker_count = requests.len().min(DEFAULT_CONCURRENCY); + let next = Arc::new(AtomicUsize::new(0)); + let (send, recv) = mpsc::unbounded(); + let mut workers = Vec::with_capacity(worker_count); + + for _ in 0..worker_count { + let file = Arc::clone(&self.file); + let allocator = Arc::clone(&self.allocator); + let requests = Arc::clone(&requests); + let next = Arc::clone(&next); + let send = send.clone(); + workers.push(self.handle.spawn_blocking(move || { + loop { + let index = next.fetch_add(1, Ordering::Relaxed); + let Some(request) = requests.get(index).copied() else { + break; + }; + let result = (|| -> VortexResult { + let mut buffer = allocator.allocate(request.length, request.alignment)?; + read_exact_at(&file, buffer.as_mut_slice(), request.offset)?; + Ok(BufferHandle::new_host(buffer.freeze())) + })(); + if send.unbounded_send((request, result)).is_err() { + break; + } + } + })); + } + drop(send); + + // Retaining task handles in the stream state aborts workers that have not started their + // next range if the consumer drops the response stream. + stream::unfold((recv, workers), |(mut recv, workers)| async move { + recv.next() + .await + .map(|response| (response, (recv, workers))) + }) + .boxed() + } } From b9462acc40b57cbd25ce2e50d0be589cfd4fc792 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 13 Aug 2026 10:28:11 +0000 Subject: [PATCH 11/17] Submit partial segment ranges as bounded batches Signed-off-by: Joe Isaacs --- vortex-file/src/read/request.rs | 11 ++ vortex-file/src/segments/source.rs | 155 +++++++++++++++++++++++---- vortex-layout/src/segments/shared.rs | 85 +++++++++++++++ 3 files changed, 233 insertions(+), 18 deletions(-) diff --git a/vortex-file/src/read/request.rs b/vortex-file/src/read/request.rs index 55aad0cf1b4..a68de2862ed 100644 --- a/vortex-file/src/read/request.rs +++ b/vortex-file/src/read/request.rs @@ -55,6 +55,17 @@ impl IoRequest { } } + /// Whether this physical request was assembled exclusively from partial segment ranges. + pub(crate) fn is_partial(&self) -> bool { + match &self.0 { + IoRequestInner::Single(request) => request.coalesce_distance.is_some(), + IoRequestInner::Coalesced(request) => request + .requests + .iter() + .all(|request| request.coalesce_distance.is_some()), + } + } + /// Resolves the request with the given result. pub fn resolve(self, result: VortexResult) { match self.0 { diff --git a/vortex-file/src/segments/source.rs b/vortex-file/src/segments/source.rs index 1f2407757d6..26689dc6956 100644 --- a/vortex-file/src/segments/source.rs +++ b/vortex-file/src/segments/source.rs @@ -7,6 +7,7 @@ use std::future::Future; use std::ops::Range; use std::pin::Pin; use std::sync::Arc; +use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Context; @@ -90,6 +91,26 @@ type SharedDriver = Shared>; /// observe completion takes the payload and re-raises it; later readers report a graceful error. type DriverPanic = Arc>>>; +const MAX_PARTIAL_SUBMISSION_REQUESTS: usize = 512; +const MAX_PARTIAL_SUBMISSION_BYTES: usize = 16 << 20; + +fn partial_submission_len(requests: &VecDeque) -> usize { + let mut count = 0usize; + let mut bytes = 0usize; + for request in requests.iter().take(MAX_PARTIAL_SUBMISSION_REQUESTS) { + if !request.is_partial() { + break; + } + let next_bytes = bytes.saturating_add(request.len()); + if count > 0 && next_bytes > MAX_PARTIAL_SUBMISSION_BYTES { + break; + } + count += 1; + bytes = next_bytes; + } + count +} + fn validate_read_result( request: &IoRequest, result: VortexResult, @@ -159,7 +180,7 @@ impl FileSegmentSource { StreamExt::boxed(recv), coalesce_config, max_alignment, - concurrency, + MAX_PARTIAL_SUBMISSION_REQUESTS, metrics.clone(), ) .boxed(); @@ -186,7 +207,12 @@ impl FileSegmentSource { } while num_active < concurrency && !pending.is_empty() { - let batch_len = (concurrency - num_active).min(pending.len()); + let batch_len = + if num_active == 0 && pending.front().is_some_and(IoRequest::is_partial) { + partial_submission_len(&pending) + } else { + (concurrency - num_active).min(pending.len()) + }; let reqs = pending.drain(..batch_len).collect::>(); num_active += batch_len; @@ -327,6 +353,27 @@ impl SegmentSource for FileSegmentSource { self.preferred_read_size.map(|size| size / 4), ) } + + fn request_ranges(&self, segment_id: SegmentId, ranges: Vec>) -> Vec { + let coalesce_distance = self.preferred_read_size.map(|size| size / 4); + let mut registered = ranges + .into_iter() + .map(|range| self.register_range(segment_id, range, coalesce_distance)) + .collect::>(); + let poll_ids: Arc<[usize]> = registered + .iter() + .filter_map(|registration| registration.as_ref().ok().map(|read| read.id)) + .collect(); + let poll_once = Arc::new(AtomicBool::new(false)); + + registered + .drain(..) + .map(|registration| match registration { + Ok(read) => self.read_future(read, Arc::clone(&poll_ids), Arc::clone(&poll_once)), + Err(error) => future::ready(Err(error)).boxed(), + }) + .collect() + } } impl FileSegmentSource { @@ -336,24 +383,36 @@ impl FileSegmentSource { range: Range, coalesce_distance: Option, ) -> SegmentFuture { + match self.register_range(segment_id, range, coalesce_distance) { + Ok(read) => { + let poll_ids = Arc::from([read.id]); + self.read_future(read, poll_ids, Arc::new(AtomicBool::new(false))) + } + Err(error) => future::ready(Err(error)).boxed(), + } + } + + fn register_range( + &self, + segment_id: SegmentId, + range: Range, + coalesce_distance: Option, + ) -> VortexResult { // We eagerly register the read request here assuming the behaviour of // [`FileSegmentSource`], where coalescing becomes effective prior to polling. let spec = *match self.segments.get(*segment_id as usize) { Some(spec) => spec, - None => { - return future::ready(Err(vortex_err!("Missing segment: {}", segment_id))).boxed(); - } + None => return Err(vortex_err!("Missing segment: {}", segment_id)), }; if range.start > range.end || range.end > u64::from(spec.length) { - return future::ready(Err(vortex_err!( + return Err(vortex_err!( "Segment {} range {}..{} is out of bounds for a {}-byte segment", segment_id, range.start, range.end, spec.length - ))) - .boxed(); + )); } let SegmentSpec { @@ -361,11 +420,10 @@ impl FileSegmentSource { } = spec; let Some(offset) = offset.checked_add(range.start) else { - return future::ready(Err(vortex_err!("Segment range offset overflow"))).boxed(); + return Err(vortex_err!("Segment range offset overflow")); }; let Ok(length) = usize::try_from(range.end - range.start) else { - return future::ready(Err(vortex_err!("Segment range length does not fit usize"))) - .boxed(); + return Err(vortex_err!("Segment range length does not fit usize")); }; let (send, recv) = oneshot::channel(); @@ -380,15 +438,28 @@ impl FileSegmentSource { }); if let Err(error) = self.events.unbounded_send(event) { - return future::ready(Err(vortex_err!("Failed to submit read request: {error}"))) - .boxed(); + return Err(vortex_err!("Failed to submit read request: {error}")); } - ReadFuture { + Ok(RegisteredRead { id, recv: recv.into_future(), + }) + } + + fn read_future( + &self, + read: RegisteredRead, + poll_ids: Arc<[usize]>, + poll_once: Arc, + ) -> SegmentFuture { + ReadFuture { + id: read.id, + recv: read.recv, polled: false, finished: false, + poll_ids, + poll_once, events: self.events.clone(), driver: self.driver.clone(), driver_panic: Arc::clone(&self.driver_panic), @@ -397,6 +468,11 @@ impl FileSegmentSource { } } +struct RegisteredRead { + id: usize, + recv: oneshot::AsyncReceiver>, +} + /// A future that resolves a read request from a [`FileSegmentSource`]. /// /// See the documentation for [`FileSegmentSource`] for details on coalescing and pre-fetching. @@ -406,6 +482,8 @@ struct ReadFuture { recv: oneshot::AsyncReceiver>, polled: bool, finished: bool, + poll_ids: Arc<[usize]>, + poll_once: Arc, events: mpsc::UnboundedSender, driver: SharedDriver, driver_panic: DriverPanic, @@ -440,11 +518,16 @@ impl Future for ReadFuture { }, Poll::Pending if !self.polled => { self.polled = true; - // Notify the I/O stream that this request has been polled. - match self.events.unbounded_send(ReadEvent::Polled(self.id)) { - Ok(()) => Poll::Pending, - Err(e) => Poll::Ready(Err(vortex_err!("ReadRequest dropped by runtime: {e}"))), + if !self.poll_once.swap(true, Ordering::AcqRel) { + for &id in self.poll_ids.iter() { + if let Err(error) = self.events.unbounded_send(ReadEvent::Polled(id)) { + return Poll::Ready(Err(vortex_err!( + "ReadRequest dropped by runtime: {error}" + ))); + } + } } + Poll::Pending } _ => Poll::Pending, } @@ -713,6 +796,7 @@ mod tests { #[derive(Clone)] struct ReadRangesOnly { calls: Arc, + max_batch: Arc, } impl VortexReadAt for ReadRangesOnly { @@ -735,6 +819,7 @@ mod tests { fn read_ranges(&self, requests: Arc<[ReadAtRequest]>) -> vortex_io::ReadAtStream { self.calls.fetch_add(1, Ordering::Relaxed); + self.max_batch.fetch_max(requests.len(), Ordering::Relaxed); let results = requests .iter() .copied() @@ -752,6 +837,7 @@ mod tests { #[tokio::test] async fn read_driver_batches_ready_requests() -> VortexResult<()> { let calls = Arc::new(AtomicUsize::new(0)); + let max_batch = Arc::new(AtomicUsize::new(0)); let segments: Arc<[SegmentSpec]> = (0..4) .map(|i| SegmentSpec { offset: i * 4, @@ -765,6 +851,7 @@ mod tests { segments, ReadRangesOnly { calls: Arc::clone(&calls), + max_batch: Arc::clone(&max_batch), }, TokioRuntime::current(), request_metrics.clone(), @@ -776,6 +863,7 @@ mod tests { assert_eq!(result?.len(), 4); } assert_eq!(calls.load(Ordering::Relaxed), 1); + assert_eq!(max_batch.load(Ordering::Relaxed), 4); 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); @@ -783,6 +871,37 @@ mod tests { Ok(()) } + #[tokio::test] + async fn read_driver_submits_partial_ranges_together() -> VortexResult<()> { + let calls = Arc::new(AtomicUsize::new(0)); + let max_batch = Arc::new(AtomicUsize::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, + ReadRangesOnly { + calls: Arc::clone(&calls), + max_batch: Arc::clone(&max_batch), + }, + TokioRuntime::current(), + RequestMetrics::new(&metrics, vec![]), + ); + + let results = source.request_ranges(SegmentId::from(0), vec![0..1, 1..2, 2..3, 3..4]); + for result in future::join_all(results).await { + assert_eq!(result?.len(), 1); + } + assert_eq!(calls.load(Ordering::Relaxed), 1); + assert_eq!(max_batch.load(Ordering::Relaxed), 4); + Ok(()) + } + #[derive(Clone)] struct ControlledReadRanges { active: Arc, diff --git a/vortex-layout/src/segments/shared.rs b/vortex-layout/src/segments/shared.rs index 2db73d2098b..b4d29ab8f42 100644 --- a/vortex-layout/src/segments/shared.rs +++ b/vortex-layout/src/segments/shared.rs @@ -6,12 +6,14 @@ use std::sync::Arc; use futures::FutureExt; use futures::TryFutureExt; +use futures::channel::oneshot; use futures::future::BoxFuture; use futures::future::WeakShared; use vortex_array::buffer::BufferHandle; use vortex_error::SharedVortexResult; use vortex_error::VortexError; use vortex_error::VortexExpect; +use vortex_error::vortex_err; use vortex_utils::aliases::dash_map::DashMap; use vortex_utils::aliases::dash_map::Entry; @@ -71,6 +73,62 @@ impl SegmentSource for SharedSegmentSource { fn request_range(&self, id: SegmentId, range: Range) -> SegmentFuture { self.request_shared(SegmentRequest::Range(id, range)) } + + fn request_ranges(&self, id: SegmentId, ranges: Vec>) -> Vec { + let mut outputs = (0..ranges.len()).map(|_| None).collect::>(); + let mut missing = Vec::new(); + + for (index, range) in ranges.into_iter().enumerate() { + let request = SegmentRequest::Range(id, range.clone()); + loop { + match self.in_flight.entry(request.clone()) { + Entry::Occupied(entry) => { + if let Some(future) = entry.get().upgrade() { + outputs[index] = Some(future.map_err(VortexError::from).boxed()); + break; + } + entry.remove(); + } + Entry::Vacant(entry) => { + let (send, receive) = oneshot::channel::(); + let guard = InFlightGuard { + in_flight: Arc::clone(&self.in_flight), + request: request.clone(), + }; + let future = async move { + let _guard = guard; + let inner = receive.await.map_err(|_| { + Arc::new(vortex_err!("Batched segment request was dropped")) + })?; + inner.await.map_err(Arc::new) + } + .boxed() + .shared(); + entry.insert( + future + .downgrade() + .vortex_expect("new shared future cannot be complete"), + ); + outputs[index] = Some(future.map_err(VortexError::from).boxed()); + missing.push((range, send)); + break; + } + } + } + } + + let inner = self + .inner + .request_ranges(id, missing.iter().map(|(range, _)| range.clone()).collect()); + for ((_, send), future) in missing.into_iter().zip(inner) { + drop(send.send(future)); + } + + outputs + .into_iter() + .map(|future| future.vortex_expect("every requested range has a future")) + .collect() + } } impl SharedSegmentSource { @@ -133,6 +191,7 @@ mod tests { segments: TestSegments, request_count: Arc, range_request_count: Arc, + range_batch_count: Arc, } impl SegmentSource for CountingSegmentSource { @@ -145,6 +204,14 @@ mod tests { self.range_request_count.fetch_add(1, Ordering::SeqCst); self.segments.request_range(id, range) } + + fn request_ranges(&self, id: SegmentId, ranges: Vec>) -> Vec { + self.range_batch_count.fetch_add(1, Ordering::SeqCst); + ranges + .into_iter() + .map(|range| self.request_range(id, range)) + .collect() + } } #[tokio::test] @@ -227,4 +294,22 @@ mod tests { assert!(shared_source.in_flight.is_empty()); Ok(()) } + + #[tokio::test] + async fn test_shared_source_forwards_missing_ranges_as_one_batch() -> VortexResult<()> { + let source = CountingSegmentSource::default(); + let data = ByteBuffer::from(vec![1, 2, 3, 4]); + let seq_id = SequenceId::root().downgrade(); + source.segments.write(seq_id, vec![data]).await?; + + let shared_source = SharedSegmentSource::new(source.clone()); + let reads = shared_source.request_ranges(SegmentId::from(0), vec![0..1, 2..4]); + let mut results = futures::future::join_all(reads).await.into_iter(); + assert_eq!(results.next().vortex_expect("first range")?.len(), 1); + assert_eq!(results.next().vortex_expect("second range")?.len(), 2); + assert_eq!(source.range_batch_count.load(Ordering::Relaxed), 1); + assert_eq!(source.range_request_count.load(Ordering::Relaxed), 2); + assert!(shared_source.in_flight.is_empty()); + Ok(()) + } } From e42dfbf07fb983236be506c5b25665a58c586573 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 13 Aug 2026 10:35:46 +0000 Subject: [PATCH 12/17] Forward grouped reads through the segment cache Signed-off-by: Joe Isaacs --- vortex-layout/src/segments/cache.rs | 30 +++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/vortex-layout/src/segments/cache.rs b/vortex-layout/src/segments/cache.rs index ced1a27cabf..8d24ed3aeb6 100644 --- a/vortex-layout/src/segments/cache.rs +++ b/vortex-layout/src/segments/cache.rs @@ -200,4 +200,34 @@ impl SegmentSource for SegmentCacheSourceAdapter { } .boxed() } + + fn request_ranges(&self, id: SegmentId, ranges: Vec>) -> Vec { + let delegates = self.source.request_ranges(id, ranges.clone()); + ranges + .into_iter() + .zip(delegates) + .map(|(range, delegate)| { + let cache = Arc::clone(&self.cache); + async move { + if let Ok(Some(segment)) = cache.get(id).await { + let start = usize::try_from(range.start)?; + let end = usize::try_from(range.end)?; + if start > end || end > segment.len() { + return Err(vortex_error::vortex_err!( + "Segment {} range {}..{} is out of bounds for cached segment length {}", + id, + range.start, + range.end, + segment.len() + )); + } + tracing::debug!("Resolved segment {} range {:?} from cache", id, range); + return Ok(BufferHandle::new_host(segment.slice(start..end))); + } + delegate.await + } + .boxed() + }) + .collect() + } } From e5bf12af35ee38e0c6979aa5a974d0335c5e34e3 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 13 Aug 2026 10:45:40 +0000 Subject: [PATCH 13/17] Keep partial range submissions batched Signed-off-by: Joe Isaacs --- vortex-file/src/segments/source.rs | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/vortex-file/src/segments/source.rs b/vortex-file/src/segments/source.rs index 26689dc6956..5730f8fc17e 100644 --- a/vortex-file/src/segments/source.rs +++ b/vortex-file/src/segments/source.rs @@ -207,6 +207,13 @@ impl FileSegmentSource { } while num_active < concurrency && !pending.is_empty() { + // A partial batch is submitted through one `read_ranges` stream. Do not + // refill individual slots from another partial batch as each range finishes: + // that turns a queued group into one syscall submission per completion. Let + // the current group drain, then submit all ready partial ranges together. + if num_active != 0 && pending.front().is_some_and(IoRequest::is_partial) { + break; + } let batch_len = if num_active == 0 && pending.front().is_some_and(IoRequest::is_partial) { partial_submission_len(&pending) From e06b4ca6928e4d8fbd547da517fa7e66e917e907 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Wed, 12 Aug 2026 10:16:02 +0000 Subject: [PATCH 14/17] Read Flat arrays from partial segment ranges Signed-off-by: Joe Isaacs --- Cargo.lock | 2 + encodings/alp/src/alp_rd/array.rs | 27 + encodings/fastlanes/src/bitpacking/mod.rs | 1 + .../fastlanes/src/bitpacking/vtable/mod.rs | 16 + vortex-array/src/arrays/list/mod.rs | 1 + vortex-array/src/arrays/list/vtable/mod.rs | 6 + vortex-array/src/mask_future.rs | 86 +- vortex-array/src/serde.rs | 66 + vortex-layout/Cargo.toml | 2 + vortex-layout/src/display.rs | 2 +- vortex-layout/src/layouts/flat/mod.rs | 3 +- vortex-layout/src/layouts/flat/partial.rs | 1326 +++++++++++++++++ vortex-layout/src/layouts/flat/reader.rs | 385 ++++- vortex-layout/src/layouts/list/writer.rs | 40 +- vortex-layout/src/layouts/table.rs | 52 +- vortex-layout/src/scan/tasks.rs | 17 +- 16 files changed, 1976 insertions(+), 56 deletions(-) create mode 100644 vortex-layout/src/layouts/flat/partial.rs diff --git a/Cargo.lock b/Cargo.lock index 6fddca6f7db..68a103e59cd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10336,11 +10336,13 @@ dependencies = [ "termtree", "tokio", "tracing", + "vortex-alp", "vortex-array", "vortex-arrow", "vortex-btrblocks", "vortex-buffer", "vortex-error", + "vortex-fastlanes", "vortex-flatbuffers", "vortex-io", "vortex-mask", diff --git a/encodings/alp/src/alp_rd/array.rs b/encodings/alp/src/alp_rd/array.rs index b6ee50d7b1f..f7c954d2222 100644 --- a/encodings/alp/src/alp_rd/array.rs +++ b/encodings/alp/src/alp_rd/array.rs @@ -71,6 +71,33 @@ pub struct ALPRDMetadata { patches: Option, } +impl ALPRDMetadata { + pub fn right_bit_width(&self) -> VortexResult { + u8::try_from(self.right_bit_width).map_err(|_| { + vortex_err!( + "right bit width {} does not fit in u8", + self.right_bit_width + ) + }) + } + + pub fn left_parts_dictionary(&self) -> VortexResult> { + self.dict + .get(..usize::try_from(self.dict_len)?) + .ok_or_else(|| vortex_err!("ALPRD dictionary length is out of bounds"))? + .iter() + .map(|&value| { + u16::try_from(value) + .map_err(|_| vortex_err!("ALPRD dictionary value {value} does not fit in u16")) + }) + .collect() + } + + pub fn patches(&self) -> Option<&PatchesMetadata> { + self.patches.as_ref() + } +} + impl ArrayHash for ALPRDData { fn array_hash(&self, state: &mut H, accuracy: EqMode) { self.left_parts_dictionary.array_hash(state, accuracy); diff --git a/encodings/fastlanes/src/bitpacking/mod.rs b/encodings/fastlanes/src/bitpacking/mod.rs index efa0677a91e..115be03d89f 100644 --- a/encodings/fastlanes/src/bitpacking/mod.rs +++ b/encodings/fastlanes/src/bitpacking/mod.rs @@ -19,6 +19,7 @@ mod vtable; pub(crate) use plugin::BitPackedPatchedPlugin; pub use vtable::BitPacked; pub use vtable::BitPackedArray; +pub use vtable::BitPackedMetadata; pub(crate) fn initialize(session: &vortex_session::VortexSession) { vtable::initialize(session); diff --git a/encodings/fastlanes/src/bitpacking/vtable/mod.rs b/encodings/fastlanes/src/bitpacking/vtable/mod.rs index 68fbf1b41d3..8be75c8e553 100644 --- a/encodings/fastlanes/src/bitpacking/vtable/mod.rs +++ b/encodings/fastlanes/src/bitpacking/vtable/mod.rs @@ -72,6 +72,22 @@ pub struct BitPackedMetadata { pub(crate) patches: Option, } +impl BitPackedMetadata { + pub fn bit_width(&self) -> VortexResult { + u8::try_from(self.bit_width) + .map_err(|_| vortex_err!("bit width {} does not fit in u8", self.bit_width)) + } + + pub fn offset(&self) -> VortexResult { + u16::try_from(self.offset) + .map_err(|_| vortex_err!("bit-packed offset {} does not fit in u16", self.offset)) + } + + pub fn patches(&self) -> Option<&PatchesMetadata> { + self.patches.as_ref() + } +} + impl ArrayHash for BitPackedData { fn array_hash(&self, state: &mut H, accuracy: EqMode) { self.offset.hash(state); diff --git a/vortex-array/src/arrays/list/mod.rs b/vortex-array/src/arrays/list/mod.rs index bfd43cdb401..6ee096e747e 100644 --- a/vortex-array/src/arrays/list/mod.rs +++ b/vortex-array/src/arrays/list/mod.rs @@ -14,6 +14,7 @@ pub(crate) mod compute; mod vtable; pub use vtable::List; +pub use vtable::ListMetadata; pub(crate) fn initialize(session: &vortex_session::VortexSession) { compute::initialize(session); diff --git a/vortex-array/src/arrays/list/vtable/mod.rs b/vortex-array/src/arrays/list/vtable/mod.rs index c55e7050351..bd508817e86 100644 --- a/vortex-array/src/arrays/list/vtable/mod.rs +++ b/vortex-array/src/arrays/list/vtable/mod.rs @@ -52,6 +52,12 @@ pub struct ListMetadata { offset_ptype: i32, } +impl ListMetadata { + pub fn elements_len(&self) -> u64 { + self.elements_len + } +} + impl ArrayHash for ListData { fn array_hash(&self, _state: &mut H, _accuracy: EqMode) {} } diff --git a/vortex-array/src/mask_future.rs b/vortex-array/src/mask_future.rs index a46107623c1..767809162cd 100644 --- a/vortex-array/src/mask_future.rs +++ b/vortex-array/src/mask_future.rs @@ -20,6 +20,9 @@ use vortex_mask::Mask; pub struct MaskFuture { inner: Shared>>, len: usize, + upper_bound: Option, + upper_bound_is_exact: bool, + partial_reads_allowed: bool, } impl MaskFuture { @@ -40,6 +43,9 @@ impl MaskFuture { .boxed() .shared(), len, + upper_bound: None, + upper_bound_is_exact: false, + partial_reads_allowed: false, } } @@ -55,7 +61,12 @@ impl MaskFuture { /// Create a MaskFuture from a ready mask. pub fn ready(mask: Mask) -> Self { - Self::new(mask.len(), async move { Ok(mask) }) + let upper_bound = mask.clone(); + let mut future = Self::new(mask.len(), async move { Ok(mask) }); + future.upper_bound = Some(upper_bound); + future.upper_bound_is_exact = true; + future.partial_reads_allowed = true; + future } /// Create a MaskFuture that resolves to a mask with all values set to true. @@ -72,7 +83,57 @@ impl MaskFuture { } let inner = self.inner.clone(); - Self::new(range.len(), async move { Ok(inner.await?.slice(range)) }) + let upper_bound = self + .upper_bound + .as_ref() + .map(|upper_bound| upper_bound.slice(range.clone())); + let mut sliced = Self::new(range.len(), async move { Ok(inner.await?.slice(range)) }); + sliced.upper_bound = upper_bound; + sliced.upper_bound_is_exact = self.upper_bound_is_exact; + sliced.partial_reads_allowed = self.partial_reads_allowed; + sliced + } + + /// Attach a conservative upper bound for the mask resolved by this future. + /// + /// Readers can use this to register I/O eagerly without waiting for filter evaluation. The + /// resolved mask must not contain a true row that is false in `upper_bound`. + pub fn with_upper_bound(mut self, upper_bound: Mask) -> Self { + assert_eq!( + upper_bound.len(), + self.len, + "MaskFuture upper bound length mismatch" + ); + self.upper_bound = Some(upper_bound); + self.upper_bound_is_exact = false; + self + } + + /// Return the conservative upper bound for this future, when one is known. + pub fn upper_bound(&self) -> Option<&Mask> { + self.upper_bound.as_ref() + } + + /// Return whether the upper bound is the exact mask returned by this future. + pub fn upper_bound_is_exact(&self) -> bool { + self.upper_bound_is_exact + } + + /// Permit readers to satisfy this selection using partial segment reads. + pub fn with_partial_reads(mut self) -> Self { + self.partial_reads_allowed = true; + self + } + + /// Prevent readers from turning this mask into partial segment reads. + pub fn without_partial_reads(mut self) -> Self { + self.partial_reads_allowed = false; + self + } + + /// Return whether readers may satisfy this selection using partial segment reads. + pub fn partial_reads_allowed(&self) -> bool { + self.partial_reads_allowed } pub fn inspect( @@ -84,6 +145,9 @@ impl MaskFuture { Self { inner: self.inner.inspect(f).boxed().shared(), len, + upper_bound: self.upper_bound, + upper_bound_is_exact: self.upper_bound_is_exact, + partial_reads_allowed: self.partial_reads_allowed, } } } @@ -119,8 +183,26 @@ mod tests { let partial = fut.slice(0..mask.len() - 1); assert_eq!(partial.len(), mask.len() - 1); + assert_eq!(partial.upper_bound(), Some(&mask.slice(0..mask.len() - 1))); assert_eq!(partial.await?, mask.slice(0..mask.len() - 1)); Ok(()) }) } + + #[test] + fn new_future_has_no_upper_bound_until_attached() { + let future = MaskFuture::new(3, async { Ok(Mask::new_false(3)) }); + assert!(future.upper_bound().is_none()); + + let upper_bound = Mask::from_indices(3, [0, 2]); + let future = future.with_upper_bound(upper_bound.clone()); + assert_eq!(future.upper_bound(), Some(&upper_bound)); + assert!(!future.upper_bound_is_exact()); + } + + #[test] + fn ready_future_has_an_exact_upper_bound() { + let future = MaskFuture::ready(Mask::from_indices(3, [0, 2])); + assert!(future.upper_bound_is_exact()); + } } diff --git a/vortex-array/src/serde.rs b/vortex-array/src/serde.rs index f84ec182269..b1d7d1fd736 100644 --- a/vortex-array/src/serde.rs +++ b/vortex-array/src/serde.rs @@ -5,6 +5,7 @@ use std::borrow::Cow; use std::fmt::Debug; use std::fmt::Formatter; use std::iter; +use std::ops::Range; use std::sync::Arc; use flatbuffers::FlatBufferBuilder; @@ -294,6 +295,31 @@ pub struct SerializedArray { buffers: Arc<[BufferHandle]>, } +/// Location and alignment of one serialized array buffer within its containing segment. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct SerializedBuffer { + index: usize, + range: Range, + alignment: Alignment, +} + +impl SerializedBuffer { + /// Return this buffer's index in the serialized array's global buffer table. + pub fn index(&self) -> usize { + self.index + } + + /// Return this buffer's byte range within the containing segment. + pub fn range(&self) -> &Range { + &self.range + } + + /// Return the alignment required when materializing this buffer independently. + pub fn alignment(&self) -> Alignment { + self.alignment + } +} + impl Debug for SerializedArray { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { f.debug_struct("SerializedArray") @@ -516,6 +542,46 @@ impl SerializedArray { .unwrap_or_default() } + /// Return the global buffer indices referenced by this array node. + pub fn buffer_indices(&self) -> Vec { + self.flatbuffer() + .buffers() + .map(|buffers| buffers.iter().map(usize::from).collect()) + .unwrap_or_default() + } + + /// Return validated locations for all data buffers in their serialized segment. + pub fn buffer_descriptors(&self) -> VortexResult> { + let fb_array = root::(self.flatbuffer.as_ref())?; + let mut offset = 0usize; + fb_array + .buffers() + .unwrap_or_default() + .iter() + .enumerate() + .map(|(index, buffer)| { + if buffer.compression() != Compression::None { + vortex_bail!( + "Partial reads do not support serialized buffer compression {:?}", + buffer.compression() + ); + } + let start = offset + .checked_add(usize::from(buffer.padding())) + .ok_or_else(|| vortex_err!("Buffer {index} padding overflows"))?; + let end = start + .checked_add(buffer.length() as usize) + .ok_or_else(|| vortex_err!("Buffer {index} length overflows"))?; + offset = end; + Ok(SerializedBuffer { + index, + range: start..end, + alignment: Alignment::try_from_untrusted_exponent(buffer.alignment_exponent())?, + }) + }) + .collect() + } + /// Validate and align the array tree flatbuffer, returning the aligned buffer and root location. fn validate_array_tree(array_tree: impl Into) -> VortexResult<(FlatBuffer, usize)> { let fb_buffer = FlatBuffer::align_from(array_tree.into()); diff --git a/vortex-layout/Cargo.toml b/vortex-layout/Cargo.toml index f772b9ab639..63a69c78bf0 100644 --- a/vortex-layout/Cargo.toml +++ b/vortex-layout/Cargo.toml @@ -40,11 +40,13 @@ termtree = { workspace = true } tokio = { workspace = true, features = ["rt"], optional = true } tracing = { workspace = true } vortex-array = { workspace = true } +vortex-alp = { workspace = true } vortex-arrow = { workspace = true } vortex-btrblocks = { workspace = true } vortex-buffer = { workspace = true } vortex-error = { workspace = true } vortex-flatbuffers = { workspace = true, features = ["layout"] } +vortex-fastlanes = { workspace = true } vortex-io = { workspace = true } vortex-mask = { workspace = true } vortex-metrics = { workspace = true } diff --git a/vortex-layout/src/display.rs b/vortex-layout/src/display.rs index c3743a1df45..0fdd28e0b00 100644 --- a/vortex-layout/src/display.rs +++ b/vortex-layout/src/display.rs @@ -346,7 +346,7 @@ vortex.struct, dtype: {numbers=i64?, strings=utf8}, children: 2, rows: 5 #[test] fn test_display_tree_with_segment_source() { if std::env::var("NEXTEST_RUN_ID").is_ok() { - temp_env::with_var("FLAT_LAYOUT_INLINE_ARRAY_NODE", None::<&str>, || { + temp_env::with_var("FLAT_LAYOUT_INLINE_ARRAY_NODE", Some("0"), || { block_on(|handle| async move { let session = new_session().with_handle(handle); let ctx = ArrayContext::empty(); diff --git a/vortex-layout/src/layouts/flat/mod.rs b/vortex-layout/src/layouts/flat/mod.rs index 6913d2fc0c4..eb2eb6a81ca 100644 --- a/vortex-layout/src/layouts/flat/mod.rs +++ b/vortex-layout/src/layouts/flat/mod.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +mod partial; mod reader; pub mod writer; @@ -34,7 +35,7 @@ use crate::segments::SegmentSource; /// Check if inline array node is enabled. pub(super) fn flat_layout_inline_array_node() -> bool { static FLAT_LAYOUT_INLINE_ARRAY_NODE: LazyLock = - LazyLock::new(|| env::var("FLAT_LAYOUT_INLINE_ARRAY_NODE").is_ok_and(|v| v == "1")); + LazyLock::new(|| env::var("FLAT_LAYOUT_INLINE_ARRAY_NODE").map_or(true, |v| v != "0")); *FLAT_LAYOUT_INLINE_ARRAY_NODE } diff --git a/vortex-layout/src/layouts/flat/partial.rs b/vortex-layout/src/layouts/flat/partial.rs new file mode 100644 index 00000000000..1745e2eb6da --- /dev/null +++ b/vortex-layout/src/layouts/flat/partial.rs @@ -0,0 +1,1326 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::collections::BTreeSet; +use std::ops::Range; +use std::sync::Arc; + +use futures::FutureExt; +use futures::future::try_join_all; +use prost::Message; +use vortex_alp::ALPRD; +use vortex_alp::ALPRDMetadata; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::IntoArray; +use vortex_array::VTable; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::FixedSizeList; +use vortex_array::arrays::FixedSizeListArray; +use vortex_array::arrays::List; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::Primitive; +use vortex_array::arrays::Struct; +use vortex_array::arrays::list::ListMetadata; +use vortex_array::buffer::BufferHandle; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::expr::stats::Stat; +use vortex_array::patches::Patches; +use vortex_array::patches::PatchesMetadata; +use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::serde::SerializedArray; +use vortex_array::serde::SerializedBuffer; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::ByteBuffer; +use vortex_error::VortexResult; +use vortex_error::vortex_err; +use vortex_fastlanes::BitPacked; +use vortex_fastlanes::BitPackedMetadata; +use vortex_mask::AllOr; +use vortex_mask::Mask; +use vortex_session::VortexSession; +use vortex_session::registry::ReadContext; + +use crate::layouts::flat::FlatLayout; +use crate::segments::SegmentFuture; +use crate::segments::SegmentId; +use crate::segments::SegmentSource; + +#[derive(Clone)] +pub(super) struct PartialReadPlan { + array_tree: ByteBuffer, + bytes_per_row: usize, + row_granularity: usize, + kind: PartialReadKind, +} + +#[derive(Clone)] +enum PartialReadKind { + Fixed(Arc<[PlannedBuffer]>), + Alprd(Box), + List(ListReadPlan), +} + +#[derive(Clone)] +struct PlannedBuffer { + descriptor: SerializedBuffer, + bytes_per_row: usize, + row_granularity: usize, + bytes_per_granule: usize, +} + +#[derive(Clone)] +struct ListReadPlan { + descriptors: Arc<[SerializedBuffer]>, + element_buffer: SerializedBuffer, + bytes_per_element: usize, + offset_buffers: Arc<[SerializedBuffer]>, + offset_dtype: DType, +} + +#[derive(Clone)] +struct ALPRDReadPlan { + descriptors: Arc<[SerializedBuffer]>, + left: BitPackedReadPlan, + right: BitPackedReadPlan, + patch_buffers: Arc<[SerializedBuffer]>, + patch_metadata: PatchesMetadata, + patch_indices_dtype: DType, + left_parts_dtype: DType, + left_parts_dictionary: Buffer, + right_bit_width: u8, + element_dtype: DType, + list_size: u32, + row_count: usize, +} + +struct PageResolveContext<'a> { + dtype: &'a DType, + row_range: &'a Range, + mask: &'a Mask, + ctx: &'a ReadContext, + session: &'a VortexSession, +} + +#[derive(Clone)] +struct BitPackedReadPlan { + descriptor: SerializedBuffer, + ptype: vortex_array::dtype::PType, + bit_width: u8, + offset: u16, +} + +pub(super) struct RegisteredPartialRead { + array_tree: ByteBuffer, + kind: RegisteredReadKind, +} + +enum RegisteredReadKind { + Fixed { + pages: Vec, + }, + Alprd { + pages: Vec, + patch_buffers: Vec<(SegmentFuture, SerializedBuffer)>, + plan: ALPRDReadPlan, + }, + List { + pages: Vec>, + offset_buffers: Vec<(SegmentFuture, SerializedBuffer)>, + plan: ListReadPlan, + source: Arc, + segment_id: SegmentId, + layout_len: usize, + }, +} + +struct RegisteredALPRDPage { + rows: Range, + left: SegmentFuture, + right: SegmentFuture, +} + +struct RegisteredPage { + rows: Range, + buffers: Vec<(SegmentFuture, SerializedBuffer)>, +} + +impl PartialReadPlan { + pub(super) fn supports_mask(mask: &Mask) -> bool { + !mask.all_true() + } + + pub(super) fn try_new(layout: &FlatLayout) -> VortexResult> { + let Some(array_tree) = layout.array_tree().cloned() else { + return Ok(None); + }; + let serialized = SerializedArray::from_array_tree(array_tree.clone())?; + let descriptors: Arc<[SerializedBuffer]> = serialized.buffer_descriptors()?.into(); + let row_count = usize::try_from(layout.row_count())?; + + if let Some((plan, bytes_per_row)) = try_alprd_plan( + &serialized, + layout.dtype(), + layout.array_ctx(), + row_count, + Arc::clone(&descriptors), + )? { + return Ok(Some(Self { + array_tree, + bytes_per_row, + row_granularity: 1, + kind: PartialReadKind::Alprd(Box::new(plan)), + })); + } + + if let Some((plan, bytes_per_row)) = try_list_plan( + &serialized, + layout.dtype(), + layout.array_ctx(), + row_count, + Arc::clone(&descriptors), + )? { + return Ok(Some(Self { + array_tree, + bytes_per_row, + row_granularity: 1, + kind: PartialReadKind::List(plan), + })); + } + + let mut planned = Vec::new(); + if !collect_raw_buffers( + &serialized, + layout.dtype(), + layout.array_ctx(), + 1, + row_count, + &descriptors, + &mut planned, + )? { + return Ok(None); + } + planned.sort_unstable_by_key(|buffer| buffer.descriptor.index()); + if planned.len() != descriptors.len() + || planned + .iter() + .enumerate() + .any(|(index, buffer)| buffer.descriptor.index() != index) + { + return Ok(None); + } + for buffer in &planned { + let expected = row_count + .div_ceil(buffer.row_granularity) + .checked_mul(buffer.bytes_per_granule) + .ok_or_else(|| vortex_err!("Partial buffer length overflow"))?; + if buffer.descriptor.range().len() != expected { + return Ok(None); + } + } + let bytes_per_row = planned.iter().try_fold(0usize, |sum, buffer| { + sum.checked_add(buffer.bytes_per_row) + .ok_or_else(|| vortex_err!("Partial row width overflow")) + })?; + if bytes_per_row == 0 { + return Ok(None); + } + let row_granularity = planned + .iter() + .map(|buffer| buffer.row_granularity) + .try_fold(1usize, checked_lcm)?; + Ok(Some(Self { + array_tree, + bytes_per_row, + row_granularity, + kind: PartialReadKind::Fixed(planned.into()), + })) + } + + pub(super) fn register( + &self, + source: &Arc, + segment_id: SegmentId, + layout_len: usize, + row_range: &Range, + mask: &Mask, + ) -> Option { + if !Self::supports_mask(mask) { + return None; + } + let preferred_read_size = usize::try_from(source.preferred_read_size()?).ok()?; + let segment_len = usize::try_from(source.segment_len(segment_id)?).ok()?; + let desired_rows = (preferred_read_size / self.bytes_per_row).max(1); + let page_rows = desired_rows + .div_ceil(self.row_granularity) + .saturating_mul(self.row_granularity); + let pages = selected_pages(page_rows, layout_len, row_range, mask)?; + let (partial_bytes, request_count) = self.estimated_partial_io(&pages, layout_len)?; + let partial_cost = partial_bytes.checked_add( + request_count + .saturating_sub(1) + .checked_mul(preferred_read_size)?, + )?; + if partial_cost >= segment_len { + tracing::trace!( + layout_len, + page_rows, + page_count = pages.len(), + partial_bytes, + request_count, + partial_cost, + segment_len, + "Flat partial read rejected by I/O cost" + ); + return None; + } + tracing::trace!( + layout_len, + page_rows, + page_count = pages.len(), + partial_bytes, + request_count, + partial_cost, + segment_len, + "Flat partial read registered" + ); + + let kind = match &self.kind { + PartialReadKind::Fixed(buffers) => { + let page_specs = pages + .into_iter() + .map(|rows| { + let ranges = buffers + .iter() + .map(|buffer| { + let start = buffer.descriptor.range().start + + (rows.start / buffer.row_granularity) + * buffer.bytes_per_granule; + let end = buffer.descriptor.range().start + + rows.end.div_ceil(buffer.row_granularity) + * buffer.bytes_per_granule; + Some(( + u64::try_from(start).ok()?..u64::try_from(end).ok()?, + buffer.descriptor.clone(), + )) + }) + .collect::>>()?; + Some((rows, ranges)) + }) + .collect::>>()?; + let requests = source.request_ranges( + segment_id, + page_specs + .iter() + .flat_map(|(_, ranges)| ranges.iter().map(|(range, _)| range.clone())) + .collect(), + ); + let mut requests = requests.into_iter(); + let pages = page_specs + .into_iter() + .map(|(rows, ranges)| { + let buffers = ranges + .into_iter() + .map(|(_, descriptor)| Some((requests.next()?, descriptor))) + .collect::>>()?; + Some(RegisteredPage { rows, buffers }) + }) + .collect::>>()?; + RegisteredReadKind::Fixed { pages } + } + PartialReadKind::Alprd(plan) => { + let values_per_row = usize::try_from(plan.list_size).ok()?; + let page_specs = pages + .into_iter() + .map(|rows| { + let inner_start = rows.start.checked_mul(values_per_row)?; + let inner_end = rows.end.checked_mul(values_per_row)?; + let left = bitpacked_range(&plan.left, inner_start..inner_end)?; + let right = bitpacked_range(&plan.right, inner_start..inner_end)?; + Some(( + rows, + u64::try_from(left.start).ok()?..u64::try_from(left.end).ok()?, + u64::try_from(right.start).ok()?..u64::try_from(right.end).ok()?, + )) + }) + .collect::>>()?; + let patch_specs = plan + .patch_buffers + .iter() + .map(|descriptor| { + Some(( + u64::try_from(descriptor.range().start).ok()? + ..u64::try_from(descriptor.range().end).ok()?, + descriptor.clone(), + )) + }) + .collect::>>()?; + let ranges = page_specs + .iter() + .flat_map(|(_, left, right)| [left.clone(), right.clone()]) + .chain(patch_specs.iter().map(|(range, _)| range.clone())) + .collect(); + let mut requests = source.request_ranges(segment_id, ranges).into_iter(); + let pages = page_specs + .into_iter() + .map(|(rows, ..)| { + Some(RegisteredALPRDPage { + rows, + left: requests.next()?, + right: requests.next()?, + }) + }) + .collect::>>()?; + let patch_buffers = patch_specs + .into_iter() + .map(|(_, descriptor)| Some((requests.next()?, descriptor))) + .collect::>>()?; + RegisteredReadKind::Alprd { + pages, + patch_buffers, + plan: plan.as_ref().clone(), + } + } + PartialReadKind::List(plan) => { + let offset_specs = plan + .offset_buffers + .iter() + .map(|descriptor| { + Some(( + u64::try_from(descriptor.range().start).ok()? + ..u64::try_from(descriptor.range().end).ok()?, + descriptor.clone(), + )) + }) + .collect::>>()?; + let offset_requests = source.request_ranges( + segment_id, + offset_specs + .iter() + .map(|(range, _)| range.clone()) + .collect(), + ); + let offset_buffers = offset_requests + .into_iter() + .zip(offset_specs) + .map(|(request, (_, descriptor))| (request, descriptor)) + .collect(); + RegisteredReadKind::List { + pages, + offset_buffers, + plan: plan.clone(), + source: Arc::clone(source), + segment_id, + layout_len, + } + } + }; + + Some(RegisteredPartialRead { + array_tree: self.array_tree.clone(), + kind, + }) + } + + fn estimated_partial_io( + &self, + pages: &[Range], + layout_len: usize, + ) -> Option<(usize, usize)> { + match &self.kind { + PartialReadKind::Fixed(buffers) => { + let bytes = pages.iter().try_fold(0usize, |total, rows| { + buffers.iter().try_fold(total, |total, buffer| { + let granules = rows + .end + .div_ceil(buffer.row_granularity) + .checked_sub(rows.start / buffer.row_granularity)?; + total.checked_add(granules.checked_mul(buffer.bytes_per_granule)?) + }) + })?; + Some((bytes, pages.len().checked_mul(buffers.len())?)) + } + PartialReadKind::Alprd(plan) => { + let values_per_row = usize::try_from(plan.list_size).ok()?; + let page_bytes = pages.iter().try_fold(0usize, |total, rows| { + let values = rows.start.checked_mul(values_per_row)? + ..rows.end.checked_mul(values_per_row)?; + let left = bitpacked_range(&plan.left, values.clone())?; + let right = bitpacked_range(&plan.right, values)?; + total.checked_add(left.len())?.checked_add(right.len()) + })?; + let patch_bytes = plan + .patch_buffers + .iter() + .try_fold(0usize, |total, buffer| { + total.checked_add(buffer.range().len()) + })?; + Some(( + page_bytes.checked_add(patch_bytes)?, + pages + .len() + .checked_mul(2)? + .checked_add(plan.patch_buffers.len())?, + )) + } + PartialReadKind::List(plan) => { + let offset_bytes = plan + .offset_buffers + .iter() + .try_fold(0usize, |total, buffer| { + total.checked_add(buffer.range().len()) + })?; + let covered_rows = pages + .iter() + .try_fold(0usize, |total, rows| total.checked_add(rows.len()))?; + let estimated_element_bytes = plan + .element_buffer + .range() + .len() + .checked_mul(covered_rows)? + .div_ceil(layout_len.max(1)); + Some(( + offset_bytes.checked_add(estimated_element_bytes)?, + plan.offset_buffers.len().checked_add(pages.len())?, + )) + } + } + } +} + +impl RegisteredPartialRead { + pub(super) async fn resolve( + self, + dtype: &DType, + row_range: &Range, + mask: &Mask, + ctx: &ReadContext, + session: &VortexSession, + ) -> VortexResult { + let chunks = match self.kind { + RegisteredReadKind::Fixed { pages } => { + resolve_fixed_pages( + self.array_tree, + pages, + PageResolveContext { + dtype, + row_range, + mask, + ctx, + session, + }, + ) + .await? + } + RegisteredReadKind::Alprd { + pages, + patch_buffers, + plan, + } => { + resolve_alprd_pages( + self.array_tree, + pages, + patch_buffers, + plan, + dtype, + row_range, + mask, + ctx, + session, + ) + .await? + } + RegisteredReadKind::List { + pages, + offset_buffers, + plan, + source, + segment_id, + layout_len, + } => { + resolve_list_pages( + self.array_tree, + pages, + offset_buffers, + plan, + source, + segment_id, + layout_len, + dtype, + row_range, + mask, + ctx, + session, + ) + .await? + } + }; + finish_chunks(chunks, dtype, session) + } +} + +async fn resolve_fixed_pages( + array_tree: ByteBuffer, + pages: Vec, + context: PageResolveContext<'_>, +) -> VortexResult> { + let mut page_futures = Vec::new(); + for page in pages { + let local_mask = page_mask(&page.rows, context.row_range, context.mask.indices())?; + if local_mask.all_false() { + continue; + } + let array_tree = array_tree.clone(); + let dtype = context.dtype.clone(); + let ctx = context.ctx.clone(); + let session = context.session.clone(); + page_futures.push(async move { + let buffers = try_join_all(page.buffers.into_iter().map( + |(future, descriptor)| async move { + future.await?.ensure_aligned(descriptor.alignment()) + }, + )) + .await?; + let array = SerializedArray::from_flatbuffer_with_buffers(array_tree, buffers)? + .decode(&dtype, page.rows.len(), &ctx, &session)?; + clear_stats(&array); + apply_page_mask(array, local_mask) + }); + } + try_join_all(page_futures).await +} + +#[allow(clippy::too_many_arguments)] +async fn resolve_alprd_pages( + array_tree: ByteBuffer, + pages: Vec, + patch_requests: Vec<(SegmentFuture, SerializedBuffer)>, + plan: ALPRDReadPlan, + dtype: &DType, + row_range: &Range, + mask: &Mask, + ctx: &ReadContext, + session: &VortexSession, +) -> VortexResult> { + let patch_handles = try_join_all(patch_requests.into_iter().map( + |(future, descriptor)| async move { + Ok::<_, vortex_error::VortexError>(( + descriptor.index(), + future.await?.ensure_aligned(descriptor.alignment())?, + )) + }, + )) + .await?; + let mut handles = empty_handles(plan.descriptors.len()); + for (index, handle) in patch_handles { + handles[index] = handle; + } + let serialized = SerializedArray::from_flatbuffer_with_buffers(array_tree, handles)?; + let alprd = serialized.child(0); + let patch_len = plan.patch_metadata.len()?; + let patch_indices = + alprd + .child(2) + .decode(&plan.patch_indices_dtype, patch_len, ctx, session)?; + let patch_values = alprd.child(3).decode( + &plan.left_parts_dtype.as_nonnullable(), + patch_len, + ctx, + session, + )?; + let full_inner_len = usize::try_from(plan.list_size)? + .checked_mul(plan.row_count) + .ok_or_else(|| vortex_err!("ALPRD inner length overflow"))?; + let full_patches = Patches::new( + full_inner_len, + plan.patch_metadata.offset()?, + patch_indices, + patch_values, + None, + )?; + + let mut page_futures = Vec::new(); + for page in pages { + let local_mask = page_mask(&page.rows, row_range, mask.indices())?; + if local_mask.all_false() { + continue; + } + let plan = plan.clone(); + let dtype = dtype.clone(); + let page_patches = full_patches.clone(); + page_futures.push( + async move { + let left = page + .left + .await? + .ensure_aligned(plan.left.descriptor.alignment())?; + let right = page + .right + .await? + .ensure_aligned(plan.right.descriptor.alignment())?; + let inner_start = page.rows.start * plan.list_size as usize; + let inner_end = page.rows.end * plan.list_size as usize; + let inner_len = inner_end - inner_start; + let left = BitPacked::try_new( + left, + plan.left.ptype, + Validity::from(plan.left_parts_dtype.nullability()), + None, + plan.left.bit_width, + inner_len, + 0, + )? + .into_array(); + let right = BitPacked::try_new( + right, + plan.right.ptype, + Validity::NonNullable, + None, + plan.right.bit_width, + inner_len, + 0, + )? + .into_array(); + let patches = page_patches.slice(inner_start..inner_end)?; + let elements = ALPRD::try_new( + plan.element_dtype.clone(), + left, + plan.left_parts_dictionary.clone(), + right, + plan.right_bit_width, + patches, + )? + .into_array(); + let array = FixedSizeListArray::try_new( + elements, + plan.list_size, + Validity::from(dtype.nullability()), + page.rows.len(), + )? + .into_array(); + clear_stats(&array); + apply_page_mask(array, local_mask) + } + .boxed(), + ); + } + try_join_all(page_futures).await +} + +#[allow(clippy::too_many_arguments)] +async fn resolve_list_pages( + array_tree: ByteBuffer, + pages: Vec>, + offset_requests: Vec<(SegmentFuture, SerializedBuffer)>, + plan: ListReadPlan, + source: Arc, + segment_id: SegmentId, + layout_len: usize, + dtype: &DType, + row_range: &Range, + mask: &Mask, + ctx: &ReadContext, + session: &VortexSession, +) -> VortexResult> { + if pages.is_empty() { + return Ok(Vec::new()); + } + + let resolved_offsets = try_join_all(offset_requests.into_iter().map( + |(future, descriptor)| async move { + Ok::<_, vortex_error::VortexError>(( + descriptor.index(), + future.await?.ensure_aligned(descriptor.alignment())?, + )) + }, + )) + .await?; + let mut offset_handles = empty_handles(plan.descriptors.len()); + for (index, handle) in resolved_offsets { + offset_handles[index] = handle; + } + let offsets = + SerializedArray::from_flatbuffer_with_buffers(array_tree.clone(), offset_handles)? + .child(1) + .decode(&plan.offset_dtype, layout_len + 1, ctx, session)?; + + let mut exec = session.create_execution_ctx(); + let mut page_specs = Vec::new(); + for rows in pages { + let local_mask = page_mask(&rows, row_range, mask.indices())?; + if local_mask.all_false() { + continue; + } + let first = offset_at(&offsets, rows.start, &mut exec)?; + let last = offset_at(&offsets, rows.end, &mut exec)?; + let start = plan.element_buffer.range().start + + first + .checked_mul(plan.bytes_per_element) + .ok_or_else(|| vortex_err!("List element offset overflow"))?; + let end = plan.element_buffer.range().start + + last + .checked_mul(plan.bytes_per_element) + .ok_or_else(|| vortex_err!("List element offset overflow"))?; + page_specs.push(( + rows, + local_mask, + first, + last, + u64::try_from(start)?..u64::try_from(end)?, + )); + } + let element_requests = source.request_ranges( + segment_id, + page_specs + .iter() + .map(|(_, _, _, _, range)| range.clone()) + .collect(), + ); + let mut page_futures = Vec::new(); + for ((rows, local_mask, first, last, _), elements_future) in + page_specs.into_iter().zip(element_requests) + { + let array_tree = array_tree.clone(); + let offsets = offsets.clone(); + let plan = plan.clone(); + let dtype = dtype.clone(); + let ctx = ctx.clone(); + let session = session.clone(); + page_futures.push( + async move { + let elements_handle = elements_future + .await? + .ensure_aligned(plan.element_buffer.alignment())?; + let mut handles = empty_handles(plan.descriptors.len()); + handles[plan.element_buffer.index()] = elements_handle; + let element_dtype = match &dtype { + DType::List(element_dtype, _) => element_dtype.as_ref(), + _ => return Err(vortex_err!("List partial plan used with non-list dtype")), + }; + let elements = SerializedArray::from_flatbuffer_with_buffers(array_tree, handles)? + .child(0) + .decode(element_dtype, last - first, &ctx, &session)?; + clear_stats(&elements); + + let page_offsets = offsets.slice(rows.start..rows.end + 1)?; + let mut exec = session.create_execution_ctx(); + let first_offset = page_offsets.execute_scalar(0, &mut exec)?; + let adjusted_offsets = page_offsets.binary( + ConstantArray::new(first_offset, page_offsets.len()).into_array(), + Operator::Sub, + )?; + let validity = Validity::from(dtype.nullability()); + let array = ListArray::try_new(elements, adjusted_offsets, validity)?.into_array(); + clear_stats(&array); + apply_page_mask(array, local_mask) + } + .boxed(), + ); + } + try_join_all(page_futures).await +} + +fn finish_chunks( + mut chunks: Vec, + dtype: &DType, + session: &VortexSession, +) -> VortexResult { + match chunks.len() { + 0 => Ok(Canonical::empty(dtype).into_array()), + 1 => Ok(chunks.remove(0)), + _ => { + let chunks = ChunkedArray::try_new(chunks, dtype.clone())?.into_array(); + let mut ctx = session.create_execution_ctx(); + Ok(chunks.execute::(&mut ctx)?.into_array()) + } + } +} + +fn apply_page_mask(array: ArrayRef, mask: Mask) -> VortexResult { + if mask.all_true() { + Ok(array) + } else if let AllOr::Some([(start, end)]) = mask.slices() { + array.slice(*start..*end) + } else { + array.filter(mask) + } +} + +fn clear_stats(array: &ArrayRef) { + for child in array.depth_first_traversal() { + for stat in Stat::all() { + child.statistics().clear(stat); + } + } +} + +fn empty_handles(len: usize) -> Vec { + (0..len) + .map(|_| BufferHandle::new_host(ByteBuffer::empty())) + .collect() +} + +fn offset_at( + offsets: &ArrayRef, + index: usize, + ctx: &mut vortex_array::ExecutionCtx, +) -> VortexResult { + offsets + .execute_scalar(index, ctx)? + .as_primitive() + .as_::() + .ok_or_else(|| vortex_err!("List offset does not fit usize")) +} + +fn selected_pages( + page_rows: usize, + layout_len: usize, + row_range: &Range, + mask: &Mask, +) -> Option>> { + let mut page_indices = BTreeSet::new(); + match mask.slices() { + AllOr::All => return None, + AllOr::None => {} + AllOr::Some(slices) => { + for &(start, end) in slices { + if start >= end { + continue; + } + let global_start = row_range.start.checked_add(start)?; + let global_end = row_range.start.checked_add(end)?; + if global_end > row_range.end || global_end > layout_len { + return None; + } + page_indices.extend(global_start / page_rows..=(global_end - 1) / page_rows); + } + } + } + Some( + page_indices + .into_iter() + .map(|page_index| { + let start = page_index * page_rows; + start..start.saturating_add(page_rows).min(layout_len) + }) + .collect(), + ) +} + +fn page_mask( + page_rows: &Range, + row_range: &Range, + selected: AllOr<&[usize]>, +) -> VortexResult { + match selected { + AllOr::None => Ok(Mask::new_false(page_rows.len())), + AllOr::All => { + let start = page_rows.start.max(row_range.start); + let end = page_rows.end.min(row_range.end); + Ok(Mask::from_indices( + page_rows.len(), + (start..end).map(|row| row - page_rows.start), + )) + } + AllOr::Some(indices) => Ok(Mask::from_indices( + page_rows.len(), + indices.iter().filter_map(|&index| { + let row = row_range.start.checked_add(index)?; + page_rows.contains(&row).then(|| row - page_rows.start) + }), + )), + } +} + +fn try_alprd_plan( + node: &SerializedArray, + dtype: &DType, + ctx: &ReadContext, + row_count: usize, + descriptors: Arc<[SerializedBuffer]>, +) -> VortexResult> { + if ctx.resolve(node.encoding_id()) != Some(FixedSizeList.id()) + || node.nbuffers() != 0 + || node.nchildren() != 1 + { + return Ok(None); + } + let DType::FixedSizeList(element_dtype, list_size, _) = dtype else { + return Ok(None); + }; + let list_size_usize = usize::try_from(*list_size)?; + if list_size_usize == 0 || !list_size_usize.is_multiple_of(1024) { + return Ok(None); + } + let DType::Primitive(element_ptype, element_nullability) = element_dtype.as_ref() else { + return Ok(None); + }; + if !matches!( + element_ptype, + vortex_array::dtype::PType::F32 | vortex_array::dtype::PType::F64 + ) { + return Ok(None); + } + + let alprd = node.child(0); + if ctx + .resolve(alprd.encoding_id()) + .is_none_or(|id| id.as_str() != "vortex.alprd") + || alprd.nbuffers() != 0 + || alprd.nchildren() != 4 + { + return Ok(None); + } + let metadata = ALPRDMetadata::decode(alprd.metadata())?; + let Some(patch_metadata) = metadata.patches().copied() else { + return Ok(None); + }; + let left_parts_dtype = DType::Primitive(metadata.left_parts_ptype(), *element_nullability); + let right_ptype = match element_ptype { + vortex_array::dtype::PType::F32 => vortex_array::dtype::PType::U32, + vortex_array::dtype::PType::F64 => vortex_array::dtype::PType::U64, + _ => unreachable!(), + }; + let inner_len = row_count + .checked_mul(list_size_usize) + .ok_or_else(|| vortex_err!("ALPRD inner length overflow"))?; + let left = try_bitpacked_plan( + &alprd.child(0), + left_parts_dtype.as_ptype(), + ctx, + inner_len, + &descriptors, + )?; + let right = try_bitpacked_plan(&alprd.child(1), right_ptype, ctx, inner_len, &descriptors)?; + let (Some(left), Some(right)) = (left, right) else { + return Ok(None); + }; + + let mut patch_indices = BTreeSet::new(); + collect_buffer_indices(&alprd.child(2), &mut patch_indices); + collect_buffer_indices(&alprd.child(3), &mut patch_indices); + let Some(patch_buffers) = patch_indices + .into_iter() + .map(|index| descriptors.get(index).cloned()) + .collect::>>() + else { + return Ok(None); + }; + if patch_buffers.is_empty() { + return Ok(None); + } + let expected_indices: BTreeSet<_> = [left.descriptor.index(), right.descriptor.index()] + .into_iter() + .chain(patch_buffers.iter().map(SerializedBuffer::index)) + .collect(); + if expected_indices.len() != descriptors.len() + || expected_indices.iter().copied().ne(0..descriptors.len()) + { + return Ok(None); + } + + let blocks_per_row = list_size_usize / 1024; + let bytes_per_row = blocks_per_row + .checked_mul(128) + .and_then(|value| value.checked_mul(left.bit_width as usize + right.bit_width as usize)) + .ok_or_else(|| vortex_err!("ALPRD row width overflow"))?; + Ok(Some(( + ALPRDReadPlan { + descriptors, + left, + right, + patch_buffers: patch_buffers.into(), + patch_metadata, + patch_indices_dtype: patch_metadata.indices_dtype()?, + left_parts_dtype, + left_parts_dictionary: metadata.left_parts_dictionary()?, + right_bit_width: metadata.right_bit_width()?, + element_dtype: element_dtype.as_ref().clone(), + list_size: *list_size, + row_count, + }, + bytes_per_row, + ))) +} + +fn try_bitpacked_plan( + node: &SerializedArray, + ptype: vortex_array::dtype::PType, + ctx: &ReadContext, + len: usize, + descriptors: &[SerializedBuffer], +) -> VortexResult> { + if ctx + .resolve(node.encoding_id()) + .is_none_or(|id| id.as_str() != "fastlanes.bitpacked") + || node.nchildren() != 0 + || node.buffer_indices().len() != 1 + { + return Ok(None); + } + let metadata = BitPackedMetadata::decode(node.metadata())?; + if metadata.patches().is_some() || metadata.offset()? != 0 { + return Ok(None); + } + let bit_width = metadata.bit_width()?; + let Some(descriptor) = descriptors.get(node.buffer_indices()[0]).cloned() else { + return Ok(None); + }; + let expected_len = len + .div_ceil(1024) + .checked_mul(128 * bit_width as usize) + .ok_or_else(|| vortex_err!("Bit-packed buffer length overflow"))?; + if descriptor.range().len() != expected_len { + return Ok(None); + } + Ok(Some(BitPackedReadPlan { + descriptor, + ptype, + bit_width, + offset: 0, + })) +} + +fn bitpacked_range(plan: &BitPackedReadPlan, values: Range) -> Option> { + if plan.offset != 0 || !values.start.is_multiple_of(1024) || !values.end.is_multiple_of(1024) { + return None; + } + let bytes_per_block = 128usize.checked_mul(plan.bit_width as usize)?; + let start = plan + .descriptor + .range() + .start + .checked_add((values.start / 1024).checked_mul(bytes_per_block)?)?; + let end = plan + .descriptor + .range() + .start + .checked_add((values.end / 1024).checked_mul(bytes_per_block)?)?; + Some(start..end) +} + +fn try_list_plan( + node: &SerializedArray, + dtype: &DType, + ctx: &ReadContext, + row_count: usize, + descriptors: Arc<[SerializedBuffer]>, +) -> VortexResult> { + if ctx.resolve(node.encoding_id()) != Some(List.id()) { + return Ok(None); + } + let DType::List(element_dtype, _) = dtype else { + return Ok(None); + }; + if node.nbuffers() != 0 || node.nchildren() != 2 { + return Ok(None); + } + let elements = node.child(0); + let DType::Primitive(element_ptype, _) = element_dtype.as_ref() else { + return Ok(None); + }; + if ctx.resolve(elements.encoding_id()) != Some(Primitive.id()) + || elements.nchildren() != 0 + || elements.buffer_indices().len() != 1 + { + return Ok(None); + } + + let metadata = ListMetadata::decode(node.metadata())?; + let element_index = elements.buffer_indices()[0]; + let Some(element_buffer) = descriptors.get(element_index).cloned() else { + return Ok(None); + }; + let elements_len = usize::try_from(metadata.elements_len())?; + let expected_elements_len = elements_len + .checked_mul(element_ptype.byte_width()) + .ok_or_else(|| vortex_err!("List elements length overflow"))?; + if element_buffer.range().len() != expected_elements_len { + return Ok(None); + } + + let offsets = node.child(1); + let mut offset_indices = BTreeSet::new(); + collect_buffer_indices(&offsets, &mut offset_indices); + if offset_indices.is_empty() || offset_indices.contains(&element_index) { + return Ok(None); + } + let Some(offset_buffers) = offset_indices + .into_iter() + .map(|index| descriptors.get(index).cloned()) + .collect::>>() + else { + return Ok(None); + }; + let offset_dtype = DType::Primitive(metadata.offset_ptype(), Nullability::NonNullable); + let offset_bytes: usize = offset_buffers + .iter() + .map(|buffer| buffer.range().len()) + .sum(); + let bytes_per_row = element_buffer + .range() + .len() + .saturating_add(offset_bytes) + .div_ceil(row_count.max(1)) + .max(1); + + Ok(Some(( + ListReadPlan { + descriptors, + element_buffer, + bytes_per_element: element_ptype.byte_width(), + offset_buffers: offset_buffers.into(), + offset_dtype, + }, + bytes_per_row, + ))) +} + +fn collect_buffer_indices(node: &SerializedArray, output: &mut BTreeSet) { + output.extend(node.buffer_indices()); + for index in 0..node.nchildren() { + collect_buffer_indices(&node.child(index), output); + } +} + +fn collect_raw_buffers( + node: &SerializedArray, + dtype: &DType, + ctx: &ReadContext, + row_multiplier: usize, + root_row_count: usize, + descriptors: &[SerializedBuffer], + output: &mut Vec, +) -> VortexResult { + let Some(id) = ctx.resolve(node.encoding_id()) else { + return Ok(false); + }; + + if id == Primitive.id() { + let DType::Primitive(ptype, _) = dtype else { + return Ok(false); + }; + if node.nchildren() != 0 || node.buffer_indices().len() != 1 { + return Ok(false); + } + let index = node.buffer_indices()[0]; + let Some(descriptor) = descriptors.get(index) else { + return Ok(false); + }; + output.push(PlannedBuffer { + descriptor: descriptor.clone(), + bytes_per_row: row_multiplier + .checked_mul(ptype.byte_width()) + .ok_or_else(|| vortex_err!("Partial primitive row width overflow"))?, + row_granularity: 1, + bytes_per_granule: row_multiplier + .checked_mul(ptype.byte_width()) + .ok_or_else(|| vortex_err!("Partial primitive row width overflow"))?, + }); + return Ok(true); + } + + if id == FixedSizeList.id() { + let DType::FixedSizeList(element_dtype, list_size, _) = dtype else { + return Ok(false); + }; + if node.nbuffers() != 0 || node.nchildren() != 1 { + return Ok(false); + } + let multiplier = row_multiplier + .checked_mul(*list_size as usize) + .ok_or_else(|| vortex_err!("Partial fixed-size-list width overflow"))?; + return collect_raw_buffers( + &node.child(0), + element_dtype, + ctx, + multiplier, + root_row_count, + descriptors, + output, + ); + } + + if id == Struct.id() { + let DType::Struct(fields, _) = dtype else { + return Ok(false); + }; + if node.nbuffers() != 0 || node.nchildren() != fields.nfields() { + return Ok(false); + } + for (index, field_dtype) in fields.fields().enumerate() { + if !collect_raw_buffers( + &node.child(index), + &field_dtype, + ctx, + row_multiplier, + root_row_count, + descriptors, + output, + )? { + return Ok(false); + } + } + return Ok(true); + } + + if id.as_str() == "vortex.alprd" { + if !matches!(dtype, DType::Primitive(_, _)) + || node.nbuffers() != 0 + || node.nchildren() != 2 + || row_multiplier == 0 + || root_row_count == 0 + { + return Ok(false); + } + let granularity = 1024 / gcd(1024, row_multiplier); + for child_index in 0..2 { + let child = node.child(child_index); + let Some(child_id) = ctx.resolve(child.encoding_id()) else { + return Ok(false); + }; + if child_id.as_str() != "fastlanes.bitpacked" + || child.nchildren() != 0 + || child.buffer_indices().len() != 1 + { + return Ok(false); + } + let index = child.buffer_indices()[0]; + let Some(descriptor) = descriptors.get(index) else { + return Ok(false); + }; + let granules = root_row_count.div_ceil(granularity); + if descriptor.range().len() % granules != 0 { + return Ok(false); + } + let bytes_per_granule = descriptor.range().len() / granules; + output.push(PlannedBuffer { + descriptor: descriptor.clone(), + bytes_per_row: bytes_per_granule.div_ceil(granularity), + row_granularity: granularity, + bytes_per_granule, + }); + } + return Ok(true); + } + + Ok(false) +} + +fn gcd(mut left: usize, mut right: usize) -> usize { + while right != 0 { + (left, right) = (right, left % right); + } + left +} + +fn checked_lcm(left: usize, right: usize) -> VortexResult { + left.checked_div(gcd(left, right)) + .and_then(|value| value.checked_mul(right)) + .ok_or_else(|| vortex_err!("Partial row granularity overflow")) +} diff --git a/vortex-layout/src/layouts/flat/reader.rs b/vortex-layout/src/layouts/flat/reader.rs index aa7609f1659..9f0709d12ad 100644 --- a/vortex-layout/src/layouts/flat/reader.rs +++ b/vortex-layout/src/layouts/flat/reader.rs @@ -4,6 +4,7 @@ use std::ops::BitAnd; use std::ops::Range; use std::sync::Arc; +use std::sync::OnceLock; use futures::FutureExt; use futures::future::BoxFuture; @@ -22,6 +23,8 @@ use vortex_session::VortexSession; use crate::layouts::SharedArrayFuture; use crate::layouts::flat::FlatLayout; +use crate::layouts::flat::partial::PartialReadPlan; +use crate::layouts::flat::partial::RegisteredPartialRead; use crate::reader::LayoutReader; use crate::reader::RowSplits; use crate::reader::SplitRange; @@ -34,11 +37,13 @@ use crate::segments::SegmentSource; // actual expression? Perhaps all expressions are given a selection mask to decide for themselves? const EXPR_EVAL_THRESHOLD: f64 = 0.2; +#[derive(Clone)] pub struct FlatReader { layout: FlatLayout, name: Arc, segment_source: Arc, session: VortexSession, + partial_plan: Arc>>, } impl FlatReader { @@ -53,9 +58,36 @@ impl FlatReader { name, segment_source, session, + partial_plan: Arc::new(OnceLock::new()), } } + fn register_partial( + &self, + row_range: &Range, + mask: &Mask, + ) -> Option { + if !PartialReadPlan::supports_mask(mask) { + return None; + } + let plan = self + .partial_plan + .get_or_init(|| match PartialReadPlan::try_new(&self.layout) { + Ok(plan) => plan, + Err(error) => { + tracing::debug!("Flat partial-read plan disabled: {error}"); + None + } + }); + plan.as_ref()?.register( + &self.segment_source, + self.layout.segment_id(), + usize::try_from(self.layout.row_count()).ok()?, + row_range, + mask, + ) + } + /// Register the segment request and return a future that would resolve into the deserialised array. fn array_future(&self) -> SharedArrayFuture { let row_count = @@ -131,18 +163,87 @@ impl LayoutReader for FlatReader { .vortex_expect("Row range begin must fit within FlatLayout size") ..usize::try_from(row_range.end) .vortex_expect("Row range end must fit within FlatLayout size"); + if !mask.partial_reads_allowed() { + let name = Arc::clone(&self.name); + let array = self.array_future(); + let expr = expr.clone(); + let session = self.session.clone(); + + return Ok(MaskFuture::new(mask.len(), async move { + let mut array = array.await?; + let mask = mask.await?; + + if row_range.start > 0 || row_range.end < array.len() { + array = array.slice(row_range.clone())?; + } + + let mask_density = mask.density(); + let array_mask = if mask_density < EXPR_EVAL_THRESHOLD { + let array = array.apply_bound(&expr)?; + let array = array.filter(mask.clone())?; + let mut ctx = session.create_execution_ctx(); + let array_mask = array.null_as_false().execute(&mut ctx)?; + mask.intersect_by_rank(&array_mask) + } else { + let array = array.apply_bound(&expr)?; + let mut ctx = session.create_execution_ctx(); + let array_mask = array.null_as_false().execute(&mut ctx)?; + mask.bitand(&array_mask) + }; + + trace!( + "Flat mask evaluation {} - {} (mask = {}) => {}", + name, + expr, + mask_density, + array_mask.density(), + ); + Ok(array_mask) + })); + } let name = Arc::clone(&self.name); - let array = self.array_future(); let expr = expr.clone(); let session = self.session.clone(); + let reader = self.clone(); + let partial_reads_allowed = mask.partial_reads_allowed(); + let registered = partial_reads_allowed + .then(|| mask.upper_bound()) + .flatten() + .and_then(|upper_bound| self.register_partial(&row_range, upper_bound)); + let eager_array = + (mask.upper_bound_is_exact() && registered.is_none()).then(|| self.array_future()); Ok(MaskFuture::new(mask.len(), async move { // 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 mask = mask.await?; + if let Some(registered) = registered.or_else(|| { + partial_reads_allowed + .then(|| reader.register_partial(&row_range, &mask)) + .flatten() + }) { + let array = registered + .resolve( + reader.layout.dtype(), + &row_range, + &mask, + reader.layout.array_ctx(), + &session, + ) + .await?; + let array = array.apply_bound(&expr)?; + let mut ctx = session.create_execution_ctx(); + let array_mask = array.null_as_false().execute(&mut ctx)?; + return Ok(mask.intersect_by_rank(&array_mask)); + } + + let mut array = match eager_array { + Some(array) => array.await?, + None => reader.array_future().await?, + }; + // Slice the array based on the row mask. if row_range.start > 0 || row_range.end < array.len() { array = array.slice(row_range.clone())?; @@ -190,16 +291,64 @@ impl LayoutReader for FlatReader { .vortex_expect("Row range begin must fit within FlatLayout size") ..usize::try_from(row_range.end) .vortex_expect("Row range end must fit within FlatLayout size"); + if !mask.partial_reads_allowed() { + let name = Arc::clone(&self.name); + let array = self.array_future(); + let expr = expr.clone(); + + return Ok(async move { + trace!("Flat array evaluation {} - {}", name, expr); + + let mut array = array.await?; + let mask = mask.await?; + + if row_range.start > 0 || row_range.end < array.len() { + array = array.slice(row_range.clone())?; + } + if !mask.all_true() { + array = array.filter(mask)?; + } + array = array.apply_bound(&expr)?; + Ok(array) + } + .boxed()); + } let name = Arc::clone(&self.name); - let array = self.array_future(); let expr = expr.clone(); + let reader = self.clone(); + let partial_reads_allowed = mask.partial_reads_allowed(); + let registered = partial_reads_allowed + .then(|| mask.upper_bound()) + .flatten() + .and_then(|upper_bound| self.register_partial(&row_range, upper_bound)); + let eager_array = ((!partial_reads_allowed || mask.upper_bound_is_exact()) + && registered.is_none()) + .then(|| self.array_future()); Ok(async move { trace!("Flat array evaluation {} - {}", name, expr); - let mut array = array.clone().await?; let mask = mask.await?; + if let Some(registered) = registered { + let mut array = registered + .resolve( + reader.layout.dtype(), + &row_range, + &mask, + reader.layout.array_ctx(), + &reader.session, + ) + .await?; + array = array.apply_bound(&expr)?; + return Ok(array); + } + + let mut array = match eager_array { + Some(array) => array.await?, + None => reader.array_future().await?, + }; + // Slice the array based on the row mask. if row_range.start > 0 || row_range.end < array.len() { array = array.slice(row_range.clone())?; @@ -228,13 +377,18 @@ impl LayoutReader for FlatReader { #[cfg(test)] mod test { + use std::ops::Range; use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use parking_lot::Mutex; use vortex_array::ArrayContext; use vortex_array::IntoArray; use vortex_array::MaskFuture; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::expr::gt; @@ -245,14 +399,46 @@ mod test { use vortex_error::VortexResult; use vortex_io::runtime::single::block_on; use vortex_io::session::RuntimeSessionExt; + use vortex_mask::Mask; use crate::LayoutStrategy; use crate::layouts::flat::writer::FlatLayoutStrategy; + use crate::segments::SegmentFuture; + use crate::segments::SegmentId; + use crate::segments::SegmentSource; + use crate::segments::SharedSegmentSource; use crate::segments::TestSegments; use crate::sequence::SequenceId; use crate::sequence::SequentialArrayStreamExt; use crate::test::new_session; + #[derive(Clone, Default)] + struct RangedTestSource { + inner: Arc, + ranges: Arc>>>, + whole_requests: Arc, + } + + impl SegmentSource for RangedTestSource { + fn preferred_read_size(&self) -> Option { + Some(16) + } + + fn segment_len(&self, id: SegmentId) -> Option { + self.inner.segment_len(id) + } + + fn request(&self, id: SegmentId) -> SegmentFuture { + self.whole_requests.fetch_add(1, Ordering::Relaxed); + self.inner.request(id) + } + + fn request_range(&self, id: SegmentId, range: Range) -> SegmentFuture { + self.ranges.lock().push(range.clone()); + self.inner.request_range(id, range) + } + } + #[test] fn flat_identity() -> VortexResult<()> { block_on(|handle| async { @@ -370,4 +556,195 @@ mod test { assert_arrays_eq!(result, expected, &mut ctx); }) } + + #[test] + fn sparse_projection_reads_only_virtual_pages() -> VortexResult<()> { + block_on(|handle| async { + let session = new_session().with_handle(handle); + let mut ctx = session.create_execution_ctx(); + let array_ctx = ArrayContext::empty(); + let source = RangedTestSource::default(); + let (ptr, eof) = SequenceId::root().split(); + let array = PrimitiveArray::from_iter(0i32..64).into_array(); + let layout = FlatLayoutStrategy::default() + .write_stream( + array_ctx.into(), + Arc::::clone(&source.inner), + array.to_array_stream().sequenced(ptr), + eof, + &session, + ) + .await?; + + let reader = layout.new_reader( + "".into(), + Arc::new(source.clone()), + &session, + &Default::default(), + )?; + let expr = root().bind(reader.dtype())?; + let result = reader + .projection_evaluation( + &(0..64), + &expr, + MaskFuture::ready(Mask::from_indices(64, [1, 10])), + )? + .await?; + + let expected = PrimitiveArray::from_iter([1i32, 10]).into_array(); + assert_arrays_eq!(result, expected, &mut ctx); + assert_eq!(source.whole_requests.load(Ordering::Relaxed), 0); + assert_eq!(*source.ranges.lock(), [0..16, 32..48]); + + let result = reader + .projection_evaluation( + &(0..64), + &expr, + MaskFuture::ready(Mask::from_indices(64, [1, 10])), + )? + .await?; + assert_arrays_eq!(result, expected, &mut ctx); + assert_eq!( + *source.ranges.lock(), + [0..16, 32..48, 0..16, 32..48], + "separate evaluations must not retain page data" + ); + Ok(()) + }) + } + + #[test] + fn dense_projection_chooses_whole_segment() -> VortexResult<()> { + block_on(|handle| async { + let session = new_session().with_handle(handle); + let mut ctx = session.create_execution_ctx(); + let array_ctx = ArrayContext::empty(); + let source = RangedTestSource::default(); + let (ptr, eof) = SequenceId::root().split(); + let array = PrimitiveArray::from_iter(0i32..64).into_array(); + let layout = FlatLayoutStrategy::default() + .write_stream( + array_ctx.into(), + Arc::::clone(&source.inner), + array.to_array_stream().sequenced(ptr), + eof, + &session, + ) + .await?; + + let reader = layout.new_reader( + "".into(), + Arc::new(source.clone()), + &session, + &Default::default(), + )?; + let expr = root().bind(reader.dtype())?; + let mask = Mask::from_indices(64, (0..64).step_by(2)); + let result = reader + .projection_evaluation(&(0..64), &expr, MaskFuture::ready(mask.clone()))? + .await?; + + assert_arrays_eq!(result, array.filter(mask)?, &mut ctx); + assert_eq!(source.whole_requests.load(Ordering::Relaxed), 1); + assert!(source.ranges.lock().is_empty()); + Ok(()) + }) + } + + #[test] + fn filter_and_projection_share_pages_while_scan_is_in_flight() -> VortexResult<()> { + block_on(|handle| async { + let session = new_session().with_handle(handle); + let mut ctx = session.create_execution_ctx(); + let array_ctx = ArrayContext::empty(); + let source = RangedTestSource::default(); + let (ptr, eof) = SequenceId::root().split(); + let array = PrimitiveArray::from_iter(0i32..64).into_array(); + let layout = FlatLayoutStrategy::default() + .write_stream( + array_ctx.into(), + Arc::::clone(&source.inner), + array.to_array_stream().sequenced(ptr), + eof, + &session, + ) + .await?; + + let reader = layout.new_reader( + "".into(), + Arc::new(SharedSegmentSource::new(source.clone())), + &session, + &Default::default(), + )?; + let projection_expr = root().bind(reader.dtype())?; + let filter_expr = gt(root(), lit(-1i32)).bind(reader.dtype())?; + let mask = Mask::from_indices(64, [1, 10]); + + let filter = reader.filter_evaluation( + &(0..64), + &filter_expr, + MaskFuture::ready(mask.clone()), + )?; + let projection = reader.projection_evaluation( + &(0..64), + &projection_expr, + MaskFuture::ready(mask), + )?; + + let (filter_mask, result) = futures::try_join!(filter, projection)?; + assert_eq!( + filter_mask.indices(), + Mask::from_indices(64, [1, 10]).indices() + ); + let expected = PrimitiveArray::from_iter([1i32, 10]).into_array(); + assert_arrays_eq!(result, expected, &mut ctx); + assert_eq!( + *source.ranges.lock(), + [0..16, 32..48], + "one in-flight request should serve filter and projection" + ); + Ok(()) + }) + } + + #[test] + fn sparse_list_projection_reads_offsets_then_selected_elements() -> VortexResult<()> { + block_on(|handle| async { + let session = new_session().with_handle(handle); + let mut ctx = session.create_execution_ctx(); + let array_ctx = ArrayContext::empty(); + let source = RangedTestSource::default(); + let (ptr, eof) = SequenceId::root().split(); + let elements = PrimitiveArray::from_iter(0i64..32).into_array(); + let offsets = PrimitiveArray::from_iter((0u32..=16).map(|v| v * 2)).into_array(); + let array = ListArray::try_new(elements, offsets, Validity::NonNullable)?.into_array(); + let layout = FlatLayoutStrategy::default() + .write_stream( + array_ctx.into(), + Arc::::clone(&source.inner), + array.to_array_stream().sequenced(ptr), + eof, + &session, + ) + .await?; + + let reader = layout.new_reader( + "".into(), + Arc::new(source.clone()), + &session, + &Default::default(), + )?; + let expr = root().bind(reader.dtype())?; + let mask = Mask::from_indices(16, [1, 10]); + let result = reader + .projection_evaluation(&(0..16), &expr, MaskFuture::ready(mask.clone()))? + .await?; + + let expected = array.filter(mask)?; + assert_arrays_eq!(result, expected, &mut ctx); + assert_eq!(source.whole_requests.load(Ordering::Relaxed), 0); + assert!(source.ranges.lock().len() >= 3); + Ok(()) + }) + } } diff --git a/vortex-layout/src/layouts/list/writer.rs b/vortex-layout/src/layouts/list/writer.rs index 4d8565fdd10..a8cecc4cdf4 100644 --- a/vortex-layout/src/layouts/list/writer.rs +++ b/vortex-layout/src/layouts/list/writer.rs @@ -399,8 +399,8 @@ mod tests { insta::assert_snapshot!(layout.display_tree(), @" vortex.list, dtype: list(i32), children: 2 - ├── elements: vortex.flat, dtype: i32, segment: 0 - └── offsets: vortex.flat, dtype: u64, segment: 1 + ├── elements: vortex.flat, dtype: i32, segment 0, buffers=[20B], total=20B + └── offsets: vortex.flat, dtype: u64, segment 1, buffers=[32B], total=32B "); Ok(()) } @@ -416,9 +416,9 @@ mod tests { insta::assert_snapshot!(layout.display_tree(), @" vortex.list, dtype: list(i32)?, children: 3 - ├── elements: vortex.flat, dtype: i32, segment: 0 - ├── offsets: vortex.flat, dtype: u64, segment: 1 - └── validity: vortex.flat, dtype: bool, segment: 2 + ├── elements: vortex.flat, dtype: i32, segment 0, buffers=[20B], total=20B + ├── offsets: vortex.flat, dtype: u64, segment 1, buffers=[32B], total=32B + └── validity: vortex.flat, dtype: bool, segment 2, buffers=[1B], total=1B "); Ok(()) } @@ -428,7 +428,7 @@ mod tests { async fn non_list_input_routes_to_fallback() -> VortexResult<()> { let primitive = buffer![1i32, 2, 3].into_array(); let layout = write(&flat_list_strategy(), primitive).await?; - insta::assert_snapshot!(layout.display_tree(), @"vortex.flat, dtype: i32, segment: 0"); + insta::assert_snapshot!(layout.display_tree(), @"vortex.flat, dtype: i32, segment 0, buffers=[12B], total=12B"); Ok(()) } @@ -478,9 +478,9 @@ mod tests { insta::assert_snapshot!(layout.display_tree(), @" vortex.list, dtype: list({a=i32, b=i32}), children: 2 ├── elements: vortex.struct, dtype: {a=i32, b=i32}, children: 2 - │ ├── a: vortex.flat, dtype: i32, segment: 1 - │ └── b: vortex.flat, dtype: i32, segment: 2 - └── offsets: vortex.flat, dtype: u64, segment: 0 + │ ├── a: vortex.flat, dtype: i32, segment 1, buffers=[20B], total=20B + │ └── b: vortex.flat, dtype: i32, segment 2, buffers=[20B], total=20B + └── offsets: vortex.flat, dtype: u64, segment 0, buffers=[32B], total=32B "); Ok(()) } @@ -506,9 +506,9 @@ mod tests { insta::assert_snapshot!(layout.display_tree(), @" vortex.list, dtype: list(list(i32)), children: 2 ├── elements: vortex.list, dtype: list(i32), children: 2 - │ ├── elements: vortex.flat, dtype: i32, segment: 1 - │ └── offsets: vortex.flat, dtype: u64, segment: 2 - └── offsets: vortex.flat, dtype: u64, segment: 0 + │ ├── elements: vortex.flat, dtype: i32, segment 1, buffers=[24B], total=24B + │ └── offsets: vortex.flat, dtype: u64, segment 2, buffers=[40B], total=40B + └── offsets: vortex.flat, dtype: u64, segment 0, buffers=[24B], total=24B "); Ok(()) } @@ -539,10 +539,10 @@ mod tests { vortex.list, dtype: list(list(list(i32))), children: 2 ├── elements: vortex.list, dtype: list(list(i32)), children: 2 │ ├── elements: vortex.list, dtype: list(i32), children: 2 - │ │ ├── elements: vortex.flat, dtype: i32, segment: 2 - │ │ └── offsets: vortex.flat, dtype: u64, segment: 3 - │ └── offsets: vortex.flat, dtype: u64, segment: 1 - └── offsets: vortex.flat, dtype: u64, segment: 0 + │ │ ├── elements: vortex.flat, dtype: i32, segment 2, buffers=[16B], total=16B + │ │ └── offsets: vortex.flat, dtype: u64, segment 3, buffers=[24B], total=24B + │ └── offsets: vortex.flat, dtype: u64, segment 1, buffers=[16B], total=16B + └── offsets: vortex.flat, dtype: u64, segment 0, buffers=[16B], total=16B "); Ok(()) } @@ -572,11 +572,11 @@ mod tests { insta::assert_snapshot!(layout.display_tree(), @" vortex.chunked, dtype: list(i32), children: 2 ├── [0]: vortex.list, dtype: list(i32), children: 2 - │ ├── elements: vortex.flat, dtype: i32, segment: 0 - │ └── offsets: vortex.flat, dtype: u64, segment: 1 + │ ├── elements: vortex.flat, dtype: i32, segment 0, buffers=[12B], total=12B + │ └── offsets: vortex.flat, dtype: u64, segment 1, buffers=[24B], total=24B └── [1]: vortex.list, dtype: list(i32), children: 2 - ├── elements: vortex.flat, dtype: i32, segment: 2 - └── offsets: vortex.flat, dtype: u64, segment: 3 + ├── elements: vortex.flat, dtype: i32, segment 2, buffers=[16B], total=16B + └── offsets: vortex.flat, dtype: u64, segment 3, buffers=[24B], total=24B "); Ok(()) } diff --git a/vortex-layout/src/layouts/table.rs b/vortex-layout/src/layouts/table.rs index 1a3c1adc524..c70b93909d9 100644 --- a/vortex-layout/src/layouts/table.rs +++ b/vortex-layout/src/layouts/table.rs @@ -406,10 +406,10 @@ mod tests { .into_array(); let layout = write(&flat_table(), struct_array).await?; - insta::assert_snapshot!(layout.display_tree(), @r" + insta::assert_snapshot!(layout.display_tree(), @" vortex.struct, dtype: {a=i32, b=i32}, children: 2 - ├── a: vortex.flat, dtype: i32, segment: 0 - └── b: vortex.flat, dtype: i32, segment: 1 + ├── a: vortex.flat, dtype: i32, segment 0, buffers=[12B], total=12B + └── b: vortex.flat, dtype: i32, segment 1, buffers=[12B], total=12B "); Ok(()) } @@ -432,12 +432,12 @@ mod tests { .into_array(); let layout = write(&flat_table().with_list_layout(), outer).await?; - insta::assert_snapshot!(layout.display_tree(), @r" + insta::assert_snapshot!(layout.display_tree(), @" vortex.list, dtype: list(list(i32)), children: 2 ├── elements: vortex.list, dtype: list(i32), children: 2 - │ ├── elements: vortex.flat, dtype: i32, segment: 1 - │ └── offsets: vortex.flat, dtype: u64, segment: 2 - └── offsets: vortex.flat, dtype: u64, segment: 0 + │ ├── elements: vortex.flat, dtype: i32, segment 1, buffers=[24B], total=24B + │ └── offsets: vortex.flat, dtype: u64, segment 2, buffers=[40B], total=40B + └── offsets: vortex.flat, dtype: u64, segment 0, buffers=[24B], total=24B "); Ok(()) } @@ -463,14 +463,14 @@ mod tests { let st = StructArray::from_fields([("items", items)].as_slice())?.into_array(); let layout = write(&flat_table().with_list_layout(), st).await?; - insta::assert_snapshot!(layout.display_tree(), @r" + insta::assert_snapshot!(layout.display_tree(), @" vortex.struct, dtype: {items=list({a=i32, b=i32})?}, children: 1 └── items: vortex.list, dtype: list({a=i32, b=i32})?, children: 3 ├── elements: vortex.struct, dtype: {a=i32, b=i32}, children: 2 - │ ├── a: vortex.flat, dtype: i32, segment: 2 - │ └── b: vortex.flat, dtype: i32, segment: 3 - ├── offsets: vortex.flat, dtype: u64, segment: 0 - └── validity: vortex.flat, dtype: bool, segment: 1 + │ ├── a: vortex.flat, dtype: i32, segment 2, buffers=[20B], total=20B + │ └── b: vortex.flat, dtype: i32, segment 3, buffers=[20B], total=20B + ├── offsets: vortex.flat, dtype: u64, segment 0, buffers=[32B], total=32B + └── validity: vortex.flat, dtype: bool, segment 1, buffers=[1B], total=1B "); Ok(()) } @@ -502,14 +502,14 @@ mod tests { ) .with_list_layout(); let layout = write(&dispatcher, chunked).await?; - insta::assert_snapshot!(layout.display_tree(), @r" + insta::assert_snapshot!(layout.display_tree(), @" vortex.list, dtype: list(i32), children: 2 ├── elements: vortex.chunked, dtype: i32, children: 2 - │ ├── [0]: vortex.flat, dtype: i32, segment: 0 - │ └── [1]: vortex.flat, dtype: i32, segment: 1 + │ ├── [0]: vortex.flat, dtype: i32, segment 0, buffers=[12B], total=12B + │ └── [1]: vortex.flat, dtype: i32, segment 1, buffers=[16B], total=16B └── offsets: vortex.chunked, dtype: u64, children: 2 - ├── [0]: vortex.flat, dtype: u64, segment: 2 - └── [1]: vortex.flat, dtype: u64, segment: 3 + ├── [0]: vortex.flat, dtype: u64, segment 2, buffers=[24B], total=24B + └── [1]: vortex.flat, dtype: u64, segment 3, buffers=[16B], total=16B "); Ok(()) } @@ -569,7 +569,7 @@ mod tests { async fn non_struct_input_uses_leaf() -> VortexResult<()> { let primitive = PrimitiveArray::from_iter([1i32, 2, 3]).into_array(); let layout = write(&flat_table(), primitive).await?; - insta::assert_snapshot!(layout.display_tree(), @"vortex.flat, dtype: i32, segment: 0"); + insta::assert_snapshot!(layout.display_tree(), @"vortex.flat, dtype: i32, segment 0, buffers=[12B], total=12B"); Ok(()) } @@ -601,14 +601,14 @@ mod tests { let chunked = ChunkedArray::try_new(vec![c0, c1], dtype)?.into_array(); let layout = write(&dispatcher, chunked).await?; - insta::assert_snapshot!(layout.display_tree(), @r" + insta::assert_snapshot!(layout.display_tree(), @" vortex.struct, dtype: {a=i32, b=i32}, children: 2 ├── a: vortex.chunked, dtype: i32, children: 2 - │ ├── [0]: vortex.flat, dtype: i32, segment: 0 - │ └── [1]: vortex.flat, dtype: i32, segment: 1 + │ ├── [0]: vortex.flat, dtype: i32, segment 0, buffers=[8B], total=8B + │ └── [1]: vortex.flat, dtype: i32, segment 1, buffers=[4B], total=4B └── b: vortex.chunked, dtype: i32, children: 2 - ├── [0]: vortex.flat, dtype: i32, segment: 2 - └── [1]: vortex.flat, dtype: i32, segment: 3 + ├── [0]: vortex.flat, dtype: i32, segment 2, buffers=[8B], total=8B + └── [1]: vortex.flat, dtype: i32, segment 3, buffers=[4B], total=4B "); Ok(()) } @@ -628,10 +628,10 @@ mod tests { let strategy = flat_table().with_field_writer(field_path!(a), Arc::new(FlatLayoutStrategy::default())); let layout = write(&strategy, struct_array).await?; - insta::assert_snapshot!(layout.display_tree(), @r" + insta::assert_snapshot!(layout.display_tree(), @" vortex.struct, dtype: {a=i32, b=i32}, children: 2 - ├── a: vortex.flat, dtype: i32, segment: 0 - └── b: vortex.flat, dtype: i32, segment: 1 + ├── a: vortex.flat, dtype: i32, segment 0, buffers=[12B], total=12B + └── b: vortex.flat, dtype: i32, segment 1, buffers=[12B], total=12B "); Ok(()) } diff --git a/vortex-layout/src/scan/tasks.rs b/vortex-layout/src/scan/tasks.rs index 218efb64a0d..d7e8342c664 100644 --- a/vortex-layout/src/scan/tasks.rs +++ b/vortex-layout/src/scan/tasks.rs @@ -69,8 +69,10 @@ pub fn split_exec( let reader = Arc::clone(&ctx.reader); let filter = Arc::clone(filter); let row_range = row_range.clone(); + let filter_upper_bound = row_mask.clone(); + let partial_reads_allowed = !row_mask.all_true(); - MaskFuture::new(row_mask.len(), async move { + let filter_mask = MaskFuture::new(row_mask.len(), async move { let mut mask = row_mask; let mut dynamic_versions = vec![None; filter.conjuncts().len()]; @@ -117,8 +119,13 @@ pub fn split_exec( return Ok(mask); } + let mask_future = if partial_reads_allowed { + MaskFuture::ready(mask) + } else { + MaskFuture::ready(mask).without_partial_reads() + }; let conjunct_mask = reader - .filter_evaluation(&row_range, conjunct, MaskFuture::ready(mask))? + .filter_evaluation(&row_range, conjunct, mask_future)? .await?; filter.report_selectivity(idx, conjunct_mask.density()); @@ -128,6 +135,12 @@ pub fn split_exec( Ok(mask) }) + .with_upper_bound(filter_upper_bound); + if partial_reads_allowed { + filter_mask.with_partial_reads() + } else { + filter_mask + } } }; From ccb455a42467bf9824d87e599f607b0a4ae4a79c Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Wed, 12 Aug 2026 21:46:10 +0000 Subject: [PATCH 15/17] Issue ALPRD partial reads in one round Signed-off-by: Joe Isaacs --- vortex-layout/src/layouts/flat/partial.rs | 142 +++++++++++----------- 1 file changed, 74 insertions(+), 68 deletions(-) diff --git a/vortex-layout/src/layouts/flat/partial.rs b/vortex-layout/src/layouts/flat/partial.rs index 1745e2eb6da..865eb7b126c 100644 --- a/vortex-layout/src/layouts/flat/partial.rs +++ b/vortex-layout/src/layouts/flat/partial.rs @@ -614,8 +614,33 @@ async fn resolve_alprd_pages( future.await?.ensure_aligned(descriptor.alignment())?, )) }, - )) - .await?; + )); + let page_handles = try_join_all(pages.into_iter().filter_map(|page| { + let local_mask = match page_mask(&page.rows, row_range, mask.indices()) { + Ok(local_mask) if !local_mask.all_false() => local_mask, + Ok(_) => return None, + Err(error) => return Some(futures::future::ready(Err(error)).left_future()), + }; + let left_alignment = plan.left.descriptor.alignment(); + let right_alignment = plan.right.descriptor.alignment(); + Some( + async move { + let (left, right) = futures::try_join!(page.left, page.right)?; + Ok::<_, vortex_error::VortexError>(( + page.rows, + local_mask, + left.ensure_aligned(left_alignment)?, + right.ensure_aligned(right_alignment)?, + )) + } + .right_future(), + ) + })); + + // Every range for this Flat layout is registered before resolution. Poll the complete set + // together so the driver can issue and coalesce patch, left-part, and right-part reads in one + // I/O round; array reconstruction starts only after that set has resolved. + let (patch_handles, page_handles) = futures::try_join!(patch_handles, page_handles)?; let mut handles = empty_handles(plan.descriptors.len()); for (index, handle) in patch_handles { handles[index] = handle; @@ -644,72 +669,53 @@ async fn resolve_alprd_pages( None, )?; - let mut page_futures = Vec::new(); - for page in pages { - let local_mask = page_mask(&page.rows, row_range, mask.indices())?; - if local_mask.all_false() { - continue; - } - let plan = plan.clone(); - let dtype = dtype.clone(); - let page_patches = full_patches.clone(); - page_futures.push( - async move { - let left = page - .left - .await? - .ensure_aligned(plan.left.descriptor.alignment())?; - let right = page - .right - .await? - .ensure_aligned(plan.right.descriptor.alignment())?; - let inner_start = page.rows.start * plan.list_size as usize; - let inner_end = page.rows.end * plan.list_size as usize; - let inner_len = inner_end - inner_start; - let left = BitPacked::try_new( - left, - plan.left.ptype, - Validity::from(plan.left_parts_dtype.nullability()), - None, - plan.left.bit_width, - inner_len, - 0, - )? - .into_array(); - let right = BitPacked::try_new( - right, - plan.right.ptype, - Validity::NonNullable, - None, - plan.right.bit_width, - inner_len, - 0, - )? - .into_array(); - let patches = page_patches.slice(inner_start..inner_end)?; - let elements = ALPRD::try_new( - plan.element_dtype.clone(), - left, - plan.left_parts_dictionary.clone(), - right, - plan.right_bit_width, - patches, - )? - .into_array(); - let array = FixedSizeListArray::try_new( - elements, - plan.list_size, - Validity::from(dtype.nullability()), - page.rows.len(), - )? - .into_array(); - clear_stats(&array); - apply_page_mask(array, local_mask) - } - .boxed(), - ); - } - try_join_all(page_futures).await + page_handles + .into_iter() + .map(|(rows, local_mask, left, right)| { + let inner_start = rows.start * plan.list_size as usize; + let inner_end = rows.end * plan.list_size as usize; + let inner_len = inner_end - inner_start; + let left = BitPacked::try_new( + left, + plan.left.ptype, + Validity::from(plan.left_parts_dtype.nullability()), + None, + plan.left.bit_width, + inner_len, + 0, + )? + .into_array(); + let right = BitPacked::try_new( + right, + plan.right.ptype, + Validity::NonNullable, + None, + plan.right.bit_width, + inner_len, + 0, + )? + .into_array(); + let patches = full_patches.slice(inner_start..inner_end)?; + let elements = ALPRD::try_new( + plan.element_dtype.clone(), + left, + plan.left_parts_dictionary.clone(), + right, + plan.right_bit_width, + patches, + )? + .into_array(); + let array = FixedSizeListArray::try_new( + elements, + plan.list_size, + Validity::from(dtype.nullability()), + rows.len(), + )? + .into_array(); + clear_stats(&array); + apply_page_mask(array, local_mask) + }) + .collect() } #[allow(clippy::too_many_arguments)] From 40bc76d2a366258249a57f0aa6b9926f004b8614 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Wed, 12 Aug 2026 22:36:44 +0000 Subject: [PATCH 16/17] Keep Flat partial reads to one I/O round Signed-off-by: Joe Isaacs --- vortex-layout/src/layouts/flat/partial.rs | 322 +--------------------- vortex-layout/src/layouts/flat/reader.rs | 42 --- 2 files changed, 1 insertion(+), 363 deletions(-) diff --git a/vortex-layout/src/layouts/flat/partial.rs b/vortex-layout/src/layouts/flat/partial.rs index 865eb7b126c..b4d8df6557b 100644 --- a/vortex-layout/src/layouts/flat/partial.rs +++ b/vortex-layout/src/layouts/flat/partial.rs @@ -16,22 +16,15 @@ use vortex_array::IntoArray; use vortex_array::VTable; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ChunkedArray; -use vortex_array::arrays::ConstantArray; use vortex_array::arrays::FixedSizeList; use vortex_array::arrays::FixedSizeListArray; -use vortex_array::arrays::List; -use vortex_array::arrays::ListArray; use vortex_array::arrays::Primitive; use vortex_array::arrays::Struct; -use vortex_array::arrays::list::ListMetadata; use vortex_array::buffer::BufferHandle; -use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; use vortex_array::expr::stats::Stat; use vortex_array::patches::Patches; use vortex_array::patches::PatchesMetadata; -use vortex_array::scalar_fn::fns::operators::Operator; use vortex_array::serde::SerializedArray; use vortex_array::serde::SerializedBuffer; use vortex_array::validity::Validity; @@ -63,7 +56,6 @@ pub(super) struct PartialReadPlan { enum PartialReadKind { Fixed(Arc<[PlannedBuffer]>), Alprd(Box), - List(ListReadPlan), } #[derive(Clone)] @@ -74,15 +66,6 @@ struct PlannedBuffer { bytes_per_granule: usize, } -#[derive(Clone)] -struct ListReadPlan { - descriptors: Arc<[SerializedBuffer]>, - element_buffer: SerializedBuffer, - bytes_per_element: usize, - offset_buffers: Arc<[SerializedBuffer]>, - offset_dtype: DType, -} - #[derive(Clone)] struct ALPRDReadPlan { descriptors: Arc<[SerializedBuffer]>, @@ -129,14 +112,6 @@ enum RegisteredReadKind { patch_buffers: Vec<(SegmentFuture, SerializedBuffer)>, plan: ALPRDReadPlan, }, - List { - pages: Vec>, - offset_buffers: Vec<(SegmentFuture, SerializedBuffer)>, - plan: ListReadPlan, - source: Arc, - segment_id: SegmentId, - layout_len: usize, - }, } struct RegisteredALPRDPage { @@ -178,21 +153,6 @@ impl PartialReadPlan { })); } - if let Some((plan, bytes_per_row)) = try_list_plan( - &serialized, - layout.dtype(), - layout.array_ctx(), - row_count, - Arc::clone(&descriptors), - )? { - return Ok(Some(Self { - array_tree, - bytes_per_row, - row_granularity: 1, - kind: PartialReadKind::List(plan), - })); - } - let mut planned = Vec::new(); if !collect_raw_buffers( &serialized, @@ -386,39 +346,6 @@ impl PartialReadPlan { plan: plan.as_ref().clone(), } } - PartialReadKind::List(plan) => { - let offset_specs = plan - .offset_buffers - .iter() - .map(|descriptor| { - Some(( - u64::try_from(descriptor.range().start).ok()? - ..u64::try_from(descriptor.range().end).ok()?, - descriptor.clone(), - )) - }) - .collect::>>()?; - let offset_requests = source.request_ranges( - segment_id, - offset_specs - .iter() - .map(|(range, _)| range.clone()) - .collect(), - ); - let offset_buffers = offset_requests - .into_iter() - .zip(offset_specs) - .map(|(request, (_, descriptor))| (request, descriptor)) - .collect(); - RegisteredReadKind::List { - pages, - offset_buffers, - plan: plan.clone(), - source: Arc::clone(source), - segment_id, - layout_len, - } - } }; Some(RegisteredPartialRead { @@ -430,7 +357,7 @@ impl PartialReadPlan { fn estimated_partial_io( &self, pages: &[Range], - layout_len: usize, + _layout_len: usize, ) -> Option<(usize, usize)> { match &self.kind { PartialReadKind::Fixed(buffers) => { @@ -468,27 +395,6 @@ impl PartialReadPlan { .checked_add(plan.patch_buffers.len())?, )) } - PartialReadKind::List(plan) => { - let offset_bytes = plan - .offset_buffers - .iter() - .try_fold(0usize, |total, buffer| { - total.checked_add(buffer.range().len()) - })?; - let covered_rows = pages - .iter() - .try_fold(0usize, |total, rows| total.checked_add(rows.len()))?; - let estimated_element_bytes = plan - .element_buffer - .range() - .len() - .checked_mul(covered_rows)? - .div_ceil(layout_len.max(1)); - Some(( - offset_bytes.checked_add(estimated_element_bytes)?, - plan.offset_buffers.len().checked_add(pages.len())?, - )) - } } } } @@ -535,30 +441,6 @@ impl RegisteredPartialRead { ) .await? } - RegisteredReadKind::List { - pages, - offset_buffers, - plan, - source, - segment_id, - layout_len, - } => { - resolve_list_pages( - self.array_tree, - pages, - offset_buffers, - plan, - source, - segment_id, - layout_len, - dtype, - row_range, - mask, - ctx, - session, - ) - .await? - } }; finish_chunks(chunks, dtype, session) } @@ -718,119 +600,6 @@ async fn resolve_alprd_pages( .collect() } -#[allow(clippy::too_many_arguments)] -async fn resolve_list_pages( - array_tree: ByteBuffer, - pages: Vec>, - offset_requests: Vec<(SegmentFuture, SerializedBuffer)>, - plan: ListReadPlan, - source: Arc, - segment_id: SegmentId, - layout_len: usize, - dtype: &DType, - row_range: &Range, - mask: &Mask, - ctx: &ReadContext, - session: &VortexSession, -) -> VortexResult> { - if pages.is_empty() { - return Ok(Vec::new()); - } - - let resolved_offsets = try_join_all(offset_requests.into_iter().map( - |(future, descriptor)| async move { - Ok::<_, vortex_error::VortexError>(( - descriptor.index(), - future.await?.ensure_aligned(descriptor.alignment())?, - )) - }, - )) - .await?; - let mut offset_handles = empty_handles(plan.descriptors.len()); - for (index, handle) in resolved_offsets { - offset_handles[index] = handle; - } - let offsets = - SerializedArray::from_flatbuffer_with_buffers(array_tree.clone(), offset_handles)? - .child(1) - .decode(&plan.offset_dtype, layout_len + 1, ctx, session)?; - - let mut exec = session.create_execution_ctx(); - let mut page_specs = Vec::new(); - for rows in pages { - let local_mask = page_mask(&rows, row_range, mask.indices())?; - if local_mask.all_false() { - continue; - } - let first = offset_at(&offsets, rows.start, &mut exec)?; - let last = offset_at(&offsets, rows.end, &mut exec)?; - let start = plan.element_buffer.range().start - + first - .checked_mul(plan.bytes_per_element) - .ok_or_else(|| vortex_err!("List element offset overflow"))?; - let end = plan.element_buffer.range().start - + last - .checked_mul(plan.bytes_per_element) - .ok_or_else(|| vortex_err!("List element offset overflow"))?; - page_specs.push(( - rows, - local_mask, - first, - last, - u64::try_from(start)?..u64::try_from(end)?, - )); - } - let element_requests = source.request_ranges( - segment_id, - page_specs - .iter() - .map(|(_, _, _, _, range)| range.clone()) - .collect(), - ); - let mut page_futures = Vec::new(); - for ((rows, local_mask, first, last, _), elements_future) in - page_specs.into_iter().zip(element_requests) - { - let array_tree = array_tree.clone(); - let offsets = offsets.clone(); - let plan = plan.clone(); - let dtype = dtype.clone(); - let ctx = ctx.clone(); - let session = session.clone(); - page_futures.push( - async move { - let elements_handle = elements_future - .await? - .ensure_aligned(plan.element_buffer.alignment())?; - let mut handles = empty_handles(plan.descriptors.len()); - handles[plan.element_buffer.index()] = elements_handle; - let element_dtype = match &dtype { - DType::List(element_dtype, _) => element_dtype.as_ref(), - _ => return Err(vortex_err!("List partial plan used with non-list dtype")), - }; - let elements = SerializedArray::from_flatbuffer_with_buffers(array_tree, handles)? - .child(0) - .decode(element_dtype, last - first, &ctx, &session)?; - clear_stats(&elements); - - let page_offsets = offsets.slice(rows.start..rows.end + 1)?; - let mut exec = session.create_execution_ctx(); - let first_offset = page_offsets.execute_scalar(0, &mut exec)?; - let adjusted_offsets = page_offsets.binary( - ConstantArray::new(first_offset, page_offsets.len()).into_array(), - Operator::Sub, - )?; - let validity = Validity::from(dtype.nullability()); - let array = ListArray::try_new(elements, adjusted_offsets, validity)?.into_array(); - clear_stats(&array); - apply_page_mask(array, local_mask) - } - .boxed(), - ); - } - try_join_all(page_futures).await -} - fn finish_chunks( mut chunks: Vec, dtype: &DType, @@ -871,18 +640,6 @@ fn empty_handles(len: usize) -> Vec { .collect() } -fn offset_at( - offsets: &ArrayRef, - index: usize, - ctx: &mut vortex_array::ExecutionCtx, -) -> VortexResult { - offsets - .execute_scalar(index, ctx)? - .as_primitive() - .as_::() - .ok_or_else(|| vortex_err!("List offset does not fit usize")) -} - fn selected_pages( page_rows: usize, layout_len: usize, @@ -1110,83 +867,6 @@ fn bitpacked_range(plan: &BitPackedReadPlan, values: Range) -> Option, -) -> VortexResult> { - if ctx.resolve(node.encoding_id()) != Some(List.id()) { - return Ok(None); - } - let DType::List(element_dtype, _) = dtype else { - return Ok(None); - }; - if node.nbuffers() != 0 || node.nchildren() != 2 { - return Ok(None); - } - let elements = node.child(0); - let DType::Primitive(element_ptype, _) = element_dtype.as_ref() else { - return Ok(None); - }; - if ctx.resolve(elements.encoding_id()) != Some(Primitive.id()) - || elements.nchildren() != 0 - || elements.buffer_indices().len() != 1 - { - return Ok(None); - } - - let metadata = ListMetadata::decode(node.metadata())?; - let element_index = elements.buffer_indices()[0]; - let Some(element_buffer) = descriptors.get(element_index).cloned() else { - return Ok(None); - }; - let elements_len = usize::try_from(metadata.elements_len())?; - let expected_elements_len = elements_len - .checked_mul(element_ptype.byte_width()) - .ok_or_else(|| vortex_err!("List elements length overflow"))?; - if element_buffer.range().len() != expected_elements_len { - return Ok(None); - } - - let offsets = node.child(1); - let mut offset_indices = BTreeSet::new(); - collect_buffer_indices(&offsets, &mut offset_indices); - if offset_indices.is_empty() || offset_indices.contains(&element_index) { - return Ok(None); - } - let Some(offset_buffers) = offset_indices - .into_iter() - .map(|index| descriptors.get(index).cloned()) - .collect::>>() - else { - return Ok(None); - }; - let offset_dtype = DType::Primitive(metadata.offset_ptype(), Nullability::NonNullable); - let offset_bytes: usize = offset_buffers - .iter() - .map(|buffer| buffer.range().len()) - .sum(); - let bytes_per_row = element_buffer - .range() - .len() - .saturating_add(offset_bytes) - .div_ceil(row_count.max(1)) - .max(1); - - Ok(Some(( - ListReadPlan { - descriptors, - element_buffer, - bytes_per_element: element_ptype.byte_width(), - offset_buffers: offset_buffers.into(), - offset_dtype, - }, - bytes_per_row, - ))) -} - fn collect_buffer_indices(node: &SerializedArray, output: &mut BTreeSet) { output.extend(node.buffer_indices()); for index in 0..node.nchildren() { diff --git a/vortex-layout/src/layouts/flat/reader.rs b/vortex-layout/src/layouts/flat/reader.rs index 9f0709d12ad..f3d89f78ef4 100644 --- a/vortex-layout/src/layouts/flat/reader.rs +++ b/vortex-layout/src/layouts/flat/reader.rs @@ -388,7 +388,6 @@ mod test { use vortex_array::MaskFuture; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; - use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; use vortex_array::expr::gt; @@ -706,45 +705,4 @@ mod test { Ok(()) }) } - - #[test] - fn sparse_list_projection_reads_offsets_then_selected_elements() -> VortexResult<()> { - block_on(|handle| async { - let session = new_session().with_handle(handle); - let mut ctx = session.create_execution_ctx(); - let array_ctx = ArrayContext::empty(); - let source = RangedTestSource::default(); - let (ptr, eof) = SequenceId::root().split(); - let elements = PrimitiveArray::from_iter(0i64..32).into_array(); - let offsets = PrimitiveArray::from_iter((0u32..=16).map(|v| v * 2)).into_array(); - let array = ListArray::try_new(elements, offsets, Validity::NonNullable)?.into_array(); - let layout = FlatLayoutStrategy::default() - .write_stream( - array_ctx.into(), - Arc::::clone(&source.inner), - array.to_array_stream().sequenced(ptr), - eof, - &session, - ) - .await?; - - let reader = layout.new_reader( - "".into(), - Arc::new(source.clone()), - &session, - &Default::default(), - )?; - let expr = root().bind(reader.dtype())?; - let mask = Mask::from_indices(16, [1, 10]); - let result = reader - .projection_evaluation(&(0..16), &expr, MaskFuture::ready(mask.clone()))? - .await?; - - let expected = array.filter(mask)?; - assert_arrays_eq!(result, expected, &mut ctx); - assert_eq!(source.whole_requests.load(Ordering::Relaxed), 0); - assert!(source.ranges.lock().len() >= 3); - Ok(()) - }) - } } From 79b2daafcd08de91e56786fe2a730d4cccd083cc Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 13 Aug 2026 10:03:01 +0000 Subject: [PATCH 17/17] Canonicalize ALPRD patch indices once Signed-off-by: Joe Isaacs --- vortex-layout/src/layouts/flat/partial.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/vortex-layout/src/layouts/flat/partial.rs b/vortex-layout/src/layouts/flat/partial.rs index b4d8df6557b..4f5ad9b3701 100644 --- a/vortex-layout/src/layouts/flat/partial.rs +++ b/vortex-layout/src/layouts/flat/partial.rs @@ -19,6 +19,7 @@ use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::FixedSizeList; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::Primitive; +use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::Struct; use vortex_array::buffer::BufferHandle; use vortex_array::dtype::DType; @@ -534,6 +535,9 @@ async fn resolve_alprd_pages( alprd .child(2) .decode(&plan.patch_indices_dtype, patch_len, ctx, session)?; + let patch_indices = patch_indices + .execute::(&mut session.create_execution_ctx())? + .into_array(); let patch_values = alprd.child(3).decode( &plan.left_parts_dtype.as_nonnullable(), patch_len,