Preparation for list-struct backend for NestedExtensionArray (#242)

* Allow NestedDtype initialization from pd.ArrowDtype

* NestedDtype.inner_dtypes

* .nest.to_flat to respect inner_dtypes

* Allow ArrowDtype in NestedDtype.from_fields

* Handle and derive inner_dtypes

* Fix a typo in variable name

* Pull changes from inner-dtypes branch

* NestedDtype.to_pandas_arrow_dtype(list_struct: bool)

* Building NestedDtype from list-struct

* NestedExtensionArray._pa_table and ._pa_struct_array

* test NestedExtensionArray.to_pyarrow_scalar

* Add to_pyarrow_scalar to API docs

Author: hombit

docs/reference/ext_array.rst

@@ -23,6 +23,7 @@ Functions
     series.ext_array.NestedExtensionArray.dropna
     series.ext_array.NestedExtensionArray.from_sequence
     series.ext_array.NestedExtensionArray.to_arrow_ext_array
+    series.ext_array.NestedExtensionArray.to_pyarrow_scalar
     series.ext_array.NestedExtensionArray.iter_field_lists
     series.ext_array.NestedExtensionArray.view_fields
     series.ext_array.NestedExtensionArray.set_flat_field

src/nested_pandas/series/dtype.py

@@ -13,14 +13,24 @@
 from pandas.core.arrays import ExtensionArray
 from pandas.core.dtypes.base import ExtensionDtype
 
-from nested_pandas.series.utils import is_pa_type_a_list
+from nested_pandas.series.utils import (
+    transpose_list_struct_type,
+    transpose_struct_list_type,
+)
 
 __all__ = ["NestedDtype"]
 
 
 @register_extension_dtype
 class NestedDtype(ExtensionDtype):
-    """Data type to handle packed time series data"""
+    """Data type to handle packed time series data
+
+    Parameters
+    ----------
+    pyarrow_dtype : pyarrow.StructType or pd.ArrowDtype
+        The pyarrow data type to use for the nested type. It must be a struct
+        type where all fields are list types.
+    """
 
     # ExtensionDtype overrides #
 
@@ -38,7 +48,12 @@ def na_value(self) -> Type[pd.NA]:
     @property
     def name(self) -> str:
         """The string representation of the nested type"""
-        fields = ", ".join([f"{field.name}: [{field.type.value_type!s}]" for field in self.pyarrow_dtype])
+        # Replace pd.ArrowDtype with pa.DataType, because it has nicer __str__
+        nice_dtypes = {
+            field: dtype.pyarrow_dtype if isinstance(dtype, pd.ArrowDtype) else dtype
+            for field, dtype in self.fields.items()
+        }
+        fields = ", ".join([f"{field}: [{dtype!s}]" for field, dtype in nice_dtypes.items()])
         return f"nested<{fields}>"
 
     @classmethod
@@ -136,7 +151,7 @@ def __from_arrow__(self, array: pa.Array | pa.ChunkedArray) -> ExtensionArray:
     pyarrow_dtype: pa.StructType
 
     def __init__(self, pyarrow_dtype: pa.DataType) -> None:
-        self.pyarrow_dtype = self._validate_dtype(pyarrow_dtype)
+        self.pyarrow_dtype, self.list_struct_pyarrow_dtype = self._validate_dtype(pyarrow_dtype)
 
     @classmethod
     def from_fields(cls, fields: Mapping[str, pa.DataType]) -> Self:  # type: ignore[name-defined] # noqa: F821
@@ -168,20 +183,33 @@ def from_fields(cls, fields: Mapping[str, pa.DataType]) -> Self:  # type: ignore
         return cls(pyarrow_dtype=pyarrow_dtype)
 
     @staticmethod
-    def _validate_dtype(pyarrow_dtype: pa.DataType) -> pa.StructType:
+    def _validate_dtype(pyarrow_dtype: pa.DataType) -> tuple[pa.StructType, pa.ListType]:
+        """Check that the given pyarrow type is castable to the nested type.
+
+        Parameters
+        ----------
+        pyarrow_dtype : pa.DataType
+            The pyarrow type to check and cast.
+
+        Returns
+        -------
+        pa.StructType
+            Struct-list pyarrow type representing the nested type.
+        pa.ListType
+            List-struct pyarrow type representing the nested type.
+        """
         if not isinstance(pyarrow_dtype, pa.DataType):
             raise TypeError(f"Expected a 'pyarrow.DataType' object, got {type(pyarrow_dtype)}")
-        if not pa.types.is_struct(pyarrow_dtype):
-            raise ValueError("NestedDtype can only be constructed with pyarrow struct type.")
-        pyarrow_dtype = cast(pa.StructType, pyarrow_dtype)
-
-        for field in pyarrow_dtype:
-            if not is_pa_type_a_list(field.type):
-                raise ValueError(
-                    "NestedDtype can only be constructed with pyarrow struct type, all fields must be list "
-                    f"type. Given struct has unsupported field {field}"
-                )
-        return pyarrow_dtype
+        if pa.types.is_struct(pyarrow_dtype):
+            struct_type = cast(pa.StructType, pyarrow_dtype)
+            return struct_type, transpose_struct_list_type(struct_type)
+        # Currently, LongList and others are not supported
+        if pa.types.is_list(pyarrow_dtype):
+            list_type = cast(pa.ListType, pyarrow_dtype)
+            return transpose_list_struct_type(list_type), list_type
+        raise ValueError(
+            f"NestedDtype can only be constructed with pa.StructType or pa.ListType only, got {pyarrow_dtype}"
+        )
 
     @property
     def fields(self) -> dict[str, pa.DataType]:
@@ -194,13 +222,14 @@ def field_names(self) -> list[str]:
         return [field.name for field in self.pyarrow_dtype]
 
     @classmethod
-    def from_pandas_arrow_dtype(cls, pandas_arrow_dtype: ArrowDtype):
+    def from_pandas_arrow_dtype(cls, pandas_arrow_dtype: ArrowDtype) -> Self:  # type: ignore[name-defined] # noqa: F821
         """Construct NestedDtype from a pandas.ArrowDtype.
 
         Parameters
         ----------
         pandas_arrow_dtype : ArrowDtype
             The pandas.ArrowDtype to construct NestedDtype from.
+            Must be struct-list or list-struct type.
 
         Returns
         -------
@@ -214,12 +243,20 @@ def from_pandas_arrow_dtype(cls, pandas_arrow_dtype: ArrowDtype):
         """
         return cls(pyarrow_dtype=pandas_arrow_dtype.pyarrow_dtype)
 
-    def to_pandas_arrow_dtype(self) -> ArrowDtype:
+    def to_pandas_arrow_dtype(self, list_struct: bool = False) -> ArrowDtype:
         """Convert NestedDtype to a pandas.ArrowDtype.
 
+        Parameters
+        ----------
+        list_struct : bool, default False
+            If False (default) use pyarrow struct-list type,
+            otherwise use pyarrow list-struct type.
+
         Returns
         -------
         ArrowDtype
             The corresponding pandas.ArrowDtype.
         """
+        if list_struct:
+            return ArrowDtype(self.list_struct_pyarrow_dtype)
         return ArrowDtype(self.pyarrow_dtype)

src/nested_pandas/series/ext_array.py

@@ -225,7 +225,7 @@ def _from_sequence(cls, scalars, *, dtype=None, copy: bool = False) -> Self:  #
         pa_array = cls._box_pa_array(scalars, pa_type=pa_type)
         return cls(pa_array)
 
-    # Tricky to implement, but required by things like pd.read_csv
+    # Tricky to implement but required by things like pd.read_csv
     @classmethod
     def _from_sequence_of_strings(cls, strings, *, dtype=None, copy: bool = False) -> Self:  # type: ignore[name-defined] # noqa: F821
         return super()._from_sequence_of_strings(strings, dtype=dtype, copy=copy)
@@ -680,6 +680,29 @@ def _list_array(self) -> pa.ChunkedArray:
             list_chunks.append(transpose_struct_list_array(struct_chunk, validate=False))
         return pa.chunked_array(list_chunks)
 
+    @property
+    def _struct_array(self) -> pa.ChunkedArray:
+        """Pyarrow chunked struct-list array representation
+
+        Returns
+        -------
+        pa.ChunkedArray
+            Pyarrow chunked-array of struct-list arrays.
+        """
+        return self._chunked_array
+
+    @property
+    def _pa_table(self) -> pa.Table:
+        """Pyarrow table representation of the extension array.
+
+        Returns
+        -------
+        pa.Table
+            Pyarrow table where each column is a list array corresponding
+            to a field of the struct array.
+        """
+        return pa.Table.from_struct_array(self._struct_array)
+
     @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
@@ -755,6 +778,23 @@ def to_arrow_ext_array(self, list_struct: bool = False) -> ArrowExtensionArray:
             return ArrowExtensionArray(self._list_array)
         return ArrowExtensionArray(self._chunked_array)
 
+    def to_pyarrow_scalar(self, list_struct: bool = False) -> pa.ListScalar:
+        """Convert to a pyarrow scalar of a list type
+
+        Parameters
+        ----------
+        list_struct : bool, optional
+            If False (default), return list-struct-list scalar,
+            otherwise list-list-struct scalar.
+
+        Returns
+        -------
+        pyarrow.ListScalar
+        """
+        pa_array = self._list_array if list_struct else self._chunked_array
+        pa_type = pa.list_(pa_array.type)
+        return cast(pa.ListScalar, pa.scalar(pa_array, type=pa_type))
+
     def _replace_chunked_array(self, pa_array: pa.ChunkedArray, *, validate: bool) -> None:
         if validate:
             self._validate(pa_array)

src/nested_pandas/series/utils.py

@@ -157,6 +157,9 @@ def transpose_list_struct_type(t: pa.ListType) -> pa.StructType:
     if not is_pa_type_a_list(t):
         raise ValueError(f"Expected a ListType, got {t}")
 
+    if not pa.types.is_struct(t.value_type):
+        raise ValueError(f"Expected a StructType as a list value type, got {t.value_type}")
+
     struct_type = cast(pa.StructType, t.value_type)
     fields = []
     for field in struct_type:

tests/nested_pandas/series/test_dtype.py

@@ -18,18 +18,38 @@
         ),
     ],
 )
