Function cast_array_to_feature

src/datasets/table.py:2011–2149

Cast an array to the arrow type that corresponds to the requested feature type. For custom features like [`Audio`] or [`Image`], it takes into account the "cast_storage" methods they defined to enable casting from other arrow types. Args: array (`pa.Array`): The PyAr

(
    array: pa.Array, feature: "FeatureType", allow_primitive_to_str: bool = True, allow_decimal_to_str: bool = True
)

@_wrap_for_chunked_arrays
def cast_array_to_feature(
    array: pa.Array, feature: "FeatureType", allow_primitive_to_str: bool = True, allow_decimal_to_str: bool = True
) -> pa.Array:
    """Cast an array to the arrow type that corresponds to the requested feature type.
    For custom features like [`Audio`] or [`Image`], it takes into account the "cast_storage" methods
    they defined to enable casting from other arrow types.

    Args:
        array (`pa.Array`):
            The PyArrow array to cast.
        feature (`datasets.features.FeatureType`):
            The target feature type.
        allow_primitive_to_str (`bool`, defaults to `True`):
            Whether to allow casting primitives to strings.
            Defaults to `True`.
        allow_decimal_to_str (`bool`, defaults to `True`):
            Whether to allow casting decimals to strings.
            Defaults to `True`.

    Raises:
        `pa.ArrowInvalidError`: if the arrow data casting fails
        `TypeError`: if the target type is not supported according, e.g.

            - if a field is missing
            - if casting from primitives and `allow_primitive_to_str` is `False`
            - if casting from decimals and `allow_decimal_to_str` is `False`

    Returns:
        array (`pyarrow.Array`): the casted array
    """
    from .features.features import LargeList, List, get_nested_type

    _c = partial(
        cast_array_to_feature,
        allow_primitive_to_str=allow_primitive_to_str,
        allow_decimal_to_str=allow_decimal_to_str,
    )

    if isinstance(array, pa.ExtensionArray):
        array = array.storage
    if hasattr(feature, "cast_storage"):
        return feature.cast_storage(array)

    if pa.types.is_struct(array.type):
        # feature must be a dict
        if isinstance(feature, dict) and (array_fields := {field.name for field in array.type}) <= set(feature):
            null_array = pa.array([None] * len(array))
            arrays = [
                _c(array.field(name) if name in array_fields else null_array, subfeature)
                for name, subfeature in feature.items()
            ]
            return pa.StructArray.from_arrays(arrays, names=list(feature), mask=array.is_null())
    elif pa.types.is_list(array.type) or pa.types.is_large_list(array.type):
        # feature must be either List(subfeature) or LargeList(subfeature)
        if isinstance(feature, LargeList):
            casted_array_values = _c(array.values, feature.feature)
            if pa.types.is_large_list(array.type) and casted_array_values.type == array.values.type:
                # Both array and feature have equal large_list type and values (within the list) type