Ensure that BSONPickler and custom constructors are registered on every Spark node (HADOOP-273).

Author: Luke Lovett

spark/src/main/java/com/mongodb/spark/PySparkBSONFileInputFormat.java

@@ -0,0 +1,16 @@
+package com.mongodb.spark;
+
+import com.mongodb.hadoop.BSONFileInputFormat;
+import com.mongodb.spark.pickle.RegisterConstructors;
+import com.mongodb.spark.pickle.RegisterPickles;
+
+public class PySparkBSONFileInputFormat extends BSONFileInputFormat {
+    private static final RegisterPickles PICKLES = new RegisterPickles();
+    private static final RegisterConstructors CONSTRUCTORS =
+      new RegisterConstructors();
+
+    static {
+        PICKLES.register();
+        CONSTRUCTORS.register();
+    }
+}

spark/src/main/java/com/mongodb/spark/PySparkBSONFileOutputFormat.java

@@ -0,0 +1,17 @@
+package com.mongodb.spark;
+
+import com.mongodb.hadoop.BSONFileOutputFormat;
+import com.mongodb.spark.pickle.RegisterConstructors;
+import com.mongodb.spark.pickle.RegisterPickles;
+
+public class PySparkBSONFileOutputFormat<K, V>
+  extends BSONFileOutputFormat<K, V> {
+    private static final RegisterPickles PICKLES = new RegisterPickles();
+    private static final RegisterConstructors CONSTRUCTORS =
+      new RegisterConstructors();
+
+    static {
+        PICKLES.register();
+        CONSTRUCTORS.register();
+    }
+}

spark/src/main/java/com/mongodb/spark/PySparkMongoInputFormat.java

@@ -0,0 +1,20 @@
+package com.mongodb.spark;
+
+import com.mongodb.hadoop.MongoInputFormat;
+import com.mongodb.spark.pickle.RegisterConstructors;
+import com.mongodb.spark.pickle.RegisterPickles;
+
+/**
+ * InputFormat that attaches custom Picklers and IObjectConstructors for
+ * reading and writing BSON types with PyMongo.
+ */
+public class PySparkMongoInputFormat extends MongoInputFormat {
+    private static final RegisterPickles PICKLES = new RegisterPickles();
+    private static final RegisterConstructors CONSTRUCTORS =
+      new RegisterConstructors();
+
+    static {
+        PICKLES.register();
+        CONSTRUCTORS.register();
+    }
+}

spark/src/main/java/com/mongodb/spark/PySparkMongoOutputFormat.java

@@ -0,0 +1,17 @@
+package com.mongodb.spark;
+
+import com.mongodb.hadoop.MongoOutputFormat;
+import com.mongodb.spark.pickle.RegisterConstructors;
+import com.mongodb.spark.pickle.RegisterPickles;
+
+public class PySparkMongoOutputFormat<K, V>
+  extends MongoOutputFormat<K, V> {
+    private static final RegisterPickles PICKLES = new RegisterPickles();
+    private static final RegisterConstructors CONSTRUCTORS =
+      new RegisterConstructors();
+
+    static {
+        PICKLES.register();
+        CONSTRUCTORS.register();
+    }
+}

spark/src/main/java/com/mongodb/spark/pickle/RegisterConstructors.java

@@ -0,0 +1,31 @@
+package com.mongodb.spark.pickle;
+
+import net.razorvine.pickle.Unpickler;
+import org.bson.BSON;
+
+public class RegisterConstructors {
+    public void register() {
+        Unpickler.registerConstructor("bson.binary", "Binary",
+          new com.mongodb.spark.pickle.BinaryConstructor());
+        Unpickler.registerConstructor("bson.code", "Code",
+          new com.mongodb.spark.pickle.CodeConstructor());
+        Unpickler.registerConstructor("bson.dbref", "DBRef",
+          new com.mongodb.spark.pickle.DBRefConstructor());
+        Unpickler.registerConstructor("bson.int64", "Int64",
+          new com.mongodb.spark.pickle.Int64Constructor());
+        Unpickler.registerConstructor("bson.max_key", "MaxKey",
+          new com.mongodb.spark.pickle.MaxKeyConstructor());
+        Unpickler.registerConstructor("bson.min_key", "MinKey",
+          new com.mongodb.spark.pickle.MinKeyConstructor());
+        Unpickler.registerConstructor("bson.timestamp", "Timestamp",
+          new com.mongodb.spark.pickle.TimestampConstructor());
+        Unpickler.registerConstructor("bson.regex", "Regex",
+          new com.mongodb.spark.pickle.RegexConstructor());
+        Unpickler.registerConstructor("bson.objectid", "ObjectId",
+          new com.mongodb.spark.pickle.ObjectIdConstructor());
+
+        BSON.addEncodingHook(
+          java.util.GregorianCalendar.class,
+          new CalendarTransformer());
+    }
+}

