arrow_array/array/
mod.rs

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

arrow_array/array/
mod.rs
