Merge pull request #139 from lincc-frameworks/fix_from_flat

add fix for from_flat

src/nested_pandas/nestedframe/core.py

@@ -211,7 +211,8 @@ def from_flat(cls, df, base_columns, nested_columns=None, index=None, name="nest
         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

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(