From 169f53a6191a2fe0a5d5b1b029debc0d8ea74f9b Mon Sep 17 00:00:00 2001 From: Gijs Burghoorn Date: Thu, 7 Nov 2024 13:04:39 +0100 Subject: [PATCH] refactor: Move Series bitops to `std::ops::Bit...` (#19673) --- crates/polars-core/src/frame/column/mod.rs | 12 +--- .../src/series/arithmetic/bitops.rs | 65 +++++++++++++++++++ .../polars-core/src/series/arithmetic/mod.rs | 1 + .../src/series/implementations/boolean.rs | 15 ----- .../src/series/implementations/mod.rs | 31 --------- crates/polars-core/src/series/series_trait.rs | 12 ---- crates/polars-python/src/series/general.rs | 16 +---- 7 files changed, 72 insertions(+), 80 deletions(-) create mode 100644 crates/polars-core/src/series/arithmetic/bitops.rs diff --git a/crates/polars-core/src/frame/column/mod.rs b/crates/polars-core/src/frame/column/mod.rs index 16eb940022b3..cd82126ecec5 100644 --- a/crates/polars-core/src/frame/column/mod.rs +++ b/crates/polars-core/src/frame/column/mod.rs @@ -1030,23 +1030,17 @@ impl Column { pub fn bitand(&self, rhs: &Self) -> PolarsResult { // @partition-opt // @scalar-opt - self.as_materialized_series() - .bitand(rhs.as_materialized_series()) - .map(Column::from) + (self.as_materialized_series() & rhs.as_materialized_series()).map(Column::from) } pub fn bitor(&self, rhs: &Self) -> PolarsResult { // @partition-opt // @scalar-opt - self.as_materialized_series() - .bitor(rhs.as_materialized_series()) - .map(Column::from) + (self.as_materialized_series() | rhs.as_materialized_series()).map(Column::from) } pub fn bitxor(&self, rhs: &Self) -> PolarsResult { // @partition-opt // @scalar-opt - self.as_materialized_series() - .bitxor(rhs.as_materialized_series()) - .map(Column::from) + (self.as_materialized_series() ^ rhs.as_materialized_series()).map(Column::from) } pub fn try_add_owned(self, other: Self) -> PolarsResult { diff --git a/crates/polars-core/src/series/arithmetic/bitops.rs b/crates/polars-core/src/series/arithmetic/bitops.rs new file mode 100644 index 000000000000..cd00e8de18db --- /dev/null +++ b/crates/polars-core/src/series/arithmetic/bitops.rs @@ -0,0 +1,65 @@ +use std::borrow::Cow; + +use polars_error::PolarsResult; + +use super::{polars_bail, BooleanChunked, ChunkedArray, DataType, IntoSeries, Series}; + +macro_rules! impl_bitop { + ($(($trait:ident, $f:ident))+) => { + $( + impl std::ops::$trait for &Series { + type Output = PolarsResult; + fn $f(self, rhs: Self) -> Self::Output { + use DataType as DT; + match self.dtype() { + DT::Boolean => { + let lhs: &BooleanChunked = self.as_ref().as_ref().as_ref(); + let rhs = lhs.unpack_series_matching_type(rhs)?; + Ok(lhs.$f(rhs).into_series()) + }, + dt if dt.is_integer() => with_match_physical_integer_polars_type!(dt, |$T| { + let lhs: &ChunkedArray<$T> = self.as_ref().as_ref().as_ref(); + + let rhs = if rhs.len() == 1 { + Cow::Owned(rhs.cast(self.dtype())?) + } else { + Cow::Borrowed(rhs) + }; + + let rhs = lhs.unpack_series_matching_type(&rhs)?; + Ok(lhs.$f(&rhs).into_series()) + }), + _ => polars_bail!(opq = $f, self.dtype()), + } + } + } + impl std::ops::$trait for Series { + type Output = PolarsResult; + #[inline(always)] + fn $f(self, rhs: Self) -> Self::Output { + <&Series as std::ops::$trait>::$f(&self, &rhs) + } + } + impl std::ops::$trait<&Series> for Series { + type Output = PolarsResult; + #[inline(always)] + fn $f(self, rhs: &Series) -> Self::Output { + <&Series as std::ops::$trait>::$f(&self, rhs) + } + } + impl std::ops::$trait for &Series { + type Output = PolarsResult; + #[inline(always)] + fn $f(self, rhs: Series) -> Self::Output { + <&Series as std::ops::$trait>::$f(self, &rhs) + } + } + )+ + }; +} + +impl_bitop! { + (BitAnd, bitand) + (BitOr, bitor) + (BitXor, bitxor) +} diff --git a/crates/polars-core/src/series/arithmetic/mod.rs b/crates/polars-core/src/series/arithmetic/mod.rs index 0a5550b7b0f3..713bd4fbece3 100644 --- a/crates/polars-core/src/series/arithmetic/mod.rs +++ b/crates/polars-core/src/series/arithmetic/mod.rs @@ -1,3 +1,4 @@ +mod bitops; mod borrowed; mod list_borrowed; mod owned; diff --git a/crates/polars-core/src/series/implementations/boolean.rs b/crates/polars-core/src/series/implementations/boolean.rs index 83bbacd12c00..b4cd48295c4c 100644 --- a/crates/polars-core/src/series/implementations/boolean.rs +++ b/crates/polars-core/src/series/implementations/boolean.rs @@ -123,21 +123,6 @@ impl SeriesTrait for SeriesWrap { Some(self.0.boxed_metadata_dyn()) } - fn bitxor(&self, other: &Series) -> PolarsResult { - let other = self.0.unpack_series_matching_type(other)?; - Ok((&self.0).bitxor(other).into_series()) - } - - fn bitand(&self, other: &Series) -> PolarsResult { - let other = self.0.unpack_series_matching_type(other)?; - Ok((&self.0).bitand(other).into_series()) - } - - fn bitor(&self, other: &Series) -> PolarsResult { - let other = self.0.unpack_series_matching_type(other)?; - Ok((&self.0).bitor(other).into_series()) - } - fn rename(&mut self, name: PlSmallStr) { self.0.rename(name); } diff --git a/crates/polars-core/src/series/implementations/mod.rs b/crates/polars-core/src/series/implementations/mod.rs index b2cb97e39b69..1b430a289bb2 100644 --- a/crates/polars-core/src/series/implementations/mod.rs +++ b/crates/polars-core/src/series/implementations/mod.rs @@ -26,7 +26,6 @@ mod time; use std::any::Any; use std::borrow::Cow; -use std::ops::{BitAnd, BitOr, BitXor}; use std::sync::RwLockReadGuard; use super::*; @@ -259,36 +258,6 @@ macro_rules! impl_dyn_series { Some(self.0.boxed_metadata_dyn()) } - fn bitand(&self, other: &Series) -> PolarsResult { - let other = if other.len() == 1 { - Cow::Owned(other.cast(self.dtype())?) - } else { - Cow::Borrowed(other) - }; - let other = self.0.unpack_series_matching_type(&other)?; - Ok(self.0.bitand(&other).into_series()) - } - - fn bitor(&self, other: &Series) -> PolarsResult { - let other = if other.len() == 1 { - Cow::Owned(other.cast(self.dtype())?) - } else { - Cow::Borrowed(other) - }; - let other = self.0.unpack_series_matching_type(&other)?; - Ok(self.0.bitor(&other).into_series()) - } - - fn bitxor(&self, other: &Series) -> PolarsResult { - let other = if other.len() == 1 { - Cow::Owned(other.cast(self.dtype())?) - } else { - Cow::Borrowed(other) - }; - let other = self.0.unpack_series_matching_type(&other)?; - Ok(self.0.bitxor(&other).into_series()) - } - fn rename(&mut self, name: PlSmallStr) { self.0.rename(name); } diff --git a/crates/polars-core/src/series/series_trait.rs b/crates/polars-core/src/series/series_trait.rs index f31637bcce3d..8f6b26733c0f 100644 --- a/crates/polars-core/src/series/series_trait.rs +++ b/crates/polars-core/src/series/series_trait.rs @@ -218,18 +218,6 @@ pub trait SeriesTrait: /// Rename the Series. fn rename(&mut self, name: PlSmallStr); - fn bitand(&self, _other: &Series) -> PolarsResult { - polars_bail!(opq = bitand, self._dtype()); - } - - fn bitor(&self, _other: &Series) -> PolarsResult { - polars_bail!(opq = bitor, self._dtype()); - } - - fn bitxor(&self, _other: &Series) -> PolarsResult { - polars_bail!(opq = bitxor, self._dtype()); - } - fn get_metadata(&self) -> Option> { None } diff --git a/crates/polars-python/src/series/general.rs b/crates/polars-python/src/series/general.rs index f65822146d2c..b14285e77aa0 100644 --- a/crates/polars-python/src/series/general.rs +++ b/crates/polars-python/src/series/general.rs @@ -168,25 +168,15 @@ impl PySeries { } fn bitand(&self, other: &PySeries) -> PyResult { - let out = self - .series - .bitand(&other.series) - .map_err(PyPolarsErr::from)?; + let out = (&self.series & &other.series).map_err(PyPolarsErr::from)?; Ok(out.into()) } - fn bitor(&self, other: &PySeries) -> PyResult { - let out = self - .series - .bitor(&other.series) - .map_err(PyPolarsErr::from)?; + let out = (&self.series | &other.series).map_err(PyPolarsErr::from)?; Ok(out.into()) } fn bitxor(&self, other: &PySeries) -> PyResult { - let out = self - .series - .bitxor(&other.series) - .map_err(PyPolarsErr::from)?; + let out = (&self.series ^ &other.series).map_err(PyPolarsErr::from)?; Ok(out.into()) }