polars_utils/
slice_enum.rs

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

polars_utils/
slice_enum.rs
