Ji/io uring read ranges

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"

benchmarks/datafusion-bench/src/main.rs

@@ -289,10 +289,14 @@ async fn register_v2_tables<B: Benchmark + ?Sized>(
             .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);
 

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<IoRequest> {
         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,13 +621,15 @@ mod tests {
             offset: 6,
             length: 5,
             alignment: Alignment::new(2),
+            coalesce_distance: None,
             callback: tx1,
         };
         let req2 = ReadRequest {
             id: 2,
             offset: 12,
             length: 1,
             alignment: Alignment::new(4),
+            coalesce_distance: None,
             callback: tx2,
         };
 
@@ -634,13 +698,15 @@ mod tests {
             offset: 0,
             length: 10,
             alignment: Alignment::none(),
+            coalesce_distance: None,
             callback: tx1,
         };
         let req2 = ReadRequest {
             id: 2,
             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));

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<BufferHandle>) {
         match self.0 {
@@ -96,6 +107,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<u64>,
     pub(crate) callback: oneshot::Sender<VortexResult<BufferHandle>>,
 }
 
@@ -106,6 +119,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()
     }