提升pandas/dask自定义聚合速度及RollingGroupby first/last报错解决

1. 报错原因：Dask的RollingGroupby对象暂未实现Pandas原生的first/last字符串别名对应的内置方法，因此直接传"first"/"last"作为聚合函数会触发AttributeError。
2. 性能变慢原因：自定义的Python聚合函数属于UDF，每一个滚动窗口都需要触发一次Python层调用，受GIL限制无法利用向量化优化，速度比C实现的内置聚合慢1~2个数量级是正常现象。

最优方案：用内置nth方法替代自定义UDF（性能和原生内置函数一致） #

滚动窗口的nth(n)方法是C实现的内置方法，支持取窗口内第n个值，完全可以替代你自定义的first/last逻辑：

• first等价于nth(0)：取窗口第一个值
• last等价于nth(-1)：取窗口最后一个值
  你可以拆分聚合逻辑实现，完全不需要自定义函数：

# 1. 先计算sum/max/min这类原生支持的聚合
builtin_agg = {"volume":"sum", "high":"max", "low":"min"}
agg_builtin = df_tx_counts_input.map_overlap(
    lambda x: x.groupby(id_col_name)[data_cols_tmp].rolling(f"{rolling_window_now}ns").agg(builtin_agg).sort_index(level=1), 
    pd.Timedelta(f"{rolling_window_now}ns"), 
    0
)
# 2. 单独计算open的first和close的last
agg_first_last = df_tx_counts_input.map_overlap(
    lambda x: x.groupby(id_col_name)[['open', 'close']].rolling(f"{rolling_window_now}ns").nth([0, -1]).sort_index(level=1),
    pd.Timedelta(f"{rolling_window_now}ns"), 
    0
)
# 3. 合并两部分结果即可
final_df_tmp = agg_builtin.join(agg_first_last)

如果你的Dask版本支持在agg中直接传nth，也可以直接写聚合规则：

grand_agg_instruction[rolling_window_now] = {
    "volume":"sum",
    "open": lambda x: x.nth(0),
    "close": lambda x: x.nth(-1),
    "high":"max",
    "low":"min"
}

备选方案：Numba编译自定义UDF（性能接近原生） #

如果后续有更复杂的自定义聚合逻辑，可以用Numba的无Python模式编译函数，消除Python层调用开销：

import numba
import numpy as np

# 加装饰器开启无Python模式编译
@numba.jit(nopython=True, fastmath=True)
def first(arr: np.ndarray):
    return arr[0]

@numba.jit(nopython=True, fastmath=True)
def last(arr: np.ndarray):
    return arr[-1]

# 聚合规则里转成numpy数组传入函数，配合numba加速
agg_operations = {
    "volume":["sum"],
    "open": [lambda x: first(x.to_numpy())],
    "close": [lambda x: last(x.to_numpy())],
    "high":["max"],
    "low":["min"]
}

该方案首次运行会有一次编译开销，后续运行速度可比纯Python UDF快50~100倍，接近原生函数性能。

• 运行前必须确认数据已经按id_col_name和时间列升序排序，否则nth/自定义的first/last取值会出错
• 优先用nth方案，不需要额外依赖，性能最高；Numba方案适合需要复杂自定义逻辑的场景

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