Spark 在列中保留与列表匹配的单词

Question

我目前有一个列表和一个 Spark 数据框：

['murder', 'violence', 'flashback', 'romantic', 'cult', 'revenge', 'psychedelic', 'comedy', 'suspenseful', 'good versus evil']

我很难想出一种方法来在数据框中创建一个新列，该列从每一行的 tags 列中获取第一个匹配的单词并将其放入为该行新创建的列。

例如，假设标签列中的第一行只有“谋杀”，我希望它显示在新列中。然后，如果下一行中有“无聊”、“愚蠢”和“邪教”，我希望它在新列中显示邪教，因为它与列表匹配。如果标签列的最后一行有“revenge”、“cult”，我希望它只显示 revenge，因为它是匹配列表的第一个词。

Answer 1

您可以使用 PySpark UDF（用户定义函数）。

首先，让我们编写一个python函数来查找列表（在本例中为您提供的列表）和字符串之间的第一个匹配项，即标签列的值：

def find_first_match(tags):
  first_match = ''
  genres= ['murder', 'violence', 'flashback', 'romantic', 'cult', 'revenge', 'psychedelic', 'comedy', 'suspenseful', 'good versus evil']
  for tag in tags.split():
    for genre in genres:
      if tag==genre:
        first_match=tag
  return first_match

然后，我们需要将此函数转换为 PySpark udf，以便我们可以将其与 .withColumn() 操作结合使用：

find_first_matchUDF = udf(lambda z:find_first_match(z))

现在我们可以应用udf函数来生成一个新的列。假设 df 是您的 DataFrame 的名称：

from pyspark.sql.functions import col

new_df = df.withColumn("first_match", find_first_matchUDF(col("tags")))

此方法仅在 tags 列中的所有标记都由空格分隔时才有效。

P.S

您可以使用注释避免第二步：

from pyspark.sql.functions import col

@udf(returnType=StringType()) 
def find_first_match(tags):
  first_match = ''
  genres= ['murder', 'violence', 'flashback', 'romantic', 'cult', 'revenge', 'psychedelic', 'comedy', 'suspenseful', 'good versus evil']
  for tag in tags.split():
    for genre in genres:
      if tag==genre:
        first_match=tag
  return first_match

new_df = df.withColumn("first_match", find_first_match(col("tags")))

Answer 2

from pyspark.sql import functions as F
df = spark.createDataFrame([('murder',), ('boring silly cult',), ('revenge cult',)], ['tags'])

mylist = ['murder', 'violence', 'flashback', 'romantic', 'cult', 'revenge', 'psychedelic', 'comedy', 'suspenseful', 'good versus evil']
pattern = '|'.join([f'({x})' for x in mylist])

df = df.withColumn('first_from_list', F.regexp_extract('tags', pattern, 0))

df.show()
# +-----------------+---------------+
# |             tags|first_from_list|
# +-----------------+---------------+
# |           murder|         murder|
# |boring silly cult|           cult|
# |     revenge cult|        revenge|
# +-----------------+---------------+

Spark 在列中保留与列表匹配的单词

Spark keeping words in column that match a list

string

list

apache-spark

apache-spark-sql

pyspark