Question

我有一个数据框，需要向其添加ID。我写了一个函数可以帮助我。我正在使用zipWithIndex来帮助我。

def add_ids(X):
    from copy import deepcopy
    from pyspark.sql.types import LongType
    schema_new = deepcopy(X.schema)
    schema_new = schema_new.add("default_id_col", LongType(), False)
    _X = X.rdd.zipWithIndex().map(lambda l: list(l[0]) + [l[1]]).toDF(schema_new)
    cols_arranged = [_X.columns[-1]] + _X.columns[0 : len(_X.columns) - 1]
    return _X.select(*cols_arranged)

代码工作正常，但是当我在具有超过一百万条记录的数据集上使用它时，它的速度确实很慢。我在3节点群集（80Gb）上使用了它，获得了500万条记录数据集，它花了大约12分钟的时间才添加了ID

有更好的方法吗？

Answer 1

只要不需要连续的ID，就可以使用monotonically_increasing_id函数。

from pyspark.sql import functions as F

df = df.withColumn("id", F.monotonically_increasing_id())

PySpark：如何将id列添加到数据框是一种快速有效的方法？

1 个答案: