python - 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

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