在Spark DataFrame中查找每组的最大行数

Question

在一个14节点的Google Dataproc集群中，我有大约600万个名字，它们被两个不同的系统转换为ID:SA和SB。每个行包含名称、ID_SA和ID_SB。我的目标是生成从ID_SA到ID_SB的映射，以便对于每个ID_SA，对应的ID_SB是附加到ID_SA的所有名称中最常见的id。

让我们试着用一个例子来澄清一下。如果我有以下行：

[Row(name='n1', id_sa='a1', id_sb='b1'),
 Row(name='n2', id_sa='a1', id_sb='b2'),
 Row(name='n3', id_sa='a1', id_sb='b2'),
 Row(name='n4', id_sa='a2', id_sb='b2')]

我的目标是生成从A1到B2的映射。实际上，与A1关联的名称是N1、N2和N3，它们分别映射到B1、B2，因此B2是与A1关联的名称中最常见的映射。同样，A2将映射到B2。假设总会有赢家是可以的：没有必要打破联系。

我希望可以在我的dataframe上使用groupby(df.id_sa)，但我不知道下一步该怎么做。我希望聚合最终能够产生以下行：

[Row(id_sa=a1, max_id_sb=b2),
 Row(id_sa=a2, max_id_sb=b2)]

但也许我试图使用错误的工具，我应该回到使用RDDS。

Answer 1

使用join（如果有联系，将导致组中有多行）：

import pyspark.sql.functions as F
from pyspark.sql.functions import count, col 

cnts = df.groupBy("id_sa", "id_sb").agg(count("*").alias("cnt")).alias("cnts")
maxs = cnts.groupBy("id_sa").agg(F.max("cnt").alias("mx")).alias("maxs")

cnts.join(maxs, 
  (col("cnt") == col("mx")) & (col("cnts.id_sa") == col("maxs.id_sa"))
).select(col("cnts.id_sa"), col("cnts.id_sb"))

使用窗口函数（将拖放领带）：

from pyspark.sql.functions import row_number
from pyspark.sql.window import Window

w = Window().partitionBy("id_sa").orderBy(col("cnt").desc())

(cnts
  .withColumn("rn", row_number().over(w))
  .where(col("rn") == 1)
  .select("id_sa", "id_sb"))

使用结构排序：

from pyspark.sql.functions import struct

(cnts
  .groupBy("id_sa")
  .agg(F.max(struct(col("cnt"), col("id_sb"))).alias("max"))
  .select(col("id_sa"), col("max.id_sb")))