arrow_cast/cast/
numeric.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
3// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements
4
5//! Cast support for numeric and boolean arrays.
6
7use crate::cast::*;
8
9/// Convert Array into a PrimitiveArray of type, and apply numeric cast
10pub(crate) fn cast_numeric_arrays<FROM, TO>(
11    from: &dyn Array,
12    cast_options: &CastOptions,
13) -> Result<ArrayRef, ArrowError>
14where
15    FROM: ArrowPrimitiveType,
16    TO: ArrowPrimitiveType,
17    FROM::Native: NumCast + NumericNative,
18    TO::Native: NumCast + NumericNative,
19{
20    let from = from.as_primitive::<FROM>();
21    let array: PrimitiveArray<TO> = if const {
22        is_infallible_numeric_cast(
23            <FROM::Native as NumericNative>::KIND,
24            <TO::Native as NumericNative>::KIND,
25        )
26    } {
27        // This cast cannot fail, so the fastest kernel, `unary`, can be used.
28        from.unary(|v| num_cast(v).expect("numeric cast is infallible"))
29    } else if cast_options.safe {
30        // If the value can't be cast to the `TO::Native`, return null
31        from.unary_opt(num_cast)
32    } else {
33        // If the value can't be cast to the `TO::Native`, return error
34        from.try_unary(|v| {
35            num_cast(v).ok_or_else(|| {
36                ArrowError::CastError(format!("Can't cast value {v:?} to type {}", TO::DATA_TYPE))
37            })
38        })?
39    };
40    Ok(Arc::new(array))
41}
42
43/// Properties of a numeric native type that decide whether [`num_cast`] to
44/// another numeric type can fail.
45#[derive(Clone, Copy)]
46pub(crate) struct NumericKind {
47    float: bool,
48    signed: bool,
49    bits: u32,
50}
51
52/// A native type with a [`NumericKind`].
53pub(crate) trait NumericNative {
54    const KIND: NumericKind;
55}
56
57macro_rules! numeric_native {
58    ($($t:ty => $float:literal, $signed:literal;)*) => {$(
59        impl NumericNative for $t {
60            const KIND: NumericKind = NumericKind {
61                float: $float,
62                signed: $signed,
63                bits: 8 * std::mem::size_of::<$t>() as u32,
64            };
65        }
66    )*};
67}
68
69numeric_native! {
70    u8 => false, false;
71    u16 => false, false;
72    u32 => false, false;
73    u64 => false, false;
74    i8 => false, true;
75    i16 => false, true;
76    i32 => false, true;
77    i64 => false, true;
78    half::f16 => true, true;
79    f32 => true, true;
80    f64 => true, true;
81}
82
83/// Returns true if [`num_cast`] succeeds for every value of `from` when casting
84/// to `to`: every cast to a float, and integer casts whose target holds every
85/// source value.
86///
87/// This must hold for every bit pattern of the source type, not only for the
88/// valid values of a particular array, because [`PrimitiveArray::unary`] applies
89/// the conversion to null slots as well, and their contents are arbitrary.
90pub(crate) const fn is_infallible_numeric_cast(from: NumericKind, to: NumericKind) -> bool {
91    if to.float {
92        true
93    } else if from.float {
94        false
95    } else if from.signed == to.signed {
96        to.bits >= from.bits
97    } else {
98        // Only an unsigned source fits a strictly wider signed target.
99        !from.signed && to.bits > from.bits
100    }
101}
102
103/// Natural cast between numeric types
104/// Return None if the input `value` can't be casted to type `O`.
105#[inline]
106pub fn num_cast<I, O>(value: I) -> Option<O>
107where
108    I: NumCast,
109    O: NumCast,
110{
111    num_traits::cast::cast::<I, O>(value)
112}
113
114/// Cast numeric types to Boolean
115///
116/// Any zero value returns `false` while non-zero returns `true`
117pub(crate) fn cast_numeric_to_bool<FROM>(from: &dyn Array) -> Result<ArrayRef, ArrowError>
118where
119    FROM: ArrowPrimitiveType,
120{
121    Ok(Arc::new(BooleanArray::from_unary(
122        from.as_primitive::<FROM>(),
123        cast_num_to_bool,
124    )))
125}
126
127/// Cast numeric types to boolean
128#[inline]
129pub fn cast_num_to_bool<I>(value: I) -> bool
130where
131    I: Default + PartialEq,
132{
133    value != I::default()
134}
135
136/// Cast Boolean types to numeric
137///
138/// `false` returns 0 while `true` returns 1
139pub(crate) fn cast_bool_to_numeric<TO>(
140    from: &dyn Array,
141    cast_options: &CastOptions,
142) -> Result<ArrayRef, ArrowError>
143where
144    TO: ArrowPrimitiveType,
145{
146    Ok(Arc::new(bool_to_numeric_cast::<TO>(
147        from.as_any().downcast_ref::<BooleanArray>().unwrap(),
148        cast_options,
149    )))
150}
151
152fn bool_to_numeric_cast<T>(from: &BooleanArray, _cast_options: &CastOptions) -> PrimitiveArray<T>
153where
154    T: ArrowPrimitiveType,
155{
156    let to_numeric = |bits: u64, i: usize| {
157        if bits & (1 << i) != 0 {
158            T::Native::ONE
159        } else {
160            T::Native::ZERO
161        }
162    };
163    // Unpack a 64-bit word at a time: the fixed-size inner loop helps the
164    // compiler vectorize the conversion.
165    let chunks = from.values().bit_chunks();
166    let mut values = Vec::with_capacity(from.len());
167    for bits in &chunks {
168        values.extend((0..64).map(|i| to_numeric(bits, i)));
169    }
170    let bits = chunks.remainder_bits();
171    values.extend((0..chunks.remainder_len()).map(|i| to_numeric(bits, i)));
172    PrimitiveArray::new(values.into(), from.nulls().cloned())
173}
174
175/// Cast single bool value to numeric value.
176#[inline]
177pub fn single_bool_to_numeric<O>(value: bool) -> Option<O>
178where
179    O: num_traits::NumCast + Default,
180{
181    if value {
182        // a workaround to cast a primitive to type O, infallible
183        num_traits::cast::cast(1)
184    } else {
185        Some(O::default())
186    }
187}
arrow_cast/cast/numeric.rs

arrow_cast/cast/
numeric.rs
