Repository: lincc-frameworks/nested-pandas

Issue:
Allow Dropna to work on nested columns as base columns
**Bug report**
Current dropna logic assumes that if the user invokes a nested layer, that it then must be trying to dropna within the rows of that nested layer. We should fix dropna to also work with nested columns as base columns, which is particularly useful when filtering out base rows that have empty data in a nested layer. Here's a reproducible sample:

```
# Test dropna on nested columns
import nested_pandas as npd
from nested_pandas.datasets import generate_data

nf = generate_data(10,100, seed=1).query("nested.t>19.75")
nf.dropna(subset=["nested"])
```

which currently fails:
```
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
/var/folders/lc/dws63_cs5gz5mf8s869hjpx40000gn/T/ipykernel_52322/3221593222.py in ?()
      2 import nested_pandas as npd
      3 
      4 from nested_pandas.datasets import generate_data
      5 nf = generate_data(10,100, seed=1).query("nested.t>19.75")
----> 6 nf.dropna(subset=["nested"])

~/lincc/nested-pandas/src/nested_pandas/nestedframe/core.py in ?(self, axis, how, thresh, on_nested, subset, inplace, ignore_index)
    335             return self
    336         # Or if not inplace
    337         new_df = self.copy()
    338         new_df[target] = packer.pack_flat(
--> 339             new_df[target]
    340             .nest.to_flat()
    341             .dropna(
    342                 axis=axis,

~/miniforge3/envs/ray310/lib/python3.10/site-packages/pandas/core/frame.py in ?(self, axis, how, thresh, subset, inplace, ignore_index)
   6653             ax = self._get_axis(agg_axis)
   6654             indices = ax.get_indexer_for(subset)
   6655             check = indices == -1
   6656             if check.any():
-> 6657                 raise KeyError(np.array(subset)[check].tolist())
   6658             agg_obj = self.take(indices, axis=agg_axis)
   6659 
   6660         if thresh is not lib.no_default:

KeyError: ['nested']
```

**Before submitting**
Please check the following:

- [x] I have described the situation in which the bug arose, including what code was executed, information about my environment, and any applicable data others will need to reproduce the problem.
- [x] I have included available evidence of the unexpected behavior (including error messages, screenshots, and/or plots) as well as a descriprion of what I expected instead.
- [ ] If I have a solution in mind, I have provided an explanation and/or pseudocode and/or task list.

Discussion:

Patch:
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 @@ class NestedFrame(pd.DataFrame):
         """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 @@ class NestedFrame(pd.DataFrame):
                 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)


Test patch:
diff --git a/tests/nested_pandas/nestedframe/test_nestedframe.py b/tests/nested_pandas/nestedframe/test_nestedframe.py
index a9c9e15..66931cc 100644
--- a/tests/nested_pandas/nestedframe/test_nestedframe.py
+++ b/tests/nested_pandas/nestedframe/test_nestedframe.py
@@ -3,6 +3,7 @@ import pandas as pd
 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"""
 
@@ -362,55 +372,3 @@ def test_reduce():
     assert result.index.name == "idx"
     for i in range(len(result)):
         assert result["offset_avg"].values[i] == expected_offset_avg[i]
-
-
-def test_reduce_duplicated_cols():
-    """Tests nf.reduce() to correctly handle duplicated column names."""
-    nf = NestedFrame(
-        data={"a": [1, 2, 3], "b": [2, 4, 6]},
-        index=pd.Index([0, 1, 2], name="idx"),
-    )
-
-    to_pack = pd.DataFrame(
-        data={
-            "time": [1, 2, 3, 1, 2, 4, 2, 1, 4],
-            "c": [0, 2, 4, 10, 4, 3, 1, 4, 1],
-            "d": [5, 4, 7, 5, 3, 1, 9, 3, 4],
-        },
-        index=pd.Index([0, 0, 0, 1, 1, 1, 2, 2, 2], name="idx"),
-    )
-
-    to_pack2 = pd.DataFrame(
-        data={
-            "time2": [
-                1,
-                2,
-                3,
-                1,
-                2,
-                3,
-                1,
-                2,
-                4,
-            ],  # TODO: fix duplicate name in join once to_list subset bug fixed
-            "e": [2, 9, 4, 1, 23, 3, 1, 4, 1],
-            "f": [5, 4, 7, 5, 3, 25, 9, 3, 4],
-        },
-        index=pd.Index([0, 0, 0, 1, 1, 1, 2, 2, 2], name="idx"),
-    )
-
-    # Add two nested layers to pack into our dataframe
-    nf = nf.add_nested(to_pack, "packed").add_nested(to_pack2, "packed2")
-
-    def cols_allclose(col1, col2):
-        return pd.Series([np.allclose(col1, col2)], index=["allclose"])
-
-    result = nf.reduce(cols_allclose, "packed.time", "packed2.f")
-    assert_frame_equal(
-        result, pd.DataFrame({"allclose": [False, False, False]}, index=pd.Index([0, 1, 2], name="idx"))
-    )
-
-    result = nf.reduce(cols_allclose, "packed.c", "packed.c")
-    assert_frame_equal(
-        result, pd.DataFrame({"allclose": [True, True, True]}, index=pd.Index([0, 1, 2], name="idx"))
-    )