Scala / Spark: Как выполнить фильтрацию и изменить значение столбца для списка столбцов?

Question

1) У меня есть DF с несколькими столбцами, но наиболее важными являются "col1", "col2" .... "col10" (10 столбцов). Меня интересуют только строки, в которых хотя бы ОДИН из этих столбцов имеет значение> 1. Теперь у меня есть:

DF.filter(col(col1) > 1 or col(col2) > 1 or col(col3) > 1 or col(col4) > 1 or col(col5) > 1   
          or col(col6) > 1 or col(col7) > 1 or col(col8) > 1 or col(col9) > 1 or col(col10) > 1)

2) Как и в # 1, у меня есть еще один DF, где мне нужно вычесть на 1 значение списка столбцов

 DF
  .withColumn("col1", col("col1") - 1)
  .withColumn("col2", col("col2") - 1)
  .withColumn("col3", col("col3") - 1)
  .withColumn("col4", col("col4") - 1)
   // and so on ..... 

Оба они работают, но есть ли более элегантный способ работы с несколькими столбцами?

Answer 1

Можно создать список интересующих столбцов и использовать reduce для сборки повторяющегося условия or в предложении where:

val df = Seq(
  (101, 0, -1, -2, -3, -4),
  (102, -1, 0, 1, 2, 3)
).toDF("id", "c1", "c2", "c3", "c4", "c5")

val colList = df.columns.filter(_.startsWith("c"))

df.where(colList.map(col(_) > 1).reduce(_ || _)).show
// +---+---+---+---+---+---+
// | id| c1| c2| c3| c4| c5|
// +---+---+---+---+---+---+
// |102| -1|  0|  1|  2|  3|
// +---+---+---+---+---+---+

Для повторяющихся преобразований withColumn рассмотрите используя foldLeft:

colList.foldLeft(df)((acc, c) => acc.withColumn(c, col(c) - 1)).show
// +---+---+---+---+---+---+
// | id| c1| c2| c3| c4| c5|
// +---+---+---+---+---+---+
// |101| -1| -2| -3| -4| -5|
// |102| -2| -1|  0|  1|  2|
// +---+---+---+---+---+---+

Однако вместо использования повторяющихся withColumn было бы более эффективно выполнять вычитание по столбцам в качестве элементов столбца ArrayType в одном преобразовании:

df.
  withColumn("arr", array(colList.map(col(_) - 1): _*)).
  select($"id" +: (0 until colList.size).map(i => $"arr"(i).as(colList(i))): _*)