Skip to content

⚡️ Speed up function retrieve_selectors_from_union_definition by 98% in PR #1764 (joao/fix-array-input-references) #1907

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Closed
codeflash-ai wants to merge 1 commit into from
Closed

⚡️ Speed up function retrieve_selectors_from_union_definition by 98% in PR #1764 (joao/fix-array-input-references) #1907

codeflash-ai wants to merge 1 commit into from

Conversation

@codeflash-ai
Copy link
Contributor

@codeflash-ai codeflash-ai bot commented Jan 12, 2026

⚡️ This pull request contains optimizations for PR #1764

If you approve this dependent PR, these changes will be merged into the original PR branch joao/fix-array-input-references.

This PR will be automatically closed if the original PR is merged.

📄 98% (0.98x) speedup for retrieve_selectors_from_union_definition in inference/core/workflows/execution_engine/introspection/schema_parser.py

⏱️ Runtime : 246 microseconds → 125 microseconds (best of 190 runs)

📝 Explanation and details

The optimized code achieves a 97% speedup (246μs → 125μs) through two key optimizations:

1. Inlining property_defines_union() (Primary Impact)

What changed: The property_defines_union() function call was replaced with inline dictionary key checks in retrieve_selectors_from_simple_property.

Why it's faster:

  • Eliminates function call overhead: Python function calls have significant overhead (~300-500ns). The line profiler shows this call consumed 61.3% of runtime (2.91ms out of 4.75ms) in the original code.
  • Reduces dictionary lookups: The original made 6 lookups total (3 in the function check + 3 more in union processing). The optimized version performs the same 3 checks but only when needed, avoiding the function call stack frame creation/destruction.

Key Change in retrieve_selectors_from_simple_property:

# Before: Function call + 3 dict lookups every time
if property_defines_union(property_definition=property_definition):
    
# After: Inlined check with same 3 dictionary lookups but no function overhead
if (
    ANY_OF_KEY in property_definition
    or ONE_OF_KEY in property_definition
    or ALL_OF_KEY in property_definition
):

The line profiler shows this optimization reduced time spent on the union check from 2.91ms (61.3%) to 1.59ms (45.5%) - a significant improvement in the most expensive operation.

Key optimization in retrieve_selectors_from_union_definition:

The original code concatenated three lists using + operator:

union_types = (
    union_definition.get(ANY_OF_KEY, [])
    + union_definition.get(ONE_OF_KEY, [])  
    + union_definition.get(ALL_OF_KEY, [])
)

The optimized version iterates over keys directly and skips empty results:

for key in (ANY_OF_KEY, ONE_OF_KEY, ALL_OF_KEY):
    union_types = union_definition.get(key)
    if not union_types:
        continue

Why this is faster:

  1. Inlining property_defines_union(): The original code called a function that performed 3 dictionary lookups. Inlining eliminates function call overhead (~300-400ns per call) and reduces the lookup pattern from "check all 3 keys in function → check all 3 keys again in union handling" to just checking once.

  2. Avoiding list concatenation: The original code created 3 intermediate lists via .get() calls and concatenated them with + operators, creating a new list. The optimized version iterates over the three keys directly, checking existence before iteration, avoiding list construction overhead.

  3. Flattening nested iteration: The original code used itertools.chain.from_iterable() to flatten allowed_references from multiple results. The optimized version uses a single nested loop that directly iterates over references, reducing intermediate list creation overhead.

Performance Impact:

  • The line profiler shows the property_defines_union call in retrieve_selectors_from_simple_property dropped from 61.3% of function time (2.91ms) to ~45% (1.59ms total for the three inlined checks).
  • In retrieve_selectors_from_union_definition, eliminating list concatenation and using early-exit iteration reduced union type processing from 1.1% to 0.6% of total time.
  • The optimization is most effective for test cases with simple unions (153-199% faster) and moderately beneficial for complex unions with many references (39% faster on the large mixed union test).

