⚡️ Speed up function create_array by 11%

[codeflash-ai]

📄 11% (0.11x) speedup for create_array in inference/core/workflows/execution_engine/v1/executor/output_constructor.py

⏱️ Runtime : 24.3 milliseconds → 21.9 milliseconds (best of 45 runs)

📝 Explanation and details

The optimization improves performance by replacing expensive repeated boolean indexing with efficient grouping. The key changes are:

What was optimized:

1. Pre-computed grouping: Uses np.unique() and builds an idx_map dictionary to group rows by their first column index upfront, avoiding repeated indices[:, 0] == idx boolean operations inside the loop.
2. Direct integer indexing: Replaces boolean masking (indices[idx_selector]) with direct integer indexing (indices[idx_indices]), which is faster for NumPy arrays.
3. Efficient empty array handling: Creates empty arrays with proper shape only when needed, rather than computing boolean masks for non-existent indices.

Why it's faster:
The original code performs boolean indexing (indices[:, 0] == idx) for every possible index in the range, creating O(n×m) operations where n is the array size and m is the max index. The optimization reduces this to O(n) by grouping once upfront, then using direct integer indexing.

Performance characteristics:

• Small arrays: Slightly slower (30-50% overhead) due to upfront np.unique() and dictionary creation costs
• Large sparse arrays: Dramatically faster (236% speedup) when max index >> number of actual indices, as it avoids unnecessary boolean operations
• Dense arrays: Modest improvements (11% overall speedup) as the grouping overhead is amortized

Impact on workloads:
Based on the function references, create_array is called in workflow execution engine output construction - a critical path for processing batch outputs. The optimization particularly benefits workflows with sparse index patterns or large batch sizes, which are common in computer vision pipelines where not all detection slots are filled.

✅ Correctness verification report:

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

⚙️ Existing Unit Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
workflows/unit_tests/execution_engine/executor/test_output_constructor.py::test_create_array_for_dimension_one	29.9μs	36.2μs	-17.6%⚠️
workflows/unit_tests/execution_engine/executor/test_output_constructor.py::test_create_array_for_dimension_three	125μs	186μs	-32.7%⚠️
workflows/unit_tests/execution_engine/executor/test_output_constructor.py::test_create_array_for_dimension_two	62.9μs	93.6μs	-32.8%⚠️

🌀 Generated Regression Tests and Runtime

import numpy as np

# imports
import pytest  # used for our unit tests
from inference.core.workflows.execution_engine.v1.executor.output_constructor import (
    create_array,
)

# unit tests

# --------------------------
# Basic Test Cases
# --------------------------


def test_empty_indices_returns_none():
    # Test that an empty array returns None
    indices = np.array([]).reshape(0, 2)
    codeflash_output = create_array(indices)  # 528ns -> 544ns (2.94% slower)


def test_single_element():
    # Test with a single index
    indices = np.array([[0, 0]])
    codeflash_output = create_array(indices)  # 23.8μs -> 42.2μs (43.7% slower)


def test_two_elements_same_row():
    # Two elements, same first index
    indices = np.array([[0, 0], [0, 1]])
    codeflash_output = create_array(indices)  # 23.2μs -> 43.4μs (46.4% slower)


def test_two_elements_different_rows():
    # Two elements, different first index
    indices = np.array([[0, 0], [1, 0]])
    codeflash_output = create_array(indices)  # 27.2μs -> 50.9μs (46.6% slower)


def test_multiple_rows_and_columns():
    # Test with a 2-level structure
    indices = np.array([[0, 0, 0], [0, 0, 1], [0, 1, 0], [1, 0, 0]])
    expected = [[[None, None], [None]], [[None]]]
    codeflash_output = create_array(indices)  # 42.5μs -> 76.7μs (44.6% slower)


# --------------------------
# Edge Test Cases
# --------------------------


def test_single_row_multiple_columns():
    # Single row, multiple columns
    indices = np.array([[0, 1, 2, 3]])
    # Should nest to depth 3, each level has one element
    codeflash_output = create_array(indices)  # 42.3μs -> 58.2μs (27.4% slower)


