Skip to content

Commit

Permalink
Merge pull request #121 from lincc-frameworks/fixed_column_ordering
Browse files Browse the repository at this point in the history
Count Nested: Fixed Column Ordering
  • Loading branch information
dougbrn authored Jul 9, 2024
2 parents 307c904 + 90e6915 commit 3cd3baa
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 4 deletions.
3 changes: 2 additions & 1 deletion src/nested_pandas/utils/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,8 +34,9 @@ def count_nested(df, nested, by=None, join=True) -> NestedFrame:
)
else:
# this may be able to be sped up using tolists() as well
counts = df[nested].apply(lambda x: x[by].value_counts())
counts = df[nested].apply(lambda x: x[by].value_counts(sort=False))
counts = counts.rename(columns={colname: f"n_{nested}_{colname}" for colname in counts.columns})
counts = counts.reindex(sorted(counts.columns), axis=1)
if join:
return df.join(counts)
# else just return the counts NestedFrame
Expand Down
12 changes: 9 additions & 3 deletions tests/nested_pandas/utils/test_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,7 @@ def test_count_nested(join):
data={
"c": [0, 2, 4, 1, np.nan, 3, 1, 4, 1],
"d": [5, 4, 7, 5, 3, 1, 9, 3, 4],
"label": ["a", "a", "b", "b", "a", "a", "b", "a", "b"],
"label": ["b", "a", "b", "b", "a", "a", "b", "a", "b"],
},
index=[0, 0, 0, 1, 1, 1, 2, 2, 2],
)
Expand All @@ -27,8 +27,14 @@ def test_count_nested(join):

# Test count by
label_counts = count_nested(base, "nested", by="label", join=join)
assert all(label_counts["n_nested_a"].values == [2, 2, 1])
assert all(label_counts["n_nested_b"].values == [1, 1, 2])

assert all(label_counts["n_nested_a"].values == [1, 2, 1])
assert all(label_counts["n_nested_b"].values == [2, 1, 2])

# Make sure the ordering is alphabetical
# https://github.com/lincc-frameworks/nested-pandas/issues/109
assert label_counts.columns[-1] == "n_nested_b"
assert label_counts.columns[-2] == "n_nested_a"

# Test join behavior
if join:
Expand Down

0 comments on commit 3cd3baa

Please sign in to comment.