⚡️ Speed up function aggregate_risk by 11%

[codeflash-ai]

📄 11% (0.11x) speedup for aggregate_risk in gs_quant/risk/core.py

⏱️ Runtime : 290 milliseconds → 261 milliseconds (best of 5 runs)

📝 Explanation and details

The optimized code achieves an 11% speedup through several key performance improvements:

Primary Optimizations:

1. Eliminated fillna(0) operation: The original code called pd.concat(dfs).fillna(0) which is expensive on large DataFrames. The optimization removes this and uses groupby().sum(min_count=1) instead, which handles NaN values correctly without the preprocessing overhead.

2. Replaced list comprehension with map(): Changed [get_df(r) for r in results] to list(map(get_df, results_list)). This reduces interpreter overhead in tight loops by avoiding repeated closure lookups.

3. Optimized threshold filtering: Changed result.value.abs() > threshold to result.value.abs().values > threshold, using numpy's underlying array for faster elementwise comparison instead of pandas Series operations.

4. Localized attribute lookups: Moved pd.DataFrame and Future.result to local variables to avoid repeated global namespace lookups in hot code paths.

Performance Impact by Test Type:

• Basic cases (2-3ms): 4-12% improvement from eliminating fillna overhead
• Large scale aggregation (20-30ms): Up to 47% improvement on tests with many unique keys, where the fillna elimination has maximum impact
• Threshold filtering: 8-16% improvement from numpy array operations
• Future handling: 4-7% improvement from localized attribute access

The optimizations are most effective for scenarios with large datasets and many grouping operations, which is typical in financial risk aggregation workflows. The changes maintain identical functionality while reducing pandas overhead in the most performance-critical code paths.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	✅ 22 Passed
🌀 Generated Regression Tests	✅ 34 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
api/test_risk.py::test_structured_calc	2.91ms	2.60ms	11.9%✅

🌀 Generated Regression Tests and Runtime

from concurrent.futures import Future
from typing import Callable, Iterable, List, Optional, Tuple, Union

# function to test (as provided)
import pandas as pd
# imports
import pytest
from gs_quant.risk.core import aggregate_risk


class DataFrameWithInfo:
    def __init__(self, df):
        self.raw_value = df
from gs_quant.risk.core import aggregate_risk

# unit tests
# Basic Test Cases

def test_single_dataframe_basic():
    # Test with a single DataFrameWithInfo, basic values
    df = pd.DataFrame([{
        'date': '2024-06-01',
        'time': '12:00',
        'mkt_type': 'TypeA',
        'mkt_asset': 'AssetA',
        'mkt_class': 'ClassA',
        'mkt_point': 'PointA',
        'value': 10.0
    }])
    obj = DataFrameWithInfo(df)
    codeflash_output = aggregate_risk([obj]); result = codeflash_output # 2.55ms -> 2.34ms (9.04% faster)

def test_multiple_dataframes_basic_sum():
    # Test with multiple DataFrameWithInfo, same bucket, values should sum
    df1 = pd.DataFrame([{
        'date': '2024-06-01',
        'time': '12:00',
        'mkt_type': 'TypeA',
        'mkt_asset': 'AssetA',
        'mkt_class': 'ClassA',
        'mkt_point': 'PointA',
        'value': 5.0
    }])
    df2 = pd.DataFrame([{
        'date': '2024-06-01',
        'time': '12:00',
        'mkt_type': 'TypeA',
        'mkt_asset': 'AssetA',
        'mkt_class': 'ClassA',
        'mkt_point': 'PointA',
        'value': 7.0
    }])
    objs = [DataFrameWithInfo(df1), DataFrameWithInfo(df2)]
    codeflash_output = aggregate_risk(objs); result = codeflash_output # 2.58ms -> 2.30ms (12.2% faster)

def test_multiple_dataframes_basic_different_buckets():
    # Test with multiple DataFrameWithInfo, different buckets
    df1 = pd.DataFrame([{
        'date': '2024-06-01',
        'time': '12:00',
        'mkt_type': 'TypeA',
        'mkt_asset': 'AssetA',
        'mkt_class': 'ClassA',
        'mkt_point': 'PointA',
        'value': 5.0
    }])
    df2 = pd.DataFrame([{
        'date': '2024-06-02',
        'time': '13:00',
        'mkt_type': 'TypeB',
        'mkt_asset': 'AssetB',
        'mkt_class': 'ClassB',
        'mkt_point': 'PointB',
        'value': 7.0
    }])
    objs = [DataFrameWithInfo(df1), DataFrameWithInfo(df2)]
    codeflash_output = aggregate_risk(objs); result = codeflash_output # 2.52ms -> 2.43ms (4.02% faster)

