cuDF apply_rows UDF计算Market Profile报TypingError

报错根因 #

• apply_rows 是逐行独立执行的UDF接口，传入内核的是单行列标量值，不存在行索引、shift() 类行操作方法，直接对标量做索引取值、调用shift必然触发类型错误
• Numba nopython/CUDA编译模式不支持Python动态特性：动态列表append、pandas对象构造、query查询这类解释执行的逻辑无法被编译，所有变量必须提前明确类型、固定长度
• datetime64类型无需做字符串/数值转换，直接在cuDF高层API做相邻行比较即可，不需要在内核中处理datetime运算

具体实现步骤 #

1. 预处理生成状态标记列 #

不要在内核中做前后行日期对比，提前在cuDF侧完成所有行级判断，把标记列直接传入内核：

import cudf
import numpy as np
from numba import jit

# 生成新交易日标记：当前行日期不等于上一行即为当日首行，首行默认标记为新交易日
df["is_new_day"] = (df["Date"] != df["Date"].shift(1)).fillna(True)

# 提前计算全局价格区间，按需求可替换为按日分组计算当日min/max实现动态区间
price_min = df["Close"].min()
price_max = df["Close"].max()
bin_width = (price_max - price_min) / 10

2. 初始化固定结构的输出列 #

提前创建30个输出列对应10个价格区间的min、max、count统计值，避免内核中动态创建列、动态扩列：

for bin_idx in range(10):
    df[f"bin_{bin_idx}_min"] = np.nan
    df[f"bin_{bin_idx}_max"] = np.nan
    df[f"bin_{bin_idx}_count"] = 0

3. 编写符合Numba编译要求的处理内核 #

禁止使用动态列表、pandas API，所有临时统计状态用固定长度numpy数组存储，通过预处理的is_new_day标记判断是否重置状态：

@jit(nopython=True)
def market_profile_kernel(is_new_day_arr, close_arr, *output_arrs):
    # 初始化固定长度的统计状态数组，长度硬编码为10对应10个价格区间
    bin_min = np.full(10, np.nan, dtype=np.float64)
    bin_max = np.full(10, np.nan, dtype=np.float64)
    bin_count = np.zeros(10, dtype=np.int32)

    row_count = len(is_new_day_arr)
    for i in range(row_count):
        # 新交易日重置所有统计状态
        if is_new_day_arr[i]:
            for b in range(10):
                bin_min[b] = np.nan
                bin_max[b] = np.nan
                bin_count[b] = 0
        
        current_price = close_arr[i]
        # 计算当前价格所属区间，做边界修正避免越界
        current_bin = int((current_price - price_min) // bin_width)
        current_bin = max(0, min(9, current_bin))

        # 更新对应区间统计值
        bin_count[current_bin] += 1
        if np.isnan(bin_min[current_bin]) or current_price < bin_min[current_bin]:
            bin_min[current_bin] = current_price
        if np.isnan(bin_max[current_bin]) or current_price > bin_max[current_bin]:
            bin_max[current_bin] = current_price
        
        # 将当前所有区间的统计值写入输出列
        for b in range(10):
            out_pos = b * 3
            output_arrs[out_pos][i] = bin_min[b]
            output_arrs[out_pos+1][i] = bin_max[b]
            output_arrs[out_pos+2][i] = bin_count[b]

4. 调用接口执行计算 #

有状态逐行累计逻辑不要使用并行的apply_rows，改用apply_chunks按数据块顺序串行处理，保证状态按行序正确更新：

# 整理输入输出列映射
input_cols = ["is_new_day", "Close"]
output_cols = []
for b in range(10):
    output_cols.extend([f"bin_{b}_min", f"bin_{b}_max", f"bin_{b}_count"])

# 执行计算
df = df.apply_chunks(
    market_profile_kernel,
    incols=input_cols,
    outcols=output_cols,
    kwargs={}
)

关键注意事项 #

• 所有跨行判断逻辑（比如日期切换）全部在cuDF高层API完成，只把布尔类型的标记列传入内核，不要在内核中尝试访问前后行值、操作datetime类型
• 内核中禁止使用任何动态Python对象：临时存储全部用固定长度、固定类型的数组，不要用list.append、pandas构造/查询这类动态逻辑
• 如果需要按当日价格动态划分10个区间，先按Date分组计算当日价格的min、max值，join回原表后替换全局price_min/price_max传入内核即可
• 有状态的累计计算必须用串行执行的接口，避免GPU多线程并行导致的状态错乱

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