PySpark中pandas的median函数返回错误结果问题咨询

问题重现 #

运行以下PySpark Pandas代码计算列b的中位数：

import pyspark.pandas as ps
psFrame = ps.DataFrame({
    'a': [24, 21, 275, 33, 26,90], 'b': [1, 2, 6,9,4, 6]}, columns=['a', 'b'])

print(psFrame['b'].median())

得到的结果是4.0，但手动计算列b的中位数应为5（列b排序后为[1,2,4,6,6,9]，中间两个数4和6的平均值为5）。

问题原因 #

PySpark Pandas的median()方法底层依赖Spark的中位数实现，Spark默认对偶数个元素采用**下中位数（lower median）**逻辑，即直接取中间位置靠左的数值，而非两个中间值的平均值，因此得到了4.0而非预期的5。

解决方法 #

方法1：使用Spark原生函数计算线性插值的中位数

将PySpark Pandas DataFrame转换为Spark DataFrame，利用approx_percentile函数并指定线性插值：

import pyspark.pandas as ps
from pyspark.sql import functions as F

psFrame = ps.DataFrame({
    'a': [24, 21, 275, 33, 26,90], 'b': [1, 2, 6,9,4, 6]}, columns=['a', 'b'])

spark_df = psFrame.to_spark()
# 计算50%分位数，设置插值方式为linear
median = spark_df.select(F.approx_percentile('b', 0.5, 0).alias('median')).collect()[0]['median']
print(median)  # 输出5.0

方法2：用PySpark Pandas的quantile方法指定插值方式

直接使用quantile方法计算0.5分位数，将interpolation参数设为'linear'，即可得到两个中间值的平均值：

import pyspark.pandas as ps

psFrame = ps.DataFrame({
    'a': [24, 21, 275, 33, 26,90], 'b': [1, 2, 6,9,4, 6]}, columns=['a', 'b'])

median = psFrame['b'].quantile(0.5, interpolation='linear')
print(median)  # 输出5.0

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