diff --git a/narwhals/_duckdb/dataframe.py b/narwhals/_duckdb/dataframe.py index 5a680e5bb..c446351f1 100644 --- a/narwhals/_duckdb/dataframe.py +++ b/narwhals/_duckdb/dataframe.py @@ -326,7 +326,13 @@ def drop_nulls(self: Self, subset: str | list[str] | None) -> Self: import duckdb rel = self._native_frame - subset_ = subset if subset is not None else rel.columns - keep_condition = " and ".join(f"{col} is not null" for col in subset_) + subset_ = ( + [subset] + if isinstance(subset, str) + else rel.columns + if subset is None + else subset + ) + keep_condition = " and ".join(f'"{col}" is not null' for col in subset_) query = f"""select * from rel where {keep_condition}""" # noqa: S608 return self._from_native_frame(duckdb.sql(query)) diff --git a/tests/frame/drop_nulls_test.py b/tests/frame/drop_nulls_test.py index 451b1cbd7..c49b17126 100644 --- a/tests/frame/drop_nulls_test.py +++ b/tests/frame/drop_nulls_test.py @@ -7,16 +7,16 @@ from tests.utils import assert_equal_data data = { - "a": [1.0, 2.0, None, 4.0], - "b": [None, 3.0, None, 5.0], + "alpha": [1.0, 2.0, None, 4.0], + "beta gamma": [None, 3.0, None, 5.0], } def test_drop_nulls(constructor: Constructor) -> None: result = nw.from_native(constructor(data)).drop_nulls() expected = { - "a": [2.0, 4.0], - "b": [3.0, 5.0], + "alpha": [2.0, 4.0], + "beta gamma": [3.0, 5.0], } assert_equal_data(result, expected) @@ -24,9 +24,9 @@ def test_drop_nulls(constructor: Constructor) -> None: @pytest.mark.parametrize( ("subset", "expected"), [ - ("a", {"a": [1, 2.0, 4.0], "b": [None, 3.0, 5.0]}), - (["a"], {"a": [1, 2.0, 4.0], "b": [None, 3.0, 5.0]}), - (["a", "b"], {"a": [2.0, 4.0], "b": [3.0, 5.0]}), + ("alpha", {"alpha": [1, 2.0, 4.0], "beta gamma": [None, 3.0, 5.0]}), + (["alpha"], {"alpha": [1, 2.0, 4.0], "beta gamma": [None, 3.0, 5.0]}), + (["alpha", "beta gamma"], {"alpha": [2.0, 4.0], "beta gamma": [3.0, 5.0]}), ], ) def test_drop_nulls_subset(