python - 如何从 Dask DataFrame 中提取前五个值而不计算整个 DataFrame？

Question

晚上好!

我有一个与下面粘贴的代码类似的代码，它有更多的数据，但前提是相同的。我必须从两个 DataFrame 中提取前五个值，但是当我处理数千万个条目时，我有时无法等待长达一个小时来计算整个 DataFrame 并返回前五个值。我也无法使用简单的 Pandas DataFrame，因为它们超出了我的内存限制。有没有办法解决这个问题？

import random
import pandas
import dask.dataframe as dd
import time

# Random list from 1 to 10,000,000.
random_pool = [random.randint(1, 1000000) for i in range(10000000)]
random.shuffle(random_pool)

df1 = dd.from_pandas(pandas.DataFrame(random_pool[:100000], columns=["ID"]), npartitions=10)
df2 = dd.from_pandas(pandas.DataFrame(random_pool, columns=["ID"]), npartitions=10)

# Sorting both dataframes.
df1 = df1.sort_values("ID", ascending=True)
df2 = df2.sort_values("ID", ascending=True)

df1_start = time.time()
df1.head(5)
print("DF1 took {:.2f}.".format(time.time() - df1_start))

df2_start = time.time()
df2.head(5)
print("DF2 took {:.2f}.".format(time.time() - df2_start))

第一个 DataFrame 大约需要 0.41 秒，而第二个 DataFrame 大约需要 1.79 秒。

Answer 1

需要记住的一件事是 dask 中的值实际上是一系列操作的序列化堆栈。许多计算都会被推迟，直到您真正要求具体化值时为止 - 例如使用 head，或者通常使用 .compute()。

本着 general advice on persist 的精神，您可以尝试在排序调用后使用 .persist() :

It is often ideal to load, filter, and shuffle data once and keep this result in memory. Afterwards, each of the several complex queries can be based off of this in-memory data rather than have to repeat the full load-filter-shuffle process each time. To do this, use the client.persist method [ The .persist() method in this case ].

花点时间考虑一下如果我们不坚持下去会发生什么 - 当您调用 head 时需要解决的 future 将包括 sort_values 调用，并且您可能会看到每次调用 head 时对所有数据进行排序的成本 - 这解释了为什么仅获取前五个项目的成本与整个数据集的大小成正比 - 因为整个数据集正在排序。

答案是 dask 获取前 5 项的速度相当快。但如果所有计算尚未在内存中，那么解决所有计算可能不会那么快。

一般来说你应该avoid whole-dataset shuffling就像这个例子一样 - 排序!