scala - 在 Spark 中使用 "when"函数填充数据框

Question

我有一个看起来像这样的数据框

 df1:
 image-id  colorList
 -------------------------
 id1       [Red,Blue]  
 id2       [White,Grey]      

现在我想使用如下所示的 df1 创建一个新的 Dataframe

 df2:
 image-id  isRed  isBlue  isWhite  isGrey
 ----------------------------------------
 id1       1      1       0        0
 id2       0      0       1        1 

我正在尝试使用以下代码，但由于类型不匹配，它无法正常工作

val df2 = df1.withColumn("image-id",$"image-id")
.withColumn("isRed", when($"colorList" contains "Red",1).otherwise(0))

我试过了

val df2 = df1.withColumn("image-id",$"image-id")
.withColumn("isRed", when($"colorList" contains Seq("Red"),1).otherwise(0))

我收到这条消息

Unsupported literal type class scala.collection.immutable.$colon$colon List(Red)

我可以选择在 df1 上展开 colorList，但这会使我的表格过于复杂。

Answer 1

你要找的是array_contains函数，而不是Column.contains(后者只适用于StringType列和检查字符串值是否包含子字符串):

df1.withColumn("isRed", when(array_contains($"colorList", "Red"),1).otherwise(0))