NestedExtensionArray._pa_table and ._pa_struct_array

Author: hombit

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

tests/nested_pandas/series/test_ext_array.py

@@ -1873,6 +1873,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):