Update ext_array.list_offsets tests

lincc-frameworks · May 29, 2024 · 18d6efc · 18d6efc
1 parent 4dcd87a
commit 18d6efc
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git 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
@@ -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():