рассчитать медианные значения с четным количеством строк в pyspark

Question

Я использую эту формулу, взятую из другого поста SO, для вычисления средних значений столбцов в pyspark:

columns = ['id', 'dogs', 'cats']
vals = [(1, 2, 0),(2, 0, 1)]
df = sqlContext.createDataFrame(vals, columns)
df.approxQuantile(list(c for c in df.columns), [0.5], 0)

Формула работает, когда в df есть нечетное количество строк, но если числоколичество строк четное, оно не принимает среднее значение между центральными элементами (я имею в виду после сортировки, конечно), а только первое.

Например, результат из приведенного выше кода:

[[1.0], [0.0], [0.0]] 

в то время как должно было быть:

[[1.5], [1.0], [0.5]] 

Как правильно рассчитать медианное значение в pyspark?

Answer 1

Вот хакерский способ получить медиану

import numpy as np
np.random.seed(2019)

def get_median(df, col):
    # if # df rows are even get the "middle" values and average
    num_rows = df.count()
    if num_rows%2==0:
        mid_pt = num_rows/2
        q1 = df.approxQuantile( col, [0.5], 0)
        q2 = df.approxQuantile( col, [(mid_pt+1)*1./num_rows], 0)
        return(np.mean(q1+q2).tolist())
    return(df.approxQuantile( col, [0.5], 0))

# now test it out
df = pd.DataFrame(data = np.random.rand(20,2), columns=['dogs', 'cats'])
df['id'] = list(range(20))

df = sqlContext.createDataFrame(df)
df.approxQuantile( "cats", [0.5], 0)
# [0.5783015865898744]