From f9ad1e895f27e2f1ac3242d40cb900c4e405b03c Mon Sep 17 00:00:00 2001
From: iabaako <iabaako@poverty-action.org>
Date: Fri, 7 Aug 2026 11:38:24 +0000
Subject: [PATCH] bugfix/invertion in prep steps for equal & not equal

---
 src/datasure/processing/prep.py | 24 +++++++++++++-----------
 tests/processing/test_prep.py   | 30 +++++++++++++++++++++++-------
 2 files changed, 36 insertions(+), 18 deletions(-)

diff --git a/src/datasure/processing/prep.py b/src/datasure/processing/prep.py
index d78c8dcf..16c4d8cc 100644
--- a/src/datasure/processing/prep.py
+++ b/src/datasure/processing/prep.py
@@ -323,18 +323,18 @@ def _filter_by_equality(
         # Ensure value is a list for is_in() to treat as literal values
         value_list = value if isinstance(value, list) else [value]
 
+        filter_expr = pl.any_horizontal(
+            [pl.col(col).is_in(value_list) for col in columns]
+        )
+
         if condition == PrepRowConditions.not_equal_to.value:
-            # Keep rows where value is NOT in the list (remove matching rows)
-            filter_expr = pl.any_horizontal(
-                [pl.col(col).is_in(value_list) for col in columns]
-            )
-            return data.filter(~filter_expr)
-        else:
-            # Keep rows where value IS in the list (remove non-matching rows)
-            filter_expr = pl.any_horizontal(
-                [pl.col(col).is_in(value_list) for col in columns]
-            )
+            # "Remove rows where value is not equal to X" - keep only the
+            # matching rows, i.e. drop everything the filter doesn't match.
             return data.filter(filter_expr)
+        else:
+            # "Remove rows where value is equal to X" - keep everything
+            # that doesn't match.
+            return data.filter(~filter_expr)
 
     def _filter_by_comparison(
         self, data: pl.DataFrame, condition: str, columns: list[str], value: Any
@@ -826,7 +826,9 @@ def _add_computed_column(
             PrepFunctions.var.value: lambda cols: pl.concat_list(cols).list.var(),
             PrepFunctions.first.value: lambda cols: pl.concat_list(cols).list.first(),
             PrepFunctions.last.value: lambda cols: pl.concat_list(cols).list.last(),
-            PrepFunctions.count.value: lambda cols: pl.concat_list(cols).list.len(),
+            PrepFunctions.count.value: lambda cols: (
+                pl.concat_list(cols).list.drop_nulls().list.len()
+            ),
             PrepFunctions.nunique.value: lambda cols: (
                 pl.concat_list(cols).list.unique().list.len()
             ),
diff --git a/tests/processing/test_prep.py b/tests/processing/test_prep.py
index d0a88589..2ace6042 100644
--- a/tests/processing/test_prep.py
+++ b/tests/processing/test_prep.py
@@ -311,8 +311,8 @@ def test_remove_by_condition_equal_to(self):
             value=[2],
         )
         result, _ = op.execute(data, prep_args)
-        # equal_to keeps matching rows (removes non-matching)
-        assert result.shape[0] == 2
+        # equal_to removes matching rows, keeps everything else
+        assert sorted(result["a"].to_list()) == [1, 3, 5]
 
     def test_remove_by_condition_not_equal_to(self):
         """Test Remove by condition not equal to."""
@@ -326,8 +326,8 @@ def test_remove_by_condition_not_equal_to(self):
             value=[2],
         )
         result, _ = op.execute(data, prep_args)
-        # not_equal_to removes matching rows
-        assert result.shape[0] == 3
+        # not_equal_to removes non-matching rows, keeps only matches
+        assert result["a"].to_list() == [2, 2]
 
     def test_remove_by_condition_greater_than(self):
         """Test Remove by condition greater than."""
@@ -516,7 +516,8 @@ def test_filter_by_equality_with_list(self):
             value=[2, 4],
         )
         result, _ = op.execute(data, prep_args)
-        assert sorted(result["a"].to_list()) == [2, 4]
+        # removes rows matching any value in the list, keeps the rest
+        assert sorted(result["a"].to_list()) == [1, 3, 5]
 
     def test_filter_by_range_single_value(self):
         """Test range filter with single value (not a list) - uses [val, val]."""
@@ -1264,7 +1265,7 @@ def test_add_last_column(self):
         assert result["l"].to_list() == [3.0, 4.0]
 
     def test_add_count_column(self):
-        """Test Add count column."""
+        """Count should tally non-null values, not just the column count."""
         op = AddNewColumnOperation()
         data = pl.DataFrame({"a": [1.0, None], "b": [3.0, 4.0]})
         prep_args = PrepActionResult(
@@ -1274,7 +1275,22 @@ def test_add_count_column(self):
             source_columns=["a", "b"],
         )
         result, _ = op.execute(data, prep_args)
-        assert result["cnt"].to_list() == [2, 2]
+        assert result["cnt"].to_list() == [2, 1]
+
+    def test_add_count_column_all_missing(self):
+        """A row with no non-null values across the source columns counts 0."""
+        op = AddNewColumnOperation()
+        data = pl.DataFrame(
+            {"a": [1.0, None, None], "b": [3.0, None, 4.0], "c": [5.0, None, None]}
+        )
+        prep_args = PrepActionResult(
+            action="add new column",
+            column_names="cnt",
+            method="count",
+            source_columns=["a", "b", "c"],
+        )
+        result, _ = op.execute(data, prep_args)
+        assert result["cnt"].to_list() == [3, 0, 1]
 
     def test_add_nunique_column(self):
         """Test Add nunique column."""