python - 按列位置子集 dask 数据框

Question

一旦有了 dask 数据框，我如何有选择地将列拉入内存中的 pandas DataFrame？假设我有一个 N x M 数据框。我如何创建一个 N x m 数据框，其中 m << M 并且是任意的。

from sklearn.datasets import load_iris
import dask.dataframe as dd

d = load_iris()
df = pd.DataFrame(d.data)
ddf = dd.from_pandas(df, chunksize=100)

我想做什么:

in_memory = ddf.iloc[:,2:4].compute()

我能做的:

ddf.map_partitions(lambda x: x.iloc[:,2:4]).compute()

map_partitions 可以工作，但在一个不是很大的文件上速度很慢。我希望我遗漏了一些非常明显的东西。

Answer 1

虽然 iloc 没有为 dask-dataframes 实现，但你可以很容易地实现索引，如下所示:

cols = list(ddf.columns[2:4])
ddf[cols].compute()

这有额外的好处，即 dask 立即知道所选列的类型，并且不需要做额外的工作。对于 map_partitions 变体，dask 至少需要检查生成的数据类型，因为您调用的函数是完全任意的。