MapType 列值上的 PySpark 杠杆函数

Question

下面是一个代表我想要完成的数据框。不过请注意，我要利用的功能比这个例子要复杂一些。

import pyspark
from pyspark.sql import SparkSession

arrayData = [
        ('1',{1:100,2:200}),
        ('1',{1:100,2:None})]

df=spark.createDataFrame(data=arrayData, schema = ['id','value'])

我想做的是利用 withColumn 创建一个新列，其中包含一个已应用函数的新地图类型对象。

假设我想对每个值进行平方。我知道我可以创建一个将值乘以 2 并使用 withColumn 的 udf...但是这似乎不适用于应用于 MapType。

我想要达到的输出是：

arrayData = [
        ('1',{1:200,2:200},{1:400,2:400}),
        ('1',{1:100,2:None},{1:200,2:None})]

df=spark.createDataFrame(data=arrayData, schema = ['id','value','newCol'])

最后，我需要保持并行性并尽量避免爆炸，因为我想将它保持在一行中。这如何实现？

Answer 1

不知道为什么它不适合你，但这个 UDF 对我来说很好用

import pyspark.sql.functions as F

def sq(m):
    return {k:e**2 if e is not None else None for k,e in m.items()}

print(sq({1:100,2:200}))  # {1: 10000, 2: 40000}
print(sq({1:100,2:None})) # {1: 10000, 2: None}

(df
    .withColumn('newCol', F.udf(sq, T.MapType(T.LongType(), T.LongType()))('value'))
    .show(10, False)
)

# Output
# +---+---------------------+------------------------+
# |id |value                |newCol                  |
# +---+---------------------+------------------------+
# |1  |{1 -> 100, 2 -> 200} |{1 -> 10000, 2 -> 40000}|
# |1  |{1 -> 100, 2 -> null}|{1 -> 10000, 2 -> null} |
# +---+---------------------+------------------------+