Применить логику / функцию к «подмножеству» столбцов Spark Dataframe

Question

Я хотел бы адаптировать следующее, передавая список столбцов для обработки, в отличие от всех, кроме ID.Независимо от того, что я делаю, есть ошибки на столбцах, строках и т. Д. Со всеми столбцами и выбором подмножества все в порядке, но не по подмножеству с точки зрения функции.

val df = sc.parallelize(Seq(
   ("r1", 1, 1),
   ("r2", 6, 4),
   ("r3", 4, 1),
   ("r4", 1, 2)
   )).toDF("ID", "a", "b")

val ones = df.columns.drop(1).map(x => when(col(x) === 1, 1).otherwise(0)).reduce(_ + _)
df.withColumn("ones", ones).show

Требуется указание.Заранее спасибо.

Answer 1

Код выше работает для меня. Я не уверен, что вы пытаетесь это не работает. Вы хотите проверить только некоторые строки? В функции?

def sumOnes(df: DataFrame, columnNames: Seq[String]): DataFrame = {
  val ones = columnNames.map(x => when(col(x) === 1, 1).otherwise(0)).reduce(_ + _)
  df.withColumn("ones", ones)
}

sumOnes(df, Seq("a")).show

+---+---+---+----+
| ID|  a|  b|ones|
+---+---+---+----+
| r1|  1|  1|   1|
| r2|  6|  4|   0|
| r3|  4|  1|   0|
| r4|  1|  2|   1|
+---+---+---+----+

sumOnes(df, Seq("b")).show

+---+---+---+----+
| ID|  a|  b|ones|
+---+---+---+----+
| r1|  1|  1|   1|
| r2|  6|  4|   0|
| r3|  4|  1|   1|
| r4|  1|  2|   0|
+---+---+---+----+

sumOnes(df, Seq("a", "b")).show

+---+---+---+----+
| ID|  a|  b|ones|
+---+---+---+----+
| r1|  1|  1|   2|
| r2|  6|  4|   0|
| r3|  4|  1|   1|
| r4|  1|  2|   1|
+---+---+---+----+