Scala Spark - 计算Dataframe列中特定字符串的出现次数

Question

如何df使用分区的Spark计算列中String的出现次数id？

例如,找到a "test"列"name"中的值df

在SQL中将是:

 SELECT
    SUM(CASE WHEN name = 'test' THEN 1 else 0 END) over window AS cnt_test
  FROM
    mytable
 WINDOW window AS (PARTITION BY id)

我试过用了 map( v => match { case "test" -> 1.. })

和类似的东西:

def getCount(df: DataFrame): DataFrame = {
    val dfCnt = df.agg(
          .withColumn("cnt_test", 
            count(col("name")==lit('test'))
)

这是一项昂贵的操作吗？什么是检查特定字符串出现然后执行操作的最佳方法(sum, max, min, etc)？

谢谢

Answer 1

你可以在火花中使用groupBy+ agg; 这里when($"name" == "test", 1)将name列转换为1if name == 'test',null否则,并count给出非空值的计数:

df.groupBy("id").agg(count(when($"name" === "test", 1)).as("cnt_test"))

示例:

val df = Seq(("a", "joe"), ("b", "test"), ("b", "john")).toDF("id", "name")
df.groupBy("id").agg(count(when($"name" === "test", 1)).as("cnt_test")).show
+---+--------+
| id|cnt_test|
+---+--------+
|  b|       1|
|  a|       0|
+---+--------+

或者类似于你的SQL查询:

df.groupBy("id").agg(sum(when($"name" === "test", 1).otherwise(0)).as("cnt_test")).show
+---+--------+
| id|cnt_test|
+---+--------+
|  b|       1|
|  a|       0|
+---+--------+