⚡️ Speed up function get_otel_attribute by 28%

[codeflash-ai]

📄 28% (0.28x) speedup for get_otel_attribute in mlflow/tracing/utils/__init__.py

⏱️ Runtime : 729 microseconds → 568 microseconds (best of 68 runs)

📝 Explanation and details

The optimization adds a string type check with selective JSON parsing that significantly reduces the number of expensive json.loads() calls.

Key Changes:

1. Added string type checking: Before attempting JSON parsing, the code now checks if the attribute value is a string
2. Heuristic-based JSON detection: For strings, it only calls json.loads() if the string appears to be JSON-formatted (starts/ends with quotes, brackets, or braces)
3. Early return for plain strings: Non-JSON-looking strings are returned directly without parsing

Why This Is Faster:

• json.loads() is computationally expensive, requiring string parsing, tokenization, and object construction
• The optimization eliminates ~53% of json.loads() calls (32 plain strings avoided out of 60 total string attributes in the profiler)
• String character checks (attribute_value[0] == '"') are orders of magnitude faster than JSON parsing
• Line profiler shows the optimization reduces time spent in json.loads() from 922ms to 633ms (31% reduction)

Performance Patterns from Tests:

• Massive speedup for non-JSON strings: 400-1600% faster (e.g., "not json" goes from 16μs to 1μs)
• Slight slowdown for actual JSON strings: 7-11% slower due to extra type checking overhead
• Excellent speedup for primitive JSON values: 300-500% faster (numbers, booleans, null)

Impact on Hot Path Usage:
Based on function references, get_otel_attribute() is called in critical tracing paths:

• mlflow.start_span() - used in trace creation workflows
• mlflow.start_span_no_context() - manual span management
• Span processors during trace lifecycle events

The 28% overall speedup will directly benefit these tracing operations, especially when spans contain many non-JSON string attributes, which appears common in OpenTelemetry attribute storage patterns.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import json
import logging

# imports
import pytest
from mlflow.tracing.utils.__init__ import get_otel_attribute


# Dummy Span class to simulate OpenTelemetry Span for testing
class DummySpan:
    def __init__(self, attributes):
        self.attributes = attributes
from mlflow.tracing.utils.__init__ import get_otel_attribute

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

# Basic Test Cases

def test_basic_string_decoding():
    # Attribute is a JSON string representing a string
    span = DummySpan({"foo": json.dumps("bar")})
    codeflash_output = get_otel_attribute(span, "foo") # 6.71μs -> 6.93μs (3.19% slower)

def test_basic_number_decoding():
    # Attribute is a JSON string representing a number
    span = DummySpan({"num": json.dumps(123)})
    codeflash_output = get_otel_attribute(span, "num") # 6.40μs -> 1.09μs (487% faster)

def test_basic_boolean_decoding():
    # Attribute is a JSON string representing a boolean
    span = DummySpan({"flag": json.dumps(True)})
    codeflash_output = get_otel_attribute(span, "flag") # 6.21μs -> 1.13μs (450% faster)

def test_basic_list_decoding():
    # Attribute is a JSON string representing a list
    span = DummySpan({"lst": json.dumps([1, 2, 3])})
    codeflash_output = get_otel_attribute(span, "lst") # 6.67μs -> 7.50μs (11.0% slower)

def test_basic_dict_decoding():
    # Attribute is a JSON string representing a dict
    span = DummySpan({"dct": json.dumps({"a": 1, "b": 2})})
    codeflash_output = get_otel_attribute(span, "dct") # 7.15μs -> 7.94μs (9.91% slower)

def test_basic_none_decoding():
    # Attribute is a JSON string representing None
    span = DummySpan({"none": json.dumps(None)})
    codeflash_output = get_otel_attribute(span, "none") # 5.95μs -> 1.15μs (416% faster)

def test_basic_missing_key():
    # Key does not exist in attributes
    span = DummySpan({"foo": json.dumps("bar")})
    codeflash_output = get_otel_attribute(span, "baz") # 704ns -> 641ns (9.83% faster)

# Edge Test Cases

def test_edge_attribute_is_none():
    # Attribute exists but value is None
    span = DummySpan({"foo": None})
    codeflash_output = get_otel_attribute(span, "foo") # 677ns -> 686ns (1.31% slower)

def test_edge_attribute_not_json():
    # Attribute value is not a valid JSON string
    span = DummySpan({"foo": "not a json"})
    codeflash_output = get_otel_attribute(span, "foo") # 15.7μs -> 1.16μs (1253% faster)

def test_edge_attribute_is_empty_string():
    # Attribute value is an empty string
    span = DummySpan({"foo": ""})
    codeflash_output = get_otel_attribute(span, "foo") # 15.0μs -> 862ns (1640% faster)

