将列添加到嵌套在数组中的结构

Question

我有一个带有结构数组的 PySpark DataFrame，包含两列（colorcode 和 name）。我想向结构添加一个新列，newcol。

回答了“如何将列添加到嵌套结构”，但我无法将其转移到我的案例中，其中结构进一步嵌套在数组中。我似乎无法 reference/recreate 数组结构模式。

我的架构：

 |-- Id: string (nullable = true)
 |-- values: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- Dep: long (nullable = true)
 |    |    |-- ABC: string (nullable = true)

应该变成：

 |-- Id: string (nullable = true)
 |-- values: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- Dep: long (nullable = true)
 |    |    |-- ABC: string (nullable = true)
 |    |    |-- newcol: string (nullable = true)

如何将解决方案转移到我的嵌套结构中？

获取上述模式的 df 的可重现代码：

data = [
    (10, [{"Dep": 10, "ABC": 1}, {"Dep": 10, "ABC": 1}]),
    (20, [{"Dep": 20, "ABC": 1}, {"Dep": 20, "ABC": 1}]),
    (30, [{"Dep": 30, "ABC": 1}, {"Dep": 30, "ABC": 1}]),
    (40, [{"Dep": 40, "ABC": 1}, {"Dep": 40, "ABC": 1}])
  ]
myschema = StructType(
[
    StructField("id", IntegerType(), True),
    StructField("values",
                ArrayType(
                    StructType([
                        StructField("Dep", StringType(), True),
                        StructField("ABC", StringType(), True)
                    ])
    ))
]
)
df = spark.createDataFrame(data=data, schema=myschema)
df.printSchema()
df.show(10, False)

Answer 1

对于spark版本>=3.1，可以使用transform函数和withField方法来实现

transform根据提供的函数对array（此处为values列）中的每个元素（此处为struct(Dep, ABC)）进行变换计算。 withField adds/replaces StructType 中的字段名称。

df = df.withColumn('values', F.transform('values', lambda x: x.withField('newcol', F.lit(1))))

Answer 2

另一种方法是使用 sql 表达式。

df = df.withColumn("values",F.expr("transform(values, x -> struct(COALESCE('1') as newcol,x.Dep,x.ABC))"))

将列添加到嵌套在数组中的结构

Add a column to a struct nested in an array

python

pyspark

apache-spark-sql