Merge pull request #254 from lincc-frameworks/list-struct-backend

List-struct backend

Author: hombit

docs/gettingstarted/quickstart.ipynb

@@ -345,7 +345,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "nested",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -359,9 +359,9 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.4"
+   "version": "3.13.3"
   }
  },
  "nbformat": 4,
- "nbformat_minor": 2
+ "nbformat_minor": 4
 }

docs/reference/ext_array.rst

@@ -15,6 +15,16 @@ Functions
 .. autosummary::
     :toctree: api/
 
+    series.ext_array.NestedExtensionArray.dtype
+    series.ext_array.NestedExtensionArray.nbytes
+    series.ext_array.NestedExtensionArray.list_array
+    series.ext_array.NestedExtensionArray.struct_array
+    series.ext_array.NestedExtensionArray.py_table
+    series.ext_array.NestedExtensionArray.list_offsets
+    series.ext_array.NestedExtensionArray.field_names
+    series.ext_array.NestedExtensionArray.list_lengths
+    series.ext_array.NestedExtensionArray.flat_length
+    series.ext_array.NestedExtensionArray.num_chunks
     series.ext_array.NestedExtensionArray.to_numpy
     series.ext_array.NestedExtensionArray.isna
     series.ext_array.NestedExtensionArray.take
@@ -24,6 +34,7 @@ Functions
     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.get_list_index
     series.ext_array.NestedExtensionArray.iter_field_lists
     series.ext_array.NestedExtensionArray.view_fields
     series.ext_array.NestedExtensionArray.set_flat_field

src/nested_pandas/series/_storage/__init__.py

@@ -0,0 +1,3 @@
+from .list_struct_storage import ListStructStorage  # noqa: F401
+from .struct_list_storage import StructListStorage  # noqa: F401
+from .table_storage import TableStorage  # noqa: F401

src/nested_pandas/series/_storage/list_struct_storage.py

@@ -0,0 +1,84 @@
+from __future__ import annotations  # Python 3.9 requires it for X | Y type hints
+
+from typing import TYPE_CHECKING, Any
+
+import pyarrow as pa
+
+from nested_pandas.series.utils import transpose_struct_list_chunked, validate_list_struct_type
+
+if TYPE_CHECKING:
+    from nested_pandas.series._storage.struct_list_storage import StructListStorage
+    from nested_pandas.series._storage.table_storage import TableStorage
+
+
+class ListStructStorage:
+    """Store nested data as a PyArrow list-struct array.
+
+    Parameters
+    ----------
+    array : pa.ListArray or pa.ChunkedArray
+        Pyarrow list-array with a struct value type. An array or a chunk-array
+    """
+
+    _data: pa.ChunkedArray
+
+    def __init__(self, array: pa.ListArray | pa.ChunkedArray) -> None:
+        if isinstance(array, pa.ListArray):
+            array = pa.chunked_array([array])
+        if not isinstance(array, pa.ChunkedArray):
+            raise ValueError("array must be of type pa.ChunkedArray")
+        validate_list_struct_type(array.type)
+        self._data = array
+
+    @property
+    def data(self) -> pa.ChunkedArray:
+        return self._data
+
+    @classmethod
+    def from_struct_list_storage(cls, struct_list_storage: StructListStorage) -> Self:  # type: ignore # noqa: F821
+        """Construct from a StructListStorage object.
+
+        Parameters
+        ----------
+        struct_list_storage : StructListStorage
+            StructListStorage object.
+        """
+        data = transpose_struct_list_chunked(struct_list_storage.data, validate=False)
+        return cls(data)
+
+    @classmethod
+    def from_table_storage(cls, table_storage: TableStorage) -> Self:  # type: ignore # noqa: F821
+        """Construct from a TableStorage object.
+
+        Parameters
+        ----------
+        table_storage : TableStorage
+            TableStorage object.
+        """
+        from nested_pandas.series._storage import StructListStorage
+
+        struct_list_storage = StructListStorage.from_table_storage(table_storage)
+        return cls.from_struct_list_storage(struct_list_storage)
+
+    def __len__(self) -> int:
+        return len(self._data)
+
+    def __eq__(self, other: Any) -> bool:
+        if not isinstance(other, type(self)):
+            return False
+        return self._data == other._data
+
+    @property
+    def nbytes(self) -> int:
+        """Number of bytes consumed by the data in memory."""
+        return self._data.nbytes
+
+    @property
+    def type(self) -> pa.ListType:
+        """Pyarrow type of the underlying array."""
+        return self._data.type
+
+    @property
+    def num_chunks(self) -> int:
+        """Number of chunks in the underlying array."""
+        return self._data.num_chunks