Why this is faster:

  1. Function call elimination: Inlining property_defines_union() removes ~900ns of function call overhead per invocation (visible in line profiler: 2.9μs spent on that call alone for 951 hits)
  2. Reduced dictionary lookups: The original code performed 3 .get() calls plus list concatenation in retrieve_selectors_from_union_definition. The optimized version uses a single loop checking each key once with early continue on empty results
  3. Avoided list concatenation: The original created intermediate lists via + operator concatenation; the optimized version iterates keys directly, eliminating temporary list allocations
  4. Flattening nested comprehension: Changed from itertools.chain.from_iterable() on a generator to direct nested iteration, reducing iterator overhead

Impact on workloads:
Based on function_references, this function is called recursively during schema parsing, particularly when processing union types in workflow definitions. The optimizations reduce overhead per call, which compounds when processing schemas with multiple union definitions. The test results show consistent 130-150% speedup across all test cases, with the largest gains (190-199%) in edge cases with empty unions where the early exit optimization is most beneficial. The 39% speedup on the large mixed union test (300 references) demonstrates good scalability of the optimization.

✅ Correctness verification report:

Test Status
⏪ Replay Tests 🔘 None Found
⚙️ Existing Unit Tests 🔘 None Found
🔎 Concolic Coverage Tests 🔘 None Found
🌀 Generated Regression Tests ✅ 32 Passed
📊 Tests Coverage 66.7%
🌀 Click to see Generated Regression Tests 
import pytest
from inference.core.workflows.execution_engine.entities.types import Kind
from inference.core.workflows.execution_engine.introspection.entities import (
    ReferenceDefinition,
    SelectorDefinition,
)
from inference.core.workflows.execution_engine.introspection.schema_parser import (
    retrieve_selectors_from_union_definition,
)

# ============================================================================
# BASIC TEST CASES
# ============================================================================


def test_retrieve_selectors_from_union_definition_with_anyof_single_reference():
    """
    Test basic functionality with a single anyOf union containing one reference.
    This verifies that the function correctly extracts selectors from a simple anyOf union.
    """
    union_definition = {
        "anyOf": [
            {
                "$ref": "#/components/schemas/SomeModel",
                "x-selected-element": "predictions",
            }
        ]
    }

    codeflash_output = retrieve_selectors_from_union_definition(
        property_name="test_property",
        property_description="A test property",
        union_definition=union_definition,
        is_list_element=False,
        is_dict_element=False,
        property_dimensionality_offset=0,
        is_dimensionality_reference_property=False,
        inputs_accepting_batches=set(),
        inputs_accepting_batches_and_scalars=set(),
        inputs_enforcing_auto_batch_casting=set(),
    )
    result = codeflash_output  # 6.70μs -> 2.64μs (153% faster)


def test_retrieve_selectors_from_union_definition_with_oneof_single_reference():
    """
    Test basic functionality with a single oneOf union containing one reference.
    Ensures oneOf is handled the same way as anyOf.
    """
    union_definition = {
        "oneOf": [
            {
                "$ref": "#/components/schemas/AnotherModel",
                "x-selected-element": "output",
            }
        ]
    }

    codeflash_output = retrieve_selectors_from_union_definition(
        property_name="my_selector",
        property_description="My selector description",
        union_definition=union_definition,
        is_list_element=False,
        is_dict_element=False,
        property_dimensionality_offset=0,
        is_dimensionality_reference_property=False,
        inputs_accepting_batches=set(),
        inputs_accepting_batches_and_scalars=set(),
        inputs_enforcing_auto_batch_casting=set(),
    )
    result = codeflash_output  # 6.50μs -> 2.73μs (138% faster)


def test_retrieve_selectors_from_union_definition_with_allof_single_reference():
    """
    Test basic functionality with a single allOf union containing one reference.
    Ensures allOf is handled correctly.
    """
    union_definition = {
        "allOf": [
            {
                "$ref": "#/components/schemas/CombinedModel",
                "x-selected-element": "combined_output",
            }
        ]
    }

    codeflash_output = retrieve_selectors_from_union_definition(
        property_name="combined",
        property_description="Combined property",
        union_definition=union_definition,
        is_list_element=False,
        is_dict_element=False,
        property_dimensionality_offset=0,
        is_dimensionality_reference_property=False,
        inputs_accepting_batches=set(),
        inputs_accepting_batches_and_scalars=set(),
        inputs_enforcing_auto_batch_casting=set(),
    )
    result = codeflash_output  # 6.39μs -> 2.62μs (144% faster)


