Class TensorArray

python/ray/data/_internal/tensor_extensions/pandas.py:647–1465

Pandas `ExtensionArray` representing a tensor column, i.e. a column consisting of ndarrays as elements. This extension supports tensors in which the elements have different shapes. However, each tensor element must be non-ragged, i.e. each tensor element must have a well-define

@PublicAPI(stability="beta")
class TensorArray(
    pd.api.extensions.ExtensionArray,
    _TensorOpsMixin,
    _TensorScalarCastMixin,
):
    """
    Pandas `ExtensionArray` representing a tensor column, i.e. a column
    consisting of ndarrays as elements.

    This extension supports tensors in which the elements have different shapes.
    However, each tensor element must be non-ragged, i.e. each tensor element must have
    a well-defined, non-ragged shape.

    Examples:
        >>> # Create a DataFrame with a list of ndarrays as a column.
        >>> import pandas as pd
        >>> import numpy as np
        >>> import ray
        >>> from ray.data.extensions import TensorArray
        >>> df = pd.DataFrame({
        ...     "one": [1, 2, 3],
        ...     "two": TensorArray(np.arange(24).reshape((3, 2, 2, 2)))})
        >>> # Note that the column dtype is TensorDtype.
        >>> df.dtypes # doctest: +SKIP
        one          int64
        two    TensorDtype(shape=(3, 2, 2, 2), dtype=int64)
        dtype: object
        >>> # Pandas is aware of this tensor column, and we can do the
        >>> # typical DataFrame operations on this column.
        >>> col = 2 * (df["two"] + 10)
        >>> # The ndarrays underlying the tensor column will be manipulated,
        >>> # but the column itself will continue to be a Pandas type.
        >>> type(col) # doctest: +SKIP
        pandas.core.series.Series
        >>> col # doctest: +SKIP
        0   [[[ 2  4]
              [ 6  8]]
             [[10 12]
               [14 16]]]
        1   [[[18 20]
              [22 24]]
             [[26 28]
              [30 32]]]
        2   [[[34 36]
              [38 40]]
             [[42 44]
              [46 48]]]
        Name: two, dtype: TensorDtype(shape=(3, 2, 2, 2), dtype=int64)
        >>> # Once you do an aggregation on that column that returns a single
        >>> # row's value, you get back our TensorArrayElement type.
        >>> tensor = col.mean() # doctest: +SKIP
        >>> type(tensor) # doctest: +SKIP
        ray.data.extensions.tensor_extension.TensorArrayElement
        >>> tensor # doctest: +SKIP
        array([[[18., 20.],
                [22., 24.]],
               [[26., 28.],
                [30., 32.]]])