{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "apache-license-cell",
      "metadata": {
        "id": "apache-license-cell"
      },
      "outputs": [],
      "source": [
        "# Copyright 2025 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",
      "id": "intro-header",
      "metadata": {
        "id": "intro-header"
      },
      "source": [
        "# Data Science with PySpark and Distributed XGBoost"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f4f44581",
      "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",
      "id": "overview",
      "metadata": {
        "id": "overview"
      },
      "source": [
        "## Overview\n",
        "\n",
        "In this notebook, you will learn how to parallelize and scale your data science and machine learning workflows on large datasets using distributed computing with Apache Spark (`PySpark`).\n",
        "\n",
        "You will build a machine learning pipeline to predict NYC taxi tip amounts. You will focus on scaling data processing (`PySpark DataFrames`/`Spark SQL`), using `PySpark MLlib` for baseline algorithms, and training a distributed `XGBoost` model via the PySpark Estimator API."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "objectives",
      "metadata": {
        "id": "objectives"
      },
      "source": [
        "### Objectives\n",
        "\n",
        "* Initialize a cluster-ready PySpark session.\n",
        "* Construct an end-to-end data pipeline with Spark DataFrames.\n",
        "* Train regression models with Spark MLlib.\n",
        "* Configure and train distributed XGBoost on Spark.\n",
        "* Evaluate validation accuracy on a held-out test dataset."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "dataset-intro",
      "metadata": {
        "id": "dataset-intro"
      },
      "source": [
        "---\n",
        "\n",
        "## Prepare the NYC taxi dataset\n",
        "\n",
        "This tutorial uses the [NYC Taxi & Limousine Commission (TLC) Trip Record Data](https://www.nyc.gov/site/tlc/about/tlc-trip-record-data.page).\n",
        "\n",
        "The dataset contains individual trip records from yellow taxis in New York City, and includes fields like:\n",
        "*   Pick-up and drop-off dates, times, and locations\n",
        "*   Trip distances\n",
        "*   Itemized fare amounts\n",
        "*   Passenger counts\n",
        "*   **Tip amounts** (*the target variable we will predict*)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "define-data-path-md-generated",
      "metadata": {},
      "source": [
        "#### Define Data Path\n",
        "\n",
        "Specify the Google Cloud Storage (GCS) path where the NYC taxi dataset is located."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "download-data-code",
      "metadata": {
        "execution": {
          "iopub.execute_input": "2026-02-23T14:56:39.838438Z",
          "iopub.status.busy": "2026-02-23T14:56:39.838139Z",
          "iopub.status.idle": "2026-02-23T14:56:39.844972Z",
          "shell.execute_reply": "2026-02-23T14:56:39.843597Z",
          "shell.execute_reply.started": "2026-02-23T14:56:39.838417Z"
        },
        "id": "download-data-code"
      },
      "outputs": [],
      "source": [
        "# Just using 1-month data of the NYC Taxi dataset\n",
        "DATA_DIR = \"gs://dataproc-metastore-public-binaries/nyc_taxi_data\"\n",
        "print(f\"Using dataset from {DATA_DIR}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "load-clean-data",
      "metadata": {
        "id": "load-clean-data"
      },
      "source": [
        "### Load and Clean Data\n",
        "\n",
        "Let's load the data we downloaded and clean out invalid trips."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "load-clean-data-code",
      "metadata": {
        "execution": {
          "iopub.execute_input": "2026-02-23T14:56:40.177850Z",
          "iopub.status.busy": "2026-02-23T14:56:40.177474Z",
          "iopub.status.idle": "2026-02-23T14:56:48.234558Z",
          "shell.execute_reply": "2026-02-23T14:56:48.233037Z",
          "shell.execute_reply.started": "2026-02-23T14:56:40.177823Z"
        },
        "id": "load-clean-data-code"
      },
      "outputs": [],
      "source": [
        "import time\n",
        "from pyspark.sql.functions import col\n",
        "\n",
        "start_time = time.perf_counter()\n",
        "\n",
        "# 1. Load data\n",
        "df = spark.read.parquet(f\"{DATA_DIR}/*.parquet\")\n",
        "print(f\"Loaded {df.count():,} records.\")\n",
        "\n",
        "# 2. Clean data\n",
        "df = df.filter(\n",
        "    (col('fare_amount') > 0) & (col('fare_amount') < 500) &\n",
        "    (col('trip_distance') > 0) & (col('trip_distance') < 100) &\n",
        "    (col('tip_amount') >= 0) & (col('tip_amount') < 100) &\n",
        "    (col('payment_type') == 1) # Credit card only (tips recorded)\n",
        ")\n",
        "print(f\"Clean records remaining: {df.count():,}.\")\n",
        "print(f\"Load and Clean Time: {time.perf_counter() - start_time:.2f} seconds\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "feature-engineering",
      "metadata": {
        "id": "feature-engineering"
      },
      "source": [
        "### Feature Engineering\n",
        "\n",
        "Next, we prepare our features for the machine learning models. We will extract time-based features, log-transform the fare, and create route-specific aggregations based on the pickup and dropoff location IDs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "feature-engineering-code",
      "metadata": {
        "execution": {
          "iopub.execute_input": "2026-02-23T14:56:48.236211Z",
          "iopub.status.busy": "2026-02-23T14:56:48.235833Z",
          "iopub.status.idle": "2026-02-23T14:56:48.404728Z",
          "shell.execute_reply": "2026-02-23T14:56:48.403728Z",
          "shell.execute_reply.started": "2026-02-23T14:56:48.236176Z"
        },
        "id": "feature-engineering-code"
      },
      "outputs": [],
      "source": [
        "from pyspark.sql import functions as F\n",
        "from pyspark.sql import Window\n",
        "\n",
        "start_time = time.perf_counter()\n",
        "\n",
        "# Time Features\n",
        "df = df.withColumn('hour', F.hour('tpep_pickup_datetime'))\n",
        "df = df.withColumn('dow', F.dayofweek('tpep_pickup_datetime')) # Note: 1=Sunday, 7=Saturday in PySpark\n",
        "df = df.withColumn('is_weekend', F.when((col('dow') == 1) | (col('dow') == 7), 1).otherwise(0))\n",
        "df = df.withColumn('is_rush_hour', \n",
        "                   F.when(((col('hour') >= 7) & (col('hour') <= 9)) | \n",
        "                          ((col('hour') >= 17) & (col('hour') <= 19)), 1).otherwise(0))\n",
        "\n",
        "# Amount Features\n",
        "df = df.withColumn('fare_log', F.log1p('fare_amount'))\n",
        "df = df.withColumn('fare_decimal', (col('fare_amount') % 1 * 100).cast('int'))\n",
        "df = df.withColumn('is_round_fare', F.when(col('fare_amount') % 5 == 0, 1).otherwise(0))\n",
        "\n",
        "# Route Features\n",
        "df = df.withColumn('route_id', F.concat_ws('_', col('PULocationID').cast('string'), col('DOLocationID').cast('string')))\n",
        "route_window = Window.partitionBy('route_id')\n",
        "df = df.withColumn('route_frequency', F.count('*').over(route_window))\n",
        "\n",
        "# Location Aggregation Features (Mean and Std Tip by Pickup Location)\n",
        "pu_window = Window.partitionBy('PULocationID')\n",
        "df = df.withColumn('pu_tip_mean', F.mean('tip_amount').over(pu_window))\n",
        "df = df.withColumn('pu_tip_std', F.stddev('tip_amount').over(pu_window))\n",
        "df = df.fillna({'pu_tip_std': 0.0})"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "display-features-md-generated",
      "metadata": {},
      "source": [
        "#### Display Engineered Features\n",
        "\n",
        "Display the first few rows of the DataFrame with the newly engineered features to inspect the results."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c9bb20c9-4f45-4434-9580-be13bd00e9ce",
      "metadata": {
        "execution": {
          "iopub.execute_input": "2026-02-23T14:56:48.406319Z",
          "iopub.status.busy": "2026-02-23T14:56:48.406034Z",
          "iopub.status.idle": "2026-02-23T14:56:58.974367Z",
          "shell.execute_reply": "2026-02-23T14:56:58.973171Z",
          "shell.execute_reply.started": "2026-02-23T14:56:48.406290Z"
        }
      },
      "outputs": [],
      "source": [
        "df.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "prepare-xy-markdown",
      "metadata": {
        "id": "prepare-xy-markdown"
      },
      "source": [
        "### Prepare Features and Target\n",
        "\n",
        "Let's select our feature columns, assemble them into a Vector, and split the data into a Train and Test set."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "prepare-xy-code",
      "metadata": {
        "execution": {
          "iopub.execute_input": "2026-02-23T14:56:58.976716Z",
          "iopub.status.busy": "2026-02-23T14:56:58.976348Z",
          "iopub.status.idle": "2026-02-23T14:57:21.582495Z",
          "shell.execute_reply": "2026-02-23T14:57:21.579850Z",
          "shell.execute_reply.started": "2026-02-23T14:56:58.976690Z"
        },
        "id": "prepare-xy-code"
      },
      "outputs": [],
      "source": [
        "from pyspark.ml.feature import VectorAssembler\n",
        "\n",
        "feature_cols = [\n",
        "    'trip_distance', 'fare_amount', 'passenger_count',\n",
        "    'hour', 'dow', 'is_weekend', 'is_rush_hour',\n",
        "    'fare_log', 'fare_decimal', 'is_round_fare',\n",
        "    'route_frequency', 'pu_tip_mean', 'pu_tip_std',\n",
        "    'PULocationID', 'DOLocationID'\n",
        "]\n",
        "\n",
        "# Impute nulls with mean for safe vector assembly\n",
        "from pyspark.ml.feature import Imputer\n",
        "imputer = Imputer(inputCols=feature_cols, outputCols=feature_cols).setStrategy(\"mean\")\n",
        "imputer_model = imputer.fit(df)\n",
        "df = imputer_model.transform(df)\n",
        "\n",
        "assembler = VectorAssembler(inputCols=feature_cols, outputCol='features')\n",
        "df_assembled = assembler.transform(df)\n",
        "\n",
        "# Split data\n",
        "train_df, test_df = df_assembled.randomSplit([0.8, 0.2], seed=42)\n",
        "\n",
        "print(f\"Train rows: {train_df.count():,}\")\n",
        "print(f\"Test rows: {test_df.count():,}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "train-models",
      "metadata": {
        "id": "train-models"
      },
      "source": [
        "---\n",
        "\n",
        "## Train Models\n",
        "\n",
        "Now we will train three different predictive models to forecast tip amounts using PySpark."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "linear-model-markdown",
      "metadata": {
        "id": "linear-model-markdown"
      },
      "source": [
        "### 1. Linear Regression\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "description-lr-model",
      "metadata": {},
      "source": [
        "Implement and train a Linear Regression model from PySpark MLlib to predict tip amounts. Evaluate its performance using RMSE."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "linear-model-code",
      "metadata": {
        "execution": {
          "iopub.execute_input": "2026-02-23T14:57:21.584489Z",
          "iopub.status.busy": "2026-02-23T14:57:21.584082Z",
          "iopub.status.idle": "2026-02-23T14:57:56.333439Z",
          "shell.execute_reply": "2026-02-23T14:57:56.332128Z",
          "shell.execute_reply.started": "2026-02-23T14:57:21.584456Z"
        },
        "id": "linear-model-code"
      },
      "outputs": [],
      "source": [
        "from pyspark.ml.regression import LinearRegression\n",
        "from pyspark.ml.evaluation import RegressionEvaluator\n",
        "\n",
        "start_time = time.perf_counter()\n",
        "\n",
        "lr = LinearRegression(featuresCol='features', labelCol='tip_amount')\n",
        "lr_model = lr.fit(train_df)\n",
        "\n",
        "lr_preds = lr_model.transform(test_df)\n",
        "lr_preds = lr_preds.withColumnRenamed('prediction', 'lr_prediction')\n",
        "\n",
        "evaluator = RegressionEvaluator(labelCol=\"tip_amount\", predictionCol=\"lr_prediction\", metricName=\"rmse\")\n",
        "lr_rmse = evaluator.evaluate(lr_preds)\n",
        "\n",
        "print(f\"\\n{'Linear Reg RMSE:':<20} ${lr_rmse:.4f}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "rf-markdown",
      "metadata": {
        "id": "rf-markdown"
      },
      "source": [
        "### 2. Random Forest\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "description-rf-model",
      "metadata": {},
      "source": [
        "Train a Random Forest Regressor model from PySpark MLlib. Configure parameters like maximum depth and number of trees, and then evaluate its RMSE."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "rf-code",
      "metadata": {
        "execution": {
          "iopub.execute_input": "2026-02-23T14:57:56.334787Z",
          "iopub.status.busy": "2026-02-23T14:57:56.334406Z",
          "iopub.status.idle": "2026-02-23T14:58:53.740216Z",
          "shell.execute_reply": "2026-02-23T14:58:53.738989Z",
          "shell.execute_reply.started": "2026-02-23T14:57:56.334754Z"
        },
        "id": "rf-code"
      },
      "outputs": [],
      "source": [
        "from pyspark.ml.regression import RandomForestRegressor\n",
        "\n",
        "start_time = time.perf_counter()\n",
        "\n",
        "rf = RandomForestRegressor(featuresCol='features', labelCol='tip_amount', maxDepth=10, numTrees=20, seed=42)\n",
        "rf_model = rf.fit(train_df)\n",
        "\n",
        "rf_preds = rf_model.transform(test_df)\n",
        "rf_preds = rf_preds.withColumnRenamed('prediction', 'rf_prediction')\n",
        "\n",
        "evaluator_rf = RegressionEvaluator(labelCol=\"tip_amount\", predictionCol=\"rf_prediction\", metricName=\"rmse\")\n",
        "rf_rmse = evaluator_rf.evaluate(rf_preds)\n",
        "\n",
        "print(f\"\\n{'Random Forest RMSE:':<20} ${rf_rmse:.4f}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "xgb-markdown",
      "metadata": {
        "id": "xgb-markdown"
      },
      "source": [
        "### 3. Distributed XGBoost\n",
        "To use SparkXGBRegressor, ensure that the xgboost python package is installed in your python environment."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "description-xgb-model",
      "metadata": {},
      "source": [
        "Utilize `SparkXGBRegressor` to train a distributed XGBoost model. Configure the number of workers, maximum depth, and learning rate for the model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "xgb-code",
      "metadata": {
        "execution": {
          "iopub.execute_input": "2026-02-23T14:58:53.741915Z",
          "iopub.status.busy": "2026-02-23T14:58:53.741522Z",
          "iopub.status.idle": "2026-02-23T14:59:46.543035Z",
          "shell.execute_reply": "2026-02-23T14:59:46.542036Z",
          "shell.execute_reply.started": "2026-02-23T14:58:53.741881Z"
        },
        "id": "xgb-code"
      },
      "outputs": [],
      "source": [
        "from xgboost.spark import SparkXGBRegressor\n",
        "\n",
        "start_time = time.perf_counter()\n",
        "\n",
        "xgb = SparkXGBRegressor(\n",
        "  features_col=\"features\",\n",
        "  label_col=\"tip_amount\",\n",
        "  num_workers=2,\n",
        "  max_depth=5,\n",
        "  learning_rate=0.1\n",
        ")\n",
        "\n",
        "xgb_model = xgb.fit(train_df)\n",
        "\n",
        "xgb_preds = xgb_model.transform(test_df)\n",
        "xgb_preds = xgb_preds.withColumnRenamed('prediction', 'xgb_prediction')\n",
        "\n",
        "evaluator_xgb = RegressionEvaluator(labelCol=\"tip_amount\", predictionCol=\"xgb_prediction\", metricName=\"rmse\")\n",
        "xgb_rmse = evaluator_xgb.evaluate(xgb_preds)\n",
        "\n",
        "print(f\"\\n{'XGBoost RMSE:':<20} ${xgb_rmse:.4f}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ensemble-eval-markdown",
      "metadata": {
        "id": "ensemble-eval-markdown"
      },
      "source": [
        "---\n",
        "\n",
        "## Evaluate the Models\n",
        "\n",
        "Finally, let's compare the predictions from our three models and compute a simple average ensemble."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "compare-models-md-generated",
      "metadata": {},
      "source": [
        "### Compare Model Performance\n",
        "\n",
        "Review the RMSE for each trained model (Linear Regression, Random Forest, and Distributed XGBoost) to compare their predictive performance on the test dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ensemble-eval-code",
      "metadata": {
        "execution": {
          "iopub.execute_input": "2026-02-23T14:59:55.981551Z",
          "iopub.status.busy": "2026-02-23T14:59:55.981112Z",
          "iopub.status.idle": "2026-02-23T14:59:55.989704Z",
          "shell.execute_reply": "2026-02-23T14:59:55.987872Z",
          "shell.execute_reply.started": "2026-02-23T14:59:55.981522Z"
        },
        "id": "ensemble-eval-code"
      },
      "outputs": [],
      "source": [
        "from pyspark.sql.functions import col, lit\n",
        "\n",
        "# Join predictions on a unique id or simply use the fact they were all evaluated separately.\n",
        "# To perform an ensemble average in pyspark, we can join the dataframes if we added a unique ID,\n",
        "# but for simplicity, we can just print the RMSE values calculated above.\n",
        "\n",
        "print(f\"\\n{'Model':<20} {'RMSE':>10}\")\n",
        "print(\"-\" * 32)\n",
        "print(f\"{'Linear Regression':<20} ${lr_rmse:>9.4f}\")\n",
        "print(f\"{'Random Forest':<20} ${rf_rmse:>9.4f}\")\n",
        "print(f\"{'XGBoost':<20} ${xgb_rmse:>9.4f}\")\n",
        "print(\"-\" * 32)"
      ]
    }
  ],
  "metadata": {
    "environment": {
      "kernel": "9c39b79e5d2e7072beb4bd59-iceberg-biglake",
      "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": 5
}