Is there a way to pass an additional/extra parameter in `foreachBatch` function in Kafka Streaming?

user_7119501 0
Viewed 399

I am trying to join a streaming dataframe with a hive table and insert the resultant dataframe into another Kafka topic. Below is the code I implemented and it is working as per the requirement.

def write_stream_batches(kafka_df: DataFrame,table_config):
    table_config = state_config
    kafka_df.writeStream \
    .format('kafka') \
    .foreachBatch(join_kafka_streams_denorm) \
    .option('checkpointLocation', table_config['checkpoint_location']) \
    .start() \
    .awaitTermination()

def join_kafka_streams_denorm(kafka_df, batch_id):
    try:
        table_config = state_config
        kafka_config = kafkaconfig

        filters = ata_filter(kafka_df=kafka_df)
        main_df = spark.sql(f'select * from db.table where {filters}')

        joined_df = join_remove_duplicate_col(kafka_df=kafka_df, denorm=main_df, table_config=table_config)
        push_to_kafka(joined_df, kafka_config, table_config, 'state')
    except Exception as error:
        print(f'Join failed with the exception: {error}')
        traceback.print_exc()
        print('Stopping the application')
        sys.exit(1)

The method write_stream_batches is receiving a streaming dataframe from kafka. I am merging this topic data to a hive table and I have my table configuration in dictionary imported from a config.py file and below is that line.

table_config = state_config

The problem here is to give the check point configuration, I am importing state_config in write_stream_batches which is fine.

But if I try to pass table_config as an argument in foreachBatch along with join_kafka_streams_denorm the code doesn't work. So I had to import it once again in the method join_kafka_streams_denorm and use the rest of the configuration there. I tried this to run into an error.

def write_stream_batches(kafka_df: DataFrame,table_config):
    table_config = state_config
    kafka_df.writeStream \
    .format('kafka') \
    .foreachBatch(join_kafka_streams_denorm, table_config) \
    .option('checkpointLocation', table_config['checkpoint_location']) \
    .start() \
    .awaitTermination()

But that syntax is not right as table_config is being considered as a batch id.

So can anyone let me know if there is any way I can pass an extra argument in foreachBatch when processing a streaming dataframe.

spark-streaming apache-kafka apache-spark
1 Answers

You can't pass arbitrary parameters in foreachBatch but you can rework your join_kafka_streams_denorm as a higher-order function that can take your config and return the actual processing function:

def join_kafka_streams_denorm(table_config):
    def _do_work(kafka_df, batch_id):
        try:
            table_config = state_config
            kafka_config = kafkaconfig

            filters = ata_filter(kafka_df=kafka_df)
            main_df = spark.sql(f'select * from db.table where {filters}')

            joined_df = join_remove_duplicate_col(kafka_df=kafka_df, denorm=main_df, table_config=table_config)
            push_to_kafka(joined_df, kafka_config, table_config, 'state')
        except Exception as error:
            print(f'Join failed with the exception: {error}')
            traceback.print_exc()
            print('Stopping the application')
            sys.exit(1)
    return _do_work

You can use such a function as:

def write_stream_batches(kafka_df: DataFrame,table_config):
    table_config = state_config
    kafka_df.writeStream \
    .format('kafka') \
    .foreachBatch(join_kafka_streams_denorm(table_config)) \
    .option('checkpointLocation', table_config['checkpoint_location']) \
    .start() \
    .awaitTermination()
user_31696690
Related