Как пройти карту в Spark Udf? - PullRequest
Новая
       7

Как пройти карту в Spark Udf?

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

У меня есть вопрос.У меня есть искровой датафрейм с несколькими столбцами, похожими на:

id Color1 красный, синий, черный2 красный, зеленый3 синий, желтый, зеленый...

У меня также есть файл карты, похожий на:Красный, 0Синий, 1Зеленый, 2Черный, 3Желтый, 4

Мне нужно сопоставить имя цвета с разными идентификаторами, такими как отображение «Красный, Синий, Черный» в массив [1,1,0,1,0].Я пишу код следующим образом: 

def mapColor(label_string:String):Array[Int]={
var labels = label_string.split(",")
var index_array = new Array[Int](COLOR_LENGTH)
for (label<-labels){
  if(COLOR_MAP.contains(label)){
    index_array(COLOR_MAP(label))=1
  }
  else{
    //dictionary does not contain the label, the last index set to be one
    index_array(COLOR_LENGTH-1)=1
  }
}
index_array 
}

COLOR_LENGTH - это длина словаря, а COLOR_MAP - это словарь, который содержит отношение string-> id.

Я называю эту функцию следующим образом: 

 val color_function = udf(mapColor:(String)=>Array[Int])
 sql.withColumn("color_idx",color_function(col("Color")))

Поскольку у меня несколько столбцов, эта операция нужна, но для разных столбцов нужны разные словари.В настоящее время я дублирую эту функцию для каждого столбца (просто измените словарь и информацию о длине).Но код выглядит утомительно.Есть ли какой-нибудь метод, я могу передать длину и словарь в функцию отображения, например 

def map(label_string:String,map:Map[String,Integer],len:Int):Array[Int]

Но как мне вызвать эту функцию в кадре данных spark?Поскольку я не могу передать параметр в объявлении 

val color_function = udf(mapColor:(String)=>Array[Int])

Ответы [ 2 ]

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

Вот полный код для краткости - 

val colrMapList = List("Red" -> 0, "Blue" -> 1, "Green" -> 2).toMap

def getColor = udf((colors: Seq[String]) => { if(!colors.isEmpty) colors.map(color => colrMapList.getOrElse(color,"0")).mkString(",") else "0"  } )

val colors = List((1, Array("Red","Blue","Black")),(2,Array("Red", "Green")))
val colrDF = sc.parallelize(colors).toDF

colrDF.withColumn("colorMap", getColor($"colors")).show

Объяснение

  1. Создание map для отображения цвета в целое число.
  2. Функция getColor извлекает соответствующие целые числа , учитывая цвета
  3. Наконец, вы применяете функцию colrDF, чтобы получить вывод
0 голосов
/ 17 мая 2018

Вы можете использовать UDF с цветовой картой в качестве базового аргумента, как в следующем примере: 

val df = Seq(
  (1, "Red, Blue, Black"),
  (2, "Red, Green"),
  (3, "Blue, Yellow, Green")
).toDF("id", "color")

val colorMap = Map("Red"-> 0, "Blue"->1, "Green"->2, "Black"->3, "Yellow"->4)

def mapColorCode(m: Map[String, Int]) = udf( (s: String) =>
  s.split("""\s*,\s*""").map(c => m.getOrElse(c, -99))
)

df.select($"id", mapColorCode(colorMap)($"color").as("colorcode")).show
// +---+----------+
// | id| colorcode|
// +---+----------+
// |  1| [0, 1, 3]|
// |  2|    [0, 2]|
// |  3| [1, 4, 2]|
// +---+----------+
...