PySpark将行合并为列StackOverFlow错误

Question

我想要什么（非常简化）：

[Input Dataset至Output dataset

我尝试过的一些代码：

def add_columns(cur_typ, target, value):
        if cur_typ == target:
            return value
        return None
  schema = T.StructType([T.StructField("name", T.StringType(), True),
                         T.StructField("typeT", T.StringType(), True),
                         T.StructField("value", T.IntegerType(), True)])
  data = [("x", "a", 3), ("x", "b", 5), ("x", "c", 7), ("y", "a", 1), ("y", "b", 2),
          ("y", "c", 4), ("z", "a", 6), ("z", "b", 2), ("z", "c", 3)]
  df = ctx.spark_session.createDataFrame(ctx.spark_session.sparkContext.parallelize(data), schema)
  targets = [i.typeT for i in df.select("typeT").distinct().collect()]
  add_columns = F.udf(add_columns)
  w = Window.partitionBy('name')
  for target in targets:
      df = df.withColumn(target, F.max(F.lit(add_columns(df["typeT"], F.lit(target), df["value"]))).over(w))
  df = df.drop("typeT", "value").dropDuplicates()

另一个版本：

targets = df.select(F.collect_set("typeT").alias("typeT")).first()["typeT"]

w = Window.partitionBy('name')

for target in targets:
  df = df.withColumn(target, F.max(F.lit(F.when(veh["typeT"] == F.lit(target), veh["value"])
                                                  .otherwise(None)).over(w)))

df = df.drop("typeT", "value").dropDuplicates()

对于小型数据集，它们都可以工作，但是我有一个具有100万行和5000个不同typeT的数据框。因此，结果应该是大约500 x 5000的表（某些名称没有某些typeT。现在我得到stackoverflow错误（py4j.protocol.Py4JJavaError：调用o7624.withColumn时发生错误。：java.lang.StackOverflowError）尝试创建此数据框。除了增加堆栈大小，我还能做什么？有没有更好的方法来获得相同的结果？

Answer 1

如果没有要添加的列数更多，则在循环中使用withColumn不好。

创建cols数组，然后选择它们，这将导致更好的性能

cols = [F.col("name")]
for target in targets:
    cols.append(F.max(F.lit(add_columns(df["typeT"], F.lit(target), df["value"]))).over(w).alias(target))
df = df.select(cols)

结果相同的输出
+----+---+---+---+
|name|  c|  b|  a|
+----+---+---+---+
|   x|  7|  5|  3|
|   z|  3|  2|  6|
|   y|  4|  2|  1|
+----+---+---+---+