Merge pull request #190 from lincc-frameworks/flat-chunked-array-fix

Fix issue with flat chunked arrays

Author: hombit

src/nested_pandas/series/packer.py

@@ -293,10 +293,17 @@ def view_sorted_series_as_list_array(
     if unique_index is None:
         unique_index = series.index[offset[:-1]]
 
+    # Input series may be represented by pyarrow.ChunkedArray, in this case pa.array(series) would fail
+    # with "TypeError: Cannot convert a 'ChunkedArray' to a 'ListArray'".
+    # https://github.com/lincc-frameworks/nested-pandas/issues/189
+    flat_array = pa.array(series, from_pandas=True)
+    if isinstance(flat_array, pa.ChunkedArray):
+        flat_array = flat_array.combine_chunks()
     list_array = pa.ListArray.from_arrays(
         offset,
-        pa.array(series, from_pandas=True),
+        flat_array,
     )
+
     return pd.Series(
         list_array,
         dtype=pd.ArrowDtype(list_array.type),

tests/nested_pandas/series/test_packer.py

@@ -524,6 +524,30 @@ def test_view_sorted_series_as_list_array_raises_when_not_sorted():
         packer.view_sorted_series_as_list_array(series)
 
 
+def test_view_sorted_series_as_list_array_chunked_input():
+    """Issue #189
+
+    https://github.com/lincc-frameworks/nested-pandas/issues/189
+    """
+    series = pd.Series(
+        pa.chunked_array([pa.array([0, 1, 2]), pa.array([None, 4])]),
+        name="a",
+        index=np.arange(5),
+        dtype=pd.ArrowDtype(pa.int64()),
+    )
+    offset = np.array([0, 2, 4, 5])
+    unique_index = ["x", "y", "z"]
+    desired = pd.Series(
+        pa.array([[0, 1], [2, None], [4]]),
+        index=unique_index,
+        dtype=pd.ArrowDtype(pa.list_(pa.int64())),
+        name="a",
+    )
+
+    actual = packer.view_sorted_series_as_list_array(series, offset, unique_index)
+    assert_series_equal(actual, desired)
+
+
 @pytest.mark.parametrize(
     "index,offsets",
     [