lance_table/
utils.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The Lance Authors

pub mod stream;

pub trait LanceIteratorExtension {
    fn exact_size(self, size: usize) -> ExactSize<Self>
    where
        Self: Sized;
}

impl<I: Iterator> LanceIteratorExtension for I {
    fn exact_size(self, size: usize) -> ExactSize<Self>
    where
        Self: Sized,
    {
        ExactSize { inner: self, size }
    }
}

/// A iterator that is tagged with a known size. This is useful when we are
/// able to pre-compute the size of the iterator but the iterator implementation
/// isn't able to itself. A common example is when using `flatten()`.
///
/// This is inspired by discussion in https://github.com/rust-lang/rust/issues/68995
pub struct ExactSize<I> {
    inner: I,
    size: usize,
}

impl<I: Iterator> Iterator for ExactSize<I> {
    type Item = I::Item;

    fn next(&mut self) -> Option<Self::Item> {
        match self.inner.next() {
            None => None,
            Some(x) => {
                self.size -= 1;
                Some(x)
            }
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.size, Some(self.size))
    }
}