PySpark 聚合操作对 MapType(*, IntegerType()) 类型的 DataFrame 列中的所有行求和答案

【问题标题】：PySpark aggregate operation that sum all rows in a DataFrame column of type MapType(*, IntegerType())PySpark 聚合操作对 MapType(*, IntegerType()) 类型的 DataFrame 列中的所有行求和
【发布时间】：2022-11-30 03:27:42
【问题描述】：

假设您创建了一个具有精确模式的 Spark DataFrame：

import pyspark.sql.functions as sf
from pyspark.sql.types import *

dfschema = StructType([
    StructField("_1", ArrayType(IntegerType())),
    StructField("_2", ArrayType(IntegerType())),
])
df = spark.createDataFrame([[[1, 2, 5], [13, 74, 1]], 
                            [[1, 2, 3], [77, 23, 15]]
                           
                           ], schema=dfschema)
df = df.select(sf.map_from_arrays("_1", "_2").alias("omap"))
df = df.withColumn("id", sf.lit(1))

上面的 DataFrame 看起来像这样：

+---------------------------+---+
|omap                       |id |
+---------------------------+---+
|{1 -> 13, 2 -> 74, 5 -> 1} |1  |
|{1 -> 77, 2 -> 23, 3 -> 15}|1  |
+---------------------------+---+

我想执行以下操作：

df.groupby("id").agg(sum_counter("omap")).show(truncate=False)

你能帮我定义一个 sum_counter 函数吗，它只使用来自 pyspark.sql.functions 的 SQL 函数（所以没有 UDF），它允许我在输出中获得这样一个 DataFrame：

+---+-----------------------------------+
|id |mapsum                             |
+---+-----------------------------------+
|1  |{1 -> 90, 2 -> 97, 5 -> 1, 3 -> 15}|
+---+-----------------------------------+

我可以使用 applyInPandas 解决这个问题：

from pyspark.sql.types import *
from collections import Counter
import pandas as pd

reschema = StructType([
    StructField("id", LongType()),
    StructField("mapsum", MapType(IntegerType(), IntegerType()))
])

def sum_counter(key: int, pdf: pd.DataFrame) -> pd.DataFrame:
    return pd.DataFrame([
        key
        + (sum([Counter(x) for x in pdf["omap"]], Counter()), )
    ])

df.groupby("id").applyInPandas(sum_counter, reschema).show(truncate=False)

+---+-----------------------------------+
|id |mapsum                             |
+---+-----------------------------------+
|1  |{1 -> 90, 2 -> 97, 5 -> 1, 3 -> 15}|
+---+-----------------------------------+

但是，出于性能原因，我想避免使用applyInPandas 或UDFs。有任何想法吗？

【问题讨论】：

标签： python-3.x pyspark apache-spark-sql

【解决方案1】：

您可以先将 omap 分解为单独的行，其中键和值将设置在单独的列中，然后像这样聚合它们：

exploded_df = df.select("*", sf.explode("omap"))
agg_df = exploded_df.groupBy("id", "key").sum("value")
agg_df.groupBy("id").agg(sf.map_from_entries(sf.collect_list(sf.struct("key","sum(value)"))).alias("mapsum")).show(truncate=False)

+---+-----------------------------------+
|id |mapsum                             |
+---+-----------------------------------+
|1  |{2 -> 97, 1 -> 90, 5 -> 1, 3 -> 15}|
+---+-----------------------------------+

【讨论】：