python - 根据另一列中的值复制 PySpark Dataframe 中的行并获得顺序

Question

例如，我想基于原始数据集创建新数据集

例如我的输入1

我的输出应该是2我引用了其他代码并得到了这个

def duplicate_function(row):
    data = []  # list of rows to return
    to_duplicate = float(row["No_of_Occ"])

    i = 0
    while i < to_duplicate:
        row_dict = row.asDict()  # convert a Spark Row object to a Python dictionary
        row_dict["No_of_Occ"] = str(i)
        new_row = Row(**row_dict)  # create a Spark Row object based on a Python dictionary
        to_return.append(new_row)  # adds this Row to the list
        i += 1

    return data  # returns the final list

但是我怎样才能得到这里的 No_of_occ？

Answer 1

一般的想法是将值复制到 No_of_Occ 的次数，然后使用 posexplode 生成更多行。

假设 df 是您的数据框。

from pyspark.sql import functions as F, types as T

output_schema = T.ArrayType(df.drop("no_of_occ").schema)

@F.udf(output_schema)
def duplicate(no_of_occ, *args):
    return list((args,) * no_of_occ)

df.select(
    "no_of_occ",
    F.posexplode(duplicate(*df.columns))
).select(
    "no_of_occ",
    (F.col("pos")+1).alias("occ_no"),
    F.col("col.*")
).show()


+---------+------+------+------+
|no_of_occ|occ_no|value1|value2|
+---------+------+------+------+
|        2|     1|     2|     3|
|        2|     2|     2|     3|
|        3|     1|     3|     4|
|        3|     2|     3|     4|
|        3|     3|     3|     4|
|        4|     1|     5|     6|
|        4|     2|     5|     6|
|        4|     3|     5|     6|
|        4|     4|     5|     6|
|        2|     1|     7|     8|
|        2|     2|     7|     8|
|        1|     1|     8|     9|
+---------+------+------+------+

需要解决no_of_occ = 0的情况。

---

编辑:如果您还需要保留带有 0 的行(如 1)，则替换 UDF:

@F.udf(output_schema)
def duplicate(no_of_occ, *args):
    dup_value = no_of_occ or 1
    return list((args,) * dup_value)