lance_file/
io.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::sync::Arc;
5
6use futures::{FutureExt, future::BoxFuture};
7use lance_encoding::EncodingsIo;
8use lance_io::scheduler::FileScheduler;
9
10use super::reader::DEFAULT_READ_CHUNK_SIZE;
11
12#[derive(Debug)]
13pub struct LanceEncodingsIo {
14    scheduler: FileScheduler,
15    /// Size of chunks when reading large pages
16    read_chunk_size: u64,
17}
18
19impl LanceEncodingsIo {
20    pub fn new(scheduler: FileScheduler) -> Self {
21        Self {
22            scheduler,
23            read_chunk_size: DEFAULT_READ_CHUNK_SIZE,
24        }
25    }
26
27    pub fn with_read_chunk_size(mut self, read_chunk_size: u64) -> Self {
28        self.read_chunk_size = read_chunk_size;
29        self
30    }
31}
32
33impl EncodingsIo for LanceEncodingsIo {
34    fn with_bypass_backpressure(&self) -> Option<Arc<dyn EncodingsIo>> {
35        Some(Arc::new(Self {
36            scheduler: self.scheduler.with_bypass_backpressure(),
37            read_chunk_size: self.read_chunk_size,
38        }))
39    }
40
41    fn with_io_stats(
42        &self,
43        stats: Arc<dyn lance_core::utils::io_stats::IoStatsRecorder>,
44    ) -> Option<Arc<dyn EncodingsIo>> {
45        Some(Arc::new(Self {
46            scheduler: self.scheduler.with_io_stats(stats),
47            read_chunk_size: self.read_chunk_size,
48        }))
49    }
50
51    fn with_coalesce_gap(&self, gap: u64) -> Option<Arc<dyn EncodingsIo>> {
52        Some(Arc::new(Self {
53            scheduler: self.scheduler.with_block_size(gap),
54            read_chunk_size: self.read_chunk_size,
55        }))
56    }
57
58    fn submit_request(
59        &self,
60        ranges: Vec<std::ops::Range<u64>>,
61        priority: u64,
62    ) -> BoxFuture<'static, lance_core::Result<Vec<bytes::Bytes>>> {
63        let mut split_ranges = Vec::new();
64        let mut split_indices = Vec::new(); // Track which original range each split came from
65        // Large ranges (above read_chunk_size) will be split into
66        // multiple reads.  Empty ranges will skip the I/O layer
67        // entirely.  If we have either of these we will need to
68        // reassemble our results, inserting empties and merging parts
69        let mut needs_reassembly = false;
70
71        // Split large ranges into smaller chunks
72        //
73        // TODO: consider read_chunk_size before submitting requests.
74        for (idx, range) in ranges.iter().enumerate() {
75            if range.start == range.end {
76                // EncodingsIo requires one result per input range. Zero-length
77                // ranges schedule no I/O, so their empty results are restored
78                // after the non-empty requests complete.
79                needs_reassembly = true;
80                continue;
81            }
82            let range_size = range.end - range.start;
83
84            if range_size > self.read_chunk_size {
85                needs_reassembly = true;
86                let num_chunks = range_size.div_ceil(self.read_chunk_size);
87                let chunk_size = range_size / num_chunks;
88
89                for i in 0..num_chunks {
90                    let start = range.start + i * chunk_size;
91                    let end = if i == num_chunks - 1 {
92                        range.end // Last chunk gets any remaining bytes
93                    } else {
94                        start + chunk_size
95                    };
96                    split_ranges.push(start..end);
97                    split_indices.push(idx);
98                }
99            } else {
100                split_ranges.push(range.clone());
101                split_indices.push(idx);
102            }
103        }
104
105        let fut = self.scheduler.submit_request(split_ranges, priority);
106
107        async move {
108            let split_results = fut.await?;
109
110            if split_results.len() != split_indices.len() {
111                return Err(lance_core::Error::internal(format!(
112                    "Encoding I/O returned {} results for {} requested range chunks",
113                    split_results.len(),
114                    split_indices.len()
115                )));
116            }
117            if !needs_reassembly {
118                return Ok(split_results);
119            }
120
121            let mut results = vec![Vec::new(); ranges.len()];
122
123            for (split_result, orig_idx) in split_results.into_iter().zip(split_indices) {
124                results[orig_idx].push(split_result);
125            }
126
127            let mut reassembled = Vec::with_capacity(ranges.len());
128            for (range, chunks) in ranges.iter().zip(results) {
129                if chunks.is_empty() {
130                    if range.start == range.end {
131                        reassembled.push(bytes::Bytes::new());
132                        continue;
133                    }
134                    return Err(lance_core::Error::internal(format!(
135                        "Encoding I/O returned no data for non-empty range {}..{}",
136                        range.start, range.end
137                    )));
138                }
139                if chunks.len() == 1 {
140                    reassembled.push(chunks[0].clone());
141                    continue;
142                }
143
144                let total_size: usize = chunks.iter().map(|c| c.len()).sum();
145                let mut combined = Vec::with_capacity(total_size);
146                for chunk in chunks {
147                    combined.extend_from_slice(&chunk);
148                }
149                reassembled.push(bytes::Bytes::from(combined));
150            }
151            Ok(reassembled)
152        }
153        .boxed()
154    }
155}
lance_file/io.rs

lance_file/
io.rs
