How to pass data into udf not from row

user_8803744 0
Viewed 35

Let's say I have the following dataframe schema:

+-------+-------+
| body  | rules |
+-------+-------+

I have a udf that takes in the body column and rule-list column for each row, and parses and evaluated the conditions for the rules based on the row (and returns a list of booleans whether it each rule matches or not). Right now, every single row in the DF has a copy of these rules because I don't know any other way to pass in these rules to the UDF. This feels very redundant and wasteful to me.

The rules are joined onto the row based on some join conditions, so each row doesn't have the exact same data but there is still a lot of redundancy (each rule is probably listed ~5000 redundant times across 1 million rows). I would prefer to join the ruleIds onto each row instead and pass a map(ruleId -> rule) into the udf. This map may be somewhat large though so however it is passed in would have to be able to handle that (ideally it would be some sort of shared variable stored at the partition level)

spark-streaming apache-spark-sql apache-spark
1 Answers

You probably want to use an external map parameter to your udf, and broadcast this map to each machine to avoid having a copy of it for each task.

For your UDF, you can do something along those lines:

def yourUDF(rulesMap: Map[String, XXX]): UserDefinedFunction = udf {
  (body: YYY, ruleId: String) => applyYourRules(body, rulesMap(ruleId))
} // XXX and YYY are the types you need, I don't know the problem you're trying to solve

And as you map is quite large, you can avoid duplicating it for every task by broadcasting the variable (you can access a broadcasted variable with variable.value):

val rulesMapBroadcasted = spark.sparkContext.broadcast(rulesMap)
df.withColumn("new", yourUDF(rulesMap = rulesMapBroadcasted.value)(col("body"), col("ruleId")))

A broadcasted variable is a read-only variable duplicated only once per machine (in comparison, a classical variable is duplicated once per task), so this is a perfect usage for large lookup table.

user_181596030
Related