向dask.dataframe添加dask.array列

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数据帧的索引列作为普通的pandas数据帧，添加Dask数组列，然后按索引列将其合并回Dask数据帧

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