Skip to content

Commit

Permalink
Add NestedExtensionArray.from_sequence
Browse files Browse the repository at this point in the history
  • Loading branch information
hombit committed May 3, 2024
1 parent d44f213 commit de45b64
Show file tree
Hide file tree
Showing 2 changed files with 86 additions and 45 deletions.
24 changes: 23 additions & 1 deletion src/nested_pandas/series/ext_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,12 +54,14 @@
__all__ = ["NestedExtensionArray"]


def to_pyarrow_dtype(dtype: NestedDtype | pd.ArrowDtype | None) -> pa.DataType | None:
def to_pyarrow_dtype(dtype: NestedDtype | pd.ArrowDtype | pa.DataType | None) -> pa.DataType | None:
"""Convert the dtype to pyarrow.DataType"""
if isinstance(dtype, NestedDtype):
return dtype.pyarrow_dtype
if isinstance(dtype, pd.ArrowDtype):
return dtype.pyarrow_dtype
if isinstance(dtype, pa.DataType):
return dtype
return None


Expand Down Expand Up @@ -532,6 +534,26 @@ def __init__(self, values: pa.Array | pa.ChunkedArray, *, validate: bool = True)
self._pa_array = values
self._dtype = NestedDtype(values.type)

@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
Parameters
----------
scalars : Sequence
The sequence of items: dictionaries (key is column name, value is array-like of nested elements),
DataFrames, None, pd.NA, pa.Array or anything convertible to PyArrow scalars of struct type with
list fields of the same lengths.
dtype : dtype or None
NestedDtype of the resulting array, or a type to infer from: pd.ArrowDtype or pa.DataType.
Returns
-------
NestedExtensionArray
The constructed extension array.
"""
return cls._from_sequence(scalars, dtype=dtype)

@property
def _pyarrow_dtype(self) -> pa.DataType:
"""PyArrow data type of the extension array"""
Expand Down
Loading

0 comments on commit de45b64

Please sign in to comment.