573 lines
19 KiB
Plaintext
573 lines
19 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"editable": true
|
|
},
|
|
"source": [
|
|
"### ETL notebook for testing the pipeline"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 1,
|
|
"metadata": {
|
|
"editable": true
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"import configparser\n",
|
|
"import pandas as pd\n",
|
|
"from datetime import datetime\n",
|
|
"from pyspark.sql.functions import dayofweek\n",
|
|
"import os\n",
|
|
"from pyspark.sql import SparkSession\n",
|
|
"from pyspark.sql.functions import udf, col, lit\n",
|
|
"from pyspark.sql.functions import year, month, dayofmonth, hour, weekofyear, date_format, to_date, upper\n",
|
|
"import logging\n",
|
|
"from pyspark.sql.types import DateType\n",
|
|
"from pyspark.sql.functions import monotonically_increasing_id"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 2,
|
|
"metadata": {
|
|
"editable": true
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"def create_spark_session():\n",
|
|
" spark = SparkSession \\\n",
|
|
" .builder \\\n",
|
|
" .config(\"spark.jars.repositories\", \"https://repos.spark-packages.org/\")\\\n",
|
|
" .config(\"spark.jars.packages\", \"org.apache.hadoop:hadoop-aws:2.7.0,saurfang:spark-sas7bdat:2.0.0-s_2.11\")\\\n",
|
|
" .enableHiveSupport().getOrCreate()\n",
|
|
" return spark"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 3,
|
|
"metadata": {
|
|
"editable": true
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"def rename_columns(table, new_columns):\n",
|
|
" for original, new in zip(table.columns, new_columns):\n",
|
|
" table = table.withColumnRenamed(original, new)\n",
|
|
" return table"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 4,
|
|
"metadata": {
|
|
"editable": true
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"def SAS_to_date(date):\n",
|
|
" if date is not None:\n",
|
|
" return pd.to_timedelta(date, unit='D') + pd.Timestamp('1960-1-1')\n",
|
|
"\n",
|
|
"SAS_to_date_udf = udf(SAS_to_date, DateType())"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 5,
|
|
"metadata": {
|
|
"editable": true
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"def process_immigration_data(spark, output_data):\n",
|
|
" \"\"\"Process immigration data to get f_immigration, d_immi_citzen and d_immi_airline tables\n",
|
|
" Arguments:\n",
|
|
" spark {object}: SparkSession object\n",
|
|
" input_data {object}: Source S3 endpoint\n",
|
|
" output_data {object}: Target S3 endpoint\n",
|
|
" Returns:\n",
|
|
" None\n",
|
|
" \"\"\"\n",
|
|
" logging.info(\"Start processing immigration\")\n",
|
|
" \n",
|
|
" # read immigration data file\n",
|
|
" df = spark.read.format(\"com.github.saurfang.sas.spark\").load(\"../../data/18-83510-I94-Data-2016/i94_apr16_sub.sas7bdat\", forceLowercaseNames=True, inferLong=True)\n",
|
|
" \n",
|
|
" logging.info(\"Start processing f_immigration\")\n",
|
|
" \n",
|
|
" # extract columns to create fact_immigration table\n",
|
|
" f_immigration = df.select('cicid', 'i94yr', 'i94mon', 'i94port', 'i94addr', 'arrdate', 'depdate', 'i94mode', 'i94visa')\n",
|
|
" f_immigration = f_immigration.distinct()\n",
|
|
" f_immigration = f_immigration.withColumn(\"immigration_id\", monotonically_increasing_id())\n",
|
|
" \n",
|
|
" # data wrangling to match data model\n",
|
|
" new_columns = ['cic_id', 'year', 'month', 'city_code', 'state_code', 'arrive_date', 'departure_date', 'mode', 'visa']\n",
|
|
" \n",
|
|
" # renaming columns using the function rename_columns()\n",
|
|
" f_immigration = rename_columns(f_immigration, new_columns)\n",
|
|
" \n",
|
|
" # add a new column to f_immigration by assigning a literal or constant value = United States\n",
|
|
" f_immigration = f_immigration.withColumn('country', lit('United States'))\n",
|
|
" \n",
|
|
" # convert column arrive_date to date type format\n",
|
|
" f_immigration = f_immigration.withColumn('arrive_date', SAS_to_date_udf(col('arrive_date')))\n",
|
|
" \n",
|
|
" # convert column departure_date to date type format\n",
|
|
" f_immigration = f_immigration.withColumn('departure_date', SAS_to_date_udf(col('departure_date')))\n",
|
|
" \n",
|
|
" logging.info(\"Start loading f_immigration parquet files partitioned by state_code\")\n",
|
|
" \n",
|
|
" # write f_immigration table to parquet files partitioned by state_code\n",
|
|
" f_immigration.write.mode(\"overwrite\").partitionBy('state_code').parquet(path=output_data + 'f_immigration') \n",
|
|
" \n",
|
|
" \n",
|
|
" \n",
|
|
" logging.info(\"Start processing d_citizen table\")\n",
|
|
" \n",
|
|
" # extract columns from immigration data file to create d_citizen table\n",
|
|
" d_citizen = df.select('cicid', 'i94cit', 'i94res', 'biryear', 'gender', 'insnum').distinct().withColumn(\"immi_citizen_id\", monotonically_increasing_id())\n",
|
|
" \n",
|
|
" # data wrangling to match data model\n",
|
|
" new_columns = ['cic_id', 'citizen_country', 'residence_country', 'birth_year', 'gender', 'ins_num']\n",
|
|
" d_citizen = rename_columns(d_citizen, new_columns)\n",
|
|
"\n",
|
|
" # write d_citizen table to parquet files\n",
|
|
" d_citizen.write.mode(\"overwrite\").parquet(path=output_data + 'd_citizen')\n",
|
|
" \n",
|
|
" \n",
|
|
" \n",
|
|
" logging.info(\"Start processing d_airline\")\n",
|
|
" \n",
|
|
" # extract columns from immigration data file to create d_airline table\n",
|
|
" d_airline = df.select('cicid', 'airline', 'admnum', 'fltno', 'visatype').distinct().withColumn(\"immi_airline_id\", monotonically_increasing_id())\n",
|
|
" \n",
|
|
" # data wrangling to match data model\n",
|
|
" new_columns = ['cic_id', 'airline', 'admin_num', 'flight_number', 'visa_type']\n",
|
|
" d_airline = rename_columns(d_airline, new_columns)\n",
|
|
"\n",
|
|
" # write d_airline table to parquet files\n",
|
|
" d_airline.write.mode(\"overwrite\").parquet(path=output_data + 'd_airline')"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 6,
|
|
"metadata": {
|
|
"editable": true
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"def process_label_descriptions(spark, input_data, output_data):\n",
|
|
" \"\"\" Parsing label desctiption file to get codes of country, city, state\n",
|
|
" Arguments:\n",
|
|
" spark {object}: SparkSession object\n",
|
|
" input_data {object}: Source S3 endpoint\n",
|
|
" output_data {object}: Target S3 endpoint\n",
|
|
" Returns:\n",
|
|
" None\n",
|
|
" \"\"\"\n",
|
|
"\n",
|
|
" logging.info(\"Start processing label descriptions\")\n",
|
|
" label_file = os.path.join(input_data + \"I94_SAS_Labels_Descriptions.SAS\")\n",
|
|
" with open(label_file) as f:\n",
|
|
" contents = f.readlines()\n",
|
|
"\n",
|
|
" country_code = {}\n",
|
|
" for countries in contents[10:245]:\n",
|
|
" pair = countries.split('=')\n",
|
|
" code, country = pair[0].strip(), pair[1].strip().strip(\"'\")\n",
|
|
" country_code[code] = country\n",
|
|
" \n",
|
|
" spark.createDataFrame(country_code.items(), ['country_code', 'country'])\\\n",
|
|
" .write.mode(\"overwrite\")\\\n",
|
|
" .parquet(path=output_data + 'country_code')\n",
|
|
"\n",
|
|
" city_code = {}\n",
|
|
" for cities in contents[302:962]:\n",
|
|
" pair = cities.split('=')\n",
|
|
" code, city = pair[0].strip(\"\\t\").strip().strip(\"'\"),\\\n",
|
|
" pair[1].strip('\\t').strip().strip(\"''\")\n",
|
|
" city_code[code] = city\n",
|
|
" spark.createDataFrame(city_code.items(), ['city_code', 'city'])\\\n",
|
|
" .write.mode(\"overwrite\")\\\n",
|
|
" .parquet(path=output_data + 'city_code')\n",
|
|
"\n",
|
|
" state_code = {}\n",
|
|
" for states in contents[981:1036]:\n",
|
|
" pair = states.split('=')\n",
|
|
" code, state = pair[0].strip('\\t').strip(\"'\"), pair[1].strip().strip(\"'\")\n",
|
|
" state_code[code] = state\n",
|
|
" spark.createDataFrame(state_code.items(), ['state_code', 'state'])\\\n",
|
|
" .write.mode(\"overwrite\")\\\n",
|
|
" .parquet(path=output_data + 'state_code')"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 7,
|
|
"metadata": {
|
|
"editable": true
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"def process_temperature_data(spark, output_data):\n",
|
|
" \"\"\" Process temperature data to get dim_temperature table\n",
|
|
" Arguments:\n",
|
|
" spark {object}: SparkSession object\n",
|
|
" input_data {object}: Source S3 endpoint\n",
|
|
" output_data {object}: Target S3 endpoint\n",
|
|
" Returns:\n",
|
|
" None\n",
|
|
" \"\"\"\n",
|
|
"\n",
|
|
" logging.info(\"Start processing d_temperature\")\n",
|
|
" # read temperature data file\n",
|
|
" tempe_data = os.path.join('../../data2/GlobalLandTemperaturesByCity.csv')\n",
|
|
" df = spark.read.csv(tempe_data, header=True)\n",
|
|
"\n",
|
|
" df = df.where(df['Country'] == 'United States')\n",
|
|
" d_temperature = df.select(['dt', 'AverageTemperature', 'AverageTemperatureUncertainty',\\\n",
|
|
" 'City', 'Country']).distinct()\n",
|
|
"\n",
|
|
" new_columns = ['dt', 'avg_temp', 'avg_temp_uncertnty', 'city', 'country']\n",
|
|
" d_temperature = rename_columns(d_temperature, new_columns)\n",
|
|
"\n",
|
|
" d_temperature = d_temperature.withColumn('dt', to_date(col('dt')))\n",
|
|
" d_temperature = d_temperature.withColumn('year', year(d_temperature['dt']))\n",
|
|
" d_temperature = d_temperature.withColumn('month', month(d_temperature['dt']))\n",
|
|
" \n",
|
|
" # write dim_temperature table to parquet files\n",
|
|
" d_temperature.write.mode(\"overwrite\")\\\n",
|
|
" .parquet(path=output_data + 'd_temperature')"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 8,
|
|
"metadata": {
|
|
"editable": true
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"def process_demography_data(spark, input_data, output_data):\n",
|
|
" \"\"\" Process demograpy data to get dim_demog_population \n",
|
|
" and d_demog_statistics table\n",
|
|
" Arguments:\n",
|
|
" spark {object}: SparkSession object\n",
|
|
" input_data {object}: Source S3 endpoint\n",
|
|
" output_data {object}: Target S3 endpoint\n",
|
|
" Returns:\n",
|
|
" None\n",
|
|
" \"\"\"\n",
|
|
"\n",
|
|
" logging.info(\"Start processing d_demog_statistics\")\n",
|
|
" # read demography data file\n",
|
|
" demog_data = os.path.join(input_data + 'us-cities-demographics.csv')\n",
|
|
" df = spark.read.format('csv').options(header=True, delimiter=';').load(demog_data)\n",
|
|
"\n",
|
|
"\n",
|
|
" d_demog_statistics = df.select(['City', 'State', 'Male Population', 'Female Population', \\\n",
|
|
" 'Number of Veterans', 'Foreign-born', 'Race']).distinct() \\\n",
|
|
" .withColumn(\"demog_pop_id\", monotonically_increasing_id())\n",
|
|
"\n",
|
|
"\n",
|
|
" new_columns = ['city', 'state', 'male_population', 'female_population', \\\n",
|
|
" 'num_vetarans', 'foreign_born', 'race']\n",
|
|
" d_demog_statistics = rename_columns(d_demog_statistics, new_columns)\n",
|
|
"\n",
|
|
" # write dim_demog_population table to parquet files\n",
|
|
" d_demog_statistics.write.mode(\"overwrite\")\\\n",
|
|
" .parquet(path=output_data + 'd_demog_statistics')\n",
|
|
"\n",
|
|
" \n",
|
|
" logging.info(\"Start processing d_demog_statistics\")\n",
|
|
" d_demog_statistics = df.select(['City', 'State', 'Median Age', 'Average Household Size'])\\\n",
|
|
" .distinct()\\\n",
|
|
" .withColumn(\"d_demog_statistics\", monotonically_increasing_id())\n",
|
|
"\n",
|
|
" new_columns = ['city', 'state', 'median_age', 'avg_household_size']\n",
|
|
" d_demog_statistics = rename_columns(d_demog_statistics, new_columns)\n",
|
|
" d_demog_statistics = d_demog_statistics.withColumn('city', upper(col('city')))\n",
|
|
" d_demog_statistics = d_demog_statistics.withColumn('state', upper(col('state')))\n",
|
|
"\n",
|
|
" # write dim_demog_statistics table to parquet files\n",
|
|
" d_demog_statistics.write.mode(\"overwrite\")\\\n",
|
|
" .parquet(path=output_data + 'd_demog_statistics')"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 9,
|
|
"metadata": {
|
|
"editable": true
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"'''Paths for local testing'''\n",
|
|
"input_data = \"./\"\n",
|
|
"output_data = \"./data/outputs/\"\n",
|
|
"\n",
|
|
"spark = create_spark_session()\n",
|
|
"\n",
|
|
"process_immigration_data(spark, output_data)\n",
|
|
"process_label_descriptions(spark, input_data, output_data)\n",
|
|
"process_temperature_data(spark, output_data)\n",
|
|
"process_demography_data(spark, input_data, output_data)\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {
|
|
"editable": true
|
|
},
|
|
"source": [
|
|
"### Data Model Schema"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 11,
|
|
"metadata": {
|
|
"editable": true
|
|
},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"root\n",
|
|
" |-- cic_id: double (nullable = true)\n",
|
|
" |-- year: double (nullable = true)\n",
|
|
" |-- month: double (nullable = true)\n",
|
|
" |-- city_code: string (nullable = true)\n",
|
|
" |-- arrive_date: date (nullable = true)\n",
|
|
" |-- departure_date: date (nullable = true)\n",
|
|
" |-- mode: double (nullable = true)\n",
|
|
" |-- visa: double (nullable = true)\n",
|
|
" |-- immigration_id: long (nullable = true)\n",
|
|
" |-- country: string (nullable = true)\n",
|
|
"\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"# f_immigration table sample\n",
|
|
"spark.read.parquet('data/outputs/f_immigration/state_code=AE/part-00000-79a43b03-8ba1-418b-bd9e-b4b3de3d441c.c000.snappy.parquet').printSchema()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 12,
|
|
"metadata": {
|
|
"editable": true
|
|
},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"root\n",
|
|
" |-- cic_id: double (nullable = true)\n",
|
|
" |-- citizen_country: double (nullable = true)\n",
|
|
" |-- residence_country: double (nullable = true)\n",
|
|
" |-- birth_year: double (nullable = true)\n",
|
|
" |-- gender: string (nullable = true)\n",
|
|
" |-- ins_num: string (nullable = true)\n",
|
|
" |-- immi_citizen_id: long (nullable = true)\n",
|
|
"\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"# d_citizen table sample\n",
|
|
"spark.read.parquet('data/outputs/d_citizen/part-00000-ddacc44c-bdd0-4694-a9fa-4a367a610df6-c000.snappy.parquet').printSchema()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 13,
|
|
"metadata": {
|
|
"editable": true
|
|
},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"root\n",
|
|
" |-- cic_id: double (nullable = true)\n",
|
|
" |-- airline: string (nullable = true)\n",
|
|
" |-- admin_num: double (nullable = true)\n",
|
|
" |-- flight_number: string (nullable = true)\n",
|
|
" |-- visa_type: string (nullable = true)\n",
|
|
" |-- immi_airline_id: long (nullable = true)\n",
|
|
"\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"# d_airline table sample\n",
|
|
"spark.read.parquet('data/outputs/d_airline/part-00000-5464a16a-e021-494e-9e76-d6df7e80c20f-c000.snappy.parquet').printSchema()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 14,
|
|
"metadata": {
|
|
"editable": true
|
|
},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"root\n",
|
|
" |-- country_code: string (nullable = true)\n",
|
|
" |-- country: string (nullable = true)\n",
|
|
"\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"# country_code table sample\n",
|
|
"spark.read.parquet('data/outputs/country_code/*.parquet').printSchema()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 15,
|
|
"metadata": {
|
|
"editable": true
|
|
},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"root\n",
|
|
" |-- city_code: string (nullable = true)\n",
|
|
" |-- city: string (nullable = true)\n",
|
|
"\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"# city_code table sample\n",
|
|
"spark.read.parquet('data/outputs/city_code/*.parquet').printSchema()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 16,
|
|
"metadata": {
|
|
"editable": true
|
|
},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"root\n",
|
|
" |-- state_code: string (nullable = true)\n",
|
|
" |-- state: string (nullable = true)\n",
|
|
"\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"# state_code table sample\n",
|
|
"spark.read.parquet('data/outputs/state_code/*.parquet').printSchema()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 17,
|
|
"metadata": {
|
|
"editable": true
|
|
},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"root\n",
|
|
" |-- dt: date (nullable = true)\n",
|
|
" |-- avg_temp: string (nullable = true)\n",
|
|
" |-- avg_temp_uncertnty: string (nullable = true)\n",
|
|
" |-- city: string (nullable = true)\n",
|
|
" |-- country: string (nullable = true)\n",
|
|
" |-- year: integer (nullable = true)\n",
|
|
" |-- month: integer (nullable = true)\n",
|
|
"\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"# d_temperature table sample\n",
|
|
"spark.read.parquet('data/outputs/d_temperature/*.parquet').printSchema()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 18,
|
|
"metadata": {
|
|
"editable": true
|
|
},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"root\n",
|
|
" |-- city: string (nullable = true)\n",
|
|
" |-- state: string (nullable = true)\n",
|
|
" |-- median_age: string (nullable = true)\n",
|
|
" |-- avg_household_size: string (nullable = true)\n",
|
|
" |-- d_demog_statistics: long (nullable = true)\n",
|
|
"\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"# d_demog_statistics table sample\n",
|
|
"spark.read.parquet('data/outputs/d_demog_statistics/*.parquet').printSchema()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"editable": true
|
|
},
|
|
"outputs": [],
|
|
"source": []
|
|
}
|
|
],
|
|
"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.6.3"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 4
|
|
}
|