diff --git a/crates/polars-arrow/src/array/binview/mod.rs b/crates/polars-arrow/src/array/binview/mod.rs index 4796a7f07c50..95e0f00fd100 100644 --- a/crates/polars-arrow/src/array/binview/mod.rs +++ b/crates/polars-arrow/src/array/binview/mod.rs @@ -315,6 +315,11 @@ impl BinaryViewArrayGeneric { pub fn total_buffer_len(&self) -> usize { self.total_buffer_len } + + #[inline(always)] + pub fn len(&self) -> usize { + self.views.len() + } } impl BinaryViewArray { @@ -370,8 +375,9 @@ impl Array for BinaryViewArrayGeneric { self } + #[inline(always)] fn len(&self) -> usize { - self.views.len() + BinaryViewArrayGeneric::len(self) } fn data_type(&self) -> &ArrowDataType { diff --git a/crates/polars-compute/src/comparisons/scalar.rs b/crates/polars-compute/src/comparisons/scalar.rs index bc338f3816a8..6ad3c4cea011 100644 --- a/crates/polars-compute/src/comparisons/scalar.rs +++ b/crates/polars-compute/src/comparisons/scalar.rs @@ -1,4 +1,6 @@ -use arrow::array::{BinaryArray, BooleanArray, PrimitiveArray, Utf8Array}; +use arrow::array::{ + BinaryArray, BinaryViewArray, BooleanArray, PrimitiveArray, Utf8Array, Utf8ViewArray, +}; use arrow::bitmap::{self, Bitmap}; use arrow::types::NativeType; use polars_utils::total_ord::{TotalEq, TotalOrd}; @@ -69,6 +71,111 @@ impl TotalOrdKernel for PrimitiveAr } } +impl TotalOrdKernel for BinaryViewArray { + type Scalar = [u8]; + + fn tot_eq_kernel(&self, other: &Self) -> Bitmap { + debug_assert!(self.len() == other.len()); + // TODO! speed-up by first comparing views + self.values_iter() + .zip(other.values_iter()) + .map(|(l, r)| l.tot_eq(&r)) + .collect() + } + + fn tot_ne_kernel(&self, other: &Self) -> Bitmap { + debug_assert!(self.len() == other.len()); + self.values_iter() + .zip(other.values_iter()) + .map(|(l, r)| l.tot_ne(&r)) + .collect() + } + + fn tot_lt_kernel(&self, other: &Self) -> Bitmap { + debug_assert!(self.len() == other.len()); + self.values_iter() + .zip(other.values_iter()) + .map(|(l, r)| l.tot_lt(&r)) + .collect() + } + + fn tot_le_kernel(&self, other: &Self) -> Bitmap { + debug_assert!(self.len() == other.len()); + self.values_iter() + .zip(other.values_iter()) + .map(|(l, r)| l.tot_le(&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() + } +} + +impl TotalOrdKernel for Utf8ViewArray { + type Scalar = str; + + fn tot_eq_kernel(&self, other: &Self) -> Bitmap { + self.to_binview().tot_eq_kernel(&other.to_binview()) + } + + fn tot_ne_kernel(&self, other: &Self) -> Bitmap { + self.to_binview().tot_ne_kernel(&other.to_binview()) + } + + fn tot_lt_kernel(&self, other: &Self) -> Bitmap { + self.to_binview().tot_lt_kernel(&other.to_binview()) + } + + fn tot_le_kernel(&self, other: &Self) -> Bitmap { + self.to_binview().tot_le_kernel(&other.to_binview()) + } + + fn tot_eq_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap { + self.to_binview().tot_eq_kernel_broadcast(other.as_bytes()) + } + + fn tot_ne_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap { + self.to_binview().tot_ne_kernel_broadcast(other.as_bytes()) + } + + fn tot_lt_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap { + self.to_binview().tot_lt_kernel_broadcast(other.as_bytes()) + } + + fn tot_le_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap { + self.to_binview().tot_le_kernel_broadcast(other.as_bytes()) + } + + fn tot_gt_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap { + self.to_binview().tot_gt_kernel_broadcast(other.as_bytes()) + } + + fn tot_ge_kernel_broadcast(&self, other: &Self::Scalar) -> Bitmap { + self.to_binview().tot_ge_kernel_broadcast(other.as_bytes()) + } +} + impl TotalOrdKernel for BinaryArray { type Scalar = [u8];