feat: add argument maintain_order in $list$unique()

pola-rs · Aug 24, 2024 · f7a3b3e · f7a3b3e
1 parent 7588ca3
commit f7a3b3e
Show file tree

Hide file tree

Showing 5 changed files with 21 additions and 10 deletions.
diff --git a/NEWS.md b/NEWS.md
@@ -56,6 +56,7 @@
 - `$glimpse()` for `DataFrame` has two new arguments `max_items_per_column` and
   `max_colname_length` (#1200).
 - New method `$list$sample()` (#1204).
+- New argument `maintain_order` in `$list$unique()` (#1207).
 
 ### Other changes
 

diff --git a/R/expr__list.R b/R/expr__list.R
@@ -67,12 +67,17 @@ ExprList_reverse = function() .pr$Expr$list_reverse(self)
 
 #' Get unique values in a list
 #'
+#' @param maintain_order Maintain order of data. This requires more work.
+#'
 #' @return Expr
 #'
 #' @examples
 #' df = pl$DataFrame(values = list(c(2, 2, NA), c(1, 2, 3), NA_real_))
 #' df$with_columns(unique = pl$col("values")$list$unique())
-ExprList_unique = function() .pr$Expr$list_unique(self)
+ExprList_unique = function(maintain_order = FALSE) {
+  .pr$Expr$list_unique(self, maintain_order) |>
+    unwrap("in $list$unique():")
+}
 
 #' Get the number of unique values in a list
 #'

diff --git a/R/extendr-wrappers.R b/R/extendr-wrappers.R
@@ -736,7 +736,7 @@ RPolarsExpr$list_sort <- function(descending) .Call(wrap__RPolarsExpr__list_sort
 
 RPolarsExpr$list_reverse <- function() .Call(wrap__RPolarsExpr__list_reverse, self)
 
-RPolarsExpr$list_unique <- function() .Call(wrap__RPolarsExpr__list_unique, self)
+RPolarsExpr$list_unique <- function(maintain_order) .Call(wrap__RPolarsExpr__list_unique, self, maintain_order)
 
 RPolarsExpr$list_n_unique <- function() .Call(wrap__RPolarsExpr__list_n_unique, self)
 

diff --git a/man/ExprList_unique.Rd b/man/ExprList_unique.Rd
diff --git a/src/rust/src/lazy/dsl.rs b/src/rust/src/lazy/dsl.rs
@@ -1156,13 +1156,15 @@ impl RPolarsExpr {
             .into()
     }
 
-    fn list_unique(&self) -> Self {
-        self.0
-            .clone()
-            .list()
-            .unique()
-            .with_fmt("list.unique")
-            .into()
+    fn list_unique(&self, maintain_order: Robj) -> RResult<Self> {
+        let e = self.0.clone();
+        let maintain_order = robj_to!(bool, maintain_order)?;
+        let out = if maintain_order {
+            e.list().unique_stable().into()
+        } else {
+            e.list().unique().into()
+        };
+        Ok(out)
     }
 
     fn list_n_unique(&self) -> Self {