首页 > 解决方案 > 将逻辑/函数应用于 Spark Dataframe 列的“子集”

问题描述

我想通过传递一个 cols 列表来调整以下内容,而不是所有,除了 ID。无论我做什么,cols、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

寻求指导。提前致谢。

标签: apache-spark

解决方案


上面的代码对我有用。我不确定你正在尝试什么不起作用。你想检查一些行中的那些吗?在函数中?

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

推荐阅读