In Spark, how to do One Hot Encoding for top N frequent values only?

问题

Let, in my dataframe df, I have a column my_category in which I have different values, and I can view the value counts using:

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

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

Now, I'd like to apply One Hot Encoding (OHE) on this column, but for the top N frequent values only (say N = 3), and put all the rest infrequent values in a dummy column (say, "default"). E.g., the output should be something like:

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

How do I do this in Spark/Scala?

Note: I know how to do this in Python, e.g., by first building a frequent-based dictionary for each unique value, and then create the OHE vector by examining the values one by one, putting the infrequent ones in a "default" column.

回答1:

A custom function could be written as, to apply One Hot Encoding (OHE) on a particular column, for the top N frequent values only (say N = 3).

It is relatively similar to Python, 1) Building a top n frequent-based Dataframe/Dictionary. 2) Pivot the top n frequent Dataframe i.e create the OHE vector. 3) Left join the given Dataframe and pivoted Dataframe, replace null with 0 i.e the default OHE vector.

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

来源：https://stackoverflow.com/questions/60236153/in-spark-how-to-do-one-hot-encoding-for-top-n-frequent-values-only

标签

scala

apache-spark

apache-spark-sql

one-hot-encoding