Interview Prep
PySpark Coding
Interview Questions
Hands-on PySpark problems: joins, window functions, aggregations, UDFs, and performance tuning scenarios.
10Questions
FreeAlways
Vishal Taneja
Q.1
You receive employee data in Spark. How would you calculate a 15% bonus column? Easy
from pyspark.sql import SparkSession
from pyspark.sql.functions import col
# Create SparkSession
spark = SparkSession.builder \
.appName("EmployeeBonus") \
.master("local[*]") \
.getOrCreate()
# Sample data
data = [
(1, "Amit", 50000),
(2, "Rahul", 60000)
]
# Column names
columns = ["id", "name", "salary"]
# Create DataFrame
df = spark.createDataFrame(data, columns)
# Display original data
print("Original Data:")
df.show()
# Calculate 15% bonus
result = df.withColumn(
"bonus",
col("salary") * 0.15
)
# Display result
print("Employee Bonus:")
result.show()
# Stop Spark
spark.stop()
▶ Run to see output
Vishal Taneja
Q.2
Find Employee whoes salary greater than 55,0000? Easy
from pyspark.sql import SparkSession
from pyspark.sql.functions import col
# 1. Create SparkSession
spark = SparkSession.builder \
.appName("EmployeeFiltering") \
.master("local[*]") \
.getOrCreate()
# 2. Create sample employee data
data = [
(1, "Amit", 50000),
(2, "Rahul", 60000),
(3, "John", 75000),
(4, "Priya", 55000),
(5, "Neha", 90000)
]
# 3. Define columns
columns = ["employee_id", "name", "salary"]
# 4. Create DataFrame
df = spark.createDataFrame(data, columns)
# 5. Display original DataFrame
print("Original Data:")
df.show()
# 6. Filter employees whose salary is greater than 55,000
result = df.filter(col("salary") > 55000)
# 7. Display filtered result
print("Employees with salary greater than 55,000:")
result.show()
# 8. Stop Spark
spark.stop()
▶ Run to see output
Vishal Taneja
Q.3
A bank wants to identify customers whose salary has significantly declined over the last 3 months. Hard
Business Requirement
A bank wants to identify customers whose salary has significantly declined over the last 3 months.
Risk Criteria
- Customer should have monthly salary credits.
- Compare the latest month's salary with the average salary of the previous 3 months.
- If the salary has dropped by more than 40%, flag the customer.
- Ignore customers having fewer than 4 salary credits.
Input DataFrame
account_transactions
| account_id | txn_date | txn_type | amount |
|---|---|---|---|
| 1001 | 2024-01-05 | CREDIT | 100000 |
| 1001 | 2024-02-05 | CREDIT | 105000 |
| 1001 | 2024-03-05 | CREDIT | 98000 |
| 1001 | 2024-04-05 | CREDIT | 50000 |
| 1002 | 2024-01-05 | CREDIT | 70000 |
| 1002 | 2024-02-05 | CREDIT | 72000 |
| 1002 | 2024-03-05 | CREDIT | 71000 |
| 1002 | 2024-04-05 | CREDIT | 70000 |
Expected Output
| account_id | avg_previous_3_months | current_salary | drop_percentage | risk_flag |
|---|---|---|---|---|
| 1001 | 101000.00 | 50000 | 50.50% | FLAG |
| 1002 | 71000.00 | 70000 | 1.41% | NOT_FLAG |
from pyspark.sql import SparkSession
from pyspark.sql.functions import (
col,
avg,
count,
round,
when,
row_number
)
from pyspark.sql.window import Window
# ============================================================
# 1. Create Spark Session
# ============================================================
spark = SparkSession.builder \
.appName("SalaryDropAnalysis") \
.master("local[*]") \
.getOrCreate()
# ============================================================
# 2. Create Source Data
# ============================================================
data = [
(1001, "2024-01-05", "CREDIT", 100000),
(1001, "2024-02-05", "CREDIT", 105000),
(1001, "2024-03-05", "CREDIT", 98000),
(1001, "2024-04-05", "CREDIT", 50000),
(1002, "2024-01-05", "CREDIT", 70000),
(1002, "2024-02-05", "CREDIT", 72000),
(1002, "2024-03-05", "CREDIT", 71000),
(1002, "2024-04-05", "CREDIT", 70000),
# Only 3 credits - should be ignored
(1003, "2024-02-05", "CREDIT", 80000),
(1003, "2024-03-05", "CREDIT", 82000),
(1003, "2024-04-05", "CREDIT", 40000)
]
columns = [
"account_id",
"txn_date",
"txn_type",
"amount"
]
# ============================================================
# 3. Create DataFrame
# ============================================================
df = spark.createDataFrame(data, columns)
print("SOURCE DATA")
df.show()
# ============================================================
# 4. Convert transaction date to date
# ============================================================
df = df.withColumn(
"txn_date",
col("txn_date").cast("date")
)
# ============================================================
# 5. Keep only salary CREDIT transactions
# ============================================================
salary_df = df.filter(
col("txn_type") == "CREDIT"
)
print("SALARY CREDIT DATA")
salary_df.show()
# ============================================================
# 6. Check customers having at least 4 salary credits
# ============================================================
valid_customers = salary_df \
.groupBy("account_id") \
.agg(
count("*").alias("credit_count")
) \
.filter(
col("credit_count") >= 4
)
print("CUSTOMERS WITH AT LEAST 4 CREDITS")
valid_customers.show()
# ============================================================
# 7. Keep only valid customers
# ============================================================
salary_df = salary_df.join(
valid_customers.select("account_id"),
on="account_id",
how="inner"
)
# ============================================================
# 8. Create Window
#
# Latest transaction gets rn = 1
# Previous transactions get rn = 2, 3, 4
# ============================================================
window_spec = Window \
.partitionBy("account_id") \
.orderBy(col("txn_date").desc())
# ============================================================
# 9. Assign row number
# ============================================================
salary_df = salary_df.withColumn(
"rn",
row_number().over(window_spec)
)
print("DATA WITH ROW NUMBER")
salary_df.show()
# ============================================================
# 10. Get latest salary
# ============================================================
current_salary_df = salary_df \
.filter(col("rn") == 1) \
.select(
"account_id",
col("amount").alias("current_salary")
)
print("CURRENT SALARY")
current_salary_df.show()
# ============================================================
# 11. Get previous 3 months average salary
# ============================================================
previous_salary_df = salary_df \
.filter(col("rn").between(2, 4)) \
.groupBy("account_id") \
.agg(
avg("amount").alias("avg_previous_3_months")
)
print("AVERAGE OF PREVIOUS 3 MONTHS")
previous_salary_df.show()
# ============================================================
# 12. Join current salary and previous salary
# ============================================================
result_df = current_salary_df.join(
previous_salary_df,
on="account_id",
how="inner"
)
# ============================================================
# 13. Calculate salary drop percentage
#
# Drop % =
# ((Previous Average - Current Salary)
# / Previous Average) * 100
# ============================================================
result_df = result_df.withColumn(
"drop_percentage",
round(
(
(
col("avg_previous_3_months")
- col("current_salary")
)
/ col("avg_previous_3_months")
) * 100,
2
)
)
# ============================================================
# 14. Apply Risk Rule
#
# Drop > 40% --> FLAG
# Otherwise --> NOT_FLAG
# ============================================================
result_df = result_df.withColumn(
"risk_flag",
when(
col("drop_percentage") > 40,
"FLAG"
).otherwise("NOT_FLAG")
)
# ============================================================
# 15. Final Result
# ============================================================
final_df = result_df.select(
"account_id",
"avg_previous_3_months",
"current_salary",
"drop_percentage",
"risk_flag"
).orderBy("account_id")
print("FINAL RESULT")
final_df.show(truncate=False)
# ============================================================
# 16. Stop Spark
# ============================================================
spark.stop()
▶ Run to see output
pyspark
Q.4
Given a DataFrame containing flight/train routes and fares, find the cheapest fare between a source and destination, including the possibility of travelling through an intermediate city? Medium
Input DataFrame
Suppose we have this route data:
| source | destination | fare |
|---|---|---|
| Delhi | Mumbai | 5000 |
| Delhi | Pune | 3000 |
| Pune | Mumbai | 1500 |
| Delhi | Bangalore | 4000 |
| Bangalore | Mumbai | 2500 |
| Delhi | Chennai | 4500 |
| Chennai | Mumbai | 3000 |
We want to find the cheapest way to travel:
Delhi → Mumbai
Possible routes
| Route | Calculation | Total Fare |
|---|---|---|
| Delhi → Mumbai | 5000 | 5000 |
| Delhi → Pune → Mumbai | 3000 + 1500 | 4500 |
| Delhi → Bangalore → Mumbai | 4000 + 2500 | 6500 |
| Delhi → Chennai → Mumbai | 4500 + 3000 | 7500 |
Therefore:
Cheapest fare = ₹4,500
Route = Delhi → Pune → Mumbai
from pyspark.sql import SparkSession
from pyspark.sql.functions import col
# Create Spark Session
spark = SparkSession.builder \
.appName("CheapestFare") \
.getOrCreate()
# Route data
data = [
("Delhi", "Mumbai", 5000),
("Delhi", "Pune", 3000),
("Pune", "Mumbai", 1500),
("Delhi", "Bangalore", 4000),
("Bangalore", "Mumbai", 2500),
("Delhi", "Chennai", 4500),
("Chennai", "Mumbai", 3000)
]
columns = ["source", "destination", "fare"]
routes_df = spark.createDataFrame(data, columns)
# Self join to find routes through an intermediate city
first_leg = routes_df.alias("a")
second_leg = routes_df.alias("b")
via_city = first_leg.join(
second_leg,
col("a.destination") == col("b.source"),
"inner"
)
# Filter source and destination
via_city = via_city.filter(
(col("a.source") == "Delhi") &
(col("b.destination") == "Mumbai")
)
# Calculate total fare
via_city = via_city.withColumn(
"total_fare",
col("a.fare") + col("b.fare")
)
# Select required columns
result = via_city.select(
col("a.source").alias("source"),
col("a.destination").alias("intermediate_city"),
col("b.destination").alias("destination"),
col("a.fare").alias("first_leg_fare"),
col("b.fare").alias("second_leg_fare"),
col("total_fare")
)
# Find cheapest route
cheapest = result.orderBy(
col("total_fare").asc()
).limit(1)
cheapest.show()
▶ Run to see output
Logic Explain:
“I have a routes DataFrame containing source, destination and fare. To find a route through an intermediate city, I self-join the DataFrame. The destination of the first leg should match the source of the second leg. Then I add the two fares to calculate the total fare and order by total fare ascending. Finally, I take the first record, which gives me the cheapest route.
Remember the logic
Self Join → Find Intermediate City → Add Fares → Sort Ascending → limit(1)
pyspark
Q.5
Given a sales DataFrame containing product_id, unit_price, and quantity, write PySpark code to calculate the total amount for each product and the overall total sales amount. Easy
1. Input DataFrame
Think of the DataFrame like this table:
| product_id | unit_price | quantity |
|---|---|---|
| 101 | 100.0 | 2 |
| 102 | 200.0 | 3 |
| 103 | 50.0 | 4 |
| 104 | 150.0 | 2 |
For each row:
Total Amount = unit_price × quantity
from pyspark.sql import SparkSession
from pyspark.sql.functions import col, sum
# Create Spark Session
spark = SparkSession.builder \
.appName("SalesCalculation") \
.getOrCreate()
# Sample sales data
data = [
(101, 100.0, 2),
(102, 200.0, 3),
(103, 50.0, 4),
(104, 150.0, 2)
]
columns = ["product_id", "unit_price", "quantity"]
# Create DataFrame
sales_df = spark.createDataFrame(data, columns)
print("Input Data:")
sales_df.show()
# Calculate total amount for each row
sales_with_amount = sales_df.withColumn(
"total_amount",
col("unit_price") * col("quantity")
)
print("Sales with Total Amount:")
sales_with_amount.show()
# Calculate overall total sales
total_sales = sales_with_amount.agg(
sum("total_amount").alias("total_sales_amount")
)
print("Overall Total Sales:")
total_sales.show()
spark.stop()
▶ Run to see output
Nitesh Kumar Sharma
Q.6
Detect Salary Drop Risk Customers Hard
Business Requirement
The bank wants to identify customers whose salary has significantly declined over the last 3 months.
Risk Criteria:
- Customer receives monthly salary credits.
- Compare latest month salary with average of previous 3 months.
- If salary drops by more than 40%, flag them.
- Ignore customers with fewer than 4 salary credits.
account_transactions txn_date txn_type amount
1001 2024-01-05 CREDIT 100000
1001 2024-02-05 CREDIT 105000
1001 2024-03-05 CREDIT 98000
1001 2024-04-05 CREDIT 50000
1002 2024-01-05 CREDIT 70000
1002 2024-02-05 CREDIT 72000
1002 2024-03-05 CREDIT 71000
1002 2024-04-05 CREDIT 70000
Output details
Customer 1001 gets flagged.
Avg Salary:
(100000+105000+98000)/3 = 101000
Current Salary:
50000
Drop = 50.49%
Vishal Taneja
Q.7
Find the number of transactions that occurred for each product. Output the product name along with the transaction count and sort by product_id Hard
You are given two datasets:
1. Inventory Table
product_id
product_name
2. Transactions Table
transaction_id
time
product_id
Output the product name along with the transaction count and sort by product_id.
Handle duplicate transactions (keep latest)
Count the number of transactions per product
Exclude products with no transactions
Sort the result by product_id
inventory_data = [
(1, 'strawberry'),
(2, 'apple_fuji'),
(3, 'orange'),
(4, 'clementines'),
(5, 'blood_orange'),
(6, 'banana') # no transactions (to test filtering)
]
inv_cols = ["product_id", "product_name"]
transaction_data = [
(153, '2016-01-06 08:57:52', 1),
(91, '2016-01-07 12:17:27', 1),
(31, '2016-01-05 13:19:25', 1),
(24, '2016-01-03 10:47:44', 3),
(4, '2016-01-06 17:57:42', 3),
(163, '2016-01-03 10:11:22', 3),
(92, '2016-01-08 12:03:20', 2),
(32, '2016-01-04 19:37:14', 4),
(253, '2016-01-06 14:15:20', 5),
(118, '2016-01-06 14:27:33', 5),
(118, '2016-01-07 14:27:33', 5), # duplicate (keep latest)
(300, '2016-01-02 09:00:00', 7) # product not in inventory (to test join)
]
trans_cols = ["transaction_id", "time", "product_id"]
from pyspark.sql import SparkSession
from pyspark.sql.functions import col, row_number, count
from pyspark.sql.window import Window
spark = SparkSession.builder.appName("InterviewSolution").getOrCreate()
# DataFrames
inv_df = spark.createDataFrame(inventory_data, ["product_id", "product_name"])
trans_df = spark.createDataFrame(transaction_data, ["transaction_id", "time", "product_id"])
# Step 1: Cast time to timestamp
trans_df = trans_df.withColumn("time", col("time").cast("timestamp"))
# Step 2: Deduplicate using window
window_spec = Window.partitionBy("transaction_id").orderBy(col("time").desc())
dedup_df = trans_df.withColumn("rn", row_number().over(window_spec)) \
.filter(col("rn") == 1) \
.drop("rn")
# Step 3: Inner join
joined_df = dedup_df.join(inv_df, on="product_id", how="inner")
# Step 4: Aggregate
result_df = joined_df.groupBy("product_id", "product_name") \
.agg(count("*").alias("transaction_count"))
# Step 5: Sort
final_df = result_df.orderBy("product_id")
final_df.show()
▶ Run to see output
Vishal Taneja
Q.8
You are given two datasets: customers(customer_id, name, city) orders(customer_id, item, amount) Easy
Write a PySpark program to perform an JOIN and return only customers who have placed at least one order.
customers = [
("C001", "Arjun", "Chennai"),
("C002", "Priya", None),
("C003", "Ravi", None)
]
orders = [
("C001", "Idli", 50),
("C002", "Coffee", 120),
("C004", "Dosa", 60)
]
from pyspark.sql import SparkSession
from pyspark.sql.functions import col
# Create Spark session
spark = SparkSession.builder.appName("CustomerOrders").getOrCreate()
# Create DataFrames
customers_df = spark.createDataFrame(customers, ["customer_id", "name", "city"])
orders_df = spark.createDataFrame(orders, ["customer_id", "item", "amount"])
# Perform INNER JOIN
result_df = customers_df.join(orders_df, on="customer_id", how="inner")
# Select only customer columns
final_df = result_df.select("customer_id", "name", "city")
# Show result
final_df.show()
▶ Run to see output
Sandip Shaw (M24DE3069)(G23AI2095)
Q.9
Give me the list of students who are coming for consecutive 3 days. Easy
| student_id | attendance_date |
|---|---|
| 1 | 01-05-2025 |
| 2 | 01-05-2025 |
| 3 | 01-05-2025 |
| 1 | 02-05-2025 |
| 2 | 02-05-2025 |
| 3 | 02-05-2025 |
| 2 | 03-05-2025 |
| 3 | 03-05-2025 |
| 1 | 04-05-2025 |
| 2 | 04-05-2025 |
CREATE TABLE attendance (
student_id INT,
attendance_date DATE
);
INSERT INTO attendance VALUES
(1,'2025-05-01'),
(2,'2025-05-01'),
(3,'2025-05-01'),
(1,'2025-05-02'),
(2,'2025-05-02'),
(3,'2025-05-02'),
(2,'2025-05-03'),
(3,'2025-05-03'),
(1,'2025-05-04'),
(2,'2025-05-04');
SELECT DISTINCT a1.student_id
FROM attendance a1
JOIN attendance a2
ON a1.student_id = a2.student_id
AND a2.attendance_date = DATE_ADD(a1.attendance_date, INTERVAL 1 DAY)
JOIN attendance a3
ON a1.student_id = a3.student_id
AND a3.attendance_date = DATE_ADD(a1.attendance_date, INTERVAL 2 DAY);
▶ Run to see output
Vishal Taneja
Q.10
Find Dupliacates records in a table? Easy
Table:Employee
| id | name | department |
|---|---|---|
| 1 | Amit | IT |
| 2 | Neha | HR |
| 3 | Amit | IT |
| 4 | Rahul | Sales |
| 5 | Neha | HR |
CREATE TABLE employees (
id INT PRIMARY KEY,
name VARCHAR(100),
department VARCHAR(100)
);
INSERT INTO employees (name, department) VALUES
('Amit', 'IT'),
('Neha', 'HR'),
('Amit', 'IT'),
('Rahul', 'Sales'),
('Neha', 'HR'),
('Sonia', 'Finance');
SELECT name, COUNT(*) AS duplicate_count
FROM employees
GROUP BY name
HAVING COUNT(*) > 1;
▶ Run to see output