在 Spark 数据框中，如果另一列中的值在广播变量数组中，则将值添加到新列答案

【问题标题】：In Spark data frame, add value to new column if value in another column is in broadcast variable array在 Spark 数据框中，如果另一列中的值在广播变量数组中，则将值添加到新列
【发布时间】：2017-12-11 04:33:26
【问题描述】：

我有一个数组作为广播变量，它包含整数：

broadcast_array.value
Array(72159153, 72159163, 72159202, 72159203, 72159238, 72159398, 72159447, 72159448, 72159455, 72159492...

我在数据集中有一列（调用是 col_id，其中包含可能在 broadcast_array 中的 IntegerType 值，但它们可能不在。

我只是想创建一个新列（称为new_col）来检查每行的col_id 值是否在broadcast_array 中。如果是这样，新的列值应该是Available，否则它可以是null

所以我有类似的东西：

val my_new_df = df.withColumn("new_col", when(broadcast_array.value.contains($"col_id"), "Available"))

但我不断收到此错误：

Name: Unknown Error
Message: <console>:45: error: type mismatch;
 found   : Boolean
 required: org.apache.spark.sql.Column
   val my_new_df = df.withColumn("new_col", when(broadcast_array.value.contains($"col_id"), "Available"))
                                                                                           ^
StackTrace:

最让我困惑的是，我认为when 语句需要一个输出一些布尔值的条件，但这里它说它需要一个列。

我应该如何根据现有列中的值是否可以在预定义的数组中找到来向新列添加值？

【问题讨论】：

标签： scala apache-spark merge apache-spark-sql

【解决方案1】：

如果你看when函数的api

def when（条件：org.apache.spark.sql.Column，值：scala.Any）：org.apache.spark.sql.Column

很明显所需的条件是列而不是布尔值。

所以你可以做复杂的lit 组合来将你的boolean 转换为column

import org.apache.spark.sql.functions._
df.withColumn("new_col", when(lit(broadcast_array.value.mkString(",")).contains($"col_id"), lit("Available"))).show(false)

或

您可以通过编写一个简单的udf 函数来实现您的目标

import org.apache.spark.sql.functions._
val broadcastContains = udf((id: Int) => broadcast_array.value.contains(id))

只需将函数调用为

df.withColumn("new_col", when(broadcastContains($"col_id"), lit("Available"))).show(false)

【讨论】：

【解决方案2】：

我在spark-daria 中添加了一个broadcastArrayContains 函数，使Ramesh 的解决方案更易于重用/访问。

def broadcastArrayContains[T](col: Column, broadcastedArray: Broadcast[Array[T]]) = {
  when(col.isNull, null)
    .when(lit(broadcastedArray.value.mkString(",")).contains(col), lit(true))
    .otherwise(lit(false))
}

假设您有以下 DataFrame (df)：

+----+
| num|
+----+
| 123|
|  hi|
|null| 
+----+

您可以按如下方式识别广播数组中的所有值：

val specialNumbers = spark.sparkContext.broadcast(Array("123", "456"))

df.withColumn(
  "is_special_number",
  functions.broadcastArrayContains[String](col("num"), specialNumbers)
)

【讨论】：