lab2_vectorized_udfs.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Lab 2 - Vectorized UDFs for Batching\n",
    "\n",
    "### Setup\n",
    "\n",
    "Below are the imports needed for this lab.  They have been included in the instantiation of the conda environment.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Importing snowpark libraries\n",
    "import snowflake.snowpark\n",
    "from snowflake.snowpark.session import Session\n",
    "from snowflake.snowpark.functions import avg, stddev, udf\n",
    "from snowflake.snowpark.types import FloatType, StringType, PandasSeries\n",
    "\n",
    "# Others\n",
    "import json\n",
    "import re\n",
    "import pandas as pd\n",
    "from datetime import datetime, timedelta"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Connect to Snowflake by opening the `credentials.json` file and creating a session."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating a session into Snowflake. This is the same as in previous sections of this lab\n",
    "with open('credentials.json') as f:\n",
    "    connection_parameters = json.load(f)\n",
    "    \n",
    "session = Session.builder.configs(connection_parameters).create()\n",
    "\n",
    "# Testing the session\n",
    "session.sql(\"SELECT current_warehouse(), current_database(), current_schema()\").show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before continuing, we will disable caching to ensure proper performance comparisons.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.sql(\"ALTER SESSION SET USE_CACHED_RESULT = FALSE\").collect()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Dataframes\n",
    "\n",
    "As mentioned earlier, our data (TPCH dataset) is available from Snowflake via an inbound data share. We are taking 4 datasets into consideration, namely, Customers & Orders with 2 different sizes.\n",
    "\n",
    "Let's create all of these datasets as dataframes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_customer_100 = session.table(\"SNOWFLAKE_SAMPLE_DATA.TPCH_SF100.customer\")\n",
    "df_customer_1000 = session.table(\"SNOWFLAKE_SAMPLE_DATA.TPCH_SF1000.customer\")\n",
    "df_order_100 = session.table(\"SNOWFLAKE_SAMPLE_DATA.TPCH_SF100.orders\")\n",
    "df_order_1000 = session.table(\"SNOWFLAKE_SAMPLE_DATA.TPCH_SF1000.orders\")\n",
    "df_customer_100.limit(2).show()\n",
    "df_order_100.limit(2).show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, let's check on the data volumes we will be working with.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"df_customer_100 = {str(df_customer_100.count()/1000000)} M, df_customer_1000 = {str(df_customer_1000.count()/1000000)} M,\\ndf_order_100 = {str(df_order_100.count()/1000000)} M, df_order_1000 = {str(df_order_1000.count()/1000000)} M\")\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup Warehouses\n",
    "As mentioned earlier, we will be using 2 different Warehouse Sizes:\n",
    "\n",
    "For *_100 datasets, we will use Small and Medium Warehouse sizes\n",
    "For *_1000 datasets, we will use Medium and Large Warehouse sizes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's change the warehouse size as required\n",
    "session.sql(\"ALTER WAREHOUSE compute_wh SET warehouse_size='Small'\").collect()\n",
    "\n",
    "# Let's check the current size of the warehouse\n",
    "session.sql(\"SHOW WAREHOUSES LIKE '%COMPUTE_WH%'\").collect()\n",
    "session.sql('SELECT \"name\", \"size\" FROM TABLE(RESULT_SCAN(LAST_QUERY_ID()))').show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use Case 1: Numeric Computation\n",
    "For this use case, we will make use of the Customer datasets. Let's first compute the mean and standard deviation for Account Balance in Customer datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## df_customer_100\n",
    "df_customer_100_mean = float(df_customer_100.agg(avg(\"C_ACCTBAL\")).to_pandas().values[0])\n",
    "df_customer_100_stddev = float(df_customer_100.agg(stddev(\"C_ACCTBAL\")).to_pandas().values[0])\n",
    "\n",
    "## df_customer_1000\n",
    "df_customer_1000_mean = float(df_customer_1000.agg(avg(\"C_ACCTBAL\")).to_pandas().values[0])\n",
    "df_customer_1000_stddev = float(df_customer_1000.agg(stddev(\"C_ACCTBAL\")).to_pandas().values[0])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create 2 sets of UDFs, one for each dataset as the means & stdevs will be different\n",
    "\n",
    "- As mentioned earlier, these are hypothetical usecases\n",
    "- In this UDF, we take the Customer Account Balance as Input, subtract with the Mean, add the Standard Deviation, and finally multiple with 10,000 Notice how we are creating a Normal and Vectorised UDF both"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## df_customer_100\n",
    "def basic_compute_100(inp):\n",
    "    return (inp - df_customer_100_mean + df_customer_100_stddev) * 10000.0\n",
    "\n",
    "# Let's create a UDF on top of that\n",
    "udf_bc_100 = udf(basic_compute_100, return_type=FloatType(), input_types=[FloatType()])\n",
    "\n",
    "# Let's vectorise the UDF\n",
    "@udf()\n",
    "def vect_udf_bc_100(inp: PandasSeries[float]) -> PandasSeries[float]:\n",
    "    return (inp - df_customer_100_mean + df_customer_100_stddev) * 10000.0"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above is an example of basic_compute_100 function created as a Temporary UDF within Snowflake. It takes a few seconds to create the UDF, the first time. Subsequent runs will take less time\n",
    "\n",
    "- You'll notice the run time parameters such as `Python version=3.8, Cloudpickle==2.0.0` package, and so on\n",
    "- More importantly towards the end, you can see the actual function definition\n",
    "\n",
    "Using @udf with a Pandas Series or Dataframe as input will automatically make use of the Python UDF Batch API or Vectorised UDF. \n",
    "The UDF definition for a Vectorised UDF looks pretty much the same as for a normal non-vectorised one However, the key differences apart from our definition for `vect_udf_bc100` is:\n",
    "\n",
    "- Importing the pandas library: this can be seen towards the end of the Vectorised UDF definition\n",
    "- The Vectorised input argument: `compute._sf_vectorized_input = pandas.DataFrame`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## df_customer_1000\n",
    "def basic_compute_1000(inp):\n",
    "    return (inp - df_customer_1000_mean + df_customer_1000_stddev) * 10000.0\n",
    "\n",
    "# Let's create a UDF on top of that\n",
    "udf_bc_1000 = udf(basic_compute_1000, return_type=FloatType(), input_types=[FloatType()])\n",
    "\n",
    "# Let's vectorise the UDF\n",
    "@udf()\n",
    "def vect_udf_bc_1000(inp: PandasSeries[float]) -> PandasSeries[float]:\n",
    "    return (inp - df_customer_1000_mean + df_customer_1000_stddev) * 10000.0"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's quickly execute these UDFs against our account balance datasets 100 & 1000. In our methodology we had discussed using different Virtual Warehouse sizes:\n",
    "\n",
    "- Small and Medium for 100 dataset\n",
    "- Medium and Large for 1000 dataset We will be re-running this accordingly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.sql(\"ALTER WAREHOUSE compute_wh SET warehouse_size='Small'\").collect()\n",
    "print(\"Use case 1.1: Numeric Computation with Small WH and 100x dataset\")\n",
    "print(\"Scalar UDF: \") \n",
    "%timeit -r 1 -n 1 df_customer_100.select(udf_bc_100(\"C_ACCTBAL\").alias(\"bal_from_mean\")).agg(avg(\"bal_from_mean\")).show()\n",
    "print(\"Vectorized UDF: \") \n",
    "%timeit -r 1 -n 1 df_customer_100.select(vect_udf_bc_100(\"C_ACCTBAL\").alias(\"bal_from_mean\")).agg(avg(\"bal_from_mean\")).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.sql(\"ALTER WAREHOUSE compute_wh SET warehouse_size='Medium'\").collect()\n",
    "print(\"Use case 1.2: Numeric Computation with Medium WH and 100x dataset\")\n",
    "print(\"Scalar UDF: \") \n",
    "%timeit -r 1 -n 1 df_customer_100.select(udf_bc_100(\"C_ACCTBAL\").alias(\"bal_from_mean\")).agg(avg(\"bal_from_mean\")).show()\n",
    "print(\"Vectorized UDF: \") \n",
    "%timeit -r 1 -n 1 df_customer_100.select(vect_udf_bc_100(\"C_ACCTBAL\").alias(\"bal_from_mean\")).agg(avg(\"bal_from_mean\")).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.sql(\"ALTER WAREHOUSE compute_wh SET warehouse_size='Medium'\").collect()\n",
    "print(\"Use case 1.3: Numeric Computation with Medium WH and 1000x dataset\")\n",
    "print(\"Scalar UDF: \") \n",
    "%timeit -r 1 -n 1 df_customer_100.select(udf_bc_1000(\"C_ACCTBAL\").alias(\"bal_from_mean\")).agg(avg(\"bal_from_mean\")).show()\n",
    "print(\"Vectorized UDF: \") \n",
    "%timeit -r 1 -n 1 df_customer_100.select(vect_udf_bc_1000(\"C_ACCTBAL\").alias(\"bal_from_mean\")).agg(avg(\"bal_from_mean\")).show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.sql(\"ALTER WAREHOUSE compute_wh SET warehouse_size='Large'\").collect()\n",
    "print(\"Use case 1.4: Numeric Computation with Large WH and 1000x dataset\")\n",
    "print(\"Scalar UDF: \") \n",
    "%timeit -r 1 -n 1 df_customer_100.select(udf_bc_1000(\"C_ACCTBAL\").alias(\"bal_from_mean\")).agg(avg(\"bal_from_mean\")).show()\n",
    "print(\"Vectorized UDF: \") \n",
    "%timeit -r 1 -n 1 df_customer_100.select(vect_udf_bc_1000(\"C_ACCTBAL\").alias(\"bal_from_mean\")).agg(avg(\"bal_from_mean\")).show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use Case 2: String Manipulation\n",
    "To test a workload based on string manipulation, let's create a single set of UDFs for Customer100 & Customer1000 datasets. Here, we take the Customer Name as input and split based on the ‘#' character. As before, we will create a Normal and Vectorised UDF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def str_manipulate(inp):\n",
    "    return inp.split('#')[1]\n",
    "\n",
    "# Let's create a UDF on top of that\n",
    "udf_sm = udf(str_manipulate, return_type=StringType(), input_types=[StringType()])\n",
    "\n",
    "# Let's vectorise the same UDF\n",
    "@udf()\n",
    "def vect_udf_sm(inp: PandasSeries[str]) -> PandasSeries[str]:\n",
    "    return inp.str.split('#', expand=True)[1]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As before, let's quickly execute these UDFs against our datasets 100 & 1000. We will re-run these UDFs based on our Dataset and VW sizing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.sql(\"ALTER WAREHOUSE compute_wh SET warehouse_size='Small'\").collect()\n",
    "print(\"Use case 2.1: String Manipulation with Small WH and 100x dataset\")\n",
    "print(\"\\rScalar UDF: \", end=\"\") \n",
    "%timeit -r 1 -n 1 df_customer_100.select(udf_sm(\"C_NAME\").alias(\"CustIDs\")).write.mode(\"overwrite\").save_as_table(\"SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.CustIDs\")\n",
    "print(\"\\rVectorized UDF: \", end=\"\") \n",
    "%timeit -r 1 -n 1 df_customer_100.select(vect_udf_sm(\"C_NAME\").alias(\"CustIDs\")).write.mode(\"overwrite\").save_as_table(\"SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.CustIDs\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.sql(\"ALTER WAREHOUSE compute_wh SET warehouse_size='Medium'\").collect()\n",
    "print(\"Use case 2.2: String Manipulation with Medium WH and 100x dataset\")\n",
    "print(\"\\rScalar UDF: \", end=\"\") \n",
    "%timeit -r 1 -n 1 df_customer_100.select(udf_sm(\"C_NAME\").alias(\"CustIDs\")).write.mode(\"overwrite\").save_as_table(\"SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.CustIDs\")\n",
    "print(\"\\rVectorized UDF: \", end=\"\") \n",
    "%timeit -r 1 -n 1 df_customer_100.select(vect_udf_sm(\"C_NAME\").alias(\"CustIDs\")).write.mode(\"overwrite\").save_as_table(\"SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.CustIDs\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.sql(\"ALTER WAREHOUSE compute_wh SET warehouse_size='Medium'\").collect()\n",
    "print(\"Use case 2.3: String Manipulation with Medium WH and 1000x dataset\")\n",
    "print(\"\\rScalar UDF: \", end=\"\") \n",
    "%timeit -r 1 -n 1 df_customer_1000.select(udf_sm(\"C_NAME\").alias(\"CustIDs\")).write.mode(\"overwrite\").save_as_table(\"SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.CustIDs\")\n",
    "print(\"\\rVectorized UDF: \", end=\"\") \n",
    "%timeit -r 1 -n 1 df_customer_1000.select(vect_udf_sm(\"C_NAME\").alias(\"CustIDs\")).write.mode(\"overwrite\").save_as_table(\"SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.CustIDs\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.sql(\"ALTER WAREHOUSE compute_wh SET warehouse_size='Large'\").collect()\n",
    "print(\"Use case 2.4: String Manipulation with Large WH and 1000x dataset\")\n",
    "print(\"\\rScalar UDF: \", end=\"\") \n",
    "%timeit -r 1 -n 1 df_customer_1000.select(udf_sm(\"C_NAME\").alias(\"CustIDs\")).write.mode(\"overwrite\").save_as_table(\"SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.CustIDs\")\n",
    "print(\"\\rVectorized UDF: \", end=\"\") \n",
    "%timeit -r 1 -n 1 df_customer_1000.select(vect_udf_sm(\"C_NAME\").alias(\"CustIDs\")).write.mode(\"overwrite\").save_as_table(\"SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.CustIDs\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Use Case 3: Regex Masking\n",
    "Another common use case is Regex Masking. To set this up, let's create a single set of UDFs for Customer100 & Customer1000 datasets. Here, we take the Customer Phone Number as input and masks the last 4 digits. As before, we will create a Normal and Vectorised UDF."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's create a slightly complex function using native python regex for string replacement/masking\n",
    "def mask_data(inp):\n",
    "    return re.sub('\\d{4}$', '****', inp)\n",
    "\n",
    "# Let's create a UDF on top of that\n",
    "udf_md = udf(mask_data, return_type=StringType(), input_types=[StringType()])\n",
    "\n",
    "# Let's vectorise the same UDF\n",
    "@udf()\n",
    "def vect_udf_md(inp: PandasSeries[str]) -> PandasSeries[str]:\n",
    "    return inp.apply(lambda x: mask_data(x))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As before, let's quickly execute these UDFs against our datasets 100 & 1000. We will re-run these UDFs based on our Dataset and VW sizing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.sql(\"ALTER WAREHOUSE compute_wh SET warehouse_size='Small'\").collect()\n",
    "print(\"Use case 3.1: Regex Masking with Small WH and 100x dataset\")\n",
    "print(\"\\rScalar UDF: \", end=\"\") \n",
    "%timeit -r 1 -n 1 df_customer_100.select(udf_md(\"C_PHONE\").alias(\"masked_phone_nums\")).write.mode(\"overwrite\").save_as_table(\"SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.masked_phone_data\")\n",
    "print(\"\\rVectorized UDF: \", end=\"\") \n",
    "%timeit -r 1 -n 1 df_customer_100.select(vect_udf_md(\"C_PHONE\").alias(\"masked_phone_nums\")).write.mode(\"overwrite\").save_as_table(\"SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.masked_phone_data\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.sql(\"ALTER WAREHOUSE compute_wh SET warehouse_size='Medium'\").collect()\n",
    "print(\"Use case 3.2: Regex Masking with Medium WH and 100x dataset\")\n",
    "print(\"\\rScalar UDF: \", end=\"\") \n",
    "%timeit -r 1 -n 1 df_customer_100.select(udf_md(\"C_PHONE\").alias(\"masked_phone_nums\")).write.mode(\"overwrite\").save_as_table(\"SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.masked_phone_data\")\n",
    "print(\"\\rVectorized UDF: \", end=\"\") \n",
    "%timeit -r 1 -n 1 df_customer_100.select(vect_udf_md(\"C_PHONE\").alias(\"masked_phone_nums\")).write.mode(\"overwrite\").save_as_table(\"SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.masked_phone_data\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.sql(\"ALTER WAREHOUSE compute_wh SET warehouse_size='Medium'\").collect()\n",
    "print(\"Use case 3.3: Regex Masking with Medium WH and 1000x dataset\")\n",
    "print(\"\\rScalar UDF: \", end=\"\") \n",
    "%timeit -r 1 -n 1 df_customer_1000.select(udf_md(\"C_PHONE\").alias(\"masked_phone_nums\")).write.mode(\"overwrite\").save_as_table(\"SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.masked_phone_data\")\n",
    "print(\"\\rVectorized UDF: \", end=\"\") \n",
    "%timeit -r 1 -n 1 df_customer_1000.select(vect_udf_md(\"C_PHONE\").alias(\"masked_phone_nums\")).write.mode(\"overwrite\").save_as_table(\"SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.masked_phone_data\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.sql(\"ALTER WAREHOUSE compute_wh SET warehouse_size='Large'\").collect()\n",
    "print(\"Use case 3.4: Regex Masking with Large WH and 1000x dataset\")\n",
    "print(\"\\rScalar UDF: \", end=\"\") \n",
    "%timeit -r 1 -n 1 df_customer_1000.select(udf_md(\"C_PHONE\").alias(\"masked_phone_nums\")).write.mode(\"overwrite\").save_as_table(\"SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.masked_phone_data\")\n",
    "print(\"\\rVectorized UDF: \", end=\"\") \n",
    "%timeit -r 1 -n 1 df_customer_1000.select(vect_udf_md(\"C_PHONE\").alias(\"masked_phone_nums\")).write.mode(\"overwrite\").save_as_table(\"SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.masked_phone_data\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To try another approach, let's vectorise the original UDF again, this time using the Pandas replace code.  After creation, check the performance against the two datasets.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@udf()\n",
    "def vect_udf_md(inp: PandasSeries[str]) -> PandasSeries[str]:\n",
    "    return inp.replace(to_replace ='\\d{4}$', value = '****', regex = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.sql(\"ALTER WAREHOUSE compute_wh SET warehouse_size='Medium'\").collect()\n",
    "print(\"Use case 3.5: Regex Masking with Pandas replace, using Medium WH and 100x then 1000x dataset\")\n",
    "print(\"\\rVectorized UDF on 100x: \", end=\"\") \n",
    "%timeit -r 1 -n 1 df_customer_100.select(vect_udf_md(\"C_PHONE\").alias(\"masked_phone_nums\")).write.mode(\"overwrite\").save_as_table(\"SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.masked_phone_data\")\n",
    "print(\"\\rVectorized UDF on 1000x: \", end=\"\") \n",
    "%timeit -r 1 -n 1 df_customer_1000.select(vect_udf_md(\"C_PHONE\").alias(\"masked_phone_nums\")).write.mode(\"overwrite\").save_as_table(\"SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.masked_phone_data\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Timestamp Manipulation\n",
    "The final example is changing a timestamp field. Let's create a single set of UDFs for Order100 & Order1000 datasets. Here, we take the Order Data as input and apply some basic Timestamp operations. As before, we will create a Normal and Vectorised UDF."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's try some timestamp manipulation\n",
    "def change_format(inp):\n",
    "    instantiate = datetime.strptime(inp, '%Y-%m-%d')\n",
    "    change_format = datetime.strptime(instantiate.strftime('%m/%d/%Y'), '%m/%d/%Y')\n",
    "    dt_add3_days = change_format + timedelta(days=10)\n",
    "    return dt_add3_days\n",
    "\n",
    "# Let's create a UDF on top of that\n",
    "udf_cf = udf(change_format, return_type=StringType(), input_types=[StringType()])\n",
    "\n",
    "# Let's vectorise the UDF - still using native python\n",
    "@udf()\n",
    "def vect_udf_cf(inp: PandasSeries[str]) -> PandasSeries[str]:\n",
    "    return inp.apply(lambda x: change_format(x))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As before, let's execute these UDFs against our datasets 100 & 1000. We will re-run these UDFs based on our Dataset and VW sizing.  \n",
    "\n",
    "**NOTE:**  **This use case is the most computationally intensive and will take the longest compared to the previous use cases.**\n",
    "The 1000 datasets take ~30 minutes to run each, so please be cognizant of credit usage and feel free to skip ahead to the analysis if desired. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.sql(\"ALTER WAREHOUSE compute_wh SET warehouse_size='Small'\").collect()\n",
    "print(\"Use case 4.1: Timestamp Manipulation with Small WH and 100x dataset\")\n",
    "print(\"\\rScalar UDF: \", end=\"\") \n",
    "%timeit -r 1 -n 1 df_order_100.select(udf_cf(\"O_ORDERDATE\").alias(\"NY_OrderDate\")).write.mode(\"overwrite\").save_as_table(\"SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.new_order_dates\")\n",
    "print(\"\\rVectorized UDF: \", end=\"\") \n",
    "%timeit -r 1 -n 1 df_order_100.select(vect_udf_cf(\"O_ORDERDATE\").alias(\"NY_OrderDate\")).write.mode(\"overwrite\").save_as_table(\"SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.new_order_dates\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.sql(\"ALTER WAREHOUSE compute_wh SET warehouse_size='Medium'\").collect()\n",
    "print(\"Use case 4.2: Timestamp Manipulation with Medium WH and 100x dataset\")\n",
    "print(\"\\rScalar UDF: \", end=\"\") \n",
    "%timeit -r 1 -n 1 df_order_100.select(udf_cf(\"O_ORDERDATE\").alias(\"NY_OrderDate\")).write.mode(\"overwrite\").save_as_table(\"SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.new_order_dates\")\n",
    "print(\"\\rVectorized UDF: \", end=\"\") \n",
    "%timeit -r 1 -n 1 df_order_100.select(vect_udf_cf(\"O_ORDERDATE\").alias(\"NY_OrderDate\")).write.mode(\"overwrite\").save_as_table(\"SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.new_order_dates\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.sql(\"ALTER WAREHOUSE compute_wh SET warehouse_size='Medium'\").collect()\n",
    "print(\"Use case 4.3: Timestamp Manipulation with Medium WH and 1000x dataset\")\n",
    "print(\"\\rScalar UDF: \", end=\"\") \n",
    "%timeit -r 1 -n 1 df_order_1000.select(udf_cf(\"O_ORDERDATE\").alias(\"NY_OrderDate\")).write.mode(\"overwrite\").save_as_table(\"SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.new_order_dates\")\n",
    "print(\"\\rVectorized UDF: \", end=\"\") \n",
    "%timeit -r 1 -n 1 df_order_1000.select(vect_udf_cf(\"O_ORDERDATE\").alias(\"NY_OrderDate\")).write.mode(\"overwrite\").save_as_table(\"SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.new_order_dates\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.sql(\"ALTER WAREHOUSE compute_wh SET warehouse_size='Large'\").collect()\n",
    "print(\"Use case 4.4: Timestamp Manipulation with Large WH and 1000x dataset\")\n",
    "print(\"\\rScalar UDF: \", end=\"\") \n",
    "%timeit -r 1 -n 1 df_order_1000.select(udf_cf(\"O_ORDERDATE\").alias(\"NY_OrderDate\")).write.mode(\"overwrite\").save_as_table(\"SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.new_order_dates\")\n",
    "print(\"\\rVectorized UDF: \", end=\"\") \n",
    "%timeit -r 1 -n 1 df_order_1000.select(vect_udf_cf(\"O_ORDERDATE\").alias(\"NY_OrderDate\")).write.mode(\"overwrite\").save_as_table(\"SNOWPARK_BEST_PRACTICES_LABS.PUBLIC.new_order_dates\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cleanup\n",
    "\n",
    "The following code block cleans up what was executed in this lab, by accomplishing the following:\n",
    " - Resetting the warehouse to SMALL\n",
    " - Drop and recreate the database to quickly remove all tables created from this lab.  \n",
    " - Close the session."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.sql(\"ALTER WAREHOUSE compute_wh SET warehouse_size='Small'\").collect()\n",
    "session.sql(\"DROP DATABASE SNOWPARK_BEST_PRACTICES_LABS\")\n",
    "session.sql(\"CREATE DATABASE SNOWPARK_BEST_PRACTICES_LABS\")\n",
    "\n",
    "if session:\n",
    "    session.close"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.15"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "1e026f6e58e7e7fcac433b20e0e98bc28079064f8b39889c46bb1bee90d36b9d"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}