Merge pull request #630 from Aarhus-Psychiatry-Research/bokajgd/issue629

feat: add new UniqueCountAggregator

requirements-dev.lock

@@ -84,7 +84,7 @@ pytest==7.2.2
     # via pytest-testmon
     # via pytest-xdist
 pytest-benchmark==4.0.0
-pytest-codspeed==2.2.0
+pytest-codspeed==3.0.1
 pytest-cov==3.0.0
 pytest-sugar==0.9.7
 pytest-testmon==2.1.0

src/timeseriesflattener/aggregators.py

@@ -23,6 +23,7 @@ def validate_compatible_fallback_type_for_aggregator(
     "bool",
     "change_per_day",
     "count",
+    "unique_count",
     "has_values",
     "max",
     "mean",
@@ -47,6 +48,7 @@ def string_to_aggregator(aggregator_name: AggregatorName, timestamp_col_name: st
         "bool": HasValuesAggregator(),
         "change_per_day": SlopeAggregator(timestamp_col_name=timestamp_col_name),
         "count": CountAggregator(),
+        "unique_count": UniqueCountAggregator(),
         "has_values": HasValuesAggregator(),
         "max": MaxAggregator(),
         "mean": MeanAggregator(),
@@ -64,8 +66,7 @@ class Aggregator(ABC):
     output_type: type[float | int | bool]
 
     @abstractmethod
-    def __call__(self, column_name: str) -> pl.Expr:
-        ...
+    def __call__(self, column_name: str) -> pl.Expr: ...
 
     def new_col_name(self, previous_col_name: str) -> str:
         return f"{previous_col_name}_{self.name}"
@@ -111,6 +112,16 @@ def __call__(self, column_name: str) -> pl.Expr:
         return pl.col(column_name).count().alias(self.new_col_name(column_name))
 
 
+class UniqueCountAggregator(Aggregator):
+    """Returns the count of non-null values in the look window."""
+
+    name: str = "unique_count"
+    output_type = int
+
+    def __call__(self, column_name: str) -> pl.Expr:
+        return pl.col(column_name).n_unique().alias(self.new_col_name(column_name))
+
+
 @dataclass(frozen=True)
 class EarliestAggregator(Aggregator):
     """Returns the earliest value in the look window."""

src/timeseriesflattener/aggregators_test.py

@@ -11,6 +11,7 @@
 from .intermediary import TimeMaskedFrame
 from .aggregators import (
     CountAggregator,
+    UniqueCountAggregator,
     EarliestAggregator,
     HasValuesAggregator,
     LatestAggregator,
@@ -85,6 +86,9 @@ def expected_output(self) -> pl.DataFrame:
         SingleVarAggregatorExample(
             aggregator=CountAggregator(), input_values=[1, 2], expected_output_values=[2]
         ),
+        SingleVarAggregatorExample(
+            aggregator=UniqueCountAggregator(), input_values=[1, 2, 1], expected_output_values=[2]
+        ),
         SingleVarAggregatorExample(
             aggregator=SumAggregator(), input_values=[1, 2], expected_output_values=[3]
         ),