使用 Spark 结构化流在 postgresql 中插入数据

Question

I am trying to run a structured streaming application using (py)spark. My data is read from a Kafka topic and then I am running windowed aggregation on event time.

# I have been able to create data frame pn_data_df after reading data from Kafka

Schema of pn_data_df
  |
   - id StringType
   - source StringType
   - source_id StringType
   - delivered_time TimeStamp

windowed_report_df = pn_data_df.filter(pn_data_df.source == 'campaign') \
    .withWatermark("delivered_time", "24 hours") \
    .groupBy('source_id', window('delivered_time', '15 minute')) \
    .count()
windowed_report_df = windowed_report_df \
    .withColumn('start_ts', unix_timestamp(windowed_report_df.window.start)) \
    .withColumn('end_ts', unix_timestamp(windowed_report_df.window.end)) \
    .selectExpr('CAST(source_id as LONG)', 'start_ts', 'end_ts', 'count')

I am writing this windowed aggregation to my postgresql database which I have already created.

CREATE TABLE pn_delivery_report(
   source_id bigint not null,
   start_ts bigint not null,
   end_ts bigint not null,
   count integer not null,
   unique(source_id, start_ts)
);

Writing to postgresql using spark jdbc allows me to either Append or Overwrite . Append mode fails if there is an existing composite key existing in the database, and Overwrite just overwrites entire table with current batch output.

def write_pn_report_to_postgres(df, epoch_id):
    df.write \
    .mode('append') \
    .format('jdbc') \
    .option("url", "jdbc:postgresql://db_endpoint/db") \
    .option("driver", "org.postgresql.Driver") \
    .option("dbtable", "pn_delivery_report") \
    .option("user", "postgres") \
    .option("password", "PASSWORD") \
    .save()

windowed_report_df.writeStream \
   .foreachBatch(write_pn_report_to_postgres) \
   .option("checkpointLocation", '/home/hadoop/campaign_report_df_windowed_checkpoint') \
   .outputMode('update') \
   .start()

How can I execute a query like

INSERT INTO pn_delivery_report (source_id, start_ts, end_ts, COUNT)
VALUES (1001, 125000000001, 125000050000, 128),
       (1002, 125000000001, 125000050000, 127) ON conflict (source_id, start_ts) DO
UPDATE
SET COUNT = excluded.count;

in foreachBatch .

Spark has a jira feature ticket open for it, but it seems that it has not been prioritised till now.

https://issues.apache.org/jira/browse/SPARK-19335

Answer 1

that's worked for me:

def _write_streaming(self,
    df,
    epoch_id
) -> None:         

    df.write \
        .mode('append') \
        .format("jdbc") \
        .option("url", f"jdbc:postgresql://localhost:5432/postgres") \
        .option("driver", "org.postgresql.Driver") \
        .option("dbtable", 'table_test') \
        .option("user", 'user') \
        .option("password", 'password') \
        .save()

df_stream.writeStream \
    .foreachBatch(_write_streaming) \
    .start() \
    .awaitTermination()

You need to add ".awaitTermination()" at the end.