如何在含pandas read_csv的函数中启用numba @njit(nopython=True)？

核心问题答案 #

不能在@njit装饰的函数里调用pd.read_csv这类IO操作。因为Numba的njit模式会把代码编译成机器码，仅支持Python和NumPy的原生数值操作，像文件读取这种依赖Python解释器、外部库的IO逻辑不在支持范围内，强行编写会直接编译失败。

你的POC失败原因 #

你测试的test_value函数报错，是因为df.to_numpy()生成了混合类型数组（包含整数和字符串），Numba对混合类型数组的支持非常有限，尤其是字符串类型。换成纯数值的DataFrame，代码就能正常运行：

import numba
import pandas as pd
import numpy as np

df = pd.DataFrame(columns=['a','c'], data=[[1,3],[2,4]])

@numba.njit 
def test_value(df_np):
  print(df_np[:,0])

test_value(df.to_numpy())

可行解决方案 #

方案1：预处理阶段一次性加载所有CSV #

把所有需要的CSV文件在@njit函数外面读好，转成Numba支持的格式（纯数值NumPy数组、结构化数组），再传入编译后的函数做计算。IO操作全在Python解释器层面完成，编译后的函数只负责纯计算逻辑，最大化利用Numba的加速效果。
示例代码：

import numba
import pandas as pd
import numpy as np

# 预处理：读取所有CSV并转成NumPy数组
csv_data = {}
list_of_value = ["file1", "file2", "file3"]
for xx in list_of_value:
    xdf = pd.read_csv(f"../corpus/{xx}.csv")
    # 假设计算只需要数值列，转成纯数值数组
    csv_data[xx] = xdf.select_dtypes(include=np.number).to_numpy()

# 用njit编译纯计算函数
@numba.njit
def calculate_prob(csv_data_values):
    results = np.empty(len(csv_data_values), dtype=np.float64)
    for i in range(len(csv_data_values)):
        data = csv_data_values[i]
        # 这里写你的概率计算逻辑，以均值为例
        pr = np.mean(data)
        results[i] = pr
    return results

# 调用函数
values = np.array(list(csv_data.values()), dtype=object)
probs = calculate_prob(values)

方案2：分文件处理，IO在外部，计算在njit函数 #

如果CSV文件太大没法一次性加载，就把单个文件的读取逻辑放在外面，每次读一个文件转成数组后，调用@njit函数处理该文件的数据，循环完成所有文件的计算：

import numba
import pandas as pd
import numpy as np

@numba.njit
def compute_single_prob(data):
    # 单个文件的计算逻辑
    pr = np.mean(data)
    return pr

list_of_value = ["file1", "file2", "file3"]
results = []
for xx in list_of_value:
    xdf = pd.read_csv(f"../corpus/{xx}.csv")
    data = xdf.select_dtypes(include=np.number).to_numpy()
    pr = compute_single_prob(data)
    results.append(pr)

方案3：优化计算逻辑适配Numba #

• 避免在@njit函数里用Python列表的append，改用NumPy数组预分配内存（比如示例里的np.empty），速度会快很多。
• 确保传入@njit函数的是纯数值类型的数组，避免混合类型、字符串类型，减少编译错误。

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