如何根据spark中的简单条件在DataFrame中创建一个额外的列

Question

我有一个数据框，我想基于一个简单的条件为它添加一个额外的列，它基本上表示另一列的值是否等于给定的字符串。我知道我可以创建一个UDF并注册然后使用它，但是我认为必须有一种更简单的方法。这是我将要做的事情的伪代码

df.withColumn("extra", if (col("a) == "str" 1 else 2))

Answer 1

你几乎就在那里：

scala> val df = Seq((1,2), (3,3), (4,5)).toDF("a", "b")
scala> df.show
+-+-+
|a|b|
+-+-+
|1|2|
|3|3|
|4|5|
+-+-+

scala> df.withColumn("New", when($"a" === $"b", "equal").otherwise("not")).show
+-+-+-----+
|a|b|  New|
+-+-+-----+
|1|2|  not|
|3|3|equal|
|4|5|  not|
+-+-+-----+

请注意，您需要为上述内容导入functions和implicits才能正常工作。