⚡️ Speed up function _generate_range_overflow_safe_signed by 91%

[codeflash-ai]

📄 91% (0.91x) speedup for _generate_range_overflow_safe_signed in pandas/core/arrays/_ranges.py

⏱️ Runtime : 401 microseconds → 210 microseconds (best of 68 runs)

📝 Explanation and details

The optimization achieves a 90% speedup by eliminating the expensive np.errstate(over="raise") context manager from the common path and reducing NumPy scalar operations.

Key optimizations:

1. Removed expensive error context: The original code wraps the main computation in np.errstate(over="raise"), which adds significant overhead (597,801ns vs 0ns in optimized version). The optimized version performs arithmetic with Python's native int type first, then uses np.int64() conversion to detect overflow.

2. Reduced NumPy scalar operations: Instead of computing np.int64(periods) * np.int64(stride) (283,455ns), the optimized version uses int(periods) * int(stride) (40,064ns) - a 7x improvement. Python's arbitrary-precision integers handle the multiplication efficiently without NumPy overhead.

3. Streamlined overflow detection: The optimized version converts the final result to np.int64 once for overflow checking, rather than creating multiple NumPy scalars during computation.

Performance impact: This function is called from _generate_range_overflow_safe, which is part of pandas' range generation machinery for datetime/timedelta arrays. The function_references show it's used in overflow-safe range calculations, making this optimization valuable for date range operations that are common in time series processing.

Test case benefits: The optimization shows consistent 80-110% speedups across all test cases, with particularly strong performance on basic cases (the most common usage patterns) and large-scale operations. Simple operations like zero periods/stride benefit most since they avoid the expensive NumPy context manager entirely.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

# imports
import pytest
from pandas.core.arrays._ranges import _generate_range_overflow_safe_signed


# Mocks for pandas._libs.lib.i8max, pandas._libs.tslibs.OutOfBoundsDatetime, pandas._libs.tslibs.iNaT
class OutOfBoundsDatetime(Exception):
    pass


i8max = 2**63 - 1
iNaT = -9223372036854775808  # pandas._libs.tslibs.iNaT

# unit tests

# 1. Basic Test Cases


def test_basic_positive_stride_start():
    # Simple case: stride=1, periods=10, endpoint=0, side='start'
    codeflash_output = _generate_range_overflow_safe_signed(
        0, 10, 1, "start"
    )  # 11.6μs -> 6.11μs (90.6% faster)


def test_basic_positive_stride_end():
    # Simple case: stride=1, periods=10, endpoint=0, side='end'
    codeflash_output = _generate_range_overflow_safe_signed(
        0, 10, 1, "end"
    )  # 9.84μs -> 5.01μs (96.2% faster)


def test_basic_negative_stride_start():
    # Simple case: stride=-2, periods=5, endpoint=10, side='start'
    codeflash_output = _generate_range_overflow_safe_signed(
        10, 5, -2, "start"
    )  # 8.78μs -> 4.44μs (97.7% faster)


def test_basic_negative_stride_end():
    # Simple case: stride=-2, periods=5, endpoint=10, side='end'
    codeflash_output = _generate_range_overflow_safe_signed(
        10, 5, -2, "end"
    )  # 8.73μs -> 4.24μs (106% faster)


def test_basic_zero_periods():
    # periods=0 should always return endpoint
    codeflash_output = _generate_range_overflow_safe_signed(
        123, 0, 10, "start"
    )  # 8.66μs -> 4.30μs (102% faster)
    codeflash_output = _generate_range_overflow_safe_signed(
        -456, 0, -10, "end"
    )  # 3.14μs -> 1.71μs (83.3% faster)


def test_basic_stride_zero():
    # stride=0 should always return endpoint
    codeflash_output = _generate_range_overflow_safe_signed(
        100, 5, 0, "start"
    )  # 7.70μs -> 3.94μs (95.5% faster)
    codeflash_output = _generate_range_overflow_safe_signed(
        -100, 5, 0, "end"
    )  # 3.00μs -> 1.67μs (79.4% faster)


