Spark udf 从现有列创建新列，以及使用 group by答案

【问题标题】：Spark udf to create new columns from existing column, along with using group bySpark udf 从现有列创建新列，以及使用 group by
【发布时间】：2017-12-17 09:19:19
【问题描述】：

我有一个包含以下列的数据框：group_id、gender 和 height。

group_id 和 height 是 Int。性别是字符串。

group_id| gender|height
 1      |  F    |  52
 1      |  F    |  53
 1      |  F    |  58 
 1      |  M    |  55
 1      |  M    |  59
 2      |  F    |  50
 2      |  M    |  60
 2      |  M    |  61
 2      |  M    |  64

我想按 group_id、gender 和 height_range 分组高度范围可以是任何东西，但都是预定义的——例如 48-50、51-58、58-64 等。为简单起见，我考虑以 5 英寸为增量：50-54、55-59、60-64。

预期输出：

group_id | gender | height_low | height_high | count
   1     | F      |    50      |    54       |   2
   1     | F      |    55      |    59       |   1
   1     | M      |    55      |    59       |   2
   2     | F      |    50      |    54       |   1
   2     | M      |    60      |    64       |   3

我尝试使用 sum(when)。这对获得低值和高值没有多大帮助。我能想到的其他方法是使用 hive udf 2 次，一次在 when 子句中使用各种条件填充 height_low，另一次填充 height_high。我想知道是否有办法同时填充这两者，因为我需要执行相同的检查。任何帮助表示赞赏。

【问题讨论】：

标签： scala apache-spark dataframe apache-spark-sql user-defined-functions

【解决方案1】：

您可以使用height / 5 的结果的floor（向下舍入）作为分组键（连同其他列），然后根据需要计算height_low 和height_high：

import org.apache.spark.sql.functions._
import spark.implicits._

val result = df.groupBy($"group_id", $"gender", floor($"height" / 5) * 5 as "height_low")
  .count()
  .withColumn("height_high", $"height_low" + 4)

编辑：更一般的情况，范围不一定具有恒定大小，确实可以使用返回表示范围的元组的 UDF 来解决：

// For a set of ranges, given as an ordered list of range lower bounds:
val rangeLowerBoundaries = List(50, 55, 60, 65) // ordered!

// create UDF for finding the range, returns a tuple of (start, end), 
// uses 0 and MaxValue as edges of first and last range, respectively:
val findMatchingRange = udf[(Int, Int), Int] { height => (
  rangeLowerBoundaries.filter(_ <= height).lastOption.getOrElse(0),
  rangeLowerBoundaries.find(_ > height).getOrElse(Int.MaxValue) - 1
)}

// group by UDF and then select and rename tuple's elements:
val result = df.groupBy($"group_id", $"gender", findMatchingRange($"height") as "range")
  .count()
  .select($"group_id", $"gender", $"range._1" as "height_low", $"range._2" as "height_high", $"count")
  .show(false)

result.show(false)
// +--------+------+----------+-----------+-----+
// |group_id|gender|height_low|height_high|count|
// +--------+------+----------+-----------+-----+
// |1       |F     |55        |59         |1    |
// |2       |M     |60        |64         |3    |
// |2       |F     |50        |54         |1    |
// |1       |M     |55        |59         |2    |
// |1       |F     |50        |54         |2    |
// +--------+------+----------+-----------+-----+

【讨论】：