Как рассчитать медиану в строке данных в 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

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

from pyspark.sql.functions import struct, udf
from pyspark.sql.types import FloatType
import numpy as np

def find_median(values_list):
    try:
        median = np.median(values_list) #get the median of values in a list in each row
        return round(float(median),2)
    except Exception:
        return None #if there is anything wrong with the given values
median_finder = udf(find_median,FloatType())

df = df.withColumn("List_abc", struct(col('a'),col('b'),col('c')))\
       .withColumn("median",median_finder("List_abc")).drop('List_abc')
df.show()
+---+---+---+------+
|  a|  b|  c|median|
+---+---+---+------+
|  1|  2|  3|   2.0|
|  1|  4|100|   4.0|
| 20| 30| 50|  30.0|
+---+---+---+------+

Answer 2

Нет встроенной функции, но вы можете легко написать ее, используя существующие компоненты.

# In Spark < 2.4  replace array_sort with sort_array
# Thanks to @RaphaelRoth for pointing that out
from pyspark.sql.functions import array, array_sort, floor, col, size
from pyspark.sql import Column

def percentile(p, *args):
    def col_(c):
        if isinstance(c, Column):
            return c
        elif isinstance(c, str):
            return col(c)
        else:
            raise TypeError("args should str or Column, got {}".format(type(c)))

    xs = array_sort(array(*[col_(x) for x in args]))
    n = size(xs)
    h = (n - 1) * p
    i = floor(h).cast("int")
    x0, x1 = xs[i], xs[i + 1]
    return x0 + (h - i) * (x1 - x0)

Пример использования:

df.withColumn("median", percentile(0.5, *df.columns)).show()

+---+---+---+------+
|  a|  b|  c|median|
+---+---+---+------+
|  1|  2|  3|   2.0|
|  1|  4|100|   4.0|
| 20| 30| 50|  30.0|
+---+---+---+------+

То же самое можно сделать в Scala :

import org.apache.spark.sql.functions._
import org.apache.spark.sql.Column

def percentile(p: Double, args: Column*) = {
    val xs = array_sort(array(args: _*))
    val n = size(xs)
    val h = (n - 1) * p
    val i = floor(h).cast("int")
    val (x0, x1) = (xs(i), xs(i + 1))
    x0 + (h - i) * (x1 - x0)
}

val df = Seq((1, 2, 3), (1, 4, 100), (20, 30, 50)).toDF("a", "b", "c")
df.withColumn("median", percentile(0.5, $"a", $"b", $"c")).show

+---+---+---+------+
|  a|  b|  c|median|
+---+---+---+------+
|  1|  2|  3|   2.0|
|  1|  4|100|   4.0|
| 20| 30| 50|  30.0|
+---+---+---+------+

Только в Python вы также можете рассмотреть векторизованный UDF - в общем, он, скорее всего, будет медленнее, чем встроенные функции, но лучше, чем векторизованный udf:

from pyspark.sql.functions import pandas_udf, PandasUDFType 
from pyspark.sql.types import DoubleType

import pandas as pd
import numpy as np 

def pandas_percentile(p=0.5):
    assert 0 <= p <= 1
    @pandas_udf(DoubleType()) 
    def _(*args): 
        return pd.Series(np.percentile(args, q = p * 100, axis = 0))
    return _


df.withColumn("median", pandas_percentile(0.5)("a", "b", "c")).show()

+---+---+---+------+                                                            
|  a|  b|  c|median|
+---+---+---+------+
|  1|  2|  3|   2.0|
|  1|  4|100|   4.0|
| 20| 30| 50|  30.0|
+---+---+---+------+

Answer 3

Я немного изменил ответ OmG, чтобы сделать UDF динамическим для числа столбцов 'n' вместо 3 *.

Код:

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

import numpy as np
from pyspark.sql.functions import udf
from pyspark.sql.types import DoubleType

def my_median(*args):
    return float(np.median(list(args)))

udf_median = udf(my_median, DoubleType())

df.withColumn('median', udf_median('a','b','c')).show()

Выход:

Answer 4

определить пользовательскую функцию с помощью 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()