如何计算 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()