Skip to main content

runlength_arrow/
numeric.rs

1//! Dispatch from an Arrow data type to the core crate's numeric types.
2
3use arrow_schema::DataType;
4
5/// Runs `$body` with `$T` bound to the Arrow type of `$values`, for the
6/// numeric types coverage supports.
7macro_rules! with_numeric_type {
8    ($values:expr, |$T:ident| $body:expr) => {
9        match $values.data_type() {
10            ::arrow_schema::DataType::Int8 => {
11                type $T = ::arrow_array::types::Int8Type;
12                $body
13            }
14            ::arrow_schema::DataType::Int16 => {
15                type $T = ::arrow_array::types::Int16Type;
16                $body
17            }
18            ::arrow_schema::DataType::Int32 => {
19                type $T = ::arrow_array::types::Int32Type;
20                $body
21            }
22            ::arrow_schema::DataType::Int64 => {
23                type $T = ::arrow_array::types::Int64Type;
24                $body
25            }
26            ::arrow_schema::DataType::UInt8 => {
27                type $T = ::arrow_array::types::UInt8Type;
28                $body
29            }
30            ::arrow_schema::DataType::UInt16 => {
31                type $T = ::arrow_array::types::UInt16Type;
32                $body
33            }
34            ::arrow_schema::DataType::UInt32 => {
35                type $T = ::arrow_array::types::UInt32Type;
36                $body
37            }
38            ::arrow_schema::DataType::UInt64 => {
39                type $T = ::arrow_array::types::UInt64Type;
40                $body
41            }
42            ::arrow_schema::DataType::Float32 => {
43                type $T = ::arrow_array::types::Float32Type;
44                $body
45            }
46            ::arrow_schema::DataType::Float64 => {
47                type $T = ::arrow_array::types::Float64Type;
48                $body
49            }
50            other => Err($crate::Error::UnsupportedType(other.clone())),
51        }
52    };
53}
54
55pub(crate) use with_numeric_type;
56
57/// The type two numeric types are converted to before combining them:
58/// themselves if equal, `Float64` if either is a float, `Int64` otherwise.
59/// `None` if either is not numeric.
60pub fn common_numeric_type(a: &DataType, b: &DataType) -> Option<DataType> {
61    if !a.is_numeric() || !b.is_numeric() {
62        return None;
63    }
64    if a == b {
65        return Some(a.clone());
66    }
67    if a.is_floating() || b.is_floating() {
68        Some(DataType::Float64)
69    } else {
70        Some(DataType::Int64)
71    }
72}