add _missing kernels

pola-rs · ritchie46 · Nov 30, 2023 · Oct 16, 2023 · Oct 16, 2023 · Oct 17, 2023
commit 74f4e3fcb44347db579950977fcfeee6f7859f8d
@@ -1,10 +1,12 @@
-use arrow::bitmap::Bitmap;
+use arrow::array::Array;
+use arrow::bitmap::{Bitmap, self};
 
 // Low-level comparison kernel.
-// Ignores validity (results for nulls are unspecified but initialized).
-pub trait TotalOrdKernel: Sized {
+pub trait TotalOrdKernel: Sized + Array {
     type Scalar: ?Sized;
 
+    // These kernels ignore validity entirely (results for nulls are unspecified
+    // but initialized).
     fn tot_eq_kernel(&self, other: &Self) -> Bitmap;
     fn tot_ne_kernel(&self, other: &Self) -> Bitmap;
     fn tot_lt_kernel(&self, other: &Self) -> Bitmap;
@@ -16,28 +18,74 @@ pub trait TotalOrdKernel: Sized {
         other.tot_le_kernel(self)
     }
 
+    // These kernels ignore validity entirely (results for nulls are unspecified
+    // but initialized).
     fn tot_eq_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap;
     fn tot_ne_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap;
     fn tot_lt_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap;
     fn tot_le_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap;
     fn tot_gt_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap;
     fn tot_ge_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap;
-}
-
-trait NotSimd {}
+
+    // These kernels treat null as any other value equal to itself but unequal
+    // to anything else.
+    fn tot_eq_missing_kernel(&self, other: &Self) -> Bitmap {
+        let q = self.tot_eq_kernel(other);
+        let combined = match (self.validity(), other.validity()) {
+            (None, None) => q,
+            (None, Some(r)) => &q & r,
+            (Some(l), None) => &q & l,
+            (Some(l), Some(r)) => {
+                bitmap::ternary(&q, l, r, |q, l, r| (q & l & r) | !(l | r))
+            },
+        };
+        combined
+    }
+
+    fn tot_ne_missing_kernel(&self, other: &Self) -> Bitmap {
+        let q = self.tot_ne_kernel(other);
+        let combined = match (self.validity(), other.validity()) {
+            (None, None) => q,
+            (None, Some(r)) => &q | &!r,
+            (Some(l), None) => &q | &!l,
+            (Some(l), Some(r)) => {
+                bitmap::ternary(&q, l, r, |q, l, r| (q & l & r) | (l ^ r))
+            },
+        };
+        combined.into()
+    }
 
-#[allow(unused)]
-macro_rules! impl_not_simd {
-    ($($T:ty,)*) => {
-        $(impl NotSimd for $T { })*
-    };
+    // These kernels treat null as any other value equal to itself but unequal
+    // to anything else. other is assumed to be non-null.
+    fn tot_eq_missing_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap {
+        let q = self.tot_eq_kernel_broadcast(other);
+        if let Some(valid) = self.validity() {
+            bitmap::binary(&q, valid, |q, v| q & v).into()
+        } else {
+            q.into()
+        }
+    }
+
+    fn tot_ne_missing_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap {
+        let q = self.tot_ne_kernel_broadcast(other);
+        if let Some(valid) = self.validity() {
+            bitmap::binary(&q, valid, |q, v| q | !v).into()
+        } else {
+            q.into()
+        }
+    }
 }
 
+// Trait to enable the scalar blanket implementation.
+trait NotSimdPrimitive {}
+
 #[cfg(not(feature = "simd"))]
-impl<T> NotSimd for T {}
+impl<T> NotSimdPrimitive for T {}
 
 #[cfg(feature = "simd")]
-impl_not_simd!(u128, i128,);
+impl NotSimdPrimitive for u128 {}
+#[cfg(feature = "simd")]
+impl NotSimdPrimitive for i128 {}
 
 mod scalar;
 

@@ -3,9 +3,9 @@ use arrow::bitmap::Bitmap;
 use arrow::types::NativeType;
 use polars_utils::total_ord::{TotalOrd, TotalEq};
 
-use super::{NotSimd, TotalOrdKernel};
+use super::{NotSimdPrimitive, TotalOrdKernel};
 
-impl<T: NativeType + NotSimd + TotalOrd> TotalOrdKernel for PrimitiveArray<T> {
+impl<T: NativeType + NotSimdPrimitive + TotalOrd> TotalOrdKernel for PrimitiveArray<T> {
     type Scalar = T;
 
     fn tot_lt_kernel(&self, other: &Self) -> Bitmap {

@@ -3,7 +3,7 @@ mod scalar;
 use std::ops::{BitOr, Not};
 
 use arrow::array::{BooleanArray, Utf8Array};
-use arrow::bitmap::{self, MutableBitmap};
+use arrow::bitmap::MutableBitmap;
 use arrow::compute;
 use arrow::compute::comparison;
 use arrow::legacy::prelude::FromData;
@@ -60,23 +60,7 @@ where
                     rhs.is_null()
                 }
             },
-            _ => arity::binary_mut_with_options(
-                self,
-                rhs,
-                |a, b| {
-                    let q = a.tot_eq_kernel(b);
-                    let combined = match (a.validity(), b.validity()) {
-                        (None, None) => q,
-                        (None, Some(r)) => &q & r,
-                        (Some(l), None) => &q & l,
-                        (Some(l), Some(r)) => {
-                            bitmap::ternary(&q, l, r, |q, l, r| (q & l & r) | !(l | r))
-                        },
-                    };
-                    combined.into()
-                },
-                "",
-            ),
+            _ => arity::binary_mut_with_options(self, rhs, |a, b| a.tot_eq_missing_kernel(b).into(), ""),
         }
     }
 
