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///         _ => u8::MAX,
64///     }
65/// }
66///
67/// assert_eq!(dictionary_key_size(&DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8))), 4);
68/// assert_eq!(dictionary_key_size(&DataType::Dictionary(Box::new(DataType::Int64), Box::new(DataType::Utf8))), 8);
69/// assert_eq!(dictionary_key_size(&DataType::Dictionary(Box::new(DataType::UInt16), Box::new(DataType::Utf8))), 2);
70/// ```
71///
72/// [`DataType`]: arrow_schema::DataType
73#[macro_export]
74macro_rules! downcast_integer {
75    ($($data_type:expr),+ => ($m:path $(, $args:tt)*), $($p:pat => $fallback:expr $(,)*)*) => {
76        match ($($data_type),+) {
77            $crate::repeat_pat!($crate::cast::__private::DataType::Int8, $($data_type),+) => {
78                $m!($crate::types::Int8Type $(, $args)*)
79            }
80            $crate::repeat_pat!($crate::cast::__private::DataType::Int16, $($data_type),+) => {
81                $m!($crate::types::Int16Type $(, $args)*)
82            }
83            $crate::repeat_pat!($crate::cast::__private::DataType::Int32, $($data_type),+) => {
84                $m!($crate::types::Int32Type $(, $args)*)
85            }
86            $crate::repeat_pat!($crate::cast::__private::DataType::Int64, $($data_type),+) => {
87                $m!($crate::types::Int64Type $(, $args)*)
88            }
89            $crate::repeat_pat!($crate::cast::__private::DataType::UInt8, $($data_type),+) => {
90                $m!($crate::types::UInt8Type $(, $args)*)
91            }
92            $crate::repeat_pat!($crate::cast::__private::DataType::UInt16, $($data_type),+) => {
93                $m!($crate::types::UInt16Type $(, $args)*)
94            }
95            $crate::repeat_pat!($crate::cast::__private::DataType::UInt32, $($data_type),+) => {
96                $m!($crate::types::UInt32Type $(, $args)*)
97            }
98            $crate::repeat_pat!($crate::cast::__private::DataType::UInt64, $($data_type),+) => {
99                $m!($crate::types::UInt64Type $(, $args)*)
100            }
101            $($p => $fallback,)*
102        }
103    };
104}
105
106/// Given one or more expressions evaluating to an integer [`DataType`] invokes the provided macro
107/// `m` with the corresponding integer [`RunEndIndexType`], followed by any additional arguments
108///
109/// ```
110/// # use std::sync::Arc;
111/// # use arrow_array::{downcast_primitive, ArrowPrimitiveType, downcast_run_end_index};
112/// # use arrow_schema::{DataType, Field};
113///
114/// macro_rules! run_end_size_helper {
115///   ($t:ty, $o:ty) => {
116///       std::mem::size_of::<<$t as ArrowPrimitiveType>::Native>() as $o
117///   };
118/// }
119///
120/// fn run_end_index_size(t: &DataType) -> u8 {
121///     match t {
122///         DataType::RunEndEncoded(k, _) => downcast_run_end_index! {
123///             k.data_type() => (run_end_size_helper, u8),
124///             _ => unreachable!(),
125///         },
126///         _ => u8::MAX,
127///     }
128/// }
129///
130/// 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);
131/// 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);
132/// 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);
133/// ```
134///
135/// [`DataType`]: arrow_schema::DataType
136#[macro_export]
137macro_rules! downcast_run_end_index {
138    ($($data_type:expr),+ => ($m:path $(, $args:tt)*), $($p:pat => $fallback:expr $(,)*)*) => {
139        match ($($data_type),+) {
140            $crate::repeat_pat!($crate::cast::__private::DataType::Int16, $($data_type),+) => {
141                $m!($crate::types::Int16Type $(, $args)*)
142            }
143            $crate::repeat_pat!($crate::cast::__private::DataType::Int32, $($data_type),+) => {
144                $m!($crate::types::Int32Type $(, $args)*)
145            }
146            $crate::repeat_pat!($crate::cast::__private::DataType::Int64, $($data_type),+) => {
147                $m!($crate::types::Int64Type $(, $args)*)
148            }
149            $($p => $fallback,)*
150        }
151    };
152}
153
154/// Given one or more expressions evaluating to primitive [`DataType`] invokes the provided macro
155/// `m` with the corresponding [`ArrowPrimitiveType`], followed by any additional arguments
156///
157/// ```
158/// # use arrow_array::{downcast_temporal, ArrowPrimitiveType};
159/// # use arrow_schema::DataType;
160///
161/// macro_rules! temporal_size_helper {
162///   ($t:ty, $o:ty) => {
163///       std::mem::size_of::<<$t as ArrowPrimitiveType>::Native>() as $o
164///   };
165/// }
166///
167/// fn temporal_size(t: &DataType) -> u8 {
168///     downcast_temporal! {
169///         t => (temporal_size_helper, u8),
170///         _ => u8::MAX
171///     }
172/// }
173///
174/// assert_eq!(temporal_size(&DataType::Date32), 4);
175/// assert_eq!(temporal_size(&DataType::Date64), 8);
176/// ```
177///
178/// [`DataType`]: arrow_schema::DataType
179#[macro_export]
180macro_rules! downcast_temporal {
181    ($($data_type:expr),+ => ($m:path $(, $args:tt)*), $($p:pat => $fallback:expr $(,)*)*) => {
182        match ($($data_type),+) {
183            $crate::repeat_pat!($crate::cast::__private::DataType::Time32($crate::cast::__private::TimeUnit::Second), $($data_type),+) => {
184                $m!($crate::types::Time32SecondType $(, $args)*)
185            }
186            $crate::repeat_pat!($crate::cast::__private::DataType::Time32($crate::cast::__private::TimeUnit::Millisecond), $($data_type),+) => {
187                $m!($crate::types::Time32MillisecondType $(, $args)*)
188            }
189            $crate::repeat_pat!($crate::cast::__private::DataType::Time64($crate::cast::__private::TimeUnit::Microsecond), $($data_type),+) => {
190                $m!($crate::types::Time64MicrosecondType $(, $args)*)
191            }
192            $crate::repeat_pat!($crate::cast::__private::DataType::Time64($crate::cast::__private::TimeUnit::Nanosecond), $($data_type),+) => {
193                $m!($crate::types::Time64NanosecondType $(, $args)*)
194            }
195            $crate::repeat_pat!($crate::cast::__private::DataType::Date32, $($data_type),+) => {
196                $m!($crate::types::Date32Type $(, $args)*)
197            }
198            $crate::repeat_pat!($crate::cast::__private::DataType::Date64, $($data_type),+) => {
199                $m!($crate::types::Date64Type $(, $args)*)
200            }
201            $crate::repeat_pat!($crate::cast::__private::DataType::Timestamp($crate::cast::__private::TimeUnit::Second, _), $($data_type),+) => {
202                $m!($crate::types::TimestampSecondType $(, $args)*)
203            }
204            $crate::repeat_pat!($crate::cast::__private::DataType::Timestamp($crate::cast::__private::TimeUnit::Millisecond, _), $($data_type),+) => {
205                $m!($crate::types::TimestampMillisecondType $(, $args)*)
206            }
207            $crate::repeat_pat!($crate::cast::__private::DataType::Timestamp($crate::cast::__private::TimeUnit::Microsecond, _), $($data_type),+) => {
208                $m!($crate::types::TimestampMicrosecondType $(, $args)*)
209            }
210            $crate::repeat_pat!($crate::cast::__private::DataType::Timestamp($crate::cast::__private::TimeUnit::Nanosecond, _), $($data_type),+) => {
211                $m!($crate::types::TimestampNanosecondType $(, $args)*)
212            }
213            $($p => $fallback,)*
214        }
215    };
216}
217
218/// Downcast an [`Array`] to a temporal [`PrimitiveArray`] based on its [`DataType`]
219/// accepts a number of subsequent patterns to match the data type
220///
221/// ```
222/// # use arrow_array::{Array, downcast_temporal_array, cast::as_string_array};
223/// # use arrow_schema::DataType;
224///
225/// fn print_temporal(array: &dyn Array) {
226///     downcast_temporal_array!(
227///         array => {
228///             for v in array {
229///                 println!("{:?}", v);
230///             }
231///         }
232///         DataType::Utf8 => {
233///             for v in as_string_array(array) {
234///                 println!("{:?}", v);
235///             }
236///         }
237///         t => println!("Unsupported datatype {}", t)
238///     )
239/// }
240/// ```
241///
242/// [`DataType`]: arrow_schema::DataType
243#[macro_export]
244macro_rules! downcast_temporal_array {
245    ($values:ident => $e:expr, $($p:pat => $fallback:expr $(,)*)*) => {
246        $crate::downcast_temporal_array!($values => {$e} $($p => $fallback)*)
247    };
248    (($($values:ident),+) => $e:expr, $($p:pat => $fallback:expr $(,)*)*) => {
249        $crate::downcast_temporal_array!($($values),+ => {$e} $($p => $fallback)*)
250    };
251    ($($values:ident),+ => $e:block $($p:pat => $fallback:expr $(,)*)*) => {
252        $crate::downcast_temporal_array!(($($values),+) => $e $($p => $fallback)*)
253    };
254    (($($values:ident),+) => $e:block $($p:pat => $fallback:expr $(,)*)*) => {
255        $crate::downcast_temporal!{
256            $($values.data_type()),+ => ($crate::downcast_primitive_array_helper, $($values),+, $e),
257            $($p => $fallback,)*
258        }
259    };
260}
261
262/// Given one or more expressions evaluating to primitive [`DataType`] invokes the provided macro
263/// `m` with the corresponding [`ArrowPrimitiveType`], followed by any additional arguments
264///
265/// ```
266/// # use arrow_array::{downcast_primitive, ArrowPrimitiveType};
267/// # use arrow_schema::DataType;
268///
269/// macro_rules! primitive_size_helper {
270///   ($t:ty, $o:ty) => {
271///       std::mem::size_of::<<$t as ArrowPrimitiveType>::Native>() as $o
272///   };
273/// }
274///
275/// fn primitive_size(t: &DataType) -> u8 {
276///     downcast_primitive! {
277///         t => (primitive_size_helper, u8),
278///         _ => u8::MAX
279///     }
280/// }
281///
282/// assert_eq!(primitive_size(&DataType::Int32), 4);
283/// assert_eq!(primitive_size(&DataType::Int64), 8);
284/// assert_eq!(primitive_size(&DataType::Float16), 2);
285/// assert_eq!(primitive_size(&DataType::Decimal128(38, 10)), 16);
286/// assert_eq!(primitive_size(&DataType::Decimal256(76, 20)), 32);
287/// ```
288///
289/// [`DataType`]: arrow_schema::DataType
290#[macro_export]
291macro_rules! downcast_primitive {
292    ($($data_type:expr),+ => ($m:path $(, $args:tt)*), $($p:pat => $fallback:expr $(,)*)*) => {
293        $crate::downcast_integer! {
294            $($data_type),+ => ($m $(, $args)*),
295            $crate::repeat_pat!($crate::cast::__private::DataType::Float16, $($data_type),+) => {
296                $m!($crate::types::Float16Type $(, $args)*)
297            }
298            $crate::repeat_pat!($crate::cast::__private::DataType::Float32, $($data_type),+) => {
299                $m!($crate::types::Float32Type $(, $args)*)
300            }
301            $crate::repeat_pat!($crate::cast::__private::DataType::Float64, $($data_type),+) => {
302                $m!($crate::types::Float64Type $(, $args)*)
303            }
304            $crate::repeat_pat!($crate::cast::__private::DataType::Decimal128(_, _), $($data_type),+) => {
305                $m!($crate::types::Decimal128Type $(, $args)*)
306            }
307            $crate::repeat_pat!($crate::cast::__private::DataType::Decimal256(_, _), $($data_type),+) => {
308                $m!($crate::types::Decimal256Type $(, $args)*)
309            }
310            $crate::repeat_pat!($crate::cast::__private::DataType::Interval($crate::cast::__private::IntervalUnit::YearMonth), $($data_type),+) => {
311                $m!($crate::types::IntervalYearMonthType $(, $args)*)
312            }
313            $crate::repeat_pat!($crate::cast::__private::DataType::Interval($crate::cast::__private::IntervalUnit::DayTime), $($data_type),+) => {
314                $m!($crate::types::IntervalDayTimeType $(, $args)*)
315            }
316            $crate::repeat_pat!($crate::cast::__private::DataType::Interval($crate::cast::__private::IntervalUnit::MonthDayNano), $($data_type),+) => {
317                $m!($crate::types::IntervalMonthDayNanoType $(, $args)*)
318            }
319            $crate::repeat_pat!($crate::cast::__private::DataType::Duration($crate::cast::__private::TimeUnit::Second), $($data_type),+) => {
320                $m!($crate::types::DurationSecondType $(, $args)*)
321            }
322            $crate::repeat_pat!($crate::cast::__private::DataType::Duration($crate::cast::__private::TimeUnit::Millisecond), $($data_type),+) => {
323                $m!($crate::types::DurationMillisecondType $(, $args)*)
324            }
325            $crate::repeat_pat!($crate::cast::__private::DataType::Duration($crate::cast::__private::TimeUnit::Microsecond), $($data_type),+) => {
326                $m!($crate::types::DurationMicrosecondType $(, $args)*)
327            }
328            $crate::repeat_pat!($crate::cast::__private::DataType::Duration($crate::cast::__private::TimeUnit::Nanosecond), $($data_type),+) => {
329                $m!($crate::types::DurationNanosecondType $(, $args)*)
330            }
331            _ => {
332                $crate::downcast_temporal! {
333                    $($data_type),+ => ($m $(, $args)*),
334                    $($p => $fallback,)*
335                }
336            }
337        }
338    };
339}
340
341#[macro_export]
342#[doc(hidden)]
343macro_rules! downcast_primitive_array_helper {
344    ($t:ty, $($values:ident),+, $e:block) => {{
345        $(let $values = $crate::cast::as_primitive_array::<$t>($values);)+
346        $e
347    }};
348}
349
350/// Downcast an [`Array`] to a [`PrimitiveArray`] based on its [`DataType`]
351/// accepts a number of subsequent patterns to match the data type
352///
353/// ```
354/// # use arrow_array::{Array, downcast_primitive_array, cast::as_string_array};
355/// # use arrow_schema::DataType;
356///
357/// fn print_primitive(array: &dyn Array) {
358///     downcast_primitive_array!(
359///         array => {
360///             for v in array {
361///                 println!("{:?}", v);
362///             }
363///         }
364///         DataType::Utf8 => {
365///             for v in as_string_array(array) {
366///                 println!("{:?}", v);
367///             }
368///         }
369///         t => println!("Unsupported datatype {}", t)
370///     )
371/// }
372/// ```
373///
374/// [`DataType`]: arrow_schema::DataType
375#[macro_export]
376macro_rules! downcast_primitive_array {
377    ($values:ident => $e:expr, $($p:pat => $fallback:expr $(,)*)*) => {
378        $crate::downcast_primitive_array!($values => {$e} $($p => $fallback)*)
379    };
380    (($($values:ident),+) => $e:expr, $($p:pat => $fallback:expr $(,)*)*) => {
381        $crate::downcast_primitive_array!($($values),+ => {$e} $($p => $fallback)*)
382    };
383    ($($values:ident),+ => $e:block $($p:pat => $fallback:expr $(,)*)*) => {
384        $crate::downcast_primitive_array!(($($values),+) => $e $($p => $fallback)*)
385    };
386    (($($values:ident),+) => $e:block $($p:pat => $fallback:expr $(,)*)*) => {
387        $crate::downcast_primitive!{
388            $($values.data_type()),+ => ($crate::downcast_primitive_array_helper, $($values),+, $e),
389            $($p => $fallback,)*
390        }
391    };
392}
393
394/// Force downcast of an [`Array`], such as an [`ArrayRef`], to
395/// [`PrimitiveArray<T>`], panic'ing on failure.
396///
397/// # Example
398///
399/// ```
400/// # use std::sync::Arc;
401/// # use arrow_array::{ArrayRef, Int32Array};
402/// # use arrow_array::cast::as_primitive_array;
403/// # use arrow_array::types::Int32Type;
404///
405/// let arr: ArrayRef = Arc::new(Int32Array::from(vec![Some(1)]));
406///
407/// // Downcast an `ArrayRef` to Int32Array / PrimitiveArray<Int32>:
408/// let primitive_array: &Int32Array = as_primitive_array(&arr);
409///
410/// // Equivalently:
411/// let primitive_array = as_primitive_array::<Int32Type>(&arr);
412///
413/// // This is the equivalent of:
414/// let primitive_array = arr
415///     .as_any()
416///     .downcast_ref::<Int32Array>()
417///     .unwrap();
418/// ```
419pub fn as_primitive_array<T>(arr: &dyn Array) -> &PrimitiveArray<T>
420where
421    T: ArrowPrimitiveType,
422{
423    arr.as_any()
424        .downcast_ref::<PrimitiveArray<T>>()
425        .expect("Unable to downcast to primitive array")
426}
427
428#[macro_export]
429#[doc(hidden)]
430macro_rules! downcast_dictionary_array_helper {
431    ($t:ty, $($values:ident),+, $e:block) => {{
432        $(let $values = $crate::cast::as_dictionary_array::<$t>($values);)+
433        $e
434    }};
435}
436
437/// Downcast an [`Array`] to a [`DictionaryArray`] based on its [`DataType`], accepts
438/// a number of subsequent patterns to match the data type
439///
440/// ```
441/// # use arrow_array::{Array, StringArray, downcast_dictionary_array, cast::as_string_array};
442/// # use arrow_schema::DataType;
443///
444/// fn print_strings(array: &dyn Array) {
445///     downcast_dictionary_array!(
446///         array => match array.values().data_type() {
447///             DataType::Utf8 => {
448///                 for v in array.downcast_dict::<StringArray>().unwrap() {
449///                     println!("{:?}", v);
450///                 }
451///             }
452///             t => println!("Unsupported dictionary value type {}", t),
453///         },
454///         DataType::Utf8 => {
455///             for v in as_string_array(array) {
456///                 println!("{:?}", v);
457///             }
458///         }
459///         t => println!("Unsupported datatype {}", t)
460///     )
461/// }
462/// ```
463///
464/// [`DataType`]: arrow_schema::DataType
465#[macro_export]
466macro_rules! downcast_dictionary_array {
467    ($values:ident => $e:expr, $($p:pat => $fallback:expr $(,)*)*) => {
468        downcast_dictionary_array!($values => {$e} $($p => $fallback)*)
469    };
470
471    ($values:ident => $e:block $($p:pat => $fallback:expr $(,)*)*) => {
472        match $values.data_type() {
473            $crate::cast::__private::DataType::Dictionary(k, _) => {
474                $crate::downcast_integer! {
475                    k.as_ref() => ($crate::downcast_dictionary_array_helper, $values, $e),
476                    k => unreachable!("unsupported dictionary key type: {}", k)
477                }
478            }
479            $($p => $fallback,)*
480        }
481    }
482}
483
484/// Force downcast of an [`Array`], such as an [`ArrayRef`] to
485/// [`DictionaryArray<T>`], panic'ing on failure.
486///
487/// # Example
488///
489/// ```
490/// # use arrow_array::{ArrayRef, DictionaryArray};
491/// # use arrow_array::cast::as_dictionary_array;
492/// # use arrow_array::types::Int32Type;
493///
494/// let arr: DictionaryArray<Int32Type> = vec![Some("foo")].into_iter().collect();
495/// let arr: ArrayRef = std::sync::Arc::new(arr);
496/// let dict_array: &DictionaryArray<Int32Type> = as_dictionary_array::<Int32Type>(&arr);
497/// ```
498pub fn as_dictionary_array<T>(arr: &dyn Array) -> &DictionaryArray<T>
499where
500    T: ArrowDictionaryKeyType,
501{
502    arr.as_any()
503        .downcast_ref::<DictionaryArray<T>>()
504        .expect("Unable to downcast to dictionary array")
505}
506
507/// Force downcast of an [`Array`], such as an [`ArrayRef`] to
508/// [`RunArray<T>`], panic'ing on failure.
509///
510/// # Example
511///
512/// ```
513/// # use arrow_array::{ArrayRef, RunArray};
514/// # use arrow_array::cast::as_run_array;
515/// # use arrow_array::types::Int32Type;
516///
517/// let arr: RunArray<Int32Type> = vec![Some("foo")].into_iter().collect();
518/// let arr: ArrayRef = std::sync::Arc::new(arr);
519/// let run_array: &RunArray<Int32Type> = as_run_array::<Int32Type>(&arr);
520/// ```
521pub fn as_run_array<T>(arr: &dyn Array) -> &RunArray<T>
522where
523    T: RunEndIndexType,
524{
525    arr.as_any()
526        .downcast_ref::<RunArray<T>>()
527        .expect("Unable to downcast to run array")
528}
529
530#[macro_export]
531#[doc(hidden)]
532macro_rules! downcast_run_array_helper {
533    ($t:ty, $($values:ident),+, $e:block) => {{
534        $(let $values = $crate::cast::as_run_array::<$t>($values);)+
535        $e
536    }};
537}
538
539/// Downcast an [`Array`] to a [`RunArray`] based on its [`DataType`], accepts
540/// a number of subsequent patterns to match the data type
541///
542/// ```
543/// # use arrow_array::{Array, StringArray, downcast_run_array, cast::as_string_array};
544/// # use arrow_schema::DataType;
545///
546/// fn print_strings(array: &dyn Array) {
547///     downcast_run_array!(
548///         array => match array.values().data_type() {
549///             DataType::Utf8 => {
550///                 for v in array.downcast::<StringArray>().unwrap() {
551///                     println!("{:?}", v);
552///                 }
553///             }
554///             t => println!("Unsupported run array value type {}", t),
555///         },
556///         DataType::Utf8 => {
557///             for v in as_string_array(array) {
558///                 println!("{:?}", v);
559///             }
560///         }
561///         t => println!("Unsupported datatype {}", t)
562///     )
563/// }
564/// ```
565///
566/// [`DataType`]: arrow_schema::DataType
567#[macro_export]
568macro_rules! downcast_run_array {
569    ($values:ident => $e:expr, $($p:pat => $fallback:expr $(,)*)*) => {
570        downcast_run_array!($values => {$e} $($p => $fallback)*)
571    };
572
573    ($values:ident => $e:block $($p:pat => $fallback:expr $(,)*)*) => {
574        match $values.data_type() {
575            $crate::cast::__private::DataType::RunEndEncoded(k, _) => {
576                $crate::downcast_run_end_index! {
577                    k.data_type() => ($crate::downcast_run_array_helper, $values, $e),
578                    k => unreachable!("unsupported run end index type: {}", k)
579                }
580            }
581            $($p => $fallback,)*
582        }
583    }
584}
585
586/// Force downcast of an [`Array`], such as an [`ArrayRef`] to
587/// [`GenericListArray<T>`], panicking on failure.
588pub fn as_generic_list_array<S: OffsetSizeTrait>(arr: &dyn Array) -> &GenericListArray<S> {
589    arr.as_any()
590        .downcast_ref::<GenericListArray<S>>()
591        .expect("Unable to downcast to list array")
592}
593
594/// Force downcast of an [`Array`], such as an [`ArrayRef`] to
595/// [`ListArray`], panicking on failure.
596#[inline]
597pub fn as_list_array(arr: &dyn Array) -> &ListArray {
598    as_generic_list_array::<i32>(arr)
599}
600
601/// Force downcast of an [`Array`], such as an [`ArrayRef`] to
602/// [`FixedSizeListArray`], panicking on failure.
603#[inline]
604pub fn as_fixed_size_list_array(arr: &dyn Array) -> &FixedSizeListArray {
605    arr.as_any()
606        .downcast_ref::<FixedSizeListArray>()
607        .expect("Unable to downcast to fixed size list array")
608}
609
610/// Force downcast of an [`Array`], such as an [`ArrayRef`] to
611/// [`LargeListArray`], panicking on failure.
612#[inline]
613pub fn as_large_list_array(arr: &dyn Array) -> &LargeListArray {
614    as_generic_list_array::<i64>(arr)
615}
616
617/// Force downcast of an [`Array`], such as an [`ArrayRef`] to
618/// [`GenericBinaryArray<S>`], panicking on failure.
619#[inline]
620pub fn as_generic_binary_array<S: OffsetSizeTrait>(arr: &dyn Array) -> &GenericBinaryArray<S> {
621    arr.as_any()
622        .downcast_ref::<GenericBinaryArray<S>>()
623        .expect("Unable to downcast to binary array")
624}
625
626/// Force downcast of an [`Array`], such as an [`ArrayRef`] to
627/// [`StringArray`], panicking on failure.
628///
629/// # Example
630///
631/// ```
632/// # use std::sync::Arc;
633/// # use arrow_array::cast::as_string_array;
634/// # use arrow_array::{ArrayRef, StringArray};
635///
636/// let arr: ArrayRef = Arc::new(StringArray::from_iter(vec![Some("foo")]));
637/// let string_array = as_string_array(&arr);
638/// ```
639pub fn as_string_array(arr: &dyn Array) -> &StringArray {
640    arr.as_any()
641        .downcast_ref::<StringArray>()
642        .expect("Unable to downcast to StringArray")
643}
644
645/// Force downcast of an [`Array`], such as an [`ArrayRef`] to
646/// [`BooleanArray`], panicking on failure.
647///
648/// # Example
649///
650/// ```
651/// # use std::sync::Arc;
652/// # use arrow_array::{ArrayRef, BooleanArray};
653/// # use arrow_array::cast::as_boolean_array;
654///
655/// let arr: ArrayRef = Arc::new(BooleanArray::from_iter(vec![Some(true)]));
656/// let boolean_array = as_boolean_array(&arr);
657/// ```
658pub fn as_boolean_array(arr: &dyn Array) -> &BooleanArray {
659    arr.as_any()
660        .downcast_ref::<BooleanArray>()
661        .expect("Unable to downcast to BooleanArray")
662}
663
664macro_rules! array_downcast_fn {
665    ($name: ident, $arrty: ty, $arrty_str:expr) => {
666        #[doc = "Force downcast of an [`Array`], such as an [`ArrayRef`] to "]
667        #[doc = $arrty_str]
668        pub fn $name(arr: &dyn Array) -> &$arrty {
669            arr.as_any().downcast_ref::<$arrty>().expect(concat!(
670                "Unable to downcast to typed array through ",
671                stringify!($name)
672            ))
673        }
674    };
675
676    // use recursive macro to generate dynamic doc string for a given array type
677    ($name: ident, $arrty: ty) => {
678        array_downcast_fn!(
679            $name,
680            $arrty,
681            concat!("[`", stringify!($arrty), "`], panicking on failure.")
682        );
683    };
684}
685
686array_downcast_fn!(as_largestring_array, LargeStringArray);
687array_downcast_fn!(as_null_array, NullArray);
688array_downcast_fn!(as_struct_array, StructArray);
689array_downcast_fn!(as_union_array, UnionArray);
690array_downcast_fn!(as_map_array, MapArray);
691
692/// Downcasts a `dyn Array` to a concrete type
693///
694/// ```
695/// # use arrow_array::{BooleanArray, Int32Array, RecordBatch, StringArray};
696/// # use arrow_array::cast::downcast_array;
697/// struct ConcreteBatch {
698///     col1: Int32Array,
699///     col2: BooleanArray,
700///     col3: StringArray,
701/// }
702///
703/// impl ConcreteBatch {
704///     fn new(batch: &RecordBatch) -> Self {
705///         Self {
706///             col1: downcast_array(batch.column(0).as_ref()),
707///             col2: downcast_array(batch.column(1).as_ref()),
708///             col3: downcast_array(batch.column(2).as_ref()),
709///         }
710///     }
711/// }
712/// ```
713///
714/// # Panics
715///
716/// Panics if array is not of the correct data type
717pub fn downcast_array<T>(array: &dyn Array) -> T
718where
719    T: From<ArrayData>,
720{
721    T::from(array.to_data())
722}
723
724mod private {
725    pub trait Sealed {}
726}
727
728/// An extension trait for `dyn Array` that provides ergonomic downcasting
729///
730/// ```
731/// # use std::sync::Arc;
732/// # use arrow_array::{ArrayRef, Int32Array};
733/// # use arrow_array::cast::AsArray;
734/// # use arrow_array::types::Int32Type;
735/// let col = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef;
736/// assert_eq!(col.as_primitive::<Int32Type>().values(), &[1, 2, 3]);
737/// ```
738pub trait AsArray: private::Sealed {
739    /// Downcast this to a [`BooleanArray`] returning `None` if not possible
740    fn as_boolean_opt(&self) -> Option<&BooleanArray>;
741
742    /// Downcast this to a [`BooleanArray`] panicking if not possible
743    fn as_boolean(&self) -> &BooleanArray {
744        self.as_boolean_opt().expect("boolean array")
745    }
746
747    /// Downcast this to a [`PrimitiveArray`] returning `None` if not possible
748    fn as_primitive_opt<T: ArrowPrimitiveType>(&self) -> Option<&PrimitiveArray<T>>;
749
750    /// Downcast this to a [`PrimitiveArray`] panicking if not possible
751    fn as_primitive<T: ArrowPrimitiveType>(&self) -> &PrimitiveArray<T> {
752        self.as_primitive_opt().expect("primitive array")
753    }
754
755    /// Downcast this to a [`GenericByteArray`] returning `None` if not possible
756    fn as_bytes_opt<T: ByteArrayType>(&self) -> Option<&GenericByteArray<T>>;
757
758    /// Downcast this to a [`GenericByteArray`] panicking if not possible
759    fn as_bytes<T: ByteArrayType>(&self) -> &GenericByteArray<T> {
760        self.as_bytes_opt().expect("byte array")
761    }
762
763    /// Downcast this to a [`GenericStringArray`] returning `None` if not possible
764    fn as_string_opt<O: OffsetSizeTrait>(&self) -> Option<&GenericStringArray<O>> {
765        self.as_bytes_opt()
766    }
767
768    /// Downcast this to a [`GenericStringArray`] panicking if not possible
769    fn as_string<O: OffsetSizeTrait>(&self) -> &GenericStringArray<O> {
770        self.as_bytes_opt().expect("string array")
771    }
772
773    /// Downcast this to a [`GenericBinaryArray`] returning `None` if not possible
774    fn as_binary_opt<O: OffsetSizeTrait>(&self) -> Option<&GenericBinaryArray<O>> {
775        self.as_bytes_opt()
776    }
777
778    /// Downcast this to a [`GenericBinaryArray`] panicking if not possible
779    fn as_binary<O: OffsetSizeTrait>(&self) -> &GenericBinaryArray<O> {
780        self.as_bytes_opt().expect("binary array")
781    }
782
783    /// Downcast this to a [`StringViewArray`] returning `None` if not possible
784    fn as_string_view_opt(&self) -> Option<&StringViewArray> {
785        self.as_byte_view_opt()
786    }
787
788    /// Downcast this to a [`StringViewArray`] panicking if not possible
789    fn as_string_view(&self) -> &StringViewArray {
790        self.as_byte_view_opt().expect("string view array")
791    }
792
793    /// Downcast this to a [`BinaryViewArray`] returning `None` if not possible
794    fn as_binary_view_opt(&self) -> Option<&BinaryViewArray> {
795        self.as_byte_view_opt()
796    }
797
798    /// Downcast this to a [`BinaryViewArray`] panicking if not possible
799    fn as_binary_view(&self) -> &BinaryViewArray {
800        self.as_byte_view_opt().expect("binary view array")
801    }
802
803    /// Downcast this to a [`GenericByteViewArray`] returning `None` if not possible
804    fn as_byte_view_opt<T: ByteViewType>(&self) -> Option<&GenericByteViewArray<T>>;
805
806    /// Downcast this to a [`GenericByteViewArray`] panicking if not possible
807    fn as_byte_view<T: ByteViewType>(&self) -> &GenericByteViewArray<T> {
808        self.as_byte_view_opt().expect("byte view array")
809    }
810
811    /// Downcast this to a [`StructArray`] returning `None` if not possible
812    fn as_struct_opt(&self) -> Option<&StructArray>;
813
814    /// Downcast this to a [`StructArray`] panicking if not possible
815    fn as_struct(&self) -> &StructArray {
816        self.as_struct_opt().expect("struct array")
817    }
818
819    /// Downcast this to a [`UnionArray`] returning `None` if not possible
820    fn as_union_opt(&self) -> Option<&UnionArray>;
821
822    /// Downcast this to a [`UnionArray`] panicking if not possible
823    fn as_union(&self) -> &UnionArray {
824        self.as_union_opt().expect("union array")
825    }
826
827    /// Downcast this to a [`GenericListArray`] returning `None` if not possible
828    fn as_list_opt<O: OffsetSizeTrait>(&self) -> Option<&GenericListArray<O>>;
829
830    /// Downcast this to a [`GenericListArray`] panicking if not possible
831    fn as_list<O: OffsetSizeTrait>(&self) -> &GenericListArray<O> {
832        self.as_list_opt().expect("list array")
833    }
834
835    /// Downcast this to a [`GenericListViewArray`] returning `None` if not possible
836    fn as_list_view_opt<O: OffsetSizeTrait>(&self) -> Option<&GenericListViewArray<O>>;
837
838    /// Downcast this to a [`GenericListViewArray`] panicking if not possible
839    fn as_list_view<O: OffsetSizeTrait>(&self) -> &GenericListViewArray<O> {
840        self.as_list_view_opt().expect("list view array")
841    }
842
843    /// Downcast this to a [`FixedSizeBinaryArray`] returning `None` if not possible
844    fn as_fixed_size_binary_opt(&self) -> Option<&FixedSizeBinaryArray>;
845
846    /// Downcast this to a [`FixedSizeBinaryArray`] panicking if not possible
847    fn as_fixed_size_binary(&self) -> &FixedSizeBinaryArray {
848        self.as_fixed_size_binary_opt()
849            .expect("fixed size binary array")
850    }
851
852    /// Downcast this to a [`FixedSizeListArray`] returning `None` if not possible
853    fn as_fixed_size_list_opt(&self) -> Option<&FixedSizeListArray>;
854
855    /// Downcast this to a [`FixedSizeListArray`] panicking if not possible
856    fn as_fixed_size_list(&self) -> &FixedSizeListArray {
857        self.as_fixed_size_list_opt()
858            .expect("fixed size list array")
859    }
860
861    /// Downcast this to a [`MapArray`] returning `None` if not possible
862    fn as_map_opt(&self) -> Option<&MapArray>;
863
864    /// Downcast this to a [`MapArray`] panicking if not possible
865    fn as_map(&self) -> &MapArray {
866        self.as_map_opt().expect("map array")
867    }
868
869    /// Downcast this to a [`DictionaryArray`] returning `None` if not possible
870    fn as_dictionary_opt<K: ArrowDictionaryKeyType>(&self) -> Option<&DictionaryArray<K>>;
871
872    /// Downcast this to a [`DictionaryArray`] panicking if not possible
873    fn as_dictionary<K: ArrowDictionaryKeyType>(&self) -> &DictionaryArray<K> {
874        self.as_dictionary_opt().expect("dictionary array")
875    }
876
877    /// Downcasts this to a [`AnyDictionaryArray`] returning `None` if not possible
878    fn as_any_dictionary_opt(&self) -> Option<&dyn AnyDictionaryArray>;
879
880    /// Downcasts this to a [`AnyDictionaryArray`] panicking if not possible
881    fn as_any_dictionary(&self) -> &dyn AnyDictionaryArray {
882        self.as_any_dictionary_opt().expect("any dictionary array")
883    }
884}
885
886impl private::Sealed for dyn Array + '_ {}
887impl AsArray for dyn Array + '_ {
888    fn as_boolean_opt(&self) -> Option<&BooleanArray> {
889        self.as_any().downcast_ref()
890    }
891
892    fn as_primitive_opt<T: ArrowPrimitiveType>(&self) -> Option<&PrimitiveArray<T>> {
893        self.as_any().downcast_ref()
894    }
895
896    fn as_bytes_opt<T: ByteArrayType>(&self) -> Option<&GenericByteArray<T>> {
897        self.as_any().downcast_ref()
898    }
899
900    fn as_byte_view_opt<T: ByteViewType>(&self) -> Option<&GenericByteViewArray<T>> {
901        self.as_any().downcast_ref()
902    }
903
904    fn as_struct_opt(&self) -> Option<&StructArray> {
905        self.as_any().downcast_ref()
906    }
907
908    fn as_union_opt(&self) -> Option<&UnionArray> {
909        self.as_any().downcast_ref()
910    }
911
912    fn as_list_opt<O: OffsetSizeTrait>(&self) -> Option<&GenericListArray<O>> {
913        self.as_any().downcast_ref()
914    }
915
916    fn as_list_view_opt<O: OffsetSizeTrait>(&self) -> Option<&GenericListViewArray<O>> {
917        self.as_any().downcast_ref()
918    }
919
920    fn as_fixed_size_binary_opt(&self) -> Option<&FixedSizeBinaryArray> {
921        self.as_any().downcast_ref()
922    }
923
924    fn as_fixed_size_list_opt(&self) -> Option<&FixedSizeListArray> {
925        self.as_any().downcast_ref()
926    }
927
928    fn as_map_opt(&self) -> Option<&MapArray> {
929        self.as_any().downcast_ref()
930    }
931
932    fn as_dictionary_opt<K: ArrowDictionaryKeyType>(&self) -> Option<&DictionaryArray<K>> {
933        self.as_any().downcast_ref()
934    }
935
936    fn as_any_dictionary_opt(&self) -> Option<&dyn AnyDictionaryArray> {
937        let array = self;
938        downcast_dictionary_array! {
939            array => Some(array),
940            _ => None
941        }
942    }
943}
944
945impl private::Sealed for ArrayRef {}
946impl AsArray for ArrayRef {
947    fn as_boolean_opt(&self) -> Option<&BooleanArray> {
948        self.as_ref().as_boolean_opt()
949    }
950
951    fn as_primitive_opt<T: ArrowPrimitiveType>(&self) -> Option<&PrimitiveArray<T>> {
952        self.as_ref().as_primitive_opt()
953    }
954
955    fn as_bytes_opt<T: ByteArrayType>(&self) -> Option<&GenericByteArray<T>> {
956        self.as_ref().as_bytes_opt()
957    }
958
959    fn as_byte_view_opt<T: ByteViewType>(&self) -> Option<&GenericByteViewArray<T>> {
960        self.as_ref().as_byte_view_opt()
961    }
962
963    fn as_struct_opt(&self) -> Option<&StructArray> {
964        self.as_ref().as_struct_opt()
965    }
966
967    fn as_union_opt(&self) -> Option<&UnionArray> {
968        self.as_any().downcast_ref()
969    }
970
971    fn as_list_opt<O: OffsetSizeTrait>(&self) -> Option<&GenericListArray<O>> {
972        self.as_ref().as_list_opt()
973    }
974
975    fn as_list_view_opt<O: OffsetSizeTrait>(&self) -> Option<&GenericListViewArray<O>> {
976        self.as_ref().as_list_view_opt()
977    }
978
979    fn as_fixed_size_binary_opt(&self) -> Option<&FixedSizeBinaryArray> {
980        self.as_ref().as_fixed_size_binary_opt()
981    }
982
983    fn as_fixed_size_list_opt(&self) -> Option<&FixedSizeListArray> {
984        self.as_ref().as_fixed_size_list_opt()
985    }
986
987    fn as_map_opt(&self) -> Option<&MapArray> {
988        self.as_any().downcast_ref()
989    }
990
991    fn as_dictionary_opt<K: ArrowDictionaryKeyType>(&self) -> Option<&DictionaryArray<K>> {
992        self.as_ref().as_dictionary_opt()
993    }
994
995    fn as_any_dictionary_opt(&self) -> Option<&dyn AnyDictionaryArray> {
996        self.as_ref().as_any_dictionary_opt()
997    }
998}
999
1000#[cfg(test)]
1001mod tests {
1002    use arrow_buffer::i256;
1003    use std::sync::Arc;
1004
1005    use super::*;
1006
1007    #[test]
1008    fn test_as_primitive_array_ref() {
1009        let array: Int32Array = vec![1, 2, 3].into_iter().map(Some).collect();
1010        assert!(!as_primitive_array::<Int32Type>(&array).is_empty());
1011
1012        // should also work when wrapped in an Arc
1013        let array: ArrayRef = Arc::new(array);
1014        assert!(!as_primitive_array::<Int32Type>(&array).is_empty());
1015    }
1016
1017    #[test]
1018    fn test_as_string_array_ref() {
1019        let array: StringArray = vec!["foo", "bar"].into_iter().map(Some).collect();
1020        assert!(!as_string_array(&array).is_empty());
1021
1022        // should also work when wrapped in an Arc
1023        let array: ArrayRef = Arc::new(array);
1024        assert!(!as_string_array(&array).is_empty())
1025    }
1026
1027    #[test]
1028    fn test_decimal128array() {
1029        let a = Decimal128Array::from_iter_values([1, 2, 4, 5]);
1030        assert!(!as_primitive_array::<Decimal128Type>(&a).is_empty());
1031    }
1032
1033    #[test]
1034    fn test_decimal256array() {
1035        let a = Decimal256Array::from_iter_values([1, 2, 4, 5].into_iter().map(i256::from_i128));
1036        assert!(!as_primitive_array::<Decimal256Type>(&a).is_empty());
1037    }
1038}
arrow_array/cast.rs

arrow_array/
cast.rs
