Per VortexReadAt coalescing window (#2241)

Introduces a new `PerformanceHint` struct that can have more knobs in
the future, and configure it for each `VortexReadAt` implementation we
have.

Author: AdamGS

vortex-file/src/io/file.rs

@@ -63,6 +63,7 @@ impl<R: VortexReadAt> IoDriver for FileIoDriver<R> {
         stream: impl Stream<Item = SegmentRequest> + 'static,
     ) -> impl Stream<Item = VortexResult<()>> + 'static {
         // We map the segment requests to their respective locations within the file.
+        let coalescing_window = self.read.performance_hint().coalescing_window();
         let segment_map = self.file_layout.segment_map().clone();
         let stream = stream.filter_map(move |request| {
             let segment_map = segment_map.clone();
@@ -127,7 +128,9 @@ impl<R: VortexReadAt> IoDriver for FileIoDriver<R> {
             .inspect(|requests| log::debug!("Processing {} segment requests", requests.len()));
 
         // Coalesce the segment requests to minimize the number of I/O operations.
-        let stream = stream.map(coalesce).flat_map(stream::iter);
+        let stream = stream
+            .map(move |r| coalesce(r, coalescing_window))
+            .flat_map(stream::iter);
 
         // Submit the coalesced requests to the I/O.
         let read = self.read.clone();
@@ -221,13 +224,15 @@ async fn evaluate<R: VortexReadAt>(
 }
 
 /// TODO(ngates): outsource coalescing to a trait
-fn coalesce(requests: Vec<FileSegmentRequest>) -> Vec<CoalescedSegmentRequest> {
-    const COALESCE: u64 = 1024 * 1024; // 1MB
+fn coalesce(
+    requests: Vec<FileSegmentRequest>,
+    coalescing_window: u64,
+) -> Vec<CoalescedSegmentRequest> {
     let fetch_ranges = merge_ranges(
         requests
             .iter()
             .map(|r| r.location.offset..r.location.offset + r.location.length as u64),
-        COALESCE,
+        coalescing_window,
     );
     let mut coalesced = fetch_ranges
         .iter()
@@ -239,7 +244,7 @@ fn coalesce(requests: Vec<FileSegmentRequest>) -> Vec<CoalescedSegmentRequest> {
 
     for req in requests {
         let idx = fetch_ranges.partition_point(|v| v.start <= req.location.offset) - 1;
-        coalesced.as_mut_slice()[idx].requests.push(req);
+        coalesced[idx].requests.push(req);
     }
 
     // Ensure we sort the requests by segment ID within the coalesced request.

vortex-io/src/offset.rs

@@ -4,7 +4,7 @@ use std::io;
 use bytes::Bytes;
 use futures::FutureExt;
 
-use crate::VortexReadAt;
+use crate::{PerformanceHint, VortexReadAt};
 
 /// An adapter that offsets all reads by a fixed amount.
 pub struct OffsetReadAt<R> {
@@ -39,7 +39,7 @@ impl<R: VortexReadAt> VortexReadAt for OffsetReadAt<R> {
         self.read.read_byte_range(pos + self.offset, len)
     }
 
-    fn performance_hint(&self) -> usize {
+    fn performance_hint(&self) -> PerformanceHint {
         self.read.performance_hint()
     }
 

vortex-io/src/read.rs

@@ -30,8 +30,8 @@ pub trait VortexReadAt: Send + Sync + Clone + 'static {
 
     // TODO(ngates): the read implementation should be able to hint at its latency/throughput
     //  allowing the caller to make better decisions about how to coalesce reads.
-    fn performance_hint(&self) -> usize {
-        0
+    fn performance_hint(&self) -> PerformanceHint {
+        PerformanceHint::default()
     }
 
     /// Asynchronously get the number of bytes of data readable.
@@ -41,6 +41,34 @@ pub trait VortexReadAt: Send + Sync + Clone + 'static {
     fn size(&self) -> impl Future<Output = io::Result<u64>> + 'static;
 }
 
+pub struct PerformanceHint {
+    coalescing_window: u64,
+}
+
+impl Default for PerformanceHint {
+    fn default() -> Self {
+        Self {
+            coalescing_window: 2 << 20, //1MB,
+        }
+    }
+}
+
+impl PerformanceHint {
+    pub fn new(coalescing_window: u64) -> Self {
+        Self { coalescing_window }
+    }
+
+    /// Creates a new instance with a profile appropriate for fast local storage, like memory or files on NVMe devices.
+    pub fn local() -> Self {
+        Self::new(0)
+    }
+
+    /// The maximum distance between two reads that should coalesced into a single operation.
+    pub fn coalescing_window(&self) -> u64 {
+        self.coalescing_window
+    }
+}
+
 impl<T: VortexReadAt> VortexReadAt for Arc<T> {
     fn read_byte_range(
         &self,
@@ -50,7 +78,7 @@ impl<T: VortexReadAt> VortexReadAt for Arc<T> {
         T::read_byte_range(self, pos, len)
     }
 
-    fn performance_hint(&self) -> usize {
+    fn performance_hint(&self) -> PerformanceHint {
         T::performance_hint(self)
     }
 
@@ -79,6 +107,10 @@ impl VortexReadAt for Bytes {
     fn size(&self) -> impl Future<Output = io::Result<u64>> + 'static {
         ready(Ok(self.len() as u64))
     }
+
+    fn performance_hint(&self) -> PerformanceHint {
+        PerformanceHint::local()
+    }
 }
 
 impl VortexReadAt for ByteBuffer {
@@ -98,6 +130,10 @@ impl VortexReadAt for ByteBuffer {
         ready(Ok(self.slice(read_start..read_end).into_inner()))
     }
 
+    fn performance_hint(&self) -> PerformanceHint {
+        PerformanceHint::local()
+    }
+
     fn size(&self) -> impl Future<Output = io::Result<u64>> + 'static {
         ready(Ok(self.len() as u64))
     }

vortex-io/src/tokio.rs

@@ -10,7 +10,7 @@ use bytes::{Bytes, BytesMut};
 use tokio::io::{AsyncWrite, AsyncWriteExt};
 use vortex_error::VortexUnwrap;
 
-use crate::{IoBuf, VortexReadAt, VortexWrite};
+use crate::{IoBuf, PerformanceHint, VortexReadAt, VortexWrite};
 
 pub struct TokioAdapter<IO>(pub IO);
 
@@ -81,6 +81,10 @@ impl VortexReadAt for TokioFile {
 
         async move { this.metadata().map(|metadata| metadata.len()) }
     }
+
+    fn performance_hint(&self) -> PerformanceHint {
+        PerformanceHint::local()
+    }
 }
 
 impl VortexWrite for tokio::fs::File {