Merge pull request #105 from lincc-frameworks/dropna_layer_as_base

allow dropna on nested columns as base columns

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)

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"""