PySpark聚合和何时条件

Question

我对PySpark有疑问。

df = (sc.parallelize([
    {"Day":1,"sensitive_id":"1234", "num":3},
    {"Day":1,"sensitive_id":"1234", "num":3}, 
    {"Day":2,"sensitive_id":"1234", "num":3},
    {"Day":3,"sensitive_id":"2345", "num":2},
    {"Day":3,"sensitive_id":"2345", "num":2},
    {"Day":3,"sensitive_id":"6789", "num":4},
    {"Day":4,"sensitive_id":"6789", "num":4},
    {"Day":4,"sensitive_id":"6789", "num":4},
    {"Day":4,"sensitive_id":"6789", "num":4}
 ]).toDF()
      )

我希望有一个新列，以具有对应的“ sensitive_id”作为“ num”列的最大值。

这是我到目前为止的结果。

from pyspark.sql import functions as F
from pyspark.sql.window import Window
(
    df 
    .withColumn(
        'top_sensitive_id',
        F.when
        (
            F.col('num') == F.max(F.col('num')).over(Window.partitionBy(F.col('Day'))),
            F.col('sensitive_id')
        )
    )
    .withColumn
    (
        'top',
        F.max(F.col('top_sensitive_id')).over(Window.partitionBy(F.col('Day')))
    )

    .show()
)

但是我相信应该有一种更优雅，更有效的方法来实现这一目标。

有人可以提供更好的方法吗？

Answer 1

您的代码几乎接近最佳方法，但是我仍然尝试添加一些内容-

1。计算一次“ top”，并将其用于比较。

2。使用单独的列定义，这将有助于提高可读性    和可维护性

 from pyspark.sql.window import Window

windowSpec = Window.partitionBy('Day')
top = max('sensitive_id').over(windowSpec).alias('top')

df.select('Day','Num','sensitive_id',top).withColumn('top_sensitive_id', expr("CASE WHEN sensitive_id = top THEN top END")).show()