pyspark根据条件从结构的数组Column中获取元素

Question

我有一个带有以下模式的spark df：

|-- col1 : string
|-- col2 : string
|-- customer: struct
|    |-- smt: string
|    |-- attributes: array (nullable = true)
|    |    |-- element: struct
|    |    |     |-- key: string
|    |    |     |-- value: string

df:

#+-------+-------+---------------------------------------------------------------------------+
#|col1   |col2   |customer                                                                   |
#+-------+-------+---------------------------------------------------------------------------+
#|col1_XX|col2_XX|"attributes":[[{"key": "A", "value": "123"},{"key": "B", "value": "456"}]  |
#+-------+-------+---------------------------------------------------------------------------+

数组的json输入如下所示：

...
"attributes": [
{
"key": "A",
"value": "123"
},
{
"key": "B",
"value": "456"
}
],

我想循环属性数组，获得具有key="B"的元素，然后选择相应的value。我不想使用爆炸，因为我想避免连接数据帧。是否可以直接使用火花"柱"执行这种操作？

预期输出为：

#+-------+-------+-----+
#|col1   |col2   |B    |                                                               |
#+-------+-------+-----+
#|col1_XX|col2_XX|456  |
#+-------+-------+-----+

如有任何帮助，将不胜感激

Answer 1

您可以使用filter函数过滤结构数组，然后获得value:

from pyspark.sql import functions as F
df2 = df.withColumn(
"B", 
F.expr("filter(customer.attributes, x -> x.key = 'B')")[0]["value"]
)