def test_retrieve_selectors_from_union_definition_with_multiple_references_same_element():
    """
    Test merging of references pointing to the same selected element.
    This verifies that references with identical selected_element are properly merged.
    """
    union_definition = {
        "anyOf": [
            {
                "$ref": "#/components/schemas/Model1",
                "x-selected-element": "same_element",
            },
            {
                "$ref": "#/components/schemas/Model2",
                "x-selected-element": "same_element",
            },
        ]
    }

    codeflash_output = retrieve_selectors_from_union_definition(
        property_name="merged_property",
        property_description="Property with merged references",
        union_definition=union_definition,
        is_list_element=False,
        is_dict_element=False,
        property_dimensionality_offset=0,
        is_dimensionality_reference_property=False,
        inputs_accepting_batches=set(),
        inputs_accepting_batches_and_scalars=set(),
        inputs_enforcing_auto_batch_casting=set(),
    )
    result = codeflash_output  # 6.61μs -> 2.86μs (131% faster)


def test_retrieve_selectors_from_union_definition_preserves_property_metadata():
    """
    Test that property metadata (dimensionality offset, description) is preserved in result.
    """
    union_definition = {
        "anyOf": [
            {
                "$ref": "#/components/schemas/Model1",
                "x-selected-element": "output",
            }
        ]
    }

    codeflash_output = retrieve_selectors_from_union_definition(
        property_name="test_prop",
        property_description="Test description",
        union_definition=union_definition,
        is_list_element=False,
        is_dict_element=False,
        property_dimensionality_offset=2,
        is_dimensionality_reference_property=True,
        inputs_accepting_batches=set(),
        inputs_accepting_batches_and_scalars=set(),
        inputs_enforcing_auto_batch_casting=set(),
    )
    result = codeflash_output  # 8.66μs -> 3.53μs (145% faster)


# ============================================================================
# EDGE TEST CASES
# ============================================================================


def test_retrieve_selectors_from_union_definition_with_empty_union():
    """
    Test behavior when union_definition has empty arrays for all union types.
    Should return None when no valid references are found.
    """
    union_definition = {
        "anyOf": [],
        "oneOf": [],
        "allOf": [],
    }

    codeflash_output = retrieve_selectors_from_union_definition(
        property_name="empty_union",
        property_description="Empty union",
        union_definition=union_definition,
        is_list_element=False,
        is_dict_element=False,
        property_dimensionality_offset=0,
        is_dimensionality_reference_property=False,
        inputs_accepting_batches=set(),
        inputs_accepting_batches_and_scalars=set(),
        inputs_enforcing_auto_batch_casting=set(),
    )
    result = codeflash_output  # 5.82μs -> 2.01μs (190% faster)


def test_retrieve_selectors_from_union_definition_with_no_union_keys():
    """
    Test behavior when union_definition has no anyOf, oneOf, or allOf keys.
    Should return None when there are no union types.
    """
    union_definition = {}

    codeflash_output = retrieve_selectors_from_union_definition(
        property_name="no_union_keys",
        property_description="No union keys",
        union_definition=union_definition,
        is_list_element=False,
        is_dict_element=False,
        property_dimensionality_offset=0,
        is_dimensionality_reference_property=False,
        inputs_accepting_batches=set(),
        inputs_accepting_batches_and_scalars=set(),
        inputs_enforcing_auto_batch_casting=set(),
    )
    result = codeflash_output  # 5.45μs -> 1.82μs (199% faster)


