dask DataFrame 中的复杂过滤

Question

我习惯于对 Pandas DataFrame 对象进行“复杂”过滤：

import numpy as np
import pandas as pd
data = pd.DataFrame(np.random.random((10000, 2)) * 512, columns=["x", "y"])
data2 = data[np.sqrt((data.x - 200)**2 + (data.y - 200)**2) < 1]

这不会产生任何问题。

但是使用 dask DataFrames 我有：

ddata = dask.dataframe.from_pandas(data, 8)
ddata2 = ddata[np.sqrt((ddata.x - 200)**2 + (ddata.y - 200)**2) < 1]
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
<ipython-input-13-c2acf73dddf6> in <module>()
----> 1 ddata2 = ddata[np.sqrt((ddata.x - 200)**2 + (ddata.y - 200)**2) < 1]

~/anaconda3/lib/python3.6/site-packages/dask/dataframe/core.py in __getitem__(self, key)
   2115             return new_dd_object(merge(self.dask, key.dask, dsk), name,
   2116                                  self, self.divisions)
-> 2117         raise NotImplementedError(key)
   2118 
   2119     def __setitem__(self, key, value):

NotImplementedError: 0       False

同时一个更简单的操作：

ddata2 = ddata[ddata.x < 200]

工作正常。

我认为问题在于，只要我进行任何“复杂”数学（即np.sqrt），结果就不再是懒惰的 dask DataFrame。

有没有解决的办法？我是否必须创建一个可以过滤的新列，或者有更好的方法吗？

Answer 1

如果你替换np.sqrt为da.sqrt那么一切正常。

import dask.array as da

您可能会注意到np.sqrtdask 系列会生成一个 numpy 数组，因此计算中的这一步并不是懒惰的，而是强制得出具体结果。使用 dask 等效函数来保持惰性并保持所有内容都符合 dask 标准。