⚡️ Speed up method Categorical._rank by 14%

[codeflash-ai]

📄 14% (0.14x) speedup for Categorical._rank in pandas/core/arrays/categorical.py

⏱️ Runtime : 2.56 milliseconds → 2.26 milliseconds (best of 79 runs)

📝 Explanation and details

The optimized code achieves a 13% speedup through two key optimizations that eliminate overhead in hot code paths:

Key Optimizations Applied

1. Streamlined rank() Function in algorithms.py

• Eliminated _ensure_data() call overhead: Replaced the expensive function call (855μs per call) with inline conversion logic using np.asarray() only when needed
• Cached attribute access: Store values.dtype and values.ndim in variables to avoid repeated property lookups
• Early branching: Check ndim directly without additional preprocessing since 1D arrays are the most common case

2. Property Caching in _values_for_rank()

• Cached repeated property access: Store self.ordered, self.codes, and self.categories in local variables upfront to avoid expensive attribute lookups in conditional branches
• Optimized numeric categorical handling: For unordered categoricals with numeric categories, return codes directly instead of calling np.array(self), which creates unnecessary array conversions

Performance Impact by Test Case

The optimizations show particularly strong gains for:

• Unordered numeric categoricals: Up to 76% faster (62.4μs → 35.6μs) due to avoiding np.array(self) conversion
• Large-scale operations: 12-15% improvements across all large datasets due to reduced overhead
• Missing value handling: 16-20% faster due to property caching eliminating repeated lookups during mask operations

The optimizations are most effective for unordered categoricals with numeric categories and large datasets, while maintaining identical correctness across all test scenarios.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import numpy as np
# imports
import pytest
from pandas.core.arrays.categorical import Categorical

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

def test_basic_ordered_average():
    # Ordered categorical, no missing, average method
    cat = Categorical(['a', 'b', 'c', 'a'], categories=['a', 'b', 'c'], ordered=True)
    codeflash_output = cat._rank(method='average'); result = codeflash_output # 40.5μs -> 36.3μs (11.6% faster)

def test_basic_unordered_numeric():
    # Unordered categorical, numeric categories
    cat = Categorical([2, 1, 3, 2], categories=[1,2,3], ordered=False)
    codeflash_output = cat._rank(method='average'); result = codeflash_output # 61.8μs -> 35.0μs (76.5% faster)

def test_basic_unordered_string():
    # Unordered categorical, string categories
    cat = Categorical(['b', 'a', 'c', 'b'], categories=['a', 'b', 'c'], ordered=False)
    codeflash_output = cat._rank(method='average'); result = codeflash_output # 201μs -> 200μs (0.454% faster)

def test_basic_dense_method():
    # Dense ranking
    cat = Categorical(['a', 'b', 'c', 'a'], categories=['a', 'b', 'c'], ordered=True)
    codeflash_output = cat._rank(method='dense'); result = codeflash_output # 40.6μs -> 35.1μs (15.5% faster)

def test_basic_first_method():
    # First ranking
    cat = Categorical(['b', 'a', 'c', 'b'], categories=['a', 'b', 'c'], ordered=True)
    codeflash_output = cat._rank(method='first'); result = codeflash_output # 40.4μs -> 35.8μs (12.8% faster)

def test_basic_pct():
    # Percentile rank
    cat = Categorical(['a', 'b', 'c', 'a'], categories=['a', 'b', 'c'], ordered=True)
    codeflash_output = cat._rank(method='average', pct=True); result = codeflash_output # 41.4μs -> 35.9μs (15.4% faster)

def test_basic_descending():
    # Descending order
    cat = Categorical(['a', 'b', 'c', 'a'], categories=['a', 'b', 'c'], ordered=True)
    codeflash_output = cat._rank(method='average', ascending=False); result = codeflash_output # 43.5μs -> 37.7μs (15.2% faster)

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

def test_edge_empty():
    # Empty input
    cat = Categorical([], categories=['a', 'b'], ordered=True)
    codeflash_output = cat._rank(); result = codeflash_output # 39.5μs -> 34.7μs (13.8% faster)

