在cuDF中使用列表和集合执行查询的方法及报错解决

在RAPIDS 22.12版本中，cuDF的query方法暂不支持直接传入Python列表作为变量进行in或isin判断，会触发Numba NRT（Numba Runtime）相关内部错误，这是该版本的已知兼容性限制。

替代方案1：直接使用布尔索引（推荐，性能更优） #

处理数十亿级规模数据时，布尔索引的性能优于query，无需经过表达式解析环节，直接利用GPU向量运算能力：

test = list(test_userid)[0:2]
# 生成筛选掩码
mask = df['src'].isin(test) | df['dst'].isin(test)
# 提取目标数据
result_df = df[mask]

dask-cudf下逻辑完全一致，框架会自动完成分块并行处理。

替代方案2：将列表转为cuDF Series后传入query #

若坚持使用query语法，可将Python列表转换为cuDF Series，避免原生列表在表达式解析中的兼容性问题：

import cudf
test = list(test_userid)[0:2]
test_series = cudf.Series(test)
result_df = df.query("src.isin(@test_series) or dst.isin(@test_series)")

注意事项 #

• 测试代码中src==@test的广播匹配写法，cuDF query完全不支持，无论版本都应避免使用。
• 针对大规模用户ID集合，优先转为cuDF Series/Index，减少CPU与GPU间的数据拷贝开销，提升处理效率。

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