From 074840f1a1012b49398de147a2a4376763ab29ca Mon Sep 17 00:00:00 2001 From: Konstantin Malanchev Date: Fri, 10 May 2024 10:39:22 -0400 Subject: [PATCH] Fix view_sorted_series_as_list_array to handle NA and series name --- src/nested_pandas/series/packer.py | 3 ++- tests/nested_pandas/series/test_packer.py | 7 +++++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/src/nested_pandas/series/packer.py b/src/nested_pandas/series/packer.py index 15b155c..a5edab7 100644 --- a/src/nested_pandas/series/packer.py +++ b/src/nested_pandas/series/packer.py @@ -285,13 +285,14 @@ def view_sorted_series_as_list_array( list_array = pa.ListArray.from_arrays( offset, - pa.array(series), + pa.array(series, from_pandas=True), ) return pd.Series( list_array, dtype=pd.ArrowDtype(list_array.type), index=unique_index, copy=False, + name=series.name, ) diff --git a/tests/nested_pandas/series/test_packer.py b/tests/nested_pandas/series/test_packer.py index 23c0698..976661d 100644 --- a/tests/nested_pandas/series/test_packer.py +++ b/tests/nested_pandas/series/test_packer.py @@ -365,8 +365,10 @@ def test_view_sorted_df_as_list_arrays(): def test_view_sorted_series_as_list_array(): """Test view_sorted_series_as_list_array().""" series = pd.Series( - data=[1, 2, 3, 4, 5, 6, 7, 8, 9], + data=[1, 2, 3, 4, 5, None, 7, 8, 9], + dtype=pd.Int64Dtype(), index=[1, 1, 2, 2, 3, 3, 4, 4, 4], + name="my_series", ) nested = packer.view_sorted_series_as_list_array(series) @@ -376,10 +378,11 @@ def test_view_sorted_series_as_list_array(): data=[ np.array([1, 2]), np.array([3, 4]), - np.array([5, 6]), + np.array([5, None]), np.array([7, 8, 9]), ], index=[1, 2, 3, 4], dtype=pd.ArrowDtype(pa.list_(pa.int64())), + name="my_series", ) assert_series_equal(nested, desired_nested)