Same Spotify queries from Module 1, but on a distributed engine. Same SQL, different execution model.
Learning goals
Run SQL queries on PySpark and compare to Postgres
Read Spark explain plans and connect them to the query plans from Module 3
Understand partitioned data lakes: read/write Parquet partitioned by key
See why distributed systems have overhead on small data but win at scale
Time: ~45 minutes
Same SQL, different engine: Run the same queries from the SQL Colab on Spark. Compare the output -- identical results, different execution.
Spark explain plans: Run explain(True) on a query. Compare to Postgres EXPLAIN from Module 3. Same concepts (hash join, scan, aggregate), different operators.
Partitioned writes: Write the Listens data partitioned by user_id. Then read only one partition. This is how data lakes at Netflix, Uber, and Airbnb work.
Scale trade-off: On our toy data, Spark is slower than Postgres (startup overhead). At what scale does Spark start winning? Think about 1TB of listens across 100 machines.
import pandas as pd
from IPython.display import display, HTML
users_data = [
{"user_id": 1, "name": "Mickey", "email": "mickey@example.com"},
{"user_id": 2, "name": "Minnie", "email": "minnie@example.com"},
{"user_id": 3, "name": "Daffy", "email": "daffy@example.com"},
{"user_id": 4, "name": "Pluto", "email": "pluto@example.com"},
]
songs_data = [
{"song_id": 1, "title": "Evermore", "artist": "Taylor Swift", "genre": "Pop"},
{"song_id": 2, "title": "Willow", "artist": "Taylor Swift", "genre": "Pop"},
{"song_id": 3, "title": "Shape of You", "artist": "Ed Sheeran", "genre": "Rock"},
{"song_id": 4, "title": "Photograph", "artist": "Ed Sheeran", "genre": "Rock"},
{"song_id": 5, "title": "Shivers", "artist": "Ed Sheeran", "genre": "Rock"},
{"song_id": 6, "title": "Yesterday", "artist": "Beatles", "genre": "Classic"},
{"song_id": 7, "title": "Yellow Submarine", "artist": "Beatles", "genre": "Classic"},
{"song_id": 8, "title": "Hey Jude", "artist": "Beatles", "genre": "Classic"},
{"song_id": 9, "title": "Bad Blood", "artist": "Taylor Swift", "genre": "Rock"},
{"song_id": 10, "title": "DJ Mix", "artist": "DJ", "genre": None},
]
listens_data = [
{"listen_id": 1, "user_id": 1, "song_id": 1, "rating": 4.5, "listen_time": "2024-08-30 14:35:00"},
{"listen_id": 2, "user_id": 1, "song_id": 2, "rating": 4.2, "listen_time": None},
{"listen_id": 3, "user_id": 1, "song_id": 6, "rating": 3.9, "listen_time": "2024-08-29 10:15:00"},
{"listen_id": 4, "user_id": 2, "song_id": 2, "rating": 4.7, "listen_time": None},
{"listen_id": 5, "user_id": 2, "song_id": 7, "rating": 4.6, "listen_time": "2024-08-28 09:20:00"},
{"listen_id": 6, "user_id": 2, "song_id": 8, "rating": 3.9, "listen_time": "2024-08-27 16:45:00"},
{"listen_id": 7, "user_id": 3, "song_id": 1, "rating": 2.9, "listen_time": None},
{"listen_id": 8, "user_id": 3, "song_id": 2, "rating": 4.9, "listen_time": "2024-08-26 12:30:00"},
{"listen_id": 9, "user_id": 3, "song_id": 6, "rating": None, "listen_time": None},
]
from pyspark.sql import SparkSession
spark = SparkSession.builder \
.master("local[*]") \
.appName("CS145 Distributed SQL") \
.getOrCreate()
print(f"Spark version: {spark.version}")
print("SparkSession created (local mode -- simulating a cluster on one machine)")
# Load data into Spark DataFrames
users_df = spark.createDataFrame(pd.DataFrame(users_data))
songs_df = spark.createDataFrame(pd.DataFrame(songs_data))
listens_clean = []
for l in listens_data:
row = l.copy()
if row["rating"] is None:
row["rating"] = float("nan")
listens_clean.append(row)
listens_df = spark.createDataFrame(pd.DataFrame(listens_clean))
users_df.createOrReplaceTempView("Users")
songs_df.createOrReplaceTempView("Songs")
listens_df.createOrReplaceTempView("Listens")
print("Tables registered")
spark.sql("SELECT * FROM Songs").show()
"""Popular songs -- exact same SQL as Postgres"""
spark.sql("""
SELECT Songs.title, Songs.artist, COUNT(*) as num_listens
FROM Songs JOIN Listens ON Songs.song_id = Listens.song_id
GROUP BY Songs.title, Songs.artist
ORDER BY num_listens DESC
""").show()
"""Average rating by genre"""
spark.sql("""
SELECT Songs.genre, ROUND(AVG(Listens.rating), 2) as avg_rating
FROM Songs JOIN Listens ON Songs.song_id = Listens.song_id
WHERE Listens.rating IS NOT NULL
GROUP BY Songs.genre
ORDER BY avg_rating DESC
""").show()
Compare this to Postgres EXPLAIN from Module 3. Same concepts -- hash join, sequential scan, aggregate -- but Spark adds exchange (shuffle) operations for distributing data across workers.
"""Explain the popular songs query"""
spark.sql("""
SELECT Songs.title, Songs.artist, COUNT(*) as num_listens
FROM Songs JOIN Listens ON Songs.song_id = Listens.song_id
GROUP BY Songs.title, Songs.artist
ORDER BY num_listens DESC
""").explain(True)
"""Explain a JOIN + GROUP BY + HAVING query"""
spark.sql("""
SELECT Songs.genre, AVG(Listens.rating) as avg_rating
FROM Songs JOIN Listens ON Songs.song_id = Listens.song_id
WHERE Listens.rating IS NOT NULL
GROUP BY Songs.genre
HAVING AVG(Listens.rating) > 4.0
""").explain(True)
Real data platforms store data as Parquet files partitioned by a key (date, user_id, genre). When you query, Spark only reads the partitions it needs -- this is called partition pruning.
"""Write Listens as partitioned Parquet files"""
import shutil, os
# Clean up previous runs
if os.path.exists("listens_partitioned"):
shutil.rmtree("listens_partitioned")
# Write partitioned by user_id
listens_df.write.partitionBy("user_id").parquet("listens_partitioned")
# Show the partition structure
for root, dirs, files in os.walk("listens_partitioned"):
level = root.replace("listens_partitioned", "").count(os.sep)
indent = " " * level
print(f"{indent}{os.path.basename(root)}/")
for f in files:
if not f.startswith("."):
print(f"{indent} {f}")
"""Read back only one partition -- Spark skips the rest"""
mickey_listens = spark.read.parquet("listens_partitioned")
mickey_listens.createOrReplaceTempView("ListensPartitioned")
# This query only reads user_id=1 partition (partition pruning)
spark.sql("""
SELECT * FROM ListensPartitioned WHERE user_id = 1
""").show()
# Look at the explain plan -- notice PartitionFilters
spark.sql("SELECT * FROM ListensPartitioned WHERE user_id = 1").explain(True)
On our 9-row Listens table, Spark is slower than Postgres. The JVM startup, query planning, and shuffle overhead dominate.
At scale (TBs of data, hundreds of machines), Spark wins because:
Data is distributed across machines -- each processes its local partition
Shuffle operations move only the data that needs to be combined
Columnar formats (Parquet) + partition pruning skip irrelevant data
This is the architecture behind data platforms at Netflix (Spark + S3), Uber (Spark + HDFS), and Airbnb (Spark + Hive).
Three-minute walkthrough.