def test_edge_all_missing():
    # All values missing
    cat = Categorical([None, None], categories=['a', 'b'], ordered=True)
    codeflash_output = cat._rank(); result = codeflash_output # 45.3μs -> 37.9μs (19.7% faster)

def test_edge_some_missing_keep():
    # Some missing, na_option='keep'
    cat = Categorical(['a', None, 'b', 'a'], categories=['a', 'b'], ordered=True)
    codeflash_output = cat._rank(na_option='keep'); result = codeflash_output # 44.4μs -> 37.4μs (18.4% faster)

def test_edge_some_missing_top():
    # Some missing, na_option='top'
    cat = Categorical(['a', None, 'b', 'a'], categories=['a', 'b'], ordered=True)
    codeflash_output = cat._rank(na_option='top'); result = codeflash_output # 45.6μs -> 38.5μs (18.3% faster)

def test_edge_some_missing_bottom():
    # Some missing, na_option='bottom'
    cat = Categorical(['a', None, 'b', 'a'], categories=['a', 'b'], ordered=True)
    codeflash_output = cat._rank(na_option='bottom'); result = codeflash_output # 44.3μs -> 36.9μs (20.1% faster)

def test_edge_duplicate_categories():
    # Duplicate categories in input values
    cat = Categorical(['a', 'a', 'b', 'c', 'b', 'c'], categories=['a', 'b', 'c'], ordered=True)
    codeflash_output = cat._rank(method='min'); result = codeflash_output # 41.2μs -> 36.3μs (13.6% faster)

def test_edge_invalid_axis():
    # axis != 0 should raise NotImplementedError
    cat = Categorical(['a','b'], categories=['a','b'], ordered=True)
    with pytest.raises(NotImplementedError):
        cat._rank(axis=1) # 1.48μs -> 1.42μs (4.15% faster)

def test_edge_invalid_method():
    # Invalid method should raise ValueError
    cat = Categorical(['a','b'], categories=['a','b'], ordered=True)
    with pytest.raises(ValueError):
        cat._rank(method='foobar')

def test_edge_invalid_na_option():
    # Invalid na_option should raise ValueError
    cat = Categorical(['a','b'], categories=['a','b'], ordered=True)
    with pytest.raises(ValueError):
        cat._rank(na_option='invalid')

def test_edge_unsorted_categories():
    # Unordered categories, input values not sorted
    cat = Categorical(['c', 'a', 'b'], categories=['a', 'b', 'c'], ordered=False)
    codeflash_output = cat._rank(); result = codeflash_output # 205μs -> 205μs (0.128% slower)

def test_edge_non_numeric_categories():
    # Categories are objects (tuples)
    cat = Categorical([(1,2), (2,3), (1,2)], categories=[(1,2), (2,3)], ordered=True)
    codeflash_output = cat._rank(); result = codeflash_output # 42.0μs -> 36.0μs (16.9% faster)

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

def test_large_unique_categories():
    # Large number of unique categories
    n = 1000
    values = [f"cat{i}" for i in range(n)]
    cat = Categorical(values, categories=values, ordered=True)
    codeflash_output = cat._rank(); result = codeflash_output # 50.1μs -> 43.9μs (14.0% faster)

def test_large_many_duplicates():
    # Large input with many duplicates
    n = 1000
    values = ['a']*500 + ['b']*500
    cat = Categorical(values, categories=['a','b'], ordered=True)
    codeflash_output = cat._rank(method='average'); result = codeflash_output # 47.2μs -> 41.5μs (13.6% faster)

def test_large_missing_values():
    # Large input with many missing values
    n = 1000
    values = ['a']*500 + [None]*500
    cat = Categorical(values, categories=['a'], ordered=True)
    codeflash_output = cat._rank(); result = codeflash_output # 52.6μs -> 45.4μs (15.9% faster)

def test_large_pct():
    # Large input, pct=True
    n = 1000
    values = ['a']*500 + ['b']*500
    cat = Categorical(values, categories=['a','b'], ordered=True)
    codeflash_output = cat._rank(method='average', pct=True); result = codeflash_output # 47.7μs -> 41.6μs (14.7% faster)

