unique counts for logical types (#3694)

pola-rs · Jun 14, 2022 · 0307a57 · 0307a57
1 parent 86167e2
commit 0307a57
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 1 deletion.
diff --git a/polars/polars-core/src/series/ops/unique.rs b/polars/polars-core/src/series/ops/unique.rs
@@ -41,7 +41,7 @@ impl Series {
     #[cfg(feature = "unique_counts")]
     #[cfg_attr(docsrs, doc(cfg(feature = "unique_counts")))]
     pub fn unique_counts(&self) -> IdxCa {
-        if self.dtype().is_numeric() {
+        if self.dtype().to_physical().is_numeric() {
             if self.bit_repr_is_large() {
                 let ca = self.bit_repr_large();
                 unique_counts(ca.into_iter())

diff --git a/py-polars/tests/test_datelike.py b/py-polars/tests/test_datelike.py
@@ -981,3 +981,26 @@ def test_datetime_instance_selection() -> None:
 
     for tu in ["ns", "us", "ms"]:
         assert df.select(pl.col([pl.Datetime(tu)])).dtypes == [pl.Datetime(tu)]
+
+
+def test_unique_counts_on_dates() -> None:
+    assert pl.DataFrame(
+        {
+            "dt_ns": pl.date_range(datetime(2020, 1, 1), datetime(2020, 3, 1), "1mo"),
+        }
+    ).with_columns(
+        [
+            pl.col("dt_ns").dt.cast_time_unit("us").alias("dt_us"),
+            pl.col("dt_ns").dt.cast_time_unit("ms").alias("dt_ms"),
+            pl.col("dt_ns").cast(pl.Date).alias("date"),
+        ]
+    ).select(
+        pl.all().unique_counts().sum()
+    ).to_dict(
+        False
+    ) == {
+        "dt_ns": [3],
+        "dt_us": [3],
+        "dt_ms": [3],
+        "date": [3],
+    }