def test_threshold_basic():
    # Test threshold parameter
    df1 = pd.DataFrame([{
        'date': '2024-06-01',
        'time': '12:00',
        'mkt_type': 'TypeA',
        'mkt_asset': 'AssetA',
        'mkt_class': 'ClassA',
        'mkt_point': 'PointA',
        'value': 0.05
    }])
    df2 = pd.DataFrame([{
        'date': '2024-06-02',
        'time': '13:00',
        'mkt_type': 'TypeB',
        'mkt_asset': 'AssetB',
        'mkt_class': 'ClassB',
        'mkt_point': 'PointB',
        'value': 2.0
    }])
    objs = [DataFrameWithInfo(df1), DataFrameWithInfo(df2)]
    codeflash_output = aggregate_risk(objs, threshold=0.1); result = codeflash_output # 2.67ms -> 2.51ms (6.33% faster)

def test_future_basic():
    # Test with Future objects
    df = pd.DataFrame([{
        'date': '2024-06-01',
        'time': '12:00',
        'mkt_type': 'TypeA',
        'mkt_asset': 'AssetA',
        'mkt_class': 'ClassA',
        'mkt_point': 'PointA',
        'value': 10.0
    }])
    fut = Future()
    fut.set_result(DataFrameWithInfo(df))
    codeflash_output = aggregate_risk([fut]); result = codeflash_output # 2.36ms -> 2.26ms (4.15% faster)

# Edge Test Cases


def test_missing_columns():
    # Test with DataFrame missing some columns (should still aggregate on present columns)
    df = pd.DataFrame([{
        'date': '2024-06-01',
        'value': 10.0
    }])
    obj = DataFrameWithInfo(df)
    codeflash_output = aggregate_risk([obj]); result = codeflash_output # 1.31ms -> 1.20ms (9.74% faster)

def test_nan_values():
    # Test with NaN values in 'value' column
    df = pd.DataFrame([{
        'date': '2024-06-01',
        'time': '12:00',
        'mkt_type': 'TypeA',
        'mkt_asset': 'AssetA',
        'mkt_class': 'ClassA',
        'mkt_point': 'PointA',
        'value': float('nan')
    }])
    obj = DataFrameWithInfo(df)
    codeflash_output = aggregate_risk([obj]); result = codeflash_output # 2.39ms -> 2.24ms (6.62% faster)

def test_negative_values_and_threshold():
    # Test with negative values and threshold
    df1 = pd.DataFrame([{
        'date': '2024-06-01',
        'value': -0.05
    }])
    df2 = pd.DataFrame([{
        'date': '2024-06-02',
        'value': -2.0
    }])
    objs = [DataFrameWithInfo(df1), DataFrameWithInfo(df2)]
    codeflash_output = aggregate_risk(objs, threshold=0.1); result = codeflash_output # 1.63ms -> 1.41ms (15.3% faster)

def test_allow_heterogeneous_types_series():
    # Test with pd.Series and allow_heterogeneous_types=True
    class SeriesWithInfo(pd.Series):
        @property
        def raw_value(self):
            return self

    s = SeriesWithInfo({'date': '2024-06-01', 'value': 5.0})
    codeflash_output = aggregate_risk([s], allow_heterogeneous_types=True); result = codeflash_output # 1.51ms -> 1.42ms (6.24% faster)

def test_unsorted_input():
    # Test that output is sorted by risk columns
    df1 = pd.DataFrame([{
        'date': '2024-06-02',
        'time': '13:00',
        'mkt_type': 'TypeB',
        'mkt_asset': 'AssetB',
        'mkt_class': 'ClassB',
        'mkt_point': 'PointB',
        'value': 7.0
    }])
    df2 = pd.DataFrame([{
        'date': '2024-06-01',
        'time': '12:00',
        'mkt_type': 'TypeA',
        'mkt_asset': 'AssetA',
        'mkt_class': 'ClassA',
        'mkt_point': 'PointA',
        'value': 5.0
    }])
    objs = [DataFrameWithInfo(df1), DataFrameWithInfo(df2)]
    codeflash_output = aggregate_risk(objs); result = codeflash_output # 2.44ms -> 2.33ms (5.07% faster)

def test_zero_values_threshold():
    # Test with zero values and threshold
    df = pd.DataFrame([{
        'date': '2024-06-01',
        'value': 0.0
    }])
    obj = DataFrameWithInfo(df)
    codeflash_output = aggregate_risk([obj], threshold=0.1); result = codeflash_output # 1.42ms -> 1.31ms (8.90% faster)