src/nested_pandas/series/_storage/struct_list_storage.py

@@ -0,0 +1,74 @@
+from __future__ import annotations  # Python 3.9 requires it for X | Y type hints
+
+from collections.abc import Iterator
+from typing import TYPE_CHECKING
+
+import pyarrow as pa
+
+from nested_pandas.series.utils import (
+    table_to_struct_array,
+    transpose_list_struct_chunked,
+    validate_struct_list_array_for_equal_lengths,
+)
+
+if TYPE_CHECKING:
+    from nested_pandas.series._storage.list_struct_storage import ListStructStorage
+    from nested_pandas.series._storage.table_storage import TableStorage
+
+
+class StructListStorage:
+    """Store nested data as a PyArrow struct-list array.
+
+    Parameters
+    ----------
+    array : pa.StructArray or pa.ChunkedArray
+        Pyarrow struct-array with all fields to be list-arrays.
+        All list-values must be "aligned", e.g., have the same length.
+    validate : bool (default True)
+        Check that all the lists have the same lengths for each struct-value.
+    """
+
+    _data: pa.ChunkedArray
+
+    def __init__(self, array: pa.StructArray | pa.ChunkedArray, *, validate: bool = True) -> None:
+        if isinstance(array, pa.StructArray):
+            array = pa.chunked_array([array])
+        if not isinstance(array, pa.ChunkedArray):
+            raise ValueError("array must be a StructArray or ChunkedArray")
+
+        if validate:
+            for chunk in array.chunks:
+                validate_struct_list_array_for_equal_lengths(chunk)
+
+        self._data = array
+
+    @property
+    def data(self) -> pa.ChunkedArray:
+        return self._data
+
+    @classmethod
+    def from_list_struct_storage(cls, list_struct_storage: ListStructStorage) -> Self:  # type: ignore # noqa: F821
+        """Construct from a ListStructStorage object.
+
+        Parameters
+        ----------
+        list_struct_storage : ListStructStorage
+            ListStructStorage object.
+        """
+        data = transpose_list_struct_chunked(list_struct_storage.data)
+        return cls(data, validate=False)
+
+    @classmethod
+    def from_table_storage(cls, table_storage: TableStorage) -> Self:  # type: ignore # noqa: F821
+        """Construct from a TableStorage object.
+
+        Parameters
+        ----------
+        table_storage : TableStorage
+            TableStorage object.
+        """
+        data = table_to_struct_array(table_storage.data)
+        return cls(data, validate=False)
+
+    def __iter__(self) -> Iterator[pa.StructScalar]:
+        return iter(self._data)

src/nested_pandas/series/_storage/table_storage.py

@@ -0,0 +1,62 @@
+from __future__ import annotations  # Python 3.9 requires it for X | Y type hints
+
+from typing import TYPE_CHECKING
+
+import pyarrow as pa
+
+from nested_pandas.series.utils import table_to_struct_array, validate_struct_list_array_for_equal_lengths
+
+if TYPE_CHECKING:
+    from nested_pandas.series._storage.list_struct_storage import ListStructStorage
+    from nested_pandas.series._storage.struct_list_storage import StructListStorage
+
+
+class TableStorage:
+    """Store nested data as a PyArrow table with list-columns.
+
+    Parameters
+    ----------
+    table : pa.Table
+        PyArrow table, all columns must be list-columns.
+        All list-values must be "aligned", e.g., have the same length.
+    """
+
+    _data: pa.Table
+
+    def __init__(self, table: pa.Table, validate: bool = True) -> None:
+        if validate:
+            struct_array = table_to_struct_array(table)
+            for chunk in struct_array.iterchunks():
+                validate_struct_list_array_for_equal_lengths(chunk)
+
+        self._data = table
+
+    @property
+    def data(self) -> pa.Table:
+        return self._data
+
+    @classmethod
+    def from_list_struct_storage(cls, list_storage: ListStructStorage) -> Self:  # type: ignore # noqa: F821
+        """Construct from a StructListStorage object.
+
+        Parameters
+        ----------
+        list_storage : ListStructStorage
+            StructListStorage object.
+        """
+        from nested_pandas.series._storage import StructListStorage
+
+        struct_list_storage = StructListStorage.from_list_struct_storage(list_storage)
+        return cls.from_struct_list_storage(struct_list_storage)
+
+    @classmethod
+    def from_struct_list_storage(cls, struct_list_storage: StructListStorage) -> Self:  # type: ignore # noqa: F821
+        """Construct from a StructListStorage object.
+
+        Parameters
+        ----------
+        struct_list_storage : StructListStorage
+            StructListStorage object.
+        """
+        table = pa.Table.from_struct_array(struct_list_storage.data)
+        return cls(table, validate=False)

