Pyspark streaming append pivoted batch to sql server

user_15463697 0
Viewed 268

I'm using pyspark streaming to stream data from a kafka server, manipulate it batch by batch (using foreachBatch), and append each batch to a Microsoft SQL server using jdbc.

here are the main relevent parts of my code:

Defining the stream

string_value_df = spark.readStream \
    .format("kafka") \
    .option("kafka.bootstrap.servers", constants.kafka_server) \
    .option("subscribe", "topics") \
    .option("maxOffsetsPerTrigger", 1000) \
    .option("startingOffsets", "earliest") \
    .load() \
    .selectExpr("CAST(value AS STRING)")

Defining schema and manipulating the df

schema = T.StructType([T.StructField('StationId', T.StringType(), False),
                            T.StructField('Date', T.StringType(), False),
                            T.StructField('Variable', T.StringType(), False),
                            T.StructField('Value', T.IntegerType(), False)])

json_df = string_value_df.select(F.from_json(F.col("value"),schema=schema).alias('json'))
streaming_df = json_df.select("json.*")

Start Streaming

query = df \
    .writeStream \
    .foreachBatch(write_to_sql) \
    .outputMode("Update") \
    .start() \
    .awaitTermination()

During each batch manipulation, I use pivot to transform multiple variable-value records to variables in columns forms, while grouping on "StationID" and "Date"

The original data form (how each batch arrives) is:

StationId Date Variable Value
A yyyyMMdd PRCP x1
A yyyyMMdd SNOW x2
A yyyyMMdd SNWD x3
A yyyyMMdd TMAX x4
A yyyyMMdd TMIN x5

And after my transformation which include pivot:

StationId Date PRCP SNOW SNWD TMAX TMIN
A yyyyMMdd x1 x2 x3 x4 x5

Here is the function applied on each batch:

Hundle a batch

def write_to_sql(df, df_id):
    df = df.groupBy("StationId", "Date").pivot("Variable").sum("Value")

    try:
        df.write \
            .format("jdbc") \
            .mode("append") \
            .option("url", constants.url) \
            .option("dbtable", constants.table_name) \
            .option("user", constants.username) \
            .option("password", constants.password) \
            .save()
    except ValueError as error:
        print("Connector write failed", error)

My problem is with appending the new batches to the SQL table on the server.
When having two rows with the same <StationId, Date> in the same batch, the pivot works fine and appeared correctly in the server.
Within, if multiple (different variables) records of specific <StationId, Date> pair, are divided between multiple batches, then when appending it to the server, it appeared not fully grouped.

StationId Date PRCP SNOW SNWD TMAX TMIN
A yyyyMMdd x1 null x3 null null
A yyyyMMdd null x2 null null null

[x1 and x3 appeared in the same batch]
Is there any efficient way to append the batches to the server, while maintaining the grouping on <StationId, Date> across different variables?
Thanks a lot

pyspark apache-spark pivot sql
1 Answers

You can do the following to match the schema ahead of time and then you won't run into trouble: This is my solution to the exercise when trying to filter out the 4 variables PRCP, TMAX, TMIN, SNWD.

batch_df = batch_df.filter(batch_df.Q_Flag.isNull()).drop("M_Flag", "Q_Flag", "S_Flag", "ObsTime")
batch_df = batch_df.withColumn("Year_Rec", F.expr("substring(Date,0,4)").cast(IntegerType())).withColumn(
        "Month_Rec", F.expr("substring(Date,5,2)").cast(IntegerType())).drop("Date")
batch_df = batch_df.filter(
        ((batch_df.Variable == "PRCP") & (batch_df.Value >= 0))
        | (batch_df.Variable == "TMAX")
        | (batch_df.Variable == "TMIN")
        | ((batch_df.Variable == "SNWD") & (batch_df.Value >= 0)))
    ger_swe_df = batch_df.groupBy("StationId", "Year_Rec", "Month_Rec", "Variable") \
        .agg(F.sum(F.col("Value")).alias("batch_sum"), F.count("*").alias("batch_count"))
    ger_swe_df = ger_swe_df.groupby("StationId", "Year_Rec", "Month_Rec").pivot("Variable") \
        .agg(F.first("batch_count").cast(LongType()).alias('batch_count'),
             F.first("batch_sum").cast(LongType()).alias('batch_sum'))
cols = {str(col)[:4] for col in ger_swe_df.columns if str(col)[:4] in vars}
missing = vars - cols
for var in missing:
    ger_swe_df = ger_swe_df.withColumn(var + '_batch_count', F.lit(None).cast(LongType())).withColumn(
            var + '_batch_sum', F.lit(None).cast(LongType()))
ger_swe_df = ger_swe_df.select("StationId", "Year_Rec", "Month_Rec", "PRCP_batch_count", 'PRCP_batch_sum',
                                   'TMAX_batch_count', 'TMAX_batch_sum', 'TMIN_batch_count', 'TMIN_batch_sum',
                                   'SNWD_batch_count', 'SNWD_batch_sum') \
        .withColumn("Season", ((F.col("Month_Rec") % 12) / 3 + 1).cast(IntegerType()))
ger_swe_df.withColumn("batchId", F.lit(batch_id)) \
        .write \
        .format("com.microsoft.sqlserver.jdbc.spark") \
        .mode("append") \
        .option("url", url) \
        .option("dbtable", table_name) \
        .option("user", username) \
        .option("password", password) \
        .save()
user_108470960
Related