01_Data_Engineering.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Citibike ML\n",
    "In this example we use the [Citibike dataset](https://ride.citibikenyc.com/system-data). Citibike is a bicycle sharing system in New York City. Everyday users choose from 20,000 bicycles at 1300 stations around New York City.\n",
    "\n",
    "To ensure customer satisfaction Citibike needs to predict how many bicycles will be needed at each station. Maintenance teams from Citibike will check each station and repair or replace bicycles. Additionally, the team will relocate bicycles between stations based on predicted demand. The business needs to be able to run reports of how many bicycles will be needed at a given station on a given day.\n",
    "\n"
   ]
  },
  {
   "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",
    "For this demo flow we will assume that the organization has the following **policies and processes** :   \n",
    "-**Dev Tools**: The data engineer can develop in their tool of choice (ie. VS Code, IntelliJ, Pycharm, Eclipse, etc.).  Snowpark Python makes it possible to use any environment where they have a python kernel.  For the sake of a demo we will use Jupyter.  \n",
    "-**Data Governance**: To preserve customer privacy no data can be stored locally.  The ingest system may store data temporarily but it must be assumed that, in production, the ingest system will not preserve intermediate data products between runs. Snowpark Python allows the user to push-down all operations to Snowflake and bring the code to the data.   \n",
    "-**Automation**: Although the data engineer can use any IDE or notebooks for development purposes the final product must be python code at the end of the work stream.  Well-documented, modularized code is necessary for good ML operations and to interface with the company's CI/CD and orchestration tools.  \n",
    "-**Compliance**: Any ML models must be traceable back to the original data set used for training.  The business needs to be able to easily remove specific user data from training datasets and retrain models.  \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": [
    "load_stage_name=state_dict['load_stage_name']\n",
    "session.sql('CREATE STAGE IF NOT EXISTS '+load_stage_name).collect()"
   ]
  },
  {
   "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 @\"+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.createDataFrame([[None]*len(load_schema1.names)], schema=load_schema1)\\\n",
    "       .na.drop()\\\n",
    "       .write\\\n",
    "       .saveAsTable(state_dict['load_table_name']+'schema1')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session.createDataFrame([[None]*len(load_schema2.names)], schema=load_schema2)\\\n",
    "       .na.drop()\\\n",
    "       .write\\\n",
    "       .saveAsTable(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_schema1)\\\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.withColumn('STARTTIME', F.regexp_replace(F.col('STARTTIME'),\n",
    "                                            F.lit(date_format_match), \n",
    "                                            F.lit(date_format_repl)))\\\n",
    "      .withColumn('STARTTIME', F.to_timestamp('STARTTIME'))\\\n",
    "      .withColumn('STOPTIME', F.regexp_replace(F.col('STOPTIME'),\n",
    "                                            F.lit(date_format_match), \n",
    "                                            F.lit(date_format_repl)))\\\n",
    "      .withColumn('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').saveAsTable(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",
    "                         .option(\"pattern\", \"'.*20.*[.]gz'\")\\\n",
    "                         .schema(load_schema1)\\\n",
    "                         .csv('@'+load_stage_names['schema1'])\\\n",
    "                         .copy_into_table(load_table_name+'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",
    "                         .option(\"pattern\", \"'.*20.*[.]gz'\")\\\n",
    "                         .schema(load_schema2)\\\n",
    "                         .csv('@'+load_stage_names['schema2'])\\\n",
    "                         .copy_into_table(load_table_name+'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",
    "                      .withColumn('STARTTIME', F.regexp_replace(F.col('STARTTIME'),\n",
    "                                                                F.lit(date_format_match), \n",
    "                                                                F.lit(date_format_repl)))\\\n",
    "                      .withColumn('STARTTIME', F.to_timestamp('STARTTIME'))\\\n",
    "                      .withColumn('STOPTIME', F.regexp_replace(F.col('STOPTIME'),\n",
    "                                                               F.lit(date_format_match), \n",
    "                                                               F.lit(date_format_repl)))\\\n",
    "                      .withColumn('STOPTIME', F.to_timestamp('STOPTIME'))\\\n",
    "                      .write.mode('overwrite').saveAsTable(trips_table_name)\n",
    "\n",
    "    return trips_table_name"
   ]
  },
  {
   "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": "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.12"
  },
  "msauthor": "trbye"
 },
 "nbformat": 4,
 "nbformat_minor": 4
}