Python + Spark Practice

pyspark-practice-problems.ipynb

Read onlyPython 3 (PySpark)Download

110 cells · 47 code · 20 saved outputs

Trusted · exact notebook preview

PySpark Transformations — 20 Exhaustive Practice Problems

Coverage: select, filter, withColumn, groupBy, agg, join, window functions, pivot, union, explode, UDFs, orderBy, distinct, dropDuplicates, cast, when/otherwise, null handling, repartition, cache, and more.

Dummy data is pre-loaded. Write your transformations in the # YOUR CODE HERE cells. Solutions are provided at the bottom of this notebook.

In [ ]

PySpark · In [ ]cell 2

from pyspark.sql import SparkSession, functions as F
from pyspark.sql.window import Window
from pyspark.sql.types import (
    DateType,
    DoubleType,
    IntegerType,
    StringType,
    StructField,
    StructType,
    TimestampType,
)

Setup — Run This First

In [ ]

PySpark · In [ ]cell 4

spark = (
    SparkSession.builder
    .appName("SparkPractice")
    .master("local[*]")
    .getOrCreate()
)

spark.sparkContext.setLogLevel("ERROR")
print(f"Spark Ready! Version: {spark.version}")

Dummy Data Setup

In [ ]

PySpark · In [ ]cell 6

# ─── employees DataFrame ───────────────────────────────────────────────────
employees_data = [
    (1, "Alice", "Engineering", "India", 95000.0, 5, "2019-03-15"),
    (2, "Bob", "Marketing", "US", 72000.0, 3, "2021-06-01"),
    (3, "Charlie", "Engineering", "US", 110000.0, 8, "2016-01-20"),
    (4, "Diana", "HR", "India", 60000.0, 2, "2022-09-10"),
    (5, "Eve", "Engineering", "Germany", 120000.0, 10, "2014-07-04"),
    (6, "Frank", "Marketing", "India", 68000.0, 4, "2020-11-30"),
    (7, "Grace", "HR", "US", 58000.0, 1, "2023-02-14"),
    (8, "Hank", "Engineering", "Germany", 98000.0, 6, "2018-05-22"),
    (9, "Ivy", "Marketing", "Germany", 75000.0, 5, "2019-08-09"),
    (10, "Jack", "HR", "US", None, 3, "2021-03-17"),
]
employeesDF = spark.createDataFrame(
    employees_data,
    ["emp_id", "name", "department", "country", "salary", "years_exp", "join_date"],
)

# ─── sales DataFrame ───────────────────────────────────────────────────────
sales_data = [
    (101, 1, "Laptop", "Electronics", 1200.0, 2, "2024-01-05"),
    (102, 2, "Phone", "Electronics", 800.0, 1, "2024-01-07"),
    (103, 1, "Desk", "Furniture", 450.0, 3, "2024-01-10"),
    (104, 3, "Laptop", "Electronics", 1200.0, 1, "2024-01-12"),
    (105, 4, "Chair", "Furniture", 250.0, 4, "2024-02-01"),
    (106, 5, "Laptop", "Electronics", 1200.0, 5, "2024-02-03"),
    (107, 2, "Monitor", "Electronics", 600.0, 2, "2024-02-15"),
    (108, 6, "Keyboard", "Electronics", 80.0, 10, "2024-03-01"),
    (109, 7, "Desk", "Furniture", 450.0, 1, "2024-03-05"),
    (110, 1, "Phone", "Electronics", 800.0, 1, "2024-03-10"),
    (111, 8, "Chair", "Furniture", 250.0, 6, "2024-03-15"),
    (112, 9, "Monitor", "Electronics", 600.0, 3, "2024-04-01"),
]
salesDF = spark.createDataFrame(
    sales_data,
    ["sale_id", "emp_id", "product", "category", "price", "quantity", "sale_date"],
)

# ─── departments DataFrame ────────────────────────────────────────────────
departments_data = [
    ("Engineering", "Alice Johnson", 500000.0),
    ("Marketing", "Bob Williams", 300000.0),
    ("HR", "Carol Smith", 150000.0),
    ("Finance", "David Brown", 200000.0),  # No employees in this dept
]
departmentsDF = spark.createDataFrame(
    departments_data, ["dept_name", "dept_head", "budget"]
)

# ─── skills DataFrame (array column) ─────────────────────────────────────
skills_data = [
    (1, ["Scala", "Spark", "SQL"]),
    (2, ["Marketing", "SEO"]),
    (3, ["Java", "Scala", "Kafka", "Spark"]),
    (4, ["HR", "Recruiting"]),
    (5, ["Scala", "Python", "ML", "Spark"]),
]
skillsDF = spark.createDataFrame(skills_data, ["emp_id", "skills"])

employeesDF.createOrReplaceTempView("employees")
salesDF.createOrReplaceTempView("sales")
departmentsDF.createOrReplaceTempView("departments")
skillsDF.createOrReplaceTempView("skills")

print("All DataFrames Ready!")
print("\n--- Employees ---")
employeesDF.show()
print("\n--- Sales ---")
salesDF.show()
print("\n--- Departments ---")
departmentsDF.show()
print("\n--- Skills ---")
skillsDF.show(truncate=False)

PROBLEMS

Q1 — select, filter, withColumn

From employeesDF:

Select name, department, salary
Filter only employees in Engineering with salary > 90000
Add a new column salary_inr = salary × 83 (USD to INR)
Show the result

In [ ]

PySpark · In [ ]cell 8

# YOUR CODE HERE

Expected output

text

+-------+-----------+--------+----------+
|   name| department|  salary|salary_inr|
+-------+-----------+--------+----------+
|  Alice|Engineering| 95000.0| 7885000.0|
|Charlie|Engineering|110000.0| 9130000.0|
|    Eve|Engineering|120000.0| 9960000.0|
|   Hank|Engineering| 98000.0| 8134000.0|
+-------+-----------+--------+----------+

Q2 — groupBy + agg (count, avg, max, min, sum)

From employeesDF, for each department compute:

total employees (count)
average salary (avg)
max salary
min salary
total salary budget (sum)

Round avg to 2 decimal places. Sort by avg salary descending.

In [ ]

PySpark · In [ ]cell 11

# YOUR CODE HERE

Expected output

text

+-----------+---------------+----------+----------+----------+------------+
| department|total_employees|avg_salary|max_salary|min_salary|total_salary|
+-----------+---------------+----------+----------+----------+------------+
|Engineering|              4|  105750.0|  120000.0|   95000.0|    423000.0|
|  Marketing|              3|  71666.67|   75000.0|   68000.0|    215000.0|
|         HR|              3|   59000.0|   60000.0|   58000.0|    118000.0|
+-----------+---------------+----------+----------+----------+------------+

Q3 — Inner, Left, Right & Anti Join

Inner Join employeesDF with salesDF on emp_id → show employee name and their sales
Left Join employeesDF with salesDF → find employees who made NO sales (anti-join pattern using left + filter null)
Anti Join — same as above but use anti join type directly
Left Join departmentsDF with employeesDF → show departments with no employees

In [ ]

PySpark · In [ ]cell 14

# YOUR CODE HERE

Expected output

text

=== Inner Join ===
+-------+-----------+--------+------+--------+----------+
|   name| department| product| price|quantity| sale_date|
+-------+-----------+--------+------+--------+----------+
|  Alice|Engineering|  Laptop|1200.0|       2|2024-01-05|
|  Alice|Engineering|    Desk| 450.0|       3|2024-01-10|
|  Alice|Engineering|   Phone| 800.0|       1|2024-03-10|
|    Bob|  Marketing|   Phone| 800.0|       1|2024-01-07|
|    Bob|  Marketing| Monitor| 600.0|       2|2024-02-15|
|Charlie|Engineering|  Laptop|1200.0|       1|2024-01-12|
|  Diana|         HR|   Chair| 250.0|       4|2024-02-01|
|    Eve|Engineering|  Laptop|1200.0|       5|2024-02-03|
|  Frank|  Marketing|Keyboard|  80.0|      10|2024-03-01|
|  Grace|         HR|    Desk| 450.0|       1|2024-03-05|
|   Hank|Engineering|   Chair| 250.0|       6|2024-03-15|
|    Ivy|  Marketing| Monitor| 600.0|       3|2024-04-01|
+-------+-----------+--------+------+--------+----------+

=== No Sales (Left + null filter) ===
+------+----+----------+
|emp_id|name|department|
+------+----+----------+
|    10|Jack|        HR|
+------+----+----------+

=== Anti Join ===
+------+----+----------+
|emp_id|name|department|
+------+----+----------+
|    10|Jack|        HR|
+------+----+----------+

=== Departments with no Employees ===
+---------+-----------+
|dept_name|  dept_head|
+---------+-----------+
|  Finance|David Brown|
+---------+-----------+

Q4 — Window Functions: rank, dense_rank, row_number

From employeesDF:

Rank employees by salary within each department (highest first)
Add columns: rank, dense_rank, row_number — all partitioned by department, ordered by salary desc
Filter to show only the top 2 employees per department by dense_rank

In [ ]

PySpark · In [ ]cell 17

# YOUR CODE HERE

Expected output

text

=== All Ranks ===
+-------+-----------+--------+----+----------+----------+
|   name| department|  salary|rank|dense_rank|row_number|
+-------+-----------+--------+----+----------+----------+
|    Eve|Engineering|120000.0|   1|         1|         1|
|Charlie|Engineering|110000.0|   2|         2|         2|
|   Hank|Engineering| 98000.0|   3|         3|         3|
|  Alice|Engineering| 95000.0|   4|         4|         4|
|  Diana|         HR| 60000.0|   1|         1|         1|
|  Grace|         HR| 58000.0|   2|         2|         2|
|   Jack|         HR|    NULL|   3|         3|         3|
|    Ivy|  Marketing| 75000.0|   1|         1|         1|
|    Bob|  Marketing| 72000.0|   2|         2|         2|
|  Frank|  Marketing| 68000.0|   3|         3|         3|
+-------+-----------+--------+----+----------+----------+

=== Top 2 per Department ===
+-------+-----------+--------+----------+
|   name| department|  salary|dense_rank|
+-------+-----------+--------+----------+
|    Eve|Engineering|120000.0|         1|
|Charlie|Engineering|110000.0|         2|
|  Diana|         HR| 60000.0|         1|
|  Grace|         HR| 58000.0|         2|
|    Ivy|  Marketing| 75000.0|         1|
|    Bob|  Marketing| 72000.0|         2|
+-------+-----------+--------+----------+

Q5 — Window Functions: lag, lead, running total

From salesDF (cast sale_date as Date, order by it per emp_id):

