python - 使用循环在 pyspark 数据框中添加多列

Question

我需要在 pyspark 的数据框中添加一些列(4000)。我正在使用 withColumn函数，但得到断言错误。

df3 = df2.withColumn("['ftr' + str(i) for i in range(0, 4000)]", [expr('ftr[' + str(x) + ']') for x in range(0, 4000)])

不知道出了什么问题。

Answer 1

我们可以使用 .select()而不是 .withColumn()使用列表作为输入来创建与链接多个 .withColumn() 类似的结果的。 ["*"]还用于选择数据框中的每个现有列。import pyspark.sql.functions as Fdf2:

+---+
|age|
+---+
| 10|
| 11|
| 13|
+---+

df3 = df2.select(["*"] + [F.lit(f"{x}").alias(f"ftr{x}") for x in range(0,10)])结果是:

+---+----+----+----+----+----+----+----+----+----+----+
|age|ftr0|ftr1|ftr2|ftr3|ftr4|ftr5|ftr6|ftr7|ftr8|ftr9|
+---+----+----+----+----+----+----+----+----+----+----+
| 10|   0|   1|   2|   3|   4|   5|   6|   7|   8|   9|
| 11|   0|   1|   2|   3|   4|   5|   6|   7|   8|   9|
| 13|   0|   1|   2|   3|   4|   5|   6|   7|   8|   9|
+---+----+----+----+----+----+----+----+----+----+----+