Skip to content

Commit

Permalink
NestedExtensionArray.chunked_array
Browse files Browse the repository at this point in the history
  • Loading branch information
hombit committed May 29, 2024
1 parent a932d18 commit d193826
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 2 deletions.
5 changes: 5 additions & 0 deletions src/nested_pandas/series/ext_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -558,6 +558,11 @@ def _pyarrow_dtype(self) -> pa.DataType:
"""PyArrow data type of the extension array"""
return self._dtype.pyarrow_dtype

@property
def chunked_array(self) -> pa.ChunkedArray:
"""The underlying PyArrow ChunkedArray"""
return self._chunked_array

@staticmethod
def _validate(array: pa.ChunkedArray) -> None:
"""Raises ValueError if the input array is not a struct array with all fields being
Expand Down
19 changes: 17 additions & 2 deletions tests/nested_pandas/series/test_ext_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -204,7 +204,7 @@ def test_from_sequence_with_arrow_array_and_dtype():
type=pa_type,
)

actual = NestedExtensionArray.from_sequence(pa_array, dtype=new_pa_type)._chunked_array
actual = NestedExtensionArray.from_sequence(pa_array, dtype=new_pa_type).chunked_array
desired = pa.chunked_array([pa_array.cast(new_pa_type)])
# pyarrow doesn't convert pandas boxed missing values to nulls in nested arrays
assert actual == desired
Expand Down Expand Up @@ -525,7 +525,7 @@ def test___setitem___series_of_dfs():
)
desired = NestedExtensionArray(desired_struct_array)

assert ext_array._chunked_array == desired._chunked_array
assert ext_array.chunked_array == desired.chunked_array
assert ext_array.equals(desired)


Expand Down Expand Up @@ -588,6 +588,21 @@ def test_series_built_raises(data):
_array = NestedExtensionArray(pa_array)


def test_chunked_array():
"""Test that the .chunked_array property is correct."""
struct_array = pa.StructArray.from_arrays(
arrays=[
pa.array([np.array([1, 2, 3]), np.array([1, 2, 1])]),
pa.array([-np.array([4.0, 5.0, 6.0]), -np.array([3.0, 4.0, 5.0])]),
],
names=["a", "b"],
)
ext_array = NestedExtensionArray(struct_array)

# pyarrow returns a single bool for ==
assert ext_array.chunked_array == pa.chunked_array(struct_array)


def test_list_offsets_single_chunk():
"""Test that the .list_offset property is correct for a single chunk."""
struct_array = pa.StructArray.from_arrays(
Expand Down

0 comments on commit d193826

Please sign in to comment.