Files
workspace/My-Data-Engineering-Portifolio/Capstone+Project/etl.ipynb
@gabriel.pereira 58602991b3 Add 'My-Data-Engineering-Portifolio/' from commit 'af9dfac5db38895dd0aacc6ea42ddede38a11ca5'
git-subtree-dir: My-Data-Engineering-Portifolio
git-subtree-mainline: d606a91bd5
git-subtree-split: af9dfac5db
2026-03-26 15:41:43 -03:00

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
}