如何使用 Spark 查找中位数和分位数

How to find median and quantiles using Spark

如何使用分布式方法 IPython 和 Spark 找到 RDD 个整数的中位数? RDD 大约有 700,000 个元素,因此太大而无法收集和找到中位数。

这个问题和这个问题类似。但是,问题的答案是使用Scala,我不知道。

使用 Scala 答案的思路,我试图在 Python 中写一个类似的答案。

我知道我首先要对RDD进行排序。我不知道怎么。我看到 sortBy(按给定的 keyfunc 对这个 RDD 进行排序)和 sortByKey(对 RDD 进行排序,假定它由(键,值)对组成。)方法。我认为两者都使用键值,而我的 RDD 只有整数元素。

  1. 首先,我想做myrdd.sortBy(lambda x: x)
  2. 接下来我会找到rdd的长度(rdd.count())。
  3. 最后,我想找到rdd中心的元素或2个元素。我也需要这个方法的帮助。

编辑:

我有个主意。也许我可以索引我的 RDD 然后 key = index 和 value = element。然后我可以尝试按值排序?我不知道这是否可能,因为只有 sortByKey 方法。

正在进行的工作

SPARK-30569 - 添加调用 percentile_approx

的 DSL 函数

Spark 2.0+:

您可以使用 approxQuantile 方法实现 Greenwald-Khanna algorithm:

Python:

df.approxQuantile("x", [0.5], 0.25)

Scala:

df.stat.approxQuantile("x", Array(0.5), 0.25)

其中最后一个参数是相对误差。数字越小,结果越准确,计算成本也越高。

自 Spark 2.2 (SPARK-14352) 起支持多列估计:

df.approxQuantile(["x", "y", "z"], [0.5], 0.25)

df.approxQuantile(Array("x", "y", "z"), Array(0.5), 0.25)

底层方法也可以使用approx_percentile函数在SQL聚合(全局和摸索)中使用:

> SELECT approx_percentile(10.0, array(0.5, 0.4, 0.1), 100);
 [10.0,10.0,10.0]
> SELECT approx_percentile(10.0, 0.5, 100);
 10.0

Spark < 2.0

Python

正如我在评论中提到的,很可能不值得大惊小怪。如果数据相对较小,就像您的情况一样,那么只需在本地收集和计算中位数即可:

import numpy as np

np.random.seed(323)
rdd = sc.parallelize(np.random.randint(1000000, size=700000))

%time np.median(rdd.collect())
np.array(rdd.collect()).nbytes

在我用了几年的旧电脑和大约 5.5MB 的内存上大约需要 0.01 秒。

如果数据大得多,排序将成为一个限制因素,因此与其获取精确值,不如在本地采样、收集和计算更好。但是如果你真的想使用 Spark,像这样的东西应该可以解决问题(如果我没有搞砸任何东西):

from numpy import floor
import time

def quantile(rdd, p, sample=None, seed=None):
    """Compute a quantile of order p ∈ [0, 1]
    :rdd a numeric rdd
    :p quantile(between 0 and 1)
    :sample fraction of and rdd to use. If not provided we use a whole dataset
    :seed random number generator seed to be used with sample
    """
    assert 0 <= p <= 1
    assert sample is None or 0 < sample <= 1

    seed = seed if seed is not None else time.time()
    rdd = rdd if sample is None else rdd.sample(False, sample, seed)

    rddSortedWithIndex = (rdd.
        sortBy(lambda x: x).
        zipWithIndex().
        map(lambda (x, i): (i, x)).
        cache())

    n = rddSortedWithIndex.count()
    h = (n - 1) * p

    rddX, rddXPlusOne = (
        rddSortedWithIndex.lookup(x)[0]
        for x in int(floor(h)) + np.array([0L, 1L]))

    return rddX + (h - floor(h)) * (rddXPlusOne - rddX)

还有一些测试:

np.median(rdd.collect()), quantile(rdd, 0.5)
## (500184.5, 500184.5)
np.percentile(rdd.collect(), 25), quantile(rdd, 0.25)
## (250506.75, 250506.75)
np.percentile(rdd.collect(), 75), quantile(rdd, 0.75)
(750069.25, 750069.25)

最后让我们定义中位数:

from functools import partial
median = partial(quantile, p=0.5)

到目前为止一切顺利,但在没有任何网络通信的本地模式下需要 4.66 秒。可能有办法改善这一点,但为什么还要费心呢?

语言无关Hive UDAF):

如果您使用 HiveContext,您还可以使用 Hive UDAF。具有整数值:

rdd.map(lambda x: (float(x), )).toDF(["x"]).registerTempTable("df")

sqlContext.sql("SELECT percentile_approx(x, 0.5) FROM df")

具有连续值:

sqlContext.sql("SELECT percentile(x, 0.5) FROM df")

percentile_approx 中,您可以传递一个额外的参数来确定要使用的记录数。

如果您只想使用 RDD 方法而不想移动到 DF,请添加解决方案。 这段代码可以让你得到一个 RDD 为 double 的百分位数。

如果您输入的百分位数为 50,您应该会获得所需的中位数。 让我知道是否有任何未解决的极端情况。

