在 PySpark 中不使用 for 循环查找 DataFrame 列中存在的相似字符串

Question

我有一个包含字符串列的 DataFrame。我想找到类似的字符串并用一些标志标记它们。我正在使用来自 python-Levenshtein 模块的函数 ratio 并希望将比率大于 0.90 的字符串标记为“相似”。以下是我拥有的 DataFrame 的示例：

sentenceDataFrame = spark.createDataFrame([
    (0, "Hi I heard about Spark"),
    (1, "I wish Java could use case classes"),
    (2, "Logistic,regression,models,are,neat"),
    (3, "Logistic,regression,model,are,neat")
], ["id", "sentence"])

想要的输出是：

+---+-----------------------------------+------------+
|id |sentence                           |similar_flag|
+---+-----------------------------------+------------+
|0  |Hi I heard about Spark             |            |
|1  |I wish Java could use case classes |            |
|2  |Logistic regression models are neat|2_0         |
|3  |Logistic regression model is neat  |2_1         |
|4  |Logistics regression model are neat|2_2         |
+---+-----------------------------------+------------+

其中“2_1”表示“2”是参考字符串的“id”（用于匹配的第一个唯一字符串），“1”表示与其匹配的第一个字符串。我想完全避免 for 循环。对于较小的数据，我使用 for-loop 在简单的 python 中实现了预期的结果，并希望在 PySpark 中也有相同的结果，因此我不想使用除 python-Levenshtein 之外的任何模块。我遇到过this 方法，但它要求我放弃 python-Levenshtein 模块。此外，我的 DataFrame 可能很大（并且预计每天都会增长），因此这种方法可能会导致内存错误。有没有更好的方法来达到预期的效果？

Answer 1

我会分三步回答。首先，您需要允许df 查看所有选项，因此您可能需要使用crossJoin 对您的数据进行Carthesian 乘积，例如：

from pyspark.sql import functions as f

df_new = (
    sentenceDataFrame.crossJoin(
                         sentenceDataFrame.select(
                             f.col('sentence').alias('second_sentence'),
                             f.col('id').alias('second_id')))
)

其次，看看pyspark.sql.functions.levehstein。一旦你的句子被一个一个地排列在一起，添加一个新的列，其中包含 Levehstein 距离，使用

df_new_with_dist = df_new.withColumn('levehstein_distance',
    f.levenshtein(f.col("sentence"), f.col("second_sentence"))
)

df_new_with_dist.show()

+---+--------------------+--------------------+---------+-------------------+
| id|            sentence|     second_sentence|second_id|levehstein_distance|
+---+--------------------+--------------------+---------+-------------------+
|  0|Hi I heard about ...|Hi I heard about ...|        0|                  0|
|  0|Hi I heard about ...|I wish Java could...|        1|                 27|
|  0|Hi I heard about ...|Logistic,regressi...|        2|                 29|
|  0|Hi I heard about ...|Logistic,regressi...|        3|                 28|
|  1|I wish Java could...|Hi I heard about ...|        0|                 27|
|  1|I wish Java could...|I wish Java could...|        1|                  0|
|  1|I wish Java could...|Logistic,regressi...|        2|                 32|
|  1|I wish Java could...|Logistic,regressi...|        3|                 31|
|  2|Logistic,regressi...|Hi I heard about ...|        0|                 29|
|  2|Logistic,regressi...|I wish Java could...|        1|                 32|
|  2|Logistic,regressi...|Logistic,regressi...|        2|                  0|
|  2|Logistic,regressi...|Logistic,regressi...|        3|                  1|
|  3|Logistic,regressi...|Hi I heard about ...|        0|                 28|
|  3|Logistic,regressi...|I wish Java could...|        1|                 31|
|  3|Logistic,regressi...|Logistic,regressi...|        2|                  1|
|  3|Logistic,regressi...|Logistic,regressi...|        3|                  0|
+---+--------------------+--------------------+---------+-------------------+

最后，过滤掉id == second_id 所在的所有行。如果你想坚持你的符号，例如，2_1，我建议你添加groupBy(f.col("id"))并在levehstein_distance上聚合f.min()。然后您可以连接您的 ID，例如使用

min_dist_df = (
    df_new_with_dist.where(f.col('id') != f.col('second_id'))
                    .groupBy(f.col('id').alias('second_id'))
                    .agg(f.min(f.col('levehstein_distance')).alias('levehstein_distance'))
)


(
    df_new_with_dist.join(min_dist_df,
                          on=['second_id', 'levehstein_distance'],
                          how='inner')
                    .withColumn('similar_flag', f.concat(f.concat(f.col('id'), f.lit('_'), f.col('second_id'))))
                    .select('id', 'sentence', 'similar_flag')
).show()

+---+--------------------+------------+
| id|            sentence|similar_flag|
+---+--------------------+------------+
|  2|Logistic,regressi...|         2_3|
|  1|I wish Java could...|         1_0|
|  0|Hi I heard about ...|         0_1|
|  3|Logistic,regressi...|         3_2|
+---+--------------------+------------+

虽然这不是完全您所要求的，但您可以过滤和调整levehstein_distance 值以获得您想要的答案。