PySpark：替换 ArrayType(String) 中的值

Question

我目前有以下代码：

def _join_intent_types(df):
  mappings = {
    'PastNews': 'ContextualInformation',
    'ContinuingNews': 'News',
    'KnownAlready': 'OriginalEvent',
    'SignificantEventChange': 'NewSubEvent',
  }
  return df.withColumn('Categories', posexplode('Categories').alias('i', 'val'))\
           .when(col('val').isin(mappings), mappings[col('i')])\
           .otherwise(col('val'))

但我不确定我的语法是否正确。我要做的是对一列列表进行操作，例如：

['EmergingThreats', 'Factoid', 'KnownAlready']

并用提供的字典中存在的映射替换该数组中的字符串，即。

['EmergingThreats', 'Factoid', 'OriginalEvent']

我尝试了多种方法，但似乎无法应用转换，任何帮助将不胜感激。我知道使用 UDF 可以做到这一点，但我担心这会如何影响性能和可扩展性。

更新： 我在下面提供了原始表结构的示例，

+------------------+-----------------------------------------------------------+
|postID            |Categories                                                 |
+------------------+-----------------------------------------------------------+
|266269932671606786|[EmergingThreats, Factoid, KnownAlready]                   |
|266804609954234369|[Donations, ServiceAvailable, ContinuingNews]              |
|266250638852243457|[EmergingThreats, Factoid, ContinuingNews]                 |
|266381928989589505|[EmergingThreats, MultimediaShare, Factoid, ContinuingNews]|
|266223346520297472|[EmergingThreats, Factoid, KnownAlready]                   |
+------------------+-----------------------------------------------------------+

如果这些数组中的字符串存在于字典中，我想要代码用它们的新映射替换它们。如果没有，请保持原样：

+------------------+-------------------------------------------------+          
|postID            |Categories                                       |
+------------------+-------------------------------------------------+
|266269932671606786|[EmergingThreats, Factoid, OriginalEvent]        |
|266804609954234369|[Donations, ServiceAvailable, News]              |
|266250638852243457|[EmergingThreats, Factoid, News]                 |
|266381928989589505|[EmergingThreats, MultimediaShare, Factoid, News]|
|266223346520297472|[EmergingThreats, Factoid, OriginalEvent]        |
+------------------+-------------------------------------------------+

Answer 1

使用explode + collect_list is expensive。这是未经测试的，但应该适用于 Spark 2.4+：

from pyspark.sql.functions import expr

for k, v in mappings.items()
    df = df.withColumn(
        'Categories', 
        expr('transform(sequence(0,size(Categories)-1), x -> replace(Categories[x], {k}, {v}))'.format(k=k, v=v))
    )

您还可以将映射转换为 CASE/WHEN 语句，然后将其应用于 SparkSQL 转换函数：

sql_epxr = "transform(Categories, x -> CASE x {} ELSE x END)".format(" ".join("WHEN '{}' THEN '{}'".format(k,v) for k,v in mappings.items()))
# this yields the following SQL expression:
# transform(Categories, x -> 
#   CASE x 
#     WHEN 'PastNews' THEN 'ContextualInformation' 
#     WHEN 'ContinuingNews' THEN 'News' 
#     WHEN 'KnownAlready' THEN 'OriginalEvent' 
#     WHEN 'SignificantEventChange' THEN 'NewSubEvent' 
#     ELSE x 
#   END
# )

df.withColumn('Categories', expr(sql_epxr)).show(truncate=False)    

对于旧版本的 spark，udf 可能是首选。

Answer 2

您可以通过一系列步骤来做到这一点，

import pandas as pd
from pyspark.sql.functions as F
from itertools import chain

df = pd.DataFrame()
df['postID'] = [266269932671606786, 266804609954234369, 266250638852243457]
df['Categories']= [
  ['EmergingThreats', 'Factoid', 'KnownAlready'],
  ['Donations', 'ServiceAvailable', 'ContinuingNews'],
  ['EmergingThreats', 'Factoid', 'ContinuingNews'] 
]

sdf = sc.createDataFrame(df)

mappings = {
    'PastNews': 'ContextualInformation',
    'ContinuingNews': 'News',
    'KnownAlready': 'OriginalEvent',
    'SignificantEventChange': 'NewSubEvent',
    'Donations': 'x'
  }

mapping_expr = F.create_map([F.lit(x) for x in chain(*mappings.items())])

sdf.select(F.col("postID"), F.explode("Categories").alias("Categories")) \
            .withColumn("Categories", F.coalesce(mapping_expr.getItem(F.col("Categories")), F.col('Categories'))) \
            .groupBy('postID').agg(F.collect_list('Categories').alias('Categories')) \
            .show(truncate=False)


+------------------+-----------------------------------------+
|postID            |Categories                               |
+------------------+-----------------------------------------+
|266250638852243457|[EmergingThreats, Factoid, News]         |
|266804609954234369|[x, ServiceAvailable, News]              |
|266269932671606786|[EmergingThreats, Factoid, OriginalEvent]|
+------------------+-----------------------------------------+

Answer 3

您可以explodeCategories 列，然后na.replace 与字典后跟groupby 并使用collect_list 聚合为数组：

import pyspark.sql.functions as F

out = (df.select(F.col("postID"),F.explode("Categories").alias("Categories"))
         .na.replace(mappings).groupby("postID")
        .agg(F.collect_list("Categories").alias("Categories")))

---

out.show(truncate=False)

+------------------+-------------------------------------------------+
|postID            |Categories                                       |
+------------------+-------------------------------------------------+
|266269932671606786|[EmergingThreats, Factoid, OriginalEvent]        |
|266250638852243457|[EmergingThreats, Factoid, News]                 |
|266381928989589505|[EmergingThreats, MultimediaShare, Factoid, News]|
|266804609954234369|[Donations, ServiceAvailable, News]              |
|266223346520297472|[EmergingThreats, Factoid, OriginalEvent]        |
+------------------+-------------------------------------------------+

---

更新：

如 cmets 中所述，您可以考虑使用 udf 考虑性能：

def fun(x):
    return [mappings.get(i,i) for i in x]
myudf = F.udf(fun)
df.withColumn("Categories",myudf(F.col("Categories"))).show(truncate=False)

+------------------+-------------------------------------------------+
|postID            |Categories                                       |
+------------------+-------------------------------------------------+
|266269932671606786|[EmergingThreats, Factoid, OriginalEvent]        |
|266804609954234369|[Donations, ServiceAvailable, News]              |
|266250638852243457|[EmergingThreats, Factoid, News]                 |
|266381928989589505|[EmergingThreats, MultimediaShare, Factoid, News]|
|266223346520297472|[EmergingThreats, Factoid, OriginalEvent]        |
+------------------+-------------------------------------------------+