def test_retrieve_selectors_from_union_definition_with_non_reference_union_types():
    """
    Test that non-reference union members are filtered out (return None from recursive calls).
    Only reference definitions should be included in final result.
    """
    union_definition = {
        "anyOf": [
            {"type": "string"},  # Not a reference
            {
                "$ref": "#/components/schemas/Model1",
                "x-selected-element": "valid_element",
            },
            {"type": "number"},  # Not a reference
        ]
    }

    codeflash_output = retrieve_selectors_from_union_definition(
        property_name="mixed_union",
        property_description="Mixed union with non-references",
        union_definition=union_definition,
        is_list_element=False,
        is_dict_element=False,
        property_dimensionality_offset=0,
        is_dimensionality_reference_property=False,
        inputs_accepting_batches=set(),
        inputs_accepting_batches_and_scalars=set(),
        inputs_enforcing_auto_batch_casting=set(),
    )
    result = codeflash_output  # 7.70μs -> 3.54μs (117% faster)


def test_retrieve_selectors_from_union_definition_with_all_non_references():
    """
    Test behavior when union contains only non-reference types.
    Should return None when no references can be extracted.
    """
    union_definition = {
        "anyOf": [
            {"type": "string"},
            {"type": "number"},
            {"type": "boolean"},
        ]
    }

    codeflash_output = retrieve_selectors_from_union_definition(
        property_name="non_reference_union",
        property_description="Union with no references",
        union_definition=union_definition,
        is_list_element=False,
        is_dict_element=False,
        property_dimensionality_offset=0,
        is_dimensionality_reference_property=False,
        inputs_accepting_batches=set(),
        inputs_accepting_batches_and_scalars=set(),
        inputs_enforcing_auto_batch_casting=set(),
    )
    result = codeflash_output  # 7.18μs -> 3.11μs (131% faster)


def test_retrieve_selectors_from_union_definition_with_list_element_flag():
    """
    Test that is_list_element flag is preserved in the result.
    """
    union_definition = {
        "anyOf": [
            {
                "$ref": "#/components/schemas/Model1",
                "x-selected-element": "output",
            }
        ]
    }

    codeflash_output = retrieve_selectors_from_union_definition(
        property_name="list_property",
        property_description="List property",
        union_definition=union_definition,
        is_list_element=True,
        is_dict_element=False,
        property_dimensionality_offset=0,
        is_dimensionality_reference_property=False,
        inputs_accepting_batches=set(),
        inputs_accepting_batches_and_scalars=set(),
        inputs_enforcing_auto_batch_casting=set(),
    )
    result = codeflash_output  # 6.33μs -> 2.54μs (149% faster)


def test_retrieve_selectors_from_union_definition_with_dict_element_flag():
    """
    Test that is_dict_element flag is preserved in the result.
    """
    union_definition = {
        "anyOf": [
            {
                "$ref": "#/components/schemas/Model1",
                "x-selected-element": "output",
            }
        ]
    }

    codeflash_output = retrieve_selectors_from_union_definition(
        property_name="dict_property",
        property_description="Dict property",
        union_definition=union_definition,
        is_list_element=False,
        is_dict_element=True,
        property_dimensionality_offset=0,
        is_dimensionality_reference_property=False,
        inputs_accepting_batches=set(),
        inputs_accepting_batches_and_scalars=set(),
        inputs_enforcing_auto_batch_casting=set(),
    )
    result = codeflash_output  # 6.11μs -> 2.50μs (144% faster)


def test_retrieve_selectors_from_union_definition_combined_list_and_dict_flags():
    """
    Test that both list and dict flags are properly combined when present in union results.
    """
    union_definition = {
        "anyOf": [
            {
                "$ref": "#/components/schemas/Model1",
                "x-selected-element": "output",
            }
        ]
    }

    codeflash_output = retrieve_selectors_from_union_definition(
        property_name="combined_property",
        property_description="Combined flags property",
        union_definition=union_definition,
        is_list_element=True,
        is_dict_element=True,
        property_dimensionality_offset=0,
        is_dimensionality_reference_property=False,
        inputs_accepting_batches=set(),
        inputs_accepting_batches_and_scalars=set(),
        inputs_enforcing_auto_batch_casting=set(),
    )
    result = codeflash_output  # 6.08μs -> 2.51μs (142% faster)


