计算并比较两列的平均值

时间:2021-02-15 14:40:52

标签: python apache-spark pyspark apache-spark-sql apache-spark-ml

我开始将我的 Pandas 实现转换为 pySpark,但是我在完成一些基本操作时遇到了麻烦。所以我有这张桌子:

+-----+-----+----+
| Col1|Col2 |Col3|
+-----+-----+----+
|  1  |[1,3]|   0|
|  44 |[2,0]|   1|
|  77 |[1,5]|   7|
+-----+-----+----+

我想要的输出是:

+-----+-----+----+----+
| Col1|Col2 |Col3|Col4|
+-----+-----+----+----+
|  1  |[1,3]|   0|2.67|
|  44 |[2,0]|   1|2.67|
|  77 |[1,5]|   7|2.67|
+-----+-----+----+----+

到达这里:

  • 我对 Col2 中每个数组的第一项求平均值,并对 Col2 中每个数组的第二项求平均值。由于第二个“子列”的平均值((3+0+5)/3)比第一个“子列”((1+2+1)/3)大,这是“获胜”条件.之后,我创建了一个新列,该列将“获胜”平均值复制到该表的行数(在本例中为 3)。 我已经能够通过“手动”选择 ta 列,对其进行平均,然后使用“点亮”来复制结果来做到这一点。我的实现的问题是 collect() 需要很多时间并且不推荐它。 你能帮我解决这个问题吗?

1 个答案:

答案 0 :(得分:1)

您可以使用 greatest 获得数组中每个(子)列的最大平均值:

from pyspark.sql import functions as F, Window

df2 = df.withColumn(
    'Col4',
    F.greatest(*[F.avg(F.udf(lambda r: [float(i) for i in r.toArray()], 'array<double>')('Col2')[i]).over(Window.orderBy()) for i in range(2)])
)

df2.show()
+----+------+----+------------------+
|Col1|  Col2|Col3|              Col4|
+----+------+----+------------------+
|   1|[1, 3]|   0|2.6666666666666665|
|  44|[2, 0]|   1|2.6666666666666665|
|  77|[1, 5]|   7|2.6666666666666665|
+----+------+----+------------------+

如果你希望数组大小是动态的,你可以这样做

arr_size = df.select(F.max(F.size(F.udf(lambda r: [float(i) for i in r.toArray()], 'array<double>')('Col2')))).head()[0]

df2 = df.withColumn(
    'Col4',
    F.greatest(*[F.avg(F.udf(lambda r: [float(i) for i in r.toArray()], 'array<double>')('Col2')[i]).over(Window.orderBy()) for i in range(arr_size)])
)