⚡️ Speed up method Styler._map by 10%

[codeflash-ai]

📄 10% (0.10x) speedup for Styler._map in pandas/io/formats/style.py

⏱️ Runtime : 19.6 milliseconds → 17.7 milliseconds (best of 13 runs)

📝 Explanation and details

The optimized code achieves a 10% performance improvement by targeting key bottlenecks in the _update_ctx method and making a minor optimization in _map.

Key Optimizations Applied:

1. Pre-computed column lookups: Instead of calling self.columns.get_loc(cn) for each column in every iteration, the optimized version pre-computes all column locations in a dictionary (columns_get_loc) at the start. This eliminates repeated expensive index lookups.

2. Direct array access: The optimization uses attrs.values and attrs.index to access the underlying NumPy array data directly, avoiding the overhead of DataFrame column extraction (attrs[cn]) which was taking significant time in the original implementation.

3. Improved null checking: The original code used not c or pd.isna(c) which could be inefficient. The optimized version checks for common falsy values (c is None or c == "") first before falling back to pd.isna(c), providing a fast path for the most common cases.

4. Conditional functools.partial: In _map, the optimization avoids creating a functools.partial wrapper when no kwargs are provided (the common case), using the function directly instead.

Why These Optimizations Work:

• Reduced function call overhead: Pre-computing lookups and using direct array access eliminates thousands of redundant method calls in typical styling operations
• Better cache locality: Accessing data through NumPy arrays is more cache-friendly than going through pandas DataFrame accessors
• Fast-path optimizations: The null checking improvements handle the most common cases (None, empty string) without expensive pandas operations

Impact on Workloads:

The optimizations are particularly effective for:

• Large DataFrames with many styled cells (e.g., 100x10 test cases show ~12-16% improvement)
• DataFrames with sparse styling where many cells return empty/null values (up to 34% improvement on empty DataFrames)
• Frequent styling operations where the reduced overhead compounds across multiple calls

The improvements maintain full backward compatibility and preserve all existing behavior while significantly reducing the computational cost of common styling workflows.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

# Function to test (minimal, correct implementation of Styler._map)
import pandas as pd

# imports
import pytest
from pandas.io.formats.style import Styler

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


def test_basic_map_entire_dataframe():
    # Map a function that sets all cells to "color: red;"
    df = pd.DataFrame([[1, 2], [3, 4]], columns=["A", "B"])
    styler = Styler(df)

    def color_red(val):
        return "color: red;"

    styler._map(color_red)  # 364μs -> 324μs (12.3% faster)
    # All cells should have "color: red;"
    for i in range(2):
        for j in range(2):
            pass


def test_basic_map_with_kwargs():
    # Map a function with kwargs
    df = pd.DataFrame([[1, 2]], columns=["A", "B"])
    styler = Styler(df)

    def color_fn(val, color):
        return f"color: {color};"

    styler._map(color_fn, color="blue")  # 343μs -> 306μs (12.0% faster)
    for key in styler.ctx:
        pass


def test_basic_map_subset_column():
    # Map only column "A"
    df = pd.DataFrame([[1, 2], [3, 4]], columns=["A", "B"])
    styler = Styler(df)

    def color_green(val):
        return "color: green;"

    styler._map(color_green, subset="A")  # 557μs -> 519μs (7.30% faster)
    # Only column "A" should be styled
    for i in range(2):
        pass


def test_basic_map_subset_tuple():
    # Map only cell (row 1, col "B")
    df = pd.DataFrame([[1, 2], [3, 4]], columns=["A", "B"])
    styler = Styler(df)

    def color_orange(val):
        return "color: orange;"

    styler._map(color_orange, subset=([1], ["B"]))  # 756μs -> 748μs (1.14% faster)


def test_basic_map_with_none_and_nan():
    # Map function should skip None and NaN
    df = pd.DataFrame([[1, None], [float("nan"), 4]], columns=["A", "B"])
    styler = Styler(df)

    def color_black(val):
        return "color: black;"

    styler._map(color_black)  # 357μs -> 328μs (8.90% faster)


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


def test_edge_empty_dataframe():
    # Should not fail on empty DataFrame
    df = pd.DataFrame([], columns=["A", "B"])
    styler = Styler(df)

    def color_red(val):
        return "color: red;"

    styler._map(color_red)  # 161μs -> 120μs (34.0% faster)


def test_edge_function_returns_empty_string():
    # Should not add anything to ctx for empty string
    df = pd.DataFrame([[1]])
    styler = Styler(df)

    def empty(val):
        return ""

    styler._map(empty)  # 264μs -> 241μs (9.49% faster)


def test_edge_function_returns_none():
    # Should not add anything to ctx for None
    df = pd.DataFrame([[1]])
    styler = Styler(df)

    def none_fn(val):
        return None

    styler._map(none_fn)  # 250μs -> 230μs (8.48% faster)


def test_edge_function_returns_invalid_css_string():
    # Should raise ValueError for invalid CSS string
    df = pd.DataFrame([[1]])
    styler = Styler(df)

    def bad_css(val):
        return "not_a_css_rule"

    with pytest.raises(ValueError):
        styler._map(bad_css)  # 255μs -> 242μs (5.25% faster)


