feat(io): make file-read coalescing configurable

rust/lance-encoding/src/lib.rs

@@ -101,6 +101,19 @@ pub trait EncodingsIo: std::fmt::Debug + Send + Sync {
     ) -> Option<Arc<dyn EncodingsIo>> {
         None
     }
+
+    /// Returns a version of this I/O service that merges the ranges of a
+    /// request that are at most `gap` bytes apart into one read, in place of
+    /// its own coalescing rule. Every range still receives exactly its bytes,
+    /// so decoded data is unchanged; only the number and size of the reads
+    /// differ.
+    ///
+    /// Returns `None` if this implementation does not coalesce requests (e.g.
+    /// in-memory or test schedulers), in which case the caller should fall
+    /// back to using self.
+    fn with_coalesce_gap(&self, _gap: u64) -> Option<Arc<dyn EncodingsIo>> {
+        None
+    }
 }
 
 /// An implementation of EncodingsIo that serves data from an in-memory buffer

rust/lance-file/src/io.rs

@@ -48,6 +48,13 @@ impl EncodingsIo for LanceEncodingsIo {
         }))
     }
 
+    fn with_coalesce_gap(&self, gap: u64) -> Option<Arc<dyn EncodingsIo>> {
+        Some(Arc::new(Self {
+            scheduler: self.scheduler.with_block_size(gap),
+            read_chunk_size: self.read_chunk_size,
+        }))
+    }
+
     fn submit_request(
         &self,
         ranges: Vec<std::ops::Range<u64>>,

rust/lance-file/src/reader.rs

@@ -625,6 +625,23 @@ impl FileReader {
         }
     }
 
+    /// Returns a clone of this reader whose I/O merges requested byte ranges
+    /// that are at most `gap` bytes apart into one request, instead of those
+    /// within the object store's block size.
+    ///
+    /// Each range is cut back out of the merged request and decoded as
+    /// before, so reads return the same data at every gap; a wider gap reads
+    /// the bytes between the ranges too, in fewer requests. All cached
+    /// metadata is shared with `self`. If the underlying I/O service does not
+    /// coalesce requests (e.g. an in-memory scheduler), the returned reader is
+    /// an ordinary clone.
+    pub fn with_coalesce_gap(&self, gap: u64) -> Self {
+        match self.core.scheduler.with_coalesce_gap(gap) {
+            Some(scheduler) => self.with_scheduler(scheduler),
+            None => self.clone(),
+        }
+    }
+
     pub fn num_rows(&self) -> u64 {
         self.core.num_rows()
     }
@@ -2436,6 +2453,7 @@ mod tests {
     };
     use arrow_buffer::{NullBuffer, OffsetBuffer, ScalarBuffer};
     use arrow_schema::{DataType, Field, Fields, Schema as ArrowSchema};
+    use arrow_select::concat::concat_batches;
     use bytes::Bytes;
     use futures::{StreamExt, prelude::stream::TryStreamExt};
     use lance_arrow::{BLOB_META_KEY, RecordBatchExt};
@@ -2450,7 +2468,7 @@ mod tests {
         encoder::{EncodedBatch, EncodingOptions, encode_batch},
         format::pb21,
     };
-    use lance_io::{stream::RecordBatchStream, utils::CachedFileSize};
+    use lance_io::{scheduler::IoStats, stream::RecordBatchStream, utils::CachedFileSize};
     use log::debug;
     use rstest::rstest;
     use tokio::sync::mpsc;
@@ -4100,6 +4118,135 @@ mod tests {
         );
     }
 
