2026-10-18 18:38:51 +01:00
{
"cells": [
{
"cell_type": "code",
"execution_count": 13,
"id": "09a05cc1",
"metadata": {},
"outputs": [],
"source": [
"from pyspark.sql import SparkSession, DataFrame\n",
"from pyspark.sql import functions as F\n",
"from pathlib import Path\n",
"import glob\n",
"import os\n",
"from dotenv import load_dotenv, find_dotenv"
]
},
{
"cell_type": "code",
2026-10-18 22:55:40 +01:00
"execution_count": null,
2026-10-18 18:38:51 +01:00
"id": "2c1aeb30",
"metadata": {},
"outputs": [],
"source": [
" ## Extract (Load in) ##\n",
"def start_spark(app_name):\n",
2026-10-18 23:26:25 +01:00
" \"\"\"Initialize and return a SparkSession with the PostgreSQL JDBC driver.\"\"\"\n",
2026-10-18 18:38:51 +01:00
" jar_path = str(Path.cwd() / \"postgresql.jar\")\n",
" spark = SparkSession.builder.appName(app_name).config(\"spark.jars\", jar_path).config(\"spark.driver.extraClassPath\", jar_path).getOrCreate()\n",
" return spark\n",
"\n",
"def get_jdbc_config():\n",
" \"\"\"Load env vars and return a Postgres URL\"\"\"\n",
" load_dotenv(find_dotenv())\n",
" \n",
" username=os.getenv(\"DATABASE_USERNAME\")\n",
" password=os.getenv(\"DATABASE_PASSWORD\")\n",
" host=os.getenv(\"DATABASE_IP\")\n",
" port=int(os.getenv(\"DATABASE_PORT\", 5432))\n",
" database=os.getenv(\"DATABASE_NAME\")\n",
" \n",
" missing = [k for k, v in {\n",
" \"DATABASE_USERNAME\": username,\n",
" \"DATABASE_PASSWORD\": password,\n",
" \"DATABASE_IP\": host,\n",
" \"DATABASE_NAME\": database,}.items() if not v]\n",
" if missing:\n",
" raise ValueError(f\"Missing env vars: {missing}\")\n",
"\n",
" url = f\"jdbc:postgresql://{host}:{port}/{database}\"\n",
" properties = {\n",
" \"user\": username,\n",
" \"password\": password,\n",
" \"driver\": \"org.postgresql.Driver\"\n",
" }\n",
" return url, properties\n",
"\n",
"def read_jdbc(spark, table_name):\n",
2026-10-18 23:26:25 +01:00
" \"\"\"Read a table from Postgres via JDBC into a Spark DataFrame.\"\"\"\n",
2026-10-18 18:38:51 +01:00
" url, properties = get_jdbc_config()\n",
2026-10-18 22:55:40 +01:00
" return spark.read.jdbc(url=url, table=table_name, properties=properties)\n",
"\n",
"def databricks_read(df, table_name):\n",
" \"\"\"Read the data from catalog.\"\"\"\n",
" catalog = \"workspace\"\n",
" schema = \"default\"\n",
" path = f\"{catalog}.{schema}.{table_name}\"\n",
" return df.read.table(path)"
2026-10-18 18:38:51 +01:00
]
},
{
"cell_type": "code",
"execution_count": 15,
"id": "9631ba07",
"metadata": {},
"outputs": [],
"source": [
" ## Transform ##"
]
},
{
"cell_type": "code",
2026-10-18 22:55:40 +01:00
"execution_count": null,
2026-10-18 18:38:51 +01:00
"id": "f7673c72",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"[Stage 0:> (0 + 1) / 1]\r"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
2026-10-18 22:25:12 +01:00
"+--------------------+--------------------+-----+\n",
"| track_name| username|count|\n",
"+--------------------+--------------------+-----+\n",
"| Mujeriego| csm3-us| 93|\n",
"| Angelito| csm3-us| 16|\n",
"|Navidad en Mi Barrio| csm3-us| 1|\n",
"| Como Perro Y Gato| csm3-us| 1|\n",
"| The Taste of Ink|fm30ajetkgrrkbvae...| 2|\n",
"| Hasta el Amanecer| csm3-us| 4|\n",
"| Tu Foto| csm3-us| 1|\n",
"|Love Is Not Over ...| csm3-us| 3|\n",
"|Like I'm Gonna Lo...|fm30ajetkgrrkbvae...| 2|\n",
"| Lo Siento BB:/| csm3-us| 5|\n",
"| Superbeast| csm3-us| 1|\n",
"| Amor de Hospital| csm3-us| 11|\n",
"| All Right Now|fm30ajetkgrrkbvae...| 2|\n",
"|Dancing On The Ce...|fm30ajetkgrrkbvae...| 4|\n",
"|Prologue / The Da...|fm30ajetkgrrkbvae...| 4|\n",
"|That Way (feat. J...| csm3-us| 11|\n",
"| Blame| csm3-us| 29|\n",
"| Speak Now|fm30ajetkgrrkbvae...| 2|\n",
"|A Little Respect ...|fm30ajetkgrrkbvae...| 3|\n",
"| Rockstar|fm30ajetkgrrkbvae...| 4|\n",
"+--------------------+--------------------+-----+\n",
"only showing top 20 rows\n"
2026-10-18 18:38:51 +01:00
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
" \r"
]
}
],
"source": [
" ## MAIN ## \n",
"def main():\n",
" app_name = \"ml_media_silver\"\n",
" spark = start_spark(app_name)\n",
"\n",
2026-10-18 22:55:40 +01:00
" bronze_table_name = \"bronze_spotify_data\"\n",
" silver_table_name = \"silver_spotify_data\"\n",
2026-10-19 08:23:52 +01:00
" bronze_df = read_jdbc(spark, bronze_table_name)\n",
2026-10-18 18:38:51 +01:00
"\n",
2026-10-19 08:23:52 +01:00
" bronze_df.groupBy(\"track_name\", \"username\").count().show()\n",
2026-10-18 22:25:12 +01:00
"\n",
" #spark.stop()\n",
2026-10-18 18:38:51 +01:00
"if __name__ == \"__main__\":\n",
2026-10-18 22:55:40 +01:00
" main()\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "c9a695ae",
"metadata": {},
"outputs": [],
"source": [
"## Databricks ##\n",
" ## MAIN ## \n",
"\n",
"# app_name = \"ml_media_silver\"\n",
"# bronze_table_name = \"bronze_spotify_data\"\n",
"# silver_table_name = \"silver_spotify_data\"\n",
"# silver_df = databricks_read(spark, bronze_table_name)\n",
"# silver_df.show(1)\n",
"\n",
"\n"
2026-10-18 18:38:51 +01:00
]
2026-10-19 08:23:52 +01:00
},
{
"cell_type": "code",
"execution_count": null,
"id": "4367ec82",
"metadata": {},
"outputs": [],
"source": [
"## --- ETL: Bronze → Silver (simple functions) --- ##\n",
"\n",
"def clean_platform(df):\n",
" \"\"\"Normalize verbose platform strings into 5 buckets.\"\"\"\n",
" return df.withColumn(\n",
" \"platform\",\n",
" F.when(F.lower(F.col(\"platform\")).contains(\"ios\"), \"ios\")\n",
" .when(F.lower(F.col(\"platform\")).contains(\"android\"), \"android\")\n",
" .when(F.lower(F.col(\"platform\")).rlike(\"windows|osx|mac|linux\"), \"desktop\")\n",
" .when(F.lower(F.col(\"platform\")).contains(\"web\"), \"web\")\n",
" .otherwise(\"other\"),\n",
" )\n",
"\n",
"def derive_skip_label(df):\n",
" \"\"\"Create a binary skip label from reason_end + ms_played.\n",
" A 'skip' = user forwarded or hit back before 30 seconds (Spotify's stream rule).\"\"\"\n",
" return df.withColumn(\n",
" \"is_skip\",\n",
" F.when(F.col(\"reason_end\").isin(\"fwdbtn\", \"backbtn\") & (F.col(\"ms_played\") < 30000), 1)\n",
" .otherwise(0),\n",
" )\n",
"\n",
"def extract_track_id(df):\n",
" \"\"\"Pull the Spotify track ID from the track_uri (spotify:track:XXXX).\"\"\"\n",
" return df.withColumn(\n",
" \"track_id\",\n",
" F.regexp_extract(F.col(\"track_uri\"), r\"spotify:track:([A-Za-z0-9]+)\", 1),\n",
" )\n",
"\n",
"def add_time_features(df):\n",
" \"\"\"Derive hour-of-day and day-of-week from the timestamp.\"\"\"\n",
" return (df\n",
" .withColumn(\"ts\", F.to_timestamp(\"ts\"))\n",
" .withColumn(\"hour\", F.hour(\"ts\"))\n",
" .withColumn(\"dow\", F.dayofweek(\"ts\"))\n",
" )\n",
"\n",
"\n",
"bronze_df = spark.table(\"workspace.default.bronze_spotify_data\")\n",
"\n",
"silver_df = (bronze_df\n",
" .transform(clean_platform)\n",
" .transform(derive_skip_label)\n",
" .transform(extract_track_id)\n",
" .transform(add_time_features)\n",
")\n",
"\n",
"print(f\"Bronze rows: {bronze_df.count()}\")\n",
"print(f\"Silver rows: {silver_df.count()}\")\n",
"silver_df.groupBy(\"is_skip\").count().show()"
]
},
{
"cell_type": "markdown",
"id": "cfe0a7c6",
"metadata": {},
"source": [
"## Feature Engineering (Silver → Gold / ML features)\n",
"\n",
"The silver table feeds into **feature engineering** — the bridge between ETL and ML.\n",
"We build per-play features that a classifier can consume."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3db7ecd1",
"metadata": {},
"outputs": [],
"source": [
"## --- Feature Engineering for Skip Prediction --- ##\n",
"\n",
"from pyspark.ml.feature import StringIndexer, OneHotEncoder, VectorAssembler\n",
"from pyspark.ml import Pipeline\n",
"\n",
"# Select features relevant to skipping behaviour\n",
"feature_cols = [\"ms_played\", \"platform\", \"reason_start\", \"shuffle\", \"hour\", \"dow\"]\n",
"\n",
"gold_df = silver_df.select(\n",
" \"is_skip\",\n",
" F.col(\"ms_played\").cast(\"double\").alias(\"ms_played\"),\n",
" \"platform\", \"reason_start\",\n",
" F.col(\"shuffle\").cast(\"string\").alias(\"shuffle\"),\n",
" \"hour\", \"dow\",\n",
").na.drop()\n",
"\n",
"# Encode categoricals\n",
"categorical_cols = [\"platform\", \"reason_start\", \"shuffle\"]\n",
"indexers = [StringIndexer(inputCol=c, outputCol=f\"{c}_idx\", handleInvalid=\"keep\") for c in categorical_cols]\n",
"encoders = [OneHotEncoder(inputCol=f\"{c}_idx\", outputCol=f\"{c}_oh\") for c in categorical_cols]\n",
"\n",
"assembler = VectorAssembler(\n",
" inputCols=[\"ms_played\", \"platform_oh\", \"reason_start_oh\", \"shuffle_oh\", \"hour\", \"dow\"],\n",
" outputCol=\"features\",\n",
")\n",
"\n",
"feature_pipeline = Pipeline(stages=indexers + encoders + [assembler])\n",
"feature_model = feature_pipeline.fit(gold_df)\n",
"ml_ready = feature_model.transform(gold_df)\n",
"\n",
"print(f\"ML-ready rows: {ml_ready.count()}\")\n",
"ml_ready.select(\"is_skip\", \"features\").show(3, truncate=False)"
]
},
{
"cell_type": "markdown",
"id": "a3ed64ca",
"metadata": {},
"source": [
"## Model Training\n",
"\n",
"**Business question:** _Will a user skip a track within the first 30 seconds?_\n",
"\n",
"This is a **binary classification** problem. We train a Logistic Regression model on the ML-ready features and evaluate with AUC-ROC."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b48e91df",
"metadata": {},
"outputs": [],
"source": [
"## --- Model Training --- ##\n",
"from pyspark.ml.classification import LogisticRegression\n",
"from pyspark.ml.evaluation import BinaryClassificationEvaluator\n",
"\n",
"# Split 80/20\n",
"train_df, test_df = ml_ready.randomSplit([0.8, 0.2], seed=42)\n",
"\n",
"lr = LogisticRegression(featuresCol=\"features\", labelCol=\"is_skip\", maxIter=20)\n",
"lr_model = lr.fit(train_df)\n",
"\n",
"predictions = lr_model.transform(test_df)\n",
"\n",
"evaluator = BinaryClassificationEvaluator(labelCol=\"is_skip\", metricName=\"areaUnderROC\")\n",
"auc = evaluator.evaluate(predictions)\n",
"\n",
"print(f\"Test AUC-ROC: {auc:.4f}\")\n",
"print(f\"Train rows: {train_df.count()} | Test rows: {test_df.count()}\")\n",
"predictions.select(\"is_skip\", \"prediction\", \"probability\").show(5, truncate=False)"
]
},
{
"cell_type": "markdown",
"id": "2c429452",
"metadata": {},
"source": [
"## Inference (Batch Scoring)\n",
"\n",
"Once trained, the model scores new plays to predict skip probability. This is **batch inference** — the same feature pipeline transforms new bronze data and the model produces predictions.\n",
"\n",
"In production this would run on a schedule (e.g., daily) or as a streaming job."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "c225f0e7",
"metadata": {},
"outputs": [],
"source": [
"## --- Inference: Score new plays --- ##\n",
"# Simulate new data: take the last 1000 rows from silver as \"new\" plays\n",
"new_plays = silver_df.orderBy(F.desc(\"ts\")).limit(1000)\n",
"\n",
"# Apply the same feature pipeline\n",
"new_features = feature_model.transform(\n",
" new_plays.select(\n",
" F.col(\"ms_played\").cast(\"double\").alias(\"ms_played\"),\n",
" \"platform\", \"reason_start\",\n",
" F.col(\"shuffle\").cast(\"string\").alias(\"shuffle\"),\n",
" \"hour\", \"dow\",\n",
" ).na.drop()\n",
")\n",
"\n",
"# Score\n",
"scored = lr_model.transform(new_features)\n",
"\n",
"from pyspark.sql.functions import udf\n",
"from pyspark.sql.types import DoubleType\n",
"skip_prob_udf = udf(lambda v: float(v[1]), DoubleType())\n",
"\n",
"display(scored.select(\n",
" \"ms_played\", \"platform\", \"reason_start\", \"shuffle\",\n",
" skip_prob_udf(\"probability\").alias(\"skip_prob\"),\n",
" \"prediction\",\n",
").orderBy(F.desc(\"skip_prob\")).limit(10))"
]
},
2026-10-19 08:27:26 +01:00
{
"cell_type": "code",
"execution_count": null,
"id": "6cd50786",
"metadata": {},
"outputs": [],
"source": [
"## --- Tests for ETL Functions --- ##\n",
"import pytest\n",
"\n",
"# Helper: collect a single column into a Python list\n",
"def col_vals(df, colname):\n",
" return [row[colname] for row in df.select(colname).collect()]\n",
"\n",
"# ── Test 1: clean_platform ──\n",
"def test_clean_platform():\n",
" \"\"\"Platform strings are normalized into 5 buckets.\"\"\"\n",
" sample = spark.createDataFrame([\n",
" (\"iOS 14.2 (iPhone)\",),\n",
" (\"Android OS 11 API 30\",),\n",
" (\"Windows 10\",),\n",
" (\"macOS 12.0\",),\n",
" (\"web player\",),\n",
" (\"Partner amazon_salmon\",),\n",
" ], [\"platform\"])\n",
" result = clean_platform(sample)\n",
" vals = col_vals(result, \"platform\")\n",
" assert vals == [\"ios\", \"android\", \"desktop\", \"desktop\", \"web\", \"other\"], f\"Got {vals}\"\n",
" print(\"[x] test_clean_platform passed\")\n",
"\n",
"# ── Test 2: derive_skip_label ──\n",
"def test_derive_skip_label():\n",
" \"\"\"Skip = fwdbtn/backbtn with ms_played < 30s; everything else is not a skip.\"\"\"\n",
" sample = spark.createDataFrame([\n",
" (\"fwdbtn\", 500.0),\n",
" (\"backbtn\", 200.0),\n",
" (\"trackdone\", 180000.0),\n",
" (\"fwdbtn\", 60000.0), # fwdbtn but >= 30s → not a skip\n",
" (\"endplay\", 100.0), # wrong reason → not a skip\n",
" ], [\"reason_end\", \"ms_played\"])\n",
" result = derive_skip_label(sample)\n",
" vals = col_vals(result, \"is_skip\")\n",
" assert vals == [1, 1, 0, 0, 0], f\"Got {vals}\"\n",
" print(\"[x] test_derive_skip_label passed\")\n",
"\n",
"# ── Test 3: extract_track_id ──\n",
"def test_extract_track_id():\n",
" \"\"\"Track ID is extracted from the spotify:track: URI.\"\"\"\n",
" sample = spark.createDataFrame([\n",
" (\"spotify:track:0cKk8BKEi7zXbdrYdyqBP5\",),\n",
" (\"spotify:track:2VgPA0WnwvE2nbltxa3BHz\",),\n",
" (\"\",), # empty URI → empty string\n",
" ], [\"track_uri\"])\n",
" result = extract_track_id(sample)\n",
" vals = col_vals(result, \"track_id\")\n",
" assert vals[0] == \"0cKk8BKEi7zXbdrYdyqBP5\"\n",
" assert vals[1] == \"2VgPA0WnwvE2nbltxa3BHz\"\n",
" assert vals[2] == \"\"\n",
" print(\"[x] test_extract_track_id passed\")\n",
"\n",
"# ── Test 4: add_time_features ──\n",
"def test_add_time_features():\n",
" \"\"\"Hour and day-of-week are derived from the ts timestamp.\"\"\"\n",
" sample = spark.createDataFrame([\n",
" (\"2021-10-27 20:17:11\",), # Wed, 8pm\n",
" (\"2022-01-03 07:30:00\",), # Mon, 7am\n",
" ], [\"ts\"])\n",
" result = add_time_features(sample)\n",
" hours = col_vals(result, \"hour\")\n",
" dows = col_vals(result, \"dow\")\n",
" assert hours == [20, 7], f\"Got hours {hours}\"\n",
" assert dows == [4, 2], f\"Got dows {dows}\" # Sun=1, Mon=2, …, Wed=4\n",
" print(\"[x] test_add_time_features passed\")\n",
"\n",
"# ── Run all tests ──\n",
"test_clean_platform()\n",
"test_derive_skip_label()\n",
"test_extract_track_id()\n",
"test_add_time_features()\n",
"print(\"\\n🎉 All 4 ETL tests passed!\")"
]
},
2026-10-19 08:23:52 +01:00
{
"cell_type": "markdown",
"id": "23c744a6",
"metadata": {},
"source": [
"## Data Framework & Business Question\n",
"\n",
"### Framework: Medallion Architecture (Lakehouse)\n",
"\n",
"This pipeline uses the **Medallion Architecture** — Databricks' multi-layer data design pattern:\n",
"\n",
"```\n",
"Raw JSON Export (Spotify)\n",
" |\n",
" v\n",
" ┌─────────┐ ┌───────────┐ ┌──────────┐ ┌──────────┐\n",
" │ RAW │→ │ BRONZE │→ │ SILVER │→ │ GOLD │→ ML Model\n",
" │ (exact) │ │ (cleaned) │ │ (conform)│ │ (aggreg) │ (train/infer)\n",
" └─────────┘ └───────────┘ └──────────┘ └──────────┘\n",
"```\n",
"\n",
"| Layer | What Happens | This Project |\n",
"| --- | --- | --- |\n",
"| **Raw** | Land the data exactly as exported | `raw_spotify_data` — 105K rows, original verbose column names, PII included |\n",
"| **Bronze** | Clean names, drop PII, add time columns | `bronze_spotify_data` — 103K rows, renamed columns, derived `hour`/`date` |\n",
"| **Silver** | Conform types, dedupe, sessionize, derive flags | `silver_df` (in-memory) — platform normalized, `is_skip` label, `track_id` extracted |\n",
"| **Gold** | Feature engineering for ML | `ml_ready` — encoded categorical features, assembled vector, ready for training |\n",
"\n",
"### How It Answers a Business Question\n",
"\n",
"**Question:** _Which tracks will a user skip, and can we predict it before it happens?_\n",
"\n",
"**Why it matters:** Skip rate is a key engagement metric for music platforms. If we can predict skips:\n",
"\n",
"* **Recommendation engines** can avoid surfaging tracks a user will skip\n",
"* **Playlist curation** can order tracks to minimize early-exit\n",
"* **A/B testing** can measure how UI changes affect skip behaviour\n",
"\n",
"**How the framework supports this:**\n",
"\n",
"1. **Bronze ETL** (`spotify_bronze_silver`) — lands and cleans the raw export so the data is queryable\n",
"2. **Silver ETL** (this notebook) — derives the `is_skip` label and normalizes features like `platform`\n",
"3. **Feature Engineering** (Gold) — encodes categoricals and assembles a feature vector\n",
"4. **Model Training** — Logistic Regression learns which feature patterns lead to skips\n",
"5. **Inference** — Scores new plays in batch; high `skip_prob` flags tracks to avoid recommending\n",
"\n",
"### ML Engineering Connection\n",
"\n",
"The ETL functions aren't just data plumbing — they are the **ML feature pipeline**:\n",
"\n",
"* `clean_platform()` → categorical feature `platform` (one-hot encoded)\n",
"* `derive_skip_label()` → **target variable** `is_skip` (what the model predicts)\n",
"* `extract_track_id()` → enables per-track aggregation and cold-start handling\n",
"* `add_time_features()` → temporal features `hour` and `dow` (when do users skip most?)\n",
"\n",
"In a production ML system, steps 2–
]
2026-10-18 18:38:51 +01:00
}
],
"metadata": {
"kernelspec": {
"display_name": ".venv",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.13.7"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
