arrow_array/
cast.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//! Defines helper functions for downcasting [`dyn Array`](Array) to concrete types
19
20use crate::array::*;
21use crate::types::*;
22use arrow_data::ArrayData;
23
24/// Re-export symbols needed for downcast macros
25///
26/// Name follows `serde` convention
27#[doc(hidden)]
28pub mod __private {
29    pub use arrow_schema::{DataType, IntervalUnit, TimeUnit};
30}
31
32/// Repeats the provided pattern based on the number of comma separated identifiers
33#[doc(hidden)]
34#[macro_export]
35macro_rules! repeat_pat {
36    ($e:pat, $v_:expr) => {
37        $e
38    };
39    ($e:pat, $v_:expr $(, $tail:expr)+) => {
40        ($e, $crate::repeat_pat!($e $(, $tail)+))
41    }
42}
43
44/// Given one or more expressions evaluating to an integer [`DataType`] invokes the provided macro
45/// `m` with the corresponding integer [`ArrowPrimitiveType`], followed by any additional arguments
46///
47/// ```
48/// # use arrow_array::{downcast_primitive, ArrowPrimitiveType, downcast_integer};
49/// # use arrow_schema::DataType;
50///
51/// macro_rules! dictionary_key_size_helper {
52///   ($t:ty, $o:ty) => {
53///       std::mem::size_of::<<$t as ArrowPrimitiveType>::Native>() as $o
54///   };
55/// }
56///
57/// fn dictionary_key_size(t: &DataType) -> u8 {
58///     match t {
59///         DataType::Dictionary(k, _) => downcast_integer! {
60///             k.as_ref() => (dictionary_key_size_helper, u8),
61///             _ => unreachable!(),
62///         },
63///         // You can also add a guard to the pattern
64///         DataType::LargeUtf8 if true => u8::MAX,
65///         _ => u8::MAX,
66///     }
67/// }
68///
69/// assert_eq!(dictionary_key_size(&DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8))), 4);
70/// assert_eq!(dictionary_key_size(&DataType::Dictionary(Box::new(DataType::Int64), Box::new(DataType::Utf8))), 8);
71/// assert_eq!(dictionary_key_size(&DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8))), 2);
72/// ```
73///
74/// [`DataType`]: arrow_schema::DataType
75#[macro_export]
76macro_rules! downcast_integer {
77    ($($data_type:expr),+ => ($m:path $(, $args:tt)*), $($p:pat $(if $pred:expr)? => $fallback:expr $(,)?)*) => {
78        match ($($data_type),+) {
79            $crate::repeat_pat!($crate::cast::__private::DataType::Int8, $($data_type),+) => {
80                $m!($crate::types::Int8Type $(, $args)*)
81            }
82            $crate::repeat_pat!($crate::cast::__private::DataType::Int16, $($data_type),+) => {
83                $m!($crate::types::Int16Type $(, $args)*)
84            }
85            $crate::repeat_pat!($crate::cast::__private::DataType::Int32, $($data_type),+) => {
86                $m!($crate::types::Int32Type $(, $args)*)
87            }
88            $crate::repeat_pat!($crate::cast::__private::DataType::Int64, $($data_type),+) => {
89                $m!($crate::types::Int64Type $(, $args)*)
90            }
91            $crate::repeat_pat!($crate::cast::__private::DataType::UInt8, $($data_type),+) => {
92                $m!($crate::types::UInt8Type $(, $args)*)
93            }
94            $crate::repeat_pat!($crate::cast::__private::DataType::UInt16, $($data_type),+) => {
95                $m!($crate::types::UInt16Type $(, $args)*)
96            }
97            $crate::repeat_pat!($crate::cast::__private::DataType::UInt32, $($data_type),+) => {
98                $m!($crate::types::UInt32Type $(, $args)*)
99            }
100            $crate::repeat_pat!($crate::cast::__private::DataType::UInt64, $($data_type),+) => {
101                $m!($crate::types::UInt64Type $(, $args)*)
102            }
103            $($p $(if $pred)? => $fallback,)*
104        }
105    };
106}
107
108/// Given one or more expressions evaluating to an integer [`PrimitiveArray`] invokes the provided macro
109/// with the corresponding array, along with match statements for any non integer array types
110///
111/// ```
112/// # use arrow_array::{Array, downcast_integer_array, cast::as_string_array, cast::as_largestring_array};
113/// # use arrow_schema::DataType;
114///
115/// fn print_integer(array: &dyn Array) {
116///     downcast_integer_array!(
117///         array => {
118///             for v in array {
119///                 println!("{:?}", v);
120///             }
121///         }
122///         DataType::Utf8 => {
123///             for v in as_string_array(array) {
124///                 println!("{:?}", v);
125///             }
126///         }
127///         // You can also add a guard to the pattern
128///         DataType::LargeUtf8 if true => {
129///             for v in as_largestring_array(array) {
130///                 println!("{:?}", v);
131///             }
132///         }
133///         t => println!("Unsupported datatype {}", t)
134///     )
135/// }
136/// ```
137///
138/// [`DataType`]: arrow_schema::DataType
139#[macro_export]
140macro_rules! downcast_integer_array {
141    ($($values:ident),+ => $e:block $($p:pat $(if $pred:expr)? => $fallback:expr $(,)?)*) => {
142        $crate::downcast_integer!{
143            $($values.data_type()),+ => ($crate::downcast_primitive_array_helper, $($values),+, $e),
144            $($p $(if $pred)? => $fallback,)*
145        }
146    };
147    // Turn $e into a block.
148    ($values:ident => $e:expr, $($p:pat $(if $pred:expr)? => $fallback:expr $(,)?)*) => {
149        $crate::downcast_integer_array!($values => {$e} $($p $(if $pred)? => $fallback,)*)
150    };
151    // Remove $values parentheses.
152    (($($values:ident),+) => $e:block $($p:pat $(if $pred:expr)? => $fallback:expr $(,)?)*) => {
153        $crate::downcast_integer_array!($($values),+ => $e $($p $(if $pred)? => $fallback,)*)
154    };
155    // Turn $e into a block & remove $values parentheses.
156    (($($values:ident),+) => $e:expr, $($p:pat $(if $pred:expr)? => $fallback:expr $(,)?)*) => {
157        $crate::downcast_integer_array!($($values),+ => {$e} $($p $(if $pred)? => $fallback,)*)
158    };
159}
160
161/// Given one or more expressions evaluating to an integer [`DataType`] invokes the provided macro
162/// `m` with the corresponding integer [`RunEndIndexType`], followed by any additional arguments
163///
164/// ```
165/// # use std::sync::Arc;
166/// # use arrow_array::{downcast_primitive, ArrowPrimitiveType, downcast_run_end_index};
167/// # use arrow_schema::{DataType, Field};
168///
169/// macro_rules! run_end_size_helper {
170///   ($t:ty, $o:ty) => {
171///       std::mem::size_of::<<$t as ArrowPrimitiveType>::Native>() as $o
172///   };
173/// }
174///
175/// fn run_end_index_size(t: &DataType) -> u8 {
176///     match t {
177///         DataType::RunEndEncoded(k, _) => downcast_run_end_index! {
178///             k.data_type() => (run_end_size_helper, u8),
179///             _ => unreachable!(),
180///         },
181///         // You can also add a guard to the pattern
182///         DataType::LargeUtf8 if true => u8::MAX,
183///         _ => u8::MAX,
184///     }
185/// }
186///
187/// assert_eq!(run_end_index_size(&DataType::RunEndEncoded(Arc::new(Field::new("a", DataType::Int32, false)), Arc::new(Field::new("b", DataType::Utf8, true)))), 4);
188/// assert_eq!(run_end_index_size(&DataType::RunEndEncoded(Arc::new(Field::new("a", DataType::Int64, false)), Arc::new(Field::new("b", DataType::Utf8, true)))), 8);
189/// assert_eq!(run_end_index_size(&DataType::RunEndEncoded(Arc::new(Field::new("a", DataType::Int16, false)), Arc::new(Field::new("b", DataType::Utf8, true)))), 2);
190/// ```
191///
192/// [`DataType`]: arrow_schema::DataType
193#[macro_export]
194macro_rules! downcast_run_end_index {
195    ($($data_type:expr),+ => ($m:path $(, $args:tt)*), $($p:pat $(if $pred:expr)? => $fallback:expr $(,)?)*) => {
196        match ($($data_type),+) {
197            $crate::repeat_pat!($crate::cast::__private::DataType::Int16, $($data_type),+) => {
198                $m!($crate::types::Int16Type $(, $args)*)
199            }
200            $crate::repeat_pat!($crate::cast::__private::DataType::Int32, $($data_type),+) => {
201                $m!($crate::types::Int32Type $(, $args)*)
202            }
203            $crate::repeat_pat!($crate::cast::__private::DataType::Int64, $($data_type),+) => {
204                $m!($crate::types::Int64Type $(, $args)*)
205            }
206            $($p $(if $pred)? => $fallback,)*
207        }
208    };
209}
210
211/// Given one or more expressions evaluating to primitive [`DataType`] invokes the provided macro
212/// `m` with the corresponding [`ArrowPrimitiveType`], followed by any additional arguments
213///
214/// ```
215/// # use arrow_array::{downcast_temporal, ArrowPrimitiveType};
216/// # use arrow_schema::DataType;
217///
218/// macro_rules! temporal_size_helper {
219///   ($t:ty, $o:ty) => {
220///       std::mem::size_of::<<$t as ArrowPrimitiveType>::Native>() as $o
221///   };
222/// }
223///
224/// fn temporal_size(t: &DataType) -> u8 {
225///     downcast_temporal! {
226///         t => (temporal_size_helper, u8),
227///         // You can also add a guard to the pattern
228///         DataType::LargeUtf8 if true => u8::MAX,
229///         _ => u8::MAX
230///     }
231/// }
232///
233/// assert_eq!(temporal_size(&DataType::Date32), 4);
234/// assert_eq!(temporal_size(&DataType::Date64), 8);
235/// ```
236///
237/// [`DataType`]: arrow_schema::DataType
238#[macro_export]
239macro_rules! downcast_temporal {
240    ($($data_type:expr),+ => ($m:path $(, $args:tt)*), $($p:pat $(if $pred:expr)? => $fallback:expr $(,)?)*) => {
241        match ($($data_type),+) {
242            $crate::repeat_pat!($crate::cast::__private::DataType::Time32($crate::cast::__private::TimeUnit::Second), $($data_type),+) => {
243                $m!($crate::types::Time32SecondType $(, $args)*)
244            }
245            $crate::repeat_pat!($crate::cast::__private::DataType::Time32($crate::cast::__private::TimeUnit::Millisecond), $($data_type),+) => {
246                $m!($crate::types::Time32MillisecondType $(, $args)*)
247            }
248            $crate::repeat_pat!($crate::cast::__private::DataType::Time64($crate::cast::__private::TimeUnit::Microsecond), $($data_type),+) => {
249                $m!($crate::types::Time64MicrosecondType $(, $args)*)
250            }
251            $crate::repeat_pat!($crate::cast::__private::DataType::Time64($crate::cast::__private::TimeUnit::Nanosecond), $($data_type),+) => {
252                $m!($crate::types::Time64NanosecondType $(, $args)*)
253            }
254            $crate::repeat_pat!($crate::cast::__private::DataType::Date32, $($data_type),+) => {
255                $m!($crate::types::Date32Type $(, $args)*)
256            }
257            $crate::repeat_pat!($crate::cast::__private::DataType::Date64, $($data_type),+) => {
258                $m!($crate::types::Date64Type $(, $args)*)
259            }
260            $crate::repeat_pat!($crate::cast::__private::DataType::Timestamp($crate::cast::__private::TimeUnit::Second, _), $($data_type),+) => {
261                $m!($crate::types::TimestampSecondType $(, $args)*)
262            }
263            $crate::repeat_pat!($crate::cast::__private::DataType::Timestamp($crate::cast::__private::TimeUnit::Millisecond, _), $($data_type),+) => {
264                $m!($crate::types::TimestampMillisecondType $(, $args)*)
265            }
266            $crate::repeat_pat!($crate::cast::__private::DataType::Timestamp($crate::cast::__private::TimeUnit::Microsecond, _), $($data_type),+) => {
267                $m!($crate::types::TimestampMicrosecondType $(, $args)*)
268            }
269            $crate::repeat_pat!($crate::cast::__private::DataType::Timestamp($crate::cast::__private::TimeUnit::Nanosecond, _), $($data_type),+) => {
270                $m!($crate::types::TimestampNanosecondType $(, $args)*)
271            }
272            $($p $(if $pred)? => $fallback,)*
273        }
274    };
275}
276
277/// Downcast an [`Array`] to a temporal [`PrimitiveArray`] based on its [`DataType`]
278/// accepts a number of subsequent patterns to match the data type
279///
280/// ```
281/// # use arrow_array::{Array, downcast_temporal_array, cast::as_string_array, cast::as_largestring_array};
282/// # use arrow_schema::DataType;
283///
284/// fn print_temporal(array: &dyn Array) {
285///     downcast_temporal_array!(
286///         array => {
287///             for v in array {
288///                 println!("{:?}", v);
289///             }
290///         }
291///         DataType::Utf8 => {
292///             for v in as_string_array(array) {
293///                 println!("{:?}", v);
294///             }
295///         }
296///         // You can also add a guard to the pattern
297///         DataType::LargeUtf8 if true => {
298///             for v in as_largestring_array(array) {
299///                 println!("{:?}", v);
300///             }
301///         }
302///         t => println!("Unsupported datatype {}", t)
303///     )
304/// }
305/// ```
306///
307/// [`DataType`]: arrow_schema::DataType
308#[macro_export]
309macro_rules! downcast_temporal_array {
310    ($($values:ident),+ => $e:block $($p:pat $(if $pred:expr)? => $fallback:expr $(,)?)*) => {
311        $crate::downcast_temporal!{
312            $($values.data_type()),+ => ($crate::downcast_primitive_array_helper, $($values),+, $e),
313            $($p $(if $pred)? => $fallback,)*
314        }
315    };
316    // Turn $e into a block.
317    ($values:ident => $e:expr, $($p:pat $(if $pred:expr)? => $fallback:expr $(,)?)*) => {
318        $crate::downcast_temporal_array!($values => {$e} $($p $(if $pred)? => $fallback,)*)
319    };
320    // Remove $values parentheses.
321    (($($values:ident),+) => $e:block $($p:pat $(if $pred:expr)? => $fallback:expr $(,)?)*) => {
322        $crate::downcast_temporal_array!($($values),+ => $e $($p $(if $pred)? => $fallback,)*)
323    };
324    // Turn $e into a block & remove $values parentheses.
325    (($($values:ident),+) => $e:expr, $($p:pat $(if $pred:expr)? => $fallback:expr $(,)?)*) => {
326        $crate::downcast_temporal_array!($($values),+ => {$e} $($p $(if $pred)? => $fallback,)*)
327    };
328}
329
330/// Given one or more expressions evaluating to primitive [`DataType`] invokes the provided macro
331/// `m` with the corresponding [`ArrowPrimitiveType`], followed by any additional arguments
332///
333/// ```
334/// # use arrow_array::{downcast_primitive, ArrowPrimitiveType};
335/// # use arrow_schema::DataType;
336///
337/// macro_rules! primitive_size_helper {
338///   ($t:ty, $o:ty) => {
339///       std::mem::size_of::<<$t as ArrowPrimitiveType>::Native>() as $o
340///   };
341/// }
342///
343/// fn primitive_size(t: &DataType) -> u8 {
344///     downcast_primitive! {
345///         t => (primitive_size_helper, u8),
346///         // You can also add a guard to the pattern
347///         DataType::LargeUtf8 if true => u8::MAX,
348///         _ => u8::MAX
349///     }
350/// }
351///
352/// assert_eq!(primitive_size(&DataType::Int32), 4);
353/// assert_eq!(primitive_size(&DataType::Int64), 8);
354/// assert_eq!(primitive_size(&DataType::Float16), 2);
355/// assert_eq!(primitive_size(&DataType::Decimal128(38, 10)), 16);
356/// assert_eq!(primitive_size(&DataType::Decimal256(76, 20)), 32);
357/// ```
358///
359/// [`DataType`]: arrow_schema::DataType
360#[macro_export]
361macro_rules! downcast_primitive {
362    ($($data_type:expr),+ => ($m:path $(, $args:tt)*), $($p:pat $(if $pred:expr)? => $fallback:expr $(,)?)*) => {
363        $crate::downcast_integer! {
364            $($data_type),+ => ($m $(, $args)*),
365            $crate::repeat_pat!($crate::cast::__private::DataType::Float16, $($data_type),+) => {
366                $m!($crate::types::Float16Type $(, $args)*)
367            }
368            $crate::repeat_pat!($crate::cast::__private::DataType::Float32, $($data_type),+) => {
369                $m!($crate::types::Float32Type $(, $args)*)
370            }
371            $crate::repeat_pat!($crate::cast::__private::DataType::Float64, $($data_type),+) => {
372                $m!($crate::types::Float64Type $(, $args)*)
373            }
374            $crate::repeat_pat!($crate::cast::__private::DataType::Decimal32(_, _), $($data_type),+) => {
375                $m!($crate::types::Decimal32Type $(, $args)*)
376            }
377            $crate::repeat_pat!($crate::cast::__private::DataType::Decimal64(_, _), $($data_type),+) => {
378                $m!($crate::types::Decimal64Type $(, $args)*)
379            }
380            $crate::repeat_pat!($crate::cast::__private::DataType::Decimal128(_, _), $($data_type),+) => {
381                $m!($crate::types::Decimal128Type $(, $args)*)
382            }
383            $crate::repeat_pat!($crate::cast::__private::DataType::Decimal256(_, _), $($data_type),+) => {
384                $m!($crate::types::Decimal256Type $(, $args)*)
385            }
386            $crate::repeat_pat!($crate::cast::__private::DataType::Interval($crate::cast::__private::IntervalUnit::YearMonth), $($data_type),+) => {
387                $m!($crate::types::IntervalYearMonthType $(, $args)*)
388            }
389            $crate::repeat_pat!($crate::cast::__private::DataType::Interval($crate::cast::__private::IntervalUnit::DayTime), $($data_type),+) => {
390                $m!($crate::types::IntervalDayTimeType $(, $args)*)
391            }
392            $crate::repeat_pat!($crate::cast::__private::DataType::Interval($crate::cast::__private::IntervalUnit::MonthDayNano), $($data_type),+) => {
393                $m!($crate::types::IntervalMonthDayNanoType $(, $args)*)
394            }
395            $crate::repeat_pat!($crate::cast::__private::DataType::Duration($crate::cast::__private::TimeUnit::Second), $($data_type),+) => {
396                $m!($crate::types::DurationSecondType $(, $args)*)
397            }
398            $crate::repeat_pat!($crate::cast::__private::DataType::Duration($crate::cast::__private::TimeUnit::Millisecond), $($data_type),+) => {
399                $m!($crate::types::DurationMillisecondType $(, $args)*)
400            }
401            $crate::repeat_pat!($crate::cast::__private::DataType::Duration($crate::cast::__private::TimeUnit::Microsecond), $($data_type),+) => {
402                $m!($crate::types::DurationMicrosecondType $(, $args)*)
403            }
404            $crate::repeat_pat!($crate::cast::__private::DataType::Duration($crate::cast::__private::TimeUnit::Nanosecond), $($data_type),+) => {
405                $m!($crate::types::DurationNanosecondType $(, $args)*)
406            }
407            _ => {
408                $crate::downcast_temporal! {
409                    $($data_type),+ => ($m $(, $args)*),
410                    $($p $(if $pred)? => $fallback,)*
411                }
412            }
413        }
414    };
415}
416
417#[macro_export]
418#[doc(hidden)]
419macro_rules! downcast_primitive_array_helper {
420    ($t:ty, $($values:ident),+, $e:block) => {{
421        $(let $values = $crate::cast::as_primitive_array::<$t>($values);)+
422        $e
423    }};
424}
425
426/// Downcast an [`Array`] to a [`PrimitiveArray`] based on its [`DataType`]
427/// accepts a number of subsequent patterns to match the data type
428///
429/// ```
430/// # use arrow_array::{Array, downcast_primitive_array, cast::as_string_array, cast::as_largestring_array};
431/// # use arrow_schema::DataType;
432///
433/// fn print_primitive(array: &dyn Array) {
434///     downcast_primitive_array!(
435///         array => {
436///             for v in array {
437///                 println!("{:?}", v);
438///             }
439///         }
440///         DataType::Utf8 => {
441///             for v in as_string_array(array) {
442///                 println!("{:?}", v);
443///             }
444///         }
445///         // You can also add a guard to the pattern
446///         DataType::LargeUtf8 if true => {
447///             for v in as_largestring_array(array) {
448///                 println!("{:?}", v);
449///             }
450///         }
451///         t => println!("Unsupported datatype {}", t)
452///     )
453/// }
454/// ```
455///
456/// [`DataType`]: arrow_schema::DataType
457#[macro_export]
458macro_rules! downcast_primitive_array {
459    ($($values:ident),+ => $e:block $($p:pat $(if $pred:expr)? => $fallback:expr $(,)?)*) => {
460        $crate::downcast_primitive!{
461            $($values.data_type()),+ => ($crate::downcast_primitive_array_helper, $($values),+, $e),
462            $($p $(if $pred)? => $fallback,)*
463        }
464    };
465    // Turn $e into a block.
466    ($values:ident => $e:expr, $($p:pat $(if $pred:expr)? => $fallback:expr $(,)?)*) => {
467        $crate::downcast_primitive_array!($values => {$e} $($p $(if $pred)? => $fallback,)*)
468    };
469    // Remove $values parentheses.
470    (($($values:ident),+) => $e:block $($p:pat $(if $pred:expr)? => $fallback:expr $(,)?)*) => {
471        $crate::downcast_primitive_array!($($values),+ => $e $($p $(if $pred)? => $fallback,)*)
472    };
473    // Turn $e into a block & remove $values parentheses.
474    (($($values:ident),+) => $e:expr, $($p:pat $(if $pred:expr)? => $fallback:expr $(,)?)*) => {
475        $crate::downcast_primitive_array!($($values),+ => {$e} $($p $(if $pred)? => $fallback,)*)
476    };
477}
478
479/// Force downcast of an [`Array`], such as an [`ArrayRef`], to
480/// [`PrimitiveArray<T>`].
481///
482/// # Example
483///
484/// ```
485/// # use std::sync::Arc;
486/// # use arrow_array::{ArrayRef, Int32Array};
487/// # use arrow_array::cast::as_primitive_array;
488/// # use arrow_array::types::Int32Type;
489///
490/// let arr: ArrayRef = Arc::new(Int32Array::from(vec![Some(1)]));
491///
492/// // Downcast an `ArrayRef` to Int32Array / PrimitiveArray<Int32>:
493/// let primitive_array: &Int32Array = as_primitive_array(&arr);
494///
495/// // Equivalently:
496/// let primitive_array = as_primitive_array::<Int32Type>(&arr);
497///
498/// // This is the equivalent of:
499/// let primitive_array = arr
500///     .as_any()
501///     .downcast_ref::<Int32Array>()
502///     .unwrap();
503/// ```
504///
505/// # Panics
506///
507/// Panics if `arr` is not a [`PrimitiveArray<T>`]
508pub fn as_primitive_array<T>(arr: &dyn Array) -> &PrimitiveArray<T>
509where
510    T: ArrowPrimitiveType,
511{
512    arr.as_any()
513        .downcast_ref::<PrimitiveArray<T>>()
514        .expect("Unable to downcast to primitive array")
515}
516
517#[macro_export]
518#[doc(hidden)]
519macro_rules! downcast_dictionary_array_helper {
520    ($t:ty, $($values:ident),+, $e:block) => {{
521        $(let $values = $crate::cast::as_dictionary_array::<$t>($values);)+
522        $e
523    }};
524}
525
526/// Downcast an [`Array`] to a [`DictionaryArray`] based on its [`DataType`], accepts
527/// a number of subsequent patterns to match the data type
528///
529/// ```
530/// # use arrow_array::{Array, StringArray, downcast_dictionary_array, cast::as_string_array, cast::as_largestring_array};
531/// # use arrow_schema::DataType;
532///
533/// fn print_strings(array: &dyn Array) {
534///     downcast_dictionary_array!(
535///         array => match array.values().data_type() {
536///             DataType::Utf8 => {
537///                 for v in array.downcast_dict::<StringArray>().unwrap() {
538///                     println!("{:?}", v);
539///                 }
540///             }
541///             t => println!("Unsupported dictionary value type {}", t),
542///         },
543///         DataType::Utf8 => {
544///             for v in as_string_array(array) {
545///                 println!("{:?}", v);
546///             }
547///         }
548///         // You can also add a guard to the pattern
549///         DataType::LargeUtf8 if true => {
550///             for v in as_largestring_array(array) {
551///                 println!("{:?}", v);
552///             }
553///         }
554///         t => println!("Unsupported datatype {}", t)
555///     )
556/// }
557/// ```
558///
559/// [`DataType`]: arrow_schema::DataType
560#[macro_export]
561macro_rules! downcast_dictionary_array {
562    ($values:ident => $e:expr, $($p:pat $(if $pred:expr)? => $fallback:expr $(,)?)*) => {
563        downcast_dictionary_array!($values => {$e} $($p $(if $pred)? => $fallback,)*)
564    };
565
566    ($values:ident => $e:block $($p:pat $(if $pred:expr)? => $fallback:expr $(,)?)*) => {
567        match $values.data_type() {
568            $crate::cast::__private::DataType::Dictionary(k, _) => {
569                $crate::downcast_integer! {
570                    k.as_ref() => ($crate::downcast_dictionary_array_helper, $values, $e),
571                    k => unreachable!("unsupported dictionary key type: {}", k)
572                }
573            }
574            $($p $(if $pred)? => $fallback,)*
575        }
576    }
577}
578
579/// Force downcast of an [`Array`], such as an [`ArrayRef`] to
580/// [`DictionaryArray<T>`].
581///
582/// # Example
583///
584/// ```
585/// # use arrow_array::{ArrayRef, DictionaryArray};
586/// # use arrow_array::cast::as_dictionary_array;
587/// # use arrow_array::types::Int32Type;
588///
589/// let arr: DictionaryArray<Int32Type> = vec![Some("foo")].into_iter().collect();
590/// let arr: ArrayRef = std::sync::Arc::new(arr);
591/// let dict_array: &DictionaryArray<Int32Type> = as_dictionary_array::<Int32Type>(&arr);
592/// ```
593///
594/// # Panics
595///
596/// Panics if `arr` is not a [`DictionaryArray<T>`]
597pub fn as_dictionary_array<T>(arr: &dyn Array) -> &DictionaryArray<T>
598where
599    T: ArrowDictionaryKeyType,
600{
601    arr.as_any()
602        .downcast_ref::<DictionaryArray<T>>()
603        .expect("Unable to downcast to dictionary array")
604}
605
606/// Force downcast of an [`Array`], such as an [`ArrayRef`] to
607/// [`RunArray<T>`].
608///
609/// # Example
610///
611/// ```
612/// # use arrow_array::{ArrayRef, RunArray};
613/// # use arrow_array::cast::as_run_array;
614/// # use arrow_array::types::Int32Type;
615///
616/// let arr: RunArray<Int32Type> = vec![Some("foo")].into_iter().collect();
617/// let arr: ArrayRef = std::sync::Arc::new(arr);
618/// let run_array: &RunArray<Int32Type> = as_run_array::<Int32Type>(&arr);
619/// ```
620///
621/// # Panics
622///
623/// Panics if `arr` is not a [`RunArray<T>`]
624pub fn as_run_array<T>(arr: &dyn Array) -> &RunArray<T>
625where
626    T: RunEndIndexType,
627{
628    arr.as_any()
629        .downcast_ref::<RunArray<T>>()
630        .expect("Unable to downcast to run array")
631}
632
633#[macro_export]
634#[doc(hidden)]
635macro_rules! downcast_run_array_helper {
636    ($t:ty, $($values:ident),+, $e:block) => {{
637        $(let $values = $crate::cast::as_run_array::<$t>($values);)+
638        $e
639    }};
640}
641
642/// Downcast an [`Array`] to a [`RunArray`] based on its [`DataType`], accepts
643/// a number of subsequent patterns to match the data type
644///
645/// ```
646/// # use arrow_array::{Array, StringArray, downcast_run_array, cast::as_string_array, cast::as_largestring_array};
647/// # use arrow_schema::DataType;
648///
649/// fn print_strings(array: &dyn Array) {
650///     downcast_run_array!(
651///         array => match array.values().data_type() {
652///             DataType::Utf8 => {
653///                 for v in array.downcast::<StringArray>().unwrap() {
654///                     println!("{:?}", v);
655///                 }
656///             }
657///             t => println!("Unsupported run array value type {}", t),
658///         },
659///         DataType::Utf8 => {
660///             for v in as_string_array(array) {
661///                 println!("{:?}", v);
662///             }
663///         }
664///         // You can also add a guard to the pattern
665///         DataType::LargeUtf8 if true => {
666///             for v in as_largestring_array(array) {
667///                 println!("{:?}", v);
668///             }
669///         }
670///         t => println!("Unsupported datatype {}", t)
671///     )
672/// }
673/// ```
674///
675/// [`DataType`]: arrow_schema::DataType
676#[macro_export]
677macro_rules! downcast_run_array {
678    ($values:ident => $e:expr, $($p:pat $(if $pred:expr)? => $fallback:expr $(,)?)*) => {
679        downcast_run_array!($values => {$e} $($p $(if $pred)? => $fallback,)*)
680    };
681
682    ($values:ident => $e:block $($p:pat $(if $pred:expr)? => $fallback:expr $(,)?)*) => {
683        match $values.data_type() {
684            $crate::cast::__private::DataType::RunEndEncoded(k, _) => {
685                $crate::downcast_run_end_index! {
686                    k.data_type() => ($crate::downcast_run_array_helper, $values, $e),
687                    k => unreachable!("unsupported run end index type: {}", k)
688                }
689            }
690            $($p $(if $pred)? => $fallback,)*
691        }
692    }
693}
694
695/// Force downcast of an [`Array`], such as an [`ArrayRef`] to
696/// [`GenericListArray<T>`].
697///
698/// # Panics
699///
700/// Panics if `arr` is not a [`GenericListArray<T>`]
701pub fn as_generic_list_array<S: OffsetSizeTrait>(arr: &dyn Array) -> &GenericListArray<S> {
702    arr.as_any()
703        .downcast_ref::<GenericListArray<S>>()
704        .expect("Unable to downcast to list array")
705}
706
707/// Force downcast of an [`Array`], such as an [`ArrayRef`] to
708/// [`ListArray`].
709///
710/// # Panics
711///
712/// Panics if `arr` is not a [`ListArray`]
713#[inline]
714pub fn as_list_array(arr: &dyn Array) -> &ListArray {
715    as_generic_list_array::<i32>(arr)
716}
717
718/// Force downcast of an [`Array`], such as an [`ArrayRef`] to
719/// [`FixedSizeListArray`].
720///
721/// # Panics
722///
723/// Panics if `arr` is not a [`FixedSizeListArray`]
724#[inline]
725pub fn as_fixed_size_list_array(arr: &dyn Array) -> &FixedSizeListArray {
726    arr.as_any()
727        .downcast_ref::<FixedSizeListArray>()
728        .expect("Unable to downcast to fixed size list array")
729}
730
731/// Force downcast of an [`Array`], such as an [`ArrayRef`] to
732/// [`LargeListArray`].
733///
734/// # Panics
735///
736/// Panics if `arr` is not a [`LargeListArray`]
737#[inline]
738pub fn as_large_list_array(arr: &dyn Array) -> &LargeListArray {
739    as_generic_list_array::<i64>(arr)
740}
741
742/// Force downcast of an [`Array`], such as an [`ArrayRef`] to
743/// [`GenericBinaryArray<S>`].
744///
745/// # Panics
746///
747/// Panics if `arr` is not a [`GenericBinaryArray<S>`]
748#[inline]
749pub fn as_generic_binary_array<S: OffsetSizeTrait>(arr: &dyn Array) -> &GenericBinaryArray<S> {
750    arr.as_any()
751        .downcast_ref::<GenericBinaryArray<S>>()
752        .expect("Unable to downcast to binary array")
753}
754
755/// Force downcast of an [`Array`], such as an [`ArrayRef`] to
756/// [`StringArray`].
757///
758/// # Example
759///
760/// ```
761/// # use std::sync::Arc;
762/// # use arrow_array::cast::as_string_array;
763/// # use arrow_array::{ArrayRef, StringArray};
764///
765/// let arr: ArrayRef = Arc::new(StringArray::from_iter(vec![Some("foo")]));
766/// let string_array = as_string_array(&arr);
767/// ```
768///
769/// # Panics
770///
771/// Panics if `arr` is not a [`StringArray`]
772pub fn as_string_array(arr: &dyn Array) -> &StringArray {
773    arr.as_any()
774        .downcast_ref::<StringArray>()
775        .expect("Unable to downcast to StringArray")
776}
777
778/// Force downcast of an [`Array`], such as an [`ArrayRef`] to
779/// [`BooleanArray`].
780///
781/// # Example
782///
783/// ```
784/// # use std::sync::Arc;
785/// # use arrow_array::{ArrayRef, BooleanArray};
786/// # use arrow_array::cast::as_boolean_array;
787///
788/// let arr: ArrayRef = Arc::new(BooleanArray::from_iter(vec![Some(true)]));
789/// let boolean_array = as_boolean_array(&arr);
790/// ```
791///
792/// # Panics
793///
794/// Panics if `arr` is not a [`BooleanArray`]
795pub fn as_boolean_array(arr: &dyn Array) -> &BooleanArray {
796    arr.as_any()
797        .downcast_ref::<BooleanArray>()
798        .expect("Unable to downcast to BooleanArray")
799}
800
801macro_rules! array_downcast_fn {
802    ($name: ident, $arrty: ty, $arrty_str:expr) => {
803        #[doc = "Force downcast of an [`Array`], such as an [`ArrayRef`] to "]
804        #[doc = $arrty_str]
805        #[doc = ""]
806        #[doc = "# Panics"]
807        #[doc = ""]
808        #[doc = "Panics if `arr` is not a "]
809        #[doc = $arrty_str]
810        pub fn $name(arr: &dyn Array) -> &$arrty {
811            arr.as_any().downcast_ref::<$arrty>().expect(concat!(
812                "Unable to downcast to typed array through ",
813                stringify!($name)
814            ))
815        }
816    };
817
818    // use recursive macro to generate dynamic doc string for a given array type
819    ($name: ident, $arrty: ty) => {
820        array_downcast_fn!($name, $arrty, concat!("[`", stringify!($arrty), "`]"));
821    };
822}
823
824array_downcast_fn!(as_largestring_array, LargeStringArray);
825array_downcast_fn!(as_null_array, NullArray);
826array_downcast_fn!(as_struct_array, StructArray);
827array_downcast_fn!(as_union_array, UnionArray);
828array_downcast_fn!(as_map_array, MapArray);
829
830/// Downcasts a `dyn Array` to a concrete type
831///
832/// ```
833/// # use arrow_array::{BooleanArray, Int32Array, RecordBatch, StringArray};
834/// # use arrow_array::cast::downcast_array;
835/// struct ConcreteBatch {
836///     col1: Int32Array,
837///     col2: BooleanArray,
838///     col3: StringArray,
839/// }
840///
841/// impl ConcreteBatch {
842///     fn new(batch: &RecordBatch) -> Self {
843///         Self {
844///             col1: downcast_array(batch.column(0).as_ref()),
845///             col2: downcast_array(batch.column(1).as_ref()),
846///             col3: downcast_array(batch.column(2).as_ref()),
847///         }
848///     }
849/// }
850/// ```
851///
852/// # Panics
853///
854/// Panics if array is not of the correct data type
855pub fn downcast_array<T>(array: &dyn Array) -> T
856where
857    T: From<ArrayData>,
858{
859    T::from(array.to_data())
860}
861
862mod private {
863    pub trait Sealed {}
864}
865
866/// An extension trait for `dyn Array` that provides ergonomic downcasting
867///
868/// ```
869/// # use std::sync::Arc;
870/// # use arrow_array::{ArrayRef, Int32Array};
871/// # use arrow_array::cast::AsArray;
872/// # use arrow_array::types::Int32Type;
873/// let col = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef;
874/// assert_eq!(col.as_primitive::<Int32Type>().values(), &[1, 2, 3]);
875/// ```
876pub trait AsArray: private::Sealed {
877    /// Downcast this to a [`BooleanArray`] returning `None` if not possible
878    fn as_boolean_opt(&self) -> Option<&BooleanArray>;
879
880    /// Downcast this to a [`BooleanArray`]
881    ///
882    /// # Panics
883    ///
884    /// Panics if this is not a [`BooleanArray`]
885    fn as_boolean(&self) -> &BooleanArray {
886        self.as_boolean_opt().expect("boolean array")
887    }
888
889    /// Downcast this to a [`PrimitiveArray`] returning `None` if not possible
890    fn as_primitive_opt<T: ArrowPrimitiveType>(&self) -> Option<&PrimitiveArray<T>>;
891
892    /// Downcast this to a [`PrimitiveArray`]
893    ///
894    /// # Panics
895    ///
896    /// Panics if this is not a [`PrimitiveArray`]
897    fn as_primitive<T: ArrowPrimitiveType>(&self) -> &PrimitiveArray<T> {
898        self.as_primitive_opt().expect("primitive array")
899    }
900
901    /// Downcast this to a [`GenericByteArray`] returning `None` if not possible
902    fn as_bytes_opt<T: ByteArrayType>(&self) -> Option<&GenericByteArray<T>>;
903
904    /// Downcast this to a [`GenericByteArray`]
905    ///
906    /// # Panics
907    ///
908    /// Panics if this is not a [`GenericByteArray`]
909    fn as_bytes<T: ByteArrayType>(&self) -> &GenericByteArray<T> {
910        self.as_bytes_opt().expect("byte array")
911    }
912
913    /// Downcast this to a [`GenericStringArray`] returning `None` if not possible
914    fn as_string_opt<O: OffsetSizeTrait>(&self) -> Option<&GenericStringArray<O>> {
915        self.as_bytes_opt()
916    }
917
918    /// Downcast this to a [`GenericStringArray`]
919    ///
920    /// # Panics
921    ///
922    /// Panics if this is not a [`GenericStringArray`]
923    fn as_string<O: OffsetSizeTrait>(&self) -> &GenericStringArray<O> {
924        self.as_bytes_opt().expect("string array")
925    }
926
927    /// Downcast this to a [`GenericBinaryArray`] returning `None` if not possible
928    fn as_binary_opt<O: OffsetSizeTrait>(&self) -> Option<&GenericBinaryArray<O>> {
929        self.as_bytes_opt()
930    }
931
932    /// Downcast this to a [`GenericBinaryArray`]
933    ///
934    /// # Panics
935    ///
936    /// Panics if this is not a [`GenericBinaryArray`]
937    fn as_binary<O: OffsetSizeTrait>(&self) -> &GenericBinaryArray<O> {
938        self.as_bytes_opt().expect("binary array")
939    }
940
941    /// Downcast this to a [`StringViewArray`] returning `None` if not possible
942    fn as_string_view_opt(&self) -> Option<&StringViewArray> {
943        self.as_byte_view_opt()
944    }
945
946    /// Downcast this to a [`StringViewArray`]
947    ///
948    /// # Panics
949    ///
950    /// Panics if this is not a [`StringViewArray`]
951    fn as_string_view(&self) -> &StringViewArray {
952        self.as_byte_view_opt().expect("string view array")
953    }
954
955    /// Downcast this to a [`BinaryViewArray`] returning `None` if not possible
956    fn as_binary_view_opt(&self) -> Option<&BinaryViewArray> {
957        self.as_byte_view_opt()
958    }
959
960    /// Downcast this to a [`BinaryViewArray`]
961    ///
962    /// # Panics
963    ///
964    /// Panics if this is not a [`BinaryViewArray`]
965    fn as_binary_view(&self) -> &BinaryViewArray {
966        self.as_byte_view_opt().expect("binary view array")
967    }
968
969    /// Downcast this to a [`GenericByteViewArray`] returning `None` if not possible
970    fn as_byte_view_opt<T: ByteViewType>(&self) -> Option<&GenericByteViewArray<T>>;
971
972    /// Downcast this to a [`GenericByteViewArray`]
973    ///
974    /// # Panics
975    ///
976    /// Panics if this is not a [`GenericByteViewArray`]
977    fn as_byte_view<T: ByteViewType>(&self) -> &GenericByteViewArray<T> {
978        self.as_byte_view_opt().expect("byte view array")
979    }
980
981    /// Downcast this to a [`StructArray`] returning `None` if not possible
982    fn as_struct_opt(&self) -> Option<&StructArray>;
983
984    /// Downcast this to a [`StructArray`]
985    ///
986    /// # Panics
987    ///
988    /// Panics if this is not a [`StructArray`]
989    fn as_struct(&self) -> &StructArray {
990        self.as_struct_opt().expect("struct array")
991    }
992
993    /// Downcast this to a [`UnionArray`] returning `None` if not possible
994    fn as_union_opt(&self) -> Option<&UnionArray>;
995
996    /// Downcast this to a [`UnionArray`]
997    ///
998    /// # Panics
999    ///
1000    /// Panics if this is not a [`UnionArray`]
1001    fn as_union(&self) -> &UnionArray {
1002        self.as_union_opt().expect("union array")
1003    }
1004
1005    /// Downcast this to a [`GenericListArray`] returning `None` if not possible
1006    fn as_list_opt<O: OffsetSizeTrait>(&self) -> Option<&GenericListArray<O>>;
1007
1008    /// Downcast this to a [`GenericListArray`]
1009    ///
1010    /// # Panics
1011    ///
1012    /// Panics if this is not a [`GenericListArray`]
1013    fn as_list<O: OffsetSizeTrait>(&self) -> &GenericListArray<O> {
1014        self.as_list_opt().expect("list array")
1015    }
1016
1017    /// Downcast this to a [`GenericListViewArray`] returning `None` if not possible
1018    fn as_list_view_opt<O: OffsetSizeTrait>(&self) -> Option<&GenericListViewArray<O>>;
1019
1020    /// Downcast this to a [`GenericListViewArray`]
1021    ///
1022    /// # Panics
1023    ///
1024    /// Panics if this is not a [`GenericListViewArray`]
1025    fn as_list_view<O: OffsetSizeTrait>(&self) -> &GenericListViewArray<O> {
1026        self.as_list_view_opt().expect("list view array")
1027    }
1028
1029    /// Downcast this to a [`FixedSizeBinaryArray`] returning `None` if not possible
1030    fn as_fixed_size_binary_opt(&self) -> Option<&FixedSizeBinaryArray>;
1031
1032    /// Downcast this to a [`FixedSizeBinaryArray`]
1033    ///
1034    /// # Panics
1035    ///
1036    /// Panics if this is not a [`FixedSizeBinaryArray`]
1037    fn as_fixed_size_binary(&self) -> &FixedSizeBinaryArray {
1038        self.as_fixed_size_binary_opt()
1039            .expect("fixed size binary array")
1040    }
1041
1042    /// Downcast this to a [`FixedSizeListArray`] returning `None` if not possible
1043    fn as_fixed_size_list_opt(&self) -> Option<&FixedSizeListArray>;
1044
1045    /// Downcast this to a [`FixedSizeListArray`]
1046    ///
1047    /// # Panics
1048    ///
1049    /// Panics if this is not a [`FixedSizeListArray`]
1050    fn as_fixed_size_list(&self) -> &FixedSizeListArray {
1051        self.as_fixed_size_list_opt()
1052            .expect("fixed size list array")
1053    }
1054
1055    /// Downcast this to a [`MapArray`] returning `None` if not possible
1056    fn as_map_opt(&self) -> Option<&MapArray>;
1057
1058    /// Downcast this to a [`MapArray`]
1059    ///
1060    /// # Panics
1061    ///
1062    /// Panics if this is not a [`MapArray`]
1063    fn as_map(&self) -> &MapArray {
1064        self.as_map_opt().expect("map array")
1065    }
1066
1067    /// Downcast this to a [`DictionaryArray`] returning `None` if not possible
1068    fn as_dictionary_opt<K: ArrowDictionaryKeyType>(&self) -> Option<&DictionaryArray<K>>;
1069
1070    /// Downcast this to a [`DictionaryArray`]
1071    ///
1072    /// # Panics
1073    ///
1074    /// Panics if this is not a [`DictionaryArray`]
1075    fn as_dictionary<K: ArrowDictionaryKeyType>(&self) -> &DictionaryArray<K> {
1076        self.as_dictionary_opt().expect("dictionary array")
1077    }
1078
1079    /// Downcast this to a [`RunArray`] returning `None` if not possible
1080    fn as_run_opt<K: RunEndIndexType>(&self) -> Option<&RunArray<K>>;
1081
1082    /// Downcast this to a [`RunArray`]
1083    ///
1084    /// # Panics
1085    ///
1086    /// Panics if this is not a [`RunArray`]
1087    fn as_run<K: RunEndIndexType>(&self) -> &RunArray<K> {
1088        self.as_run_opt().expect("run array")
1089    }
1090
1091    /// Downcasts this to a [`AnyDictionaryArray`] returning `None` if not possible
1092    fn as_any_dictionary_opt(&self) -> Option<&dyn AnyDictionaryArray>;
1093
1094    /// Downcasts this to a [`AnyDictionaryArray`]
1095    ///
1096    /// # Panics
1097    ///
1098    /// Panics if this is not a [`AnyDictionaryArray`]
1099    fn as_any_dictionary(&self) -> &dyn AnyDictionaryArray {
1100        self.as_any_dictionary_opt().expect("any dictionary array")
1101    }
1102
1103    /// Downcasts this to a [`AnyRunEndArray`] returning `None` if not possible
1104    fn as_any_ree_opt(&self) -> Option<&dyn AnyRunEndArray>;
1105
1106    /// Downcasts this to a [`AnyRunEndArray`]
1107    ///
1108    /// # Panics
1109    ///
1110    /// Panics if this is not a [`AnyRunEndArray`]
1111    fn as_any_ree(&self) -> &dyn AnyRunEndArray {
1112        self.as_any_ree_opt().expect("any run end array")
1113    }
1114}
1115
1116impl private::Sealed for dyn Array + '_ {}
1117impl AsArray for dyn Array + '_ {
1118    fn as_boolean_opt(&self) -> Option<&BooleanArray> {
1119        self.as_any().downcast_ref()
1120    }
1121
1122    fn as_primitive_opt<T: ArrowPrimitiveType>(&self) -> Option<&PrimitiveArray<T>> {
1123        self.as_any().downcast_ref()
1124    }
1125
1126    fn as_bytes_opt<T: ByteArrayType>(&self) -> Option<&GenericByteArray<T>> {
1127        self.as_any().downcast_ref()
1128    }
1129
1130    fn as_byte_view_opt<T: ByteViewType>(&self) -> Option<&GenericByteViewArray<T>> {
1131        self.as_any().downcast_ref()
1132    }
1133
1134    fn as_struct_opt(&self) -> Option<&StructArray> {
1135        self.as_any().downcast_ref()
1136    }
1137
1138    fn as_union_opt(&self) -> Option<&UnionArray> {
1139        self.as_any().downcast_ref()
1140    }
1141
1142    fn as_list_opt<O: OffsetSizeTrait>(&self) -> Option<&GenericListArray<O>> {
1143        self.as_any().downcast_ref()
1144    }
1145
1146    fn as_list_view_opt<O: OffsetSizeTrait>(&self) -> Option<&GenericListViewArray<O>> {
1147        self.as_any().downcast_ref()
1148    }
1149
1150    fn as_fixed_size_binary_opt(&self) -> Option<&FixedSizeBinaryArray> {
1151        self.as_any().downcast_ref()
1152    }
1153
1154    fn as_fixed_size_list_opt(&self) -> Option<&FixedSizeListArray> {
1155        self.as_any().downcast_ref()
1156    }
1157
1158    fn as_map_opt(&self) -> Option<&MapArray> {
1159        self.as_any().downcast_ref()
1160    }
1161
1162    fn as_dictionary_opt<K: ArrowDictionaryKeyType>(&self) -> Option<&DictionaryArray<K>> {
1163        self.as_any().downcast_ref()
1164    }
1165
1166    fn as_run_opt<K: RunEndIndexType>(&self) -> Option<&RunArray<K>> {
1167        self.as_any().downcast_ref()
1168    }
1169
1170    fn as_any_dictionary_opt(&self) -> Option<&dyn AnyDictionaryArray> {
1171        let array = self;
1172        downcast_dictionary_array! {
1173            array => Some(array),
1174            _ => None
1175        }
1176    }
1177
1178    fn as_any_ree_opt(&self) -> Option<&dyn AnyRunEndArray> {
1179        let array = self;
1180        downcast_run_array! {
1181            array => Some(array),
1182            _ => None
1183        }
1184    }
1185}
1186
1187impl private::Sealed for ArrayRef {}
1188impl AsArray for ArrayRef {
1189    fn as_boolean_opt(&self) -> Option<&BooleanArray> {
1190        self.as_ref().as_boolean_opt()
1191    }
1192
1193    fn as_primitive_opt<T: ArrowPrimitiveType>(&self) -> Option<&PrimitiveArray<T>> {
1194        self.as_ref().as_primitive_opt()
1195    }
1196
1197    fn as_bytes_opt<T: ByteArrayType>(&self) -> Option<&GenericByteArray<T>> {
1198        self.as_ref().as_bytes_opt()
1199    }
1200
1201    fn as_byte_view_opt<T: ByteViewType>(&self) -> Option<&GenericByteViewArray<T>> {
1202        self.as_ref().as_byte_view_opt()
1203    }
1204
1205    fn as_struct_opt(&self) -> Option<&StructArray> {
1206        self.as_ref().as_struct_opt()
1207    }
1208
1209    fn as_union_opt(&self) -> Option<&UnionArray> {
1210        self.as_any().downcast_ref()
1211    }
1212
1213    fn as_list_opt<O: OffsetSizeTrait>(&self) -> Option<&GenericListArray<O>> {
1214        self.as_ref().as_list_opt()
1215    }
1216
1217    fn as_list_view_opt<O: OffsetSizeTrait>(&self) -> Option<&GenericListViewArray<O>> {
1218        self.as_ref().as_list_view_opt()
1219    }
1220
1221    fn as_fixed_size_binary_opt(&self) -> Option<&FixedSizeBinaryArray> {
1222        self.as_ref().as_fixed_size_binary_opt()
1223    }
1224
1225    fn as_fixed_size_list_opt(&self) -> Option<&FixedSizeListArray> {
1226        self.as_ref().as_fixed_size_list_opt()
1227    }
1228
1229    fn as_map_opt(&self) -> Option<&MapArray> {
1230        self.as_any().downcast_ref()
1231    }
1232
1233    fn as_dictionary_opt<K: ArrowDictionaryKeyType>(&self) -> Option<&DictionaryArray<K>> {
1234        self.as_ref().as_dictionary_opt()
1235    }
1236
1237    fn as_any_dictionary_opt(&self) -> Option<&dyn AnyDictionaryArray> {
1238        self.as_ref().as_any_dictionary_opt()
1239    }
1240
1241    fn as_any_ree_opt(&self) -> Option<&dyn AnyRunEndArray> {
1242        self.as_ref().as_any_ree_opt()
1243    }
1244
1245    fn as_run_opt<K: RunEndIndexType>(&self) -> Option<&RunArray<K>> {
1246        self.as_ref().as_run_opt()
1247    }
1248
1249    fn as_string_opt<O: OffsetSizeTrait>(&self) -> Option<&GenericStringArray<O>> {
1250        self.as_ref().as_string_opt()
1251    }
1252}
1253
1254#[cfg(test)]
1255mod tests {
1256    use super::*;
1257    use arrow_buffer::i256;
1258    use arrow_schema::DataType;
1259    use std::sync::Arc;
1260
1261    #[test]
1262    fn test_as_primitive_array_ref() {
1263        let array: Int32Array = vec![1, 2, 3].into_iter().map(Some).collect();
1264        assert!(!as_primitive_array::<Int32Type>(&array).is_empty());
1265
1266        // should also work when wrapped in an Arc
1267        let array: ArrayRef = Arc::new(array);
1268        assert!(!as_primitive_array::<Int32Type>(&array).is_empty());
1269    }
1270
1271    #[test]
1272    fn test_as_string_array_ref() {
1273        let array: StringArray = vec!["foo", "bar"].into_iter().map(Some).collect();
1274        assert!(!as_string_array(&array).is_empty());
1275
1276        // should also work when wrapped in an Arc
1277        let array: ArrayRef = Arc::new(array);
1278        assert!(!as_string_array(&array).is_empty())
1279    }
1280
1281    #[test]
1282    fn test_decimal32array() {
1283        let a = Decimal32Array::from_iter_values([1, 2, 4, 5]);
1284        assert!(!as_primitive_array::<Decimal32Type>(&a).is_empty());
1285    }
1286
1287    #[test]
1288    fn test_decimal64array() {
1289        let a = Decimal64Array::from_iter_values([1, 2, 4, 5]);
1290        assert!(!as_primitive_array::<Decimal64Type>(&a).is_empty());
1291    }
1292
1293    #[test]
1294    fn test_decimal128array() {
1295        let a = Decimal128Array::from_iter_values([1, 2, 4, 5]);
1296        assert!(!as_primitive_array::<Decimal128Type>(&a).is_empty());
1297    }
1298
1299    #[test]
1300    fn test_decimal256array() {
1301        let a = Decimal256Array::from_iter_values([1, 2, 4, 5].into_iter().map(i256::from_i128));
1302        assert!(!as_primitive_array::<Decimal256Type>(&a).is_empty());
1303    }
1304
1305    #[test]
1306    fn downcast_integer_array_should_match_only_integers() {
1307        let i32_array: ArrayRef = Arc::new(Int32Array::new_null(1));
1308        let i32_array_ref = &i32_array;
1309        downcast_integer_array!(
1310            i32_array_ref => {
1311                assert_eq!(i32_array_ref.null_count(), 1);
1312            },
1313            _ => panic!("unexpected data type")
1314        );
1315    }
1316
1317    #[test]
1318    fn downcast_integer_array_should_not_match_primitive_that_are_not_integers() {
1319        let array: ArrayRef = Arc::new(Float32Array::new_null(1));
1320        let array_ref = &array;
1321        downcast_integer_array!(
1322            array_ref => {
1323                panic!("unexpected data type {}", array_ref.data_type())
1324            },
1325            DataType::Float32 => {
1326                assert_eq!(array_ref.null_count(), 1);
1327            },
1328            _ => panic!("unexpected data type")
1329        );
1330    }
1331
1332    #[test]
1333    fn downcast_integer_array_should_not_match_non_primitive() {
1334        let array: ArrayRef = Arc::new(StringArray::new_null(1));
1335        let array_ref = &array;
1336        downcast_integer_array!(
1337            array_ref => {
1338                panic!("unexpected data type {}", array_ref.data_type())
1339            },
1340            DataType::Utf8 => {
1341                assert_eq!(array_ref.null_count(), 1);
1342            },
1343            _ => panic!("unexpected data type")
1344        );
1345    }
1346}
arrow_array/cast.rs

arrow_array/
cast.rs
