MongoDB Spark Connector使用

背景

对现有的存储在MongoDB中的海量数据进行清洗操作。为了提供数据的处理速度，决定使用spark作为引擎对数据进行操作。

先决条件

• MongoDB 2.6+
• Spark 2.2.x
• Scala 2.11.x
• Java 1.8+

MonogoDB Spark Connector介绍

1. 支持数据的读取和写入。
2. 将查询拆分为n个子任务, 如Connector会将一次match,拆分为多个子任务交给spark来处理, 减少数据的全量读取。
3. 支持Aggregation，即在执行任务时，Spark可以将Aggregation操作下发到MongoDB去执行聚合后再读取相应的数据。
4. 支持通过Spark SQL操作数据集。

示例

Maven依赖

  <dependency>
    <groupId>org.mongodb.spark</groupId>
    <artifactId>mongo-spark-connector_2.11</artifactId>
    <version>2.4.0</version> <!-- 2.2.0会有Mongo游标异常的bug-->
  </dependency>
  <dependency>
    <groupId>org.apache.spark</groupId>
    <artifactId>spark-core_2.11</artifactId>
    <version>2.2.0</version>
  </dependency>
  <dependency>
    <groupId>org.apache.spark</groupId>
    <artifactId>spark-sql_2.11</artifactId>
    <version>2.2.0</version>
  </dependency>

创建SparkSession

SparkSession spark = SparkSession.builder()
                .master("local")
                .appName("MongoSparkDataClean")
                            //设置读取的mongo集合信息。mongodb://127.0.0.1/maoke.kg_entries_address
                .config("spark.mongodb.input.uri", inputUri)
                                //设置写入的mongo集合信息。 mongodb://127.0.0.1/maoke.kg_entries_address
                .config("spark.mongodb.output.uri", outputUri)
                .getOrCreate();

处理MongoDB数据

JavaSparkContext jsc = new JavaSparkContext(spark.sparkContext());
//读取MongoDB数据，可以通过ReadConfig配置读取操作的特殊配置。
JavaMongoRDD<Document> mongoRDD = MongoSpark.load(jsc);
//处理MongoDB数据
Dataset<KgEntriesAddress> kgEntriesAddressDataset = mongoRDD.toDS(KgEntriesAddress.class)
                .map((MapFunction<KgEntriesAddress, KgEntriesAddress>) value -> {
                    boolean checkAddress = checkData(ADDRESS_URL, value.getValue());
                    value.setFlag(checkAddress);
                    return value;
                }, Encoders.bean(KgEntriesAddress.class));
//写入MongoDB，可以通过WriteConfig配置写入操作的特殊配置。
MongoSpark.save(kgEntriesAddressDataset);
LOGGER.info("Total number of clean address processed data: {}", mongoRDD.count());

⚠️注意

作者在使用MongoSpark.save()写入数据时，一开始MongoDB会抛出异常:

com.mongodb.MongoBulkWriteException: Bulk write operation error on server 127.0.0.1:27017. Write errors: [BulkWriteError{index=0, code=11000, message='E11000 duplicate key error collection: maoke.kg_entries_address index: _id_ dup key: { : "116d2839b7c5aee6758961ea4875ee65" }', details={ }}]. 
.....  

查看MongoSpark.save()源码后发现当参数类型为RDD时，由于spark无法得知数据schema，不能指定_id字段,采用的是insertMany对数据进行写入。所以当写入的数据与原始的数据_id相同时，抛出了E11000 duplicate key error异常。

所以，如果你是对同一个MongoDB Collection进行读取和写入操作，需要将写入的数据类型转化为DataSet。

/**
   * Save data to MongoDB
   *
   * @param rdd the RDD data to save to MongoDB
   * @param writeConfig the writeConfig
   * @tparam D the type of the data in the RDD
   */
  def save[D: ClassTag](rdd: RDD[D], writeConfig: WriteConfig): Unit = {
    val mongoConnector = MongoConnector(writeConfig.asOptions)
    rdd.foreachPartition(iter => if (iter.nonEmpty) {
      mongoConnector.withCollectionDo(writeConfig, { collection: MongoCollection[D] =>
        iter.grouped(writeConfig.maxBatchSize).foreach(batch => collection.insertMany(batch.toList.asJava))
      })
    })
  }

/**
   * Save data to MongoDB
   *
   * '''Note:''' If the dataFrame contains an `_id` field the data will upserted and replace any existing documents in the collection.
   *
   * @param dataset the dataset to save to MongoDB
   * @param writeConfig the writeConfig
   * @tparam D
   * @since 1.1.0
   */
  def save[D](dataset: Dataset[D], writeConfig: WriteConfig): Unit = {
    val mongoConnector = MongoConnector(writeConfig.asOptions)
    val documentRdd: RDD[BsonDocument] = dataset.toDF().rdd.map(row => rowToDocument(row))
    if (dataset.schema.fields.exists(_.name == "_id")) {
      documentRdd.foreachPartition(iter => if (iter.nonEmpty) {
        mongoConnector.withCollectionDo(writeConfig, { collection: MongoCollection[BsonDocument] =>
          iter.grouped(writeConfig.maxBatchSize).foreach(batch => {
            val updateOptions = new UpdateOptions().upsert(true)
            val requests = batch.map(doc =>
              Option(doc.get("_id")) match {
                case Some(_id) =>
                  if (writeConfig.replaceDocument) {
                    new ReplaceOneModel[BsonDocument](new BsonDocument("_id", _id), doc, updateOptions)
                  } else {
                    doc.remove("_id")
                    new UpdateOneModel[BsonDocument](new BsonDocument("_id", _id), new BsonDocument("$set", doc), updateOptions)
                  }
                case None => new InsertOneModel[BsonDocument](doc)
              })
            collection.bulkWrite(requests.toList.asJava)
          })
        })
      })
    } else {
      MongoSpark.save(documentRdd, writeConfig)
    }
  }

参考链接

• MongoDB Spark Connector官方文档
• MongoDB Spark Connector GitHub