Transmitir una clase definida por el usuario en Spark

Estoy tratando de transmitir una variable definida por el usuario en una aplicación PySpark pero siempre tengo el siguiente error:

File "/usr/local/spark-2.1.0-bin-hadoop2.7/python/lib/pyspark.zip/pyspark/worker.py", line 174, in main process() File "/usr/local/spark-2.1.0-bin-hadoop2.7/python/lib/pyspark.zip/pyspark/worker.py", line 169, in process serializer.dump_stream(func(split_index, iterator), outfile) File "/usr/local/spark-2.1.0-bin-hadoop2.7/python/lib/pyspark.zip/pyspark/serializers.py", line 268, in dump_stream vs = list(itertools.islice(iterator, batch)) File "/home/.../sparkbroad.py", line 29, in  output = input_.map(lambda item: b.value.map(item)) File "/usr/local/spark-2.1.0-bin-hadoop2.7/python/lib/pyspark.zip/pyspark/broadcast.py", line 106, in value self._value = self.load(self._path) File "/usr/local/spark-2.1.0-bin-hadoop2.7/python/lib/pyspark.zip/pyspark/broadcast.py", line 97, in load return pickle.load(f) AttributeError: 'module' object has no attribute 'FooMap' 

El código, en el módulo sparkbrad.py es el siguiente:

 import random import pyspark as spark class FooMap(object): def __init__(self): keys = list(range(10)) values = [2 * key for key in keys] self._map = dict(zip(keys, values)) def map(self, value): if value not in self._map: return -1 return self._map[value] class FooMapJob(object): def __init__(self, inputs): self._inputs = inputs self._foomap = FooMap() def run(self): sc = spark.SparkContext('local', 'FooMap') input_ = sc.parallelize(self._inputs, 4) b = sc.broadcast(self._foomap) output = input_.map(lambda item: b.value.map(item)) b.unpersist() result = list(output.toLocalIterator()) sc.stop() return result def main(): inputs = [random.randint(0, 10) for _ in range(10)] job = FooMapJob(inputs) print(job.run()) if __name__ == '__main__': main() 

y lo estoy ejecutando a través de:

 :~$ spark-submit --master local[4] --py-files sparkbroad.py sparkbroad.py 

donde he agregado el argumento --py-files pero parece que no cambia mucho. Desafortunadamente, no pude encontrar ningún ejemplo en línea relacionado con la transmisión de clases complejas (solo listas o diccionarios). Cualquier sugerencia es apreciada. Gracias por adelantado.

ACTUALIZACIÓN : colocando la clase FooMap en un módulo separado, todo parece funcionar bien, incluso sin la directiva --py-files .