Сумма нескольких столбцов в зависимости от условия

Question

Я хочу суммировать массив столбцов в строке на основе условия.Я нашел способ сделать это, но это не очень хорошая идея, когда у меня есть более 20 столбцов для суммирования, так как это сгенерирует дополнительный столбец для каждого.

Требуемый результат: Сумма значений длявсе столбцы, заканчивающиеся на «_val», где значение равно 0 или 1 (или <2, я просто хочу исключить значение 3 прямо сейчас) </p>

 val df1 = Seq(
  ("id1", 1, 0, 3),
  ("id2", 0, 0, 3),
  ("id3", 1, 1, 3))
  .toDF("id", "bla_val", "blub_val", "bli_val")

Мое решение с требуемым результатом в сумме столбцов

val channelNames = df1.schema.fieldNames.filter(_.endsWith("_val"))
val ch = channelNames.map(x => col(x+"_redval"))

val df2 = df1.select(col("*") +: (channelNames.map(c =>
  when(col(c) === 1, lit(1))
    .otherwise(lit(0)).as(c+"_redval"))): _*) 

val df3 = df2.withColumn("sum", ch.reduce(_+_))
df3.show()

Пример вывода:

+---+-------+--------+-------+--------------+---------------+--------------+---+ | id|bla_val|blub_val|bli_val|bla_val_redval|blub_val_redval|bli_val_redval|sum| +---+-------+--------+-------+--------------+---------------+--------------+---+ |id1| 1| 0| 3| 1| 0| 0| 1| |id2| 0| 0| 3| 0| 0| 0| 0| |id3| 1| 1| 3| 1| 1| 0| 2| +---+-------+--------+-------+--------------+---------------+--------------+---+

Answer 1

Вы можете объединить выражение, используя операцию Redu (). Проверьте это:

val df1 = Seq(
  ("id1", 1, 0, 3),
  ("id2", 0, 0, 3),
  ("id3", 1, 1, 3))
  .toDF("id", "bla_val", "blub_val", "bli_val")

val newcols= df1.columns.filter(_.endsWith("_val")).map( x=> when(col(x)===1, lit(1)).otherwise(lit(0))).reduce(_+_)
df1.withColumn("redval_count",newcols).show(false)

Выход:

+---+-------+--------+-------+------------+
|id |bla_val|blub_val|bli_val|redval_count|
+---+-------+--------+-------+------------+
|id1|1      |0       |3      |1           |
|id2|0      |0       |3      |0           |
|id3|1      |1       |3      |2           |
+---+-------+--------+-------+------------+

Answer 2

def sumNot3(s: Seq[Int]): Int = {
  s.filter(_ != 3).sum
}

val sumNot3Udf = udf(sumNot3(_: Seq[Int]))

val channelNameCols = df1.schema.fieldNames.filter(_.endsWith("_val")).map(c => col(c))

df1.select(sumNot3Udf(array(channelNameCols: _*)).as("sum"))