def test_edge_attribute_is_json_empty_string():
    # Attribute value is a JSON empty string
    span = DummySpan({"foo": json.dumps("")})
    codeflash_output = get_otel_attribute(span, "foo") # 6.75μs -> 7.53μs (10.3% slower)

def test_edge_attribute_is_json_empty_list():
    # Attribute value is a JSON empty list
    span = DummySpan({"foo": json.dumps([])})
    codeflash_output = get_otel_attribute(span, "foo") # 6.12μs -> 7.08μs (13.6% slower)

def test_edge_attribute_is_json_empty_dict():
    # Attribute value is a JSON empty dict
    span = DummySpan({"foo": json.dumps({})})
    codeflash_output = get_otel_attribute(span, "foo") # 6.23μs -> 6.92μs (9.94% slower)

def test_edge_attribute_is_json_false():
    # Attribute value is a JSON false boolean
    span = DummySpan({"foo": json.dumps(False)})
    codeflash_output = get_otel_attribute(span, "foo") # 6.06μs -> 1.13μs (435% faster)

def test_edge_attribute_is_json_zero():
    # Attribute value is a JSON zero
    span = DummySpan({"foo": json.dumps(0)})
    codeflash_output = get_otel_attribute(span, "foo") # 6.22μs -> 1.13μs (449% faster)

def test_edge_attribute_is_json_float():
    # Attribute value is a JSON float
    span = DummySpan({"foo": json.dumps(3.1415)})
    codeflash_output = get_otel_attribute(span, "foo") # 6.80μs -> 1.19μs (469% faster)

def test_edge_attribute_is_json_nested():
    # Attribute value is a JSON nested structure
    data = {"a": [1, {"b": "c"}], "d": None}
    span = DummySpan({"foo": json.dumps(data)})
    codeflash_output = get_otel_attribute(span, "foo") # 7.64μs -> 8.25μs (7.33% slower)

def test_edge_attribute_is_json_with_unicode():
    # Attribute value is a JSON string with unicode characters
    span = DummySpan({"foo": json.dumps("こんにちは")})
    codeflash_output = get_otel_attribute(span, "foo") # 7.01μs -> 7.71μs (9.03% slower)

def test_edge_attribute_is_json_with_escaped_characters():
    # Attribute value is a JSON string with escaped characters
    span = DummySpan({"foo": json.dumps("line1\nline2")})
    codeflash_output = get_otel_attribute(span, "foo") # 7.41μs -> 7.79μs (4.89% slower)

def test_edge_span_has_no_attributes():
    # Span has no attributes dictionary
    class EmptySpan:
        def __init__(self):
            self.attributes = {}
    span = EmptySpan()
    codeflash_output = get_otel_attribute(span, "foo") # 721ns -> 707ns (1.98% faster)



def test_large_scale_many_attributes():
    # Span with 1000 attributes, each a JSON string of its index
    attributes = {str(i): json.dumps(i) for i in range(1000)}
    span = DummySpan(attributes)
    # Check random sampling of keys
    for i in [0, 10, 500, 999]:
        codeflash_output = get_otel_attribute(span, str(i)) # 11.0μs -> 2.65μs (313% faster)
    # Check missing key
    codeflash_output = get_otel_attribute(span, "1000") # 294ns -> 274ns (7.30% faster)

def test_large_scale_large_json_object():
    # Attribute is a large JSON object (dict of 1000 items)
    large_dict = {str(i): i for i in range(1000)}
    span = DummySpan({"big": json.dumps(large_dict)})
    codeflash_output = get_otel_attribute(span, "big"); result = codeflash_output # 109μs -> 108μs (0.289% faster)
    for i in [0, 10, 500, 999]:
        pass

def test_large_scale_large_json_array():
    # Attribute is a large JSON array (list of 1000 items)
    large_list = list(range(1000))
    span = DummySpan({"biglist": json.dumps(large_list)})
    codeflash_output = get_otel_attribute(span, "biglist"); result = codeflash_output # 35.8μs -> 35.8μs (0.008% faster)
    for i in [0, 10, 500, 999]:
        pass

def test_large_scale_large_string():
    # Attribute is a large JSON string (length 1000)
    large_string = "a" * 1000
    span = DummySpan({"bigstr": json.dumps(large_string)})
    codeflash_output = get_otel_attribute(span, "bigstr"); result = codeflash_output # 7.52μs -> 8.10μs (7.24% slower)

def test_large_scale_performance():
    # Ensure function doesn't take too long with large input
    import time
    large_list = list(range(1000))
    span = DummySpan({"biglist": json.dumps(large_list)})
    start = time.time()
    codeflash_output = get_otel_attribute(span, "biglist"); result = codeflash_output # 35.2μs -> 36.4μs (3.28% slower)
    duration = time.time() - start
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import json
import logging

