Repository: ray-project/ray
File: python/ray/data/_internal/planner/plan_expression/expression_evaluator.py (Python)

Code before:
        #   -- indices: [0, 1]
        # Output: regular string array ["a", "b"]
    Args:
        x: The input array to convert.
    Returns:
        The converted string array.
    """
    if pa.types.is_dictionary(x.type) and _is_pa_string_type(x.type.value_type):
        return pc.cast(x, pa.string())
    return x


def _to_pa_string_input(x: Any) -> Any:
    if isinstance(x, str):
        return pa.scalar(x)
    if isinstance(x, (pa.Array, pa.ChunkedArray)) and _is_pa_string_like(x):
        return _pa_decode_dict_string_array(x)
    if isinstance(x, (pa.Array, pa.ChunkedArray)):
        actual_type = str(x.type)
    else:
        actual_type = type(x).__name__
    raise TypeError(
        "Expected string or string-like pyarrow Array/ChunkedArray for string "
        f"concatenation, got {actual_type}."
    )


def _pa_add_or_concat(left: Any, right: Any) -> Any:
    if isinstance(left, pa.Scalar):
        left = left.as_py()
    if isinstance(right, pa.Scalar):
        right = right.as_py()
    # If either side is string-like, perform string concatenation.
    if (
        isinstance(left, str)
        or isinstance(right, str)
        or (isinstance(left, (pa.Array, pa.ChunkedArray)) and _is_pa_string_like(left))
        or (
            isinstance(right, (pa.Array, pa.ChunkedArray)) and _is_pa_string_like(right)
        )
    ):
        left_input = _to_pa_string_input(left)
        right_input = _to_pa_string_input(right)
        return pc.binary_join_element_wise(left_input, right_input, "")
    return pc.add(left, right)


_ARROW_EXPR_OPS_MAP: Dict[Operation, Callable[..., Any]] = {
    Operation.ADD: _pa_add_or_concat,
    Operation.SUB: pc.subtract,
    Operation.MUL: pc.multiply,

Review comment (suggestion):
Suggested cleanup:
```
    if isinstance(x, (pa.Array, pa.ChunkedArray)) and _is_pa_string_like(x):
        return _pa_decode_dict_string_array(x)
    actual_type = str(x.type) if isinstance(x, (pa.Array, pa.ChunkedArray)) else type(x).__name__
    raise TypeError(
        "Expected string or string-like pyarrow Array/ChunkedArray for string "
        f"concatenation, got {actual_type}."
    )
```

Code after:
        # Input: pa.array(['a', 'b']).dictionary_encode()
        #   -- dictionary: ["a", "b"]
        #   -- indices: [0, 1]
        # Output: regular string array ["a", "b"]
    Args:
        x: The input array to convert.
    Returns:
        The converted string array.
    """
    if pa.types.is_dictionary(x.type) and _is_pa_string_type(x.type.value_type):
        return pc.cast(x, pa.string())
    return x


def _to_pa_string_input(x: Any) -> Any:
    if isinstance(x, str):
        return pa.scalar(x)
    if isinstance(x, (pa.Array, pa.ChunkedArray)) and _is_pa_string_like(x):
        return _pa_decode_dict_string_array(x)
    actual_type = (
        str(x.type) if isinstance(x, (pa.Array, pa.ChunkedArray)) else type(x).__name__
    )
    raise TypeError(
        "Expected string or string-like pyarrow Array/ChunkedArray for string "
        f"concatenation, got {actual_type}."
    )


def _pa_add_or_concat(left: Any, right: Any) -> Any:
    if isinstance(left, pa.Scalar):
        left = left.as_py()
    if isinstance(right, pa.Scalar):
        right = right.as_py()
    # If either side is string-like, perform string concatenation.
    if (
        isinstance(left, str)
        or isinstance(right, str)
        or (isinstance(left, (pa.Array, pa.ChunkedArray)) and _is_pa_string_like(left))
        or (
            isinstance(right, (pa.Array, pa.ChunkedArray)) and _is_pa_string_like(right)
        )
    ):
        left_input = _to_pa_string_input(left)
        right_input = _to_pa_string_input(right)
        return pc.binary_join_element_wise(left_input, right_input, "")
    return pc.add(left, right)


_ARROW_EXPR_OPS_MAP: Dict[Operation, Callable[..., Any]] = {
    Operation.ADD: _pa_add_or_concat,
    Operation.SUB: pc.subtract,