将新的键/值对添加到 Spark MapType 列

Question

我有一个带有 MapType 字段的数据框。

>>> from pyspark.sql.functions import *
>>> from pyspark.sql.types import *
>>> fields = StructType([
...         StructField('timestamp',      TimestampType(), True),
...         StructField('other_field',    StringType(), True),
...         StructField('payload',        MapType(
...                                         keyType=StringType(),
...                                         valueType=StringType()),
...                                                     True),   ])
>>> import datetime
>>> rdd = sc.parallelize([[datetime.datetime.now(), 'this should be in', {'akey': 'aValue'}]])
>>> df = rdd.toDF(fields)
>>> df.show()
+--------------------+-----------------+-------------------+
|           timestamp|      other_field|            payload|
+--------------------+-----------------+-------------------+
|2018-01-10 12:56:...|this should be in|Map(akey -> aValue)|
+--------------------+-----------------+-------------------+

我想添加other_field作为字段中的键payload。

我知道我可以使用 udf：

>>> def _add_to_map(name, value, map_field):
...     map_field[name] = value
...     return map_field
...
>>> add_to_map = udf(_add_to_map, MapType(StringType(),StringType()))
>>> df.select(add_to_map(lit('other_field'), 'other_field', 'payload')).show(1, False)
+------------------------------------------------------+
|PythonUDF#_add_to_map(other_field,other_field,payload)|
+------------------------------------------------------+
|Map(other_field -> this should be in, akey -> aValue) |
+------------------------------------------------------+

有没有办法在没有 的情况下做到这一点udf？

Answer 1

使用map_concat和create_map（pyspark 2.4+）：

(
    df.withColumn(
        "new_map",
        F.map_concat(
            "old_map",
            F.create_map(F.lit("key"), F.lit("val"))
        )
    )
)

借助 F.create_map，您可以一次添加多个键，但 F.map_concat 不会替换旧键。

Answer 2

udf如果您提前知道密钥，这里有一种方法可以做到这一点。使用该create_map功能。至于这是否更有效，我不知道。

from pyspark.sql.functions import col, lit, create_map

df.select(
    create_map(
        lit('other_field'),
        col('other_field'),
        lit('akey'),
        col('payload')['akey']
    )
).show(n=1, truncate=False)

输出：

+-----------------------------------------------------+
|map(other_field, other_field, akey, payload['akey']) |
+-----------------------------------------------------+
|Map(other_field -> this should be in, akey -> aValue)|
+-----------------------------------------------------+

更新

这是一种无需对字典键进行硬编码的方法。不幸的是，它涉及一项collect()操作。

模拟一些数据

首先，让我们修改原始数据框以在字段中再包含一个键值对MapType()。

from pyspark.sql.functions import col, lit, create_map
import datetime
rdd = sc.parallelize(
    [
        [
            datetime.datetime.now(),
            'this should be in',
            {'akey': 'aValue', 'bkey': 'bValue'}
        ]
    ]
)
df = rdd.toDF(fields)
df.show(n=1, truncate=False)

这会创建以下内容：

+--------------------------+-----------------+-----------------------------------+
|timestamp                 |other_field      |payload                            |
+--------------------------+-----------------+-----------------------------------+
|2018-01-10 17:37:58.859603|this should be in|Map(bkey -> bValue, akey -> aValue)|
+--------------------------+-----------------+-----------------------------------+

获取地图的钥匙

使用explode()和collect()，您可以获得密钥，如下所示：

from pyspark.sql.functions import explode

keys = [
    x['key'] for x in (df.select(explode("payload"))
                        .select("key")
                        .distinct()
                        .collect())
]

创建包含所有字段的新地图

现在如上所述使用create_map，但使用来自的信息keys动态创建键值对。我使用它是reduce(add, ...)因为create_map期望输入按顺序为键值对 - 我想不出另一种方法来压平列表。

from operator import add
df.select(
    create_map
    (
        *([lit('other_field'), col('other_field')] + reduce(add, [[lit(k), col('payload').getItem(k)] for k in keys]))
    )
).show(n=1, truncate=False)

最后结果：

+---------------------------------------------------------------------------+
|map(other_field, other_field, akey, payload['akey'], bkey, payload['bkey'])|
+---------------------------------------------------------------------------+
|Map(other_field -> this should be in, akey -> aValue, bkey -> bValue)      |
+---------------------------------------------------------------------------+

参考

1. pyspark：从现有列创建 MapType 列

2. PySpark 将“map”类型的列转换为数据框中的多列