polars_utils/
slice_enum.rs

1use std::num::TryFromIntError;
2use std::ops::Range;
3
4#[derive(Debug, Clone, PartialEq, Eq, Hash)]
5#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
6#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
7pub enum Slice {
8    /// Or zero
9    Positive {
10        offset: usize,
11        len: usize,
12    },
13    Negative {
14        offset_from_end: usize,
15        len: usize,
16    },
17}
18
19impl Slice {
20    #[allow(clippy::len_without_is_empty)]
21    pub fn len(&self) -> usize {
22        match self {
23            Slice::Positive { len, .. } => *len,
24            Slice::Negative { len, .. } => *len,
25        }
26    }
27
28    pub fn len_mut(&mut self) -> &mut usize {
29        match self {
30            Slice::Positive { len, .. } => len,
31            Slice::Negative { len, .. } => len,
32        }
33    }
34
35    /// Returns the offset of a positive slice.
36    ///
37    /// # Panics
38    /// Panics if `self` is [`Slice::Negative`]
39    pub fn positive_offset(&self) -> usize {
40        let Slice::Positive { offset, len: _ } = self.clone() else {
41            panic!("cannot use positive_offset() on a negative slice");
42        };
43
44        offset
45    }
46
47    /// Returns the end position of the slice (offset + len).
48    ///
49    /// # Panics
50    /// Panics if self is negative.
51    pub fn end_position(&self) -> usize {
52        let Slice::Positive { offset, len } = self.clone() else {
53            panic!("cannot use end_position() on a negative slice");
54        };
55
56        offset.saturating_add(len)
57    }
58
59    /// Returns the equivalent slice to apply from an offsetted position.
60    ///
61    /// # Panics
62    /// Panics if self is negative.
63    pub fn offsetted(self, position: usize) -> Self {
64        let Slice::Positive { offset, len } = self else {
65            panic!("cannot use offsetted() on a negative slice");
66        };
67
68        let (offset, len) = if position <= offset {
69            (offset - position, len)
70        } else {
71            let n_past_offset = position - offset;
72            (0, len.saturating_sub(n_past_offset))
73        };
74
75        Slice::Positive { offset, len }
76    }
77
78    /// Restricts the bounds of the slice to within a number of rows. Negative slices will also
79    /// be translated to the positive equivalent.
80    pub fn restrict_to_bounds(self, n_rows: usize) -> Self {
81        match self {
82            Slice::Positive { offset, len } => {
83                let offset = offset.min(n_rows);
84                let len = len.min(n_rows - offset);
85                Slice::Positive { offset, len }
86            },
87            Slice::Negative {
88                offset_from_end,
89                len,
90            } => {
91                if n_rows >= offset_from_end {
92                    // Trim extra starting rows
93                    let offset = n_rows - offset_from_end;
94                    let len = len.min(n_rows - offset);
95                    Slice::Positive { offset, len }
96                } else {
97                    // Slice offset goes past start of data.
98                    let stop_at_n_from_end = offset_from_end.saturating_sub(len);
99                    let len = n_rows.saturating_sub(stop_at_n_from_end);
100
101                    Slice::Positive { offset: 0, len }
102                }
103            },
104        }
105    }
106}
107
108impl From<(usize, usize)> for Slice {
109    fn from((offset, len): (usize, usize)) -> Self {
110        Slice::Positive { offset, len }
111    }
112}
113
114impl From<(i64, usize)> for Slice {
115    fn from((offset, len): (i64, usize)) -> Self {
116        if offset >= 0 {
117            Slice::Positive {
118                offset: usize::try_from(offset).unwrap(),
119                len,
120            }
121        } else {
122            Slice::Negative {
123                offset_from_end: usize::try_from(-offset).unwrap(),
124                len,
125            }
126        }
127    }
128}
129
130impl TryFrom<Slice> for (i64, usize) {
131    type Error = TryFromIntError;
132
133    fn try_from(value: Slice) -> Result<Self, Self::Error> {
134        match value {
135            Slice::Positive { offset, len } => Ok((i64::try_from(offset)?, len)),
136            Slice::Negative {
137                offset_from_end,
138                len,
139            } => Ok((-i64::try_from(offset_from_end)?, len)),
140        }
141    }
142}
143
144impl From<Slice> for Range<usize> {
145    fn from(value: Slice) -> Self {
146        match value {
147            Slice::Positive { offset, len } => offset..offset.checked_add(len).unwrap(),
148            Slice::Negative { .. } => panic!("cannot convert negative slice into range"),
149        }
150    }
151}
152
153#[cfg(test)]
154mod tests {
155    use super::Slice;
156
157    #[test]
158    fn test_slice_offset() {
159        assert_eq!(
160            Slice::Positive { offset: 3, len: 10 }.offsetted(1),
161            Slice::Positive { offset: 2, len: 10 }
162        );
163        assert_eq!(
164            Slice::Positive { offset: 3, len: 10 }.offsetted(5),
165            Slice::Positive { offset: 0, len: 8 }
166        );
167    }
168
169    #[test]
170    fn test_slice_restrict_to_bounds() {
171        assert_eq!(
172            Slice::Positive { offset: 3, len: 10 }.restrict_to_bounds(7),
173            Slice::Positive { offset: 3, len: 4 },
174        );
175        assert_eq!(
176            Slice::Positive { offset: 3, len: 10 }.restrict_to_bounds(0),
177            Slice::Positive { offset: 0, len: 0 },
178        );
179        assert_eq!(
180            Slice::Positive { offset: 3, len: 10 }.restrict_to_bounds(1),
181            Slice::Positive { offset: 1, len: 0 },
182        );
183        assert_eq!(
184            Slice::Positive { offset: 2, len: 0 }.restrict_to_bounds(10),
185            Slice::Positive { offset: 2, len: 0 },
186        );
187        assert_eq!(
188            Slice::Negative {
189                offset_from_end: 3,
190                len: 1
191            }
192            .restrict_to_bounds(4),
193            Slice::Positive { offset: 1, len: 1 },
194        );
195        assert_eq!(
196            Slice::Negative {
197                offset_from_end: 3,
198                len: 1
199            }
200            .restrict_to_bounds(1),
201            Slice::Positive { offset: 0, len: 0 },
202        );
203    }
204}
polars_utils/slice_enum.rs

polars_utils/
slice_enum.rs