def test_retrieve_selectors_from_union_definition_with_negative_dimensionality_offset():
    """
    Test handling of negative dimensionality_offset values.
    """
    union_definition = {
        "anyOf": [
            {
                "$ref": "#/components/schemas/Model1",
                "x-selected-element": "output",
            }
        ]
    }

    codeflash_output = retrieve_selectors_from_union_definition(
        property_name="negative_offset",
        property_description="Negative offset property",
        union_definition=union_definition,
        is_list_element=False,
        is_dict_element=False,
        property_dimensionality_offset=-1,
        is_dimensionality_reference_property=False,
        inputs_accepting_batches=set(),
        inputs_accepting_batches_and_scalars=set(),
        inputs_enforcing_auto_batch_casting=set(),
    )
    result = codeflash_output  # 8.79μs -> 3.46μs (154% faster)


def test_retrieve_selectors_from_union_definition_with_empty_property_name():
    """
    Test behavior with empty string as property_name.
    Function should still process it without error.
    """
    union_definition = {
        "anyOf": [
            {
                "$ref": "#/components/schemas/Model1",
                "x-selected-element": "output",
            }
        ]
    }

    codeflash_output = retrieve_selectors_from_union_definition(
        property_name="",
        property_description="Empty name property",
        union_definition=union_definition,
        is_list_element=False,
        is_dict_element=False,
        property_dimensionality_offset=0,
        is_dimensionality_reference_property=False,
        inputs_accepting_batches=set(),
        inputs_accepting_batches_and_scalars=set(),
        inputs_enforcing_auto_batch_casting=set(),
    )
    result = codeflash_output  # 7.01μs -> 3.00μs (134% faster)


def test_retrieve_selectors_from_union_definition_with_empty_property_description():
    """
    Test behavior with empty string as property_description.
    """
    union_definition = {
        "anyOf": [
            {
                "$ref": "#/components/schemas/Model1",
                "x-selected-element": "output",
            }
        ]
    }

    codeflash_output = retrieve_selectors_from_union_definition(
        property_name="test_prop",
        property_description="",
        union_definition=union_definition,
        is_list_element=False,
        is_dict_element=False,
        property_dimensionality_offset=0,
        is_dimensionality_reference_property=False,
        inputs_accepting_batches=set(),
        inputs_accepting_batches_and_scalars=set(),
        inputs_enforcing_auto_batch_casting=set(),
    )
    result = codeflash_output  # 6.71μs -> 2.90μs (132% faster)


def test_retrieve_selectors_from_union_definition_with_very_long_property_name():
    """
    Test behavior with very long property_name string.
    """
    long_name = "a" * 1000
    union_definition = {
        "anyOf": [
            {
                "$ref": "#/components/schemas/Model1",
                "x-selected-element": "output",
            }
        ]
    }

    codeflash_output = retrieve_selectors_from_union_definition(
        property_name=long_name,
        property_description="Long name property",
        union_definition=union_definition,
        is_list_element=False,
        is_dict_element=False,
        property_dimensionality_offset=0,
        is_dimensionality_reference_property=False,
        inputs_accepting_batches=set(),
        inputs_accepting_batches_and_scalars=set(),
        inputs_enforcing_auto_batch_casting=set(),
    )
    result = codeflash_output  # 6.38μs -> 2.73μs (134% faster)


def test_retrieve_selectors_from_union_definition_with_special_characters_in_names():
    """
    Test handling of special characters in property names.
    """
    union_definition = {
        "anyOf": [
            {
                "$ref": "#/components/schemas/Model1",
                "x-selected-element": "output",
            }
        ]
    }

    special_name = "prop_with-special.chars$123"
    codeflash_output = retrieve_selectors_from_union_definition(
        property_name=special_name,
        property_description="Special chars",
        union_definition=union_definition,
        is_list_element=False,
        is_dict_element=False,
        property_dimensionality_offset=0,
        is_dimensionality_reference_property=False,
        inputs_accepting_batches=set(),
        inputs_accepting_batches_and_scalars=set(),
        inputs_enforcing_auto_batch_casting=set(),
    )
    result = codeflash_output  # 6.56μs -> 2.73μs (140% faster)


