如何将dask.dataframe预缓存到所有工作程序和分区,以减少通信需求

时间:2019-07-30 14:47:31

标签: python pandas dask rapids cudf

有时使用dask.dataframe.map_partitions进行合并等操作很有吸引力。在某些情况下,当使用left_dfright_dfmap_partitions之间进行合并时,我想在执行合并之前实质上预缓存right_df以减少网络开销/本地改组。有没有明确的方法可以做到这一点?感觉应该可以使用client.scatter(the_df)client.run(func_to_cache_the_df)之一或组合或其他一些智能广播来实现。

在大型left_df和较小的right_df上进行左联接的情况下,这尤其明显,这实际上是一个查找表。感觉right_df应该能够读入内存,并持久地/分散到所有合并前的工作程序/分区中,以减少跨分区通信的需求,直到最后。如何分散right_df才能成功做到这一点?

以下是使用cuDF和Dask进行的这种不平衡合并的一个较小示例(但从概念上讲,与pandas和Dask相同)

import pandas as pd
import cudf
import dask_cudf
import numpy as np
from dask.distributed import Client
from dask_cuda import LocalCUDACluster

# create a local CUDA cluster
cluster = LocalCUDACluster()
client = Client(cluster)

np.random.seed(12)

nrows_left = 1000000
nrows_right = 1000

left = cudf.DataFrame({'a': np.random.randint(0,nrows_right,nrows_left), 'left_value':np.arange(nrows_left)})
right = cudf.DataFrame({'a': np.arange(nrows_right), 'lookup_val': np.random.randint(0,1000,nrows_right)})

print(left.shape, right.shape) # (1000000, 2) (1000, 2)

ddf_left = dask_cudf.from_cudf(left, npartitions=500)
ddf_right = dask_cudf.from_cudf(right, npartitions=2)

def dask_merge(L, R):
    return L.merge(R, how='left', on='a')

result = ddf_left.map_partitions(dask_merge, R=ddf_right).compute()
result.head()
<cudf.DataFrame ncols=3 nrows=5 >
     a  left_value  lookup_val
0  219        1952         822
1  873        1953         844
2  908        1954         142
3  290        1955         810
4  863        1956         910

1 个答案:

答案 0 :(得分:1)

如果您执行以下任何一项操作,则应该没事:

  • 与单分区dask数据帧合并
  • 与非黄昏数据框(例如Pandas或cuDF)合并
  • 具有非黄昏数据框(例如Pandas或cuDF)的map_partitions

这是怎么回事:

  1. 将单个分区推送给单个工作人员
  2. 在执行过程中,一些工作人员将复制该数据,然后其他人将从这些工作人员复制,依此类推,以树的形式传达数据。
  3. 工人将按预期进行合并

这大约和预期的一样快。但是,如果您正在执行基准测试之类的操作,并且想要将步骤1,2和3分开,则可以使用client.replicate

left = ... # multi-partition dataframe
right = ... # single-partition dataframe
right = right.persist()  # make sure it exists in one worker
client.replicate(right)  # replicate it across many workers

... proceed as normal

这不会更快,但是步骤1,2将被拉入复制步骤。

在您的示例中,看来right有两个分区。您可能需要将其更改为一个。在这种情况下,Dask采用了不同的代码路径,该路径实际上只是map_partitions