Snowpark_For_Python_DE.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "06060903-e0b5-4970-be4d-799d5e7ba52f",
   "metadata": {
    "collapsed": false,
    "name": "cell18"
   },
   "source": [
    "# Project Overview\n",
    "\n",
    "Perform data analysis and data preparation tasks to train a Linear Regression model to predict future ROI (Return On Investment) of variable ad spend budgets across multiple channels including search, video, social media, and email using Snowpark for Python, Snowpark ML and Streamlit. By the end of the session, you will have an interactive web application deployed visualizing the ROI of different allocated advertising spend budgets.\n",
    "\n",
    "## Data Engineering -- Data Analysis and Data Preparation\n",
    "\n",
    "In this Notebook, we will focus on Data Engineering in Snowflake using Snowpark for Python.\n",
    "\n",
    "- Load data from Snowflake tables into Snowpark DataFrames\n",
    "- Perform Exploratory Data Analysis on Snowpark DataFrames\n",
    "- Pivot and Join data from multiple tables using Snowpark DataFrames\n",
    "- Demostrate how to automate data preparation using Snowflake Tasks\n",
    "- For environment setup including loading data into Snowflake tables, and step-by-step instructions, please refer to the [QuickStart Guide](https://quickstarts.snowflake.com/guide/getting_started_with_dataengineering_ml_using_snowpark_python/index.html#0)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75978241-4905-406c-a0c2-208bbcdbf2ad",
   "metadata": {
    "collapsed": false,
    "name": "cell19"
   },
   "source": [
    "## Import Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3775908f-ca36-4846-8f38-5adca39217f2",
   "metadata": {
    "codeCollapsed": false,
    "collapsed": false,
    "language": "python",
    "name": "cell1"
   },
   "outputs": [],
   "source": [
    "from snowflake.snowpark.session import Session\n",
    "from snowflake.snowpark.context import get_active_session\n",
    "from snowflake.snowpark.functions import month,year,col,sum\n",
    "from snowflake.snowpark.version import VERSION\n",
    "from snowflake.core import Root\n",
    "from snowflake.core.task import Task, StoredProcedureCall\n",
    "from snowflake.core.task.dagv1 import DAG, DAGTask, DAGOperation\n",
    "from snowflake.core import CreateMode\n",
    "\n",
    "# Misc\n",
    "from datetime import timedelta\n",
    "import json\n",
    "import logging \n",
    "logger = logging.getLogger(\"snowflake.snowpark.session\")\n",
    "logger.setLevel(logging.ERROR)\n",
    "\n",
    "session = get_active_session()\n",
    "# Add a query tag to the session. This helps with monitoring and troubleshooting.\n",
    "session.query_tag = {\"origin\":\"sf_sit-is\", \n",
    "                     \"name\":\"aiml_notebooks_ad_spend_roi\", \n",
    "                     \"version\":{\"major\":1, \"minor\":0},\n",
    "                     \"attributes\":{\"is_quickstart\":1, \"source\":\"notebook\"}}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acf35526-e031-44fe-8ced-5725a051f154",
   "metadata": {
    "collapsed": false,
    "name": "cell20"
   },
   "source": [
    "## Load Aggregated Campaign Spend Data from Snowflake table into Snowpark DataFrame\n",
    "\n",
    "Let's first load the campaign spend data. This table contains ad click data that has been aggregated to show daily spend across digital ad channels including search engines, social media, email and video.\n",
    "\n",
    "Note: Some other ways to load data in a Snowpark DataFrame\n",
    "\n",
    "- session.sql(\"select col1, col2... from tableName\")\n",
    "- session.read.options({\"field_delimiter\": \",\", \"skip_header\": 1}).schema(user_schema).csv(\"@mystage/testCSV.csv\")\n",
    "- session.read.parquet(\"@stageName/path/to/file\")\n",
    "- session.create_dataframe([1,2,3], schema=[\"col1\"])\n",
    "\n",
    "TIP: Learn more about [Snowpark DataFrames](https://docs.snowflake.com/en/developer-guide/snowpark/reference/python/latest/dataframe)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d50cbf4-0c8d-4950-86cb-114990437ac9",
   "metadata": {
    "codeCollapsed": false,
    "collapsed": false,
    "language": "python",
    "name": "cell2"
   },
   "outputs": [],
   "source": [
    "snow_df_spend = session.table('campaign_spend')\n",
    "snow_df_spend.queries"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "429243f4-d59a-43fd-b787-7bed1207729a",
   "metadata": {
    "collapsed": false,
    "name": "cell21"
   },
   "source": [
    "Actions like show(), collect(), count() send the DataFrame SQL for execution on the server\n",
    "\n",
    "Note: History object provides the query ID which can be helpful for debugging as well as the SQL query executed on the server."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c695373e-ac74-4b62-a1f1-08206cbd5c81",
   "metadata": {
    "codeCollapsed": false,
    "collapsed": false,
    "language": "python",
    "name": "cell3"
   },
   "outputs": [],
   "source": [
    "with session.query_history() as history:\n",
    "    snow_df_spend.show(20)\n",
    "history.queries"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "388f6cae-b2ea-4db5-b948-62f745e5853d",
   "metadata": {
    "collapsed": false,
    "name": "cell22"
   },
   "source": [
    "## Total Spend per Year and Month For All Channels\n",
    "\n",
    "Let's transform the data so we can see total cost per year/month per channel using group_by() and agg() Snowpark DataFrame functions.\n",
    "\n",
    "TIP: For a full list of functions, refer to the [documentation](https://docs.snowflake.com/en/developer-guide/snowpark/reference/python/latest/functions)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc22a0e7-83ca-470c-8430-60b5f9e1e4c1",
   "metadata": {
    "codeCollapsed": false,
    "collapsed": false,
    "language": "python",
    "name": "cell4"
   },
   "outputs": [],
   "source": [
    "# Stats per Month per Channel\n",
    "snow_df_spend_per_channel = snow_df_spend.group_by(year('DATE'), month('DATE'),'CHANNEL').agg(sum('TOTAL_COST').as_('TOTAL_COST')).\\\n",
    "    with_column_renamed('\"YEAR(DATE)\"',\"YEAR\").with_column_renamed('\"MONTH(DATE)\"',\"MONTH\").sort('YEAR','MONTH')\n",
    "\n",
    "snow_df_spend_per_channel.show(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b9a17f7-f557-4321-b318-4068c98cc68e",
   "metadata": {
    "collapsed": false,
    "name": "cell23"
   },
   "source": [
    "## Pivot on Channel: Total Spend Across All Channels\n",
    "\n",
    "Let's further transform the campaign spend data so that each row will represent total cost across all channels per year/month using pivot() and sum() Snowpark DataFrame functions. This transformation will enable us to join with the revenue table such that we will have our input features and target variable in a single table for model training.\n",
    "\n",
    "TIP: For a full list of functions, refer to the [documentation](https://docs.snowflake.com/en/developer-guide/snowpark/reference/python/latest/functions)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "813b07ec-cf73-44dd-bbf0-261268f73d8e",
   "metadata": {
    "collapsed": false,
    "language": "python",
    "name": "cell5"
   },
   "outputs": [],
   "source": [
    "snow_df_spend_per_month = snow_df_spend_per_channel.pivot('CHANNEL',['search_engine','social_media','video','email']).sum('TOTAL_COST').sort('YEAR','MONTH')\n",
    "snow_df_spend_per_month = snow_df_spend_per_month.select(\n",
    "    col(\"YEAR\"),\n",
    "    col(\"MONTH\"),\n",
    "    col(\"'search_engine'\").as_(\"SEARCH_ENGINE\"),\n",
    "    col(\"'social_media'\").as_(\"SOCIAL_MEDIA\"),\n",
    "    col(\"'video'\").as_(\"VIDEO\"),\n",
    "    col(\"'email'\").as_(\"EMAIL\")\n",
    ")\n",
    "snow_df_spend_per_month.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16abb4e0-4a66-415b-99b0-a12863906fbb",
   "metadata": {
    "collapsed": false,
    "name": "cell24"
   },
   "source": [
    "## Save Transformed Data into Snowflake Table\n",
    "\n",
    "Let's save the transformed data into a Snowflake table SPEND_PER_MONTH."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd9caea2-4825-4f86-bec7-80b3d9d07103",
   "metadata": {
    "codeCollapsed": false,
    "collapsed": false,
    "language": "python",
    "name": "cell6"
   },
   "outputs": [],
   "source": [
    "snow_df_spend_per_month.write.mode('overwrite').save_as_table('SPEND_PER_MONTH')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8a2e164-b2ba-44a0-9ace-fcb2503f27ac",
   "metadata": {
    "collapsed": false,
    "name": "cell25"
   },
   "source": [
    "## Automation: Run Campaign Spend Data Transformations As a Snowflake Task\n",
    "\n",
    "Note: Optionally you can run all these transformations as an automated task by deploying the code to Snowflake as a Snowpark Stored Procedure and executing it as a Snowflake Task.\n",
    "\n",
    "TIP: Learn more about [Stored Procedures](https://docs.snowflake.com/en/sql-reference/stored-procedures-python) and [Snowflake Tasks](https://docs.snowflake.com/en/sql-reference/sql/create-task)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f01b5c4-3f5b-4384-9051-e15ded939c48",
   "metadata": {
    "codeCollapsed": false,
    "collapsed": false,
    "language": "python",
    "name": "cell7"
   },
   "outputs": [],
   "source": [
    "def campaign_spend_data_pipeline(session: Session) -> str:\n",
    "  # DATA TRANSFORMATIONS\n",
    "  # Perform the following actions to transform the data\n",
    "\n",
    "  # Load the campaign spend data\n",
    "  snow_df_spend_t = session.table('campaign_spend')\n",
    "\n",
    "  # Transform the data so we can see total cost per year/month per channel using group_by() and agg() Snowpark DataFrame functions\n",
    "  snow_df_spend_per_channel_t = snow_df_spend_t.group_by(year('DATE'), month('DATE'),'CHANNEL').agg(sum('TOTAL_COST').as_('TOTAL_COST')).\\\n",
    "      with_column_renamed('\"YEAR(DATE)\"',\"YEAR\").with_column_renamed('\"MONTH(DATE)\"',\"MONTH\").sort('YEAR','MONTH')\n",
    "\n",
    "  # Transform the data so that each row will represent total cost across all channels per year/month using pivot() and sum() Snowpark DataFrame functions\n",
    "  snow_df_spend_per_month_t = snow_df_spend_per_channel_t.pivot('CHANNEL',['search_engine','social_media','video','email']).sum('TOTAL_COST').sort('YEAR','MONTH')\n",
    "  snow_df_spend_per_month_t = snow_df_spend_per_month_t.select(\n",
    "      col(\"YEAR\"),\n",
    "      col(\"MONTH\"),\n",
    "      col(\"'search_engine'\").as_(\"SEARCH_ENGINE\"),\n",
    "      col(\"'social_media'\").as_(\"SOCIAL_MEDIA\"),\n",
    "      col(\"'video'\").as_(\"VIDEO\"),\n",
    "      col(\"'email'\").as_(\"EMAIL\")\n",
    "  )\n",
    "\n",
    "  # Save transformed data\n",
    "  snow_df_spend_per_month_t.write.mode('overwrite').save_as_table('SPEND_PER_MONTH')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c7035a7-3609-402b-8da6-8104cb7b748e",
   "metadata": {
    "collapsed": false,
    "language": "python",
    "name": "cell8"
   },
   "outputs": [],
   "source": [
    "# Register data pipeline function as a task\n",
    "root = Root(session)\n",
    "my_task = Task(name='campaign_spend_data_pipeline_task'\n",
    "               , definition=StoredProcedureCall(\n",
    "                   campaign_spend_data_pipeline, stage_location='@dash_sprocs'\n",
    "               )\n",
    "               , warehouse='DASH_S'\n",
    "               , schedule=timedelta(minutes=3))\n",
    "\n",
    "tasks = root.databases[session.get_current_database()].schemas[session.get_current_schema()].tasks\n",
    "task_res = tasks.create(my_task,mode=CreateMode.or_replace)\n",
    "\n",
    "# By default a Task is suspended and we need to resume it if we want it run based on the schema. Note that we can still execute a task by calling the execute method.\n",
    "task_res.execute()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "155beaf3-16bd-489f-a1ce-538181dbf0b6",
   "metadata": {
    "collapsed": false,
    "name": "cell26"
   },
   "source": [
    "## Total Revenue per Year And Month\n",
    "\n",
    "Now let's load revenue table and transform the data into revenue per year/month using group_by() and agg() functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d4578a9-3330-4e03-9b4f-d2eed87b052f",
   "metadata": {
    "codeCollapsed": false,
    "collapsed": false,
    "language": "python",
    "name": "cell9"
   },
   "outputs": [],
   "source": [
    "snow_df_revenue = session.table('monthly_revenue')\n",
    "snow_df_revenue_per_month = snow_df_revenue.group_by('YEAR','MONTH').agg(sum('REVENUE')).sort('YEAR','MONTH').with_column_renamed('SUM(REVENUE)','REVENUE')\n",
    "snow_df_revenue_per_month.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "630e468c-4781-4412-9f51-b894690d39e7",
   "metadata": {
    "collapsed": false,
    "name": "cell27"
   },
   "source": [
    "## Join Total Spend and Total Revenue per Year and Month Across All Channels\n",
    "\n",
    "Next let's join this revenue data with the transformed campaign spend data so that our input features (i.e. cost per channel) and target variable (i.e. revenue) can be loaded into a single table for model training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23705e14-a843-4b92-8b96-b14f9bdf2296",
   "metadata": {
    "collapsed": false,
    "language": "python",
    "name": "cell10"
   },
   "outputs": [],
   "source": [
    "snow_df_spend_and_revenue_per_month = snow_df_spend_per_month.join(snow_df_revenue_per_month, [\"YEAR\",\"MONTH\"])\n",
    "snow_df_spend_and_revenue_per_month.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a19f2ef7-c6d7-4d73-9d31-987da24b0f36",
   "metadata": {
    "collapsed": false,
    "name": "cell28"
   },
   "source": [
    "## Examine Snowpark DataFrame Query and Execution Plan\n",
    "\n",
    "Snowpark makes is really convenient to look at the DataFrame query and execution plan using explain() Snowpark DataFrame function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "186edbc9-d878-4ec9-a03f-d694ab9b0494",
   "metadata": {
    "codeCollapsed": false,
    "collapsed": false,
    "language": "python",
    "name": "cell11"
   },
   "outputs": [],
   "source": [
    "snow_df_spend_and_revenue_per_month.explain()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d18128c0-1c16-4710-be99-909329d418d3",
   "metadata": {
    "collapsed": false,
    "name": "cell29"
   },
   "source": [
    "## Save Transformed Data into Snowflake Table\n",
    "\n",
    "Let's save the transformed data into a Snowflake table SPEND_AND_REVENUE_PER_MONTH."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8d3d424-ec65-4120-9dfc-0da0a21b5289",
   "metadata": {
    "codeCollapsed": false,
    "collapsed": false,
    "language": "python",
    "name": "cell12"
   },
   "outputs": [],
   "source": [
    "snow_df_spend_and_revenue_per_month.write.mode('overwrite').save_as_table('SPEND_AND_REVENUE_PER_MONTH')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37b86074-4332-4d6b-9314-44e93964a282",
   "metadata": {
    "collapsed": false,
    "name": "cell30"
   },
   "source": [
    "## Automation: Run Monthly Revenue Data Transformations As a Snowflake Task (DAG)\n",
    "\n",
    "Note: Optionally you can run all these transformations as an automated task by deploying the code to Snowflake as a Snowpark Stored Procedure and executing it as a Snowflake Task. By using a DAG we can run it AFTER campaign_spend_data_pipeline_task.\n",
    "\n",
    "TIP: Learn more about [Stored Procedures](https://docs.snowflake.com/en/sql-reference/stored-procedures-python) and [Snowflake Tasks](https://docs.snowflake.com/en/sql-reference/sql/create-task)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db9ffcd9-ce7d-425a-86e4-21b42f9655cc",
   "metadata": {
    "codeCollapsed": false,
    "collapsed": false,
    "language": "python",
    "name": "cell13"
   },
   "outputs": [],
   "source": [
    "def monthly_revenue_data_pipeline(session: Session) -> str:\n",
    "  # Load revenue table and transform the data into revenue per year/month using group_by and agg() functions\n",
    "  snow_df_spend_per_month_t = session.table('spend_per_month')\n",
    "  snow_df_revenue_t = session.table('monthly_revenue')\n",
    "  snow_df_revenue_per_month_t = snow_df_revenue_t.group_by('YEAR','MONTH').agg(sum('REVENUE')).sort('YEAR','MONTH').with_column_renamed('SUM(REVENUE)','REVENUE')\n",
    "\n",
    "  # Join revenue data with the transformed campaign spend data so that our input features (i.e. cost per channel) and target variable (i.e. revenue) can be loaded into a single table for model training\n",
    "  snow_df_spend_and_revenue_per_month_t = snow_df_spend_per_month_t.join(snow_df_revenue_per_month_t, [\"YEAR\",\"MONTH\"])\n",
    "\n",
    "  # SAVE in a new table for the next task\n",
    "  snow_df_spend_and_revenue_per_month_t.write.mode('overwrite').save_as_table('SPEND_AND_REVENUE_PER_MONTH')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcb56324-fe90-47e5-bf8c-330da47e38cf",
   "metadata": {
    "collapsed": false,
    "name": "cell31"
   },
   "source": [
    "Note: Since monthly_revenue_data_pipeline is depened on that campaign_spend_data_pipeline is executed first we want to create a DAG to make sure they run in the correct order."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6892443-055b-4cc3-9b12-5ea58a15f3d8",
   "metadata": {
    "codeCollapsed": false,
    "collapsed": false,
    "language": "python",
    "name": "cell14"
   },
   "outputs": [],
   "source": [
    "# Delete the previous task\n",
    "task_res.delete()\n",
    "\n",
    "with DAG(\"de_pipeline_dag\", schedule=timedelta(minutes=3)) as dag:\n",
    "    # Create a task that runs our first pipleine\n",
    "    dag_spend_task = DAGTask(name='campaign_spend_data_pipeline_task'\n",
    "                        , definition=StoredProcedureCall(\n",
    "                                    campaign_spend_data_pipeline, stage_location='@dash_sprocs'\n",
    "                                )\n",
    "                        ,warehouse='DASH_S'\n",
    "                        )\n",
    "    # Create a task that runs our second pipleine\n",
    "    dag_revenue_task = DAGTask(name='monthly_revenue_data_pipeline'\n",
    "                          , definition=StoredProcedureCall(\n",
    "                                monthly_revenue_data_pipeline, stage_location='@dash_sprocs'\n",
    "                            )\n",
    "                        ,warehouse='DASH_S'\n",
    "                        )\n",
    "# Shift right and left operators can specify task relationships.\n",
    "dag_spend_task >> dag_revenue_task  # dag_spend_task is a predecessor of dag_revenue_task\n",
    "\n",
    "schema = root.databases[session.get_current_database()].schemas[session.get_current_schema()]\n",
    "dag_op = DAGOperation(schema)\n",
    "\n",
    "dag_op.deploy(dag,mode=CreateMode.or_replace)\n",
    "\n",
    "# A DAG is not suspended by default so we will suspend the root task that will suspend the full DAG\n",
    "root_task = tasks[\"DE_PIPELINE_DAG\"]\n",
    "root_task.suspend()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7aac0fb4-3a27-4800-bbd2-72a39fb264ea",
   "metadata": {
    "collapsed": false,
    "name": "cell32"
   },
   "source": [
    "## Run DAG\n",
    "\n",
    "Note that we can manually run DAGs even if they are suspended"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab79fde1-28e2-41e2-a3bd-4377d392c13e",
   "metadata": {
    "codeCollapsed": false,
    "language": "python",
    "name": "cell15"
   },
   "outputs": [],
   "source": [
    "# dag_op.run(dag)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd0fa7d5-b026-43e4-aefd-d4ae259e8b23",
   "metadata": {
    "collapsed": false,
    "name": "cell33"
   },
   "source": [
    "## Resume DAG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d53b663-04bd-4f3f-9adb-7f722087c6b5",
   "metadata": {
    "language": "python",
    "name": "cell16"
   },
   "outputs": [],
   "source": [
    "# root_task = tasks[\"DE_PIPELINE_DAG\"]\n",
    "# root_task.resume()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8d165ce-eb80-4e71-ac82-4618e8ab7f37",
   "metadata": {
    "collapsed": false,
    "name": "cell35"
   },
   "source": [
    "## Suspend Tasks\n",
    "\n",
    "Note: For the sake of this lab, if you resume the above tasks, suspend them to avoid unecessary resource utilization by executing the following commands."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbde842a-562d-430f-9ded-44fd2c7d501d",
   "metadata": {
    "codeCollapsed": false,
    "language": "python",
    "name": "cell17"
   },
   "outputs": [],
   "source": [
    "# root_task = tasks[\"DE_PIPELINE_DAG\"]\n",
    "# root_task.suspend()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Streamlit Notebook",
   "name": "streamlit"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}