arrow_array/array/
run_array.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
5use std::any::Any;
6use std::sync::Arc;
7
8use arrow_buffer::{ArrowNativeType, BooleanBufferBuilder, NullBuffer, RunEndBuffer, ScalarBuffer};
9use arrow_data::{ArrayData, ArrayDataBuilder};
10use arrow_schema::{ArrowError, DataType, Field, FieldRef};
11
12use crate::{
13    Array, ArrayAccessor, ArrayRef, PrimitiveArray,
14    builder::StringRunBuilder,
15    make_array,
16    run_iterator::RunArrayIter,
17    types::{Int16Type, Int32Type, Int64Type, RunEndIndexType},
18};
19
20/// Recursively applies a function to the values of a RunEndEncoded array, preserving the run structure.
21///
22/// # Example
23///
24/// ```ignore
25/// let result = ree_recurse!(array, Int32Type, my_function)?;
26/// ```
27///
28/// This macro is useful for implementing functions that should work on the logical values
29/// of a REE array while preserving the run-end encoding structure.
30#[macro_export]
31macro_rules! ree_map {
32    ($array:expr, $run_type:ty, $func:expr) => {{
33        let ree = $array.as_run_opt::<$run_type>().unwrap();
34        let inner_values = $func(ree.values().as_ref())?;
35        Ok(std::sync::Arc::new(ree.with_values(inner_values)))
36    }};
37}
38
39/// An array of [run-end encoded values].
40///
41/// This encoding is variation on [run-length encoding (RLE)] and is good for representing
42/// data containing the same values repeated consecutively.
43///
44/// A [`RunArray`] consists of a `run_ends` buffer and a `values` array of equivalent
45/// lengths. The `run_ends` buffer stores the indexes at which the run ends. The
46/// `values` array stores the corresponding value of each run. The below example
47/// illustrates how a logical array is represented by a [`RunArray`]:
48///
49/// ```text
50/// ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─┐
51///   ┌─────────────────┐  ┌─────────┐       ┌─────────────────┐
52/// │ │        A        │  │    2    │ │     │        A        │
53///   ├─────────────────┤  ├─────────┤       ├─────────────────┤
54/// │ │        D        │  │    3    │ │     │        A        │    run length of 'A' = run_ends[0] - 0 = 2
55///   ├─────────────────┤  ├─────────┤       ├─────────────────┤
56/// │ │        B        │  │    6    │ │     │        D        │    run length of 'D' = run_ends[1] - run_ends[0] = 1
57///   └─────────────────┘  └─────────┘       ├─────────────────┤
58/// │        values          run_ends  │     │        B        │
59///                                          ├─────────────────┤
60/// └ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─┘     │        B        │
61///                                          ├─────────────────┤
62///                RunArray                  │        B        │    run length of 'B' = run_ends[2] - run_ends[1] = 3
63///          offset = 0, len = 6             └─────────────────┘
64///
65///                                             Logical array
66///                                                Contents
67/// ```
68///
69/// # `offset` and `len`
70///
71/// The [`Array::len`] and [`Array::offset`] for a `RunArray` refer to the
72/// **logical** length, not the `run_ends` and `values` child arrays. In the
73/// example above the `RunArray` has length `6` even though `run_ends` and
74/// `values` each contain `3` elements.
75///
76/// The values in `run_ends` represent indexes of the *unsliced* logical array.
77/// This allows zero-copy [`RunArray::slice`]: `run_ends` and `values` are left
78/// untouched and only `offset` and `len` are adjusted.
79///
80/// A [`RunArray`] must satisfy the following invariants:
81///
82/// * every run end is greater than zero and `run_ends` is strictly increasing
83/// * `run_ends` has no null values
84/// * `run_ends` and `values` have the same length
85/// * if `len > 0`, the last run end is greater than or equal to `offset + len`.
86///
87/// Note an empty `RunArray` (`len == 0`) is valid even when `run_ends` is empty
88/// or its last run end is less than `offset`. Note this is more permissive than
89/// the Arrow C++ implementation.
90///
91/// The example below shows the same array as above after calling `slice(1, 3)`
92///
93/// ```text
94/// ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─┐
95///   ┌─────────────────┐  ┌─────────┐       ┌ ─ ─ ─ ─ ─ ─ ─ ─ ┐
96/// │ │        A        │  │    2    │ │              A             logical index 0: before offset, not visible
97///   ├─────────────────┤  ├─────────┤       ┌─────────────────┐ ◄── offset = 1
98/// │ │        D        │  │    3    │ │     │        A        │    logical index 1 < run_ends[0] = 2, so 'A'
99///   ├─────────────────┤  ├─────────┤       ├─────────────────┤
100/// │ │        B        │  │    6    │ │     │        D        │    logical index 2 < run_ends[1] = 3, so 'D'
101///   └─────────────────┘  └─────────┘       ├─────────────────┤
102/// │        values          run_ends  │     │        B        │    logical index 3 < run_ends[2] = 6, so 'B'
103///                                          └─────────────────┘ ◄── offset + len = 4
104/// └ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─┘              B             logical index 4: after offset + len, not visible
105///                                          ├ ─ ─ ─ ─ ─ ─ ─ ─ ┤
106///                RunArray                           B             logical index 5: after offset + len, not visible
107///          offset = 1, len = 3             └ ─ ─ ─ ─ ─ ─ ─ ─ ┘
108///
109///                                             Logical array
110///                                                Contents
111/// ```
112///
113/// # Sliced child arrays
114///
115/// [`RunArray::slice`] doesn't slice child arrays, but a `RunArray` may be
116/// built (for example from [`ArrayData`]) from child arrays that are already
117/// sliced. In that case each child's own `offset` and `len` are respected,
118/// independently of the logical `offset` and `len` of the `RunArray`:
119///
120/// * Run `i` ends at `run_ends.value(i)`, the `i`th (logical) element
121/// * The value of run `i` is `values.value(i)`, the `i`th (logical) element
122///
123/// [run-end encoded values]: https://arrow.apache.org/docs/format/Columnar.html#run-end-encoded-layout
124/// [run-length encoding (RLE)]: https://en.wikipedia.org/wiki/Run-length_encoding
125pub struct RunArray<R: RunEndIndexType> {
126    data_type: DataType,
127    run_ends: RunEndBuffer<R::Native>,
128    values: ArrayRef,
129}
130
131impl<R: RunEndIndexType> Clone for RunArray<R> {
132    fn clone(&self) -> Self {
133        Self {
134            data_type: self.data_type.clone(),
135            run_ends: self.run_ends.clone(),
136            values: self.values.clone(),
137        }
138    }
139}
140
141impl<R: RunEndIndexType> RunArray<R> {
142    /// Calculates the logical length of the array encoded by treating the `run_ends`
143    /// array as if it were a [`RunEndBuffer`].
144    pub fn logical_len(run_ends: &PrimitiveArray<R>) -> usize {
145        let len = run_ends.len();
146        if len == 0 {
147            return 0;
148        }
149        run_ends.value(len - 1).as_usize()
150    }
151
152    /// Attempts to create a [`RunArray`] using the given `run_ends` and `values`.
153    ///
154    /// The resulting array is unsliced: its logical offset is `0` and its logical
155    /// length is the last value of `run_ends` (or `0` if `run_ends` is empty).
156    ///
157    /// # Errors
158    ///
159    /// Returns an error if the [`RunArray` invariants] are not satisfied:
160    ///
161    /// - If `run_ends` and `values` have different lengths
162    /// - If `run_ends` has any null values
163    /// - If `run_ends` doesn't consist of strictly increasing positive integers
164    ///
165    /// [`RunArray` invariants]: RunArray#offset-and-len
166    pub fn try_new(run_ends: &PrimitiveArray<R>, values: &dyn Array) -> Result<Self, ArrowError> {
167        // 1. run_ends and values must have the same (physical) length.
168        if run_ends.len() != values.len() {
169            return Err(ArrowError::InvalidArgumentError(format!(
170                "The run_ends array length should be the same as values array length. Run_ends array length is {}, values array length is {}",
171                run_ends.len(),
172                values.len()
173            )));
174        }
175
176        // 2. run_ends must not contain null values.
177        if run_ends.nulls().is_some() {
178            return Err(ArrowError::InvalidArgumentError(
179                "Found null values in run_ends array. The run_ends array should not have null values."
180                    .to_string(),
181            ));
182        }
183
184        // 3. run_ends must be strictly increasing, strictly positive integers.
185        if let Some(first) = run_ends.values().first() {
186            let mut prev_value = *first;
187            if prev_value <= R::Native::usize_as(0) {
188                return Err(ArrowError::InvalidArgumentError(format!(
189                    "The values in run_ends array should be strictly positive. Found value {prev_value:?} at index 0 that does not match the criteria."
190                )));
191            }
192
193            for (ix, &run_end) in run_ends.values().iter().enumerate().skip(1) {
194                if run_end <= prev_value {
195                    return Err(ArrowError::InvalidArgumentError(format!(
196                        "The values in run_ends array should be strictly increasing. Found value {run_end:?} at index {ix} with previous value {prev_value:?} that does not match the criteria."
197                    )));
198                }
199                prev_value = run_end;
200            }
201        }
202
203        let data_type = DataType::RunEndEncoded(
204            Arc::new(Field::new(
205                Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
206                run_ends.data_type().clone(),
207                false,
208            )),
209            Arc::new(Field::new(
210                Field::REE_VALUES_FIELD_DEFAULT_NAME,
211                values.data_type().clone(),
212                true,
213            )),
214        );
215
216        let logical_len = RunArray::logical_len(run_ends);
217        // Safety: validated above that the run ends are strictly increasing, strictly
218        // positive integers, so the last value equals the logical length.
219        let run_ends_buffer =
220            unsafe { RunEndBuffer::new_unchecked(run_ends.values().clone(), 0, logical_len) };
221
222        Ok(Self {
223            data_type,
224            run_ends: run_ends_buffer,
225            values: values.slice(0, values.len()),
226        })
227    }
228
229    /// Create a new [`RunArray`] from the provided parts, without validation
230    ///
231    /// # Safety
232    ///
233    /// Safe if [`Self::try_new`] would not error, that is if `run_ends` and
234    /// `values` satisfy the [`RunArray` invariants]
235    ///
236    /// [`RunArray` invariants]: RunArray#offset-and-len
237    pub unsafe fn new_unchecked(
238        data_type: DataType,
239        run_ends: RunEndBuffer<R::Native>,
240        values: ArrayRef,
241    ) -> Self {
242        if cfg!(feature = "force_validate") {
243            match &data_type {
244                DataType::RunEndEncoded(run_ends, values_field) => {
245                    assert!(!run_ends.is_nullable(), "run_ends should not be nullable");
246                    assert_eq!(
247                        run_ends.data_type(),
248                        &R::DATA_TYPE,
249                        "Incorrect run ends type"
250                    );
251                    assert_eq!(
252                        values_field.data_type(),
253                        values.data_type(),
254                        "Incorrect values type"
255                    );
256                }
257                _ => {
258                    panic!(
259                        "Invalid data type {data_type:?} for RunArray. Should be DataType::RunEndEncoded"
260                    );
261                }
262            }
263
264            let run_array = Self {
265                data_type,
266                run_ends,
267                values,
268            };
269
270            // Safety: `validate_data` checks below
271            //    1. The given array data has exactly two child arrays.
272            //    2. The first child array (run_ends) has valid data type.
273            //    3. run_ends array does not have null values
274            //    4. run_ends array has non-zero and strictly increasing values.
275            //    5. The length of run_ends array and values array are the same.
276            run_array
277                .to_data()
278                .validate_data()
279                .expect("RunArray data should be valid");
280
281            return run_array;
282        }
283
284        Self {
285            data_type,
286            run_ends,
287            values,
288        }
289    }
290
291    /// Deconstruct this array into its constituent parts
292    pub fn into_parts(self) -> (DataType, RunEndBuffer<R::Native>, ArrayRef) {
293        (self.data_type, self.run_ends, self.values)
294    }
295
296    /// The field that describes the run ends of this array.
297    pub fn run_ends_field(&self) -> &FieldRef {
298        match &self.data_type {
299            DataType::RunEndEncoded(f, _) => f,
300            _ => unreachable!(),
301        }
302    }
303
304    /// The field that describes the values of this array.
305    pub fn values_field(&self) -> &FieldRef {
306        match &self.data_type {
307            DataType::RunEndEncoded(_, f) => f,
308            _ => unreachable!(),
309        }
310    }
311
312    /// Returns a reference to the [`RunEndBuffer`].
313    pub fn run_ends(&self) -> &RunEndBuffer<R::Native> {
314        &self.run_ends
315    }
316
317    /// Returns a reference to the values array.
318    ///
319    /// Any slicing of this [`RunArray`] array is **not** applied to the returned
320    /// values here and must be handled separately.
321    pub fn values(&self) -> &ArrayRef {
322        &self.values
323    }
324
325    /// Returns a new [`RunArray`] with the same `run_ends` and the supplied `values`.
326    ///
327    /// # Panics
328    ///
329    /// Panics if `values.len()` does not equal `self.values().len()`.
330    ///
331    /// # Example
332    ///
333    /// ```
334    /// # use std::sync::Arc;
335    /// # use arrow_array::{RunArray, Int32Array, StringArray, ArrayRef,Array};
336    /// # use arrow_array::types::Int32Type;
337    /// // A RunArray logically representing ["a", "a", "b", "c", "c"]
338    /// let run_ends = Int32Array::from(vec![2, 3, 5]);
339    /// let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c"]));
340    /// let run_array = RunArray::<Int32Type>::try_new(&run_ends, &values).unwrap();
341    ///
342    /// // Swap in new values while keeping the same run pattern.
343    /// // The result logically represents ["x", "x", "y", "z", "z"].
344    /// let new_values: ArrayRef = Arc::new(StringArray::from(vec!["x", "y", "z"]));
345    /// let new_run_array = run_array.with_values(new_values);
346    ///
347    /// assert_eq!(new_run_array.len(), 5);
348    /// assert_eq!(new_run_array.run_ends().values(), &[2, 3, 5]);
349    /// ```
350    pub fn with_values(&self, values: ArrayRef) -> Self {
351        assert_eq!(values.len(), self.values.len());
352        let (run_ends_field, values_field) = match &self.data_type {
353            DataType::RunEndEncoded(r, v) => {
354                let new_v = Arc::new(Field::new(
355                    v.name(),
356                    values.data_type().clone(),
357                    v.is_nullable(),
358                ));
359                (r, new_v)
360            }
361            _ => unreachable!("RunArray should have type RunEndEncoded"),
362        };
363        let data_type = DataType::RunEndEncoded(Arc::clone(run_ends_field), values_field);
364
365        Self {
366            data_type,
367            run_ends: self.run_ends.clone(),
368            values,
369        }
370    }
371
372    /// Similar to [`values`] but accounts for logical slicing, returning only the values
373    /// that are part of the logical slice of this array.
374    ///
375    /// [`values`]: Self::values
376    pub fn values_slice(&self) -> ArrayRef {
377        if self.is_empty() {
378            return self.values.slice(0, 0);
379        }
380        let start = self.get_start_physical_index();
381        let end = self.get_end_physical_index();
382        self.values.slice(start, end - start + 1)
383    }
384
385    /// Returns the physical index at which the array slice starts.
386    ///
387    /// See [`RunEndBuffer::get_start_physical_index`].
388    pub fn get_start_physical_index(&self) -> usize {
389        self.run_ends.get_start_physical_index()
390    }
391
392    /// Returns the physical index at which the array slice ends.
393    ///
394    /// See [`RunEndBuffer::get_end_physical_index`].
395    pub fn get_end_physical_index(&self) -> usize {
396        self.run_ends.get_end_physical_index()
397    }
398
399    /// Downcast this [`RunArray`] to a [`TypedRunArray`]
400    ///
401    /// ```
402    /// use arrow_array::{Array, ArrayAccessor, RunArray, StringArray, types::Int32Type};
403    ///
404    /// let orig = [Some("a"), Some("b"), None];
405    /// let run_array = RunArray::<Int32Type>::from_iter(orig);
406    /// let typed = run_array.downcast::<StringArray>().unwrap();
407    /// assert_eq!(typed.value(0), "a");
408    /// assert_eq!(typed.value(1), "b");
409    /// assert!(typed.values().is_null(2));
410    /// ```
411    pub fn downcast<V: 'static>(&self) -> Option<TypedRunArray<'_, R, V>> {
412        let values = self.values.as_any().downcast_ref()?;
413        Some(TypedRunArray {
414            run_array: self,
415            values,
416        })
417    }
418
419    /// Calls [`RunEndBuffer::get_physical_index`].
420    ///
421    /// The result is arbitrary if `logical_index >= self.len()`
422    pub fn get_physical_index(&self, logical_index: usize) -> usize {
423        self.run_ends.get_physical_index(logical_index)
424    }
425
426    /// Returns the physical indices corresponding to the provided logical indices.
427    ///
428    /// See [`RunEndBuffer::get_physical_indices`] for more details.
429    #[inline]
430    pub fn get_physical_indices<I>(&self, logical_indices: &[I]) -> Result<Vec<usize>, ArrowError>
431    where
432        I: ArrowNativeType,
433    {
434        self.run_ends()
435            .get_physical_indices(logical_indices)
436            .map_err(|index| {
437                ArrowError::InvalidArgumentError(format!(
438                    "Logical index {} is out of bounds for RunArray of length {}",
439                    index.as_usize(),
440                    self.len()
441                ))
442            })
443    }
444
445    /// Returns a zero-copy slice of this array with the indicated offset and length.
446    ///
447    /// # Panics
448    ///
449    /// - Specified slice (`offset` + `length`) exceeds existing length
450    pub fn slice(&self, offset: usize, length: usize) -> Self {
451        Self {
452            data_type: self.data_type.clone(),
453            run_ends: self.run_ends.slice(offset, length),
454            values: self.values.clone(),
455        }
456    }
457}
458
459/// Converts an [`ArrayData`] into a [`RunArray`].
460///
461/// The `len` and `offset` of `data` become the logical length and offset of the
462/// [`RunArray`]; see the [`RunArray`] documentation for how they relate to the
463/// run ends.
464///
465/// The `offset` and `len` of the `run_ends` child are respected: only the `len`
466/// elements starting at `offset` are used as run ends. See [Sliced child
467/// arrays].
468///
469/// [Sliced child arrays]: RunArray#sliced-child-arrays
470///
471/// # Panics
472///
473/// Panics if `data` is not a [`DataType::RunEndEncoded`] array with a `run_ends`
474/// child of type `R`, or does not have exactly two children.
475impl<R: RunEndIndexType> From<ArrayData> for RunArray<R> {
476    // The method assumes the caller already validated the data using `ArrayData::validate_data()`
477    fn from(data: ArrayData) -> Self {
478        let (data_type, len, _nulls, offset, _buffers, child_data) = data.into_parts();
479
480        match &data_type {
481            DataType::RunEndEncoded(_, _) => {}
482            _ => {
483                panic!(
484                    "Invalid data type {data_type:?} for RunArray. Should be DataType::RunEndEncoded"
485                );
486            }
487        }
488
489        let [run_end_child, values_child]: [ArrayData; 2] = child_data
490            .try_into()
491            .expect("RunArray data should have exactly two child arrays");
492
493        // deconstruct the run ends child array
494        let (
495            run_end_data_type,
496            run_end_len,
497            _run_end_nulls,
498            run_end_offset,
499            run_end_buffers,
500            _run_end_child_data,
501        ) = run_end_child.into_parts();
502        assert_eq!(run_end_data_type, R::DATA_TYPE, "Incorrect run ends type");
503        let [run_end_buffer]: [arrow_buffer::Buffer; 1] = run_end_buffers
504            .try_into()
505            .expect("Run ends should have exactly one buffer");
506        let scalar = ScalarBuffer::new(run_end_buffer, run_end_offset, run_end_len);
507        let run_ends = unsafe { RunEndBuffer::new_unchecked(scalar, offset, len) };
508
509        let values = make_array(values_child);
510
511        Self {
512            data_type,
513            run_ends,
514            values,
515        }
516    }
517}
518
519impl<R: RunEndIndexType> From<RunArray<R>> for ArrayData {
520    fn from(array: RunArray<R>) -> Self {
521        let len = array.run_ends.len();
522        let offset = array.run_ends.offset();
523
524        let run_ends = ArrayDataBuilder::new(R::DATA_TYPE)
525            .len(array.run_ends.values().len())
526            .buffers(vec![array.run_ends.into_inner().into_inner()]);
527
528        let run_ends = unsafe { run_ends.build_unchecked() };
529
530        let builder = ArrayDataBuilder::new(array.data_type)
531            .len(len)
532            .offset(offset)
533            .child_data(vec![run_ends, array.values.to_data()]);
534
535        unsafe { builder.build_unchecked() }
536    }
537}
538
539/// SAFETY: Correctly implements the contract of Arrow Arrays
540unsafe impl<T: RunEndIndexType> Array for RunArray<T> {
541    fn as_any(&self) -> &dyn Any {
542        self
543    }
544
545    fn to_data(&self) -> ArrayData {
546        self.clone().into()
547    }
548
549    fn into_data(self) -> ArrayData {
550        self.into()
551    }
552
553    fn data_type(&self) -> &DataType {
554        &self.data_type
555    }
556
557    fn slice(&self, offset: usize, length: usize) -> ArrayRef {
558        Arc::new(self.slice(offset, length))
559    }
560
561    fn len(&self) -> usize {
562        self.run_ends.len()
563    }
564
565    fn is_empty(&self) -> bool {
566        self.run_ends.is_empty()
567    }
568
569    fn shrink_to_fit(&mut self) {
570        self.run_ends.shrink_to_fit();
571        self.values.shrink_to_fit();
572    }
573
574    fn offset(&self) -> usize {
575        self.run_ends.offset()
576    }
577
578    fn nulls(&self) -> Option<&NullBuffer> {
579        None
580    }
581
582    fn logical_nulls(&self) -> Option<NullBuffer> {
583        let len = self.len();
584        let nulls = self.values.logical_nulls()?;
585        let mut out = BooleanBufferBuilder::new(len);
586        let offset = self.run_ends.offset();
587        let mut valid_start = 0;
588        let mut last_end = 0;
589        for (idx, end) in self.run_ends.values().iter().enumerate() {
590            let end = end.as_usize();
591            if end < offset {
592                continue;
593            }
594            let end = (end - offset).min(len);
595            if nulls.is_null(idx) {
596                if valid_start < last_end {
597                    out.append_n(last_end - valid_start, true);
598                }
599                out.append_n(end - last_end, false);
600                valid_start = end;
601            }
602            last_end = end;
603            if end == len {
604                break;
605            }
606        }
607        if valid_start < len {
608            out.append_n(len - valid_start, true)
609        }
610        // Sanity check
611        assert_eq!(out.len(), len);
612        Some(out.finish().into())
613    }
614
615    fn is_nullable(&self) -> bool {
616        !self.is_empty() && self.values.is_nullable()
617    }
618
619    fn get_buffer_memory_size(&self) -> usize {
620        self.run_ends.inner().inner().capacity() + self.values.get_buffer_memory_size()
621    }
622
623    fn get_array_memory_size(&self) -> usize {
624        std::mem::size_of::<Self>()
625            + self.run_ends.inner().inner().capacity()
626            + self.values.get_array_memory_size()
627    }
628
629    #[cfg(feature = "pool")]
630    fn claim(&self, pool: &dyn arrow_buffer::MemoryPool) {
631        self.run_ends.claim(pool);
632        self.values.claim(pool);
633    }
634}
635
636impl<R: RunEndIndexType> std::fmt::Debug for RunArray<R> {
637    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
638        writeln!(
639            f,
640            "RunArray {{run_ends: {:?}, values: {:?}}}",
641            self.run_ends.values(),
642            self.values
643        )
644    }
645}
646
647/// Constructs a `RunArray` from an iterator of optional strings.
648///
649/// # Example:
650/// ```
651/// use arrow_array::{RunArray, PrimitiveArray, StringArray, types::Int16Type};
652///
653/// let test = vec!["a", "a", "b", "c", "c"];
654/// let array: RunArray<Int16Type> = test
655///     .iter()
656///     .map(|&x| if x == "b" { None } else { Some(x) })
657///     .collect();
658/// assert_eq!(
659///     "RunArray {run_ends: [2, 3, 5], values: StringArray\n[\n  \"a\",\n  null,\n  \"c\",\n]}\n",
660///     format!("{:?}", array)
661/// );
662/// ```
663impl<'a, T: RunEndIndexType> FromIterator<Option<&'a str>> for RunArray<T> {
664    fn from_iter<I: IntoIterator<Item = Option<&'a str>>>(iter: I) -> Self {
665        let it = iter.into_iter();
666        let (lower, _) = it.size_hint();
667        let mut builder = StringRunBuilder::with_capacity(lower, 256);
668        it.for_each(|i| {
669            builder.append_option(i);
670        });
671
672        builder.finish()
673    }
674}
675
676/// Constructs a `RunArray` from an iterator of strings.
677///
678/// # Example:
679///
680/// ```
681/// use arrow_array::{RunArray, PrimitiveArray, StringArray, types::Int16Type};
682///
683/// let test = vec!["a", "a", "b", "c"];
684/// let array: RunArray<Int16Type> = test.into_iter().collect();
685/// assert_eq!(
686///     "RunArray {run_ends: [2, 3, 4], values: StringArray\n[\n  \"a\",\n  \"b\",\n  \"c\",\n]}\n",
687///     format!("{:?}", array)
688/// );
689/// ```
690impl<'a, T: RunEndIndexType> FromIterator<&'a str> for RunArray<T> {
691    fn from_iter<I: IntoIterator<Item = &'a str>>(iter: I) -> Self {
692        let it = iter.into_iter();
693        let (lower, _) = it.size_hint();
694        let mut builder = StringRunBuilder::with_capacity(lower, 256);
695        it.for_each(|i| {
696            builder.append_value(i);
697        });
698
699        builder.finish()
700    }
701}
702
703///
704/// A [`RunArray`] with `i16` run ends
705///
706/// # Example: Using `collect`
707/// ```
708/// # use arrow_array::{Array, Int16RunArray, Int16Array, StringArray};
709/// # use std::sync::Arc;
710///
711/// let array: Int16RunArray = vec!["a", "a", "b", "c", "c"].into_iter().collect();
712/// let values: Arc<dyn Array> = Arc::new(StringArray::from(vec!["a", "b", "c"]));
713/// assert_eq!(array.run_ends().values(), &[2, 3, 5]);
714/// assert_eq!(array.values(), &values);
715/// ```
716pub type Int16RunArray = RunArray<Int16Type>;
717
718///
719/// A [`RunArray`] with `i32` run ends
720///
721/// # Example: Using `collect`
722/// ```
723/// # use arrow_array::{Array, Int32RunArray, Int32Array, StringArray};
724/// # use std::sync::Arc;
725///
726/// let array: Int32RunArray = vec!["a", "a", "b", "c", "c"].into_iter().collect();
727/// let values: Arc<dyn Array> = Arc::new(StringArray::from(vec!["a", "b", "c"]));
728/// assert_eq!(array.run_ends().values(), &[2, 3, 5]);
729/// assert_eq!(array.values(), &values);
730/// ```
731pub type Int32RunArray = RunArray<Int32Type>;
732
733///
734/// A [`RunArray`] with `i64` run ends
735///
736/// # Example: Using `collect`
737/// ```
738/// # use arrow_array::{Array, Int64RunArray, Int64Array, StringArray};
739/// # use std::sync::Arc;
740///
741/// let array: Int64RunArray = vec!["a", "a", "b", "c", "c"].into_iter().collect();
742/// let values: Arc<dyn Array> = Arc::new(StringArray::from(vec!["a", "b", "c"]));
743/// assert_eq!(array.run_ends().values(), &[2, 3, 5]);
744/// assert_eq!(array.values(), &values);
745/// ```
746pub type Int64RunArray = RunArray<Int64Type>;
747
748/// A [`RunArray`] typed on its child values array
749///
750/// Implements [`ArrayAccessor`] and [`IntoIterator`] allowing fast access to its elements
751///
752/// ```
753/// use arrow_array::{RunArray, StringArray, types::Int32Type};
754///
755/// let orig = ["a", "b", "a", "b"];
756/// let ree_array = RunArray::<Int32Type>::from_iter(orig);
757///
758/// // `TypedRunArray` allows you to access the values directly
759/// let typed = ree_array.downcast::<StringArray>().unwrap();
760///
761/// for (maybe_val, orig) in typed.into_iter().zip(orig) {
762///     assert_eq!(maybe_val.unwrap(), orig)
763/// }
764/// ```
765pub struct TypedRunArray<'a, R: RunEndIndexType, V> {
766    /// The run array
767    run_array: &'a RunArray<R>,
768
769    /// The values of the run_array
770    values: &'a V,
771}
772
773// Manually implement `Clone` to avoid `V: Clone` type constraint
774impl<R: RunEndIndexType, V> Clone for TypedRunArray<'_, R, V> {
775    fn clone(&self) -> Self {
776        *self
777    }
778}
779
780impl<R: RunEndIndexType, V> Copy for TypedRunArray<'_, R, V> {}
781
782impl<R: RunEndIndexType, V> std::fmt::Debug for TypedRunArray<'_, R, V> {
783    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
784        writeln!(f, "TypedRunArray({:?})", self.run_array)
785    }
786}
787
788impl<'a, R: RunEndIndexType, V> TypedRunArray<'a, R, V> {
789    /// Returns the run_ends of this [`TypedRunArray`]
790    pub fn run_ends(&self) -> &'a RunEndBuffer<R::Native> {
791        self.run_array.run_ends()
792    }
793
794    /// Returns the values of this [`TypedRunArray`]
795    pub fn values(&self) -> &'a V {
796        self.values
797    }
798
799    /// Returns the run array of this [`TypedRunArray`]
800    pub fn run_array(&self) -> &'a RunArray<R> {
801        self.run_array
802    }
803}
804
805/// SAFETY: Correctly implements the contract of Arrow Arrays
806unsafe impl<R: RunEndIndexType, V: Sync> Array for TypedRunArray<'_, R, V> {
807    fn as_any(&self) -> &dyn Any {
808        self.run_array
809    }
810
811    fn to_data(&self) -> ArrayData {
812        self.run_array.to_data()
813    }
814
815    fn into_data(self) -> ArrayData {
816        self.run_array.into_data()
817    }
818
819    fn data_type(&self) -> &DataType {
820        self.run_array.data_type()
821    }
822
823    fn slice(&self, offset: usize, length: usize) -> ArrayRef {
824        Arc::new(self.run_array.slice(offset, length))
825    }
826
827    fn len(&self) -> usize {
828        self.run_array.len()
829    }
830
831    fn is_empty(&self) -> bool {
832        self.run_array.is_empty()
833    }
834
835    fn offset(&self) -> usize {
836        self.run_array.offset()
837    }
838
839    fn nulls(&self) -> Option<&NullBuffer> {
840        self.run_array.nulls()
841    }
842
843    fn logical_nulls(&self) -> Option<NullBuffer> {
844        self.run_array.logical_nulls()
845    }
846
847    fn logical_null_count(&self) -> usize {
848        self.run_array.logical_null_count()
849    }
850
851    fn is_nullable(&self) -> bool {
852        self.run_array.is_nullable()
853    }
854
855    fn get_buffer_memory_size(&self) -> usize {
856        self.run_array.get_buffer_memory_size()
857    }
858
859    fn get_array_memory_size(&self) -> usize {
860        self.run_array.get_array_memory_size()
861    }
862
863    #[cfg(feature = "pool")]
864    fn claim(&self, pool: &dyn arrow_buffer::MemoryPool) {
865        self.run_array.claim(pool);
866    }
867}
868
869// Array accessor converts the index of logical array to the index of the physical array
870// using binary search. The time complexity is O(log N) where N is number of runs.
871impl<'a, R, V> ArrayAccessor for TypedRunArray<'a, R, V>
872where
873    R: RunEndIndexType,
874    V: Sync + Send,
875    &'a V: ArrayAccessor,
876    <&'a V as ArrayAccessor>::Item: Default,
877{
878    type Item = <&'a V as ArrayAccessor>::Item;
879
880    fn value(&self, logical_index: usize) -> Self::Item {
881        assert!(
882            logical_index < self.len(),
883            "Trying to access an element at index {} from a TypedRunArray of length {}",
884            logical_index,
885            self.len()
886        );
887        unsafe { self.value_unchecked(logical_index) }
888    }
889
890    unsafe fn value_unchecked(&self, logical_index: usize) -> Self::Item {
891        let physical_index = self.run_array.get_physical_index(logical_index);
892        unsafe { self.values().value_unchecked(physical_index) }
893    }
894}
895
896impl<'a, R, V> IntoIterator for TypedRunArray<'a, R, V>
897where
898    R: RunEndIndexType,
899    V: Sync + Send,
900    &'a V: ArrayAccessor,
901    <&'a V as ArrayAccessor>::Item: Default,
902{
903    type Item = Option<<&'a V as ArrayAccessor>::Item>;
904    type IntoIter = RunArrayIter<'a, R, V>;
905
906    fn into_iter(self) -> Self::IntoIter {
907        RunArrayIter::new(self)
908    }
909}
910/// An array that can be downcast to a [`RunArray`] of any run end type and any value type.
911///
912/// This can be used to efficiently implement kernels for all possible run end
913/// types without needing to create specialized implementations for each key type.
914pub trait AnyRunEndArray: Array {
915    /// Returns the values of this array.
916    fn values(&self) -> &Arc<dyn Array>;
917
918    /// Returns a new run-end encoded array with the given values, preserving the
919    /// existing run ends.
920    fn with_values(&self, values: ArrayRef) -> ArrayRef;
921}
922
923impl<R: RunEndIndexType> AnyRunEndArray for RunArray<R> {
924    fn values(&self) -> &Arc<dyn Array> {
925        &self.values
926    }
927
928    fn with_values(&self, values: ArrayRef) -> ArrayRef {
929        Arc::new(RunArray::<R>::with_values(self, values))
930    }
931}
932
933#[cfg(test)]
934mod tests {
935    use rand::RngExt;
936    use rand::rng;
937    use rand::seq::SliceRandom;
938
939    use super::*;
940    use crate::Float32Array;
941    use crate::Int64Array;
942    use crate::builder::PrimitiveRunBuilder;
943    use crate::cast::AsArray;
944    use crate::new_empty_array;
945    use crate::types::{Int8Type, UInt32Type};
946    use crate::{Int16Array, Int32Array, StringArray};
947
948    fn build_input_array(size: usize) -> Vec<Option<i32>> {
949        // The input array is created by shuffling and repeating
950        // the seed values random number of times.
951        let mut seed: Vec<Option<i32>> = vec![
952            None,
953            None,
954            None,
955            Some(1),
956            Some(2),
957            Some(3),
958            Some(4),
959            Some(5),
960            Some(6),
961            Some(7),
962            Some(8),
963            Some(9),
964        ];
965        let mut result: Vec<Option<i32>> = Vec::with_capacity(size);
966        let mut ix = 0;
967        let mut rng = rng();
968        // run length can go up to 8. Cap the max run length for smaller arrays to size / 2.
969        let max_run_length = 8_usize.min(1_usize.max(size / 2));
970        while result.len() < size {
971            // shuffle the seed array if all the values are iterated.
972            if ix == 0 {
973                seed.shuffle(&mut rng);
974            }
975            // repeat the items between 1 and 8 times. Cap the length for smaller sized arrays
976            let num = max_run_length.min(rng.random_range(1..=max_run_length));
977            for _ in 0..num {
978                result.push(seed[ix]);
979            }
980            ix += 1;
981            if ix == seed.len() {
982                ix = 0
983            }
984        }
985        result.resize(size, None);
986        result
987    }
988
989    // Asserts that `logical_array[logical_indices[*]] == physical_array[physical_indices[*]]`
990    fn compare_logical_and_physical_indices(
991        logical_indices: &[u32],
992        logical_array: &[Option<i32>],
993        physical_indices: &[usize],
994        physical_array: &PrimitiveArray<Int32Type>,
995    ) {
996        assert_eq!(logical_indices.len(), physical_indices.len());
997
998        // check value in logical index in the logical_array matches physical index in physical_array
999        logical_indices
1000            .iter()
1001            .map(|f| f.as_usize())
1002            .zip(physical_indices.iter())
1003            .for_each(|(logical_ix, physical_ix)| {
1004                let expected = logical_array[logical_ix];
1005                match expected {
1006                    Some(val) => {
1007                        assert!(physical_array.is_valid(*physical_ix));
1008                        let actual = physical_array.value(*physical_ix);
1009                        assert_eq!(val, actual);
1010                    }
1011                    None => {
1012                        assert!(physical_array.is_null(*physical_ix))
1013                    }
1014                }
1015            });
1016    }
1017    #[test]
1018    fn test_run_array() {
1019        // Construct a value array
1020        let value_data =
1021            PrimitiveArray::<Int8Type>::from_iter_values([10_i8, 11, 12, 13, 14, 15, 16, 17]);
1022
1023        // Construct a run_ends array:
1024        let run_ends_values = [4_i16, 6, 7, 9, 13, 18, 20, 22];
1025        let run_ends_data =
1026            PrimitiveArray::<Int16Type>::from_iter_values(run_ends_values.iter().copied());
1027
1028        // Construct a run ends encoded array from the above two
1029        let ree_array = RunArray::<Int16Type>::try_new(&run_ends_data, &value_data).unwrap();
1030
1031        assert_eq!(ree_array.len(), 22);
1032        assert_eq!(ree_array.null_count(), 0);
1033
1034        let values = ree_array.values();
1035        assert_eq!(value_data.into_data(), values.to_data());
1036        assert_eq!(&DataType::Int8, values.data_type());
1037
1038        let run_ends = ree_array.run_ends();
1039        assert_eq!(run_ends.values(), &run_ends_values);
1040    }
1041
1042    #[test]
1043    fn test_run_array_empty() {
1044        let runs = new_empty_array(&DataType::Int16);
1045        let runs = runs.as_primitive::<Int16Type>();
1046        let values = new_empty_array(&DataType::Int64);
1047        let array = RunArray::try_new(runs, &values).unwrap();
1048
1049        fn assertions(array: &RunArray<Int16Type>) {
1050            assert!(array.is_empty());
1051            assert_eq!(array.get_start_physical_index(), 0);
1052            assert_eq!(array.get_end_physical_index(), 0);
1053            assert!(array.get_physical_indices::<i16>(&[]).unwrap().is_empty());
1054            assert!(array.run_ends().is_empty());
1055            assert_eq!(array.run_ends().sliced_values().count(), 0);
1056        }
1057
1058        assertions(&array);
1059        assertions(&array.slice(0, 0));
1060    }
1061
1062    #[test]
1063    fn test_run_array_fmt_debug() {
1064        let mut builder = PrimitiveRunBuilder::<Int16Type, UInt32Type>::with_capacity(3);
1065        builder.append_value(12345678);
1066        builder.append_null();
1067        builder.append_value(22345678);
1068        let array = builder.finish();
1069        assert_eq!(
1070            "RunArray {run_ends: [1, 2, 3], values: PrimitiveArray<UInt32>\n[\n  12345678,\n  null,\n  22345678,\n]}\n",
1071            format!("{array:?}")
1072        );
1073
1074        let mut builder = PrimitiveRunBuilder::<Int16Type, UInt32Type>::with_capacity(20);
1075        for _ in 0..20 {
1076            builder.append_value(1);
1077        }
1078        let array = builder.finish();
1079
1080        assert_eq!(array.len(), 20);
1081        assert_eq!(array.null_count(), 0);
1082        assert_eq!(array.logical_null_count(), 0);
1083
1084        assert_eq!(
1085            "RunArray {run_ends: [20], values: PrimitiveArray<UInt32>\n[\n  1,\n]}\n",
1086            format!("{array:?}")
1087        );
1088    }
1089
1090    #[test]
1091    fn test_run_array_from_iter() {
1092        let test = vec!["a", "a", "b", "c"];
1093        let array: RunArray<Int16Type> = test
1094            .iter()
1095            .map(|&x| if x == "b" { None } else { Some(x) })
1096            .collect();
1097        assert_eq!(
1098            "RunArray {run_ends: [2, 3, 4], values: StringArray\n[\n  \"a\",\n  null,\n  \"c\",\n]}\n",
1099            format!("{array:?}")
1100        );
1101
1102        assert_eq!(array.len(), 4);
1103        assert_eq!(array.null_count(), 0);
1104        assert_eq!(array.logical_null_count(), 1);
1105
1106        let array: RunArray<Int16Type> = test.into_iter().collect();
1107        assert_eq!(
1108            "RunArray {run_ends: [2, 3, 4], values: StringArray\n[\n  \"a\",\n  \"b\",\n  \"c\",\n]}\n",
1109            format!("{array:?}")
1110        );
1111    }
1112
1113    #[test]
1114    fn test_run_array_run_ends_as_primitive_array() {
1115        let test = vec!["a", "b", "c", "a"];
1116        let array: RunArray<Int16Type> = test.into_iter().collect();
1117
1118        assert_eq!(array.len(), 4);
1119        assert_eq!(array.null_count(), 0);
1120        assert_eq!(array.logical_null_count(), 0);
1121
1122        let run_ends = array.run_ends();
1123        assert_eq!(&[1, 2, 3, 4], run_ends.values());
1124    }
1125
1126    #[test]
1127    fn test_run_array_as_primitive_array_with_null() {
1128        let test = vec![Some("a"), None, Some("b"), None, None, Some("a")];
1129        let array: RunArray<Int32Type> = test.into_iter().collect();
1130
1131        assert_eq!(array.len(), 6);
1132        assert_eq!(array.null_count(), 0);
1133        assert_eq!(array.logical_null_count(), 3);
1134
1135        let run_ends = array.run_ends();
1136        assert_eq!(&[1, 2, 3, 5, 6], run_ends.values());
1137
1138        let values_data = array.values();
1139        assert_eq!(2, values_data.null_count());
1140        assert_eq!(5, values_data.len());
1141    }
1142
1143    #[test]
1144    fn test_run_array_all_nulls() {
1145        let test = vec![None, None, None];
1146        let array: RunArray<Int32Type> = test.into_iter().collect();
1147
1148        assert_eq!(array.len(), 3);
1149        assert_eq!(array.null_count(), 0);
1150        assert_eq!(array.logical_null_count(), 3);
1151
1152        let run_ends = array.run_ends();
1153        assert_eq!(3, run_ends.len());
1154        assert_eq!(&[3], run_ends.values());
1155
1156        let values_data = array.values();
1157        assert_eq!(1, values_data.null_count());
1158    }
1159
1160    #[test]
1161    fn test_run_array_try_new() {
1162        let values: StringArray = [Some("foo"), Some("bar"), None, Some("baz")]
1163            .into_iter()
1164            .collect();
1165        let run_ends: Int32Array = [Some(1), Some(2), Some(3), Some(4)].into_iter().collect();
1166
1167        let array = RunArray::<Int32Type>::try_new(&run_ends, &values).unwrap();
1168        assert_eq!(array.values().data_type(), &DataType::Utf8);
1169
1170        assert_eq!(array.null_count(), 0);
1171        assert_eq!(array.logical_null_count(), 1);
1172        assert_eq!(array.len(), 4);
1173        assert_eq!(array.values().null_count(), 1);
1174
1175        assert_eq!(
1176            "RunArray {run_ends: [1, 2, 3, 4], values: StringArray\n[\n  \"foo\",\n  \"bar\",\n  null,\n  \"baz\",\n]}\n",
1177            format!("{array:?}")
1178        );
1179    }
1180
1181    #[test]
1182    fn test_run_array_int16_type_definition() {
1183        let array: Int16RunArray = vec!["a", "a", "b", "c", "c"].into_iter().collect();
1184        let values: Arc<dyn Array> = Arc::new(StringArray::from(vec!["a", "b", "c"]));
1185        assert_eq!(array.run_ends().values(), &[2, 3, 5]);
1186        assert_eq!(array.values(), &values);
1187    }
1188
1189    #[test]
1190    fn test_run_array_empty_string() {
1191        let array: Int16RunArray = vec!["a", "a", "", "", "c"].into_iter().collect();
1192        let values: Arc<dyn Array> = Arc::new(StringArray::from(vec!["a", "", "c"]));
1193        assert_eq!(array.run_ends().values(), &[2, 4, 5]);
1194        assert_eq!(array.values(), &values);
1195    }
1196
1197    #[test]
1198    fn test_run_array_length_mismatch() {
1199        let values: StringArray = [Some("foo"), Some("bar"), None, Some("baz")]
1200            .into_iter()
1201            .collect();
1202        let run_ends: Int32Array = [Some(1), Some(2), Some(3)].into_iter().collect();
1203
1204        let actual = RunArray::<Int32Type>::try_new(&run_ends, &values);
1205        let expected = ArrowError::InvalidArgumentError("The run_ends array length should be the same as values array length. Run_ends array length is 3, values array length is 4".to_string());
1206        assert_eq!(expected.to_string(), actual.err().unwrap().to_string());
1207    }
1208    #[test]
1209    fn test_run_array_with_values_changes_value_type() {
1210        let values = StringArray::from(vec!["foo", "bar", "baz"]);
1211        let run_ends: Int32Array = [Some(1), Some(2), Some(3)].into_iter().collect();
1212        let ree = RunArray::<Int32Type>::try_new(&run_ends, &values).unwrap();
1213
1214        let new_values = Int64Array::from(vec![10, 20, 30]);
1215        let result = ree.with_values(Arc::new(new_values));
1216
1217        match result.data_type() {
1218            DataType::RunEndEncoded(_, v) => {
1219                assert_eq!(v.data_type(), &DataType::Int64);
1220            }
1221            other => panic!("expected RunEndEncoded, got {other:?}"),
1222        }
1223
1224        assert_eq!(result.values().data_type(), &DataType::Int64);
1225        assert_eq!(result.values().len(), 3);
1226    }
1227
1228    #[test]
1229    fn test_run_array_run_ends_with_null() {
1230        let values: StringArray = [Some("foo"), Some("bar"), Some("baz")]
1231            .into_iter()
1232            .collect();
1233        let run_ends: Int32Array = [Some(1), None, Some(3)].into_iter().collect();
1234
1235        let actual = RunArray::<Int32Type>::try_new(&run_ends, &values);
1236        let expected = ArrowError::InvalidArgumentError(
1237            "Found null values in run_ends array. The run_ends array should not have null values."
1238                .to_string(),
1239        );
1240        assert_eq!(expected.to_string(), actual.err().unwrap().to_string());
1241    }
1242
1243    #[test]
1244    fn test_run_array_run_ends_with_zeroes() {
1245        let values: StringArray = [Some("foo"), Some("bar"), Some("baz")]
1246            .into_iter()
1247            .collect();
1248        let run_ends: Int32Array = [Some(0), Some(1), Some(3)].into_iter().collect();
1249
1250        let actual = RunArray::<Int32Type>::try_new(&run_ends, &values);
1251        let expected = ArrowError::InvalidArgumentError("The values in run_ends array should be strictly positive. Found value 0 at index 0 that does not match the criteria.".to_string());
1252        assert_eq!(expected.to_string(), actual.err().unwrap().to_string());
1253    }
1254
1255    #[test]
1256    fn test_run_array_run_ends_non_increasing() {
1257        let values: StringArray = [Some("foo"), Some("bar"), Some("baz")]
1258            .into_iter()
1259            .collect();
1260        let run_ends: Int32Array = [Some(1), Some(4), Some(4)].into_iter().collect();
1261
1262        let actual = RunArray::<Int32Type>::try_new(&run_ends, &values);
1263        let expected = ArrowError::InvalidArgumentError("The values in run_ends array should be strictly increasing. Found value 4 at index 2 with previous value 4 that does not match the criteria.".to_string());
1264        assert_eq!(expected.to_string(), actual.err().unwrap().to_string());
1265    }
1266
1267    #[test]
1268    #[should_panic(expected = "Incorrect run ends type")]
1269    fn test_run_array_run_ends_data_type_mismatch() {
1270        let a = RunArray::<Int32Type>::from_iter(["32"]);
1271        let _ = RunArray::<Int64Type>::from(a.into_data());
1272    }
1273
1274    #[test]
1275    fn test_ree_array_accessor() {
1276        let input_array = build_input_array(256);
1277
1278        // Encode the input_array to ree_array
1279        let mut builder =
1280            PrimitiveRunBuilder::<Int16Type, Int32Type>::with_capacity(input_array.len());
1281        builder.extend(input_array.iter().copied());
1282        let run_array = builder.finish();
1283        let typed = run_array.downcast::<PrimitiveArray<Int32Type>>().unwrap();
1284
1285        // Access every index and check if the value in the input array matches returned value.
1286        for (i, inp_val) in input_array.iter().enumerate() {
1287            if let Some(val) = inp_val {
1288                let actual = typed.value(i);
1289                assert_eq!(*val, actual)
1290            } else {
1291                let physical_ix = run_array.get_physical_index(i);
1292                assert!(typed.values().is_null(physical_ix));
1293            }
1294        }
1295    }
1296
1297    #[test]
1298    #[cfg_attr(miri, ignore)] // Takes too long
1299    fn test_get_physical_indices() {
1300        // Test for logical lengths starting from 10 to 250 increasing by 10
1301        for logical_len in (0..250).step_by(10) {
1302            let input_array = build_input_array(logical_len);
1303
1304            // create run array using input_array
1305            let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
1306            builder.extend(input_array.clone());
1307
1308            let run_array = builder.finish();
1309            let physical_values_array = run_array.values().as_primitive::<Int32Type>();
1310
1311            // create an array consisting of all the indices repeated twice and shuffled.
1312            let mut logical_indices: Vec<u32> = (0_u32..(logical_len as u32)).collect();
1313            // add same indices once more
1314            logical_indices.append(&mut logical_indices.clone());
1315            let mut rng = rng();
1316            logical_indices.shuffle(&mut rng);
1317
1318            let physical_indices = run_array.get_physical_indices(&logical_indices).unwrap();
1319
1320            assert_eq!(logical_indices.len(), physical_indices.len());
1321
1322            // check value in logical index in the input_array matches physical index in typed_run_array
1323            compare_logical_and_physical_indices(
1324                &logical_indices,
1325                &input_array,
1326                &physical_indices,
1327                physical_values_array,
1328            );
1329        }
1330    }
1331
1332    #[test]
1333    #[cfg_attr(miri, ignore)] // Takes too long
1334    fn test_get_physical_indices_sliced() {
1335        let total_len = 80;
1336        let input_array = build_input_array(total_len);
1337
1338        // Encode the input_array to run array
1339        let mut builder =
1340            PrimitiveRunBuilder::<Int16Type, Int32Type>::with_capacity(input_array.len());
1341        builder.extend(input_array.iter().copied());
1342        let run_array = builder.finish();
1343        let physical_values_array = run_array.values().as_primitive::<Int32Type>();
1344
1345        // test for all slice lengths.
1346        for slice_len in 1..=total_len {
1347            // create an array consisting of all the indices repeated twice and shuffled.
1348            let mut logical_indices: Vec<u32> = (0_u32..(slice_len as u32)).collect();
1349            // add same indices once more
1350            logical_indices.append(&mut logical_indices.clone());
1351            let mut rng = rng();
1352            logical_indices.shuffle(&mut rng);
1353
1354            // test for offset = 0 and slice length = slice_len
1355            // slice the input array using which the run array was built.
1356            let sliced_input_array = &input_array[0..slice_len];
1357
1358            // slice the run array
1359            let sliced_run_array: RunArray<Int16Type> =
1360                run_array.slice(0, slice_len).into_data().into();
1361
1362            // Get physical indices.
1363            let physical_indices = sliced_run_array
1364                .get_physical_indices(&logical_indices)
1365                .unwrap();
1366
1367            compare_logical_and_physical_indices(
1368                &logical_indices,
1369                sliced_input_array,
1370                &physical_indices,
1371                physical_values_array,
1372            );
1373
1374            // test for offset = total_len - slice_len and slice length = slice_len
1375            // slice the input array using which the run array was built.
1376            let sliced_input_array = &input_array[total_len - slice_len..total_len];
1377
1378            // slice the run array
1379            let sliced_run_array: RunArray<Int16Type> = run_array
1380                .slice(total_len - slice_len, slice_len)
1381                .into_data()
1382                .into();
1383
1384            // Get physical indices
1385            let physical_indices = sliced_run_array
1386                .get_physical_indices(&logical_indices)
1387                .unwrap();
1388
1389            compare_logical_and_physical_indices(
1390                &logical_indices,
1391                sliced_input_array,
1392                &physical_indices,
1393                physical_values_array,
1394            );
1395        }
1396    }
1397
1398    #[test]
1399    fn test_logical_nulls() {
1400        let run = Int32Array::from(vec![3, 6, 9, 12]);
1401        let values = Int32Array::from(vec![Some(0), None, Some(1), None]);
1402        let array = RunArray::try_new(&run, &values).unwrap();
1403
1404        let expected = [
1405            true, true, true, false, false, false, true, true, true, false, false, false,
1406        ];
1407
1408        let n = array.logical_nulls().unwrap();
1409        assert_eq!(n.null_count(), 6);
1410
1411        let slices = [(0, 12), (0, 2), (2, 5), (3, 0), (3, 3), (3, 4), (4, 8)];
1412        for (offset, length) in slices {
1413            let a = array.slice(offset, length);
1414            let n = a.logical_nulls().unwrap();
1415            let n = n.into_iter().collect::<Vec<_>>();
1416            assert_eq!(&n, &expected[offset..offset + length], "{offset} {length}");
1417        }
1418    }
1419
1420    #[test]
1421    fn test_run_array_eq_identical() {
1422        let run_ends1 = Int32Array::from(vec![2, 4, 6]);
1423        let values1 = StringArray::from(vec!["a", "b", "c"]);
1424        let array1 = RunArray::<Int32Type>::try_new(&run_ends1, &values1).unwrap();
1425
1426        let run_ends2 = Int32Array::from(vec![2, 4, 6]);
1427        let values2 = StringArray::from(vec!["a", "b", "c"]);
1428        let array2 = RunArray::<Int32Type>::try_new(&run_ends2, &values2).unwrap();
1429
1430        assert_eq!(array1, array2);
1431    }
1432
1433    #[test]
1434    fn test_run_array_ne_different_run_ends() {
1435        let run_ends1 = Int32Array::from(vec![2, 4, 6]);
1436        let values1 = StringArray::from(vec!["a", "b", "c"]);
1437        let array1 = RunArray::<Int32Type>::try_new(&run_ends1, &values1).unwrap();
1438
1439        let run_ends2 = Int32Array::from(vec![1, 4, 6]);
1440        let values2 = StringArray::from(vec!["a", "b", "c"]);
1441        let array2 = RunArray::<Int32Type>::try_new(&run_ends2, &values2).unwrap();
1442
1443        assert_ne!(array1, array2);
1444    }
1445
1446    #[test]
1447    fn test_run_array_ne_different_values() {
1448        let run_ends1 = Int32Array::from(vec![2, 4, 6]);
1449        let values1 = StringArray::from(vec!["a", "b", "c"]);
1450        let array1 = RunArray::<Int32Type>::try_new(&run_ends1, &values1).unwrap();
1451
1452        let run_ends2 = Int32Array::from(vec![2, 4, 6]);
1453        let values2 = StringArray::from(vec!["a", "b", "d"]);
1454        let array2 = RunArray::<Int32Type>::try_new(&run_ends2, &values2).unwrap();
1455
1456        assert_ne!(array1, array2);
1457    }
1458
1459    #[test]
1460    fn test_run_array_eq_with_nulls() {
1461        let run_ends1 = Int32Array::from(vec![2, 4, 6]);
1462        let values1 = StringArray::from(vec![Some("a"), None, Some("c")]);
1463        let array1 = RunArray::<Int32Type>::try_new(&run_ends1, &values1).unwrap();
1464
1465        let run_ends2 = Int32Array::from(vec![2, 4, 6]);
1466        let values2 = StringArray::from(vec![Some("a"), None, Some("c")]);
1467        let array2 = RunArray::<Int32Type>::try_new(&run_ends2, &values2).unwrap();
1468
1469        assert_eq!(array1, array2);
1470    }
1471
1472    #[test]
1473    fn test_run_array_eq_different_run_end_types() {
1474        let run_ends_i16_1 = Int16Array::from(vec![2_i16, 4, 6]);
1475        let values_i16_1 = StringArray::from(vec!["a", "b", "c"]);
1476        let array_i16_1 = RunArray::<Int16Type>::try_new(&run_ends_i16_1, &values_i16_1).unwrap();
1477
1478        let run_ends_i16_2 = Int16Array::from(vec![2_i16, 4, 6]);
1479        let values_i16_2 = StringArray::from(vec!["a", "b", "c"]);
1480        let array_i16_2 = RunArray::<Int16Type>::try_new(&run_ends_i16_2, &values_i16_2).unwrap();
1481
1482        assert_eq!(array_i16_1, array_i16_2);
1483    }
1484
1485    #[test]
1486    fn test_run_array_values_slice() {
1487        // 0, 0, 1, 1, 1, 2...2 (15 2s)
1488        let run_ends: PrimitiveArray<Int32Type> = vec![2, 5, 20].into();
1489        let values: PrimitiveArray<Int32Type> = vec![0, 1, 2].into();
1490        let array = RunArray::<Int32Type>::try_new(&run_ends, &values).unwrap();
1491
1492        let slice = array.slice(1, 4); // 0 | 1, 1, 1 |
1493        // logical indices: 1, 2, 3, 4
1494        // physical indices: 0, 1, 1, 1
1495        // values at 0 is 0
1496        // values at 1 is 1
1497        // values slice should be [0, 1]
1498        assert_eq!(slice.get_start_physical_index(), 0);
1499        assert_eq!(slice.get_end_physical_index(), 1);
1500
1501        let values_slice = slice.values_slice();
1502        let values_slice = values_slice.as_primitive::<Int32Type>();
1503        assert_eq!(values_slice.values(), &[0, 1]);
1504
1505        let slice2 = array.slice(2, 3); // 1, 1, 1
1506        // logical indices: 2, 3, 4
1507        // physical indices: 1, 1, 1
1508        assert_eq!(slice2.get_start_physical_index(), 1);
1509        assert_eq!(slice2.get_end_physical_index(), 1);
1510
1511        let values_slice2 = slice2.values_slice();
1512        let values_slice2 = values_slice2.as_primitive::<Int32Type>();
1513        assert_eq!(values_slice2.values(), &[1]);
1514    }
1515
1516    #[test]
1517    fn test_run_array_values_slice_empty() {
1518        let run_ends = Int32Array::from(vec![2, 5, 10]);
1519        let values = StringArray::from(vec!["a", "b", "c"]);
1520        let array = RunArray::<Int32Type>::try_new(&run_ends, &values).unwrap();
1521
1522        let slice = array.slice(0, 0);
1523        assert_eq!(slice.len(), 0);
1524
1525        let values_slice = slice.values_slice();
1526        assert_eq!(values_slice.len(), 0);
1527        assert_eq!(values_slice.data_type(), &DataType::Utf8);
1528    }
1529
1530    #[test]
1531    fn test_run_array_eq_empty() {
1532        let run_ends = Int32Array::from(vec![2, 5, 10]);
1533        let values = StringArray::from(vec!["a", "b", "c"]);
1534        let array = RunArray::<Int32Type>::try_new(&run_ends, &values).unwrap();
1535
1536        let slice1 = array.slice(0, 0);
1537        let slice2 = array.slice(1, 0);
1538        let slice3 = array.slice(10, 0);
1539
1540        assert_eq!(slice1, slice2);
1541        assert_eq!(slice2, slice3);
1542
1543        let empty_array = new_empty_array(array.data_type());
1544        let empty_array = crate::cast::as_run_array::<Int32Type>(empty_array.as_ref());
1545
1546        assert_eq!(&slice1, empty_array);
1547    }
1548
1549    #[test]
1550    fn test_run_array_eq_diff_physical_same_logical() {
1551        let run_ends1 = Int32Array::from(vec![1, 3, 6]);
1552        let values1 = StringArray::from(vec!["a", "b", "c"]);
1553        let array1 = RunArray::<Int32Type>::try_new(&run_ends1, &values1).unwrap();
1554
1555        let run_ends2 = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
1556        let values2 = StringArray::from(vec!["a", "b", "b", "c", "c", "c"]);
1557        let array2 = RunArray::<Int32Type>::try_new(&run_ends2, &values2).unwrap();
1558
1559        assert_eq!(array1, array2);
1560    }
1561
1562    #[test]
1563    fn test_run_array_eq_sliced() {
1564        let run_ends1 = Int32Array::from(vec![2, 5, 10]);
1565        let values1 = StringArray::from(vec!["a", "b", "c"]);
1566        let array1 = RunArray::<Int32Type>::try_new(&run_ends1, &values1).unwrap();
1567        // Logical: a, a, b, b, b, c, c, c, c, c
1568
1569        let slice1 = array1.slice(1, 6);
1570        // Logical: a, b, b, b, c, c
1571
1572        let run_ends2 = Int32Array::from(vec![1, 4, 6]);
1573        let values2 = StringArray::from(vec!["a", "b", "c"]);
1574        let array2 = RunArray::<Int32Type>::try_new(&run_ends2, &values2).unwrap();
1575        // Logical: a, b, b, b, c, c
1576
1577        assert_eq!(slice1, array2);
1578
1579        let slice2 = array1.slice(2, 3);
1580        // Logical: b, b, b
1581        let run_ends3 = Int32Array::from(vec![3]);
1582        let values3 = StringArray::from(vec!["b"]);
1583        let array3 = RunArray::<Int32Type>::try_new(&run_ends3, &values3).unwrap();
1584        assert_eq!(slice2, array3);
1585    }
1586
1587    #[test]
1588    fn test_run_array_eq_sliced_different_offsets() {
1589        let run_ends1 = Int32Array::from(vec![2, 5, 10]);
1590        let values1 = StringArray::from(vec!["a", "b", "c"]);
1591        let array1 = RunArray::<Int32Type>::try_new(&run_ends1, &values1).unwrap();
1592        let array2 = array1.clone();
1593        assert_eq!(array1, array2);
1594
1595        let slice1 = array1.slice(1, 4); // a, b, b, b
1596        let slice2 = array1.slice(1, 4);
1597        assert_eq!(slice1, slice2);
1598
1599        let slice3 = array1.slice(0, 4); // a, a, b, b
1600        assert_ne!(slice1, slice3);
1601    }
1602
1603    #[test]
1604    #[cfg(not(feature = "force_validate"))]
1605    fn allow_to_create_invalid_array_using_new_unchecked() {
1606        let valid = RunArray::<Int32Type>::from_iter(["32"]);
1607        let (_, buffer, values) = valid.into_parts();
1608
1609        let _ = unsafe {
1610            // mismatch data type
1611            RunArray::<Int32Type>::new_unchecked(DataType::Int64, buffer, values)
1612        };
1613    }
1614
1615    #[test]
1616    #[should_panic(
1617        expected = "Invalid data type Int64 for RunArray. Should be DataType::RunEndEncoded"
1618    )]
1619    #[cfg(feature = "force_validate")]
1620    fn should_not_be_able_to_create_invalid_array_using_new_unchecked_when_force_validate_is_enabled()
1621     {
1622        let valid = RunArray::<Int32Type>::from_iter(["32"]);
1623        let (_, buffer, values) = valid.into_parts();
1624
1625        let _ = unsafe {
1626            // mismatch data type
1627            RunArray::<Int32Type>::new_unchecked(DataType::Int64, buffer, values)
1628        };
1629    }
1630
1631    #[test]
1632    fn test_run_array_roundtrip() {
1633        let run = Int32Array::from(vec![3, 6, 9, 12]);
1634        let values = Int32Array::from(vec![Some(0), None, Some(1), None]);
1635        let array = RunArray::try_new(&run, &values).unwrap();
1636
1637        let (dt, buffer, values) = array.clone().into_parts();
1638        let created_from_parts =
1639            unsafe { RunArray::<Int32Type>::new_unchecked(dt, buffer, values) };
1640        assert_eq!(array, created_from_parts);
1641    }
1642
1643    #[test]
1644    fn test_from_array_data_offset() {
1645        let dt = DataType::RunEndEncoded(
1646            Arc::new(Field::new(
1647                Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
1648                DataType::Int32,
1649                false,
1650            )),
1651            Arc::new(Field::new(
1652                Field::REE_VALUES_FIELD_DEFAULT_NAME,
1653                DataType::Float32,
1654                true,
1655            )),
1656        );
1657
1658        let runs = ArrayData::builder(DataType::Int32)
1659            .len(2)
1660            .offset(2)
1661            .add_buffer(vec![-10_i32, -20, 1_i32, 4, 0, 100].into())
1662            .build()
1663            .unwrap();
1664        let values = Float32Array::from(vec![1.5, 2.5]);
1665
1666        let data = ArrayDataBuilder::new(dt)
1667            .len(3)
1668            .child_data(vec![runs, values.to_data()])
1669            .build()
1670            .unwrap();
1671        let array = RunArray::<Int32Type>::from(data);
1672        let array = array.downcast::<Float32Array>().unwrap();
1673
1674        assert_eq!(array.value(0), 1.5);
1675        assert_eq!(array.value(1), 2.5);
1676        assert_eq!(array.value(2), 2.5);
1677    }
1678}
arrow_array/array/run_array.rs

arrow_array/array/
run_array.rs
