加权中值Pypark数据帧

Question

为了计算加权中值，我编写了以下代码。我哪里出错了，生成的值是空的？列A是值，列B是与这些值相关联的权重

代码：

def get_median(values,weights):
    return np.median(np.repeat(values,weights))    # function created to calculate wt. median

wimedian = F.udf(get_median,DoubleType())    # registering as udf here

myview = df.groupBy('category').agg(
    F.collect_list(F.col('col_A')),
    F.collect_list(F.col('col_B'))
).withColumn('Weighted_median',wimedian(F.col('col_A'),F.col('col_B')))

myview.show(3)

输出表：

+-----------+--------+-------+---------------+
|category   |col_A   |col_B  |Weighted_median|
+-----------+--------+-------+---------------+
|001        |[69]    |[8]    |null           |
|002        |[69]    |[14]   |null           |
|003        |[28, 21]|[3, 1] |null           |
+-----------+--------+-------+---------------+

仅供参考，此表第3行的正确输出应为median of [28,28,28,21] = 28。 这就是为什么np.median和np.repeat在那里的原因

Answer 1

问题可能是返回类型，因为dataframe不理解numpy类型，而且withColumn语句中的列引用也不正确

我将类型转换为float，它现在正在运行

def get_median(values,weights):
    return float(np.median(np.repeat(values,weights)))

wimedian = F.udf(get_median,DoubleType())
df = sc.parallelize([["001",69,8],["002",69,14],["003",28,3],["003",21,1]]).toDF(["category","col_A","col_B"])

myview = df.groupBy('category').agg(
    F.collect_list(F.col('col_A')),
    F.collect_list(F.col('col_B'))).withColumn('Weighted_median',wimedian(F.col("collect_list(col_A)"),F.col("collect_list(col_B)"))).show()