Add prev_sale_price = lag of price by 1, partitioned by emp_id
Add next_sale_price = lead of price by 1, partitioned by emp_id
Add running_revenue = cumulative sum of price * quantity ordered by sale_date

In [ ]

PySpark · In [ ]cell 20

# YOUR CODE HERE

Expected output

text

+------+--------+------+--------+----------+---------------+---------------+---------------+
|emp_id| product| price|quantity| sale_date|prev_sale_price|next_sale_price|running_revenue|
+------+--------+------+--------+----------+---------------+---------------+---------------+
|     1|  Laptop|1200.0|       2|2024-01-05|           NULL|          450.0|         2400.0|
|     1|    Desk| 450.0|       3|2024-01-10|         1200.0|          800.0|         3750.0|
|     1|   Phone| 800.0|       1|2024-03-10|          450.0|           NULL|         4550.0|
|     2|   Phone| 800.0|       1|2024-01-07|           NULL|          600.0|          800.0|
|     2| Monitor| 600.0|       2|2024-02-15|          800.0|           NULL|         2000.0|
|     3|  Laptop|1200.0|       1|2024-01-12|           NULL|           NULL|         1200.0|
|     4|   Chair| 250.0|       4|2024-02-01|           NULL|           NULL|         1000.0|
|     5|  Laptop|1200.0|       5|2024-02-03|           NULL|           NULL|         6000.0|
|     6|Keyboard|  80.0|      10|2024-03-01|           NULL|           NULL|          800.0|
|     7|    Desk| 450.0|       1|2024-03-05|           NULL|           NULL|          450.0|
|     8|   Chair| 250.0|       6|2024-03-15|           NULL|           NULL|         1500.0|
|     9| Monitor| 600.0|       3|2024-04-01|           NULL|           NULL|         1800.0|
+------+--------+------+--------+----------+---------------+---------------+---------------+

Q6 — when / otherwise (conditional column)

From employeesDF, create a new column salary_band:

"High" if salary >= 100000
"Mid" if salary >= 70000
"Low" if salary >= 0
"Unknown" if salary is null

Also create experience_level:

"Senior" if years_exp >= 7
"Mid" if years_exp >= 4
"Junior" otherwise

In [ ]

PySpark · In [ ]cell 23

# YOUR CODE HERE

Expected output

text

+-------+--------+-----------+---------+----------------+
|   name|  salary|salary_band|years_exp|experience_level|
+-------+--------+-----------+---------+----------------+
|  Alice| 95000.0|        Mid|        5|             Mid|
|    Bob| 72000.0|        Mid|        3|          Junior|
|Charlie|110000.0|       High|        8|          Senior|
|  Diana| 60000.0|        Low|        2|          Junior|
|    Eve|120000.0|       High|       10|          Senior|
|  Frank| 68000.0|        Low|        4|             Mid|
|  Grace| 58000.0|        Low|        1|          Junior|
|   Hank| 98000.0|        Mid|        6|             Mid|
|    Ivy| 75000.0|        Mid|        5|             Mid|
|   Jack|    NULL|    Unknown|        3|          Junior|
+-------+--------+-----------+---------+----------------+

Q7 — Null Handling: isNull, isNotNull, fillna, dropna, coalesce

From employeesDF:

Count rows where salary is null
Fill null salaries with the average salary of the dataset
Use coalesce to replace null salary with 60000.0 as fallback
Drop rows where any column is null
Drop rows where salary and country are both null (subset dropna)

In [ ]

PySpark · In [ ]cell 26

# YOUR CODE HERE

Expected output

text

Null salary count: 1
=== Filled with avg ===
+-------+--------+
|   name|  salary|
+-------+--------+
|  Alice| 95000.0|
|    Bob| 72000.0|
|Charlie|110000.0|
|  Diana| 60000.0|
|    Eve|120000.0|
|  Frank| 68000.0|
|  Grace| 58000.0|
|   Hank| 98000.0|
|    Ivy| 75000.0|
|   Jack| 84000.0|
+-------+--------+

=== Coalesce ===
+-------+--------+
|   name|  salary|
+-------+--------+
|  Alice| 95000.0|
|    Bob| 72000.0|
|Charlie|110000.0|
|  Diana| 60000.0|
|    Eve|120000.0|
|  Frank| 68000.0|
|  Grace| 58000.0|
|   Hank| 98000.0|
|    Ivy| 75000.0|
|   Jack| 60000.0|
+-------+--------+

After na.drop (any): 9 rows
After subset na.drop (all): 10 rows

Q8 — String Functions

From employeesDF:

Uppercase name, lowercase department
Concatenate name and department as "Alice (Engineering)"
Extract first 3 chars of name as name_code
Check if name contains the letter "a" (case-insensitive) → boolean column
Replace "India" with "IN" and "US" with "USA" in the country column using regexp_replace

In [ ]

PySpark · In [ ]cell 29

# YOUR CODE HERE

In [ ]

PySpark · In [ ]cell 30

# YOUR CODE HERE

In [ ]

PySpark · In [ ]cell 31

# YOUR CODE HERE

Expected output

text

