emr_files/Customer_Income_EMR.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4f7bb117",
   "metadata": {},
   "source": [
    "# Running Spark ML Pipelines on Amazon EMR Cluster"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f00d213b",
   "metadata": {},
   "source": [
    "Notebook [reference](https://github.com/aws/amazon-sagemaker-examples/blob/master/sagemaker-python-sdk/sparkml_serving_emr_mleap_abalone/sparkml_serving_emr_mleap_abalone.ipynb). "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49dba8fe",
   "metadata": {},
   "source": [
    "Please ensure you use the following configuration when setting "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9a2c3c3d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "Current session configs: <tt>{'driverMemory': '1000M', 'executorCores': 2, 'kind': 'pyspark'}</tt><br>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<table>\n",
       "<tr><th>ID</th><th>YARN Application ID</th><th>Kind</th><th>State</th><th>Spark UI</th><th>Driver log</th><th>User</th><th>Current session?</th></tr><tr><td>2</td><td>application_1631487410552_0003</td><td>pyspark</td><td>busy</td><td><a target=\"_blank\" href=\"http://ip-172-31-4-247.us-east-2.compute.internal:20888/proxy/application_1631487410552_0003/\">Link</a></td><td><a target=\"_blank\" href=\"http://ip-172-31-13-187.us-east-2.compute.internal:8042/node/containerlogs/container_1631487410552_0003_01_000001/livy\">Link</a></td><td>None</td><td></td></tr><tr><td>3</td><td>application_1631487410552_0004</td><td>pyspark</td><td>idle</td><td><a target=\"_blank\" href=\"http://ip-172-31-4-247.us-east-2.compute.internal:20888/proxy/application_1631487410552_0004/\">Link</a></td><td><a target=\"_blank\" href=\"http://ip-172-31-13-187.us-east-2.compute.internal:8042/node/containerlogs/container_1631487410552_0004_01_000001/livy\">Link</a></td><td>None</td><td></td></tr></table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5a613198",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting Spark application\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table>\n",
       "<tr><th>ID</th><th>YARN Application ID</th><th>Kind</th><th>State</th><th>Spark UI</th><th>Driver log</th><th>User</th><th>Current session?</th></tr><tr><td>5</td><td>application_1631487410552_0006</td><td>pyspark</td><td>idle</td><td><a target=\"_blank\" href=\"http://ip-172-31-4-247.us-east-2.compute.internal:20888/proxy/application_1631487410552_0006/\">Link</a></td><td><a target=\"_blank\" href=\"http://ip-172-31-13-187.us-east-2.compute.internal:8042/node/containerlogs/container_1631487410552_0006_01_000001/livy\">Link</a></td><td>None</td><td>✔</td></tr></table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cdababccb52e454aa4313b4375009606",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SparkSession available as 'spark'.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2fcc5cb0a88c4f72ab34e14b7e083576",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from __future__ import print_function\n",
    "\n",
    "import os\n",
    "import shutil\n",
    "import boto3\n",
    "import numpy as np\n",
    "import time\n",
    "\n",
    "import pyspark\n",
    "import pyspark.sql.functions as func\n",
    "from pyspark.sql import SparkSession\n",
    "from pyspark.ml import Pipeline\n",
    "from pyspark.sql.functions import *\n",
    "from pyspark.sql.types import (\n",
    "    StructField, \n",
    "    StructType, \n",
    "    StringType, \n",
    "    DoubleType,\n",
    "    FloatType\n",
    ")\n",
    "\n",
    "from pyspark.ml.feature import (\n",
    "    StringIndexer,\n",
    "    VectorIndexer,\n",
    "    OneHotEncoderEstimator, #MLeap =<1.17.0 only supports OneHotEncoderEstimator for OHE\n",
    "    VectorAssembler,\n",
    "    IndexToString,\n",
    "    MinMaxScaler\n",
    ")\n",
    "\n",
    "from pyspark.ml.classification import (\n",
    "    LogisticRegression,\n",
    "    DecisionTreeClassifier,\n",
    "    RandomForestClassifier,\n",
    "    GBTClassifier\n",
    ")\n",
    "\n",
    "from pyspark.ml.tuning import ParamGridBuilder, CrossValidator\n",
    "from pyspark.ml.evaluation import MulticlassClassificationEvaluator\n",
    "from mleap.pyspark.spark_support import SimpleSparkSerializer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ee00ed4f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fa3a4cfa957040feb3c2ab76a2546e8c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "spark = SparkSession \\\n",
    "    .builder \\\n",
    "    .appName(\"Classifying if a customer's income is >$50k or lesser.\") \\\n",
    "    .getOrCreate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0b863f6a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7f6f5fea904444b3af1f9ac2113d5f5c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Load data from a public S3 bucket using the object URI\n",
    "# The data does not include an header\n",
    "# Also, the leading whitespace for each value needs to be removed\n",
    "\n",
    "rawData = spark.read.format('csv') \\\n",
    "                .option('header', 'false') \\\n",
    "                .option('ignoreLeadingWhiteSpace', 'true') \\\n",
    "                .load('s3://layer-guide/dataset/adult-dataset.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ba4165a8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e5c46a1d6ced40768478cf3296f7a93b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "## Another way to read the dataset--simpler.\n",
    "\n",
    "# rawData = spark.read.csv(\n",
    "#     \"s3://layer-guide/dataset/adult-dataset.csv\", header=False, schema=schema\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b917da82",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f6eacc0a90b54cf89940d9541ca59162",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Specify the features and label you want to use for the dataset\n",
    "\n",
    "dataset = rawData.toDF('Age',\n",
    "               'WorkClass',\n",
    "               'FnlWgt',\n",
    "               'Education',\n",
    "               'EducationNum',\n",
    "               'MaritalStatus',\n",
    "               'Occupation',\n",
    "               'Relationship',\n",
    "               'Race',\n",
    "               'Gender',\n",
    "               'CapitalGain',\n",
    "               'CapitalLoss',\n",
    "               'HoursPerWeek',\n",
    "               'NativeCountry',\n",
    "               'Income'\n",
    "                )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f01230a3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ab1b119f6fe9467a81e5640427db37f9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total Columns: 15\n",
      "root\n",
      " |-- Age: string (nullable = true)\n",
      " |-- WorkClass: string (nullable = true)\n",
      " |-- FnlWgt: string (nullable = true)\n",
      " |-- Education: string (nullable = true)\n",
      " |-- EducationNum: string (nullable = true)\n",
      " |-- MaritalStatus: string (nullable = true)\n",
      " |-- Occupation: string (nullable = true)\n",
      " |-- Relationship: string (nullable = true)\n",
      " |-- Race: string (nullable = true)\n",
      " |-- Gender: string (nullable = true)\n",
      " |-- CapitalGain: string (nullable = true)\n",
      " |-- CapitalLoss: string (nullable = true)\n",
      " |-- HoursPerWeek: string (nullable = true)\n",
      " |-- NativeCountry: string (nullable = true)\n",
      " |-- Income: string (nullable = true)"
     ]
    }
   ],
   "source": [
    "# Validate data structure\n",
    "\n",
    "print(f'Total Columns: {len(dataset.dtypes)}')\n",
    "dataset.printSchema()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "56dd485c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3877f23ed1d041e4a113f9dcab2a33e4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "48842"
     ]
    }
   ],
   "source": [
    "# WARNING: This can be inefficient for large datasets\n",
    "dataset.count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ea1ea01a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d8753f3e7d8e4099a3f216ab957f6192",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Drop the column that's not useful\n",
    "\n",
    "dataset = dataset.drop('FnlWgt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "efb0faac",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "89dad30d44c54aa0bce449480001dd4d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---+----------------+---------+------------+--------------------+-----------------+-------------+-----+------+-----------+-----------+------------+-------------+------+\n",
      "|Age|       WorkClass|Education|EducationNum|       MaritalStatus|       Occupation| Relationship| Race|Gender|CapitalGain|CapitalLoss|HoursPerWeek|NativeCountry|Income|\n",
      "+---+----------------+---------+------------+--------------------+-----------------+-------------+-----+------+-----------+-----------+------------+-------------+------+\n",
      "| 39|       State-gov|Bachelors|          13|       Never-married|     Adm-clerical|Not-in-family|White|  Male|       2174|          0|          40|United-States| <=50K|\n",
      "| 50|Self-emp-not-inc|Bachelors|          13|  Married-civ-spouse|  Exec-managerial|      Husband|White|  Male|          0|          0|          13|United-States| <=50K|\n",
      "| 38|         Private|  HS-grad|           9|            Divorced|Handlers-cleaners|Not-in-family|White|  Male|          0|          0|          40|United-States| <=50K|\n",
      "| 53|         Private|     11th|           7|  Married-civ-spouse|Handlers-cleaners|      Husband|Black|  Male|          0|          0|          40|United-States| <=50K|\n",
      "| 28|         Private|Bachelors|          13|  Married-civ-spouse|   Prof-specialty|         Wife|Black|Female|          0|          0|          40|         Cuba| <=50K|\n",
      "| 37|         Private|  Masters|          14|  Married-civ-spouse|  Exec-managerial|         Wife|White|Female|          0|          0|          40|United-States| <=50K|\n",
      "| 49|         Private|      9th|           5|Married-spouse-ab...|    Other-service|Not-in-family|Black|Female|          0|          0|          16|      Jamaica| <=50K|\n",
      "| 52|Self-emp-not-inc|  HS-grad|           9|  Married-civ-spouse|  Exec-managerial|      Husband|White|  Male|          0|          0|          45|United-States|  >50K|\n",
      "| 31|         Private|  Masters|          14|       Never-married|   Prof-specialty|Not-in-family|White|Female|      14084|          0|          50|United-States|  >50K|\n",
      "| 42|         Private|Bachelors|          13|  Married-civ-spouse|  Exec-managerial|      Husband|White|  Male|       5178|          0|          40|United-States|  >50K|\n",
      "+---+----------------+---------+------------+--------------------+-----------------+-------------+-----+------+-----------+-----------+------------+-------------+------+"
     ]
    }
   ],
   "source": [
    "dataset.limit(10).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "58a132e4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6d3b3f4ade8d45d08c4e32a1d21c9c8a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Replace `?` with NaN val\n",
    "\n",
    "dataset = dataset.replace('?', None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "b4cfd5b4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "78ae44c77b4842069032fb1139448bf6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Drop NaN  val\n",
    "\n",
    "dataset = dataset.dropna(how='any')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "7c07a483",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1eb94940c3c641f68b4cbd5043ac1873",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---+----------------+---------+------------+--------------------+-----------------+-------------+-----+------+-----------+-----------+------------+-------------+------+\n",
      "|Age|       WorkClass|Education|EducationNum|       MaritalStatus|       Occupation| Relationship| Race|Gender|CapitalGain|CapitalLoss|HoursPerWeek|NativeCountry|Income|\n",
      "+---+----------------+---------+------------+--------------------+-----------------+-------------+-----+------+-----------+-----------+------------+-------------+------+\n",
      "| 39|       State-gov|Bachelors|          13|       Never-married|     Adm-clerical|Not-in-family|White|  Male|       2174|          0|          40|United-States| <=50K|\n",
      "| 50|Self-emp-not-inc|Bachelors|          13|  Married-civ-spouse|  Exec-managerial|      Husband|White|  Male|          0|          0|          13|United-States| <=50K|\n",
      "| 38|         Private|  HS-grad|           9|            Divorced|Handlers-cleaners|Not-in-family|White|  Male|          0|          0|          40|United-States| <=50K|\n",
      "| 53|         Private|     11th|           7|  Married-civ-spouse|Handlers-cleaners|      Husband|Black|  Male|          0|          0|          40|United-States| <=50K|\n",
      "| 28|         Private|Bachelors|          13|  Married-civ-spouse|   Prof-specialty|         Wife|Black|Female|          0|          0|          40|         Cuba| <=50K|\n",
      "| 37|         Private|  Masters|          14|  Married-civ-spouse|  Exec-managerial|         Wife|White|Female|          0|          0|          40|United-States| <=50K|\n",
      "| 49|         Private|      9th|           5|Married-spouse-ab...|    Other-service|Not-in-family|Black|Female|          0|          0|          16|      Jamaica| <=50K|\n",
      "| 52|Self-emp-not-inc|  HS-grad|           9|  Married-civ-spouse|  Exec-managerial|      Husband|White|  Male|          0|          0|          45|United-States|  >50K|\n",
      "| 31|         Private|  Masters|          14|       Never-married|   Prof-specialty|Not-in-family|White|Female|      14084|          0|          50|United-States|  >50K|\n",
      "| 42|         Private|Bachelors|          13|  Married-civ-spouse|  Exec-managerial|      Husband|White|  Male|       5178|          0|          40|United-States|  >50K|\n",
      "+---+----------------+---------+------------+--------------------+-----------------+-------------+-----+------+-----------+-----------+------------+-------------+------+"
     ]
    }
   ],
   "source": [
    "# Observe the target column \"Income\" is a string\n",
    "\n",
    "dataset.limit(10).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "7cda2de7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "90d3667553a64f29943135ec6c714670",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Convert target column (Better to do this w/o a transformer to avoid schema complications)\n",
    "mapping= {\n",
    "        '<=50K': '0',\n",
    "        '>50K': '1'\n",
    "    }\n",
    "dataset = dataset.replace(to_replace=mapping, subset=['Income'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "1594c4c3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "adb5cc0da602450294a58122b404f8f6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "root\n",
      " |-- Age: string (nullable = true)\n",
      " |-- WorkClass: string (nullable = true)\n",
      " |-- Education: string (nullable = true)\n",
      " |-- EducationNum: string (nullable = true)\n",
      " |-- MaritalStatus: string (nullable = true)\n",
      " |-- Occupation: string (nullable = true)\n",
      " |-- Relationship: string (nullable = true)\n",
      " |-- Race: string (nullable = true)\n",
      " |-- Gender: string (nullable = true)\n",
      " |-- CapitalGain: string (nullable = true)\n",
      " |-- CapitalLoss: string (nullable = true)\n",
      " |-- HoursPerWeek: string (nullable = true)\n",
      " |-- NativeCountry: string (nullable = true)\n",
      " |-- Income: string (nullable = true)"
     ]
    }
   ],
   "source": [
    "dataset.printSchema()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "edcaa683",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e48140d966ea423ca8694721b4a5a6fe",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----+----------------+---------+------------+--------------------+-----------------+-------------+-----+------+-----------+-----------+------------+-------------+------+\n",
      "| Age|       WorkClass|Education|EducationNum|       MaritalStatus|       Occupation| Relationship| Race|Gender|CapitalGain|CapitalLoss|HoursPerWeek|NativeCountry|Income|\n",
      "+----+----------------+---------+------------+--------------------+-----------------+-------------+-----+------+-----------+-----------+------------+-------------+------+\n",
      "|39.0|       State-gov|Bachelors|        13.0|       Never-married|     Adm-clerical|Not-in-family|White|  Male|     2174.0|        0.0|        40.0|United-States|   0.0|\n",
      "|50.0|Self-emp-not-inc|Bachelors|        13.0|  Married-civ-spouse|  Exec-managerial|      Husband|White|  Male|        0.0|        0.0|        13.0|United-States|   0.0|\n",
      "|38.0|         Private|  HS-grad|         9.0|            Divorced|Handlers-cleaners|Not-in-family|White|  Male|        0.0|        0.0|        40.0|United-States|   0.0|\n",
      "|53.0|         Private|     11th|         7.0|  Married-civ-spouse|Handlers-cleaners|      Husband|Black|  Male|        0.0|        0.0|        40.0|United-States|   0.0|\n",
      "|28.0|         Private|Bachelors|        13.0|  Married-civ-spouse|   Prof-specialty|         Wife|Black|Female|        0.0|        0.0|        40.0|         Cuba|   0.0|\n",
      "|37.0|         Private|  Masters|        14.0|  Married-civ-spouse|  Exec-managerial|         Wife|White|Female|        0.0|        0.0|        40.0|United-States|   0.0|\n",
      "|49.0|         Private|      9th|         5.0|Married-spouse-ab...|    Other-service|Not-in-family|Black|Female|        0.0|        0.0|        16.0|      Jamaica|   0.0|\n",
      "|52.0|Self-emp-not-inc|  HS-grad|         9.0|  Married-civ-spouse|  Exec-managerial|      Husband|White|  Male|        0.0|        0.0|        45.0|United-States|   1.0|\n",
      "|31.0|         Private|  Masters|        14.0|       Never-married|   Prof-specialty|Not-in-family|White|Female|    14084.0|        0.0|        50.0|United-States|   1.0|\n",
      "|42.0|         Private|Bachelors|        13.0|  Married-civ-spouse|  Exec-managerial|      Husband|White|  Male|     5178.0|        0.0|        40.0|United-States|   1.0|\n",
      "+----+----------------+---------+------------+--------------------+-----------------+-------------+-----+------+-----------+-----------+------------+-------------+------+"
     ]
    }
   ],
   "source": [
    "# Convert string-encoded Num columns to numeric features\n",
    "\n",
    "dataset = dataset.withColumn('Age', \n",
    "                             dataset['Age'].cast(DoubleType()))\n",
    "dataset = dataset.withColumn('EducationNum', \n",
    "                             dataset['EducationNum'].cast(DoubleType()))\n",
    "dataset = dataset.withColumn('CapitalGain', \n",
    "                             dataset['CapitalGain'].cast(DoubleType()))\n",
    "dataset = dataset.withColumn('CapitalLoss', \n",
    "                             dataset['CapitalLoss'].cast(DoubleType()))\n",
    "dataset = dataset.withColumn('HoursPerWeek', \n",
    "                             dataset['HoursPerWeek'].cast(DoubleType()))\n",
    "dataset = dataset.withColumn('Income', \n",
    "                             dataset['Income'].cast(DoubleType()))\n",
    "\n",
    "dataset.limit(10).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "4bbacf10",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "46241c94bb684d6e8972d79abd4985ca",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "root\n",
      " |-- Age: double (nullable = true)\n",
      " |-- WorkClass: string (nullable = true)\n",
      " |-- Education: string (nullable = true)\n",
      " |-- EducationNum: double (nullable = true)\n",
      " |-- MaritalStatus: string (nullable = true)\n",
      " |-- Occupation: string (nullable = true)\n",
      " |-- Relationship: string (nullable = true)\n",
      " |-- Race: string (nullable = true)\n",
      " |-- Gender: string (nullable = true)\n",
      " |-- CapitalGain: double (nullable = true)\n",
      " |-- CapitalLoss: double (nullable = true)\n",
      " |-- HoursPerWeek: double (nullable = true)\n",
      " |-- NativeCountry: string (nullable = true)\n",
      " |-- Income: double (nullable = true)"
     ]
    }
   ],
   "source": [
    "# Validate the data types...\n",
    "\n",
    "#dataset.schema.fields\n",
    "\n",
    "# You can also do this with:\n",
    "#dataset.dtypes\n",
    "\n",
    "dataset.printSchema()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "d441dfce",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e565f56fdd9247348ae5c118dbe1124b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-RECORD 0----------------------\n",
      " Age           | 39.0          \n",
      " WorkClass     | State-gov     \n",
      " Education     | Bachelors     \n",
      " EducationNum  | 13.0          \n",
      " MaritalStatus | Never-married \n",
      " Occupation    | Adm-clerical  \n",
      " Relationship  | Not-in-family \n",
      " Race          | White         \n",
      " Gender        | Male          \n",
      " CapitalGain   | 2174.0        \n",
      " CapitalLoss   | 0.0           \n",
      " HoursPerWeek  | 40.0          \n",
      " NativeCountry | United-States \n",
      " Income        | 0.0           \n",
      "only showing top 1 row"
     ]
    }
   ],
   "source": [
    "dataset.limit(10).show(n=1, truncate=False, vertical=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "ebb37027",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e4f3c921d0164522be4a0df6e3a2766f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Data segragation\n",
    "\n",
    "(trainingData, testData) = dataset.randomSplit([0.8,0.2], seed=42)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74af30a8",
   "metadata": {},
   "source": [
    "## Feature Transformation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "00c2b529",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1c9299b321964e3eb939d3e6699d475e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['WorkClass', 'Education', 'MaritalStatus', 'Occupation', 'Relationship', 'Race', 'Gender', 'NativeCountry']"
     ]
    }
   ],
   "source": [
    "# Select categorical features\n",
    "\n",
    "categoricalFeatures = [item[0] for item in dataset.dtypes if \\\n",
    "                     item[1].startswith('string')]\n",
    "\n",
    "print(categoricalFeatures)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "32fcb9a1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "29a4f107b8e94719b70e5db32966fe4c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Create an array of StringIndexers to convert the categorical values to indices\n",
    "\n",
    "indexers = [StringIndexer(\n",
    "    inputCol=column, \n",
    "    outputCol=column + '_index', \n",
    "    handleInvalid='keep') for column in categoricalFeatures]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "58f37192",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e39ac9e397db4bf780675a83e3f39504",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Create an array of OneHotEncoders to encode the categorical values\n",
    "# Using OneHotEncoderEstimator here because OneHotEncoder transformer is...\n",
    "# ... not compatible with MLeap (model serialization) as of this time.\n",
    "encoders = [OneHotEncoderEstimator(\n",
    "    inputCols=[column + '_index'], \n",
    "    outputCols= [column + '_encoded']) for column in categoricalFeatures]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "41dc99a5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a1464c8db569484783cdd43e88ae2897",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Age', 'EducationNum', 'CapitalGain', 'CapitalLoss', 'HoursPerWeek']"
     ]
    }
   ],
   "source": [
    "# Select numeric features\n",
    "\n",
    "numFeatures = [item[0] for item in dataset.dtypes \\\n",
    "               if item[1].startswith('double')]\n",
    "\n",
    "# Only select features for scaling; not the label/target \"Income\"\n",
    "numFeatures = numFeatures[:-1] \n",
    "\n",
    "print(numFeatures)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "29ca4034",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "09c95d972c2f4d008f42c20eb34f302c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Num features must be converted to single feature vector before they can be scaled.\n",
    "\n",
    "vectorFeatures = [VectorAssembler(inputCols=[col], \\\n",
    "                                  outputCol=col + \"_vec\") for col in numFeatures]\n",
    "\n",
    "# Scale Num Vectors (better to also train models with and without scaling and observer performance)\n",
    "numScaler = [MinMaxScaler(\n",
    "    inputCol= column + \"_vec\", \n",
    "    outputCol= column + '_scaled') for column in numFeatures]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "e7c07d39",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bbbeb28b26024909afcfe8e2cf90f446",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Create a pipeline to observe transformed features\n",
    "\n",
    "pipeline = Pipeline(stages=indexers + encoders + vectorFeatures \\\n",
    "                    + numScaler)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "36af3294",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f446a64043ef4975afae58e62c005d10",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "transformedDF = pipeline.fit(trainingData).transform(trainingData)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "c5378057",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "abaf61cc89d04a789893be08e05a9d02",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-RECORD 0-------------------------------------\n",
      " Age                   | 17.0                 \n",
      " WorkClass             | Federal-gov          \n",
      " Education             | 11th                 \n",
      " EducationNum          | 7.0                  \n",
      " MaritalStatus         | Never-married        \n",
      " Occupation            | Adm-clerical         \n",
      " Relationship          | Not-in-family        \n",
      " Race                  | Black                \n",
      " Gender                | Female               \n",
      " CapitalGain           | 0.0                  \n",
      " CapitalLoss           | 1602.0               \n",
      " HoursPerWeek          | 40.0                 \n",
      " NativeCountry         | United-States        \n",
      " Income                | 0.0                  \n",
      " WorkClass_index       | 5.0                  \n",
      " Education_index       | 5.0                  \n",
      " MaritalStatus_index   | 1.0                  \n",
      " Occupation_index      | 3.0                  \n",
      " Relationship_index    | 1.0                  \n",
      " Race_index            | 1.0                  \n",
      " Gender_index          | 1.0                  \n",
      " NativeCountry_index   | 0.0                  \n",
      " WorkClass_encoded     | (7,[5],[1.0])        \n",
      " Education_encoded     | (16,[5],[1.0])       \n",
      " MaritalStatus_encoded | (7,[1],[1.0])        \n",
      " Occupation_encoded    | (14,[3],[1.0])       \n",
      " Relationship_encoded  | (6,[1],[1.0])        \n",
      " Race_encoded          | (5,[1],[1.0])        \n",
      " Gender_encoded        | (2,[1],[1.0])        \n",
      " NativeCountry_encoded | (41,[0],[1.0])       \n",
      " Age_vec               | [17.0]               \n",
      " EducationNum_vec      | [7.0]                \n",
      " CapitalGain_vec       | [0.0]                \n",
      " CapitalLoss_vec       | [1602.0]             \n",
      " HoursPerWeek_vec      | [40.0]               \n",
      " Age_scaled            | [0.0]                \n",
      " EducationNum_scaled   | [0.4]                \n",
      " CapitalGain_scaled    | [0.0]                \n",
      " CapitalLoss_scaled    | [0.3677685950413223] \n",
      " HoursPerWeek_scaled   | [0.3979591836734694] \n",
      "only showing top 1 row"
     ]
    }
   ],
   "source": [
    "transformedDF.limit(10).show(n=1, truncate=False, vertical=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "6ff84dd6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6778dd30d1134b4b9f6265929412b86c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Select required features\n",
    "\n",
    "requiredFeatures = [\n",
    "    'Age_scaled',\n",
    "    'EducationNum_scaled',\n",
    "    'CapitalGain_scaled',\n",
    "    'CapitalLoss_scaled',\n",
    "    'HoursPerWeek_scaled',\n",
    "    'WorkClass_encoded',\n",
    "    'Education_encoded',\n",
    "    'MaritalStatus_encoded',\n",
    "    'Occupation_encoded',\n",
    "    'Relationship_encoded',\n",
    "    'Race_encoded',\n",
    "    'Gender_encoded',\n",
    "    'NativeCountry_encoded'\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "4ac0c248",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ccacb8c97872496ebc7649ac96145b9c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-RECORD 0----------------------------------------------------------------------------------------------------------------------------------\n",
      " Age                   | 17.0                                                                                                              \n",
      " WorkClass             | Federal-gov                                                                                                       \n",
      " Education             | 11th                                                                                                              \n",
      " EducationNum          | 7.0                                                                                                               \n",
      " MaritalStatus         | Never-married                                                                                                     \n",
      " Occupation            | Adm-clerical                                                                                                      \n",
      " Relationship          | Not-in-family                                                                                                     \n",
      " Race                  | Black                                                                                                             \n",
      " Gender                | Female                                                                                                            \n",
      " CapitalGain           | 0.0                                                                                                               \n",
      " CapitalLoss           | 1602.0                                                                                                            \n",
      " HoursPerWeek          | 40.0                                                                                                              \n",
      " NativeCountry         | United-States                                                                                                     \n",
      " Income                | 0.0                                                                                                               \n",
      " WorkClass_index       | 5.0                                                                                                               \n",
      " Education_index       | 5.0                                                                                                               \n",
      " MaritalStatus_index   | 1.0                                                                                                               \n",
      " Occupation_index      | 3.0                                                                                                               \n",
      " Relationship_index    | 1.0                                                                                                               \n",
      " Race_index            | 1.0                                                                                                               \n",
      " Gender_index          | 1.0                                                                                                               \n",
      " NativeCountry_index   | 0.0                                                                                                               \n",
      " WorkClass_encoded     | (7,[5],[1.0])                                                                                                     \n",
      " Education_encoded     | (16,[5],[1.0])                                                                                                    \n",
      " MaritalStatus_encoded | (7,[1],[1.0])                                                                                                     \n",
      " Occupation_encoded    | (14,[3],[1.0])                                                                                                    \n",
      " Relationship_encoded  | (6,[1],[1.0])                                                                                                     \n",
      " Race_encoded          | (5,[1],[1.0])                                                                                                     \n",
      " Gender_encoded        | (2,[1],[1.0])                                                                                                     \n",
      " NativeCountry_encoded | (41,[0],[1.0])                                                                                                    \n",
      " Age_vec               | [17.0]                                                                                                            \n",
      " EducationNum_vec      | [7.0]                                                                                                             \n",
      " CapitalGain_vec       | [0.0]                                                                                                             \n",
      " CapitalLoss_vec       | [1602.0]                                                                                                          \n",
      " HoursPerWeek_vec      | [40.0]                                                                                                            \n",
      " Age_scaled            | [0.0]                                                                                                             \n",
      " EducationNum_scaled   | [0.4]                                                                                                             \n",
      " CapitalGain_scaled    | [0.0]                                                                                                             \n",
      " CapitalLoss_scaled    | [0.3677685950413223]                                                                                              \n",
      " HoursPerWeek_scaled   | [0.3979591836734694]                                                                                              \n",
      " features              | (103,[1,3,4,10,17,29,38,50,56,61,62],[0.4,0.3677685950413223,0.3979591836734694,1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0]) \n",
      "only showing top 1 row"
     ]
    }
   ],
   "source": [
    "# Create a Vec assembler\n",
    "\n",
    "assembler = VectorAssembler(inputCols=requiredFeatures, outputCol='features')\n",
    "\n",
    "transformedDF = assembler.transform(transformedDF)\n",
    "\n",
    "transformedDF.limit(10).show(n=1, truncate=False, vertical=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "8577df15",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d6e52bd8ee064c31bea562e9c166698d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------+\n",
      "|            features|\n",
      "+--------------------+\n",
      "|(103,[1,3,4,10,17...|\n",
      "|(103,[1,4,7,17,29...|\n",
      "|(103,[1,4,7,17,29...|\n",
      "|(103,[1,4,7,17,29...|\n",
      "|(103,[1,4,7,17,29...|\n",
      "|(103,[1,4,7,17,29...|\n",
      "|(103,[1,4,7,17,29...|\n",
      "|(103,[1,4,7,17,29...|\n",
      "|(103,[1,4,7,17,29...|\n",
      "|(103,[1,4,7,17,29...|\n",
      "+--------------------+"
     ]
    }
   ],
   "source": [
    "transformedDF.select('features').limit(10).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0769644d",
   "metadata": {},
   "source": [
    "## Selecting Model Estimators"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "495f9788",
   "metadata": {},
   "source": [
    "- [Binomial logistic regression](https://spark.apache.org/docs/latest/ml-classification-regression.html#binomial-logistic-regression) --> For predicting a binary class problem.\n",
    "\n",
    " - Assumption for selecting this estimator: To establish a baseline performance for the classifiers on the dataset.\n",
    " - Learn more about LR in the [docs](https://spark.apache.org/docs/latest/api/python/reference/api/pyspark.ml.classification.LogisticRegression.html).\n",
    "\n",
    "- [Decision tree classifier](https://spark.apache.org/docs/latest/ml-classification-regression.html#decision-tree-classifier): Decision trees are great for giving some boost in performance but also risk overfitting and are mostly unstable for production usage.\n",
    "\n",
    " - Learn more about this estimator in the [docs](https://spark.apache.org/docs/latest/api/python/reference/api/pyspark.ml.classification.DecisionTreeClassifier.html#pyspark.ml.classification.DecisionTreeClassifier).\n",
    "\n",
    "\n",
    "- [Random forest classifier](https://spark.apache.org/docs/latest/ml-classification-regression.html#random-forest-classifier): To curb any potential for overfitting on the dataset, and it's a really good candidate for an optimal production model.\n",
    " - If we deploy this to production, it can be able to handle missing values from new data it needs to predict on.\n",
    " - Learn more about this estimator in the [docs](https://spark.apache.org/docs/3.1.1/api/python/reference/api/pyspark.ml.classification.RandomForestClassifier.html).\n",
    "\n",
    "\n",
    "- (Optional) [Gradient-boosted tree classifier](https://spark.apache.org/docs/latest/ml-classification-regression.html#gradient-boosted-tree-classifier): An alternative to random forest, if it performs better.\n",
    " - Learn more about this estimator in the [docs](https://spark.apache.org/docs/latest/api/python/reference/api/pyspark.ml.classification.GBTClassifier.html).\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "04d3b48b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fc4f639fe095445c804c72585e541a87",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# lr = LogisticRegression(labelCol='Income', \\\n",
    "#                         featuresCol='features')\n",
    "\n",
    "# dt = DecisionTreeClassifier(labelCol=\"Income\", \\\n",
    "#                             featuresCol=\"features\")\n",
    "\n",
    "# rf = RandomForestClassifier(labelCol='Income', \n",
    "#                              featuresCol='features')\n",
    "\n",
    "gbt = GBTClassifier(labelCol=\"Income\", featuresCol=\"features\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "974574de",
   "metadata": {},
   "source": [
    "Combine the transformed features and model estimator into a single pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "2a60bf3b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "655c0963b9b741eba686dcd0a899450e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "\n",
    "# pipeline_lr = Pipeline(\n",
    "#     stages=indexers + encoders + vectorFeatures + numScaler \\\n",
    "#      + [assembler, lr]\n",
    "# )\n",
    "\n",
    "# pipeline_dt = Pipeline(\n",
    "#     stages=indexers + encoders + vectorFeatures + numScaler \\\n",
    "#      + [assembler, dt]\n",
    "# )\n",
    "\n",
    "# pipeline_rf = Pipeline(\n",
    "#     stages=indexers + encoders + vectorFeatures + numScaler \\\n",
    "#      + [assembler, rf]\n",
    "# )\n",
    "\n",
    "pipeline_gbt = Pipeline(\n",
    "    stages=indexers + encoders + vectorFeatures + numScaler \\\n",
    "     + [assembler, gbt]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7988e3cd",
   "metadata": {},
   "source": [
    "Define a set of parameters to be used for each estimator. If you want to save time, you may want to use only one model, as running the set of parameters below can get pretty computer intensive and may cost you some more cents.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "b278d637",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "09c7f3eca64c4210b18f2108446bef6d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "# paramGrid_lr = ParamGridBuilder().addGrid(\n",
    "#     lr.maxIter, [10,50,100]).addGrid(\n",
    "#     lr.regParam, [0.1, 0.3, 1.0]).addGrid(\n",
    "#     lr.elasticNetParam, [0.0, 0.7, 1.0]).addGrid(\n",
    "#     lr.threshold, [0.5, 0.75]\n",
    "#     ).build()\n",
    "\n",
    "\n",
    "# paramGrid_dt = ParamGridBuilder().addGrid(\n",
    "#     dt.maxDepth, [5, 10, 15, 20]).build()\n",
    "\n",
    "\n",
    "# paramGrid_rf = ParamGridBuilder().addGrid(\n",
    "#     rf.maxDepth, [5, 10, 15]).addGrid(\n",
    "#     rf.numTrees, [50, 100, 150]).build()\n",
    "\n",
    "\n",
    "# paramGrid_gbt = ParamGridBuilder().addGrid(\n",
    "#     gbt.maxDepth, [5, 10, 15]).addGrid(\n",
    "#     gbt.maxIter, [20, 25, 50]).build()\n",
    "\n",
    "\n",
    "# For the purposes of demonstration, and to save time,\n",
    "# ... the following parameters are used instead.\n",
    "# You may want to commit this out.\n",
    "paramGrid_gbt = ParamGridBuilder().addGrid(\n",
    "    gbt.maxDepth, [10]).addGrid(\n",
    "    gbt.maxIter, [50]).build()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3dc340be",
   "metadata": {},
   "source": [
    "The class imbalance is not severe with this data, so accuracy can be used to evaluate the estimators."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "d8009b6a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ef40aeb1fe4349de9fbeb43e035acbe4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Define the estimator evaluators\n",
    "\n",
    "evaluator = MulticlassClassificationEvaluator(\n",
    "    labelCol='Income', \n",
    "    predictionCol='prediction', \n",
    "    metricName='accuracy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "76bcd900",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "81bfe88ff07342398695ac3fad20c711",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "crossVal = CrossValidator(estimator=pipeline_gbt,\n",
    "                          estimatorParamMaps=paramGrid_gbt,\n",
    "                          evaluator=evaluator,\n",
    "                          numFolds=3) # You can choose n>5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5407fc54",
   "metadata": {},
   "source": [
    "Train the model with the CrossValidator estimator..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "4191ed00",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aae9b04c3e7f43ffa2bfcd1b7cd96cb9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "587.4561674594879"
     ]
    }
   ],
   "source": [
    "# Fit CrossVal estimator with the pipeline, parameters, and evaluator\n",
    "start = time.time()\n",
    "model_gbt = crossVal.fit(trainingData)\n",
    "end = time.time()\n",
    "print(end - start)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ea5b0c0",
   "metadata": {},
   "source": [
    "Alternatively, the parameters can be defined directly within the model estimator `gbt = GBTClassifier(labelCol=\"Income\", featuresCol=\"features\", maxDepth=10, maxIter=50)` and the pipeline fitted directly on the training data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72218c9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "## IF YOU ARE NOT USING CROSSVALIDATOR ##\n",
    "# start = time.time()\n",
    "# model = pipeline.fit(trainingData)\n",
    "# end = time.time()\n",
    "# print(end - start)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5725847",
   "metadata": {},
   "source": [
    "## Make predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "fffc02d3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8697648565664e6c9aa7c521f0d24ecb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# TRANSFORM WILL APPEND A NEW \"prediction\" column to the training/test data dataframe\n",
    "predictions_train = model_gbt.transform(trainingData)\n",
    "predictions = model_gbt.transform(testData)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "a3c60c2c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f99a21bb06f847deb5fdf97a99ccde49",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Accuracy =  0.892994122678734\n",
      "Test Accuracy =  0.8584790112459637"
     ]
    }
   ],
   "source": [
    "# Evaluate training and test performance for overfitting\n",
    "\n",
    "accuracy_train = evaluator.evaluate(predictions_train)\n",
    "accuracy_test = evaluator.evaluate(predictions)\n",
    "print('Train Accuracy = ', accuracy_train)\n",
    "print('Test Accuracy = ', accuracy_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77084c54",
   "metadata": {},
   "source": [
    "The model overfits on this run. Can be improved furthe, of course to combat overfitting. Perhaps reducing MaxDepth, adding regularizers, and other techniques."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "00c9e60c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e989d85f431e4491a84eff152257110d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Age', 'WorkClass', 'Education', 'EducationNum', 'MaritalStatus', 'Occupation', 'Relationship', 'Race', 'Gender', 'CapitalGain', 'CapitalLoss', 'HoursPerWeek', 'NativeCountry', 'Income', 'WorkClass_index', 'Education_index', 'MaritalStatus_index', 'Occupation_index', 'Relationship_index', 'Race_index', 'Gender_index', 'NativeCountry_index', 'WorkClass_encoded', 'Education_encoded', 'MaritalStatus_encoded', 'Occupation_encoded', 'Relationship_encoded', 'Race_encoded', 'Gender_encoded', 'NativeCountry_encoded', 'Age_vec', 'EducationNum_vec', 'CapitalGain_vec', 'CapitalLoss_vec', 'HoursPerWeek_vec', 'Age_scaled', 'EducationNum_scaled', 'CapitalGain_scaled', 'CapitalLoss_scaled', 'HoursPerWeek_scaled', 'features', 'rawPrediction', 'probability', 'prediction']"
     ]
    }
   ],
   "source": [
    "# View new columns that have been appended.\n",
    "predictions.columns "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d5ba48c",
   "metadata": {},
   "source": [
    "## (Optional) Save model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2c03173",
   "metadata": {},
   "source": [
    "Optionally save your model to a local directory or upload to an S3 bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "207b57b1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b366193f29d94f78a0b53dd4eaed07a4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# import tempfile\n",
    "\n",
    "# model = model_gbt.bestModel\n",
    "\n",
    "# path = tempfile.mkdtemp()\n",
    "# model_path = path + \"/model\"\n",
    "# model.write().save(model_path)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "887890e5",
   "metadata": {},
   "source": [
    "## Serialize the pipeline for deployment"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "013dded0",
   "metadata": {},
   "source": [
    "Use [MLeap](https://github.com/combust/mleap) to serialize entire pipeline as it provides an easy-to-use Pipeline serialization format & execution engine for our online inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "beb33a37",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cd62dc60dd234d2b8025dbae6a5de608",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "SimpleSparkSerializer().serializeToBundle(\n",
    "    model, \"jar:file:/tmp/model.zip\", predictions\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d1fc9ae",
   "metadata": {},
   "source": [
    "To deploy the model to SageMaker, it has to be present in tar.gz format, but MLeap produces the model zip format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "bf4e33eb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "547ecf75c07c4edab244c3671a2af026",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Unzip the model artifacts and store it in tar.gz format.\n",
    "\n",
    "import zipfile\n",
    "\n",
    "with zipfile.ZipFile(\"/tmp/model.zip\") as zf:\n",
    "    zf.extractall(\"/tmp/model\")\n",
    "\n",
    "import tarfile\n",
    "\n",
    "with tarfile.open(\"/tmp/model.tar.gz\", \"w:gz\") as tar:\n",
    "    tar.add(\"/tmp/model/bundle.json\", arcname=\"bundle.json\")\n",
    "    tar.add(\"/tmp/model/root\", arcname=\"root\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d8815c1",
   "metadata": {},
   "source": [
    "Upload to an S3 bucket... Change folder and file namesif needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "5e063c57",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8e00e6defeae46d9a3847d125ab330ff",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "s3 = boto3.resource(\"s3\")\n",
    "file_name = os.path.join(\"emr/sparkml-layer/mleap\", \"model.tar.gz\")\n",
    "s3.Bucket(\"layer-bootstrap-script\").upload_file(\"/tmp/model.tar.gz\", file_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b382271",
   "metadata": {},
   "source": [
    "Optionally, delete the model from local storage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "de87b329",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "97fcc7127a794d7793f74538e48302bc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# os.remove(\"/tmp/model.zip\")\n",
    "# os.remove(\"/tmp/model.tar.gz\")\n",
    "# shutil.rmtree(\"/tmp/model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1225c2c",
   "metadata": {},
   "source": [
    "# Deploying the pipeline on AWS SageMaker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "cba7aa54",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3423a2eb567a4ea59b92ea12dec02d92",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "root\n",
      " |-- Age: double (nullable = true)\n",
      " |-- WorkClass: string (nullable = true)\n",
      " |-- Education: string (nullable = true)\n",
      " |-- EducationNum: double (nullable = true)\n",
      " |-- MaritalStatus: string (nullable = true)\n",
      " |-- Occupation: string (nullable = true)\n",
      " |-- Relationship: string (nullable = true)\n",
      " |-- Race: string (nullable = true)\n",
      " |-- Gender: string (nullable = true)\n",
      " |-- CapitalGain: double (nullable = true)\n",
      " |-- CapitalLoss: double (nullable = true)\n",
      " |-- HoursPerWeek: double (nullable = true)\n",
      " |-- NativeCountry: string (nullable = true)\n",
      " |-- Income: double (nullable = true)\n",
      " |-- WorkClass_index: double (nullable = false)\n",
      " |-- Education_index: double (nullable = false)\n",
      " |-- MaritalStatus_index: double (nullable = false)\n",
      " |-- Occupation_index: double (nullable = false)\n",
      " |-- Relationship_index: double (nullable = false)\n",
      " |-- Race_index: double (nullable = false)\n",
      " |-- Gender_index: double (nullable = false)\n",
      " |-- NativeCountry_index: double (nullable = false)\n",
      " |-- WorkClass_encoded: vector (nullable = true)\n",
      " |-- Education_encoded: vector (nullable = true)\n",
      " |-- MaritalStatus_encoded: vector (nullable = true)\n",
      " |-- Occupation_encoded: vector (nullable = true)\n",
      " |-- Relationship_encoded: vector (nullable = true)\n",
      " |-- Race_encoded: vector (nullable = true)\n",
      " |-- Gender_encoded: vector (nullable = true)\n",
      " |-- NativeCountry_encoded: vector (nullable = true)\n",
      " |-- Age_vec: vector (nullable = true)\n",
      " |-- EducationNum_vec: vector (nullable = true)\n",
      " |-- CapitalGain_vec: vector (nullable = true)\n",
      " |-- CapitalLoss_vec: vector (nullable = true)\n",
      " |-- HoursPerWeek_vec: vector (nullable = true)\n",
      " |-- Age_scaled: vector (nullable = true)\n",
      " |-- EducationNum_scaled: vector (nullable = true)\n",
      " |-- CapitalGain_scaled: vector (nullable = true)\n",
      " |-- CapitalLoss_scaled: vector (nullable = true)\n",
      " |-- HoursPerWeek_scaled: vector (nullable = true)\n",
      " |-- features: vector (nullable = true)\n",
      " |-- rawPrediction: vector (nullable = true)\n",
      " |-- probability: vector (nullable = true)\n",
      " |-- prediction: double (nullable = false)"
     ]
    }
   ],
   "source": [
    "# Optional\n",
    "# predictions.printSchema()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "edeeffc9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "90d37bbd49d94a91b60d0e6a67ed472e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\"input\": [{\"name\": \"Age\", \"type\": \"double\"}, {\"name\": \"WorkClass\", \"type\": \"string\"}, {\"name\": \"Education\", \"type\": \"string\"}, {\"name\": \"EducationNum\", \"type\": \"double\"}, {\"name\": \"MaritalStatus\", \"type\": \"string\"}, {\"name\": \"Occupation\", \"type\": \"string\"}, {\"name\": \"Relationship\", \"type\": \"string\"}, {\"name\": \"Race\", \"type\": \"string\"}, {\"name\": \"Gender\", \"type\": \"string\"}, {\"name\": \"CapitalGain\", \"type\": \"double\"}, {\"name\": \"CapitalLoss\", \"type\": \"double\"}, {\"name\": \"HoursPerWeek\", \"type\": \"double\"}, {\"name\": \"NativeCountry\", \"type\": \"string\"}], \"output\": {\"name\": \"prediction\", \"type\": \"double\"}}"
     ]
    }
   ],
   "source": [
    "import json\n",
    "\n",
    "# Define the raw data schema (which will be tranformed and fitted \\\n",
    "# by the pipeline in production)\n",
    "\n",
    "schema = {\n",
    "    \"input\": [\n",
    "        {\"name\": \"Age\", \"type\": \"double\"},\n",
    "        {\"name\": \"WorkClass\", \"type\": \"string\"},\n",
    "        {\"name\": \"Education\", \"type\": \"string\"},\n",
    "        {\"name\": \"EducationNum\", \"type\": \"double\"},\n",
    "        {\"name\": \"MaritalStatus\", \"type\": \"string\"},\n",
    "        {\"name\": \"Occupation\", \"type\": \"string\"},\n",
    "        {\"name\": \"Relationship\", \"type\": \"string\"},\n",
    "        {\"name\": \"Race\", \"type\": \"string\"},\n",
    "        {\"name\": \"Gender\", \"type\": \"string\"},\n",
    "        {\"name\": \"CapitalGain\", \"type\": \"double\"},\n",
    "        {\"name\": \"CapitalLoss\", \"type\": \"double\"},\n",
    "        {\"name\": \"HoursPerWeek\", \"type\": \"double\"},\n",
    "        {\"name\": \"NativeCountry\", \"type\": \"string\"},\n",
    "    ],\n",
    "    \"output\": {\"name\": \"prediction\", \"type\": \"double\"},\n",
    "}\n",
    "\n",
    "# Convert to JSON object\n",
    "schema_json = json.dumps(schema)\n",
    "print(schema_json)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "f0aefb30",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6eb174eccc234d2788deeac7e6d2fefe",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Set the region to deploy the pipeline on AWS\n",
    "\n",
    "os.environ[\"AWS_DEFAULT_REGION\"] = 'us-east-2'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "9e7f7be1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "28e5d426afd443ffaeacff69227b3115",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-------------!<sagemaker.sparkml.model.SparkMLPredictor object at 0x7ff611792518>"
     ]
    }
   ],
   "source": [
    "from time import gmtime, strftime\n",
    "import time\n",
    "\n",
    "timestamp_prefix = strftime(\"%Y-%m-%d-%H-%M-%S\", gmtime())\n",
    "\n",
    "import sagemaker\n",
    "from sagemaker import get_execution_role\n",
    "from sagemaker.sparkml.model import SparkMLModel\n",
    "\n",
    "sess = sagemaker.Session()\n",
    "role = get_execution_role()\n",
    "\n",
    "# S3 location of where you uploaded your trained and serialized SparkML model\n",
    "sparkml_data = \"s3://{}/{}/{}\".format(\n",
    "    \"layer-bootstrap-script\", \"emr/sparkml-layer/mleap\", \"model.tar.gz\"\n",
    ")\n",
    "model_name = \"sparkml-layer-\" + timestamp_prefix\n",
    "sparkml_model = SparkMLModel(\n",
    "    model_data=sparkml_data,\n",
    "    role=role,\n",
    "    sagemaker_session=sess,\n",
    "    name=model_name,\n",
    "    # passing the schema defined above by using an environment\n",
    "    # variable that sagemaker-sparkml-serving understands\n",
    "    env={\"SAGEMAKER_SPARKML_SCHEMA\": schema_json},\n",
    ")\n",
    "\n",
    "\n",
    "endpoint_name = \"sparkml-layer-ep-\" + timestamp_prefix\n",
    "\n",
    "# You can deploy to multiple instances for high availability and distributed ML\n",
    "sparkml_model.deploy(\n",
    "    initial_instance_count=1, instance_type=\"ml.c4.xlarge\", endpoint_name=endpoint_name\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "909b89ea",
   "metadata": {},
   "source": [
    "## Test the deployed pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bc68bb2",
   "metadata": {},
   "source": [
    "Send payload (data) to the pipeline endpoint in CSV format..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "f786e5c7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "545170cb9a0f4e9b85eca69a93c56631",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "b'0.0'"
     ]
    }
   ],
   "source": [
    "from sagemaker.predictor import (\n",
    "    json_serializer,\n",
    "    csv_serializer,\n",
    "    json_deserializer,\n",
    "    Predictor,\n",
    ")\n",
    "\n",
    "payload = \"17.0, Local-gov, 10th, 6.0, Never-married, Protective-serv, \\\n",
    "            Own-child, White, Female, 0.0, 1602.0, 40.0, United-States\"\n",
    "\n",
    "predictor = Predictor(\n",
    "    endpoint_name=endpoint_name,\n",
    "    sagemaker_session=sess,\n",
    "    serializer=csv_serializer,\n",
    "#     deserializer=csv_serializer\n",
    "#     content_type=\"text/csv\",\n",
    "#     accept=\"text/csv\",\n",
    ")\n",
    "print(predictor.predict(payload))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95ef7ce7",
   "metadata": {},
   "source": [
    "Send payload (data) to the pipeline as a JSON object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "34822a5c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "de06c1d69e914bd6af30d7c9e1469027",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "b'0.0'"
     ]
    }
   ],
   "source": [
    "payload = {\"data\": [17.0, \"Local-gov\", \"10th\", 6.0, \"Never-married\", \"Protective-serv\", \\\n",
    "                    \"Own-child\", \"White\", \"Female\", 0.0, 1602.0, 40.0, \"United-States\"]}\n",
    "predictor = Predictor(\n",
    "    endpoint_name=endpoint_name,\n",
    "    sagemaker_session=sess,\n",
    "    serializer=json_serializer\n",
    ")\n",
    "\n",
    "print(predictor.predict(payload))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00fd9037",
   "metadata": {},
   "source": [
    "Send payload with schema changes to the data... The endpoint can overwrite existing schema and use the data schema from the client application."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "26a26962",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "299b4b796fb24af982b6dde4dba6108c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatProgress(value=0.0, bar_style='info', description='Progress:', layout=Layout(height='25px', width='50%'),…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "b'0.0'"
     ]
    }
   ],
   "source": [
    "payload = {\n",
    "    \"schema\" : {\n",
    "        \"input\": [\n",
    "            {\"name\": \"Age\", \"type\": \"double\"},\n",
    "            {\"name\": \"WorkClass\", \"type\": \"string\"},\n",
    "            {\"name\": \"Education\", \"type\": \"string\"},\n",
    "            {\"name\": \"EducationNum\", \"type\": \"double\"},\n",
    "            {\"name\": \"MaritalStatus\", \"type\": \"string\"},\n",
    "            {\"name\": \"Occupation\", \"type\": \"string\"},\n",
    "            {\"name\": \"Relationship\", \"type\": \"string\"},\n",
    "            {\"name\": \"Race\", \"type\": \"string\"},\n",
    "            {\"name\": \"Gender\", \"type\": \"string\"},\n",
    "            {\"name\": \"CapitalGain\", \"type\": \"double\"},\n",
    "            {\"name\": \"CapitalLoss\", \"type\": \"double\"},\n",
    "            {\"name\": \"HoursPerWeek\", \"type\": \"double\"},\n",
    "            {\"name\": \"NativeCountry\", \"type\": \"string\"},\n",
    "        ],\n",
    "        \"output\": {\"name\": \"prediction\", \"type\": \"double\"},\n",
    "    },\n",
    "    \"data\": [17.0, \"Local-gov\", \"10th\", 6.0, \"Never-married\", \"Protective-serv\", \\\n",
    "                    \"Own-child\", \"White\", \"Female\", 0.0, 1602.0, 40.0, \"United-States\"],\n",
    "}\n",
    "\n",
    "predictor = Predictor(\n",
    "    endpoint_name=endpoint_name,\n",
    "    sagemaker_session=sess,\n",
    "    serializer=json_serializer\n",
    ")\n",
    "\n",
    "print(predictor.predict(payload))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "319c89ee",
   "metadata": {},
   "source": [
    "Learn how to invoke the endpoint from your application, backend server, or computer in [this reference](https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_runtime_InvokeEndpoint.html). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d26318f9",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Sparkmagic (PySpark)",
   "language": "python",
   "name": "pysparkkernel"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "python",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "pyspark",
   "pygments_lexer": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}