Dask：在大型 dataframe 上設置索引會導致處理期間的磁盤空間使用率很高

Question

我正在處理一個大型數據集（220 000 000 行，~25Gb 作為 csv 文件），它存儲為少數 csv 文件。

我已經設法用 Dask 閱讀了這些 csv 並將數據保存為鑲木地板文件，其中包含以下內容：

import pandas as pd
from dask.distributed import Client
import dask.dataframe as dd
client = Client()

init_fields = {
# definition of csv fields
}

raw_data_paths = [
# filenames with their path
]

read_csv_kwargs = dict(
    sep=";",
    header=None,
    names=list(init_fields.keys()),      
    dtype=init_fields, 
    parse_dates=['date'],    
)

ddf = dd.read_csv(
    raw_data_paths,
    **read_csv_kwargs,
)
ddf.to_parquet(persist_path / 'raw_data.parquet')

它就像一個魅力，並在幾分鍾內完成。 我得到一個鑲木地板文件，里面有一個 Dask Dataframe，有 455 個分區，我完全可以使用。

但是，此 dataframe 包含大量客戶訂單，我想按日期對其進行索引以便進一步處理。

當我嘗試通過以下調整運行代碼時：

ddf = dd.read_csv(
    raw_data_paths,
    **read_csv_kwargs,
).set_index('date')
ddf.to_parquet(persist_path / 'raw_data.parquet')

處理變得非常長，有 26 000 多個任務（我可以理解，這是要排序的大量數據），但是工作人員在使用 memory 一段時間后開始死亡。

每個工人死亡，都會失去一些進展，似乎處理永遠不會完成。

我注意到工人死亡與我機器的磁盤達到極限有關，每當工人死亡時，都會釋放一些空間。 在處理開始時，我有大約 37 Gb 的可用磁盤空間。

我對 Dask 很陌生，所以有幾個問題：

• 在轉儲到鑲木地板文件之前設置索引是個好主意嗎？ 在接下來的步驟中，我有幾個分組日期，根據 Dask 文檔，使用這個字段作為索引在我看來是個好主意。
• 如果我在轉儲為 parquet 文件之前設法設置索引，parquet 文件是否會被排序並且我的進一步處理不需要更多的改組？
• 上述行為（memory 錯誤中的磁盤使用率高）看起來正常還是在我的設置或使用 Dask 時有些奇怪？ 有一些我可以調整的參數嗎？
• 或者我真的需要更多的磁盤空間，因為排序這么多數據需要它？ 估計所需的總磁盤空間是多少？

提前致謝！

編輯：我終於設法通過以下方式設置索引：

• 在我的機器上添加磁盤空間
• 調整客戶端參數以使每個工作人員擁有更多 memory

我使用的參數是：

client = Client(
    n_workers=1,
    threads_per_worker=8,
    processes=True,
    memory_limit='31GB'
)

我不太堅持認為磁盤使用是我的工人死於缺少 memory 的根本原因，因為僅增加磁盤空間並不能完成處理。 它還要求擴展每個工作人員的 memory，這是我通過使用機器的整個 memory 創建一個工作人員來實現的。

但是，我很驚訝需要這么多 memory。 我認為 Dask（和其他大數據工具）的目標之一是啟用“核心外處理”。 我在這里做錯了什么還是設置索引需要大量的 memory，無論如何？

問候，

Answer 1

這是我理解事物的方式，但我可能會遺漏一些重要的觀點。

讓我們從一個很好的索引數據集開始，以獲得一個可重現的示例。

import dask
import dask.dataframe as dd

df = dask.datasets.timeseries(start='2000-01-01', end='2000-01-2', freq='2h', partition_freq='12h')

print(len(df), df.npartitions)
# 12 2

所以我們正在處理一個很小的數據集，只有 12 行，分成 2 個分區。 由於此 dataframe 已編入索引，因此對其進行合並將非常快，因為 dask 知道哪些分區包含哪些（索引）值。

%%time
_ = df.merge(df, how='outer', left_index=True, right_index=True).compute()
#CPU times: user 25.7 ms, sys: 4.23 ms, total: 29.9 ms
#Wall time: 27.7 ms

現在，如果我們嘗試在非索引列上進行合並，dask 將不知道哪個分區包含哪些值，因此它必須在工作人員之間交換信息並在工作人員之間傳輸數據位。

%%time
_ = df.merge(df, how='outer', on=['name']).compute()
#CPU times: user 82.3 ms, sys: 8.19 ms, total: 90.4 ms
#Wall time: 85.4 ms

這在這個小數據集上可能看起來不多，但將其與pandas所需的時間進行比較：

%%time
_ = df.compute().merge(df.compute(), how='outer', on=['name'])
#CPU times: user 18.9 ms, sys: 3.39 ms, total: 22.3 ms
#Wall time: 19.7 ms

另一種查看方式是使用 DAG，將用於合並索引列的 DAG 與用於合並非索引列的 DAG 進行比較。 第一個很好地平行：

第二個（使用非索引列）要復雜得多：

那么隨着數據大小的增長會發生什么，使用非索引列執行操作會變得更加昂貴。 對於包含許多唯一值（例如字符串）的列尤其如此。 您可以嘗試增加上面構建的 dataframe df中的分區數量，您將觀察到非索引情況如何變得越來越復雜，而索引數據的 DAG 保持可擴展性。

回到你的具體情況，你從一個非索引的 dataframe 開始，在索引之后它將是一個非常復雜的實體。 您可以使用.visualize()查看索引 dataframe 的 DAG，根據經驗，我猜它看起來不太漂亮。

因此，當您保存到鑲木地板（或啟動數據幀的其他計算）時，工作人員開始對數據進行混洗，這將很快耗盡 memory（特別是如果有很多列和/或許多分區和/或列有很多唯一值）。 一旦工人 memory 限制接近，工人將開始將數據溢出到磁盤（如果允許的話），這就是為什么您能夠通過增加 memory 和可用磁盤空間來完成任務的原因。

在這些選項都不可能的情況下，您可能需要實現使用delayed的 API（或動態圖的futures ）的自定義工作流，以便此工作流使用一些不明確提供給 dask 的信息。 例如，如果原始 csv 文件按感興趣的列分區，則您可能希望分批處理這些 csv 文件，而不是將它們攝取到單個 dask Z6A8064B5DF4794555500553C47DZC 中。