python - Spark RDD - 使用额外参数进行映射

Question

是否可以将额外的参数传递给 pySpark 中的映射函数？具体来说，我有以下代码配方:

raw_data_rdd = sc.textFile("data.json", use_unicode=True)
json_data_rdd = raw_data_rdd.map(lambda line: json.loads(line))
mapped_rdd = json_data_rdd.flatMap(processDataLine)

processDataLine 函数除了 JSON 对象外还接受额外的参数，如:

def processDataLine(dataline, arg1, arg2)

如何将额外的参数 arg1 和 arg2 传递给 flaMap 函数？

Answer 1

1. 您可以直接在 flatMap

   中使用匿名函数

   json_data_rdd.flatMap(lambda j: processDataLine(j, arg1, arg2))
   

   或柯里化(Currying)processDataLine

   f = lambda j: processDataLine(dataline, arg1, arg2)
   json_data_rdd.flatMap(f)
   

2. 你可以像这样生成processDataLine:

   def processDataLine(arg1, arg2):
       def _processDataLine(dataline):
           return ... # Do something with dataline, arg1, arg2
       return _processDataLine
   
   json_data_rdd.flatMap(processDataLine(arg1, arg2))
   

3. toolz库提供了有用的 curry 装饰器:

   from toolz.functoolz import curry
   
   @curry
   def processDataLine(arg1, arg2, dataline): 
       return ... # Do something with dataline, arg1, arg2
   
   json_data_rdd.flatMap(processDataLine(arg1, arg2))
   

   请注意，我已将 dataline 参数推到最后一个位置。这不是必需的，但这样我们就不必使用关键字 args。

4. 终于有了functools.partial Avihoo Mamka 已经提到过在评论中。