Skip to content

Commit

Permalink
Fix ndf.add_nested() for empty frames
Browse files Browse the repository at this point in the history
  • Loading branch information
hombit committed May 29, 2024
1 parent 41cce93 commit d4eb8fc
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 1 deletion.
4 changes: 3 additions & 1 deletion src/nested_pandas/nestedframe/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,7 +99,9 @@ def add_nested(
# Add sources to objects
packed = packer.pack(obj, name=name, dtype=dtype)
label = packed.name
return self.assign(**{f"{label}": packed})
new_df = self.copy()
new_df[label] = packed
return new_df

def _split_query(self, expr) -> dict:
"""Splits a pandas query into multiple subqueries for nested and base layers"""
Expand Down
14 changes: 14 additions & 0 deletions tests/nested_pandas/nestedframe/test_nestedframe.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import numpy as np
import pandas as pd
import pyarrow as pa
import pytest
from nested_pandas import NestedFrame
from pandas.testing import assert_frame_equal
Expand Down Expand Up @@ -130,6 +131,19 @@ def test_add_nested_with_series_and_mismatched_index():
assert pd.isna(base.loc[1]["nested"])


def test_add_nested_for_empty_df():
"""Test that .add_nested() works for empty frame and empty input"""
base = NestedFrame(data={"a": [], "b": []}, index=[])
nested = pd.DataFrame(data={"c": []}, index=[])
new_base = base.add_nested(nested, "nested")

# Check original frame is unchanged
assert_frame_equal(base, NestedFrame(data={"a": [], "b": []}, index=[]))

assert "nested" in new_base.columns
assert_frame_equal(new_base.nested.nest.to_flat(), nested.astype(pd.ArrowDtype(pa.float64())))


def test_query():
"""Test that NestedFrame.query handles nested queries correctly"""

Expand Down

0 comments on commit d4eb8fc

Please sign in to comment.