我有一个带有MapType字段的Dataframe。
>>> 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
?
答案 0 :(得分: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) |
+---------------------------------------------------------------------------+