内置函数（数组，何时点亮）

Question

我有一个包含两列的数据框，

+---+-------+
| id|  fruit|
+---+-------+
|  0|  apple|
|  1| banana|
|  2|coconut|
|  1| banana|
|  2|coconut|
+---+-------+

我还有一个包含所有项目的通用列表，

fruitList: Seq[String] = WrappedArray(apple, coconut, banana)

现在我想在数据框中创建一个新的列，该列的数组为1、0，其中1表示该项目存在，而0表示该行不存在该项目。

所需的输出

    +---+-----------+
    | id|  fruitlist|
    +---+-----------+
    |  0|  [1,0,0]  |
    |  1| [0,1,0]   |
    |  2|[0,0,1]    |
    |  1| [0,1,0]   |
    |  2|[0,0,1]    |
    +---+-----------+

这是我尝试过的事情，

import org.apache.spark.ml.feature.{OneHotEncoder, StringIndexer}

val df = spark.createDataFrame(Seq(
  (0, "apple"),
  (1, "banana"),
  (2, "coconut"),
  (1, "banana"),
  (2, "coconut")
)).toDF("id", "fruit")

df.show
import org.apache.spark.sql.functions._
val fruitList = df.select(collect_set("fruit")).first().getAs[Seq[String]](0)
print(fruitList)

我试图用OneHotEncoder解决这个问题，但是转换为密集向量后，结果却是这样，这不是我所需要的。

    +---+-------+----------+-------------+---------+
| id|  fruit|fruitIndex|     fruitVec|       vd|
+---+-------+----------+-------------+---------+
|  0|  apple|       2.0|    (2,[],[])|[0.0,0.0]|
|  1| banana|       1.0|(2,[1],[1.0])|[0.0,1.0]|
|  2|coconut|       0.0|(2,[0],[1.0])|[1.0,0.0]|
|  1| banana|       1.0|(2,[1],[1.0])|[0.0,1.0]|
|  2|coconut|       0.0|(2,[0],[1.0])|[1.0,0.0]|
+---+-------+----------+-------------+---------+

Answer 1

如果您有一个收藏夹

val fruitList: Seq[String] = Array("apple", "coconut", "banana")

然后您可以使用内置函数或 udf函数

内置函数（数组，何时点亮）

import org.apache.spark.sql.functions._
df.withColumn("fruitList", array(fruitList.map(x => when(lit(x) === col("fruit"),1).otherwise(0)): _*)).show(false)

udf函数

import org.apache.spark.sql.functions._
def containedUdf = udf((fruit: String) => fruitList.map(x => if(x == fruit) 1 else 0))

df.withColumn("fruitList", containedUdf(col("fruit"))).show(false)

应该给您

+---+-------+---------+
|id |fruit  |fruitList|
+---+-------+---------+
|0  |apple  |[1, 0, 0]|
|1  |banana |[0, 0, 1]|
|2  |coconut|[0, 1, 0]|
|1  |banana |[0, 0, 1]|
|2  |coconut|[0, 1, 0]|
+---+-------+---------+

udf函数简单易懂，可以处理原始数据类型，但是如果可以使用经过优化的快速内置函数来完成相同任务，则应避免使用

我希望答案会有所帮助

由OneHotEncoder

1 个答案:

内置函数（数组，何时点亮）

udf函数