Series#

@staticmethod
def from_arrow(array: pa.Array | pa.ChunkedArray, name: str = "arrow_series") -> Series:
    """Construct a Series from an pyarrow array or chunked array.

    Args:
        array: The pyarrow (chunked) array whose data we wish to put in the Series.
        name: The name associated with the Series; this is usually the column name.
    """
    _ensure_registered_super_ext_type()
    if DataType.from_arrow_type(array.type) == DataType.python():
        # If the Arrow type is not natively supported, go through the Python list path.
        return Series.from_pylist(array.to_pylist(), name=name, pyobj="force")
    elif isinstance(array, pa.Array):
        array = ensure_array(array)
        if isinstance(array.type, getattr(pa, "FixedShapeTensorType", ())):
            series = Series.from_arrow(array.storage, name=name)
            return series.cast(DataType.from_arrow_type(array.type))
        else:
            pys = PySeries.from_arrow(name, array)
            return Series._from_pyseries(pys)
    elif isinstance(array, pa.ChunkedArray):
        array = ensure_chunked_array(array)
        arr_type = array.type
        if isinstance(arr_type, pa.BaseExtensionType):
            combined_storage_array = array.cast(arr_type.storage_type).combine_chunks()
            combined_array = arr_type.wrap_array(combined_storage_array)
        else:
            combined_array = array.combine_chunks()
        return Series.from_arrow(combined_array)
    else:
        raise TypeError(f"expected either PyArrow Array or Chunked Array, got {type(array)}")

@classmethod
def from_numpy(cls, data: np.ndarray[Any, Any], name: str = "numpy_series") -> Series:
    """Construct a Series from a NumPy ndarray.

    If the provided NumPy ndarray is 1-dimensional, Daft will attempt to store the ndarray
    in a pyarrow Array. If the ndarray has more than 1 dimension OR storing the 1D array in Arrow failed,
    Daft will store the ndarray data as a Python list of NumPy ndarrays.

    Args:
        data: The NumPy ndarray whose data we wish to put in the Series.
        name: The name associated with the Series; this is usually the column name.
    """
    if not isinstance(data, np.ndarray):
        raise TypeError(f"Expected a NumPy ndarray, got {type(data)}")
    if data.ndim <= 1:
        try:
            arrow_array = pa.array(data)
        except pa.ArrowInvalid:
            pass
        else:
            return cls.from_arrow(arrow_array, name=name)
    # TODO(Clark): Represent the tensor series with an Arrow extension type in order
    # to keep the series data contiguous.
    list_ndarray = [np.asarray(item) for item in data]
    return cls.from_pylist(list_ndarray, name=name, pyobj="allow")

@classmethod
def from_pandas(cls, data: pd.Series[Any], name: str = "pd_series") -> Series:
    """Construct a Series from a pandas Series.

    This will first try to convert the series into a pyarrow array, then will fall
    back to converting the series to a NumPy ndarray and going through that construction path,
    and will finally fall back to converting the series to a Python list and going through that
    path.

    Args:
        data: The pandas Series whose data we wish to put in the Daft Series.
        name: The name associated with the Series; this is usually the column name.
    """
    if not isinstance(data, pd.Series):
        raise TypeError(f"expected a pandas Series, got {type(data)}")
    # First, try Arrow path.
    try:
        arrow_arr = pa.Array.from_pandas(data)
    except pa.ArrowInvalid:
        pass
    else:
        return cls.from_arrow(arrow_arr, name=name)
    # Second, fall back to NumPy path. Note that .from_numpy() does _not_ fall back to
    # the pylist representation for 1D arrays and instead raises an error that we can catch.
    # We do the pylist representation fallback ourselves since the pd.Series.to_list()
    # preserves more type information for types that are not natively representable in Python.
    try:
        ndarray = data.to_numpy()
        return cls.from_numpy(ndarray, name=name)
    except Exception:
        pass
    # Finally, fall back to pylist path.
    # NOTE: For element types that don't have a native Python representation,
    # a Pandas scalar object will be returned.
    return cls.from_pylist(data.to_list(), name=name, pyobj="force")

@staticmethod
def from_pylist(
    data: list[Any],
    name: str = "list_series",
    pyobj: Literal["allow", "disallow", "force"] = "allow",
) -> Series:
    """Construct a Series from a Python list.

    The resulting type depends on the setting of pyobjects:
        - ``"allow"``: Arrow-backed types if possible, else PyObject;
        - ``"disallow"``: Arrow-backed types only, raising error if not convertible;
        - ``"force"``: Store as PyObject types.

    Args:
        data: The Python list whose data we wish to put in the Series.
        name: The name associated with the Series; this is usually the column name.
        pyobj: Whether we want to ``"allow"`` coercion to Arrow types, ``"disallow"``
            falling back to Python type representation, or ``"force"`` the data to only
            have a Python type representation. Default is ``"allow"``.
    """
    if not isinstance(data, list):
        raise TypeError(f"expected a python list, got {type(data)}")

    if pyobj not in {"allow", "disallow", "force"}:
        raise ValueError(f"pyobj: expected either 'allow', 'disallow', or 'force', but got {pyobj})")

    if pyobj == "force":
        pys = PySeries.from_pylist(name, data, pyobj=pyobj)
        return Series._from_pyseries(pys)

    try:
        # Workaround: wrap list of np.datetime64 in an np.array
        #   - https://github.com/apache/arrow/issues/40580
        #   - https://github.com/Eventual-Inc/Daft/issues/3826
        if data and np.module_available() and isinstance(data[0], np.datetime64):  # type: ignore[attr-defined]
            np_arr = np.array(data)
            arrow_array = pa.array(np_arr)
        else:
            arrow_array = pa.array(data)
        return Series.from_arrow(arrow_array, name=name)
    except pa.lib.ArrowInvalid:
        if pyobj == "disallow":
            raise
        pys = PySeries.from_pylist(name, data, pyobj=pyobj)
        return Series._from_pyseries(pys)

def to_arrow(self) -> pa.Array:
    """Convert this Series to an pyarrow array."""
    _ensure_registered_super_ext_type()

    dtype = self.datatype()
    arrow_arr = self._series.to_arrow()

    # Special-case for PyArrow FixedShapeTensor if it is supported by the version of PyArrow
    # TODO: Push this down into self._series.to_arrow()?
    if dtype.is_fixed_shape_tensor() and pyarrow_supports_fixed_shape_tensor():
        pyarrow_dtype = dtype.to_arrow_dtype()
        arrow_series = self._series.to_arrow()
        return pa.ExtensionArray.from_storage(pyarrow_dtype, arrow_series.storage)

    return arrow_arr

def to_pylist(self) -> list[Any]:
    """Convert this Series to a Python list."""
    if self.datatype().is_python():
        return self._series.to_pylist()
    elif self.datatype()._should_cast_to_python():
        return self._series.cast(DataType.python()._dtype).to_pylist()
    else:
        return self._series.to_arrow().to_pylist()