Module 4: Working with Data in Databricks

customers_csv = spark.read.option("header", True).csv(f"{base}/customers_csv")
products_json = spark.read.json(f"{base}/products_json")
orders_parquet = spark.read.parquet(f"{base}/orders_parquet")

for name, df in [("customers", customers_csv), ("products", products_json), ("orders", orders_parquet)]:
    print(name, df.count(), "rows")

orders_parquet.write.mode("overwrite").saveAsTable("orders_raw")
spark.read.table("orders_raw").printSchema()

spark.read.option("header", True).csv(f"{base}/customers_csv").printSchema()

customers = (spark.read
    .schema("customer_id INT, name STRING, country STRING, signup_date DATE")
    .option("header", True)
    .csv(f"{base}/customers_csv"))
customers.printSchema()

customers = (spark.read
    .schema("customer_id INT, name STRING, country STRING, signup_date DATE")
    .option("header", True).csv(f"{base}/customers_csv"))
products = (spark.read
    .schema("product_id INT, product_name STRING, category STRING, unit_price DECIMAL(8,2)")
    .json(f"{base}/products_json"))
orders = spark.read.parquet(f"{base}/orders_parquet")

from pyspark.sql import Row
import pandas as pd

df_a = spark.createDataFrame([(1, "Ada", "NL"), (2, "Grace", "DE")], "id INT, name STRING, country STRING")
df_b = spark.createDataFrame([Row(id=3, name="Linus", country="FI")])    # schema inferred
df_c = spark.createDataFrame(pd.DataFrame({"id": [4], "name": ["Margaret"], "country": ["US"]}))
df_d = spark.range(5)
df_e = spark.sql("SELECT 1 AS one, 'a' AS letter")
df_f = spark.read.table("orders_raw")

df_a.unionByName(df_b).show()

from pyspark.sql import functions as F

orders_t = (
    orders
    .withColumn("order_date", F.to_date("order_ts"))
    .withColumn("size", F.when(F.col("quantity") >= 4, "bulk").otherwise("regular"))
    .withColumnRenamed("status", "order_status")
    .select("order_id", "customer_id", "product_id", "quantity",
            "size", "order_status", "order_date")
)
display(orders_t.limit(10))

completed = orders.filter(F.col("status") == "completed")

bulk_or_shipped = orders.where(
    (F.col("quantity") >= 4) & (F.col("status").isin("completed", "shipped"))
)

recent = orders.filter(F.col("order_ts") >= "2025-06-01").orderBy(F.col("order_ts").desc()).limit(5)
display(recent)

summary = (
    orders.groupBy("status")
          .agg(F.count("*").alias("orders"),
               F.round(F.avg("quantity"), 2).alias("avg_quantity"),
               F.countDistinct("customer_id").alias("customers"))
          .orderBy(F.desc("orders"))
)
summary.show()

# dedupe first so revenue isn't counted twice
order_lines = (
    orders.dropDuplicates(["order_id"])
          .join(products, "product_id", "left")
          .withColumn("revenue", F.col("quantity") * F.col("unit_price"))
)

by_category = (order_lines.groupBy("category")
                 .agg(F.round(F.sum("revenue"), 2).alias("revenue"))
                 .orderBy(F.desc("revenue")))
by_category.show()

no_orders = customers.join(orders.select("customer_id").distinct(), "customer_id", "left_anti")
print(no_orders.count())    # may be 0: every customer ordered at least once

from pyspark.sql.functions import broadcast
joined = orders.join(broadcast(products), "product_id")    # small table sent to every executor

from pyspark.sql.window import Window

product_revenue = (order_lines
    .groupBy("category", "product_id", "product_name")
    .agg(F.sum("revenue").alias("revenue")))

w_rank = Window.partitionBy("category").orderBy(F.desc("revenue"))
top3 = (product_revenue
    .withColumn("rank", F.row_number().over(w_rank))
    .filter(F.col("rank") <= 3))
display(top3.orderBy("category", "rank"))

