From 18d6efc8f288e3d524e69e13a0e92564f30e099b Mon Sep 17 00:00:00 2001 From: Konstantin Malanchev Date: Wed, 29 May 2024 11:39:29 -0400 Subject: [PATCH] Update ext_array.list_offsets tests --- tests/nested_pandas/series/test_ext_array.py | 8 +++++--- tests/nested_pandas/series/test_packer.py | 4 +++- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/tests/nested_pandas/series/test_ext_array.py b/tests/nested_pandas/series/test_ext_array.py index 1ea0193..f25077b 100644 --- a/tests/nested_pandas/series/test_ext_array.py +++ b/tests/nested_pandas/series/test_ext_array.py @@ -599,8 +599,9 @@ def test_list_offsets_single_chunk(): ) ext_array = NestedExtensionArray(struct_array) - desired = pa.chunked_array([pa.array([0, 3, 6])]) - assert_array_equal(ext_array.list_offsets, desired) + desired = pa.array([0, 3, 6], type=pa.int32()) + # pyarrow returns a single bool for == + assert ext_array.list_offsets == desired def test_list_offsets_multiple_chunks(): @@ -616,7 +617,8 @@ def test_list_offsets_multiple_chunks(): ext_array = NestedExtensionArray(chunked_arrray) desired = chunked_arrray.combine_chunks().field("a").offsets - assert_array_equal(ext_array.list_offsets, desired) + # pyarrow returns a single bool for == + assert ext_array.list_offsets == desired def test___getitem___with_integer(): diff --git a/tests/nested_pandas/series/test_packer.py b/tests/nested_pandas/series/test_packer.py index a77edf4..204cace 100644 --- a/tests/nested_pandas/series/test_packer.py +++ b/tests/nested_pandas/series/test_packer.py @@ -402,7 +402,9 @@ def test_view_sorted_series_as_list_array_raises_when_not_sorted(): ) def test_calculate_sorted_index_offsets(index, offsets): """Test calculate_sorted_index_offsets().""" - assert_array_equal(packer.calculate_sorted_index_offsets(index), offsets) + actual = packer.calculate_sorted_index_offsets(index) + assert actual.dtype == np.int32 + assert_array_equal(actual, offsets) def test_calculate_sorted_index_offsets_raises_when_not_sorted():