PySpark practice

PySpark groupBy practice questions

8 problems, ordered from easiest. They run in your browser, and 2 are free.

The idea in a minute

groupBy collects rows into groups, and agg says what to calculate for each group. You can ask for several results in one pass, such as a count, a sum and an average.

The questions here are the PySpark twin of GROUP BY in SQL, so they are a good place to compare the two.

Example: Orders and revenue per city

from pyspark.sql import SparkSession, functions as F
spark = SparkSession.builder.master("local[2]").config("spark.ui.enabled","false").getOrCreate()
spark.sparkContext.setLogLevel("ERROR")
orders = spark.createDataFrame([(1,"Asha","Pune",450),(2,"Ravi","Pune",300),(3,"Meera","Delhi",900),(4,"Asha","Delhi",150),(5,"Ravi","Pune",250)],["order_id","customer","city","amount"])
(orders.groupBy("city")
       .agg(F.count("*").alias("orders"), F.sum("amount").alias("revenue"))
       .orderBy("city")
       .show())

What it prints

+-----+------+-------+
| city|orders|revenue|
+-----+------+-------+
|Delhi|     2|   1050|
| Pune|     3|   1000|
+-----+------+-------+

What to watch for

Give each result a name with alias, or the column will be called sum(amount).
groupBy returns a grouped object, not a DataFrame. You get a DataFrame back only after agg.
Sort with orderBy after the aggregation if you need a fixed order.

Aggregation problems

Want a guided order across every topic? Open the practice sheets or the PySpark hub.