python - 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 更快地完成这项工作和 .values Series 的属性对象，但想针对 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 中指定返回类型争论，但我似乎无法弄清楚。

Answer 1

问题是numba不容易支持字符串( see here 和 see here )。

解决方案 是处理 bool 逻辑 if x=='a'外numba装饰功能。如下修改您的示例(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)  

返回:

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

如预期的。

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