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 } => {
164                offset
165                    ..offset
166                        .checked_add(len)
167                        // Infer if no len was specified, in which case we set to range end to the max possible.
168                        .or((len == usize::MAX).then_some(usize::MAX))
169                        .unwrap()
170            },
171            Slice::Negative { .. } => panic!("cannot convert negative slice into range"),
172        }
173    }
174}
175
176#[cfg(test)]
177mod tests {
178    use super::Slice;
179
180    #[test]
181    fn test_slice_offset() {
182        assert_eq!(
183            Slice::Positive { offset: 3, len: 10 }.offsetted(1),
184            Slice::Positive { offset: 2, len: 10 }
185        );
186        assert_eq!(
187            Slice::Positive { offset: 3, len: 10 }.offsetted(5),
188            Slice::Positive { offset: 0, len: 8 }
189        );
190    }
191
192    #[test]
193    fn test_slice_restrict_to_bounds() {
194        assert_eq!(
195            Slice::Positive { offset: 3, len: 10 }.restrict_to_bounds(7),
196            Slice::Positive { offset: 3, len: 4 },
197        );
198        assert_eq!(
199            Slice::Positive { offset: 3, len: 10 }.restrict_to_bounds(0),
200            Slice::Positive { offset: 0, len: 0 },
201        );
202        assert_eq!(
203            Slice::Positive { offset: 3, len: 10 }.restrict_to_bounds(1),
204            Slice::Positive { offset: 1, len: 0 },
205        );
206        assert_eq!(
207            Slice::Positive { offset: 2, len: 0 }.restrict_to_bounds(10),
208            Slice::Positive { offset: 2, len: 0 },
209        );
210        assert_eq!(
211            Slice::Negative {
212                offset_from_end: 3,
213                len: 2
214            }
215            .restrict_to_bounds(4),
216            Slice::Positive { offset: 1, len: 2 },
217        );
218        assert_eq!(
219            Slice::Negative {
220                offset_from_end: 3,
221                len: 2
222            }
223            .restrict_to_bounds(3),
224            Slice::Positive { offset: 0, len: 2 },
225        );
226        assert_eq!(
227            Slice::Negative {
228                offset_from_end: 3,
229                len: 2
230            }
231            .restrict_to_bounds(2),
232            Slice::Positive { offset: 0, len: 1 },
233        );
234        assert_eq!(
235            Slice::Negative {
236                offset_from_end: 3,
237                len: 2
238            }
239            .restrict_to_bounds(1),
240            Slice::Positive { offset: 0, len: 0 },
241        );
242    }
243}
polars_utils/slice_enum.rs

polars_utils/
slice_enum.rs
