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};
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/// An array of [run-end encoded values].
34///
35/// This encoding is variation on [run-length encoding (RLE)] and is good for representing
36/// data containing the same values repeated consecutively.
37///
38/// A [`RunArray`] consists of a `run_ends` buffer and a `values` array of equivalent
39/// lengths. The `run_ends` buffer stores the indexes at which the run ends. The
40/// `values` array stores the corresponding value of each run. The below example
41/// illustrates how a logical array is represented by a [`RunArray`]:
42///
43/// ```text
44/// ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─┐
45///   ┌─────────────────┐  ┌─────────┐       ┌─────────────────┐
46/// │ │        A        │  │    2    │ │     │        A        │
47///   ├─────────────────┤  ├─────────┤       ├─────────────────┤
48/// │ │        D        │  │    3    │ │     │        A        │    run length of 'A' = runs_ends[0] - 0 = 2
49///   ├─────────────────┤  ├─────────┤       ├─────────────────┤
50/// │ │        B        │  │    6    │ │     │        D        │    run length of 'D' = run_ends[1] - run_ends[0] = 1
51///   └─────────────────┘  └─────────┘       ├─────────────────┤
52/// │        values          run_ends  │     │        B        │
53///                                          ├─────────────────┤
54/// └ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─┘     │        B        │
55///                                          ├─────────────────┤
56///                RunArray                  │        B        │    run length of 'B' = run_ends[2] - run_ends[1] = 3
57///               length = 3                 └─────────────────┘
58///
59///                                             Logical array
60///                                                Contents
61/// ```
62///
63/// [run-end encoded values]: https://arrow.apache.org/docs/format/Columnar.html#run-end-encoded-layout
64/// [run-length encoding (RLE)]: https://en.wikipedia.org/wiki/Run-length_encoding
65pub struct RunArray<R: RunEndIndexType> {
66    data_type: DataType,
67    run_ends: RunEndBuffer<R::Native>,
68    values: ArrayRef,
69}
70
71impl<R: RunEndIndexType> Clone for RunArray<R> {
72    fn clone(&self) -> Self {
73        Self {
74            data_type: self.data_type.clone(),
75            run_ends: self.run_ends.clone(),
76            values: self.values.clone(),
77        }
78    }
79}
80
81impl<R: RunEndIndexType> RunArray<R> {
82    /// Calculates the logical length of the array encoded by treating the `run_ends`
83    /// array as if it were a [`RunEndBuffer`].
84    pub fn logical_len(run_ends: &PrimitiveArray<R>) -> usize {
85        let len = run_ends.len();
86        if len == 0 {
87            return 0;
88        }
89        run_ends.value(len - 1).as_usize()
90    }
91
92    /// Attempts to create a [`RunArray`] using the given `run_ends` and `values`.
93    ///
94    /// # Errors
95    ///
96    /// - If `run_ends` and `values` have different lengths
97    /// - If `run_ends` has any null values
98    /// - If `run_ends` doesn't consist of strictly increasing positive integers
99    pub fn try_new(run_ends: &PrimitiveArray<R>, values: &dyn Array) -> Result<Self, ArrowError> {
100        let run_ends_type = run_ends.data_type().clone();
101        let values_type = values.data_type().clone();
102        let ree_array_type = DataType::RunEndEncoded(
103            Arc::new(Field::new("run_ends", run_ends_type, false)),
104            Arc::new(Field::new("values", values_type, true)),
105        );
106        let len = RunArray::logical_len(run_ends);
107        let builder = ArrayDataBuilder::new(ree_array_type)
108            .len(len)
109            .add_child_data(run_ends.to_data())
110            .add_child_data(values.to_data());
111
112        // `build_unchecked` is used to avoid recursive validation of child arrays.
113        let array_data = unsafe { builder.build_unchecked() };
114
115        // Safety: `validate_data` checks below
116        //    1. The given array data has exactly two child arrays.
117        //    2. The first child array (run_ends) has valid data type.
118        //    3. run_ends array does not have null values
119        //    4. run_ends array has non-zero and strictly increasing values.
120        //    5. The length of run_ends array and values array are the same.
121        array_data.validate_data()?;
122
123        Ok(array_data.into())
124    }
125
126    /// Returns a reference to the [`RunEndBuffer`].
127    pub fn run_ends(&self) -> &RunEndBuffer<R::Native> {
128        &self.run_ends
129    }
130
131    /// Returns a reference to the values array.
132    ///
133    /// Any slicing of this [`RunArray`] array is **not** applied to the returned
134    /// values here and must be handled separately.
135    pub fn values(&self) -> &ArrayRef {
136        &self.values
137    }
138
139    /// Returns the physical index at which the array slice starts.
140    ///
141    /// See [`RunEndBuffer::get_start_physical_index`].
142    pub fn get_start_physical_index(&self) -> usize {
143        self.run_ends.get_start_physical_index()
144    }
145
146    /// Returns the physical index at which the array slice ends.
147    ///
148    /// See [`RunEndBuffer::get_end_physical_index`].
149    pub fn get_end_physical_index(&self) -> usize {
150        self.run_ends.get_end_physical_index()
151    }
152
153    /// Downcast this [`RunArray`] to a [`TypedRunArray`]
154    ///
155    /// ```
156    /// use arrow_array::{Array, ArrayAccessor, RunArray, StringArray, types::Int32Type};
157    ///
158    /// let orig = [Some("a"), Some("b"), None];
159    /// let run_array = RunArray::<Int32Type>::from_iter(orig);
160    /// let typed = run_array.downcast::<StringArray>().unwrap();
161    /// assert_eq!(typed.value(0), "a");
162    /// assert_eq!(typed.value(1), "b");
163    /// assert!(typed.values().is_null(2));
164    /// ```
165    pub fn downcast<V: 'static>(&self) -> Option<TypedRunArray<'_, R, V>> {
166        let values = self.values.as_any().downcast_ref()?;
167        Some(TypedRunArray {
168            run_array: self,
169            values,
170        })
171    }
172
173    /// Calls [`RunEndBuffer::get_physical_index`].
174    ///
175    /// The result is arbitrary if `logical_index >= self.len()`
176    pub fn get_physical_index(&self, logical_index: usize) -> usize {
177        self.run_ends.get_physical_index(logical_index)
178    }
179
180    /// Returns the physical indices corresponding to the provided logical indices.
181    ///
182    /// See [`RunEndBuffer::get_physical_indices`] for more details.
183    #[inline]
184    pub fn get_physical_indices<I>(&self, logical_indices: &[I]) -> Result<Vec<usize>, ArrowError>
185    where
186        I: ArrowNativeType,
187    {
188        self.run_ends()
189            .get_physical_indices(logical_indices)
190            .map_err(|index| {
191                ArrowError::InvalidArgumentError(format!(
192                    "Logical index {} is out of bounds for RunArray of length {}",
193                    index.as_usize(),
194                    self.len()
195                ))
196            })
197    }
198
199    /// Returns a zero-copy slice of this array with the indicated offset and length.
200    ///
201    /// # Panics
202    ///
203    /// - Specified slice (`offset` + `length`) exceeds existing length
204    pub fn slice(&self, offset: usize, length: usize) -> Self {
205        Self {
206            data_type: self.data_type.clone(),
207            run_ends: self.run_ends.slice(offset, length),
208            values: self.values.clone(),
209        }
210    }
211}
212
213impl<R: RunEndIndexType> From<ArrayData> for RunArray<R> {
214    // The method assumes the caller already validated the data using `ArrayData::validate_data()`
215    fn from(data: ArrayData) -> Self {
216        match data.data_type() {
217            DataType::RunEndEncoded(_, _) => {}
218            _ => {
219                panic!(
220                    "Invalid data type for RunArray. The data type should be DataType::RunEndEncoded"
221                );
222            }
223        }
224
225        // Safety
226        // ArrayData is valid
227        let child = &data.child_data()[0];
228        assert_eq!(child.data_type(), &R::DATA_TYPE, "Incorrect run ends type");
229        let run_ends = unsafe {
230            let scalar = ScalarBuffer::new(child.buffers()[0].clone(), child.offset(), child.len());
231            RunEndBuffer::new_unchecked(scalar, data.offset(), data.len())
232        };
233
234        let values = make_array(data.child_data()[1].clone());
235        Self {
236            data_type: data.data_type().clone(),
237            run_ends,
238            values,
239        }
240    }
241}
242
243impl<R: RunEndIndexType> From<RunArray<R>> for ArrayData {
244    fn from(array: RunArray<R>) -> Self {
245        let len = array.run_ends.len();
246        let offset = array.run_ends.offset();
247
248        let run_ends = ArrayDataBuilder::new(R::DATA_TYPE)
249            .len(array.run_ends.values().len())
250            .buffers(vec![array.run_ends.into_inner().into_inner()]);
251
252        let run_ends = unsafe { run_ends.build_unchecked() };
253
254        let builder = ArrayDataBuilder::new(array.data_type)
255            .len(len)
256            .offset(offset)
257            .child_data(vec![run_ends, array.values.to_data()]);
258
259        unsafe { builder.build_unchecked() }
260    }
261}
262
263/// SAFETY: Correctly implements the contract of Arrow Arrays
264unsafe impl<T: RunEndIndexType> Array for RunArray<T> {
265    fn as_any(&self) -> &dyn Any {
266        self
267    }
268
269    fn to_data(&self) -> ArrayData {
270        self.clone().into()
271    }
272
273    fn into_data(self) -> ArrayData {
274        self.into()
275    }
276
277    fn data_type(&self) -> &DataType {
278        &self.data_type
279    }
280
281    fn slice(&self, offset: usize, length: usize) -> ArrayRef {
282        Arc::new(self.slice(offset, length))
283    }
284
285    fn len(&self) -> usize {
286        self.run_ends.len()
287    }
288
289    fn is_empty(&self) -> bool {
290        self.run_ends.is_empty()
291    }
292
293    fn shrink_to_fit(&mut self) {
294        self.run_ends.shrink_to_fit();
295        self.values.shrink_to_fit();
296    }
297
298    fn offset(&self) -> usize {
299        self.run_ends.offset()
300    }
301
302    fn nulls(&self) -> Option<&NullBuffer> {
303        None
304    }
305
306    fn logical_nulls(&self) -> Option<NullBuffer> {
307        let len = self.len();
308        let nulls = self.values.logical_nulls()?;
309        let mut out = BooleanBufferBuilder::new(len);
310        let offset = self.run_ends.offset();
311        let mut valid_start = 0;
312        let mut last_end = 0;
313        for (idx, end) in self.run_ends.values().iter().enumerate() {
314            let end = end.as_usize();
315            if end < offset {
316                continue;
317            }
318            let end = (end - offset).min(len);
319            if nulls.is_null(idx) {
320                if valid_start < last_end {
321                    out.append_n(last_end - valid_start, true);
322                }
323                out.append_n(end - last_end, false);
324                valid_start = end;
325            }
326            last_end = end;
327            if end == len {
328                break;
329            }
330        }
331        if valid_start < len {
332            out.append_n(len - valid_start, true)
333        }
334        // Sanity check
335        assert_eq!(out.len(), len);
336        Some(out.finish().into())
337    }
338
339    fn is_nullable(&self) -> bool {
340        !self.is_empty() && self.values.is_nullable()
341    }
342
343    fn get_buffer_memory_size(&self) -> usize {
344        self.run_ends.inner().inner().capacity() + self.values.get_buffer_memory_size()
345    }
346
347    fn get_array_memory_size(&self) -> usize {
348        std::mem::size_of::<Self>()
349            + self.run_ends.inner().inner().capacity()
350            + self.values.get_array_memory_size()
351    }
352}
353
354impl<R: RunEndIndexType> std::fmt::Debug for RunArray<R> {
355    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
356        writeln!(
357            f,
358            "RunArray {{run_ends: {:?}, values: {:?}}}",
359            self.run_ends.values(),
360            self.values
361        )
362    }
363}
364
365/// Constructs a `RunArray` from an iterator of optional strings.
366///
367/// # Example:
368/// ```
369/// use arrow_array::{RunArray, PrimitiveArray, StringArray, types::Int16Type};
370///
371/// let test = vec!["a", "a", "b", "c", "c"];
372/// let array: RunArray<Int16Type> = test
373///     .iter()
374///     .map(|&x| if x == "b" { None } else { Some(x) })
375///     .collect();
376/// assert_eq!(
377///     "RunArray {run_ends: [2, 3, 5], values: StringArray\n[\n  \"a\",\n  null,\n  \"c\",\n]}\n",
378///     format!("{:?}", array)
379/// );
380/// ```
381impl<'a, T: RunEndIndexType> FromIterator<Option<&'a str>> for RunArray<T> {
382    fn from_iter<I: IntoIterator<Item = Option<&'a str>>>(iter: I) -> Self {
383        let it = iter.into_iter();
384        let (lower, _) = it.size_hint();
385        let mut builder = StringRunBuilder::with_capacity(lower, 256);
386        it.for_each(|i| {
387            builder.append_option(i);
388        });
389
390        builder.finish()
391    }
392}
393
394/// Constructs a `RunArray` from an iterator of strings.
395///
396/// # Example:
397///
398/// ```
399/// use arrow_array::{RunArray, PrimitiveArray, StringArray, types::Int16Type};
400///
401/// let test = vec!["a", "a", "b", "c"];
402/// let array: RunArray<Int16Type> = test.into_iter().collect();
403/// assert_eq!(
404///     "RunArray {run_ends: [2, 3, 4], values: StringArray\n[\n  \"a\",\n  \"b\",\n  \"c\",\n]}\n",
405///     format!("{:?}", array)
406/// );
407/// ```
408impl<'a, T: RunEndIndexType> FromIterator<&'a str> for RunArray<T> {
409    fn from_iter<I: IntoIterator<Item = &'a str>>(iter: I) -> Self {
410        let it = iter.into_iter();
411        let (lower, _) = it.size_hint();
412        let mut builder = StringRunBuilder::with_capacity(lower, 256);
413        it.for_each(|i| {
414            builder.append_value(i);
415        });
416
417        builder.finish()
418    }
419}
420
421///
422/// A [`RunArray`] with `i16` run ends
423///
424/// # Example: Using `collect`
425/// ```
426/// # use arrow_array::{Array, Int16RunArray, Int16Array, StringArray};
427/// # use std::sync::Arc;
428///
429/// let array: Int16RunArray = vec!["a", "a", "b", "c", "c"].into_iter().collect();
430/// let values: Arc<dyn Array> = Arc::new(StringArray::from(vec!["a", "b", "c"]));
431/// assert_eq!(array.run_ends().values(), &[2, 3, 5]);
432/// assert_eq!(array.values(), &values);
433/// ```
434pub type Int16RunArray = RunArray<Int16Type>;
435
436///
437/// A [`RunArray`] with `i32` run ends
438///
439/// # Example: Using `collect`
440/// ```
441/// # use arrow_array::{Array, Int32RunArray, Int32Array, StringArray};
442/// # use std::sync::Arc;
443///
444/// let array: Int32RunArray = vec!["a", "a", "b", "c", "c"].into_iter().collect();
445/// let values: Arc<dyn Array> = Arc::new(StringArray::from(vec!["a", "b", "c"]));
446/// assert_eq!(array.run_ends().values(), &[2, 3, 5]);
447/// assert_eq!(array.values(), &values);
448/// ```
449pub type Int32RunArray = RunArray<Int32Type>;
450
451///
452/// A [`RunArray`] with `i64` run ends
453///
454/// # Example: Using `collect`
455/// ```
456/// # use arrow_array::{Array, Int64RunArray, Int64Array, StringArray};
457/// # use std::sync::Arc;
458///
459/// let array: Int64RunArray = vec!["a", "a", "b", "c", "c"].into_iter().collect();
460/// let values: Arc<dyn Array> = Arc::new(StringArray::from(vec!["a", "b", "c"]));
461/// assert_eq!(array.run_ends().values(), &[2, 3, 5]);
462/// assert_eq!(array.values(), &values);
463/// ```
464pub type Int64RunArray = RunArray<Int64Type>;
465
466/// A [`RunArray`] typed typed on its child values array
467///
468/// Implements [`ArrayAccessor`] and [`IntoIterator`] allowing fast access to its elements
469///
470/// ```
471/// use arrow_array::{RunArray, StringArray, types::Int32Type};
472///
473/// let orig = ["a", "b", "a", "b"];
474/// let ree_array = RunArray::<Int32Type>::from_iter(orig);
475///
476/// // `TypedRunArray` allows you to access the values directly
477/// let typed = ree_array.downcast::<StringArray>().unwrap();
478///
479/// for (maybe_val, orig) in typed.into_iter().zip(orig) {
480///     assert_eq!(maybe_val.unwrap(), orig)
481/// }
482/// ```
483pub struct TypedRunArray<'a, R: RunEndIndexType, V> {
484    /// The run array
485    run_array: &'a RunArray<R>,
486
487    /// The values of the run_array
488    values: &'a V,
489}
490
491// Manually implement `Clone` to avoid `V: Clone` type constraint
492impl<R: RunEndIndexType, V> Clone for TypedRunArray<'_, R, V> {
493    fn clone(&self) -> Self {
494        *self
495    }
496}
497
498impl<R: RunEndIndexType, V> Copy for TypedRunArray<'_, R, V> {}
499
500impl<R: RunEndIndexType, V> std::fmt::Debug for TypedRunArray<'_, R, V> {
501    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
502        writeln!(f, "TypedRunArray({:?})", self.run_array)
503    }
504}
505
506impl<'a, R: RunEndIndexType, V> TypedRunArray<'a, R, V> {
507    /// Returns the run_ends of this [`TypedRunArray`]
508    pub fn run_ends(&self) -> &'a RunEndBuffer<R::Native> {
509        self.run_array.run_ends()
510    }
511
512    /// Returns the values of this [`TypedRunArray`]
513    pub fn values(&self) -> &'a V {
514        self.values
515    }
516
517    /// Returns the run array of this [`TypedRunArray`]
518    pub fn run_array(&self) -> &'a RunArray<R> {
519        self.run_array
520    }
521}
522
523/// SAFETY: Correctly implements the contract of Arrow Arrays
524unsafe impl<R: RunEndIndexType, V: Sync> Array for TypedRunArray<'_, R, V> {
525    fn as_any(&self) -> &dyn Any {
526        self.run_array
527    }
528
529    fn to_data(&self) -> ArrayData {
530        self.run_array.to_data()
531    }
532
533    fn into_data(self) -> ArrayData {
534        self.run_array.into_data()
535    }
536
537    fn data_type(&self) -> &DataType {
538        self.run_array.data_type()
539    }
540
541    fn slice(&self, offset: usize, length: usize) -> ArrayRef {
542        Arc::new(self.run_array.slice(offset, length))
543    }
544
545    fn len(&self) -> usize {
546        self.run_array.len()
547    }
548
549    fn is_empty(&self) -> bool {
550        self.run_array.is_empty()
551    }
552
553    fn offset(&self) -> usize {
554        self.run_array.offset()
555    }
556
557    fn nulls(&self) -> Option<&NullBuffer> {
558        self.run_array.nulls()
559    }
560
561    fn logical_nulls(&self) -> Option<NullBuffer> {
562        self.run_array.logical_nulls()
563    }
564
565    fn logical_null_count(&self) -> usize {
566        self.run_array.logical_null_count()
567    }
568
569    fn is_nullable(&self) -> bool {
570        self.run_array.is_nullable()
571    }
572
573    fn get_buffer_memory_size(&self) -> usize {
574        self.run_array.get_buffer_memory_size()
575    }
576
577    fn get_array_memory_size(&self) -> usize {
578        self.run_array.get_array_memory_size()
579    }
580}
581
582// Array accessor converts the index of logical array to the index of the physical array
583// using binary search. The time complexity is O(log N) where N is number of runs.
584impl<'a, R, V> ArrayAccessor for TypedRunArray<'a, R, V>
585where
586    R: RunEndIndexType,
587    V: Sync + Send,
588    &'a V: ArrayAccessor,
589    <&'a V as ArrayAccessor>::Item: Default,
590{
591    type Item = <&'a V as ArrayAccessor>::Item;
592
593    fn value(&self, logical_index: usize) -> Self::Item {
594        assert!(
595            logical_index < self.len(),
596            "Trying to access an element at index {} from a TypedRunArray of length {}",
597            logical_index,
598            self.len()
599        );
600        unsafe { self.value_unchecked(logical_index) }
601    }
602
603    unsafe fn value_unchecked(&self, logical_index: usize) -> Self::Item {
604        let physical_index = self.run_array.get_physical_index(logical_index);
605        unsafe { self.values().value_unchecked(physical_index) }
606    }
607}
608
609impl<'a, R, V> IntoIterator for TypedRunArray<'a, R, V>
610where
611    R: RunEndIndexType,
612    V: Sync + Send,
613    &'a V: ArrayAccessor,
614    <&'a V as ArrayAccessor>::Item: Default,
615{
616    type Item = Option<<&'a V as ArrayAccessor>::Item>;
617    type IntoIter = RunArrayIter<'a, R, V>;
618
619    fn into_iter(self) -> Self::IntoIter {
620        RunArrayIter::new(self)
621    }
622}
623
624#[cfg(test)]
625mod tests {
626    use rand::Rng;
627    use rand::rng;
628    use rand::seq::SliceRandom;
629
630    use super::*;
631    use crate::Float32Array;
632    use crate::builder::PrimitiveRunBuilder;
633    use crate::cast::AsArray;
634    use crate::types::{Int8Type, UInt32Type};
635    use crate::{Int16Array, Int32Array, StringArray};
636
637    fn build_input_array(size: usize) -> Vec<Option<i32>> {
638        // The input array is created by shuffling and repeating
639        // the seed values random number of times.
640        let mut seed: Vec<Option<i32>> = vec![
641            None,
642            None,
643            None,
644            Some(1),
645            Some(2),
646            Some(3),
647            Some(4),
648            Some(5),
649            Some(6),
650            Some(7),
651            Some(8),
652            Some(9),
653        ];
654        let mut result: Vec<Option<i32>> = Vec::with_capacity(size);
655        let mut ix = 0;
656        let mut rng = rng();
657        // run length can go up to 8. Cap the max run length for smaller arrays to size / 2.
658        let max_run_length = 8_usize.min(1_usize.max(size / 2));
659        while result.len() < size {
660            // shuffle the seed array if all the values are iterated.
661            if ix == 0 {
662                seed.shuffle(&mut rng);
663            }
664            // repeat the items between 1 and 8 times. Cap the length for smaller sized arrays
665            let num = max_run_length.min(rng.random_range(1..=max_run_length));
666            for _ in 0..num {
667                result.push(seed[ix]);
668            }
669            ix += 1;
670            if ix == seed.len() {
671                ix = 0
672            }
673        }
674        result.resize(size, None);
675        result
676    }
677
678    // Asserts that `logical_array[logical_indices[*]] == physical_array[physical_indices[*]]`
679    fn compare_logical_and_physical_indices(
680        logical_indices: &[u32],
681        logical_array: &[Option<i32>],
682        physical_indices: &[usize],
683        physical_array: &PrimitiveArray<Int32Type>,
684    ) {
685        assert_eq!(logical_indices.len(), physical_indices.len());
686
687        // check value in logical index in the logical_array matches physical index in physical_array
688        logical_indices
689            .iter()
690            .map(|f| f.as_usize())
691            .zip(physical_indices.iter())
692            .for_each(|(logical_ix, physical_ix)| {
693                let expected = logical_array[logical_ix];
694                match expected {
695                    Some(val) => {
696                        assert!(physical_array.is_valid(*physical_ix));
697                        let actual = physical_array.value(*physical_ix);
698                        assert_eq!(val, actual);
699                    }
700                    None => {
701                        assert!(physical_array.is_null(*physical_ix))
702                    }
703                };
704            });
705    }
706    #[test]
707    fn test_run_array() {
708        // Construct a value array
709        let value_data =
710            PrimitiveArray::<Int8Type>::from_iter_values([10_i8, 11, 12, 13, 14, 15, 16, 17]);
711
712        // Construct a run_ends array:
713        let run_ends_values = [4_i16, 6, 7, 9, 13, 18, 20, 22];
714        let run_ends_data =
715            PrimitiveArray::<Int16Type>::from_iter_values(run_ends_values.iter().copied());
716
717        // Construct a run ends encoded array from the above two
718        let ree_array = RunArray::<Int16Type>::try_new(&run_ends_data, &value_data).unwrap();
719
720        assert_eq!(ree_array.len(), 22);
721        assert_eq!(ree_array.null_count(), 0);
722
723        let values = ree_array.values();
724        assert_eq!(value_data.into_data(), values.to_data());
725        assert_eq!(&DataType::Int8, values.data_type());
726
727        let run_ends = ree_array.run_ends();
728        assert_eq!(run_ends.values(), &run_ends_values);
729    }
730
731    #[test]
732    fn test_run_array_fmt_debug() {
733        let mut builder = PrimitiveRunBuilder::<Int16Type, UInt32Type>::with_capacity(3);
734        builder.append_value(12345678);
735        builder.append_null();
736        builder.append_value(22345678);
737        let array = builder.finish();
738        assert_eq!(
739            "RunArray {run_ends: [1, 2, 3], values: PrimitiveArray<UInt32>\n[\n  12345678,\n  null,\n  22345678,\n]}\n",
740            format!("{array:?}")
741        );
742
743        let mut builder = PrimitiveRunBuilder::<Int16Type, UInt32Type>::with_capacity(20);
744        for _ in 0..20 {
745            builder.append_value(1);
746        }
747        let array = builder.finish();
748
749        assert_eq!(array.len(), 20);
750        assert_eq!(array.null_count(), 0);
751        assert_eq!(array.logical_null_count(), 0);
752
753        assert_eq!(
754            "RunArray {run_ends: [20], values: PrimitiveArray<UInt32>\n[\n  1,\n]}\n",
755            format!("{array:?}")
756        );
757    }
758
759    #[test]
760    fn test_run_array_from_iter() {
761        let test = vec!["a", "a", "b", "c"];
762        let array: RunArray<Int16Type> = test
763            .iter()
764            .map(|&x| if x == "b" { None } else { Some(x) })
765            .collect();
766        assert_eq!(
767            "RunArray {run_ends: [2, 3, 4], values: StringArray\n[\n  \"a\",\n  null,\n  \"c\",\n]}\n",
768            format!("{array:?}")
769        );
770
771        assert_eq!(array.len(), 4);
772        assert_eq!(array.null_count(), 0);
773        assert_eq!(array.logical_null_count(), 1);
774
775        let array: RunArray<Int16Type> = test.into_iter().collect();
776        assert_eq!(
777            "RunArray {run_ends: [2, 3, 4], values: StringArray\n[\n  \"a\",\n  \"b\",\n  \"c\",\n]}\n",
778            format!("{array:?}")
779        );
780    }
781
782    #[test]
783    fn test_run_array_run_ends_as_primitive_array() {
784        let test = vec!["a", "b", "c", "a"];
785        let array: RunArray<Int16Type> = test.into_iter().collect();
786
787        assert_eq!(array.len(), 4);
788        assert_eq!(array.null_count(), 0);
789        assert_eq!(array.logical_null_count(), 0);
790
791        let run_ends = array.run_ends();
792        assert_eq!(&[1, 2, 3, 4], run_ends.values());
793    }
794
795    #[test]
796    fn test_run_array_as_primitive_array_with_null() {
797        let test = vec![Some("a"), None, Some("b"), None, None, Some("a")];
798        let array: RunArray<Int32Type> = test.into_iter().collect();
799
800        assert_eq!(array.len(), 6);
801        assert_eq!(array.null_count(), 0);
802        assert_eq!(array.logical_null_count(), 3);
803
804        let run_ends = array.run_ends();
805        assert_eq!(&[1, 2, 3, 5, 6], run_ends.values());
806
807        let values_data = array.values();
808        assert_eq!(2, values_data.null_count());
809        assert_eq!(5, values_data.len());
810    }
811
812    #[test]
813    fn test_run_array_all_nulls() {
814        let test = vec![None, None, None];
815        let array: RunArray<Int32Type> = test.into_iter().collect();
816
817        assert_eq!(array.len(), 3);
818        assert_eq!(array.null_count(), 0);
819        assert_eq!(array.logical_null_count(), 3);
820
821        let run_ends = array.run_ends();
822        assert_eq!(3, run_ends.len());
823        assert_eq!(&[3], run_ends.values());
824
825        let values_data = array.values();
826        assert_eq!(1, values_data.null_count());
827    }
828
829    #[test]
830    fn test_run_array_try_new() {
831        let values: StringArray = [Some("foo"), Some("bar"), None, Some("baz")]
832            .into_iter()
833            .collect();
834        let run_ends: Int32Array = [Some(1), Some(2), Some(3), Some(4)].into_iter().collect();
835
836        let array = RunArray::<Int32Type>::try_new(&run_ends, &values).unwrap();
837        assert_eq!(array.values().data_type(), &DataType::Utf8);
838
839        assert_eq!(array.null_count(), 0);
840        assert_eq!(array.logical_null_count(), 1);
841        assert_eq!(array.len(), 4);
842        assert_eq!(array.values().null_count(), 1);
843
844        assert_eq!(
845            "RunArray {run_ends: [1, 2, 3, 4], values: StringArray\n[\n  \"foo\",\n  \"bar\",\n  null,\n  \"baz\",\n]}\n",
846            format!("{array:?}")
847        );
848    }
849
850    #[test]
851    fn test_run_array_int16_type_definition() {
852        let array: Int16RunArray = vec!["a", "a", "b", "c", "c"].into_iter().collect();
853        let values: Arc<dyn Array> = Arc::new(StringArray::from(vec!["a", "b", "c"]));
854        assert_eq!(array.run_ends().values(), &[2, 3, 5]);
855        assert_eq!(array.values(), &values);
856    }
857
858    #[test]
859    fn test_run_array_empty_string() {
860        let array: Int16RunArray = vec!["a", "a", "", "", "c"].into_iter().collect();
861        let values: Arc<dyn Array> = Arc::new(StringArray::from(vec!["a", "", "c"]));
862        assert_eq!(array.run_ends().values(), &[2, 4, 5]);
863        assert_eq!(array.values(), &values);
864    }
865
866    #[test]
867    fn test_run_array_length_mismatch() {
868        let values: StringArray = [Some("foo"), Some("bar"), None, Some("baz")]
869            .into_iter()
870            .collect();
871        let run_ends: Int32Array = [Some(1), Some(2), Some(3)].into_iter().collect();
872
873        let actual = RunArray::<Int32Type>::try_new(&run_ends, &values);
874        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());
875        assert_eq!(expected.to_string(), actual.err().unwrap().to_string());
876    }
877
878    #[test]
879    fn test_run_array_run_ends_with_null() {
880        let values: StringArray = [Some("foo"), Some("bar"), Some("baz")]
881            .into_iter()
882            .collect();
883        let run_ends: Int32Array = [Some(1), None, Some(3)].into_iter().collect();
884
885        let actual = RunArray::<Int32Type>::try_new(&run_ends, &values);
886        let expected = ArrowError::InvalidArgumentError(
887            "Found null values in run_ends array. The run_ends array should not have null values."
888                .to_string(),
889        );
890        assert_eq!(expected.to_string(), actual.err().unwrap().to_string());
891    }
892
893    #[test]
894    fn test_run_array_run_ends_with_zeroes() {
895        let values: StringArray = [Some("foo"), Some("bar"), Some("baz")]
896            .into_iter()
897            .collect();
898        let run_ends: Int32Array = [Some(0), Some(1), Some(3)].into_iter().collect();
899
900        let actual = RunArray::<Int32Type>::try_new(&run_ends, &values);
901        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());
902        assert_eq!(expected.to_string(), actual.err().unwrap().to_string());
903    }
904
905    #[test]
906    fn test_run_array_run_ends_non_increasing() {
907        let values: StringArray = [Some("foo"), Some("bar"), Some("baz")]
908            .into_iter()
909            .collect();
910        let run_ends: Int32Array = [Some(1), Some(4), Some(4)].into_iter().collect();
911
912        let actual = RunArray::<Int32Type>::try_new(&run_ends, &values);
913        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());
914        assert_eq!(expected.to_string(), actual.err().unwrap().to_string());
915    }
916
917    #[test]
918    #[should_panic(expected = "Incorrect run ends type")]
919    fn test_run_array_run_ends_data_type_mismatch() {
920        let a = RunArray::<Int32Type>::from_iter(["32"]);
921        let _ = RunArray::<Int64Type>::from(a.into_data());
922    }
923
924    #[test]
925    fn test_ree_array_accessor() {
926        let input_array = build_input_array(256);
927
928        // Encode the input_array to ree_array
929        let mut builder =
930            PrimitiveRunBuilder::<Int16Type, Int32Type>::with_capacity(input_array.len());
931        builder.extend(input_array.iter().copied());
932        let run_array = builder.finish();
933        let typed = run_array.downcast::<PrimitiveArray<Int32Type>>().unwrap();
934
935        // Access every index and check if the value in the input array matches returned value.
936        for (i, inp_val) in input_array.iter().enumerate() {
937            if let Some(val) = inp_val {
938                let actual = typed.value(i);
939                assert_eq!(*val, actual)
940            } else {
941                let physical_ix = run_array.get_physical_index(i);
942                assert!(typed.values().is_null(physical_ix));
943            };
944        }
945    }
946
947    #[test]
948    #[cfg_attr(miri, ignore)] // Takes too long
949    fn test_get_physical_indices() {
950        // Test for logical lengths starting from 10 to 250 increasing by 10
951        for logical_len in (0..250).step_by(10) {
952            let input_array = build_input_array(logical_len);
953
954            // create run array using input_array
955            let mut builder = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
956            builder.extend(input_array.clone().into_iter());
957
958            let run_array = builder.finish();
959            let physical_values_array = run_array.values().as_primitive::<Int32Type>();
960
961            // create an array consisting of all the indices repeated twice and shuffled.
962            let mut logical_indices: Vec<u32> = (0_u32..(logical_len as u32)).collect();
963            // add same indices once more
964            logical_indices.append(&mut logical_indices.clone());
965            let mut rng = rng();
966            logical_indices.shuffle(&mut rng);
967
968            let physical_indices = run_array.get_physical_indices(&logical_indices).unwrap();
969
970            assert_eq!(logical_indices.len(), physical_indices.len());
971
972            // check value in logical index in the input_array matches physical index in typed_run_array
973            compare_logical_and_physical_indices(
974                &logical_indices,
975                &input_array,
976                &physical_indices,
977                physical_values_array,
978            );
979        }
980    }
981
982    #[test]
983    #[cfg_attr(miri, ignore)] // Takes too long
984    fn test_get_physical_indices_sliced() {
985        let total_len = 80;
986        let input_array = build_input_array(total_len);
987
988        // Encode the input_array to run array
989        let mut builder =
990            PrimitiveRunBuilder::<Int16Type, Int32Type>::with_capacity(input_array.len());
991        builder.extend(input_array.iter().copied());
992        let run_array = builder.finish();
993        let physical_values_array = run_array.values().as_primitive::<Int32Type>();
994
995        // test for all slice lengths.
996        for slice_len in 1..=total_len {
997            // create an array consisting of all the indices repeated twice and shuffled.
998            let mut logical_indices: Vec<u32> = (0_u32..(slice_len as u32)).collect();
999            // add same indices once more
1000            logical_indices.append(&mut logical_indices.clone());
1001            let mut rng = rng();
1002            logical_indices.shuffle(&mut rng);
1003
1004            // test for offset = 0 and slice length = slice_len
1005            // slice the input array using which the run array was built.
1006            let sliced_input_array = &input_array[0..slice_len];
1007
1008            // slice the run array
1009            let sliced_run_array: RunArray<Int16Type> =
1010                run_array.slice(0, slice_len).into_data().into();
1011
1012            // Get physical indices.
1013            let physical_indices = sliced_run_array
1014                .get_physical_indices(&logical_indices)
1015                .unwrap();
1016
1017            compare_logical_and_physical_indices(
1018                &logical_indices,
1019                sliced_input_array,
1020                &physical_indices,
1021                physical_values_array,
1022            );
1023
1024            // test for offset = total_len - slice_len and slice length = slice_len
1025            // slice the input array using which the run array was built.
1026            let sliced_input_array = &input_array[total_len - slice_len..total_len];
1027
1028            // slice the run array
1029            let sliced_run_array: RunArray<Int16Type> = run_array
1030                .slice(total_len - slice_len, slice_len)
1031                .into_data()
1032                .into();
1033
1034            // Get physical indices
1035            let physical_indices = sliced_run_array
1036                .get_physical_indices(&logical_indices)
1037                .unwrap();
1038
1039            compare_logical_and_physical_indices(
1040                &logical_indices,
1041                sliced_input_array,
1042                &physical_indices,
1043                physical_values_array,
1044            );
1045        }
1046    }
1047
1048    #[test]
1049    fn test_logical_nulls() {
1050        let run = Int32Array::from(vec![3, 6, 9, 12]);
1051        let values = Int32Array::from(vec![Some(0), None, Some(1), None]);
1052        let array = RunArray::try_new(&run, &values).unwrap();
1053
1054        let expected = [
1055            true, true, true, false, false, false, true, true, true, false, false, false,
1056        ];
1057
1058        let n = array.logical_nulls().unwrap();
1059        assert_eq!(n.null_count(), 6);
1060
1061        let slices = [(0, 12), (0, 2), (2, 5), (3, 0), (3, 3), (3, 4), (4, 8)];
1062        for (offset, length) in slices {
1063            let a = array.slice(offset, length);
1064            let n = a.logical_nulls().unwrap();
1065            let n = n.into_iter().collect::<Vec<_>>();
1066            assert_eq!(&n, &expected[offset..offset + length], "{offset} {length}");
1067        }
1068    }
1069
1070    #[test]
1071    fn test_run_array_eq_identical() {
1072        let run_ends1 = Int32Array::from(vec![2, 4, 6]);
1073        let values1 = StringArray::from(vec!["a", "b", "c"]);
1074        let array1 = RunArray::<Int32Type>::try_new(&run_ends1, &values1).unwrap();
1075
1076        let run_ends2 = Int32Array::from(vec![2, 4, 6]);
1077        let values2 = StringArray::from(vec!["a", "b", "c"]);
1078        let array2 = RunArray::<Int32Type>::try_new(&run_ends2, &values2).unwrap();
1079
1080        assert_eq!(array1, array2);
1081    }
1082
1083    #[test]
1084    fn test_run_array_ne_different_run_ends() {
1085        let run_ends1 = Int32Array::from(vec![2, 4, 6]);
1086        let values1 = StringArray::from(vec!["a", "b", "c"]);
1087        let array1 = RunArray::<Int32Type>::try_new(&run_ends1, &values1).unwrap();
1088
1089        let run_ends2 = Int32Array::from(vec![1, 4, 6]);
1090        let values2 = StringArray::from(vec!["a", "b", "c"]);
1091        let array2 = RunArray::<Int32Type>::try_new(&run_ends2, &values2).unwrap();
1092
1093        assert_ne!(array1, array2);
1094    }
1095
1096    #[test]
1097    fn test_run_array_ne_different_values() {
1098        let run_ends1 = Int32Array::from(vec![2, 4, 6]);
1099        let values1 = StringArray::from(vec!["a", "b", "c"]);
1100        let array1 = RunArray::<Int32Type>::try_new(&run_ends1, &values1).unwrap();
1101
1102        let run_ends2 = Int32Array::from(vec![2, 4, 6]);
1103        let values2 = StringArray::from(vec!["a", "b", "d"]);
1104        let array2 = RunArray::<Int32Type>::try_new(&run_ends2, &values2).unwrap();
1105
1106        assert_ne!(array1, array2);
1107    }
1108
1109    #[test]
1110    fn test_run_array_eq_with_nulls() {
1111        let run_ends1 = Int32Array::from(vec![2, 4, 6]);
1112        let values1 = StringArray::from(vec![Some("a"), None, Some("c")]);
1113        let array1 = RunArray::<Int32Type>::try_new(&run_ends1, &values1).unwrap();
1114
1115        let run_ends2 = Int32Array::from(vec![2, 4, 6]);
1116        let values2 = StringArray::from(vec![Some("a"), None, Some("c")]);
1117        let array2 = RunArray::<Int32Type>::try_new(&run_ends2, &values2).unwrap();
1118
1119        assert_eq!(array1, array2);
1120    }
1121
1122    #[test]
1123    fn test_run_array_eq_different_run_end_types() {
1124        let run_ends_i16_1 = Int16Array::from(vec![2_i16, 4, 6]);
1125        let values_i16_1 = StringArray::from(vec!["a", "b", "c"]);
1126        let array_i16_1 = RunArray::<Int16Type>::try_new(&run_ends_i16_1, &values_i16_1).unwrap();
1127
1128        let run_ends_i16_2 = Int16Array::from(vec![2_i16, 4, 6]);
1129        let values_i16_2 = StringArray::from(vec!["a", "b", "c"]);
1130        let array_i16_2 = RunArray::<Int16Type>::try_new(&run_ends_i16_2, &values_i16_2).unwrap();
1131
1132        assert_eq!(array_i16_1, array_i16_2);
1133    }
1134
1135    #[test]
1136    fn test_from_array_data_offset() {
1137        let dt = DataType::RunEndEncoded(
1138            Arc::new(Field::new("run_ends", DataType::Int32, false)),
1139            Arc::new(Field::new("values", DataType::Float32, true)),
1140        );
1141
1142        let runs = ArrayData::builder(DataType::Int32)
1143            .len(2)
1144            .offset(2)
1145            .add_buffer(vec![-10_i32, -20, 1_i32, 4, 0, 100].into())
1146            .build()
1147            .unwrap();
1148        let values = Float32Array::from(vec![1.5, 2.5]);
1149
1150        let data = ArrayDataBuilder::new(dt)
1151            .len(3)
1152            .child_data(vec![runs, values.to_data()])
1153            .build()
1154            .unwrap();
1155        let array = RunArray::<Int32Type>::from(data);
1156        let array = array.downcast::<Float32Array>().unwrap();
1157
1158        assert_eq!(array.value(0), 1.5);
1159        assert_eq!(array.value(1), 2.5);
1160        assert_eq!(array.value(2), 2.5);
1161    }
1162}
arrow_array/array/run_array.rs

arrow_array/array/
run_array.rs
