numba.vectorize-不支持的数组dtype

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它不容易支持字符串（请参阅此处和此处）。

解决方案是处理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上面的装饰器中指定了类型。是的，这有点烦人，但是我认为这是最佳做法，因为它使您完全省却了这种麻烦：如果指定了类型，则将无法找到字符串类型，那么就可以解决它。