add n_unique aggregation for floats

pola-rs · Oct 23, 2021 · 834214b · 834214b
1 parent 486a86b
commit 834214b
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 4 deletions.
diff --git a/polars/polars-core/src/frame/groupby/aggregations.rs b/polars/polars-core/src/frame/groupby/aggregations.rs
@@ -418,9 +418,16 @@ where
     }
 }
 
-// todo! could use mantissa method here
-impl AggNUnique for Float32Chunked {}
-impl AggNUnique for Float64Chunked {}
+impl AggNUnique for Float32Chunked {
+    fn agg_n_unique(&self, groups: &[(u32, Vec<u32>)]) -> Option<UInt32Chunked> {
+        self.bit_repr_small().agg_n_unique(groups)
+    }
+}
+impl AggNUnique for Float64Chunked {
+    fn agg_n_unique(&self, groups: &[(u32, Vec<u32>)]) -> Option<UInt32Chunked> {
+        self.bit_repr_large().agg_n_unique(groups)
+    }
+}
 impl AggNUnique for ListChunked {}
 #[cfg(feature = "dtype-categorical")]
 impl AggNUnique for CategoricalChunked {

diff --git a/polars/polars-core/src/frame/groupby/mod.rs b/polars/polars-core/src/frame/groupby/mod.rs
@@ -1323,7 +1323,7 @@ mod test {
         let gb = df.groupby("date").unwrap().n_unique().unwrap();
         println!("{:?}", df.groupby("date").unwrap().n_unique().unwrap());
         // check the group by column is filtered out.
-        assert_eq!(gb.width(), 2);
+        assert_eq!(gb.width(), 3);
         println!(
             "{:?}",
             df.groupby("date")

diff --git a/py-polars/tests/test_df.py b/py-polars/tests/test_df.py
@@ -1141,3 +1141,14 @@ def test_asof_join():
         720.92,
         51.95,
     ]
+
+
+def test_groupby_agg_n_unique_floats():
+    # tests proper dispatch
+    df = pl.DataFrame({"a": [1, 1, 3], "b": [1.0, 2.0, 2.0]})
+
+    for dtype in [pl.Float32, pl.Float64]:
+        out = df.groupby("a", maintain_order=True).agg(
+            [pl.col("b").cast(dtype).n_unique()]
+        )
+        out["b_n_unique"].to_list() == [2, 1]