使用PySpark DataFrame统计一列作为另一列子串的出现次数

核心思路 #

将两个DataFrame的字符串列拆分为数组，通过判断df1的水果组合是否是df2中items集合的子集，再统计符合条件的匹配次数。

完整代码实现 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import split, size, array_intersect, count, col

# 初始化SparkSession
spark = SparkSession.builder.appName("fruit_frequency_count").getOrCreate()

# 创建示例DataFrame
# df1
data_df1 = [("apple, orange",), ("apple, banana",), ("banana, grape",)]
df1 = spark.createDataFrame(data_df1, ["fruit"])

# df2
data_df2 = [("apple, orange, banana",), ("apple, orange, grape",), ("apple, banana, grape",)]
df2 = spark.createDataFrame(data_df2, ["items"])

# 1. 将逗号分隔的字符串转为数组，便于后续集合操作
df1 = df1.withColumn("fruit_list", split(col("fruit"), ", "))
df2 = df2.withColumn("items_list", split(col("items"), ", "))

# 2. 交叉连接并判断匹配关系，统计频率
result_df = df1.crossJoin(df2) \
    .withColumn("is_match", size(array_intersect(col("fruit_list"), col("items_list"))) == size(col("fruit_list"))) \
    .filter(col("is_match")) \
    .groupBy("fruit") \
    .agg(count("*").alias("frequency"))

# 查看结果
result_df.show()

代码细节解释 #

1. 字符串转数组：使用split函数将水果组合字符串拆分为数组，为后续的集合交集计算做准备。
2. 匹配判断逻辑：通过array_intersect计算两个数组的交集，若交集长度等于df1水果数组的长度，说明该组合完全包含在df2的items中。
3. 分组统计：过滤出所有匹配的行后，按fruit字段分组并计数，得到每个组合的出现频率。

性能优化提示 #

如果数据量较大，直接crossJoin会产生大量中间数据，可通过广播df2减少数据传输开销：

from pyspark.sql.functions import broadcast

result_df = df1.crossJoin(broadcast(df2)) \
    .withColumn("is_match", size(array_intersect(col("fruit_list"), col("items_list"))) == size(col("fruit_list"))) \
    .filter(col("is_match")) \
    .groupBy("fruit") \
    .agg(count("*").alias("frequency"))

内容的提问来源于stack exchange，提问作者Dipie