⚡️ Speed up function _get_level_lengths by 8%

[codeflash-ai]

📄 8% (0.08x) speedup for _get_level_lengths in pandas/io/formats/style_render.py

⏱️ Runtime : 10.6 milliseconds → 9.87 milliseconds (best of 25 runs)

📝 Explanation and details

The optimization improves performance by replacing list membership checks with set membership checks for the hidden_elements parameter, which provides O(1) average-case lookup time instead of O(n) for lists.

Key Optimizations:

1. Set-based membership testing: Converts hidden_elements from a list to a set at function entry, changing if j not in hidden_elements checks from O(n) to O(1) operations. This is particularly impactful since these checks occur within nested loops that can iterate thousands of times.

2. Safe dictionary access: Replaces direct dictionary key access lengths[(i, last_label)] with lengths.get((i, last_label), None) to avoid potential KeyError exceptions and improve robustness.

3. Variable initialization: Adds last_label = None initialization to ensure the variable is always defined before use in the inner loop.

Performance Impact:
The optimization delivers a 7% speedup overall with the most significant gains appearing in large-scale test cases:

• test_large_multiindex_with_hidden_and_trimming: 16.6% faster
• test_large_index_with_all_hidden: 299% faster (most dramatic improvement)
• Small overhead (1-3% slower) on very small datasets due to set creation cost

Hot Path Context:
Based on the function references, _get_level_lengths is called from _translate() - a core styling method that processes DataFrames for HTML rendering. This function runs in pandas' styling pipeline where large MultiIndex DataFrames are common, making the O(1) membership optimization particularly valuable for real-world performance.

The optimization is most effective for cases with larger hidden element lists and higher iteration counts, which aligns well with typical pandas styling workloads involving complex hierarchical indexes.

✅ Correctness verification report:

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

⚙️ Existing Unit Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
io/formats/style/test_style.py::TestStyler.test_get_level_lengths	273μs	270μs	1.21%✅
io/formats/style/test_style.py::TestStyler.test_get_level_lengths_un_sorted	265μs	261μs	1.25%✅
io/formats/style/test_style.py::test_get_level_lengths_mi_hidden	176μs	176μs	0.359%✅

🌀 Generated Regression Tests and Runtime

# function to test (copied from pandas/io/formats/style_render.py)

# We'll use the real pandas Index and MultiIndex for realistic test cases
import pandas as pd

# imports
from pandas.io.formats.style_render import _get_level_lengths

# --------------------------
# Unit Tests for _get_level_lengths
# --------------------------

# 1. BASIC TEST CASES


def test_single_level_index_no_hidden():
    # Simple Index, no hidden elements
    idx = pd.Index(["a", "b", "c"])
    codeflash_output = _get_level_lengths(idx, sparsify=True, max_index=10)
    result = codeflash_output  # 72.2μs -> 71.3μs (1.31% faster)


def test_single_level_index_with_hidden():
    # Index with one hidden element
    idx = pd.Index(["a", "b", "c"])
    codeflash_output = _get_level_lengths(
        idx, sparsify=True, max_index=10, hidden_elements=[1]
    )
    result = codeflash_output  # 68.6μs -> 70.1μs (2.04% slower)


def test_multiindex_sparsify():
    # MultiIndex, sparsify=True (simulate sparsification)
    idx = pd.MultiIndex.from_tuples([("A", 1), ("A", 2), ("B", 1)], names=["x", "y"])
    # The _format_multi method is not public, but we can check that the function doesn't error
    codeflash_output = _get_level_lengths(idx, sparsify=True, max_index=10)
    result = codeflash_output  # 154μs -> 155μs (0.674% slower)
    # All visible elements should have length >=1, keys are (level, pos)
    for k, v in result.items():
        pass


def test_empty_index():
    # Empty Index
    idx = pd.Index([])
    codeflash_output = _get_level_lengths(idx, sparsify=True, max_index=10)
    result = codeflash_output  # 60.7μs -> 62.6μs (3.07% slower)


def test_single_element_index():
    # Single element Index
    idx = pd.Index(["z"])
    codeflash_output = _get_level_lengths(idx, sparsify=True, max_index=10)
    result = codeflash_output  # 74.6μs -> 75.4μs (1.05% slower)


# 2. EDGE TEST CASES


def test_hidden_all_elements():
    # All elements hidden
    idx = pd.Index(["a", "b", "c"])
    codeflash_output = _get_level_lengths(
        idx, sparsify=True, max_index=10, hidden_elements=[0, 1, 2]
    )
    result = codeflash_output  # 71.7μs -> 71.4μs (0.443% faster)


def test_hidden_some_multiindex():
    # Hide some elements in MultiIndex
    idx = pd.MultiIndex.from_tuples([("A", 1), ("A", 2), ("B", 1)])
    codeflash_output = _get_level_lengths(
        idx, sparsify=False, max_index=10, hidden_elements=[1]
    )
    result = codeflash_output  # 157μs -> 157μs (0.105% slower)


