arrow_buffer/buffer/
run.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use crate::ArrowNativeType;
19use crate::buffer::ScalarBuffer;
20
21/// A buffer of monotonically increasing, positive integers used to store run-ends.
22///
23/// Used to compactly represent runs of the same value. Values being represented
24/// are stored in a separate buffer from this struct. See [`RunArray`] for an example
25/// of how this is used with a companion array to represent the values.
26///
27/// # Logical vs Physical
28///
29/// Physically, each value in the `run_ends` buffer is the cumulative length of
30/// all runs in the logical representation, up to that physical index. Consider
31/// the following example:
32///
33/// ```text
34///           physical                        logical
35///     ┌─────────┬─────────┐           ┌─────────┬─────────┐
36///     │    3    │    0    │ ◄──────┬─ │    A    │    0    │
37///     ├─────────┼─────────┤        │  ├─────────┼─────────┤
38///     │    4    │    1    │ ◄────┐ ├─ │    A    │    1    │
39///     ├─────────┼─────────┤      │ │  ├─────────┼─────────┤
40///     │    6    │    2    │ ◄──┐ │ └─ │    A    │    2    │
41///     └─────────┴─────────┘    │ │    ├─────────┼─────────┤
42///      run-ends    index       │ └─── │    B    │    3    │
43///                              │      ├─────────┼─────────┤
44///      logical_offset = 0      ├───── │    C    │    4    │
45///      logical_length = 6      │      ├─────────┼─────────┤
46///                              └───── │    C    │    5    │
47///                                     └─────────┴─────────┘
48///                                       values     index
49/// ```
50///
51/// A [`RunEndBuffer`] is physically the buffer and offset with length on the left.
52/// In this case, the offset and length represent the whole buffer, so it is essentially
53/// unsliced. See the section below on slicing for more details on how this buffer
54/// handles slicing.
55///
56/// This means that multiple logical values are represented in the same physical index,
57/// and multiple logical indices map to the same physical index. The [`RunEndBuffer`]
58/// containing `[3, 4, 6]` is essentially the physical indices `[0, 0, 0, 1, 2, 2]`,
59/// and having a separately stored buffer of values such as `[A, B, C]` can turn
60/// this into a representation of `[A, A, A, B, C, C]`.
61///
62/// # Slicing
63///
64/// In order to provide zero-copy slicing, this struct stores a separate **logical**
65/// offset and length. Consider the following example:
66///
67/// ```text
68///           physical                        logical
69///     ┌─────────┬─────────┐           ┌ ─ ─ ─ ─ ┬ ─ ─ ─ ─ ┐
70///     │    3    │    0    │ ◄──────┐       A         0
71///     ├─────────┼─────────┤        │  ├── ─ ─ ─ ┼ ─ ─ ─ ─ ┤
72///     │    4    │    1    │ ◄────┐ │       A         1
73///     ├─────────┼─────────┤      │ │  ├─────────┼─────────┤
74///     │    6    │    2    │ ◄──┐ │ └─ │    A    │    2    │◄─── logical_offset
75///     └─────────┴─────────┘    │ │    ├─────────┼─────────┤
76///      run-ends    index       │ └─── │    B    │    3    │
77///                              │      ├─────────┼─────────┤
78///      logical_offset = 2      └───── │    C    │    4    │
79///      logical_length = 3             ├─────────┼─────────┤
80///                                          C         5     ◄─── logical_offset + logical_length
81///                                     └ ─ ─ ─ ─ ┴ ─ ─ ─ ─ ┘
82///                                       values     index
83/// ```
84///
85/// The physical `run_ends` [`ScalarBuffer`] remains unchanged, in order to facilitate
86/// zero-copy. However, we now offset into the **logical** representation with an
87/// accompanying length. This allows us to represent values `[A, B, C]` using physical
88/// indices `0, 1, 2` with the same underlying physical buffer, at the cost of two
89/// extra `usize`s to represent the logical slice that was taken.
90///
91/// (A [`RunEndBuffer`] is considered unsliced when `logical_offset` is `0` and
92/// `logical_length` is equal to the last value in `run_ends`)
93///
94/// [`RunArray`]: https://docs.rs/arrow/latest/arrow/array/struct.RunArray.html
95/// [Run-End encoded layout]: https://arrow.apache.org/docs/format/Columnar.html#run-end-encoded-layout
96#[derive(Debug, Clone)]
97pub struct RunEndBuffer<E: ArrowNativeType> {
98    run_ends: ScalarBuffer<E>,
99    logical_length: usize,
100    logical_offset: usize,
101}
102
103impl<E> RunEndBuffer<E>
104where
105    E: ArrowNativeType,
106{
107    /// Create a new [`RunEndBuffer`] from a [`ScalarBuffer`], `logical_offset`
108    /// and `logical_length`.
109    ///
110    /// # Panics
111    ///
112    /// Panics if the [`RunArray` invariants] are not satisfied:
113    ///
114    /// - `run_ends` does not contain strictly increasing values greater than zero
115    /// - `logical_length` is not 0 and the last value of `run_ends` is less than
116    ///   `logical_offset + logical_length` (an empty slice is always accepted)
117    ///
118    /// [`RunArray` invariants]: https://docs.rs/arrow/latest/arrow/array/struct.RunArray.html#offset-and-len
119    pub fn new(run_ends: ScalarBuffer<E>, logical_offset: usize, logical_length: usize) -> Self {
120        assert!(
121            run_ends.windows(2).all(|w| w[0] < w[1]),
122            "run-ends not strictly increasing"
123        );
124
125        if logical_length != 0 {
126            assert!(!run_ends.is_empty(), "non-empty slice but empty run-ends");
127            let end = E::from_usize(logical_offset.saturating_add(logical_length)).unwrap();
128            assert!(
129                *run_ends.first().unwrap() > E::usize_as(0),
130                "run-ends not greater than 0"
131            );
132            assert!(
133                *run_ends.last().unwrap() >= end,
134                "slice beyond bounds of run-ends"
135            );
136        }
137
138        Self {
139            run_ends,
140            logical_length,
141            logical_offset,
142        }
143    }
144
145    /// Create a new [`RunEndBuffer`] from a [`ScalarBuffer`], `logical_offset`
146    /// and `logical_length`.
147    ///
148    /// # Safety
149    ///
150    /// `run_ends` must satisfy the [`RunArray` invariants]:
151    ///
152    /// - `run_ends` must contain strictly increasing values greater than zero
153    /// - If `logical_length` is not 0, the last value of `run_ends` must be
154    ///   greater than or equal to `logical_offset + logical_length`
155    ///
156    /// [`RunArray` invariants]: https://docs.rs/arrow/latest/arrow/array/struct.RunArray.html#offset-and-len
157    pub unsafe fn new_unchecked(
158        run_ends: ScalarBuffer<E>,
159        logical_offset: usize,
160        logical_length: usize,
161    ) -> Self {
162        Self {
163            run_ends,
164            logical_length,
165            logical_offset,
166        }
167    }
168
169    /// Returns the logical offset into the run-ends stored by this buffer.
170    #[inline]
171    pub fn offset(&self) -> usize {
172        self.logical_offset
173    }
174
175    /// Returns the logical length of the run-ends stored by this buffer.
176    #[inline]
177    pub fn len(&self) -> usize {
178        self.logical_length
179    }
180
181    /// Returns true if this buffer is logically empty.
182    #[inline]
183    pub fn is_empty(&self) -> bool {
184        self.logical_length == 0
185    }
186
187    /// Free up unused memory.
188    pub fn shrink_to_fit(&mut self) {
189        // TODO(emilk): we could shrink even more in the case where we are a small sub-slice of the full buffer
190        self.run_ends.shrink_to_fit();
191    }
192
193    /// Returns the physical (**unsliced**) run ends of this buffer.
194    ///
195    /// Take care when operating on these values as it doesn't take into account
196    /// any logical slicing that may have occurred.
197    #[inline]
198    pub fn values(&self) -> &[E] {
199        &self.run_ends
200    }
201
202    /// Returns an iterator yielding run ends adjusted for the logical slice.
203    ///
204    /// Each yielded value is subtracted by the [`logical_offset`] and capped
205    /// at the [`logical_length`].
206    ///
207    /// [`logical_offset`]: Self::offset
208    /// [`logical_length`]: Self::len
209    pub fn sliced_values(&self) -> impl Iterator<Item = E> + '_ {
210        let offset = self.logical_offset;
211        let len = self.logical_length;
212        // Doing this roundabout way since the iterator type we return must be
213        // the same (i.e. cannot use std::iter::empty())
214        let physical_slice = if self.is_empty() {
215            &self.run_ends[0..0]
216        } else {
217            let start = self.get_start_physical_index();
218            let end = self.get_end_physical_index();
219            &self.run_ends[start..=end]
220        };
221        physical_slice.iter().map(move |&val| {
222            // `len` is at most the largest run end, so it always fits in `E`
223            let val = val.as_usize().saturating_sub(offset).min(len);
224            E::usize_as(val)
225        })
226    }
227
228    /// Returns the maximum run-end encoded in the underlying buffer; that is, the
229    /// last physical run of the buffer. This does not take into account any logical
230    /// slicing that may have occurred.
231    #[inline]
232    pub fn max_value(&self) -> usize {
233        self.values().last().copied().unwrap_or_default().as_usize()
234    }
235
236    /// Performs a binary search to find the physical index for the given logical
237    /// index.
238    ///
239    /// Useful for extracting the corresponding physical `run_ends` when this buffer
240    /// is logically sliced.
241    ///
242    /// The result is arbitrary if `logical_index >= self.len()`.
243    pub fn get_physical_index(&self, logical_index: usize) -> usize {
244        let logical_index = self.logical_offset + logical_index;
245        let cmp = |p: &E| p.as_usize().cmp(&logical_index);
246
247        match self.run_ends.binary_search_by(cmp) {
248            Ok(idx) => idx + 1,
249            Err(idx) => idx,
250        }
251    }
252
253    /// Returns the physical index at which the logical array starts.
254    ///
255    /// The same as calling `get_physical_index(0)` but with a fast path if the
256    /// buffer is not logically sliced, in which case it always returns `0`.
257    pub fn get_start_physical_index(&self) -> usize {
258        if self.logical_offset == 0 || self.logical_length == 0 {
259            return 0;
260        }
261        // Fallback to binary search
262        self.get_physical_index(0)
263    }
264
265    /// Returns the physical index at which the logical array ends.
266    ///
267    /// The same as calling `get_physical_index(length - 1)` but with a fast path
268    /// if the buffer is not logically sliced, in which case it returns `length - 1`.
269    pub fn get_end_physical_index(&self) -> usize {
270        if self.logical_length == 0 {
271            return 0;
272        }
273        if self.max_value() == self.logical_offset + self.logical_length {
274            return self.values().len() - 1;
275        }
276        // Fallback to binary search
277        self.get_physical_index(self.logical_length - 1)
278    }
279
280    /// Slices this [`RunEndBuffer`] by the provided `logical_offset` and `logical_length`.
281    ///
282    /// # Panics
283    ///
284    /// - Specified slice (`logical_offset` + `logical_length`) exceeds existing
285    ///   logical length
286    pub fn slice(&self, logical_offset: usize, logical_length: usize) -> Self {
287        assert!(
288            logical_offset.saturating_add(logical_length) <= self.logical_length,
289            "the length + offset of the sliced RunEndBuffer cannot exceed the existing length"
290        );
291        Self {
292            run_ends: self.run_ends.clone(),
293            logical_offset: self.logical_offset + logical_offset,
294            logical_length,
295        }
296    }
297
298    /// Returns the inner [`ScalarBuffer`].
299    pub fn inner(&self) -> &ScalarBuffer<E> {
300        &self.run_ends
301    }
302
303    /// Returns the inner [`ScalarBuffer`], consuming self.
304    pub fn into_inner(self) -> ScalarBuffer<E> {
305        self.run_ends
306    }
307
308    /// Claim memory used by this buffer in the provided memory pool.
309    #[cfg(feature = "pool")]
310    pub fn claim(&self, pool: &dyn crate::MemoryPool) {
311        self.run_ends.claim(pool);
312    }
313
314    /// Returns the physical indices corresponding to the provided logical indices.
315    ///
316    /// Given a slice of logical indices, this method returns a `Vec` containing the
317    /// corresponding physical indices into the run-ends buffer.
318    ///
319    /// This method operates by iterating the logical indices in sorted order, instead of
320    /// finding the physical index for each logical index using binary search via
321    /// the function [`RunEndBuffer::get_physical_index`].
322    ///
323    /// Running benchmarks on both approaches showed that the approach used here
324    /// scaled well for larger inputs.
325    ///
326    /// See <https://github.com/apache/arrow-rs/pull/3622#issuecomment-1407753727> for more details.
327    ///
328    /// # Errors
329    ///
330    /// If any logical index is out of bounds (>= self.len()), returns an error containing the invalid index.
331    #[inline]
332    pub fn get_physical_indices<I>(&self, logical_indices: &[I]) -> Result<Vec<usize>, I>
333    where
334        I: ArrowNativeType,
335    {
336        let len = self.len();
337        let offset = self.offset();
338
339        let indices_len = logical_indices.len();
340
341        if indices_len == 0 {
342            return Ok(vec![]);
343        }
344
345        // `ordered_indices` store index into `logical_indices` and can be used
346        // to iterate `logical_indices` in sorted order.
347        let mut ordered_indices: Vec<usize> = (0..indices_len).collect();
348
349        // Instead of sorting `logical_indices` directly, sort the `ordered_indices`
350        // whose values are index of `logical_indices`
351        ordered_indices.sort_unstable_by_key(|&idx| logical_indices[idx].as_usize());
352
353        // Return early if all the logical indices cannot be converted to physical indices.
354        // `ordered_indices` has `indices_len` entries, and the empty case returned above.
355        let largest_logical_index = logical_indices[ordered_indices[indices_len - 1]];
356        if largest_logical_index.as_usize() >= len {
357            return Err(largest_logical_index);
358        }
359
360        // Skip some physical indices based on offset.
361        let skip_value = self.get_start_physical_index();
362
363        let mut physical_indices = vec![0; indices_len];
364
365        let mut ordered_index = 0_usize;
366        for (physical_index, run_end) in self.values().iter().enumerate().skip(skip_value) {
367            // Get the run end index (relative to offset) of current physical index
368            let run_end_value = run_end.as_usize() - offset;
369
370            // All the `logical_indices` that are less than current run end index
371            // belongs to current physical index.
372            while ordered_index < indices_len
373                && logical_indices[ordered_indices[ordered_index]].as_usize() < run_end_value
374            {
375                physical_indices[ordered_indices[ordered_index]] = physical_index;
376                ordered_index += 1;
377            }
378        }
379
380        // If there are input values >= run_ends.last_value then we'll not be able to convert
381        // all logical indices to physical indices.
382        if ordered_index < logical_indices.len() {
383            return Err(logical_indices[ordered_indices[ordered_index]]);
384        }
385        Ok(physical_indices)
386    }
387}
388
389#[cfg(test)]
390mod tests {
391    use crate::buffer::RunEndBuffer;
392
393    #[test]
394    fn test_zero_length_slice() {
395        let buffer = RunEndBuffer::new(vec![1_i32, 4_i32].into(), 0, 4);
396        assert_eq!(buffer.get_start_physical_index(), 0);
397        assert_eq!(buffer.get_end_physical_index(), 1);
398        assert_eq!(buffer.get_physical_index(3), 1);
399
400        for offset in 0..4 {
401            let sliced = buffer.slice(offset, 0);
402            assert_eq!(sliced.get_start_physical_index(), 0);
403            assert_eq!(sliced.get_end_physical_index(), 0);
404        }
405
406        let buffer = RunEndBuffer::new(Vec::<i32>::new().into(), 0, 0);
407        assert_eq!(buffer.get_start_physical_index(), 0);
408        assert_eq!(buffer.get_end_physical_index(), 0);
409    }
410
411    #[test]
412    fn test_sliced_values() {
413        // [0, 0, 1, 2, 2, 2]
414        let buffer = RunEndBuffer::new(vec![2i32, 3, 6].into(), 0, 6);
415
416        // Slice: [0, 1, 2, 2] start: 1, len: 4
417        // Logical indices: 1, 2, 3, 4
418        // Original run ends: [2, 3, 6]
419        // Adjusted: [2-1, 3-1, 6-1] capped at 4 -> [1, 2, 4]
420        let sliced = buffer.slice(1, 4);
421        let sliced_values: Vec<i32> = sliced.sliced_values().collect();
422        assert_eq!(sliced_values, &[1, 2, 4]);
423
424        // Slice: [2, 2] start: 4, len: 2
425        // Original run ends: [2, 3, 6]
426        // Slicing at 4 means we only have the last run (physical index 2, which ends at 6)
427        // Adjusted: [6-4] capped at 2 -> [2]
428        let sliced = buffer.slice(4, 2);
429        let sliced_values: Vec<i32> = sliced.sliced_values().collect();
430        assert_eq!(sliced_values, &[2]);
431    }
432}
arrow_buffer/buffer/run.rs

arrow_buffer/buffer/
run.rs
