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

Question

我需要将许多列（4000）添加到 pyspark 的数据框中。我正在使用该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 F

df2：

+---+
|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|
+---+----+----+----+----+----+----+----+----+----+----+

Answer 2

尝试做这样的事情：

df2 = df3
for i in range(0, 4000):
  df2 = df2.withColumn(f"ftr{i}", lit(f"frt{i}"))