Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 13 additions & 11 deletions src/datasure/processing/prep.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()
),
Expand Down
30 changes: 23 additions & 7 deletions tests/processing/test_prep.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand All @@ -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."""
Expand Down Expand Up @@ -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]."""
Expand Down Expand Up @@ -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(
Expand All @@ -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."""
Expand Down
Loading