def test_large_scale_map_100x10():
    # DataFrame with 100 rows, 10 columns
    df = pd.DataFrame([[i * j for j in range(10)] for i in range(100)])
    styler = Styler(df)

    def color_fn(val):
        return f"color: #{val % 10}00;"

    styler._map(color_fn)  # 2.21ms -> 2.02ms (9.67% faster)


def test_large_scale_map_only_subset():
    # Only style a subset in a large DataFrame
    df = pd.DataFrame([[i * j for j in range(10)] for i in range(100)])
    styler = Styler(df)

    def color_fn(val):
        return "color: magenta;"

    # Only style first 5 rows, columns 0 and 1
    styler._map(color_fn, subset=(range(5), [0, 1]))  # 765μs -> 727μs (5.22% faster)
    for i in range(5):
        for j in [0, 1]:
            pass
    # Other cells should not be styled
    for i in range(5, 10):
        for j in [0, 1]:
            pass
    for i in range(5):
        for j in range(2, 10):
            pass


def test_large_scale_map_with_nan_and_none():
    # Large DataFrame with NaN/None scattered
    import numpy as np

    df = pd.DataFrame(
        [
            [
                np.nan if (i + j) % 17 == 0 else None if (i + j) % 19 == 0 else i + j
                for j in range(20)
            ]
            for i in range(50)
        ]
    )
    styler = Styler(df)

    def color_fn(val):
        return "color: navy;"

    styler._map(color_fn)  # 2.39ms -> 2.07ms (15.4% faster)
    # Only non-NaN, non-None cells should be styled
    for i in range(50):
        for j in range(20):
            val = df.iloc[i, j]
            if val is None or (isinstance(val, float) and pd.isna(val)):
                pass
            else:
                pass


def test_large_scale_map_performance():
    # Should finish in reasonable time for 1000 cells
    import time

    df = pd.DataFrame([[i + j for j in range(20)] for i in range(50)])
    styler = Styler(df)

    def color_fn(val):
        return "color: olive;"

    t0 = time.time()
    styler._map(color_fn)  # 2.41ms -> 2.07ms (16.0% faster)
    t1 = time.time()


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

import pandas as pd

# imports
import pytest

# function to test (Styler._map is defined in the provided code above)

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


def test_basic_apply_single_column():
    # Test mapping a function to a single column
    df = pd.DataFrame({"A": [1, 2], "B": [3, 4]})
    styler = pd.io.formats.style.Styler(df.copy())

    def func(val):
        return "color: red;" if val % 2 == 0 else "color: blue;"

    styler._map(func, subset=["A"])  # 591μs -> 534μs (10.6% faster)
    # Check that the ctx for column "A" is updated correctly
    i0 = df.index.get_loc(df.index[0])
    i1 = df.index.get_loc(df.index[1])
    jA = df.columns.get_loc("A")
    jB = df.columns.get_loc("B")


def test_basic_apply_all_columns():
    # Test mapping a function to all columns (default subset)
    df = pd.DataFrame({"A": [1, 2], "B": [3, 4]})
    styler = pd.io.formats.style.Styler(df.copy())

    def func(val):
        return "font-weight: bold;" if val > 2 else ""

    styler._map(func)  # 335μs -> 302μs (11.1% faster)
    for i in range(2):
        for j in range(2):
            if df.iloc[i, j] > 2:
                pass
            else:
                pass


def test_basic_apply_with_kwargs():
    # Test that kwargs are passed to the function via partial
    df = pd.DataFrame({"A": [1, 2]})
    styler = pd.io.formats.style.Styler(df.copy())

    def func(val, color):
        return f"color: {color};"

    styler._map(func, color="green")  # 285μs -> 260μs (9.66% faster)
    for i in range(2):
        pass


def test_edge_empty_dataframe():
    # Test with an empty DataFrame
    df = pd.DataFrame(columns=["A", "B"])
    styler = pd.io.formats.style.Styler(df.copy())

    def func(val):
        return "color: red;"

    styler._map(func)  # 161μs -> 113μs (41.8% faster)


def test_edge_nan_values():
    # Test with NaN values in the DataFrame
    import numpy as np

    df = pd.DataFrame({"A": [1, np.nan], "B": [np.nan, 4]})
    styler = pd.io.formats.style.Styler(df.copy())

    def func(val):
        if pd.isna(val):
            return "background-color: yellow;"
        return "color: black;"

    styler._map(func)  # 343μs -> 312μs (10.1% faster)


def test_edge_function_returns_none_or_empty():
    # Test when function returns None or empty string
    df = pd.DataFrame({"A": [1, 2]})
    styler = pd.io.formats.style.Styler(df.copy())

    def func(val):
        return None if val == 1 else ""

    styler._map(func)  # 264μs -> 250μs (5.37% faster)


