Как найти медиану в Apache Spark с API Python Dataframe?

Question

Pyspark API предоставляет множество агрегатных функций, кроме медианы. Spark 2 поставляется с приблизительно Quantile, который дает приблизительные квантили, но точная медиана очень дорога для вычисления. Есть ли способ Pyspark для вычисления медианы для столбца значений в кадре данных Spark?

Answer 1

Вот пример реализации с использованием Dataframe API в Python (Spark 1.6 +).

import pyspark.sql.functions as F
import numpy as np
from pyspark.sql.types import FloatType

Предположим, у нас есть ежемесячные зарплаты для клиентов в кадре данных "зарплаты", например:

месяц | customer_id | оплата труда

и мы хотели бы найти среднюю зарплату на одного клиента в течение всех месяцев

Шаг 1: написать пользовательскую функцию для вычисления медианы

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 = F.udf(find_median,FloatType())

Шаг 2: Сгруппировать в столбце зарплаты, собрав их в список зарплат в каждом ряду:

salaries_list = salaries.groupBy("customer_id").agg(F.collect_list("salary").alias("salaries"))

Шаг 3: Вызовите udf median_finder в столбце зарплаты и добавьте медианные значения в качестве нового столбца.

salaries_list = salaries_list.withColumn("median",median_finder("salaries")) 

Answer 2

Для точной медианы (для кадров данных малого и среднего размера), начиная с Spark 2.1 , можно использоватьpercentileфункция, завернутая в expr:

      F.expr('percentile(c2, 0.5)')

      df = spark.createDataFrame(
    [(1, 10),
     (1, 20),
     (2, 50)],
    ['c1', 'c2'])
df.groupby('c1').agg(F.expr('percentile(c2, 0.5)').alias('median')).show()
#  +---+------+
#  | c1|median|
#  +---+------+
#  |  1|  15.0|
#  |  2|  50.0|
#  +---+------+
df.withColumn('median', F.expr('percentile(c2, 0.5)').over(W.partitionBy('c1'))).show()
#  +---+---+------+
#  | c1| c2|median|
#  +---+---+------+
#  |  1| 10|  15.0|
#  |  1| 20|  15.0|
#  |  2| 50|  50.0|
#  +---+---+------+

Приблизительная медиана часто может быть лучшим выбором для кадров данных среднего размера.

Искра 2.1 реализуетapprox_percentileа такжеpercentile_approx:

      F.expr('percentile_approx(c2, 0.5)')

Начиная с Spark 3.1 его можно использовать напрямую в PySpark API:

      F.percentile_approx('c2', 0.5)