Merge pull request #203 from lincc-frameworks/extarr-list-struct

Use both struct-list and list-struct representations in the extension array

Author: hombit

src/nested_pandas/series/ext_array.py

@@ -51,7 +51,14 @@
 from pandas.io.formats.format import format_array
 
 from nested_pandas.series.dtype import NestedDtype
-from nested_pandas.series.utils import enumerate_chunks, is_pa_type_a_list
+from nested_pandas.series.utils import (
+    enumerate_chunks,
+    is_pa_type_a_list,
+    transpose_list_struct_array,
+    transpose_struct_list_array,
+    transpose_struct_list_type,
+    validate_struct_list_array_for_equal_lengths,
+)
 
 __all__ = ["NestedExtensionArray"]
 
@@ -549,6 +556,8 @@ def __arrow_array__(self, type=None):
         """Convert the extension array to a PyArrow array."""
         if type is None:
             return self._chunked_array
+        if isinstance(type, pa.ListType):
+            return self._list_array.cast(type)
         return self._chunked_array.cast(type)
 
     def __array__(self, dtype=None):
@@ -650,12 +659,27 @@ def __init__(self, values: pa.Array | pa.ChunkedArray, *, validate: bool = True)
         if isinstance(values, pa.Array):
             values = pa.chunked_array([values])
 
-        if validate:
+        # Convert list-struct array to struct-list array
+        if is_pa_type_a_list(values.type):
+            struct_chunks = []
+            for list_chunk in values.iterchunks():
+                struct_chunks.append(transpose_list_struct_array(list_chunk))
+            values = pa.chunked_array(struct_chunks)
+        # Validate struct-array with list fields
+        elif validate:
             self._validate(values)
 
         self._chunked_array = values
         self._dtype = NestedDtype(values.type)
 
+    @property
+    def _list_array(self) -> pa.ChunkedArray:
+        """Pyarrow chunked list-struct array representation"""
+        list_chunks = []
+        for struct_chunk in self._chunked_array.iterchunks():
+            list_chunks.append(transpose_struct_list_array(struct_chunk, validate=False))
+        return pa.chunked_array(list_chunks)
+
     @classmethod
     def from_sequence(cls, scalars, *, dtype: NestedDtype | pd.ArrowDtype | pa.DataType = None) -> Self:  # type: ignore[name-defined] # noqa: F821
         """Construct a NestedExtensionArray from a sequence of items
@@ -677,53 +701,65 @@ def from_sequence(cls, scalars, *, dtype: NestedDtype | pd.ArrowDtype | pa.DataT
         return cls._from_sequence(scalars, dtype=dtype)
 
     @property
-    def _pyarrow_dtype(self) -> pa.DataType:
+    def _pyarrow_dtype(self) -> pa.StructType:
         """PyArrow data type of the extension array"""
         return self._dtype.pyarrow_dtype
 
+    @property
+    def _pyarrow_list_struct_dtype(self) -> pa.ListType:
+        """PyArrow data type of the list-struct view over the ext. array"""
+        return transpose_struct_list_type(self._pyarrow_dtype)
+
     @property
     def chunked_array(self) -> pa.ChunkedArray:
         """The underlying PyArrow ChunkedArray"""
         return self._chunked_array
 
+    @property
+    def chunked_list_struct_array(self) -> pa.ChunkedArray:
+        """Chunked list-struct view over the extension array"""
+        return self._list_array
+
     @staticmethod
     def _validate(array: pa.ChunkedArray) -> None:
         """Raises ValueError if the input array is not a struct array with all fields being
         list arrays of the same lengths.
+
+        Parameters
+        ----------
+        array : pa.ChunkedArray
+            The array to validate.
+
+        Raises
+        ------
+        ValueError
         """
         for chunk in array.iterchunks():
-            if not pa.types.is_struct(chunk.type):
-                raise ValueError(f"Expected a StructArray, got {chunk.type}")
-            struct_array = cast(pa.StructArray, chunk)
-
-            first_list_array: pa.ListArray | None = None
-            for field in struct_array.type:
-                inner_array = struct_array.field(field.name)
-                if not is_pa_type_a_list(inner_array.type):
-                    raise ValueError(f"Expected a ListArray, got {inner_array.type}")
-                list_array = cast(pa.ListArray, inner_array)
-
-                if first_list_array is None:
-                    first_list_array = list_array
-                    continue
-                # compare offsets from the first list array with the current one
-                if not first_list_array.offsets.equals(list_array.offsets):
-                    raise ValueError("Offsets of all ListArrays must be the same")
+            validate_struct_list_array_for_equal_lengths(chunk)
 
     @classmethod
     def from_arrow_ext_array(cls, array: ArrowExtensionArray) -> Self:  # type: ignore[name-defined] # noqa: F821
         """Create a NestedExtensionArray from pandas' ArrowExtensionArray"""
         return cls(array._pa_array)
 