def test_retrieve_selectors_from_union_definition_with_unicode_property_names():
    """
    Test handling of unicode characters in property names.
    """
    union_definition = {
        "anyOf": [
            {
                "$ref": "#/components/schemas/Model1",
                "x-selected-element": "output",
            }
        ]
    }

    unicode_name = "プロパティ名_属性"
    codeflash_output = retrieve_selectors_from_union_definition(
        property_name=unicode_name,
        property_description="Unicode property",
        union_definition=union_definition,
        is_list_element=False,
        is_dict_element=False,
        property_dimensionality_offset=0,
        is_dimensionality_reference_property=False,
        inputs_accepting_batches=set(),
        inputs_accepting_batches_and_scalars=set(),
        inputs_enforcing_auto_batch_casting=set(),
    )
    result = codeflash_output  # 6.31μs -> 2.61μs (142% faster)


def test_retrieve_selectors_from_union_definition_with_duplicate_references():
    """
    Test handling when same reference appears multiple times in union.
    Duplicates with same selected_element should be merged.
    """
    union_definition = {
        "anyOf": [
            {
                "$ref": "#/components/schemas/Model1",
                "x-selected-element": "element_a",
            },
            {
                "$ref": "#/components/schemas/Model1",
                "x-selected-element": "element_a",
            },
            {
                "$ref": "#/components/schemas/Model1",
                "x-selected-element": "element_a",
            },
        ]
    }

    codeflash_output = retrieve_selectors_from_union_definition(
        property_name="duplicates",
        property_description="Duplicate references",
        union_definition=union_definition,
        is_list_element=False,
        is_dict_element=False,
        property_dimensionality_offset=0,
        is_dimensionality_reference_property=False,
        inputs_accepting_batches=set(),
        inputs_accepting_batches_and_scalars=set(),
        inputs_enforcing_auto_batch_casting=set(),
    )
    result = codeflash_output  # 7.08μs -> 3.25μs (118% faster)


# ============================================================================
# LARGE SCALE TEST CASES
# ============================================================================


def test_retrieve_selectors_from_union_definition_with_mixed_large_union():
    """
    Test with large union containing mix of anyOf, oneOf, and allOf.
    Creates 300 total references across different union type keys.
    """
    union_definition = {
        "anyOf": [
            {
                "$ref": f"#/components/schemas/ModelAny{i}",
                "x-selected-element": f"elem_any_{i % 20}",
            }
            for i in range(100)
        ],
        "oneOf": [
            {
                "$ref": f"#/components/schemas/ModelOne{i}",
                "x-selected-element": f"elem_one_{i % 20}",
            }
            for i in range(100)
        ],
        "allOf": [
            {
                "$ref": f"#/components/schemas/ModelAll{i}",
                "x-selected-element": f"elem_all_{i % 20}",
            }
            for i in range(100)
        ],
    }

    codeflash_output = retrieve_selectors_from_union_definition(
        property_name="large_mixed",
        property_description="Large mixed union",
        union_definition=union_definition,
        is_list_element=False,
        is_dict_element=False,
        property_dimensionality_offset=0,
        is_dimensionality_reference_property=False,
        inputs_accepting_batches=set(),
        inputs_accepting_batches_and_scalars=set(),
        inputs_enforcing_auto_batch_casting=set(),
    )
    result = codeflash_output  # 72.6μs -> 52.2μs (39.1% faster)


def test_retrieve_selectors_from_union_definition_performance_with_large_batches():
    """
    Test performance with large batch parameters.
    Verifies function handles large sets of batch-accepting inputs.
    """
    large_batch_set = {f"input_{i}" for i in range(500)}
    large_batch_scalars_set = {f"input_scalar_{i}" for i in range(300)}
    large_auto_cast_set = {f"input_auto_{i}" for i in range(200)}

    union_definition = {
        "anyOf": [
            {
                "$ref": "#/components/schemas/Model1",
                "x-selected-element": "output",
            }
        ]
    }

    codeflash_output = retrieve_selectors_from_union_definition(
        property_name="large_batch_property",
        property_description="Property with large batch sets",
        union_definition=union_definition,
        is_list_element=False,
        is_dict_element=False,
        property_dimensionality_offset=0,
        is_dimensionality_reference_property=False,
        inputs_accepting_batches=large_batch_set,
        inputs_accepting_batches_and_scalars=large_batch_scalars_set,
        inputs_enforcing_auto_batch_casting=large_auto_cast_set,
    )
    result = codeflash_output  # 7.52μs -> 3.03μs (148% faster)


