arrow_array/
cast.rs

Help
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/// Force downcast of an Array, such as an ArrayRef to Decimal128Array, panic’ing on failure.
693#[deprecated(note = "please use `as_primitive_array::<Decimal128Type>` instead")]
694pub fn as_decimal_array(arr: &dyn Array) -> &PrimitiveArray<Decimal128Type> {
695    as_primitive_array::<Decimal128Type>(arr)
696}
697
698/// Downcasts a `dyn Array` to a concrete type
699///
700/// ```
701/// # use arrow_array::{BooleanArray, Int32Array, RecordBatch, StringArray};
702/// # use arrow_array::cast::downcast_array;
703/// struct ConcreteBatch {
704///     col1: Int32Array,
705///     col2: BooleanArray,
706///     col3: StringArray,
707/// }
708///
709/// impl ConcreteBatch {
710///     fn new(batch: &RecordBatch) -> Self {
711///         Self {
712///             col1: downcast_array(batch.column(0).as_ref()),
713///             col2: downcast_array(batch.column(1).as_ref()),
714///             col3: downcast_array(batch.column(2).as_ref()),
715///         }
716///     }
717/// }
718/// ```
719///
720/// # Panics
721///
722/// Panics if array is not of the correct data type
723pub fn downcast_array<T>(array: &dyn Array) -> T
724where
725    T: From<ArrayData>,
726{
727    T::from(array.to_data())
728}
729
730mod private {
731    pub trait Sealed {}
732}
733
734/// An extension trait for `dyn Array` that provides ergonomic downcasting
735///
736/// ```
737/// # use std::sync::Arc;
738/// # use arrow_array::{ArrayRef, Int32Array};
739/// # use arrow_array::cast::AsArray;
740/// # use arrow_array::types::Int32Type;
741/// let col = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef;
742/// assert_eq!(col.as_primitive::<Int32Type>().values(), &[1, 2, 3]);
743/// ```
744pub trait AsArray: private::Sealed {
745    /// Downcast this to a [`BooleanArray`] returning `None` if not possible
746    fn as_boolean_opt(&self) -> Option<&BooleanArray>;
747
748    /// Downcast this to a [`BooleanArray`] panicking if not possible
749    fn as_boolean(&self) -> &BooleanArray {
750        self.as_boolean_opt().expect("boolean array")
751    }
752
753    /// Downcast this to a [`PrimitiveArray`] returning `None` if not possible
754    fn as_primitive_opt<T: ArrowPrimitiveType>(&self) -> Option<&PrimitiveArray<T>>;
755
756    /// Downcast this to a [`PrimitiveArray`] panicking if not possible
757    fn as_primitive<T: ArrowPrimitiveType>(&self) -> &PrimitiveArray<T> {
758        self.as_primitive_opt().expect("primitive array")
759    }
760
761    /// Downcast this to a [`GenericByteArray`] returning `None` if not possible
762    fn as_bytes_opt<T: ByteArrayType>(&self) -> Option<&GenericByteArray<T>>;
763
764    /// Downcast this to a [`GenericByteArray`] panicking if not possible
765    fn as_bytes<T: ByteArrayType>(&self) -> &GenericByteArray<T> {
766        self.as_bytes_opt().expect("byte array")
767    }
768
769    /// Downcast this to a [`GenericStringArray`] returning `None` if not possible
770    fn as_string_opt<O: OffsetSizeTrait>(&self) -> Option<&GenericStringArray<O>> {
771        self.as_bytes_opt()
772    }
773
774    /// Downcast this to a [`GenericStringArray`] panicking if not possible
775    fn as_string<O: OffsetSizeTrait>(&self) -> &GenericStringArray<O> {
776        self.as_bytes_opt().expect("string array")
777    }
778
779    /// Downcast this to a [`GenericBinaryArray`] returning `None` if not possible
780    fn as_binary_opt<O: OffsetSizeTrait>(&self) -> Option<&GenericBinaryArray<O>> {
781        self.as_bytes_opt()
782    }
783
784    /// Downcast this to a [`GenericBinaryArray`] panicking if not possible
785    fn as_binary<O: OffsetSizeTrait>(&self) -> &GenericBinaryArray<O> {
786        self.as_bytes_opt().expect("binary array")
787    }
788
789    /// Downcast this to a [`StringViewArray`] returning `None` if not possible
790    fn as_string_view_opt(&self) -> Option<&StringViewArray> {
791        self.as_byte_view_opt()
792    }
793
794    /// Downcast this to a [`StringViewArray`] panicking if not possible
795    fn as_string_view(&self) -> &StringViewArray {
796        self.as_byte_view_opt().expect("string view array")
797    }
798
799    /// Downcast this to a [`BinaryViewArray`] returning `None` if not possible
800    fn as_binary_view_opt(&self) -> Option<&BinaryViewArray> {
801        self.as_byte_view_opt()
802    }
803
804    /// Downcast this to a [`BinaryViewArray`] panicking if not possible
805    fn as_binary_view(&self) -> &BinaryViewArray {
806        self.as_byte_view_opt().expect("binary view array")
807    }
808
809    /// Downcast this to a [`GenericByteViewArray`] returning `None` if not possible
810    fn as_byte_view_opt<T: ByteViewType>(&self) -> Option<&GenericByteViewArray<T>>;
811
812    /// Downcast this to a [`GenericByteViewArray`] panicking if not possible
813    fn as_byte_view<T: ByteViewType>(&self) -> &GenericByteViewArray<T> {
814        self.as_byte_view_opt().expect("byte view array")
815    }
816
817    /// Downcast this to a [`StructArray`] returning `None` if not possible
818    fn as_struct_opt(&self) -> Option<&StructArray>;
819
820    /// Downcast this to a [`StructArray`] panicking if not possible
821    fn as_struct(&self) -> &StructArray {
822        self.as_struct_opt().expect("struct array")
823    }
824
825    /// Downcast this to a [`UnionArray`] returning `None` if not possible
826    fn as_union_opt(&self) -> Option<&UnionArray>;
827
828    /// Downcast this to a [`UnionArray`] panicking if not possible
829    fn as_union(&self) -> &UnionArray {
830        self.as_union_opt().expect("union array")
831    }
832
833    /// Downcast this to a [`GenericListArray`] returning `None` if not possible
834    fn as_list_opt<O: OffsetSizeTrait>(&self) -> Option<&GenericListArray<O>>;
835
836    /// Downcast this to a [`GenericListArray`] panicking if not possible
837    fn as_list<O: OffsetSizeTrait>(&self) -> &GenericListArray<O> {
838        self.as_list_opt().expect("list array")
839    }
840
841    /// Downcast this to a [`FixedSizeBinaryArray`] returning `None` if not possible
842    fn as_fixed_size_binary_opt(&self) -> Option<&FixedSizeBinaryArray>;
843
844    /// Downcast this to a [`FixedSizeBinaryArray`] panicking if not possible
845    fn as_fixed_size_binary(&self) -> &FixedSizeBinaryArray {
846        self.as_fixed_size_binary_opt()
847            .expect("fixed size binary array")
848    }
849
850    /// Downcast this to a [`FixedSizeListArray`] returning `None` if not possible
851    fn as_fixed_size_list_opt(&self) -> Option<&FixedSizeListArray>;
852
853    /// Downcast this to a [`FixedSizeListArray`] panicking if not possible
854    fn as_fixed_size_list(&self) -> &FixedSizeListArray {
855        self.as_fixed_size_list_opt()
856            .expect("fixed size list array")
857    }
858
859    /// Downcast this to a [`MapArray`] returning `None` if not possible
860    fn as_map_opt(&self) -> Option<&MapArray>;
861
862    /// Downcast this to a [`MapArray`] panicking if not possible
863    fn as_map(&self) -> &MapArray {
864        self.as_map_opt().expect("map array")
865    }
866
867    /// Downcast this to a [`DictionaryArray`] returning `None` if not possible
868    fn as_dictionary_opt<K: ArrowDictionaryKeyType>(&self) -> Option<&DictionaryArray<K>>;
869
870    /// Downcast this to a [`DictionaryArray`] panicking if not possible
871    fn as_dictionary<K: ArrowDictionaryKeyType>(&self) -> &DictionaryArray<K> {
872        self.as_dictionary_opt().expect("dictionary array")
873    }
874
875    /// Downcasts this to a [`AnyDictionaryArray`] returning `None` if not possible
876    fn as_any_dictionary_opt(&self) -> Option<&dyn AnyDictionaryArray>;
877
878    /// Downcasts this to a [`AnyDictionaryArray`] panicking if not possible
879    fn as_any_dictionary(&self) -> &dyn AnyDictionaryArray {
880        self.as_any_dictionary_opt().expect("any dictionary array")
881    }
882}
883
884impl private::Sealed for dyn Array + '_ {}
885impl AsArray for dyn Array + '_ {
886    fn as_boolean_opt(&self) -> Option<&BooleanArray> {
887        self.as_any().downcast_ref()
888    }
889
890    fn as_primitive_opt<T: ArrowPrimitiveType>(&self) -> Option<&PrimitiveArray<T>> {
891        self.as_any().downcast_ref()
892    }
893
894    fn as_bytes_opt<T: ByteArrayType>(&self) -> Option<&GenericByteArray<T>> {
895        self.as_any().downcast_ref()
896    }
897
898    fn as_byte_view_opt<T: ByteViewType>(&self) -> Option<&GenericByteViewArray<T>> {
899        self.as_any().downcast_ref()
900    }
901
902    fn as_struct_opt(&self) -> Option<&StructArray> {
903        self.as_any().downcast_ref()
904    }
905
906    fn as_union_opt(&self) -> Option<&UnionArray> {
907        self.as_any().downcast_ref()
908    }
909
910    fn as_list_opt<O: OffsetSizeTrait>(&self) -> Option<&GenericListArray<O>> {
911        self.as_any().downcast_ref()
912    }
913
914    fn as_fixed_size_binary_opt(&self) -> Option<&FixedSizeBinaryArray> {
915        self.as_any().downcast_ref()
916    }
917
918    fn as_fixed_size_list_opt(&self) -> Option<&FixedSizeListArray> {
919        self.as_any().downcast_ref()
920    }
921
922    fn as_map_opt(&self) -> Option<&MapArray> {
923        self.as_any().downcast_ref()
924    }
925
926    fn as_dictionary_opt<K: ArrowDictionaryKeyType>(&self) -> Option<&DictionaryArray<K>> {
927        self.as_any().downcast_ref()
928    }
929
930    fn as_any_dictionary_opt(&self) -> Option<&dyn AnyDictionaryArray> {
931        let array = self;
932        downcast_dictionary_array! {
933            array => Some(array),
934            _ => None
935        }
936    }
937}
938
939impl private::Sealed for ArrayRef {}
940impl AsArray for ArrayRef {
941    fn as_boolean_opt(&self) -> Option<&BooleanArray> {
942        self.as_ref().as_boolean_opt()
943    }
944
945    fn as_primitive_opt<T: ArrowPrimitiveType>(&self) -> Option<&PrimitiveArray<T>> {
946        self.as_ref().as_primitive_opt()
947    }
948
949    fn as_bytes_opt<T: ByteArrayType>(&self) -> Option<&GenericByteArray<T>> {
950        self.as_ref().as_bytes_opt()
951    }
952
953    fn as_byte_view_opt<T: ByteViewType>(&self) -> Option<&GenericByteViewArray<T>> {
954        self.as_ref().as_byte_view_opt()
955    }
956
957    fn as_struct_opt(&self) -> Option<&StructArray> {
958        self.as_ref().as_struct_opt()
959    }
960
961    fn as_union_opt(&self) -> Option<&UnionArray> {
962        self.as_any().downcast_ref()
963    }
964
965    fn as_list_opt<O: OffsetSizeTrait>(&self) -> Option<&GenericListArray<O>> {
966        self.as_ref().as_list_opt()
967    }
968
969    fn as_fixed_size_binary_opt(&self) -> Option<&FixedSizeBinaryArray> {
970        self.as_ref().as_fixed_size_binary_opt()
971    }
972
973    fn as_fixed_size_list_opt(&self) -> Option<&FixedSizeListArray> {
974        self.as_ref().as_fixed_size_list_opt()
975    }
976
977    fn as_map_opt(&self) -> Option<&MapArray> {
978        self.as_any().downcast_ref()
979    }
980
981    fn as_dictionary_opt<K: ArrowDictionaryKeyType>(&self) -> Option<&DictionaryArray<K>> {
982        self.as_ref().as_dictionary_opt()
983    }
984
985    fn as_any_dictionary_opt(&self) -> Option<&dyn AnyDictionaryArray> {
986        self.as_ref().as_any_dictionary_opt()
987    }
988}
989
990#[cfg(test)]
991mod tests {
992    use arrow_buffer::i256;
993    use std::sync::Arc;
994
995    use super::*;
996
997    #[test]
998    fn test_as_primitive_array_ref() {
999        let array: Int32Array = vec![1, 2, 3].into_iter().map(Some).collect();
1000        assert!(!as_primitive_array::<Int32Type>(&array).is_empty());
1001
1002        // should also work when wrapped in an Arc
1003        let array: ArrayRef = Arc::new(array);
1004        assert!(!as_primitive_array::<Int32Type>(&array).is_empty());
1005    }
1006
1007    #[test]
1008    fn test_as_string_array_ref() {
1009        let array: StringArray = vec!["foo", "bar"].into_iter().map(Some).collect();
1010        assert!(!as_string_array(&array).is_empty());
1011
1012        // should also work when wrapped in an Arc
1013        let array: ArrayRef = Arc::new(array);
1014        assert!(!as_string_array(&array).is_empty())
1015    }
1016
1017    #[test]
1018    fn test_decimal128array() {
1019        let a = Decimal128Array::from_iter_values([1, 2, 4, 5]);
1020        assert!(!as_primitive_array::<Decimal128Type>(&a).is_empty());
1021    }
1022
1023    #[test]
1024    fn test_decimal256array() {
1025        let a = Decimal256Array::from_iter_values([1, 2, 4, 5].into_iter().map(i256::from_i128));
1026        assert!(!as_primitive_array::<Decimal256Type>(&a).is_empty());
1027    }
1028}
arrow_array/cast.rs

arrow_array/
cast.rs