spark/src/main/java/com/mongodb/spark/pickle/RegisterPickles.java

@@ -0,0 +1,23 @@
+package com.mongodb.spark.pickle;
+
+import net.razorvine.pickle.Pickler;
+
+public class RegisterPickles {
+    private static final BSONPickler PICKLER = new BSONPickler();
+
+    public void register() {
+        Pickler.registerCustomPickler(org.bson.types.ObjectId.class, PICKLER);
+        Pickler.registerCustomPickler(org.bson.types.Binary.class, PICKLER);
+        Pickler.registerCustomPickler(org.bson.types.Code.class, PICKLER);
+        Pickler.registerCustomPickler(org.bson.types.CodeWScope.class, PICKLER);
+        Pickler.registerCustomPickler(
+          org.bson.types.CodeWithScope.class, PICKLER);
+        Pickler.registerCustomPickler(org.bson.types.MaxKey.class, PICKLER);
+        Pickler.registerCustomPickler(org.bson.types.MinKey.class, PICKLER);
+        Pickler.registerCustomPickler(
+          org.bson.types.BSONTimestamp.class, PICKLER);
+        Pickler.registerCustomPickler(com.mongodb.DBRef.class, PICKLER);
+        Pickler.registerCustomPickler(java.util.regex.Pattern.class, PICKLER);
+        Pickler.registerCustomPickler(java.util.Date.class, PICKLER);
+    }
+}

spark/src/main/python/README.rst

@@ -23,6 +23,9 @@ Installation
      cd mongo-hadoop/spark/src/main/python
      python setup.py install
 
+3. Install `pymongo <https://pypi.python.org/pypi/pymongo>`_ on each machine in
+   your Spark cluster.
+
 You'll also need to put ``mongo-hadoop-spark.jar`` (see above for instructions
 on how to obtain this) somewhere on Spark's ``CLASSPATH`` prior to using this
 package.

spark/src/main/python/pymongo_spark.py

@@ -14,89 +14,9 @@
 
 __version__ = '0.1'
 
-import sys
-
-import py4j
 import pyspark
 
 