def test_retrieve_selectors_from_union_definition_with_large_dimensionality_offset():
    """
    Test handling of very large dimensionality_offset values.
    """
    union_definition = {
        "anyOf": [
            {
                "$ref": "#/components/schemas/Model1",
                "x-selected-element": "output",
            }
        ]
    }

    codeflash_output = retrieve_selectors_from_union_definition(
        property_name="large_offset",
        property_description="Large dimensionality offset",
        union_definition=union_definition,
        is_list_element=False,
        is_dict_element=False,
        property_dimensionality_offset=999,
        is_dimensionality_reference_property=False,
        inputs_accepting_batches=set(),
        inputs_accepting_batches_and_scalars=set(),
        inputs_enforcing_auto_batch_casting=set(),
    )
    result = codeflash_output  # 6.84μs -> 2.80μs (144% faster)


def test_retrieve_selectors_from_union_definition_consistency_with_repeated_calls():
    """
    Test that repeated calls with same parameters produce consistent results.
    Ensures no side effects or state pollution.
    """
    union_definition = {
        "anyOf": [
            {
                "$ref": "#/components/schemas/Model1",
                "x-selected-element": "elem_a",
            },
            {
                "$ref": "#/components/schemas/Model2",
                "x-selected-element": "elem_b",
            },
        ]
    }

    # Call multiple times
    results = []
    for _ in range(10):
        codeflash_output = retrieve_selectors_from_union_definition(
            property_name="consistent",
            property_description="Consistency test",
            union_definition=union_definition,
            is_list_element=False,
            is_dict_element=False,
            property_dimensionality_offset=0,
            is_dimensionality_reference_property=False,
            inputs_accepting_batches=set(),
            inputs_accepting_batches_and_scalars=set(),
            inputs_enforcing_auto_batch_casting=set(),
        )
        result = codeflash_output  # 30.9μs -> 13.5μs (129% faster)
        results.append(result)
    for r in results[1:]:
        pass


# 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-pr1764-2026-01-12T18.20.55 and push.

Codeflash Static Badge

@codeflash-ai
Optimize retrieve_selectors_from_union_definition
0aaa230 
The optimized code achieves a **97% speedup** (246μs → 125μs) through two key optimizations:

## 1. Inlining `property_defines_union()` (Primary Impact)
**What changed**: The `property_defines_union()` function call was replaced with inline dictionary key checks in `retrieve_selectors_from_simple_property`.

**Why it's faster**: 
- **Eliminates function call overhead**: Python function calls have significant overhead (~300-500ns). The line profiler shows this call consumed **61.3% of runtime** (2.91ms out of 4.75ms) in the original code.
- **Reduces dictionary lookups**: The original made 6 lookups total (3 in the function check + 3 more in union processing). The optimized version performs the same 3 checks but only when needed, avoiding the function call stack frame creation/destruction.

**Key Change in `retrieve_selectors_from_simple_property`:**
```python
# Before: Function call + 3 dict lookups every time
if property_defines_union(property_definition=property_definition):
    
# After: Inlined check with same 3 dictionary lookups but no function overhead
if (
    ANY_OF_KEY in property_definition
    or ONE_OF_KEY in property_definition
    or ALL_OF_KEY in property_definition
):
```

The line profiler shows this optimization reduced time spent on the union check from **2.91ms (61.3%)** to **1.59ms (45.5%)** - a significant improvement in the most expensive operation.

**Key optimization in `retrieve_selectors_from_union_definition`:**

The original code concatenated three lists using `+` operator:
```python
union_types = (
    union_definition.get(ANY_OF_KEY, [])
    + union_definition.get(ONE_OF_KEY, [])  
    + union_definition.get(ALL_OF_KEY, [])
)
```