def test_max_index_trimming():
    # max_index trims the output
    idx = pd.Index(["a", "b", "c", "d", "e"])
    # max_index=2, only first 3 elements processed (since visible_row_count > max_index breaks)
    codeflash_output = _get_level_lengths(idx, sparsify=True, max_index=2)
    result = codeflash_output  # 75.3μs -> 75.0μs (0.392% faster)


def test_max_index_zero():
    # max_index=0, only first element processed
    idx = pd.Index(["a", "b", "c"])
    codeflash_output = _get_level_lengths(idx, sparsify=True, max_index=0)
    result = codeflash_output  # 69.0μs -> 69.9μs (1.41% slower)


def test_hidden_elements_out_of_bounds():
    # hidden_elements includes out-of-bounds index
    idx = pd.Index(["a", "b"])
    # hidden_elements=[2] is out of bounds, should be ignored
    codeflash_output = _get_level_lengths(
        idx, sparsify=True, max_index=10, hidden_elements=[2]
    )
    result = codeflash_output  # 66.6μs -> 65.7μs (1.32% faster)


def test_multiindex_with_duplicate_labels():
    # MultiIndex with duplicate labels
    idx = pd.MultiIndex.from_tuples([("A", 1), ("A", 1), ("B", 2)])
    codeflash_output = _get_level_lengths(idx, sparsify=False, max_index=10)
    result = codeflash_output  # 156μs -> 152μs (1.99% faster)


def test_multiindex_with_hidden_and_sparsify():
    # MultiIndex with hidden elements and sparsify=True
    idx = pd.MultiIndex.from_tuples([("A", 1), ("A", 2), ("B", 1), ("B", 1)])
    codeflash_output = _get_level_lengths(
        idx, sparsify=True, max_index=10, hidden_elements=[1, 3]
    )
    result = codeflash_output  # 158μs -> 154μs (2.88% faster)
    # Only visible elements should be present
    for k in result.keys():
        pass


# 3. LARGE SCALE TEST CASES


def test_large_single_index():
    # Large Index, 1000 elements
    idx = pd.Index(range(1000))
    codeflash_output = _get_level_lengths(idx, sparsify=True, max_index=999)
    result = codeflash_output  # 274μs -> 273μs (0.267% faster)


def test_large_multiindex():
    # Large MultiIndex, 100 x 10 = 1000 rows
    tuples = [(i, j) for i in range(100) for j in range(10)]
    idx = pd.MultiIndex.from_tuples(tuples, names=["level0", "level1"])
    codeflash_output = _get_level_lengths(idx, sparsify=False, max_index=999)
    result = codeflash_output  # 2.13ms -> 2.12ms (0.184% faster)


def test_large_multiindex_with_hidden_and_trimming():
    # Large MultiIndex, hide every 10th element, trim at 500
    tuples = [(i, j) for i in range(50) for j in range(20)]  # 1000 rows
    idx = pd.MultiIndex.from_tuples(tuples)
    hidden = list(range(0, 1000, 10))
    codeflash_output = _get_level_lengths(
        idx, sparsify=False, max_index=500, hidden_elements=hidden
    )
    result = codeflash_output  # 2.36ms -> 2.02ms (16.6% faster)


def test_large_index_with_all_hidden():
    # Large Index, all elements hidden
    idx = pd.Index(range(500))
    codeflash_output = _get_level_lengths(
        idx, sparsify=True, max_index=499, hidden_elements=list(range(500))
    )
    result = codeflash_output  # 496μs -> 124μs (299% faster)


def test_large_multiindex_sparsify_true():
    # Large MultiIndex, sparsify=True
    tuples = [(i, 0) for i in range(1000)]
    idx = pd.MultiIndex.from_tuples(tuples)
    codeflash_output = _get_level_lengths(idx, sparsify=True, max_index=999)
    result = codeflash_output  # 2.17ms -> 2.13ms (2.07% faster)


# 4. ADDITIONAL EDGE CASES


def test_index_with_non_string_labels():
    # Index with integer and float labels
    idx = pd.Index([1, 2.5, 3])
    codeflash_output = _get_level_lengths(idx, sparsify=True, max_index=10)
    result = codeflash_output  # 153μs -> 156μs (2.33% slower)


def test_multiindex_with_mixed_types():
    # MultiIndex with mixed types
    idx = pd.MultiIndex.from_tuples([("A", 1), ("B", 2.5), ("C", "foo")])
    codeflash_output = _get_level_lengths(idx, sparsify=False, max_index=10)
    result = codeflash_output  # 167μs -> 166μs (0.799% faster)


def test_index_with_duplicate_labels():
    # Index with duplicates
    idx = pd.Index(["x", "x", "y"])
    codeflash_output = _get_level_lengths(idx, sparsify=True, max_index=10)
    result = codeflash_output  # 66.1μs -> 67.6μs (2.14% slower)


