Фильтр Dataframe по столбцам, определенным в списке

Question

У меня есть фрейм данных df

+----------+----+----+----+---+---+----+---+---+-------+-------+
|      WEEK|DIM1|DIM2|  T1| T2| T3|  T1| T2| T3|T1_diff|T2_diff|
+----------+----+----+----+---+---+----+---+---+-------+-------+
|2016-04-02|  14|NULL|9874|880| 23|9879|820| 45|     -5|     60|
|2016-04-30|  14|  FR|9875| 13| 34|9785|  9| 67|     90|      4|
+----------+----+----+----+---+---+----+---+---+-------+-------+

Я хочу сделать две вещи на этом фрейме данных:

1. Выбрать только WEEK, DIM1, DIM2, T1_diff, T2_diff
2. Фильтр T1_diff или T2_diff> 3

В настоящее время я делаю это так -

val selectColumns = Seq("WEEK", "DIM1", "DIM2","T1_diff","T2_diff")
df.select(selectColumns.head, selectColumns.tail: _*).filter($"T1_diff" > 3 or $"T2_diff" > 3).show()

У меня есть сценарий использования, в котором мои targetColumns определены так-

val targetColumns = Seq("T1_diff", "T2_diff")

Мне нужно использовать приведенную выше последовательность, чтобы применить ее в фильтре.Это последовательно, потому что в список targetColumns можно добавить больше столбцов.Я пытался что-то вроде этого -

df.filter(r => !targetColumns.map(x => col(x) > 3).isEmpty).show()

Это не похоже на работу.Кто-нибудь может сказать мне, как лучше всего это сделать?

Answer 1

Вы можете создать строку, используя targetColumns List, а затем передать эту строку в функцию where.

val targetColumns = List("T1_diff", "T2_diff")
val selectColumns = Seq("WEEK", "DIM1", "DIM2", "T1_diff", "T2_diff")

//create the where condition to filter the columns
val condition = targetColumns.map(c => s"$c>3").mkString(" OR ")

//select the columns and apply filter using where function.
df.select(selectColumns.head, selectColumns.tail: _*).where(condition).show(false)

Answer 2

Вы можете просто сделать следующее, сделав string query

val targetColumns = Seq("T1_diff", "T2_diff")

df.filter(targetColumns.map(x => s"$x > 3").mkString(" or ")).show()

, и вы можете добавить столько столбцов в targetColumns, сколько захотите

Answer 3

Вы можете использовать reduce в последовательности целевых столбцов после того, как вы отобразили каждый из них в условие (col(name) > 3), используя or, чтобы «объединить» их в одно условие:

import org.apache.spark.sql.functions._

val selectColumns = Seq("id", "type", "DIM2","T1_diff","T2_diff")
val targetColumns = Seq("T1_diff", "T2_diff")

df.select(selectColumns.head, selectColumns.tail: _*)
  .filter(targetColumns.map(name => col(name) > 3).reduce(_ or _))
  .show()