UDF根據不同的概率隨機分配值

Question

我想創建一個UDF以根據不同的概率隨機分配值。

在下面的示例取決於由蘭特返回的值：

• 0〜0.5的值應爲A（50％的概率）
• 0.8至1的值應爲B（20％的概率）
• 什麼都值應爲C（30％概率）

val names = Array("A", "B", "C")

val allocate = udf((p: Double) => { 
    if(p < 0.5) names(0) 
    else if (p > 0.8) names(1) 
    else names(2)}) 

val test = sqlContext.range(0, 100).select(($"id"),(round(abs(rand),2)).alias("val"), allocate(abs(rand)).alias("name")) 
` 

但是，當我打印結果時，不會根據UDF中定義的規則分配名稱。

+---+----+----+ 
| id| val|name| 
+---+----+----+ 
| 0|0.17| C| => should be A 
| 1|0.12| A| 
| 2|0.36| A| 
| 3|0.56| B| 
| 4|0.82| A|=> should be C 

Answer 1

這裏沒有什麼意外的事情發生。您可以調用rand函數兩次，以獲得兩個不同的隨機值。

要麼的用於呼叫提供相同的種子：

sqlContext.range(0, 100) 
    .select(
    $"id", 
    abs(rand(1)).alias("val"), 
    allocate(abs(rand(1))).alias("name") 
) 

或重用值：

sqlContext.range(0, 100) 
    .withColumn("val", abs(rand)) 
    .withColumn("name", allocate($"val"))