@@ -118,23 +102,7 @@ where
                     rhs.is_not_null()
                 }
             },
-            _ => arity::binary_mut_with_options(
-                self,
-                rhs,
-                |a, b| {
-                    let q = a.tot_ne_kernel(b);
-                    let combined = match (a.validity(), b.validity()) {
-                        (None, None) => q,
-                        (None, Some(r)) => &q | &!r,
-                        (Some(l), None) => &q | &!l,
-                        (Some(l), Some(r)) => {
-                            bitmap::ternary(&q, l, r, |q, l, r| (q & l & r) | (l ^ r))
-                        },
-                    };
-                    combined.into()
-                },
-                "",
-            ),
+            _ => arity::binary_mut_with_options(self, rhs, |a, b| a.tot_ne_missing_kernel(b).into(), ""),
         }
     }
 
@@ -458,23 +426,7 @@ impl ChunkCompare<&BinaryChunked> for BinaryChunked {
                     rhs.is_null()
                 }
             },
-            _ => arity::binary_mut_with_options(
-                self,
-                rhs,
-                |a, b| {
-                    let q = a.tot_eq_kernel(b);
-                    let combined = match (a.validity(), b.validity()) {
-                        (None, None) => q,
-                        (None, Some(r)) => &q & r,
-                        (Some(l), None) => &q & l,
-                        (Some(l), Some(r)) => {
-                            bitmap::ternary(&q, l, r, |q, l, r| (q & l & r) | !(l | r))
-                        },
-                    };
-                    combined.into()
-                },
-                "",
-            ),
+            _ => arity::binary_mut_with_options(self, rhs, |a, b| a.tot_eq_missing_kernel(b).into(), ""),
         }
     }
 
@@ -516,23 +468,7 @@ impl ChunkCompare<&BinaryChunked> for BinaryChunked {
                     rhs.is_not_null()
                 }
             },
-            _ => arity::binary_mut_with_options(
-                self,
-                rhs,
-                |a, b| {
-                    let q = a.tot_ne_kernel(b);
-                    let combined = match (a.validity(), b.validity()) {
-                        (None, None) => q,
-                        (None, Some(r)) => &q | &!r,
-                        (Some(l), None) => &q | &!l,
-                        (Some(l), Some(r)) => {
-                            bitmap::ternary(&q, l, r, |q, l, r| (q & l & r) | (l ^ r))
-                        },
-                    };
-                    combined.into()
-                },
-                "",
-            ),
+            _ => arity::binary_mut_with_options(self, rhs, |a, b| a.tot_ne_missing_kernel(b).into(), ""),
         }
     }
 

@@ -46,15 +46,7 @@ where
 
     fn equal_missing(&self, rhs: Rhs) -> BooleanChunked {
         let rhs: T::Native = NumCast::from(rhs).unwrap();
-        let iter = self.downcast_iter().map(|arr| {
-            let eq = arr.tot_eq_kernel_broadcast(&rhs);
-            if let Some(valid) = arr.validity() {
-                bitmap::binary(&eq, valid, |e, v| e & v).into()
-            } else {
-                eq.into()
-            }
-        });
-        ChunkedArray::from_chunk_iter(self.name(), iter)
+        arity::unary_mut_with_options(self, |arr| arr.tot_eq_missing_kernel_broadcast(&rhs).into())
     }
 
     fn not_equal(&self, rhs: Rhs) -> BooleanChunked {
@@ -64,15 +56,7 @@ where
 
     fn not_equal_missing(&self, rhs: Rhs) -> BooleanChunked {
         let rhs: T::Native = NumCast::from(rhs).unwrap();
-        let iter = self.downcast_iter().map(|arr| {
-            let ne = arr.tot_ne_kernel_broadcast(&rhs);
-            if let Some(valid) = arr.validity() {
-                bitmap::binary(&ne, valid, |n, v| n | !v).into()
-            } else {
-                ne.into()
-            }
-        });
-        ChunkedArray::from_chunk_iter(self.name(), iter)
+        arity::unary_mut_with_options(self, |arr| arr.tot_ne_missing_kernel_broadcast(&rhs).into())
     }
 
     fn gt(&self, rhs: Rhs) -> BooleanChunked {

@@ -45,6 +45,19 @@ where
     ChunkedArray::from_chunk_iter(ca.name(), iter)
 }
 
+/// Applies a kernel that produces `Array` types.
+#[inline]
+pub fn unary_mut_with_options<T, V, F, Arr>(ca: &ChunkedArray<T>, mut op: F) -> ChunkedArray<V>
+where
+    T: PolarsDataType,
+    V: PolarsDataType<Array = Arr>,
+    Arr: Array + StaticArray,
+    F: FnMut(&T::Array) -> Arr,
+{
+    let iter = ca.downcast_iter().map(|arr| op(arr));
+    ChunkedArray::from_chunk_iter(ca.name(), iter)
+}
+
 #[inline]
 pub fn binary_elementwise<T, U, V, F>(
     lhs: &ChunkedArray<T>,