Comment calculer la médiane par ligne dans un Spark DataFrame

définir une fonction définie par l'utilisateur en utilisant udf , puis en utilisant withColumn pour ajouter la colonne spécifiée au bloc de données:

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()

p>

Il n'y a pas de fonction intégrée, mais vous pouvez facilement en écrire une, en utilisant des composants existants.

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

Exemple d'utilisation:

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|
+---+---+---+------+

La même chose peut être faite dans 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|
+---+---+---+------+

Dans Python uniquement, vous pouvez également envisager l'UDF vectorisé - en général, il est susceptible d'être plus lent que les fonctions intégrées, mais supérieur par rapport aux udf non vectorisés:

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

# 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)

J'ai légèrement modifié la réponse d'OmG pour rendre l'UDF dynamique pour un nombre 'n' de colonnes au lieu de seulement 3.

Code:

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()

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|
+---+---+---+------+