55

from pyspark.sql import SparkSession
from pyspark.sql import Column
from pyspark.sql.functions import upper, col, concat, lit, current_timestamp, current_date, count
from datetime import datetime
import sys

source = sys.argv[1]
accountId = sys.argv[2]

def create_redshift_context(database):
    url_redshift = f"jdbc:redshift:iam://redshift-spark-wg.{accountId}.us-east-1.redshift-serverless.amazonaws.com:5439/{database}"
    aws_role = f"arn:aws:iam::{accountId}:role/AmazonRedshiftCommandsAccessRole"
    temp_dir= f"s3://emrserverless-workshop-{accountId}/redshift-spark-demo-temp-dir/"
    
    redshiftOptions = {
        "url": url_redshift,
        "tempdir": temp_dir,
        "aws_iam_role": aws_role
    }
    
    return redshiftOptions

def read_from_s3(spark):

    startTime = datetime.now()
    print("Process started at: ", startTime)

    as_of_date = startTime.strftime("%Y-%m-%d")
    
    rawDF = spark.read.parquet(f"s3://emrserverless-workshop-{accountId}/input/")
    rawDF.printSchema()

    print("Total number of records: " + str(rawDF.count()))

    return rawDF

def read_from_redshift(spark, redshiftOptions):
    customerDF = (
        spark.read
            .format("io.github.spark_redshift_community.spark.redshift")
            .options(**redshiftOptions)
            .option("dbtable","tpcds.customer")
            .load()
    )
    
    storeSalesDF = (
        spark.read
            .format("io.github.spark_redshift_community.spark.redshift")
            .options(**redshiftOptions)
            .option("dbtable","tpcds.store_sales")
            .load()
    )
    
    webSalesDF = (
        spark.read
            .format("io.github.spark_redshift_community.spark.redshift")
            .options(**redshiftOptions)
            .option("dbtable","tpcds.web_Sales")
            .load()
    )
    
    catalogSalesDF = (
        spark.read
            .format("io.github.spark_redshift_community.spark.redshift")
            .options(**redshiftOptions)
            .option("dbtable","tpcds.catalog_sales")
            .load()
            )
            
    aggregateDF = customerDF \
            .join(storeSalesDF, customerDF.c_customer_sk == storeSalesDF.ss_customer_sk, "left") \
            .join(webSalesDF, customerDF.c_customer_sk == webSalesDF.ws_bill_customer_sk,"left") \
            .join(catalogSalesDF, customerDF.c_customer_sk == catalogSalesDF.cs_bill_customer_sk, "left") \
            .groupBy("c_customer_id") \
            .agg(count("ss_customer_sk"),count("cs_bill_customer_sk"),count("ws_bill_customer_sk")) \
            .select(col("c_customer_id"),\
            col("count(ss_customer_sk)"), \
            col("count(cs_bill_customer_sk)"), \
            col("count(ws_bill_customer_sk)")) \
            .withColumnRenamed("count(ss_customer_sk)", "total_store_sales") \
            .withColumnRenamed("count(cs_bill_customer_sk)", "total_catalog_sales") \
            .withColumnRenamed("count(ws_bill_customer_sk)", "total_web_sales")
    return aggregateDF
       
def write_to_redshift (spark, dataframe, redshiftOptions, tableName, schemaName):
    
    dataframe.write.format("io.github.spark_redshift_community.spark.redshift")\
    .options(**redshiftOptions)\
    .option("dbtable", f"{schemaName}.{tableName}")\
    .option("tempformat", "CSV")\
    .mode("append")\
    .save()
    
if __name__ == "__main__":
	spark = SparkSession.builder.appName("Redshift Spark Connector").config('spark.sql.debug.maxToStringFields',2000).getOrCreate()
	redshiftOptionsWrite = create_redshift_context("test")
	redshiftOptionsRead = create_redshift_context("sample_data_dev")
	sourceDF = read_from_s3(spark) if source.upper() == "S3" else read_from_redshift(spark, redshiftOptionsRead)
	write_to_redshift(spark, sourceDF, redshiftOptionsWrite, "nyc_taxi_trips","public") if source.upper() == "S3" else write_to_redshift(spark, sourceDF, redshiftOptionsWrite, "total_by_customer","public")
	spark.stop()

1
2
3
4
export JOB_ROLE_ARN=<<EmrServerlessS3RuntimeRoleARN>>
export S3_BUCKET=s3://<<S3Bucket>>
export APPLICATION_ID=<<application_id>>
export ACCOUNT_ID=<<accountId>> 

aws emr-serverless start-job-run --application-id ${APPLICATION_ID} --region us-east-1 --execution-role-arn ${JOB_ROLE_ARN} --name "Redshift-Redshift-Job" --job-driver '{
        "sparkSubmit": {
            "entryPoint": "'"$S3_BUCKET"'/script/load_to_redshift.py",
            "entryPointArguments": ["redshift","'"$ACCOUNT_ID"'"]
        }
    }' --configuration-overrides '{
    "monitoringConfiguration": {
        "s3MonitoringConfiguration": {
            "logUri": "'"${S3_BUCKET}"'/logs/"
        }
    }
}'

1
aws emr-serverless get-job-run --application-id ${APPLICATION_ID} --job-run-id jobRunId --region us-east-1