如何在单独的组中存储列并保留空值

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'

我查阅了一些文档，其中说Python3不允许比较空值的数字。但是有没有一种方法可以让我把这些空值扔进一个单独的组中，这样我就不会扔掉数据了。它们不是坏数据

谢谢

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

---

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

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