lance_table/
utils.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4pub mod stream;
5
6pub trait LanceIteratorExtension {
7    fn exact_size(self, size: usize) -> ExactSize<Self>
8    where
9        Self: Sized;
10}
11
12impl<I: Iterator> LanceIteratorExtension for I {
13    fn exact_size(self, size: usize) -> ExactSize<Self>
14    where
15        Self: Sized,
16    {
17        ExactSize { inner: self, size }
18    }
19}
20
21/// A iterator that is tagged with a known size. This is useful when we are
22/// able to pre-compute the size of the iterator but the iterator implementation
23/// isn't able to itself. A common example is when using `flatten()`.
24///
25/// This is inspired by discussion in https://github.com/rust-lang/rust/issues/68995
26pub struct ExactSize<I> {
27    inner: I,
28    size: usize,
29}
30
31impl<I: Iterator> Iterator for ExactSize<I> {
32    type Item = I::Item;
33
34    fn next(&mut self) -> Option<Self::Item> {
35        match self.inner.next() {
36            None => None,
37            Some(x) => {
38                self.size -= 1;
39                Some(x)
40            }
41        }
42    }
43
44    fn size_hint(&self) -> (usize, Option<usize>) {
45        (self.size, Some(self.size))
46    }
47}