-    def to_arrow_ext_array(self) -> ArrowExtensionArray:
-        """Convert the extension array to pandas' ArrowExtensionArray"""
+    def to_arrow_ext_array(self, list_struct: bool = False) -> ArrowExtensionArray:
+        """Convert the extension array to pandas' ArrowExtensionArray
+
+        Parameters
+        ----------
+        list_struct : bool, optional
+            If False (default), return struct-list array, otherwise return
+            list-struct array.
+        """
+        if list_struct:
+            return ArrowExtensionArray(self._list_array)
         return ArrowExtensionArray(self._chunked_array)
 
     def _replace_chunked_array(self, pa_array: pa.ChunkedArray, *, validate: bool) -> None:
         if validate:
             self._validate(pa_array)
         self._chunked_array = pa_array
-        self._dtype = NestedDtype(pa_array.chunk(0).type)
+        self._dtype = NestedDtype(pa_array.type)
 
     @property
     def list_offsets(self) -> pa.Array:
@@ -737,48 +773,32 @@ def list_offsets(self) -> pa.Array:
         pa.ChunkedArray
             The list offsets of the field arrays.
         """
-        # Quick and cheap path for a single chunk
+        # Cheap path for a single chunk
         if self._chunked_array.num_chunks == 1:
             struct_array = cast(pa.StructArray, self._chunked_array.chunk(0))
             return cast(pa.ListArray, struct_array.field(0)).offsets
 
-        chunks = []
-        # The offset of the current chunk in the flat array.
-        # Offset arrays use int32 type, so we cast to it
-        chunk_offset = pa.scalar(0, type=pa.int32())
-        for chunk in self._chunked_array.iterchunks():
-            list_array = cast(pa.ListArray, chunk.field(0))
-            if chunk_offset.equals(pa.scalar(0, type=pa.int32())):
-                offsets = list_array.offsets
-            else:
-                offsets = pa.compute.add(list_array.offsets[1:], chunk_offset)
-            chunks.append(offsets)
-            chunk_offset = offsets[-1]
-        return pa.concat_arrays(chunks)
+        zero_and_lengths = pa.chunked_array(
+            [pa.array([0], type=pa.int32()), pa.array(self.list_lengths, type=pa.int32())]
+        )
+        offsets = pa.compute.cumulative_sum(zero_and_lengths)
+        return offsets.chunk(0) if offsets.num_chunks == 1 else offsets.combine_chunks()
 
     @property
     def field_names(self) -> list[str]:
         """Names of the nested columns"""
         return [field.name for field in self._chunked_array.chunk(0).type]
 
-    def _iter_list_lengths(self) -> Generator[int, None, None]:
-        """Iterate over the lengths of the list arrays"""
-        for chunk in self._chunked_array.iterchunks():
-            for length in chunk.field(0).value_lengths():
-                if length.is_valid:
-                    yield length.as_py()
-                else:
-                    yield 0
-
     @property
-    def list_lengths(self) -> list[int]:
+    def list_lengths(self) -> np.ndarray:
         """Lengths of the list arrays"""
-        return list(self._iter_list_lengths())
+        list_lengths = pa.compute.list_value_length(self._list_array)
+        return np.asarray(list_lengths)
 
     @property
     def flat_length(self) -> int:
         """Length of the flat arrays"""
-        return sum(self._iter_list_lengths())
+        return pa.compute.sum(self.list_lengths).as_py()
 
     @property
     def num_chunks(self) -> int:
@@ -790,8 +810,8 @@ def get_list_index(self) -> np.ndarray:
         if len(self) == 0:
             # Since we have no list offsets, return an empty array
             return np.array([], dtype=int)
-        list_index = np.arange(len(self))
-        return np.repeat(list_index, np.diff(self.list_offsets))
+        list_index = np.arange(len(self), dtype=int)
+        return np.repeat(list_index, self.list_lengths)
 
     def iter_field_lists(self, field: str) -> Generator[np.ndarray, None, None]:
         """Iterate over single field nested lists, as numpy arrays
@@ -813,7 +833,7 @@ def iter_field_lists(self, field: str) -> Generator[np.ndarray, None, None]:
                 yield np.asarray(list_scalar.values)
 
     def view_fields(self, fields: str | list[str]) -> Self:  # type: ignore[name-defined] # noqa: F821
-        """Get a view of the extension array with only the specified fields
+        """Get a view of the extension array with the specified fields only
 
         Parameters
         ----------
@@ -842,7 +862,7 @@ def view_fields(self, fields: str | list[str]) -> Self:  # type: ignore[name-def
             chunks.append(struct_array)
         pa_array = pa.chunked_array(chunks)
 
-        return self.__class__(pa_array, validate=False)
+        return type(self)(pa_array, validate=False)
 
     def set_flat_field(self, field: str, value: ArrayLike, *, keep_dtype: bool = False) -> None:
         """Set the field from flat-array of values

src/nested_pandas/series/utils.py

@@ -1,9 +1,12 @@
+from __future__ import annotations  # Python 3.9 requires it for X | Y type hints
+
 from collections.abc import Generator
+from typing import cast
 
 import pyarrow as pa
 
 
-def is_pa_type_a_list(pa_type: type[pa.Array]) -> bool:
+def is_pa_type_a_list(pa_type: pa.DataType) -> bool:
     """Check if the given pyarrow type is a list type.
 
     I.e. one of the following types: ListArray, LargeListArray,
