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

arrow_array/array/
run_array.rs