src/nested_pandas/series/accessor.py

@@ -13,6 +13,7 @@
 
 from nested_pandas.series.dtype import NestedDtype
 from nested_pandas.series.packer import pack_sorted_df_into_struct
+from nested_pandas.series.utils import nested_types_mapper
 
 __all__ = ["NestSeriesAccessor"]
 
@@ -70,25 +71,10 @@ def to_lists(self, fields: list[str] | None = None) -> pd.DataFrame:
         if len(fields) == 0:
             raise ValueError("Cannot convert a struct with no fields to lists")
 
-        list_chunks = defaultdict(list)
-        for chunk in self._series.array._chunked_array.iterchunks():
-            struct_array = cast(pa.StructArray, chunk)
-            for field in fields:
-                list_array = cast(pa.ListArray, struct_array.field(field))
-                list_chunks[field].append(list_array)
+        list_df = self._series.array.pa_table.select(fields).to_pandas(types_mapper=nested_types_mapper)
+        list_df.index = self._series.index
 
-        list_series = {}
-        for field, chunks in list_chunks.items():
-            chunked_array = pa.chunked_array(chunks)
-            list_series[field] = pd.Series(
-                chunked_array,
-                dtype=pd.ArrowDtype(chunked_array.type),
-                index=self._series.index,
-                name=field,
-                copy=False,
-            )
-
-        return pd.DataFrame(list_series)
+        return list_df
 
     def to_flat(self, fields: list[str] | None = None) -> pd.DataFrame:
         """Convert nested series into dataframe of flat arrays
@@ -130,7 +116,7 @@ def to_flat(self, fields: list[str] | None = None) -> pd.DataFrame:
         index = pd.Series(self.get_flat_index(), name=self._series.index.name)
 
         flat_chunks = defaultdict(list)
-        for chunk in self._series.array._chunked_array.iterchunks():
+        for chunk in self._series.array.struct_array.iterchunks():
             struct_array = cast(pa.StructArray, chunk)
             for field in fields:
                 list_array = cast(pa.ListArray, struct_array.field(field))
@@ -439,7 +425,7 @@ def get_flat_series(self, field: str) -> pd.Series:
         """
 
         flat_chunks = []
-        for nested_chunk in self._series.array._chunked_array.iterchunks():
+        for nested_chunk in self._series.array.struct_array.iterchunks():
             struct_array = cast(pa.StructArray, nested_chunk)
             list_array = cast(pa.ListArray, struct_array.field(field))
             flat_array = list_array.flatten()
@@ -483,12 +469,7 @@ def get_list_series(self, field: str) -> pd.Series:
         Name: flux, dtype: list<item: double>[pyarrow]
         """
 
-        list_chunks = []
-        for nested_chunk in self._series.array._chunked_array.iterchunks():
-            struct_array = cast(pa.StructArray, nested_chunk)
-            list_array = struct_array.field(field)
-            list_chunks.append(list_array)
-        list_chunked_array = pa.chunked_array(list_chunks)
+        list_chunked_array = self._series.array.pa_table[field]
         return pd.Series(
             list_chunked_array,
             dtype=pd.ArrowDtype(list_chunked_array.type),

src/nested_pandas/series/dtype.py

@@ -154,7 +154,12 @@ 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.list_struct_pyarrow_dtype = self._validate_dtype(pyarrow_dtype)
+        self.pyarrow_dtype, self.list_struct_pa_dtype = self._validate_dtype(pyarrow_dtype)
+
+    @property
+    def struct_list_pa_dtype(self) -> pa.StructType:
+        """Struct-list pyarrow type representing the nested type."""
+        return self.pyarrow_dtype
 
     @classmethod
     def from_fields(cls, fields: Mapping[str, pa.DataType]) -> Self:  # type: ignore[name-defined] # noqa: F821
@@ -261,5 +266,5 @@ def to_pandas_arrow_dtype(self, list_struct: bool = False) -> ArrowDtype:
             The corresponding pandas.ArrowDtype.
         """
         if list_struct:
-            return ArrowDtype(self.list_struct_pyarrow_dtype)
+            return ArrowDtype(self.list_struct_pa_dtype)
         return ArrowDtype(self.pyarrow_dtype)