spark_regression.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "!pip install PySpark"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ykLmOfJtkDG7",
        "outputId": "2845c483-9a1b-421c-ec34-196daf164fa2"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "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.1)\n",
            "Requirement already satisfied: py4j==0.10.9.5 in /usr/local/lib/python3.8/dist-packages (from PySpark) (0.10.9.5)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "WHh74H3Kf7IB"
      },
      "outputs": [],
      "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"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "spark = SparkSession.builder.appName(\"test_regression\").getOrCreate()"
      ],
      "metadata": {
        "id": "05lGSWi1jYkG"
      },
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "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)"
      ],
      "metadata": {
        "id": "w3R-k49ngGKV"
      },
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "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",
        "  # Create model\n",
        "  lr = LogisticRegression(featuresCol='features', labelCol='target')\n",
        "  model = lr.fit(train_spark)\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}\")"
      ],
      "metadata": {
        "id": "7loXDkWlgGEm"
      },
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Test Statsmodel\n",
        "vars = ['y', 'z', 'w']\n",
        "\n",
        "# Run regression\n",
        "for var in vars:\n",
        "  logit_statmodels(train, test, var)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wYXRG4qkmYqU",
        "outputId": "0fd7d913-f76e-4eff-fdac-9f648978d3d9"
      },
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Model with y as the independent variable has AUC of 1.00\n",
            "logit_statmodels took 0.61 seconds to run.\n",
            "Model with z as the independent variable has AUC of 0.47\n",
            "logit_statmodels took 0.43 seconds to run.\n",
            "Model with w as the independent variable has AUC of 0.51\n",
            "logit_statmodels took 0.29 seconds to run.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Test PySpark\n",
        "vars = ['y', 'z', 'w']\n",
        "\n",
        "# Run regression\n",
        "for var in vars:\n",
        "  logit_pyspark(var, spark_train, spark_test)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "G6fyVTaCmYlh",
        "outputId": "4c424403-3535-4d7b-ff2a-de2ed67327bb"
      },
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Model with y as the independent variable has AUC of 1.00\n",
            "logit_pyspark took 20.79 seconds to run.\n",
            "Model with z as the independent variable has AUC of 0.53\n",
            "logit_pyspark took 11.27 seconds to run.\n",
            "Model with w as the independent variable has AUC of 0.51\n",
            "logit_pyspark took 8.37 seconds to run.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "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)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "E_ZH6rHcoAOW",
        "outputId": "ec1b032d-6868-4c9d-f6c7-ba777beb1181"
      },
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Model with y as the independent variable has AUC of 1.00\n",
            "logit_pyspark took 9.79 seconds to run.\n",
            "Model with z as the independent variable has AUC of 0.53\n",
            "logit_pyspark took 7.93 seconds to run.\n",
            "Model with w as the independent variable has AUC of 0.51\n",
            "logit_pyspark took 7.21 seconds to run.\n"
          ]
        }
      ]
    }
  ]
}