def test_basic_one_period():
    # periods=1 should return endpoint + stride or endpoint - stride
    codeflash_output = _generate_range_overflow_safe_signed(
        50, 1, 5, "start"
    )  # 7.68μs -> 3.73μs (106% faster)
    codeflash_output = _generate_range_overflow_safe_signed(
        50, 1, 5, "end"
    )  # 2.98μs -> 1.54μs (92.9% faster)


def test_basic_side_assertion():
    # Invalid side should raise AssertionError
    with pytest.raises(AssertionError):
        _generate_range_overflow_safe_signed(
            0, 1, 1, "foo"
        )  # 868ns -> 919ns (5.55% slower)


# 2. Edge Test Cases


def test_edge_int64_max_start():
    # endpoint at int64 max, stride positive, periods=0
    codeflash_output = _generate_range_overflow_safe_signed(
        i8max, 0, 1, "start"
    )  # 10.5μs -> 5.86μs (80.1% faster)


def test_edge_int64_max_end():
    # endpoint at int64 max, stride positive, periods=1, side='end'
    codeflash_output = _generate_range_overflow_safe_signed(
        i8max, 1, 1, "end"
    )  # 9.00μs -> 4.77μs (88.5% faster)


def test_edge_int64_min_end():
    # endpoint at int64 min, stride negative, periods=1, side='end'
    codeflash_output = _generate_range_overflow_safe_signed(
        iNaT, 1, -1, "end"
    )  # 13.8μs -> 7.75μs (78.1% faster)


def test_edge_stride_opposite_sign_endpoint():
    # If stride and endpoint have opposite signs, should not overflow
    # Should not raise
    codeflash_output = _generate_range_overflow_safe_signed(
        -10, 2, 5, "start"
    )  # 13.5μs -> 7.21μs (86.7% faster)
    codeflash_output = _generate_range_overflow_safe_signed(
        10, 2, -5, "start"
    )  # 2.80μs -> 1.51μs (85.4% faster)


def test_edge_stride_zero_large_endpoint():
    # stride=0, large endpoint, should just return endpoint
    codeflash_output = _generate_range_overflow_safe_signed(
        i8max, 10, 0, "start"
    )  # 8.48μs -> 4.46μs (90.3% faster)


def test_edge_just_over_i8max():
    # endpoint + addend just overflows i8max, but within i8max + stride
    endpoint = i8max - 1
    stride = 2
    periods = 1
    # uresult = i8max - 1 + 2 = i8max + 1
    codeflash_output = _generate_range_overflow_safe_signed(
        endpoint, periods, stride, "start"
    )  # 13.9μs -> 8.09μs (72.2% faster)


def test_edge_just_under_i8max():
    # endpoint + addend just below i8max, should not overflow
    endpoint = i8max - 2
    stride = 1
    periods = 1
    codeflash_output = _generate_range_overflow_safe_signed(
        endpoint, periods, stride, "start"
    )  # 8.52μs -> 4.14μs (106% faster)


def test_edge_negative_stride_large_negative_endpoint():
    # Large negative endpoint, negative stride, but no overflow
    endpoint = -(2**62)
    stride = -2
    periods = 2
    codeflash_output = _generate_range_overflow_safe_signed(
        endpoint, periods, stride, "start"
    )  # 8.47μs -> 4.31μs (96.7% faster)


# 3. Large Scale Test Cases


def test_large_scale_positive_stride():
    # Large periods, positive stride, positive endpoint
    endpoint = 100
    stride = 5
    periods = 1000
    expected = endpoint + stride * periods
    codeflash_output = _generate_range_overflow_safe_signed(
        endpoint, periods, stride, "start"
    )  # 8.34μs -> 4.10μs (103% faster)


