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