UDF_UDTF_Examples.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f2d0f473-3e84-410e-b9f8-d8cc86dd946c",
   "metadata": {},
   "source": [
    "# UDF/UDTF Examples\n",
    "\n",
    "This notebook contains different examples of how to create UDF/UDTF using the Snowpark API. A lot of this notebook comes from Mats Stellwall in his repo at https://github.com/mstellwa/snowpark_examples. Other places for help include the Snowflake documentation at https://docs.snowflake.com/en/developer-guide/udf/python/udf-python-batch and Snowflake API documentation at https://docs.snowflake.com/ko/developer-guide/snowpark/reference/python/1.8.0/api/snowflake.snowpark.udf.UDFRegistration.\n",
    "\n",
    "\n",
    "There is threes diffrent types of UDFs  in this notebook:\n",
    "* UDF (Scalar User Defined Function)\n",
    "    * Is a scalar function that returns one output row for each input row. \n",
    "    * The returned row consists of a single column/value.\n",
    "    * Python UDF batch API enables defining UDFs that receive batches of input rows as Pandas DataFrames and return batches of results as Pandas arrays or Series\n",
    "* A Vectorized UDF - \n",
    "* UDTF (User Defined Tabular Function)\n",
    "    * A tabular function, also called a table function, returns zero, one, or multiple rows for each input row."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcd9b33e",
   "metadata": {},
   "source": [
    "## 1. Setup Environment \n",
    "I have not put all the imports for the notebook in this section, so it's easier to see what is needed for the UDF/UDTF examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0a241ea1-cfd9-41b1-bf59-ee17aad33860",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style>pre { white-space: pre !important; }</style>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using Snowpark: 1.11.1\n"
     ]
    }
   ],
   "source": [
    "# Make sure we do not get line breaks when doing show on wide dataframes\n",
    "from IPython.core.display import HTML\n",
    "display(HTML(\"<style>pre { white-space: pre !important; }</style>\"))\n",
    "\n",
    "# Snowpark imports \n",
    "import snowflake.snowpark as S\n",
    "from snowflake.snowpark import Session\n",
    "from snowflake.snowpark import functions as F\n",
    "from snowflake.snowpark import types as T\n",
    "from snowflake.snowpark import Window\n",
    "\n",
    "# Used for reading creds.json\n",
    "import json\n",
    "\n",
    "# Print the version of Snowpark we are using\n",
    "print(f\"Using Snowpark: {S.__version__}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b255d52c-2f3f-4978-a308-f1433e8f37b3",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Current role: \"RAJIV\", Current schema: \"RAJIV\".\"PUBLIC\", Current WH: \"RAJIV\"\n"
     ]
    }
   ],
   "source": [
    "with open('../creds.json') as f:\n",
    "    connection_parameters = json.load(f)\n",
    "\n",
    "session = Session.builder.configs(connection_parameters).create()\n",
    "print(\"Current role: \" + session.get_current_role() + \", Current schema: \" + session.get_fully_qualified_current_schema() + \", Current WH: \" + session.get_current_warehouse())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ede898e",
   "metadata": {},
   "source": [
    "## 2. UDF - Hello World Example\n",
    "Start by creating a UDF that returns a string, by setting *is_permanent=False* the UDF will only be avalible for our user and also only until the active Snowflake session is closed. The Function will be called for each input row ie it is not using the Batch API. By using **session.clear_imports()** and **session.clear_packages()** we make sure that old imports and packages are not included for the creation.\n",
    "\n",
    "A UDF can be created using the **@udf** decorator, the **udf** function or the **udf.register** method ofthe session object. It can be permanent or temporary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5b9288e1-9f74-4208-a1b6-75e886761042",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "session.clear_imports()\n",
    "session.clear_packages()\n",
    "@F.udf(name=\"hello_udf\", is_permanent=False, replace=True, session=session)\n",
    "def hello_udf(name: str) -> str:\n",
    "    return f'Hello {name}!'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3851e2f3-ea13-42f3-8a46-5d19e4340e32",
   "metadata": {},
   "source": [
    "Create a Snowpark DataFrame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "16477727-2e4e-4ed6-9559-179958ca0288",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------\n",
      "|\"NAME\"  |\n",
      "----------\n",
      "|Mats    |\n",
      "|Pia     |\n",
      "----------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "test_name_df = session.create_dataframe([['Mats'], ['Pia']], schema=[\"name\"])\n",
    "test_name_df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fd8c454",
   "metadata": {},
   "source": [
    "Test the function on the Snowpark DataFrame. Yes, your UDF is running inside Snowflake!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "dc189eea",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-------------------------\n",
      "|\"HELLO_UDF(\"\"NAME\"\")\"  |\n",
      "-------------------------\n",
      "|Hello Pia!             |\n",
      "|Hello Mats!            |\n",
      "-------------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "test_name_df.select(F.call_function(\"hello_udf\", F.col(\"name\"))).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f73f0595-9c48-462d-aee5-83169ae3698f",
   "metadata": {},
   "source": [
    "If a NULL value is provided to a UDF, it will be converted into a *None* value for the Python function **How big a deal is this???**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 162,
   "id": "fde03c58-e99b-45ce-ba10-cb48fb9864d4",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------\n",
      "|\"NAME\"  |\n",
      "----------\n",
      "|Mats    |\n",
      "|NULL    |\n",
      "|Pia     |\n",
      "----------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "test_name_with_null_df = session.create_dataframe([['Mats'],[None], ['Pia']], schema=[\"name\"])\n",
    "test_name_with_null_df.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 163,
   "id": "cedc061a-8ae5-425a-b479-83e5bca27f2b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-------------------------\n",
      "|\"HELLO_UDF(\"\"NAME\"\")\"  |\n",
      "-------------------------\n",
      "|Hello Mats!            |\n",
      "|Hello None!            |\n",
      "|Hello Pia!             |\n",
      "-------------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "test_name_with_null_df.select(F.call_function(\"hello_udf\", F.col(\"name\"))).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c25cbc2",
   "metadata": {},
   "source": [
    "### Batch Example of Hello World"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "545468f0-c5a4-457d-816b-0ee6d1c4d3b9",
   "metadata": {},
   "source": [
    "Create the same function again using the Python UDF batch API, this is done by changing the parameter to **PandasDataframe** or **PandasSeries** and the return to **PandasSeries**. The benfit of using the Python UDF batch API is that the function will not be called for each input row , but for a batches of rows instead."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "de22105e-725a-4087-aab1-db1ce80e8591",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "session.clear_imports()\n",
    "session.clear_packages()\n",
    "@F.udf(name=\"hello_batch_udf\", is_permanent=False, replace=True, session=session)\n",
    "def hello_batch_udf(ds: T.PandasSeries[str]) -> T.PandasSeries[str]:\n",
    "    n = len(ds)\n",
    "    return ds.apply(lambda x: f'Hello {x}, we got {n} rows')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 165,
   "id": "b7d642a3-454a-42dc-95ea-5697496a61b5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-------------------------------\n",
      "|\"HELLO_BATCH_UDF(\"\"NAME\"\")\"  |\n",
      "-------------------------------\n",
      "|Hello Mats, we got 1 rows    |\n",
      "|Hello Pia, we got 1 rows     |\n",
      "-------------------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "test_name_df.select(F.call_function(\"hello_batch_udf\", F.col(\"name\"))).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "630310d2-8567-4ab9-a50c-b8046bc1500f",
   "metadata": {},
   "source": [
    "Use a larger dataset for testing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 166,
   "id": "45b68c42-b2a7-473a-b4bf-0ed890e69598",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Nbr of customers: 100,000,000\n",
      "---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------\n",
      "|\"C_CUSTOMER_SK\"  |\"C_CUSTOMER_ID\"   |\"C_CURRENT_CDEMO_SK\"  |\"C_CURRENT_HDEMO_SK\"  |\"C_CURRENT_ADDR_SK\"  |\"C_FIRST_SHIPTO_DATE_SK\"  |\"C_FIRST_SALES_DATE_SK\"  |\"C_SALUTATION\"  |\"C_FIRST_NAME\"  |\"C_LAST_NAME\"  |\"C_PREFERRED_CUST_FLAG\"  |\"C_BIRTH_DAY\"  |\"C_BIRTH_MONTH\"  |\"C_BIRTH_YEAR\"  |\"C_BIRTH_COUNTRY\"  |\"C_LOGIN\"  |\"C_EMAIL_ADDRESS\"                |\"C_LAST_REVIEW_DATE\"  |\n",
      "---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------\n",
      "|93375407         |AAAAAAAAPKLMAJFA  |1205001               |4398                  |8856948              |2449648                   |2449618                  |Dr.             |Fred            |Randall        |N                        |22             |2                |1977            |MONGOLIA           |NULL       |Fred.Randall@IxNRsHK.edu         |2452532               |\n",
      "|93375408         |AAAAAAAAALLMAJFA  |390804                |6418                  |17606032             |2451686                   |2451656                  |Mr.             |Michael         |Kirby          |Y                        |5              |12               |1951            |GERMANY            |NULL       |Michael.Kirby@hF4G.com           |2452550               |\n",
      "|93375409         |AAAAAAAABLLMAJFA  |369656                |6884                  |3882865              |2451902                   |2451872                  |Sir             |Preston         |Clark          |Y                        |29             |6                |1990            |GEORGIA            |NULL       |Preston.Clark@1ed.com            |2452362               |\n",
      "|93375410         |AAAAAAAACLLMAJFA  |1469488               |3120                  |47381105             |2451956                   |2451926                  |Dr.             |David           |Morrison       |N                        |27             |12               |1978            |SIERRA LEONE       |NULL       |David.Morrison@jUvFTpsEedsi.com  |2452285               |\n",
      "|93375411         |AAAAAAAADLLMAJFA  |NULL                  |NULL                  |14271181             |NULL                      |NULL                     |Dr.             |NULL            |Holmes         |Y                        |8              |11               |1927            |TIMOR-LESTE        |NULL       |Launa.Holmes@gjtzIiPi.edu        |NULL                  |\n",
      "---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "customers_df = session.table(\"SFSALESSHARED_SFC_SAMPLES_PROD3_SAMPLE_DATA.tpcds_sf100tcl.customer\")# Name of Snowflake Sample Database might be different...\n",
    "print(f\"Nbr of customers: {customers_df.count():,}\")\n",
    "customers_df.show(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f19ad88-2fc1-4689-b79c-fc677bcfffd7",
   "metadata": {
    "tags": []
   },
   "source": [
    "If we test this using **show** we will see that it is only providing 15 rows since that is the limit we are setting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 167,
   "id": "fda5c5b8-5be5-48b6-bb6c-8c11d6bc7ea6",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------------------------------------\n",
      "|\"HELLO_BATCH_UDF(\"\"C_FIRST_NAME\"\")\"  |\n",
      "---------------------------------------\n",
      "|Hello Brian, we got 15 rows          |\n",
      "|Hello Margurite, we got 15 rows      |\n",
      "|Hello Keith, we got 15 rows          |\n",
      "|Hello Jessica, we got 15 rows        |\n",
      "|Hello Susan, we got 15 rows          |\n",
      "|Hello Orpha, we got 15 rows          |\n",
      "|Hello Matthew, we got 15 rows        |\n",
      "|Hello Jared, we got 15 rows          |\n",
      "|Hello Brenda, we got 15 rows         |\n",
      "|Hello William, we got 15 rows        |\n",
      "|Hello Dennis, we got 15 rows         |\n",
      "|Hello Donna, we got 15 rows          |\n",
      "|Hello Kevin, we got 15 rows          |\n",
      "|Hello Adam, we got 15 rows           |\n",
      "|Hello John, we got 15 rows           |\n",
      "---------------------------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "customers_df.select(F.col(\"C_FIRST_NAME\")).select(F.call_function(\"hello_batch_udf\", F.col(\"C_FIRST_NAME\"))).show(15)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cb0ac85-92ad-407b-a3d0-382522d95975",
   "metadata": {
    "tags": []
   },
   "source": [
    "By using **cache_result** we can temprary store the result of the query generated by the DataFrame and then seee that each call to the function does provide more rows.  You can see we can run more than 100M in less than 1 minute!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 168,
   "id": "a6f4a916-3f51-4780-abc0-e21dc14bd395",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------------------------------------\n",
      "|\"HELLO_BATCH_UDF(\"\"C_FIRST_NAME\"\")\"  |\n",
      "---------------------------------------\n",
      "|Hello <NA>, we got 256 rows          |\n",
      "|Hello Joseph, we got 256 rows        |\n",
      "|Hello Gregory, we got 256 rows       |\n",
      "|Hello Melvin, we got 256 rows        |\n",
      "|Hello Joe, we got 256 rows           |\n",
      "|Hello Thomas, we got 256 rows        |\n",
      "|Hello Pearl, we got 256 rows         |\n",
      "|Hello Ellen, we got 256 rows         |\n",
      "|Hello Trenton, we got 256 rows       |\n",
      "|Hello Cynthia, we got 256 rows       |\n",
      "|Hello Irene, we got 256 rows         |\n",
      "|Hello Craig, we got 256 rows         |\n",
      "|Hello Vincent, we got 256 rows       |\n",
      "|Hello Alan, we got 256 rows          |\n",
      "|Hello Tony, we got 256 rows          |\n",
      "---------------------------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "batch_udf_df = customers_df.select(F.call_function(\"hello_batch_udf\", F.col(\"C_FIRST_NAME\"))).cache_result()\n",
    "batch_udf_df.show(15)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8e37088-d430-46ba-a9f6-789f68f0554d",
   "metadata": {},
   "source": [
    "## 3. Create a UDF that uses as saved Python object"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e96a2c44-ea04-47ea-9b6b-e914a3f2f01c",
   "metadata": {},
   "source": [
    "Creating a UDF that uses as saved Python object. In this case a fitted scikit-learn pipline. **I am using native SKLearn to show how to use a any sort of python package.**  \n",
    "For many packages, Snowpark has a distributed version that is much easier to use.  Those versions also allow for easy use of the model registry and predict function.  \n",
    "The goal here is when you are using Python packages that are not optimized for Snowpark. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d437535f",
   "metadata": {},
   "source": [
    "### Build the pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "bb62aba2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import random\n",
    "import string\n",
    "\n",
    "from sklearn.datasets import make_regression\n",
    "from sklearn.preprocessing import MinMaxScaler, OrdinalEncoder\n",
    "from sklearn.pipeline import Pipeline\n",
    "from sklearn.ensemble import RandomForestRegressor\n",
    "from sklearn.compose import ColumnTransformer\n",
    "\n",
    "\n",
    "\n",
    "NUMERICAL_COLS = [\"X1\", \"X2\", \"X3\"]\n",
    "CATEGORICAL_COLS = [\"C1\", \"C2\", \"C3\"]\n",
    "FEATURE_COLS = NUMERICAL_COLS + CATEGORICAL_COLS\n",
    "\n",
    "# Create a dataset with numerical and categorical features\n",
    "X, y = make_regression(\n",
    "    n_samples=1000,\n",
    "    n_features=3,\n",
    "    noise=0.1,\n",
    "    random_state=0,\n",
    ")\n",
    "\n",
    "def generate_random_string(length):\n",
    "    return \"\".join(random.choices(string.ascii_uppercase, k=length))\n",
    "\n",
    "num_categorical_cols, categorical_feature_length = 3, 2\n",
    "categorical_features = []\n",
    "for _ in range(num_categorical_cols):\n",
    "    categorical_column = [generate_random_string(categorical_feature_length) for _ in range(X.shape[0])]\n",
    "    categorical_features.append(categorical_column)\n",
    "X = np.column_stack((X, *categorical_features))\n",
    "X = pd.DataFrame(X, columns=FEATURE_COLS)\n",
    "X[NUMERICAL_COLS] = X[NUMERICAL_COLS].astype(float)  ##Change\n",
    "\n",
    "preprocessor = ColumnTransformer(\n",
    "    transformers=[\n",
    "        ('num', MinMaxScaler(), NUMERICAL_COLS),\n",
    "        ('cat', OrdinalEncoder(handle_unknown=\"use_encoded_value\", unknown_value=99999), CATEGORICAL_COLS)\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Create the pipeline with the ColumnTransformer\n",
    "pipeline = Pipeline(steps=[('preprocessor', preprocessor),\n",
    "                       ('classifier',RandomForestRegressor())])\n",
    "\n",
    "rc_pipeline = pipeline.fit(X,y)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54b1a9cf-a60a-4784-aee2-ddb8034e3591",
   "metadata": {},
   "source": [
    "Remember what the attributes are of the object you are using. Often for models they will have a predict function that we later want to use.\n",
    "Save the fitted pipeline as a file locally using joblib."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 339,
   "id": "84937bd7-9e86-4fc6-a241-f40e165e3d64",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['rc_pipeline.joblib']"
      ]
     },
     "execution_count": 339,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "joblib.dump(rc_pipeline, \"rc_pipeline.joblib\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bc064fb-500d-4601-8aa7-03c77659e9c4",
   "metadata": {},
   "source": [
    "Upload the file to the Snowflake stage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c2a87063-29c2-49e6-8d86-4a58b6e54e5e",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[PutResult(source='rc_pipeline.joblib', target='rc_pipeline.joblib', source_size=9161831, target_size=9161840, source_compression='NONE', target_compression='NONE', status='UPLOADED', message='')]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "udf_stage_name =\"RAJ_MODELS\" ##Put your stage name here\n",
    "session.file.put(\"rc_pipeline.joblib\", f\"@{udf_stage_name}\", auto_compress=False, overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0856d004",
   "metadata": {},
   "source": [
    "Verify the object is in the stage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 438,
   "id": "1b77df23-e914-4b05-b94a-78d85f9ed1d6",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------------------------------------------------------------------------------------------------\n",
      "|\"name\"                         |\"size\"   |\"md5\"                             |\"last_modified\"                |\n",
      "--------------------------------------------------------------------------------------------------------------\n",
      "|raj_models/model.joblib.gz     |13216    |abe8081ad8ca6dea523caac2518f7c11  |Tue, 6 Feb 2024 02:39:52 GMT   |\n",
      "|raj_models/rc_pipeline.joblib  |9161840  |57c818aac2f62e96685e603d2a880d70  |Sun, 11 Feb 2024 03:19:51 GMT  |\n",
      "--------------------------------------------------------------------------------------------------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "session.sql(f\"ls @{udf_stage_name}\").show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9cde4470",
   "metadata": {},
   "source": [
    "### Creating the UDF"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ce6962a-4f27-41f8-80c8-17a4ac5204a1",
   "metadata": {
    "tags": []
   },
   "source": [
    "Create a function to load the file using joblib, use cachetools so the read from stage is only done once"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "083c4332-2b75-459c-8562-e4d565f2ffba",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "\n",
    "import cachetools\n",
    "\n",
    "@cachetools.cached(cache={})\n",
    "def load_joblib_file(filename):\n",
    "    import joblib\n",
    "    import sys\n",
    "    import os\n",
    "\n",
    "    import_dir = sys._xoptions.get(\"snowflake_import_directory\")\n",
    "    if import_dir:\n",
    "        with open(os.path.join(import_dir, filename), 'rb') as file:\n",
    "            m = joblib.load(file)\n",
    "            return m\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8bc65f4-7171-4b6c-8d11-fd0908227517",
   "metadata": {
    "tags": []
   },
   "source": [
    "Create the UDF, it is important that the *imports* parameter is refering the stage and file. Also, only the filename is needed for the *load_joblib_file* function.\n",
    "\n",
    "Since the function is depended on **Pandas**, **scikit-learn** and **cachetools** we need to add them to the *packages* parameter.\n",
    "\n",
    "We will also make sure UDF scikit-learn version matches the local one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "31f112d2-027a-4457-b16a-0273cc680049",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1.3.0'"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from sklearn import __version__ as sk_version\n",
    "sk_version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "36183140",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1.5.3'"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from pandas import __version__ as pd_version\n",
    "pd_version"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "129f9077",
   "metadata": {},
   "source": [
    "If you get an error here like: `Cannot add package scikit-learn==1.3.2 because it is not available in Snowflake ` - Welcome to production data science where version numbers matter - go ahead and install a supported version of scikit-learn in your local account. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "17ad2cc4-95a6-48fb-b872-0b94e2dc3b0b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "session.clear_imports()\n",
    "session.clear_packages()\n",
    "@F.udf(name = \"predict_model_udf\", is_permanent = False, imports = [f\"@{udf_stage_name}/rc_pipeline.joblib\"]  \n",
    "       , packages = [f'pandas=={pd_version}', f'scikit-learn=={sk_version}', 'cachetools'], replace = True, session = session)\n",
    "def predict_survive(pd_df: T.PandasDataFrame[float, float, float, str, str, str]) -> T.PandasSeries[int]:\n",
    "    \n",
    "    pd_df.columns = [FEATURE_COLS]\n",
    "    model = load_joblib_file('rc_pipeline.joblib') # Only call with the file name!\n",
    "\n",
    "    return model.predict(pd_df)    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95aa59a2-8553-414a-ae1f-124ed182638c",
   "metadata": {},
   "source": [
    "Test that the UDF works.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "fa4d5d2a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-------------------------------------------------------------------------------------------\n",
      "|\"X1\"                 |\"X2\"                   |\"X3\"                  |\"C1\"  |\"C2\"  |\"C3\"  |\n",
      "-------------------------------------------------------------------------------------------\n",
      "|-1.1802856063511906  |1.012168295174788      |-0.28044778146745414  |HY    |OH    |XT    |\n",
      "|0.7335948682293696   |2.0118642631265615     |0.3030110504561574    |NN    |CB    |LR    |\n",
      "|-0.5139029502799155  |-0.7688491596748099    |0.9882405737426969    |NV    |BG    |DV    |\n",
      "|1.037585667050634    |0.018791791774257802   |1.3925184494342724    |PB    |WM    |EM    |\n",
      "|-0.3550287310553741  |-1.8923618933173414    |-0.3004787855854223   |RZ    |GR    |FV    |\n",
      "|1.183119557331707    |0.7188971655282916     |1.4969109935208271    |QX    |LH    |AA    |\n",
      "|0.2390336012467649   |-1.0003303489537054    |-0.04932407014757259  |OS    |YG    |KK    |\n",
      "|0.46637957438197886  |-0.09439250641118496   |-0.47728627040322935  |NO    |VE    |DK    |\n",
      "|0.2456164279017758   |-0.032996480264666096  |0.8425887964569804    |RR    |ZG    |SD    |\n",
      "|0.5898798207345195   |-0.3638588099707899    |0.7679024077327037    |LA    |DI    |DC    |\n",
      "-------------------------------------------------------------------------------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "features_df = session.create_dataframe(X)\n",
    "features_df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72103672",
   "metadata": {},
   "source": [
    "Now you are using the UDF on your Snowflake data!!!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "8fa73c0d-5503-416c-b508-226200250feb",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------------------------------------------------\n",
      "|\"X1\"                 |\"X2\"                   |\"X3\"                  |\"C1\"  |\"C2\"  |\"C3\"  |\"PREDICTION\"  |\n",
      "----------------------------------------------------------------------------------------------------------\n",
      "|-1.1802856063511906  |1.012168295174788      |-0.28044778146745414  |HY    |OH    |XT    |-15           |\n",
      "|0.7335948682293696   |2.0118642631265615     |0.3030110504561574    |NN    |CB    |LR    |257           |\n",
      "|-0.5139029502799155  |-0.7688491596748099    |0.9882405737426969    |NV    |BG    |DV    |-54           |\n",
      "|1.037585667050634    |0.018791791774257802   |1.3925184494342724    |PB    |WM    |EM    |168           |\n",
      "|-0.3550287310553741  |-1.8923618933173414    |-0.3004787855854223   |RZ    |GR    |FV    |-222          |\n",
      "|1.183119557331707    |0.7188971655282916     |1.4969109935208271    |QX    |LH    |AA    |237           |\n",
      "|0.2390336012467649   |-1.0003303489537054    |-0.04932407014757259  |OS    |YG    |KK    |-82           |\n",
      "|0.46637957438197886  |-0.09439250641118496   |-0.47728627040322935  |NO    |VE    |DK    |-5            |\n",
      "|0.2456164279017758   |-0.032996480264666096  |0.8425887964569804    |RR    |ZG    |SD    |64            |\n",
      "|0.5898798207345195   |-0.3638588099707899    |0.7679024077327037    |LA    |DI    |DC    |51            |\n",
      "----------------------------------------------------------------------------------------------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "input_cols = [F.col(col) for col in FEATURE_COLS ]\n",
    "features_df.with_column(\"PREDICTION\", F.call_function(\"predict_model_udf\",  *input_cols)).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc7bfbc9-40ab-4975-a82f-acbe09dede50",
   "metadata": {},
   "source": [
    "### Batch UDF\n",
    "A batch UDF can also be called by providing the inputs as a array for example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "7027b99b-919d-4072-bc4c-7d90da9fef8d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "session.clear_imports()\n",
    "session.clear_packages()\n",
    "\n",
    "@F.udf(name = \"predict_model_array_udf\", is_permanent = False, imports = [f\"@{udf_stage_name}/rc_pipeline.joblib\"]\n",
    "       , packages = [f'pandas=={pd_version}', f'scikit-learn=={sk_version}', 'cachetools'], replace = True, session = session)\n",
    "def predict_survive_array(pd_s: T.PandasSeries[list]) -> T.PandasSeries[int]:\n",
    "    pd_df = pd.DataFrame.from_dict(dict(zip(pd_s.index, pd_s.values))).T\n",
    "    pd_df.columns = [FEATURE_COLS]\n",
    "    model = load_joblib_file('rc_pipeline.joblib') # Only call with the file name!\n",
    "\n",
    "    return model.predict(pd_df)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "32453411-2540-4d25-a11d-eb30be6c567a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------------------------------------------------\n",
      "|\"X1\"                 |\"X2\"                   |\"X3\"                  |\"C1\"  |\"C2\"  |\"C3\"  |\"PREDICTION\"  |\n",
      "----------------------------------------------------------------------------------------------------------\n",
      "|-1.1802856063511906  |1.012168295174788      |-0.28044778146745414  |HY    |OH    |XT    |-15           |\n",
      "|0.7335948682293696   |2.0118642631265615     |0.3030110504561574    |NN    |CB    |LR    |257           |\n",
      "|-0.5139029502799155  |-0.7688491596748099    |0.9882405737426969    |NV    |BG    |DV    |-54           |\n",
      "|1.037585667050634    |0.018791791774257802   |1.3925184494342724    |PB    |WM    |EM    |168           |\n",
      "|-0.3550287310553741  |-1.8923618933173414    |-0.3004787855854223   |RZ    |GR    |FV    |-222          |\n",
      "|1.183119557331707    |0.7188971655282916     |1.4969109935208271    |QX    |LH    |AA    |237           |\n",
      "|0.2390336012467649   |-1.0003303489537054    |-0.04932407014757259  |OS    |YG    |KK    |-82           |\n",
      "|0.46637957438197886  |-0.09439250641118496   |-0.47728627040322935  |NO    |VE    |DK    |-5            |\n",
      "|0.2456164279017758   |-0.032996480264666096  |0.8425887964569804    |RR    |ZG    |SD    |64            |\n",
      "|0.5898798207345195   |-0.3638588099707899    |0.7679024077327037    |LA    |DI    |DC    |51            |\n",
      "----------------------------------------------------------------------------------------------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "input_cols = [F.col(col) for col in FEATURE_COLS ]\n",
    "features_df.with_column(\"PREDICTION\", F.call_function(\"predict_model_array_udf\", F.array_construct(*input_cols))).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a36d090-f2a2-465b-9fee-e27d76214bbe",
   "metadata": {},
   "source": [
    "### Dict as an input into the UDF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "c7cb487c-5a9c-4f2e-b36b-0b03feb48a5f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "session.clear_imports()\n",
    "session.clear_packages()\n",
    "\n",
    "@F.udf(name = \"predict_model_dict_udf\", is_permanent = False, imports = [f\"@{udf_stage_name}/rc_pipeline.joblib\"]\n",
    "       , packages = [f'pandas=={pd_version}', f'scikit-learn=={sk_version}', 'cachetools'], replace = True, session = session)\n",
    "def predict_survive_dict(pd_s: T.PandasSeries[dict]) -> T.PandasSeries[int]:\n",
    "    pd_df = pd.json_normalize(pd_s)[['X1', 'X2', 'X3', 'C1', 'C2', 'C3']]\n",
    "    model = load_joblib_file('rc_pipeline.joblib') # Only call with the file name!\n",
    "\n",
    "    return model.predict(pd_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "17ec4e68-cc95-4a12-be21-38e640391fd3",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------------------------------------------------\n",
      "|\"X1\"                 |\"X2\"                   |\"X3\"                  |\"C1\"  |\"C2\"  |\"C3\"  |\"PREDICTION\"  |\n",
      "----------------------------------------------------------------------------------------------------------\n",
      "|-1.1802856063511906  |1.012168295174788      |-0.28044778146745414  |HY    |OH    |XT    |-15           |\n",
      "|0.7335948682293696   |2.0118642631265615     |0.3030110504561574    |NN    |CB    |LR    |257           |\n",
      "|-0.5139029502799155  |-0.7688491596748099    |0.9882405737426969    |NV    |BG    |DV    |-54           |\n",
      "|1.037585667050634    |0.018791791774257802   |1.3925184494342724    |PB    |WM    |EM    |168           |\n",
      "|-0.3550287310553741  |-1.8923618933173414    |-0.3004787855854223   |RZ    |GR    |FV    |-222          |\n",
      "|1.183119557331707    |0.7188971655282916     |1.4969109935208271    |QX    |LH    |AA    |237           |\n",
      "|0.2390336012467649   |-1.0003303489537054    |-0.04932407014757259  |OS    |YG    |KK    |-82           |\n",
      "|0.46637957438197886  |-0.09439250641118496   |-0.47728627040322935  |NO    |VE    |DK    |-5            |\n",
      "|0.2456164279017758   |-0.032996480264666096  |0.8425887964569804    |RR    |ZG    |SD    |64            |\n",
      "|0.5898798207345195   |-0.3638588099707899    |0.7679024077327037    |LA    |DI    |DC    |51            |\n",
      "----------------------------------------------------------------------------------------------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "features_df.with_column(\"PREDICTION\", F.call_function(\"predict_model_dict_udf\", F.object_construct('*'))).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cd1be76-4884-4989-930a-e02bffcffc03",
   "metadata": {},
   "source": [
    "It is possible to return multiple values with a UDF a list or Dict is needed. The code below doesn't work for my model, since I bult a regressor. But by definining the output as a list, you are able to return multiple values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "a1c38181-dd5f-42c1-b806-2b2e1c0ef689",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "##Doesn't work -- My model does't support predict_proba\n",
    "session.clear_imports()\n",
    "session.clear_packages()\n",
    "\n",
    "@F.udf(name = \"predict_model_array_return_udf\", is_permanent = False, imports = [f\"@{udf_stage_name}/rc_pipeline.joblib\"]\n",
    "       , packages = [f'pandas=={pd_version}', f'scikit-learn=={sk_version}', 'cachetools'], replace = True, session = session)\n",
    "def predict_survive_array_return(pd_df: T.PandasDataFrame[str, str, str, float, float]) -> T.PandasSeries[list]:\n",
    "    \n",
    "    pd_df.columns = [FEATURE_COLS]\n",
    "    model = load_joblib_file('rc_pipeline.joblib') # Only call with the file name!\n",
    "    prediction_proba = model.predict_proba(pd_df)\n",
    "        \n",
    "    # Get the label for the highest probablility\n",
    "    predicted_classes_idx = np.argmax(prediction_proba, axis=1)\n",
    "    classes = model.classes_\n",
    "    predicted_classes = classes[predicted_classes_idx]\n",
    "\n",
    "    # Create a list with return values\n",
    "    return_array = np.column_stack((prediction_proba, predicted_classes))\n",
    "\n",
    "    return return_array\n",
    "\n",
    "# Doesn't work -- My model does't support predict_proba\n",
    "#features_df.with_column(\"RETURN_ARRAY\", F.call_function(\"predict_survive_array_return_udf\", *input_cols)).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4eb3dfdd",
   "metadata": {},
   "source": [
    "## 4. Vectorized UDFs\n",
    "Vectorized UDFs and standard udfs are both parallelized the exact same way. The difference between them is that a vectorized udf relies on passing numpy arrays of data as input versus rowsets which can be much more efficient when your udf is utilizing numpy based functions for computation (which is the case with most many ML libraries)\n",
    "You then run the UDF with batches, each Batch must finish within 180 seconds.  \n",
    "\n",
    "To use a vectorized UDF, you just need to set the proper return variable types. Snowpark will then automatically use the vectorized UDFs. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "5e9f0a20",
   "metadata": {},
   "outputs": [],
   "source": [
    "session.clear_imports()\n",
    "session.clear_packages()\n",
    "@F.udf(name = \"predict_survive_udf\", is_permanent = False, imports = [f\"@{udf_stage_name}/rc_pipeline.joblib\"],max_batch_size= 100\n",
    "       , packages = [f'pandas=={pd_version}', f'scikit-learn=={sk_version}', 'cachetools'], replace = True, session = session)\n",
    "def predict_survive(pd_df: T.PandasDataFrame[float, float, float, str, str, str]) -> T.PandasSeries[int]:\n",
    "    \n",
    "    pd_df.columns = [FEATURE_COLS]\n",
    "    model = load_joblib_file('rc_pipeline.joblib') # Only call with the file name!\n",
    "\n",
    "    return model.predict(pd_df)    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "82e6dea5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------------------------------------------------\n",
      "|\"X1\"                 |\"X2\"                   |\"X3\"                  |\"C1\"  |\"C2\"  |\"C3\"  |\"PREDICTION\"  |\n",
      "----------------------------------------------------------------------------------------------------------\n",
      "|-1.1802856063511906  |1.012168295174788      |-0.28044778146745414  |HY    |OH    |XT    |-15           |\n",
      "|0.7335948682293696   |2.0118642631265615     |0.3030110504561574    |NN    |CB    |LR    |257           |\n",
      "|-0.5139029502799155  |-0.7688491596748099    |0.9882405737426969    |NV    |BG    |DV    |-54           |\n",
      "|1.037585667050634    |0.018791791774257802   |1.3925184494342724    |PB    |WM    |EM    |168           |\n",
      "|-0.3550287310553741  |-1.8923618933173414    |-0.3004787855854223   |RZ    |GR    |FV    |-222          |\n",
      "|1.183119557331707    |0.7188971655282916     |1.4969109935208271    |QX    |LH    |AA    |237           |\n",
      "|0.2390336012467649   |-1.0003303489537054    |-0.04932407014757259  |OS    |YG    |KK    |-82           |\n",
      "|0.46637957438197886  |-0.09439250641118496   |-0.47728627040322935  |NO    |VE    |DK    |-5            |\n",
      "|0.2456164279017758   |-0.032996480264666096  |0.8425887964569804    |RR    |ZG    |SD    |64            |\n",
      "|0.5898798207345195   |-0.3638588099707899    |0.7679024077327037    |LA    |DI    |DC    |51            |\n",
      "----------------------------------------------------------------------------------------------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "input_cols = [F.col(col) for col in FEATURE_COLS ]\n",
    "features_df.with_column(\"PREDICTION\", F.call_function(\"predict_survive_udf\",  *input_cols)).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a57f1468-447d-421f-bbd8-eb9170980e48",
   "metadata": {},
   "source": [
    "## 5. UDTF\n",
    "\n",
    "User Defined Table Functions (UDTF) is a function that returns zero, one, or multiple rows for each input row.\n",
    "\n",
    "When creating a UDTF a Python class has to be used as the handler\n",
    "\n",
    "A UDTF handler class implements the following, which Snowflake invokes at run time:\n",
    "* An **__init__** method. Optional. Invoked to initialize stateful processing of input partitions.\n",
    "* A **process** method. Required. Invoked for each input row. The method returns a tabular value as tuples.\n",
    "* An **end_partition** method. Optional. Invoked to finalize processing of input partitions.\n",
    "\n",
    "A UDTF can be created using the **@udtf** decorator, the **udtf** function or the **udtf.register** method ofthe session object. It can be permanent or temporary.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7805c6ff",
   "metadata": {},
   "source": [
    "### Python model as UDTF\n",
    "Use our model and move it to a UDTF. The key here is setting the input and output types. For a UDTF the output is typically multiple rows, so dataframe is used here. This is getting to more code, so I broke up registration of the UTDF and the code for the UTDF. This allows me to ensure the code is running correctly locally, before I try to register the UDTF inside Snowflake."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "a0491331",
   "metadata": {},
   "outputs": [],
   "source": [
    "from snowflake.snowpark.functions import udtf, lit\n",
    "from snowflake.snowpark.types import PandasDataFrameType, IntegerType, StringType, FloatType\n",
    "session.clear_imports()\n",
    "session.clear_packages()\n",
    "\n",
    "class predict_model:\n",
    "    # We load the model from stage at the start of each partition\n",
    "    def __init__(self) -> None:\n",
    "        import os\n",
    "        import joblib\n",
    "        import sys\n",
    "        import pandas as pd\n",
    "        #self.model = joblib.load('rc_pipeline.joblib')\n",
    "        import_dir = sys._xoptions.get(\"snowflake_import_directory\")\n",
    "        with open(os.path.join(import_dir, 'rc_pipeline.joblib'), 'rb') as file:\n",
    "            self.model = joblib.load(file)\n",
    "        \n",
    "    # Called after the last row in a partion has been processed\n",
    "    def end_partition(self, pd_df):\n",
    "        # Set the column name of the provided Pandas DataFrame\n",
    "        pd_df.columns=[\"X1\",\"X2\",\"X3\",\"C1\",\"C2\",\"C3\"]\n",
    "        results = self.model.predict(pd_df) # Get the name of the 7th column\n",
    "        print (results)\n",
    "        df_seventh_column = pd.DataFrame(results)\n",
    "        print (df_seventh_column)\n",
    "        yield df_seventh_column\n",
    "\n",
    "predict_udtf = session.udtf.register(predict_model,is_permanent=False, replace=True, imports = [f\"@{udf_stage_name}/rc_pipeline.joblib\"]\n",
    "        , input_types=[T.PandasDataFrameType([T.FloatType(), T.FloatType(),T.FloatType(), T.StringType(), T.StringType(), T.StringType()])]\n",
    "        , output_schema=PandasDataFrameType([FloatType()], [\"PREDICTIONS\"])\n",
    "        , packages = [f'pandas=={pd_version}', f'scikit-learn=={sk_version}', 'cachetools'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb8a249c",
   "metadata": {},
   "source": [
    "Run it locally like: \n",
    "```\n",
    "predictor = predict_udtf()\n",
    "final = predictor.end_partition(X)\n",
    "final\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "a7d2c3b1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-----------------------\n",
      "|\"PREDICTIONS\"        |\n",
      "-----------------------\n",
      "|53.91033785203392    |\n",
      "|137.7003877591698    |\n",
      "|-195.21377894882605  |\n",
      "|-4.268479029560137   |\n",
      "|2.2400981635396415   |\n",
      "|-160.1945418373434   |\n",
      "|103.51438829189482   |\n",
      "|-109.63587457610556  |\n",
      "|-152.69719097490136  |\n",
      "|-6.194146150918908   |\n",
      "-----------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "results = features_df.select(predict_udtf(\"X1\",\"X2\",\"X3\",\"C1\",\"C2\",\"C3\").over(partition_by=[\"X1\"]))\n",
    "results.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b47040d",
   "metadata": {},
   "source": [
    "### Word Count UDTF\n",
    "\n",
    "Start with a simple UDTF that splits a string into words and fore each unique word it returns a row with it and the number of ocurrances in the string of it. We need to provide the output schema ie the columns of the returning rows. If only names are provided the data types are inheried from the process parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import Counter\n",
    "from typing import Iterable, Tuple\n",
    "\n",
    "session.clear_imports()\n",
    "session.clear_packages()\n",
    "\n",
    "@F.udtf(name=\"word_count_udtf\", output_schema=[\"word\", \"count\"], is_permanent=False, replace=True, session=session)\n",
    "class MyWordCount:\n",
    "    # Called once for each partition\n",
    "    def __init__(self):\n",
    "        self._total_per_partition = 0\n",
    "    \n",
    "    # Called for each row in a partition\n",
    "    def process(self, s1: str) -> Iterable[Tuple[str, int]]:\n",
    "        words = s1.split()\n",
    "        self._total_per_partition = len(words)\n",
    "        # Counter will return a dict with the uinique words as keys and the number ocurrances as the values\n",
    "        counter = Counter(words) \n",
    "        yield from counter.items()\n",
    "    \n",
    "    # Called after the last row in a partion has been processed\n",
    "    def end_partition(self):\n",
    "        yield (\"partition_total\", self._total_per_partition)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "712f853a-f5cb-4449-9192-b856a880e551",
   "metadata": {},
   "source": [
    "Test the UDTF, by using session.table_function we will get a new DataFrame with the data generated by the UDTF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "af33476c-a11b-410f-b4e8-685d8fa3e852",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-----------------------------\n",
      "|\"WORD\"           |\"COUNT\"  |\n",
      "-----------------------------\n",
      "|w1               |1        |\n",
      "|w2               |2        |\n",
      "|w3               |3        |\n",
      "|partition_total  |6        |\n",
      "-----------------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df_udtf = session.table_function(\"word_count_udtf\", F.lit(\"w1 w2 w2 w3 w3 w3\"))\n",
    "df_udtf.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff11d9aa-d7ac-4bc1-afcc-c4bcab82e104",
   "metadata": {},
   "source": [
    "We can also use it with a DataFrame, using call_table_function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "bd36b562-aa76-46b5-b585-8050e576519c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---------------------\n",
      "|\"TEXT\"             |\n",
      "---------------------\n",
      "|w1 w2 w2 w3 w3 w3  |\n",
      "---------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df_udtf_data = session.create_dataframe([[\"w1 w2 w2 w3 w3 w3\"]], schema=[\"text\"])\n",
    "df_udtf_data.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "11a7572a-028d-479c-84b3-8fe5a1f10025",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-----------------------------\n",
      "|\"WORD\"           |\"COUNT\"  |\n",
      "-----------------------------\n",
      "|w1               |1        |\n",
      "|w2               |2        |\n",
      "|w3               |3        |\n",
      "|partition_total  |6        |\n",
      "-----------------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df_udtf_data.select(F.call_table_function(\"word_count_udtf\", F.col(\"TEXT\"))).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1830123b-d953-4e84-a48c-ec45ffa7d291",
   "metadata": {},
   "source": [
    "If we want to do the split/count by a column, the partition_by parameter can be used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "debeed02-0cde-4462-94f3-ebc4128b8b97",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-----------------------------------\n",
      "|\"PARTITION\"  |\"TEXT\"             |\n",
      "-----------------------------------\n",
      "|1            |w1 w2 w2 w3 w3 w3  |\n",
      "|2            |w4 w4 w4 w4 w1     |\n",
      "-----------------------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df_udtf_part_data = session.create_dataframe([[\"1\", \"w1 w2 w2 w3 w3 w3\"], [\"2\", \"w4 w4 w4 w4 w1\"]], schema=[\"partition\",\"text\"])\n",
    "df_udtf_part_data.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "164ac085-0081-448a-a9ce-6afdc2d76621",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-------------------------------------------\n",
      "|\"PARTITION\"  |\"WORD\"           |\"COUNT\"  |\n",
      "-------------------------------------------\n",
      "|1            |w1               |1        |\n",
      "|1            |w2               |2        |\n",
      "|1            |w3               |3        |\n",
      "|1            |partition_total  |6        |\n",
      "|2            |w4               |4        |\n",
      "|2            |w1               |1        |\n",
      "|2            |partition_total  |5        |\n",
      "-------------------------------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df_udtf_part_data.select(\"partition\", F.call_table_function(\"word_count_udtf\", F.col(\"TEXT\")).over(partition_by=\"partition\")).show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}