finish

pola-rs · Sep 30, 2024 · 0c45d90 · 0c45d90
1 parent f367006
commit 0c45d90
Show file tree

Hide file tree

Showing 40 changed files with 356 additions and 147 deletions.
diff --git a/crates/polars-compute/src/bitwise/mod.rs b/crates/polars-compute/src/bitwise/mod.rs
@@ -8,14 +8,14 @@ use bytemuck::Zeroable;
 pub trait BitwiseKernel {
     type Scalar;
 
-    fn count_ones(&self) -> PrimitiveArray<u8>;
-    fn count_zeros(&self) -> PrimitiveArray<u8>;
+    fn count_ones(&self) -> PrimitiveArray<u32>;
+    fn count_zeros(&self) -> PrimitiveArray<u32>;
 
-    fn leading_ones(&self) -> PrimitiveArray<u8>;
-    fn leading_zeros(&self) -> PrimitiveArray<u8>;
+    fn leading_ones(&self) -> PrimitiveArray<u32>;
+    fn leading_zeros(&self) -> PrimitiveArray<u32>;
 
-    fn trailing_ones(&self) -> PrimitiveArray<u8>;
-    fn trailing_zeros(&self) -> PrimitiveArray<u8>;
+    fn trailing_ones(&self) -> PrimitiveArray<u32>;
+    fn trailing_zeros(&self) -> PrimitiveArray<u32>;
 
     fn reduce_and(&self) -> Option<Self::Scalar>;
     fn reduce_or(&self) -> Option<Self::Scalar>;
@@ -33,77 +33,77 @@ macro_rules! impl_bitwise_kernel {
             type Scalar = $T;
 
             #[inline(never)]
-            fn count_ones(&self) -> PrimitiveArray<u8> {
+            fn count_ones(&self) -> PrimitiveArray<u32> {
                 PrimitiveArray::new(
-                    ArrowDataType::UInt8,
+                    ArrowDataType::UInt32,
                     self.values()
                         .iter()
-                        .map(|&v| ($to_bits(v).count_ones() & 0xFF) as u8)
+                        .map(|&v| $to_bits(v).count_ones())
                         .collect_trusted::<Vec<_>>()
                         .into(),
                     self.validity().cloned(),
                 )
             }
 
             #[inline(never)]
-            fn count_zeros(&self) -> PrimitiveArray<u8> {
+            fn count_zeros(&self) -> PrimitiveArray<u32> {
                 PrimitiveArray::new(
-                    ArrowDataType::UInt8,
+                    ArrowDataType::UInt32,
                     self
                         .values()
                         .iter()
-                        .map(|&v| ($to_bits(v).count_zeros() & 0xFF) as u8)
+                        .map(|&v| $to_bits(v).count_zeros())
                         .collect_trusted::<Vec<_>>()
                         .into(),
                     self.validity().cloned(),
                 )
             }
 
             #[inline(never)]
-            fn leading_ones(&self) -> PrimitiveArray<u8> {
+            fn leading_ones(&self) -> PrimitiveArray<u32> {
                 PrimitiveArray::new(
-                    ArrowDataType::UInt8,
+                    ArrowDataType::UInt32,
                     self.values()
                         .iter()
-                        .map(|&v| ($to_bits(v).leading_ones() & 0xFF) as u8)
+                        .map(|&v| $to_bits(v).leading_ones())
                         .collect_trusted::<Vec<_>>()
                         .into(),
                     self.validity().cloned(),
                 )
             }
 
             #[inline(never)]
-            fn leading_zeros(&self) -> PrimitiveArray<u8> {
+            fn leading_zeros(&self) -> PrimitiveArray<u32> {
                 PrimitiveArray::new(
-                    ArrowDataType::UInt8,
+                    ArrowDataType::UInt32,
                     self.values()
                         .iter()
-                        .map(|&v| ($to_bits(v).leading_zeros() & 0xFF) as u8)
+                        .map(|&v| $to_bits(v).leading_zeros())
                         .collect_trusted::<Vec<_>>()
                         .into(),
                     self.validity().cloned(),
                 )
             }
 
             #[inline(never)]
-            fn trailing_ones(&self) -> PrimitiveArray<u8> {
+            fn trailing_ones(&self) -> PrimitiveArray<u32> {
                 PrimitiveArray::new(
-                    ArrowDataType::UInt8,
+                    ArrowDataType::UInt32,
                     self.values()
                         .iter()
-                        .map(|&v| ($to_bits(v).trailing_ones() & 0xFF) as u8)
+                        .map(|&v| $to_bits(v).trailing_ones())
                         .collect_trusted::<Vec<_>>()
                         .into(),
                     self.validity().cloned(),
                 )
             }
 
             #[inline(never)]
-            fn trailing_zeros(&self) -> PrimitiveArray<u8> {
+            fn trailing_zeros(&self) -> PrimitiveArray<u32> {
                 PrimitiveArray::new(
-                    ArrowDataType::UInt8,
+                    ArrowDataType::UInt32,
                     self.values().iter()
-                        .map(|&v| ($to_bits(v).trailing_zeros() & 0xFF) as u8)
+                        .map(|&v| $to_bits(v).trailing_zeros())
                         .collect_trusted::<Vec<_>>()
                         .into(),
                     self.validity().cloned(),
@@ -186,48 +186,48 @@ impl BitwiseKernel for BooleanArray {
     type Scalar = bool;
 
     #[inline(never)]
-    fn count_ones(&self) -> PrimitiveArray<u8> {
+    fn count_ones(&self) -> PrimitiveArray<u32> {
         PrimitiveArray::new(
-            ArrowDataType::UInt8,
+            ArrowDataType::UInt32,
             self.values()
                 .iter()
-                .map(u8::from)
+                .map(u32::from)
                 .collect_trusted::<Vec<_>>()
                 .into(),
             self.validity().cloned(),
         )
     }
 
     #[inline(never)]
-    fn count_zeros(&self) -> PrimitiveArray<u8> {
+    fn count_zeros(&self) -> PrimitiveArray<u32> {
         PrimitiveArray::new(
-            ArrowDataType::UInt8,
+            ArrowDataType::UInt32,
             self.values()
                 .iter()
-                .map(|v| u8::from(!v))
+                .map(|v| u32::from(!v))
                 .collect_trusted::<Vec<_>>()
                 .into(),
             self.validity().cloned(),
         )
     }
 
     #[inline(always)]
-    fn leading_ones(&self) -> PrimitiveArray<u8> {
+    fn leading_ones(&self) -> PrimitiveArray<u32> {
         self.count_ones()
     }
 
     #[inline(always)]
-    fn leading_zeros(&self) -> PrimitiveArray<u8> {
+    fn leading_zeros(&self) -> PrimitiveArray<u32> {
         self.count_zeros()
     }
 
     #[inline(always)]
-    fn trailing_ones(&self) -> PrimitiveArray<u8> {
+    fn trailing_ones(&self) -> PrimitiveArray<u32> {
         self.count_ones()
     }
 
     #[inline(always)]
-    fn trailing_zeros(&self) -> PrimitiveArray<u8> {
+    fn trailing_zeros(&self) -> PrimitiveArray<u32> {
         self.count_zeros()
     }
 

diff --git a/crates/polars-core/Cargo.toml b/crates/polars-core/Cargo.toml
@@ -75,6 +75,7 @@ fmt_no_tty = ["comfy-table"]
 rows = []
 
 # operations
+bitwise = ["algorithm_group_by"]
 zip_with = []
 round_series = []
 checked_arithmetic = []

diff --git a/crates/polars-core/src/chunked_array/ops/mod.rs b/crates/polars-core/src/chunked_array/ops/mod.rs
@@ -9,6 +9,7 @@ pub(crate) mod append;
 mod apply;
 pub mod arity;
 mod bit_repr;
+#[cfg(feature = "bitwise")]
 mod bitwise_reduce;
 pub(crate) mod chunkops;
 pub(crate) mod compare_inner;
@@ -297,6 +298,7 @@ pub trait ChunkVar {
 }
 
 /// Bitwise Reduction Operations.
+#[cfg(feature = "bitwise")]
 pub trait ChunkBitwiseReduce {
     type Physical;
 

diff --git a/crates/polars-core/src/frame/column/mod.rs b/crates/polars-core/src/frame/column/mod.rs
@@ -525,6 +525,7 @@ impl Column {
     /// # Safety
     ///
     /// Does no bounds checks, groups must be correct.
+    #[cfg(feature = "algorithm_group_by")]
     pub unsafe fn agg_first(&self, groups: &GroupsProxy) -> Self {
         // @scalar-opt
         unsafe { self.as_materialized_series().agg_first(groups) }.into()
@@ -533,6 +534,7 @@ impl Column {
     /// # Safety
     ///
     /// Does no bounds checks, groups must be correct.
+    #[cfg(feature = "algorithm_group_by")]
     pub unsafe fn agg_last(&self, groups: &GroupsProxy) -> Self {
         // @scalar-opt
         unsafe { self.as_materialized_series().agg_last(groups) }.into()
@@ -541,6 +543,7 @@ impl Column {
     /// # Safety
     ///
     /// Does no bounds checks, groups must be correct.
+    #[cfg(feature = "algorithm_group_by")]
     pub unsafe fn agg_n_unique(&self, groups: &GroupsProxy) -> Self {
         // @scalar-opt
         unsafe { self.as_materialized_series().agg_n_unique(groups) }.into()
@@ -549,6 +552,7 @@ impl Column {
     /// # Safety
     ///
     /// Does no bounds checks, groups must be correct.
+    #[cfg(feature = "algorithm_group_by")]
     pub unsafe fn agg_quantile(
         &self,
         groups: &GroupsProxy,

diff --git a/crates/polars-core/src/frame/group_by/aggregations/boolean.rs b/crates/polars-core/src/frame/group_by/aggregations/boolean.rs
@@ -17,6 +17,7 @@ where
     ca.into_series()
 }
 
+#[cfg(feature = "bitwise")]
 unsafe fn bitwise_agg(
     ca: &BooleanChunked,
     groups: &GroupsProxy,
@@ -51,6 +52,7 @@ unsafe fn bitwise_agg(
     }
 }
 
+#[cfg(feature = "bitwise")]
 impl BooleanChunked {
     pub(crate) unsafe fn agg_and(&self, groups: &GroupsProxy) -> Series {
         bitwise_agg(self, groups, ChunkBitwiseReduce::and_reduce)
@@ -63,7 +65,9 @@ impl BooleanChunked {
     pub(crate) unsafe fn agg_xor(&self, groups: &GroupsProxy) -> Series {
         bitwise_agg(self, groups, ChunkBitwiseReduce::xor_reduce)
     }
+}
 
+impl BooleanChunked {
     pub(crate) unsafe fn agg_min(&self, groups: &GroupsProxy) -> Series {
         // faster paths
         match (self.is_sorted_flag(), self.null_count()) {

diff --git a/crates/polars-core/src/frame/group_by/aggregations/mod.rs b/crates/polars-core/src/frame/group_by/aggregations/mod.rs
@@ -458,6 +458,7 @@ where
 /// # Safety
 ///
 /// No bounds checks on `groups`.
+#[cfg(feature = "bitwise")]
 unsafe fn bitwise_agg<T: PolarsNumericType>(
     ca: &ChunkedArray<T>,
     groups: &GroupsProxy,
@@ -496,6 +497,7 @@ where
     }
 }
 
+#[cfg(feature = "bitwise")]
 impl<T> ChunkedArray<T>
 where
     T: PolarsNumericType,

diff --git a/crates/polars-core/src/frame/group_by/mod.rs b/crates/polars-core/src/frame/group_by/mod.rs
@@ -869,13 +869,17 @@ pub enum GroupByMethod {
     Groups,
     NUnique,
     Quantile(f64, QuantileInterpolOptions),
-    Count { include_nulls: bool },
+    Count {
+        include_nulls: bool,
+    },
     Implode,
     Std(u8),
     Var(u8),
+    #[cfg(feature = "bitwise")]
     Bitwise(GroupByBitwiseMethod),
 }
 
+#[cfg(feature = "bitwise")]
 #[derive(Copy, Clone, Debug)]
 pub enum GroupByBitwiseMethod {
     And,
@@ -903,6 +907,7 @@ impl Display for GroupByMethod {
             Implode => "list",
             Std(_) => "std",
             Var(_) => "var",
+            #[cfg(feature = "bitwise")]
             Bitwise(t) => {
                 f.write_str("bitwise_")?;
                 return Display::fmt(t, f);
@@ -912,6 +917,7 @@ impl Display for GroupByMethod {
     }
 }
 
+#[cfg(feature = "bitwise")]
 impl Display for GroupByBitwiseMethod {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         match self {
@@ -942,7 +948,8 @@ pub fn fmt_group_by_column(name: &str, method: GroupByMethod) -> PlSmallStr {
         Quantile(quantile, _interpol) => format_pl_smallstr!("{name}_quantile_{quantile:.2}"),
         Std(_) => format_pl_smallstr!("{name}_agg_std"),
         Var(_) => format_pl_smallstr!("{name}_agg_var"),
-        Bitwise(_) => format_pl_smallstr!("{name}_agg_var"),
+        #[cfg(feature = "bitwise")]
+        Bitwise(f) => format_pl_smallstr!("{name}_agg_bitwise_{f}"),
     }
 }
 

diff --git a/crates/polars-core/src/series/implementations/boolean.rs b/crates/polars-core/src/series/implementations/boolean.rs
@@ -84,12 +84,15 @@ impl private::PrivateSeries for SeriesWrap<BooleanChunked> {
             .agg_var(groups, _ddof)
     }
 
+    #[cfg(feature = "bitwise")]
     unsafe fn agg_and(&self, groups: &GroupsProxy) -> Series {
         self.0.agg_and(groups)
     }
+    #[cfg(feature = "bitwise")]
     unsafe fn agg_or(&self, groups: &GroupsProxy) -> Series {
         self.0.agg_or(groups)
     }
+    #[cfg(feature = "bitwise")]
     unsafe fn agg_xor(&self, groups: &GroupsProxy) -> Series {
         self.0.agg_xor(groups)
     }