web-dev-qa-db-ja.com

Spark DataFrameでグループごとの最大行を見つける

RDDの代わりにSparkデータフレームを使用しようとしています。それらはRDDよりも高レベルであり、より読みやすいコードを生成する傾向があるためです。

14ノードのGoogle Dataprocクラスターには、約600万の名前があり、2つの異なるシステムsaとsbによってIDに変換されます。各Rowには、name、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にマップされているため、b2がa1に関連付けられた名前の中で最も頻繁にマッピングされます。同様に、a2はb2にマップされます。常に勝者がいると仮定しても構いません。関係を壊す必要はありません。

データフレームでgroupBy(df.id_sa)を使用できることを望んでいましたが、次に何をすべきかわかりません。最終的に次の行を生成できる集計を期待していました。

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

しかし、間違ったツールを使用しようとしているため、RDDの使用に戻る必要があります。

apache-sparkpysparkapache-spark-sql
42
Quentin Pradet

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"))

struct順序の使用：

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")))

各グループの最初の行を選択する方法 も参照してください。

52
zero323

あなたが探しているのはウィンドウ関数だと思います： http://spark.Apache.org/docs/latest/api/python/pyspark.sql.html?highlight=window#pyspark.sql.Window

https://databricks.com/blog/2015/07/15/introducing-window-functions-in-spark-sql.html

Scalaの例を次に示します（現在、Hiveを使用できるSpark Shellがないため、コードをテストできませんでしたが、動作するはずです） ：

case class MyRow(name: String, id_sa: String, id_sb: String)

val myDF = sc.parallelize(Array(
    MyRow("n1", "a1", "b1"),
    MyRow("n2", "a1", "b2"),
    MyRow("n3", "a1", "b2"),
    MyRow("n1", "a2", "b2")
)).toDF("name", "id_sa", "id_sb")

import org.Apache.spark.sql.expressions.Window

val windowSpec = Window.partitionBy(myDF("id_sa")).orderBy(myDF("id_sb").desc)

myDF.withColumn("max_id_b", first(myDF("id_sb")).over(windowSpec).as("max_id_sb")).filter("id_sb = max_id_sb")

Window関数で同じ結果を達成するためのより効率的な方法はおそらくありますが、これが正しい方向を示していることを願っています。

9
alghimo