def test_zero_shape_but_not_empty():
    # Shape (0, n) should return None
    indices = np.empty((0, 3), dtype=int)
    codeflash_output = create_array(indices)  # 540ns -> 546ns (1.10% slower)


def test_nonzero_first_col_but_empty():
    # If max index is not 0, should fill up to max index
    indices = np.array([[2, 0]])
    # Should have 3 outer lists, only the last one has [None]
    codeflash_output = create_array(indices)  # 29.3μs -> 43.8μs (33.1% slower)


def test_gapped_indices():
    # Indices skip some values in first column
    indices = np.array([[0, 0], [2, 0]])
    # Should have 3 outer lists: [ [None], None, [None] ]
    codeflash_output = create_array(indices)  # 29.0μs -> 52.8μs (45.0% slower)


def test_high_dimension():
    # Test with 4D indices
    indices = np.array([[0, 0, 0, 0], [0, 0, 0, 1], [0, 1, 0, 0], [1, 0, 0, 0]])
    expected = [[[[None, None]], [[None]]], [[[None]]]]
    codeflash_output = create_array(indices)  # 51.6μs -> 99.6μs (48.2% slower)


def test_all_indices_same():
    # All indices are the same
    indices = np.array([[0, 0, 0], [0, 0, 0]])
    # Should produce [[ [None, None] ]]
    codeflash_output = create_array(indices)  # 22.0μs -> 45.9μs (52.1% slower)


def test_single_row_single_column():
    # 1D index, should return [None, None, None]
    indices = np.array([[0], [1], [2]])
    codeflash_output = create_array(indices)  # 18.7μs -> 30.2μs (38.2% slower)


def test_non_consecutive_indices():
    # Indices with gaps in first column
    indices = np.array([[0, 0], [2, 1]])
    expected = [[None], None, [None]]
    codeflash_output = create_array(indices)  # 32.8μs -> 52.0μs (37.0% slower)


# --------------------------
# Large Scale Test Cases
# --------------------------


def test_large_flat_array():
    # Large 1D array
    n = 1000
    indices = np.arange(n).reshape(n, 1)
    codeflash_output = create_array(indices)
    result = codeflash_output  # 4.01ms -> 1.76ms (128% faster)


def test_large_2d_array():
    # Large 2D array, 1000 rows, 2 columns
    n = 1000
    indices = np.stack([np.arange(n), np.zeros(n, dtype=int)], axis=1)
    codeflash_output = create_array(indices)
    result = codeflash_output  # 7.51ms -> 8.45ms (11.2% slower)
    for sub in result:
        pass


def test_large_sparse_array():
    # Large array with only a few nonzero entries
    n = 1000
    indices = np.array([[0, 0], [n - 1, 0]])
    codeflash_output = create_array(indices)
    result = codeflash_output  # 2.02ms -> 601μs (236% faster)
    for i in range(1, n - 1):
        pass


def test_large_deep_array():
    # Test with large depth, but only one path
    depth = 10
    indices = np.zeros((1, depth), dtype=int)
    codeflash_output = create_array(indices)
    result = codeflash_output  # 48.5μs -> 98.7μs (50.9% slower)
    # Should be nested lists, depth-1, ending in [None]
    node = result
    for _ in range(depth - 1):
        node = node[0]


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

import numpy as np

# imports
import pytest  # used for our unit tests
from inference.core.workflows.execution_engine.v1.executor.output_constructor import (
    create_array,
)

# unit tests

# ---------------- Basic Test Cases ----------------


def test_empty_input_returns_none():
    # Test with completely empty array
    arr = np.empty((0, 2), dtype=int)
    codeflash_output = create_array(arr)  # 573ns -> 587ns (2.39% slower)


def test_single_element():
    # Test with a single 2D index
    arr = np.array([[0, 0]])
    codeflash_output = create_array(arr)  # 26.0μs -> 46.2μs (43.7% slower)


def test_single_row_multiple_columns():
    # Test with a single row, more columns
    arr = np.array([[0, 1, 2]])
    # Should produce [[[None]]] (3 levels)
    codeflash_output = create_array(arr)  # 34.2μs -> 51.0μs (33.0% slower)


