PySpark结构化流：含列表字段的Kafka DataFrame与MongoDB DataFrame关联问题

问题背景 #

我正在使用PySpark结构化流从Kafka Topic读取数据，其中key为整数类型，value是逗号分隔的整数列表。需要基于value中的列表值，过滤MongoDB DataFrame中id列匹配这些值的数据，得到指定结果。

示例数据 #

Kafka流DataFrame：

MongoDB DataFrame：

期望结果：

解决方案 #

核心思路是先将Kafka流中的value字段拆分为单个整数，再通过关联操作匹配MongoDB中的数据，完全不需要遍历列表。

步骤1：处理Kafka流数据

读取Kafka流后，将二进制类型的key和value转换为对应类型，再把value按逗号拆分并展开为多行数据：

from pyspark.sql import SparkSession
from pyspark.sql.functions import split, explode, col, cast

# 初始化SparkSession
spark = SparkSession.builder \
    .appName("KafkaMongoJoin") \
    .config("spark.mongodb.input.uri", "mongodb://localhost:27017/db.collection") \
    .getOrCreate()

# 读取Kafka流数据
kafka_df = spark.readStream \
    .format("kafka") \
    .option("kafka.bootstrap.servers", "localhost:9092") \
    .option("subscribe", "your_topic") \
    .load()

# 处理key和value：转换类型，拆分value为列表并展开
processed_kafka_df = kafka_df \
    .select(
        col("key").cast("integer").alias("kafka_key"),
        split(col("value").cast("string"), ",").alias("value_list")
    ) \
    .withColumn("target_id", explode(col("value_list")).cast("integer"))

步骤2：读取MongoDB数据

读取MongoDB中的静态DataFrame（若MongoDB数据动态更新，可改用结构化流读取MongoDB变更流）：

# 读取MongoDB数据
mongo_df = spark.read \
    .format("mongo") \
    .option("uri", "mongodb://localhost:27017/db.collection") \
    .load() \
    .select("name", col("id").cast("integer"))

步骤3：关联并获取结果

将处理后的Kafka流数据与MongoDB数据通过target_id和id关联，得到期望结果：

# 关联两个DataFrame
result_df = processed_kafka_df.join(
    mongo_df,
    processed_kafka_df.target_id == mongo_df.id,
    "inner"
).select("name", "id")

# 输出结果（按需选择输出模式和sink）
query = result_df.writeStream \
    .format("console") \
    .outputMode("append") \
    .start()

query.awaitTermination()

备选方案：使用isin过滤（适合小批量数据）

若Kafka流每个批次的value列表数据量较小，可收集为静态集合后用isin过滤，但需注意该操作会将数据拉取到Driver端：

from pyspark.sql.functions import collect_set

# 获取当前批次的目标id集合
target_ids = processed_kafka_df.select(collect_set("target_id")).first()[0]

# 过滤MongoDB数据
filtered_mongo_df = mongo_df.filter(col("id").isin(target_ids))

注意事项 #

• 确保Kafka和MongoDB的连接配置（如bootstrap.servers、mongodb.uri）与实际环境一致。
• 结构化流中inner join在append模式下可正常工作，若需保留更多数据可调整连接类型和输出模式。
• 若MongoDB数据动态更新，建议用结构化流读取MongoDB变更流，实现全实时关联。

内容的提问来源于stack exchange，提问作者Rodrigo Alarcón