def test_large_descending():
    # Large input, descending order
    n = 1000
    values = ['a']*500 + ['b']*500
    cat = Categorical(values, categories=['a','b'], ordered=True)
    codeflash_output = cat._rank(method='average', ascending=False); result = codeflash_output # 48.8μs -> 43.3μs (12.7% faster)
# 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
from pandas.core.arrays.categorical import Categorical

# --------------------------
# Unit tests for _rank
# --------------------------

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

def test_rank_simple_ordered_average():
    # Ordered categorical, no missing, method=average
    cat = Categorical(['a', 'b', 'c', 'a', 'b'], categories=['a', 'b', 'c'], ordered=True)
    codeflash_output = cat._rank(); ranks = codeflash_output # 43.2μs -> 37.9μs (14.1% faster)

def test_rank_unordered_numeric_categories():
    # Unordered, numeric categories
    cat = Categorical([2, 1, 3, 2, 1], categories=[1, 2, 3], ordered=False)
    codeflash_output = cat._rank(); ranks = codeflash_output # 62.4μs -> 35.6μs (75.2% faster)

def test_rank_unordered_string_categories():
    # Unordered, string categories
    cat = Categorical(['b', 'a', 'c', 'b', 'a'], categories=['a', 'b', 'c'], ordered=False)
    codeflash_output = cat._rank(); ranks = codeflash_output # 201μs -> 201μs (0.100% faster)

def test_rank_ascending_false():
    # Ascending False
    cat = Categorical(['a', 'b', 'c', 'a', 'b'], categories=['a', 'b', 'c'], ordered=True)
    codeflash_output = cat._rank(ascending=False); ranks = codeflash_output # 42.5μs -> 37.7μs (12.7% faster)

def test_rank_method_min_max_first_dense():
    # method=min
    cat = Categorical(['a', 'a', 'b', 'c', 'b'], categories=['a', 'b', 'c'], ordered=True)
    codeflash_output = cat._rank(method='min'); r_min = codeflash_output # 41.4μs -> 36.0μs (15.0% faster)
    codeflash_output = cat._rank(method='max'); r_max = codeflash_output # 20.7μs -> 17.0μs (21.7% faster)
    codeflash_output = cat._rank(method='first'); r_first = codeflash_output # 17.1μs -> 13.3μs (28.9% faster)
    codeflash_output = cat._rank(method='dense'); r_dense = codeflash_output # 15.1μs -> 11.8μs (28.4% faster)

def test_rank_pct_true():
    cat = Categorical(['a', 'b', 'c', 'a', 'b'], categories=['a', 'b', 'c'], ordered=True)
    codeflash_output = cat._rank(pct=True); ranks = codeflash_output # 38.1μs -> 33.9μs (12.6% faster)

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

def test_rank_with_missing_values():
    # Missing values present
    cat = Categorical(['a', None, 'b', np.nan, 'a'], categories=['a', 'b'], ordered=True)
    codeflash_output = cat._rank(); ranks = codeflash_output # 44.4μs -> 37.8μs (17.5% faster)

def test_rank_na_option_top_bottom():
    cat = Categorical(['a', None, 'b', np.nan, 'a'], categories=['a', 'b'], ordered=True)
    codeflash_output = cat._rank(na_option='top'); r_top = codeflash_output # 44.8μs -> 38.4μs (16.5% faster)
    codeflash_output = cat._rank(na_option='bottom'); r_bottom = codeflash_output # 23.2μs -> 18.9μs (22.8% faster)

def test_rank_empty():
    cat = Categorical([], categories=['a', 'b'], ordered=True)
    codeflash_output = cat._rank(); ranks = codeflash_output # 38.4μs -> 33.4μs (15.1% faster)

def test_rank_all_missing():
    cat = Categorical([None, np.nan, None], categories=['a', 'b'], ordered=True)
    codeflash_output = cat._rank(); ranks = codeflash_output # 42.8μs -> 36.6μs (16.8% faster)