def test_multiple_elements_simple():
    # Test with two indices in first dimension
    arr = np.array([[0, 0], [1, 0]])
    # Should produce [[None], [None]]
    codeflash_output = create_array(arr)  # 25.4μs -> 51.5μs (50.6% slower)


def test_multiple_elements_multiple_levels():
    # Test with multiple indices in two dimensions
    arr = np.array([[0, 0], [0, 1], [1, 0]])
    # Should produce [[None, None], [None]]
    codeflash_output = create_array(arr)  # 29.3μs -> 53.3μs (45.0% slower)


def test_nonzero_start_index():
    # Test with indices starting at nonzero (should fill with None for missing indices)
    arr = np.array([[1, 0], [2, 0]])
    # Should produce [None, [None], [None]]
    codeflash_output = create_array(arr)  # 28.8μs -> 49.6μs (41.8% slower)


# ---------------- Edge Test Cases ----------------


def test_empty_second_dimension():
    # Test with shape (N, 0): should return None for any N
    arr = np.empty((3, 0), dtype=int)
    codeflash_output = create_array(arr)  # 547ns -> 523ns (4.59% faster)


def test_single_column():
    # Test with only one column (should return list of None)
    arr = np.array([[0], [1], [2]])
    # Should produce [None, None, None]
    codeflash_output = create_array(arr)  # 23.1μs -> 35.8μs (35.4% slower)


def test_missing_indices():
    # Test with missing indices in the first column (should fill with None)
    arr = np.array([[0, 0], [2, 0]])
    # Should produce [[None], None, [None]]
    codeflash_output = create_array(arr)  # 31.4μs -> 54.5μs (42.5% slower)


def test_deeply_nested():
    # Test with 3D indices
    arr = np.array([[0, 0, 0], [0, 1, 0], [1, 0, 0]])
    # Should produce [[[None], [None]], [[None]]]
    codeflash_output = create_array(arr)  # 38.8μs -> 74.3μs (47.7% slower)


def test_irregular_indices():
    # Test with indices not sorted, and with gaps
    arr = np.array([[2, 1], [0, 0], [2, 0]])
    # Should produce [[None], None, [None, None]]
    codeflash_output = create_array(arr)  # 31.5μs -> 53.4μs (41.0% slower)


def test_zero_level():
    # Test with an array that has only one column and one row
    arr = np.array([[0]])
    # Should produce [None]
    codeflash_output = create_array(arr)  # 13.4μs -> 22.7μs (40.9% slower)


# ---------------- Large Scale Test Cases ----------------


def test_large_flat_array():
    # Test with a large number of flat indices
    arr = np.stack([np.arange(1000)], axis=1)
    arr = np.hstack([arr, np.zeros((1000, 1), dtype=int)])
    # Should produce [None, None, ..., None] (1000 Nones)
    codeflash_output = create_array(arr)
    result = codeflash_output  # 7.50ms -> 8.47ms (11.5% slower)


def test_large_nested_array():
    # Test with a large number of indices in two dimensions
    arr = np.array([[i, j] for i in range(10) for j in range(10)])
    # Should produce a 10x10 nested list of None
    codeflash_output = create_array(arr)
    result = codeflash_output  # 245μs -> 295μs (16.8% slower)
    for sublist in result:
        pass


def test_large_sparse_array():
    # Test with a large first dimension but only a few indices filled
    arr = np.array([[0, 0], [999, 0]])
    # Should produce [ [None], None, ..., None, [None] ] (1000 elements)
    codeflash_output = create_array(arr)
    result = codeflash_output  # 2.02ms -> 600μs (236% faster)
    for i in range(1, 999):
        pass


def test_large_deeply_nested_array():
    # Test with 3D indices, 10x10x10
    arr = np.array([[i, j, k] for i in range(3) for j in range(3) for k in range(3)])
    codeflash_output = create_array(arr)
    result = codeflash_output  # 109μs -> 176μs (38.3% slower)
    for l2 in result:
        for l3 in l2:
            for item in l3:
                pass


# 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

To edit these changes git checkout codeflash/optimize-create_array-miqiwooe and push.