try read_ranges IO

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 }

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<T: VortexReadAt + Clone> VortexReadAt for CopyDeviceReadAt<T> {
         self.read.size()
     }
 
+    fn read_ranges(
+        &self,
+        ranges: Vec<ReadOp>,
+    ) -> BoxFuture<'static, VortexResult<Vec<BufferHandle>>> {
+        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,

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<ReadOp>,
+    ) -> BoxFuture<'static, VortexResult<Vec<BufferHandle>>> {
+        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<ReadOp>) -> BoxFuture<'static, VortexResult<Vec<BufferHandle>>> {
+        let read_futures = ops
+            .into_iter()
+            .map(|op| self.read_at(op.offset, op.length, op.alignment))
+            .collect::<Vec<_>>();
+        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<ReadOp>) -> BoxFuture<'static, VortexResult<Vec<BufferHandle>>> {
+        let read_futures = ops
+            .into_iter()
+            .map(|op| self.read_at(op.offset, op.length, op.alignment))
+            .collect::<Vec<_>>();
+        future::try_join_all(read_futures).boxed()
+    }
+
     fn read_at(
         &self,
         offset: u64,

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<BufferHandle>> {
-            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<ReadOp>,
+        ) -> BoxFuture<'static, VortexResult<Vec<BufferHandle>>> {
+            let total = ops.iter().map(|op| op.length).sum::<usize>();
+            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 {

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;
@@ -35,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 {
@@ -107,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),
@@ -118,28 +123,51 @@ impl FileSegmentSource {
 
         let drive_fut = async move {
             stream
-                .map(move |req| {
+                .ready_chunks(read_batch_size)
+                .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(read_batch_concurrency)
                 .collect::<()>()
                 .await
         };

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<R: VortexReadAt> VortexReadAt for Compat<R> {
         Compat::new(self.inner().size()).boxed()
     }
 
+    fn read_ranges(
+        &self,
+        ranges: Vec<ReadOp>,
+    ) -> BoxFuture<'static, VortexResult<Vec<BufferHandle>>> {
+        Compat::new(self.inner().read_ranges(ranges)).boxed()
+    }
+
     fn read_at(
         &self,
         offset: u64,