在 Spark 中，如何只对前 N 个频繁值进行一次热编码？

Question

让，在我的数据框 df 中，我有一个列 my_category，其中我有不同的值，我可以使用以下方法查看值计数：

df.groupBy("my_category").count().show()

value   count
a    197
b    166
c    210
d      5
e      2
f      9
g      3

现在，我想在此列上应用 One Hot Encoding (OHE)，但仅针对顶部的 N 频繁值（例如 N = 3），并将所有其余不频繁值放在一个虚拟列中（说，“默认”）。例如，输出应该是这样的：

a  b  c  default
0  0  1  0
1  0  0  0
0  1  0  0
1  0  0  0
...
0  0  0  1
0  0  0  1
...

如何在 Spark/Scala 中执行此操作？

注意： 我知道如何在 Python 中做到这一点，例如，首先为每个唯一值构建一个基于频繁的字典，然后通过逐个检查值来创建 OHE 向量，将不频繁的值放在“默认”列中。

Answer 1

自定义函数可以写成，在特定列上应用一个热编码 (OHE)，仅用于前 N 个频繁值（比如 N = 3）。

它与Python比较相似，1）构建一个基于top n频繁的Dataframe/Dictionary。 2）旋转前 n 个频繁的 Dataframe，即创建 OHE 向量。 3）左连接给定的Dataframe和pivot Dataframe，用0替换null，即默认的OHE向量。

import org.apache.spark.sql.DataFrame
import org.apache.spark.sql.functions.{col, lit, when}
import org.apache.spark.sql.Column

import spark.implicits._
val df = spark
  .sparkContext
  .parallelize(Seq("a", "b", "c", "a", "b", "c", "d", "e", "a", "b", "f", "a", "g", "a", "b", "c", "a", "d", "e", "f", "a", "b", "g", "b", "c", "f", "a", "b", "c"))
  .toDF("value")

val oheEncodedDF = oheEncoding(df, "value", 3)


def oheEncoding(df: DataFrame, colName: String, n: Int): DataFrame = {
  df.createOrReplaceTempView("data")
  val topNDF = spark.sql(s"select $colName, count(*) as count from data group by $colName order by count desc limit $n")

  val pivotTopNDF = topNDF
    .groupBy(colName)
    .pivot(colName)
    .count()
    .withColumn("default", lit(1))

  val joinedTopNDF = df.join(pivotTopNDF, Seq(colName), "left").drop(colName)

  val oheEncodedDF = joinedTopNDF
    .na.fill(0, joinedTopNDF.columns)
    .withColumn("default", flip(col("default")))

   oheEncodedDF
}

def flip(col: Column): Column = when(col === 1, lit(0)).otherwise(lit(1))