Question

我想要一个函数，一个列和一个bucket作为参数列表，并返回相应的存储桶。我想用Spark API解决这个问题，不想使用UDF。

假设我们从这个DataFrame（df）开始：

+--------+
|some_num|
+--------+
|       3|
|      24|
|      45|
|    null|
+--------+

这是函数的期望行为：

df.withColumn(
  "bucket",
  bucketFinder(
    col("some_num"),
    Array(
      (0, 10),
      (10, 20),
      (20, 30),
      (30, 70)
    )
  )
).show()

+--------+------+
|some_num|bucket|
+--------+------+
|       3|  0-10|
|      24| 20-30|
|      45| 30-70|
|    null|  null|
+--------+------+

以下是我尝试过的不起作用的代码：

def bucketFinder(col: Column, buckets: Array[(Any, Any)]): Column = {

  buckets.foreach { res: (Any, Any) =>
    when(col.between(res._1, res._2), lit(s"$res._1 - $res._2"))
  }

}

使用UDF编写此代码非常容易，但仅限于Spark API时很难。

Answer 1

您可以将列除以10，然后该列的floor和ceil应该成为您需要的存储桶：

val bucket_size = 10    
val floor_col = floor(df("some_num") / bucket_size) * bucket_size

df.withColumn("bucket", concat_ws("-", floor_col, floor_col + bucket_size)).show
+--------+------+
|some_num|bucket|
+--------+------+
|       3|  0-10|
|      24| 20-30|

铲斗尺寸为5：

val bucket_size1 = 5
val floor_col = floor(df("some_num") / bucket_size1) * bucket_size1

df.withColumn("bucket", concat_ws("-", floor_col, floor_col + bucket_size1)).show
+--------+------+
|some_num|bucket|
+--------+------+
|       3|   0-5|
|      24| 20-25|

Answer 2

这是一个纯粹的Spark解决方案：

def bucketFinder(col: Column, buckets: Array[(Any, Any)]): Column = {

  val b = buckets.map { res: (Any, Any) =>
    when(col.between(res._1, res._2), lit(s"${res._1}-${res._2}"))
  }

  coalesce(b: _*)

}

我会稍微留下这个问题，看看别人是否有更优雅的解决方案。

Spark代码用于对存储桶中的整数进行分组

2 个答案: