PySpark使用pivot时出现非预期隐式类型转换，如何保留原类型？

Spark中sum聚合函数对Short/Integer类型的输入，默认会自动提升为Long类型返回，这是框架内置的防溢出机制——即使是在pivot操作中调用sum，也会遵循这个逻辑，所以你看到结果列变成了LongType。遗憾的是，Spark没有提供全局配置直接禁用这种隐式类型提升。

要彻底避免生成Long类型的中间数据，你可以自定义针对ShortType的聚合函数（UDAF），直接在聚合阶段保持Short类型计算，无需转换为Long。

自定义Short类型求和UDAF实现 #

from pyspark.sql import SparkSession
from pyspark.sql.types import *
from pyspark.sql.functions import *
from pyspark.sql.expressions import Aggregator

# 自定义聚合器：针对Short类型求和，保持结果为Short
class ShortSum(Aggregator):
    # 输入类型：ShortType
    def inputSchema(self):
        return StructType([StructField("value", ShortType())])
    
    # 缓冲区类型：用Short存储累加值
    def bufferSchema(self):
        return StructType([StructField("sum", ShortType())])
    
    # 输出类型：ShortType
    def outputDataType(self):
        return ShortType()
    
    # 是否确定性运算
    def deterministic(self):
        return True
    
    # 初始化缓冲区
    def initialize(self, buffer):
        buffer[0] = 0
    
    # 累加输入值到缓冲区
    def update(self, buffer, input):
        if input[0] is not None:
            buffer[0] = buffer[0] + input[0]
    
    # 合并两个缓冲区
    def merge(self, buffer1, buffer2):
        buffer1[0] = buffer1[0] + buffer2[0]
    
    # 生成最终结果
    def evaluate(self, buffer):
        return buffer[0]

# 将自定义聚合器注册为可调用的函数
short_sum = udf(ShortSum(), ShortType())

使用自定义UDAF完成pivot聚合 #

# 初始化SparkSession
spark = SparkSession.builder.appName("ShortPivotSum").getOrCreate()

# 构建原始DataFrame（与你的示例一致）
my_schema = StructType([
   StructField("key1", IntegerType(), False),
   StructField("key2", IntegerType(), False),
   StructField("value", ShortType(), False)
])

df1 = spark.createDataFrame(
    [[1,1,1],[1,1,1],[1,2,1],[2,1,1],[2,2,1]],
    schema=my_schema
)

# 使用自定义short_sum进行pivot聚合
df2 = df1.groupby('key1').pivot('key2').agg(short_sum('value').alias('sum'))
# 重命名列（可选，与原示例列名保持一致）
df2 = df2.withColumnRenamed("sum_1", "1").withColumnRenamed("sum_2", "2")

# 查看最终schema
df2.printSchema()

执行后df2的schema会是：

root
 |-- key1: integer (nullable = false)
 |-- 1: short (nullable = true)
 |-- 2: short (nullable = true)

方案优势 #

这个方法全程在聚合阶段使用Short类型计算，不会生成Long类型的中间数据，完全避免了不必要的内存和计算资源浪费，符合你数据量可控、无溢出风险的场景需求。

内容的提问来源于stack exchange，提问作者user344577