Pyspark DataFrame中两列逗号分隔字符串的相似度计算方法咨询

核心思路 #

直接使用PySpark内置的数组操作函数实现需求，无需引入第三方依赖，性能优于自定义UDF方案，计算逻辑完全匹配你给出的公式：匹配段落数 / col2总段落数。

内置函数实现代码（推荐，Spark 2.4+） #

# 导入依赖
from pyspark.sql import SparkSession
from pyspark.sql.functions import split, size, array_intersect, when, trim, lower, transform

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

# 示例数据，可替换为你的实际DataFrame
test_data = [
    ("苹果,香蕉,橙子,葡萄", "香蕉,橙子,西瓜"),
    ("Java,Python,Go", "Java,Python,Go"),
    ("123,456,789", "abc,def"),
    ("测试文本", "")
]
df = spark.createDataFrame(test_data, schema=["col1", "col2"])

# 计算相似度
result_df = df.withColumn("col1_segs", split("col1", ",")) \
              .withColumn("col2_segs", split("col2", ",")) \
              # 可选：如果需要忽略大小写和前后空格，可加上下面两行对数组元素做标准化处理
              # .withColumn("col1_segs", transform("col1_segs", lambda x: lower(trim(x)))) \
              # .withColumn("col2_segs", transform("col2_segs", lambda x: lower(trim(x)))) \
              .withColumn("similarity", 
                          when(size("col2_segs") == 0, 0.0)  # 处理col2为空的情况，避免除以0报错
                          .otherwise(size(array_intersect("col1_segs", "col2_segs")) / size("col2_segs"))
              )
# 查看结果
result_df.select("col1", "col2", "similarity").show(truncate=False)

# 可选：删除中间生成的数组列
# result_df = result_df.drop("col1_segs", "col2_segs")

低版本Spark兼容方案（Spark 2.4以下，自定义UDF） #

如果你的Spark版本不支持array_intersect函数，可以用自定义UDF实现相同逻辑：

from pyspark.sql.types import DoubleType
from pyspark.sql.functions import udf

def calc_sim(col1_str: str, col2_str: str) -> float:
    if not col2_str:
        return 0.0
    col1_seg_set = set(col1_str.split(","))
    col2_seg_list = col2_str.split(",")
    match_count = sum(1 for seg in col2_seg_list if seg in col1_seg_set)
    return match_count / len(col2_seg_list)

sim_udf = udf(calc_sim, DoubleType())
result_df = df.withColumn("similarity", sim_udf("col1", "col2"))

注意事项 #

• 如果你的场景中段落存在重复值，需要去重后统计匹配数，可以先对拆分后的数组调用array_distinct函数去重再计算
• 内置函数方案没有序列化和Python进程通信开销，大数据量下性能远高于自定义UDF，优先使用
• 可根据业务需要调整col2为空时的默认返回值，比如将0.0改为None

内容的提问来源于stack exchange，提问作者ABDELOUAHAB AMINE TAFAT