lance_io/
lib.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3#![recursion_limit = "512"]
4use std::{
5    ops::{Range, RangeFrom, RangeFull, RangeTo},
6    sync::Arc,
7};
8
9use arrow::datatypes::UInt32Type;
10use arrow_array::{PrimitiveArray, UInt32Array};
11
12use lance_core::{Error, Result};
13
14pub mod encodings;
15pub mod ffi;
16pub mod local;
17pub mod object_reader;
18pub mod object_store;
19pub mod object_writer;
20pub mod scheduler;
21pub mod spill;
22pub mod stream;
23#[cfg(test)]
24pub mod testing;
25pub mod traits;
26#[cfg(target_os = "linux")]
27pub mod uring;
28pub mod utils;
29
30pub use scheduler::{bytes_read_counter, iops_counter};
31
32/// Defines a selection of rows to read from a file/batch
33#[derive(Debug, Clone, PartialEq, Default)]
34pub enum ReadBatchParams {
35    /// Select a contiguous range of rows
36    Range(Range<usize>),
37    /// Select multiple contiguous ranges of rows
38    Ranges(Arc<[Range<u64>]>),
39    /// Select all rows (this is the default)
40    #[default]
41    RangeFull,
42    /// Select all rows up to a given index
43    RangeTo(RangeTo<usize>),
44    /// Select all rows starting at a given index
45    RangeFrom(RangeFrom<usize>),
46    /// Select scattered non-contiguous rows
47    Indices(UInt32Array),
48}
49
50impl std::fmt::Display for ReadBatchParams {
51    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
52        match self {
53            Self::Range(r) => write!(f, "Range({}..{})", r.start, r.end),
54            Self::Ranges(ranges) => {
55                let mut ranges_str = ranges.iter().fold(String::new(), |mut acc, r| {
56                    acc.push_str(&format!("{}..{}", r.start, r.end));
57                    acc.push(',');
58                    acc
59                });
60                // Remove the trailing comma
61                if !ranges_str.is_empty() {
62                    ranges_str.pop();
63                }
64                write!(f, "Ranges({})", ranges_str)
65            }
66            Self::RangeFull => write!(f, "RangeFull"),
67            Self::RangeTo(r) => write!(f, "RangeTo({})", r.end),
68            Self::RangeFrom(r) => write!(f, "RangeFrom({})", r.start),
69            Self::Indices(indices) => {
70                let mut indices_str = indices.values().iter().fold(String::new(), |mut acc, v| {
71                    acc.push_str(&v.to_string());
72                    acc.push(',');
73                    acc
74                });
75                if !indices_str.is_empty() {
76                    indices_str.pop();
77                }
78                write!(f, "Indices({})", indices_str)
79            }
80        }
81    }
82}
83
84impl From<&[u32]> for ReadBatchParams {
85    fn from(value: &[u32]) -> Self {
86        Self::Indices(UInt32Array::from_iter_values(value.iter().copied()))
87    }
88}
89
90impl From<UInt32Array> for ReadBatchParams {
91    fn from(value: UInt32Array) -> Self {
92        Self::Indices(value)
93    }
94}
95
96impl From<RangeFull> for ReadBatchParams {
97    fn from(_: RangeFull) -> Self {
98        Self::RangeFull
99    }
100}
101
102impl From<Range<usize>> for ReadBatchParams {
103    fn from(r: Range<usize>) -> Self {
104        Self::Range(r)
105    }
106}
107
108impl From<RangeTo<usize>> for ReadBatchParams {
109    fn from(r: RangeTo<usize>) -> Self {
110        Self::RangeTo(r)
111    }
112}
113
114impl From<RangeFrom<usize>> for ReadBatchParams {
115    fn from(r: RangeFrom<usize>) -> Self {
116        Self::RangeFrom(r)
117    }
118}
119
120impl From<&Self> for ReadBatchParams {
121    fn from(params: &Self) -> Self {
122        params.clone()
123    }
124}
125
126impl ReadBatchParams {
127    /// Validate that the selection is valid given the length of the batch
128    pub fn valid_given_len(&self, len: usize) -> bool {
129        match self {
130            Self::Indices(indices) => indices.iter().all(|i| i.unwrap_or(0) < len as u32),
131            Self::Range(r) => r.start < len && r.end <= len,
132            Self::Ranges(ranges) => ranges.iter().all(|r| r.end <= len as u64),
133            Self::RangeFull => true,
134            Self::RangeTo(r) => r.end <= len,
135            Self::RangeFrom(r) => r.start < len,
136        }
137    }
138
139    /// Slice the selection
140    ///
141    /// For example, given ReadBatchParams::RangeFull and slice(10, 20), the output will be
142    /// ReadBatchParams::Range(10..20)
143    ///
144    /// Given ReadBatchParams::Range(10..20) and slice(5, 3), the output will be
145    /// ReadBatchParams::Range(15..18)
146    ///
147    /// Given ReadBatchParams::RangeTo(20) and slice(10, 5), the output will be
148    /// ReadBatchParams::Range(10..15)
149    ///
150    /// Given ReadBatchParams::RangeFrom(20) and slice(10, 5), the output will be
151    /// ReadBatchParams::Range(30..35)
152    ///
153    /// Given ReadBatchParams::Indices([1, 3, 5, 7, 9]) and slice(1, 3), the output will be
154    /// ReadBatchParams::Indices([3, 5, 7])
155    ///
156    /// You cannot slice beyond the bounds of the selection and an attempt to do so will
157    /// return an error.
158    pub fn slice(&self, start: usize, length: usize) -> Result<Self> {
159        let out_of_bounds = |size: usize| {
160            Err(Error::invalid_input_source(
161                format!(
162                    "Cannot slice from {} with length {} given a selection of size {}",
163                    start, length, size
164                )
165                .into(),
166            ))
167        };
168
169        match self {
170            Self::Indices(indices) => {
171                if start + length > indices.len() {
172                    return out_of_bounds(indices.len());
173                }
174                Ok(Self::Indices(indices.slice(start, length)))
175            }
176            Self::Range(r) => {
177                if (r.start + start + length) > r.end {
178                    return out_of_bounds(r.end - r.start);
179                }
180                Ok(Self::Range((r.start + start)..(r.start + start + length)))
181            }
182            Self::Ranges(ranges) => {
183                let mut new_ranges = Vec::with_capacity(ranges.len());
184                let mut to_skip = start as u64;
185                let mut to_take = length as u64;
186                let mut total_num_rows = 0;
187                for r in ranges.as_ref() {
188                    let num_rows = r.end - r.start;
189                    total_num_rows += num_rows;
190                    if to_skip > num_rows {
191                        to_skip -= num_rows;
192                        continue;
193                    }
194                    let new_start = r.start + to_skip;
195                    let to_take_this_range = (num_rows - to_skip).min(to_take);
196                    new_ranges.push(new_start..(new_start + to_take_this_range));
197                    to_skip = 0;
198                    to_take -= to_take_this_range;
199                    if to_take == 0 {
200                        break;
201                    }
202                }
203                if to_take > 0 {
204                    out_of_bounds(total_num_rows as usize)
205                } else {
206                    Ok(Self::Ranges(new_ranges.into()))
207                }
208            }
209            Self::RangeFull => Ok(Self::Range(start..(start + length))),
210            Self::RangeTo(range) => {
211                if start + length > range.end {
212                    return out_of_bounds(range.end);
213                }
214                Ok(Self::Range(start..(start + length)))
215            }
216            Self::RangeFrom(r) => {
217                // No way to validate out_of_bounds, assume caller will do so
218                Ok(Self::Range((r.start + start)..(r.start + start + length)))
219            }
220        }
221    }
222
223    /// Convert a read range into a vector of row offsets
224    ///
225    /// RangeFull and RangeFrom are unbounded and cannot be converted into row offsets
226    /// and any attempt to do so will return an error.  Call slice first
227    pub fn to_offsets(&self) -> Result<PrimitiveArray<UInt32Type>> {
228        match self {
229            Self::Indices(indices) => Ok(indices.clone()),
230            Self::Range(r) => Ok(UInt32Array::from(Vec::from_iter(
231                r.start as u32..r.end as u32,
232            ))),
233            Self::Ranges(ranges) => {
234                let num_rows = ranges
235                    .iter()
236                    .map(|r| (r.end - r.start) as usize)
237                    .sum::<usize>();
238                let mut offsets = Vec::with_capacity(num_rows);
239                for r in ranges.as_ref() {
240                    offsets.extend(r.start as u32..r.end as u32);
241                }
242                Ok(UInt32Array::from(offsets))
243            }
244            Self::RangeFull => Err(Error::invalid_input("cannot materialize RangeFull")),
245            Self::RangeTo(r) => Ok(UInt32Array::from(Vec::from_iter(0..r.end as u32))),
246            Self::RangeFrom(_) => Err(Error::invalid_input("cannot materialize RangeFrom")),
247        }
248    }
249
250    pub fn iter_offset_ranges<'a>(
251        &'a self,
252    ) -> Result<Box<dyn Iterator<Item = Range<u32>> + Send + 'a>> {
253        match self {
254            Self::Indices(indices) => Ok(Box::new(indices.values().iter().map(|i| *i..(*i + 1)))),
255            Self::Range(r) => Ok(Box::new(std::iter::once(r.start as u32..r.end as u32))),
256            Self::Ranges(ranges) => Ok(Box::new(
257                ranges.iter().map(|r| r.start as u32..r.end as u32),
258            )),
259            Self::RangeFull => Err(Error::invalid_input("cannot materialize RangeFull")),
260            Self::RangeTo(r) => Ok(Box::new(std::iter::once(0..r.end as u32))),
261            Self::RangeFrom(_) => Err(Error::invalid_input("cannot materialize RangeFrom")),
262        }
263    }
264
265    /// Convert a read range into a vector of row ranges
266    pub fn to_ranges(&self) -> Result<Vec<Range<u64>>> {
267        match self {
268            Self::Indices(indices) => Ok(indices
269                .values()
270                .iter()
271                .map(|i| *i as u64..(*i + 1) as u64)
272                .collect()),
273            Self::Range(r) => Ok(vec![r.start as u64..r.end as u64]),
274            Self::Ranges(ranges) => Ok(ranges.to_vec()),
275            Self::RangeFull => Err(Error::invalid_input("cannot materialize RangeFull")),
276            Self::RangeTo(r) => Ok(vec![0..r.end as u64]),
277            Self::RangeFrom(_) => Err(Error::invalid_input("cannot materialize RangeFrom")),
278        }
279    }
280
281    /// Same thing as to_offsets but the caller knows the total number of rows in the file
282    ///
283    /// This makes it possible to materialize RangeFull / RangeFrom
284    pub fn to_offsets_total(&self, total: u32) -> PrimitiveArray<UInt32Type> {
285        match self {
286            Self::Indices(indices) => indices.clone(),
287            Self::Range(r) => UInt32Array::from_iter_values(r.start as u32..r.end as u32),
288            Self::Ranges(ranges) => {
289                let num_rows = ranges
290                    .iter()
291                    .map(|r| (r.end - r.start) as usize)
292                    .sum::<usize>();
293                let mut offsets = Vec::with_capacity(num_rows);
294                for r in ranges.as_ref() {
295                    offsets.extend(r.start as u32..r.end as u32);
296                }
297                UInt32Array::from(offsets)
298            }
299            Self::RangeFull => UInt32Array::from_iter_values(0_u32..total),
300            Self::RangeTo(r) => UInt32Array::from_iter_values(0..r.end as u32),
301            Self::RangeFrom(r) => UInt32Array::from_iter_values(r.start as u32..total),
302        }
303    }
304}
305
306#[cfg(test)]
307mod test {
308    use std::ops::{RangeFrom, RangeTo};
309
310    use arrow_array::UInt32Array;
311
312    use crate::ReadBatchParams;
313
314    #[test]
315    fn test_params_slice() {
316        let params = ReadBatchParams::Ranges(vec![0..15, 20..40].into());
317        let sliced = params.slice(10, 10).unwrap();
318        assert_eq!(sliced, ReadBatchParams::Ranges(vec![10..15, 20..25].into()));
319    }
320
321    #[test]
322    fn test_params_to_offsets() {
323        let check = |params: ReadBatchParams, base_offset, length, expected: Vec<u32>| {
324            let offsets = params
325                .slice(base_offset, length)
326                .unwrap()
327                .to_offsets()
328                .unwrap();
329            let expected = UInt32Array::from(expected);
330            assert_eq!(offsets, expected);
331        };
332
333        check(ReadBatchParams::RangeFull, 0, 100, (0..100).collect());
334        check(ReadBatchParams::RangeFull, 50, 100, (50..150).collect());
335        check(
336            ReadBatchParams::RangeFrom(RangeFrom { start: 500 }),
337            0,
338            100,
339            (500..600).collect(),
340        );
341        check(
342            ReadBatchParams::RangeFrom(RangeFrom { start: 500 }),
343            100,
344            100,
345            (600..700).collect(),
346        );
347        check(
348            ReadBatchParams::RangeTo(RangeTo { end: 800 }),
349            0,
350            100,
351            (0..100).collect(),
352        );
353        check(
354            ReadBatchParams::RangeTo(RangeTo { end: 800 }),
355            200,
356            100,
357            (200..300).collect(),
358        );
359        check(
360            ReadBatchParams::Indices(UInt32Array::from(vec![1, 3, 5, 7, 9])),
361            0,
362            2,
363            vec![1, 3],
364        );
365        check(
366            ReadBatchParams::Indices(UInt32Array::from(vec![1, 3, 5, 7, 9])),
367            2,
368            2,
369            vec![5, 7],
370        );
371
372        let check_error = |params: ReadBatchParams, base_offset, length| {
373            assert!(params.slice(base_offset, length).is_err());
374        };
375
376        check_error(ReadBatchParams::Indices(UInt32Array::from(vec![1])), 0, 2);
377        check_error(ReadBatchParams::Indices(UInt32Array::from(vec![1])), 1, 1);
378        check_error(ReadBatchParams::Range(0..10), 5, 6);
379        check_error(ReadBatchParams::RangeTo(RangeTo { end: 10 }), 5, 6);
380
381        assert!(ReadBatchParams::RangeFull.to_offsets().is_err());
382        assert!(
383            ReadBatchParams::RangeFrom(RangeFrom { start: 10 })
384                .to_offsets()
385                .is_err()
386        );
387    }
388}
lance_io/lib.rs

lance_io/
lib.rs