def test_edge_non_unique_index_raises():
    # Test that non-unique index raises KeyError
    df = pd.DataFrame({"A": [1, 2]}, index=[0, 0])
    styler = pd.io.formats.style.Styler(df.copy())

    def func(val):
        return "color: red;"

    with pytest.raises(KeyError):
        styler._map(func)  # 229μs -> 235μs (2.58% slower)


def test_edge_non_unique_columns_raises():
    # Test that non-unique columns raises KeyError
    df = pd.DataFrame([[1, 2]], columns=["A", "A"])
    styler = pd.io.formats.style.Styler(df.copy())

    def func(val):
        return "color: red;"

    with pytest.raises(KeyError):
        styler._map(func)  # 275μs -> 274μs (0.517% faster)


def test_edge_function_raises_exception():
    # Test that if the function raises, the error is propagated
    df = pd.DataFrame({"A": [1, 2]})
    styler = pd.io.formats.style.Styler(df.copy())

    def func(val):
        raise ValueError("test error")

    with pytest.raises(ValueError):
        styler._map(func)  # 102μs -> 102μs (0.546% faster)


def test_edge_subset_as_scalar():
    # Test that a scalar subset works (single column)
    df = pd.DataFrame({"A": [1, 2], "B": [3, 4]})
    styler = pd.io.formats.style.Styler(df.copy())

    def func(val):
        return "color: red;"

    styler._map(func, subset="A")  # 556μs -> 514μs (8.30% faster)
    # Only column A should be styled
    for i in range(2):
        jA = df.columns.get_loc("A")
        jB = df.columns.get_loc("B")


def test_edge_subset_as_tuple():
    # Test that a tuple subset works (row, column)
    df = pd.DataFrame({"A": [1, 2], "B": [3, 4]}, index=["x", "y"])
    styler = pd.io.formats.style.Styler(df.copy())

    def func(val):
        return "color: green;"

    styler._map(func, subset=("x", "A"))  # 678μs -> 647μs (4.82% faster)
    # Only cell (x, A) should be styled
    i = df.index.get_loc("x")
    j = df.columns.get_loc("A")
    # All other cells should not be styled
    for idx in range(2):
        for col in range(2):
            if not (idx == i and col == j):
                pass


def test_edge_func_returns_invalid_css_string():
    # Test that an invalid CSS string raises ValueError
    df = pd.DataFrame({"A": [1]})
    styler = pd.io.formats.style.Styler(df.copy())

    def func(val):
        return "not_a_css_rule"

    with pytest.raises(ValueError):
        styler._map(func)  # 275μs -> 258μs (6.53% faster)


def test_edge_func_returns_mixed_types():
    # Test that a function returning a mix of valid and invalid CSS rules raises
    df = pd.DataFrame({"A": [1, 2]})
    styler = pd.io.formats.style.Styler(df.copy())

    def func(val):
        return "color: red;" if val == 1 else "invalid_css"

    with pytest.raises(ValueError):
        styler._map(func)  # 276μs -> 259μs (6.67% faster)


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


def test_large_scale_100x10():
    # Test with a 100x10 DataFrame
    rows, cols = 100, 10
    df = pd.DataFrame({f"C{j}": range(rows) for j in range(cols)})
    styler = pd.io.formats.style.Styler(df.copy())

    def func(val):
        # Style even numbers
        return "background-color: #eee;" if val % 2 == 0 else ""

    styler._map(func)  # 1.47ms -> 1.31ms (12.8% faster)
    for i in range(rows):
        for j in range(cols):
            if df.iloc[i, j] % 2 == 0:
                pass
            else:
                pass


def test_large_scale_all_nan():
    # Test with a 50x5 DataFrame of all NaN values
    import numpy as np

    df = pd.DataFrame(np.nan, index=range(50), columns=[f"C{i}" for i in range(5)])
    styler = pd.io.formats.style.Styler(df.copy())

    def func(val):
        return "opacity: 0.5;" if pd.isna(val) else ""

    styler._map(func)  # 896μs -> 802μs (11.7% faster)
    for i in range(50):
        for j in range(5):
            pass


def test_large_scale_subset():
    # Test with a subset on a large DataFrame
    rows, cols = 200, 6
    df = pd.DataFrame({f"C{j}": range(rows) for j in range(cols)})
    styler = pd.io.formats.style.Styler(df.copy())

    # Only style columns C1 and C3
    def func(val):
        return "border: 1px solid black;"

    styler._map(func, subset=["C1", "C3"])  # 1.21ms -> 1.12ms (8.37% faster)
    for i in range(rows):
        for j, col in enumerate(df.columns):
            if col in ["C1", "C3"]:
                pass
            else:
                pass


def test_large_scale_performance():
    # This test is mainly to ensure it doesn't take too long or error on 500x2
    df = pd.DataFrame({"A": range(500), "B": range(500, 1000)})
    styler = pd.io.formats.style.Styler(df.copy())

    def func(val):
        return "color: #123456;" if val % 100 == 0 else ""

    styler._map(func)  # 532μs -> 495μs (7.51% faster)


# 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-Styler._map-mj9xqj6m and push.