mostly fix broadcasting comparisons

pola-rs · ritchie46 · Nov 30, 2023 · Oct 16, 2023 · Oct 16, 2023 · Oct 17, 2023
commit d37d9c0c427d5da83f6e6afe9d0011136ded2102
@@ -1,19 +1,27 @@
 use arrow::bitmap::Bitmap;
 
+// Low-level comparison kernel.
+// Ignores validity (results for nulls are unspecified but initialized).
 pub trait TotalOrdKernel : Sized {
+    type Scalar;
+
+    fn tot_eq_kernel(&self, other: &Self) -> Bitmap;
+    fn tot_ne_kernel(&self, other: &Self) -> Bitmap;
     fn tot_lt_kernel(&self, other: &Self) -> Bitmap;
     fn tot_le_kernel(&self, other: &Self) -> Bitmap;
-
     fn tot_gt_kernel(&self, other: &Self) -> Bitmap {
         other.tot_lt_kernel(self)
     }
-
     fn tot_ge_kernel(&self, other: &Self) -> Bitmap {
         other.tot_le_kernel(self)
     }
-
-    fn tot_eq_kernel(&self, other: &Self) -> Bitmap;
-    fn tot_ne_kernel(&self, other: &Self) -> Bitmap;
+
+    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 { }

@@ -1,18 +1,19 @@
-// fn to_bitmap(
-
 use arrow::array::PrimitiveArray;
 use arrow::bitmap::Bitmap;
 use arrow::types::NativeType;
+use polars_utils::total_ord::TotalOrd;
 
 use super::{TotalOrdKernel, NotSimd};
 
-impl<T: NativeType + NotSimd> TotalOrdKernel for PrimitiveArray<T> {
+impl<T: NativeType + NotSimd + TotalOrd> TotalOrdKernel for PrimitiveArray<T> {
+    type Scalar = T;
+
     fn tot_lt_kernel(&self, other: &Self) -> Bitmap {
         assert!(self.len() == other.len());
         self.values()
             .iter()
             .zip(other.values().iter())
-            .map(|(a, b)| a.tot_lt(b))
+            .map(|(l, r)| l.tot_lt(r))
             .collect()
     }
 
@@ -21,7 +22,7 @@ impl<T: NativeType + NotSimd> TotalOrdKernel for PrimitiveArray<T> {
         self.values()
             .iter()
             .zip(other.values().iter())
-            .map(|(a, b)| a.tot_le(b))
+            .map(|(l, r)| l.tot_le(r))
             .collect()
     }
 
@@ -30,7 +31,7 @@ impl<T: NativeType + NotSimd> TotalOrdKernel for PrimitiveArray<T> {
         self.values()
             .iter()
             .zip(other.values().iter())
-            .map(|(a, b)| a.tot_eq(b))
+            .map(|(l, r)| l.tot_eq(r))
             .collect()
     }
 
@@ -39,7 +40,50 @@ impl<T: NativeType + NotSimd> TotalOrdKernel for PrimitiveArray<T> {
         self.values()
             .iter()
             .zip(other.values().iter())
-            .map(|(a, b)| a.tot_ne(b))
+            .map(|(l, r)| l.tot_ne(r))
+            .collect()
+    }
+
+    fn tot_eq_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap {
+        self.values()
+            .iter()
+            .map(|l| l.tot_eq(other))
+            .collect()
+    }
+
+    fn tot_ne_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap {
+        self.values()
+            .iter()
+            .map(|l| l.tot_ne(other))
+            .collect()
+    }
+
+    fn tot_lt_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap {
+        self.values()
+            .iter()
+            .map(|l| l.tot_lt(other))
             .collect()
     }
+
+    fn tot_le_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap {
+        self.values()
+            .iter()
+            .map(|l| l.tot_le(other))
+            .collect()
+    }
+
+    fn tot_gt_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap {
+        self.values()
+            .iter()
+            .map(|l| l.tot_gt(other))
+            .collect()
+    }
+
+    fn tot_ge_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap {
+        self.values()
+            .iter()
+            .map(|l| l.tot_ge(other))
+            .collect()
+    }
+
 }
@@ -8,7 +8,7 @@ use bytemuck::Pod;
 use super::TotalOrdKernel;
 
 
-fn apply_comp_kernel<const N: usize, M: Pod, T, F>(
+fn apply_binary_kernel<const N: usize, M: Pod, T, F>(
     lhs: &PrimitiveArray<T>,
     rhs: &PrimitiveArray<T>,
     mut f: F,
@@ -45,30 +45,105 @@ where
     Bitmap::from_u8_vec(bytemuck::cast_vec(v), n)
 }
 
+fn apply_unary_kernel<const N: usize, M: Pod, T, F>(
+    arg: &PrimitiveArray<T>,
+    mut f: F,
+) -> Bitmap
+where
+    T: NativeType,
+    F: FnMut(&[T; N]) -> M,
+{
+    assert!(std::mem::size_of::<M>() == N);
+    let n = arg.len();
+
+    let arg_buf = arg.values().as_slice();
+    let mut arg_chunks = arg_buf.chunks_exact(N);
+    let mut v = Vec::with_capacity(n.div_ceil(N));
+    v.extend(
+        arg_chunks.by_ref()
+            .map(|l| unsafe {
+                f(l.try_into().unwrap_unchecked())
+            })
+    );
+
+    if n % N > 0 {
+        let mut l: [T; N] = [T::zeroed(); N];
+        l.copy_from_slice(arg_chunks.remainder());
+        v.push(f(&l));
+    }
+
+    Bitmap::from_u8_vec(bytemuck::cast_vec(v), n)
+}
+
+
 macro_rules! impl_int_total_ord_kernel {
     ($T: ty, $width: literal, $mask: ty) => {
         impl TotalOrdKernel for PrimitiveArray<$T> {
+            type Scalar = $T;
+
+            fn tot_eq_kernel(&self, other: &Self) -> Bitmap {
+                apply_binary_kernel::<$width, $mask, _, _>(self, other, |l, r| {
+                    Simd::from(*l).simd_eq(Simd::from(*r)).to_bitmask()
+                })
+            }
+
+            fn tot_ne_kernel(&self, other: &Self) -> Bitmap {
+                apply_binary_kernel::<$width, $mask, _, _>(self, other, |l, r| {
+                    Simd::from(*l).simd_ne(Simd::from(*r)).to_bitmask()
+                })
+            }
+
             fn tot_lt_kernel(&self, other: &Self) -> Bitmap {
-                apply_comp_kernel::<$width, $mask, _, _>(self, other, |l, r| {
+                apply_binary_kernel::<$width, $mask, _, _>(self, other, |l, r| {
                     Simd::from(*l).simd_lt(Simd::from(*r)).to_bitmask()
                 })
             }
 
             fn tot_le_kernel(&self, other: &Self) -> Bitmap {
-                apply_comp_kernel::<$width, $mask, _, _>(self, other, |l, r| {
+                apply_binary_kernel::<$width, $mask, _, _>(self, other, |l, r| {
                     Simd::from(*l).simd_le(Simd::from(*r)).to_bitmask()
                 })
             }
 
-            fn tot_eq_kernel(&self, other: &Self) -> Bitmap {
-                apply_comp_kernel::<$width, $mask, _, _>(self, other, |l, r| {
-                    Simd::from(*l).simd_eq(Simd::from(*r)).to_bitmask()
+            fn tot_eq_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap {
+                let r = Simd::splat(*other);
+                apply_unary_kernel::<$width, $mask, _, _>(self, |l| {
+                    Simd::from(*l).simd_eq(r).to_bitmask()
                 })
             }
 
-            fn tot_ne_kernel(&self, other: &Self) -> Bitmap {
-                apply_comp_kernel::<$width, $mask, _, _>(self, other, |l, r| {
-                    Simd::from(*l).simd_ne(Simd::from(*r)).to_bitmask()
+            fn tot_ne_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap {
+                let r = Simd::splat(*other);
+                apply_unary_kernel::<$width, $mask, _, _>(self, |l| {
+                    Simd::from(*l).simd_ne(r).to_bitmask()
+                })
+            }
+
+            fn tot_lt_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap {
+                let r = Simd::splat(*other);
+                apply_unary_kernel::<$width, $mask, _, _>(self, |l| {
+                    Simd::from(*l).simd_lt(r).to_bitmask()
+                })
+            }
+
+            fn tot_le_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap {
+                let r = Simd::splat(*other);
+                apply_unary_kernel::<$width, $mask, _, _>(self, |l| {
+                    Simd::from(*l).simd_le(r).to_bitmask()
+                })
+            }
+
+            fn tot_gt_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap {
+                let r = Simd::splat(*other);
+                apply_unary_kernel::<$width, $mask, _, _>(self, |l| {
+                    Simd::from(*l).simd_gt(r).to_bitmask()
+                })
+            }
+
+            fn tot_ge_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap {
+                let r = Simd::splat(*other);
+                apply_unary_kernel::<$width, $mask, _, _>(self, |l| {
+                    Simd::from(*l).simd_ge(r).to_bitmask()
                 })
             }
         }
@@ -78,8 +153,30 @@ macro_rules! impl_int_total_ord_kernel {
 macro_rules! impl_float_total_ord_kernel {
     ($T: ty, $width: literal, $mask: ty) => {
         impl TotalOrdKernel for PrimitiveArray<$T> {
+            type Scalar = $T;
+
+            fn tot_eq_kernel(&self, other: &Self) -> Bitmap {
+                apply_binary_kernel::<$width, $mask, _, _>(self, other, |l, r| {
+                    let ls = Simd::from(*l);
+                    let rs = Simd::from(*r);
+                    let lhs_is_nan = ls.simd_ne(ls);
+                    let rhs_is_nan = rs.simd_ne(rs);
+                    ((lhs_is_nan & rhs_is_nan) | ls.simd_eq(rs)).to_bitmask()
+                })
+            }
+
+            fn tot_ne_kernel(&self, other: &Self) -> Bitmap {
+                apply_binary_kernel::<$width, $mask, _, _>(self, other, |l, r| {
+                    let ls = Simd::from(*l);
+                    let rs = Simd::from(*r);
+                    let lhs_is_nan = ls.simd_ne(ls);
+                    let rhs_is_nan = rs.simd_ne(rs);
+                    (!((lhs_is_nan & rhs_is_nan) | ls.simd_eq(rs))).to_bitmask()
+                })
+            }
+
             fn tot_lt_kernel(&self, other: &Self) -> Bitmap {
-                apply_comp_kernel::<$width, $mask, _, _>(self, other, |l, r| {
+                apply_binary_kernel::<$width, $mask, _, _>(self, other, |l, r| {
                     let ls = Simd::from(*l);
                     let rs = Simd::from(*r);
                     let lhs_is_nan = ls.simd_ne(ls);
@@ -88,33 +185,69 @@ macro_rules! impl_float_total_ord_kernel {
             }
 
             fn tot_le_kernel(&self, other: &Self) -> Bitmap {
-                apply_comp_kernel::<$width, $mask, _, _>(self, other, |l, r| {
+                apply_binary_kernel::<$width, $mask, _, _>(self, other, |l, r| {
                     let ls = Simd::from(*l);
                     let rs = Simd::from(*r);
                     let rhs_is_nan = rs.simd_ne(rs);
                     (rhs_is_nan | ls.simd_le(rs)).to_bitmask()
                 })
             }
 
-            fn tot_eq_kernel(&self, other: &Self) -> Bitmap {
-                apply_comp_kernel::<$width, $mask, _, _>(self, other, |l, r| {
+            fn tot_eq_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap {
+                let rs = Simd::splat(*other);
+                apply_unary_kernel::<$width, $mask, _, _>(self, |l| {
                     let ls = Simd::from(*l);
-                    let rs = Simd::from(*r);
                     let lhs_is_nan = ls.simd_ne(ls);
                     let rhs_is_nan = rs.simd_ne(rs);
                     ((lhs_is_nan & rhs_is_nan) | ls.simd_eq(rs)).to_bitmask()
                 })
             }
 
-            fn tot_ne_kernel(&self, other: &Self) -> Bitmap {
-                apply_comp_kernel::<$width, $mask, _, _>(self, other, |l, r| {
+            fn tot_ne_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap {
+                let rs = Simd::splat(*other);
+                apply_unary_kernel::<$width, $mask, _, _>(self, |l| {
                     let ls = Simd::from(*l);
-                    let rs = Simd::from(*r);
                     let lhs_is_nan = ls.simd_ne(ls);
                     let rhs_is_nan = rs.simd_ne(rs);
                     (!((lhs_is_nan & rhs_is_nan) | ls.simd_eq(rs))).to_bitmask()
                 })
             }
+
+            fn tot_lt_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap {
+                let rs = Simd::splat(*other);
+                apply_unary_kernel::<$width, $mask, _, _>(self, |l| {
+                    let ls = Simd::from(*l);
+                    let lhs_is_nan = ls.simd_ne(ls);
+                    (!(lhs_is_nan | ls.simd_ge(rs))).to_bitmask()
+                })
+            }
+
+            fn tot_le_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap {
+                let rs = Simd::splat(*other);
+                apply_unary_kernel::<$width, $mask, _, _>(self, |l| {
+                    let ls = Simd::from(*l);
+                    let rhs_is_nan = rs.simd_ne(rs);
+                    (rhs_is_nan | ls.simd_le(rs)).to_bitmask()
+                })
+            }
+
+            fn tot_gt_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap {
+                let rs = Simd::splat(*other);
+                apply_unary_kernel::<$width, $mask, _, _>(self, |l| {
+                    let ls = Simd::from(*l);
+                    let rhs_is_nan = rs.simd_ne(rs);
+                    (!(rhs_is_nan | rs.simd_ge(ls))).to_bitmask()
+                })
+            }
+
+            fn tot_ge_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap {
+                let rs = Simd::splat(*other);
+                apply_unary_kernel::<$width, $mask, _, _>(self, |l| {
+                    let ls = Simd::from(*l);
+                    let lhs_is_nan = ls.simd_ne(ls);
+                    (lhs_is_nan | rs.simd_le(ls)).to_bitmask()
+                })
+            }
         }
     }
 }