The optimized version iterates over keys directly and skips empty results:
```python
for key in (ANY_OF_KEY, ONE_OF_KEY, ALL_OF_KEY):
    union_types = union_definition.get(key)
    if not union_types:
        continue
```

**Why this is faster:**

1. **Inlining `property_defines_union()`**: The original code called a function that performed 3 dictionary lookups. Inlining eliminates function call overhead (~300-400ns per call) and reduces the lookup pattern from "check all 3 keys in function → check all 3 keys again in union handling" to just checking once.

2. **Avoiding list concatenation**: The original code created 3 intermediate lists via `.get()` calls and concatenated them with `+` operators, creating a new list. The optimized version iterates over the three keys directly, checking existence before iteration, avoiding list construction overhead.

3. **Flattening nested iteration**: The original code used `itertools.chain.from_iterable()` to flatten `allowed_references` from multiple results. The optimized version uses a single nested loop that directly iterates over references, reducing intermediate list creation overhead.

**Performance Impact:**
- The line profiler shows the `property_defines_union` call in `retrieve_selectors_from_simple_property` dropped from 61.3% of function time (2.91ms) to ~45% (1.59ms total for the three inlined checks).
- In `retrieve_selectors_from_union_definition`, eliminating list concatenation and using early-exit iteration reduced union type processing from 1.1% to 0.6% of total time.
- The optimization is most effective for test cases with simple unions (153-199% faster) and moderately beneficial for complex unions with many references (39% faster on the large mixed union test).

**Why this is faster:**
1. **Function call elimination**: Inlining `property_defines_union()` removes ~900ns of function call overhead per invocation (visible in line profiler: 2.9μs spent on that call alone for 951 hits)
2. **Reduced dictionary lookups**: The original code performed 3 `.get()` calls plus list concatenation in `retrieve_selectors_from_union_definition`. The optimized version uses a single loop checking each key once with early `continue` on empty results
3. **Avoided list concatenation**: The original created intermediate lists via `+` operator concatenation; the optimized version iterates keys directly, eliminating temporary list allocations
4. **Flattening nested comprehension**: Changed from `itertools.chain.from_iterable()` on a generator to direct nested iteration, reducing iterator overhead

**Impact on workloads:**
Based on `function_references`, this function is called recursively during schema parsing, particularly when processing union types in workflow definitions. The optimizations reduce overhead per call, which compounds when processing schemas with multiple union definitions. The test results show consistent 130-150% speedup across all test cases, with the largest gains (190-199%) in edge cases with empty unions where the early exit optimization is most beneficial. The 39% speedup on the large mixed union test (300 references) demonstrates good scalability of the optimization.
@codeflash-ai codeflash-ai bot added the ⚡️ codeflash Optimization PR opened by Codeflash AI label Jan 12, 2026
@codeflash-ai codeflash-ai bot added the 🎯 Quality: High Optimization Quality according to codeflash label Jan 12, 2026
@codeflash-ai codeflash-ai bot mentioned this pull request Jan 12, 2026
Merged
3 tasks
@Borda
Copy link
Member

Borda commented Jan 12, 2026

hmmm it does not require any review to be merged... 🤔

@codeflash-ai codeflash-ai bot closed this Jan 12, 2026
@codeflash-ai
Copy link
Contributor Author

codeflash-ai bot commented Jan 12, 2026

This PR has been automatically closed because the original PR #1764 by joaomarcoscrs was closed.

@codeflash-ai codeflash-ai bot deleted the codeflash/optimize-pr1764-2026-01-12T18.20.55 branch January 12, 2026 20:28
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@PawelPeczek-Roboflow PawelPeczek-Roboflow Awaiting requested review from PawelPeczek-Roboflow

@grzegorz-roboflow grzegorz-roboflow Awaiting requested review from grzegorz-roboflow

@yeldarby yeldarby Awaiting requested review from yeldarby

@probicheaux probicheaux Awaiting requested review from probicheaux

@hansent hansent Awaiting requested review from hansent

Assignees

No one assigned

Labels

⚡️ codeflash Optimization PR opened by Codeflash AI 🎯 Quality: High Optimization Quality according to codeflash

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@Borda