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

Question

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

例如，在df的"name"列中查找值为"test"的字符串。

在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'))
)

这是一个昂贵的操作吗？如何最好地检查特定字符串的出现次数，然后执行相应的操作（求和、最大值、最小值等）？

谢谢。

- Lou_Ds

---

有任何答案对您有帮助吗？如果有，请接受它。 - Raphael Roth

Answer 1

你可以在Spark中使用groupBy+agg; 这里的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|
+---+--------+

- Psidom

Answer 2

如果您想要翻译您的SQL，您也可以在Spark中使用窗口函数：

def getCount(df: DataFrame): DataFrame = {
  import org.apache.spark.sql.expressions.Window

  df.withColumn("cnt_test",
      sum(when($"name" === "test", 1).otherwise(0)).over(Window.partitionBy($"id"))
    )
}

- Raphael Roth