⚡️ Speed up function maybe_coerce_to_str by 6%

[codeflash-ai]

📄 6% (0.06x) speedup for maybe_coerce_to_str in xarray/core/utils.py

⏱️ Runtime : 1.62 milliseconds → 1.53 milliseconds (best of 45 runs)

📝 Explanation and details

The optimization improves the result_type function in xarray/core/dtypes.py by replacing nested any() calls with a single-pass algorithm that uses early termination.

Key optimization: The original code executed any(issubclass(t, left) for t in types) and any(issubclass(t, right) for t in types) for each promotion rule, resulting in potentially 2×N×M issubclass calls (where N is the number of types and M is the number of promotion rules). The optimized version uses a single loop per promotion rule with boolean flags and early return, reducing this to at most N×M calls with frequent early termination.

Why this speeds up the code:

• Eliminates redundant iterations: Instead of scanning the entire type set twice per promotion rule, it scans once and tracks both conditions simultaneously
• Early termination: Returns immediately when both left and right types are found, avoiding unnecessary issubclass checks
• Reduces function call overhead: Eliminates the any() generator expressions and their associated overhead

Impact on workloads: The function reference shows maybe_coerce_to_str is called in IndexVariable.concat, which processes pandas Index objects during concatenation operations. This is likely in data processing hot paths where arrays are frequently combined. The 5% overall speedup becomes valuable when processing large datasets or performing many concatenation operations.

Test performance patterns: The optimization shows consistent 8-27% improvements across most test cases, with particularly strong gains (19-27%) on cases involving type promotion decisions (mixed bytes/unicode, bool/string, number/string), where the early termination logic provides maximum benefit.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	✅ 35 Passed
🌀 Generated Regression Tests	✅ 44 Passed
⏪ Replay Tests	✅ 8 Passed
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

⚙️ Existing Unit Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_utils.py::test_maybe_coerce_to_str	20.4μs	18.5μs	10.3%✅
test_utils.py::test_maybe_coerce_to_str_minimal_str_dtype	12.1μs	11.0μs	9.76%✅

🌀 Generated Regression Tests and Runtime

import numpy as np

# imports
import pytest
from xarray.core.utils import maybe_coerce_to_str

# ------------------ UNIT TESTS ------------------

# Basic Test Cases


def test_basic_str_dtype():
    # All coords are str, index is object dtype, should coerce to str dtype
    index = np.array(["a", "b", "c"], dtype=object)
    coords = [np.array(["a", "b", "c"], dtype="U")]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 12.3μs -> 11.0μs (12.0% faster)


def test_basic_bytes_dtype():
    # All coords are bytes, index is object dtype, should coerce to bytes dtype
    index = np.array([b"a", b"b", b"c"], dtype=object)
    coords = [np.array([b"a", b"b", b"c"], dtype="S")]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 12.8μs -> 11.1μs (15.2% faster)


def test_basic_object_dtype_no_str():
    # Index is object dtype, coords are int, should not coerce
    index = np.array([1, 2, 3], dtype=object)
    coords = [np.array([1, 2, 3], dtype=int)]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 8.32μs -> 6.54μs (27.3% faster)


def test_basic_mixed_str_and_object():
    # Index is object dtype, coords are mixed str and object, should coerce to str
    index = np.array(["a", "b", "c"], dtype=object)
    coords = [
        np.array(["a", "b", "c"], dtype="U"),
        np.array(["d", "e", "f"], dtype=object),
    ]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 8.14μs -> 7.52μs (8.19% faster)


def test_basic_already_str_dtype():
    # Index is already str dtype, should remain unchanged
    index = np.array(["a", "b", "c"], dtype="U")
    coords = [np.array(["a", "b", "c"], dtype="U")]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 8.41μs -> 7.41μs (13.5% faster)


# Edge Test Cases


def test_coords_typeerror():
    # coords that cannot be used in np.result_type, triggers TypeError
    index = np.array(["a", "b", "c"], dtype=object)
    coords = [None]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 9.32μs -> 7.53μs (23.8% faster)


def test_mixed_bytes_and_unicode_coords():
    # Should promote to object dtype (bytes + unicode)
    index = np.array(["a", "b", "c"], dtype=object)
    coords = [
        np.array(["a", "b", "c"], dtype="U"),
        np.array([b"a", b"b", b"c"], dtype="S"),
    ]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 9.08μs -> 7.61μs (19.3% faster)


def test_mixed_bool_and_str_coords():
    # Should promote to object dtype (bool + str)
    index = np.array(["True", "False"], dtype=object)
    coords = [
        np.array([True, False], dtype=bool),
        np.array(["True", "False"], dtype="U"),
    ]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 9.14μs -> 7.65μs (19.5% faster)


def test_mixed_number_and_str_coords():
    # Should promote to object dtype (number + str)
    index = np.array(["1", "2", "3"], dtype=object)
    coords = [np.array([1, 2, 3], dtype=int), np.array(["1", "2", "3"], dtype="U")]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 8.47μs -> 7.41μs (14.4% faster)