+    /// Read `ranges` of every column through a clone of `reader` that merges
+    /// requested byte ranges at most `gap` bytes apart (the store's block size
+    /// when `None`), returning the batch and the requests the read made.
+    async fn read_ranges_with_gap(
+        reader: &FileReader,
+        ranges: &[std::ops::Range<u64>],
+        gap: Option<u64>,
+    ) -> (RecordBatch, u64) {
+        let stats = IoStats::new();
+        let mut reader = reader.with_io_stats(stats.recorder());
+        if let Some(gap) = gap {
+            reader = reader.with_coalesce_gap(gap);
+        }
+        let batches = reader
+            .read_stream(
+                lance_io::ReadBatchParams::Ranges(ranges.into()),
+                u32::MAX,
+                1,
+                FilterExpression::no_filter(),
+            )
+            .await
+            .unwrap()
+            .try_collect::<Vec<_>>()
+            .await
+            .unwrap();
+        let batch = concat_batches(&batches[0].schema(), &batches).unwrap();
+        (batch, stats.snapshot().iops)
+    }
+
+    /// A reader's coalescing gap changes only the requests that a read of
+    /// scattered rows makes: every gap returns the same batch, a wider gap
+    /// makes no more requests than a narrower one, and a gap that covers the
+    /// file makes at most one request per column page.
+    #[rstest]
+    #[case::v2_0(ConcreteFileVersion::V2_0)]
+    #[case::v2_2(ConcreteFileVersion::V2_2)]
+    #[tokio::test]
+    async fn test_coalesce_gap_keeps_ranges_reads(#[case] version: ConcreteFileVersion) {
+        const PAGES: u32 = 4;
+        const PAGE_ROWS: u64 = 1024;
+        let fs = FsFixture::default();
+        let codes =
+            DataType::FixedSizeList(Arc::new(Field::new("item", DataType::UInt8, true)), 64);
+        let reader = gen_batch()
+            .col("x", array::step::<Int32Type>())
+            .col("codes", array::rand_type(&codes))
+            .into_reader_rows(RowCount::from(PAGE_ROWS), BatchCount::from(PAGES));
+        // Without a cache, every written batch is a page of each column.
+        let WrittenFile { data, .. } = write_lance_file(
+            reader,
+            &fs,
+            version,
+            FileWriterOptions {
+                data_cache_bytes: Some(1),
+                ..Default::default()
+            },
+        )
+        .await;
+        let file_scheduler = fs
+            .scheduler
+            .open_file(&fs.tmp_path, &CachedFileSize::unknown())
+            .await
+            .unwrap();
+        let file_reader = FileReader::try_open(
+            file_scheduler,
+            None,
+            Arc::<DecoderPlugins>::default(),
+            &test_cache(),
+            FileReaderOptions::default(),
+        )
+        .await
+        .unwrap();
+        let num_pages = file_reader
+            .metadata()
+            .column_infos
+            .iter()
+            .map(|column| column.page_infos.len() as u64)
+            .sum::<u64>();
+        assert!(
+            num_pages >= 2 * u64::from(PAGES),
+            "the writer must cut a page of each column per batch, got {num_pages}"
+        );
+
+        // Single rows 97 apart, a run across the end of the first page and
+        // a run inside a page.
+        let rows = (0..PAGE_ROWS * u64::from(PAGES))
+            .filter(|row| {
+                row.is_multiple_of(97) || (1020..1030).contains(row) || (3000..3003).contains(row)
+            })
+            .collect::<Vec<_>>();
+        let mut ranges: Vec<std::ops::Range<u64>> = Vec::new();
+        for &row in &rows {
+            match ranges.last_mut() {
+                Some(last) if last.end == row => last.end += 1,
+                _ => ranges.push(row..row + 1),
+            }
+        }
+        let indices = UInt32Array::from_iter_values(rows.iter().map(|&row| row as u32));
+        let expected = concat_batches(&data[0].schema(), &data)
+            .unwrap()
+            .take(&indices)
+            .unwrap();
+
+        // The first read also initializes the page metadata; the cache keeps
+        // it, so the reads below are data I/O only.
+        let (baseline, _) = read_ranges_with_gap(&file_reader, &ranges, None).await;
+        assert_eq!(baseline.columns(), expected.columns());
+        let mut requests = Vec::new();
+        for gap in [0, 64 * 1024, 256 * 1024, 1024 * 1024, u64::MAX] {
+            let (batch, gap_requests) =
+                read_ranges_with_gap(&file_reader, &ranges, Some(gap)).await;
+            assert_eq!(batch, baseline, "gap={gap}");
+            requests.push(gap_requests);
+        }
+        assert!(
+            requests.windows(2).all(|pair| pair[1] <= pair[0]),
+            "{requests:?}"
+        );
+        let covering = requests[requests.len() - 1];
+        assert!(
+            covering <= num_pages,
+            "{covering} requests for {num_pages} column pages"
+        );
+        if version == ConcreteFileVersion::V2_0 {
+            // Each run of each column is its own request at gap 0.
+            assert!(requests[0] > covering, "{requests:?}");
+        }
+    }
+
     #[rstest]
     #[tokio::test]
     async fn test_blocking_take(

rust/lance-io/src/scheduler.rs

@@ -1171,7 +1171,8 @@ pub struct FileScheduler {
 
 fn is_close_together(range1: &Range<u64>, range2: &Range<u64>, block_size: u64) -> bool {
     // Note that range1.end <= range2.start is possible (e.g. when decoding string arrays)
-    range2.start <= (range1.end + block_size)
+    // A gap from `FileScheduler::with_block_size` may be as large as `u64::MAX`.
+    range2.start <= range1.end.saturating_add(block_size)
 }
 
 fn is_overlapping(range1: &Range<u64>, range2: &Range<u64>) -> bool {
@@ -1340,6 +1341,21 @@ impl FileScheduler {
         }
     }
 
+    /// Returns a copy of this scheduler that merges the ranges of a request
+    /// that are at most `block_size` bytes apart into one I/O, instead of
+    /// using the object store's block size.
+    ///
+    /// Each requested range is still cut back out of the merged read, so the
+    /// bytes returned are the same for every gap; only the number and size of
+    /// the I/Os change. A wider gap trades extra bytes for fewer requests,
+    /// which pays off where every request is a round trip to remote storage.
+    pub fn with_block_size(&self, block_size: u64) -> Self {
+        Self {
+            block_size,
+            ..self.clone()
+        }
+    }
+
     /// Returns a copy of this scheduler that bypasses backpressure for all requests.
     ///
     /// This should be used for indirect I/O (e.g. fetching items after decoding offsets) where
@@ -1713,6 +1729,39 @@ mod tests {
             );
         }
         assert_eq!(11, scheduler.stats().iops);
+
+        // A gap set on the file scheduler replaces the store's 4KiB block
+        // size: gap 0 merges only touching ranges, a wider gap merges every
+        // range it reaches, and `u64::MAX` merges them all. Each range still
+        // gets exactly its own bytes.
+        let reads = vec![
+            50_000..51_000,
+            52_000..53_000,
+            150_000..151_000,
+            151_000..152_000,
+        ];
+        for (gap, expected_iops) in [(0, 3), (4 * 1024, 2), (100_000, 1), (u64::MAX, 1)] {
+            let iops_before = scheduler.stats().iops;
+            let bytes = file_scheduler
+                .with_block_size(gap)
+                .submit_request(reads.clone(), 0)
+                .await
+                .unwrap();
+            assert_eq!(bytes.len(), reads.len(), "gap={gap}");
+            for (range, bytes) in reads.iter().zip(&bytes) {
+                let expected = &some_data[range.start as usize..range.end as usize];
+                assert_eq!(bytes, expected, "gap={gap} range={range:?}");
+            }
+            assert_eq!(
+                scheduler.stats().iops - iops_before,
+                expected_iops,
+                "gap={gap}"
+            );
+        }
+        // The scheduler the gap was set on keeps the store's block size.
+        let iops_before = scheduler.stats().iops;
+        file_scheduler.submit_request(reads, 0).await.unwrap();
+        assert_eq!(scheduler.stats().iops - iops_before, 2);
     }
 
     #[rstest]