Apache spark 将逻辑/函数应用于Spark数据帧列的“子集”

我想通过传递一个col列表来处理以下内容，而不是传递除ID之外的all。无论我做什么，col、Row等都会有一些错误。对于所有列，选择一个子集可以很好地工作，但对于函数方面的子集则不行

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

---

寻求指导。提前谢谢。

上面的代码对我很有用。我不确定你在尝试什么，那是行不通的。您想只在某些行中检查一个吗？在函数中

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)
}

---

上面的代码对我有用。你能提供一个你遇到麻烦的代码样本吗？是的，我知道它是有效的。但我只想说a列。这个答案中的第一个例子是a列。看起来不错，我想知道这将如何与列一起工作？即每行我对我的问题有什么想法吗？

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|
+---+---+---+----+