/**
  * Gets the nth percentile entry for an RDD of doubles
  *
  * @param inputScore : Input scores consisting of a RDD of doubles
  * @param percentile : The percentile cutoff required (between 0 to 100), e.g 90%ile of [1,4,5,9,19,23,44] = ~23.
  *                     It prefers the higher value when the desired quantile lies between two data points
  * @return : The number best representing the percentile in the Rdd of double
  */    
  def getRddPercentile(inputScore: RDD[Double], percentile: Double): Double = {
    val numEntries = inputScore.count().toDouble
    val retrievedEntry = (percentile * numEntries / 100.0 ).min(numEntries).max(0).toInt


    inputScore
      .sortBy { case (score) => score }
      .zipWithIndex()
      .filter { case (score, index) => index == retrievedEntry }
      .map { case (score, index) => score }
      .collect()(0)
  }

这是我使用 window 函数(使用 pyspark 2.2.0)的方法。

from pyspark.sql import DataFrame

class median():
    """ Create median class with over method to pass partition """
    def __init__(self, df, col, name):
        assert col
        self.column=col
        self.df = df
        self.name = name

    def over(self, window):
        from pyspark.sql.functions import percent_rank, pow, first

        first_window = window.orderBy(self.column)                                  # first, order by column we want to compute the median for
        df = self.df.withColumn("percent_rank", percent_rank().over(first_window))  # add percent_rank column, percent_rank = 0.5 coressponds to median
        second_window = window.orderBy(pow(df.percent_rank-0.5, 2))                 # order by (percent_rank - 0.5)^2 ascending
        return df.withColumn(self.name, first(self.column).over(second_window))     # the first row of the window corresponds to median

def addMedian(self, col, median_name):
    """ Method to be added to spark native DataFrame class """
    return median(self, col, median_name)

# Add method to DataFrame class
DataFrame.addMedian = addMedian

然后调用addMedian方法计算col2的中位数:

from pyspark.sql import Window

median_window = Window.partitionBy("col1")
df = df.addMedian("col2", "median").over(median_window)

终于可以根据需要分组了。

df.groupby("col1", "median")

我写了一个函数,它将数据帧作为输入,returns 一个数据帧,它在一个分区上将中位数作为输出,order_col 是我们要计算中位数的列part_col 是我们要计算 :

中位数的水平
from pyspark.sql import Window
import pyspark.sql.functions as F

def calculate_median(dataframe, part_col, order_col):
    win = Window.partitionBy(*part_col).orderBy(order_col)
#     count_row = dataframe.groupby(*part_col).distinct().count()
    dataframe.persist()
    dataframe.count()
    temp = dataframe.withColumn("rank", F.row_number().over(win))
    temp = temp.withColumn(
        "count_row_part",
        F.count(order_col).over(Window.partitionBy(part_col))
    )
    temp = temp.withColumn(
        "even_flag",
        F.when(
            F.col("count_row_part") %2 == 0,
            F.lit(1)
        ).otherwise(
            F.lit(0)
        )
    ).withColumn(
        "mid_value",
        F.floor(F.col("count_row_part")/2)
    )

    temp = temp.withColumn(
        "avg_flag",
        F.when(
            (F.col("even_flag")==1) &
            (F.col("rank") == F.col("mid_value"))|
            ((F.col("rank")-1) == F.col("mid_value")),
            F.lit(1)
        ).otherwise(
        F.when(
            F.col("rank") == F.col("mid_value")+1,
            F.lit(1)
            )
        )
    )
    temp.show(10)
    return temp.filter(
        F.col("avg_flag") == 1
    ).groupby(
        part_col + ["avg_flag"]
    ).agg(
        F.avg(F.col(order_col)).alias("median")
    ).drop("avg_flag")

有两种方法可以使用。一种是使用 approxQuantile method and the other percentile_approx 方法。但是,当记录数为偶数时,这两种方法都可能无法给出准确的结果。

importpyspark.sql.functions.percentile_approx as F
# df.select(F.percentile_approx("COLUMN_NAME_FOR_WHICH_MEDIAN_TO_BE_COMPUTED", 0.5).alias("MEDIAN)) # might not give proper results when there are even number of records

((
df.select(F.percentile_approx("COLUMN_NAME_FOR_WHICH_MEDIAN_TO_BE_COMPUTED", 0.5) + df.select(F.percentile_approx("COLUMN_NAME_FOR_WHICH_MEDIAN_TO_BE_COMPUTED", 0.51)
)*.5).alias("MEDIAN))

对于精确的中位数计算,您可以使用以下函数并将其与 PySpark DataFrame 一起使用 API:

def median_exact(col: Union[Column, str]) -> Column:
    """
    For grouped aggregations, Spark provides a way via pyspark.sql.functions.percentile_approx("col", .5) function,
    since for large datasets, computing the median is computationally expensive.
    This function manually computes the median and should only be used for small to mid sized datasets / groupings.
    :param col: Column to compute the median for.
    :return: A pyspark `Column` containing the median calculation expression
    """
    list_expr = F.filter(F.collect_list(col), lambda x: x.isNotNull())
    sorted_list_expr = F.sort_array(list_expr)
    size_expr = F.size(sorted_list_expr)

    even_num_elements = (size_expr % 2) == 0
    odd_num_elements = ~even_num_elements

    return F.when(size_expr == 0, None).otherwise(
        F.when(odd_num_elements, sorted_list_expr[F.floor(size_expr / 2)]).otherwise(
            (
                sorted_list_expr[(size_expr / 2 - 1).cast("long")]
                + sorted_list_expr[(size_expr / 2).cast("long")]
            )
            / 2
        )
    )

像这样应用它:

output_df = input_spark_df.groupby("group").agg(
    median_exact("elems").alias("elems_median")
)