python - 是否可以使用 dask 获取集合的交集？

Question

我有一个大数据集(5000 万行)，我需要在其中进行一些按行计算，例如获取两个集合的交集(每个集合在不同的列中)

例如

col_1:{1587004, 1587005, 1587006, 1587007}
col_2:{1587004, 1587005}
col_1.intersection(col_2) = {1587004, 1587005}

这适用于我的虚拟数据集 (100 000) 行。但是，当我尝试对实际内存进行相同操作时，内存用完了

我的编码工作使用 Pandas 1:1 移植到 dask 不起作用NotImplementedError: Series getitem in only supported for other series objects with matching partition structure

到目前为止，尝试使用 map_partitions 是行不通的

工作代码:

df["intersection"] = [col_1.intersection(col_2) for col_1,col2 in zip(df.col_1,df.col_2)]

用 dask df 替换 pandas df 在未实现的错误中运行:

ddf["intersection"] = [col_1.intersection(col_2) for col_1,col2 in zip(df.col_1,df.col_2)]

使用 map_partions “有效”，但我不知道如何将结果分配给现有的 ddf

def intersect_sets(df, col_1, col_2):
    result = df[col_1].intersection(df[col_2])
    return result

newCol = ddf.map_partitions(lambda df : df.apply(lambda series: intersect_sets(series,"col_1","col_2"),axis=1),meta=str).compute()

只是在做:

ddf['result'] = newCol

导致:ValueError:并非所有分区都已知，无法对齐分区。请使用 set_index 设置索引。

更新:重置索引会消除错误，但是包含交叉点的列不再与其他两列匹配。好像顺序乱了...

ddf2 = ddf.reset_index().set_index('index')
ddf2 ['result'] = result

我希望有一个包含以下列的 dask 数据框

col_1:{1587004, 1587005, 1587006, 1587007}
col_2:{1587004, 1587005}
col_3:{1587004, 1587005}

不仅欣赏完美工作的解决方案，而且对 map_partitions 如何工作的一些见解也会对我有很大帮助:)

更新:感谢 M.Rocklin，我弄明白了。对于 future 我或其他人绊倒这个问题:

ddf = ddf.assign(
       new_col = ddf.map_partitions(
           lambda df : df.apply(
                        lambda series:intersect_sets(
                                series,"col_1","col_2"),axis=1),meta=str)
 )
 df = ddf.compute()

Answer 1

如果您有一个适用于 pandas 数据帧的函数:

def f(df: pandas.DataFrame) -> pandas.Series:
    return df.apply(...)

然后您可以将此函数映射到您的分区

df['new'] = df.map_partitions(f)

我认为你的问题是你在这里不必要地调用了计算，所以你试图将 pandas 数据帧插入 dask 数据帧。

# Don't do this
new = df.map_partitions(f).compute() 
df['new'] = new  # tries to put a pandas dataframe into a dask dataframe