Skip to content

Commit

Permalink
feat: implement binview comparison kernels (pola-rs#13715)
Browse files Browse the repository at this point in the history
  • Loading branch information
ritchie46 authored and r-brink committed Jan 24, 2024
1 parent 6c8c213 commit 8a1d186
Show file tree
Hide file tree
Showing 2 changed files with 115 additions and 2 deletions.
8 changes: 7 additions & 1 deletion crates/polars-arrow/src/array/binview/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -315,6 +315,11 @@ impl<T: ViewType + ?Sized> BinaryViewArrayGeneric<T> {
pub fn total_buffer_len(&self) -> usize {
self.total_buffer_len
}

#[inline(always)]
pub fn len(&self) -> usize {
self.views.len()
}
}

impl BinaryViewArray {
Expand Down Expand Up @@ -370,8 +375,9 @@ impl<T: ViewType + ?Sized> Array for BinaryViewArrayGeneric<T> {
self
}

#[inline(always)]
fn len(&self) -> usize {
self.views.len()
BinaryViewArrayGeneric::len(self)
}

fn data_type(&self) -> &ArrowDataType {
Expand Down
109 changes: 108 additions & 1 deletion crates/polars-compute/src/comparisons/scalar.rs
Original file line number Diff line number Diff line change
@@ -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};
Expand Down Expand Up @@ -69,6 +71,111 @@ impl<T: NativeType + NotSimdPrimitive + TotalOrd> 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<i64> {
type Scalar = [u8];

Expand Down

0 comments on commit 8a1d186

Please sign in to comment.