# Large Scale Test Cases

def test_large_scale_sum():
    # Test aggregation over many rows (up to 1000)
    dfs = []
    for i in range(1000):
        df = pd.DataFrame([{
            'date': f'2024-06-{i%30+1:02d}',
            'time': '12:00',
            'mkt_type': 'TypeA',
            'mkt_asset': 'AssetA',
            'mkt_class': 'ClassA',
            'mkt_point': 'PointA',
            'value': 1.0
        }])
        dfs.append(DataFrameWithInfo(df))
    codeflash_output = aggregate_risk(dfs); result = codeflash_output # 20.7ms -> 20.0ms (3.46% faster)
    # Each date should have 1000/30 = 33 or 34 rows summed to that value
    for day in range(1, 31):
        date = f'2024-06-{day:02d}'
        count = sum(1 for i in range(1000) if i%30+1 == day)

def test_large_scale_threshold():
    # Test aggregation with threshold on large scale
    dfs = []
    for i in range(1000):
        val = 0.05 if i % 2 == 0 else 2.0
        df = pd.DataFrame([{
            'date': f'2024-06-{i%30+1:02d}',
            'value': val
        }])
        dfs.append(DataFrameWithInfo(df))
    codeflash_output = aggregate_risk(dfs, threshold=0.1); result = codeflash_output # 18.3ms -> 17.9ms (2.10% faster)
    # Only rows with value 2.0 should remain
    for day in range(1, 31):
        date = f'2024-06-{day:02d}'
        count = sum(1 for i in range(1000) if i%30+1 == day and 2.0 == (0.05 if i%2==0 else 2.0))
        if count:
            pass
        else:
            pass

def test_large_scale_future():
    # Test with many Future objects
    dfs = []
    for i in range(100):
        df = pd.DataFrame([{
            'date': f'2024-06-{i%10+1:02d}',
            'value': 1.0
        }])
        fut = Future()
        fut.set_result(DataFrameWithInfo(df))
        dfs.append(fut)
    codeflash_output = aggregate_risk(dfs); result = codeflash_output # 3.07ms -> 2.96ms (3.69% faster)
    # Each date should have 10 rows summed to 10.0
    for day in range(1, 11):
        date = f'2024-06-{day:02d}'

def test_large_scale_heterogeneous_types():
    # Test with mix of DataFrameWithInfo and Series, allow_heterogeneous_types=True
    class SeriesWithInfo(pd.Series):
        @property
        def raw_value(self):
            return self
    objs = []
    for i in range(500):
        df = pd.DataFrame([{
            'date': f'2024-06-{i%10+1:02d}',
            'value': 1.0
        }])
        objs.append(DataFrameWithInfo(df))
    for i in range(500):
        s = SeriesWithInfo({'date': f'2024-06-{i%10+1:02d}', 'value': 2.0})
        objs.append(s)
    codeflash_output = aggregate_risk(objs, allow_heterogeneous_types=True); result = codeflash_output # 73.7ms -> 72.0ms (2.29% faster)
    # Each date should have 1000/10 = 100 rows, 500 with value 1.0 and 500 with value 2.0, sum = 500*1 + 500*2 = 1500.0
    for day in range(1, 11):
        date = f'2024-06-{day:02d}'
        count1 = sum(1 for i in range(500) if i%10+1 == day)
        count2 = sum(1 for i in range(500) if i%10+1 == day)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from concurrent.futures import Future
from typing import Callable, Iterable, List, Optional, Tuple, Union

# function to test (copied verbatim from user block)
import pandas as pd
# imports
import pytest
from gs_quant.risk.core import aggregate_risk


class DataFrameWithInfo:
    def __init__(self, df):
        self.raw_value = df
from gs_quant.risk.core import aggregate_risk

# unit tests

# Helper to create a DataFrameWithInfo for testing
def make_dfw(**kwargs):
    # All risk columns plus 'value'
    cols = ['date', 'time', 'mkt_type', 'mkt_asset', 'mkt_class', 'mkt_point', 'value']
    row = {col: kwargs.get(col, None) for col in cols}
    df = pd.DataFrame([row])
    return DataFrameWithInfo(df)

# Helper to create a Future that returns a DataFrameWithInfo
def make_future(dfw):
    f = Future()
    f.set_result(dfw)
    return f

# 1. BASIC TEST CASES

