时间戳

首页 > 解决方案 > Scala spark，输入数据框，返回所有值等于 1 的列

scala - Scala spark，输入数据框，返回所有值等于 1 的列

问题描述

给定一个数据框，假设它包含 4 列和 3 行。我想编写一个函数来返回该列中所有值都等于 1 的列。

这是一个 Scala 代码。我想使用一些火花转换来转换或过滤数据框输入。这个过滤器应该在一个函数中实现。

case class Grade(c1: Integral, c2: Integral, c3: Integral, c4: Integral)
val example = Seq(
      Grade(1,3,1,1),
      Grade(1,1,null,1),
      Grade(1,10,2,1)
    )

    val dfInput = spark.createDataFrame(example)

在我调用函数 filterColumns()

val dfOutput = dfInput.filterColumns()

它应该返回 3 行 2 列数据框，其值全部为 1。

标签： scaladataframeapache-sparkfilterapache-spark-sql

解决方案

使用更具可读性的方法Dataset[Grade]

import org.apache.spark.sql.functions.col
import scala.collection.mutable
import org.apache.spark.sql.Column

val tmp = dfInput.map(grade => grade.dropWhenNotEqualsTo(1))
val rowsCount = dfInput.count()

val colsToRetain = mutable.Set[Column]()
 for (column <- tmp.columns) {
   val withoutNullsCount = tmp.select(column).na.drop().count()
   if (rowsCount == withoutNullsCount) colsToRetain += col(column)
}

dfInput.select(colsToRetain.toArray:_*).show()

+---+---+
| c4| c1|
+---+---+
|  1|  1|
|  1|  1|
|  1|  1|
+---+---+

和案例对象

case class Grade(c1: Integer, c2: Integer, c3: Integer, c4: Integer) {
  def dropWhenNotEqualsTo(n: Integer): Grade = {
    Grade(nullOrValue(c1, n), nullOrValue(c2, n), nullOrValue(c3, n), nullOrValue(c4, n))
  }
  def nullOrValue(c: Integer, n: Integer) = if (c == n) c else null
}
  1. grade.dropWhenNotEqualsTo(1)-> 返回一个新的成绩，其中不满足条件的值被替换为空
+---+----+----+---+
| c1|  c2|  c3| c4|
+---+----+----+---+
|  1|null|   1|  1|
|  1|   1|null|  1|
|  1|null|null|  1|
+---+----+----+---+

  1. (column <- tmp.columns) -> 遍历列

  2. tmp.select(column).na.drop()-> 删除带有空值的行，例如，c2这将返回

+---+
| c2|
+---+
|  1|
+---+
  1. if (rowsCount == withoutNullsCount) colsToRetain += col(column)-> 如果列包含空值，则删除它

推荐阅读