coalesced request align all segments (#6158)

We align the coalesced range start down to the max alignment of all
requests. That makes `coalesced_start % max_alignment == 0`, just like
the file start. Because segment offsets are aligned in the file, their
offsets relative to the coalesced start remain aligned, so each buffer
stays properly aligned without per‑buffer realignment.

example:
bytes in the file:
[x, x, x, x, x, x, A, A, A, A, A, x, B]
A aligned to 2
B aligned to 4

coalesced read both to a 2 aligned buffer:
[ A, A, A, A, A, x, B]
A aligned to 2
B is not aligned to 4

with this change we read this to the coalesced buffer:
[x, x, A, A, A, A, A, x, B]
A aligned to 2
B aligned to 4

---------

Signed-off-by: Onur Satici <onur@spiraldb.com>

Author: onursatici

vortex-file/src/read/driver.rs

@@ -10,6 +10,7 @@ use std::task::Poll;
 
 use futures::Stream;
 use pin_project_lite::pin_project;
+use vortex_buffer::Alignment;
 use vortex_error::VortexExpect;
 use vortex_io::CoalesceConfig;
 use vortex_metrics::Counter;
@@ -47,6 +48,7 @@ impl<S> IoRequestStream<S> {
     pub(crate) fn new(
         events: S,
         coalesce_window: Option<CoalesceConfig>,
+        coalesced_buffer_alignment: Alignment,
         metrics: VortexMetrics,
     ) -> Self
     where
@@ -56,7 +58,7 @@ impl<S> IoRequestStream<S> {
             events,
             inner_done: false,
             coalesce_window,
-            state: State::new(metrics),
+            state: State::new(metrics, coalesced_buffer_alignment),
         }
     }
 }
