CoCalc -- dyn_array.rs

GitHub Repository: pola-rs/polars
Path: blob/main/crates/polars-compute/src/min_max/dyn_array.rs
⁶⁹³⁹ views
1
use arrow::array::{
2
    Array, BinaryArray, BinaryViewArray, BooleanArray, PrimitiveArray, Utf8Array, Utf8ViewArray,
3
};
4
use arrow::scalar::{BinaryScalar, BinaryViewScalar, BooleanScalar, PrimitiveScalar, Scalar};
5

6
use crate::min_max::MinMaxKernel;
7

8
macro_rules! call_op {
9
    ($T:ty, $scalar:ty, $arr:expr, $op:path) => {{
10
        let arr: &$T = $arr.as_any().downcast_ref().unwrap();
11
        $op(arr).map(|v| Box::new(<$scalar>::new(Some(v))) as Box<dyn Scalar>)
12
    }};
13
    (dt: $T:ty, $scalar:ty, $arr:expr, $op:path) => {{
14
        let arr: &$T = $arr.as_any().downcast_ref().unwrap();
15
        $op(arr).map(|v| Box::new(<$scalar>::new(arr.dtype().clone(), Some(v))) as Box<dyn Scalar>)
16
    }};
17
    ($T:ty, $scalar:ty, $arr:expr, $op:path, ret_two) => {{
18
        let arr: &$T = $arr.as_any().downcast_ref().unwrap();
19
        $op(arr).map(|(l, r)| {
20
            (
21
                Box::new(<$scalar>::new(Some(l))) as Box<dyn Scalar>,
22
                Box::new(<$scalar>::new(Some(r))) as Box<dyn Scalar>,
23
            )
24
        })
25
    }};
26
    (dt: $T:ty, $scalar:ty, $arr:expr, $op:path, ret_two) => {{
27
        let arr: &$T = $arr.as_any().downcast_ref().unwrap();
28
        $op(arr).map(|(l, r)| {
29
            (
30
                Box::new(<$scalar>::new(arr.dtype().clone(), Some(l))) as Box<dyn Scalar>,
31
                Box::new(<$scalar>::new(arr.dtype().clone(), Some(r))) as Box<dyn Scalar>,
32
            )
33
        })
34
    }};
35
}
36

37
macro_rules! call {
38
    ($arr:expr, $op:path$(, $variant:ident)?) => {{
39
        let arr = $arr;
40

41
        use arrow::datatypes::{PhysicalType as PH, PrimitiveType as PR};
42
        use PrimitiveArray as PArr;
43
        use PrimitiveScalar as PScalar;
44
        match arr.dtype().to_physical_type() {
45
            PH::Boolean => call_op!(BooleanArray, BooleanScalar, arr, $op$(, $variant)?),
46
            PH::Primitive(PR::Int8) => call_op!(dt: PArr<i8>, PScalar<i8>, arr, $op$(, $variant)?),
47
            PH::Primitive(PR::Int16) => call_op!(dt: PArr<i16>, PScalar<i16>, arr, $op$(, $variant)?),
48
            PH::Primitive(PR::Int32) => call_op!(dt: PArr<i32>, PScalar<i32>, arr, $op$(, $variant)?),
49
            PH::Primitive(PR::Int64) => call_op!(dt: PArr<i64>, PScalar<i64>, arr, $op$(, $variant)?),
50
            PH::Primitive(PR::Int128) => call_op!(dt: PArr<i128>, PScalar<i128>, arr, $op$(, $variant)?),
51
            PH::Primitive(PR::UInt8) => call_op!(dt: PArr<u8>, PScalar<u8>, arr, $op$(, $variant)?),
52
            PH::Primitive(PR::UInt16) => call_op!(dt: PArr<u16>, PScalar<u16>, arr, $op$(, $variant)?),
53
            PH::Primitive(PR::UInt32) => call_op!(dt: PArr<u32>, PScalar<u32>, arr, $op$(, $variant)?),
54
            PH::Primitive(PR::UInt64) => call_op!(dt: PArr<u64>, PScalar<u64>, arr, $op$(, $variant)?),
55
            PH::Primitive(PR::UInt128) => call_op!(dt: PArr<u128>, PScalar<u128>, arr, $op$(, $variant)?),
56
            PH::Primitive(PR::Float32) => call_op!(dt: PArr<f32>, PScalar<f32>, arr, $op$(, $variant)?),
57
            PH::Primitive(PR::Float64) => call_op!(dt: PArr<f64>, PScalar<f64>, arr, $op$(, $variant)?),
58

59
            PH::BinaryView => call_op!(BinaryViewArray, BinaryViewScalar<[u8]>, arr, $op$(, $variant)?),
60
            PH::Utf8View => call_op!(Utf8ViewArray, BinaryViewScalar<str>, arr, $op$(, $variant)?),
61

62
            PH::Binary => call_op!(BinaryArray<i32>, BinaryScalar<i32>, arr, $op$(, $variant)?),
63
            PH::LargeBinary => call_op!(BinaryArray<i64>, BinaryScalar<i64>, arr, $op$(, $variant)?),
64
            PH::Utf8 => call_op!(Utf8Array<i32>, BinaryScalar<i32>, arr, $op$(, $variant)?),
65
            PH::LargeUtf8 => call_op!(Utf8Array<i64>, BinaryScalar<i64>, arr, $op$(, $variant)?),
66

67
            _ => todo!("Dynamic MinMax is not yet implemented for {:?}", arr.dtype()),
68
        }
69
    }};
70
}
71

72
pub fn dyn_array_min_ignore_nan(arr: &dyn Array) -> Option<Box<dyn Scalar>> {
73
    call!(arr, MinMaxKernel::min_ignore_nan_kernel)
74
}
75

76
pub fn dyn_array_max_ignore_nan(arr: &dyn Array) -> Option<Box<dyn Scalar>> {
77
    call!(arr, MinMaxKernel::max_ignore_nan_kernel)
78
}
79

80
pub fn dyn_array_min_propagate_nan(arr: &dyn Array) -> Option<Box<dyn Scalar>> {
81
    call!(arr, MinMaxKernel::min_propagate_nan_kernel)
82
}
83

84
pub fn dyn_array_max_propagate_nan(arr: &dyn Array) -> Option<Box<dyn Scalar>> {
85
    call!(arr, MinMaxKernel::max_propagate_nan_kernel)
86
}
87

88
pub fn dyn_array_min_max_propagate_nan(
89
    arr: &dyn Array,
90
) -> Option<(Box<dyn Scalar>, Box<dyn Scalar>)> {
91
    call!(arr, MinMaxKernel::min_max_propagate_nan_kernel, ret_two)
92
}
93

94
Product

Resources

Company
