Scala udf 检查 df 列值是不是在列表中

Posted 2023-04-13

【中文标题】Scala udf 检查 df 列值是不是在列表中

【英文标题】：Scala udf check if df column value is in a listScala udf 检查 df 列值是否在列表中

【发布时间】：2019-04-15 12:36:58

【问题描述】：

我正在编写一个 scala UDF，它从列中获取值并检查列表中的特定值是否存在然后执行某些操作，如果存在其他值执行某些操作等等。例如：

val listOfValues = List("001", "002", "003", "004", "005")

if ($"column".isin(listOfValues: _*) || (logic 2) && (logic 3)) "value 1"
else if ($"column".isin(listOfValues: _*) || (logic 3) || (logic 4)) "value 2"
else if ($"column".isin(listOfValues: _*) && (logic 4) && (logic 5)) "value 3"
else "value 4"

现在的问题是$"column".isin(listOfValues: _*) 返回一列，但没有返回单个真/假值。我需要每行都有一个真/假值才能正确使用条件。

关于如何检查 df 列是否包含列表中的这些值作为真/假有什么建议吗？

【参考方案1】：

假设你有一个像下面这样的 DF：

+---+----+-----+
|ID |Type|Value|
+---+----+-----+
|ID1|001 |1    |
|ID1|002 |5    |
|ID2|A   |12   |
|ID3|A   |3    |
|ID3|B   |3    |
|ID3|002 |5    |
|ID4|A   |10   |
+---+----+-----+

您可以创建一个 UDF 来检查列值是否在列表中。例如：

val listOfValues = List("001", "002", "003", "004", "005")
def isInDef(p1: String): String = if (listOfValues.contains(p1) || (logic 2) && (logic 3)) "value 1"
                                  else if (listOfValues.contains(p1) || (logic 3) || (logic 4)) "value 2"
                                  else if (listOfValues.contains(p1) && (logic 4) && (logic 5)) "value 3"
                                  else "value 4"
val isIn = udf[String, String](isInDef)

然后您可以使用 UDF 从条件中创建一个具有正确值的新列：

df = df.withColumn("contain", isIn($"Type"))

+---+----+-----+-------+
|ID |Type|Value|contain|
+---+----+-----+-------+
|ID1|001 |1    |value 1|
|ID1|002 |5    |value 2|
|ID2|A   |12   |value 4|
|ID3|A   |3    |value 4|
|ID3|B   |3    |value 4|
|ID3|002 |5    |value 3|
|ID4|A   |10   |value 4|
+---+----+-----+-------+