Support fixed-size and large lists (#418)

* Support fixed-size lists

* Apply suggestions from code review

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Support fixed-size lists

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: hombit

docs/reference/ext_array.rst

@@ -41,3 +41,4 @@ Functions
     series.ext_array.NestedExtensionArray.set_list_field
     series.ext_array.NestedExtensionArray.fill_field_lists
     series.ext_array.NestedExtensionArray.pop_fields
+    series.ext_array.NestedExtensionArray.is_input_pa_type_supported

src/nested_pandas/nestedframe/core.py

@@ -24,9 +24,9 @@
     _subexprs_by_nest,
 )
 from nested_pandas.series.dtype import NestedDtype
+from nested_pandas.series.ext_array import NestedExtensionArray
 from nested_pandas.series.nestedseries import NestedSeries
 from nested_pandas.series.packer import pack, pack_lists, pack_sorted_df_into_struct
-from nested_pandas.series.utils import is_pa_type_a_list
 
 pd.set_option("display.max_rows", 30)
 pd.set_option("display.min_rows", 5)
@@ -68,13 +68,11 @@ def _cast_cols_to_nested(self, *, struct_list: bool) -> None:
             if not isinstance(dtype, pd.ArrowDtype):
                 continue
             pa_type = dtype.pyarrow_dtype
-            if not is_pa_type_a_list(pa_type) and not (struct_list and pa.types.is_struct(pa_type)):
+            if pa.types.is_struct(pa_type) and not struct_list:
                 continue
-            try:
-                nested_dtype = NestedDtype(pa_type)
-            except (TypeError, ValueError):
+            if not NestedExtensionArray.is_input_pa_type_supported(pa_type):
                 continue
-            self[column] = self[column].astype(nested_dtype)
+            self[column] = NestedExtensionArray(pa.array(self[column]))
 
     @property
     def _constructor(self) -> Self:  # type: ignore[name-defined] # noqa: F821

src/nested_pandas/nestedframe/io.py

@@ -11,9 +11,9 @@
 from pyarrow.lib import ArrowInvalid
 from upath import UPath
 
-from ..series.dtype import NestedDtype
+from ..series.ext_array import NestedExtensionArray
 from ..series.packer import pack_lists
-from ..series.utils import table_to_struct_array
+from ..series.utils import is_pa_type_a_list, table_to_struct_array
 from .core import NestedFrame
 
 # Use smaller block size for these FSSPEC filesystems.
