arrow_array/
run_iterator.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! Idiomatic iterator for [`RunArray`](crate::RunArray)
6
7use crate::{Array, TypedRunArray, array::ArrayAccessor, types::RunEndIndexType};
8use arrow_buffer::ArrowNativeType;
9
10/// The [`RunArrayIter`] provides an idiomatic way to iterate over the run array.
11/// It returns Some(T) if there is a value or None if the value is null.
12///
13/// The iterator comes with a cost as it has to iterate over three arrays to determine
14/// the value to be returned. The run_ends array is used to determine the index of the value.
15/// The nulls array is used to determine if the value is null and the values array is used to
16/// get the value.
17///
18/// Unlike other iterators in this crate, [`RunArrayIter`] does not use [`ArrayAccessor`]
19/// because the run array accessor does binary search to access each value which is too slow.
20/// The run array iterator can determine the next value in constant time.
21///
22#[derive(Debug)]
23pub struct RunArrayIter<'a, R, V>
24where
25    R: RunEndIndexType,
26    V: Sync + Send,
27    &'a V: ArrayAccessor,
28    <&'a V as ArrayAccessor>::Item: Default,
29{
30    array: TypedRunArray<'a, R, V>,
31    current_front_logical: usize,
32    current_front_physical: usize,
33    current_back_logical: usize,
34    current_back_physical: usize,
35}
36
37impl<'a, R, V> RunArrayIter<'a, R, V>
38where
39    R: RunEndIndexType,
40    V: Sync + Send,
41    &'a V: ArrayAccessor,
42    <&'a V as ArrayAccessor>::Item: Default,
43{
44    /// create a new iterator
45    pub fn new(array: TypedRunArray<'a, R, V>) -> Self {
46        let current_front_physical = array.run_array().get_start_physical_index();
47        let current_back_physical = array.run_array().get_end_physical_index() + 1;
48        RunArrayIter {
49            array,
50            current_front_logical: array.offset(),
51            current_front_physical,
52            current_back_logical: array.offset() + array.len(),
53            current_back_physical,
54        }
55    }
56}
57
58impl<'a, R, V> Iterator for RunArrayIter<'a, R, V>
59where
60    R: RunEndIndexType,
61    V: Sync + Send,
62    &'a V: ArrayAccessor,
63    <&'a V as ArrayAccessor>::Item: Default,
64{
65    type Item = Option<<&'a V as ArrayAccessor>::Item>;
66
67    #[inline]
68    fn next(&mut self) -> Option<Self::Item> {
69        if self.current_front_logical == self.current_back_logical {
70            return None;
71        }
72
73        // If current logical index is greater than current run end index then increment
74        // the physical index.
75        let run_ends = self.array.run_ends().values();
76        if self.current_front_logical >= run_ends[self.current_front_physical].as_usize() {
77            // As the run_ends is expected to be strictly increasing, there
78            // should be at least one logical entry in one physical entry. Because of this
79            // reason the next value can be accessed by incrementing physical index once.
80            self.current_front_physical += 1;
81        }
82        if self.array.values().is_null(self.current_front_physical) {
83            self.current_front_logical += 1;
84            Some(None)
85        } else {
86            self.current_front_logical += 1;
87            // Safety:
88            // The self.current_physical is kept within bounds of self.current_logical.
89            // The self.current_logical will not go out of bounds because of the check
90            // `self.current_logical = self.current_end_logical` above.
91            unsafe {
92                Some(Some(
93                    self.array
94                        .values()
95                        .value_unchecked(self.current_front_physical),
96                ))
97            }
98        }
99    }
100
101    fn size_hint(&self) -> (usize, Option<usize>) {
102        (
103            self.current_back_logical - self.current_front_logical,
104            Some(self.current_back_logical - self.current_front_logical),
105        )
106    }
107}
108
109impl<'a, R, V> DoubleEndedIterator for RunArrayIter<'a, R, V>
110where
111    R: RunEndIndexType,
112    V: Sync + Send,
113    &'a V: ArrayAccessor,
114    <&'a V as ArrayAccessor>::Item: Default,
115{
116    fn next_back(&mut self) -> Option<Self::Item> {
117        if self.current_back_logical == self.current_front_logical {
118            return None;
119        }
120
121        self.current_back_logical -= 1;
122
123        let run_ends = self.array.run_ends().values();
124        if self.current_back_physical > 0
125            && self.current_back_logical < run_ends[self.current_back_physical - 1].as_usize()
126        {
127            // As the run_ends is expected to be strictly increasing, there
128            // should be at least one logical entry in one physical entry. Because of this
129            // reason the next value can be accessed by decrementing physical index once.
130            self.current_back_physical -= 1;
131        }
132        Some(if self.array.values().is_null(self.current_back_physical) {
133            None
134        } else {
135            // Safety:
136            // The check `self.current_end_physical > 0` ensures the value will not underflow.
137            // Also self.current_end_physical starts with array.len() and
138            // decrements based on the bounds of self.current_end_logical.
139            unsafe {
140                Some(
141                    self.array
142                        .values()
143                        .value_unchecked(self.current_back_physical),
144                )
145            }
146        })
147    }
148}
149
150/// all arrays have known size.
151impl<'a, R, V> ExactSizeIterator for RunArrayIter<'a, R, V>
152where
153    R: RunEndIndexType,
154    V: Sync + Send,
155    &'a V: ArrayAccessor,
156    <&'a V as ArrayAccessor>::Item: Default,
157{
158}
159
160#[cfg(test)]
161mod tests {
162    use rand::{RngExt, rng, seq::SliceRandom};
163
164    use crate::{
165        Array, Int64RunArray, PrimitiveArray, RunArray,
166        array::{Int32Array, StringArray},
167        builder::PrimitiveRunBuilder,
168        types::{Int16Type, Int32Type},
169    };
170
171    fn build_input_array(size: usize) -> Vec<Option<i32>> {
172        // The input array is created by shuffling and repeating
173        // the seed values random number of times.
174        let mut seed: Vec<Option<i32>> = vec![
175            None,
176            None,
177            None,
178            Some(1),
179            Some(2),
180            Some(3),
181            Some(4),
182            Some(5),
183            Some(6),
184            Some(7),
185            Some(8),
186            Some(9),
187        ];
188        let mut result: Vec<Option<i32>> = Vec::with_capacity(size);
189        let mut ix = 0;
190        let mut rng = rng();
191        // run length can go up to 8. Cap the max run length for smaller arrays to size / 2.
192        let max_run_length = 8_usize.min(1_usize.max(size / 2));
193        while result.len() < size {
194            // shuffle the seed array if all the values are iterated.
195            if ix == 0 {
196                seed.shuffle(&mut rng);
197            }
198            // repeat the items between 1 and 8 times. Cap the length for smaller sized arrays
199            let num = max_run_length.min(rng.random_range(1..=max_run_length));
200            for _ in 0..num {
201                result.push(seed[ix]);
202            }
203            ix += 1;
204            if ix == seed.len() {
205                ix = 0
206            }
207        }
208        result.resize(size, None);
209        result
210    }
211
212    #[test]
213    fn test_primitive_array_iter_round_trip() {
214        let mut input_vec = vec![
215            Some(32),
216            Some(32),
217            None,
218            Some(64),
219            Some(64),
220            Some(64),
221            Some(72),
222        ];
223        let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
224        builder.extend(input_vec.iter().copied());
225        let ree_array = builder.finish();
226        let ree_array = ree_array.downcast::<Int32Array>().unwrap();
227
228        let output_vec: Vec<Option<i32>> = ree_array.into_iter().collect();
229        assert_eq!(input_vec, output_vec);
230
231        let rev_output_vec: Vec<Option<i32>> = ree_array.into_iter().rev().collect();
232        input_vec.reverse();
233        assert_eq!(input_vec, rev_output_vec);
234    }
235
236    #[test]
237    fn test_double_ended() {
238        let input_vec = vec![
239            Some(32),
240            Some(32),
241            None,
242            Some(64),
243            Some(64),
244            Some(64),
245            Some(72),
246        ];
247        let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
248        builder.extend(input_vec);
249        let ree_array = builder.finish();
250        let ree_array = ree_array.downcast::<Int32Array>().unwrap();
251
252        let mut iter = ree_array.into_iter();
253        assert_eq!(Some(Some(32)), iter.next());
254        assert_eq!(Some(Some(72)), iter.next_back());
255        assert_eq!(Some(Some(32)), iter.next());
256        assert_eq!(Some(Some(64)), iter.next_back());
257        assert_eq!(Some(None), iter.next());
258        assert_eq!(Some(Some(64)), iter.next_back());
259        assert_eq!(Some(Some(64)), iter.next());
260        assert_eq!(None, iter.next_back());
261        assert_eq!(None, iter.next());
262    }
263
264    #[test]
265    fn test_run_iterator_comprehensive() {
266        // Test forward and backward iterator for different array lengths.
267        let logical_lengths = vec![1_usize, 2, 3, 4, 15, 16, 17, 63, 64, 65];
268
269        for logical_len in logical_lengths {
270            let input_array = build_input_array(logical_len);
271
272            let mut run_array_builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
273            run_array_builder.extend(input_array.iter().copied());
274            let run_array = run_array_builder.finish();
275            let typed_array = run_array.downcast::<Int32Array>().unwrap();
276
277            // test forward iterator
278            let mut input_iter = input_array.iter().copied();
279            let mut run_array_iter = typed_array.into_iter();
280            for _ in 0..logical_len {
281                assert_eq!(input_iter.next(), run_array_iter.next());
282            }
283            assert_eq!(None, run_array_iter.next());
284
285            // test reverse iterator
286            let mut input_iter = input_array.iter().rev().copied();
287            let mut run_array_iter = typed_array.into_iter().rev();
288            for _ in 0..logical_len {
289                assert_eq!(input_iter.next(), run_array_iter.next());
290            }
291            assert_eq!(None, run_array_iter.next());
292        }
293    }
294
295    #[test]
296    fn test_string_array_iter_round_trip() {
297        let input_vec = vec!["ab", "ab", "ba", "cc", "cc"];
298        let input_ree_array: Int64RunArray = input_vec.into_iter().collect();
299        let string_ree_array = input_ree_array.downcast::<StringArray>().unwrap();
300
301        // to and from iter, with a +1
302        let result: Vec<Option<String>> = string_ree_array
303            .into_iter()
304            .map(|e| {
305                e.map(|e| {
306                    let mut a = e.to_string();
307                    a.push('b');
308                    a
309                })
310            })
311            .collect();
312
313        let result_asref: Vec<Option<&str>> = result.iter().map(|f| f.as_deref()).collect();
314
315        let expected_vec = vec![
316            Some("abb"),
317            Some("abb"),
318            Some("bab"),
319            Some("ccb"),
320            Some("ccb"),
321        ];
322
323        assert_eq!(expected_vec, result_asref);
324    }
325
326    #[test]
327    #[cfg_attr(miri, ignore)] // Takes too long
328    fn test_sliced_run_array_iterator() {
329        let total_len = 80;
330        let input_array = build_input_array(total_len);
331
332        // Encode the input_array to run array
333        let mut builder =
334            PrimitiveRunBuilder::<Int16Type, Int32Type>::with_capacity(input_array.len());
335        builder.extend(input_array.iter().copied());
336        let run_array = builder.finish();
337
338        // test for all slice lengths.
339        for slice_len in 1..=total_len {
340            // test for offset = 0, slice length = slice_len
341            let sliced_run_array: RunArray<Int16Type> =
342                run_array.slice(0, slice_len).into_data().into();
343            let sliced_typed_run_array = sliced_run_array
344                .downcast::<PrimitiveArray<Int32Type>>()
345                .unwrap();
346
347            // Iterate on sliced typed run array
348            let actual: Vec<Option<i32>> = sliced_typed_run_array.into_iter().collect();
349            let expected: Vec<Option<i32>> = input_array.iter().take(slice_len).copied().collect();
350            assert_eq!(expected, actual);
351
352            // test for offset = total_len - slice_len, length = slice_len
353            let sliced_run_array: RunArray<Int16Type> = run_array
354                .slice(total_len - slice_len, slice_len)
355                .into_data()
356                .into();
357            let sliced_typed_run_array = sliced_run_array
358                .downcast::<PrimitiveArray<Int32Type>>()
359                .unwrap();
360
361            // Iterate on sliced typed run array
362            let actual: Vec<Option<i32>> = sliced_typed_run_array.into_iter().collect();
363            let expected: Vec<Option<i32>> = input_array
364                .iter()
365                .skip(total_len - slice_len)
366                .copied()
367                .collect();
368            assert_eq!(expected, actual);
369        }
370    }
371}
arrow_array/run_iterator.rs

arrow_array/
run_iterator.rs
