Scala-spark: как случайным образом выбрать элемент в столбце массивов разных размеров - PullRequest
Новая
       8

Scala-spark: как случайным образом выбрать элемент в столбце массивов разных размеров

0 голосов
/ 22 мая 2018

С учетом кадра данных со столбцом массивов целых чисел разных размеров: 

scala> sampleDf.show()
+------------+
|      arrays|
+------------+
|[15, 16, 17]|
|[15, 16, 17]|
|        [14]|
|        [11]|
|        [11]|
+------------+

scala> sampleDf.printSchema()
root
 |-- arrays: array (nullable = true)
 |    |-- element: integer (containsNull = true)

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

  import scala.util.Random

  def getRandomElement(arr: Array[Int]): Int = {
    arr(Random.nextInt(arr.size))
  }

  val getRandomElementUdf = udf{arr: Array[Int] => getRandomElement(arr)}

  sampleDf.withColumn("randomItem", getRandomElementUdf('arrays)).show

вылетает в последней строке с длинным сообщением об ошибке: (извлекает) 

...
Caused by: org.apache.spark.SparkException: Failed to execute user defined function($anonfun$1: (array<int>) => int)
...
Caused by: java.lang.ClassCastException: scala.collection.mutable.WrappedArray$ofRef cannot be cast to [I

Имеюпробовал с альтернативным определением udf: 

 val getRandomElementUdf = udf[Int, Array[Int]] (getRandomElement)

но я получаю ту же ошибку.
2. Второй метод путем создания промежуточных столбцов со случайным индексом в диапазоне соответствующего массива: 

 // Return a dataframe with a column with random index from column of Arrays with different sizes
 def choice(df: DataFrame, colName: String): DataFrame = {
    df.withColumn("array_size", size(col(colName)))
      .withColumn("random_idx", least('array_size, floor(rand * 'array_size))) 
 }

 choice(sampleDf, "arrays").show

выводит: 

+------------+----------+----------+
|      arrays|array_size|random_idx|
+------------+----------+----------+
|[15, 16, 17]|         3|         2|
|[15, 16, 17]|         3|         1|
|        [14]|         1|         0|
|        [11]|         1|         0|
|        [11]|         1|         0|
+------------+----------+----------+

и в идеале мы хотели бы использовать столбец random_idx для выбора элемента в столбце arrays, вид: 

sampleDf.withColumn("choosen_item", 'arrays.getItem('random_idx))

К сожалению, getItem не может принять столбец в качестве аргумента.

Любые предложения приветствуются.

Ответы [ 2 ]

0 голосов
/ 22 мая 2018

Если вы хотите остаться без udf, вот такая возможность: сначала добавьте ключ к фрейму данных, выданный choice (предположим, его имя choiceDf) 

val myDf = choiceDf.withColumn("key", monotonically_increasing_id())

, затем создайтепромежуточный фрейм данных, который разбивает столбец массивов и сохраняет индекс значений 

val tmp = myDf.select('key, posexplode('arrays))

, наконец, объединяются с использованием ключа и random_idx

myDf.join(tmp.withColumnRenamed("pos", "random_idx"), Seq("key", "random_idx", "left")

искомый элемент сохраняется в столбцеcol 

+---+----------+------------+----------+---+
|key|random_idx|      arrays|array_size|col|
+---+----------+------------+----------+---+
|  0|         2|[15, 16, 17]|         3| 17|
|  1|         1|[15, 16, 17]|         3| 16|
|  2|         0|        [14]|         1| 14|
+---+----------+------------+----------+---+
0 голосов
/ 22 мая 2018

Вы можете использовать ниже udf, чтобы выбрать элемент random из массива как 

val getRandomElement = udf ((array: Seq[Integer]) => {
  array(Random.nextInt(array.size))
})


df.withColumn("c1", getRandomElement($"arrays"))
  .withColumn("c2", getRandomElement($"arrays"))
  .withColumn("c3", getRandomElement($"arrays"))
  .withColumn("c4", getRandomElement($"arrays"))
  .withColumn("c5", getRandomElement($"arrays"))

  .show(false)

. Вы можете видеть элемент random, выбранный при каждом использовании в качестве нового столбца. 

+------------+---+---+---+---+---+
|arrays      |c1 |c2 |c3 |c4 |c5 |
+------------+---+---+---+---+---+
|[15, 16, 17]|15 |16 |16 |17 |16 |
|[15, 16, 17]|16 |16 |17 |15 |15 |
|[14]        |14 |14 |14 |14 |14 |
|[11]        |11 |11 |11 |11 |11 |
|[11]        |11 |11 |11 |11 |11 |
+------------+---+---+---+---+---+
...