Skip to content

Commit

Permalink
Merge pull request #139 from lincc-frameworks/fix_from_flat
Browse files Browse the repository at this point in the history
add fix for from_flat
  • Loading branch information
dougbrn authored Aug 20, 2024
2 parents 6e2f433 + 0556592 commit 0d8c2f3
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 1 deletion.
3 changes: 2 additions & 1 deletion src/nested_pandas/nestedframe/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
18 changes: 18 additions & 0 deletions tests/nested_pandas/nestedframe/test_nestedframe.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down

0 comments on commit 0d8c2f3

Please sign in to comment.