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

arrow_array/array/
run_array.rs
