从dask dataframe保存多个地块文件

Question

我想从Dask dataframe中保存多个拼图文件，一个拼图文件用于特定列中的所有唯一值。因此，拼图文件的数量应该等于该列中唯一值的数量。

例如，给定以下数据帧，我想保存四个拼图文件，因为在列"A“中有四个唯一的值。

import pandas as pd
from dask import dataframe as dd

df = pd.DataFrame(
    {
        "A": [1, 1, 2, 3, 1, 3, 6, 6],
        "B": ["A", "L", "C", "D", "A", "B", "A", "B"],
        "C": [1, 2, 3, 4, 5, 6, 7, 8],
    }
)
ddf = dd.from_pandas(df, npartitions=2)

for i in ddf["A"].unique().compute():
    ddf.loc[ddf["A"] == i].to_parquet(f"file_{i}.parquet", schema="infer")

我不确定循环遍历Dask dataframe是否是扩大规模的正确方法(可能unique().compute()可能比我的内存大)。此外，我不确定我是否必须提前订购。

如果你有一些建议，如何正确地实现这一点或考虑事项，我将很高兴！

Answer 1

这并不完全是您想要的，但是可以使用.to_parquet的partition_on选项

ddf.to_parquet("file_parquet", schema="infer", partition_on="A")

请注意，这并不能保证每个分区都有一个您想要的文件，相反，file_parquet中将有多个子文件夹，其中可能包含多个文件。

Answer 2

您可以通过将索引设置为感兴趣的列并将divisions设置为遵循该列中的唯一值来实现这一点。

这应该能起到作用：

import dask.dataframe as dd
import pandas as pd
import numpy as np

# create dummy dataset with 3 partitions
df = pd.DataFrame(
    {"letter": ["a", "b", "c", "a", "a", "d", "d", "b", "c", "b", "a", "b", "c", "e", "e", "e"], "number": np.arange(0,16)}
)

ddf = dd.from_pandas(df, npartitions=3)

# set index to column of interest
ddf = ddf.set_index('letter').persist()

# generate list of divisions (last value needs to be repeated)
index_values = list(df.letter.unique())
divisions = index_values.append(df.letter.unique()[-1])

# repartition 
ddf = ddf.repartition(divisions=divisions).persist()

# write out partitions as separate parquet files
for i in range(ddf.npartitions):
    ddf.partitions[i].to_parquet(f"file_{i}.parquet", engine='pyarrow')

请注意，值'e'在分区列表中出现了两次。根据the Dask docs：“分区包括每个分区的索引的最小值和最后一个分区的索引的最大值。”这意味着最后一个值需要包含两次，因为它既是最后一个分区索引的开始，也是最后一个分区索引的结束。