def test_single_dataframe_basic():
    # Test aggregating a single DataFrameWithInfo
    dfw = make_dfw(date='2024-06-01', time='10:00', mkt_type='IR', mkt_asset='GBP', mkt_class='Swap', mkt_point='5y', value=100)
    codeflash_output = aggregate_risk([dfw]); result = codeflash_output # 3.01ms -> 2.86ms (5.40% faster)

def test_multiple_dataframes_basic():
    # Test aggregating two DataFrameWithInfo with different mkt_point
    dfw1 = make_dfw(date='2024-06-01', time='10:00', mkt_type='IR', mkt_asset='GBP', mkt_class='Swap', mkt_point='5y', value=100)
    dfw2 = make_dfw(date='2024-06-01', time='10:00', mkt_type='IR', mkt_asset='GBP', mkt_class='Swap', mkt_point='10y', value=200)
    codeflash_output = aggregate_risk([dfw1, dfw2]); result = codeflash_output # 2.60ms -> 2.56ms (1.37% faster)

def test_sum_on_duplicate_keys_basic():
    # Test aggregation sums values when all keys except value are the same
    dfw1 = make_dfw(date='2024-06-01', time='10:00', mkt_type='IR', mkt_asset='GBP', mkt_class='Swap', mkt_point='5y', value=100)
    dfw2 = make_dfw(date='2024-06-01', time='10:00', mkt_type='IR', mkt_asset='GBP', mkt_class='Swap', mkt_point='5y', value=50)
    codeflash_output = aggregate_risk([dfw1, dfw2]); result = codeflash_output # 2.49ms -> 2.33ms (6.86% faster)

def test_future_basic():
    # Test that a Future containing a DataFrameWithInfo is handled
    dfw = make_dfw(date='2024-06-01', time='10:00', mkt_type='IR', mkt_asset='GBP', mkt_class='Swap', mkt_point='5y', value=123)
    fut = make_future(dfw)
    codeflash_output = aggregate_risk([fut]); result = codeflash_output # 2.40ms -> 2.24ms (7.22% faster)

def test_threshold_basic():
    # Test threshold filters out values below threshold
    dfw1 = make_dfw(date='2024-06-01', time='10:00', mkt_type='IR', mkt_asset='GBP', mkt_class='Swap', mkt_point='5y', value=0.05)
    dfw2 = make_dfw(date='2024-06-01', time='10:00', mkt_type='IR', mkt_asset='GBP', mkt_class='Swap', mkt_point='10y', value=0.2)
    codeflash_output = aggregate_risk([dfw1, dfw2], threshold=0.1); result = codeflash_output # 2.79ms -> 2.57ms (8.66% faster)

def test_threshold_negative_values_basic():
    # Test threshold with negative values
    dfw1 = make_dfw(date='2024-06-01', time='10:00', mkt_type='IR', mkt_asset='GBP', mkt_class='Swap', mkt_point='5y', value=-0.05)
    dfw2 = make_dfw(date='2024-06-01', time='10:00', mkt_type='IR', mkt_asset='GBP', mkt_class='Swap', mkt_point='10y', value=-0.2)
    codeflash_output = aggregate_risk([dfw1, dfw2], threshold=0.1); result = codeflash_output # 2.79ms -> 2.49ms (12.1% faster)

# 2. EDGE TEST CASES


def test_all_filtered_by_threshold_edge():
    # All values below threshold, should return empty DataFrame
    dfw1 = make_dfw(date='2024-06-01', time='10:00', mkt_type='IR', mkt_asset='GBP', mkt_class='Swap', mkt_point='5y', value=0.05)
    dfw2 = make_dfw(date='2024-06-01', time='10:00', mkt_type='IR', mkt_asset='GBP', mkt_class='Swap', mkt_point='10y', value=-0.09)
    codeflash_output = aggregate_risk([dfw1, dfw2], threshold=0.1); result = codeflash_output # 2.61ms -> 2.49ms (4.86% faster)

def test_missing_columns_edge():
    # DataFrame missing some risk columns
    df = pd.DataFrame([{'date': '2024-06-01', 'value': 42}])
    dfw = DataFrameWithInfo(df)
    codeflash_output = aggregate_risk([dfw]); result = codeflash_output # 1.27ms -> 1.23ms (3.25% faster)

def test_nan_values_edge():
    # DataFrame with NaN value, should be filled with 0
    df = pd.DataFrame([{'date': '2024-06-01', 'value': float('nan')}])
    dfw = DataFrameWithInfo(df)
    codeflash_output = aggregate_risk([dfw]); result = codeflash_output # 1.27ms -> 1.19ms (6.81% faster)

