Merge pull request #267 from lincc-frameworks/empty-data-storages

Fix .nest.to_flatten_inner for empty inputs

Author: hombit

src/nested_pandas/series/_storage/table_storage.py

@@ -4,7 +4,11 @@
 
 import pyarrow as pa
 
-from nested_pandas.series.utils import table_to_struct_array, validate_struct_list_array_for_equal_lengths
+from nested_pandas.series.utils import (
+    table_from_struct_array,
+    table_to_struct_array,
+    validate_struct_list_array_for_equal_lengths,
+)
 
 if TYPE_CHECKING:
     from nested_pandas.series._storage.list_struct_storage import ListStructStorage
@@ -58,5 +62,5 @@ def from_struct_list_storage(cls, struct_list_storage: StructListStorage) -> Sel
         struct_list_storage : StructListStorage
             StructListStorage object.
         """
-        table = pa.Table.from_struct_array(struct_list_storage.data)
+        table = table_from_struct_array(struct_list_storage.data)
         return cls(table, validate=False)

src/nested_pandas/series/accessor.py

@@ -1,7 +1,6 @@
 # Python 3.9 doesn't support "|" for types
 from __future__ import annotations
 
-from collections import defaultdict
 from collections.abc import Generator, Mapping
 from typing import cast
 
@@ -115,7 +114,7 @@ def to_flat(self, fields: list[str] | None = None) -> pd.DataFrame:
 
         index = pd.Series(self.get_flat_index(), name=self._series.index.name)
 
-        flat_chunks = defaultdict(list)
+        flat_chunks: dict[str, list[pa.Array]] = {field: [] for field in fields}
         for chunk in self._series.array.struct_array.iterchunks():
             struct_array = cast(pa.StructArray, chunk)
             for field in fields:
@@ -125,13 +124,16 @@ def to_flat(self, fields: list[str] | None = None) -> pd.DataFrame:
 
         flat_series = {}
         for field, chunks in flat_chunks.items():
+            dtype = self._series.dtype.field_dtype(field)
+            if len(chunks) == 0:
+                chunks = [pa.array([])]
             chunked_array = pa.chunked_array(chunks)
             flat_series[field] = pd.Series(
                 chunked_array,
                 index=index,
                 name=field,
                 copy=False,
-                dtype=self._series.dtype.field_dtype(field),
+                dtype=dtype,
             )
 
         return pd.DataFrame(flat_series)

src/nested_pandas/series/packer.py

@@ -209,8 +209,8 @@ def pack_lists(df: pd.DataFrame, name: str | None = None, *, validate: bool = Tr
     chunk_lengths = pa.array([[len(chunk) for chunk in arr.chunks] for arr in pa_chunked_arrays.values()])
     if all(chunk_length == chunk_lengths[0] for chunk_length in chunk_lengths):
         chunks = []
-        numpy_chunks = next(iter(pa_chunked_arrays.values())).num_chunks
-        for i in range(numpy_chunks):
+        num_chunks = next(iter(pa_chunked_arrays.values())).num_chunks
+        for i in range(num_chunks):
             chunks.append(
                 pa.StructArray.from_arrays(
                     [arr.chunk(i) for arr in pa_chunked_arrays.values()],

src/nested_pandas/series/utils.py

@@ -307,3 +307,10 @@ def table_to_struct_array(table: pa.Table) -> pa.ChunkedArray:
     if len(table) == 0:
         return pa.chunked_array([], type=pa.struct(table.schema))
     return table.to_struct_array()
+
+
+def table_from_struct_array(array: pa.ChunkedArray | pa.array) -> pa.Table:
+    """pa.Table.from_struct_array, but working with chunkless input"""
+    if isinstance(array, pa.ChunkedArray) and array.num_chunks == 0:
+        array = pa.array([], type=array.type)
+    return pa.Table.from_struct_array(array)

tests/nested_pandas/series/test_accessor.py

@@ -1098,3 +1098,23 @@ def test_to_flatten_outer_wrong_field():
     nf = generate_data(10, 2)
     with pytest.raises(ValueError):
         nf.nested.nest.to_flatten_inner("t")
+
+
+def test_issue266():
+    """Test .nest.to_flatten_inner() with empty series.
+
+    https://github.com/lincc-frameworks/nested-pandas/issues/266
+    """
+
+    nf = generate_data(10, 2)
+    nf = nf.assign(id=np.repeat(np.r_[0:5], 2))
+    nf = nf.rename(columns={"nested": "inner"})
+    nnf = NestedFrame.from_flat(nf, base_columns=[], on="id", name="outer")
+
+    empty_nnf = nnf.iloc[0:0]
+
+    empty_outer_flatten = empty_nnf["outer"].nest.to_flatten_inner("inner")
+
+    assert empty_outer_flatten.dtype == NestedDtype.from_fields(
+        {"t": pa.float64(), "flux": pa.float64(), "band": pa.string(), "a": pa.float64(), "b": pa.float64()}
+    )