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

Question

如何计算 df 中字符串的出现次数使用 Spark 分区的列 id ?

例如找到值 "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在 Spark 中；这里when($"name" == "test", 1)转换 name列到 1如果 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|
+---+--------+