+-------+----------+-----------+--------------------+---------+-----+------------+
|   name|name_upper| dept_lower|           name_dept|name_code|has_a|country_code|
+-------+----------+-----------+--------------------+---------+-----+------------+
|  Alice|     ALICE|engineering| Alice (Engineering)|      Ali| true|          IN|
|    Bob|       BOB|  marketing|     Bob (Marketing)|      Bob|false|         USA|
|Charlie|   CHARLIE|engineering|Charlie (Engineer...|      Cha| true|         USA|
|  Diana|     DIANA|         hr|          Diana (HR)|      Dia| true|          IN|
|    Eve|       EVE|engineering|   Eve (Engineering)|      Eve|false|     Germany|
|  Frank|     FRANK|  marketing|   Frank (Marketing)|      Fra| true|          IN|
|  Grace|     GRACE|         hr|          Grace (HR)|      Gra| true|         USA|
|   Hank|      HANK|engineering|  Hank (Engineering)|      Han| true|     Germany|
|    Ivy|       IVY|  marketing|     Ivy (Marketing)|      Ivy|false|     Germany|
|   Jack|      JACK|         hr|           Jack (HR)|      Jac| true|         USA|
+-------+----------+-----------+--------------------+---------+-----+------------+

Q9 — Date Functions

From employeesDF (treat join_date as a date):

Cast join_date to DateType
Extract year, month, day as separate columns
Calculate tenure_days = today's date minus join_date
Calculate tenure_years (rounded to 1 decimal)
Add 30 days to join_date as probation_end

In [ ]

PySpark · In [ ]cell 34

# YOUR CODE HERE

In [ ]

PySpark · In [ ]cell 35

# YOUR CODE HERE

Expected output

text

+-------+----------+---------+----------+--------+-----------+------------+-------------+
|   name| join_date|join_year|join_month|join_day|tenure_days|tenure_years|probation_end|
+-------+----------+---------+----------+--------+-----------+------------+-------------+
|  Alice|2019-03-15|     2019|         3|      15|       2764|         7.6|   2019-04-14|
|    Bob|2021-06-01|     2021|         6|       1|       1955|         5.4|   2021-07-01|
|Charlie|2016-01-20|     2016|         1|      20|       3914|        10.7|   2016-02-19|
|  Diana|2022-09-10|     2022|         9|      10|       1489|         4.1|   2022-10-10|
|    Eve|2014-07-04|     2014|         7|       4|       4479|        12.3|   2014-08-03|
|  Frank|2020-11-30|     2020|        11|      30|       2138|         5.9|   2020-12-30|
|  Grace|2023-02-14|     2023|         2|      14|       1332|         3.6|   2023-03-16|
|   Hank|2018-05-22|     2018|         5|      22|       3061|         8.4|   2018-06-21|
|    Ivy|2019-08-09|     2019|         8|       9|       2617|         7.2|   2019-09-08|
|   Jack|2021-03-17|     2021|         3|      17|       2031|         5.6|   2021-04-16|
+-------+----------+---------+----------+--------+-----------+------------+-------------+

Q10 — explode + array functions

From skillsDF:

explode the skills array → one row per skill per employee
Find how many skills each employee has (size of array)
Find all employees who have "Spark" in their skills (array_contains)
Join with employeesDF to show employee name alongside their individual skills
collect_list — after explode and join, re-aggregate skills back per employee

In [ ]

PySpark · In [ ]cell 38

# YOUR CODE HERE

In [ ]

PySpark · In [ ]cell 39

# YOUR CODE HERE

Expected output

text

=== Exploded ===
+------+----------+
|emp_id|     skill|
+------+----------+
|     1|     Scala|
|     1|     Spark|
|     1|       SQL|
|     2| Marketing|
|     2|       SEO|
|     3|      Java|
|     3|     Scala|
|     3|     Kafka|
|     3|     Spark|
|     4|        HR|
|     4|Recruiting|
|     5|     Scala|
|     5|    Python|
|     5|        ML|
|     5|     Spark|
+------+----------+

=== Skill Count ===
+------+--------------------+-----------+
|emp_id|              skills|skill_count|
+------+--------------------+-----------+
|     1| [Scala, Spark, SQL]|          3|
|     2|    [Marketing, SEO]|          2|
|     3|[Java, Scala, Kaf...|          4|
|     4|    [HR, Recruiting]|          2|
|     5|[Scala, Python, M...|          4|
+------+--------------------+-----------+

=== Spark Developers ===
+------+--------------------+
|emp_id|              skills|
+------+--------------------+
|     1| [Scala, Spark, SQL]|
|     3|[Java, Scala, Kaf...|
|     5|[Scala, Python, M...|
+------+--------------------+

=== Name + Skills ===
+------+-------+----------+
|emp_id|   name|     skill|
+------+-------+----------+
|     1|  Alice|     Scala|
|     1|  Alice|     Spark|
|     1|  Alice|       SQL|
|     2|    Bob| Marketing|
|     2|    Bob|       SEO|
|     3|Charlie|      Java|
|     3|Charlie|     Scala|
|     3|Charlie|     Kafka|
|     3|Charlie|     Spark|
|     4|  Diana|        HR|
|     4|  Diana|Recruiting|
|     5|    Eve|     Scala|
|     5|    Eve|    Python|
|     5|    Eve|        ML|
|     5|    Eve|     Spark|
+------+-------+----------+

=== Re-aggregated ===
+------+-------+---------------------------+
|emp_id|name   |skills_list                |
+------+-------+---------------------------+
|1     |Alice  |[Scala, Spark, SQL]        |
|2     |Bob    |[Marketing, SEO]           |
|3     |Charlie|[Java, Scala, Kafka, Spark]|
|4     |Diana  |[HR, Recruiting]           |
|5     |Eve    |[Scala, Python, ML, Spark] |
+------+-------+---------------------------+

Q11 — Pivot Table

From salesDF:

Pivot on category (columns: Electronics, Furniture)
Rows = emp_id
Values = total revenue (sum of price * quantity)
Fill nulls with 0
Show the pivoted result

In [ ]

PySpark · In [ ]cell 42

# YOUR CODE HERE

Expected output

text

+------+-----------+---------+
|emp_id|Electronics|Furniture|
+------+-----------+---------+
|     7|        0.0|    450.0|
|     6|      800.0|      0.0|
|     9|     1800.0|      0.0|
|     5|     6000.0|      0.0|
|     1|     3200.0|   1350.0|
|     3|     1200.0|      0.0|
|     8|        0.0|   1500.0|
|     2|     2000.0|      0.0|
|     4|        0.0|   1000.0|
+------+-----------+---------+

Q12 — Union, UnionByName, Intersect, Except

Create two subsets from employeesDF: engDF (Engineering) and mktDF (Marketing)
union them and show result
unionByName — demonstrate with a reordered mktDF
Create highSalaryDF (salary > 80000) and seniorDF (years_exp >= 5)
Use intersect → employees both high-salary AND senior
Use except → high-salary employees who are NOT senior

In [ ]

PySpark · In [ ]cell 45

# YOUR CODE HERE

Expected output

text

=== Union ===
+------+-------+-----------+-------+--------+---------+----------+
|emp_id|   name| department|country|  salary|years_exp| join_date|
+------+-------+-----------+-------+--------+---------+----------+
|     1|  Alice|Engineering|  India| 95000.0|        5|2019-03-15|
|     3|Charlie|Engineering|     US|110000.0|        8|2016-01-20|
|     5|    Eve|Engineering|Germany|120000.0|       10|2014-07-04|
|     8|   Hank|Engineering|Germany| 98000.0|        6|2018-05-22|
|     2|    Bob|  Marketing|     US| 72000.0|        3|2021-06-01|
|     6|  Frank|  Marketing|  India| 68000.0|        4|2020-11-30|
|     9|    Ivy|  Marketing|Germany| 75000.0|        5|2019-08-09|
+------+-------+-----------+-------+--------+---------+----------+

=== unionByName ===
+------+-------+-----------+-------+--------+---------+----------+
|emp_id|   name| department|country|  salary|years_exp| join_date|
+------+-------+-----------+-------+--------+---------+----------+
|     1|  Alice|Engineering|  India| 95000.0|        5|2019-03-15|
|     3|Charlie|Engineering|     US|110000.0|        8|2016-01-20|
|     5|    Eve|Engineering|Germany|120000.0|       10|2014-07-04|
|     8|   Hank|Engineering|Germany| 98000.0|        6|2018-05-22|
|     2|    Bob|  Marketing|     US| 72000.0|        3|2021-06-01|
|     6|  Frank|  Marketing|  India| 68000.0|        4|2020-11-30|
|     9|    Ivy|  Marketing|Germany| 75000.0|        5|2019-08-09|
+------+-------+-----------+-------+--------+---------+----------+

=== Intersect: High Salary AND Senior ===
+-------+--------+---------+
|   name|  salary|years_exp|
+-------+--------+---------+
|    Eve|120000.0|       10|
|   Hank| 98000.0|        6|
|Charlie|110000.0|        8|
|  Alice| 95000.0|        5|
+-------+--------+---------+

=== Except: High Salary but NOT Senior ===
+----+------+---------+
|name|salary|years_exp|
+----+------+---------+
+----+------+---------+

Q13 — UDF (User Defined Function)

Write a UDF gradeEmployee that takes salary and returns:

"A" if >= 100000, "B" if >= 75000, "C" if >= 50000, "D" otherwise

Register and apply it to employeesDF as a new column grade
Write a second UDF maskName that returns first char + "***" (e.g., "Alice" → "A***")
Apply maskName to the name column

In [ ]

PySpark · In [ ]cell 48

# YOUR CODE HERE

Expected output

text

+-------+-----------+--------+-----+
|   name|masked_name|  salary|grade|
+-------+-----------+--------+-----+
|  Alice|       A***| 95000.0|    B|
|    Bob|       B***| 72000.0|    C|
|Charlie|       C***|110000.0|    A|
|  Diana|       D***| 60000.0|    C|
|    Eve|       E***|120000.0|    A|
|  Frank|       F***| 68000.0|    C|
|  Grace|       G***| 58000.0|    C|
|   Hank|       H***| 98000.0|    B|
|    Ivy|       I***| 75000.0|    B|
|   Jack|       J***|    NULL|    D|
+-------+-----------+--------+-----+

Q14 — distinct, dropDuplicates

From salesDF, find all unique product names using distinct
Create a DataFrame with intentional duplicates by union-ing salesDF with itself
Use dropDuplicates() to remove full row duplicates
Use dropDuplicates(Seq("emp_id", "product")) to keep only one sale per employee-product combo
Confirm counts at each step

In [ ]

PySpark · In [ ]cell 51

# YOUR CODE HERE

Expected output

text

=== Distinct Products ===
+--------+
| product|
+--------+
|  Laptop|
|    Desk|
|   Phone|
|   Chair|
| Monitor|
|Keyboard|
+--------+

With duplicates: 24 rows
After dropDuplicates(): 12 rows
After dropDuplicates(emp_id, product): 12 rows
+-------+------+--------+-----------+------+--------+----------+
|sale_id|emp_id| product|   category| price|quantity| sale_date|
+-------+------+--------+-----------+------+--------+----------+
|    103|     1|    Desk|  Furniture| 450.0|       3|2024-01-10|
|    101|     1|  Laptop|Electronics|1200.0|       2|2024-01-05|
|    110|     1|   Phone|Electronics| 800.0|       1|2024-03-10|
|    107|     2| Monitor|Electronics| 600.0|       2|2024-02-15|
|    102|     2|   Phone|Electronics| 800.0|       1|2024-01-07|
|    104|     3|  Laptop|Electronics|1200.0|       1|2024-01-12|
|    105|     4|   Chair|  Furniture| 250.0|       4|2024-02-01|
|    106|     5|  Laptop|Electronics|1200.0|       5|2024-02-03|
|    108|     6|Keyboard|Electronics|  80.0|      10|2024-03-01|
|    109|     7|    Desk|  Furniture| 450.0|       1|2024-03-05|
|    111|     8|   Chair|  Furniture| 250.0|       6|2024-03-15|
|    112|     9| Monitor|Electronics| 600.0|       3|2024-04-01|
+-------+------+--------+-----------+------+--------+----------+

Q15 — cast, schema manipulation, printSchema

Print the schema of employeesDF and salesDF
Cast join_date (String) to DateType in employeesDF
Cast salary to IntegerType
Cast sale_date (String) to TimestampType in salesDF
Define a manual schema using StructType + StructField and re-read the employees data with it

In [ ]

PySpark · In [ ]cell 54

# YOUR CODE HERE

Expected output

text

=== Employees Schema ===
root
 |-- emp_id: long (nullable = true)
 |-- name: string (nullable = true)
 |-- department: string (nullable = true)
 |-- country: string (nullable = true)
 |-- salary: double (nullable = true)
 |-- years_exp: long (nullable = true)
 |-- join_date: string (nullable = true)

=== Sales Schema ===
root
 |-- sale_id: long (nullable = true)
 |-- emp_id: long (nullable = true)
 |-- product: string (nullable = true)
 |-- category: string (nullable = true)
 |-- price: double (nullable = true)
 |-- quantity: long (nullable = true)
 |-- sale_date: string (nullable = true)

=== Casted Employees Schema ===
root
 |-- emp_id: long (nullable = true)
 |-- name: string (nullable = true)
 |-- department: string (nullable = true)
 |-- country: string (nullable = true)
 |-- salary: integer (nullable = true)
 |-- years_exp: long (nullable = true)
 |-- join_date: date (nullable = true)

root
 |-- sale_id: long (nullable = true)
 |-- emp_id: long (nullable = true)
 |-- product: string (nullable = true)
 |-- category: string (nullable = true)
 |-- price: double (nullable = true)
 |-- quantity: long (nullable = true)
 |-- sale_date: timestamp (nullable = true)

=== Manual Schema ===
root
 |-- emp_id: integer (nullable = false)
 |-- name: string (nullable = true)
 |-- department: string (nullable = true)
 |-- country: string (nullable = true)
 |-- salary: double (nullable = true)
 |-- years_exp: integer (nullable = true)
 |-- join_date: string (nullable = true)

Q16 — orderBy, sort, limit, sample

Sort employeesDF by salary descending, then by name ascending
Get the top 3 highest paid employees
Get the bottom 3 lowest paid employees (handle nulls last)
Take a 50% random sample (with seed=42, without replacement)
Sort salesDF by sale_date ascending and show the first 5 records

In [ ]

PySpark · In [ ]cell 57

# YOUR CODE HERE

Expected output

text

=== Top 3 Highest Paid ===
+------+-------+-----------+-------+--------+---------+----------+
|emp_id|   name| department|country|  salary|years_exp| join_date|
+------+-------+-----------+-------+--------+---------+----------+
|     5|    Eve|Engineering|Germany|120000.0|       10|2014-07-04|
|     3|Charlie|Engineering|     US|110000.0|        8|2016-01-20|
|     8|   Hank|Engineering|Germany| 98000.0|        6|2018-05-22|
+------+-------+-----------+-------+--------+---------+----------+

=== Bottom 3 Lowest Paid ===
+------+-----+----------+-------+-------+---------+----------+
|emp_id| name|department|country| salary|years_exp| join_date|
+------+-----+----------+-------+-------+---------+----------+
|     7|Grace|        HR|     US|58000.0|        1|2023-02-14|
|     4|Diana|        HR|  India|60000.0|        2|2022-09-10|
|     6|Frank| Marketing|  India|68000.0|        4|2020-11-30|
+------+-----+----------+-------+-------+---------+----------+

=== 50% Sample ===
+------+-----+-----------+-------+--------+---------+----------+
|emp_id| name| department|country|  salary|years_exp| join_date|
+------+-----+-----------+-------+--------+---------+----------+
|     5|  Eve|Engineering|Germany|120000.0|       10|2014-07-04|
|     6|Frank|  Marketing|  India| 68000.0|        4|2020-11-30|
|     8| Hank|Engineering|Germany| 98000.0|        6|2018-05-22|
|     9|  Ivy|  Marketing|Germany| 75000.0|        5|2019-08-09|
|    10| Jack|         HR|     US|    NULL|        3|2021-03-17|
+------+-----+-----------+-------+--------+---------+----------+

=== First 5 Sales by Date ===
+-------+------+-------+-----------+------+--------+----------+
|sale_id|emp_id|product|   category| price|quantity| sale_date|
+-------+------+-------+-----------+------+--------+----------+
|    101|     1| Laptop|Electronics|1200.0|       2|2024-01-05|
|    102|     2|  Phone|Electronics| 800.0|       1|2024-01-07|
|    103|     1|   Desk|  Furniture| 450.0|       3|2024-01-10|
|    104|     3| Laptop|Electronics|1200.0|       1|2024-01-12|
|    105|     4|  Chair|  Furniture| 250.0|       4|2024-02-01|
+-------+------+-------+-----------+------+--------+----------+
only showing top 5 rows

Q17 — repartition, coalesce, cache, persist

Check the current number of partitions in employeesDF
repartition(4) — repartition to 4 partitions (full shuffle)
coalesce(2) — reduce to 2 partitions without full shuffle
repartition(col("department")) — partition by department column
cache() the employeesDF and trigger materialization with an action
Unpersist/uncache it after use

In [ ]

PySpark · In [ ]cell 60

# YOUR CODE HERE

Expected output

text

Default partitions: 8
After repartition(4): 4
After coalesce(2): 2
After repartition(department): 1
DataFrame cached
DataFrame unpersisted

Q18 — Advanced Aggregations: countDistinct, approx_count_distinct, collect_list, collect_set, first, last

From salesDF:

Count distinct products sold per category
Use approx_count_distinct for approximate distinct count
collect_list of all product names per emp_id
collect_set of unique products per emp_id (no duplicates)
Get the first and last sale per emp_id (by sale_date)

In [ ]

PySpark · In [ ]cell 63

# YOUR CODE HERE

Expected output

text

=== countDistinct ===
+-----------+-----------------+
|   category|distinct_products|
+-----------+-----------------+
|Electronics|                4|
|  Furniture|                2|
+-----------+-----------------+

=== approx_count_distinct ===
+-----------+---------------+
|   category|approx_distinct|
+-----------+---------------+
|Electronics|              4|
|  Furniture|              2|
+-----------+---------------+

=== collect_list ===
+------+---------------------+
|emp_id|all_products         |
+------+---------------------+
|1     |[Laptop, Desk, Phone]|
|2     |[Phone, Monitor]     |
|3     |[Laptop]             |
|5     |[Laptop]             |
|4     |[Chair]              |
|7     |[Desk]               |
|6     |[Keyboard]           |
|9     |[Monitor]            |
|8     |[Chair]              |
+------+---------------------+

=== collect_set ===
+------+---------------------+
|emp_id|unique_products      |
+------+---------------------+
|1     |[Laptop, Phone, Desk]|
|2     |[Phone, Monitor]     |
|3     |[Laptop]             |
|5     |[Laptop]             |
|4     |[Chair]              |
|7     |[Desk]               |
|6     |[Keyboard]           |
|9     |[Monitor]            |
|8     |[Chair]              |
+------+---------------------+

=== First & Last Sale ===
+------+-------------+------------+
|emp_id|first_product|last_product|
+------+-------------+------------+
|     1|       Laptop|       Phone|
|     2|        Phone|     Monitor|
|     3|       Laptop|      Laptop|
|     4|        Chair|       Chair|
|     5|       Laptop|      Laptop|
|     6|     Keyboard|    Keyboard|
|     7|         Desk|        Desk|
|     8|        Chair|       Chair|
|     9|      Monitor|     Monitor|
+------+-------------+------------+

Q19 — Complex Pipeline (Multi-step)

Build a complete employee performance report combining multiple transformations:

Join employeesDF with salesDF on emp_id (left join to keep all employees)
Calculate revenue_per_emp = sum(price × quantity) per employee
Fill null revenue (employees with no sales) with 0
Join with departmentsDF on department = dept_name
Add salary_pct_of_budget = (salary / budget) × 100
Add window-based dept_salary_rank = rank within department by salary
Add perf_label using when/otherwise: "Star" if revenue > 3000, "Good" if > 1000, "Low" otherwise
Select: name, department, salary, revenue_per_emp, salary_pct_of_budget, dept_salary_rank, perf_label
Sort by department, then dept_salary_rank

In [ ]

PySpark · In [ ]cell 66

# YOUR CODE HERE

Expected output

text

+-------+-----------+--------+---------------+--------------------+----------------+----------+
|   name| department|  salary|revenue_per_emp|salary_pct_of_budget|dept_salary_rank|perf_label|
+-------+-----------+--------+---------------+--------------------+----------------+----------+
|    Eve|Engineering|120000.0|         6000.0|                24.0|               1|      Star|
|Charlie|Engineering|110000.0|         1200.0|                22.0|               2|      Good|
|   Hank|Engineering| 98000.0|         1500.0|                19.6|               3|      Good|
|  Alice|Engineering| 95000.0|         4550.0|                19.0|               4|      Star|
|  Diana|         HR| 60000.0|         1000.0|                40.0|               1|       Low|
|  Grace|         HR| 58000.0|          450.0|               38.67|               2|       Low|
|   Jack|         HR|    NULL|            0.0|                NULL|               3|       Low|
|    Ivy|  Marketing| 75000.0|         1800.0|                25.0|               1|      Good|
|    Bob|  Marketing| 72000.0|         2000.0|                24.0|               2|      Good|
|  Frank|  Marketing| 68000.0|          800.0|               22.67|               3|       Low|
+-------+-----------+--------+---------------+--------------------+----------------+----------+

Q20 — Spark SQL + createOrReplaceTempView

Use Spark SQL (not DataFrame API) to:

Find the department with the highest average salary
Find employees who earn more than their department's average salary
Write a SQL query to get the top 2 earners per department using ROW_NUMBER() OVER (PARTITION BY ...)
Join all three tables (employees, sales, departments) in one SQL query to compute total revenue per department
Use a CTE (WITH clause) in SQL to find departments where total salary > 50% of budget

In [ ]

PySpark · In [ ]cell 69

# YOUR CODE HERE

Expected output

text

=== Highest Avg Salary Dept ===
+-----------+----------+
| department|avg_salary|
+-----------+----------+
|Engineering|  105750.0|
+-----------+----------+

=== Above Dept Average ===
+-------+-----------+--------+--------+
|   name| department|  salary|dept_avg|
+-------+-----------+--------+--------+
|    Bob|  Marketing| 72000.0|71666.67|
|Charlie|Engineering|110000.0|105750.0|
|    Eve|Engineering|120000.0|105750.0|
|  Diana|         HR| 60000.0| 59000.0|
|    Ivy|  Marketing| 75000.0|71666.67|
+-------+-----------+--------+--------+

=== Top 2 per Dept (SQL Window) ===
+-------+-----------+--------+---+
|   name| department|  salary| rn|
+-------+-----------+--------+---+
|    Eve|Engineering|120000.0|  1|
|Charlie|Engineering|110000.0|  2|
|  Diana|         HR| 60000.0|  1|
|  Grace|         HR| 58000.0|  2|
|    Ivy|  Marketing| 75000.0|  1|
|    Bob|  Marketing| 72000.0|  2|
+-------+-----------+--------+---+

=== Revenue per Dept ===
+-----------+-------------+-------------+
| department|    dept_head|total_revenue|
+-----------+-------------+-------------+
|Engineering|Alice Johnson|      13250.0|
|  Marketing| Bob Williams|       4600.0|
|         HR|  Carol Smith|       1450.0|
+-----------+-------------+-------------+

=== Salary > 50% of Budget (CTE) ===
+-----------+------------+--------+-------------+
| department|total_salary|  budget|pct_of_budget|
+-----------+------------+--------+-------------+
|Engineering|    423000.0|500000.0|         84.6|
|         HR|    118000.0|150000.0|         78.7|
|  Marketing|    215000.0|300000.0|         71.7|
+-----------+------------+--------+-------------+

SOLUTIONS

Solution 1 — select, filter, withColumn

In [ ]

PySpark · In [ ]cell 72

q1 = (
    employeesDF
    .select("name", "department", "salary")
    .filter((F.col("department") == "Engineering") & (F.col("salary") > 90000))
    .withColumn("salary_inr", F.col("salary") * 83)
)

q1.show()

Output

+-------+-----------+--------+----------+
|   name| department|  salary|salary_inr|
+-------+-----------+--------+----------+
|  Alice|Engineering| 95000.0| 7885000.0|
|Charlie|Engineering|110000.0| 9130000.0|
|    Eve|Engineering|120000.0| 9960000.0|
|   Hank|Engineering| 98000.0| 8134000.0|
+-------+-----------+--------+----------+

Solution 2 — groupBy + agg

In [ ]

PySpark · In [ ]cell 74

q2 = (
    employeesDF
    .groupBy("department")
    .agg(
        F.count("emp_id").alias("total_employees"),
        F.round(F.avg("salary"), 2).alias("avg_salary"),
        F.max("salary").alias("max_salary"),
        F.min("salary").alias("min_salary"),
        F.sum("salary").alias("total_salary"),
    )
    .orderBy(F.col("avg_salary").desc())
)

q2.show()

Output

+-----------+---------------+----------+----------+----------+------------+
| department|total_employees|avg_salary|max_salary|min_salary|total_salary|
+-----------+---------------+----------+----------+----------+------------+
|Engineering|              4|  105750.0|  120000.0|   95000.0|    423000.0|
|  Marketing|              3|  71666.67|   75000.0|   68000.0|    215000.0|
|         HR|              3|   59000.0|   60000.0|   58000.0|    118000.0|
+-----------+---------------+----------+----------+----------+------------+

Solution 3 — Joins

In [ ]

PySpark · In [ ]cell 76

# 1. Inner Join
inner_join = (
    employeesDF
    .join(salesDF, "emp_id", "inner")
    .select("name", "department", "product", "price", "quantity", "sale_date")
)
print("=== Inner Join ===")
inner_join.show()

# 2. Left Join + null filter (employees with no sales)
no_sales = (
    employeesDF
    .join(salesDF, ["emp_id"], "left")
    .filter(F.col("sale_id").isNull())
    .select("emp_id", "name", "department")
)
print("=== No Sales (Left + null filter) ===")
no_sales.show()

# 3. Anti Join
anti_join = (
    employeesDF
    .join(salesDF, ["emp_id"], "left_anti")
    .select("emp_id", "name", "department")
)
print("=== Anti Join ===")
anti_join.show()

# 4. Departments with no employees
dept_no_emp = (
    departmentsDF.alias("d")
    .join(
        employeesDF.alias("e"),
        F.col("d.dept_name") == F.col("e.department"),
        "left",
    )
    .filter(F.col("e.emp_id").isNull())
    .select(F.col("d.dept_name"), F.col("d.dept_head"))
)
print("=== Departments with no Employees ===")
dept_no_emp.show()

Output

=== Inner Join ===
+-------+-----------+--------+------+--------+----------+
|   name| department| product| price|quantity| sale_date|
+-------+-----------+--------+------+--------+----------+
|  Alice|Engineering|  Laptop|1200.0|       2|2024-01-05|
|  Alice|Engineering|    Desk| 450.0|       3|2024-01-10|
|  Alice|Engineering|   Phone| 800.0|       1|2024-03-10|
|    Bob|  Marketing|   Phone| 800.0|       1|2024-01-07|
|    Bob|  Marketing| Monitor| 600.0|       2|2024-02-15|
|Charlie|Engineering|  Laptop|1200.0|       1|2024-01-12|
|  Diana|         HR|   Chair| 250.0|       4|2024-02-01|
|    Eve|Engineering|  Laptop|1200.0|       5|2024-02-03|
|  Frank|  Marketing|Keyboard|  80.0|      10|2024-03-01|
|  Grace|         HR|    Desk| 450.0|       1|2024-03-05|
|   Hank|Engineering|   Chair| 250.0|       6|2024-03-15|
|    Ivy|  Marketing| Monitor| 600.0|       3|2024-04-01|
+-------+-----------+--------+------+--------+----------+

=== No Sales (Left + null filter) ===
+------+----+----------+
|emp_id|name|department|
+------+----+----------+
|    10|Jack|        HR|
+------+----+----------+

=== Anti Join ===
+------+----+----------+
|emp_id|name|department|
+------+----+----------+
|    10|Jack|        HR|
+------+----+----------+

=== Departments with no Employees ===
+---------+-----------+
|dept_name|  dept_head|
+---------+-----------+
|  Finance|David Brown|
+---------+-----------+

Solution 4 — rank, dense_rank, row_number

In [ ]

PySpark · In [ ]cell 78

window_spec = Window.partitionBy("department").orderBy(F.col("salary").desc())

ranked = (
    employeesDF
    .withColumn("rank", F.rank().over(window_spec))
    .withColumn("dense_rank", F.dense_rank().over(window_spec))
    .withColumn("row_number", F.row_number().over(window_spec))
)

print("=== All Ranks ===")
ranked.select(
    "name", "department", "salary", "rank", "dense_rank", "row_number"
).show()

print("=== Top 2 per Department ===")
(
    ranked
    .filter(F.col("dense_rank") <= 2)
    .select("name", "department", "salary", "dense_rank")
    .show()
)

Output

=== All Ranks ===
+-------+-----------+--------+----+----------+----------+
|   name| department|  salary|rank|dense_rank|row_number|
+-------+-----------+--------+----+----------+----------+
|    Eve|Engineering|120000.0|   1|         1|         1|
|Charlie|Engineering|110000.0|   2|         2|         2|
|   Hank|Engineering| 98000.0|   3|         3|         3|
|  Alice|Engineering| 95000.0|   4|         4|         4|
|  Diana|         HR| 60000.0|   1|         1|         1|
|  Grace|         HR| 58000.0|   2|         2|         2|
|   Jack|         HR|    NULL|   3|         3|         3|
|    Ivy|  Marketing| 75000.0|   1|         1|         1|
|    Bob|  Marketing| 72000.0|   2|         2|         2|
|  Frank|  Marketing| 68000.0|   3|         3|         3|
+-------+-----------+--------+----+----------+----------+

=== Top 2 per Department ===
+-------+-----------+--------+----------+
|   name| department|  salary|dense_rank|
+-------+-----------+--------+----------+
|    Eve|Engineering|120000.0|         1|
|Charlie|Engineering|110000.0|         2|
|  Diana|         HR| 60000.0|         1|
|  Grace|         HR| 58000.0|         2|
|    Ivy|  Marketing| 75000.0|         1|
|    Bob|  Marketing| 72000.0|         2|
+-------+-----------+--------+----------+

Solution 5 — lag, lead, running total

In [ ]

PySpark · In [ ]cell 80

sales_with_date = salesDF.withColumn("sale_date", F.to_date("sale_date"))

employee_window = Window.partitionBy("emp_id").orderBy("sale_date")
running_window = employee_window.rowsBetween(
    Window.unboundedPreceding, Window.currentRow
)

q5 = (
    sales_with_date
    .withColumn("prev_sale_price", F.lag("price", 1).over(employee_window))
    .withColumn("next_sale_price", F.lead("price", 1).over(employee_window))
    .withColumn("revenue", F.col("price") * F.col("quantity"))
    .withColumn("running_revenue", F.sum("revenue").over(running_window))
)

q5.select(
    "emp_id", "product", "price", "quantity", "sale_date",
    "prev_sale_price", "next_sale_price", "running_revenue",
).show()

Output

+------+--------+------+--------+----------+---------------+---------------+---------------+
|emp_id| product| price|quantity| sale_date|prev_sale_price|next_sale_price|running_revenue|
+------+--------+------+--------+----------+---------------+---------------+---------------+
|     1|  Laptop|1200.0|       2|2024-01-05|           NULL|          450.0|         2400.0|
|     1|    Desk| 450.0|       3|2024-01-10|         1200.0|          800.0|         3750.0|
|     1|   Phone| 800.0|       1|2024-03-10|          450.0|           NULL|         4550.0|
|     2|   Phone| 800.0|       1|2024-01-07|           NULL|          600.0|          800.0|
|     2| Monitor| 600.0|       2|2024-02-15|          800.0|           NULL|         2000.0|
|     3|  Laptop|1200.0|       1|2024-01-12|           NULL|           NULL|         1200.0|
|     4|   Chair| 250.0|       4|2024-02-01|           NULL|           NULL|         1000.0|
|     5|  Laptop|1200.0|       5|2024-02-03|           NULL|           NULL|         6000.0|
|     6|Keyboard|  80.0|      10|2024-03-01|           NULL|           NULL|          800.0|
|     7|    Desk| 450.0|       1|2024-03-05|           NULL|           NULL|          450.0|
|     8|   Chair| 250.0|       6|2024-03-15|           NULL|           NULL|         1500.0|
|     9| Monitor| 600.0|       3|2024-04-01|           NULL|           NULL|         1800.0|
+------+--------+------+--------+----------+---------------+---------------+---------------+

Solution 6 — when / otherwise

In [ ]

PySpark · In [ ]cell 82

q6 = (
    employeesDF
    .withColumn(
        "salary_band",
        F.when(F.col("salary").isNull(), "Unknown")
        .when(F.col("salary") >= 100000, "High")
        .when(F.col("salary") >= 70000, "Mid")
        .otherwise("Low"),
    )
    .withColumn(
        "experience_level",
        F.when(F.col("years_exp") >= 7, "Senior")
        .when(F.col("years_exp") >= 4, "Mid")
        .otherwise("Junior"),
    )
)

q6.select(
    "name", "salary", "salary_band", "years_exp", "experience_level"
).show()

Output

+-------+--------+-----------+---------+----------------+
|   name|  salary|salary_band|years_exp|experience_level|
+-------+--------+-----------+---------+----------------+
|  Alice| 95000.0|        Mid|        5|             Mid|
|    Bob| 72000.0|        Mid|        3|          Junior|
|Charlie|110000.0|       High|        8|          Senior|
|  Diana| 60000.0|        Low|        2|          Junior|
|    Eve|120000.0|       High|       10|          Senior|
|  Frank| 68000.0|        Low|        4|             Mid|
|  Grace| 58000.0|        Low|        1|          Junior|
|   Hank| 98000.0|        Mid|        6|             Mid|
|    Ivy| 75000.0|        Mid|        5|             Mid|
|   Jack|    NULL|    Unknown|        3|          Junior|
+-------+--------+-----------+---------+----------------+

Solution 7 — Null Handling

In [ ]

PySpark · In [ ]cell 84

# 1. Count null salaries
null_count = employeesDF.filter(F.col("salary").isNull()).count()
print(f"Null salary count: {null_count}")

# 2. Fill null with the average salary
avg_salary = employeesDF.agg(F.avg("salary")).first()[0]
filled = employeesDF.na.fill({"salary": avg_salary})
print("=== Filled with avg ===")
filled.select("name", "salary").show()

# 3. coalesce
with_coalesce = employeesDF.withColumn(
    "salary", F.coalesce(F.col("salary"), F.lit(60000.0))
)
print("=== Coalesce ===")
with_coalesce.select("name", "salary").show()

# 4. Drop rows where any column is null
drop_all = employeesDF.na.drop()
print(f"After na.drop (any): {drop_all.count()} rows")

# 5. Drop rows only when both salary and country are null
drop_subset = employeesDF.na.drop(how="all", subset=["salary", "country"])
print(f"After subset na.drop (all): {drop_subset.count()} rows")

Output

Null salary count: 1
=== Filled with avg ===
+-------+--------+
|   name|  salary|
+-------+--------+
|  Alice| 95000.0|
|    Bob| 72000.0|
|Charlie|110000.0|
|  Diana| 60000.0|
|    Eve|120000.0|
|  Frank| 68000.0|
|  Grace| 58000.0|
|   Hank| 98000.0|
|    Ivy| 75000.0|
|   Jack| 84000.0|
+-------+--------+

=== Coalesce ===
+-------+--------+
|   name|  salary|
+-------+--------+
|  Alice| 95000.0|
|    Bob| 72000.0|
|Charlie|110000.0|
|  Diana| 60000.0|
|    Eve|120000.0|
|  Frank| 68000.0|
|  Grace| 58000.0|
|   Hank| 98000.0|
|    Ivy| 75000.0|
|   Jack| 60000.0|
+-------+--------+

After na.drop (any): 9 rows
After subset na.drop (all): 10 rows

Solution 8 — String Functions

In [ ]

PySpark · In [ ]cell 86

q8 = (
    employeesDF
    .withColumn("name_upper", F.upper("name"))
    .withColumn("dept_lower", F.lower("department"))
    .withColumn(
        "name_dept",
        F.concat(F.col("name"), F.lit(" ("), F.col("department"), F.lit(")")),
    )
    .withColumn("name_code", F.substring("name", 1, 3))
    .withColumn("has_a", F.lower("name").contains("a"))
    .withColumn(
        "country_code",
        F.regexp_replace(
            F.regexp_replace(F.col("country"), "India", "IN"), "^US$", "USA"
        ),
    )
)

q8.select(
    "name", "name_upper", "dept_lower", "name_dept",
    "name_code", "has_a", "country_code",
).show()

Output

+-------+----------+-----------+--------------------+---------+-----+------------+
|   name|name_upper| dept_lower|           name_dept|name_code|has_a|country_code|
+-------+----------+-----------+--------------------+---------+-----+------------+
|  Alice|     ALICE|engineering| Alice (Engineering)|      Ali| true|          IN|
|    Bob|       BOB|  marketing|     Bob (Marketing)|      Bob|false|         USA|
|Charlie|   CHARLIE|engineering|Charlie (Engineer...|      Cha| true|         USA|
|  Diana|     DIANA|         hr|          Diana (HR)|      Dia| true|          IN|
|    Eve|       EVE|engineering|   Eve (Engineering)|      Eve|false|     Germany|
|  Frank|     FRANK|  marketing|   Frank (Marketing)|      Fra| true|          IN|
|  Grace|     GRACE|         hr|          Grace (HR)|      Gra| true|         USA|
|   Hank|      HANK|engineering|  Hank (Engineering)|      Han| true|     Germany|
|    Ivy|       IVY|  marketing|     Ivy (Marketing)|      Ivy|false|     Germany|
|   Jack|      JACK|         hr|           Jack (HR)|      Jac| true|         USA|
+-------+----------+-----------+--------------------+---------+-----+------------+

Solution 9 — Date Functions

In [ ]

PySpark · In [ ]cell 88

q9 = (
    employeesDF
    .withColumn("join_date", F.to_date("join_date", "yyyy-MM-dd"))
    .withColumn("join_year", F.year("join_date"))
    .withColumn("join_month", F.month("join_date"))
    .withColumn("join_day", F.dayofmonth("join_date"))
    .withColumn("tenure_days", F.datediff(F.current_date(), F.col("join_date")))
    .withColumn("tenure_years", F.round(F.col("tenure_days") / 365.0, 1))
    .withColumn("probation_end", F.date_add(F.col("join_date"), 30))
)

q9.select(
    "name", "join_date", "join_year", "join_month", "join_day",
    "tenure_days", "tenure_years", "probation_end",
).show()

Output

+-------+----------+---------+----------+--------+-----------+------------+-------------+
|   name| join_date|join_year|join_month|join_day|tenure_days|tenure_years|probation_end|
+-------+----------+---------+----------+--------+-----------+------------+-------------+
|  Alice|2019-03-15|     2019|         3|      15|       2764|         7.6|   2019-04-14|
|    Bob|2021-06-01|     2021|         6|       1|       1955|         5.4|   2021-07-01|
|Charlie|2016-01-20|     2016|         1|      20|       3914|        10.7|   2016-02-19|
|  Diana|2022-09-10|     2022|         9|      10|       1489|         4.1|   2022-10-10|
|    Eve|2014-07-04|     2014|         7|       4|       4479|        12.3|   2014-08-03|
|  Frank|2020-11-30|     2020|        11|      30|       2138|         5.9|   2020-12-30|
|  Grace|2023-02-14|     2023|         2|      14|       1332|         3.6|   2023-03-16|
|   Hank|2018-05-22|     2018|         5|      22|       3061|         8.4|   2018-06-21|
|    Ivy|2019-08-09|     2019|         8|       9|       2617|         7.2|   2019-09-08|
|   Jack|2021-03-17|     2021|         3|      17|       2031|         5.6|   2021-04-16|
+-------+----------+---------+----------+--------+-----------+------------+-------------+

Solution 10 — explode + array functions

In [ ]

PySpark · In [ ]cell 90

# 1. Explode
exploded_df = skillsDF.withColumn("skill", F.explode("skills"))
print("=== Exploded ===")
exploded_df.select("emp_id", "skill").show()

# 2. Array size
with_size = skillsDF.withColumn("skill_count", F.size("skills"))
print("=== Skill Count ===")
with_size.show()

# 3. array_contains
spark_devs = skillsDF.filter(F.array_contains("skills", "Spark"))
print("=== Spark Developers ===")
spark_devs.show()

# 4. Join with employees
with_names = (
    exploded_df
    .join(employeesDF, "emp_id")
    .select("emp_id", "name", "skill")
)
print("=== Name + Skills ===")
with_names.show()

# 5. collect_list
reaggregated = with_names.groupBy("emp_id", "name").agg(
    F.collect_list("skill").alias("skills_list")
)
print("=== Re-aggregated ===")
reaggregated.show(truncate=False)

Output

=== Exploded ===
+------+----------+
|emp_id|     skill|
+------+----------+
|     1|     Scala|
|     1|     Spark|
|     1|       SQL|
|     2| Marketing|
|     2|       SEO|
|     3|      Java|
|     3|     Scala|
|     3|     Kafka|
|     3|     Spark|
|     4|        HR|
|     4|Recruiting|
|     5|     Scala|
|     5|    Python|
|     5|        ML|
|     5|     Spark|
+------+----------+

=== Skill Count ===
+------+--------------------+-----------+
|emp_id|              skills|skill_count|
+------+--------------------+-----------+
|     1| [Scala, Spark, SQL]|          3|
|     2|    [Marketing, SEO]|          2|
|     3|[Java, Scala, Kaf...|          4|
|     4|    [HR, Recruiting]|          2|
|     5|[Scala, Python, M...|          4|
+------+--------------------+-----------+

=== Spark Developers ===
+------+--------------------+
|emp_id|              skills|
+------+--------------------+
|     1| [Scala, Spark, SQL]|
|     3|[Java, Scala, Kaf...|
|     5|[Scala, Python, M...|
+------+--------------------+

=== Name + Skills ===
+------+-------+----------+
|emp_id|   name|     skill|
+------+-------+----------+
|     1|  Alice|     Scala|
|     1|  Alice|     Spark|
|     1|  Alice|       SQL|
|     2|    Bob| Marketing|
|     2|    Bob|       SEO|
|     3|Charlie|      Java|
|     3|Charlie|     Scala|
|     3|Charlie|     Kafka|
|     3|Charlie|     Spark|
|     4|  Diana|        HR|
|     4|  Diana|Recruiting|
|     5|    Eve|     Scala|
|     5|    Eve|    Python|
|     5|    Eve|        ML|
|     5|    Eve|     Spark|
+------+-------+----------+

=== Re-aggregated ===
+------+-------+---------------------------+
|emp_id|name   |skills_list                |
+------+-------+---------------------------+
|1     |Alice  |[Scala, Spark, SQL]        |
|2     |Bob    |[Marketing, SEO]           |
|3     |Charlie|[Java, Scala, Kafka, Spark]|
|4     |Diana  |[HR, Recruiting]           |
|5     |Eve    |[Scala, Python, ML, Spark] |
+------+-------+---------------------------+

Solution 11 — Pivot Table

In [ ]

PySpark · In [ ]cell 92

q11 = (
    salesDF
    .withColumn("revenue", F.col("price") * F.col("quantity"))
    .groupBy("emp_id")
    .pivot("category", ["Electronics", "Furniture"])
    .agg(F.round(F.sum("revenue"), 2))
    .na.fill(0)
)

q11.show()

Output

+------+-----------+---------+
|emp_id|Electronics|Furniture|
+------+-----------+---------+
|     7|        0.0|    450.0|
|     6|      800.0|      0.0|
|     9|     1800.0|      0.0|
|     5|     6000.0|      0.0|
|     1|     3200.0|   1350.0|
|     3|     1200.0|      0.0|
|     8|        0.0|   1500.0|
|     2|     2000.0|      0.0|
|     4|        0.0|   1000.0|
+------+-----------+---------+

Solution 12 — Union, Intersect, Except

In [ ]

PySpark · In [ ]cell 94

eng_df = employeesDF.filter(F.col("department") == "Engineering")
mkt_df = employeesDF.filter(F.col("department") == "Marketing")

# Union
print("=== Union ===")
eng_df.union(mkt_df).show()

# unionByName with reordered columns
mkt_reordered = mkt_df.select(
    "department", "name", "emp_id", "country", "salary", "years_exp", "join_date"
)
print("=== unionByName ===")
eng_df.unionByName(mkt_reordered).show()

# Intersect and set difference
high_salary_df = employeesDF.filter(F.col("salary") > 80000)
senior_df = employeesDF.filter(F.col("years_exp") >= 5)

print("=== Intersect: High Salary AND Senior ===")
high_salary_df.intersect(senior_df).select("name", "salary", "years_exp").show()

print("=== Except: High Salary but NOT Senior ===")
high_salary_df.subtract(senior_df).select("name", "salary", "years_exp").show()

Output

=== Union ===
+------+-------+-----------+-------+--------+---------+----------+
|emp_id|   name| department|country|  salary|years_exp| join_date|
+------+-------+-----------+-------+--------+---------+----------+
|     1|  Alice|Engineering|  India| 95000.0|        5|2019-03-15|
|     3|Charlie|Engineering|     US|110000.0|        8|2016-01-20|
|     5|    Eve|Engineering|Germany|120000.0|       10|2014-07-04|
|     8|   Hank|Engineering|Germany| 98000.0|        6|2018-05-22|
|     2|    Bob|  Marketing|     US| 72000.0|        3|2021-06-01|
|     6|  Frank|  Marketing|  India| 68000.0|        4|2020-11-30|
|     9|    Ivy|  Marketing|Germany| 75000.0|        5|2019-08-09|
+------+-------+-----------+-------+--------+---------+----------+

=== unionByName ===
+------+-------+-----------+-------+--------+---------+----------+
|emp_id|   name| department|country|  salary|years_exp| join_date|
+------+-------+-----------+-------+--------+---------+----------+
|     1|  Alice|Engineering|  India| 95000.0|        5|2019-03-15|
|     3|Charlie|Engineering|     US|110000.0|        8|2016-01-20|
|     5|    Eve|Engineering|Germany|120000.0|       10|2014-07-04|
|     8|   Hank|Engineering|Germany| 98000.0|        6|2018-05-22|
|     2|    Bob|  Marketing|     US| 72000.0|        3|2021-06-01|
|     6|  Frank|  Marketing|  India| 68000.0|        4|2020-11-30|
|     9|    Ivy|  Marketing|Germany| 75000.0|        5|2019-08-09|
+------+-------+-----------+-------+--------+---------+----------+

=== Intersect: High Salary AND Senior ===
+-------+--------+---------+
|   name|  salary|years_exp|
+-------+--------+---------+
|    Eve|120000.0|       10|
|   Hank| 98000.0|        6|
|Charlie|110000.0|        8|
|  Alice| 95000.0|        5|
+-------+--------+---------+

=== Except: High Salary but NOT Senior ===
+----+------+---------+
|name|salary|years_exp|
+----+------+---------+
+----+------+---------+

Solution 13 — UDFs

In [ ]

PySpark · In [ ]cell 96

def grade_employee(salary):
    if salary is None:
        return "D"
    if salary >= 100000:
        return "A"
    if salary >= 75000:
        return "B"
    if salary >= 50000:
        return "C"
    return "D"


def mask_name(name):
    if not name:
        return ""
    return name[0] + "***"


grade_employee_udf = F.udf(grade_employee, StringType())
mask_name_udf = F.udf(mask_name, StringType())

spark.udf.register("gradeEmployee", grade_employee, StringType())
spark.udf.register("maskName", mask_name, StringType())

q13 = (
    employeesDF
    .withColumn("grade", grade_employee_udf("salary"))
    .withColumn("masked_name", mask_name_udf("name"))
)

q13.select("name", "masked_name", "salary", "grade").show()

Output

+-------+-----------+--------+-----+
|   name|masked_name|  salary|grade|
+-------+-----------+--------+-----+
|  Alice|       A***| 95000.0|    B|
|    Bob|       B***| 72000.0|    C|
|Charlie|       C***|110000.0|    A|
|  Diana|       D***| 60000.0|    C|
|    Eve|       E***|120000.0|    A|
|  Frank|       F***| 68000.0|    C|
|  Grace|       G***| 58000.0|    C|
|   Hank|       H***| 98000.0|    B|
|    Ivy|       I***| 75000.0|    B|
|   Jack|       J***|    NULL|    D|
+-------+-----------+--------+-----+

Solution 14 — distinct, dropDuplicates

In [ ]

PySpark · In [ ]cell 98

# 1. Distinct products
print("=== Distinct Products ===")
salesDF.select("product").distinct().show()

# 2. Create duplicates
dupe_df = salesDF.union(salesDF)
print(f"With duplicates: {dupe_df.count()} rows")

# 3. Remove full-row duplicates
deduped = dupe_df.dropDuplicates()
print(f"After dropDuplicates(): {deduped.count()} rows")

# 4. Remove duplicates using a column subset
one_per_emp_product = dupe_df.dropDuplicates(["emp_id", "product"])
print(
    f"After dropDuplicates(emp_id, product): {one_per_emp_product.count()} rows"
)
one_per_emp_product.show()

Output

=== Distinct Products ===
+--------+
| product|
+--------+
|  Laptop|
|    Desk|
|   Phone|
|   Chair|
| Monitor|
|Keyboard|
+--------+

With duplicates: 24 rows
After dropDuplicates(): 12 rows
After dropDuplicates(emp_id, product): 12 rows
+-------+------+--------+-----------+------+--------+----------+
|sale_id|emp_id| product|   category| price|quantity| sale_date|
+-------+------+--------+-----------+------+--------+----------+
|    103|     1|    Desk|  Furniture| 450.0|       3|2024-01-10|
|    101|     1|  Laptop|Electronics|1200.0|       2|2024-01-05|
|    110|     1|   Phone|Electronics| 800.0|       1|2024-03-10|
|    107|     2| Monitor|Electronics| 600.0|       2|2024-02-15|
|    102|     2|   Phone|Electronics| 800.0|       1|2024-01-07|
|    104|     3|  Laptop|Electronics|1200.0|       1|2024-01-12|
|    105|     4|   Chair|  Furniture| 250.0|       4|2024-02-01|
|    106|     5|  Laptop|Electronics|1200.0|       5|2024-02-03|
|    108|     6|Keyboard|Electronics|  80.0|      10|2024-03-01|
|    109|     7|    Desk|  Furniture| 450.0|       1|2024-03-05|
|    111|     8|   Chair|  Furniture| 250.0|       6|2024-03-15|
|    112|     9| Monitor|Electronics| 600.0|       3|2024-04-01|
+-------+------+--------+-----------+------+--------+----------+

Solution 15 — cast, schema, StructType

In [ ]

PySpark · In [ ]cell 100

# 1. Print schemas
print("=== Employees Schema ===")
employeesDF.printSchema()
print("=== Sales Schema ===")
salesDF.printSchema()

# 2 and 3. Cast employee columns
emp_casted = (
    employeesDF
    .withColumn("join_date", F.col("join_date").cast(DateType()))
    .withColumn("salary", F.col("salary").cast(IntegerType()))
)
print("=== Casted Employees Schema ===")
emp_casted.printSchema()

# 4. Cast sale_date to timestamp
sales_casted = salesDF.withColumn(
    "sale_date", F.col("sale_date").cast(TimestampType())
)
sales_casted.printSchema()

# 5. Define and apply a manual schema
manual_schema = StructType([
    StructField("emp_id", IntegerType(), nullable=False),
    StructField("name", StringType(), nullable=True),
    StructField("department", StringType(), nullable=True),
    StructField("country", StringType(), nullable=True),
    StructField("salary", DoubleType(), nullable=True),
    StructField("years_exp", IntegerType(), nullable=True),
    StructField("join_date", StringType(), nullable=True),
])

manual_employees_df = spark.createDataFrame(employees_data, schema=manual_schema)
print("=== Manual Schema ===")
manual_employees_df.printSchema()

Output

=== Employees Schema ===
root
 |-- emp_id: long (nullable = true)
 |-- name: string (nullable = true)
 |-- department: string (nullable = true)
 |-- country: string (nullable = true)
 |-- salary: double (nullable = true)
 |-- years_exp: long (nullable = true)
 |-- join_date: string (nullable = true)

=== Sales Schema ===
root
 |-- sale_id: long (nullable = true)
 |-- emp_id: long (nullable = true)
 |-- product: string (nullable = true)
 |-- category: string (nullable = true)
 |-- price: double (nullable = true)
 |-- quantity: long (nullable = true)
 |-- sale_date: string (nullable = true)

=== Casted Employees Schema ===
root
 |-- emp_id: long (nullable = true)
 |-- name: string (nullable = true)
 |-- department: string (nullable = true)
 |-- country: string (nullable = true)
 |-- salary: integer (nullable = true)
 |-- years_exp: long (nullable = true)
 |-- join_date: date (nullable = true)

root
 |-- sale_id: long (nullable = true)
 |-- emp_id: long (nullable = true)
 |-- product: string (nullable = true)
 |-- category: string (nullable = true)
 |-- price: double (nullable = true)
 |-- quantity: long (nullable = true)
 |-- sale_date: timestamp (nullable = true)

=== Manual Schema ===
root
 |-- emp_id: integer (nullable = false)
 |-- name: string (nullable = true)
 |-- department: string (nullable = true)
 |-- country: string (nullable = true)
 |-- salary: double (nullable = true)
 |-- years_exp: integer (nullable = true)
 |-- join_date: string (nullable = true)

Solution 16 — orderBy, limit, sample

In [ ]

PySpark · In [ ]cell 102

# 1 and 2. Sort and get the top 3
print("=== Top 3 Highest Paid ===")
(
    employeesDF
    .orderBy(F.col("salary").desc_nulls_last(), F.col("name").asc())
    .limit(3)
    .show()
)

# 3. Bottom 3
print("=== Bottom 3 Lowest Paid ===")
employeesDF.orderBy(F.col("salary").asc_nulls_last()).limit(3).show()

# 4. Sample 50%
print("=== 50% Sample ===")
employeesDF.sample(withReplacement=False, fraction=0.5, seed=42).show()

# 5. Earliest five sales
print("=== First 5 Sales by Date ===")
salesDF.orderBy(F.to_date("sale_date").asc()).show(5)

Output

=== Top 3 Highest Paid ===
+------+-------+-----------+-------+--------+---------+----------+
|emp_id|   name| department|country|  salary|years_exp| join_date|
+------+-------+-----------+-------+--------+---------+----------+
|     5|    Eve|Engineering|Germany|120000.0|       10|2014-07-04|
|     3|Charlie|Engineering|     US|110000.0|        8|2016-01-20|
|     8|   Hank|Engineering|Germany| 98000.0|        6|2018-05-22|
+------+-------+-----------+-------+--------+---------+----------+

=== Bottom 3 Lowest Paid ===
+------+-----+----------+-------+-------+---------+----------+
|emp_id| name|department|country| salary|years_exp| join_date|
+------+-----+----------+-------+-------+---------+----------+
|     7|Grace|        HR|     US|58000.0|        1|2023-02-14|
|     4|Diana|        HR|  India|60000.0|        2|2022-09-10|
|     6|Frank| Marketing|  India|68000.0|        4|2020-11-30|
+------+-----+----------+-------+-------+---------+----------+

=== 50% Sample ===
+------+-----+-----------+-------+--------+---------+----------+
|emp_id| name| department|country|  salary|years_exp| join_date|
+------+-----+-----------+-------+--------+---------+----------+
|     5|  Eve|Engineering|Germany|120000.0|       10|2014-07-04|
|     6|Frank|  Marketing|  India| 68000.0|        4|2020-11-30|
|     8| Hank|Engineering|Germany| 98000.0|        6|2018-05-22|
|     9|  Ivy|  Marketing|Germany| 75000.0|        5|2019-08-09|
|    10| Jack|         HR|     US|    NULL|        3|2021-03-17|
+------+-----+-----------+-------+--------+---------+----------+

=== First 5 Sales by Date ===
+-------+------+-------+-----------+------+--------+----------+
|sale_id|emp_id|product|   category| price|quantity| sale_date|
+-------+------+-------+-----------+------+--------+----------+
|    101|     1| Laptop|Electronics|1200.0|       2|2024-01-05|
|    102|     2|  Phone|Electronics| 800.0|       1|2024-01-07|
|    103|     1|   Desk|  Furniture| 450.0|       3|2024-01-10|
|    104|     3| Laptop|Electronics|1200.0|       1|2024-01-12|
|    105|     4|  Chair|  Furniture| 250.0|       4|2024-02-01|
+-------+------+-------+-----------+------+--------+----------+
only showing top 5 rows

Solution 17 — repartition, coalesce, cache

In [ ]

PySpark · In [ ]cell 104

# 1. Current partitions
print(f"Default partitions: {employeesDF.rdd.getNumPartitions()}")

# 2. Repartition to 4 (full shuffle)
repartitioned = employeesDF.repartition(4)
print(f"After repartition(4): {repartitioned.rdd.getNumPartitions()}")

# 3. Coalesce to 2
coalesced = repartitioned.coalesce(2)
print(f"After coalesce(2): {coalesced.rdd.getNumPartitions()}")

# 4. Repartition by column
by_dept = employeesDF.repartition(F.col("department"))
print(f"After repartition(department): {by_dept.rdd.getNumPartitions()}")

# 5. Cache and materialize
cached_df = employeesDF.cache()
cached_df.count()
print("DataFrame cached")

# 6. Unpersist
cached_df.unpersist()
print("DataFrame unpersisted")

Output

Default partitions: 8
After repartition(4): 4
After coalesce(2): 2
After repartition(department): 1
DataFrame cached
DataFrame unpersisted

Solution 18 — Advanced Aggregations

In [ ]

PySpark · In [ ]cell 106

# 1. countDistinct per category
print("=== countDistinct ===")
salesDF.groupBy("category").agg(
    F.countDistinct("product").alias("distinct_products")
).show()

# 2. approx_count_distinct
print("=== approx_count_distinct ===")
salesDF.groupBy("category").agg(
    F.approx_count_distinct("product").alias("approx_distinct")
).show()

# 3. collect_list
print("=== collect_list ===")
salesDF.groupBy("emp_id").agg(
    F.collect_list("product").alias("all_products")
).show(truncate=False)

# 4. collect_set
print("=== collect_set ===")
salesDF.groupBy("emp_id").agg(
    F.collect_set("product").alias("unique_products")
).show(truncate=False)

# 5. Deterministic first and last sale by date
sales_ordered = salesDF.withColumn("sale_date", F.to_date("sale_date"))
sale_window = (
    Window.partitionBy("emp_id")
    .orderBy("sale_date")
    .rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing)
)

first_last = (
    sales_ordered
    .withColumn("first_product", F.first("product").over(sale_window))
    .withColumn("last_product", F.last("product").over(sale_window))
    .select("emp_id", "first_product", "last_product")
    .dropDuplicates(["emp_id"])
)
print("=== First & Last Sale ===")
first_last.show()

Output

=== countDistinct ===
+-----------+-----------------+
|   category|distinct_products|
+-----------+-----------------+
|Electronics|                4|
|  Furniture|                2|
+-----------+-----------------+

=== approx_count_distinct ===
+-----------+---------------+
|   category|approx_distinct|
+-----------+---------------+
|Electronics|              4|
|  Furniture|              2|
+-----------+---------------+

=== collect_list ===
+------+---------------------+
|emp_id|all_products         |
+------+---------------------+
|1     |[Laptop, Desk, Phone]|
|2     |[Phone, Monitor]     |
|3     |[Laptop]             |
|5     |[Laptop]             |
|4     |[Chair]              |
|7     |[Desk]               |
|6     |[Keyboard]           |
|9     |[Monitor]            |
|8     |[Chair]              |
+------+---------------------+

=== collect_set ===
+------+---------------------+
|emp_id|unique_products      |
+------+---------------------+
|1     |[Laptop, Phone, Desk]|
|2     |[Phone, Monitor]     |
|3     |[Laptop]             |
|5     |[Laptop]             |
|4     |[Chair]              |
|7     |[Desk]               |
|6     |[Keyboard]           |
|9     |[Monitor]            |
|8     |[Chair]              |
+------+---------------------+

=== First & Last Sale ===
+------+-------------+------------+
|emp_id|first_product|last_product|
+------+-------------+------------+
|     1|       Laptop|       Phone|
|     2|        Phone|     Monitor|
|     3|       Laptop|      Laptop|
|     4|        Chair|       Chair|
|     5|       Laptop|      Laptop|
|     6|     Keyboard|    Keyboard|
|     7|         Desk|        Desk|
|     8|        Chair|       Chair|
|     9|      Monitor|     Monitor|
+------+-------------+------------+

Solution 19 — Complex Pipeline

In [ ]

PySpark · In [ ]cell 108

dept_window = Window.partitionBy("department").orderBy(
    F.col("salary").desc_nulls_last()
)

# Revenue per employee
revenue_df = (
    salesDF
    .withColumn("revenue", F.col("price") * F.col("quantity"))
    .groupBy("emp_id")
    .agg(F.round(F.sum("revenue"), 2).alias("revenue_per_emp"))
)

pipeline = (
    employeesDF
    .join(revenue_df, ["emp_id"], "left")
    .na.fill({"revenue_per_emp": 0.0})
    .join(
        departmentsDF,
        employeesDF["department"] == departmentsDF["dept_name"],
        "left",
    )
    .withColumn(
        "salary_pct_of_budget",
        F.round((F.col("salary") / F.col("budget")) * 100, 2),
    )
    .withColumn("dept_salary_rank", F.rank().over(dept_window))
    .withColumn(
        "perf_label",
        F.when(F.col("revenue_per_emp") > 3000, "Star")
        .when(F.col("revenue_per_emp") > 1000, "Good")
        .otherwise("Low"),
    )
    .select(
        "name", "department", "salary", "revenue_per_emp",
        "salary_pct_of_budget", "dept_salary_rank", "perf_label",
    )
    .orderBy("department", "dept_salary_rank")
)

pipeline.show()

Output

+-------+-----------+--------+---------------+--------------------+----------------+----------+
|   name| department|  salary|revenue_per_emp|salary_pct_of_budget|dept_salary_rank|perf_label|
+-------+-----------+--------+---------------+--------------------+----------------+----------+
|    Eve|Engineering|120000.0|         6000.0|                24.0|               1|      Star|
|Charlie|Engineering|110000.0|         1200.0|                22.0|               2|      Good|
|   Hank|Engineering| 98000.0|         1500.0|                19.6|               3|      Good|
|  Alice|Engineering| 95000.0|         4550.0|                19.0|               4|      Star|
|  Diana|         HR| 60000.0|         1000.0|                40.0|               1|       Low|
|  Grace|         HR| 58000.0|          450.0|               38.67|               2|       Low|
|   Jack|         HR|    NULL|            0.0|                NULL|               3|       Low|
|    Ivy|  Marketing| 75000.0|         1800.0|                25.0|               1|      Good|
|    Bob|  Marketing| 72000.0|         2000.0|                24.0|               2|      Good|
|  Frank|  Marketing| 68000.0|          800.0|               22.67|               3|       Low|
+-------+-----------+--------+---------------+--------------------+----------------+----------+

Solution 20 — Spark SQL

In [ ]

PySpark · In [ ]cell 110

# 1. Department with the highest average salary
print("=== Highest Avg Salary Dept ===")
spark.sql("""
  SELECT department, ROUND(AVG(salary), 2) AS avg_salary
  FROM employees
  GROUP BY department
  ORDER BY avg_salary DESC
  LIMIT 1
""").show()

# 2. Employees earning more than their department average
print("=== Above Dept Average ===")
spark.sql("""
  SELECT e.name, e.department, e.salary, ROUND(d.avg_sal, 2) AS dept_avg
  FROM employees e
  JOIN (
    SELECT department, AVG(salary) AS avg_sal
    FROM employees
    GROUP BY department
  ) d ON e.department = d.department
  WHERE e.salary > d.avg_sal
""").show()

# 3. Top two earners per department
print("=== Top 2 per Dept (SQL Window) ===")
spark.sql("""
  SELECT name, department, salary, rn FROM (
    SELECT name, department, salary,
           ROW_NUMBER() OVER (
             PARTITION BY department ORDER BY salary DESC
           ) AS rn
    FROM employees
  ) ranked
  WHERE rn <= 2
""").show()

# 4. Total revenue per department using all three tables
print("=== Revenue per Dept ===")
spark.sql("""
  SELECT e.department, d.dept_head,
         ROUND(SUM(s.price * s.quantity), 2) AS total_revenue
  FROM employees e
  JOIN sales s ON e.emp_id = s.emp_id
  JOIN departments d ON e.department = d.dept_name
  GROUP BY e.department, d.dept_head
  ORDER BY total_revenue DESC
""").show()

# 5. CTE: departments where salary exceeds 50% of budget
print("=== Salary > 50% of Budget (CTE) ===")
spark.sql("""
  WITH dept_salary AS (
    SELECT e.department, SUM(e.salary) AS total_salary, d.budget
    FROM employees e
    JOIN departments d ON e.department = d.dept_name
    GROUP BY e.department, d.budget
  )
  SELECT department,
         ROUND(total_salary, 0) AS total_salary,
         budget,
         ROUND(total_salary / budget * 100, 1) AS pct_of_budget
  FROM dept_salary
  WHERE total_salary > 0.5 * budget
  ORDER BY pct_of_budget DESC
""").show()

Output

=== Highest Avg Salary Dept ===
+-----------+----------+
| department|avg_salary|
+-----------+----------+
|Engineering|  105750.0|
+-----------+----------+

=== Above Dept Average ===
+-------+-----------+--------+--------+
|   name| department|  salary|dept_avg|
+-------+-----------+--------+--------+
|    Bob|  Marketing| 72000.0|71666.67|
|Charlie|Engineering|110000.0|105750.0|
|    Eve|Engineering|120000.0|105750.0|
|  Diana|         HR| 60000.0| 59000.0|
|    Ivy|  Marketing| 75000.0|71666.67|
+-------+-----------+--------+--------+

=== Top 2 per Dept (SQL Window) ===
+-------+-----------+--------+---+
|   name| department|  salary| rn|
+-------+-----------+--------+---+
|    Eve|Engineering|120000.0|  1|
|Charlie|Engineering|110000.0|  2|
|  Diana|         HR| 60000.0|  1|
|  Grace|         HR| 58000.0|  2|
|    Ivy|  Marketing| 75000.0|  1|
|    Bob|  Marketing| 72000.0|  2|
+-------+-----------+--------+---+

=== Revenue per Dept ===
+-----------+-------------+-------------+
| department|    dept_head|total_revenue|
+-----------+-------------+-------------+
|Engineering|Alice Johnson|      13250.0|
|  Marketing| Bob Williams|       4600.0|
|         HR|  Carol Smith|       1450.0|
+-----------+-------------+-------------+

=== Salary > 50% of Budget (CTE) ===
+-----------+------------+--------+-------------+
| department|total_salary|  budget|pct_of_budget|
+-----------+------------+--------+-------------+
|Engineering|    423000.0|500000.0|         84.6|
|         HR|    118000.0|150000.0|         78.7|
|  Marketing|    215000.0|300000.0|         71.7|
+-----------+------------+--------+-------------+

Preparing chapter

Loading your data design lesson

The curriculum shell is ready while the requested chapter is being prepared. You can wait here or return to the design library.

Open design library

Continue learning

NetflixStreaming analytics platform

UberMarketplace data platform

YouTubeVideo analytics platform