def test_object_dtype_with_non_str_like():
    # Index is object dtype, but contains ints, coords are str
    index = np.array([1, 2, 3], dtype=object)
    coords = [np.array(["a", "b", "c"], dtype="U")]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 12.7μs -> 11.6μs (9.63% faster)


def test_object_dtype_with_mixed_types():
    # Index contains mixed types, coords are str
    index = np.array(["a", 1, None], dtype=object)
    coords = [np.array(["a", "b", "c"], dtype="U")]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 13.2μs -> 12.1μs (8.95% faster)


def test_object_dtype_with_bytes_and_str_coords():
    # Index is object dtype, coords are bytes and str, should promote to object
    index = np.array(["a", b"b", "c"], dtype=object)
    coords = [
        np.array(["a", "b", "c"], dtype="U"),
        np.array([b"a", b"b", b"c"], dtype="S"),
    ]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 8.89μs -> 7.46μs (19.2% faster)


def test_object_dtype_with_bool_and_str_coords():
    # Index is object dtype, coords are bool and str, should promote to object
    index = np.array(["True", "False"], dtype=object)
    coords = [
        np.array([True, False], dtype=bool),
        np.array(["True", "False"], dtype="U"),
    ]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 8.48μs -> 7.44μs (13.9% faster)


# Large Scale Test Cases


def test_large_str_array():
    # Large array of strings, should coerce to str dtype
    large_size = 1000
    index = np.array([str(i) for i in range(large_size)], dtype=object)
    coords = [np.array([str(i) for i in range(large_size)], dtype="U")]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 77.2μs -> 75.8μs (1.88% faster)


def test_large_bytes_array():
    # Large array of bytes, should coerce to bytes dtype
    large_size = 1000
    index = np.array([bytes([i % 256]) for i in range(large_size)], dtype=object)
    coords = [np.array([bytes([i % 256]) for i in range(large_size)], dtype="S")]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 59.0μs -> 56.8μs (3.90% faster)


def test_large_mixed_str_and_number_coords():
    # Large array, coords are mixed str and number, should promote to object
    large_size = 1000
    index = np.array([str(i) for i in range(large_size)], dtype=object)
    coords = [
        np.array([i for i in range(large_size)], dtype=int),
        np.array([str(i) for i in range(large_size)], dtype="U"),
    ]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 8.76μs -> 7.16μs (22.4% faster)


def test_large_object_dtype_with_non_str_like():
    # Large object dtype index with ints, coords are str, should coerce to str
    large_size = 1000
    index = np.array([i for i in range(large_size)], dtype=object)
    coords = [np.array([str(i) for i in range(large_size)], dtype="U")]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 118μs -> 103μs (14.6% faster)


def test_large_object_dtype_with_mixed_types():
    # Large object dtype index with mixed types, coords are str
    large_size = 1000
    index = np.array(
        [str(i) if i % 2 == 0 else i for i in range(large_size)], dtype=object
    )
    coords = [np.array([str(i) for i in range(large_size)], dtype="U")]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 99.2μs -> 90.7μs (9.37% faster)
    expected = np.array([str(i) for i in range(large_size)], dtype="U")


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import sys
import types

# function to test
from typing import Any

import numpy as np

# imports
import pytest
from xarray.core.utils import maybe_coerce_to_str

# ------------------ UNIT TESTS ------------------

# 1. BASIC TEST CASES


def test_already_str_index_and_coords():
    # index and coords are both str, should remain unchanged but dtype may be promoted to <U
    index = ["a", "b", "c"]
    coords = ["a", "b", "c"]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 19.2μs -> 18.4μs (4.48% faster)


def test_object_str_index_and_coords():
    # index is object dtype but contains all str
    index = np.array(["x", "y", "z"], dtype=object)
    coords = ["x", "y", "z"]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 6.42μs -> 6.14μs (4.60% faster)


def test_int_index_and_coords():
    # index and coords are int, should remain unchanged
    index = np.array([1, 2, 3], dtype=int)
    coords = [1, 2, 3]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 11.8μs -> 9.77μs (20.4% faster)


def test_mixed_int_and_str_promotes_to_object():
    # index is int, coords are str, should promote to object
    index = np.array([1, 2, 3], dtype=int)
    coords = ["1", "2", "3"]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 24.2μs -> 23.7μs (2.01% faster)


def test_bytes_index_and_coords():
    # index and coords are bytes, should return numpy array of bytes
    index = np.array([b"a", b"b", b"c"], dtype="S1")
    coords = [b"a", b"b", b"c"]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 17.9μs -> 16.7μs (7.67% faster)


# 2. EDGE TEST CASES


def test_single_element_str():
    # Single element str
    index = np.array(["alpha"], dtype=object)
    coords = ["alpha"]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 6.66μs -> 6.59μs (1.03% faster)


