從 dask dataframe 保存多個鑲木地板文件

Question

我想保存來自 Dask dataframe 的多個鑲木地板文件，一個用於特定列中所有唯一值的鑲木地板文件。 因此，parquet 文件的數量應等於該列中唯一值的數量。

例如，給定以下 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'的兩次出現。 根據Dask 文檔：“分區包括每個分區索引的最小值和最后一個分區索引的最大值。” 這意味着最后一個值需要包含兩次，因為它既作為最后一個分區索引的開始又作為結束。