# imports
import pytest
from mlflow.tracing.utils.__init__ import get_otel_attribute


# Helper: A minimal Span mock for testing purposes
class DummySpan:
    def __init__(self, attributes):
        self.attributes = attributes

# ------------------ BASIC TEST CASES ------------------

def test_basic_valid_json_string():
    # Attribute value is a valid JSON string for a string
    span = DummySpan({"foo": json.dumps("bar")})
    codeflash_output = get_otel_attribute(span, "foo") # 6.51μs -> 7.24μs (10.1% slower)

def test_basic_valid_json_int():
    # Attribute value is a valid JSON string for an int
    span = DummySpan({"num": json.dumps(123)})
    codeflash_output = get_otel_attribute(span, "num") # 6.27μs -> 1.12μs (460% faster)

def test_basic_valid_json_list():
    # Attribute value is a valid JSON string for a list
    span = DummySpan({"lst": json.dumps([1, 2, 3])})
    codeflash_output = get_otel_attribute(span, "lst") # 6.62μs -> 7.45μs (11.1% slower)

def test_basic_valid_json_dict():
    # Attribute value is a valid JSON string for a dict
    span = DummySpan({"dct": json.dumps({"a": 1, "b": 2})})
    codeflash_output = get_otel_attribute(span, "dct") # 7.10μs -> 7.81μs (9.05% slower)

def test_basic_none_key():
    # Attribute key does not exist
    span = DummySpan({"foo": json.dumps("bar")})
    codeflash_output = get_otel_attribute(span, "baz") # 707ns -> 656ns (7.77% faster)

# ------------------ EDGE TEST CASES ------------------

def test_edge_attribute_value_is_none():
    # Attribute value is None
    span = DummySpan({"foo": None})
    codeflash_output = get_otel_attribute(span, "foo") # 692ns -> 663ns (4.37% faster)

def test_edge_attribute_value_not_json():
    # Attribute value is not a valid JSON string
    span = DummySpan({"foo": "not json"})
    codeflash_output = get_otel_attribute(span, "foo") # 16.0μs -> 1.06μs (1402% faster)

def test_edge_attribute_value_empty_string():
    # Attribute value is an empty string
    span = DummySpan({"foo": ""})
    codeflash_output = get_otel_attribute(span, "foo") # 15.2μs -> 848ns (1695% faster)

def test_edge_attribute_value_json_null():
    # Attribute value is the JSON string "null"
    span = DummySpan({"foo": "null"})
    codeflash_output = get_otel_attribute(span, "foo") # 6.82μs -> 1.11μs (512% faster)

def test_edge_attribute_value_json_true_false():
    # Attribute value is the JSON string "true" and "false"
    span = DummySpan({"foo": "true", "bar": "false"})
    codeflash_output = get_otel_attribute(span, "foo") # 6.79μs -> 1.14μs (495% faster)
    codeflash_output = get_otel_attribute(span, "bar") # 1.76μs -> 437ns (303% faster)

def test_edge_attribute_value_json_number_as_string():
    # Attribute value is a JSON string representing a number
    span = DummySpan({"foo": "123"})
    # Not valid JSON, should return None
    codeflash_output = get_otel_attribute(span, "foo") # 7.12μs -> 1.11μs (542% faster)

def test_edge_attribute_value_json_string_with_spaces():
    # Attribute value is a valid JSON string with spaces
    span = DummySpan({"foo": json.dumps("   spaced   ")})
    codeflash_output = get_otel_attribute(span, "foo") # 6.71μs -> 7.34μs (8.62% slower)

def test_edge_attribute_key_is_empty_string():
    # Attribute key is an empty string
    span = DummySpan({"": json.dumps("emptykey")})
    codeflash_output = get_otel_attribute(span, "") # 6.64μs -> 7.26μs (8.64% slower)

def test_edge_attribute_key_is_none():
    # Attribute key is None
    span = DummySpan({None: json.dumps("nonekey")})
    codeflash_output = get_otel_attribute(span, None) # 6.66μs -> 7.12μs (6.49% slower)

def test_edge_attribute_value_is_json_array_of_dicts():
    # Attribute value is a JSON array of dicts
    value = json.dumps([{"a": 1}, {"b": 2}])
    span = DummySpan({"foo": value})
    codeflash_output = get_otel_attribute(span, "foo") # 7.67μs -> 8.12μs (5.61% slower)

def test_edge_attribute_value_is_json_nested():
    # Attribute value is a deeply nested JSON structure
    value = json.dumps({"a": {"b": {"c": [1, 2, 3]}}})
    span = DummySpan({"foo": value})
    codeflash_output = get_otel_attribute(span, "foo") # 7.83μs -> 8.30μs (5.77% slower)

