dask DataFrame 中的复杂过滤

Question

我习惯于对 pandas DataFrame 对象进行 "complex" 过滤：

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]

工作正常。

我认为问题是一旦我做了任何 "complex" 数学运算（即 np.sqrt），结果就不再是一个懒惰的 dask DataFrame。

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

Answer 1

如果将 np.sqrt 替换为 da.sqrt，则一切正常。

import dask.array as da

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

dask DataFrame 中的复杂过滤

Complex filtering in dask DataFrame

python

numpy

dataframe

pandas

dask