如何对一列进行装箱并在一个单独的组中保留空值

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")))

---

但是，更通用的解决方案将允许您传递存储桶列表并动态返回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]))