def test_single_element_int():
    # Single element int
    index = np.array([42], dtype=int)
    coords = [42]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 10.2μs -> 8.53μs (19.3% faster)


def test_object_mixed_str_and_int():
    # index is object dtype with mixed str and int, coords are str
    index = np.array(["a", 2, "c"], dtype=object)
    coords = ["a", "2", "c"]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 32.4μs -> 32.3μs (0.121% faster)


def test_coords_with_non_iterable():
    # coords is not iterable (should raise TypeError in result_type)
    index = np.array([1, 2, 3])
    coords = None
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 3.53μs -> 3.44μs (2.88% faster)


def test_coords_with_bytes_and_str():
    # bytes and unicode, should promote to object
    index = np.array([b"a", b"b"], dtype="S1")
    coords = ["a", "b"]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 17.5μs -> 15.3μs (13.9% faster)


def test_bool_and_str_promotes_to_object():
    # bool and str, should promote to object
    index = np.array([True, False], dtype=bool)
    coords = ["True", "False"]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 6.47μs -> 6.17μs (4.90% faster)


def test_index_is_python_list():
    # index is a python list, coords are str
    index = ["foo", "bar"]
    coords = ["foo", "bar"]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 6.97μs -> 7.06μs (1.26% slower)


def test_coords_are_numpy_array():
    # coords is a numpy array of str
    index = ["x", "y"]
    coords = np.array(["x", "y"], dtype="U1")
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 20.0μs -> 18.1μs (10.5% faster)


def test_index_and_coords_with_unicode():
    # index and coords with unicode
    index = np.array(["α", "β", "γ"], dtype=object)
    coords = ["α", "β", "γ"]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 7.28μs -> 7.38μs (1.48% slower)


def test_index_object_with_bytes():
    # index is object dtype but contains bytes
    index = np.array([b"a", b"b"], dtype=object)
    coords = [b"a", b"b"]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 17.5μs -> 16.0μs (9.39% faster)


def test_index_object_with_mixed_bytes_and_str():
    # index is object dtype with mixed bytes and str, coords are str
    index = np.array([b"a", "b"], dtype=object)
    coords = ["a", "b"]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 16.7μs -> 15.4μs (8.48% faster)


def test_index_object_with_bool_and_str():
    # index is object dtype with bool and str, coords are str
    index = np.array([True, "False"], dtype=object)
    coords = ["True", "False"]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 6.40μs -> 6.28μs (1.89% faster)


def test_index_object_with_only_bool():
    # index is object dtype with only bool, coords are bool
    index = np.array([True, False], dtype=object)
    coords = [True, False]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 12.1μs -> 10.3μs (18.1% faster)


# 3. LARGE SCALE TEST CASES


def test_large_str_array():
    # Large array of str, should coerce to numpy array of str dtype (not object)
    N = 1000
    index = [str(i) for i in range(N)]
    coords = [str(i) for i in range(N)]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 29.1μs -> 28.9μs (0.547% faster)


def test_large_object_str_array():
    # Large object array of str, should coerce to numpy array of str dtype
    N = 1000
    index = np.array([str(i) for i in range(N)], dtype=object)
    coords = [str(i) for i in range(N)]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 28.7μs -> 27.9μs (3.19% faster)


def test_large_int_array():
    # Large int array, should remain int
    N = 1000
    index = np.arange(N)
    coords = np.arange(N)
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 648μs -> 634μs (2.13% faster)


def test_large_mixed_int_and_str():
    # Large mixed int and str, should promote to object
    N = 1000
    index = np.arange(N)
    coords = [str(i) for i in range(N)]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 28.5μs -> 28.4μs (0.327% faster)


def test_large_bytes_array():
    # Large bytes array, should remain bytes dtype
    N = 1000
    index = np.array([b"x"] * N, dtype="S1")
    coords = [b"x"] * N
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 9.93μs -> 10.1μs (1.69% slower)


def test_large_object_mixed_bytes_and_str():
    # Large mixed bytes and str, should promote to object
    N = 1000
    index = np.array([b"x" if i % 2 == 0 else "y" for i in range(N)], dtype=object)
    coords = ["x" if i % 2 == 0 else "y" for i in range(N)]
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 10.8μs -> 10.7μs (0.851% faster)


def test_large_object_with_unicode():
    # Large unicode object array, should coerce to numpy unicode dtype
    N = 1000
    index = np.array(["α"] * N, dtype=object)
    coords = ["α"] * N
    codeflash_output = maybe_coerce_to_str(index, coords)
    result = codeflash_output  # 11.0μs -> 12.2μs (9.62% slower)


# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

⏪ Replay Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
test_pytest_xarrayteststest_concat_py_xarrayteststest_computation_py_xarrayteststest_formatting_py_xarray__replay_test_0.py::test_xarray_core_utils_maybe_coerce_to_str	72.9μs	65.8μs	10.8%✅

To edit these changes git checkout codeflash/optimize-maybe_coerce_to_str-mj9tj1t3 and push.