ray.data.extensions.tensor_extension.ArrowTensorArray#

class ray.data.extensions.tensor_extension.ArrowTensorArray[source]#

Bases: ray.air.util.tensor_extensions.arrow._ArrowTensorScalarIndexingMixin, pyarrow.lib.ExtensionArray

An array of fixed-shape, homogeneous-typed tensors.

This is the Arrow side of TensorArray.

See Arrow docs for customizing extension arrays: https://arrow.apache.org/docs/python/extending_types.html#custom-extension-array-class

PublicAPI (beta): This API is in beta and may change before becoming stable.

OFFSET_DTYPE#

alias of numpy.int32

classmethod from_numpy(arr: Union[numpy.ndarray, Iterable[numpy.ndarray]]) → Union[ray.air.util.tensor_extensions.arrow.ArrowTensorArray, ray.air.util.tensor_extensions.arrow.ArrowVariableShapedTensorArray][source]#

Convert an ndarray or an iterable of ndarrays to an array of homogeneous-typed tensors. If given fixed-shape tensor elements, this will return an ArrowTensorArray; if given variable-shape tensor elements, this will return an ArrowVariableShapedTensorArray.

Parameters

arr – An ndarray or an iterable of ndarrays.

Returns

  • If fixed-shape tensor elements, an ArrowTensorArray containing len(arr) tensors of fixed shape.

  • If variable-shaped tensor elements, an ArrowVariableShapedTensorArray containing len(arr) tensors of variable shape.

  • If scalar elements, a pyarrow.Array.

to_numpy(zero_copy_only: bool = True)[source]#

Convert the entire array of tensors into a single ndarray.

Parameters

zero_copy_only – If True, an exception will be raised if the conversion to a NumPy array would require copying the underlying data (e.g. in presence of nulls, or for non-primitive types). This argument is currently ignored, so zero-copy isn’t enforced even if this argument is true.

Returns

A single ndarray representing the entire array of tensors.

to_variable_shaped_tensor_array() → ray.air.util.tensor_extensions.arrow.ArrowVariableShapedTensorArray[source]#

Convert this tensor array to a variable-shaped tensor array.

This is primarily used when concatenating multiple chunked tensor arrays where at least one chunked array is already variable-shaped and/or the shapes of the chunked arrays differ, in which case the resulting concatenated tensor array will need to be in the variable-shaped representation.