Merge pull request #260 from lincc-frameworks/nf-autocast-list-struct

Autocast list-struct columns to nested

Author: hombit

src/nested_pandas/nestedframe/core.py

@@ -18,6 +18,7 @@
 )
 from nested_pandas.series.dtype import NestedDtype
 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)
@@ -41,6 +42,32 @@ class NestedFrame(pd.DataFrame):
     # contains those aliases, keyed by the cleaned name.
     _metadata = ["_aliases"]
 
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self._cast_cols_to_nested(struct_list=False)
+
+    def _cast_cols_to_nested(self, *, struct_list: bool) -> None:
+        """Cast arrow columns to nested.
+
+        Parameters
+        ----------
+        struct_list : bool
+            If `False` cast list-struct columns only. If `True`, also
+            try to cast struct-list columns validating if they have
+            valid nested structure.
+        """
+        for column, dtype in self.dtypes.items():
+            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)):
+                continue
+            try:
+                nested_dtype = NestedDtype(pa_type)
+            except (TypeError, ValueError):
+                continue
+            self[column] = self[column].astype(nested_dtype)
+
     @property
     def _constructor(self) -> Self:  # type: ignore[name-defined] # noqa: F821
         return NestedFrame
@@ -224,7 +251,8 @@ def __setitem__(self, key, value):
             self._update_inplace(new_df)
             return None
 
-        return super().__setitem__(key, value)
+        super().__setitem__(key, value)
+        self._cast_cols_to_nested(struct_list=False)
 
     def add_nested(
         self,

src/nested_pandas/series/ext_array.py

@@ -240,11 +240,11 @@ def _from_sequence(cls, scalars, *, dtype=None, copy: bool = False) -> Self:  #
         """
         del copy
 
-        if isinstance(dtype, NestedDtype) and isinstance(scalars, (cls, pa.Array, pa.ChunkedArray)):
-            if is_pa_type_a_list(scalars.type):
-                return cls(scalars.cast(dtype.list_struct_pa_dtype))
-            elif pa.types.is_struct(scalars.type):
-                return cls(scalars.cast(dtype.pyarrow_dtype))
+        if isinstance(dtype, NestedDtype):
+            try:
+                return cls._from_arrow_like(scalars, dtype=dtype)
+            except ValueError:
+                pass
 
         pa_type = to_pyarrow_dtype(dtype)
         pa_array = cls._box_pa_array(scalars, pa_type=pa_type)
@@ -663,6 +663,29 @@ def _box_pa_array(cls, value, *, pa_type: pa.DataType | None) -> pa.Array | pa.C
 
         return pa_array
 
+    @classmethod
+    def _from_arrow_like(cls, arraylike, dtype: NestedDtype | None = None) -> Self:  # type: ignore[name-defined] # noqa: F821
+        if isinstance(arraylike, cls):
+            if dtype is None or dtype == arraylike.dtype:
+                return arraylike
+            array = arraylike.list_array
+        elif isinstance(arraylike, (pa.Array, pa.ChunkedArray)):
+            array = arraylike
+        else:
+            array = pa.array(arraylike)
+
+        if dtype is None:
+            return cls(array)
+
+        try:
+            cast_array = array.cast(dtype.pyarrow_dtype)
+        except (ValueError, TypeError, KeyError, pa.ArrowNotImplementedError):
+            try:
+                cast_array = array.cast(dtype.list_struct_pa_dtype)
+            except (ValueError, TypeError, KeyError, pa.ArrowNotImplementedError):
+                raise ValueError(f"Cannot cast input to {dtype}") from None
+        return cls(cast_array)
+
     @classmethod
     def _convert_struct_scalar_to_df(cls, value: pa.StructScalar, *, copy: bool, na_value: Any = None) -> Any:
         """Converts a struct scalar of equal-length list scalars to a pd.DataFrame

tests/nested_pandas/nestedframe/test_nestedframe.py

@@ -2,7 +2,7 @@
 import pandas as pd
 import pyarrow as pa
 import pytest
-from nested_pandas import NestedFrame
+from nested_pandas import NestedDtype, NestedFrame
 from nested_pandas.datasets import generate_data
 from nested_pandas.nestedframe.core import _SeriesFromNest
 from nested_pandas.series.packer import pack_lists
@@ -14,6 +14,17 @@ def test_nestedframe_construction():
     base = NestedFrame(data={"a": [1, 2, 3], "b": [2, 4, 6]}, index=[0, 1, 2])
 
     assert isinstance(base, NestedFrame)
+    assert_frame_equal(base, pd.DataFrame({"a": [1, 2, 3], "b": [2, 4, 6]}, index=[0, 1, 2]))
+
+    list_struct_array = pa.array(
+        [[{"x": 1, "y": 1.0}], [{"x": 2, "y": 2.0}], [{"x": 3, "y": 3.0}, {"x": 4, "y": 4.0}]]
+    )
+    list_struct_series = pd.Series(list_struct_array, dtype=pd.ArrowDtype(list_struct_array.type))
+    nested_series = pd.Series(list_struct_series, dtype=NestedDtype(list_struct_array.type))
+
+    nf = NestedFrame(base.to_dict(orient="series") | {"list_struct": list_struct_series})
+    # Test auto-cast to nested
+    assert_frame_equal(nf, base.assign(list_struct=nested_series))
 
 
 def test_nestedseries_construction():
@@ -223,6 +234,22 @@ def test_set_new_nested_col():
     )
 
 
+def test_set_list_struct_col():
+    """Test that __setitem__ would cast list-struct columns to nested."""
+    nf = generate_data(10, 3)
+    nf["a"] = nf["a"].astype(pd.ArrowDtype(pa.float64()))
+    nf["b"] = nf["b"].astype(pd.ArrowDtype(pa.float64()))
+
+    list_struct_array = pa.array(nf.nested)
+    list_struct_series = pd.Series(list_struct_array, dtype=pd.ArrowDtype(list_struct_array.type))
+
+    nf["nested2"] = list_struct_series
+    assert_frame_equal(nf.nested.nest.to_flat(), nf.nested2.nest.to_flat())
+
+    nf = nf.assign(nested3=list_struct_series)
+    assert_frame_equal(nf.nested.nest.to_flat(), nf.nested3.nest.to_flat())
+
+
 def test_get_dot_names():
     """Test the ability to still work with column names with '.' characters outside of nesting"""
     nf = NestedFrame.from_flat(