如何优化Python Polars指标计算代码的运行效率？

问题背景 #

我有一个Python的IndicatorManager类，接收Polars DataFrame并计算60个指标（部分基于TA-Lib实现，其余基于Polars代码开发）。发现当以210行数据为限制运行时，受限数据的计算结果与全量历史数据的结果存在差异。为解决此问题，我尝试以210行窗口滑动计算并提取每个窗口的最后一行，但当前实现效率极低，需要优化。

数据为通过python-binance库从Binance API获取的加密货币K线数据（例如BTCUSDT的5分钟线，实际包含751583行）。

现有代码分析 #

当前run函数实现 #

def run(self, data: pl.DataFrame) -> pl.DataFrame:
    return pl.concat(
        [
            self._calculate(data.slice(i, 210)).tail(-1)
            for i in range(len(data) - 209)
        ]
    )

该实现通过循环逐个切片210行数据计算，每个窗口与前一个窗口有209行重叠，导致大量重复计算，效率极低。

_calculate方法实现 #

def _calculate(self, data: pl.DataFrame) -> pl.DataFrame:
    indicators = [
        self._ao, self._cci, self._cfo, self._macd, self._rsi, self._tsi, self._uo, self._willr,
        *[lambda d: self._ema(d, p) for p in (10, 20, 50, 100, 200)],
        *[lambda d: self._sma(d, p) for p in (10, 20, 50, 100, 200)], self._supertrend,
        *[lambda d: self._wma(d, p) for p in (10, 20, 50, 100, 200)],
        self._adx, self._aroon, self._bbands, self._cmf, self._mfi, self._obv, self._kama, 
        self._ppo, self._pvo, self._roc, self._stochrsi, self._stoch, self._ad, self._eom, 
        self._efi, self._nvi, self._pvt, self._vwap, self._donchian, self._kc, self._ui, 
        self._ichimoku, self._kst, self._massi, self._psar, self._stc, self._trix, self._vortex, 
        self._cmo, self._wvf, self._crsi, self._dmi, self._ttm, self._vwmacd, self._ci,
        self._ft, self._atr
    ]

    for indicator in indicators:
        data = indicator(data)

    return data

逐个调用指标函数并生成新的DataFrame，存在不必要的内存复制开销。

示例指标与静态方法 #

def _ao(self, data: pl.DataFrame) -> pl.DataFrame:
    fast = 5
    slow = 34

    median_price = 0.5 * (data["high"] + data["low"])
    return data.with_columns(
        (self.__sma(median_price, fast) - self.__sma(median_price, slow)).alias(f"AO_{fast}_{slow}")
    )

@staticmethod
def __sma(
        close: pl.Series,
        length: int
) -> pl.Series:
    return talib.SMA(close, length)

失败的LazyDataFrame尝试 #

def run(self, data: pl.DataFrame) -> pl.DataFrame:
    window_size = 210
    
    data = data.with_row_index()
    
    result = (
        data
        .lazy()
        .filter(pl.col("index") >= window_size - 1)
        .with_columns(
            (pl.col("index") - (window_size - 1)).alias("window_start")
        )
        .group_by("window_start", maintain_order=True)
        .apply(lambda df: self._calculate(
            data.slice(int(df["window_start"][0]), window_size)
        ).tail(1))
    ).collect()

    return result

Polars LazyDataFrame的group_by不支持apply方法，导致该方案无法运行。

优化方案 #

1. 利用滚动计算替代逐窗口切片 #

核心思路是避免重复计算重叠窗口数据，直接在全量数据上使用Polars或TA-Lib的滚动计算能力，最后只保留窗口完整的行（从第210行开始）。

2. 重构指标计算逻辑，批量生成列 #

将逐个调用指标函数改为一次性生成所有指标的表达式，通过with_columns批量添加，减少DataFrame复制开销。

3. 优化TA-Lib指标计算 #

TA-Lib的多数函数支持直接传入全量序列和窗口长度，自动返回滚动计算结果，无需逐窗口切片。

4. 自定义指标改用Polars窗口函数 #

对于非TA-Lib实现的指标，用Polars原生窗口函数（如rolling_mean）替代，利用Polars的向量化运算提升效率。

优化后的代码示例 #

优化后的run函数 #

def run(self, data: pl.DataFrame) -> pl.DataFrame:
    # 直接在全量数据上计算所有指标
    calculated_data = self._calculate(data)
    # 只保留窗口完整的行（从第210行开始，索引从0开始即slice(209)）
    return calculated_data.slice(209)

重构后的_calculate方法 #

def _calculate(self, data: pl.DataFrame) -> pl.DataFrame:
    exprs = []
    
    # AO指标：用Polars滚动均值替代TA-Lib，实现全量滚动计算
    median_price = 0.5 * (pl.col("high") + pl.col("low"))
    exprs.append(
        (median_price.rolling_mean(window_size=5) - median_price.rolling_mean(window_size=34)).alias("AO_5_34")
    )
    
    # TA-Lib SMA指标：直接传入全量序列计算滚动值
    close_series = data["close"]
    for period in (10, 20, 50, 100, 200):
        exprs.append(
            pl.Series(talib.SMA(close_series, period)).alias(f"SMA_{period}")
        )
    
    # EMA指标同理，用TA-Lib直接计算全量滚动值
    for period in (10, 20, 50, 100, 200):
        exprs.append(
            pl.Series(talib.EMA(close_series, period)).alias(f"EMA_{period}")
        )
    
    # 其他指标按此逻辑批量添加表达式
    # ...（补充剩余50+个指标的表达式）
    
    # 一次性添加所有指标列，减少内存复制
    return data.with_columns(exprs)

特殊指标处理（如Supertrend） #

对于需要递进计算的指标（如Supertrend），可使用Polars的rolling_apply或自定义函数实现滚动窗口计算：

def _supertrend(self, data: pl.DataFrame, period=10, multiplier=3) -> pl.Expr:
    def calculate_supertrend(window_df: pl.DataFrame) -> float:
        # 实现Supertrend的计算逻辑，返回窗口最后一行的结果
        atr = talib.ATR(window_df["high"], window_df["low"], window_df["close"], period)[-1]
        basic_upper = (window_df["high"] + window_df["low"]) / 2 + multiplier * atr
        basic_lower = (window_df["high"] + window_df["low"]) / 2 - multiplier * atr
        # 完成剩余Supertrend逻辑计算
        # ...
        return final_supertrend_value
    
    # 使用rolling_apply对全量数据滚动计算
    return pl.struct(["high", "low", "close"]).rolling_apply(
        window_size=210,
        function=calculate_supertrend,
        dtype=pl.Float64
    ).alias("Supertrend")

注意：rolling_apply的性能略低于原生窗口函数，但远优于逐窗口切片计算。

优化效果说明 #

• 避免了重叠窗口的重复计算，计算量从O(N*210)降至O(N)
• 利用Polars向量化运算和TA-Lib的底层优化，大幅提升计算速度
• 减少DataFrame复制次数，降低内存开销

内容的提问来源于stack exchange，提问作者YenilmezHunt