Retrieve top n in each group of a DataFrame in pyspark

by

I believe you need to use window functions to attain the rank of each row based on user_id and score, and subsequently filter your results to only keep the first two values.

from pyspark.sql.window import Window
from pyspark.sql.functions import rank, col

window = Window.partitionBy(df['user_id']).orderBy(df['score'].desc())

df.select('*', rank().over(window).alias('rank')) 
  .filter(col('rank') <= 2) 
  .show() 
#+-------+---------+-----+----+
#|user_id|object_id|score|rank|
#+-------+---------+-----+----+
#| user_1| object_1|    3|   1|
#| user_1| object_2|    2|   2|
#| user_2| object_2|    6|   1|
#| user_2| object_1|    5|   2|
#+-------+---------+-----+----+

In general, the official programming guide is a good place to start learning Spark.

Data

rdd = sc.parallelize([("user_1",  "object_1",  3), 
                      ("user_1",  "object_2",  2), 
                      ("user_2",  "object_1",  5), 
                      ("user_2",  "object_2",  2), 
                      ("user_2",  "object_2",  6)])
df = sqlContext.createDataFrame(rdd, ["user_id", "object_id", "score"])

More Related Contents:

  1. How to add a constant column in a Spark DataFrame?
  2. Spark Dataframe distinguish columns with duplicated name
  3. Pyspark: Split multiple array columns into rows
  4. How do I add a new column to a Spark DataFrame (using PySpark)?
  5. How to change a dataframe column from String type to Double type in PySpark?
  6. Count number of non-NaN entries in each column of Spark dataframe with Pyspark
  7. Updating a dataframe column in spark
  8. Dividing complex rows of dataframe to simple rows in Pyspark
  9. Filter Pyspark dataframe column with None value
  10. PySpark converting a column of type ‘map’ to multiple columns in a dataframe
  11. How to explode multiple columns of a dataframe in pyspark
  12. Spark add new column to dataframe with value from previous row
  13. PySpark: multiple conditions in when clause
  14. Pyspark: Replacing value in a column by searching a dictionary
  15. Count number of non-NaN entries in each column of Spark dataframe in PySpark
  16. Create Spark DataFrame. Can not infer schema for type
  17. Pivot String column on Pyspark Dataframe
  18. Filtering DataFrame using the length of a column
  19. Multiple condition filter on dataframe
  20. How to return a “Tuple type” in a UDF in PySpark?
  21. How to change dataframe column names in pyspark?
  22. How to export a table dataframe in PySpark to csv?
  23. Pyspark: explode json in column to multiple columns
  24. Pyspark ‘NoneType’ object has no attribute ‘_jvm’ error
  25. Cast column containing multiple string date formats to DateTime in Spark
  26. Create single row dataframe from list of list PySpark
  27. Concatenate two PySpark dataframes
  28. How do I convert an array (i.e. list) column to Vector
  29. How to drop all columns with null values in a PySpark DataFrame?
  30. how to add Row id in pySpark dataframes [duplicate]

Leave a Comment