From 1c5c215cf2263a25ef6541ff7a9e367b15f4663e Mon Sep 17 00:00:00 2001 From: Yang Cen <159225399+BubbleCal@users.noreply.github.com> Date: Fri, 9 Oct 2026 22:13:01 +0800 Subject: [PATCH] feat(io): add sparse column read and coalescing controls --- rust/lance-encoding/src/lib.rs | 13 +++ rust/lance-file/src/io.rs | 7 ++ rust/lance-file/src/reader.rs | 149 ++++++++++++++++++++++++++++++++- rust/lance-io/src/scheduler.rs | 51 ++++++++++- 4 files changed, 218 insertions(+), 2 deletions(-) diff --git a/rust/lance-encoding/src/lib.rs b/rust/lance-encoding/src/lib.rs index 338028b79d0..ff5f0ad4ad7 100644 --- a/rust/lance-encoding/src/lib.rs +++ b/rust/lance-encoding/src/lib.rs @@ -101,6 +101,19 @@ pub trait EncodingsIo: std::fmt::Debug + Send + Sync { ) -> Option> { 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> { + None + } } /// An implementation of EncodingsIo that serves data from an in-memory buffer diff --git a/rust/lance-file/src/io.rs b/rust/lance-file/src/io.rs index 86e5189a81d..1383b045295 100644 --- a/rust/lance-file/src/io.rs +++ b/rust/lance-file/src/io.rs @@ -48,6 +48,13 @@ impl EncodingsIo for LanceEncodingsIo { })) } + fn with_coalesce_gap(&self, gap: u64) -> Option> { + Some(Arc::new(Self { + scheduler: self.scheduler.with_block_size(gap), + read_chunk_size: self.read_chunk_size, + })) + } + fn submit_request( &self, ranges: Vec>, diff --git a/rust/lance-file/src/reader.rs b/rust/lance-file/src/reader.rs index 65aa19939a7..49528de72fd 100644 --- a/rust/lance-file/src/reader.rs +++ b/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], + gap: Option, + ) -> (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::>() + .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::()) + .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::::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::(); + 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::>(); + let mut ranges: Vec> = 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( diff --git a/rust/lance-io/src/scheduler.rs b/rust/lance-io/src/scheduler.rs index 1876fe3feba..0e0095b73a4 100644 --- a/rust/lance-io/src/scheduler.rs +++ b/rust/lance-io/src/scheduler.rs @@ -1171,7 +1171,8 @@ pub struct FileScheduler { fn is_close_together(range1: &Range, range2: &Range, 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, range2: &Range) -> 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]