@@ -117,6 +119,7 @@ struct State {
 
     // Metrics for tracking I/O request patterns
     metrics: StateMetrics,
+    coalesced_buffer_alignment: Alignment,
 }
 
 struct StateMetrics {
@@ -136,12 +139,13 @@ impl StateMetrics {
 }
 
 impl State {
-    fn new(metrics: VortexMetrics) -> Self {
+    fn new(metrics: VortexMetrics, coalesced_buffer_alignment: Alignment) -> Self {
         Self {
             requests: BTreeMap::new(),
             polled_requests: BTreeMap::new(),
             requests_by_offset: BTreeSet::new(),
             metrics: StateMetrics::new(metrics),
+            coalesced_buffer_alignment,
         }
     }
 
@@ -206,14 +210,23 @@ impl State {
         None
     }
 
+    /// Coalesce nearby requests into a single range while aligning the range start down to the
+    /// global maximum segment alignment.
+    ///
+    /// Example (file offsets):
+    /// [x, x, x, x, x, x, A, A, A, A, A, x, B]
+    /// A aligned to 2, B aligned to 4
+    /// Coalesced range starts at 4, so the buffer is:
+    /// [x, x, A, A, A, A, A, x, B]
+    /// A stays 2-aligned, B stays 4-aligned
     fn next_coalesced(&mut self, window: &CoalesceConfig) -> Option<CoalescedRequest> {
         // Find the next valid request in priority order
         let first_req = self.next_uncoalesced()?;
 
         let mut requests = vec![first_req];
         let mut current_start = requests[0].offset;
         let mut current_end = requests[0].offset + requests[0].length as u64;
-        let alignment = requests[0].alignment;
+        let align = *self.coalesced_buffer_alignment as u64;
 
         let mut keys_to_remove = Vec::new();
         let mut found_new_requests = true;
@@ -256,25 +269,25 @@ impl State {
                     // Calculate what the new range would be if we include this request
                     let new_start = current_start.min(req_offset);
                     let new_end = current_end.max(req_end);
-                    let new_total_size = new_end - new_start;
-
-                    // Check if the coalesced request would exceed max_size
-                    if new_total_size <= window.max_size {
-                        current_start = new_start;
-                        current_end = new_end;
-
-                        let req = self
-                            .polled_requests
-                            .remove(&req_id)
-                            .or_else(|| self.requests.remove(&req_id))
-                            .vortex_expect("Missing request in requests_by_offset");
-
-                        requests.push(req);
-                        keys_to_remove.push((req_offset, req_id));
-                        found_new_requests = true;
+                    let aligned_start = new_start - (new_start % align);
+                    let new_total_size = new_end - aligned_start;
+
+                    if new_total_size > window.max_size {
+                        // Skip it but keep it available for future coalescing operations.
+                        continue;
                     }
-                    // If adding this request would exceed max_size, we skip it but don't remove it
-                    // so it can be considered for future coalescing operations
+
+                    current_start = new_start;
+                    current_end = new_end;
+                    let req = self
+                        .polled_requests
+                        .remove(&req_id)
+                        .or_else(|| self.requests.remove(&req_id))
+                        .vortex_expect("Missing request in requests_by_offset");
+
+                    requests.push(req);
+                    keys_to_remove.push((req_offset, req_id));
+                    found_new_requests = true;
                 }
             }
         }
@@ -290,17 +303,19 @@ impl State {
         // Sort requests by offset for correct slicing in resolve
         requests.sort_unstable_by_key(|r| r.offset);
 
+        let aligned_start = current_start - (current_start % align);
+
         tracing::debug!(
             "Coalesced {} requests into range {}..{} (len={})",
             requests.len(),
-            current_start,
+            aligned_start,
             current_end,
-            current_end - current_start,
+            current_end - aligned_start,
         );
 
         Some(CoalescedRequest {
-            range: current_start..current_end,
-            alignment,
+            range: aligned_start..current_end,
+            alignment: self.coalesced_buffer_alignment,
             requests,
         })
     }
@@ -338,10 +353,23 @@ mod tests {
     async fn collect_outputs(
         events: Vec<ReadEvent>,
         coalesce_window: Option<CoalesceConfig>,
+    ) -> Vec<IoRequest> {
+        collect_outputs_with_alignment(events, coalesce_window, Alignment::none()).await
+    }
+
+    async fn collect_outputs_with_alignment(
+        events: Vec<ReadEvent>,
+        coalesce_window: Option<CoalesceConfig>,
+        coalesced_buffer_alignment: Alignment,
     ) -> Vec<IoRequest> {
         let event_stream = stream::iter(events);
         let metrics = VortexMetrics::default();
-        let io_stream = IoRequestStream::new(event_stream, coalesce_window, metrics);
+        let io_stream = IoRequestStream::new(
+            event_stream,
+            coalesce_window,
+            coalesced_buffer_alignment,
+            metrics,
+        );
         io_stream.collect().await
     }
 
@@ -442,6 +470,57 @@ mod tests {
         }
     }
 
+    #[tokio::test]
+    async fn test_coalesce_alignment_adjustment() {
+        let (tx1, _rx1) = oneshot::channel();
+        let (tx2, _rx2) = oneshot::channel();
+
+        let req1 = ReadRequest {
+            id: 1,
+            offset: 6,
+            length: 5,
+            alignment: Alignment::new(2),
+            callback: tx1,
+        };
+        let req2 = ReadRequest {
+            id: 2,
+            offset: 12,
+            length: 1,
+            alignment: Alignment::new(4),
+            callback: tx2,
+        };
+
+        let events = vec![
+            ReadEvent::Request(req1),
+            ReadEvent::Request(req2),
+            ReadEvent::Polled(1),
+            ReadEvent::Polled(2),
+        ];
+
+        let outputs = collect_outputs_with_alignment(
+            events,
+            Some(CoalesceConfig {
+                distance: 1,
+                max_size: 1024,
+            }),
+            Alignment::new(4),
+        )
+        .await;
+
+        assert_eq!(outputs.len(), 1);
+        match outputs[0].inner() {
+            IoRequestInner::Coalesced(coalesced) => {
+                assert_eq!(coalesced.range.start, 4);
+                assert_eq!(coalesced.alignment, Alignment::new(4));
+                for req in &coalesced.requests {
+                    let rel = req.offset - coalesced.range.start;
+                    assert_eq!(rel % *req.alignment as u64, 0);
+                }
+            }
+            _ => panic!("Expected coalesced request"),
+        }
+    }
+
     #[tokio::test]
     async fn test_dropped_requests() {
         let (req1, _rx1) = create_request(1, 0, 10);
@@ -609,6 +688,7 @@ mod tests {
                 distance: 5,
                 max_size: 1024,
             }),
+            Alignment::none(),
             metrics.clone(),
         );
 
@@ -662,7 +742,8 @@ mod tests {
         let event_stream = stream::iter(events);
         let metrics = VortexMetrics::default();
         // No coalescing window - should be individual requests
-        let io_stream = IoRequestStream::new(event_stream, None, metrics.clone());
+        let io_stream =
+            IoRequestStream::new(event_stream, None, Alignment::none(), metrics.clone());
 
         let outputs: Vec<IoRequest> = io_stream.collect().await;
         assert_eq!(outputs.len(), 2);

vortex-file/src/read/request.rs

@@ -116,7 +116,7 @@ impl ReadRequest {
 /// A set of I/O requests that have been coalesced into a single larger request.
 pub(crate) struct CoalescedRequest {
     pub(crate) range: Range<u64>,
-    pub(crate) alignment: Alignment, // The alignment of the first request in the coalesced range.
+    pub(crate) alignment: Alignment, // Global max segment alignment used for the coalesced range.
     pub(crate) requests: Vec<ReadRequest>, // TODO(ngates): we could have enum of Single/Many to avoid Vec.
 }
 

vortex-file/src/segments/source.rs

@@ -13,6 +13,7 @@ use futures::FutureExt;
 use futures::StreamExt;
 use futures::channel::mpsc;
 use vortex_array::buffer::BufferHandle;
+use vortex_buffer::Alignment;
 use vortex_buffer::ByteBuffer;
 use vortex_error::VortexResult;
 use vortex_error::vortex_err;
@@ -73,12 +74,28 @@ impl FileSegmentSource {
     ) -> Self {
         let (send, recv) = mpsc::unbounded();
 
-        let coalesce_config = reader.coalesce_config();
+        let max_alignment = segments
+            .iter()
+            .map(|segment| segment.alignment)
+            .max()
+            .unwrap_or_else(Alignment::none);
+        let coalesce_config = reader.coalesce_config().map(|mut config| {
+            // Aligning the coalesced start down can add up to (alignment - 1) bytes.
+            // Increase max_size to keep the effective payload window consistent.
+            let extra = (*max_alignment as u64).saturating_sub(1);
+            config.max_size = config.max_size.saturating_add(extra);
+            config
+        });
         let concurrency = reader.concurrency();
 
         let drive_fut = async move {
-            let stream =
-                IoRequestStream::new(StreamExt::boxed(recv), coalesce_config, metrics).boxed();
+            let stream = IoRequestStream::new(
+                StreamExt::boxed(recv),
+                coalesce_config,
+                max_alignment,
+                metrics,
+            )
+            .boxed();
 
             stream
                 .map(move |req| {