polars_arrow/types/
index.rs

1use super::NativeType;
2use crate::trusted_len::TrustedLen;
3
4/// Sealed trait describing the subset of [`NativeType`] (`i32`, `i64`, `u32` and `u64`)
5/// that can be used to index a slot of an array.
6pub trait Index:
7    NativeType
8    + std::ops::AddAssign
9    + std::ops::Sub<Output = Self>
10    + num_traits::One
11    + num_traits::Num
12    + num_traits::CheckedAdd
13    + PartialOrd
14    + Ord
15{
16    /// Convert itself to [`usize`].
17    fn to_usize(&self) -> usize;
18    /// Convert itself from [`usize`].
19    fn from_usize(index: usize) -> Option<Self>;
20
21    /// Convert itself from [`usize`].
22    fn from_as_usize(index: usize) -> Self;
23
24    /// An iterator from (inclusive) `start` to (exclusive) `end`.
25    fn range(start: usize, end: usize) -> Option<IndexRange<Self>> {
26        let start = Self::from_usize(start);
27        let end = Self::from_usize(end);
28        match (start, end) {
29            (Some(start), Some(end)) => Some(IndexRange::new(start, end)),
30            _ => None,
31        }
32    }
33}
34
35macro_rules! index {
36    ($t:ty) => {
37        impl Index for $t {
38            #[inline]
39            fn to_usize(&self) -> usize {
40                *self as usize
41            }
42
43            #[inline]
44            fn from_usize(value: usize) -> Option<Self> {
45                Self::try_from(value).ok()
46            }
47
48            #[inline]
49            fn from_as_usize(value: usize) -> Self {
50                value as $t
51            }
52        }
53    };
54}
55
56index!(i8);
57index!(i16);
58index!(i32);
59index!(i64);
60index!(u8);
61index!(u16);
62index!(u32);
63index!(u64);
64
65/// Range of [`Index`], equivalent to `(a..b)`.
66/// `Step` is unstable in Rust, which does not allow us to implement (a..b) for [`Index`].
67pub struct IndexRange<I: Index> {
68    start: I,
69    end: I,
70}
71
72impl<I: Index> IndexRange<I> {
73    /// Returns a new [`IndexRange`].
74    pub fn new(start: I, end: I) -> Self {
75        assert!(end >= start);
76        Self { start, end }
77    }
78}
79
80impl<I: Index> Iterator for IndexRange<I> {
81    type Item = I;
82
83    #[inline]
84    fn next(&mut self) -> Option<Self::Item> {
85        if self.start == self.end {
86            return None;
87        }
88        let old = self.start;
89        self.start += I::one();
90        Some(old)
91    }
92
93    #[inline]
94    fn size_hint(&self) -> (usize, Option<usize>) {
95        let len = (self.end - self.start).to_usize();
96        (len, Some(len))
97    }
98}
99
100/// # Safety
101///
102/// A range is always of known length.
103unsafe impl<I: Index> TrustedLen for IndexRange<I> {}
polars_arrow/types/index.rs

polars_arrow/types/
index.rs
