polars_compute/min_max/
dyn_array.rs

1use arrow::array::{
2    Array, BinaryArray, BinaryViewArray, BooleanArray, PrimitiveArray, Utf8Array, Utf8ViewArray,
3};
4use arrow::scalar::{BinaryScalar, BinaryViewScalar, BooleanScalar, PrimitiveScalar, Scalar};
5
6use crate::min_max::MinMaxKernel;
7
8macro_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
37macro_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
72pub fn dyn_array_min_ignore_nan(arr: &dyn Array) -> Option<Box<dyn Scalar>> {
73    call!(arr, MinMaxKernel::min_ignore_nan_kernel)
74}
75
76pub fn dyn_array_max_ignore_nan(arr: &dyn Array) -> Option<Box<dyn Scalar>> {
77    call!(arr, MinMaxKernel::max_ignore_nan_kernel)
78}
79
80pub fn dyn_array_min_propagate_nan(arr: &dyn Array) -> Option<Box<dyn Scalar>> {
81    call!(arr, MinMaxKernel::min_propagate_nan_kernel)
82}
83
84pub fn dyn_array_max_propagate_nan(arr: &dyn Array) -> Option<Box<dyn Scalar>> {
85    call!(arr, MinMaxKernel::max_propagate_nan_kernel)
86}
87
88pub 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}
polars_compute/min_max/dyn_array.rs

polars_compute/min_max/
dyn_array.rs
