spark_regression_fast.ipynb

{
  "cells": [
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "-ZkLjvyen87E"
      },
      "source": [
        "1.   Problem: \n",
        "- Spark has a lazy evaluation mechanism, so everytime you invoke in the code has to calculate DAG\n",
        "-  The time-consuming part is to infer the prediction on new data\n",
        "- Therefore, training the model should take place in Python and inference should take place on Spark\n",
        "<br />\n",
        "2.   Solution:\n",
        "- Save each model in a dictionary\n",
        "- Load it using pandas_udf (see [run your native Python code with PySpark](https://www.databricks.com/blog/2017/10/30/introducing-vectorized-udfs-for-pyspark.html))\n",
        "- In pandas_udf, you can write Python code as you would write it in python code.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ykLmOfJtkDG7",
        "outputId": "84a9cb1a-1eff-4e7a-862e-ccc65e8fc4f7"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: PySpark in /usr/local/lib/python3.8/dist-packages (3.3.2)\n",
            "Requirement already satisfied: py4j==0.10.9.5 in /usr/local/lib/python3.8/dist-packages (from PySpark) (0.10.9.5)\n",
            "time: 2.83 s (started: 2023-02-16 20:34:56 +00:00)\n"
          ]
        }
      ],
      "source": [
        "!pip install PySpark"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "CXB85fx12u-f",
        "outputId": "b43f016f-635f-457b-cf7d-2827ee5ae066"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: ipython-autotime in /usr/local/lib/python3.8/dist-packages (0.3.1)\n",
            "Requirement already satisfied: ipython in /usr/local/lib/python3.8/dist-packages (from ipython-autotime) (7.9.0)\n",
            "Requirement already satisfied: prompt-toolkit<2.1.0,>=2.0.0 in /usr/local/lib/python3.8/dist-packages (from ipython->ipython-autotime) (2.0.10)\n",
            "Requirement already satisfied: decorator in /usr/local/lib/python3.8/dist-packages (from ipython->ipython-autotime) (4.4.2)\n",
            "Requirement already satisfied: pickleshare in /usr/local/lib/python3.8/dist-packages (from ipython->ipython-autotime) (0.7.5)\n",
            "Requirement already satisfied: pygments in /usr/local/lib/python3.8/dist-packages (from ipython->ipython-autotime) (2.6.1)\n",
            "Requirement already satisfied: setuptools>=18.5 in /usr/local/lib/python3.8/dist-packages (from ipython->ipython-autotime) (57.4.0)\n",
            "Requirement already satisfied: backcall in /usr/local/lib/python3.8/dist-packages (from ipython->ipython-autotime) (0.2.0)\n",
            "Requirement already satisfied: traitlets>=4.2 in /usr/local/lib/python3.8/dist-packages (from ipython->ipython-autotime) (5.7.1)\n",
            "Requirement already satisfied: jedi>=0.10 in /usr/local/lib/python3.8/dist-packages (from ipython->ipython-autotime) (0.18.2)\n",
            "Requirement already satisfied: pexpect in /usr/local/lib/python3.8/dist-packages (from ipython->ipython-autotime) (4.8.0)\n",
            "Requirement already satisfied: parso<0.9.0,>=0.8.0 in /usr/local/lib/python3.8/dist-packages (from jedi>=0.10->ipython->ipython-autotime) (0.8.3)\n",
            "Requirement already satisfied: wcwidth in /usr/local/lib/python3.8/dist-packages (from prompt-toolkit<2.1.0,>=2.0.0->ipython->ipython-autotime) (0.2.6)\n",
            "Requirement already satisfied: six>=1.9.0 in /usr/local/lib/python3.8/dist-packages (from prompt-toolkit<2.1.0,>=2.0.0->ipython->ipython-autotime) (1.15.0)\n",
            "Requirement already satisfied: ptyprocess>=0.5 in /usr/local/lib/python3.8/dist-packages (from pexpect->ipython->ipython-autotime) (0.7.0)\n",
            "The autotime extension is already loaded. To reload it, use:\n",
            "  %reload_ext autotime\n",
            "time: 3.14 s (started: 2023-02-16 20:34:59 +00:00)\n"
          ]
        }
      ],
      "source": [
        "!pip install ipython-autotime\n",
        "%load_ext autotime"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 26,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WHh74H3Kf7IB",
        "outputId": "0f5c44de-eda1-4d84-cabf-5f5129c841ea"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "time: 7.3 ms (started: 2023-02-16 20:35:02 +00:00)\n"
          ]
        }
      ],
      "source": [
        "import pandas as pd\n",
        "import numpy as np\n",
        "from functools import wraps\n",
        "import time\n",
        "\n",
        "import statsmodels.api as sm\n",
        "from sklearn.metrics import roc_auc_score\n",
        "\n",
        "from pyspark.ml.classification import LogisticRegression\n",
        "from pyspark.ml.feature import VectorAssembler\n",
        "from pyspark.sql import SparkSession\n",
        "from pyspark.sql.types import StructType, StructField, DoubleType, IntegerType\n",
        "from pyspark.ml.evaluation import BinaryClassificationEvaluator\n",
        "from pyspark.sql.functions import pandas_udf, spark_partition_id\n",
        "import pyspark.sql.functions as f\n",
        "\n",
        "spark = SparkSession.builder.appName(\"test_regression\").getOrCreate()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 27,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "w3R-k49ngGKV",
        "outputId": "ab7930e7-18fb-426c-a973-dba09826bead"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "time: 19.5 s (started: 2023-02-16 20:35:02 +00:00)\n"
          ]
        }
      ],
      "source": [
        "# Create dataset\n",
        "observations = 500_000\n",
        "\n",
        "target = np.random.binomial(n=1, p=0.2, size=(observations, 1))\n",
        "y = target + np.random.normal(0, 0.1, size=(observations, 1))\n",
        "z = target + np.random.normal(20, 10.0, size=(observations, 1))\n",
        "w = np.random.normal(3, 1.0, size=(observations, 1))\n",
        "df = pd.DataFrame(np.hstack([target, y, z, w]), \n",
        "                  columns=['target', 'y', 'z', 'w'])\n",
        "\n",
        "# Create train test datasets\n",
        "train = df.sample(frac=0.8, random_state=1)\n",
        "test = df.drop(train.index)\n",
        "\n",
        "# Convert to pyspark dataset\n",
        "spark_train = spark.createDataFrame(train)\n",
        "spark_test = spark.createDataFrame(test)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 28,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OK893jun_mHC",
        "outputId": "4bd32593-ff3a-45a9-9960-538fabdb00f3"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "time: 519 µs (started: 2023-02-16 20:35:21 +00:00)\n"
          ]
        }
      ],
      "source": [
        "# a dict to store 3 models\n",
        "models = {}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 29,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7loXDkWlgGEm",
        "outputId": "02eeb428-dd57-4b1e-f262-0d5d308b5b4a"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "time: 5.37 ms (started: 2023-02-16 20:35:21 +00:00)\n"
          ]
        }
      ],
      "source": [
        "# Create functions for timing, pyspark regression and statsmodels regression\n",
        "def timeit(func):\n",
        "    @wraps(func)\n",
        "    def wrapper(*args, **kwargs):\n",
        "        start_time = time.time()\n",
        "        result = func(*args, **kwargs)\n",
        "        end_time = time.time()\n",
        "        elapsed_time = end_time - start_time\n",
        "        print(f\"{func.__name__} took {elapsed_time:.2f} seconds to run.\")\n",
        "        return result\n",
        "    return wrapper\n",
        "\n",
        "@timeit\n",
        "def logit_pyspark(var, train_data, test_data):\n",
        "  # create a VectorAssembler to combine the independent variables\n",
        "  assembler = VectorAssembler(inputCols=[var], outputCol='features')\n",
        "  train_spark = assembler.transform(train_data)\n",
        "  test_spark = assembler.transform(test_data)\n",
        "\n",
        "\n",
        "  # Create model\n",
        "  lr = LogisticRegression(featuresCol='features', labelCol='target')\n",
        "  model = lr.fit(train_spark)\n",
        "\n",
        "\n",
        "  # Test model\n",
        "  predictions = model.transform(test_spark)\n",
        "  evaluator = BinaryClassificationEvaluator(labelCol='target')\n",
        "  auc = evaluator.evaluate(predictions)\n",
        "  print(f\"Model with {var} as the independent variable has AUC of {auc:.2f}\")\n",
        "\n",
        "@timeit\n",
        "def logit_statmodels(train_df, test_df, var):\n",
        "  # Create model\n",
        "  logit_model = sm.Logit(train_df[['target']], train_df[[var]])\n",
        "  result = logit_model.fit(disp=0)  \n",
        "\n",
        "  # Create AUC\n",
        "  test['y_pred'] = result.predict(test_df[[var]])\n",
        "  auc = roc_auc_score(test_df['target'], test_df['y_pred'])\n",
        "  print(f\"Model with {var} as the independent variable has AUC of {auc:.2f}\")\n",
        "  return result"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AVNOeMjS5a5x"
      },
      "source": [
        "# Stats model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 30,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wYXRG4qkmYqU",
        "outputId": "d4e333c5-8034-41fe-be22-d850f2f0cc55"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Model with y as the independent variable has AUC of 1.00\n",
            "logit_statmodels took 0.37 seconds to run.\n",
            "Model with z as the independent variable has AUC of 0.47\n",
            "logit_statmodels took 0.19 seconds to run.\n",
            "Model with w as the independent variable has AUC of 0.51\n",
            "logit_statmodels took 0.31 seconds to run.\n",
            "time: 870 ms (started: 2023-02-16 20:35:21 +00:00)\n"
          ]
        }
      ],
      "source": [
        "# Test Statsmodel\n",
        "vars = ['y', 'z', 'w']\n",
        "\n",
        "# Run regression\n",
        "for var in vars:\n",
        "  models[var] = logit_statmodels(train, test, var)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VT5dpOS05luk"
      },
      "source": [
        "# Pyspark test 1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 31,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "G6fyVTaCmYlh",
        "outputId": "e84714a8-1e9c-42b0-88dc-f45b0aebf234"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Model with w as the independent variable has AUC of 0.49\n",
            "logit_pyspark took 5.27 seconds to run.\n",
            "Model with z as the independent variable has AUC of 0.53\n",
            "logit_pyspark took 8.10 seconds to run.\n",
            "Model with y as the independent variable has AUC of 1.00\n",
            "logit_pyspark took 7.52 seconds to run.\n",
            "time: 20.9 s (started: 2023-02-16 20:35:22 +00:00)\n"
          ]
        }
      ],
      "source": [
        "# Test PySpark\n",
        "vars = ['w', 'z', 'y']\n",
        "\n",
        "# Run regression\n",
        "for var in vars:\n",
        "  logit_pyspark(var, spark_train, spark_test)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "CpsVlW205nlR"
      },
      "source": [
        "# Pyspark test 2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 32,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "E_ZH6rHcoAOW",
        "outputId": "fd46eb5c-e518-4df1-b836-a9764b00c73e"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Model with y as the independent variable has AUC of 1.00\n",
            "logit_pyspark took 7.32 seconds to run.\n",
            "Model with z as the independent variable has AUC of 0.53\n",
            "logit_pyspark took 3.80 seconds to run.\n",
            "Model with w as the independent variable has AUC of 0.49\n",
            "logit_pyspark took 4.76 seconds to run.\n",
            "time: 15.9 s (started: 2023-02-16 20:35:43 +00:00)\n"
          ]
        }
      ],
      "source": [
        "# Test PySpark\n",
        "vars = ['y', 'z', 'w']\n",
        "\n",
        "spark_train_part_1 = spark_train.coalesce(1)\n",
        "spark_test_part_1 = spark_test.coalesce(1)\n",
        "\n",
        "# Run regression\n",
        "for var in vars:\n",
        "  logit_pyspark(var, spark_train_part_1, spark_test_part_1)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "Ggv6ECBD5o_R"
      },
      "source": [
        "# Final version\n",
        "1.   pred(df): a python function, it will be used as pandas_udf later\n",
        "2.   spark_test.withColumn(\"id\", f.lit(1)).groupBy('id'): since all rows has to feed into the model, this column is just a dummy column\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 33,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EJX7LD8Frucn",
        "outputId": "c20a227c-f8d0-4e3c-fc1b-93cd0dba5598"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "time: 1.02 ms (started: 2023-02-16 20:35:59 +00:00)\n"
          ]
        }
      ],
      "source": [
        "\n",
        "def pred(df):\n",
        "    df['y_pred'] = models['y'].predict(df['y'])\n",
        "    df['z_pred'] = models['z'].predict(df['z'])\n",
        "    df['w_pred'] = models['w'].predict(df['w'])\n",
        "    df['auc_y'] = roc_auc_score(df['target'], df['y_pred'])\n",
        "    df['auc_z'] = roc_auc_score(df['target'], df['z_pred'])\n",
        "    df['auc_w'] = roc_auc_score(df['target'], df['w_pred'])\n",
        "    \n",
        "    return df[['auc_y', 'auc_z', 'auc_w']].iloc[[1]]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "vV7WvpxUBevh",
        "outputId": "f4d652a6-fb96-4142-a4de-63a285a2c610"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "+-----+-------------------+------------------+\n",
            "|auc_y|              auc_z|             auc_w|\n",
            "+-----+-------------------+------------------+\n",
            "|  1.0|0.47158886270789624|0.5066707479655908|\n",
            "+-----+-------------------+------------------+\n",
            "\n",
            "time: 1.47 s (started: 2023-02-16 20:35:59 +00:00)\n"
          ]
        }
      ],
      "source": [
        "spark_test.withColumn(\"id\", f.lit(1)).groupBy('id').applyInPandas(pred, schema='auc_y double, auc_z double, auc_w double').show()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.9.13 (tags/v3.9.13:6de2ca5, May 17 2022, 16:36:42) [MSC v.1929 64 bit (AMD64)]"
    },
    "vscode": {
      "interpreter": {
        "hash": "eabca979b0553fa6d87e9a00c352604d3b703d4afc9641643dd42376492b80f6"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}