Repository: lincc-frameworks/nested-pandas

Issue:
NestedFrame.from_flat treats all non-base columns as nested columns
**Bug report**
Realized this when trying to implement on the dask side, the from_flat logic is flawed in that it will always define nested_columns as all columns not in base_columns, regardless of whether the user specifies a subset of nested_columns to use. This will limit the ability to omit columns from the result, with base columns being especially problematic as they'll be looped in with the nested structure.

All this should need is a `if nested_columns is None:` line before this: https://github.com/lincc-frameworks/nested-pandas/blob/main/src/nested_pandas/nestedframe/core.py#L214

**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.
- [ ] 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.
- [x] 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 cb743fc..c7f2b38 100644
--- a/src/nested_pandas/nestedframe/core.py
+++ b/src/nested_pandas/nestedframe/core.py
@@ -211,7 +211,8 @@ class NestedFrame(pd.DataFrame):
         out_df = df[base_columns][~df.index.duplicated(keep="first")]
 
         # add nested
-        nested_columns = [col for col in df.columns if col not in base_columns]
+        if nested_columns is None:
+            nested_columns = [col for col in df.columns if col not in base_columns]
         return out_df.add_nested(df[nested_columns], name=name)
 
     @classmethod


Test patch:
diff --git a/tests/nested_pandas/nestedframe/test_nestedframe.py b/tests/nested_pandas/nestedframe/test_nestedframe.py
index 7c4d2fc..097ce1b 100644
--- a/tests/nested_pandas/nestedframe/test_nestedframe.py
+++ b/tests/nested_pandas/nestedframe/test_nestedframe.py
@@ -320,6 +320,24 @@ def test_recover_from_flat():
     assert nf2.equals(nf)
 
 
+def test_from_flat_omitting_columns():
+    """test that from_flat successfully produces subsets"""
+    flat = NestedFrame(
+        {"a": [1, 1, 1, 2, 2], "b": [2, 2, 2, 4, 4], "c": [1, 2, 3, 4, 5], "d": [2, 4, 6, 8, 10]},
+        index=[0, 0, 0, 1, 1],
+    )
+
+    # omit a base column
+    nf = NestedFrame.from_flat(flat, base_columns=["b"], nested_columns=["c", "d"])
+    assert list(nf.columns) == ["b", "nested"]
+    assert list(nf.nested.nest.fields) == ["c", "d"]
+
+    # omit a nested column
+    nf = NestedFrame.from_flat(flat, base_columns=["a", "b"], nested_columns=["c"])
+    assert list(nf.columns) == ["a", "b", "nested"]
+    assert list(nf.nested.nest.fields) == ["c"]
+
+
 def test_from_lists():
     """Test NestedFrame.from_lists behavior"""
     nf = NestedFrame(