All Downloads are FREE. Search and download functionalities are using the official Maven repository.

com.mongodb.hadoop.splitter.MongoCollectionSplitter Maven / Gradle / Ivy

/**
 * Copyright 2010-2013 10gen Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.mongodb.hadoop.splitter;

import com.mongodb.BasicDBObject;
import com.mongodb.DB;
import com.mongodb.DBCollection;
import com.mongodb.DBCursor;
import com.mongodb.DBObject;
import com.mongodb.Mongo;
import com.mongodb.MongoClientURI;
import com.mongodb.MongoURI;
import com.mongodb.hadoop.input.MongoInputSplit;
import com.mongodb.hadoop.util.MongoConfigUtil;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.mapreduce.InputSplit;
import org.bson.BSONObject;
import org.bson.types.MaxKey;
import org.bson.types.MinKey;

import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

public abstract class MongoCollectionSplitter extends MongoSplitter {

    private static final Log LOG = LogFactory.getLog(MongoCollectionSplitter.class);

    public static final MinKey MIN_KEY_TYPE = new MinKey();
    public static final MaxKey MAX_KEY_TYPE = new MaxKey();
    //CHECKSTYLE:OFF
//    protected Mongo mongo;
//    protected DBCollection inputCollection;

    //A reference to the database that we authenticated against, if there was
    //some authURI provided in the config.
    protected DB authDB;
    //CHECKSTYLE:ON


    public MongoCollectionSplitter() {
    }

    public MongoCollectionSplitter(final Configuration conf) {
        super(conf);
    }

    @Override
    public abstract List calculateSplits() throws SplitFailedException;

    /**
     * Contacts the config server and builds a map of each shard's name to its
     * host(s) by examining config.shards.
     * @return a Map of shard name onto shard hostnames
     */
    protected Map> getShardsMap() {
        DBCursor cur = null;
        HashMap> shardsMap = new HashMap>();
        DB configDB;
        try {
            configDB = getConfigDB();
            DBCollection shardsCollection = configDB.getCollection("shards");
            cur = shardsCollection.find();
            while (cur.hasNext()) {
                final BasicDBObject row = (BasicDBObject) cur.next();
                String host = row.getString("host");
                // for replica sets host will look like: "setname/localhost:20003,localhost:20004"
                int slashIndex = host.indexOf('/');
                if (slashIndex > 0) {
                    host = host.substring(slashIndex + 1);
                }
                shardsMap.put(row.getString("_id"), Arrays.asList(host.split(",")));
            }
        } finally {
            if (cur != null) {
                cur.close();
            }
        }
        return shardsMap;
    }

    protected DB getConfigDB() {
        Mongo mongo;
        MongoClientURI inputURI = MongoConfigUtil.getInputURI(getConfiguration());
        MongoClientURI authURI = MongoConfigUtil.getAuthURI(getConfiguration());

        final DBCollection inputCollection;
        if (authURI != null
                && authURI.getUsername() != null
                && authURI.getPassword() != null) {
            inputCollection = MongoConfigUtil.getCollectionWithAuth(inputURI, authURI);
        } else {
            inputCollection = MongoConfigUtil.getCollection(inputURI);
        }

        DB db = inputCollection.getDB();
        mongo = db.getMongo();
        if (authURI != null) {
            if (authURI.getUsername() != null
                && authURI.getPassword() != null) {
                authDB = mongo.getDB(authURI.getDatabase());
            }
        }

        return mongo.getDB("config");
    }

    /**
     * Takes an existing {@link MongoClientURI} and returns a new modified URI which replaces the original's server host + port with a
     * supplied new server host + port, but maintaining all the same original options. This is useful for generating distinct URIs for each
     * mongos instance so that large batch reads can all target them separately, distributing the load more evenly. It can also be used to
     * force a block of data to be read directly from the shard's servers directly, bypassing mongos entirely.
     *
     * @param originalUri  the URI to rewrite
     * @param newServerUris the new host(s) to target, e.g. server1:port1[,
     *                      server2:port2,...]
     * @return the rewritten URI
     */
    protected static MongoClientURI rewriteURI(
      final MongoClientURI originalUri, final List newServerUris) {
        String originalUriString = originalUri.toString();
        originalUriString = originalUriString.substring(MongoURI.MONGODB_PREFIX.length());

        // uris look like: mongodb://fred:foobar@server1[,server2]/path?options
        //

        //Locate the last character of the original hostname
        int serverEnd;
        int idx = originalUriString.lastIndexOf("/");
        serverEnd = idx < 0 ? originalUriString.length() : idx;

        //Locate the first character of the original hostname
        idx = originalUriString.indexOf("@");
        int serverStart = idx > 0 ? idx + 1 : 0;

        StringBuilder sb = new StringBuilder(originalUriString);
        StringBuilder joinedHosts = new StringBuilder();
        for (String host : newServerUris) {
            joinedHosts.append(host).append(',');
        }
        sb.replace(serverStart, serverEnd,
          joinedHosts.substring(0, joinedHosts.length() - 1));
        return new MongoClientURI(MongoURI.MONGODB_PREFIX + sb);
    }

    protected static MongoClientURI rewriteURI(
      final MongoClientURI originalURI, final String newURI) {
        return rewriteURI(originalURI, Collections.singletonList(newURI));
    }

    /**
     * Create an instance of MongoInputSplit that represents a view of this
     * splitter's input URI between the given lower/upper bounds. If this
     * splitter has range queries enabled, it will attempt to use $gte/$lt
     * clauses in the query construct to create the split, otherwise it will use
     * min/max index boundaries (default behavior).
     *
     * @param lowerBound the lower bound of the collection
     * @param upperBound the upper bound of the collection
     * @return a MongoInputSplit in the given bounds
     * @throws SplitFailedException if the split could not be created
     */
    public MongoInputSplit createSplitFromBounds(final BasicDBObject lowerBound, final BasicDBObject upperBound)
        throws SplitFailedException {
        LOG.info("Created split: min=" + (lowerBound != null ? lowerBound.toString() : "null") + ", max= " + (upperBound != null
                                                                                                              ? upperBound.toString()
                                                                                                              : "null"));
        //Objects to contain upper/lower bounds for each split
        DBObject splitMin = new BasicDBObject();
        DBObject splitMax = new BasicDBObject();
        if (lowerBound != null) {
            for (Entry entry : lowerBound.entrySet()) {
                String key = entry.getKey();
                Object val = entry.getValue();

                if (!val.equals(MIN_KEY_TYPE)) {
                    splitMin.put(key, val);
                }
            }
        }

        if (upperBound != null) {
            for (Entry entry : upperBound.entrySet()) {
                String key = entry.getKey();
                Object val = entry.getValue();

                if (!val.equals(MAX_KEY_TYPE)) {
                    splitMax.put(key, val);
                }
            }
        }

        MongoInputSplit split = null;

        // If enabled, attempt to build the split using $gte/$lt.
        if (MongoConfigUtil.isRangeQueryEnabled(getConfiguration())) {
            try {
                DBObject query = MongoConfigUtil.getQuery(getConfiguration());
                split = createRangeQuerySplit(lowerBound, upperBound, query);
            } catch (Exception e) {
                throw new SplitFailedException("Couldn't use range query to create split: " + e.getMessage());
            }
        }
        if (split == null) {
            split = new MongoInputSplit(getConfiguration());
            split.setMin(splitMin);
            split.setMax(splitMax);
        }
        return split;
    }

    /**
     * Creates an instance of {@link MongoInputSplit} whose upper and lower
     * bounds are restricted by adding $gte/$lt clauses to the query
     * filter. This requires that the boundaries are not compound keys, and that
     * the query does not contain any keys used in the split key.
     *
     * @param chunkLowerBound the lower bound of the chunk (min)
     * @param chunkUpperBound the upper bound of the chunk (max)
     * @param query a query filtering the documents within the split
     * @return a MongoInputSplit from a range query
     * @throws IllegalArgumentException if the query conflicts with the chunk bounds, or the either of the bounds are compound keys.
     */
    public MongoInputSplit createRangeQuerySplit(final BasicDBObject chunkLowerBound, final BasicDBObject chunkUpperBound,
                                                 final BSONObject query) {

        //If the boundaries are actually empty, just return
        //a split without boundaries.
        if (chunkLowerBound == null && chunkUpperBound == null) {
            DBObject splitQuery = new BasicDBObject();
            splitQuery.putAll(query);
            MongoInputSplit split = new MongoInputSplit(getConfiguration());
            split.setQuery(splitQuery);
            return split;
        }

        // The boundaries are not empty, so try to build a split using $gte/$lt.

        //First check that the split contains no compound keys.
        // e.g. this is valid: { _id : "foo" }
        // but this is not {_id : "foo", name : "bar"}
        Entry minKey = chunkLowerBound != null && chunkLowerBound.keySet().size() == 1
                                           ? chunkLowerBound.entrySet().iterator().next() : null;
        Entry maxKey = chunkUpperBound != null && chunkUpperBound.keySet().size() == 1
                                           ? chunkUpperBound.entrySet().iterator().next() : null;
        if (minKey == null && maxKey == null) {
            throw new IllegalArgumentException("Range query is enabled but one or more split boundaries contains a compound key:\n"
                                               + "min:  " + chunkLowerBound + "\nmax:  " + chunkUpperBound);
        }

        //Now, check that the lower and upper bounds don't have any keys
        //which overlap with the query.
        if (minKey != null && query.containsField(minKey.getKey()) || maxKey != null && query.containsField(maxKey.getKey())) {
            throw new IllegalArgumentException("Range query is enabled but split key conflicts with query filter:\n"
                                               + "min:  " + chunkLowerBound + "\nmax:  " + chunkUpperBound + "\nquery:  " + query);
        }

        String key = null;
        BasicDBObject rangeObj = new BasicDBObject();
        if (minKey != null) {
            key = minKey.getKey();
            rangeObj.put("$gte", minKey.getValue());
        }

        if (maxKey != null) {
            key = maxKey.getKey();
            rangeObj.put("$lt", maxKey.getValue());
        }

        DBObject splitQuery = new BasicDBObject();
        splitQuery.putAll(query);
        splitQuery.put(key, rangeObj);
        MongoInputSplit split = new MongoInputSplit(getConfiguration());
        split.setQuery(splitQuery);
        return split;
    }

}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy