arrow_array/array/
mod.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
18//! The concrete array definitions
19
20mod binary_array;
21
22use crate::types::*;
23use arrow_buffer::{ArrowNativeType, Buffer, NullBuffer, OffsetBuffer, ScalarBuffer};
24use arrow_data::ArrayData;
25use arrow_schema::{DataType, IntervalUnit, TimeUnit};
26use std::any::Any;
27use std::sync::Arc;
28
29pub use binary_array::*;
30
31mod boolean_array;
32pub use boolean_array::*;
33
34mod byte_array;
35pub use byte_array::*;
36
37mod dictionary_array;
38pub use dictionary_array::*;
39
40mod fixed_size_binary_array;
41pub use fixed_size_binary_array::*;
42
43mod fixed_size_list_array;
44pub use fixed_size_list_array::*;
45
46mod list_array;
47pub use list_array::*;
48
49mod map_array;
50pub use map_array::*;
51
52mod null_array;
53pub use null_array::*;
54
55mod primitive_array;
56pub use primitive_array::*;
57
58mod string_array;
59pub use string_array::*;
60
61mod struct_array;
62pub use struct_array::*;
63
64mod union_array;
65pub use union_array::*;
66
67mod run_array;
68
69pub use run_array::*;
70
71mod byte_view_array;
72
73pub use byte_view_array::*;
74
75mod list_view_array;
76
77pub use list_view_array::*;
78
79use crate::iterator::ArrayIter;
80
81/// An array in the [Arrow Columnar Format](https://arrow.apache.org/docs/format/Columnar.html)
82///
83/// # Safety
84///
85/// Implementations of this trait must ensure that all methods implementations comply with
86/// the Arrow specification. No safety guards are placed and failing to comply with it can
87/// translate into panics or undefined behavior. For example, a value computed based on `len`
88/// may be used as a direct index into memory regions without checks.
89///
90/// Note that it is likely impossible to correctly implement the trait for a
91/// third party type, as substantial arrow-rs functionality is based on the
92/// return values of [`Array::data_type`] and third party types cannot extend
93/// the [`DataType`] enum. So any code that attempts casting based on data type
94/// (including internal arrow library code) risks a panic or undefined behavior.
95/// See [this discussion] for more details.
96///
97/// This trait might be sealed in the future. Use at your own risk.
98///
99/// [this discussion]: https://github.com/apache/arrow-rs/pull/9234#pullrequestreview-3708950936
100pub unsafe trait Array: std::fmt::Debug + Send + Sync {
101    /// Returns the array as [`Any`] so that it can be
102    /// downcasted to a specific implementation.
103    ///
104    /// # Example:
105    ///
106    /// ```
107    /// # use std::sync::Arc;
108    /// # use arrow_array::{Int32Array, RecordBatch};
109    /// # use arrow_schema::{Schema, Field, DataType, ArrowError};
110    ///
111    /// let id = Int32Array::from(vec![1, 2, 3, 4, 5]);
112    /// let batch = RecordBatch::try_new(
113    ///     Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])),
114    ///     vec![Arc::new(id)]
115    /// ).unwrap();
116    ///
117    /// let int32array = batch
118    ///     .column(0)
119    ///     .as_any()
120    ///     .downcast_ref::<Int32Array>()
121    ///     .expect("Failed to downcast");
122    /// ```
123    fn as_any(&self) -> &dyn Any;
124
125    /// Returns the underlying data of this array
126    fn to_data(&self) -> ArrayData;
127
128    /// Returns the underlying data of this array
129    ///
130    /// Unlike [`Array::to_data`] this consumes self, allowing it avoid unnecessary clones
131    fn into_data(self) -> ArrayData;
132
133    /// Returns a reference to the [`DataType`] of this array.
134    ///
135    /// # Example:
136    ///
137    /// ```
138    /// use arrow_schema::DataType;
139    /// use arrow_array::{Array, Int32Array};
140    ///
141    /// let array = Int32Array::from(vec![1, 2, 3, 4, 5]);
142    ///
143    /// assert_eq!(*array.data_type(), DataType::Int32);
144    /// ```
145    fn data_type(&self) -> &DataType;
146
147    /// Returns a zero-copy slice of this array with the indicated offset and length.
148    ///
149    /// # Example:
150    ///
151    /// ```
152    /// use arrow_array::{Array, Int32Array};
153    ///
154    /// let array = Int32Array::from(vec![1, 2, 3, 4, 5]);
155    /// // Make slice over the values [2, 3, 4]
156    /// let array_slice = array.slice(1, 3);
157    ///
158    /// assert_eq!(&array_slice, &Int32Array::from(vec![2, 3, 4]));
159    /// ```
160    ///
161    /// # Panics
162    /// Panics if `offset + length > self.len()`
163    fn slice(&self, offset: usize, length: usize) -> ArrayRef;
164
165    /// Returns the length (i.e., number of elements) of this array.
166    ///
167    /// # Example:
168    ///
169    /// ```
170    /// use arrow_array::{Array, Int32Array};
171    ///
172    /// let array = Int32Array::from(vec![1, 2, 3, 4, 5]);
173    ///
174    /// assert_eq!(array.len(), 5);
175    /// ```
176    fn len(&self) -> usize;
177
178    /// Returns whether this array is empty.
179    ///
180    /// # Example:
181    ///
182    /// ```
183    /// use arrow_array::{Array, Int32Array};
184    ///
185    /// let array = Int32Array::from(vec![1, 2, 3, 4, 5]);
186    ///
187    /// assert_eq!(array.is_empty(), false);
188    /// ```
189    fn is_empty(&self) -> bool;
190
191    /// Shrinks the capacity of any exclusively owned buffer as much as possible
192    ///
193    /// Shared or externally allocated buffers will be ignored, and
194    /// any buffer offsets will be preserved.
195    fn shrink_to_fit(&mut self) {}
196
197    /// Returns the offset into the underlying data used by this array(-slice).
198    /// Note that the underlying data can be shared by many arrays.
199    /// This defaults to `0`.
200    ///
201    /// # Example:
202    ///
203    /// ```
204    /// use arrow_array::{Array, BooleanArray};
205    ///
206    /// let array = BooleanArray::from(vec![false, false, true, true]);
207    /// let array_slice = array.slice(1, 3);
208    ///
209    /// assert_eq!(array.offset(), 0);
210    /// assert_eq!(array_slice.offset(), 1);
211    /// ```
212    fn offset(&self) -> usize;
213
214    /// Returns the null buffer of this array if any.
215    ///
216    /// The null buffer contains the "physical" nulls of an array, that is how
217    /// the nulls are represented in the underlying arrow format.
218    ///
219    /// The physical representation is efficient, but is sometimes non intuitive
220    /// for certain array types such as those with nullable child arrays like
221    /// [`DictionaryArray::values`], [`RunArray::values`] or [`UnionArray`], or without a
222    /// null buffer, such as [`NullArray`].
223    ///
224    /// To determine if each element of such an array is "logically" null,
225    /// use the slower [`Array::logical_nulls`] to obtain a computed mask.
226    fn nulls(&self) -> Option<&NullBuffer>;
227
228    /// Returns a potentially computed [`NullBuffer`] that represents the logical
229    /// null values of this array, if any.
230    ///
231    /// Logical nulls represent the values that are null in the array,
232    /// regardless of the underlying physical arrow representation.
233    ///
234    /// For most array types, this is equivalent to the "physical" nulls
235    /// returned by [`Array::nulls`]. It is different for the following cases, because which
236    /// elements are null is not encoded in a single null buffer:
237    ///
238    /// * [`DictionaryArray`] where [`DictionaryArray::values`] contains nulls
239    /// * [`RunArray`] where [`RunArray::values`] contains nulls
240    /// * [`NullArray`] where all indices are nulls
241    /// * [`UnionArray`] where the selected values contains nulls
242    ///
243    /// In these cases a logical [`NullBuffer`] will be computed, encoding the
244    /// logical nullability of these arrays, beyond what is encoded in
245    /// [`Array::nulls`]
246    fn logical_nulls(&self) -> Option<NullBuffer> {
247        self.nulls().cloned()
248    }
249
250    /// Returns whether the element at `index` is null according to [`Array::nulls`]
251    ///
252    /// Note: For performance reasons, this method returns nullability solely as determined by the
253    /// null buffer. This difference can lead to surprising results, for example, [`NullArray::is_null`] always
254    /// returns `false` as the array lacks a null buffer. Similarly [`DictionaryArray`], [`RunArray`] and [`UnionArray`] may
255    /// encode nullability in their children. See [`Self::logical_nulls`] for more information.
256    ///
257    /// # Example:
258    ///
259    /// ```
260    /// use arrow_array::{Array, Int32Array, NullArray};
261    ///
262    /// let array = Int32Array::from(vec![Some(1), None]);
263    /// assert_eq!(array.is_null(0), false);
264    /// assert_eq!(array.is_null(1), true);
265    ///
266    /// // NullArrays do not have a null buffer, and therefore always
267    /// // return false for is_null.
268    /// let array = NullArray::new(1);
269    /// assert_eq!(array.is_null(0), false);
270    /// ```
271    ///
272    /// # Panics
273    ///
274    /// Panics if `index >= self.len()`.
275    ///
276    /// Note: arrays without a null buffer currently return `false` instead of
277    /// panicking, but callers must not rely on this.
278    fn is_null(&self, index: usize) -> bool {
279        self.nulls().is_some_and(|n| n.is_null(index))
280    }
281
282    /// Returns whether the element at `index` is *not* null, the
283    /// opposite of [`Self::is_null`].
284    ///
285    /// # Example:
286    ///
287    /// ```
288    /// use arrow_array::{Array, Int32Array};
289    ///
290    /// let array = Int32Array::from(vec![Some(1), None]);
291    ///
292    /// assert_eq!(array.is_valid(0), true);
293    /// assert_eq!(array.is_valid(1), false);
294    /// ```
295    ///
296    /// # Panics
297    ///
298    /// Panics if `index >= self.len()`.
299    ///
300    /// Note: arrays without a null buffer currently return `true` instead of
301    /// panicking, but callers must not rely on this.
302    fn is_valid(&self, index: usize) -> bool {
303        !self.is_null(index)
304    }
305
306    /// Returns the total number of physical null values in this array.
307    ///
308    /// Note: this method returns the physical null count, i.e. that encoded in [`Array::nulls`],
309    /// see [`Array::logical_nulls`] for logical nullability
310    ///
311    /// # Example:
312    ///
313    /// ```
314    /// use arrow_array::{Array, Int32Array};
315    ///
316    /// // Construct an array with values [1, NULL, NULL]
317    /// let array = Int32Array::from(vec![Some(1), None, None]);
318    ///
319    /// assert_eq!(array.null_count(), 2);
320    /// ```
321    fn null_count(&self) -> usize {
322        self.nulls().map(|n| n.null_count()).unwrap_or_default()
323    }
324
325    /// Returns the total number of logical null values in this array.
326    ///
327    /// Note: this method returns the logical null count, i.e. that encoded in
328    /// [`Array::logical_nulls`]. In general this is equivalent to [`Array::null_count`] but may differ in the
329    /// presence of logical nullability, see [`Array::nulls`] and [`Array::logical_nulls`].
330    ///
331    /// # Example:
332    ///
333    /// ```
334    /// use arrow_array::{Array, Int32Array};
335    ///
336    /// // Construct an array with values [1, NULL, NULL]
337    /// let array = Int32Array::from(vec![Some(1), None, None]);
338    ///
339    /// assert_eq!(array.logical_null_count(), 2);
340    /// ```
341    fn logical_null_count(&self) -> usize {
342        self.logical_nulls()
343            .map(|n| n.null_count())
344            .unwrap_or_default()
345    }
346
347    /// Returns `false` if the array is guaranteed to not contain any logical nulls
348    ///
349    /// This is generally equivalent to `Array::logical_null_count() != 0` unless determining
350    /// the logical nulls is expensive, in which case this method can return true even for an
351    /// array without nulls.
352    ///
353    /// This is also generally equivalent to `Array::null_count() != 0` but may differ in the
354    /// presence of logical nullability, see [`Array::logical_null_count`] and [`Array::null_count`].
355    ///
356    /// Implementations will return `true` unless they can cheaply prove no logical nulls
357    /// are present. For example a [`DictionaryArray`] with nullable values will still return true,
358    /// even if the nulls present in [`DictionaryArray::values`] are not referenced by any key,
359    /// and therefore would not appear in [`Array::logical_nulls`].
360    fn is_nullable(&self) -> bool {
361        self.logical_null_count() != 0
362    }
363
364    /// Returns the total number of bytes of memory pointed to by this array.
365    /// The buffers store bytes in the Arrow memory format, and include the data as well as the validity map.
366    /// Note that this does not always correspond to the exact memory usage of an array,
367    /// since multiple arrays can share the same buffers or slices thereof.
368    fn get_buffer_memory_size(&self) -> usize;
369
370    /// Returns the total number of bytes of memory occupied physically by this array.
371    /// This value will always be greater than returned by `get_buffer_memory_size()` and
372    /// includes the overhead of the data structures that contain the pointers to the various buffers.
373    fn get_array_memory_size(&self) -> usize;
374
375    /// Claim memory used by this array in the provided memory pool.
376    ///
377    /// This recursively claims memory for:
378    /// - All data buffers in this array
379    /// - All child arrays (for nested types like List, Struct, etc.)
380    /// - The null bitmap buffer if present
381    ///
382    /// This method guarantees that the memory pool will only compute occupied memory
383    /// exactly once. For example, if this array is derived from operations like `slice`,
384    /// calling `claim` on it would not change the memory pool's usage if the underlying buffers
385    /// are already counted before.
386    ///
387    /// # Example
388    /// ```
389    /// # use arrow_array::{Int32Array, Array};
390    /// # use arrow_buffer::TrackingMemoryPool;
391    /// # use arrow_buffer::MemoryPool;
392    ///
393    /// let pool = TrackingMemoryPool::default();
394    ///
395    /// let small_array = Int32Array::from(vec![1, 2, 3, 4, 5]);
396    /// let small_array_size = small_array.get_buffer_memory_size();
397    ///
398    /// // Claim the array's memory in the pool
399    /// small_array.claim(&pool);
400    ///
401    /// // Create and claim slices of `small_array`; should not increase memory usage
402    /// let slice1 = small_array.slice(0, 2);
403    /// let slice2 = small_array.slice(2, 2);
404    /// slice1.claim(&pool);
405    /// slice2.claim(&pool);
406    ///
407    /// assert_eq!(pool.used(), small_array_size);
408    ///
409    /// // Create a `large_array` which does not derive from the original `small_array`
410    ///
411    /// let large_array = Int32Array::from((0..1000).collect::<Vec<i32>>());
412    /// let large_array_size = large_array.get_buffer_memory_size();
413    ///
414    /// large_array.claim(&pool);
415    ///
416    /// // Trying to claim more than once is a no-op
417    /// large_array.claim(&pool);
418    /// large_array.claim(&pool);
419    ///
420    /// assert_eq!(pool.used(), small_array_size + large_array_size);
421    ///
422    /// let sum_of_all_sizes = small_array_size + large_array_size + slice1.get_buffer_memory_size() + slice2.get_buffer_memory_size();
423    ///
424    /// // `get_buffer_memory_size` works independently of the memory pool, so a sum of all the
425    /// // arrays in scope will always be >= the memory used reported by the memory pool.
426    /// assert_ne!(pool.used(), sum_of_all_sizes);
427    ///
428    /// // Until the final claim is dropped the buffer size remains accounted for
429    /// drop(small_array);
430    /// drop(slice1);
431    ///
432    /// assert_eq!(pool.used(), small_array_size + large_array_size);
433    ///
434    /// // Dropping this finally releases the buffer that was backing `small_array`
435    /// drop(slice2);
436    ///
437    /// assert_eq!(pool.used(), large_array_size);
438    /// ```
439    #[cfg(feature = "pool")]
440    fn claim(&self, pool: &dyn arrow_buffer::MemoryPool) {
441        self.to_data().claim(pool)
442    }
443}
444
445/// A reference-counted reference to a generic `Array`
446pub type ArrayRef = Arc<dyn Array>;
447
448/// Ergonomics: Allow use of an ArrayRef as an `&dyn Array`
449unsafe impl Array for ArrayRef {
450    fn as_any(&self) -> &dyn Any {
451        self.as_ref().as_any()
452    }
453
454    fn to_data(&self) -> ArrayData {
455        self.as_ref().to_data()
456    }
457
458    fn into_data(self) -> ArrayData {
459        self.to_data()
460    }
461
462    fn data_type(&self) -> &DataType {
463        self.as_ref().data_type()
464    }
465
466    fn slice(&self, offset: usize, length: usize) -> ArrayRef {
467        self.as_ref().slice(offset, length)
468    }
469
470    fn len(&self) -> usize {
471        self.as_ref().len()
472    }
473
474    fn is_empty(&self) -> bool {
475        self.as_ref().is_empty()
476    }
477
478    /// For shared buffers, this is a no-op.
479    fn shrink_to_fit(&mut self) {
480        if let Some(slf) = Arc::get_mut(self) {
481            slf.shrink_to_fit();
482        } else {
483            // We ignore shared buffers.
484        }
485    }
486
487    fn offset(&self) -> usize {
488        self.as_ref().offset()
489    }
490
491    fn nulls(&self) -> Option<&NullBuffer> {
492        self.as_ref().nulls()
493    }
494
495    fn logical_nulls(&self) -> Option<NullBuffer> {
496        self.as_ref().logical_nulls()
497    }
498
499    fn is_null(&self, index: usize) -> bool {
500        self.as_ref().is_null(index)
501    }
502
503    fn is_valid(&self, index: usize) -> bool {
504        self.as_ref().is_valid(index)
505    }
506
507    fn null_count(&self) -> usize {
508        self.as_ref().null_count()
509    }
510
511    fn logical_null_count(&self) -> usize {
512        self.as_ref().logical_null_count()
513    }
514
515    fn is_nullable(&self) -> bool {
516        self.as_ref().is_nullable()
517    }
518
519    fn get_buffer_memory_size(&self) -> usize {
520        self.as_ref().get_buffer_memory_size()
521    }
522
523    fn get_array_memory_size(&self) -> usize {
524        self.as_ref().get_array_memory_size()
525    }
526
527    #[cfg(feature = "pool")]
528    fn claim(&self, pool: &dyn arrow_buffer::MemoryPool) {
529        self.as_ref().claim(pool)
530    }
531}
532
533unsafe impl<T: Array> Array for &T {
534    fn as_any(&self) -> &dyn Any {
535        T::as_any(self)
536    }
537
538    fn to_data(&self) -> ArrayData {
539        T::to_data(self)
540    }
541
542    fn into_data(self) -> ArrayData {
543        self.to_data()
544    }
545
546    fn data_type(&self) -> &DataType {
547        T::data_type(self)
548    }
549
550    fn slice(&self, offset: usize, length: usize) -> ArrayRef {
551        T::slice(self, offset, length)
552    }
553
554    fn len(&self) -> usize {
555        T::len(self)
556    }
557
558    fn is_empty(&self) -> bool {
559        T::is_empty(self)
560    }
561
562    fn offset(&self) -> usize {
563        T::offset(self)
564    }
565
566    fn nulls(&self) -> Option<&NullBuffer> {
567        T::nulls(self)
568    }
569
570    fn logical_nulls(&self) -> Option<NullBuffer> {
571        T::logical_nulls(self)
572    }
573
574    fn is_null(&self, index: usize) -> bool {
575        T::is_null(self, index)
576    }
577
578    fn is_valid(&self, index: usize) -> bool {
579        T::is_valid(self, index)
580    }
581
582    fn null_count(&self) -> usize {
583        T::null_count(self)
584    }
585
586    fn logical_null_count(&self) -> usize {
587        T::logical_null_count(self)
588    }
589
590    fn is_nullable(&self) -> bool {
591        T::is_nullable(self)
592    }
593
594    fn get_buffer_memory_size(&self) -> usize {
595        T::get_buffer_memory_size(self)
596    }
597
598    fn get_array_memory_size(&self) -> usize {
599        T::get_array_memory_size(self)
600    }
601
602    #[cfg(feature = "pool")]
603    fn claim(&self, pool: &dyn arrow_buffer::MemoryPool) {
604        T::claim(self, pool)
605    }
606}
607
608/// A generic trait for accessing the values of an [`Array`]
609///
610/// This trait helps write specialized implementations of algorithms for
611/// different array types. Specialized implementations allow the compiler
612/// to optimize the code for the specific array type, which can lead to
613/// significant performance improvements.
614///
615/// # Example
616/// For example, to write three different implementations of a string length function
617/// for [`StringArray`], [`LargeStringArray`], and [`StringViewArray`], you can write
618///
619/// ```
620/// # use std::sync::Arc;
621/// # use arrow_array::{ArrayAccessor, ArrayRef, ArrowPrimitiveType, OffsetSizeTrait, PrimitiveArray};
622/// # use arrow_buffer::ArrowNativeType;
623/// # use arrow_array::cast::AsArray;
624/// # use arrow_array::iterator::ArrayIter;
625/// # use arrow_array::types::{Int32Type, Int64Type};
626/// # use arrow_schema::{ArrowError, DataType};
627/// /// This function takes a dynamically typed `ArrayRef` and calls
628/// /// calls one of three specialized implementations
629/// fn character_length(arg: ArrayRef) -> Result<ArrayRef, ArrowError> {
630///     match arg.data_type() {
631///         DataType::Utf8 => {
632///             // downcast the ArrayRef to a StringArray and call the specialized implementation
633///             let string_array = arg.as_string::<i32>();
634///             character_length_general::<Int32Type, _>(string_array)
635///         }
636///         DataType::LargeUtf8 => {
637///             character_length_general::<Int64Type, _>(arg.as_string::<i64>())
638///         }
639///         DataType::Utf8View => {
640///             character_length_general::<Int32Type, _>(arg.as_string_view())
641///         }
642///         _ => Err(ArrowError::InvalidArgumentError("Unsupported data type".to_string())),
643///     }
644/// }
645///
646/// /// A generic implementation of the character_length function
647/// /// This function uses the `ArrayAccessor` trait to access the values of the array
648/// /// so the compiler can generated specialized implementations for different array types
649/// ///
650/// /// Returns a new array with the length of each string in the input array
651/// /// * Int32Array for Utf8 and Utf8View arrays (lengths are 32-bit integers)
652/// /// * Int64Array for LargeUtf8 arrays (lengths are 64-bit integers)
653/// ///
654/// /// This is generic on the type of the primitive array (different string arrays have
655/// /// different lengths) and the type of the array accessor (different string arrays
656/// /// have different ways to access the values)
657/// fn character_length_general<'a, T: ArrowPrimitiveType, V: ArrayAccessor<Item = &'a str>>(
658///     array: V,
659/// ) -> Result<ArrayRef, ArrowError>
660/// where
661///     T::Native: OffsetSizeTrait,
662/// {
663///     let iter = ArrayIter::new(array);
664///     // Create a Int32Array / Int64Array with the length of each string
665///     let result = iter
666///         .map(|string| {
667///             string.map(|string: &str| {
668///                 T::Native::from_usize(string.chars().count())
669///                     .expect("should not fail as string.chars will always return integer")
670///             })
671///         })
672///         .collect::<PrimitiveArray<T>>();
673///
674///     /// Return the result as a new ArrayRef (dynamically typed)
675///     Ok(Arc::new(result) as ArrayRef)
676/// }
677/// ```
678///
679/// # Validity
680///
681/// An [`ArrayAccessor`] must always return a well-defined value for an index
682/// that is within the bounds `0..Array::len`, including for null indexes where
683/// [`Array::is_null`] is true.
684///
685/// The value at null indexes is unspecified, and implementations must not rely
686/// on a specific value such as [`Default::default`] being returned, however, it
687/// must not be undefined
688pub trait ArrayAccessor: Array {
689    /// The Arrow type of the element being accessed.
690    type Item: Send + Sync;
691
692    /// Returns the element at index `i`
693    /// # Panics
694    /// Panics if the value is outside the bounds of the array
695    fn value(&self, index: usize) -> Self::Item;
696
697    /// Returns the element at index `i`
698    /// # Safety
699    /// Caller is responsible for ensuring that the index is within the bounds of the array
700    unsafe fn value_unchecked(&self, index: usize) -> Self::Item;
701}
702
703/// A trait for Arrow String Arrays, currently three types are supported:
704/// - `StringArray`
705/// - `LargeStringArray`
706/// - `StringViewArray`
707///
708/// This trait helps to abstract over the different types of string arrays
709/// so that we don't need to duplicate the implementation for each type.
710pub trait StringArrayType<'a>: ArrayAccessor<Item = &'a str> + Sized {
711    /// Returns true if all data within this string array is ASCII
712    fn is_ascii(&self) -> bool;
713
714    /// Constructs a new iterator
715    fn iter(&self) -> ArrayIter<Self>;
716}
717
718impl<'a, O: OffsetSizeTrait> StringArrayType<'a> for &'a GenericStringArray<O> {
719    fn is_ascii(&self) -> bool {
720        GenericStringArray::<O>::is_ascii(self)
721    }
722
723    fn iter(&self) -> ArrayIter<Self> {
724        GenericStringArray::<O>::iter(self)
725    }
726}
727impl<'a> StringArrayType<'a> for &'a StringViewArray {
728    fn is_ascii(&self) -> bool {
729        StringViewArray::is_ascii(self)
730    }
731
732    fn iter(&self) -> ArrayIter<Self> {
733        StringViewArray::iter(self)
734    }
735}
736
737/// A trait for Arrow Binary Arrays, currently four types are supported:
738/// - `BinaryArray`
739/// - `LargeBinaryArray`
740/// - `BinaryViewArray`
741/// - `FixedSizeBinaryArray`
742///
743/// This trait helps to abstract over the different types of binary arrays
744/// so that we don't need to duplicate the implementation for each type.
745pub trait BinaryArrayType<'a>: ArrayAccessor<Item = &'a [u8]> + Sized {
746    /// Constructs a new iterator
747    fn iter(&self) -> ArrayIter<Self>;
748}
749
750impl<'a, O: OffsetSizeTrait> BinaryArrayType<'a> for &'a GenericBinaryArray<O> {
751    fn iter(&self) -> ArrayIter<Self> {
752        GenericBinaryArray::<O>::iter(self)
753    }
754}
755impl<'a> BinaryArrayType<'a> for &'a BinaryViewArray {
756    fn iter(&self) -> ArrayIter<Self> {
757        BinaryViewArray::iter(self)
758    }
759}
760impl<'a> BinaryArrayType<'a> for &'a FixedSizeBinaryArray {
761    fn iter(&self) -> ArrayIter<Self> {
762        FixedSizeBinaryArray::iter(self)
763    }
764}
765
766/// A trait for Arrow list-like arrays, abstracting over
767/// [`GenericListArray`], [`GenericListViewArray`], and [`FixedSizeListArray`].
768///
769/// This trait provides a uniform interface for accessing the child values and
770/// computing the element range for a given index, regardless of the underlying
771/// list layout (offsets, offsets+sizes, or fixed-size).
772pub trait ListLikeArray: Array {
773    /// Returns the child values array.
774    fn values(&self) -> &ArrayRef;
775
776    /// Returns the start and end indices into the values array for the list
777    /// element at `index`.
778    fn element_range(&self, index: usize) -> std::ops::Range<usize>;
779}
780
781impl PartialEq for dyn Array + '_ {
782    fn eq(&self, other: &Self) -> bool {
783        self.to_data().eq(&other.to_data())
784    }
785}
786
787impl<T: Array> PartialEq<T> for dyn Array + '_ {
788    fn eq(&self, other: &T) -> bool {
789        self.to_data().eq(&other.to_data())
790    }
791}
792
793impl PartialEq for NullArray {
794    fn eq(&self, other: &NullArray) -> bool {
795        self.to_data().eq(&other.to_data())
796    }
797}
798
799impl<T: ArrowPrimitiveType> PartialEq for PrimitiveArray<T> {
800    fn eq(&self, other: &PrimitiveArray<T>) -> bool {
801        self.to_data().eq(&other.to_data())
802    }
803}
804
805impl<K: ArrowDictionaryKeyType> PartialEq for DictionaryArray<K> {
806    fn eq(&self, other: &Self) -> bool {
807        self.to_data().eq(&other.to_data())
808    }
809}
810
811impl PartialEq for BooleanArray {
812    fn eq(&self, other: &BooleanArray) -> bool {
813        self.to_data().eq(&other.to_data())
814    }
815}
816
817impl<OffsetSize: OffsetSizeTrait> PartialEq for GenericStringArray<OffsetSize> {
818    fn eq(&self, other: &Self) -> bool {
819        self.to_data().eq(&other.to_data())
820    }
821}
822
823impl<OffsetSize: OffsetSizeTrait> PartialEq for GenericBinaryArray<OffsetSize> {
824    fn eq(&self, other: &Self) -> bool {
825        self.to_data().eq(&other.to_data())
826    }
827}
828
829impl PartialEq for FixedSizeBinaryArray {
830    fn eq(&self, other: &Self) -> bool {
831        self.to_data().eq(&other.to_data())
832    }
833}
834
835impl<OffsetSize: OffsetSizeTrait> PartialEq for GenericListArray<OffsetSize> {
836    fn eq(&self, other: &Self) -> bool {
837        self.to_data().eq(&other.to_data())
838    }
839}
840
841impl<OffsetSize: OffsetSizeTrait> PartialEq for GenericListViewArray<OffsetSize> {
842    fn eq(&self, other: &Self) -> bool {
843        self.to_data().eq(&other.to_data())
844    }
845}
846
847impl PartialEq for MapArray {
848    fn eq(&self, other: &Self) -> bool {
849        self.to_data().eq(&other.to_data())
850    }
851}
852
853impl PartialEq for FixedSizeListArray {
854    fn eq(&self, other: &Self) -> bool {
855        self.to_data().eq(&other.to_data())
856    }
857}
858
859impl PartialEq for StructArray {
860    fn eq(&self, other: &Self) -> bool {
861        self.to_data().eq(&other.to_data())
862    }
863}
864
865impl<T: ByteViewType + ?Sized> PartialEq for GenericByteViewArray<T> {
866    fn eq(&self, other: &Self) -> bool {
867        self.to_data().eq(&other.to_data())
868    }
869}
870
871impl<R: RunEndIndexType> PartialEq for RunArray<R> {
872    fn eq(&self, other: &Self) -> bool {
873        self.to_data().eq(&other.to_data())
874    }
875}
876
877/// Constructs an [`ArrayRef`] from an [`ArrayData`].
878///
879/// # Notes:
880///
881/// It is more efficient to directly construct the concrete array type rather
882/// than using this function as creating an `ArrayData` requires at least one
883/// additional allocation (the Vec of buffers).
884///
885/// # Example:
886/// ```
887/// # use std::sync::Arc;
888/// # use arrow_data::ArrayData;
889/// # use arrow_array::{make_array, ArrayRef, Int32Array};
890/// # use arrow_buffer::{Buffer, ScalarBuffer};
891/// # use arrow_schema::DataType;
892/// // Create an Int32Array with values [1, 2, 3]
893/// let values_buffer = Buffer::from_slice_ref(&[1, 2, 3]);
894/// // ArrayData can be constructed using ArrayDataBuilder
895///  let builder = ArrayData::builder(DataType::Int32)
896///    .len(3)
897///    .add_buffer(values_buffer.clone());
898/// let array_data = builder.build().unwrap();
899/// // Create the ArrayRef from the ArrayData
900/// let array = make_array(array_data);
901///
902/// // It is equivalent to directly constructing the Int32Array
903/// let scalar_buffer = ScalarBuffer::from(values_buffer);
904/// let int32_array: ArrayRef = Arc::new(Int32Array::new(scalar_buffer, None));
905/// assert_eq!(&array, &int32_array);
906/// ```
907pub fn make_array(data: ArrayData) -> ArrayRef {
908    match data.data_type() {
909        DataType::Boolean => Arc::new(BooleanArray::from(data)) as ArrayRef,
910        DataType::Int8 => Arc::new(Int8Array::from(data)) as ArrayRef,
911        DataType::Int16 => Arc::new(Int16Array::from(data)) as ArrayRef,
912        DataType::Int32 => Arc::new(Int32Array::from(data)) as ArrayRef,
913        DataType::Int64 => Arc::new(Int64Array::from(data)) as ArrayRef,
914        DataType::UInt8 => Arc::new(UInt8Array::from(data)) as ArrayRef,
915        DataType::UInt16 => Arc::new(UInt16Array::from(data)) as ArrayRef,
916        DataType::UInt32 => Arc::new(UInt32Array::from(data)) as ArrayRef,
917        DataType::UInt64 => Arc::new(UInt64Array::from(data)) as ArrayRef,
918        DataType::Float16 => Arc::new(Float16Array::from(data)) as ArrayRef,
919        DataType::Float32 => Arc::new(Float32Array::from(data)) as ArrayRef,
920        DataType::Float64 => Arc::new(Float64Array::from(data)) as ArrayRef,
921        DataType::Date32 => Arc::new(Date32Array::from(data)) as ArrayRef,
922        DataType::Date64 => Arc::new(Date64Array::from(data)) as ArrayRef,
923        DataType::Time32(TimeUnit::Second) => Arc::new(Time32SecondArray::from(data)) as ArrayRef,
924        DataType::Time32(TimeUnit::Millisecond) => {
925            Arc::new(Time32MillisecondArray::from(data)) as ArrayRef
926        }
927        DataType::Time64(TimeUnit::Microsecond) => {
928            Arc::new(Time64MicrosecondArray::from(data)) as ArrayRef
929        }
930        DataType::Time64(TimeUnit::Nanosecond) => {
931            Arc::new(Time64NanosecondArray::from(data)) as ArrayRef
932        }
933        DataType::Timestamp(TimeUnit::Second, _) => {
934            Arc::new(TimestampSecondArray::from(data)) as ArrayRef
935        }
936        DataType::Timestamp(TimeUnit::Millisecond, _) => {
937            Arc::new(TimestampMillisecondArray::from(data)) as ArrayRef
938        }
939        DataType::Timestamp(TimeUnit::Microsecond, _) => {
940            Arc::new(TimestampMicrosecondArray::from(data)) as ArrayRef
941        }
942        DataType::Timestamp(TimeUnit::Nanosecond, _) => {
943            Arc::new(TimestampNanosecondArray::from(data)) as ArrayRef
944        }
945        DataType::Interval(IntervalUnit::YearMonth) => {
946            Arc::new(IntervalYearMonthArray::from(data)) as ArrayRef
947        }
948        DataType::Interval(IntervalUnit::DayTime) => {
949            Arc::new(IntervalDayTimeArray::from(data)) as ArrayRef
950        }
951        DataType::Interval(IntervalUnit::MonthDayNano) => {
952            Arc::new(IntervalMonthDayNanoArray::from(data)) as ArrayRef
953        }
954        DataType::Duration(TimeUnit::Second) => {
955            Arc::new(DurationSecondArray::from(data)) as ArrayRef
956        }
957        DataType::Duration(TimeUnit::Millisecond) => {
958            Arc::new(DurationMillisecondArray::from(data)) as ArrayRef
959        }
960        DataType::Duration(TimeUnit::Microsecond) => {
961            Arc::new(DurationMicrosecondArray::from(data)) as ArrayRef
962        }
963        DataType::Duration(TimeUnit::Nanosecond) => {
964            Arc::new(DurationNanosecondArray::from(data)) as ArrayRef
965        }
966        DataType::Binary => Arc::new(BinaryArray::from(data)) as ArrayRef,
967        DataType::LargeBinary => Arc::new(LargeBinaryArray::from(data)) as ArrayRef,
968        DataType::FixedSizeBinary(_) => Arc::new(FixedSizeBinaryArray::from(data)) as ArrayRef,
969        DataType::BinaryView => Arc::new(BinaryViewArray::from(data)) as ArrayRef,
970        DataType::Utf8 => Arc::new(StringArray::from(data)) as ArrayRef,
971        DataType::LargeUtf8 => Arc::new(LargeStringArray::from(data)) as ArrayRef,
972        DataType::Utf8View => Arc::new(StringViewArray::from(data)) as ArrayRef,
973        DataType::List(_) => Arc::new(ListArray::from(data)) as ArrayRef,
974        DataType::LargeList(_) => Arc::new(LargeListArray::from(data)) as ArrayRef,
975        DataType::ListView(_) => Arc::new(ListViewArray::from(data)) as ArrayRef,
976        DataType::LargeListView(_) => Arc::new(LargeListViewArray::from(data)) as ArrayRef,
977        DataType::Struct(_) => Arc::new(StructArray::from(data)) as ArrayRef,
978        DataType::Map(_, _) => Arc::new(MapArray::from(data)) as ArrayRef,
979        DataType::Union(_, _) => Arc::new(UnionArray::from(data)) as ArrayRef,
980        DataType::FixedSizeList(_, _) => Arc::new(FixedSizeListArray::from(data)) as ArrayRef,
981        DataType::Dictionary(key_type, _) => match key_type.as_ref() {
982            DataType::Int8 => Arc::new(DictionaryArray::<Int8Type>::from(data)) as ArrayRef,
983            DataType::Int16 => Arc::new(DictionaryArray::<Int16Type>::from(data)) as ArrayRef,
984            DataType::Int32 => Arc::new(DictionaryArray::<Int32Type>::from(data)) as ArrayRef,
985            DataType::Int64 => Arc::new(DictionaryArray::<Int64Type>::from(data)) as ArrayRef,
986            DataType::UInt8 => Arc::new(DictionaryArray::<UInt8Type>::from(data)) as ArrayRef,
987            DataType::UInt16 => Arc::new(DictionaryArray::<UInt16Type>::from(data)) as ArrayRef,
988            DataType::UInt32 => Arc::new(DictionaryArray::<UInt32Type>::from(data)) as ArrayRef,
989            DataType::UInt64 => Arc::new(DictionaryArray::<UInt64Type>::from(data)) as ArrayRef,
990            dt => unimplemented!("Unexpected dictionary key type {dt}"),
991        },
992        DataType::RunEndEncoded(run_ends_type, _) => match run_ends_type.data_type() {
993            DataType::Int16 => Arc::new(RunArray::<Int16Type>::from(data)) as ArrayRef,
994            DataType::Int32 => Arc::new(RunArray::<Int32Type>::from(data)) as ArrayRef,
995            DataType::Int64 => Arc::new(RunArray::<Int64Type>::from(data)) as ArrayRef,
996            dt => unimplemented!("Unexpected data type for run_ends array {dt}"),
997        },
998        DataType::Null => Arc::new(NullArray::from(data)) as ArrayRef,
999        DataType::Decimal32(_, _) => Arc::new(Decimal32Array::from(data)) as ArrayRef,
1000        DataType::Decimal64(_, _) => Arc::new(Decimal64Array::from(data)) as ArrayRef,
1001        DataType::Decimal128(_, _) => Arc::new(Decimal128Array::from(data)) as ArrayRef,
1002        DataType::Decimal256(_, _) => Arc::new(Decimal256Array::from(data)) as ArrayRef,
1003        dt => unimplemented!("Unexpected data type {dt}"),
1004    }
1005}
1006
1007/// Creates a new empty array
1008///
1009/// ```
1010/// use std::sync::Arc;
1011/// use arrow_schema::DataType;
1012/// use arrow_array::{ArrayRef, Int32Array, new_empty_array};
1013///
1014/// let empty_array = new_empty_array(&DataType::Int32);
1015/// let array: ArrayRef = Arc::new(Int32Array::from(vec![] as Vec<i32>));
1016///
1017/// assert_eq!(&array, &empty_array);
1018/// ```
1019pub fn new_empty_array(data_type: &DataType) -> ArrayRef {
1020    let data = ArrayData::new_empty(data_type);
1021    make_array(data)
1022}
1023
1024/// Creates a new array of `data_type` of length `length` filled
1025/// entirely of `NULL` values
1026///
1027/// ```
1028/// use std::sync::Arc;
1029/// use arrow_schema::DataType;
1030/// use arrow_array::{ArrayRef, Int32Array, new_null_array};
1031///
1032/// let null_array = new_null_array(&DataType::Int32, 3);
1033/// let array: ArrayRef = Arc::new(Int32Array::from(vec![None, None, None]));
1034///
1035/// assert_eq!(&array, &null_array);
1036/// ```
1037pub fn new_null_array(data_type: &DataType, length: usize) -> ArrayRef {
1038    make_array(ArrayData::new_null(data_type, length))
1039}
1040
1041/// Helper function that creates an [`OffsetBuffer`] from a buffer and array offset/ length
1042///
1043/// # Safety
1044///
1045/// - buffer must contain valid arrow offsets ( [`OffsetBuffer`] ) for the
1046///   given length and offset.
1047unsafe fn get_offsets_from_buffer<O: ArrowNativeType>(
1048    buffer: Buffer,
1049    offset: usize,
1050    len: usize,
1051) -> OffsetBuffer<O> {
1052    if len == 0 && buffer.is_empty() {
1053        return OffsetBuffer::new_empty();
1054    }
1055
1056    let scalar_buffer = ScalarBuffer::new(buffer, offset, len + 1);
1057    // Safety:
1058    // Arguments were valid
1059    unsafe { OffsetBuffer::new_unchecked(scalar_buffer) }
1060}
1061
1062/// Helper function for printing potentially long arrays.
1063fn print_long_array<A, F>(array: &A, f: &mut std::fmt::Formatter, print_item: F) -> std::fmt::Result
1064where
1065    A: Array,
1066    F: Fn(&A, usize, &mut std::fmt::Formatter) -> std::fmt::Result,
1067{
1068    let head = std::cmp::min(10, array.len());
1069
1070    for i in 0..head {
1071        if array.is_null(i) {
1072            writeln!(f, "  null,")?;
1073        } else {
1074            write!(f, "  ")?;
1075            print_item(array, i, f)?;
1076            writeln!(f, ",")?;
1077        }
1078    }
1079    if array.len() > 10 {
1080        if array.len() > 20 {
1081            writeln!(f, "  ...{} elements...,", array.len() - 20)?;
1082        }
1083
1084        let tail = std::cmp::max(head, array.len() - 10);
1085
1086        for i in tail..array.len() {
1087            if array.is_null(i) {
1088                writeln!(f, "  null,")?;
1089            } else {
1090                write!(f, "  ")?;
1091                print_item(array, i, f)?;
1092                writeln!(f, ",")?;
1093            }
1094        }
1095    }
1096    Ok(())
1097}
1098
1099#[cfg(test)]
1100mod tests {
1101    use super::*;
1102    use crate::cast::{as_union_array, downcast_array};
1103    use crate::downcast_run_array;
1104    use arrow_buffer::MutableBuffer;
1105    use arrow_schema::{Field, Fields, UnionFields, UnionMode};
1106
1107    #[test]
1108    fn test_empty_primitive() {
1109        let array = new_empty_array(&DataType::Int32);
1110        let a = array.as_any().downcast_ref::<Int32Array>().unwrap();
1111        assert_eq!(a.len(), 0);
1112        let expected: &[i32] = &[];
1113        assert_eq!(a.values(), expected);
1114    }
1115
1116    #[test]
1117    fn test_empty_variable_sized() {
1118        let array = new_empty_array(&DataType::Utf8);
1119        let a = array.as_any().downcast_ref::<StringArray>().unwrap();
1120        assert_eq!(a.len(), 0);
1121        assert_eq!(a.value_offsets()[0], 0i32);
1122    }
1123
1124    #[test]
1125    fn test_empty_list_primitive() {
1126        let data_type = DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false)));
1127        let array = new_empty_array(&data_type);
1128        let a = array.as_any().downcast_ref::<ListArray>().unwrap();
1129        assert_eq!(a.len(), 0);
1130        assert_eq!(a.value_offsets()[0], 0i32);
1131    }
1132
1133    #[test]
1134    fn test_null_boolean() {
1135        let array = new_null_array(&DataType::Boolean, 9);
1136        let a = array.as_any().downcast_ref::<BooleanArray>().unwrap();
1137        assert_eq!(a.len(), 9);
1138        for i in 0..9 {
1139            assert!(a.is_null(i));
1140        }
1141    }
1142
1143    #[test]
1144    fn test_null_primitive() {
1145        let array = new_null_array(&DataType::Int32, 9);
1146        let a = array.as_any().downcast_ref::<Int32Array>().unwrap();
1147        assert_eq!(a.len(), 9);
1148        for i in 0..9 {
1149            assert!(a.is_null(i));
1150        }
1151    }
1152
1153    #[test]
1154    fn test_null_struct() {
1155        // It is possible to create a null struct containing a non-nullable child
1156        // see https://github.com/apache/arrow-rs/pull/3244 for details
1157        let struct_type = DataType::Struct(vec![Field::new("data", DataType::Int64, false)].into());
1158        let array = new_null_array(&struct_type, 9);
1159
1160        let a = array.as_any().downcast_ref::<StructArray>().unwrap();
1161        assert_eq!(a.len(), 9);
1162        assert_eq!(a.column(0).len(), 9);
1163        for i in 0..9 {
1164            assert!(a.is_null(i));
1165        }
1166
1167        // Make sure we can slice the resulting array.
1168        a.slice(0, 5);
1169    }
1170
1171    #[test]
1172    fn test_null_variable_sized() {
1173        let array = new_null_array(&DataType::Utf8, 9);
1174        let a = array.as_any().downcast_ref::<StringArray>().unwrap();
1175        assert_eq!(a.len(), 9);
1176        assert_eq!(a.value_offsets()[9], 0i32);
1177        for i in 0..9 {
1178            assert!(a.is_null(i));
1179        }
1180    }
1181
1182    #[test]
1183    fn test_null_list_primitive() {
1184        let data_type = DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true)));
1185        let array = new_null_array(&data_type, 9);
1186        let a = array.as_any().downcast_ref::<ListArray>().unwrap();
1187        assert_eq!(a.len(), 9);
1188        assert_eq!(a.value_offsets()[9], 0i32);
1189        for i in 0..9 {
1190            assert!(a.is_null(i));
1191        }
1192    }
1193
1194    #[test]
1195    fn test_null_map() {
1196        let data_type = DataType::Map(
1197            Arc::new(Field::new(
1198                Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
1199                DataType::Struct(Fields::from(vec![
1200                    Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
1201                    Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Int32, true),
1202                ])),
1203                false,
1204            )),
1205            false,
1206        );
1207        let array = new_null_array(&data_type, 9);
1208        let a = array.as_any().downcast_ref::<MapArray>().unwrap();
1209        assert_eq!(a.len(), 9);
1210        assert_eq!(a.value_offsets()[9], 0i32);
1211        for i in 0..9 {
1212            assert!(a.is_null(i));
1213        }
1214    }
1215
1216    #[test]
1217    fn test_null_dictionary() {
1218        let values =
1219            vec![None, None, None, None, None, None, None, None, None] as Vec<Option<&str>>;
1220
1221        let array: DictionaryArray<Int8Type> = values.into_iter().collect();
1222        let array = Arc::new(array) as ArrayRef;
1223
1224        let null_array = new_null_array(array.data_type(), 9);
1225        assert_eq!(&array, &null_array);
1226        assert_eq!(
1227            array.to_data().buffers()[0].len(),
1228            null_array.to_data().buffers()[0].len()
1229        );
1230    }
1231
1232    #[test]
1233    fn test_null_union() {
1234        for mode in [UnionMode::Sparse, UnionMode::Dense] {
1235            let data_type = DataType::Union(
1236                UnionFields::try_new(
1237                    vec![2, 1],
1238                    vec![
1239                        Field::new("foo", DataType::Int32, true),
1240                        Field::new("bar", DataType::Int64, true),
1241                    ],
1242                )
1243                .unwrap(),
1244                mode,
1245            );
1246            let array = new_null_array(&data_type, 4);
1247
1248            let array = as_union_array(array.as_ref());
1249            assert_eq!(array.len(), 4);
1250            assert_eq!(array.null_count(), 0);
1251            assert_eq!(array.logical_null_count(), 4);
1252
1253            for i in 0..4 {
1254                let a = array.value(i);
1255                assert_eq!(a.len(), 1);
1256                assert_eq!(a.null_count(), 1);
1257                assert_eq!(a.logical_null_count(), 1);
1258                assert!(a.is_null(0))
1259            }
1260
1261            array.to_data().validate_full().unwrap();
1262        }
1263    }
1264
1265    #[test]
1266    fn test_null_runs() {
1267        for r in [DataType::Int16, DataType::Int32, DataType::Int64] {
1268            let data_type = DataType::RunEndEncoded(
1269                Arc::new(Field::new("run_ends", r, false)),
1270                Arc::new(Field::new("values", DataType::Utf8, true)),
1271            );
1272
1273            let array = new_null_array(&data_type, 4);
1274            let array = array.as_ref();
1275
1276            downcast_run_array! {
1277                array => {
1278                    assert_eq!(array.len(), 4);
1279                    assert_eq!(array.null_count(), 0);
1280                    assert_eq!(array.logical_null_count(), 4);
1281                    assert_eq!(array.values().len(), 1);
1282                    assert_eq!(array.values().null_count(), 1);
1283                    assert_eq!(array.run_ends().len(), 4);
1284                    assert_eq!(array.run_ends().values(), &[4]);
1285
1286                    let idx = array.get_physical_indices(&[0, 1, 2, 3]).unwrap();
1287                    assert_eq!(idx, &[0,0,0,0]);
1288                }
1289                d => unreachable!("{d}")
1290            }
1291        }
1292    }
1293
1294    #[test]
1295    fn test_null_fixed_size_binary() {
1296        for size in [1, 2, 7] {
1297            let array = new_null_array(&DataType::FixedSizeBinary(size), 6);
1298            let array = array
1299                .as_ref()
1300                .as_any()
1301                .downcast_ref::<FixedSizeBinaryArray>()
1302                .unwrap();
1303
1304            assert_eq!(array.len(), 6);
1305            assert_eq!(array.null_count(), 6);
1306            assert_eq!(array.logical_null_count(), 6);
1307            array.iter().for_each(|x| assert!(x.is_none()));
1308        }
1309    }
1310
1311    #[test]
1312    fn test_memory_size_null() {
1313        let null_arr = NullArray::new(32);
1314
1315        assert_eq!(0, null_arr.get_buffer_memory_size());
1316        assert_eq!(
1317            std::mem::size_of::<usize>(),
1318            null_arr.get_array_memory_size()
1319        );
1320    }
1321
1322    #[test]
1323    fn test_memory_size_primitive() {
1324        let arr = PrimitiveArray::<Int64Type>::from_iter_values(0..128);
1325        let empty = PrimitiveArray::<Int64Type>::from(ArrayData::new_empty(arr.data_type()));
1326
1327        // subtract empty array to avoid magic numbers for the size of additional fields
1328        assert_eq!(
1329            arr.get_array_memory_size() - empty.get_array_memory_size(),
1330            128 * std::mem::size_of::<i64>()
1331        );
1332    }
1333
1334    #[test]
1335    fn test_memory_size_primitive_sliced() {
1336        let arr = PrimitiveArray::<Int64Type>::from_iter_values(0..128);
1337        let slice1 = arr.slice(0, 64);
1338        let slice2 = arr.slice(64, 64);
1339
1340        // both slices report the full buffer memory usage, even though the buffers are shared
1341        assert_eq!(slice1.get_array_memory_size(), arr.get_array_memory_size());
1342        assert_eq!(slice2.get_array_memory_size(), arr.get_array_memory_size());
1343    }
1344
1345    #[test]
1346    fn test_memory_size_primitive_nullable() {
1347        let arr: PrimitiveArray<Int64Type> = (0..128)
1348            .map(|i| if i % 20 == 0 { Some(i) } else { None })
1349            .collect();
1350        let empty_with_bitmap = PrimitiveArray::<Int64Type>::from(
1351            ArrayData::builder(arr.data_type().clone())
1352                .add_buffer(MutableBuffer::new(0).into())
1353                .null_bit_buffer(Some(MutableBuffer::new_null(0).into()))
1354                .build()
1355                .unwrap(),
1356        );
1357
1358        // expected size is the size of the PrimitiveArray struct,
1359        // which includes the optional validity buffer
1360        // plus one buffer on the heap
1361        assert_eq!(
1362            std::mem::size_of::<PrimitiveArray<Int64Type>>(),
1363            empty_with_bitmap.get_array_memory_size()
1364        );
1365
1366        // subtract empty array to avoid magic numbers for the size of additional fields
1367        // the size of the validity bitmap is rounded up to 64 bytes
1368        assert_eq!(
1369            arr.get_array_memory_size() - empty_with_bitmap.get_array_memory_size(),
1370            128 * std::mem::size_of::<i64>() + 64
1371        );
1372    }
1373
1374    #[test]
1375    fn test_memory_size_dictionary() {
1376        let values = PrimitiveArray::<Int64Type>::from_iter_values(0..16);
1377        let keys = PrimitiveArray::<Int16Type>::from_iter_values(
1378            (0..256).map(|i| (i % values.len()) as i16),
1379        );
1380
1381        let dict_data_type = DataType::Dictionary(
1382            Box::new(keys.data_type().clone()),
1383            Box::new(values.data_type().clone()),
1384        );
1385        let dict_data = keys
1386            .into_data()
1387            .into_builder()
1388            .data_type(dict_data_type)
1389            .child_data(vec![values.into_data()])
1390            .build()
1391            .unwrap();
1392
1393        let empty_data = ArrayData::new_empty(&DataType::Dictionary(
1394            Box::new(DataType::Int16),
1395            Box::new(DataType::Int64),
1396        ));
1397
1398        let arr = DictionaryArray::<Int16Type>::from(dict_data);
1399        let empty = DictionaryArray::<Int16Type>::from(empty_data);
1400
1401        let expected_keys_size = 256 * std::mem::size_of::<i16>();
1402        assert_eq!(
1403            arr.keys().get_array_memory_size() - empty.keys().get_array_memory_size(),
1404            expected_keys_size
1405        );
1406
1407        let expected_values_size = 16 * std::mem::size_of::<i64>();
1408        assert_eq!(
1409            arr.values().get_array_memory_size() - empty.values().get_array_memory_size(),
1410            expected_values_size
1411        );
1412
1413        let expected_size = expected_keys_size + expected_values_size;
1414        assert_eq!(
1415            arr.get_array_memory_size() - empty.get_array_memory_size(),
1416            expected_size
1417        );
1418    }
1419
1420    /// Test function that takes an &dyn Array
1421    fn compute_my_thing(arr: &dyn Array) -> bool {
1422        !arr.is_empty()
1423    }
1424
1425    #[test]
1426    fn test_array_ref_as_array() {
1427        let arr: Int32Array = vec![1, 2, 3].into_iter().map(Some).collect();
1428
1429        // works well!
1430        assert!(compute_my_thing(&arr));
1431
1432        // Should also work when wrapped as an ArrayRef
1433        let arr: ArrayRef = Arc::new(arr);
1434        assert!(compute_my_thing(&arr));
1435        assert!(compute_my_thing(arr.as_ref()));
1436    }
1437
1438    #[test]
1439    fn test_downcast_array() {
1440        let array: Int32Array = vec![1, 2, 3].into_iter().map(Some).collect();
1441
1442        let boxed: ArrayRef = Arc::new(array);
1443        let array: Int32Array = downcast_array(&boxed);
1444
1445        let expected: Int32Array = vec![1, 2, 3].into_iter().map(Some).collect();
1446        assert_eq!(array, expected);
1447    }
1448}
arrow_array/array/mod.rs

arrow_array/array/
mod.rs
