Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
pola-rs
GitHub Repository: pola-rs/polars
 Path: blob/main/crates/polars-compute/src/comparisons/dyn_array.rs
8415 views
1
use arrow::array::{

2
    Array, BinaryArray, BinaryViewArray, BooleanArray, DictionaryArray, FixedSizeBinaryArray,

3
    ListArray, NullArray, PrimitiveArray, StructArray, Utf8Array, Utf8ViewArray,

4
};

5
use arrow::bitmap::Bitmap;

6
use arrow::types::{days_ms, i256, months_days_ns};

7
use polars_utils::float16::pf16;

8


9
use crate::comparisons::TotalEqKernel;

10


11
macro_rules! call_binary {

12
    ($T:ty, $lhs:expr, $rhs:expr, $op:path) => {{

13
        let lhs: &$T = $lhs.as_any().downcast_ref().unwrap();

14
        let rhs: &$T = $rhs.as_any().downcast_ref().unwrap();

15
        $op(lhs, rhs)

16
    }};

17
}

18


19
macro_rules! compare {

20
    ($lhs:expr, $rhs:expr, $op:path) => {{

21
        let lhs = $lhs;

22
        let rhs = $rhs;

23


24
        assert_eq!(lhs.dtype(), rhs.dtype());

25


26
        use arrow::datatypes::{IntegerType as I, PhysicalType as PH, PrimitiveType as PR};

27
        match lhs.dtype().to_physical_type() {

28
            PH::Boolean => call_binary!(BooleanArray, lhs, rhs, $op),

29
            PH::BinaryView => call_binary!(BinaryViewArray, lhs, rhs, $op),

30
            PH::Utf8View => call_binary!(Utf8ViewArray, lhs, rhs, $op),

31
            PH::Primitive(PR::Int8) => call_binary!(PrimitiveArray<i8>, lhs, rhs, $op),

32
            PH::Primitive(PR::Int16) => call_binary!(PrimitiveArray<i16>, lhs, rhs, $op),

33
            PH::Primitive(PR::Int32) => call_binary!(PrimitiveArray<i32>, lhs, rhs, $op),

34
            PH::Primitive(PR::Int64) => call_binary!(PrimitiveArray<i64>, lhs, rhs, $op),

35
            PH::Primitive(PR::Int128) => call_binary!(PrimitiveArray<i128>, lhs, rhs, $op),

36
            PH::Primitive(PR::UInt8) => call_binary!(PrimitiveArray<u8>, lhs, rhs, $op),

37
            PH::Primitive(PR::UInt16) => call_binary!(PrimitiveArray<u16>, lhs, rhs, $op),

38
            PH::Primitive(PR::UInt32) => call_binary!(PrimitiveArray<u32>, lhs, rhs, $op),

39
            PH::Primitive(PR::UInt64) => call_binary!(PrimitiveArray<u64>, lhs, rhs, $op),

40
            PH::Primitive(PR::UInt128) => call_binary!(PrimitiveArray<u128>, lhs, rhs, $op),

41
            PH::Primitive(PR::Float16) => call_binary!(PrimitiveArray<pf16>, lhs, rhs, $op),

42
            PH::Primitive(PR::Float32) => call_binary!(PrimitiveArray<f32>, lhs, rhs, $op),

43
            PH::Primitive(PR::Float64) => call_binary!(PrimitiveArray<f64>, lhs, rhs, $op),

44
            PH::Primitive(PR::Int256) => call_binary!(PrimitiveArray<i256>, lhs, rhs, $op),

45
            PH::Primitive(PR::DaysMs) => call_binary!(PrimitiveArray<days_ms>, lhs, rhs, $op),

46
            PH::Primitive(PR::MonthDayNano) => {

47
                call_binary!(PrimitiveArray<months_days_ns>, lhs, rhs, $op)

48
            },

49
            PH::Primitive(PR::MonthDayMillis) => unimplemented!(),

50


51
            #[cfg(feature = "dtype-array")]

52
            PH::FixedSizeList => call_binary!(arrow::array::FixedSizeListArray, lhs, rhs, $op),

53
            #[cfg(not(feature = "dtype-array"))]

54
            PH::FixedSizeList => todo!(

55
                "Comparison of FixedSizeListArray is not supported without dtype-array feature"

56
            ),

57


58
            PH::Null => call_binary!(NullArray, lhs, rhs, $op),

59
            PH::FixedSizeBinary => call_binary!(FixedSizeBinaryArray, lhs, rhs, $op),

60
            PH::Binary => call_binary!(BinaryArray<i32>, lhs, rhs, $op),

61
            PH::LargeBinary => call_binary!(BinaryArray<i64>, lhs, rhs, $op),

62
            PH::Utf8 => call_binary!(Utf8Array<i32>, lhs, rhs, $op),

63
            PH::LargeUtf8 => call_binary!(Utf8Array<i64>, lhs, rhs, $op),

64
            PH::List => call_binary!(ListArray<i32>, lhs, rhs, $op),

65
            PH::LargeList => call_binary!(ListArray<i64>, lhs, rhs, $op),

66
            PH::Struct => call_binary!(StructArray, lhs, rhs, $op),

67
            PH::Union => todo!("Comparison of UnionArrays is not yet supported"),

68
            PH::Map => todo!("Comparison of MapArrays is not yet supported"),

69
            PH::Dictionary(I::Int8) => call_binary!(DictionaryArray<i8>, lhs, rhs, $op),

70
            PH::Dictionary(I::Int16) => call_binary!(DictionaryArray<i16>, lhs, rhs, $op),

71
            PH::Dictionary(I::Int32) => call_binary!(DictionaryArray<i32>, lhs, rhs, $op),

72
            PH::Dictionary(I::Int64) => call_binary!(DictionaryArray<i64>, lhs, rhs, $op),

73
            PH::Dictionary(I::Int128) => call_binary!(DictionaryArray<i128>, lhs, rhs, $op),

74
            PH::Dictionary(I::UInt8) => call_binary!(DictionaryArray<u8>, lhs, rhs, $op),

75
            PH::Dictionary(I::UInt16) => call_binary!(DictionaryArray<u16>, lhs, rhs, $op),

76
            PH::Dictionary(I::UInt32) => call_binary!(DictionaryArray<u32>, lhs, rhs, $op),

77
            PH::Dictionary(I::UInt64) => call_binary!(DictionaryArray<u64>, lhs, rhs, $op),

78
            PH::Dictionary(I::UInt128) => call_binary!(DictionaryArray<u128>, lhs, rhs, $op),

79
        }

80
    }};

81
}

82


83
pub fn array_tot_eq_missing_kernel(lhs: &dyn Array, rhs: &dyn Array) -> Bitmap {

84
    compare!(lhs, rhs, TotalEqKernel::tot_eq_missing_kernel)

85
}

86


87
pub fn array_tot_ne_missing_kernel(lhs: &dyn Array, rhs: &dyn Array) -> Bitmap {

88
    compare!(lhs, rhs, TotalEqKernel::tot_ne_missing_kernel)

89
}

90


91