def test_edge_attribute_value_is_json_empty_list():
    # Attribute value is a JSON empty list
    span = DummySpan({"foo": "[]"})
    codeflash_output = get_otel_attribute(span, "foo") # 7.08μs -> 7.67μs (7.69% slower)

def test_edge_attribute_value_is_json_empty_dict():
    # Attribute value is a JSON empty dict
    span = DummySpan({"foo": "{}"})
    codeflash_output = get_otel_attribute(span, "foo") # 7.21μs -> 7.84μs (7.99% slower)

def test_edge_span_attributes_is_empty():
    # Span attributes is empty dict
    span = DummySpan({})
    codeflash_output = get_otel_attribute(span, "foo") # 706ns -> 709ns (0.423% slower)


def test_edge_span_attributes_is_not_dict():
    # Span attributes is not a dict
    class DummySpanNotDict:
        attributes = "not a dict"
    span = DummySpanNotDict()
    codeflash_output = get_otel_attribute(span, "foo") # 5.48μs -> 5.25μs (4.42% faster)

def test_edge_span_attributes_has_non_string_keys():
    # Span attributes has non-string keys
    span = DummySpan({1: json.dumps("one"), (2, 3): json.dumps("tuple")})
    codeflash_output = get_otel_attribute(span, 1) # 6.46μs -> 7.27μs (11.1% slower)
    codeflash_output = get_otel_attribute(span, (2, 3)) # 2.03μs -> 2.13μs (4.64% slower)

def test_edge_attribute_value_is_json_float():
    # Attribute value is a JSON float
    span = DummySpan({"foo": json.dumps(1.234)})
    codeflash_output = get_otel_attribute(span, "foo") # 6.86μs -> 1.14μs (500% faster)

# ------------------ LARGE SCALE TEST CASES ------------------

def test_large_scale_many_attributes():
    # Span with 1000 attributes, all valid JSON strings
    attributes = {f"key{i}": json.dumps(i) for i in range(1000)}
    span = DummySpan(attributes)
    # Test a few random keys
    codeflash_output = get_otel_attribute(span, "key0") # 6.27μs -> 1.34μs (366% faster)
    codeflash_output = get_otel_attribute(span, "key999") # 1.95μs -> 452ns (332% faster)
    codeflash_output = get_otel_attribute(span, "key500") # 1.24μs -> 298ns (317% faster)
    # Test a missing key
    codeflash_output = get_otel_attribute(span, "key1000") # 341ns -> 300ns (13.7% faster)

def test_large_scale_many_invalid_json_attributes():
    # Span with 1000 attributes, all invalid JSON strings
    attributes = {f"badkey{i}": "not json" for i in range(1000)}
    span = DummySpan(attributes)
    # All should return None
    for i in [0, 499, 999]:
        codeflash_output = get_otel_attribute(span, f"badkey{i}") # 23.8μs -> 2.13μs (1021% faster)

def test_large_scale_mixed_valid_invalid_json():
    # Span with 500 valid and 500 invalid JSON attributes
    attributes = {f"good{i}": json.dumps(i) for i in range(500)}
    attributes.update({f"bad{i}": "not json" for i in range(500)})
    span = DummySpan(attributes)
    # Valid keys
    for i in [0, 250, 499]:
        codeflash_output = get_otel_attribute(span, f"good{i}") # 9.69μs -> 2.14μs (352% faster)
    # Invalid keys
    for i in [0, 250, 499]:
        codeflash_output = get_otel_attribute(span, f"bad{i}") # 17.5μs -> 987ns (1676% faster)

def test_large_scale_large_json_object():
    # Attribute value is a large JSON object
    large_dict = {f"k{i}": i for i in range(1000)}
    span = DummySpan({"large": json.dumps(large_dict)})
    codeflash_output = get_otel_attribute(span, "large") # 107μs -> 110μs (2.27% slower)

def test_large_scale_large_json_array():
    # Attribute value is a large JSON array
    large_list = [i for i in range(1000)]
    span = DummySpan({"large_list": json.dumps(large_list)})
    codeflash_output = get_otel_attribute(span, "large_list") # 35.5μs -> 37.0μs (3.96% slower)

def test_large_scale_large_nested_json():
    # Attribute value is a large, deeply nested JSON structure
    nested = {"level1": {"level2": {"level3": [i for i in range(1000)]}}}
    span = DummySpan({"nested": json.dumps(nested)})
    codeflash_output = get_otel_attribute(span, "nested") # 36.3μs -> 36.8μs (1.46% slower)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-get_otel_attribute-mhwy4jyu and push.