w_run = (Window.partitionBy("customer_id").orderBy("order_ts")
               .rowsBetween(Window.unboundedPreceding, Window.currentRow))
running = order_lines.withColumn("running_revenue", F.sum("revenue").over(w_run))

null_counts = orders.select(
    [F.count(F.when(F.col(c).isNull(), c)).alias(c) for c in orders.columns]
)
null_counts.show()

print(orders.count(), orders.select("order_id").distinct().count())
print(orders.distinct().count(), orders.dropDuplicates(["order_id"]).count())

from pyspark.sql.window import Window

w = Window.partitionBy("order_id").orderBy(F.desc("order_ts"))
latest = (orders.withColumn("rn", F.row_number().over(w))
                .filter("rn = 1")
                .drop("rn"))

display(
    customers.select(
        "customer_id",
        F.upper("country").alias("country_uc"),
        F.lpad(F.col("customer_id").cast("string"), 5, "0").alias("customer_code"),
        F.concat_ws(" | ", "name", "country").alias("label"),
        F.year("signup_date").alias("signup_year"),
        F.date_format("signup_date", "yyyy-MM").alias("signup_month"),
        F.datediff(F.current_date(), "signup_date").alias("days_since_signup"),
    ).limit(10)
)

tags = spark.createDataFrame([(1, "a,b,c"), (2, "d")], "id INT, tags STRING")
display(tags.select("id", F.explode(F.split("tags", ",")).alias("tag")))     # list -> rows

monthly = (orders.withColumn("month", F.date_trunc("month", "order_ts"))
                 .groupBy("month").count().orderBy("month"))                  # day -> month
display(monthly)

orders.createOrReplaceTempView("orders_v")
customers.createOrReplaceTempView("customers_v")

%sql
SELECT c.country,
       COUNT(DISTINCT o.order_id) AS orders
FROM orders_v o
JOIN customers_v c USING (customer_id)
GROUP BY c.country
ORDER BY orders DESC

%sql
CREATE OR REPLACE TEMP VIEW big_orders AS
SELECT * FROM orders_v WHERE quantity >= 4

big = spark.table("big_orders")      # a DataFrame again

clean = orders.dropDuplicates(["order_id"])

clean.write.mode("overwrite").saveAsTable("orders")
print("after overwrite:", spark.table("orders").count())

clean.filter("order_id <= 10").write.mode("append").saveAsTable("orders")
print("after append:   ", spark.table("orders").count())

clean.write.mode("overwrite").saveAsTable("orders")        # restore
print("after restore:  ", spark.table("orders").count())

try:
    clean.write.mode("errorifexists").saveAsTable("orders")
except Exception as e:
    print("table exists:", str(e).splitlines()[0])

# later topics and modules query these
customers.write.mode("overwrite").saveAsTable("customers")
products.write.mode("overwrite").saveAsTable("products")

# plain files: a folder of part files
clean.write.mode("overwrite").parquet(f"{base}/orders_clean_parquet")

%sql
SELECT o.customer_id,
       ROUND(SUM(o.quantity * p.unit_price), 2) AS revenue
FROM orders o
JOIN products p USING (product_id)
WHERE o.customer_id IS NOT NULL
GROUP BY o.customer_id
ORDER BY revenue DESC
LIMIT 5

top5 = (spark.table("orders")
        .join(spark.table("products"), "product_id")
        .where(F.col("customer_id").isNotNull())
        .groupBy("customer_id")
        .agg(F.round(F.sum(F.col("quantity") * F.col("unit_price")), 2).alias("revenue"))
        .orderBy(F.desc("revenue"))
        .limit(5))
display(top5)

top5.createOrReplaceTempView("top5")

%sql
SELECT t.customer_id, c.name, c.country, t.revenue
FROM top5 t
JOIN customers c USING (customer_id)
ORDER BY t.revenue DESC

display(spark.sql(
    "SELECT * FROM customers WHERE country = :country",
    args={"country": "NL"},
))