Koalas新建中位数列报PandasNotImplementedError如何解决

报错原因 #

• 第一种写法存在两个问题：一是np.median要求第一个参数传入可迭代的数组对象，你直接传入x.col1、x.col2两个Koalas Series对象时，Numpy会尝试迭代Series，而Koalas为了避免分布式场景下全量数据拉取到单节点，默认禁止了Series的迭代操作，因此触发PandasNotImplementedError；二是np.median的第二个参数默认是axis，你直接传入x.col2会被识别为轴参数，本身逻辑就不正确。
• 第二种apply写法的问题：Koalas的行式apply（axis=1）默认开启严格模式，不支持直接调用Numpy方法处理单行字段值，同时也存在上述np.median参数传错的问题。

正确实现方案 #

方案1：两列场景高性能写法 #

两个数值的中位数等价于两个数的平均值，直接用向量化运算实现，完全兼容Koalas分布式执行逻辑，性能最优：

df = df.assign(newcol = (df.col1 + df.col2) / 2)

如果需要兼容空值，可以加上分支判断：

from databricks.koalas import when

df = df.assign(
    newcol = when(df.col1.isNotNull() & df.col2.isNotNull(), (df.col1 + df.col2)/2)
            .otherwise(None)
)

方案2：通用多列行中位数写法 #

如果后续需要扩展到更多列计算行中位数，可以拼接目标列后按行调用median方法：

df = df.assign(
    newcol = ks.concat([df.col1, df.col2], axis=1).median(axis=1)
)

方案3：修正后的apply写法 #

如果一定要使用apply实现，需要先开启Koalas的unsafe类型允许配置，同时修正np.median的参数：

ks.set_option("compute.default_index_type", "distributed")
ks.set_option("apply.allow_unsafe_types", True)

df['newcol'] = df.apply(lambda row: np.median([row.col1, row.col2]), axis=1)

内容的提问来源于stack exchange，提问作者Fluxy