Repository: lincc-frameworks/nested-pandas

Issue:
Struct_array casting sometimes produces a ChunkedArray
**Bug report**
There are cases we see on the nested-dask side currently where the accessor code fails. Take for example this:

```
# Load some ZTF data
catalogs_dir = "https://epyc.astro.washington.edu/~lincc-frameworks/half_degree_surveys/ztf/"

object_ndf = npd.read_parquet(
    f"{catalogs_dir}/ztf_object/Norder=3/Dir=0/Npix=432.parquet",
    columns=["ra", "dec", "ps1_objid"],
).set_index("ps1_objid")

source_ndf = npd.read_parquet(
    f"{catalogs_dir}/ztf_source/Norder=6/Dir=20000/Npix=27711.parquet",
    columns=["mjd", "mag", "magerr", "band", "ps1_objid", "catflags"],
).set_index("ps1_objid")

object_ndf = object_ndf.add_nested(source_ndf, "ztf_source")
object_ndf

nf = object_ndf
nf.reduce(np.mean, "ztf_source.mjd")
```
In this case the code fails with:
```
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[2], line 21
     18 object_ndf
     20 nf = object_ndf.compute()
---> 21 nf.reduce(np.mean, "ztf_source.mjd")

File ~/lincc/nested-pandas/src/nested_pandas/nestedframe/core.py:426, in NestedFrame.reduce(self, func, *args, **kwargs)
    422         apply_df = apply_df.join(self[columns], how="outer")
    423     else:
    424         # TODO: It should be faster to pass these columns to to_lists, but its 20x slower
    425         # columns = [col[1] for col in requested_columns if col[0] == layer]
--> 426         apply_df = apply_df.join(self[layer].nest.to_lists(), how="outer")
    428 # Translates the requested columns into the scalars or arrays we pass to func.
    429 def translate_cols(frame, layer, col):

File ~/lincc/nested-pandas/src/nested_pandas/series/accessor.py:61, in NestSeriesAccessor.to_lists(self, fields)
     59 list_series = {}
     60 for field in fields:
---> 61     list_array = cast(pa.ListArray, struct_array.field(field))
     62     list_series[field] = pd.Series(
     63         list_array,
     64         dtype=pd.ArrowDtype(list_array.type),
   (...)
     67         copy=False,
     68     )
     70 return pd.DataFrame(list_series)

AttributeError: 'pyarrow.lib.ChunkedArray' object has no attribute 'field'
```

The problem line appears to be a cast call in the accessor code, where it tries to cast it to a struct array but we instead get a ChunkedArray object with mismatched API.
```struct_array = cast(pa.StructArray, pa.array(self._series))```



**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 descriprion 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/accessor.py b/src/nested_pandas/series/accessor.py
index 3bd967b..85ca6dc 100644
--- a/src/nested_pandas/series/accessor.py
+++ b/src/nested_pandas/series/accessor.py
@@ -1,6 +1,7 @@
 # Python 3.9 doesn't support "|" for types
 from __future__ import annotations
 
+from collections import defaultdict
 from collections.abc import Generator, MutableMapping
 from typing import cast
 
@@ -53,14 +54,19 @@ class NestSeriesAccessor(MutableMapping):
         if len(fields) == 0:
             raise ValueError("Cannot convert a struct with no fields to lists")
 
-        struct_array = cast(pa.StructArray, pa.array(self._series))
+        list_chunks = defaultdict(list)
+        for chunk in self._series.array._pa_array.iterchunks():
+            struct_array = cast(pa.StructArray, chunk)
+            for field in fields:
+                list_array = cast(pa.ListArray, struct_array.field(field))
+                list_chunks[field].append(list_array)
 
         list_series = {}
