feat[rust]: make clip expr inplace when owned (#4530)

pola-rs · Aug 22, 2022 · 0c9086c · 0c9086c
1 parent a92d61d
commit 0c9086c
Show file tree

Hide file tree

Showing 2 changed files with 73 additions and 31 deletions.
diff --git a/polars/polars-core/src/series/ops/round.rs b/polars/polars-core/src/series/ops/round.rs
@@ -60,37 +60,23 @@ impl Series {
 
     #[cfg_attr(docsrs, doc(cfg(feature = "round_series")))]
     /// Ceil underlying floating point array to the highest integers smaller or equal to the float value.
-    pub fn clip(&self, min: f64, max: f64) -> Result<Self> {
-        if let Ok(ca) = self.f32() {
-            let min = min as f32;
-            let max = max as f32;
-            let s = ca.apply(|val| val.clamp(min, max)).into_series();
-            return Ok(s);
-        }
-        if let Ok(ca) = self.f64() {
-            let s = ca.apply(|val| val.clamp(min, max)).into_series();
-            return Ok(s);
-        }
-        if let Ok(ca) = self.i64() {
-            let min = min as i64;
-            let max = max as i64;
-            let s = ca.apply(|val| val.clamp(min, max)).into_series();
-            return Ok(s);
-        }
-        if let Ok(ca) = self.i32() {
-            let min = min as i32;
-            let max = max as i32;
-            let s = ca.apply(|val| val.clamp(min, max)).into_series();
-            return Ok(s);
-        }
-        if let Ok(ca) = self.u32() {
-            let min = min as u32;
-            let max = max as u32;
-            let s = ca.apply(|val| val.clamp(min, max)).into_series();
-            return Ok(s);
+    pub fn clip(mut self, min: f64, max: f64) -> Result<Self> {
+        if self.dtype().is_numeric() {
+            macro_rules! apply_clip {
+                ($pl_type:ty, $ca:expr, $min:expr, $max: expr) => {{
+                    let min = min as <$pl_type as PolarsNumericType>::Native;
+                    let max = max as <$pl_type as PolarsNumericType>::Native;
+
+                    $ca.apply_mut(|val| val.clamp(min, max));
+                }};
+            }
+            let mutable = self._get_inner_mut();
+            downcast_as_macro_arg_physical_mut!(mutable, apply_clip, min, max);
+            Ok(self)
+        } else {
+            Err(PolarsError::SchemaMisMatch(
+                format!("Cannot use 'clip' on dtype {:?}, consider using a when -> then -> otherwise expression", self.dtype()).into(),
+            ))
         }
-        Err(PolarsError::SchemaMisMatch(
-            format!("{:?} is not one of {{Float32, Float64, Int32, Int64, UInt32}} consider using a when -> then -> otherwise", self.dtype()).into(),
-        ))
     }
 }
diff --git a/polars/polars-core/src/utils/mod.rs b/polars/polars-core/src/utils/mod.rs
@@ -327,6 +327,62 @@ macro_rules! downcast_as_macro_arg_physical {
     }};
 }
 
+/// Apply a macro on the Downcasted ChunkedArray's of DataTypes that are logical numerics.
+/// So no logical.
+#[macro_export]
+macro_rules! downcast_as_macro_arg_physical_mut {
+    ($self:expr, $macro:ident $(, $opt_args:expr)*) => {{
+        // clone so that we do not borrow
+        match $self.dtype().clone() {
+            #[cfg(feature = "dtype-u8")]
+            DataType::UInt8 => {
+                let ca: &mut UInt8Chunked = $self.as_mut();
+                $macro!(UInt8Type, ca $(, $opt_args)*)
+            },
+            #[cfg(feature = "dtype-u16")]
+            DataType::UInt16 => {
+                let ca: &mut UInt16Chunked = $self.as_mut();
+                $macro!(UInt16Type, ca $(, $opt_args)*)
+            },
+            DataType::UInt32 => {
+                let ca: &mut UInt32Chunked = $self.as_mut();
+                $macro!(UInt32Type, ca $(, $opt_args)*)
+            },
+            DataType::UInt64 => {
+                let ca: &mut UInt64Chunked = $self.as_mut();
+                $macro!(UInt64Type, ca $(, $opt_args)*)
+            },
+            #[cfg(feature = "dtype-i8")]
+            DataType::Int8 => {
+                let ca: &mut Int8Chunked = $self.as_mut();
+                $macro!(Int8Type, ca $(, $opt_args)*)
+            },
+            #[cfg(feature = "dtype-i16")]
+            DataType::Int16 => {
+                let ca: &mut Int16Chunked = $self.as_mut();
+                $macro!(Int16Type, ca $(, $opt_args)*)
+            },
+            DataType::Int32 => {
+                let ca: &mut Int32Chunked = $self.as_mut();
+                $macro!(Int32Type, ca $(, $opt_args)*)
+            },
+            DataType::Int64 => {
+                let ca: &mut Int64Chunked = $self.as_mut();
+                $macro!(Int64Type, ca $(, $opt_args)*)
+            },
+            DataType::Float32 => {
+                let ca: &mut Float32Chunked = $self.as_mut();
+                $macro!(Float32Type, ca $(, $opt_args)*)
+            },
+            DataType::Float64 => {
+                let ca: &mut Float64Chunked = $self.as_mut();
+                $macro!(Float64Type, ca $(, $opt_args)*)
+            },
+            dt => panic!("not implemented for {:?}", dt),
+        }
+    }};
+}
+
 #[macro_export]
 macro_rules! apply_method_all_arrow_series {
     ($self:expr, $method:ident, $($args:expr),*) => {