Spark Binary Encode Array

Question

Я пытался закодировать массив значений в Spark 2.3.0 в Scala.Я попытался использовать Spark OneHotEncoderEstimator, который отлично работает для отдельных String или Int, но я не смог расширить его для работы с массивом (один из способов, которым я пытался, был explode столбец - который работал,но тогда я не смог вернуть его обратно в массив. Наоборот, collect_set, просто дал массив с горячим кодированием Vectors).

Я нашел этот вопрос:

Одна горячая кодировка в RDD в scala

, которую я использовал в качестве основы для разработки функции для работы смассив.Я изменил функцию, опубликованную в решении:

def encode(x: String) = {
var encodeArray = Array.fill(21)(0)
encodeArray(indexed_user.get(x).get.toInt)=1
encodeArray}

так:

def encodeSeq(x:Seq[String]) = {
var encodeArray = Array.fill(21)(0)
x.foreach(s => encodeArray(indexed.get(s).get.toInt)=1)
encodeArr}

Моя indexed - это карта String -> Int.Я протестировал эту функцию на одной последовательности, и она возвращает ожидаемые результаты:

val testSeq = Seq("a", "b")
encodeSeq(testSeq)
res97: Array[Int] = Array(1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0)

Но при применении к RDD я просто получаю RDD[Ints] назад, а не ожидаемый RDD[Array]:

val encoded = unencoded.rdd.flatMap(x=> encodeSeq(x.getSeq[String](0)))
encoded: org.apache.spark.rdd.RDD[Int] = MapPartitionsRDD[336] at flatMap at <console>:44

Здесь unencoded равно org.apache.spark.sql.DataFrame = [col1: array<string>, col2: array<string> ... 3 more fields]

Проверка результата кодирования дает:

encoded.toDF().show(5, false)
+-----+
|value|
+-----+
|0    |
|0    |
|0    |
|0    |
|0    |
+-----+

Я все еще новичок в Spark и Scala,Я делаю что-то не так на карте или где-то еще?Или есть лучший способ добиться этого?