Objective

In this problem, you will be working with a DataFrame that represents a set of architectural structures. Each row in the dataset represents a unique building and includes various details about its architecture.

Task

Write a PySpark function that calculates the average height per floor for each building, with the result rounded to two decimal places.

If the number of floors is zero, the average height must gracefully default to zero to avoid mathematical errors. The result should include the id, name, city, country, and the calculated average height per floor (named avg_height_per_floor). Order the final output in ascending order by the avg_height_per_floor. If there is a tie, break it by sorting id in ascending order.

File Path

• Buildings Dataset: /home/interview/buildings.csv
• Starter script: /home/interview/building_heights.py

Schema

buildings.csv

Expected Output Schema

Example

Given this sample input:

buildings

| id | name | city | country | height_m | floors |
|----+---------------------------+----------+--------------+----------+--------|
| 1 | One World Trade Center | New York | USA | 541.3 | 104 |
| 2 | Willis Tower | Chicago | USA | 442.1 | 108 |
| 3 | Burj Khalifa | Dubai | UAE | 828.0 | 163 |
| 4 | The Shard | London | UK | 309.6 | 72 |
| 5 | Abraj Al-Bait | Mecca | Saudi Arabia | 601.0 | 120 |

The expected output would be:

| id | name | city | country | avg_height_per_floor |
|----+---------------------------+----------+--------------+----------------------|
| 2 | Willis Tower | Chicago | USA | 4.09 |
| 4 | The Shard | London | UK | 4.30 |
| 5 | Abraj Al-Bait | Mecca | Saudi Arabia | 5.01 |
| 3 | Burj Khalifa | Dubai | UAE | 5.08 |
| 1 | One World Trade Center | New York | USA | 5.20 |

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

spark = SparkSession.builder.appName("PrepareshSpark").getOrCreate()
spark.sparkContext.setLogLevel("ERROR")

buildings = spark.read.csv("/home/interview/buildings.csv", header=True, inferSchema=True)

# Step 1: Calculate the average height per floor, checking for zero division
result_df = buildings.withColumn(
    "avg_height_per_floor",
    F.when(F.col("floors") == 0, 0.0)
     .otherwise(F.round(F.col("height_m") / F.col("floors"), 2))
)

# Step 2: Select the required columns and order deterministically
result_df = result_df.select(
    "id", "name", "city", "country", "avg_height_per_floor"
).orderBy("avg_height_per_floor", "id")

# --- Do not edit below this line ---
result_df.coalesce(1).write.csv("/home/interview/output", header=True, mode="overwrite")
spark.stop()

Explanation

Step 1: Handling Division by Zero

F.when(F.col("floors") == 0, 0.0)

When dealing with real-world data, you must always anticipate edge cases. If a building is incomplete or an error occurred in data entry resulting in floors = 0, directly dividing height_m by floors will crash the PySpark job with a division-by-zero error. We use F.when() to catch this exact scenario and safely return 0.0.

Step 2: Division and Rounding

.otherwise(F.round(F.col("height_m") / F.col("floors"), 2))

If the floors are greater than 0, the .otherwise() block executes. PySpark allows standard mathematical operators (/, *, +, -) directly between columns. We divide the height by the number of floors, and then wrap that entire calculation in F.round(..., 2) to ensure the decimal precision is strictly truncated to two places as requested by the prompt.

Step 3: Output Formatting and Sorting

result_df.select("id", "name", "city", "country", "avg_height_per_floor") \
         .orderBy("avg_height_per_floor", "id")

Finally, we use .select() to arrange the columns exactly as requested by the Output Schema, automatically dropping the now-unnecessary height_m and floors columns. We chain .orderBy("avg_height_per_floor", "id") to sort the buildings from shortest floor average to tallest, using the building's ID as a tie-breaker to guarantee the order is deterministic.