-        for field in fields:
-            list_array = cast(pa.ListArray, struct_array.field(field))
+        for field, chunks in list_chunks.items():
+            chunked_array = pa.chunked_array(chunks)
             list_series[field] = pd.Series(
-                list_array,
-                dtype=pd.ArrowDtype(list_array.type),
+                chunked_array,
+                dtype=pd.ArrowDtype(chunked_array.type),
                 index=self._series.index,
                 name=field,
                 copy=False,
@@ -85,21 +91,25 @@ class NestSeriesAccessor(MutableMapping):
         if len(fields) == 0:
             raise ValueError("Cannot flatten a struct with no fields")
 
-        struct_array = cast(pa.StructArray, pa.array(self._series))
+        index = pd.Series(self.get_flat_index(), name=self._series.index.name)
+
+        flat_chunks = defaultdict(list)
+        for chunk in self._series.array._pa_array.iterchunks():
+            struct_array = cast(pa.StructArray, chunk)
+            for field in fields:
+                list_array = cast(pa.ListArray, struct_array.field(field))
+                flat_array = list_array.flatten()
+                flat_chunks[field].append(flat_array)
 
         flat_series = {}
-        index = None
-        for field in fields:
-            list_array = cast(pa.ListArray, struct_array.field(field))
-            flat_array = list_array.flatten()
-            if index is None:
-                index = self.get_flat_index()
+        for field, chunks in flat_chunks.items():
+            chunked_array = pa.chunked_array(chunks)
             flat_series[field] = pd.Series(
-                flat_array,
+                chunked_array,
                 index=pd.Series(index, name=self._series.index.name),
                 name=field,
                 copy=False,
-                dtype=pd.ArrowDtype(flat_array.type),
+                dtype=pd.ArrowDtype(chunked_array.type),
             )
 
         return pd.DataFrame(flat_series)
diff --git a/src/nested_pandas/series/ext_array.py b/src/nested_pandas/series/ext_array.py
index 71f1777..e915cdd 100644
--- a/src/nested_pandas/series/ext_array.py
+++ b/src/nested_pandas/series/ext_array.py
@@ -598,7 +598,7 @@ class NestedExtensionArray(ExtensionArray):
         self._dtype = NestedDtype(pa_array.chunk(0).type)
 
     @property
-    def list_offsets(self) -> pa.ChunkedArray:
+    def list_offsets(self) -> pa.Array:
         """The list offsets of the field arrays.
 
         It is a chunk array of list offsets of the first field array.
@@ -609,7 +609,25 @@ class NestedExtensionArray(ExtensionArray):
         pa.ChunkedArray
             The list offsets of the field arrays.
         """
-        return pa.chunked_array([chunk.field(0).offsets for chunk in self._pa_array.iterchunks()])
+        # Quick and cheap path for a single chunk
+        if self._pa_array.num_chunks == 1:
+            struct_array = cast(pa.StructArray, self._pa_array.chunk(0))
+            return cast(pa.ListArray, struct_array.field(0)).offsets
+
+        chunks = []
+        # The offset of the current chunk in the flat array.
+        # It is 0 for the first chunk, and the last offset of the previous chunk for the next chunks,
+        # as a pa.Scalar.
+        chunk_offset: pa.Scalar | int = 0
+        for chunk in self._pa_array.iterchunks():
+            list_array = cast(pa.ListArray, chunk.field(0))
+            if chunk_offset == 0:
+                offsets = list_array.offsets
+            else:
+                offsets = pa.compute.add(list_array.offsets[1:], chunk_offset)
+            chunks.append(offsets)
+            chunk_offset = offsets[-1]
+        return pa.concat_arrays(chunks)
 
     @property
     def field_names(self) -> list[str]:
@@ -621,6 +639,11 @@ class NestedExtensionArray(ExtensionArray):
         """Length of the flat arrays"""
         return sum(chunk.field(0).value_lengths().sum().as_py() for chunk in self._pa_array.iterchunks())
 
+    @property
+    def num_chunks(self) -> int:
+        """Number of chunks in underlying pyarrow.ChunkedArray"""
+        return self._pa_array.num_chunks
+
     def view_fields(self, fields: str | list[str]) -> Self:  # type: ignore[name-defined] # noqa: F821
         """Get a view of the series with only the specified fields
 
@@ -673,8 +696,7 @@ class NestedExtensionArray(ExtensionArray):
         if len(pa_array) != self.flat_length:
             raise ValueError("The input must be a struct_scalar or have the same length as the flat arrays")
 
-        offsets = self.list_offsets.combine_chunks()
-        list_array = pa.ListArray.from_arrays(values=pa_array, offsets=offsets)
+        list_array = pa.ListArray.from_arrays(values=pa_array, offsets=self.list_offsets)
 
         return self.set_list_field(field, list_array)
 


Test patch:
diff --git a/tests/nested_pandas/e2e_tests/test_issue89.py b/tests/nested_pandas/e2e_tests/test_issue89.py
new file mode 100644
index 0000000..6919647
--- /dev/null
+++ b/tests/nested_pandas/e2e_tests/test_issue89.py
@@ -0,0 +1,29 @@
+"""Based on https://github.com/lincc-frameworks/nested-pandas/issues/89"""
+
+import nested_pandas as npd
+import numpy as np
+
+
+def test_issue89():
+    """Check that code snippet from issue 89 works as expected
+
+    https://github.com/lincc-frameworks/nested-pandas/issues/89
+    """
+
+    # Load some ZTF data
+    catalogs_dir = "https://epyc.astro.washington.edu/~lincc-frameworks/half_degree_surveys/ztf/"
+
+    object_ndf = npd.read_parquet(
+        f"{catalogs_dir}/ztf_object/Norder=3/Dir=0/Npix=432.parquet",
+        columns=["ra", "dec", "ps1_objid"],
+    ).set_index("ps1_objid")
+
+    source_ndf = npd.read_parquet(
+        f"{catalogs_dir}/ztf_source/Norder=6/Dir=20000/Npix=27711.parquet",
+        columns=["mjd", "mag", "magerr", "band", "ps1_objid", "catflags"],
+    ).set_index("ps1_objid")
+
+    object_ndf = object_ndf.add_nested(source_ndf, "ztf_source")
+
+    nf = object_ndf
+    nf.reduce(np.mean, "ztf_source.mjd")
diff --git a/tests/nested_pandas/series/test_accessor.py b/tests/nested_pandas/series/test_accessor.py
index 5ee617d..fcb6ee5 100644
--- a/tests/nested_pandas/series/test_accessor.py
+++ b/tests/nested_pandas/series/test_accessor.py
@@ -70,6 +70,39 @@ def test_to_lists():
     assert_frame_equal(lists, desired)
 
 
+def test_to_lists_for_chunked_array():
+    """ ""Test that the .nest.to_lists() when underlying array is chunked"""
+    struct_array = pa.StructArray.from_arrays(
+        arrays=[
+            [np.array([1.0, 2.0, 3.0]), -np.array([1.0, 2.0, 1.0])],
+            [np.array([4.0, 5.0, 6.0]), -np.array([3.0, 4.0, 5.0])],
+        ],
+        names=["a", "b"],
+    )
+    chunked_array = pa.chunked_array([struct_array] * 3)
+    assert chunked_array.length() == 6
+    series = pd.Series(chunked_array, dtype=NestedDtype(chunked_array.type), index=[0, 1, 2, 3, 4, 5])
+    assert series.array.num_chunks == 3
+
+    lists = series.nest.to_lists()
+
+    desired = pd.DataFrame(
+        data={
+            "a": pd.Series(
+                data=[np.array([1.0, 2.0, 3.0]), -np.array([1.0, 2.0, 1.0])] * 3,
+                dtype=pd.ArrowDtype(pa.list_(pa.float64())),
+                index=[0, 1, 2, 3, 4, 5],
+            ),
+            "b": pd.Series(
+                data=[np.array([4.0, 5.0, 6.0]), -np.array([3.0, 4.0, 5.0])] * 3,
+                dtype=pd.ArrowDtype(pa.list_(pa.float64())),
+                index=[0, 1, 2, 3, 4, 5],
+            ),
+        },
+    )
+    assert_frame_equal(lists, desired)
+
+
 def test_to_lists_with_fields():
     """Test that the .nest.to_lists(fields=...) method works."""
     struct_array = pa.StructArray.from_arrays(
@@ -154,6 +187,43 @@ def test_to_flat():
         assert_array_equal(flat[column], desired[column])
 
 
+def test_to_flat_for_chunked_array():
+    """Test that the .nest.to_flat() when underlying array is pa.ChunkedArray."""
+    struct_array = pa.StructArray.from_arrays(
+        arrays=[
+            [np.array([1.0, 2.0, 3.0]), -np.array([1.0, 2.0, 1.0])],
+            [np.array([4.0, 5.0, 6.0]), -np.array([3.0, 4.0, 5.0])],
+        ],
+        names=["a", "b"],
+    )
+    chunked_array = pa.chunked_array([struct_array] * 3)
+    assert chunked_array.length() == 6
+    series = pd.Series(chunked_array, dtype=NestedDtype(chunked_array.type), index=[0, 1, 2, 3, 4, 5])
+    assert series.array.num_chunks == 3
+
+    flat = series.nest.to_flat()
+
+    desired = pd.DataFrame(
+        data={
+            "a": pd.Series(
+                data=[1.0, 2.0, 3.0, -1.0, -2.0, -1.0] * 3,
+                name="a",
+                index=np.repeat([0, 1, 2, 3, 4, 5], 3),
+                dtype=pd.ArrowDtype(pa.float64()),
+            ),
+            "b": pd.Series(
+                data=[4.0, 5.0, 6.0, -3.0, -4.0, -5.0] * 3,
+                name="b",
+                index=np.repeat([0, 1, 2, 3, 4, 5], 3),
+                dtype=pd.ArrowDtype(pa.float64()),
+            ),
+        },
+        index=pd.Index(np.repeat([0, 1, 2, 3, 4, 5], 3)),
+    )
+
+    assert_frame_equal(flat, desired)
+
+
 def test_to_flat_with_fields():
     """Test that the .nest.to_flat(fields=...) method works."""
     struct_array = pa.StructArray.from_arrays(
diff --git a/tests/nested_pandas/series/test_ext_array.py b/tests/nested_pandas/series/test_ext_array.py
index b263240..1ea0193 100644
--- a/tests/nested_pandas/series/test_ext_array.py
+++ b/tests/nested_pandas/series/test_ext_array.py
@@ -588,8 +588,8 @@ def test_series_built_raises(data):
         _array = NestedExtensionArray(pa_array)
 
 
-def test_list_offsets():
-    """Test that the list offsets are correct."""
+def test_list_offsets_single_chunk():
+    """Test that the .list_offset property is correct for a single chunk."""
     struct_array = pa.StructArray.from_arrays(
         arrays=[
             pa.array([np.array([1, 2, 3]), np.array([1, 2, 1])], type=pa.list_(pa.uint8())),
@@ -603,6 +603,22 @@ def test_list_offsets():
     assert_array_equal(ext_array.list_offsets, desired)
 
 
+def test_list_offsets_multiple_chunks():
+    """Test that the .list_offset property is correct for multiple chunks."""
+    struct_array = pa.StructArray.from_arrays(
+        arrays=[
+            pa.array([np.array([1, 2, 3]), np.array([1, 2, 1])], type=pa.list_(pa.uint8())),
+            pa.array([-np.array([4.0, 5.0, 6.0]), -np.array([3.0, 4.0, 5.0])]),
+        ],
+        names=["a", "b"],
+    )
+    chunked_arrray = pa.chunked_array([struct_array, struct_array[:1], struct_array])
+    ext_array = NestedExtensionArray(chunked_arrray)
+
+    desired = chunked_arrray.combine_chunks().field("a").offsets
+    assert_array_equal(ext_array.list_offsets, desired)
+
+
 def test___getitem___with_integer():
     """Test [i] is a valid DataFrame."""
     struct_array = pa.StructArray.from_arrays(
@@ -1207,6 +1223,21 @@ def test_flat_length():
     assert ext_array.flat_length == 7
 
 
+def test_num_chunks():
+    """Tests .num_chunks property."""
+    struct_array = pa.StructArray.from_arrays(
+        arrays=[
+            pa.array([np.array([1.0, 2.0, 3.0]), np.array([1.0, 2.0, 1.0, 2.0])]),
+            pa.array([-np.array([4.0, 5.0, 6.0]), -np.array([3.0, 4.0, 5.0, 6.0])]),
+        ],
+        names=["a", "b"],
+    )
+    chunked_array = pa.chunked_array([struct_array] * 7)
+    ext_array = NestedExtensionArray(chunked_array)
+
+    assert ext_array.num_chunks == 7
+
+
 def test_view_fields_with_single_field():
     """Tests ext_array.view("field")"""
     arrays = [