arg_min/max to dsl

pola-rs · Jan 10, 2022 · 599c589 · 599c589
1 parent 2bbc851
commit 599c589
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 14 deletions.
diff --git a/polars/polars-lazy/src/dsl.rs b/polars/polars-lazy/src/dsl.rs
@@ -838,6 +838,22 @@ impl Expr {
         )
     }
 
+    /// Get the index value that has the minumum value
+    pub fn arg_min(self) -> Self {
+        self.apply(
+            |s| Ok(Series::new(s.name(), &[s.arg_min().map(|idx| idx as u32)])),
+            GetOutput::from_type(DataType::UInt32),
+        )
+    }
+
+    /// Get the index value that has the maximum value
+    pub fn arg_max(self) -> Self {
+        self.apply(
+            |s| Ok(Series::new(s.name(), &[s.arg_max().map(|idx| idx as u32)])),
+            GetOutput::from_type(DataType::UInt32),
+        )
+    }
+
     /// Get the index values that would sort this expression.
     pub fn arg_sort(self, reverse: bool) -> Self {
         assert!(

diff --git a/py-polars/src/lazy/dsl.rs b/py-polars/src/lazy/dsl.rs
@@ -179,22 +179,10 @@ impl PyExpr {
         self.clone().inner.arg_sort(reverse).into()
     }
     pub fn arg_max(&self) -> PyExpr {
-        self.clone()
-            .inner
-            .apply(
-                |s| Ok(Series::new(s.name(), &[s.arg_max().map(|idx| idx as u32)])),
-                GetOutput::from_type(DataType::UInt32),
-            )
-            .into()
+        self.clone().inner.arg_max().into()
     }
     pub fn arg_min(&self) -> PyExpr {
-        self.clone()
-            .inner
-            .apply(
-                |s| Ok(Series::new(s.name(), &[s.arg_min().map(|idx| idx as u32)])),
-                GetOutput::from_type(DataType::UInt32),
-            )
-            .into()
+        self.clone().inner.arg_min().into()
     }
     pub fn take(&self, idx: PyExpr) -> PyExpr {
         self.clone().inner.take(idx.inner).into()