web-dev-qa-db-ja.com

PySpark-グループ内の各行の行番号を取得します

Pysparkを使用して、sparkデータフレームをグループ化し、グループをソートして、行番号を提供できるようにしたいと考えています。したがって、

Group    Date
  A      2000
  A      2002
  A      2007
  B      1999
  B      2015

なるだろう

Group    Date    row_num
  A      2000      0
  A      2002      1
  A      2007      2
  B      1999      0
  B      2015      1
apache-sparkpysparkapache-spark-sqlspark-dataframepyspark-sql
7
Luke

ウィンドウ関数を使用：

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

df.withColumn("row_num", row_number().over(Window.partitionBy("Group").orderBy("Date")))
18
user8419108

受け入れられた解決策はそれをほぼ正しくしています。質問で要求された出力に基づく解決策は次のとおりです。

df = spark.createDataFrame([("A", 2000), ("A", 2002), ("A", 2007), ("B", 1999), ("B", 2015)], ["Group", "Date"])

+-----+----+
|Group|Date|
+-----+----+
|    A|2000|
|    A|2002|
|    A|2007|
|    B|1999|
|    B|2015|
+-----+----+

# accepted solution above


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

df.withColumn("row_num", row_number().over(Window.partitionBy("Group").orderBy("Date")))


# accepted solution above output


+-----+----+-------+
|Group|Date|row_num|
+-----+----+-------+
|    B|1999|      1|
|    B|2015|      2|
|    A|2000|      1|
|    A|2002|      2|
|    A|2007|      3|
+-----+----+-------+

ご覧のとおり、関数row_numberは0ではなく1から始まり、リクエストされた質問はrow_numを0から始めたいと思っていました。以下のような簡単な変更：

df.withColumn("row_num", row_number().over(Window.partitionBy("Group").orderBy("Date"))-1).show()

出力：

+-----+----+-------+
|Group|Date|row_num|
+-----+----+-------+
|    B|1999|      0|
|    B|2015|      1|
|    A|2000|      0|
|    A|2002|      1|
|    A|2007|      2|
+-----+----+-------+

次に、「グループ」列を任意の順序で並べ替えることができます。上記の解決策はほとんどそれを持っていますが、row_numberは0ではなく1で始まることを覚えておくことが重要です。

1
Rahul P