capstone-dataeng/main.py at main · ccaloian/capstone-dataeng · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
from pyspark.sql import SparkSession

import os

import dataio
import etl


DATA_DIR = "./data/processed"

def create_session():
    """Create SparkSession.

    Returns:
        SparkSession
    """
    spark = (SparkSession.builder
                .master("local[*]")
                .appName("ETL")
                .config("spark.jars.packages", "org.apache.hadoop:hadoop-aws:2.7.0")
                .config("spark.driver.memory", "15g")
                .config("spark.sql.legacy.parquet.datetimeRebaseModeInWrite", "CORRECTED")
                .getOrCreate()
    )
    return spark

def main():
    """Run ETL pipeline and save files.

    Output directory is './data/processed/'.
    """

    if not os.path.exists(DATA_DIR):
        os.makedirs(DATA_DIR)

    # dem_etl_pipeline = etl.DemographicsPipeline()
    # dem_df = dem_etl_pipeline.run()
    # dem_df.write\
    #     .mode("overwrite")\
    #     .parquet(f"{DATA_DIR}/demographics/demographics.parquet")

    imm_etl_pipeline = etl.ImmigrationPipeline()
    imm_df = imm_etl_pipeline.run()
    imm_df.write\
        .mode("overwrite")\
        .partitionBy("i94mon")\
        .parquet(f"{DATA_DIR}/immigration/immigration.parquet")

if __name__ == "__main__":
    main()