diff --git a/src/nested_pandas/nestedframe/core.py b/src/nested_pandas/nestedframe/core.py index 6142238..b3be03e 100644 --- a/src/nested_pandas/nestedframe/core.py +++ b/src/nested_pandas/nestedframe/core.py @@ -201,7 +201,6 @@ def _resolve_dropna_target(self, on_nested, subset): """resolves the target layer for a given set of dropna kwargs""" nested_cols = self.nested_columns - columns = self.columns # first check the subset kwarg input subset_target = [] @@ -210,13 +209,15 @@ def _resolve_dropna_target(self, on_nested, subset): subset = [subset] for col in subset: - col = col.split(".")[0] - if col in nested_cols: - subset_target.append(col) - elif col in columns: + # Without a ".", always assume base layer + if "." not in col: subset_target.append("base") else: - raise ValueError(f"Column name {col} not found in any base or nested columns") + layer, col = col.split(".") + if layer in nested_cols: + subset_target.append(layer) + else: + raise ValueError(f"layer '{layer}' not found in the base columns") # Check for 1 target subset_target = np.unique(subset_target) diff --git a/tests/nested_pandas/nestedframe/test_nestedframe.py b/tests/nested_pandas/nestedframe/test_nestedframe.py index a9c9e15..fffbbe9 100644 --- a/tests/nested_pandas/nestedframe/test_nestedframe.py +++ b/tests/nested_pandas/nestedframe/test_nestedframe.py @@ -3,6 +3,7 @@ import pyarrow as pa import pytest from nested_pandas import NestedFrame +from nested_pandas.datasets import generate_data from pandas.testing import assert_frame_equal @@ -203,6 +204,15 @@ def test_dropna(): assert len(dn_hierarchical["nested"].nest.to_flat() == 8) +def test_dropna_layer_as_base_column(): + """Test that a nested column still works as a top level column for dropna""" + nf = generate_data(10, 100, seed=1).query("nested.t>19.75") + nf = nf.dropna(subset=["nested"]) + + # make sure rows have been dropped as expected + assert len(nf) == 6 + + def test_dropna_inplace_base(): """Test in-place behavior of dropna"""