def test_allow_heterogeneous_types_series_edge():
    # Test allow_heterogeneous_types with pd.Series
    class SeriesWithInfo(pd.Series):
        @property
        def raw_value(self):
            return self

    s = SeriesWithInfo({'date': '2024-06-01', 'value': 99})
    codeflash_output = aggregate_risk([s], allow_heterogeneous_types=True); result = codeflash_output # 1.54ms -> 1.45ms (6.28% faster)

def test_heterogeneous_types_not_allowed_edge():
    # Should raise if pd.Series is passed and allow_heterogeneous_types is False
    class SeriesWithInfo(pd.Series):
        @property
        def raw_value(self):
            return self

    s = SeriesWithInfo({'date': '2024-06-01', 'value': 99})
    with pytest.raises(AttributeError):
        aggregate_risk([s], allow_heterogeneous_types=False) # 206μs -> 140μs (47.0% faster)

def test_sorting_order_edge():
    # Test that sorting is correct for mkt_point using point_sort_order
    dfw1 = make_dfw(date='2024-06-01', time='10:00', mkt_type='IR', mkt_asset='GBP', mkt_class='Swap', mkt_point='10y', value=1)
    dfw2 = make_dfw(date='2024-06-01', time='10:00', mkt_type='IR', mkt_asset='GBP', mkt_class='Swap', mkt_point='2y', value=2)
    dfw3 = make_dfw(date='2024-06-01', time='10:00', mkt_type='IR', mkt_asset='GBP', mkt_class='Swap', mkt_point='5y', value=3)
    codeflash_output = aggregate_risk([dfw1, dfw2, dfw3]); result = codeflash_output # 2.51ms -> 2.48ms (1.34% faster)

def test_multiple_dates_edge():
    # Test aggregation with multiple dates
    dfw1 = make_dfw(date='2024-06-01', value=1)
    dfw2 = make_dfw(date='2024-06-02', value=2)
    codeflash_output = aggregate_risk([dfw1, dfw2]); result = codeflash_output # 2.47ms -> 2.26ms (9.50% faster)

def test_large_number_of_duplicate_keys_edge():
    # Test aggregation with many duplicate keys
    dfws = [make_dfw(date='2024-06-01', mkt_point='5y', value=1) for _ in range(10)]
    codeflash_output = aggregate_risk(dfws); result = codeflash_output # 2.68ms -> 2.32ms (15.1% faster)

# 3. LARGE SCALE TEST CASES

def test_large_scale_aggregation():
    # Test with 1000 DataFrameWithInfo objects with different mkt_point
    dfws = [make_dfw(date='2024-06-01', mkt_point=f'{i}y', value=i) for i in range(1000)]
    codeflash_output = aggregate_risk(dfws); result = codeflash_output # 29.1ms -> 19.7ms (47.7% faster)

def test_large_scale_duplicate_keys():
    # Test with 1000 DataFrameWithInfo objects, all with same keys
    dfws = [make_dfw(date='2024-06-01', mkt_point='5y', value=1) for _ in range(1000)]
    codeflash_output = aggregate_risk(dfws); result = codeflash_output # 20.2ms -> 18.1ms (11.2% faster)

def test_large_scale_threshold():
    # Test with 1000 DataFrameWithInfo, half above threshold, half below
    dfws = [make_dfw(date='2024-06-01', mkt_point=f'{i}y', value=i) for i in range(1000)]
    codeflash_output = aggregate_risk(dfws, threshold=500); result = codeflash_output # 22.4ms -> 19.2ms (16.8% faster)
    # Only values > 500 should remain
    expected_points = [f'{i}y' for i in range(501, 1000)]

def test_large_scale_multiple_dates():
    # Test with 1000 DataFrameWithInfo, 10 unique dates, 100 per date
    dfws = []
    for d in range(10):
        date = f'2024-06-{d+1:02d}'
        for i in range(100):
            dfws.append(make_dfw(date=date, mkt_point=f'{i}y', value=i))
    codeflash_output = aggregate_risk(dfws); result = codeflash_output # 22.2ms -> 18.8ms (17.9% faster)

def test_large_scale_future():
    # Test with 1000 Future objects
    dfws = [make_dfw(date='2024-06-01', mkt_point=f'{i}y', value=i) for i in range(1000)]
    futs = [make_future(dfw) for dfw in dfws]
    codeflash_output = aggregate_risk(futs); result = codeflash_output # 23.9ms -> 19.5ms (22.7% 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-aggregate_risk-mhayhrxk and push.