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

arrow_array/array/
run_array.rs