@@ -39,3 +42,148 @@ def enumerate_chunks(array: pa.ChunkedArray) -> Generator[tuple[slice, pa.Array]
         index_stop = index_start + len(chunk)
         yield slice(index_start, index_stop), chunk
         index_start = index_stop
+
+
+def validate_struct_list_array_for_equal_lengths(array: pa.StructArray) -> None:
+    """Check if the given struct array has lists of equal length.
+
+    Parameters
+    ----------
+    array : pa.StructArray
+        Input struct array.
+
+    Raises
+    ------
+    ValueError
+        If the struct array has lists of unequal length or type of the input
+        array is not a StructArray or fields are not ListArrays.
+    """
+    if not pa.types.is_struct(array.type):
+        raise ValueError(f"Expected a StructArray, got {array.type}")
+
+    first_list_array: pa.ListArray | None = None
+    for field in array.type:
+        inner_array = array.field(field.name)
+        if not is_pa_type_a_list(inner_array.type):
+            raise ValueError(f"Expected a ListArray, got {inner_array.type}")
+        list_array = cast(pa.ListArray, inner_array)
+
+        if first_list_array is None:
+            first_list_array = list_array
+            continue
+        # compare offsets from the first list array with the current one
+        if not first_list_array.offsets.equals(list_array.offsets):
+            raise ValueError("Offsets of all ListArrays must be the same")
+
+
+def transpose_struct_list_type(t: pa.StructType) -> pa.ListType:
+    """Converts a type of struct-list array into a type of list-struct array.
+
+    Parameters
+    ----------
+    t : pa.DataType
+        Input type of struct-list array.
+
+    Returns
+    -------
+    pa.DataType
+        Type of list-struct array.
+
+    Raises
+    ------
+    ValueError
+        If the input type is not a struct-list type.
+    """
+    if not pa.types.is_struct(t):
+        raise ValueError(f"Expected a StructType, got {t}")
+
+    fields = []
+    for field in t:
+        if not is_pa_type_a_list(field.type):
+            raise ValueError(f"Expected a ListType, got {field.type}")
+        list_type = cast(pa.ListType, field.type)
+        fields.append(pa.field(field.name, list_type.value_type))
+
+    list_type = cast(pa.ListType, pa.list_(pa.struct(fields)))
+    return list_type
+
+
+def transpose_struct_list_array(array: pa.StructArray, validate: bool = True) -> pa.ListArray:
+    """Converts a struct-array of lists into a list-array of structs.
+
+    Parameters
+    ----------
+    array : pa.StructArray
+        Input struct array, each scalar must have lists of equal length.
+    validate : bool, default True
+        Whether to validate the input array for list lengths. Raises ValueError
+        if something is wrong.
+
+    Returns
+    -------
+    pa.ListArray
+        List array of structs.
+    """
+    if validate:
+        validate_struct_list_array_for_equal_lengths(array)
+
+    # Since we know that all lists have the same length, we can use the first list to get offsets
+    offsets = array.field(0).offsets
+    struct_flat_array = pa.StructArray.from_arrays(
+        [field.values for field in array.flatten()],
+        names=array.type.names,
+    )
+    return pa.ListArray.from_arrays(offsets, struct_flat_array)
+
+
+def transpose_list_struct_type(t: pa.ListType) -> pa.StructType:
+    """Converts a type of list-struct array into a type of struct-list array.
+
+    Parameters
+    ----------
+    t : pa.DataType
+        Input type of list-struct array.
+
+    Returns
+    -------
+    pa.DataType
+        Type of struct-list array.
+
+    Raises
+    ------
+    ValueError
+        If the input type is not a list-struct type.
+    """
+    if not is_pa_type_a_list(t):
+        raise ValueError(f"Expected a ListType, got {t}")
+
+    struct_type = cast(pa.StructType, t.value_type)
+    fields = []
+    for field in struct_type:
+        fields.append(pa.field(field.name, pa.list_(field.type)))
+
+    struct_type = cast(pa.StructType, pa.struct(fields))
+    return struct_type
+
+
+def transpose_list_struct_array(array: pa.ListArray) -> pa.StructArray:
+    """Converts a list-array of structs into a struct-array of lists.
+
+    Parameters
+    ----------
+    array : pa.ListArray
+        Input list array of structs.
+
+    Returns
+    -------
+    pa.StructArray
+        Struct array of lists.
+    """
+    offsets, values = array.offsets, array.values
+
+    fields = []
+    for field_values in values.flatten():
+        list_array = pa.ListArray.from_arrays(offsets, field_values)
+        fields.append(list_array)
+
+    return pa.StructArray.from_arrays(fields, names=array.type.value_type.names)