將新列附加到 dask 數據框

Question

這是在 dask 中混洗數據的后續問題。

我有一個現有的 dask 數據框df ，我希望在其中執行以下操作：

df['rand_index'] = np.random.permutation(len(df))

但是，這會產生錯誤， Column assignment doesn't support type ndarray 。 我嘗試使用df.assign(rand_index = np.random.permutation(len(df)) ，它給出了相同的錯誤。

這是一個最小（非）工作示例：

import pandas as pd
import dask.dataframe as dd
import numpy as np

df = dd.from_pandas(pd.DataFrame({'A':[1,2,3]*10, 'B':[3,2,1]*10}), npartitions=10)
df['rand_index'] = np.random.permutation(len(df))

筆記：

上一個問題提到使用df = df.map_partitions(add_random_column_to_pandas_dataframe, ...)但我不確定這是否與這種特殊情況有關。

編輯 1

我嘗試了df['rand_index'] = dd.from_array(np.random.permutation(len_df)) ，它沒有問題地執行。 當我檢查df.head() ，似乎新列創建得很好。 但是，當我查看df.tail() ， rand_index是一堆NaN 。

事實上只是為了確認我檢查了df.rand_index.max().compute()結果比len(df)-1 。 所以這可能是df.map_partitions發揮作用的地方，因為我懷疑這是 dask 被分區的問題。 在我的特殊情況下，我有 80 個分區（不是指示例案例）。

Answer 1

您需要將np.random.permutation(len(df))轉換為 dask 理解的類型：

permutations = dd.from_array(np.random.permutation(len(df)))
df['rand_index'] = permutations
df

這將產生：

Dask DataFrame Structure:
                    A      B rand_index
npartitions=10                         
0               int64  int64      int32
3                 ...    ...        ...
...               ...    ...        ...
27                ...    ...        ...
29                ...    ...        ...
Dask Name: assign, 61 tasks

因此，現在是否要.compute()計算實際結果.compute()您。

Answer 2

要分配一列，您應該使用df.assign

Answer 3

遇到與編輯 1 中相同的問題。

我的解決方法是從現有數據框中獲取一個唯一的列，並將其輸入到要附加的數據框中。

import dask.dataframe as dd
import dask.array as da
import numpy as np
import panda as pd

df = dd.from_pandas(pd.DataFrame({'A':[1,2,3]*2, 'B':[3,2,1]*2, 'idx':[0,1,2,3,4,5]}), npartitions=10)
chunks = tuple(df.map_partitions(len).compute())
size = sum(chunks)
permutations = da.from_array(np.random.permutation(len(df)), chunks=chunks)
idx = da.from_array(df['idx'].compute(), chunks=chunks)
ddf = dd.concat([dd.from_dask_array(c) for c in [idx,permutations]], axis = 1)
ddf.columns = ['idx','rand_idx']
df = df.merge(ddf, on='idx')
df = df.set_index('rand_idx')
df.compute().head()