根据规则在 Spark 数据框中添加列

Question

我有一个数据框 df，其中包含以下数据：

**customers**   **product**   **Val_id**
     1               A            1
     2               B            X
     3               C               
     4               D            Z

我提供了两条规则，如下：

**rule_id**   **rule_name**  **product value**  **priority**
   123              ABC             A,B               1
   456              DEF             A,B,D             2

要求是按优先级顺序将这些规则应用于数据帧 df，通过规则 1 的客户不应考虑使用规则 2，并在最终数据帧中添加另外两列 rule_id 和 rule_name，我已经编写了下面的代码来实现它：

val rule_name = when(col("product").isin("A","B"), "ABC").otherwise(when(col("product").isin("A","B","D"), "DEF").otherwise(""))
val rule_id = when(col("product").isin("A","B"), "123").otherwise(when(col("product").isin("A","B","D"), "456").otherwise(""))
val df1 = df_customers.withColumn("rule_name" , rule_name).withColumn("rule_id" , rule_id)
df1.show()

最终输出如下：

**customers**   **product**   **Val_id**  **rule_name**  **rule_id**
     1               A            1           ABC            123
     2               B            X           ABC            123
     3               C               
     4               D            Z           DEF            456

有没有更好的方法来实现它，只需遍历整个数据集一次而不是遍历整个数据集两次来添加两列？

Answer 1

问题：有没有更好的方法来实现它，添加两列 只需遍历整个数据集一次，而不是遍历 整个数据集两次？

答案：你可以在scala中有一个Map返回类型...

限制：例如，如果您使用 With Column，则此 udf 列名是 ruleIDandRuleName 那么你可以使用单个函数 使用 Map 数据类型或任何可接受的 spark sql 列数据类型。 否则你不能使用下面提到的方法

如下例sn-p所示

      def ruleNameAndruleId = udf((product : String) => {  
 if(Seq("A", "B").contains(product)) Map("ruleName"->"ABC","ruleId"->"123")   
    else if(Seq("A", "B", "D").contains(product)) (Map("ruleName"->"DEF","ruleId"->"456")   
    else (Map("ruleName"->"","ruleId"->"") })

来电者将是

df.withColumn("ruleIDandRuleName",ruleNameAndruleId(product here) ) // ruleNameAndruleId  will return a map containing rulename and rule id

Answer 2

您的解决方案的替代方法是使用udf 函数。它几乎类似于when 功能，因为两者都需要serialization 和deserialization。由您来测试哪个更快更高效。

def rule_name = udf((product : String) => {
  if(Seq("A", "B").contains(product)) "ABC"
  else if(Seq("A", "B", "D").contains(product)) "DEF"
  else ""
})

def rule_id = udf((product : String) => {
  if(Seq("A", "B").contains(product)) "123"
  else if(Seq("A", "B", "D").contains(product)) "456"
  else ""
})

val df1 = df_customers.withColumn("rule_name" , rule_name(col("product"))).withColumn("rule_id" , rule_id(col("product")))
df1.show()