Pyspark 中 groupby 的替代方案以提高 Pyspark 代码的性能

Question

我的 Pyspark 数据框如下所示。我必须从 pyspark 代码中删除按功能分组以提高代码的性能。我必须对 100k 数据执行操作。

[初始数据]

创建数据框

df = spark.createDataFrame([
(0, ['-9.53', '-9.35', '0.18']),
(1, ['-7.77', '-7.61', '0.16']),
(2, ['-5.80', '-5.71', '0.10']),
(0, ['1', '2', '3']),
(1, ['4', '5', '6']),
(2, ['8', '98', '32'])
 ], ["id", "Array"])

使用此代码产生预期的输出。

import pyspark.sql.functions as f

df.groupBy('id').agg(f.collect_list(f.col("Array")).alias('Array')).\
select("id",f.flatten("Array")).show()

我必须实现这种格式的输出。上面的代码给了我这个输出。我必须通过删除 groupby 函数来达到同样的目的。

+---+-------------------------------+
|id |flatten(Array)                 |
+---+-------------------------------+
|0  |[-9.53, -9.35, 0.18, 1, 2, 3]  |
|1  |[-7.77, -7.61, 0.16, 4, 5, 6]  |
|2  |[-5.80, -5.71, 0.10, 8, 98, 32]|
+---+-------------------------------+

Answer 1

如果不想分组，可以使用窗口函数：

import pyspark.sql.functions as f
from pyspark.sql.window import Window

df2 = df.select(
    "id",
    f.flatten(f.collect_list(f.col("Array")).over(Window.partitionBy("id"))).alias("Array")
).distinct()

df2.show(truncate=False)
+---+-------------------------------+
|id |Array                          |
+---+-------------------------------+
|0  |[-9.53, -9.35, 0.18, 1, 2, 3]  |
|1  |[-7.77, -7.61, 0.16, 4, 5, 6]  |
|2  |[-5.80, -5.71, 0.10, 8, 98, 32]|
+---+-------------------------------+

你也可以试试

df.select(
    'id', 
    f.explode('Array').alias('Array')
).groupBy('id').agg(f.collect_list('Array').alias('Array'))

虽然我不确定它是否会更快。