@@ -148,7 +148,7 @@ def read_parquet(
                 # if any of the columns are not list type, reject the cast
                 # and remove the column from the list of nested structures if
                 # it was added
-                if not pa.types.is_list(table.schema[i].type):
+                if not is_pa_type_a_list(table.schema[i].type):
                     reject_nesting.append(nested_col)
                     if nested_col in nested_structures:
                         # remove the column from the list of nested structures
@@ -455,28 +455,25 @@ def _cast_struct_cols_to_nested(df, reject_nesting):
     """cast struct columns to nested dtype"""
     # Attempt to cast struct columns to NestedDTypes
     for col, dtype in df.dtypes.items():
-        # First validate the dtype
-        # will return valueerror when not a struct-list
-        valid_dtype = True
+        if col in reject_nesting:
+            continue
+
+        if not NestedExtensionArray.is_input_pa_type_supported(dtype.pyarrow_dtype):
+            continue
+
         try:
-            NestedDtype._validate_dtype(dtype.pyarrow_dtype)
-        except ValueError:
-            valid_dtype = False
-
-        if valid_dtype and col not in reject_nesting:
-            try:
-                # Attempt to cast Struct to NestedDType
-                df = df.astype({col: NestedDtype(dtype.pyarrow_dtype)})
-            except ValueError as err:
-                # If cast fails, the struct likely does not fit nested-pandas
-                # criteria for a valid nested column
-                raise ValueError(
-                    f"Column '{col}' is a Struct, but an attempt to cast it to a NestedDType failed. "
-                    "This is likely due to the struct not meeting the requirements for a nested column "
-                    "(all fields should be equal length). To proceed, you may add the column to the "
-                    "`reject_nesting` argument of the read_parquet function to skip the cast attempt:"
-                    f" read_parquet(..., reject_nesting=['{col}'])"
-                ) from err
+            # Attempt to cast Struct to NestedDType
+            df[col] = NestedExtensionArray(pa.array(df[col]))
+        except ValueError as err:
+            # If cast fails, the struct likely does not fit nested-pandas
+            # criteria for a valid nested column
+            raise ValueError(
+                f"Column '{col}' is a Struct, but an attempt to cast it to a NestedDType failed. "
+                "This is likely due to the struct not meeting the requirements for a nested column "
+                "(all fields should be equal length). To proceed, you may add the column to the "
+                "`reject_nesting` argument of the read_parquet function to skip the cast attempt:"
+                f" read_parquet(..., reject_nesting=['{col}'])"
+            ) from err
     return df
 
 

src/nested_pandas/series/_storage/list_struct_storage.py

@@ -4,7 +4,11 @@
 
 import pyarrow as pa
 
-from nested_pandas.series.utils import transpose_struct_list_chunked, validate_list_struct_type
+from nested_pandas.series.utils import (
+    normalize_list_array,
+    transpose_struct_list_chunked,
+    validate_list_struct_type,
+)
 
 if TYPE_CHECKING:
     from nested_pandas.series._storage.struct_list_storage import StructListStorage
@@ -22,11 +26,14 @@ class ListStructStorage:
 
     _data: pa.ChunkedArray
 
-    def __init__(self, array: pa.ListArray | pa.ChunkedArray) -> None:
+    def __init__(
+        self, array: pa.ListArray | pa.FixedSizeListArray | pa.LargeListArray | 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")
+        array = normalize_list_array(array)
         validate_list_struct_type(array.type)
         self._data = array
 

src/nested_pandas/series/_storage/struct_list_storage.py

@@ -7,6 +7,7 @@
 
 from nested_pandas.series.utils import (
     align_chunked_struct_list_offsets,
+    normalize_struct_list_array,
     table_to_struct_array,
     transpose_list_struct_chunked,
 )
@@ -39,6 +40,7 @@ def __init__(self, array: pa.StructArray | pa.ChunkedArray, *, validate: bool =
             raise ValueError("array must be a StructArray or ChunkedArray")
 
         if validate:
+            array = normalize_struct_list_array(array)
             array = align_chunked_struct_list_offsets(array)
 
         self._data = array

src/nested_pandas/series/dtype.py

@@ -186,7 +186,7 @@ def __from_arrow__(self, array: pa.Array | pa.ChunkedArray) -> ExtensionArray:
 
     pyarrow_dtype: pa.StructType
 
-    def __init__(self, pyarrow_dtype: pa.DataType) -> None:
+    def __init__(self, pyarrow_dtype: pa.DataType | Mapping) -> None:
         # Allow pd.ArrowDtypes on init
         if isinstance(pyarrow_dtype, pd.ArrowDtype):
             pyarrow_dtype = pyarrow_dtype.pyarrow_dtype

src/nested_pandas/series/ext_array.py

@@ -65,6 +65,8 @@
 from nested_pandas.series.utils import (
     chunk_lengths,
     is_pa_type_a_list,
+    normalize_list_array,
+    normalize_struct_list_type,
     rechunk,
     struct_field_names,
     transpose_struct_list_type,
@@ -253,7 +255,7 @@ def __init__(self, values: pa.Array | pa.ChunkedArray, *, validate: bool = True)
             list_struct_storage = ListStructStorage.from_struct_list_storage(struct_list_storage)
 
         self._storage = list_struct_storage
-        self._dtype = NestedDtype(values.type)
+        self._dtype = NestedDtype(self._storage.type)
 
     # End of Constructor and initialized attributes #
 
@@ -669,6 +671,53 @@ def __getstate__(self):
 
     # End of Additional magic methods #
 
+    @classmethod
+    def is_input_pa_type_supported(cls, pa_type: pa.DataType) -> bool:
+        """Check whether a pyarrow data type is supported by the constructor.
+
+        Calling this method is cheaper than trying to construct the array,
+        because data transformations are avoided.
+
+        Parameters
+        ----------
+        pa_type : pyarrow.DataType
+            The pyarrow data type to check for compatibility with
+            NestedExtensionArray.
+
+        Returns
+        -------
+        bool
+            ``True`` if ``pa_type`` is a supported input type for
+            NestedExtensionArray, ``False`` otherwise.
+
+        Examples
+        --------
+        Check support for a list-of-structs type::
+
+            >>> import pyarrow as pa
+            >>> value_type = pa.struct([("a", pa.int64())])
+            >>> pa_type = pa.list_(value_type)
+            >>> NestedExtensionArray.is_input_pa_type_supported(pa_type)
+            True
+
+        Check support for an unsupported primitive type::
+
+            >>> NestedExtensionArray.is_input_pa_type_supported(pa.int64())
+            False
+        """
+        if is_pa_type_a_list(pa_type):
+            list_type = cast(
+                pa.ListType | pa.LargeListType | pa.FixedSizeListType,
+                pa_type,
+            )
+            return pa.types.is_struct(list_type.value_type)
+
+        try:
+            _ = normalize_struct_list_type(pa_type)
+        except ValueError:
+            return False
+        return True
+
     @classmethod
     def _box_pa_scalar(cls, value, *, pa_type: pa.DataType | None) -> pa.Scalar:
         """Convert a value to a PyArrow scalar with the specified type."""
@@ -1098,8 +1147,7 @@ def set_list_field(self, field: str, value: ArrayLike, *, keep_dtype: bool = Fal
                 "or NestedExtensionArray.set_list_field(..., keep_dtype=False) instead."
             ) from e
 
-        if not is_pa_type_a_list(pa_array.type):
-            raise ValueError(f"Expected a list array, got {pa_array.type}")
+        pa_array = normalize_list_array(pa_array)
 
         if len(pa_array) != len(self):
             raise ValueError("The length of the list-array must be equal to the length of the series")