python - 如何在 DASK_CUDF 中创建唯一 ID 列

Question

如何在所有分区的 dsak cudf 数据帧中创建唯一 id 列到目前为止，我正在使用以下技术，但如果我将数据增加到超过 10cr 行，则会出现内存错误。

def unique_id(df):
    rag = cupy.arrange(len(df))
    df['unique_id']=rag
    return df
    
part = data.npartitions
data = data.repartitions(npartitions=1)
cols_meta={c:str(data[c].dtype) for c in data.columns}
data = data.map_partitions(lambda df:unique_id(df), meta={**cols_meta,'unique_id'})
data = data.repartitions(npartitions=part)

如果有其他方法，或者代码有任何修改，请提出建议。谢谢你的帮助

Answer 1

您遇到内存错误的原因是这一步：

data = data.repartitions(npartitions=1)

通过使用单个分区，您将所有数据强制在单个工作人员上，随着数据集大小的增加，这将导致内存问题。您想要做的是在维护每个分区的同时分配一个唯一标识符，请参阅此答案。

Answer 2

我这样做是因为想按顺序创建 id，直到长度数据。

其他建议可能会奏效。但是，最简单的方法之一是创建一个值为 1 的临时列并使用cumsum，如下所示：

import cudf
import dask_cudf
​
df = cudf.DataFrame({
    "a": ["dog"]*10
})
ddf = dask_cudf.from_cudf(df, 3)
​
ddf["temp"] = 1
ddf["monotonic_id"] = ddf["temp"].cumsum()
del ddf["temp"]
​
print(ddf.partitions[2].compute())
     a  monotonic_id
8  dog             9
9  dog            10

正如预期的那样，分区索引 2 中的两行 ID 为 9 和 10。如果您需要索引从 0 开始，则可以减去 1。