def test_large_scale_negative_stride():
    # Large periods, negative stride, negative endpoint
    endpoint = -100
    stride = -5
    periods = 1000
    expected = endpoint + stride * periods
    codeflash_output = _generate_range_overflow_safe_signed(
        endpoint, periods, stride, "start"
    )  # 8.01μs -> 4.10μs (95.2% faster)


def test_large_scale_end_side():
    # Large periods, positive stride, side='end'
    endpoint = 100
    stride = 5
    periods = 1000
    expected = endpoint - stride * periods
    codeflash_output = _generate_range_overflow_safe_signed(
        endpoint, periods, stride, "end"
    )  # 8.73μs -> 4.31μs (102% faster)


def test_large_scale_large_numbers_no_overflow():
    # Large numbers, but within int64 bounds
    endpoint = 2**62
    stride = 1
    periods = 1000
    expected = endpoint + stride * periods
    codeflash_output = _generate_range_overflow_safe_signed(
        endpoint, periods, stride, "start"
    )  # 8.57μs -> 4.00μs (114% faster)

from __future__ import annotations

# imports
import pytest
from pandas.core.arrays._ranges import _generate_range_overflow_safe_signed

# Minimal stand-in for pandas._libs.lib.i8max and pandas._libs.tslibs.iNaT, OutOfBoundsDatetime
i8max = 2**63 - 1
iNaT = -(2**63)


class OutOfBoundsDatetime(Exception):
    pass


# unit tests

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


def test_basic_start_positive_stride():
    # Simple case: start at 10, 5 periods, stride 2, side 'start'
    # Should return 10 + 5*2 = 20
    codeflash_output = _generate_range_overflow_safe_signed(
        10, 5, 2, "start"
    )  # 13.4μs -> 7.25μs (85.5% faster)


def test_basic_start_negative_stride():
    # Simple case: start at 10, 5 periods, stride -2, side 'start'
    # Should return 10 + 5*(-2) = 0
    codeflash_output = _generate_range_overflow_safe_signed(
        10, 5, -2, "start"
    )  # 9.35μs -> 4.61μs (103% faster)


def test_basic_end_positive_stride():
    # Simple case: end at 20, 5 periods, stride 2, side 'end'
    # Should return 20 + 5*(-2) = 10
    codeflash_output = _generate_range_overflow_safe_signed(
        20, 5, 2, "end"
    )  # 9.10μs -> 4.50μs (102% faster)


def test_basic_end_negative_stride():
    # Simple case: end at 0, 5 periods, stride -2, side 'end'
    # Should return 0 + 5*2 = 10
    codeflash_output = _generate_range_overflow_safe_signed(
        0, 5, -2, "end"
    )  # 8.81μs -> 4.19μs (110% faster)


def test_basic_zero_stride():
    # Stride of zero: should always return endpoint
    codeflash_output = _generate_range_overflow_safe_signed(
        123, 5, 0, "start"
    )  # 8.56μs -> 4.28μs (99.8% faster)
    codeflash_output = _generate_range_overflow_safe_signed(
        -456, 10, 0, "end"
    )  # 3.21μs -> 1.77μs (81.3% faster)


def test_basic_zero_periods():
    # Zero periods: should always return endpoint
    codeflash_output = _generate_range_overflow_safe_signed(
        789, 0, 3, "start"
    )  # 7.93μs -> 4.24μs (87.1% faster)
    codeflash_output = _generate_range_overflow_safe_signed(
        -321, 0, -7, "end"
    )  # 3.06μs -> 1.65μs (84.8% faster)


def test_basic_negative_periods():
    # Negative periods: should work as negative multiplication
    # start at 10, -2 periods, stride 3: 10 + (-2)*3 = 4
    codeflash_output = _generate_range_overflow_safe_signed(
        10, -2, 3, "start"
    )  # 7.55μs -> 4.05μs (86.7% faster)


def test_basic_side_assertion():
    # Invalid side string should assert
    with pytest.raises(AssertionError):
        _generate_range_overflow_safe_signed(
            10, 5, 2, "middle"
        )  # 880ns -> 918ns (4.14% slower)


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


