Создайте столбец ArrayType [StructType] из Dataframe в UDF - PullRequest
Новая
       18

Создайте столбец ArrayType [StructType] из Dataframe в UDF

0 голосов
/ 10 апреля 2019

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

def approxNearestNeighbors(
      dataset: Dataset[_],
      key: Vector,
      numNearestNeighbors: Int): Dataset[_]

Чтобы запустить ее на каждой строке кадра данных, моя идея состоит в том, чтобы создать некоторый udf-файл, который вызывает эту функцию, и преобразовать полученный набор данных в столбец ArrayType [StructType].

Предположим, что моя начальная схема 

root
 |-- genderIndex: double (nullable = false)
 |-- genderIndexVec: vector (nullable = true)
 |-- categoryIndex: double (nullable = false)
 |-- categoryIndexVec: vector (nullable = true)
 |-- features: vector (nullable = true)
 |-- featureStdDev: vector (nullable = true)

Моя целевая схема (после вызова .withColumn ($ "featureStdDev", udf ...)) равна 

root
 |-- genderIndex: double (nullable = false)
 |-- genderIndexVec: vector (nullable = true)
 |-- categoryIndex: double (nullable = false)
 |-- categoryIndexVec: vector (nullable = true)
 |-- features: vector (nullable = true)
 |-- featureStdDev: vector (nullable = true)
 |-- neighbours: array(nullable = true)
      |-- elem: struct
           |-- genderIndex: double (nullable = false)
           |-- genderIndexVec: vector (nullable = true)
           |-- categoryIndex: double (nullable = false)
           |-- categoryIndexVec: vector (nullable = true)
           |-- features: vector (nullable = true)
           |-- featureStdDev: vector (nullable = true)

Пожалуйста, помогите с моим UDF, так как я не уверен, как заставить его работать. 

val model = // BucketedRandomProjectionLSHModel definition
val inputDF = // Input definition
val nn = udf{ (featureVector: SparseVector, k: Int) =>
      model.approxNearestNeighbors(inputDF, featureVector, k)
      // What now...
    }
...