Add NestedExtensionArray.list_lengths

Author: hombit

src/nested_pandas/series/ext_array.py

@@ -761,12 +761,24 @@ 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]:
+        """Lengths of the list arrays"""
+        return list(self._iter_list_lengths())
+
     @property
     def flat_length(self) -> int:
         """Length of the flat arrays"""
-        return sum(
-            chunk.field(0).value_lengths().sum().as_py() or 0 for chunk in self._chunked_array.iterchunks()
-        )
+        return sum(self._iter_list_lengths())
 
     @property
     def num_chunks(self) -> int:

tests/nested_pandas/series/test_ext_array.py

@@ -1226,6 +1226,24 @@ def test_field_names():
     assert ext_array.field_names == ["a", "b"]
 
 
+def test_list_lengths():
+    """Tests that the list lengths of the extension array are correct."""
+    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"],
+    )
+    empty_struct_array = pa.array([], type=struct_array.type)
+    null_struct_array = pa.array([None], type=struct_array.type)
+    ext_array = NestedExtensionArray(
+        pa.chunked_array([struct_array, empty_struct_array, struct_array, null_struct_array])
+    )
+
+    assert ext_array.list_lengths == [3, 4, 3, 4, 0]
+
+
 def test_flat_length():
     """Tests that the flat length of the extension array is correct."""
     struct_array = pa.StructArray.from_arrays(