def test_mutation_wrong_lengths_fail():
    # If the function returns wrong lengths, this test should fail.
    idx = pd.Index(["a", "b"])
    wrong = {(0, 0): 2, (0, 1): 1}
    codeflash_output = _get_level_lengths(idx, sparsify=True, max_index=10)
    result = codeflash_output  # 66.6μs -> 66.2μs (0.518% faster)


def test_mutation_wrong_keys_fail():
    # If the function returns wrong keys, this test should fail.
    idx = pd.Index(["a", "b"])
    wrong = {(1, 0): 1, (1, 1): 1}
    codeflash_output = _get_level_lengths(idx, sparsify=True, max_index=10)
    result = codeflash_output  # 66.2μs -> 66.5μs (0.483% slower)


# 6. DETERMINISM TEST


def test_determinism():
    # Multiple calls should return the same result
    idx = pd.Index(["a", "b", "c"])
    codeflash_output = _get_level_lengths(idx, sparsify=True, max_index=10)
    r1 = codeflash_output  # 67.0μs -> 67.0μs (0.046% slower)
    codeflash_output = _get_level_lengths(idx, sparsify=True, max_index=10)
    r2 = codeflash_output  # 37.9μs -> 36.2μs (4.62% faster)


# 7. TEST WITH NAMED INDEX


def test_named_index():
    idx = pd.Index(["a", "b", "c"], name="letters")
    codeflash_output = _get_level_lengths(idx, sparsify=True, max_index=10)
    result = codeflash_output  # 63.7μs -> 65.3μs (2.50% slower)


# 8. TEST WITH NAMED MULTIINDEX


def test_named_multiindex():
    idx = pd.MultiIndex.from_tuples([("A", 1), ("B", 2)], names=["foo", "bar"])
    codeflash_output = _get_level_lengths(idx, sparsify=False, max_index=10)
    result = codeflash_output  # 154μs -> 154μs (0.179% faster)


# 9. TEST WITH NON-SEQUENCE HIDDEN_ELEMENTS


def test_hidden_elements_none():
    idx = pd.Index(["a", "b", "c"])
    codeflash_output = _get_level_lengths(
        idx, sparsify=True, max_index=10, hidden_elements=None
    )
    result = codeflash_output  # 68.6μs -> 68.0μs (0.770% faster)


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

# imports
from pandas.io.formats.style_render import _get_level_lengths


# Minimal MultiIndex and Index mocks for testing
class MockMultiIndex:
    def __init__(self, levels, codes, names=None, format_multi=None):
        self.levels = levels
        self.codes = codes
        self.names = names or [None] * len(levels)
        # format_multi is a list of lists of strings, one per level
        self._format_multi = lambda sparsify, include_names: format_multi


class MockIndex:
    def __init__(self, values, name=None, format_flat=None):
        self.values = values
        self.name = name
        # format_flat is a list of strings
        self._format_flat = lambda include_name: format_flat


# Dummy lib.no_default for sparsified cells
class DummyNoDefault:
    pass


no_default = DummyNoDefault()

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

# 1. BASIC TEST CASES


def test_flat_index_no_hidden():
    # Flat index, no hidden elements
    idx = MockIndex(values=["a", "b", "c"], format_flat=["a", "b", "c"])
    codeflash_output = _get_level_lengths(idx, sparsify=True, max_index=10)
    result = codeflash_output  # 2.76μs -> 3.14μs (12.1% slower)


def test_flat_index_with_hidden():
    # Flat index, hide one element
    idx = MockIndex(values=["a", "b", "c"], format_flat=["a", "b", "c"])
    codeflash_output = _get_level_lengths(
        idx, sparsify=True, max_index=10, hidden_elements=[1]
    )
    result = codeflash_output  # 2.51μs -> 3.01μs (16.7% slower)


def test_flat_index_empty():
    # Flat index, empty
    idx = MockIndex(values=[], format_flat=[])
    codeflash_output = _get_level_lengths(idx, sparsify=True, max_index=10)
    result = codeflash_output  # 1.78μs -> 1.89μs (6.02% slower)


def test_flat_index_all_hidden():
    # Flat index, all hidden
    idx = MockIndex(values=["a", "b", "c"], format_flat=["a", "b", "c"])
    codeflash_output = _get_level_lengths(
        idx, sparsify=True, max_index=10, hidden_elements=[0, 1, 2]
    )
    result = codeflash_output  # 2.41μs -> 2.93μs (17.7% slower)


def test_large_flat_index():
    # Large flat index (1000 elements)
    n = 1000
    idx = MockIndex(values=list(range(n)), format_flat=[str(i) for i in range(n)])
    codeflash_output = _get_level_lengths(idx, sparsify=True, max_index=1000)
    result = codeflash_output  # 71.8μs -> 72.5μs (0.922% slower)


def test_flat_index_with_max_index_trim():
    # Flat index, max_index trims output
    idx = MockIndex(values=["a", "b", "c", "d"], format_flat=["a", "b", "c", "d"])
    codeflash_output = _get_level_lengths(idx, sparsify=True, max_index=1)
    result = codeflash_output  # 2.89μs -> 3.12μs (7.31% 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_level_lengths-mio5nuyd and push.