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
6939 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, f16, i256, months_days_ns};

7


8
use crate::comparisons::TotalEqKernel;

9


10
macro_rules! call_binary {

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

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

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

14
        $op(lhs, rhs)

15
    }};

16
}

17


18
macro_rules! compare {

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

20
        let lhs = $lhs;

21
        let rhs = $rhs;

22


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

24


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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

47
            },

48


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

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

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

52
            PH::FixedSizeList => todo!(

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

54
            ),

55


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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

76
        }

77
    }};

78
}

79


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

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

82
}

83


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

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

86
}

87


88