web-dev-qa-db-ja.com

Pysparkの最大値を持つGroupBy列とフィルター行

これは以前に尋ねられたことはほぼ確実ですが、 stackoverflowを介した検索 は私の質問に答えませんでした。 [2] の重複ではありません。なぜなら、最も頻繁なアイテムではなく、最大値が必要だからです。私はpysparkが初めてで、本当に簡単なことをしようとしています。列 "A"をgroupByし、列 "B"に最大値を持つ各グループの行のみを保持します。このような:

df_cleaned = df.groupBy("A").agg(F.max("B"))

残念ながら、これは他のすべての列を破棄します-df_cleanedには列「A」とBの最大値のみが含まれます。代わりに行を保持するにはどうすればよいですか? (「A」、「B」、「C」...)

16
Thomas

udfを使用してWindowなしでこれを行うことができます。

次の例を考えてみましょう。

import pyspark.sql.functions as f
data = [
    ('a', 5),
    ('a', 8),
    ('a', 7),
    ('b', 1),
    ('b', 3)
]
df = sqlCtx.createDataFrame(data, ["A", "B"])
df.show()
#+---+---+
#|  A|  B|
#+---+---+
#|  a|  5|
#|  a|  8|
#|  a|  7|
#|  b|  1|
#|  b|  3|
#+---+---+

Windowを作成して列Aで分割し、これを使用して各グループの最大値を計算します。次に、列Bの値が最大値に等しくなるように行をフィルタリングします。

from pyspark.sql import Window
w = Window.partitionBy('A')
df.withColumn('maxB', f.max('B').over(w))\
    .where(f.col('B') == f.col('maxB'))\
    .drop('maxB')\
    .show()
#+---+---+
#|  A|  B|
#+---+---+
#|  a|  8|
#|  b|  3|
#+---+---+

または同等にpyspark-sqlを使用:

df.registerTempTable('table')
q = "SELECT A, B FROM (SELECT *, MAX(B) OVER (PARTITION BY A) AS maxB FROM table) M WHERE B = maxB"
sqlCtx.sql(q).show()
#+---+---+
#|  A|  B|
#+---+---+
#|  b|  3|
#|  a|  8|
#+---+---+
21
pault

もう1つの可能なアプローチは、「leftsemi」を指定するデータフレーム自体を結合することです。この種類の結合には、左側のデータフレームのすべての列が含まれ、右側の列は含まれません。

例えば:

_import pyspark.sql.functions as f
data = [
    ('a', 5, 'c'),
    ('a', 8, 'd'),
    ('a', 7, 'e'),
    ('b', 1, 'f'),
    ('b', 3, 'g')
]
df = sqlContext.createDataFrame(data, ["A", "B", "C"])
df.show()
+---+---+---+
|  A|  B|  C|
+---+---+---+
|  a|  5|  c|
|  a|  8|  d|
|  a|  7|  e|
|  b|  1|  f|
|  b|  3|  g|
+---+---+---+
_

列Aごとの列Bの最大値は、次のようにして選択できます。

_df.groupBy('A').agg(f.max('B')
+---+---+
|  A|  B|
+---+---+
|  a|  8|
|  b|  3|
+---+---+
_

この式を左半結合の右側として使用し、取得した列max(B)の名前を元の名前Bに戻すと、必要な結果を取得できます。

_df.join(df.groupBy('A').agg(f.max('B').alias('B')),on='B',how='leftsemi').show()
+---+---+---+
|  B|  A|  C|
+---+---+---+
|  3|  b|  g|
|  8|  a|  d|
+---+---+---+
_

このソリューションの背後にある物理的な計画と受け入れられた答えからのものは異なり、大きなデータフレームでどのパフォーマンスが優れているかはまだわかりません。

spark SQL構文を実行すると、同じ結果が得られます。

_df.registerTempTable('table')
q = '''SELECT *
FROM table a LEFT SEMI
JOIN (
    SELECT 
        A,
        max(B) as max_B
    FROM table
    GROUP BY A
    ) t
ON a.A=t.A AND a.B=t.max_B
'''
sqlContext.sql(q).show()
+---+---+---+
|  A|  B|  C|
+---+---+---+
|  b|  3|  g|
|  a|  8|  d|
+---+---+---+
_
4
ndricca