PySpark 数组列答案

【问题标题】：PySpark array columnPySpark 数组列
【发布时间】：2021-08-05 16:21:19
【问题描述】：

我有一个 PySpark DataFrame，其中包含一组书籍，其中每本书可以有一个或多个标题。每个标题都被归类为原始标题OT 或替代标题AT。为简单起见，我省略了其他标题类型。我的验证需要确保每本书都有一个 OT 标题，可以有任意数量的 AT 标题。

我要做的是清理数据，以便：

如果一本书有多个OT 标题，请保留第一个并将其余更改为AT
如果一本书没有OT 标题，请将第一个AT 标题更改为OT

from pyspark.sql.types import StructType, StructField, IntegerType, StringType
from pyspark.sql.functions import collect_list, col, struct

data = ([
  (1, 'Title 1', 'OT'),
  (1, 'Title 2', 'OT'),
  (2, 'Title 3', 'AT'),
  (2, 'Title 4', 'OT'),
  (3, 'Title 5', 'AT'),
])

schema = StructType([ 
    StructField("BookID", IntegerType(), False),
    StructField("Title", StringType(), True), 
    StructField("Type", StringType(), True),
  ])

df = spark.createDataFrame(data, schema)
df = df.groupby('BookID').agg(collect_list(struct(col('Title'), col('Type'))).alias('Titles'))

display(df)

听起来应该很容易，但我不知道该怎么做。任何帮助将不胜感激。

我尝试过使用如下所示的 udf，但到目前为止，这种方法行不通。我收到一条错误消息，提示 lambda cannot contain assignment。

def process_titles(titles):
  x = list(filter(lambda t: t.Type == 'OT', titles))[1::]
  map(lambda t: t.Type = 'AT', x)
  
  return x

process_titles_udf = udf(lambda x: process_titles(x), titles)

df = df.withColumn('test', process_titles_udf('Titles'))

udf 返回一个类型的对象：

titles = ArrayType(StructType([ 
    StructField("Title", StringType(), True), 
    StructField("Type", StringType(), True)
  ]))

【问题讨论】：

titles 中的process_titles_udf = udf(lambda x: process_titles(x), titles) 是什么？
抱歉，这是复制粘贴错误。那将是包含标题和类型的结构的 ArrayType()。我现在会更新问题。

标签： apache-spark pyspark

【解决方案1】：

首先，当您说“保持第一”时，您必须知道collect_list 是不确定的。所以根据你的跑步，你可能有不同的“第一个”OT。

如果您想继续这种不确定的行为，这里是您的 UDF：

@udf(titles)
def process_titles(titles):
    OTs = [x for x in titles if x["Type"] == "OT"]  # Collect all OT types
    if OTs:
        OT = OTs[0]  # Keep the first OT as only OT if it exists
    else:
        OT = {
            "Title": titles[0]["Title"],
            "Type": "OT",
        }  # otherwise, use the first AT as OT

    ATs = [
        {"Title": x["Title"], "Type": "AT"} for x in titles if x["Title"] != OT["Title"]
    ]  # Transform all other titles as AT
    return [OT] + ATs


df.select("titles", process_titles(F.col("Titles"))).show(truncate=False)
+------------------------------+------------------------------+                 
|titles                        |process_titles(Titles)        |
+------------------------------+------------------------------+
|[[Title 1, OT], [Title 2, OT]]|[[Title 1, OT], [Title 2, AT]]|
|[[Title 5, AT]]               |[[Title 5, OT]]               |
|[[Title 3, AT], [Title 4, OT]]|[[Title 4, OT], [Title 3, AT]]|
+------------------------------+------------------------------+

【讨论】：