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

arrow_array/array/
run_array.rs
