如何在 if-else 条件下的列中使用 Spark 值 - Scala

Posted 2023-04-15

【中文标题】如何在 if-else 条件下的列中使用 Spark 值 - Scala

【英文标题】：How to use Spark value in a column in if-else conditions - Scala

【发布时间】：2021-03-26 02:00:03

【问题描述】：

我有以下数据框，我希望将列中的值与字符串进行比较，并将其用作 if-else 语句中的条件。

数据框：

id   type       name
1    fruit     apple
2    toy       football
3   fruit      orange 

我希望实现的是：

if(df("type") == "fruit")
    //do something to the df
else if ( df("type") == "toy")
    //something else to the df

我尝试使用val type= df.select("type").collectAsList().getString(0)，但这不正确。有人可以帮忙吗？非常感谢。我不认为这与这个问题重复，因为我不想添加新列。 Spark: Add column to dataframe conditionally 我不想使用withColumn

【问题讨论】：

df("type") 是一列...您无法将其与字符串进行比较。例如如果你把它比作“水果”，它应该是真的还是假的？第 1 行和第 3 行返回 true，而第 2 行返回 false。

问题是你想用 if else 语句中的 df 做什么。

您在 df 的 if..else 语句中尝试做的任何事情都发生在整个 df 上，而不是只有水果和玩具的行，除非您将这些行的 df 过滤为新的 df .

【参考方案1】：

Dataframe 的隐式类应该可以解决问题，下面的代码（忽略我的疯狂导入?）

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

implicit class typeTransform(df: DataFrame) 
    def tranformByType(typeCol: Column) = 
      if(df.filter(typeCol === lit("fruit")).count > 0) df // do something to df 
      //........ more if Statements
      else df// do something with df 
    

用法可以是这样的

val someDf = Seq(("fruit", "1"), ("toy", "2")).toDF("type", "id").tranformByType(col("type"))