如何对列进行分箱并将空值保留在单独的组中

Question

我有一个包含连续变量的列，想将其装箱以进行绘图。但是，此列也包含空值。

我使用了下面的代码来装箱：

def a(b):
  if b<= 20: return "<= 20"
  elif b<= 40: return "20 < <= 40"
  elif b<= 45: return "40 < <= 45"
  else: return "> 45"
audf = udf(a, StringType())
data= data.withColumn("a_bucket", audf("b"))

我在 Python 3 上运行并向我抛出以下错误：

TypeError: '<=' not supported between instances of 'NoneType' and 'int'

我查阅了一些文档说 Python 3 不允许比较具有空值的数字。但是有没有办法让我将这些空值放入一个单独的组中，这样我就不会丢弃数据。它们不是坏数据。

谢谢。

Answer 1

您可以在没有 udf 的情况下执行此操作。这是重写代码的一种方法，并为 null 个值设置一个特殊的存储桶：

from pyspark.sql.functions import col, when

def a(b):
    return when(b.isNull(), "Other")\
        .when(b <= 20, "<= 20")\
        .when(b <= 40, "20 < <= 40")\
        .when(b <= 45, "40 < <= 45")\
        .otherwise("> 45")

data = data.withColumn("a_bucket", a(col("b")))

但是，更通用的解决方案将允许您传入存储桶列表并动态 return bin 输出（未经测试）：

from functools import reduce

def aNew(b, buckets):
    """assumes buckets are sorted"""
    if not buckets:
        raise ValueError("buckets can not be empty")
    return reduce(
        lambda w, i: w.when(
            b.between(buckets[i-1], buckets[i]), 
            "{low} < <= {high}".format(low=buckets[i-1], high=buckets[i]))
        ),
        range(1, len(buckets)),
        when(
            b.isNull(), 
            "Other"
        ).when(
            b <= buckets[0], 
            "<= {first}".format(first=buckets[0])
        )
    ).otherwise("> {last}".format(last=buckets[-1]))

data = data.withColumn("a_bucket", aNew(col("b"), buckets=[20, 40, 45]))

如何对列进行分箱并将空值保留在单独的组中

How to bin a column and keep null values in a separate group

python

pyspark

pyspark-sql

pyspark-dataframes