Predict-salary-python-spark.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Predicting Salary with Python and Pyspark"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this project, I attempt to predict whether an adult's income exceeds $50,000 per year based upon census data. The data is in the Data folder. You can also download it from the <a href=\"http://archive.ics.uci.edu/ml/datasets/Adult\">UC Irvine Machine Learning Repository</a>."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The header is not included, so we defined the names ourselves as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "column_names = ['age','workclass','fnlwgt','education','education-num','marital-status','occupation',\n",
    "                'relationship','race','sex','capital-gain','capital-loss','hours-per-week','native-country',\n",
    "                'salary']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load and prepare the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df = pd.read_csv('Data/adult.data', names = column_names)\n",
    "test_df = pd.read_csv('Data/adult.test', names = column_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Drop first row which is not useful and reset the index\n",
    "test_df = test_df.drop(test_df.index[0])\n",
    "test_df = test_df.reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'pandas.core.frame.DataFrame'>\n",
      "RangeIndex: 32561 entries, 0 to 32560\n",
      "Data columns (total 15 columns):\n",
      " #   Column          Non-Null Count  Dtype \n",
      "---  ------          --------------  ----- \n",
      " 0   age             32561 non-null  int64 \n",
      " 1   workclass       32561 non-null  object\n",
      " 2   fnlwgt          32561 non-null  int64 \n",
      " 3   education       32561 non-null  object\n",
      " 4   education-num   32561 non-null  int64 \n",
      " 5   marital-status  32561 non-null  object\n",
      " 6   occupation      32561 non-null  object\n",
      " 7   relationship    32561 non-null  object\n",
      " 8   race            32561 non-null  object\n",
      " 9   sex             32561 non-null  object\n",
      " 10  capital-gain    32561 non-null  int64 \n",
      " 11  capital-loss    32561 non-null  int64 \n",
      " 12  hours-per-week  32561 non-null  int64 \n",
      " 13  native-country  32561 non-null  object\n",
      " 14  salary          32561 non-null  object\n",
      "dtypes: int64(6), object(9)\n",
      "memory usage: 3.7+ MB\n"
     ]
    }
   ],
   "source": [
    "train_df.info()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'pandas.core.frame.DataFrame'>\n",
      "RangeIndex: 16281 entries, 0 to 16280\n",
      "Data columns (total 15 columns):\n",
      " #   Column          Non-Null Count  Dtype  \n",
      "---  ------          --------------  -----  \n",
      " 0   age             16281 non-null  object \n",
      " 1   workclass       16281 non-null  object \n",
      " 2   fnlwgt          16281 non-null  float64\n",
      " 3   education       16281 non-null  object \n",
      " 4   education-num   16281 non-null  float64\n",
      " 5   marital-status  16281 non-null  object \n",
      " 6   occupation      16281 non-null  object \n",
      " 7   relationship    16281 non-null  object \n",
      " 8   race            16281 non-null  object \n",
      " 9   sex             16281 non-null  object \n",
      " 10  capital-gain    16281 non-null  float64\n",
      " 11  capital-loss    16281 non-null  float64\n",
      " 12  hours-per-week  16281 non-null  float64\n",
      " 13  native-country  16281 non-null  object \n",
      " 14  salary          16281 non-null  object \n",
      "dtypes: float64(5), object(10)\n",
      "memory usage: 1.9+ MB\n"
     ]
    }
   ],
   "source": [
    "test_df.info()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The test data contains columns with data types not matching the training data. We have to make the data types match."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Change the affected columns to integer type\n",
    "test_df['age'] = test_df['age'].astype(int)\n",
    "\n",
    "float_cols = test_df.select_dtypes(include=['float64']).columns\n",
    "test_df[float_cols] = test_df[float_cols].astype(int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "object    9\n",
       "int64     6\n",
       "dtype: int64"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_df.dtypes.value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "object    9\n",
       "int64     6\n",
       "dtype: int64"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_df.dtypes.value_counts()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Also, the testing set does not contain Holand-Netherlands which is present in the training set. We will therefore remove it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df = train_df.apply(lambda x: x.astype(str).str.strip() if x.dtype == 'object' else x)\n",
    "\n",
    "# Make a copy\n",
    "train_df_cp = train_df.copy()\n",
    "\n",
    "# Filter out the rows with 'Holand-Netherlands'\n",
    "train_df = train_df.loc[train_df_cp['native-country'] != 'Holand-Netherlands']\n",
    "\n",
    "# Save to a CSV file\n",
    "train_df.to_csv('Data/train.csv', index=False, header=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also save the test data to a CSV file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_df = test_df.apply(lambda x: x.astype(str).str.strip() if x.dtype == 'object' else x)\n",
    "\n",
    "test_df.to_csv('Data/test.csv', index=False, header=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Visualize the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training data shape:  (32560, 15)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>age</th>\n",
       "      <th>workclass</th>\n",
       "      <th>fnlwgt</th>\n",
       "      <th>education</th>\n",
       "      <th>education-num</th>\n",
       "      <th>marital-status</th>\n",
       "      <th>occupation</th>\n",
       "      <th>relationship</th>\n",
       "      <th>race</th>\n",
       "      <th>sex</th>\n",
       "      <th>capital-gain</th>\n",
       "      <th>capital-loss</th>\n",
       "      <th>hours-per-week</th>\n",
       "      <th>native-country</th>\n",
       "      <th>salary</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>39</td>\n",
       "      <td>State-gov</td>\n",
       "      <td>77516</td>\n",
       "      <td>Bachelors</td>\n",
       "      <td>13</td>\n",
       "      <td>Never-married</td>\n",
       "      <td>Adm-clerical</td>\n",
       "      <td>Not-in-family</td>\n",
       "      <td>White</td>\n",
       "      <td>Male</td>\n",
       "      <td>2174</td>\n",
       "      <td>0</td>\n",
       "      <td>40</td>\n",
       "      <td>United-States</td>\n",
       "      <td>&lt;=50K</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>50</td>\n",
       "      <td>Self-emp-not-inc</td>\n",
       "      <td>83311</td>\n",
       "      <td>Bachelors</td>\n",
       "      <td>13</td>\n",
       "      <td>Married-civ-spouse</td>\n",
       "      <td>Exec-managerial</td>\n",
       "      <td>Husband</td>\n",
       "      <td>White</td>\n",
       "      <td>Male</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>13</td>\n",
       "      <td>United-States</td>\n",
       "      <td>&lt;=50K</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>38</td>\n",
       "      <td>Private</td>\n",
       "      <td>215646</td>\n",
       "      <td>HS-grad</td>\n",
       "      <td>9</td>\n",
       "      <td>Divorced</td>\n",
       "      <td>Handlers-cleaners</td>\n",
       "      <td>Not-in-family</td>\n",
       "      <td>White</td>\n",
       "      <td>Male</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>40</td>\n",
       "      <td>United-States</td>\n",
       "      <td>&lt;=50K</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   age         workclass  fnlwgt  education  education-num  \\\n",
       "0   39         State-gov   77516  Bachelors             13   \n",
       "1   50  Self-emp-not-inc   83311  Bachelors             13   \n",
       "2   38           Private  215646    HS-grad              9   \n",
       "\n",
       "       marital-status         occupation   relationship   race   sex  \\\n",
       "0       Never-married       Adm-clerical  Not-in-family  White  Male   \n",
       "1  Married-civ-spouse    Exec-managerial        Husband  White  Male   \n",
       "2            Divorced  Handlers-cleaners  Not-in-family  White  Male   \n",
       "\n",
       "   capital-gain  capital-loss  hours-per-week native-country salary  \n",
       "0          2174             0              40  United-States  <=50K  \n",
       "1             0             0              13  United-States  <=50K  \n",
       "2             0             0              40  United-States  <=50K  "
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print('Training data shape: ', train_df.shape)\n",
    "train_df.head(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing data shape:  (16281, 15)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>age</th>\n",
       "      <th>workclass</th>\n",
       "      <th>fnlwgt</th>\n",
       "      <th>education</th>\n",
       "      <th>education-num</th>\n",
       "      <th>marital-status</th>\n",
       "      <th>occupation</th>\n",
       "      <th>relationship</th>\n",
       "      <th>race</th>\n",
       "      <th>sex</th>\n",
       "      <th>capital-gain</th>\n",
       "      <th>capital-loss</th>\n",
       "      <th>hours-per-week</th>\n",
       "      <th>native-country</th>\n",
       "      <th>salary</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>25</td>\n",
       "      <td>Private</td>\n",
       "      <td>226802</td>\n",
       "      <td>11th</td>\n",
       "      <td>7</td>\n",
       "      <td>Never-married</td>\n",
       "      <td>Machine-op-inspct</td>\n",
       "      <td>Own-child</td>\n",
       "      <td>Black</td>\n",
       "      <td>Male</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>40</td>\n",
       "      <td>United-States</td>\n",
       "      <td>&lt;=50K.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>38</td>\n",
       "      <td>Private</td>\n",
       "      <td>89814</td>\n",
       "      <td>HS-grad</td>\n",
       "      <td>9</td>\n",
       "      <td>Married-civ-spouse</td>\n",
       "      <td>Farming-fishing</td>\n",
       "      <td>Husband</td>\n",
       "      <td>White</td>\n",
       "      <td>Male</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>50</td>\n",
       "      <td>United-States</td>\n",
       "      <td>&lt;=50K.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>28</td>\n",
       "      <td>Local-gov</td>\n",
       "      <td>336951</td>\n",
       "      <td>Assoc-acdm</td>\n",
       "      <td>12</td>\n",
       "      <td>Married-civ-spouse</td>\n",
       "      <td>Protective-serv</td>\n",
       "      <td>Husband</td>\n",
       "      <td>White</td>\n",
       "      <td>Male</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>40</td>\n",
       "      <td>United-States</td>\n",
       "      <td>&gt;50K.</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   age  workclass  fnlwgt   education  education-num      marital-status  \\\n",
       "0   25    Private  226802        11th              7       Never-married   \n",
       "1   38    Private   89814     HS-grad              9  Married-civ-spouse   \n",
       "2   28  Local-gov  336951  Assoc-acdm             12  Married-civ-spouse   \n",
       "\n",
       "          occupation relationship   race   sex  capital-gain  capital-loss  \\\n",
       "0  Machine-op-inspct    Own-child  Black  Male             0             0   \n",
       "1    Farming-fishing      Husband  White  Male             0             0   \n",
       "2    Protective-serv      Husband  White  Male             0             0   \n",
       "\n",
       "   hours-per-week native-country  salary  \n",
       "0              40  United-States  <=50K.  \n",
       "1              50  United-States  <=50K.  \n",
       "2              40  United-States   >50K.  "
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print('Testing data shape: ', test_df.shape)\n",
    "test_df.head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check for missing values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "age               0\n",
       "workclass         0\n",
       "fnlwgt            0\n",
       "education         0\n",
       "education-num     0\n",
       "marital-status    0\n",
       "occupation        0\n",
       "relationship      0\n",
       "race              0\n",
       "sex               0\n",
       "capital-gain      0\n",
       "capital-loss      0\n",
       "hours-per-week    0\n",
       "native-country    0\n",
       "salary            0\n",
       "dtype: int64"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_df.isnull().sum()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's check if the training and testing columns match"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "workclass          9\n",
       "education         16\n",
       "marital-status     7\n",
       "occupation        15\n",
       "relationship       6\n",
       "race               5\n",
       "sex                2\n",
       "native-country    41\n",
       "salary             2\n",
       "dtype: int64"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_df.select_dtypes('object').apply(pd.Series.nunique, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "workclass          9\n",
       "education         16\n",
       "marital-status     7\n",
       "occupation        15\n",
       "relationship       6\n",
       "race               5\n",
       "sex                2\n",
       "native-country    41\n",
       "salary             2\n",
       "dtype: int64"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_df.select_dtypes('object').apply(pd.Series.nunique, axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We manually encode the target column to avoid creating two columns during one-hot encoding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df['salary'] = train_df['salary'].apply(lambda x: 0 if x == '<=50K' else 1)\n",
    "test_df['salary'] = test_df['salary'].apply(lambda x: 0 if x == '<=50K.' else 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df = pd.get_dummies(train_df)\n",
    "test_df = pd.get_dummies(test_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training features shape:  (32560, 108)\n",
      "Testing features shape:  (16281, 108)\n"
     ]
    }
   ],
   "source": [
    "print('Training features shape: ', train_df.shape)\n",
    "print('Testing features shape: ', test_df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Keep only features with columns in both train and test\n",
    "train_df,test_df = train_df.align(test_df, join = 'inner', axis = 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training features shape : (32560, 108)\n",
      "Testing features shape : (16281, 108)\n"
     ]
    }
   ],
   "source": [
    "print('Training features shape :', train_df.shape)\n",
    "print('Testing features shape :', test_df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>age</th>\n",
       "      <th>fnlwgt</th>\n",
       "      <th>education-num</th>\n",
       "      <th>capital-gain</th>\n",
       "      <th>capital-loss</th>\n",
       "      <th>hours-per-week</th>\n",
       "      <th>salary</th>\n",
       "      <th>workclass_?</th>\n",
       "      <th>workclass_Federal-gov</th>\n",
       "      <th>workclass_Local-gov</th>\n",
       "      <th>...</th>\n",
       "      <th>native-country_Portugal</th>\n",
       "      <th>native-country_Puerto-Rico</th>\n",
       "      <th>native-country_Scotland</th>\n",
       "      <th>native-country_South</th>\n",
       "      <th>native-country_Taiwan</th>\n",
       "      <th>native-country_Thailand</th>\n",
       "      <th>native-country_Trinadad&amp;Tobago</th>\n",
       "      <th>native-country_United-States</th>\n",
       "      <th>native-country_Vietnam</th>\n",
       "      <th>native-country_Yugoslavia</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>39</td>\n",
       "      <td>77516</td>\n",
       "      <td>13</td>\n",
       "      <td>2174</td>\n",
       "      <td>0</td>\n",
       "      <td>40</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>50</td>\n",
       "      <td>83311</td>\n",
       "      <td>13</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>13</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>38</td>\n",
       "      <td>215646</td>\n",
       "      <td>9</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>40</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>...</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>3 rows × 108 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "   age  fnlwgt  education-num  capital-gain  capital-loss  hours-per-week  \\\n",
       "0   39   77516             13          2174             0              40   \n",
       "1   50   83311             13             0             0              13   \n",
       "2   38  215646              9             0             0              40   \n",
       "\n",
       "   salary  workclass_?  workclass_Federal-gov  workclass_Local-gov  ...  \\\n",
       "0       0            0                      0                    0  ...   \n",
       "1       0            0                      0                    0  ...   \n",
       "2       0            0                      0                    0  ...   \n",
       "\n",
       "   native-country_Portugal  native-country_Puerto-Rico  \\\n",
       "0                        0                           0   \n",
       "1                        0                           0   \n",
       "2                        0                           0   \n",
       "\n",
       "   native-country_Scotland  native-country_South  native-country_Taiwan  \\\n",
       "0                        0                     0                      0   \n",
       "1                        0                     0                      0   \n",
       "2                        0                     0                      0   \n",
       "\n",
       "   native-country_Thailand  native-country_Trinadad&Tobago  \\\n",
       "0                        0                               0   \n",
       "1                        0                               0   \n",
       "2                        0                               0   \n",
       "\n",
       "   native-country_United-States  native-country_Vietnam  \\\n",
       "0                             1                       0   \n",
       "1                             1                       0   \n",
       "2                             1                       0   \n",
       "\n",
       "   native-country_Yugoslavia  \n",
       "0                          0  \n",
       "1                          0  \n",
       "2                          0  \n",
       "\n",
       "[3 rows x 108 columns]"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_df.head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Separate dependent and independent variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train = train_df.drop('salary', axis=1)\n",
    "y_train = train_df['salary']\n",
    "\n",
    "X_test = test_df.drop('salary', axis=1)\n",
    "y_test = test_df['salary']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sclaing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.preprocessing import MinMaxScaler\n",
    "\n",
    "scaler = MinMaxScaler(feature_range = (0,1))\n",
    "\n",
    "scaler.fit(X_train)\n",
    "X_train = scaler.transform(X_train)\n",
    "X_test = scaler.transform(X_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Logistic Regression model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.linear_model import LogisticRegression\n",
    "\n",
    "lr = LogisticRegression(max_iter=1000)\n",
    "lr.fit(X_train, y_train)\n",
    "\n",
    "lr_pred = lr.predict(X_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.8507462686567164"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from sklearn.metrics import accuracy_score\n",
    "\n",
    "accuracy_score(y_test, lr_pred)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using Spark"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import spark libraries\n",
    "from pyspark import SparkConf, SparkContext\n",
    "from pyspark.sql import SparkSession\n",
    "from pyspark.ml.classification import LogisticRegression\n",
    "from pyspark.ml.feature import OneHotEncoder, StringIndexer, VectorAssembler\n",
    "from pyspark.ml import Pipeline\n",
    "from pyspark.sql.types import StructType, StructField, IntegerType, StringType"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize our session\n",
    "\n",
    "spark = SparkSession.builder.appName('Predict Adult Salary').getOrCreate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the schema\n",
    "\n",
    "schema = StructType([\n",
    "    StructField(\"age\", IntegerType(), True),\n",
    "    StructField(\"workclass\", StringType(), True),\n",
    "    StructField(\"fnlwgt\", IntegerType(), True),\n",
    "    StructField(\"education\", StringType(), True),\n",
    "    StructField(\"education-num\", IntegerType(), True),\n",
    "    StructField(\"marital-status\", StringType(), True),\n",
    "    StructField(\"occupation\", StringType(), True),\n",
    "    StructField(\"relationship\", StringType(), True),\n",
    "    StructField(\"race\", StringType(), True),\n",
    "    StructField(\"sex\", StringType(), True),\n",
    "    StructField(\"capital-gain\", IntegerType(), True),\n",
    "    StructField(\"capital-loss\", IntegerType(), True),\n",
    "    StructField(\"hours-per-week\", IntegerType(), True),\n",
    "    StructField(\"native-country\", StringType(), True),\n",
    "    StructField(\"salary\", StringType(), True)\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the data\n",
    "\n",
    "train = spark.read.csv('Data/train.csv',header=False, schema=schema)\n",
    "\n",
    "test = spark.read.csv('Data/test.csv', header=False, schema=schema)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Row(age=39, workclass='State-gov', fnlwgt=77516, education='Bachelors', education-num=13, marital-status='Never-married', occupation='Adm-clerical', relationship='Not-in-family', race='White', sex='Male', capital-gain=2174, capital-loss=0, hours-per-week=40, native-country='United-States', salary='<=50K'),\n",
       " Row(age=50, workclass='Self-emp-not-inc', fnlwgt=83311, education='Bachelors', education-num=13, marital-status='Married-civ-spouse', occupation='Exec-managerial', relationship='Husband', race='White', sex='Male', capital-gain=0, capital-loss=0, hours-per-week=13, native-country='United-States', salary='<=50K'),\n",
       " Row(age=38, workclass='Private', fnlwgt=215646, education='HS-grad', education-num=9, marital-status='Divorced', occupation='Handlers-cleaners', relationship='Not-in-family', race='White', sex='Male', capital-gain=0, capital-loss=0, hours-per-week=40, native-country='United-States', salary='<=50K')]"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train.head(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>age</th>\n",
       "      <th>workclass</th>\n",
       "      <th>fnlwgt</th>\n",
       "      <th>education</th>\n",
       "      <th>education-num</th>\n",
       "      <th>marital-status</th>\n",
       "      <th>occupation</th>\n",
       "      <th>relationship</th>\n",
       "      <th>race</th>\n",
       "      <th>sex</th>\n",
       "      <th>capital-gain</th>\n",
       "      <th>capital-loss</th>\n",
       "      <th>hours-per-week</th>\n",
       "      <th>native-country</th>\n",
       "      <th>salary</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>39</td>\n",
       "      <td>State-gov</td>\n",
       "      <td>77516</td>\n",
       "      <td>Bachelors</td>\n",
       "      <td>13</td>\n",
       "      <td>Never-married</td>\n",
       "      <td>Adm-clerical</td>\n",
       "      <td>Not-in-family</td>\n",
       "      <td>White</td>\n",
       "      <td>Male</td>\n",
       "      <td>2174</td>\n",
       "      <td>0</td>\n",
       "      <td>40</td>\n",
       "      <td>United-States</td>\n",
       "      <td>&lt;=50K</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>50</td>\n",
       "      <td>Self-emp-not-inc</td>\n",
       "      <td>83311</td>\n",
       "      <td>Bachelors</td>\n",
       "      <td>13</td>\n",
       "      <td>Married-civ-spouse</td>\n",
       "      <td>Exec-managerial</td>\n",
       "      <td>Husband</td>\n",
       "      <td>White</td>\n",
       "      <td>Male</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>13</td>\n",
       "      <td>United-States</td>\n",
       "      <td>&lt;=50K</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>38</td>\n",
       "      <td>Private</td>\n",
       "      <td>215646</td>\n",
       "      <td>HS-grad</td>\n",
       "      <td>9</td>\n",
       "      <td>Divorced</td>\n",
       "      <td>Handlers-cleaners</td>\n",
       "      <td>Not-in-family</td>\n",
       "      <td>White</td>\n",
       "      <td>Male</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>40</td>\n",
       "      <td>United-States</td>\n",
       "      <td>&lt;=50K</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   age         workclass  fnlwgt  education  education-num  \\\n",
       "0   39         State-gov   77516  Bachelors             13   \n",
       "1   50  Self-emp-not-inc   83311  Bachelors             13   \n",
       "2   38           Private  215646    HS-grad              9   \n",
       "\n",
       "       marital-status         occupation   relationship   race   sex  \\\n",
       "0       Never-married       Adm-clerical  Not-in-family  White  Male   \n",
       "1  Married-civ-spouse    Exec-managerial        Husband  White  Male   \n",
       "2            Divorced  Handlers-cleaners  Not-in-family  White  Male   \n",
       "\n",
       "   capital-gain  capital-loss  hours-per-week native-country salary  \n",
       "0          2174             0              40  United-States  <=50K  \n",
       "1             0             0              13  United-States  <=50K  \n",
       "2             0             0              40  United-States  <=50K  "
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train.limit(3).toPandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>age</th>\n",
       "      <th>workclass</th>\n",
       "      <th>fnlwgt</th>\n",
       "      <th>education</th>\n",
       "      <th>education-num</th>\n",
       "      <th>marital-status</th>\n",
       "      <th>occupation</th>\n",
       "      <th>relationship</th>\n",
       "      <th>race</th>\n",
       "      <th>sex</th>\n",
       "      <th>capital-gain</th>\n",
       "      <th>capital-loss</th>\n",
       "      <th>hours-per-week</th>\n",
       "      <th>native-country</th>\n",
       "      <th>salary</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>25</td>\n",
       "      <td>Private</td>\n",
       "      <td>226802</td>\n",
       "      <td>11th</td>\n",
       "      <td>7</td>\n",
       "      <td>Never-married</td>\n",
       "      <td>Machine-op-inspct</td>\n",
       "      <td>Own-child</td>\n",
       "      <td>Black</td>\n",
       "      <td>Male</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>40</td>\n",
       "      <td>United-States</td>\n",
       "      <td>&lt;=50K.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>38</td>\n",
       "      <td>Private</td>\n",
       "      <td>89814</td>\n",
       "      <td>HS-grad</td>\n",
       "      <td>9</td>\n",
       "      <td>Married-civ-spouse</td>\n",
       "      <td>Farming-fishing</td>\n",
       "      <td>Husband</td>\n",
       "      <td>White</td>\n",
       "      <td>Male</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>50</td>\n",
       "      <td>United-States</td>\n",
       "      <td>&lt;=50K.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>28</td>\n",
       "      <td>Local-gov</td>\n",
       "      <td>336951</td>\n",
       "      <td>Assoc-acdm</td>\n",
       "      <td>12</td>\n",
       "      <td>Married-civ-spouse</td>\n",
       "      <td>Protective-serv</td>\n",
       "      <td>Husband</td>\n",
       "      <td>White</td>\n",
       "      <td>Male</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>40</td>\n",
       "      <td>United-States</td>\n",
       "      <td>&gt;50K.</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   age  workclass  fnlwgt   education  education-num      marital-status  \\\n",
       "0   25    Private  226802        11th              7       Never-married   \n",
       "1   38    Private   89814     HS-grad              9  Married-civ-spouse   \n",
       "2   28  Local-gov  336951  Assoc-acdm             12  Married-civ-spouse   \n",
       "\n",
       "          occupation relationship   race   sex  capital-gain  capital-loss  \\\n",
       "0  Machine-op-inspct    Own-child  Black  Male             0             0   \n",
       "1    Farming-fishing      Husband  White  Male             0             0   \n",
       "2    Protective-serv      Husband  White  Male             0             0   \n",
       "\n",
       "   hours-per-week native-country  salary  \n",
       "0              40  United-States  <=50K.  \n",
       "1              50  United-States  <=50K.  \n",
       "2              40  United-States   >50K.  "
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test.limit(3).toPandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['age',\n",
       " 'workclass',\n",
       " 'fnlwgt',\n",
       " 'education',\n",
       " 'education-num',\n",
       " 'marital-status',\n",
       " 'occupation',\n",
       " 'relationship',\n",
       " 'race',\n",
       " 'sex',\n",
       " 'capital-gain',\n",
       " 'capital-loss',\n",
       " 'hours-per-week',\n",
       " 'native-country',\n",
       " 'salary']"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train.columns"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Encode the categorical data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "cat_cols = ['workclass','education','marital-status','occupation','relationship','race','sex','native-country']\n",
    "\n",
    "indexers = [StringIndexer(inputCol=column, outputCol=column+\"-index\") for column in cat_cols]\n",
    "\n",
    "encoder = OneHotEncoder(\n",
    "    inputCols=[indexer.getOutputCol() for indexer in indexers], \n",
    "    outputCols=[\"{0}-encoded\".format(indexer.getOutputCol()) for indexer in indexers]\n",
    ")\n",
    "\n",
    "assembler = VectorAssembler(inputCols=encoder.getOutputCols(),outputCol=\"cat_features\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipeline = Pipeline(stages=indexers + [encoder, assembler])\n",
    "\n",
    "train = pipeline.fit(train).transform(train)\n",
    "\n",
    "test = pipeline.fit(test).transform(test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "root\n",
      " |-- age: integer (nullable = true)\n",
      " |-- workclass: string (nullable = true)\n",
      " |-- fnlwgt: integer (nullable = true)\n",
      " |-- education: string (nullable = true)\n",
      " |-- education-num: integer (nullable = true)\n",
      " |-- marital-status: string (nullable = true)\n",
      " |-- occupation: string (nullable = true)\n",
      " |-- relationship: string (nullable = true)\n",
      " |-- race: string (nullable = true)\n",
      " |-- sex: string (nullable = true)\n",
      " |-- capital-gain: integer (nullable = true)\n",
      " |-- capital-loss: integer (nullable = true)\n",
      " |-- hours-per-week: integer (nullable = true)\n",
      " |-- native-country: string (nullable = true)\n",
      " |-- salary: string (nullable = true)\n",
      " |-- workclass-index: double (nullable = false)\n",
      " |-- education-index: double (nullable = false)\n",
      " |-- marital-status-index: double (nullable = false)\n",
      " |-- occupation-index: double (nullable = false)\n",
      " |-- relationship-index: double (nullable = false)\n",
      " |-- race-index: double (nullable = false)\n",
      " |-- sex-index: double (nullable = false)\n",
      " |-- native-country-index: double (nullable = false)\n",
      " |-- workclass-index-encoded: vector (nullable = true)\n",
      " |-- occupation-index-encoded: vector (nullable = true)\n",
      " |-- marital-status-index-encoded: vector (nullable = true)\n",
      " |-- native-country-index-encoded: vector (nullable = true)\n",
      " |-- education-index-encoded: vector (nullable = true)\n",
      " |-- race-index-encoded: vector (nullable = true)\n",
      " |-- relationship-index-encoded: vector (nullable = true)\n",
      " |-- sex-index-encoded: vector (nullable = true)\n",
      " |-- cat_features: vector (nullable = true)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "train.printSchema()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---------------------------------------------------------------+\n",
      "|cat_features                                                   |\n",
      "+---------------------------------------------------------------+\n",
      "|(93,[4,10,24,32,44,48,52,53],[1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0])|\n",
      "|(93,[1,10,23,31,43,48,52,53],[1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0])|\n",
      "|(93,[0,8,25,38,44,48,52,53],[1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0]) |\n",
      "|(93,[0,13,23,38,43,49,52,53],[1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0])|\n",
      "|(93,[0,10,23,29,47,49,62],[1.0,1.0,1.0,1.0,1.0,1.0,1.0])       |\n",
      "|(93,[0,11,23,31,47,48,53],[1.0,1.0,1.0,1.0,1.0,1.0,1.0])       |\n",
      "|(93,[0,18,28,34,44,49,64],[1.0,1.0,1.0,1.0,1.0,1.0,1.0])       |\n",
      "|(93,[1,8,23,31,43,48,52,53],[1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0]) |\n",
      "|(93,[0,11,24,29,44,48,53],[1.0,1.0,1.0,1.0,1.0,1.0,1.0])       |\n",
      "|(93,[0,10,23,31,43,48,52,53],[1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0])|\n",
      "|(93,[0,9,23,31,43,49,52,53],[1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0]) |\n",
      "|(93,[4,10,23,29,43,50,52,61],[1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0])|\n",
      "|(93,[0,10,24,32,45,48,53],[1.0,1.0,1.0,1.0,1.0,1.0,1.0])       |\n",
      "|(93,[0,14,24,33,44,49,52,53],[1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0])|\n",
      "|(93,[0,12,23,30,43,50,52,55],[1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0])|\n",
      "|(93,[0,16,23,37,43,51,52,54],[1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0])|\n",
      "|(93,[1,8,24,39,45,48,52,53],[1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0]) |\n",
      "|(93,[0,8,24,35,46,48,52,53],[1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0]) |\n",
      "|(93,[0,13,23,33,43,48,52,53],[1.0,1.0,1.0,1.0,1.0,1.0,1.0,1.0])|\n",
      "|(93,[1,11,25,31,46,48,53],[1.0,1.0,1.0,1.0,1.0,1.0,1.0])       |\n",
      "+---------------------------------------------------------------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "train.select('cat_features').show(truncate=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### We combine our continuous variables with our categorical variables into a single column."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_cols = ['age', 'fnlwgt', 'education-num', 'capital-gain', 'capital-loss', 'hours-per-week']\n",
    "\n",
    "assembler = VectorAssembler(inputCols=['cat_features', *num_cols],outputCol='features')\n",
    "\n",
    "train = assembler.transform(train)\n",
    "\n",
    "test = assembler.transform(test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SparseVector(99, {4: 1.0, 10: 1.0, 24: 1.0, 32: 1.0, 44: 1.0, 48: 1.0, 52: 1.0, 53: 1.0, 93: 39.0, 94: 77516.0, 95: 13.0, 96: 2174.0, 98: 40.0})"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train.limit(5).toPandas()['features'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0    (1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, ...\n",
       "1    (1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, ...\n",
       "2    (0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, ...\n",
       "3    (1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, ...\n",
       "4    (0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, ...\n",
       "Name: features, dtype: object"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test.limit(5).toPandas()['features']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### We encode our target label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "indexer = StringIndexer(inputCol='salary', outputCol='label')\n",
    "\n",
    "train = indexer.fit(train).transform(train)\n",
    "\n",
    "test = indexer.fit(test).transform(test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------+-----+\n",
      "|            features|label|\n",
      "+--------------------+-----+\n",
      "|(99,[4,10,24,32,4...|  0.0|\n",
      "|(99,[1,10,23,31,4...|  0.0|\n",
      "|(99,[0,8,25,38,44...|  0.0|\n",
      "|(99,[0,13,23,38,4...|  0.0|\n",
      "|(99,[0,10,23,29,4...|  0.0|\n",
      "|(99,[0,11,23,31,4...|  0.0|\n",
      "|(99,[0,18,28,34,4...|  0.0|\n",
      "|(99,[1,8,23,31,43...|  1.0|\n",
      "|(99,[0,11,24,29,4...|  1.0|\n",
      "|(99,[0,10,23,31,4...|  1.0|\n",
      "|(99,[0,9,23,31,43...|  1.0|\n",
      "|(99,[4,10,23,29,4...|  1.0|\n",
      "|(99,[0,10,24,32,4...|  0.0|\n",
      "|(99,[0,14,24,33,4...|  0.0|\n",
      "|(99,[0,12,23,30,4...|  1.0|\n",
      "|(99,[0,16,23,37,4...|  0.0|\n",
      "|(99,[1,8,24,39,45...|  0.0|\n",
      "|(99,[0,8,24,35,46...|  0.0|\n",
      "|(99,[0,13,23,33,4...|  0.0|\n",
      "|(99,[1,11,25,31,4...|  1.0|\n",
      "+--------------------+-----+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "train.select(['features','label']).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Logistic Regression model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fit and train the model\n",
    "\n",
    "lr = LogisticRegression(featuresCol='features', labelCol='label')\n",
    "\n",
    "model = lr.fit(train)\n",
    "\n",
    "pred = model.transform(test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>label</th>\n",
       "      <th>prediction</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   label  prediction\n",
       "0    0.0         0.0\n",
       "1    0.0         0.0\n",
       "2    1.0         0.0\n",
       "3    1.0         1.0\n",
       "4    0.0         0.0\n",
       "5    0.0         0.0\n",
       "6    0.0         0.0\n",
       "7    1.0         1.0\n",
       "8    0.0         0.0\n",
       "9    0.0         0.0"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pred.limit(10).toPandas()[['label', 'prediction']]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Evaluate model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.ml.evaluation import BinaryClassificationEvaluator\n",
    "\n",
    "result = BinaryClassificationEvaluator(rawPredictionCol='prediction',labelCol='label')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "AUC = result.evaluate(pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AUC\n",
      "0.7442772306790945\n"
     ]
    }
   ],
   "source": [
    "print(\"AUC\")\n",
    "print(AUC)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}