def test_edge_i8max_no_overflow():
    # Should not overflow: endpoint at i8max, stride 0, periods 0
    codeflash_output = _generate_range_overflow_safe_signed(
        i8max, 0, 0, "start"
    )  # 10.6μs -> 5.86μs (80.4% faster)


def test_edge_i8max_negative_stride():
    # Should not overflow: endpoint at i8max, stride -1, periods 1
    # i8max + 1*-1 = i8max - 1
    codeflash_output = _generate_range_overflow_safe_signed(
        i8max, 1, -1, "start"
    )  # 13.5μs -> 7.51μs (80.2% faster)


def test_edge_negative_endpoint_positive_stride():
    # endpoint negative, stride positive, sum does not overflow
    codeflash_output = _generate_range_overflow_safe_signed(
        -100, 5, 10, "start"
    )  # 13.6μs -> 7.47μs (81.9% faster)


def test_edge_large_stride_small_periods():
    # Large stride, small periods, no overflow
    codeflash_output = _generate_range_overflow_safe_signed(
        100, 2, 2**30, "start"
    )  # 13.6μs -> 7.62μs (78.8% faster)


def test_edge_large_negative_stride_small_periods():
    # Large negative stride, small periods, no overflow
    codeflash_output = _generate_range_overflow_safe_signed(
        -100, 2, -(2**30), "start"
    )  # 9.68μs -> 4.91μs (97.2% faster)


def test_edge_large_negative_periods():
    # Large negative periods, negative stride
    codeflash_output = _generate_range_overflow_safe_signed(100, -10, -2, "start")
    result = codeflash_output  # 13.6μs -> 7.35μs (85.5% faster)


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


def test_large_scale_many_periods():
    # Large number of periods, moderate stride, should not overflow
    endpoint = 0
    periods = 999
    stride = 1000000
    expected = endpoint + periods * stride
    codeflash_output = _generate_range_overflow_safe_signed(
        endpoint, periods, stride, "start"
    )  # 9.38μs -> 4.41μs (112% faster)


def test_large_scale_large_negative_periods():
    # Large negative periods, negative stride, should not overflow
    endpoint = 1000000
    periods = -999
    stride = -1000000
    expected = endpoint + periods * stride
    codeflash_output = _generate_range_overflow_safe_signed(
        endpoint, periods, stride, "start"
    )  # 8.65μs -> 4.33μs (99.7% faster)


def test_large_scale_max_stride():
    # Stride is very large, but periods is 1, should not overflow
    endpoint = 0
    stride = i8max
    periods = 1
    expected = endpoint + stride * periods
    codeflash_output = _generate_range_overflow_safe_signed(
        endpoint, periods, stride, "start"
    )  # 13.4μs -> 7.44μs (80.5% faster)


def test_large_scale_max_negative_stride():
    # Stride is very large negative, periods is 1, should not overflow
    endpoint = 0
    stride = -i8max
    periods = 1
    expected = endpoint + stride * periods
    codeflash_output = _generate_range_overflow_safe_signed(
        endpoint, periods, stride, "start"
    )  # 9.53μs -> 4.88μs (95.4% faster)


def test_large_scale_high_periods_small_stride():
    # Large periods, small stride, does not overflow
    endpoint = 0
    stride = 1
    periods = 999
    expected = endpoint + stride * periods
    codeflash_output = _generate_range_overflow_safe_signed(
        endpoint, periods, stride, "start"
    )  # 8.87μs -> 4.17μs (113% faster)


def test_large_scale_high_periods_negative_stride():
    # Large periods, negative stride, does not overflow
    endpoint = 0
    stride = -1
    periods = 999
    expected = endpoint + stride * periods
    codeflash_output = _generate_range_overflow_safe_signed(
        endpoint, periods, stride, "start"
    )  # 8.96μs -> 4.39μs (104% 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-_generate_range_overflow_safe_signed-mir42oii and push.