From 65fd5a941de58546c2ee90f51684924b541a8555 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Tue, 7 Jul 2026 14:13:40 +0100 Subject: [PATCH 1/4] try read_ranges IO Signed-off-by: Adam Gutglick --- Cargo.lock | 1 - vortex-array/Cargo.toml | 1 - vortex-cuda/src/device_read_at.rs | 23 ++++ vortex-cuda/src/pooled_read_at.rs | 58 +++++++++ vortex-file/src/open.rs | 28 ++-- vortex-file/src/segments/source.rs | 60 ++++++--- vortex-io/src/compat/read_at.rs | 8 ++ vortex-io/src/object_store/read_at.rs | 131 ++++++++++--------- vortex-io/src/read_at.rs | 180 +++++++++++++++++++------- vortex-io/src/runtime/tests.rs | 34 ++--- vortex-io/src/std_file/read_at.rs | 100 ++++++++++++++ vortex-web/crate/src/wasm.rs | 10 ++ 12 files changed, 471 insertions(+), 163 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d4a9f942e97..d1848943e71 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9803,7 +9803,6 @@ dependencies = [ "arbitrary", "arc-swap", "arcref", - "arrow-arith 58.3.0", "arrow-array 58.3.0", "arrow-buffer 58.3.0", "arrow-cast 58.3.0", diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index 4d8ef4ea2bb..78070c8f82c 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -23,7 +23,6 @@ workspace = true arbitrary = { workspace = true, optional = true } arc-swap = { workspace = true } arcref = { workspace = true } -arrow-arith = { workspace = true } arrow-array = { workspace = true, features = ["ffi"] } arrow-buffer = { workspace = true } arrow-cast = { workspace = true } diff --git a/vortex-cuda/src/device_read_at.rs b/vortex-cuda/src/device_read_at.rs index 47500935da5..58ba43946db 100644 --- a/vortex-cuda/src/device_read_at.rs +++ b/vortex-cuda/src/device_read_at.rs @@ -9,6 +9,7 @@ use vortex::array::buffer::BufferHandle; use vortex::buffer::Alignment; use vortex::error::VortexResult; use vortex::io::CoalesceConfig; +use vortex::io::ReadOp; use vortex::io::VortexReadAt; use crate::stream::VortexCudaStream; @@ -43,6 +44,28 @@ impl VortexReadAt for CopyDeviceReadAt { self.read.size() } + fn read_ranges( + &self, + ranges: Vec, + ) -> BoxFuture<'static, VortexResult>> { + let read = self.read.clone(); + let stream = self.stream.clone(); + async move { + let handles = read.read_ranges(ranges).await?; + let mut buffers = Vec::with_capacity(handles.len()); + for handle in handles { + if handle.is_on_device() { + buffers.push(handle); + } else { + let host_buffer = handle.as_host().clone(); + buffers.push(stream.copy_to_device(host_buffer)?.await?); + } + } + Ok(buffers) + } + .boxed() + } + fn read_at( &self, offset: u64, diff --git a/vortex-cuda/src/pooled_read_at.rs b/vortex-cuda/src/pooled_read_at.rs index 1b5ec4fb31a..ef11fd37691 100644 --- a/vortex-cuda/src/pooled_read_at.rs +++ b/vortex-cuda/src/pooled_read_at.rs @@ -8,6 +8,7 @@ use std::sync::Arc; use futures::FutureExt; use futures::StreamExt; +use futures::future; use futures::future::BoxFuture; use object_store::GetOptions; use object_store::GetRange; @@ -23,6 +24,7 @@ use vortex::error::VortexResult; use vortex::error::vortex_ensure; use vortex::error::vortex_err; use vortex::io::CoalesceConfig; +use vortex::io::ReadOp; use vortex::io::VortexReadAt; use vortex::io::runtime::Handle; use vortex::io::std_file::read_exact_at; @@ -118,6 +120,46 @@ impl VortexReadAt for PooledFileReadAt { } .boxed() } + + fn read_ranges( + &self, + ranges: Vec, + ) -> BoxFuture<'static, VortexResult>> { + if ranges.is_empty() { + return async { Ok(Vec::new()) }.boxed(); + } + + let file = Arc::clone(&self.file); + let handle = self.handle.clone(); + let stream = self.stream.clone(); + let pool = Arc::clone(&self.pool); + + async move { + let mut targets = Vec::with_capacity(ranges.len()); + for range in ranges { + targets.push((range.offset, pool.get(range.length)?)); + } + + let targets = handle + .spawn_blocking(move || { + let mut targets = targets; + for (offset, target) in &mut targets { + read_exact_at(&file, target.as_mut_slice(), *offset)?; + } + Ok::<_, io::Error>(targets) + }) + .await + .map_err(VortexError::from)?; + + let mut buffers = Vec::with_capacity(targets.len()); + for (_, target) in targets { + let cuda_buf = target.transfer_to_device(&stream)?; + buffers.push(BufferHandle::new_device(Arc::new(cuda_buf))); + } + Ok(buffers) + } + .boxed() + } } /// Object store reader that uses CUDA pinned host memory for I/O buffers and @@ -198,6 +240,14 @@ impl VortexReadAt for PooledObjectStoreReadAt { .boxed() } + fn read_ranges(&self, ops: Vec) -> BoxFuture<'static, VortexResult>> { + let read_futures = ops + .into_iter() + .map(|op| self.read_at(op.offset, op.length, op.alignment)) + .collect::>(); + future::try_join_all(read_futures).boxed() + } + fn read_at( &self, offset: u64, @@ -323,6 +373,14 @@ impl VortexReadAt for PooledByteBufferReadAt { async move { Ok(len) }.boxed() } + fn read_ranges(&self, ops: Vec) -> BoxFuture<'static, VortexResult>> { + let read_futures = ops + .into_iter() + .map(|op| self.read_at(op.offset, op.length, op.alignment)) + .collect::>(); + future::try_join_all(read_futures).boxed() + } + fn read_at( &self, offset: u64, diff --git a/vortex-file/src/open.rs b/vortex-file/src/open.rs index 1bd42c7be41..8a0ba7c7632 100644 --- a/vortex-file/src/open.rs +++ b/vortex-file/src/open.rs @@ -396,6 +396,7 @@ mod tests { use vortex_buffer::Buffer; use vortex_buffer::ByteBuffer; use vortex_buffer::ByteBufferMut; + use vortex_io::ReadOp; use vortex_io::session::RuntimeSession; use vortex_layout::session::LayoutSession; @@ -415,20 +416,21 @@ mod tests { self.inner.size() } - fn read_at( + fn read_ranges( &self, - offset: u64, - length: usize, - alignment: Alignment, - ) -> BoxFuture<'static, VortexResult> { - self.total_read.fetch_add(length, Ordering::Relaxed); - let _ = self.first_read_len.compare_exchange( - 0, - length, - Ordering::Relaxed, - Ordering::Relaxed, - ); - self.inner.read_at(offset, length, alignment) + ops: Vec, + ) -> BoxFuture<'static, VortexResult>> { + let total = ops.iter().map(|op| op.length).sum::(); + self.total_read.fetch_add(total, Ordering::Relaxed); + if let Some(first) = ops.first() { + let _ = self.first_read_len.compare_exchange( + 0, + first.length, + Ordering::Relaxed, + Ordering::Relaxed, + ); + } + self.inner.read_ranges(ops) } fn concurrency(&self) -> usize { diff --git a/vortex-file/src/segments/source.rs b/vortex-file/src/segments/source.rs index e7c9dc3b222..488a656c7c6 100644 --- a/vortex-file/src/segments/source.rs +++ b/vortex-file/src/segments/source.rs @@ -15,10 +15,11 @@ use futures::future; 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::ReadOp; use vortex_io::VortexReadAt; use vortex_io::runtime::Handle; use vortex_layout::segments::SegmentFuture; @@ -118,28 +119,51 @@ impl FileSegmentSource { let drive_fut = async move { stream - .map(move |req| { + .ready_chunks(concurrency) + .map(move |requests| { 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 - ) - } - Ok(buffer) - }); + let ranges = requests + .iter() + .map(|req| ReadOp::aligned(req.offset(), req.len(), req.alignment())) + .collect(); - req.resolve(result); + match reader.read_ranges(ranges).await { + Ok(buffers) if buffers.len() == requests.len() => { + for (req, buffer) in requests.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(buffers) => { + let err = Arc::new(vortex_err!( + "FileSegmentSource: expected {} buffers but received {}", + requests.len(), + buffers.len() + )); + for req in requests { + req.resolve(Err(VortexError::from(Arc::clone(&err)))); + } + } + Err(e) => { + let err = Arc::new(e); + for req in requests { + req.resolve(Err(VortexError::from(Arc::clone(&err)))); + } + } + } } }) - .buffer_unordered(concurrency) + .buffer_unordered(1) .collect::<()>() .await }; diff --git a/vortex-io/src/compat/read_at.rs b/vortex-io/src/compat/read_at.rs index 4fc49785d28..e3838486cc7 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::ReadOp; use crate::VortexReadAt; use crate::compat::Compat; @@ -32,6 +33,13 @@ impl VortexReadAt for Compat { Compat::new(self.inner().size()).boxed() } + fn read_ranges( + &self, + ranges: Vec, + ) -> BoxFuture<'static, VortexResult>> { + Compat::new(self.inner().read_ranges(ranges)).boxed() + } + fn read_at( &self, offset: u64, diff --git a/vortex-io/src/object_store/read_at.rs b/vortex-io/src/object_store/read_at.rs index 086d70c1bcf..57a250f4ada 100644 --- a/vortex-io/src/object_store/read_at.rs +++ b/vortex-io/src/object_store/read_at.rs @@ -6,6 +6,7 @@ use std::sync::Arc; use futures::FutureExt; use futures::StreamExt; +use futures::future; use futures::future::BoxFuture; use object_store::GetOptions; use object_store::GetRange; @@ -16,12 +17,12 @@ use object_store::path::Path as ObjectPath; use vortex_array::buffer::BufferHandle; use vortex_array::memory::DefaultHostAllocator; use vortex_array::memory::HostAllocatorRef; -use vortex_buffer::Alignment; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use crate::CoalesceConfig; +use crate::ReadOp; use crate::VortexReadAt; use crate::runtime::Handle; #[cfg(not(target_arch = "wasm32"))] @@ -105,79 +106,82 @@ impl VortexReadAt for ObjectStoreReadAt { .boxed() } - fn read_at( - &self, - offset: u64, - length: usize, - alignment: Alignment, - ) -> BoxFuture<'static, VortexResult> { + fn read_ranges(&self, ops: Vec) -> 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); - let range = offset..(offset + length as u64); // Requires to deal with borrowed lifetimes 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)? + future::try_join_all(ops.into_iter().map(|op| { + let store = Arc::clone(&store); + let path = path.clone(); + let io_handle = io_handle.clone(); + let allocator = Arc::clone(&allocator); + async move { + let range = op.byte_range()?; + let mut buffer = allocator.allocate(op.length, op.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 <= op.length, + "Object store stream returned too many bytes: {} > expected {} (range: {:?})", + end, + op.length, + range + ); + buffer.as_mut_slice()[written..end].copy_from_slice(&bytes); + written = end; + } + + vortex_ensure!( + written == op.length, + "Object store stream returned {} bytes but expected {} bytes (range: {:?})", + written, + op.length, + range + ); + + buffer + } + }; + + Ok(BufferHandle::new_host(buffer.freeze())) } - #[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())) + })).await }) .boxed() } @@ -191,6 +195,7 @@ mod tests { use object_store::PutPayload; use object_store::memory::InMemory; + use vortex_buffer::Alignment; use super::*; use crate::runtime::AbortHandle; diff --git a/vortex-io/src/read_at.rs b/vortex-io/src/read_at.rs index aa9a8a03abf..b5d08765687 100644 --- a/vortex-io/src/read_at.rs +++ b/vortex-io/src/read_at.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; @@ -11,6 +12,8 @@ use vortex_buffer::ByteBuffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; use vortex_metrics::Counter; use vortex_metrics::Histogram; use vortex_metrics::Label; @@ -49,6 +52,77 @@ impl CoalesceConfig { } } +/// A positional read request against a [`VortexReadAt`] source. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ReadOp { + /// Starting byte offset in the source. + pub offset: u64, + /// Number of bytes to read. + pub length: usize, + /// Alignment required by the returned buffer. + pub alignment: Alignment, +} + +impl ReadOp { + /// Create a new read operation with no alignment requirement. + pub const fn new(offset: u64, length: usize) -> Self { + Self::aligned(offset, length, Alignment::none()) + } + + /// Create a new read operation with an explicit alignment requirement. + pub const fn aligned(offset: u64, length: usize, alignment: Alignment) -> Self { + Self { + offset, + length, + alignment, + } + } + + /// Return a copy of this read operation with a different alignment. + pub const fn with_alignment(mut self, alignment: Alignment) -> Self { + self.alignment = alignment; + self + } + + /// Starting byte offset in the source. + pub const fn offset(&self) -> u64 { + self.offset + } + + /// Number of bytes to read. + pub const fn len(&self) -> usize { + self.length + } + + /// Whether this operation reads zero bytes. + pub const fn is_empty(&self) -> bool { + self.length == 0 + } + + /// Alignment required by the returned buffer. + pub const fn alignment(&self) -> Alignment { + self.alignment + } + + /// Exclusive end offset for this read operation. + pub fn end(&self) -> VortexResult { + let length = u64::try_from(self.length) + .map_err(|_| vortex_err!("ReadOp length exceeds u64::MAX: {}", self.length))?; + self.offset.checked_add(length).ok_or_else(|| { + vortex_err!( + "ReadOp range overflow: offset={}, length={}", + self.offset, + self.length + ) + }) + } + + /// Byte range covered by this read operation. + pub fn byte_range(&self) -> VortexResult> { + Ok(self.offset..self.end()?) + } +} + /// The unified read trait for Vortex I/O sources. /// /// This trait provides async positional reads to underlying storage and is used by the vortex-file @@ -64,11 +138,12 @@ pub trait VortexReadAt: Send + Sync + 'static { None } - /// Maximum number of concurrent I/O requests for that should be pulled from this source. + /// Maximum number of physical read operations the driver should pull for this source. /// - /// This value is used to control how many [`VortexReadAt::read_at`] calls can - /// be in-flight simultaneously. Higher values allow more parallelism but consume - /// more resources (memory, file descriptors, network connections). + /// This value controls the largest batch passed to [`VortexReadAt::read_ranges`]. + /// Implementations may execute the operations concurrently or serially depending on the + /// underlying storage system. Higher values allow more coalescing and internal parallelism but + /// consume more resources (memory, file descriptors, network connections). /// /// Implementations should choose a value appropriate for their underlying storage /// characteristics. Low-latency sources benefit less from high concurrency, while @@ -79,6 +154,11 @@ pub trait VortexReadAt: Send + Sync + 'static { /// Asynchronously get the number of bytes of the underlying source. fn size(&self) -> BoxFuture<'static, VortexResult>; + /// Request asynchronous positional reads. + /// + /// Results must be returned in the same order as `ops`. + fn read_ranges(&self, ops: Vec) -> BoxFuture<'static, VortexResult>>; + /// Request an asynchronous positional read. Results will be returned as a [`BufferHandle`]. /// /// If the reader does not have the requested number of bytes, the returned Future will complete @@ -88,7 +168,21 @@ pub trait VortexReadAt: Send + Sync + 'static { offset: u64, length: usize, alignment: Alignment, - ) -> BoxFuture<'static, VortexResult>; + ) -> BoxFuture<'static, VortexResult> { + let read_fut = self.read_ranges(vec![ReadOp::aligned(offset, length, alignment)]); + async move { + let mut buffers = read_fut.await?; + vortex_ensure!( + buffers.len() == 1, + "VortexReadAt::read_ranges returned {} buffers for one read operation", + buffers.len() + ); + Ok(buffers + .pop() + .vortex_expect("single read operation returns one buffer")) + } + .boxed() + } } impl VortexReadAt for Arc { @@ -108,13 +202,8 @@ impl VortexReadAt for Arc { self.as_ref().size() } - fn read_at( - &self, - offset: u64, - length: usize, - alignment: Alignment, - ) -> BoxFuture<'static, VortexResult> { - self.as_ref().read_at(offset, length, alignment) + fn read_ranges(&self, ops: Vec) -> BoxFuture<'static, VortexResult>> { + self.as_ref().read_ranges(ops) } } @@ -135,13 +224,8 @@ impl VortexReadAt for Arc { self.as_ref().size() } - fn read_at( - &self, - offset: u64, - length: usize, - alignment: Alignment, - ) -> BoxFuture<'static, VortexResult> { - self.as_ref().read_at(offset, length, alignment) + fn read_ranges(&self, ops: Vec) -> BoxFuture<'static, VortexResult>> { + self.as_ref().read_ranges(ops) } } @@ -155,28 +239,26 @@ impl VortexReadAt for ByteBuffer { 16 } - fn read_at( - &self, - offset: u64, - length: usize, - alignment: Alignment, - ) -> BoxFuture<'static, VortexResult> { + fn read_ranges(&self, ops: Vec) -> BoxFuture<'static, VortexResult>> { let buffer = self.clone(); async move { - let start = usize::try_from(offset).vortex_expect("start too big for usize"); - let end = - usize::try_from(offset + length as u64).vortex_expect("end too big for usize"); - if end > buffer.len() { - vortex_bail!( - "Requested range {}..{} out of bounds for buffer of length {}", - start, - end, - buffer.len() - ); - } - Ok(BufferHandle::new_host( - buffer.slice_unaligned(start..end).aligned(alignment), - )) + ops.into_iter() + .map(|op| { + let start = usize::try_from(op.offset).vortex_expect("start too big for usize"); + let end = usize::try_from(op.end()?).vortex_expect("end too big for usize"); + if end > buffer.len() { + vortex_bail!( + "Requested range {}..{} out of bounds for buffer of length {}", + start, + end, + buffer.len() + ); + } + Ok(BufferHandle::new_host( + buffer.slice_unaligned(start..end).aligned(op.alignment), + )) + }) + .collect() } .boxed() } @@ -296,23 +378,21 @@ impl VortexReadAt for InstrumentedReadAt { self.read.size() } - fn read_at( - &self, - offset: u64, - length: usize, - alignment: Alignment, - ) -> BoxFuture<'static, VortexResult> { + fn read_ranges(&self, ops: Vec) -> BoxFuture<'static, VortexResult>> { let durations = self.metrics.durations.clone(); let sizes = self.metrics.sizes.clone(); let total_size = self.metrics.total_size.clone(); + let lengths = ops.iter().map(|op| op.length).collect::>(); - let read_fut = self.read.read_at(offset, length, alignment); + let read_fut = self.read.read_ranges(ops); async move { let _timer = durations.time(); - let buf = read_fut.await; - sizes.update(length as f64); - total_size.add(length as u64); - buf + let buffers = read_fut.await; + for length in lengths { + sizes.update(length as f64); + total_size.add(length as u64); + } + buffers } .boxed() } diff --git a/vortex-io/src/runtime/tests.rs b/vortex-io/src/runtime/tests.rs index e35d9db6f78..f23abda44a4 100644 --- a/vortex-io/src/runtime/tests.rs +++ b/vortex-io/src/runtime/tests.rs @@ -17,6 +17,7 @@ use vortex_buffer::ByteBuffer; use vortex_buffer::ByteBufferMut; use vortex_error::VortexResult; +use crate::ReadOp; use crate::VortexReadAt; use crate::runtime::single::block_on; use crate::runtime::tokio::TokioRuntime; @@ -248,25 +249,24 @@ impl VortexReadAt for CountingReadAt { 16 } - fn read_at( - &self, - offset: u64, - length: usize, - alignment: Alignment, - ) -> BoxFuture<'static, VortexResult> { - self.read_count.fetch_add(1, Ordering::SeqCst); + fn read_ranges(&self, ops: Vec) -> BoxFuture<'static, VortexResult>> { + self.read_count.fetch_add(ops.len(), Ordering::SeqCst); let data = self.data.clone(); async move { - let start = offset as usize; - if start + length > data.len() { - return Err(vortex_error::vortex_err!("Read out of bounds")); - } - let mut buffer = ByteBufferMut::with_capacity_aligned(length, alignment); - unsafe { buffer.set_len(length) }; - buffer - .as_mut_slice() - .copy_from_slice(&data.as_slice()[start..start + length]); - Ok(BufferHandle::new_host(buffer.freeze())) + ops.into_iter() + .map(|op| { + let start = op.offset as usize; + if start + op.length > data.len() { + return Err(vortex_error::vortex_err!("Read out of bounds")); + } + let mut buffer = ByteBufferMut::with_capacity_aligned(op.length, op.alignment); + unsafe { buffer.set_len(op.length) }; + buffer + .as_mut_slice() + .copy_from_slice(&data.as_slice()[start..start + op.length]); + Ok(BufferHandle::new_host(buffer.freeze())) + }) + .collect() } .boxed() } diff --git a/vortex-io/src/std_file/read_at.rs b/vortex-io/src/std_file/read_at.rs index 3d59a595f70..c75393823b7 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::ReadOp; use crate::VortexReadAt; use crate::runtime::Handle; @@ -135,4 +136,103 @@ impl VortexReadAt for FileReadAt { } .boxed() } + + fn read_ranges( + &self, + ranges: Vec, + ) -> BoxFuture<'static, VortexResult>> { + if ranges.is_empty() { + return async { Ok(Vec::new()) }.boxed(); + } + + let file = Arc::clone(&self.file); + let handle = self.handle.clone(); + let allocator = Arc::clone(&self.allocator); + async move { + handle + .spawn_blocking(move || { + let mut buffers = Vec::with_capacity(ranges.len()); + for range in ranges { + let mut buffer = allocator.allocate(range.length, range.alignment)?; + read_exact_at(&file, buffer.as_mut_slice(), range.offset)?; + buffers.push(BufferHandle::new_host(buffer.freeze())); + } + Ok(buffers) + }) + .await + } + .boxed() + } +} + +#[cfg(test)] +mod tests { + use std::io::Write; + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + + use futures::future::BoxFuture; + use tempfile::NamedTempFile; + + use super::*; + use crate::runtime::AbortHandle; + use crate::runtime::AbortHandleRef; + use crate::runtime::Executor; + use crate::runtime::Handle; + + #[derive(Default)] + struct CountingExecutor { + blocking_count: AtomicUsize, + } + + impl Executor for CountingExecutor { + fn spawn(&self, fut: BoxFuture<'static, ()>) -> AbortHandleRef { + TokioAbortHandle::new_handle(tokio::spawn(fut).abort_handle()) + } + + fn spawn_cpu(&self, task: Box) -> AbortHandleRef { + TokioAbortHandle::new_handle(tokio::spawn(async move { task() }).abort_handle()) + } + + fn spawn_blocking_io(&self, task: Box) -> AbortHandleRef { + self.blocking_count.fetch_add(1, Ordering::SeqCst); + TokioAbortHandle::new_handle(tokio::task::spawn_blocking(task).abort_handle()) + } + } + + struct TokioAbortHandle(tokio::task::AbortHandle); + + impl TokioAbortHandle { + fn new_handle(handle: tokio::task::AbortHandle) -> AbortHandleRef { + Box::new(Self(handle)) + } + } + + impl AbortHandle for TokioAbortHandle { + fn abort(self: Box) { + self.0.abort(); + } + } + + #[tokio::test] + async fn read_ranges_uses_one_blocking_task() -> VortexResult<()> { + let mut file = NamedTempFile::new()?; + file.write_all(b"0123456789abcdef")?; + + let executor = Arc::new(CountingExecutor::default()); + let runtime = Arc::clone(&executor) as Arc; + let handle = Handle::new(Arc::downgrade(&runtime)); + let reader = FileReadAt::open(file.path(), handle)?; + + let buffers = reader + .read_ranges(vec![ReadOp::new(1, 4), ReadOp::new(10, 3)]) + .await?; + + assert_eq!(buffers[0].to_host().await.as_slice(), b"1234"); + assert_eq!(buffers[1].to_host().await.as_slice(), b"abc"); + assert_eq!(executor.blocking_count.load(Ordering::SeqCst), 1); + + Ok(()) + } } diff --git a/vortex-web/crate/src/wasm.rs b/vortex-web/crate/src/wasm.rs index 04abe8ee6d3..5da2d83e9d7 100644 --- a/vortex-web/crate/src/wasm.rs +++ b/vortex-web/crate/src/wasm.rs @@ -16,6 +16,7 @@ use arrow_schema::Field; use arrow_schema::Schema; use futures::FutureExt; use futures::TryStreamExt; +use futures::future; use futures::future::BoxFuture; use serde::Serialize; use vortex::array::ArrayRef; @@ -33,6 +34,7 @@ use vortex::file::OpenOptionsSessionExt; use vortex::file::VERSION; use vortex::file::VortexFile; use vortex::io::CoalesceConfig; +use vortex::io::ReadOp; use vortex::io::VortexReadAt; use vortex::layout::LayoutChildType; use vortex::layout::LayoutRef; @@ -92,6 +94,14 @@ impl VortexReadAt for BlobReadAt { async move { Ok(size) }.boxed() } + fn read_ranges(&self, ops: Vec) -> BoxFuture<'static, VortexResult>> { + let read_futures = ops + .into_iter() + .map(|op| self.read_at(op.offset, op.length, op.alignment)) + .collect::>(); + future::try_join_all(read_futures).boxed() + } + fn read_at( &self, offset: u64, From 840ec326a4c0b669a7609ea2650c754be5411205 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Tue, 7 Jul 2026 15:20:42 +0100 Subject: [PATCH 2/4] file/layout changes Signed-off-by: Adam Gutglick --- vortex-file/src/segments/source.rs | 8 ++++++-- vortex-layout/src/layouts/struct_/reader.rs | 15 +++++++++------ 2 files changed, 15 insertions(+), 8 deletions(-) diff --git a/vortex-file/src/segments/source.rs b/vortex-file/src/segments/source.rs index 488a656c7c6..fb4ed555f0f 100644 --- a/vortex-file/src/segments/source.rs +++ b/vortex-file/src/segments/source.rs @@ -36,6 +36,8 @@ use crate::read::IoRequestStream; use crate::read::ReadRequest; use crate::read::RequestId; +const READ_RANGES_BATCH_SIZE: usize = 16; + #[derive(Debug)] /// Events sent from segment futures to the coalescing read driver. pub enum ReadEvent { @@ -108,6 +110,8 @@ impl FileSegmentSource { reader.uri() ); } + let read_batch_size = concurrency.min(READ_RANGES_BATCH_SIZE); + let read_batch_concurrency = concurrency.div_ceil(read_batch_size); let stream = IoRequestStream::new( StreamExt::boxed(recv), @@ -119,7 +123,7 @@ impl FileSegmentSource { let drive_fut = async move { stream - .ready_chunks(concurrency) + .ready_chunks(read_batch_size) .map(move |requests| { let reader = reader.clone(); async move { @@ -163,7 +167,7 @@ impl FileSegmentSource { } } }) - .buffer_unordered(1) + .buffer_unordered(read_batch_concurrency) .collect::<()>() .await }; diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index ab03a4624fa..14eecdbc54a 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -157,15 +157,18 @@ impl StructReader { /// Utility for partitioning an expression over the fields of a struct. fn partition_expr(&self, expr: Expression) -> VortexResult { let key = ExactExpr(expr.clone()); - let binding = self - .partitioned_expr_cache - .entry(key) - .or_insert_with(|| Arc::new(OnceLock::new())); - let entry = binding.value(); - if let Some(value) = entry.get() { + if let Some(entry) = self.partitioned_expr_cache.get(&key) + && let Some(value) = entry.value().get() + { return Ok(value.clone()); } + let result = self.compute_partitioned_expr(expr)?; + let entry = self + .partitioned_expr_cache + .entry(key) + .or_insert_with(|| Arc::new(OnceLock::new())) + .clone(); let result = entry.get_or_init(|| result); Ok(result.clone()) } From a01c3910b422b335be56e66076eb523ae174016c Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Tue, 7 Jul 2026 15:27:34 +0100 Subject: [PATCH 3/4] better Signed-off-by: Adam Gutglick --- vortex-layout/src/layouts/struct_/reader.rs | 27 +++++++++------------ 1 file changed, 12 insertions(+), 15 deletions(-) diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index 14eecdbc54a..3526bf015ba 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -3,10 +3,10 @@ use std::ops::Range; use std::sync::Arc; -use std::sync::OnceLock; use futures::try_join; use itertools::Itertools; +use once_cell::sync::OnceCell; use vortex_array::ArrayRef; use vortex_array::IntoArray; use vortex_array::MaskFuture; @@ -59,7 +59,7 @@ pub struct StructReader { expanded_root_expr: Expression, field_lookup: Option>, - partitioned_expr_cache: DashMap>>, + partitioned_expr_cache: DashMap>>, } impl StructReader { @@ -157,19 +157,16 @@ impl StructReader { /// Utility for partitioning an expression over the fields of a struct. fn partition_expr(&self, expr: Expression) -> VortexResult { let key = ExactExpr(expr.clone()); - if let Some(entry) = self.partitioned_expr_cache.get(&key) - && let Some(value) = entry.value().get() - { - return Ok(value.clone()); - } - - let result = self.compute_partitioned_expr(expr)?; - let entry = self - .partitioned_expr_cache - .entry(key) - .or_insert_with(|| Arc::new(OnceLock::new())) - .clone(); - let result = entry.get_or_init(|| result); + let entry = match self.partitioned_expr_cache.get(&key) { + Some(entry) => Arc::clone(entry.value()), + None => Arc::clone( + self.partitioned_expr_cache + .entry(key) + .or_insert_with(|| Arc::new(OnceCell::new())) + .value(), + ), + }; + let result = entry.get_or_try_init(|| self.compute_partitioned_expr(expr))?; Ok(result.clone()) } From 235c03a12014e5ef0021fd8d3073b68cc21220d4 Mon Sep 17 00:00:00 2001 From: Adam Gutglick Date: Tue, 7 Jul 2026 15:54:05 +0100 Subject: [PATCH 4/4] minimize dashmap contention Signed-off-by: Adam Gutglick --- vortex-layout/src/layouts/dict/reader.rs | 20 ++-- vortex-layout/src/layouts/row_idx/mod.rs | 30 +++--- vortex-layout/src/layouts/zoned/pruning.rs | 103 ++++++++++++--------- 3 files changed, 85 insertions(+), 68 deletions(-) diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs index 489d8823d39..ff4431c99a9 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -53,7 +53,7 @@ pub struct DictReader { /// Cached dict values array values_array: OnceLock, /// Cache of expression evaluation results on the values array by expression - values_evals: DashMap, + values_evals: DashMap>>, values: LayoutReaderRef, codes: LayoutReaderRef, @@ -139,14 +139,18 @@ impl DictReader { // shouldn't. // TODO(joe): fixme - // Check cache first with read-only lock - if let Some(fut) = self.values_evals.get(&expr) { - return fut.clone(); - } + let entry = match self.values_evals.get(&expr) { + Some(entry) => Arc::clone(entry.value()), + None => Arc::clone( + self.values_evals + .entry(expr.clone()) + .or_insert_with(|| Arc::new(OnceLock::new())) + .value(), + ), + }; - self.values_evals - .entry(expr.clone()) - .or_insert_with(|| { + entry + .get_or_init(|| { self.values_array_uncanonical() .map(move |array| { let array = array?.apply(&expr)?; diff --git a/vortex-layout/src/layouts/row_idx/mod.rs b/vortex-layout/src/layouts/row_idx/mod.rs index 4e5e8e494fa..b330969b84e 100644 --- a/vortex-layout/src/layouts/row_idx/mod.rs +++ b/vortex-layout/src/layouts/row_idx/mod.rs @@ -8,12 +8,12 @@ use std::fmt::Formatter; use std::ops::BitAnd; use std::ops::Range; use std::sync::Arc; -use std::sync::OnceLock; use Nullability::NonNullable; pub use expr::*; use futures::FutureExt; use futures::future::BoxFuture; +use once_cell::sync::OnceCell; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::IntoArray; @@ -50,7 +50,7 @@ pub struct RowIdxLayoutReader { name: Arc, row_offset: u64, child: Arc, - partition_cache: DashMap>>, + partition_cache: DashMap>>, session: VortexSession, } @@ -67,22 +67,18 @@ impl RowIdxLayoutReader { fn partition_expr(&self, expr: &Expression) -> VortexResult { let key = ExactExpr(expr.clone()); + let entry = match self.partition_cache.get(&key) { + Some(entry) => Arc::clone(entry.value()), + None => Arc::clone( + self.partition_cache + .entry(key) + .or_insert_with(|| Arc::new(OnceCell::new())) + .value(), + ), + }; - // Check cache first with read-only lock. - if let Some(entry) = self.partition_cache.get(&key) - && let Some(partitioning) = entry.value().get() - { - return Ok(partitioning.clone()); - } - - let result = self.compute_partitioning(expr)?; - - self.partition_cache - .entry(key) - .or_insert_with(|| Arc::new(OnceLock::new())) - .get_or_init(|| result.clone()); - - Ok(result) + let partitioning = entry.get_or_try_init(|| self.compute_partitioning(expr))?; + Ok(partitioning.clone()) } fn compute_partitioning(&self, expr: &Expression) -> VortexResult { diff --git a/vortex-layout/src/layouts/zoned/pruning.rs b/vortex-layout/src/layouts/zoned/pruning.rs index b517df985c9..783f1ec2c3c 100644 --- a/vortex-layout/src/layouts/zoned/pruning.rs +++ b/vortex-layout/src/layouts/zoned/pruning.rs @@ -35,6 +35,7 @@ use crate::layouts::zoned::zone_map::ZoneMap; type SharedZoneMap = Shared>>; pub(super) type SharedPruningResult = Shared>>>; +type PruningResultCache = Arc>>; type PredicateCache = Arc>>; pub(super) struct PruningState { @@ -46,7 +47,7 @@ pub(super) struct PruningState { lazy_children: Arc, session: VortexSession, - pruning_result: LazyLock>>, + pruning_result: LazyLock>, zone_map: OnceLock, pruning_predicates: LazyLock>>, } @@ -73,53 +74,69 @@ impl PruningState { } pub(super) fn pruning_mask_future(&self, expr: Expression) -> Option { - if let Some(result) = self.pruning_result.get(&expr) { - return result.value().clone(); - } + let entry = match self.pruning_result.get(&expr) { + Some(entry) => Arc::clone(entry.value()), + None => Arc::clone( + self.pruning_result + .entry(expr.clone()) + .or_insert_with(|| Arc::new(OnceLock::new())) + .value(), + ), + }; - self.pruning_result - .entry(expr.clone()) - .or_insert_with(|| match self.pruning_predicate(expr.clone()) { - None => { - trace!(%expr, "no pruning predicate"); - None - } - Some(predicate) => { - trace!(%expr, ?predicate, "constructed pruning predicate"); - let zone_map = self.zone_map(); - let dynamic_updates = DynamicExprUpdates::new(&expr); - let session = self.session.clone(); - - Some( - async move { - let zone_map = zone_map.await?; - let initial_mask = - zone_map.prune(&predicate, &session).map_err(|err| { - err.with_context(format!( - "While evaluating pruning predicate {} (derived from {})", - predicate, expr - )) - })?; - Ok(Arc::new(PruningResult { - zone_map, - predicate, - dynamic_updates, - latest_result: RwLock::new((0, initial_mask)), - session, - })) - } - .boxed() - .shared(), - ) - } - }) + entry + .get_or_init(|| self.compute_pruning_mask_future(expr)) .clone() } + fn compute_pruning_mask_future(&self, expr: Expression) -> Option { + match self.pruning_predicate(expr.clone()) { + None => { + trace!(%expr, "no pruning predicate"); + None + } + Some(predicate) => { + trace!(%expr, ?predicate, "constructed pruning predicate"); + let zone_map = self.zone_map(); + let dynamic_updates = DynamicExprUpdates::new(&expr); + let session = self.session.clone(); + + Some( + async move { + let zone_map = zone_map.await?; + let initial_mask = zone_map.prune(&predicate, &session).map_err(|err| { + err.with_context(format!( + "While evaluating pruning predicate {} (derived from {})", + predicate, expr + )) + })?; + Ok(Arc::new(PruningResult { + zone_map, + predicate, + dynamic_updates, + latest_result: RwLock::new((0, initial_mask)), + session, + })) + } + .boxed() + .shared(), + ) + } + } + } + fn pruning_predicate(&self, expr: Expression) -> Option { - self.pruning_predicates - .entry(expr.clone()) - .or_default() + let entry = match self.pruning_predicates.get(&expr) { + Some(entry) => Arc::clone(entry.value()), + None => Arc::clone( + self.pruning_predicates + .entry(expr.clone()) + .or_insert_with(|| Arc::new(OnceLock::new())) + .value(), + ), + }; + + entry .get_or_init(move || match expr.falsify(&self.dtype, &self.session) { Ok(predicate) => predicate, Err(error) => {