-def test_from_pyarrow_dtype(pyarrow_dtype):
+def test_from_pyarrow_dtype_struct_list(pyarrow_dtype):
     """Test that we can construct NestedDtype from pyarrow struct type."""
     dtype = NestedDtype(pyarrow_dtype)
     assert dtype.pyarrow_dtype == pyarrow_dtype
 
 
+@pytest.mark.parametrize(
+    "pyarrow_dtype",
+    [
+        pa.list_(pa.struct([pa.field("a", pa.int64())])),
+        pa.list_(pa.struct([pa.field("a", pa.int64()), pa.field("b", pa.float64())])),
+        pa.list_(
+            pa.struct(
+                [
+                    pa.field("a", pa.list_(pa.int64())),
+                    pa.field("b", pa.list_(pa.float64())),
+                ]
+            )
+        ),
+    ],
+)
+def test_from_pyarrow_dtype_list_struct(pyarrow_dtype):
+    """Test that we can construct NestedDtype from pyarrow list type."""
+    dtype = NestedDtype(pyarrow_dtype)
+    assert dtype.list_struct_pyarrow_dtype == pyarrow_dtype
+
+
 @pytest.mark.parametrize(
     "pyarrow_dtype",
     [
         pa.int64(),
         pa.list_(pa.int64()),
-        pa.list_(pa.struct([pa.field("a", pa.int64())])),
         pa.struct([pa.field("a", pa.int64())]),
         pa.struct([pa.field("a", pa.int64()), pa.field("b", pa.float64())]),
         pa.struct([pa.field("a", pa.list_(pa.int64())), pa.field("b", pa.float64())]),
@@ -49,6 +69,26 @@ def test_to_pandas_arrow_dtype():
     )
 
 
+def test_from_pandas_arrow_dtype():
+    """Test that we can construct NestedDtype from pandas.ArrowDtype."""
+    dtype_from_struct = NestedDtype.from_pandas_arrow_dtype(
+        pd.ArrowDtype(pa.struct([pa.field("a", pa.list_(pa.int64()))]))
+    )
+    assert dtype_from_struct.pyarrow_dtype == pa.struct([pa.field("a", pa.list_(pa.int64()))])
+    dtype_from_list = NestedDtype.from_pandas_arrow_dtype(
+        pd.ArrowDtype(pa.list_(pa.struct([pa.field("a", pa.int64())])))
+    )
+    assert dtype_from_list.pyarrow_dtype == pa.struct([pa.field("a", pa.list_(pa.int64()))])
+
+
+def test_to_pandas_list_struct_arrow_dtype():
+    """Test that NestedDtype.to_pandas_arrow_dtype(list_struct=True) returns the correct pyarrow type."""
+    dtype = NestedDtype.from_fields({"a": pa.list_(pa.int64()), "b": pa.float64()})
+    assert dtype.to_pandas_arrow_dtype(list_struct=True) == pd.ArrowDtype(
+        pa.list_(pa.struct([pa.field("a", pa.list_(pa.int64())), pa.field("b", pa.float64())]))
+    )
+
+
 def test_from_fields():
     """Test NestedDtype.from_fields()."""
     fields = {"a": pa.int64(), "b": pa.float64()}

tests/nested_pandas/series/test_ext_array.py

@@ -626,6 +626,38 @@ def test_chunked_list_struct_array():
     assert ext_array.chunked_list_struct_array.type == ext_array._pyarrow_list_struct_dtype
 
 
+def test_to_pyarrow_scalar():
+    """Test .to_pyarrow_scalar 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)
+
+    desired_struct_list = pa.scalar(
+        [
+            {"a": [1, 2, 3], "b": [-4.0, -5.0, -6.0]},
+            {"a": [1, 2, 1], "b": [-3.0, -4.0, -5.0]},
+        ],
+        type=pa.list_(
+            pa.struct([pa.field("a", pa.list_(pa.int64())), pa.field("b", pa.list_(pa.float64()))])
+        ),
+    )
+    desired_list_struct = pa.scalar(
+        [
+            [{"a": 1, "b": -4.0}, {"a": 2, "b": -5.0}, {"a": 3, "b": -6.0}],
+            [{"a": 1, "b": -3.0}, {"a": 2, "b": -4.0}, {"a": 1, "b": -5.0}],
+        ],
+        type=pa.list_(pa.list_(pa.struct([pa.field("a", pa.int64()), pa.field("b", pa.float64())]))),
+    )
+    # pyarrow returns a single bool for ==
+    assert ext_array.to_pyarrow_scalar(list_struct=False) == desired_struct_list
+    assert ext_array.to_pyarrow_scalar(list_struct=True) == desired_list_struct
+
+
 def test_list_offsets_single_chunk():
     """Test that the .list_offset property is correct for a single chunk."""
     struct_array = pa.StructArray.from_arrays(
@@ -1873,6 +1905,36 @@ def test___init___with_list_struct_array():
     assert pa.array(ext_array) == struct_array
 
 
+def test__struct_array():
+    """Test ._struct_array property"""
+    struct_array = pa.StructArray.from_arrays(
+        arrays=[
+            pa.array([np.array([1.0, 2.0, 3.0]), np.array([1.0, 2.0, 1.0, 2.0])]),
+            pa.array([-np.array([4.0, 5.0, 6.0]), -np.array([3.0, 4.0, 5.0, 6.0])]),
+        ],
+        names=["a", "b"],
+    )
+    ext_array = NestedExtensionArray(struct_array)
+
+    assert ext_array._struct_array.combine_chunks() == struct_array
+
+
+def test__pa_table():
+    """Test ._pa_table property"""
+    struct_array = pa.StructArray.from_arrays(
+        arrays=[
+            pa.array([np.array([1.0, 2.0, 3.0]), np.array([1.0, 2.0, 1.0, 2.0])]),
+            pa.array([-np.array([4.0, 5.0, 6.0]), -np.array([3.0, 4.0, 5.0, 6.0])]),
+        ],
+        names=["a", "b"],
+    )
+    ext_array = NestedExtensionArray(struct_array)
+
+    assert ext_array._pa_table == pa.Table.from_arrays(
+        arrays=[struct_array.field("a"), struct_array.field("b")], names=["a", "b"]
+    )
+
+
 def test__from_sequence_of_strings():
     """We do not support from_sequence_of_strings() which would apply things like pd.read_csv()"""
     with pytest.raises(NotImplementedError):