如何将 SparkSession 对象传递给 Kafka-Spark 流的 foreachBatch 方法？

Question

我有一个 python 脚本loader.py，它由创建sparkSession 对象的主类组成​​，如下所示，并调用各种方法来执行不同的操作。 from utils import extract_kafka_data, do_some_transformation

def main():
    try:
        spark = SparkSession.builder.appName(config['kafka_transformations']).enableHiveSupport().getOrCreate()
        kafka_df = extract_kafka_data(spark=spark, config=config, topic_name=topic_name)
        do_some_transformation(kafka_df, spark)
    except Exception as exc:
        print(f'Failed with Exception:{exc}')
        traceback.print_exc()
        print('Stopping the application')
        sys.exit(1)


if __name__ == '__main__':
    main()

extract_kafka_data、do_some_transformation 方法存在于不同的 Python 脚本中：utils.py 我的 utils.py 文件中还有很多其他方法可以执行各种转换。以下是该场景中需要解决的几种方法。

def extract_kafka_data(spark: SparkSession, config: dict, topic_name: str):
    jass_config = config['jaas_config'] + " oauth.token.endpoint.uri=" + '"' + config['endpoint_uri'] + '"' + " oauth.client.id=" + '"' + config['client_id'] + '"' + " oauth.client.secret=" + '"' + config['client_secret'] + '" ;'
    stream_df = spark.readStream \
        .format('kafka') \
        .option('kafka.bootstrap.servers', config['kafka_broker']) \
        .option('subscribe', topic_name) \
        .option('kafka.security.protocol', config['kafka_security_protocol']) \
        .option('kafka.sasl.mechanism', config['kafka_sasl_mechanism']) \
        .option('kafka.sasl.jaas.config', jass_config) \
        .option('kafka.sasl.login.callback.handler.class', config['kafka_sasl_login_callback_handler_class']) \
        .option('startingOffsets', 'earliest') \
        .option('fetchOffset.retryIntervalMs', config['kafka_fetch_offset_retry_intervalms']) \
        .option('fetchOffset.numRetries', config['retries']) \
        .option('failOnDataLoss', 'False') \
        .option('checkpointLocation', checkpoint_location) \
        .load() \
        .select(from_json(col('value').cast('string'), schema).alias("json_dta")).selectExpr('json_dta.*')
    return stream_df

def do_some_transformation(spark: SparkSession, kafka_df: Dataframe):
    kafka_df.writeStream \
        .format('kafka') \
        .foreachBatch(my_transformation_method) \
        .option('checkpointLocation', checkpoint_location) \
        .trigger(processingTime='10 minutes') \
        .start()
        .awaitTermination()

def my_transformation_method(kafka_df: Dataframe, batch_id: int):
    base_delta = DeltaTable.forPath(spark, config['delta_path'])
    base_delta.alias("base") \
        .merge(source=kafka_df.alias("inc"), condition=build_update_condition(config['merge_keys'], config['inc_keys'])) \
        .whenMatchedUpdateAll() \
        .whenNotMatchedInsertAll() \
        .execute()

我在这里面临的问题是方法：my_transformation_method。 内部方法：my_transformation_method 我正在将我的 kafka 数据帧与我的增量表合并。 为了读取基表数据，我需要运行以下语句： base_delta = DeltaTable.forPath(spark, config['delta_path']) 但是这里的问题是方法：foreachBatch 在do_some_transformation 方法中调用的方法：my_transformation_method 只能接收两个方法参数：1. Dataframe 2. batch_id 按照火花流的语法。

我可以将 spark 会话对象设为全局，但我不想这样做，因为它似乎不是标准方式。 当我从 do_some_transformation 调用它时，有什么方法可以使 sparkSession 对象 spark 可用于方法 my_transformation_method ？ 非常感谢任何帮助。

Answer 1

DataFrame API ptovides 可以使用的 sparkSession 方法：

spark = kafka_df.sparkSession()