{ "cells": [ { "cell_type": "code", "execution_count": null, "metadata": { "id": "eBQvliwsLiZM" }, "outputs": [], "source": [ "# Copyright 2026 Google LLC\n", "#\n", "# Licensed under the Apache License, Version 2.0 (the \"License\");\n", "# you may not use this file except in compliance with the License.\n", "# You may obtain a copy of the License at\n", "#\n", "# https://www.apache.org/licenses/LICENSE-2.0\n", "#\n", "# Unless required by applicable law or agreed to in writing, software\n", "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", "# See the License for the specific language governing permissions and\n", "# limitations under the License." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Propensity Modeling & Churn Predictions" ] }, { "cell_type": "markdown", "metadata": { "id": "QCcazsyjJvXK" }, "source": [ "## Overview\n", "\n", "TheLook, a hypothetical eCommerce clothing retailer, stores data on customers, products, orders, logistics, web events, and digital marketing campaigns in BigQuery. The company wants to leverage the team's existing SQL and PySpark expertise to analyze this data using Apache Spark.\n", "\n", "To avoid manual infrastructure provisioning or tuning for Spark, TheLook seeks an auto-scaling solution that allows them to focus on workloads rather than cluster management. Additionally, they want to minimize the effort required to integrate Spark and BigQuery.\n", "\n", "In this use case, we will demonstrate how to build a logistic regression classification model using PySpark to predict whether a user will make a purchase. The entire workflow is executed within a notebook in their IDE (Antigravity/VSCode), taking advantage of the serverless Spark engine running remotely. This approach allows our data science team to use familiar PySpark libraries for data exploration and model training directly on data stored in BigQuery, creating a seamless experience from data to model within a single, integrated environment.\n", "\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Environment Setup (Antigravity / Data Cloud)\n", "\n", "If you are running this notebook in **Antigravity** (with the Google Data Cloud extension), the Spark session (`spark` variable) is automatically initialized once you select a remote kernel. \n", "\n", "**Note:** When creating the Dataproc Serverless runtime template for this notebook, ensure you set `spark.dynamicAllocation.enabled` to `false` in the Spark properties.\n", "\n", "If you are running **locally** or in another environment, you will need to manually initialize the `SparkSession` and install any missing dependencies." ] }, { "cell_type": "markdown", "metadata": { "id": "RtgdfcrNWNkU" }, "source": [ "## Setup\n", "\n", "The following steps create resources that will be used throughout the tutorial." ] }, { "cell_type": "markdown", "metadata": { "id": "YHIbYWxZWxk1" }, "source": [ "Configure a project id and location." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "SzD0y_2iBsoc", "tags": [ "parameters" ] }, "outputs": [], "source": [ "PROJECT_ID = \"\" # @param {type:\"string\"}\n", "LOCATION = \"\" # @param {type:\"string\"}\n", "\n", "if not PROJECT_ID: \n", " PROJECT_ID = input(\"Enter project id\")\n", "if not LOCATION: \n", " LOCATION = input(\"Enter location\")" ] }, { "cell_type": "markdown", "metadata": { "id": "ZJCWG4vrK8dq" }, "source": [ "## Load Data\n" ] }, { "cell_type": "markdown", "metadata": { "id": "MxS70d6xVAYN" }, "source": [ "Load each table into Spark and register them as SparkSQL tables." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "ClRTDEzzLB8v" }, "outputs": [], "source": [ "# Read thelook_ecommerce.users from BigQuery and create a temporary view\n", "users = spark.read.format(\"bigquery\").option(\"table\", \"bigquery-public-data.thelook_ecommerce.users\").load()\n", "users.createOrReplaceTempView(\"users\")\n", "\n", "# Read thelook_ecommerce.order_items from BigQuery\n", "order_items = spark.read.format(\"bigquery\").option(\"table\", \"bigquery-public-data.thelook_ecommerce.order_items\").load()\n", "order_items.createOrReplaceTempView(\"order_items\")" ] }, { "cell_type": "markdown", "metadata": { "id": "3orI5dVOMNdH" }, "source": [ "## Data Exploration\n", "\n", "Bigquery Studio can leverage Gemini for [advanced code completion capabilities](https://cloud.google.com/bigquery/docs/write-sql-gemini#generate_python_code?utm_campaign=CDR_0x225cfd13_default_b407565440&utm_source=external&utm_medium=web) which can use Natual Language to perform exploratory analysis using SQL and even generate PySpark Code for Feature Engineering.\n", "\n", "Try the following examples.\n", "\n", "* **Prompt 1**: Use Spark to explore the users table and show the first 10 rows.\n", "* **Prompt 2**: Use Spark to explore the order_items table and show the first 10 rows.\n", "* **Prompt 3**: Generate PySpark code to show the top 5 most frequent countries in the users table. Display the country and the number of users from each country.\n", "* **Prompt 4**: Generate PySpark code to find the average sale price of items in the order_items table.\n", "* **Prompt 5**: Using the table \"users\", generate code to plot country vs traffic source using a suitable plotting library.\n", "* **Prompt 6:** Create a histogram showing the distribution of \"age\", \"country_hash\", \"gender_hash\", \"traffic_source_hash\"" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Display the first 10 rows of the users table." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "U_B1PuIuW4J9" }, "outputs": [], "source": [ "# prompt: Use Spark to explore the users table and show the first 10 rows.\n", "\n", "users.show(10)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Display the first 10 rows of the order_items table." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "ekdFabYZXbnx" }, "outputs": [], "source": [ "# prompt: Use Spark to explore the order_items table and show the first 10 rows.\n", "\n", "order_items.show(10)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Show the top 5 most frequent countries and their user counts." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "Oy9y1tbFLsba" }, "outputs": [], "source": [ "# prompt: Generate PySpark code to show the top 5 most frequent countries in the users table. Display the country and the number of users from each country.\n", "\n", "from pyspark.sql.functions import col, count\n", "\n", "users.groupBy(\"country\").agg(count(\"*\").alias(\"user_count\")).orderBy(col(\"user_count\").desc()).limit(5).show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Calculate the average sale price of items." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "wBeMtOJ3Lzn8" }, "outputs": [], "source": [ "# prompt: Generate code to find the average sale price of items in the order_items table.\n", "\n", "order_items.agg({\"sale_price\": \"avg\"}).show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Visualize the distribution of key user attributes." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "bHOXovKiL8sS" }, "outputs": [], "source": [ "# prompt: Create a histogram showing the distribution of \"age\", \"country_hash\", \"gender_hash\", \"traffic_source_hash\"\n", "\n", "import matplotlib.pyplot as plt\n", "\n", "# Convert Spark DataFrame to Pandas DataFrame for plotting\n", "users_pd = users.toPandas()\n", "\n", "# Create histograms for the specified columns\n", "fig, axes = plt.subplots(2, 2, figsize=(12, 10))\n", "fig.suptitle('Distribution of Age, Country, Gender, and Traffic Source')\n", "\n", "# Age histogram\n", "axes[0, 0].hist(users_pd['age'].dropna(), bins=20, edgecolor='black')\n", "axes[0, 0].set_title('Age Distribution')\n", "axes[0, 0].set_xlabel('Age')\n", "axes[0, 0].set_ylabel('Frequency')\n", "\n", "# Country histogram (using value counts for categorical data)\n", "users_pd['country'].value_counts().plot(kind='bar', ax=axes[0, 1])\n", "axes[0, 1].set_title('Country Distribution')\n", "axes[0, 1].set_xlabel('Country')\n", "axes[0, 1].set_ylabel('Frequency')\n", "axes[0, 1].tick_params(axis='x', rotation=45)\n", "\n", "# Gender histogram (using value counts for categorical data)\n", "users_pd['gender'].value_counts().plot(kind='bar', ax=axes[1, 0])\n", "axes[1, 0].set_title('Gender Distribution')\n", "axes[1, 0].set_xlabel('Gender')\n", "axes[1, 0].set_ylabel('Frequency')\n", "axes[1, 0].tick_params(axis='x', rotation=45)\n", "\n", "# Traffic Source histogram (using value counts for categorical data)\n", "users_pd['traffic_source'].value_counts().plot(kind='bar', ax=axes[1, 1])\n", "axes[1, 1].set_title('Traffic Source Distribution')\n", "axes[1, 1].set_xlabel('Traffic Source')\n", "axes[1, 1].set_ylabel('Frequency')\n", "axes[1, 1].tick_params(axis='x', rotation=45)\n", "\n", "plt.tight_layout(rect=[0, 0.03, 1, 0.95])\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": { "id": "pvg1I2Y5QC7g" }, "source": [ "## Feature Engineering\n", "\n", "In this step, we derive two key columns from the input data:\n", "\n", "**Creation of features column**:\n", "It combines user attributes (age, hashed categorical features) into a numerical array, preparing them for a machine learning model.\n", "\n", "**Generation of label column:**\n", "It creates a binary target variable indicating whether a user has made a purchase or not, derived from order information." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "Le1JbxZPpEH4" }, "outputs": [], "source": [ "# Load BigQuery dataset with feature engineering in SQL\n", "features = spark.sql(\"\"\"\n", "SELECT\n", " CAST(u.age AS DOUBLE),\n", " CAST(hash(u.country) AS BIGINT) * 1.0 AS country_hash,\n", " CAST(hash(u.gender) AS BIGINT) * 1.0 AS gender_hash,\n", " CAST(hash(u.traffic_source) AS BIGINT) * 1.0 AS traffic_source_hash,\n", " CASE WHEN COALESCE(SUM(oi.sale_price), 0) > 0 THEN 1 ELSE 0 END AS label\n", "FROM users AS u\n", "LEFT JOIN order_items AS oi\n", "ON u.id = oi.user_id\n", "GROUP BY u.id, u.age, u.country, u.gender, u.traffic_source;\n", "\"\"\")\n", "features.show()" ] }, { "cell_type": "markdown", "metadata": { "id": "-Wj-0bASnK8V" }, "source": [ "## Perform ML Task\n", "\n", "This code trains a logistic regression model to predict user purchase behavior, with these steps:\n", "\n", "**Feature Scaling:** StandardScaler scales the \"features\" column.\n", "\n", "**Model Initialization:** LogisticRegression is set up to predict the binary \"label\" (purchase/no purchase), with hyperparameters for training.\n", "\n", "**Pipeline Definition:** A Pipeline chains StandardScaler and LogisticRegression for streamlined scaling and training.\n", "\n", "**Model Training:** `pipeline.fit(dataset)` trains the pipeline (scaling and then the model).\n", "\n", "**Prediction:** `pipeline_model.transform(dataset)` generates predictions, and `transformed_dataset.show()` displays the results.\n", "\n", "In short, this step scales features, trains a logistic regression model within a pipeline, and produces purchase predictions." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "_4UXH0rLm6fd" }, "outputs": [], "source": [ "from pyspark.ml.classification import LogisticRegression, LogisticRegressionModel\n", "from pyspark.ml.evaluation import BinaryClassificationEvaluator\n", "from pyspark.ml.feature import StandardScaler, VectorAssembler\n", "from pyspark.ml.pipeline import Pipeline\n", "\n", "# Split Train and Test Data (90:10)\n", "train_data, test_data = features.randomSplit([0.9, 0.1], seed=42)\n", "\n", "# Define VectorAssembler to combine individual columns into a 'features' vector\n", "assembler = VectorAssembler(\n", " inputCols=[\"age\", \"country_hash\", \"gender_hash\", \"traffic_source_hash\"], \n", " outputCol=\"features\"\n", ")\n", "\n", "# Initialize StandardScaler\n", "scaler = StandardScaler(inputCol=\"features\", outputCol=\"scaled_features\")\n", "\n", "# Initialize Logistic Regression model\n", "lr = LogisticRegression(maxIter=30, featuresCol=\"scaled_features\", labelCol=\"label\")\n", "\n", "# Define pipeline\n", "pipeline = Pipeline(stages=[assembler, scaler, lr])" ] }, { "cell_type": "markdown", "metadata": { "id": "5dzndxJYOPsk" }, "source": [ "Train the model." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "G2T4T64rm_97" }, "outputs": [], "source": [ "# Fit the model\n", "pipeline_model = pipeline.fit(train_data)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Generate and display predictions on the test dataset." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "tV2JJ-t1ONES" }, "outputs": [], "source": [ "# Transform the dataset using the trained model\n", "transformed_dataset = pipeline_model.transform(test_data)\n", "\n", "# Print the new data\n", "transformed_dataset.show()" ] }, { "cell_type": "markdown", "metadata": { "id": "7z9R-zu7nWtO" }, "source": [ "## Evaluation\n", "\n", "This code evaluates the trained model's performance by:\n", "\n", "**Initializing an Evaluator:** A BinaryClassificationEvaluator is set up to calculate the Area Under the Precision-Recall Curve (AUC-PR).\n", "\n", "**Calculating AUC-PR:** The evaluate() method calculates the AUC-PR score using the model's predictions.\n", "\n", "This step quantifies the model's ability to distinguish between the two classes (e.g., purchase/no purchase).\n", "\n", "\n", "Further we will use NLP2SQL code generation to visualize the output\n", "\n", "**Prompt 1:** Generate code to plot the Precision-Recall (PR) curve. Calculate precision and recall from the model's predictions and display the PR curve using a suitable plotting library.\n", "\n", "**Prompt 2:** Generate code to create a confusion matrix visualization. Calculate the confusion matrix from the model's predictions and display it as a heatmap or a table with counts of true positives (TP), true negatives (TN), false positives (FP), and false negatives (FN)." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "1PIlkb-Unatt" }, "outputs": [], "source": [ "# Model evaluation\n", "eva = BinaryClassificationEvaluator(metricName=\"areaUnderPR\")\n", "aucPR = eva.evaluate(transformed_dataset)\n", "print(f\"AUC PR: {aucPR}\")" ] }, { "cell_type": "markdown", "metadata": { "id": "m5wsbhsAncHV" }, "source": [ "## Visualization\n", "\n", "Let's visualize the results to see how our model performs, and how it has predicted.\n", "\n", "**Prompt 1:** Generate code to plot the Precision-Recall (PR) curve using a suitable plotting library.\n", "\n", "\n", "**Prompt 2:** Generate code to create a confusion matrix visualization. Calculate the confusion matrix from the model's predictions and display it as a heat map or a table with counts of true positives (TP), true negatives (TN), false positives (FP), and false negatives (FN).\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Plot the Precision-Recall (PR) Curve." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "juXEzTJIRWDB" }, "outputs": [], "source": [ "# prompt: Generate code to plot the Precision-Recall (PR) curve using a suitable plotting library.\n", "\n", "from sklearn.metrics import precision_recall_curve, auc\n", "import matplotlib.pyplot as plt\n", "\n", "# Get predictions and labels from the transformed dataset\n", "predictions_and_labels = transformed_dataset.select(\"prediction\", \"label\").collect()\n", "y_scores = [row.prediction for row in predictions_and_labels]\n", "y_true = [row.label for row in predictions_and_labels]\n", "\n", "# Calculate precision and recall\n", "precision, recall, _ = precision_recall_curve(y_true, y_scores)\n", "pr_auc = auc(recall, precision)\n", "\n", "# Plot the PR curve\n", "plt.figure(figsize=(8, 6))\n", "plt.plot(recall, precision, color='blue', lw=2, label=f'PR curve (AUC = {pr_auc:.2f})')\n", "plt.xlabel('Recall')\n", "plt.ylabel('Precision')\n", "plt.title('Precision-Recall Curve')\n", "plt.legend(loc='lower left')\n", "plt.grid(True)\n", "plt.show()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Visualize the Confusion Matrix." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "SOvDhaEkSI2e" }, "outputs": [], "source": [ "# prompt: Generate code to create a confusion matrix visualization. Calculate the confusion matrix from the model's predictions and display it as a heatmap or a table with counts of true positives (TP), true negatives (TN), false positives (FP), and false negatives (FN).\n", "\n", "from sklearn.metrics import confusion_matrix\n", "import seaborn as sns\n", "import matplotlib.pyplot as plt\n", "\n", "# Get predictions and labels from the transformed dataset\n", "predictions_and_labels = transformed_dataset.select(\"prediction\", \"label\").collect()\n", "y_pred = [row.prediction for row in predictions_and_labels]\n", "y_true = [row.label for row in predictions_and_labels]\n", "\n", "# Calculate the confusion matrix\n", "cm = confusion_matrix(y_true, y_pred)\n", "\n", "# Display the confusion matrix as a heatmap\n", "plt.figure(figsize=(8, 6))\n", "sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', cbar=False,\n", " xticklabels=['Predicted Negative', 'Predicted Positive'],\n", " yticklabels=['Actual Negative', 'Actual Positive'])\n", "plt.xlabel('Predicted Label')\n", "plt.ylabel('True Label')\n", "plt.title('Confusion Matrix')\n", "plt.show()\n", "\n", "# Display the confusion matrix as a table\n", "tn, fp, fn, tp = cm.ravel()\n", "print(f\"True Negatives (TN): {tn}\")\n", "print(f\"False Positives (FP): {fp}\")\n", "print(f\"False Negatives (FN): {fn}\")\n", "print(f\"True Positives (TP): {tp}\")" ] }, { "cell_type": "markdown", "metadata": { "id": "ncZ55Z0sxF7Q" }, "source": [ "## Write Predictions to BigQuery\n", "\n", "Use Gemini to write predictions to BigQuery.\n", "\n", "**Prompt:** Use Python to create a new BigQuery dataset called predictions and then use Spark to write the new prediction data to a table in this dataset." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Create BigQuery Dataset and Write Predictions." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "IWd29JJ8Sgcj" }, "outputs": [], "source": [ "# prompt: Use Python to create a new BigQuery dataset called predictions, and then use Spark to write the new prediction data to a table in this dataset.\n", "\n", "from google.cloud import bigquery\n", "\n", "# Construct a BigQuery client object.\n", "client = bigquery.Client(project=PROJECT_ID)\n", "\n", "# Set dataset ID\n", "dataset_id = f\"{PROJECT_ID}.predictions\"\n", "\n", "# Construct a full Dataset object to send to the API.\n", "dataset = bigquery.Dataset(dataset_id)\n", "\n", "# Specify the geographic location where the dataset should reside.\n", "dataset.location = LOCATION\n", "\n", "# Send the dataset to the API for creation, with an explicit timeout.\n", "# Raises google.api_core.exceptions.Conflict if the dataset already exists.\n", "try:\n", " dataset = client.create_dataset(dataset, timeout=30) # Make an API request.\n", " print(f\"Created dataset {client.project}.{dataset.dataset_id}\")\n", "except Exception as e:\n", " print(f\"Dataset {client.project}.{dataset.dataset_id} already exists or an error occurred: {e}\")\n", "\n", "# Write the transformed_dataset to a new BigQuery table\n", "transformed_dataset.write \\\n", " .format(\"bigquery\") \\\n", " .option(\"table\", f\"{dataset_id}.user_purchase_predictions\") \\\n", " .option(\"writeMethod\", \"direct\") \\\n", " .mode(\"overwrite\") \\\n", " .save()\n", "\n", "print(f\"Predictions written to BigQuery table: {dataset_id}.user_purchase_predictions\")" ] } ], "metadata": { "colab": { "cell_execution_strategy": "setup", "name": "Purchase_Predictions_Spark (2)", "provenance": [] }, "environment": { "kernel": "9c39b79e5d2e7072beb4bd59-runtime-iceberg", "name": "workbench-notebooks.m139", "type": "gcloud", "uri": "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/workbench-notebooks:m139" }, "kernelspec": { "display_name": "runtime-iceberg on Serverless Spark", "language": "python", "name": "9c39b79e5d2e7072beb4bd59-runtime-iceberg" }, "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.11.12" } }, "nbformat": 4, "nbformat_minor": 0 }