01_Data_Engineering.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Engineering\n",
    "We begin where all ML use cases do: data engineering. In this section of the demo, we will utilize Snowpark's Python client-side Dataframe API to build an **ELT pipeline**.  We will extract the data from the source system (s3), load it into snowflake and add transformations to clean the data before analysis. \n",
    "\n",
    "The data engineer has been told that there is historical data going back to 2013 and new data will be made available at the end of each month. \n",
    "\n",
    "Input: Historical bulk data at `https://s3.amazonaws.com/tripdata/`. Incremental data to be loaded one month at a time.  \n",
    "Output: `trips` table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import snowflake.snowpark as snp\n",
    "from snowflake.snowpark import functions as F\n",
    "from snowflake.snowpark import types as T\n",
    "\n",
    "import pandas as pd\n",
    "from datetime import datetime\n",
    "import requests\n",
    "from zipfile import ZipFile\n",
    "from io import BytesIO\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Load  credentials and connect to Snowflake\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will utilize a simple json file to store our credentials. This should **never** be done in production and is for demo purposes only."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dags.snowpark_connection import snowpark_connect\n",
    "session, state_dict = snowpark_connect('./include/state.json')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Create a stage for loading data to Snowflake"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "state_dict['load_stage_name']='LOAD_STAGE' \n",
    "state_dict['download_base_url']='https://s3.amazonaws.com/tripdata/'\n",
    "state_dict['trips_table_name']='TRIPS'\n",
    "state_dict['load_table_name'] = 'RAW_'\n",
    "\n",
    "import json\n",
    "with open('./include/state.json', 'w') as sdf:\n",
    "    json.dump(state_dict, sdf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def reset_database(session, state_dict:dict, prestaged=False):\n",
    "    _ = session.sql('CREATE OR REPLACE DATABASE '+state_dict['connection_parameters']['database']).collect()\n",
    "    _ = session.sql('CREATE SCHEMA '+state_dict['connection_parameters']['schema']).collect() \n",
    "\n",
    "    if prestaged:\n",
    "        sql_cmd = 'CREATE OR REPLACE STAGE '+state_dict['load_stage_name']+\\\n",
    "                  ' url='+state_dict['connection_parameters']['download_base_url']\n",
    "        _ = session.sql(sql_cmd).collect()\n",
    "    else: \n",
    "        _ = session.sql('CREATE STAGE IF NOT EXISTS '+state_dict['load_stage_name']).collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reset_database(session, state_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Extract:  \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a list of files to download and upload to stage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from datetime import datetime\n",
    "\n",
    "#For files like 201306-citibike-tripdata.zip\n",
    "date_range1 = pd.period_range(start=datetime.strptime(\"201306\", \"%Y%m\"), \n",
    "                             end=datetime.strptime(\"201612\", \"%Y%m\"), \n",
    "                             freq='M').strftime(\"%Y%m\")\n",
    "file_name_end1 = '-citibike-tripdata.zip'\n",
    "files_to_download = [date+file_name_end1 for date in date_range1.to_list()]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Starting in January 2017 Citibike changed the format of the file name."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#For files like 201701-citibike-tripdata.csv.zip\n",
    "date_range2 = pd.period_range(start=datetime.strptime(\"201701\", \"%Y%m\"), \n",
    "                             end=datetime.strptime(\"202112\", \"%Y%m\"), \n",
    "                             freq='M').strftime(\"%Y%m\")\n",
    "file_name_end2 = '-citibike-tripdata.csv.zip'\n",
    "files_to_download = files_to_download + [date+file_name_end2 for date in date_range2.to_list()]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For development purposes we will start with loading just a couple of files.  We will create a bulk load process afterwards."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "files_to_download = [files_to_download[i] for i in [0,102]] #19,50,100,102]]\n",
    "files_to_download"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.use_warehouse(state_dict['compute_parameters']['fe_warehouse'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "schema1_download_files = list()\n",
    "schema2_download_files = list()\n",
    "schema2_start_date = datetime.strptime('202102', \"%Y%m\")\n",
    "\n",
    "for file_name in files_to_download:\n",
    "    file_start_date = datetime.strptime(file_name.split(\"-\")[0], \"%Y%m\")\n",
    "    if file_start_date < schema2_start_date:\n",
    "        schema1_download_files.append(file_name)\n",
    "    else:\n",
    "        schema2_download_files.append(file_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "schema1_download_files, schema2_download_files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "schema1_load_stage = state_dict['load_stage_name']+'/schema1/'\n",
    "schema2_load_stage = state_dict['load_stage_name']+'/schema2/'\n",
    "\n",
    "schema1_files_to_load = list()\n",
    "for zip_file_name in schema1_download_files:\n",
    "    \n",
    "    url = state_dict['download_base_url']+zip_file_name\n",
    "    \n",
    "    print('Downloading and unzipping: '+url)\n",
    "    r = requests.get(url)\n",
    "    file = ZipFile(BytesIO(r.content))\n",
    "    csv_file_name=file.namelist()[0]\n",
    "    file.extract(csv_file_name)\n",
    "    file.close()\n",
    "    \n",
    "    print('Putting '+csv_file_name+' to stage: '+schema1_load_stage)\n",
    "    session.file.put(local_file_name=csv_file_name, \n",
    "                     stage_location=schema1_load_stage, \n",
    "                     source_compression='NONE', \n",
    "                     overwrite=True)\n",
    "    schema1_files_to_load.append(csv_file_name)\n",
    "    os.remove(csv_file_name)\n",
    "    \n",
    "schema2_files_to_load = list()\n",
    "for zip_file_name in schema2_download_files:\n",
    "    \n",
    "    url = state_dict['download_base_url']+zip_file_name\n",
    "    \n",
    "    print('Downloading and unzipping: '+url)\n",
    "    r = requests.get(url)\n",
    "    file = ZipFile(BytesIO(r.content))\n",
    "    csv_file_name=file.namelist()[0]\n",
    "    file.extract(csv_file_name)\n",
    "    file.close()\n",
    "    \n",
    "    print('Putting '+csv_file_name+' to stage: '+schema2_load_stage)\n",
    "    session.file.put(local_file_name=csv_file_name, \n",
    "                     stage_location=schema2_load_stage, \n",
    "                     source_compression='NONE', \n",
    "                     overwrite=True)\n",
    "    schema2_files_to_load.append(csv_file_name)\n",
    "    os.remove(csv_file_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.sql(\"list @\"+state_dict['load_stage_name']+\" pattern='.*20.*[.]gz'\").collect()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. Load: \n",
    "Load raw as all string type.  We will fix data types in the transform stage.\n",
    "\n",
    "There are two schema types so we will create two ingest tables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Upper case fields are common to both schemas.\n",
    "#Schema from 2013 to 2021\n",
    "load_schema1 = T.StructType([T.StructField(\"tripduration\", T.StringType()),\n",
    "                             T.StructField(\"STARTTIME\", T.StringType()), \n",
    "                             T.StructField(\"STOPTIME\", T.StringType()), \n",
    "                             T.StructField(\"START_STATION_ID\", T.StringType()),\n",
    "                             T.StructField(\"START_STATION_NAME\", T.StringType()), \n",
    "                             T.StructField(\"START_STATION_LATITUDE\", T.StringType()),\n",
    "                             T.StructField(\"START_STATION_LONGITUDE\", T.StringType()),\n",
    "                             T.StructField(\"END_STATION_ID\", T.StringType()),\n",
    "                             T.StructField(\"END_STATION_NAME\", T.StringType()), \n",
    "                             T.StructField(\"END_STATION_LATITUDE\", T.StringType()),\n",
    "                             T.StructField(\"END_STATION_LONGITUDE\", T.StringType()),\n",
    "                             T.StructField(\"bike_id\", T.StringType()),\n",
    "                             T.StructField(\"USERTYPE\", T.StringType()), \n",
    "                             T.StructField(\"birth_year\", T.StringType()),\n",
    "                             T.StructField(\"gender\", T.StringType())])\n",
    "\n",
    "#starting in February 2021 the schema changed\n",
    "load_schema2 = T.StructType([T.StructField(\"ride_id\", T.StringType()), \n",
    "                             T.StructField(\"rideable_type\", T.StringType()), \n",
    "                             T.StructField(\"STARTTIME\", T.StringType()), \n",
    "                             T.StructField(\"STOPTIME\", T.StringType()), \n",
    "                             T.StructField(\"START_STATION_NAME\", T.StringType()), \n",
    "                             T.StructField(\"START_STATION_ID\", T.StringType()),\n",
    "                             T.StructField(\"END_STATION_NAME\", T.StringType()), \n",
    "                             T.StructField(\"END_STATION_ID\", T.StringType()),\n",
    "                             T.StructField(\"START_STATION_LATITUDE\", T.StringType()),\n",
    "                             T.StructField(\"START_STATION_LONGITUDE\", T.StringType()),\n",
    "                             T.StructField(\"END_STATION_LATITUDE\", T.StringType()),\n",
    "                             T.StructField(\"END_STATION_LONGITUDE\", T.StringType()),\n",
    "                             T.StructField(\"USERTYPE\", T.StringType())])\n",
    "\n",
    "trips_table_schema = T.StructType([T.StructField(\"STARTTIME\", T.StringType()), \n",
    "                             T.StructField(\"STOPTIME\", T.StringType()), \n",
    "                             T.StructField(\"START_STATION_NAME\", T.StringType()), \n",
    "                             T.StructField(\"START_STATION_ID\", T.StringType()),\n",
    "                             T.StructField(\"END_STATION_NAME\", T.StringType()), \n",
    "                             T.StructField(\"END_STATION_ID\", T.StringType()),\n",
    "                             T.StructField(\"START_STATION_LATITUDE\", T.StringType()),\n",
    "                             T.StructField(\"START_STATION_LONGITUDE\", T.StringType()),\n",
    "                             T.StructField(\"END_STATION_LATITUDE\", T.StringType()),\n",
    "                             T.StructField(\"END_STATION_LONGITUDE\", T.StringType()),\n",
    "                             T.StructField(\"USERTYPE\", T.StringType())])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create empty tables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.create_dataframe([[None]*len(load_schema1.names)], schema=load_schema1)\\\n",
    "       .na.drop()\\\n",
    "       .write\\\n",
    "       .save_as_table(state_dict['load_table_name']+'schema1')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.create_dataframe([[None]*len(load_schema2.names)], schema=load_schema2)\\\n",
    "       .na.drop()\\\n",
    "       .write\\\n",
    "       .save_as_table(state_dict['load_table_name']+'schema2')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load schema1 data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "csv_file_format_options = {\"FIELD_OPTIONALLY_ENCLOSED_BY\": \"'\\\"'\", \"skip_header\": 1}\n",
    "\n",
    "loaddf = session.read.option(\"SKIP_HEADER\", 1)\\\n",
    "                     .option(\"FIELD_OPTIONALLY_ENCLOSED_BY\", \"\\042\")\\\n",
    "                     .option(\"COMPRESSION\", \"GZIP\")\\\n",
    "                     .option(\"NULL_IF\", \"\\\\\\\\N\")\\\n",
    "                     .option(\"NULL_IF\", \"NULL\")\\\n",
    "                     .option(\"pattern\", \"'.*20.*[.]gz'\")\\\n",
    "                     .schema(load_schema1)\\\n",
    "                     .csv('@'+schema1_load_stage)\\\n",
    "                     .copy_into_table(state_dict['load_table_name']+str('schema1'), \n",
    "                                      format_type_options=csv_file_format_options)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load schema2 data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loaddf = session.read.option(\"SKIP_HEADER\", 1)\\\n",
    "                     .option(\"FIELD_OPTIONALLY_ENCLOSED_BY\", \"\\042\")\\\n",
    "                     .option(\"COMPRESSION\", \"GZIP\")\\\n",
    "                     .option(\"NULL_IF\", \"\\\\\\\\N\")\\\n",
    "                     .option(\"NULL_IF\", \"NULL\")\\\n",
    "                     .option(\"pattern\", \"'.*20.*[.]gz'\")\\\n",
    "                     .schema(load_schema2)\\\n",
    "                     .csv('@'+schema2_load_stage)\\\n",
    "                     .copy_into_table(state_dict['load_table_name']+str('schema2'), \n",
    "                                      format_type_options=csv_file_format_options)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5. Transform:\n",
    "We have the raw data loaded. Now let's transform this data and clean it up. This will push the data to a final \\\"transformed\\\" table to be consumed by our Data Science team.  First we start by combining the two tables with the common columns.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trips_table_schema_names = [field.name for field in trips_table_schema.fields]\n",
    "transdf1 = session.table(state_dict['load_table_name']+'schema1')[trips_table_schema_names]\n",
    "transdf2 = session.table(state_dict['load_table_name']+'schema2')[trips_table_schema_names]\n",
    "transdf = transdf1.union_by_name(transdf2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are three different date formats \"2014-08-10 15:21:22\", \"1/1/2015 1:30\" and \"12/1/2014 02:04:53\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "date_format_2 = \"1/1/2015 [0-9]:.*$\"      #1/1/2015 1:30 -> #M*M/D*D/YYYY H*H:M*M(:SS)*\n",
    "date_format_3 = \"1/1/2015 [0-9][0-9]:.*$\" #1/1/2015 10:30 -> #M*M/D*D/YYYY H*H:M*M(:SS)*\n",
    "date_format_4 = \"12/1/2014.*\"             #12/1/2014 02:04:53 -> M*M/D*D/YYYY \n",
    "\n",
    "#Change all dates to YYYY-MM-DD HH:MI:SS format\n",
    "date_format_match = \"^([0-9]?[0-9])/([0-9]?[0-9])/([0-9][0-9][0-9][0-9]) ([0-9]?[0-9]):([0-9][0-9])(:[0-9][0-9])?.*$\"\n",
    "date_format_repl = \"\\\\3-\\\\1-\\\\2 \\\\4:\\\\5\\\\6\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transdf.with_column('STARTTIME', F.regexp_replace(F.col('STARTTIME'),\n",
    "                                            F.lit(date_format_match), \n",
    "                                            F.lit(date_format_repl)))\\\n",
    "      .with_column('STARTTIME', F.to_timestamp('STARTTIME'))\\\n",
    "      .with_column('STOPTIME', F.regexp_replace(F.col('STOPTIME'),\n",
    "                                            F.lit(date_format_match), \n",
    "                                            F.lit(date_format_repl)))\\\n",
    "      .with_column('STOPTIME', F.to_timestamp('STOPTIME'))\\\n",
    "      .select(F.col('STARTTIME'), \n",
    "              F.col('STOPTIME'), \n",
    "              F.col('START_STATION_ID'), \n",
    "              F.col('START_STATION_NAME'), \n",
    "              F.col('START_STATION_LATITUDE'), \n",
    "              F.col('START_STATION_LONGITUDE'), \n",
    "              F.col('END_STATION_ID'), \n",
    "              F.col('END_STATION_NAME'), F.col('END_STATION_LATITUDE'), \n",
    "              F.col('END_STATION_LONGITUDE'), \n",
    "              F.col('USERTYPE'))\\\n",
    "      .write.mode('overwrite').save_as_table(state_dict['trips_table_name'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "testdf = session.table(state_dict['trips_table_name'])\n",
    "testdf.schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "testdf.count()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6. Export code in functional modules for MLOps and orchestration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile dags/elt.py\n",
    "def schema1_definition():\n",
    "    from snowflake.snowpark import types as T\n",
    "    load_schema1 = T.StructType([T.StructField(\"TRIPDURATION\", T.StringType()),\n",
    "                             T.StructField(\"STARTTIME\", T.StringType()), \n",
    "                             T.StructField(\"STOPTIME\", T.StringType()), \n",
    "                             T.StructField(\"START_STATION_ID\", T.StringType()),\n",
    "                             T.StructField(\"START_STATION_NAME\", T.StringType()), \n",
    "                             T.StructField(\"START_STATION_LATITUDE\", T.StringType()),\n",
    "                             T.StructField(\"START_STATION_LONGITUDE\", T.StringType()),\n",
    "                             T.StructField(\"END_STATION_ID\", T.StringType()),\n",
    "                             T.StructField(\"END_STATION_NAME\", T.StringType()), \n",
    "                             T.StructField(\"END_STATION_LATITUDE\", T.StringType()),\n",
    "                             T.StructField(\"END_STATION_LONGITUDE\", T.StringType()),\n",
    "                             T.StructField(\"BIKEID\", T.StringType()),\n",
    "                             T.StructField(\"USERTYPE\", T.StringType()), \n",
    "                             T.StructField(\"BIRTH_YEAR\", T.StringType()),\n",
    "                             T.StructField(\"GENDER\", T.StringType())])\n",
    "    return load_schema1\n",
    "\n",
    "def schema2_definition():\n",
    "    from snowflake.snowpark import types as T\n",
    "    load_schema2 = T.StructType([T.StructField(\"ride_id\", T.StringType()), \n",
    "                             T.StructField(\"rideable_type\", T.StringType()), \n",
    "                             T.StructField(\"STARTTIME\", T.StringType()), \n",
    "                             T.StructField(\"STOPTIME\", T.StringType()), \n",
    "                             T.StructField(\"START_STATION_NAME\", T.StringType()), \n",
    "                             T.StructField(\"START_STATION_ID\", T.StringType()),\n",
    "                             T.StructField(\"END_STATION_NAME\", T.StringType()), \n",
    "                             T.StructField(\"END_STATION_ID\", T.StringType()),\n",
    "                             T.StructField(\"START_STATION_LATITUDE\", T.StringType()),\n",
    "                             T.StructField(\"START_STATION_LONGITUDE\", T.StringType()),\n",
    "                             T.StructField(\"END_STATION_LATITUDE\", T.StringType()),\n",
    "                             T.StructField(\"END_STATION_LONGITUDE\", T.StringType()),\n",
    "                             T.StructField(\"USERTYPE\", T.StringType())])\n",
    "    return load_schema2\n",
    "\n",
    "def conformed_schema():\n",
    "    from snowflake.snowpark import types as T\n",
    "    trips_table_schema = T.StructType([T.StructField(\"STARTTIME\", T.StringType()), \n",
    "                             T.StructField(\"STOPTIME\", T.StringType()), \n",
    "                             T.StructField(\"START_STATION_NAME\", T.StringType()), \n",
    "                             T.StructField(\"START_STATION_ID\", T.StringType()),\n",
    "                             T.StructField(\"END_STATION_NAME\", T.StringType()), \n",
    "                             T.StructField(\"END_STATION_ID\", T.StringType()),\n",
    "                             T.StructField(\"START_STATION_LATITUDE\", T.StringType()),\n",
    "                             T.StructField(\"START_STATION_LONGITUDE\", T.StringType()),\n",
    "                             T.StructField(\"END_STATION_LATITUDE\", T.StringType()),\n",
    "                             T.StructField(\"END_STATION_LONGITUDE\", T.StringType()),\n",
    "                             T.StructField(\"USERTYPE\", T.StringType())])\n",
    "    return trips_table_schema\n",
    "\n",
    "def extract_trips_to_stage(session, files_to_download: list, download_base_url: str, load_stage_name:str):\n",
    "    import os \n",
    "    import requests\n",
    "    from zipfile import ZipFile\n",
    "    import gzip\n",
    "    from datetime import datetime\n",
    "    from io import BytesIO\n",
    "    \n",
    "    schema1_download_files = list()\n",
    "    schema2_download_files = list()\n",
    "    schema2_start_date = datetime.strptime('202102', \"%Y%m\")\n",
    "\n",
    "    for file_name in files_to_download:\n",
    "        file_start_date = datetime.strptime(file_name.split(\"-\")[0], \"%Y%m\")\n",
    "        if file_start_date < schema2_start_date:\n",
    "            schema1_download_files.append(file_name)\n",
    "        else:\n",
    "            schema2_download_files.append(file_name)\n",
    "         \n",
    "        \n",
    "    schema1_load_stage = load_stage_name+'/schema1/'\n",
    "    schema1_files_to_load = list()\n",
    "    for zip_file_name in schema1_download_files:\n",
    "\n",
    "        url = download_base_url+zip_file_name\n",
    "\n",
    "        print('Downloading and unzipping: '+url)\n",
    "        r = requests.get(url)\n",
    "        file = ZipFile(BytesIO(r.content))\n",
    "        csv_file_name=file.namelist()[0]\n",
    "        file.extract(csv_file_name)\n",
    "        file.close()\n",
    "\n",
    "        print('Putting '+csv_file_name+' to stage: '+schema1_load_stage)\n",
    "        session.file.put(local_file_name=csv_file_name, \n",
    "                         stage_location=schema1_load_stage, \n",
    "                         source_compression='NONE', \n",
    "                         overwrite=True)\n",
    "        schema1_files_to_load.append(csv_file_name)\n",
    "        os.remove(csv_file_name)\n",
    "\n",
    "        \n",
    "    schema2_load_stage = load_stage_name+'/schema2/'\n",
    "    schema2_files_to_load = list()\n",
    "    for zip_file_name in schema2_download_files:\n",
    "\n",
    "        url = download_base_url+zip_file_name\n",
    "\n",
    "        print('Downloading and unzipping: '+url)\n",
    "        r = requests.get(url)\n",
    "        file = ZipFile(BytesIO(r.content))\n",
    "        csv_file_name=file.namelist()[0]\n",
    "        file.extract(csv_file_name)\n",
    "        file.close()\n",
    "\n",
    "        print('Putting '+csv_file_name+' to stage: '+schema2_load_stage)\n",
    "        session.file.put(local_file_name=csv_file_name, \n",
    "                         stage_location=schema2_load_stage, \n",
    "                         source_compression='NONE', \n",
    "                         overwrite=True)\n",
    "        schema2_files_to_load.append(csv_file_name)\n",
    "        os.remove(csv_file_name)\n",
    "        \n",
    "    load_stage_names = {'schema1' : schema1_load_stage, 'schema2' : schema2_load_stage}\n",
    "    files_to_load = {'schema1': schema1_files_to_load, 'schema2': schema2_files_to_load}\n",
    "\n",
    "    return load_stage_names, files_to_load\n",
    "    \n",
    "def load_trips_to_raw(session, files_to_load:dict, load_stage_names:dict, load_table_name:str):\n",
    "    from snowflake.snowpark import functions as F\n",
    "    from snowflake.snowpark import types as T\n",
    "    from datetime import datetime\n",
    "    \n",
    "    csv_file_format_options = {\"FIELD_OPTIONALLY_ENCLOSED_BY\": \"'\\\"'\", \"skip_header\": 1}\n",
    "\n",
    "    if len(files_to_load['schema1']) > 0:\n",
    "        load_schema1 = schema1_definition()\n",
    "        loaddf = session.read.option(\"SKIP_HEADER\", 1)\\\n",
    "                         .option(\"FIELD_OPTIONALLY_ENCLOSED_BY\", \"\\042\")\\\n",
    "                         .option(\"COMPRESSION\", \"GZIP\")\\\n",
    "                         .option(\"NULL_IF\", \"\\\\\\\\N\")\\\n",
    "                         .option(\"NULL_IF\", \"NULL\")\\\n",
    "                         .schema(load_schema1)\\\n",
    "                         .csv('@'+load_stage_names['schema1'])\\\n",
    "                         .copy_into_table(load_table_name+'schema1', \n",
    "                                          files=files_to_load['schema1'],\n",
    "                                          format_type_options=csv_file_format_options)\n",
    "                              \n",
    "    if len(files_to_load['schema2']) > 0:\n",
    "        load_schema2 = schema2_definition()\n",
    "        loaddf = session.read.option(\"SKIP_HEADER\", 1)\\\n",
    "                         .option(\"FIELD_OPTIONALLY_ENCLOSED_BY\", \"\\042\")\\\n",
    "                         .option(\"COMPRESSION\", \"GZIP\")\\\n",
    "                         .option(\"NULL_IF\", \"\\\\\\\\N\")\\\n",
    "                         .option(\"NULL_IF\", \"NULL\")\\\n",
    "                         .schema(load_schema2)\\\n",
    "                         .csv('@'+load_stage_names['schema2'])\\\n",
    "                         .copy_into_table(load_table_name+'schema2', \n",
    "                                          files=files_to_load['schema2'],\n",
    "                                          format_type_options=csv_file_format_options)\n",
    "        \n",
    "    load_table_names = {'schema1' : load_table_name+str('schema1'), \n",
    "                         'schema2' : load_table_name+str('schema2')}\n",
    "                         \n",
    "    return load_table_names\n",
    "    \n",
    "def transform_trips(session, stage_table_names:dict, trips_table_name:str):\n",
    "    from snowflake.snowpark import functions as F\n",
    "        \n",
    "    #Change all dates to YYYY-MM-DD HH:MI:SS format\n",
    "    date_format_match = \"^([0-9]?[0-9])/([0-9]?[0-9])/([0-9][0-9][0-9][0-9]) ([0-9]?[0-9]):([0-9][0-9])(:[0-9][0-9])?.*$\"\n",
    "    date_format_repl = \"\\\\3-\\\\1-\\\\2 \\\\4:\\\\5\\\\6\"\n",
    "\n",
    "    trips_table_schema = conformed_schema()\n",
    "                         \n",
    "    trips_table_schema_names = [field.name for field in trips_table_schema.fields]\n",
    "                         \n",
    "    transdf1 = session.table(stage_table_names['schema1'])[trips_table_schema_names]\n",
    "    transdf2 = session.table(stage_table_names['schema2'])[trips_table_schema_names]\n",
    "                         \n",
    "    transdf = transdf1.union_by_name(transdf2)\\\n",
    "                      .with_column('STARTTIME', F.regexp_replace(F.col('STARTTIME'),\n",
    "                                                                F.lit(date_format_match), \n",
    "                                                                F.lit(date_format_repl)))\\\n",
    "                      .with_column('STARTTIME', F.to_timestamp('STARTTIME'))\\\n",
    "                      .with_column('STOPTIME', F.regexp_replace(F.col('STOPTIME'),\n",
    "                                                               F.lit(date_format_match), \n",
    "                                                               F.lit(date_format_repl)))\\\n",
    "                      .with_column('STOPTIME', F.to_timestamp('STOPTIME'))\\\n",
    "                      .write.mode('overwrite').save_as_table(trips_table_name)\n",
    "\n",
    "    return trips_table_name\n",
    "\n",
    "def reset_database(session, state_dict:dict, prestaged=False):\n",
    "    _ = session.sql('CREATE OR REPLACE DATABASE '+state_dict['connection_parameters']['database']).collect()\n",
    "    _ = session.sql('CREATE SCHEMA '+state_dict['connection_parameters']['schema']).collect() \n",
    "\n",
    "    if prestaged:\n",
    "        sql_cmd = 'CREATE OR REPLACE STAGE '+state_dict['load_stage_name']+\\\n",
    "                  ' url='+state_dict['connection_parameters']['download_base_url']\n",
    "        _ = session.sql(sql_cmd).collect()\n",
    "    else: \n",
    "        _ = session.sql('CREATE STAGE IF NOT EXISTS '+state_dict['load_stage_name']).collect()\n",
    "\n",
    "    load_schema1=schema1_definition()\n",
    "    session.create_dataframe([[None]*len(load_schema1.names)], schema=load_schema1)\\\n",
    "           .na.drop()\\\n",
    "           .write\\\n",
    "           .save_as_table(state_dict['load_table_name']+'schema1')\n",
    "\n",
    "    load_schema2=schema2_definition()\n",
    "    session.create_dataframe([[None]*len(load_schema2.names)], schema=load_schema2)\\\n",
    "           .na.drop()\\\n",
    "           .write\\\n",
    "           .save_as_table(state_dict['load_table_name']+'schema2')\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "authors": [
   {
    "name": "cforbe"
   }
  ],
  "kernelspec": {
   "display_name": "snowpark_0110:Python",
   "language": "python",
   "name": "conda-env-snowpark_0110-py"
  },
  "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.13"
  },
  "msauthor": "trbye"
 },
 "nbformat": 4,
 "nbformat_minor": 4
}