fix(rust, python): unique include null (#5112)

pola-rs · Oct 5, 2022 · 2991849 · 2991849
1 parent 32bcacc
commit 2991849
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 2 deletions.
diff --git a/polars/polars-core/src/chunked_array/ops/unique/mod.rs b/polars/polars-core/src/chunked_array/ops/unique/mod.rs
@@ -185,8 +185,38 @@ where
         }
         match self.is_sorted2() {
             IsSorted::Ascending | IsSorted::Descending => {
-                let mask = self.not_equal(&self.shift(1));
-                self.filter(&mask)
+                // TODO! optimize this branch
+                if self.null_count() > 0 {
+                    let mut arr = MutablePrimitiveArray::with_capacity(self.len());
+                    let mut iter = self.into_iter();
+                    let mut last = None;
+
+                    if let Some(val) = iter.next() {
+                        last = val;
+                        arr.push(val)
+                    };
+
+                    #[allow(clippy::unnecessary_filter_map)]
+                    let to_extend = iter.filter_map(|opt_val| {
+                        if opt_val != last {
+                            last = opt_val;
+                            Some(opt_val)
+                        } else {
+                            None
+                        }
+                    });
+
+                    arr.extend(to_extend);
+                    let arr: PrimitiveArray<T::Native> = arr.into();
+
+                    Ok(ChunkedArray::from_chunks(
+                        self.name(),
+                        vec![Box::new(arr) as ArrayRef],
+                    ))
+                } else {
+                    let mask = self.not_equal(&self.shift(1));
+                    self.filter(&mask)
+                }
             }
             IsSorted::Not => {
                 let sorted = self.sort(false);

diff --git a/py-polars/tests/unit/test_lazy.py b/py-polars/tests/unit/test_lazy.py
@@ -1221,6 +1221,9 @@ def test_unique() -> None:
         .collect()
         .frame_equal(expected)
     )
+    s0 = pl.Series("a", [1, 2, None, 2])
+    # test if the null is included
+    assert s0.unique().to_list() == [None, 1, 2]
 
 
 def test_lazy_concat(df: pl.DataFrame) -> None: