How to check all values in columns efficiently using Spark? - scala

Question

I'm wondering how to make dynamic filter given unknown columns in Spark.
For example, the dataframe is like below:
+-------+-------+-------+-------+-------+-------+
| colA | colB | colC | colD | colE | colF |
+-------+-------+-------+-------+-------+-------+
| Red | Red | Red | Red | Red | Red |
| Red | Red | Red | Red | Red | Red |
| Red | Blue | Red | Red | Red | Red |
| Red | Red | Red | Red | Red | Red |
| Red | Red | Red | Red | Blue | Red |
| Red | Red | White | Red | Red | Red |
+-------+-------+-------+-------+-------+-------+
The columns can only be known at runtime, meaning it can have colG, H ..
I need to check if the whole column's value is Red, then get a count, in above case is 3 as colA, colD and ColF columns are all Red.
What I am doing is something like below, and it is SLOW..
val allColumns = df.columns
df.foldLeft(allColumns) {
(df, column) =>
val tmpDf = df.filter(df(column) === "Red")
if (tmpDf.rdd.isEmpty) {
count += 1
}
df
}
I am wondering if there is a better way. Many thanks!

Answer 1

you got N RDD scans where N is number of columns. You can scan all of them at once and reducing in parallel. For example this way:
df.reduce((a, r) => Row.fromSeq(a.toSeq.zip(r.toSeq)
.map { case (a, r) =>
if (a == "Red" && r == "Red") "Red" else "Not"
}
))
res11: org.apache.spark.sql.Row = [Red,Not,Not]
This code will do one RDD scan, and then iterated Row columns inside reduce. Row.toSeq get Seq from Row. fromSeq restore Row to return the same object.
Edit: for count just add: .toSeq.filter(_ == "Red").size

Answer 2

Why not simply do df.filter + df.count using only DataFrame API?
val filter_expr = df.columns.map(c => col(c) === lit("Red")).reduce(_ and _)
val count = df.filter(filter_expr).count
//count: Long = 3