apache-spark - 如何计算 Spark DataFrame 中的行中位数

Question

我有一个以下格式的 Spark 数据框。

df = spark.createDataFrame([(1, 2, 3), (1, 4, 100), (20, 30, 50)],['a', 'b', 'c'])
df.show()

输入:

我想添加一个新列“中位数”作为“a”、“b”、“c”列的中位数。如何在 PySpark 中做到这一点。

预期输出:

我使用的是 Spark 2.3.1

Answer 1

使用udf定义一个用户自定义函数，然后使用withColumn将指定的列添加到数据框中:

from numpy import median
from pyspark.sql.functions import col, udf
from pyspark.sql.types import IntegerType

def my_median(a, b, c):
    return int(median([int(a),int(b),int(c)]))

udf_median = udf(my_median, IntegerType())

df_t = df.withColumn('median', udf_median(df['a'], df['b'], df['c']))
df_t.show()