Repository: lincc-frameworks/nested-pandas

Issue:
Error when nesting lists
**Bug report**

Something is wrong with packing code, error from lsdb:

```python
import lsdb

cat = lsdb.read_hats('/data3/epyc/data3/hats/catalogs/ztf_dr22/ztf_lc')

lc_columns = ['hmjd', 'mag', 'magerr', 'clrcoeff', 'catflags']
cat = cat.nest_lists(
    base_columns=[col for col in cat.columns if col not in lc_columns],
    list_columns=lc_columns,
    name='lc',
)
cat.head()
```

```
...
File ~/.virtualenvs/default/lib/python3.10/site-packages/nested_pandas/series/packer.py:199, in pack_lists(df, name, validate)
    166 def pack_lists(df: pd.DataFrame, name: str | None = None, *, validate: bool = True) -> pd.Series:
    167     """Make a series of arrow structures from a dataframe with nested arrays.
    168 
    169     For the input dataframe with repeated indexes, make a pandas.Series,
   (...)
    197     nested_pandas.series.packer.pack_flat : Pack a "flat" dataframe with repeated indexes.
    198     """
--> 199     struct_array = pa.StructArray.from_arrays(
    200         [df[column] for column in df.columns],
    201         names=df.columns,
    202     )
    203     ext_array = NestedExtensionArray(struct_array, validate=validate)
    204     return pd.Series(
    205         ext_array,
    206         index=df.index,
    207         copy=False,
    208         name=name,
    209     )

File ~/.virtualenvs/default/lib/python3.10/site-packages/pyarrow/array.pxi:4079, in pyarrow.lib.StructArray.from_arrays()

TypeError: Expected Array, got <class 'pyarrow.lib.ChunkedArray'>
```

**Before submitting**
Please check the following:

- [x] I have described the situation in which the bug arose, including what code was executed, information about my environment, and any applicable data others will need to reproduce the problem.
- [x] I have included available evidence of the unexpected behavior (including error messages, screenshots, and/or plots) as well as a description of what I expected instead.
- [ ] If I have a solution in mind, I have provided an explanation and/or pseudocode and/or task list.


Discussion:

Patch:
diff --git a/src/nested_pandas/series/packer.py b/src/nested_pandas/series/packer.py
index 1ea8abb..dedaaca 100644
--- a/src/nested_pandas/series/packer.py
+++ b/src/nested_pandas/series/packer.py
@@ -196,10 +196,34 @@ def pack_lists(df: pd.DataFrame, name: str | None = None, *, validate: bool = Tr
     nested_pandas.series.dtype.NestedDtype : The dtype of the output series.
     nested_pandas.series.packer.pack_flat : Pack a "flat" dataframe with repeated indexes.
     """
-    struct_array = pa.StructArray.from_arrays(
-        [df[column] for column in df.columns],
-        names=df.columns,
-    )
+    # When series is converted to pa.array it may be both Array and ChunkedArray
+    # We convert it to chunked for the sake of consistency
+    pa_arrays_maybe_chunked = {column: pa.array(df[column]) for column in df.columns}
+    pa_chunked_arrays = {
+        column: arr if isinstance(arr, pa.ChunkedArray) else pa.chunked_array([arr])
+        for column, arr in pa_arrays_maybe_chunked.items()
+    }
+
+    # If all chunk arrays have the same chunk lengths, we can build a chunked struct array with no
+    # data copying.
+    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):
+            chunks.append(
+                pa.StructArray.from_arrays(
+                    [arr.chunk(i) for arr in pa_chunked_arrays.values()],
+                    names=pa_chunked_arrays.keys(),
+                )
+            )
+        struct_array = pa.chunked_array(chunks)
+    else:  # "flatten" the chunked arrays
+        struct_array = pa.StructArray.from_arrays(
+            [arr.combine_chunks() for arr in pa_chunked_arrays.values()],
+            names=pa_chunked_arrays.keys(),
+        )
+
     ext_array = NestedExtensionArray(struct_array, validate=validate)
     return pd.Series(
         ext_array,


Test patch:
diff --git a/tests/nested_pandas/series/test_packer.py b/tests/nested_pandas/series/test_packer.py
index d06f8d2..b022a17 100644
--- a/tests/nested_pandas/series/test_packer.py
+++ b/tests/nested_pandas/series/test_packer.py
@@ -286,6 +286,42 @@ def test_pack_lists():
         assert_series_equal(series.nest.get_list_series(field_name), packed_df[field_name])
 
 
+def test_pack_lists_with_chunked_arrays():
+    """Issue https://github.com/lincc-frameworks/nested-pandas/issues/180"""
+    chunked_a = pd.Series(
+        pa.chunked_array([pa.array([[1, 2, 3], [4, 5]])] * 3),
+        dtype=pd.ArrowDtype(pa.list_(pa.int64())),
+        name="a",
+    )
+    chunked_b = pd.Series(
+        pa.chunked_array([pa.array([[0.0, 1.0, 2.0], [3.0, 4.0]])] * 3),
+        dtype=pd.ArrowDtype(pa.list_(pa.float64())),
+        name="b",
+    )
+    list_df = pd.DataFrame({"a": chunked_a, "b": chunked_b}, index=[0, 1, 2, 3, 4, 5])
+    series = packer.pack_lists(list_df)
+    assert_series_equal(series.nest.get_list_series("a"), chunked_a)
+    assert_series_equal(series.nest.get_list_series("b"), chunked_b)
+
+
+def test_pack_lists_with_uneven_chunked_arrays():
+    """Issue https://github.com/lincc-frameworks/nested-pandas/issues/180"""
+    chunked_a = pd.Series(
+        pa.chunked_array([pa.array([[1, 2, 3], [4, 5]])] * 3),
+        dtype=pd.ArrowDtype(pa.list_(pa.int64())),
+        name="a",
+    )
+    chunked_b = pd.Series(
+        pa.array([[0.0, 1.0, 2.0], [3.0, 4.0]] * 3),
+        dtype=pd.ArrowDtype(pa.list_(pa.float64())),
+        name="b",
+    )
+    list_df = pd.DataFrame({"a": chunked_a, "b": chunked_b}, index=[0, 1, 2, 3, 4, 5])
+    series = packer.pack_lists(list_df)
+    assert_series_equal(series.nest.get_list_series("a"), chunked_a)
+    assert_series_equal(series.nest.get_list_series("b"), chunked_b)
+
+
 def test_pack_seq_with_dfs_and_index():
     """Test pack_seq()."""
     dfs = [