etl.py

import configparser
from datetime import datetime
import os
from pyspark.sql import SparkSession
from pyspark.sql.functions import udf, col
from pyspark.sql.functions import year, month, dayofmonth, hour, weekofyear, dayofweek, monotonically_increasing_id
from pyspark.sql.types import StructType as R, StructField as Fld, DoubleType as Dbl, StringType as Str, IntegerType as Int, DateType as Date

config = configparser.ConfigParser()
config.read('dl.cfg')

os.environ['AWS_ACCESS_KEY_ID']=config['AWS']['AWS_ACCESS_KEY_ID']
os.environ['AWS_SECRET_ACCESS_KEY']=config['AWS']['AWS_SECRET_ACCESS_KEY']


def create_spark_session():
    '''Creates Spark session.
    
    Returns currently running Spark Session, or creates a new session and returns it.
    
    Parameters:
        None
    
    Returns:
        SparkSession: spark
    
    '''
    spark = SparkSession \
        .builder \
        .config("spark.jars.packages", "org.apache.hadoop:hadoop-aws:2.7.0") \
        .getOrCreate()
    return spark


def process_song_data(spark, input_data, output_data):
    '''Processes song data.
    
    Creates songs_table and artists_table using song data.
    
    Parameters:
        spark (SparkSession): Current Spark session
        input_data (str): String of directory of song_data and log_data
        output_data (str): String of directory where output data is written

    Returns:
        None
        
    '''
    
    print('process_song_data starting     (0/2)')
    
    # get filepath to song data file
    song_data = os.path.join(input_data, 'song_data/*/*/*/*.json')

    # read song data file
    df = spark.read.json(song_data) 

    print('Creating songs_table           (1/2)')
    # extract columns to create songs table
    songs_table = df.select('song_id', 'title', 'artist_id', 'year', 'duration').drop_duplicates()
    
    # write songs table to parquet files partitioned by year and artist
    songs_table.write.partitionBy("year", "artist_id").mode('overwrite').parquet(os.path.join(output_data, 'songs.parquet'))
    
    print('Creating artists_table         (2/2)')
    
    # extract columns to create artists table
    artists_table = df.select('artist_id', 'artist_name', 'artist_location', 'artist_latitude', 'artist_longitude').drop_duplicates()
    
    # write artists table to parquet files
    artists_table.write.mode('overwrite').parquet(os.path.join(output_data, 'artists.parquet'))
    
    print('process_song_data completed')

    
def process_log_data(spark, input_data, output_data):
    '''Processes log data.
    
    Creates users_table and time_table using log data. The song and log datasets are joined to create the songplays_table.
    
    Parameters:
        spark (SparkSession): Current Spark session
        input_data (str): String of directory of song_data and log_data
        output_data (str): String of directory where output data is written 
    
    Returns:
        None
    
    '''
    
    print('process_log_data starting      (0/3)')
    # get filepath to log data file
    log_data = os.path.join(input_data, 'log_data/*/*/*.json')

    # read log data file
    df = spark.read.json(log_data) 
    
    # filter by actions for song plays
    df = df.filter(df.page == 'NextSong')

    print('Creating users_table           (1/3)')    
    
    # extract columns for users table    
    users_table = df.select('userId', 'firstName', 'lastName', 'gender', 'level').drop_duplicates()
    
    # write users table to parquet files
    users_table.write.mode('overwrite').parquet(os.path.join(output_data, 'users.parquet'))
    
    # create timestamp column from original timestamp column 
    get_timestamp = udf(lambda x: str(int(x)/1000))
    df = df.withColumn("timestamp", get_timestamp(col('ts')))
    
    # create datetime column from original timestamp column
    get_datetime = udf(lambda x: datetime.fromtimestamp(int(x)/1000)) 
    df = df.withColumn("datetime", get_datetime(col('ts')))

    print('Creating time_table            (2/3)')
    
    # extract columns to create time table
    time_table = df.select(col('datetime').alias('start_time'), 
                           hour(col('datetime')).alias('hour'),
                           dayofmonth(col('datetime')).alias('day'),
                           weekofyear(col('datetime')).alias('week'),
                           month(col('datetime')).alias('month'),
                           year(col('datetime')).alias('year'),
                           dayofweek(col('datetime')).alias('weekday'))
    
    # write time table to parquet files partitioned by year and month
    time_table.write.partitionBy("year", "month").mode('overwrite').parquet(os.path.join(output_data, 'time.parquet'))
    
    # read in song data to use for songplays table
    song_df = spark.read.json(os.path.join(input_data, 'song_data/*/*/*/*.json'))

    print('Creating songplays_table       (3/3)')

    # extract columns from joined song and log datasets to create songplays table 
    songplays_table = df.join(song_df, (df.song == song_df.title) & (df.artist == song_df.artist_name) & (df.length == song_df.duration), 'left_outer') \
        .select(
            df.timestamp,
            col("userId").alias('user_id'),
            df.level,
            song_df.song_id,
            song_df.artist_id,
            col("sessionId").alias("session_id"),
            df.location,
            col("useragent").alias("user_agent"),
            year('datetime').alias('year'),
            month('datetime').alias('month')) \
            .withColumn('songplay_id', monotonically_increasing_id())
    
    # write songplays table to parquet files partitioned by year and month
    songplays_table.write.partitionBy("year", "month").mode('overwrite').parquet(os.path.join(output_data, 'songplays.parquet'))
    
    print('process_log_data completed')
    
    
def main():
    spark = create_spark_session()
    input_data = "s3a://udacity-dend/"
    output_data = "output"
    
    process_song_data(spark, input_data, output_data)    
    process_log_data(spark, input_data, output_data)


if __name__ == "__main__":
    main()