{ "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", "execution_count": null, "id": "2c1aeb30", "metadata": {}, "outputs": [], "source": [ " ## Extract (Load in) ##\n", "def start_spark(app_name):\n", " \"\"\"Initialize and return a SparkSession with the PostgreSQL JDBC driver.\"\"\"\n", " 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", " \"\"\"Read a table from Postgres via JDBC into a Spark DataFrame.\"\"\"\n", " url, properties = get_jdbc_config()\n", " 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)" ] }, { "cell_type": "code", "execution_count": 15, "id": "9631ba07", "metadata": {}, "outputs": [], "source": [ " ## Transform ##" ] }, { "cell_type": "code", "execution_count": null, "id": "f7673c72", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "[Stage 0:> (0 + 1) / 1]\r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "+--------------------+--------------------+-----+\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" ] }, { "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", " bronze_table_name = \"bronze_spotify_data\"\n", " silver_table_name = \"silver_spotify_data\"\n", " bronze_df = read_jdbc(spark, bronze_table_name)\n", "\n", " bronze_df.groupBy(\"track_name\", \"username\").count().show()\n", "\n", " #spark.stop()\n", "if __name__ == \"__main__\":\n", " 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" ] }, { "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))" ] }, { "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!\")" ] }, { "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–3 would be packaged as a **feature pipeline** (e.g., via Databricks Feature Store or a scheduled job), the model in step 4 would be registered in **MLflow**, and step 5 would run as a **Model Serving endpoint** for real-time scoring." ] } ], "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 }