如何将python dask延迟的输出（来自xarray）保存到数据帧中

Question

我对尝试并行化python代码非常陌生。我试图对xarray执行一些分析，然后用结果填充pandas数据框。数据帧的列是独立的，所以我认为使用daskdelayed进行并行化应该很简单，但无法解决如何进行并行化。我的xarray相当大，所以这个循环需要一段时间，而且内存很大。如果更容易的话，它也可以按时间分块（这可能有助于记忆）

以下是未并行化的版本：

from time import sleep
import time
import pandas as pd
import dask.dataframe as dd

data1 = np.random.rand(4, 3,3)
data2=np.random.randint(4,size=(3,3))

locs1 = ["IA", "IL", "IN"]
locs2 = ['a', 'b', 'c']
times = pd.date_range("2000-01-01", periods=4)

xarray1 = xr.DataArray(data1, coords=[times, locs1, locs2], dims=["time", "space1", "space2"])
xarray2= xr.DataArray(data2, coords=[locs1, locs2], dims=[ "space1", "space2"])

def delayed_where(xarray1,xarray2,id):
    sleep(1)
    return xarray1.where(xarray2==id).mean(axis=(1,2)).to_dataframe(id)

final_df=pd.DataFrame(columns=range(4),index=times)

for column in final_df:
    final_df[column]=delayed_where(xarray1,xarray2,column)

我想并行化for循环，但已尝试：

final_df_delayed=pd.DataFrame(columns=range(4),index=times)

for column in final_df:
    final_df_delayed[column]=delayed(delayed_where)(xarray1,xarray2,column)

final_df.compute()

或者dask的事耽搁了

final_df_dd=dd.from_pandas(final_df, npartitions=2)
for column in final_df:
    final_df_dd[column]=delayed(delayed_where)(xarray1,xarray2,column)

final_df_dd.compute()

但这些都不管用。有人能帮忙吗

Answer 1

您正确地使用了delayed，但无法按照指定的方式构造dask数据帧

from dask import delayed
import dask

@delayed
def delayed_where(xarray1,xarray2,id):
    sleep(1)
    return xarray1.where(xarray2==id).mean(axis=(1,2)).to_dataframe(id)

@delayed
def form_df(list_col_results):
    final_df=pd.DataFrame(columns=range(4),index=times)

    for n, column in enumerate(final_df):
        final_df[column]=list_col_results[n]
    return final_df

delayed_cols = [delayed_where(xarray1,xarray2, col) for col in final_df.columns]

delayed_df = form_df(delayed_cols)

delayed_df.compute()

请注意，枚举是获取列的正确顺序的一种笨拙方法，但实际问题可能会引导您找到更好的方法来指定它（例如，通过将每个列显式指定为单个参数）