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        Ok(match value {
135            Slice::Positive { offset, len } => (i64::try_from(offset)?, len),
136            Slice::Negative {
137                offset_from_end,
138                len,
139            } => (-i64::try_from(offset_from_end)?, len),
140        })
141    }
142}
143
144impl From<Slice> for (i128, i128) {
145    fn from(value: Slice) -> Self {
146        match value {
147            Slice::Positive { offset, len } => (
148                i128::try_from(offset).unwrap(),
149                i128::try_from(len).unwrap(),
150            ),
151            Slice::Negative {
152                offset_from_end,
153                len,
154            } => (
155                -i128::try_from(offset_from_end).unwrap(),
156                i128::try_from(len).unwrap(),
157            ),
158        }
159    }
160}
161
162impl From<Slice> for Range<usize> {
163    fn from(value: Slice) -> Self {
164        match value {
165            Slice::Positive { offset, len } => offset..offset.checked_add(len).unwrap(),
166            Slice::Negative { .. } => panic!("cannot convert negative slice into range"),
167        }
168    }
169}
170
171#[cfg(test)]
172mod tests {
173    use super::Slice;
174
175    #[test]
176    fn test_slice_offset() {
177        assert_eq!(
178            Slice::Positive { offset: 3, len: 10 }.offsetted(1),
179            Slice::Positive { offset: 2, len: 10 }
180        );
181        assert_eq!(
182            Slice::Positive { offset: 3, len: 10 }.offsetted(5),
183            Slice::Positive { offset: 0, len: 8 }
184        );
185    }
186
187    #[test]
188    fn test_slice_restrict_to_bounds() {
189        assert_eq!(
190            Slice::Positive { offset: 3, len: 10 }.restrict_to_bounds(7),
191            Slice::Positive { offset: 3, len: 4 },
192        );
193        assert_eq!(
194            Slice::Positive { offset: 3, len: 10 }.restrict_to_bounds(0),
195            Slice::Positive { offset: 0, len: 0 },
196        );
197        assert_eq!(
198            Slice::Positive { offset: 3, len: 10 }.restrict_to_bounds(1),
199            Slice::Positive { offset: 1, len: 0 },
200        );
201        assert_eq!(
202            Slice::Positive { offset: 2, len: 0 }.restrict_to_bounds(10),
203            Slice::Positive { offset: 2, len: 0 },
204        );
205        assert_eq!(
206            Slice::Negative {
207                offset_from_end: 3,
208                len: 1
209            }
210            .restrict_to_bounds(4),
211            Slice::Positive { offset: 1, len: 1 },
212        );
213        assert_eq!(
214            Slice::Negative {
215                offset_from_end: 3,
216                len: 1
217            }
218            .restrict_to_bounds(1),
219            Slice::Positive { offset: 0, len: 0 },
220        );
221    }
222}
polars_utils/slice_enum.rs

polars_utils/
slice_enum.rs
