将 dask.array 列添加到 dask.dataframe

Question

我有一个 dask 数据框和一个 dask 数组，它们的行数相同，逻辑顺序相同。 数据帧行由字符串索引。 我正在尝试将数组列之一添加到数据框中。 我尝试了几种方法，但都以它们特定的方式失败了。

df['col'] = da.col
# TypeError: Column assignment doesn't support type Array

df['col'] = da.to_frame(columns='col')
# TypeError: '<' not supported between instances of 'str' and 'int'

df['col'] = da.to_frame(columns=['col']).set_index(df.col).col
# TypeError: '<' not supported between instances of 'str' and 'int'

df = df.reset_index()
df['col'] = da.to_frame(columns='col')
# ValueError: Not all divisions are known, can't align partitions. Please use `set_index` to set the index.

和其他一些变体。

当结构在逻辑上兼容时，将 dask 数组列添加到 dask 数据帧的正确方法是什么？

Answer 1

这似乎从 dask 版本2021.4.0 ，可能更早。 只需确保数据帧分区的数量与数组块的数量相匹配。

import dask.array as da
import dask.dataframe as dd
import numpy as np
import pandas as pd
ddf = dd.from_pandas(pd.DataFrame({'z': np.arange(100, 104)}),
                     npartitions=2)
ddf['a'] = da.arange(200,204, chunks=2)
print(ddf.compute())

输出：

     z    a
0  100  200
1  101  201
2  102  202
3  103  203

Answer 2

解决办法是将原Dask dataframe的index列取出为plain pandas dataframe，在其上添加Dask数组列，然后通过index列将其合并回Dask dataframe

index_col = df['index'].compute()
index_col['new_col'] = da.col.compute()
df = df.merge(index_col, 'left', on='index')