diff --git a/src/nested_pandas/series/packer.py b/src/nested_pandas/series/packer.py index dedaacaf..73dc58ce 100644 --- a/src/nested_pandas/series/packer.py +++ b/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), diff --git a/tests/nested_pandas/series/test_packer.py b/tests/nested_pandas/series/test_packer.py index b022a173..df59c0f8 100644 --- a/tests/nested_pandas/series/test_packer.py +++ b/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", [