Spark: одновременный запрос нескольких условий

Question

Итак, Dataframe.where можно использовать для фильтрации фрейма данных для строк, заданных выражением, например:

df.where(($"group_id" == 1234) || ($"group_id" == 4434))

или для получения более сложного примера

df.where(($"group_id" == 1234 && $"country" === "PL") || ($"group_id" == 4434 $"country" === "FR"))

Меня интересует, могу ли я каким-то образом представить эти условия в виде списка, поэтому предположим, что у меня есть список group_id, List((1234, "PL"), (4434, "FR"), ....), тогда я хотел бы эффективно отфильтровать фрейм данных.

Answer 1

Вы можете попробовать что-то вроде этого:

val df = Seq((1,"a"),(2,"b"),(3,"c")).toDF()
df.show()

+---+---+
| _1| _2|
+---+---+
|  1|  a|
|  2|  b|
|  3|  c|
+---+---+

val list = List((1,"a"),(3,"c"))
val cols = List("_1","_2")

def mkCol(values: List[(Any,Any)], columns: List[String]) = list.map(s=>col(columns.apply(0)) === s._1 && col(columns.apply(1)) === s._2)
                                                                .reduce((a,b)=>a.or(b))

val col = mkCol(list,cols)
col.explain(true)

((('_1 = 1) && ('_2 = a)) || (('_1 = 3) && ('_2 = c)))

df.where(col).show()

+---+---+
| _1| _2|
+---+---+
|  1|  a|
|  3|  c|
+---+---+