numba.vectorize - 不支持的数组数据类型

Question

我是 numba 的新手，似乎无法理解要传递给 vectorize 的参数。这是我正在尝试做的事情：

test = [x for x in range(10)]
test2 = ['a', 'a', 'a', 'b', 'b', 'c', 'c', 'c', 'c', 'c']
test_df = pd.DataFrame({'test': test, 'test2': test2})
test_df['test3'] = np.where(test_df['test'].values % 2 == 0,
                            test_df['test'].values, 
                            np.nan)


  test  test2   test3   test4
0    0      a     0.0     0.0
1    1      a     NaN     NaN
2    2      a     2.0     4.0
3    3      b     NaN     NaN
4    4      b     4.0    16.0
5    5      c     NaN     NaN
6    6      c     6.0    36.0
7    7      c     NaN     NaN
8    8      c     8.0    64.0
9    9      c     NaN     NaN

任务是根据以下逻辑创建一个新列，首先是根据标准 pandas:

def nonnumba_test(row):
    if row['test2'] == 'a':
        return row['test'] * row['test3']
    else:
        return np.nan

使用apply；我知道我可以使用 np.where 和 Series 对象的 .values 属性更快地完成此操作，但想针对 numba.

进行测试

test_df.apply(nonnumba_test, axis=1)

0    0.0
1    NaN
2    4.0
3    NaN
4    NaN
5    NaN
6    NaN
7    NaN
8    NaN
9    NaN
dtype: float64

接下来，当我尝试使用 numba.vectorize 装饰器时

@numba.vectorize()
def numba_test(x, y, z):
    if x == 'a':
        return y * z
    else:
        return np.nan

我收到以下错误

numba_test(test_df['test2'].values, 
           test_df['test'].values, 
           test_df['test3'].values)

ValueError: Unsupported array dtype: object

我想我需要在 signature 参数中指定 return 类型，但我似乎无法弄清楚。

Answer 1

问题是numba不容易支持字符串（see here and ）。

解决方案是在numba修饰函数外处理布尔逻辑if x=='a'。如下修改示例（numba_test 和输入参数）会产生所需的输出（示例中最后两个块上方的所有内容均未更改）：

from numba import vectorize, float64, int64, boolean

#@vectorize() will also work here, but I think it's best practice with numba to specify types.
@vectorize([float64(boolean, int64, float64)])
def numba_test(x, y, z):
    if x:
        return y * z
    else:
        return np.nan

# now test it...
# NOTICE the boolean argument, **not** string!
numba_test(test_df['test2'].values =='a', 
           test_df['test'].values, 
           test_df['test3'].values)

Returns:

array([  0.,  nan,   4.,  nan,  nan,  nan,  nan,  nan,  nan,  nan])

随心所欲。

最后说明：您会看到我在上面的 vectorize 装饰器中指定了类型。是的，这有点烦人，但我认为这是最佳实践，因为它可以让您免于头痛正是就像这样：如果您指定了类型，您将无法找到字符串类型, 这样就解决了。

numba.vectorize - 不支持的数组数据类型

numba.vectorize - Unsupported array dtype

python

pandas

numba