-# These types need special pickling to work correctly with PyMongo.
-_PICKLE_BSON_TYPES = (
-    'org.bson.types.ObjectId',
-    'org.bson.types.Binary',
-    'org.bson.types.Code',
-    'org.bson.types.CodeWScope',
-    'org.bson.types.CodeWithScope',
-    'org.bson.types.MaxKey',
-    'org.bson.types.MinKey',
-    'org.bson.types.BSONTimestamp',
-    'com.mongodb.DBRef',
-    'java.util.regex.Pattern',
-    'java.util.Date'
-)
-
-
-# Register Constructors for unpickling.
-# (module, class)
-_UNPICKLE_CONSTRUCTORS = (
-    ('bson.binary', 'Binary'),
-    ('bson.code', 'Code'),
-    ('bson.dbref', 'DBRef'),
-    ('bson.int64', 'Int64'),
-    ('bson.max_key', 'MaxKey'),
-    ('bson.min_key', 'MinKey'),
-    ('bson.timestamp', 'Timestamp'),
-    ('bson.regex', 'Regex'),
-    ('bson.objectid', 'ObjectId')
-)
-
-
-def _ensure_pickles(self):
-    if not getattr(self, '__registered_picklers', False):
-        try:
-            jvm = self._jvm
-            pickler = jvm.net.razorvine.pickle.Pickler
-            bson_pickler = jvm.com.mongodb.spark.pickle.BSONPickler()
-
-            for pbt in _PICKLE_BSON_TYPES:
-                pickler.registerCustomPickler(
-                    jvm.java.lang.Class.forName(pbt), bson_pickler)
-
-            unpickler = jvm.net.razorvine.pickle.Unpickler
-            for unpc in _UNPICKLE_CONSTRUCTORS:
-                unpickler.registerConstructor(
-                    unpc[0], unpc[1],
-                    jvm.java.lang.Class.forName(
-                        'com.mongodb.spark.pickle.%sConstructor' % unpc[1])
-                    .newInstance())
-
-            # Register CalendarTransformer with the Java driver so that we can
-            # encode java.util.GregorianCalendar objects. GregorianCalendar is
-            # what is constructed out of pickled datetime objects.
-            # We can't create a custom IObjectConstructor like we do for other
-            # BSON types, because the Razorvine library already has a
-            # constructor for datetimes.
-            jvm.org.bson.BSON.addEncodingHook(
-                # SyntaxError to access ".class" attribute.
-                jvm.java.lang.Class.forName('java.util.GregorianCalendar'),
-                jvm.java.lang.Class.forName(
-                    'com.mongodb.spark.pickle.CalendarTransformer')
-                .newInstance())
-            self.__registered_picklers = True
-        except py4j.protocol.Py4JError:
-            orig_t, orig_v, orig_tb = sys.exc_info()
-            try:
-                # Try to guess most common cause of failure.
-                raise (py4j.protocol.Py4JError,
-                       "Error while communicating with the JVM. "
-                       "Is the MongoDB Spark jar on Spark's CLASSPATH? : " +
-                       str(orig_v),
-                       orig_tb)
-            finally:
-                # Avoid circular reference with traceback.
-                del orig_tb
-
-
 def saveToMongoDB(self, connection_string, config=None):
     """Save this RDD to MongoDB."""
     conf = {'mongo.output.uri': connection_string}
@@ -112,7 +32,7 @@ def saveToMongoDB(self, connection_string, config=None):
         keyClass = 'org.apache.hadoop.io.NullWritable'
     to_save.saveAsNewAPIHadoopFile(
         'file:///this-is-unused',
-        outputFormatClass='com.mongodb.hadoop.MongoOutputFormat',
+        outputFormatClass='com.mongodb.spark.PySparkMongoOutputFormat',
         keyClass=keyClass,
         valueClass='com.mongodb.hadoop.io.BSONWritable',
         keyConverter='com.mongodb.spark.pickle.NoopConverter',
@@ -132,7 +52,7 @@ def saveToBSON(self, file_path, config=None):
         keyClass = 'org.apache.hadoop.io.NullWritable'
     to_save.saveAsNewAPIHadoopFile(
         file_path,
-        outputFormatClass='com.mongodb.hadoop.BSONFileOutputFormat',
+        outputFormatClass='com.mongodb.spark.PySparkBSONFileOutputFormat',
         keyClass=keyClass,
         valueClass='com.mongodb.hadoop.io.BSONWritable',
         keyConverter='com.mongodb.spark.pickle.NoopConverter',
@@ -143,23 +63,21 @@ def saveToBSON(self, file_path, config=None):
 
 def BSONFilePairRDD(self, file_path, config=None):
     """Create a pair RDD backed by a BSON file."""
-    _ensure_pickles(self)
     return self.newAPIHadoopFile(
         file_path,
-        inputFormatClass='com.mongodb.hadoop.BSONFileInputFormat',
+        inputFormatClass='com.mongodb.spark.PySparkBSONFileInputFormat',
         keyClass='com.mongodb.hadoop.io.BSONWritable',
         valueClass='com.mongodb.hadoop.io.BSONWritable',
         conf=config)
 
 
 def mongoPairRDD(self, connection_string, config=None):
     """Create a pair RDD backed by MongoDB."""
-    _ensure_pickles(self)
     conf = {'mongo.input.uri': connection_string}
     if config:
         conf.update(config)
     return self.newAPIHadoopRDD(
-        inputFormatClass='com.mongodb.hadoop.MongoInputFormat',
+        inputFormatClass='com.mongodb.spark.PySparkMongoInputFormat',
         keyClass='com.mongodb.hadoop.io.BSONWritable',
         valueClass='com.mongodb.hadoop.io.BSONWritable',
         conf=conf)