Support null dimension categories in simple.AdditiveAction

kausaltech · Feb 24, 2025 · d2a49a9 · d2a49a9
1 parent e78f2f2
commit d2a49a9
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 2 deletions.
diff --git a/nodes/actions/action.py b/nodes/actions/action.py
@@ -79,6 +79,7 @@ class ActionNode(Node):
     allowed_parameters: ClassVar[Sequence[Parameter]] = [
         ENABLED_PARAM,
         NumberParameter(local_id='multiplier', label='Multiplies the output', is_customizable=True),
+        BoolParameter(local_id='allow_null_categories', description='Allow null dimension categories', is_customizable=False),
     ]
 
     def __init_subclass__(cls) -> None:

diff --git a/nodes/actions/simple.py b/nodes/actions/simple.py
@@ -23,6 +23,9 @@ class AdditiveAction(ActionNode):
     def compute_effect(self):
         df = self.get_input_dataset_pl()
 
+        if self.get_parameter_value('allow_null_categories', required=False):
+            self.allow_null_categories = True
+
         for m in self.output_metrics.values():
             if not self.is_enabled():
                 df = df.with_columns(pl.when(pl.col(m.column_id).is_null()).then(None)

diff --git a/nodes/node.py b/nodes/node.py
@@ -866,7 +866,7 @@ def _get_output_for_target(self, df: ppl.PathsDataFrame, target_node: Node) -> p
                     raise NodeError(self, 'Dimension %s not in output df' % dim_id)
                 filter_expr = pl.col(dim_id).is_in(cat_ids)
                 if edge_dim.exclude:
-                    filter_expr = ~filter_expr
+                    filter_expr = pl.col(dim_id).is_null() | ~filter_expr
                 df = df.filter(filter_expr)
                 if len(df) == 0:
                     raise NodeError(self, 'No rows left after filtering by %s' % dim_id)
@@ -990,7 +990,7 @@ def validate_output(self, df: ppl.PathsDataFrame) -> None:  # noqa: C901, PLR091
 
             cats = set(df[dim_id].cast(pl.Utf8).unique())
             if getattr(self, 'allow_null_categories', None):
-                cats -= {''}
+                cats -= {'', None}
 
             dim_cats = dim.get_cat_ids()
             diff = cats - dim_cats