def test_rank_single_element():
    cat = Categorical(['a'], categories=['a', 'b'], ordered=True)
    codeflash_output = cat._rank(); ranks = codeflash_output # 38.2μs -> 33.9μs (12.7% faster)

def test_rank_not_implemented_axis():
    cat = Categorical(['a', 'b'], categories=['a', 'b'], ordered=True)
    with pytest.raises(NotImplementedError):
        cat._rank(axis=1) # 1.45μs -> 1.44μs (0.347% faster)

def test_rank_invalid_method():
    cat = Categorical(['a', 'b'], categories=['a', 'b'], ordered=True)
    with pytest.raises(ValueError):
        cat._rank(method='invalid')

def test_rank_invalid_na_option():
    cat = Categorical(['a', 'b'], categories=['a', 'b'], ordered=True)
    with pytest.raises(ValueError):
        cat._rank(na_option='invalid')

def test_rank_unordered_non_numeric_categories():
    # Unordered, non-numeric categories, e.g. strings
    cat = Categorical(['dog', 'cat', 'dog', 'bird'], categories=['bird', 'cat', 'dog'], ordered=False)
    codeflash_output = cat._rank(); ranks = codeflash_output # 205μs -> 206μs (0.596% slower)

def test_rank_with_duplicate_categories():
    # Categories with duplicates in values
    cat = Categorical(['a', 'b', 'a', 'c', 'c', 'b'], categories=['a', 'b', 'c'], ordered=True)
    codeflash_output = cat._rank(); ranks = codeflash_output # 41.6μs -> 36.4μs (14.5% faster)

def test_rank_with_unsorted_categories():
    # Categories not sorted
    cat = Categorical(['b', 'a', 'c'], categories=['c', 'b', 'a'], ordered=True)
    codeflash_output = cat._rank(); ranks = codeflash_output # 41.6μs -> 36.4μs (14.2% faster)

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

def test_rank_large_ordered():
    # Large ordered categorical
    vals = ['a'] * 300 + ['b'] * 400 + ['c'] * 300
    cat = Categorical(vals, categories=['a', 'b', 'c'], ordered=True)
    codeflash_output = cat._rank(); ranks = codeflash_output # 46.0μs -> 41.2μs (11.5% faster)

def test_rank_large_unordered_numeric():
    # Large unordered numeric
    vals = [i % 10 for i in range(1000)]
    cat = Categorical(vals, categories=list(range(10)), ordered=False)
    codeflash_output = cat._rank(); ranks = codeflash_output # 89.5μs -> 41.3μs (117% faster)
    # Each value appears 100 times: ranks for 0 = 50.5, for 1 = 150.5, ..., for 9 = 950.5
    for i in range(10):
        idx = [j for j, v in enumerate(vals) if v == i]
        expected = (i * 100 + 50.5)

def test_rank_large_with_missing():
    # Large with missing values
    vals = ['a'] * 500 + [None] * 100 + ['b'] * 400
    cat = Categorical(vals, categories=['a', 'b'], ordered=True)
    codeflash_output = cat._rank(); ranks = codeflash_output # 53.5μs -> 46.3μs (15.5% faster)

def test_rank_large_pct_true():
    vals = ['a'] * 500 + ['b'] * 500
    cat = Categorical(vals, categories=['a', 'b'], ordered=True)
    codeflash_output = cat._rank(pct=True); ranks = codeflash_output # 47.5μs -> 42.1μs (12.9% faster)

def test_rank_large_all_missing():
    vals = [None] * 1000
    cat = Categorical(vals, categories=['a', 'b'], ordered=True)
    codeflash_output = cat._rank(); ranks = codeflash_output # 52.2μs -> 46.3μs (12.7% faster)

def test_rank_large_all_same_category():
    vals = ['a'] * 1000
    cat = Categorical(vals, categories=['a', 'b'], ordered=True)
    codeflash_output = cat._rank(); ranks = codeflash_output # 46.3μs -> 40.6μs (14.0% 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-Categorical._rank-mhbwq18h and push.