mlops_demo_sparkml.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "142e3c27-8676-4154-91c3-9bb92bbcfb92",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Import Packages ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "collapsed": true,
     "inputWidgets": {},
     "nuid": "1cd79563-86d0-4f55-871c-fa16b4d3f66d",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001B[43mNote: you may need to restart the kernel using dbutils.library.restartPython() to use updated packages.\u001B[0m\nRequirement already satisfied: mlflow-skinny[databricks]>=2.4.1 in /databricks/python3/lib/python3.11/site-packages (2.11.1)\nRequirement already satisfied: importlib-metadata!=4.7.0,<8,>=3.7.0 in /databricks/python3/lib/python3.11/site-packages (from mlflow-skinny[databricks]>=2.4.1) (6.0.0)\nRequirement already satisfied: click<9,>=7.0 in /databricks/python3/lib/python3.11/site-packages (from mlflow-skinny[databricks]>=2.4.1) (8.0.4)\nRequirement already satisfied: pytz<2025 in /databricks/python3/lib/python3.11/site-packages (from mlflow-skinny[databricks]>=2.4.1) (2022.7)\nRequirement already satisfied: packaging<24 in /databricks/python3/lib/python3.11/site-packages (from mlflow-skinny[databricks]>=2.4.1) (23.2)\nRequirement already satisfied: entrypoints<1 in /databricks/python3/lib/python3.11/site-packages (from mlflow-skinny[databricks]>=2.4.1) (0.4)\nRequirement already satisfied: cloudpickle<4 in /databricks/python3/lib/python3.11/site-packages (from mlflow-skinny[databricks]>=2.4.1) (2.2.1)\nRequirement already satisfied: gitpython<4,>=3.1.9 in /databricks/python3/lib/python3.11/site-packages (from mlflow-skinny[databricks]>=2.4.1) (3.1.27)\nRequirement already satisfied: protobuf<5,>=3.12.0 in /databricks/python3/lib/python3.11/site-packages (from mlflow-skinny[databricks]>=2.4.1) (4.24.1)\nRequirement already satisfied: requests<3,>=2.17.3 in /databricks/python3/lib/python3.11/site-packages (from mlflow-skinny[databricks]>=2.4.1) (2.31.0)\nRequirement already satisfied: sqlparse<1,>=0.4.0 in /databricks/python3/lib/python3.11/site-packages (from mlflow-skinny[databricks]>=2.4.1) (0.4.2)\nRequirement already satisfied: pyyaml<7,>=5.1 in /databricks/python3/lib/python3.11/site-packages (from mlflow-skinny[databricks]>=2.4.1) (6.0)\nRequirement already satisfied: azure-storage-file-datalake>12 in /databricks/python3/lib/python3.11/site-packages (from mlflow-skinny[databricks]>=2.4.1) (12.14.0)\nRequirement already satisfied: botocore in /databricks/python3/lib/python3.11/site-packages (from mlflow-skinny[databricks]>=2.4.1) (1.34.39)\nRequirement already satisfied: boto3>1 in /databricks/python3/lib/python3.11/site-packages (from mlflow-skinny[databricks]>=2.4.1) (1.34.39)\nRequirement already satisfied: google-cloud-storage>=1.30.0 in /databricks/python3/lib/python3.11/site-packages (from mlflow-skinny[databricks]>=2.4.1) (2.11.0)\nRequirement already satisfied: typing-extensions>=4.3.0 in /databricks/python3/lib/python3.11/site-packages (from azure-storage-file-datalake>12->mlflow-skinny[databricks]>=2.4.1) (4.7.1)\nRequirement already satisfied: isodate>=0.6.1 in /databricks/python3/lib/python3.11/site-packages (from azure-storage-file-datalake>12->mlflow-skinny[databricks]>=2.4.1) (0.6.1)\nRequirement already satisfied: azure-core<2.0.0,>=1.28.0 in /databricks/python3/lib/python3.11/site-packages (from azure-storage-file-datalake>12->mlflow-skinny[databricks]>=2.4.1) (1.30.1)\nRequirement already satisfied: azure-storage-blob<13.0.0,>=12.19.0 in /databricks/python3/lib/python3.11/site-packages (from azure-storage-file-datalake>12->mlflow-skinny[databricks]>=2.4.1) (12.19.1)\nRequirement already satisfied: jmespath<2.0.0,>=0.7.1 in /databricks/python3/lib/python3.11/site-packages (from boto3>1->mlflow-skinny[databricks]>=2.4.1) (0.10.0)\nRequirement already satisfied: s3transfer<0.11.0,>=0.10.0 in /databricks/python3/lib/python3.11/site-packages (from boto3>1->mlflow-skinny[databricks]>=2.4.1) (0.10.0)\nRequirement already satisfied: python-dateutil<3.0.0,>=2.1 in /databricks/python3/lib/python3.11/site-packages (from botocore->mlflow-skinny[databricks]>=2.4.1) (2.8.2)\nRequirement already satisfied: urllib3<2.1,>=1.25.4 in /databricks/python3/lib/python3.11/site-packages (from botocore->mlflow-skinny[databricks]>=2.4.1) (1.26.16)\nRequirement already satisfied: gitdb<5,>=4.0.1 in /databricks/python3/lib/python3.11/site-packages (from gitpython<4,>=3.1.9->mlflow-skinny[databricks]>=2.4.1) (4.0.11)\nRequirement already satisfied: google-resumable-media>=2.6.0 in /databricks/python3/lib/python3.11/site-packages (from google-cloud-storage>=1.30.0->mlflow-skinny[databricks]>=2.4.1) (2.7.0)\nRequirement already satisfied: google-api-core!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.0,<3.0.0dev,>=1.31.5 in /databricks/python3/lib/python3.11/site-packages (from google-cloud-storage>=1.30.0->mlflow-skinny[databricks]>=2.4.1) (2.17.1)\nRequirement already satisfied: google-auth<3.0dev,>=1.25.0 in /databricks/python3/lib/python3.11/site-packages (from google-cloud-storage>=1.30.0->mlflow-skinny[databricks]>=2.4.1) (2.21.0)\nRequirement already satisfied: google-cloud-core<3.0dev,>=2.3.0 in /databricks/python3/lib/python3.11/site-packages (from google-cloud-storage>=1.30.0->mlflow-skinny[databricks]>=2.4.1) (2.4.1)\nRequirement already satisfied: zipp>=0.5 in /databricks/python3/lib/python3.11/site-packages (from importlib-metadata!=4.7.0,<8,>=3.7.0->mlflow-skinny[databricks]>=2.4.1) (3.11.0)\nRequirement already satisfied: idna<4,>=2.5 in /databricks/python3/lib/python3.11/site-packages (from requests<3,>=2.17.3->mlflow-skinny[databricks]>=2.4.1) (3.4)\nRequirement already satisfied: certifi>=2017.4.17 in /databricks/python3/lib/python3.11/site-packages (from requests<3,>=2.17.3->mlflow-skinny[databricks]>=2.4.1) (2023.7.22)\nRequirement already satisfied: charset-normalizer<4,>=2 in /databricks/python3/lib/python3.11/site-packages (from requests<3,>=2.17.3->mlflow-skinny[databricks]>=2.4.1) (2.0.4)\nRequirement already satisfied: six>=1.11.0 in /usr/lib/python3/dist-packages (from azure-core<2.0.0,>=1.28.0->azure-storage-file-datalake>12->mlflow-skinny[databricks]>=2.4.1) (1.16.0)\nRequirement already satisfied: cryptography>=2.1.4 in /databricks/python3/lib/python3.11/site-packages (from azure-storage-blob<13.0.0,>=12.19.0->azure-storage-file-datalake>12->mlflow-skinny[databricks]>=2.4.1) (41.0.3)\nRequirement already satisfied: smmap<6,>=3.0.1 in /databricks/python3/lib/python3.11/site-packages (from gitdb<5,>=4.0.1->gitpython<4,>=3.1.9->mlflow-skinny[databricks]>=2.4.1) (5.0.0)\nRequirement already satisfied: googleapis-common-protos<2.0.dev0,>=1.56.2 in /databricks/python3/lib/python3.11/site-packages (from google-api-core!=2.0.*,!=2.1.*,!=2.2.*,!=2.3.0,<3.0.0dev,>=1.31.5->google-cloud-storage>=1.30.0->mlflow-skinny[databricks]>=2.4.1) (1.63.0)\nRequirement already satisfied: pyasn1-modules>=0.2.1 in /databricks/python3/lib/python3.11/site-packages (from google-auth<3.0dev,>=1.25.0->google-cloud-storage>=1.30.0->mlflow-skinny[databricks]>=2.4.1) (0.2.8)\nRequirement already satisfied: cachetools<6.0,>=2.0.0 in /databricks/python3/lib/python3.11/site-packages (from google-auth<3.0dev,>=1.25.0->google-cloud-storage>=1.30.0->mlflow-skinny[databricks]>=2.4.1) (5.3.3)\nRequirement already satisfied: rsa<5,>=3.1.4 in /databricks/python3/lib/python3.11/site-packages (from google-auth<3.0dev,>=1.25.0->google-cloud-storage>=1.30.0->mlflow-skinny[databricks]>=2.4.1) (4.9)\nRequirement already satisfied: google-crc32c<2.0dev,>=1.0 in /databricks/python3/lib/python3.11/site-packages (from google-resumable-media>=2.6.0->google-cloud-storage>=1.30.0->mlflow-skinny[databricks]>=2.4.1) (1.5.0)\nRequirement already satisfied: cffi>=1.12 in /databricks/python3/lib/python3.11/site-packages (from cryptography>=2.1.4->azure-storage-blob<13.0.0,>=12.19.0->azure-storage-file-datalake>12->mlflow-skinny[databricks]>=2.4.1) (1.15.1)\nRequirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /databricks/python3/lib/python3.11/site-packages (from pyasn1-modules>=0.2.1->google-auth<3.0dev,>=1.25.0->google-cloud-storage>=1.30.0->mlflow-skinny[databricks]>=2.4.1) (0.4.8)\nRequirement already satisfied: pycparser in /databricks/python3/lib/python3.11/site-packages (from cffi>=1.12->cryptography>=2.1.4->azure-storage-blob<13.0.0,>=12.19.0->azure-storage-file-datalake>12->mlflow-skinny[databricks]>=2.4.1) (2.21)\n\u001B[43mNote: you may need to restart the kernel using dbutils.library.restartPython() to use updated packages.\u001B[0m\n"
     ]
    }
   ],
   "source": [
    "# Comment code below after running the mlflow update\n",
    "%pip install \"mlflow-skinny[databricks]>=2.4.1\"\n",
    "dbutils.library.restartPython()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "cb1051bc-ea06-41c5-abd1-cad201f3953c",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from pyspark.ml.classification import RandomForestClassifier\n",
    "from pyspark.ml.evaluation import MulticlassClassificationEvaluator\n",
    "from pyspark.ml.feature import VectorAssembler\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "import mlflow\n",
    "from mlflow.models.signature import infer_signature\n",
    "from pyspark.ml import Pipeline\n",
    "from pyspark.sql.functions import col\n",
    "from mlflow import MlflowClient"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "52c29ff7-4183-4b71-89af-411901fb08a9",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Define Catalog and Schema in Unity Catalog ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "a12b5cab-945a-448f-a7d9-3d625ec756b2",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "catalog = \"data_science\"\n",
    "schema = \"models\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "d3040d68-06c1-48ce-8e8a-9b98e25b2042",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Load Kaggle Data Set ###"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "e7f278b7-a4be-49bc-92af-07eec9608aec",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "https://www.kaggle.com/datasets/fedesoriano/stroke-prediction-dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "89529bfd-720e-4869-bc2e-2062c6921dc7",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "data = pd.read_csv('/Workspace/Users/awnish.choudhary@anthology.ai/healthcare-dataset-stroke-data.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "1a3b8d88-5a58-40aa-8aee-d1b3e4d7268b",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Clean the Data ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "d975d001-f203-4b95-a777-4e9b79285bf1",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [
    {
     "output_type": "execute_result",
     "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>id</th>\n",
       "      <th>gender</th>\n",
       "      <th>age</th>\n",
       "      <th>hypertension</th>\n",
       "      <th>heart_disease</th>\n",
       "      <th>ever_married</th>\n",
       "      <th>work_type</th>\n",
       "      <th>Residence_type</th>\n",
       "      <th>avg_glucose_level</th>\n",
       "      <th>bmi</th>\n",
       "      <th>smoking_status</th>\n",
       "      <th>stroke</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>9046</td>\n",
       "      <td>0</td>\n",
       "      <td>67.0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>228.69</td>\n",
       "      <td>36.6</td>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>51676</td>\n",
       "      <td>1</td>\n",
       "      <td>61.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>202.21</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>31112</td>\n",
       "      <td>0</td>\n",
       "      <td>80.0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>105.92</td>\n",
       "      <td>32.5</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>60182</td>\n",
       "      <td>1</td>\n",
       "      <td>49.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>171.23</td>\n",
       "      <td>34.4</td>\n",
       "      <td>3</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1665</td>\n",
       "      <td>1</td>\n",
       "      <td>79.0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>174.12</td>\n",
       "      <td>24.0</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5105</th>\n",
       "      <td>18234</td>\n",
       "      <td>1</td>\n",
       "      <td>80.0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>83.75</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5106</th>\n",
       "      <td>44873</td>\n",
       "      <td>1</td>\n",
       "      <td>81.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>125.20</td>\n",
       "      <td>40.0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5107</th>\n",
       "      <td>19723</td>\n",
       "      <td>1</td>\n",
       "      <td>35.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>82.99</td>\n",
       "      <td>30.6</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5108</th>\n",
       "      <td>37544</td>\n",
       "      <td>0</td>\n",
       "      <td>51.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>166.29</td>\n",
       "      <td>25.6</td>\n",
       "      <td>2</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5109</th>\n",
       "      <td>44679</td>\n",
       "      <td>1</td>\n",
       "      <td>44.0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "      <td>0</td>\n",
       "      <td>85.28</td>\n",
       "      <td>26.2</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5110 rows × 12 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "         id  gender   age  ...   bmi  smoking_status  stroke\n",
       "0      9046       0  67.0  ...  36.6               2       1\n",
       "1     51676       1  61.0  ...   NaN               1       1\n",
       "2     31112       0  80.0  ...  32.5               1       1\n",
       "3     60182       1  49.0  ...  34.4               3       1\n",
       "4      1665       1  79.0  ...  24.0               1       1\n",
       "...     ...     ...   ...  ...   ...             ...     ...\n",
       "5105  18234       1  80.0  ...   NaN               1       0\n",
       "5106  44873       1  81.0  ...  40.0               1       0\n",
       "5107  19723       1  35.0  ...  30.6               1       0\n",
       "5108  37544       0  51.0  ...  25.6               2       0\n",
       "5109  44679       1  44.0  ...  26.2               0       0\n",
       "\n",
       "[5110 rows x 12 columns]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Convert string columns to numeric values in data\n",
    "data['gender'] = data['gender'].map({'Male': 0, 'Female': 1, 'Other': 2})\n",
    "data['ever_married'] = data['ever_married'].map({'No': 0, 'Yes': 1})\n",
    "data['work_type'] = data['work_type'].map({'Private': 0, 'Self-employed': 1, 'Govt_job': 2, 'children': 3, 'Never_worked': 4})\n",
    "data['Residence_type'] = data['Residence_type'].map({'Urban': 0, 'Rural': 1})\n",
    "data['smoking_status'] = data['smoking_status'].map({'Unknown': 0, 'never smoked': 1, 'formerly smoked': 2, 'smokes': 3})\n",
    "\n",
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "fbed0b81-a6ca-4e53-8795-c442c2109a35",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "# Drop rows with missing values\n",
    "data.dropna(inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "3f2c52fd-7361-4aa5-a951-c545eed73a54",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "data = spark.createDataFrame(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "a27de549-1f4c-438b-92c9-aef2a5fcdf38",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "# Select features and target\n",
    "featureCols = [col for col in data.columns if col != \"stroke\"]\n",
    "assembler = VectorAssembler(inputCols=featureCols, outputCol=\"features\")\n",
    "data_prepared = assembler.transform(data).select(col(\"features\"), col(\"stroke\").alias(\"label\"))\n",
    "\n",
    "# Split the data\n",
    "(train_data, test_data) = data_prepared.randomSplit([0.7, 0.3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "acde63bf-5651-4494-9d22-c3d17b460b06",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Model training ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "3ae4064a-5c93-43b6-9fc0-e79d3a114b7b",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [
    {
     "output_type": "display_data",
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "79e922e247c34a62962c5c147a6f0cfb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading artifacts:   0%|          | 0/20 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "output_type": "display_data",
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5ab9df31cb0b4ba08e0f798edac17bc9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Uploading artifacts:   0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "output_type": "stream",
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Accuracy = 0.9591973244147157\n"
     ]
    }
   ],
   "source": [
    "# Train a RandomForest model\n",
    "rf = RandomForestClassifier(labelCol=\"label\", featuresCol=\"features\", numTrees=100)\n",
    "model = rf.fit(train_data)\n",
    "model_name = 'random_forest_classifier'\n",
    "# Make predictions\n",
    "predictions = model.transform(test_data)\n",
    "# Evaluate the model\n",
    "evaluator = MulticlassClassificationEvaluator(labelCol=\"label\", predictionCol=\"prediction\", metricName=\"accuracy\")\n",
    "accuracy = evaluator.evaluate(predictions)\n",
    "print(f\"Test Accuracy = {accuracy}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "bd95b153-89bc-4ee4-9f89-f6b4904e5fa1",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Add model signature which is required to register the model ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "4ce10a34-b3d3-48b2-a047-dc1d3e9aa1c4",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "# Sample Data\n",
    "sample = train_data.limit(5)\n",
    "# Convert Spark DataFrame to Pandas DataFrame\n",
    "train_data_pd = sample.toPandas()\n",
    "# Infer model signature using the Pandas DataFrame\n",
    "signature = infer_signature(train_data_pd, model.transform(sample).toPandas())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "ed5a9ccf-708b-436a-bba7-7ed786fdda90",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Log and Register model to MLflow ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "7f4ffd2d-d880-4e36-87e7-fdd658dd6342",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024/09/17 14:09:06 INFO mlflow.spark: Inferring pip requirements by reloading the logged model from the databricks artifact repository, which can be time-consuming. To speed up, explicitly specify the conda_env or pip_requirements when calling log_model().\n"
     ]
    },
    {
     "output_type": "display_data",
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2ddb8ce6db044698873ad55405118855",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading artifacts:   0%|          | 0/20 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024/09/17 14:09:08 INFO mlflow.store.artifact.artifact_repo: The progress bar can be disabled by setting the environment variable MLFLOW_ENABLE_ARTIFACTS_PROGRESS_BAR to false\n2024/09/17 14:09:35 WARNING mlflow.utils.environment: Encountered an unexpected error while inferring pip requirements (model URI: dbfs:/databricks/mlflow-tracking/2302940123439732/d153c62e099046188ceb9d7eaf82f303/artifacts/random_forest_classifier/sparkml, flavor: spark). Fall back to return ['pyspark==3.5.1']. Set logging level to DEBUG to see the full traceback. \n"
     ]
    },
    {
     "output_type": "display_data",
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "be81ad50b8884f089f24fbad2c7046e7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Uploading artifacts:   0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Successfully registered model 'data_science.models.random_forest_classifier'.\n2024/09/17 14:09:37 WARNING mlflow.store._unity_catalog.registry.rest_store: Unable to get model version source run's workspace ID from request headers. No run link will be recorded for the model version\n"
     ]
    },
    {
     "output_type": "display_data",
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "495d52602a0c4777aec39bb3be7220d0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading artifacts:   0%|          | 0/24 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024/09/17 14:09:40 INFO mlflow.store.artifact.artifact_repo: The progress bar can be disabled by setting the environment variable MLFLOW_ENABLE_ARTIFACTS_PROGRESS_BAR to false\n"
     ]
    },
    {
     "output_type": "display_data",
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "73a1d07ee15e446688090fe7f07836ff",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Uploading artifacts:   0%|          | 0/24 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024/09/17 14:09:41 INFO mlflow.store.artifact.cloud_artifact_repo: The progress bar can be disabled by setting the environment variable MLFLOW_ENABLE_ARTIFACTS_PROGRESS_BAR to false\nCreated version '1' of model 'data_science.models.random_forest_classifier'.\n"
     ]
    }
   ],
   "source": [
    "with mlflow.start_run():\n",
    "    mlflow.spark.log_model(model, model_name, signature=signature)\n",
    "    uri = mlflow.get_artifact_uri(model_name)\n",
    "    # Log metrics\n",
    "    mlflow.log_metric(\"accuracy_score\", accuracy)\n",
    "    mlflow.set_registry_uri(\"databricks-uc\")\n",
    "    # Register Model\n",
    "    mlflow.register_model(\n",
    "        model_uri=uri,\n",
    "        name=f\"{catalog}.{schema}.{model_name}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "b20c1f72-1cf7-4482-9a4e-3a7ba7c33b85",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Create \"Champion\" alias for latest version of the model currently in production  ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "c05772ef-6521-41db-bc89-d9801f7b7e14",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "# Initialize the MLflow client\n",
    "client = MlflowClient()\n",
    "\n",
    "# Search for all versions of the model and fetch the latest one\n",
    "model_version_infos = client.search_model_versions(f\"name='{catalog}.{schema}.{model_name}'\")\n",
    "new_model_version = max(model_version_info.version for model_version_info in model_version_infos)\n",
    "\n",
    "# Set the alias for the latest model version\n",
    "client.set_registered_model_alias(\n",
    "    name=f\"{catalog}.{schema}.{model_name}\",\n",
    "    alias=\"Champion\",\n",
    "    version=new_model_version\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "729e7426-4294-4d06-82de-e8e8d8fc2ca8",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Load Registered Model for inference ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "e158bb8a-aae3-4414-83e3-b6ec3fe58ecb",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024/09/17 14:09:43 INFO mlflow.spark: 'models:/data_science.models.random_forest_classifier@Champion' resolved as 's3://caden-os-prod-databricks-metastore-storage/7adaa701-ad6e-4968-af13-eddf695e53bc/models/a36649dc-65ca-4b5b-9d16-f23d100dc545/versions/e4350f3e-96f7-48cd-a6e1-a8188a955911'\n"
     ]
    },
    {
     "output_type": "display_data",
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "725624ec9db64dfb894b2a1438883db1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading artifacts:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "output_type": "display_data",
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6e29213467eb40be81ecde5af41e1742",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading artifacts:   0%|          | 0/24 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "output_type": "stream",
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024/09/17 14:09:46 INFO mlflow.store.artifact.artifact_repo: The progress bar can be disabled by setting the environment variable MLFLOW_ENABLE_ARTIFACTS_PROGRESS_BAR to false\n"
     ]
    }
   ],
   "source": [
    "model_version_uri = 'models:/'+f\"{catalog}.{schema}.{model_name}@Champion\"\n",
    "champion_version = mlflow.spark.load_model(model_version_uri)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "56c02c1d-3ef8-4741-b7d7-d9b5cb860e67",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [
    {
     "output_type": "display_data",
     "data": {
      "text/html": [
       "<style scoped>\n",
       "  .table-result-container {\n",
       "    max-height: 300px;\n",
       "    overflow: auto;\n",
       "  }\n",
       "  table, th, td {\n",
       "    border: 1px solid black;\n",
       "    border-collapse: collapse;\n",
       "  }\n",
       "  th, td {\n",
       "    padding: 5px;\n",
       "  }\n",
       "  th {\n",
       "    text-align: left;\n",
       "  }\n",
       "</style><div class='table-result-container'><table class='table-result'><thead style='background-color: white'><tr><th>features</th><th>label</th><th>rawPrediction</th><th>probability</th><th>prediction</th></tr></thead><tbody><tr><td>Map(vectorType -> sparse, length -> 11, indices -> List(0, 1, 2, 5, 8, 9), values -> List(15525.0, 1.0, 63.0, 1.0, 96.26, 31.8))</td><td>0</td><td>Map(vectorType -> dense, length -> 2, values -> List(96.84831067226501, 3.151689327734992))</td><td>Map(vectorType -> dense, length -> 2, values -> List(0.9684831067226501, 0.031516893277349925))</td><td>0.0</td></tr><tr><td>Map(vectorType -> sparse, length -> 11, indices -> List(0, 1, 2, 5, 8, 9), values -> List(26134.0, 1.0, 28.0, 1.0, 111.22, 25.5))</td><td>0</td><td>Map(vectorType -> dense, length -> 2, values -> List(97.8523946574489, 2.1476053425511097))</td><td>Map(vectorType -> dense, length -> 2, values -> List(0.9785239465744889, 0.021476053425511097))</td><td>0.0</td></tr><tr><td>Map(vectorType -> sparse, length -> 11, indices -> List(0, 1, 2, 5, 8, 9), values -> List(32257.0, 1.0, 47.0, 1.0, 210.95, 50.1))</td><td>0</td><td>Map(vectorType -> dense, length -> 2, values -> List(96.8129161270661, 3.1870838729338957))</td><td>Map(vectorType -> dense, length -> 2, values -> List(0.9681291612706611, 0.031870838729338956))</td><td>0.0</td></tr><tr><td>Map(vectorType -> sparse, length -> 11, indices -> List(0, 1, 2, 5, 8, 9), values -> List(32689.0, 1.0, 48.0, 1.0, 84.38, 27.1))</td><td>0</td><td>Map(vectorType -> dense, length -> 2, values -> List(97.74741060448416, 2.2525893955158467))</td><td>Map(vectorType -> dense, length -> 2, values -> List(0.9774741060448416, 0.022525893955158468))</td><td>0.0</td></tr><tr><td>Map(vectorType -> sparse, length -> 11, indices -> List(0, 1, 2, 5, 8, 9), values -> List(62783.0, 1.0, 76.0, 1.0, 198.02, 38.7))</td><td>0</td><td>Map(vectorType -> dense, length -> 2, values -> List(86.25338464331385, 13.746615356686148))</td><td>Map(vectorType -> dense, length -> 2, values -> List(0.8625338464331385, 0.13746615356686148))</td><td>0.0</td></tr></tbody></table></div>"
      ]
     },
     "metadata": {
      "application/vnd.databricks.v1+output": {
       "addedWidgets": {},
       "aggData": [],
       "aggError": "",
       "aggOverflow": false,
       "aggSchema": [],
       "aggSeriesLimitReached": false,
       "aggType": "",
       "arguments": {},
       "columnCustomDisplayInfos": {},
       "data": [
        [
         {
          "indices": [
           0,
           1,
           2,
           5,
           8,
           9
          ],
          "length": 11,
          "values": [
           15525.0,
           1.0,
           63.0,
           1.0,
           96.26,
           31.8
          ],
          "vectorType": "sparse"
         },
         0,
         {
          "length": 2,
          "values": [
           96.84831067226501,
           3.151689327734992
          ],
          "vectorType": "dense"
         },
         {
          "length": 2,
          "values": [
           0.9684831067226501,
           0.031516893277349925
          ],
          "vectorType": "dense"
         },
         0.0
        ],
        [
         {
          "indices": [
           0,
           1,
           2,
           5,
           8,
           9
          ],
          "length": 11,
          "values": [
           26134.0,
           1.0,
           28.0,
           1.0,
           111.22,
           25.5
          ],
          "vectorType": "sparse"
         },
         0,
         {
          "length": 2,
          "values": [
           97.8523946574489,
           2.1476053425511097
          ],
          "vectorType": "dense"
         },
         {
          "length": 2,
          "values": [
           0.9785239465744889,
           0.021476053425511097
          ],
          "vectorType": "dense"
         },
         0.0
        ],
        [
         {
          "indices": [
           0,
           1,
           2,
           5,
           8,
           9
          ],
          "length": 11,
          "values": [
           32257.0,
           1.0,
           47.0,
           1.0,
           210.95,
           50.1
          ],
          "vectorType": "sparse"
         },
         0,
         {
          "length": 2,
          "values": [
           96.8129161270661,
           3.1870838729338957
          ],
          "vectorType": "dense"
         },
         {
          "length": 2,
          "values": [
           0.9681291612706611,
           0.031870838729338956
          ],
          "vectorType": "dense"
         },
         0.0
        ],
        [
         {
          "indices": [
           0,
           1,
           2,
           5,
           8,
           9
          ],
          "length": 11,
          "values": [
           32689.0,
           1.0,
           48.0,
           1.0,
           84.38,
           27.1
          ],
          "vectorType": "sparse"
         },
         0,
         {
          "length": 2,
          "values": [
           97.74741060448416,
           2.2525893955158467
          ],
          "vectorType": "dense"
         },
         {
          "length": 2,
          "values": [
           0.9774741060448416,
           0.022525893955158468
          ],
          "vectorType": "dense"
         },
         0.0
        ],
        [
         {
          "indices": [
           0,
           1,
           2,
           5,
           8,
           9
          ],
          "length": 11,
          "values": [
           62783.0,
           1.0,
           76.0,
           1.0,
           198.02,
           38.7
          ],
          "vectorType": "sparse"
         },
         0,
         {
          "length": 2,
          "values": [
           86.25338464331385,
           13.746615356686148
          ],
          "vectorType": "dense"
         },
         {
          "length": 2,
          "values": [
           0.8625338464331385,
           0.13746615356686148
          ],
          "vectorType": "dense"
         },
         0.0
        ]
       ],
       "datasetInfos": [],
       "dbfsResultPath": null,
       "isJsonSchema": true,
       "metadata": {},
       "overflow": false,
       "plotOptions": {
        "customPlotOptions": {},
        "displayType": "table",
        "pivotAggregation": null,
        "pivotColumns": null,
        "xColumns": null,
        "yColumns": null
       },
       "removedWidgets": [],
       "schema": [
        {
         "metadata": "{\"ml_attr\":{\"attrs\":{\"numeric\":[{\"idx\":0,\"name\":\"id\"},{\"idx\":1,\"name\":\"gender\"},{\"idx\":2,\"name\":\"age\"},{\"idx\":3,\"name\":\"hypertension\"},{\"idx\":4,\"name\":\"heart_disease\"},{\"idx\":5,\"name\":\"ever_married\"},{\"idx\":6,\"name\":\"work_type\"},{\"idx\":7,\"name\":\"Residence_type\"},{\"idx\":8,\"name\":\"avg_glucose_level\"},{\"idx\":9,\"name\":\"bmi\"},{\"idx\":10,\"name\":\"smoking_status\"}]},\"num_attrs\":11}}",
         "name": "features",
         "type": "{\"type\":\"udt\",\"class\":\"org.apache.spark.ml.linalg.VectorUDT\",\"pyClass\":\"pyspark.ml.linalg.VectorUDT\",\"sqlType\":{\"type\":\"struct\",\"fields\":[{\"name\":\"type\",\"type\":\"byte\",\"nullable\":false,\"metadata\":{}},{\"name\":\"size\",\"type\":\"integer\",\"nullable\":true,\"metadata\":{}},{\"name\":\"indices\",\"type\":{\"type\":\"array\",\"elementType\":\"integer\",\"containsNull\":false},\"nullable\":true,\"metadata\":{}},{\"name\":\"values\",\"type\":{\"type\":\"array\",\"elementType\":\"double\",\"containsNull\":false},\"nullable\":true,\"metadata\":{}}]}}"
        },
        {
         "metadata": "{}",
         "name": "label",
         "type": "\"long\""
        },
        {
         "metadata": "{\"ml_attr\":{\"num_attrs\":2}}",
         "name": "rawPrediction",
         "type": "{\"type\":\"udt\",\"class\":\"org.apache.spark.ml.linalg.VectorUDT\",\"pyClass\":\"pyspark.ml.linalg.VectorUDT\",\"sqlType\":{\"type\":\"struct\",\"fields\":[{\"name\":\"type\",\"type\":\"byte\",\"nullable\":false,\"metadata\":{}},{\"name\":\"size\",\"type\":\"integer\",\"nullable\":true,\"metadata\":{}},{\"name\":\"indices\",\"type\":{\"type\":\"array\",\"elementType\":\"integer\",\"containsNull\":false},\"nullable\":true,\"metadata\":{}},{\"name\":\"values\",\"type\":{\"type\":\"array\",\"elementType\":\"double\",\"containsNull\":false},\"nullable\":true,\"metadata\":{}}]}}"
        },
        {
         "metadata": "{\"ml_attr\":{\"num_attrs\":2}}",
         "name": "probability",
         "type": "{\"type\":\"udt\",\"class\":\"org.apache.spark.ml.linalg.VectorUDT\",\"pyClass\":\"pyspark.ml.linalg.VectorUDT\",\"sqlType\":{\"type\":\"struct\",\"fields\":[{\"name\":\"type\",\"type\":\"byte\",\"nullable\":false,\"metadata\":{}},{\"name\":\"size\",\"type\":\"integer\",\"nullable\":true,\"metadata\":{}},{\"name\":\"indices\",\"type\":{\"type\":\"array\",\"elementType\":\"integer\",\"containsNull\":false},\"nullable\":true,\"metadata\":{}},{\"name\":\"values\",\"type\":{\"type\":\"array\",\"elementType\":\"double\",\"containsNull\":false},\"nullable\":true,\"metadata\":{}}]}}"
        },
        {
         "metadata": "{\"ml_attr\":{\"type\":\"nominal\",\"num_vals\":2}}",
         "name": "prediction",
         "type": "\"double\""
        }
       ],
       "type": "table"
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "display(champion_version.transform(test_data).limit(5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "1d7eb041-ca2e-4863-ba65-ee7b538e87cc",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "application/vnd.databricks.v1+notebook": {
   "dashboards": [],
   "environmentMetadata": null,
   "language": "python",
   "notebookMetadata": {
    "pythonIndentUnit": 4
   },
   "notebookName": "mlops_demo_sparkml",
   "widgets": {}
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}