aName
aName

Reputation: 3043

spark aggregation count on condition

I'm trying to group a data frame, then when aggregating rows, with a count, I want to apply a condition on rows before counting.
here is an example :

val test=Seq(("A","X"),("A","X"),("B","O"),("B","O"),("c","O"),("c","X"),("d","X"),("d","O")).toDF
test.show
+---+---+
| _1| _2|
+---+---+
|  A|  X|
|  A|  X|
|  B|  O|
|  B|  O|
|  c|  O|
|  c|  X|
|  d|  X|
|  d|  O|
+---+---+

in this example I want to group by column _1 on count on column _2 when the value ='X'
here is the expected result :

+---+-----------+
| _1| count(_2) |
+---+-----------+
|  A|  2        |
|  B|  0        |
|  c|  1        |
|  d|  1        |
+---+-----------+

Upvotes: 3

Views: 10300

Answers (3)

Dorren Chen
Dorren Chen

Reputation: 316

import spark.implicits._

val test=Seq(("A","X"),("A","X"),("B","O"),("B","O"),("c","O"),("c","X"),("d","X"),("d","O")).toDF

test.groupBy("_1").agg(count(when($"_2"==="X", 1)).as("count")).orderBy("_1").show
+---+-----+
| _1|count|
+---+-----+
|  A|    2|
|  B|    0|
|  c|    1|
|  d|    1|
+---+-----+

Upvotes: 4

igr
igr

Reputation: 3499

As alternative, in Scala, it can be:

val counter1 = test.select( col("_1"), 
      when(col("_2") === lit("X"), lit(1)).otherwise(lit(0)).as("_2"))

val agg1 = counter1.groupBy("_1").agg(sum("_2")).orderBy("_1")

agg1.show

gives result:

+---+-------+
| _1|sum(_2)|
+---+-------+
|  A|      2|
|  B|      0|
|  c|      1|
|  d|      1|
+---+-------+

Upvotes: -1

Vamsi Prabhala
Vamsi Prabhala

Reputation: 49260

Use when to get this aggregation. PySpark solution shown here.

from pyspark.sql.functions import when,count
test.groupBy(col("col_1")).agg(count(when(col("col_2") == 'X',1))).show()

Upvotes: 5

Related Questions