<jupyter_start><jupyter_text>Spark ETL Demo PythonThis demo written in Python for Watson Data Studio illustrates the use of a Spark cluster to perform ETL. It imports data in flat files into Spark DataFrames, manipulates the data, aggregates it  and then writes the result out to a relational  database. The advantage of using Spark for this is scalability  (by using a larger cluster one can achieve close to linear scalability) and simplified error recovery (a failed attempt at running this ETL job can be repeated at any stage and the final result will be the same). Step 1 Read in the source dataWe read two CSV files. One has statistics about Social Security payments for the state of Texas by zipcode and the other maps US zipcodes to US counties so we can aggregate the Social Security data by county rather than zipcode. Grab the input data files from Github and stick them in in gpfs using wget<jupyter_code># Install wget if you don't already have it.
!pip install wget
import wget

link_to_ssdata = 'https://raw.githubusercontent.com/djccarew/sparketldemo/master/data/oasdi-tx-clean.csv'
link_to_zipdata = 'https://raw.githubusercontent.com/djccarew/sparketldemo/master/data/zip_codes_states.csv'
social_security_data_file = wget.download(link_to_ssdata)

print(social_security_data_file)

zipcode_data_file = wget.download(link_to_zipdata)

print(zipcode_data_file)<jupyter_output><empty_output><jupyter_text>Read in the Social Security data file into a DataFrame using a schema. Note the schema can be inferred but the inferred schema typically converts various numeric types to string so it's better to specify the schema so you know what you end up with<jupyter_code>from pyspark.sql.types import StructType, StructField, IntegerType, StringType, DoubleType
from pyspark.sql import SparkSession
spark = SparkSession.builder.getOrCreate()
ssdata_schema = StructType([
    StructField("Zip", StringType(), False),
    StructField("NumTotal", IntegerType(), False),
    StructField("NumRetired", IntegerType(), False),
    StructField("NumDisabled", IntegerType(), False),
    StructField("NumWidowerOrParent", IntegerType(), False),
    StructField("NumSpouses", IntegerType(), False),
    StructField("NumChildren", IntegerType(), False),
    StructField("BenTotal", IntegerType(), False),
    StructField("BenRetired", IntegerType(), False),
    StructField("BenWidowerOrParent", IntegerType(), False),
    StructField("NumSeniors", IntegerType(), False)])

df_ssdata_raw = spark.read\
  .format('org.apache.spark.sql.execution.datasources.csv.CSVFileFormat')\
  .option('header', 'true')\
  .load(social_security_data_file, schema=ssdata_schema)
    
df_ssdata_raw.printSchema()<jupyter_output><empty_output><jupyter_text>Repeat for zipcode data file<jupyter_code>zipdata_schema = StructType([
    StructField("Zip", StringType(), False),
    StructField("Latitude", DoubleType(), False),
    StructField("Longitude", DoubleType(), False),
    StructField("City", StringType(), False),
    StructField("State", StringType(), False),
    StructField("County", StringType(), False)])
    

df_zipdata_raw = spark.read\
  .format('org.apache.spark.sql.execution.datasources.csv.CSVFileFormat')\
  .option('header', 'true')\
  .load(zipcode_data_file, schema=zipdata_schema)
df_zipdata_raw.printSchema()<jupyter_output><empty_output><jupyter_text>Step 2 Transform raw source data Only need County name and zip code columns for this demo so we don't use the other columns in the zipcode data<jupyter_code>df_counties = df_zipdata_raw.select('Zip','County')
df_counties.printSchema()<jupyter_output><empty_output><jupyter_text>Join Social Security data with zipcode data to add a County column to Social Security data<jupyter_code>df_ssdata_counties = df_ssdata_raw.join(df_counties, "Zip")
df_ssdata_counties.printSchema()<jupyter_output><empty_output><jupyter_text>Don't need the zipcode column anymore since we'll be aggregating by County instead<jupyter_code>df_ssdata_counties = df_ssdata_counties.drop("Zip")
df_ssdata_counties.printSchema()<jupyter_output><empty_output><jupyter_text>Create a temp view so we can do the "by county" aggregation via SQL rather than using the Spark SQL DataFrame API. (Doing it via SQL is usually easier)<jupyter_code>df_ssdata_counties.createOrReplaceTempView("aggregated_by_county")<jupyter_output><empty_output><jupyter_text>Spark SQL query to aggregate Social Security data by county and sort by county name<jupyter_code>df_ssdata_data_by_county = spark.sql("select County, sum(NumTotal) as NumTotal, sum(NumRetired) as NumRetired, sum(NumDisabled) as NumDisabled, sum(NumWidowerOrParent) as NumWidowerOrParent, sum(NumSpouses) as NumSpouses, sum(NumChildren) as NumChildren, sum(BenTotal) as BenTotal, sum(BenRetired) as BenRetired, sum(BenWidowerOrParent) as BenWidowerOrParent, sum(NumSeniors) as NumSeniors from aggregated_by_county group by County order by County")
df_ssdata_data_by_county.take(5)<jupyter_output><empty_output><jupyter_text>Step 3 Write modified data to target database We use the jdbc method of the DataFrameWriter to write the modified data to the target db. Appropriate credentials for the target db need to be set up first. Modify the code below with the approparaite values for your database<jupyter_code>jdbc_url = 'your-jdbc-url'
dest_table = 'your-table-name'
jdbc_properties = {
    'driver': 'com.ibm.db2.jcc.DB2Driver',
    'user': 'your-db-user',
    'password': '<PASSWORD>'
}

df_ssdata_data_by_county.write.jdbc(jdbc_url, table=dest_table, mode='overwrite', properties=jdbc_properties)<jupyter_output><empty_output>
StarcoderdataJupyterStructuredCleanDedup