如何使用 Numba + Dask 正确并行化通用代码

Question

我是使用Dask和Numba加速代码的新手，我希望这对用户来说是一个有价值的问题，可以帮助他们获得有关如何并行化代码的最佳实践的答案。我已经制作了一个包含 3 列的数据框的通用测试用例pandas。

将在框架中的 3 个向量上实现通用函数，以表示数据分析中可能进行的转换类型：前两列进行平方、相加，然后取平方根，然后是布尔值将结果与第三列进行比较计算。

我实现了 4 个测试用例：(a) a pandasapply、(b) Dask、(c)Numba和 (d)Dask以及Numba在一起。

Numba效果很好。我所有的问题都与Dask. 以下是我遇到的问题：

1. Dask，无论我制作什么大小的向量，都会更慢。我可能不完全理解如何以及何时计算数据帧的某些部分，或者如何使其正确并行化。它比常规申请慢。
2. 如何正确使用Dask进行并行化？我将其设计为 4 个分区，并且有 2 个核心处理器，但是您实际上如何决定如何格式化它？

# Practice parallelizing
from dask import dataframe as dd
from numba import jit
import pandas as pd
import numpy as np
import time

# df is going to be the regular dataframe
df = pd.DataFrame(np.random.random(size=(1000000,3))*100,columns=['col1','col2','col3'])

# ddf is the dask dataframe
ddf = dd.from_pandas(df,npartitions=4)

# Check the distance regular (probably wouldn't write like this but doing for symmetry)
def check_dist(col1,col2,col3):
    dist = np.sqrt(col1**2+col2**2)
    check = dist < col3
    return check

# Jit
@jit(nopython=True)
def check_dist_fast(col1,col2,col3):
    dist = np.sqrt(col1**2+col2**2)
    check = dist < col3
    return check

#####################################
# Regular Python Apply
#####################################
t0 = time.time()
df['col4'] = df.apply(lambda x: check_dist(x.col1,x.col2,x.col3),axis=1)
t1 = time.time()-t0
print("Regular pandas took",t1)
df = df.drop('col4',axis=1)

#####################################
# Dask Apply
#####################################
t0 = time.time()
ddf['col4'] = ddf.map_partitions(lambda d: d.apply(
                                    lambda x: check_dist(x.col1,x.col2,x.col3),axis=1)
                                ).compute()
t1 = time.time()-t0
print("Dask pandas took",t1)
ddf = ddf.drop('col4',axis=1)


#####################################
# Numba Pandas
#####################################
t0 = time.time()
df['col4'] = check_dist_fast(df.col1.to_numpy(),df.col2.to_numpy(),df.col3.to_numpy())
t1 = time.time()-t0
print("Numba pandas took",t1)
df = df.drop('col4',axis=1)


#####################################
# Numba + Jit Pandas
#####################################
t0 = time.time()
t0 = time.time()

ddf['col4'] = ddf.map_partitions(lambda d: d.apply(lambda x:
                    check_dist_fast(x.col1,x.col2,x.col3),axis=1)).compute()
t1 = time.time()-t0
print("Numba Dask pandas took",t1)
ddf = ddf.drop('col4',axis=1)

最后，我们还应该了解哪些其他最佳实践。这个想法是将其发送到某种具有许多节点的集群。

时间是：

• 普通熊猫花了 150.6191689968109
• Dask 熊猫拿了 153.70575094223022
• Numba 熊猫花了 0.710655927658081
• Numba Dask 熊猫花了 139.57402181625366

Answer 1

我认为 dask 太慢了，因为您正在使用以下方法计算一系列：

ddf.map_partitions(
    lambda d: d.apply(lambda x: check_dist(x.col1,x.col2,x.col3), axis=1)
    ).compute()

然后将其分配给一个新列，这样 dask 就无法并行化该过程。以下代码执行相同的操作，但运行时间为 0.06 秒：

#####################################
# Dask Assign
#####################################
t0 = time.time()
ddf = ddf.assign(col4=lambda x: check_dist(x.col1,x.col2,x.col3))
ddf.compute()
t1 = time.time()-t0
print("Dask using Assign took",t1)
ddf = ddf.drop('col4',axis=1)

我建议查看 dask 文档中的最佳实践部分。

希望这可以帮助！