contrib/storage-mongo/src/main/java/org/apache/drill/exec/store/mongo/MongoGroupScan.java - drill - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you 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 org.apache.drill.exec.store.mongo;

 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Comparator;
 import java.util.Iterator;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Objects;
 import java.util.PriorityQueue;
 import java.util.Queue;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;

 import com.fasterxml.jackson.annotation.JacksonInject;
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.annotation.JsonTypeName;
 import com.mongodb.client.MongoClient;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.drill.common.PlanStringBuilder;
 import org.apache.drill.common.exceptions.DrillRuntimeException;
 import org.apache.drill.common.expression.SchemaPath;
 import org.apache.drill.exec.physical.EndpointAffinity;
 import org.apache.drill.exec.physical.base.AbstractGroupScan;
 import org.apache.drill.exec.physical.base.GroupScan;
 import org.apache.drill.exec.physical.base.PhysicalOperator;
 import org.apache.drill.exec.physical.base.ScanStats;
 import org.apache.drill.exec.physical.base.ScanStats.GroupScanProperty;
 import org.apache.drill.exec.proto.CoordinationProtos.DrillbitEndpoint;
 import org.apache.drill.exec.store.StoragePluginRegistry;
 import org.apache.drill.exec.store.mongo.MongoSubScan.MongoSubScanSpec;
 import org.apache.drill.exec.store.mongo.MongoSubScan.ShardedMongoSubScanSpec;
 import org.apache.drill.exec.store.mongo.common.ChunkInfo;
 import org.bson.Document;
 import org.bson.codecs.BsonTypeClassMap;
 import org.bson.codecs.DocumentCodec;
 import org.bson.conversions.Bson;
 import org.bson.types.MaxKey;
 import org.bson.types.MinKey;

 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Joiner;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Stopwatch;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
 import com.mongodb.ReadPreference;
 import com.mongodb.ServerAddress;
 import com.mongodb.client.FindIterable;
 import com.mongodb.client.MongoCollection;
 import com.mongodb.client.MongoCursor;
 import com.mongodb.client.MongoDatabase;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 @JsonTypeName("mongo-scan")
 public class MongoGroupScan extends AbstractGroupScan implements
     DrillMongoConstants {

   private static final int SELECT = 1;

   private static final Comparator<List<BaseMongoSubScanSpec>> LIST_SIZE_COMPARATOR = Comparator.comparingInt(List::size);

   private static final Comparator<List<BaseMongoSubScanSpec>> LIST_SIZE_COMPARATOR_REV = Collections.reverseOrder(LIST_SIZE_COMPARATOR);

   private static final Logger logger = LoggerFactory.getLogger(MongoGroupScan.class);

   private MongoStoragePlugin storagePlugin;

   private MongoStoragePluginConfig storagePluginConfig;

   private MongoScanSpec scanSpec;

   private List<SchemaPath> columns;

   private Map<Integer, List<BaseMongoSubScanSpec>> endpointFragmentMapping;

   // Sharding with replica sets contains all the replica server addresses for
   // each chunk.
   private Map<String, Set<ServerAddress>> chunksMapping;

   private Map<String, List<ChunkInfo>> chunksInverseMapping;

   private final Stopwatch watch = Stopwatch.createUnstarted();

   private boolean useAggregate;

   @JsonCreator
   public MongoGroupScan(
       @JsonProperty("userName") String userName,
       @JsonProperty("mongoScanSpec") MongoScanSpec scanSpec,
       @JsonProperty("storage") MongoStoragePluginConfig storagePluginConfig,
       @JsonProperty("columns") List<SchemaPath> columns,
       @JsonProperty("useAggregate") boolean useAggregate,
       @JacksonInject StoragePluginRegistry pluginRegistry) {
     this(userName,
         pluginRegistry.resolve(storagePluginConfig, MongoStoragePlugin.class),
         scanSpec, columns, useAggregate);
   }

   public MongoGroupScan(String userName, MongoStoragePlugin storagePlugin,
       MongoScanSpec scanSpec, List<SchemaPath> columns, boolean useAggregate) {
     super(userName);
     this.storagePlugin = storagePlugin;
     this.storagePluginConfig = storagePlugin.getConfig();
     this.scanSpec = scanSpec;
     this.columns = columns;
     this.useAggregate = useAggregate;
     init();
   }


   /**
    * Private constructor, used for cloning.
    * @param that
    *          The MongoGroupScan to clone
    */
   private MongoGroupScan(MongoGroupScan that) {
     super(that);
     this.scanSpec = that.scanSpec;
     this.columns = that.columns;
     this.storagePlugin = that.storagePlugin;
     this.storagePluginConfig = that.storagePluginConfig;
     this.chunksMapping = that.chunksMapping;
     this.chunksInverseMapping = that.chunksInverseMapping;
     this.endpointFragmentMapping = that.endpointFragmentMapping;
     this.useAggregate = that.useAggregate;
   }

   private boolean isShardedCluster(MongoClient client) {
     MongoDatabase db = client.getDatabase(scanSpec.getDbName());
     String msg = db.runCommand(new Document("isMaster", 1)).getString("msg");
     return msg != null && msg.equals("isdbgrid");
   }

   @SuppressWarnings({ "rawtypes" })
   private void init() {

     List<String> h = storagePluginConfig.getHosts();
     List<ServerAddress> addresses = Lists.newArrayList();
     for (String host : h) {
       addresses.add(new ServerAddress(host));
     }
     MongoClient client = storagePlugin.getClient();
     chunksMapping = Maps.newHashMap();
     chunksInverseMapping = Maps.newLinkedHashMap();
     if (useAggregate && isShardedCluster(client)) {
       handleUnshardedCollection(getPrimaryShardInfo());
     } else if (isShardedCluster(client)) {
       MongoDatabase db = client.getDatabase(CONFIG);
       MongoCollection<Document> chunksCollection = db.getCollection(CHUNKS);
       Document filter = new Document();
       filter
           .put(
               NS,
               this.scanSpec.getDbName() + "."
                   + this.scanSpec.getCollectionName());

       Document projection = new Document();
       projection.put(SHARD, SELECT);
       projection.put(MIN, SELECT);
       projection.put(MAX, SELECT);

       FindIterable<Document> chunkCursor = chunksCollection.find(filter).projection(projection);
       MongoCursor<Document> iterator = chunkCursor.iterator();

       MongoCollection<Document> shardsCollection = db.getCollection(SHARDS);

       projection = new Document();
       projection.put(HOST, SELECT);

       boolean hasChunks = false;
       while (iterator.hasNext()) {
         Document chunkObj = iterator.next();
         String shardName = (String) chunkObj.get(SHARD);
         // creates hexadecimal string representation of ObjectId
         String chunkId = chunkObj.get(ID).toString();
         filter = new Document(ID, shardName);
         FindIterable<Document> hostCursor = shardsCollection.find(filter).projection(projection);
         for (Document hostObj : hostCursor) {
           String hostEntry = (String) hostObj.get(HOST);
           String[] tagAndHost = StringUtils.split(hostEntry, '/');
           String[] hosts = tagAndHost.length > 1 ? StringUtils.split(
               tagAndHost[1], ',') : StringUtils.split(tagAndHost[0], ',');
           Set<ServerAddress> addressList = getPreferredHosts(storagePlugin.getClient(addresses));
           if (addressList == null) {
             addressList = Sets.newHashSet();
             for (String host : hosts) {
               addressList.add(new ServerAddress(host));
             }
           }
           chunksMapping.put(chunkId, addressList);
           ServerAddress address = addressList.iterator().next();
           List<ChunkInfo> chunkList = chunksInverseMapping.computeIfAbsent(address.getHost(), k -> new ArrayList<>());
           List<String> chunkHostsList = new ArrayList<>();
           for (ServerAddress serverAddr : addressList) {
             chunkHostsList.add(serverAddr.toString());
           }
           ChunkInfo chunkInfo = new ChunkInfo(chunkHostsList, chunkId);
           Document minMap = (Document) chunkObj.get(MIN);

           Map<String, Object> minFilters = Maps.newHashMap();
           Set keySet = minMap.keySet();
           for (Object keyObj : keySet) {
             Object object = minMap.get(keyObj);
             if (!(object instanceof MinKey)) {
               minFilters.put(keyObj.toString(), object);
             }
           }
           chunkInfo.setMinFilters(minFilters);

           Map<String, Object> maxFilters = Maps.newHashMap();
           Map maxMap = (Document) chunkObj.get(MAX);
           keySet = maxMap.keySet();
           for (Object keyObj : keySet) {
             Object object = maxMap.get(keyObj);
             if (!(object instanceof MaxKey)) {
               maxFilters.put(keyObj.toString(), object);
             }
           }

           chunkInfo.setMaxFilters(maxFilters);
           chunkList.add(chunkInfo);
         }
         hasChunks = true;
       }
       // In a sharded environment, if a collection doesn't have any chunks, it is considered as an
       // unsharded collection and it will be stored in the primary shard of that database.
       if (!hasChunks) {
         handleUnshardedCollection(getPrimaryShardInfo());
       }
     } else {
       handleUnshardedCollection(storagePluginConfig.getHosts());
     }

   }

   private void handleUnshardedCollection(List<String> hosts) {
     String chunkName = Joiner.on('.').join(scanSpec.getDbName(), scanSpec.getCollectionName());
     Set<ServerAddress> addressList = Sets.newHashSet();

     for (String host : hosts) {
       addressList.add(new ServerAddress(host));
     }
     chunksMapping.put(chunkName, addressList);

     String host = hosts.get(0);
     ServerAddress address = new ServerAddress(host);
     ChunkInfo chunkInfo = new ChunkInfo(hosts, chunkName);
     chunkInfo.setMinFilters(Collections.emptyMap());
     chunkInfo.setMaxFilters(Collections.emptyMap());
     List<ChunkInfo> chunksList = Lists.newArrayList();
     chunksList.add(chunkInfo);
     chunksInverseMapping.put(address.getHost(), chunksList);
   }

   private List<String> getPrimaryShardInfo() {
     MongoDatabase database = storagePlugin.getClient().getDatabase(CONFIG);
     //Identify the primary shard of the queried database.
     MongoCollection<Document> collection = database.getCollection(DATABASES);
     Bson filter = new Document(ID, this.scanSpec.getDbName());
     Bson projection = new Document(PRIMARY, SELECT);
     Document document = Objects.requireNonNull(collection.find(filter).projection(projection).first());
     String shardName = document.getString(PRIMARY);
     Preconditions.checkNotNull(shardName);

     //Identify the host(s) on which this shard resides.
     MongoCollection<Document> shardsCol = database.getCollection(SHARDS);
     filter = new Document(ID, shardName);
     projection = new Document(HOST, SELECT);
     Document hostInfo = Objects.requireNonNull(shardsCol.find(filter).projection(projection).first());
     String hostEntry = hostInfo.getString(HOST);
     Preconditions.checkNotNull(hostEntry);

     String[] tagAndHost = StringUtils.split(hostEntry, '/');
     String[] hosts = tagAndHost.length > 1 ? StringUtils.split(tagAndHost[1],
         ',') : StringUtils.split(tagAndHost[0], ',');
     return Lists.newArrayList(hosts);
   }

   @SuppressWarnings("unchecked")
   private Set<ServerAddress> getPreferredHosts(MongoClient client) {
     Set<ServerAddress> addressList = Sets.newHashSet();
     MongoDatabase db = client.getDatabase(scanSpec.getDbName());
     ReadPreference readPreference = db.getReadPreference();
     Document command = db.runCommand(new Document("isMaster", 1));

     final String primaryHost = command.getString("primary");
     final List<String> hostsList = (List<String>) command.get("hosts");

     switch (readPreference.getName().toUpperCase()) {
     case "PRIMARY":
     case "PRIMARYPREFERRED":
       if (primaryHost == null) {
         return null;
       }
       addressList.add(new ServerAddress(primaryHost));
       return addressList;
     case "SECONDARY":
     case "SECONDARYPREFERRED":
       if (primaryHost == null || hostsList == null) {
         return null;
       }
       hostsList.remove(primaryHost);
       for (String host : hostsList) {
         addressList.add(new ServerAddress(host));
       }
       return addressList;
     case "NEAREST":
       if (hostsList == null) {
         return null;
       }
       for (String host : hostsList) {
         addressList.add(new ServerAddress(host));
       }
       return addressList;
     default:
       return null;
     }
   }

   @Override
   public GroupScan clone(List<SchemaPath> columns) {
     MongoGroupScan clone = new MongoGroupScan(this);
     clone.columns = columns;
     return clone;
   }

   public GroupScan clone(int maxRecords) {
     MongoGroupScan clone = new MongoGroupScan(this);
     clone.useAggregate = true;
     clone.getScanSpec().getOperations().add(new Document("$limit", maxRecords).toJson());
     return clone;
   }

   @Override
   public boolean canPushdownProjects(List<SchemaPath> columns) {
     return true;
   }

   @Override
   public void applyAssignments(List<DrillbitEndpoint> endpoints) {
     logger.debug("Incoming endpoints :" + endpoints);
     watch.reset();
     watch.start();

     final int numSlots = endpoints.size();
     int totalAssignmentsTobeDone = chunksMapping.size();

     Preconditions.checkArgument(numSlots <= totalAssignmentsTobeDone, String
         .format("Incoming endpoints %d is greater than number of chunks %d",
             numSlots, totalAssignmentsTobeDone));

     final int minPerEndpointSlot = (int) Math
         .floor((double) totalAssignmentsTobeDone / numSlots);
     final int maxPerEndpointSlot = (int) Math
         .ceil((double) totalAssignmentsTobeDone / numSlots);

     endpointFragmentMapping = Maps.newHashMapWithExpectedSize(numSlots);
     Map<String, Queue<Integer>> endpointHostIndexListMap = Maps.newHashMap();

     for (int i = 0; i < numSlots; ++i) {
       endpointFragmentMapping.put(i, new ArrayList<>(
           maxPerEndpointSlot));
       String hostname = endpoints.get(i).getAddress();
       Queue<Integer> hostIndexQueue = endpointHostIndexListMap.computeIfAbsent(hostname, k -> new LinkedList<>());
       hostIndexQueue.add(i);
     }

     Set<Entry<String, List<ChunkInfo>>> chunksToAssignSet = Sets
         .newHashSet(chunksInverseMapping.entrySet());

     for (Iterator<Entry<String, List<ChunkInfo>>> chunksIterator = chunksToAssignSet
         .iterator(); chunksIterator.hasNext();) {
       Entry<String, List<ChunkInfo>> chunkEntry = chunksIterator.next();
       Queue<Integer> slots = endpointHostIndexListMap.get(chunkEntry.getKey());
       if (slots != null) {
         for (ChunkInfo chunkInfo : chunkEntry.getValue()) {
           Integer slotIndex = slots.poll();
           List<BaseMongoSubScanSpec> subScanSpecList = endpointFragmentMapping
               .get(slotIndex);
           subScanSpecList.add(buildSubScanSpecAndGet(chunkInfo));
           slots.offer(slotIndex);
         }
         chunksIterator.remove();
       }
     }

     PriorityQueue<List<BaseMongoSubScanSpec>> minHeap = new PriorityQueue<>(
         numSlots, LIST_SIZE_COMPARATOR);
     PriorityQueue<List<BaseMongoSubScanSpec>> maxHeap = new PriorityQueue<>(
         numSlots, LIST_SIZE_COMPARATOR_REV);
     for (List<BaseMongoSubScanSpec> listOfScan : endpointFragmentMapping.values()) {
       if (listOfScan.size() < minPerEndpointSlot) {
         minHeap.offer(listOfScan);
       } else if (listOfScan.size() > minPerEndpointSlot) {
         maxHeap.offer(listOfScan);
       }
     }

     if (chunksToAssignSet.size() > 0) {
       for (Entry<String, List<ChunkInfo>> chunkEntry : chunksToAssignSet) {
         for (ChunkInfo chunkInfo : chunkEntry.getValue()) {
           List<BaseMongoSubScanSpec> smallestList = minHeap.poll();
           smallestList.add(buildSubScanSpecAndGet(chunkInfo));
           minHeap.offer(smallestList);
         }
       }
     }

     while (minHeap.peek() != null && minHeap.peek().size() < minPerEndpointSlot) {
       List<BaseMongoSubScanSpec> smallestList = minHeap.poll();
       List<BaseMongoSubScanSpec> largestList = maxHeap.poll();
       smallestList.add(largestList.remove(largestList.size() - 1));
       if (largestList.size() > minPerEndpointSlot) {
         maxHeap.offer(largestList);
       }
       if (smallestList.size() < minPerEndpointSlot) {
         minHeap.offer(smallestList);
       }
     }

     logger.debug(
         "Built assignment map in {} µs.\nEndpoints: {}.\nAssignment Map: {}",
         watch.elapsed(TimeUnit.NANOSECONDS) / 1000, endpoints,
         endpointFragmentMapping.toString());
   }

   private BaseMongoSubScanSpec buildSubScanSpecAndGet(ChunkInfo chunkInfo) {
     if (useAggregate) {
       return MongoSubScanSpec.builder()
           .operations(scanSpec.getOperations())
           .dbName(scanSpec.getDbName())
           .collectionName(scanSpec.getCollectionName())
           .hosts(chunkInfo.getChunkLocList())
           .build();
     }
     return ShardedMongoSubScanSpec.builder()
         .minFilters(chunkInfo.getMinFilters())
         .maxFilters(chunkInfo.getMaxFilters())
         .filter(scanSpec.getFilters())
         .dbName(scanSpec.getDbName())
         .collectionName(scanSpec.getCollectionName())
         .hosts(chunkInfo.getChunkLocList())
         .build();
   }

   @Override
   public MongoSubScan getSpecificScan(int minorFragmentId) {
     return new MongoSubScan(getUserName(), storagePlugin, storagePluginConfig,
         endpointFragmentMapping.get(minorFragmentId), columns);
   }

   @Override
   public int getMaxParallelizationWidth() {
     return chunksMapping.size();
   }

   @Override
   public String getDigest() {
     return toString();
   }

   @Override
   public ScanStats getScanStats() {
     try{
       MongoClient client = storagePlugin.getClient();
       MongoDatabase db = client.getDatabase(scanSpec.getDbName());
       MongoCollection<Document> collection = db.getCollection(scanSpec.getCollectionName());
       long recordCount = collection.estimatedDocumentCount();

       float approxDiskCost = 0;
       if (recordCount != 0) {
         //toJson should use client's codec, otherwise toJson could fail on
         // some types not known to DocumentCodec, e.g. DBRef.
         DocumentCodec codec = new DocumentCodec(db.getCodecRegistry(), new BsonTypeClassMap());
         String json = collection.find().first().toJson(codec);
         approxDiskCost = json.getBytes(StandardCharsets.UTF_8).length * recordCount;
       }
       return new ScanStats(GroupScanProperty.ESTIMATED_TOTAL_COST, recordCount, 1, approxDiskCost);
     } catch (Exception e) {
       throw new DrillRuntimeException(e.getMessage(), e);
     }
   }

   @Override
   public PhysicalOperator getNewWithChildren(List<PhysicalOperator> children) {
     Preconditions.checkArgument(children.isEmpty());
     return new MongoGroupScan(this);
   }

   @Override
   public List<EndpointAffinity> getOperatorAffinity() {
     watch.reset();
     watch.start();

     Map<String, DrillbitEndpoint> endpointMap = Maps.newHashMap();
     for (DrillbitEndpoint endpoint : storagePlugin.getContext().getBits()) {
       endpointMap.put(endpoint.getAddress(), endpoint);
       logger.debug("Endpoint address: {}", endpoint.getAddress());
     }

     Map<DrillbitEndpoint, EndpointAffinity> affinityMap = Maps.newHashMap();
     // As of now, considering only the first replica, though there may be
     // multiple replicas for each chunk.
     for (Set<ServerAddress> addressList : chunksMapping.values()) {
       // Each replica can be on multiple machines, take the first one, which
       // meets affinity.
       for (ServerAddress address : addressList) {
         DrillbitEndpoint ep = endpointMap.get(address.getHost());
         if (ep != null) {
           EndpointAffinity affinity = affinityMap.get(ep);
           if (affinity == null) {
             affinityMap.put(ep, new EndpointAffinity(ep, 1));
           } else {
             affinity.addAffinity(1);
           }
           break;
         }
       }
     }
     logger.debug("Took {} µs to get operator affinity",
         watch.elapsed(TimeUnit.NANOSECONDS) / 1000);
     logger.debug("Affined drillbits : " + affinityMap.values());
     return Lists.newArrayList(affinityMap.values());
   }

   @Override
   public boolean supportsLimitPushdown() {
     return true;
   }

   @Override
   public GroupScan applyLimit(int maxRecords) {
     return clone(maxRecords);
   }

   @Override
   @JsonProperty
   public List<SchemaPath> getColumns() {
     return columns;
   }

   @JsonProperty("mongoScanSpec")
   public MongoScanSpec getScanSpec() {
     return scanSpec;
   }

   @JsonProperty("storage")
   public MongoStoragePluginConfig getStorageConfig() {
     return storagePluginConfig;
   }

   @JsonIgnore
   public MongoStoragePlugin getStoragePlugin() {
     return storagePlugin;
   }

   @JsonProperty("useAggregate")
   public void setUseAggregate(boolean useAggregate) {
     this.useAggregate = useAggregate;
   }

   @JsonProperty("useAggregate")
   public boolean isUseAggregate() {
     return useAggregate;
   }

   @Override
   public String toString() {
     return new PlanStringBuilder(this)
       .field("MongoScanSpec", scanSpec)
       .field("columns", columns)
       .field("useAggregate", useAggregate)
       .toString();
   }

   @VisibleForTesting
   MongoGroupScan() {
     super((String) null);
   }

   @JsonIgnore
   @VisibleForTesting
   void setChunksMapping(Map<String, Set<ServerAddress>> chunksMapping) {
     this.chunksMapping = chunksMapping;
   }

   @JsonIgnore
   @VisibleForTesting
   void setScanSpec(MongoScanSpec scanSpec) {
     this.scanSpec = scanSpec;
   }

   @JsonIgnore
   @VisibleForTesting
   void setInverseChunksMapping(Map<String, List<ChunkInfo>> chunksInverseMapping) {
     this.chunksInverseMapping = chunksInverseMapping;
   }
 }
	/*
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you 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 org.apache.drill.exec.store.mongo;

	import java.nio.charset.StandardCharsets;
	import java.util.ArrayList;
	import java.util.Collections;
	import java.util.Comparator;
	import java.util.Iterator;
	import java.util.LinkedList;
	import java.util.List;
	import java.util.Map;
	import java.util.Map.Entry;
	import java.util.Objects;
	import java.util.PriorityQueue;
	import java.util.Queue;
	import java.util.Set;
	import java.util.concurrent.TimeUnit;

	import com.fasterxml.jackson.annotation.JacksonInject;
	import com.fasterxml.jackson.annotation.JsonCreator;
	import com.fasterxml.jackson.annotation.JsonIgnore;
	import com.fasterxml.jackson.annotation.JsonProperty;
	import com.fasterxml.jackson.annotation.JsonTypeName;
	import com.mongodb.client.MongoClient;
	import org.apache.commons.lang3.StringUtils;
	import org.apache.drill.common.PlanStringBuilder;
	import org.apache.drill.common.exceptions.DrillRuntimeException;
	import org.apache.drill.common.expression.SchemaPath;
	import org.apache.drill.exec.physical.EndpointAffinity;
	import org.apache.drill.exec.physical.base.AbstractGroupScan;
	import org.apache.drill.exec.physical.base.GroupScan;
	import org.apache.drill.exec.physical.base.PhysicalOperator;
	import org.apache.drill.exec.physical.base.ScanStats;
	import org.apache.drill.exec.physical.base.ScanStats.GroupScanProperty;
	import org.apache.drill.exec.proto.CoordinationProtos.DrillbitEndpoint;
	import org.apache.drill.exec.store.StoragePluginRegistry;
	import org.apache.drill.exec.store.mongo.MongoSubScan.MongoSubScanSpec;
	import org.apache.drill.exec.store.mongo.MongoSubScan.ShardedMongoSubScanSpec;
	import org.apache.drill.exec.store.mongo.common.ChunkInfo;
	import org.bson.Document;
	import org.bson.codecs.BsonTypeClassMap;
	import org.bson.codecs.DocumentCodec;
	import org.bson.conversions.Bson;
	import org.bson.types.MaxKey;
	import org.bson.types.MinKey;

	import com.google.common.annotations.VisibleForTesting;
	import com.google.common.base.Joiner;
	import com.google.common.base.Preconditions;
	import com.google.common.base.Stopwatch;
	import com.google.common.collect.Lists;
	import com.google.common.collect.Maps;
	import com.google.common.collect.Sets;
	import com.mongodb.ReadPreference;
	import com.mongodb.ServerAddress;
	import com.mongodb.client.FindIterable;
	import com.mongodb.client.MongoCollection;
	import com.mongodb.client.MongoCursor;
	import com.mongodb.client.MongoDatabase;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	@JsonTypeName("mongo-scan")
	public class MongoGroupScan extends AbstractGroupScan implements
	DrillMongoConstants {

	private static final int SELECT = 1;

	private static final Comparator<List<BaseMongoSubScanSpec>> LIST_SIZE_COMPARATOR = Comparator.comparingInt(List::size);

	private static final Comparator<List<BaseMongoSubScanSpec>> LIST_SIZE_COMPARATOR_REV = Collections.reverseOrder(LIST_SIZE_COMPARATOR);

	private static final Logger logger = LoggerFactory.getLogger(MongoGroupScan.class);

	private MongoStoragePlugin storagePlugin;

	private MongoStoragePluginConfig storagePluginConfig;

	private MongoScanSpec scanSpec;

	private List<SchemaPath> columns;

	private Map<Integer, List<BaseMongoSubScanSpec>> endpointFragmentMapping;

	// Sharding with replica sets contains all the replica server addresses for
	// each chunk.
	private Map<String, Set<ServerAddress>> chunksMapping;

	private Map<String, List<ChunkInfo>> chunksInverseMapping;

	private final Stopwatch watch = Stopwatch.createUnstarted();

	private boolean useAggregate;

	@JsonCreator
	public MongoGroupScan(
	@JsonProperty("userName") String userName,
	@JsonProperty("mongoScanSpec") MongoScanSpec scanSpec,
	@JsonProperty("storage") MongoStoragePluginConfig storagePluginConfig,
	@JsonProperty("columns") List<SchemaPath> columns,
	@JsonProperty("useAggregate") boolean useAggregate,
	@JacksonInject StoragePluginRegistry pluginRegistry) {
	this(userName,
	pluginRegistry.resolve(storagePluginConfig, MongoStoragePlugin.class),
	scanSpec, columns, useAggregate);
	}

	public MongoGroupScan(String userName, MongoStoragePlugin storagePlugin,
	MongoScanSpec scanSpec, List<SchemaPath> columns, boolean useAggregate) {
	super(userName);
	this.storagePlugin = storagePlugin;
	this.storagePluginConfig = storagePlugin.getConfig();
	this.scanSpec = scanSpec;
	this.columns = columns;
	this.useAggregate = useAggregate;
	init();
	}


	/**
	* Private constructor, used for cloning.
	* @param that
	* The MongoGroupScan to clone
	*/
	private MongoGroupScan(MongoGroupScan that) {
	super(that);
	this.scanSpec = that.scanSpec;
	this.columns = that.columns;
	this.storagePlugin = that.storagePlugin;
	this.storagePluginConfig = that.storagePluginConfig;
	this.chunksMapping = that.chunksMapping;
	this.chunksInverseMapping = that.chunksInverseMapping;
	this.endpointFragmentMapping = that.endpointFragmentMapping;
	this.useAggregate = that.useAggregate;
	}

	private boolean isShardedCluster(MongoClient client) {
	MongoDatabase db = client.getDatabase(scanSpec.getDbName());
	String msg = db.runCommand(new Document("isMaster", 1)).getString("msg");
	return msg != null && msg.equals("isdbgrid");
	}

	@SuppressWarnings({ "rawtypes" })
	private void init() {

	List<String> h = storagePluginConfig.getHosts();
	List<ServerAddress> addresses = Lists.newArrayList();
	for (String host : h) {
	addresses.add(new ServerAddress(host));
	}
	MongoClient client = storagePlugin.getClient();
	chunksMapping = Maps.newHashMap();
	chunksInverseMapping = Maps.newLinkedHashMap();
	if (useAggregate && isShardedCluster(client)) {
	handleUnshardedCollection(getPrimaryShardInfo());
	} else if (isShardedCluster(client)) {
	MongoDatabase db = client.getDatabase(CONFIG);
	MongoCollection<Document> chunksCollection = db.getCollection(CHUNKS);
	Document filter = new Document();
	filter
	.put(
	NS,
	this.scanSpec.getDbName() + "."
	+ this.scanSpec.getCollectionName());

	Document projection = new Document();
	projection.put(SHARD, SELECT);
	projection.put(MIN, SELECT);
	projection.put(MAX, SELECT);

	FindIterable<Document> chunkCursor = chunksCollection.find(filter).projection(projection);
	MongoCursor<Document> iterator = chunkCursor.iterator();

	MongoCollection<Document> shardsCollection = db.getCollection(SHARDS);

	projection = new Document();
	projection.put(HOST, SELECT);

	boolean hasChunks = false;
	while (iterator.hasNext()) {
	Document chunkObj = iterator.next();
	String shardName = (String) chunkObj.get(SHARD);
	// creates hexadecimal string representation of ObjectId
	String chunkId = chunkObj.get(ID).toString();
	filter = new Document(ID, shardName);
	FindIterable<Document> hostCursor = shardsCollection.find(filter).projection(projection);
	for (Document hostObj : hostCursor) {
	String hostEntry = (String) hostObj.get(HOST);
	String[] tagAndHost = StringUtils.split(hostEntry, '/');
	String[] hosts = tagAndHost.length > 1 ? StringUtils.split(
	tagAndHost[1], ',') : StringUtils.split(tagAndHost[0], ',');
	Set<ServerAddress> addressList = getPreferredHosts(storagePlugin.getClient(addresses));
	if (addressList == null) {
	addressList = Sets.newHashSet();
	for (String host : hosts) {
	addressList.add(new ServerAddress(host));
	}
	}
	chunksMapping.put(chunkId, addressList);
	ServerAddress address = addressList.iterator().next();
	List<ChunkInfo> chunkList = chunksInverseMapping.computeIfAbsent(address.getHost(), k -> new ArrayList<>());
	List<String> chunkHostsList = new ArrayList<>();
	for (ServerAddress serverAddr : addressList) {
	chunkHostsList.add(serverAddr.toString());
	}
	ChunkInfo chunkInfo = new ChunkInfo(chunkHostsList, chunkId);
	Document minMap = (Document) chunkObj.get(MIN);

	Map<String, Object> minFilters = Maps.newHashMap();
	Set keySet = minMap.keySet();
	for (Object keyObj : keySet) {
	Object object = minMap.get(keyObj);
	if (!(object instanceof MinKey)) {
	minFilters.put(keyObj.toString(), object);
	}
	}
	chunkInfo.setMinFilters(minFilters);

	Map<String, Object> maxFilters = Maps.newHashMap();
	Map maxMap = (Document) chunkObj.get(MAX);
	keySet = maxMap.keySet();
	for (Object keyObj : keySet) {
	Object object = maxMap.get(keyObj);
	if (!(object instanceof MaxKey)) {
	maxFilters.put(keyObj.toString(), object);
	}
	}

	chunkInfo.setMaxFilters(maxFilters);
	chunkList.add(chunkInfo);
	}
	hasChunks = true;
	}
	// In a sharded environment, if a collection doesn't have any chunks, it is considered as an
	// unsharded collection and it will be stored in the primary shard of that database.
	if (!hasChunks) {
	handleUnshardedCollection(getPrimaryShardInfo());
	}
	} else {
	handleUnshardedCollection(storagePluginConfig.getHosts());
	}

	}

	private void handleUnshardedCollection(List<String> hosts) {
	String chunkName = Joiner.on('.').join(scanSpec.getDbName(), scanSpec.getCollectionName());
	Set<ServerAddress> addressList = Sets.newHashSet();

	for (String host : hosts) {
	addressList.add(new ServerAddress(host));
	}
	chunksMapping.put(chunkName, addressList);

	String host = hosts.get(0);
	ServerAddress address = new ServerAddress(host);
	ChunkInfo chunkInfo = new ChunkInfo(hosts, chunkName);
	chunkInfo.setMinFilters(Collections.emptyMap());
	chunkInfo.setMaxFilters(Collections.emptyMap());
	List<ChunkInfo> chunksList = Lists.newArrayList();
	chunksList.add(chunkInfo);
	chunksInverseMapping.put(address.getHost(), chunksList);
	}

	private List<String> getPrimaryShardInfo() {
	MongoDatabase database = storagePlugin.getClient().getDatabase(CONFIG);
	//Identify the primary shard of the queried database.
	MongoCollection<Document> collection = database.getCollection(DATABASES);
	Bson filter = new Document(ID, this.scanSpec.getDbName());
	Bson projection = new Document(PRIMARY, SELECT);
	Document document = Objects.requireNonNull(collection.find(filter).projection(projection).first());
	String shardName = document.getString(PRIMARY);
	Preconditions.checkNotNull(shardName);

	//Identify the host(s) on which this shard resides.
	MongoCollection<Document> shardsCol = database.getCollection(SHARDS);
	filter = new Document(ID, shardName);
	projection = new Document(HOST, SELECT);
	Document hostInfo = Objects.requireNonNull(shardsCol.find(filter).projection(projection).first());
	String hostEntry = hostInfo.getString(HOST);
	Preconditions.checkNotNull(hostEntry);

	String[] tagAndHost = StringUtils.split(hostEntry, '/');
	String[] hosts = tagAndHost.length > 1 ? StringUtils.split(tagAndHost[1],
	',') : StringUtils.split(tagAndHost[0], ',');
	return Lists.newArrayList(hosts);
	}

	@SuppressWarnings("unchecked")
	private Set<ServerAddress> getPreferredHosts(MongoClient client) {
	Set<ServerAddress> addressList = Sets.newHashSet();
	MongoDatabase db = client.getDatabase(scanSpec.getDbName());
	ReadPreference readPreference = db.getReadPreference();
	Document command = db.runCommand(new Document("isMaster", 1));

	final String primaryHost = command.getString("primary");
	final List<String> hostsList = (List<String>) command.get("hosts");

	switch (readPreference.getName().toUpperCase()) {
	case "PRIMARY":
	case "PRIMARYPREFERRED":
	if (primaryHost == null) {
	return null;
	}
	addressList.add(new ServerAddress(primaryHost));
	return addressList;
	case "SECONDARY":
	case "SECONDARYPREFERRED":
	if (primaryHost == null \|\| hostsList == null) {
	return null;
	}
	hostsList.remove(primaryHost);
	for (String host : hostsList) {
	addressList.add(new ServerAddress(host));
	}
	return addressList;
	case "NEAREST":
	if (hostsList == null) {
	return null;
	}
	for (String host : hostsList) {
	addressList.add(new ServerAddress(host));
	}
	return addressList;
	default:
	return null;
	}
	}

	@Override
	public GroupScan clone(List<SchemaPath> columns) {
	MongoGroupScan clone = new MongoGroupScan(this);
	clone.columns = columns;
	return clone;
	}

	public GroupScan clone(int maxRecords) {
	MongoGroupScan clone = new MongoGroupScan(this);
	clone.useAggregate = true;
	clone.getScanSpec().getOperations().add(new Document("$limit", maxRecords).toJson());
	return clone;
	}

	@Override
	public boolean canPushdownProjects(List<SchemaPath> columns) {
	return true;
	}

	@Override
	public void applyAssignments(List<DrillbitEndpoint> endpoints) {
	logger.debug("Incoming endpoints :" + endpoints);
	watch.reset();
	watch.start();

	final int numSlots = endpoints.size();
	int totalAssignmentsTobeDone = chunksMapping.size();

	Preconditions.checkArgument(numSlots <= totalAssignmentsTobeDone, String
	.format("Incoming endpoints %d is greater than number of chunks %d",
	numSlots, totalAssignmentsTobeDone));

	final int minPerEndpointSlot = (int) Math
	.floor((double) totalAssignmentsTobeDone / numSlots);
	final int maxPerEndpointSlot = (int) Math
	.ceil((double) totalAssignmentsTobeDone / numSlots);

	endpointFragmentMapping = Maps.newHashMapWithExpectedSize(numSlots);
	Map<String, Queue<Integer>> endpointHostIndexListMap = Maps.newHashMap();

	for (int i = 0; i < numSlots; ++i) {
	endpointFragmentMapping.put(i, new ArrayList<>(
	maxPerEndpointSlot));
	String hostname = endpoints.get(i).getAddress();
	Queue<Integer> hostIndexQueue = endpointHostIndexListMap.computeIfAbsent(hostname, k -> new LinkedList<>());
	hostIndexQueue.add(i);
	}

	Set<Entry<String, List<ChunkInfo>>> chunksToAssignSet = Sets
	.newHashSet(chunksInverseMapping.entrySet());

	for (Iterator<Entry<String, List<ChunkInfo>>> chunksIterator = chunksToAssignSet
	.iterator(); chunksIterator.hasNext();) {
	Entry<String, List<ChunkInfo>> chunkEntry = chunksIterator.next();
	Queue<Integer> slots = endpointHostIndexListMap.get(chunkEntry.getKey());
	if (slots != null) {
	for (ChunkInfo chunkInfo : chunkEntry.getValue()) {
	Integer slotIndex = slots.poll();
	List<BaseMongoSubScanSpec> subScanSpecList = endpointFragmentMapping
	.get(slotIndex);
	subScanSpecList.add(buildSubScanSpecAndGet(chunkInfo));
	slots.offer(slotIndex);
	}
	chunksIterator.remove();
	}
	}

	PriorityQueue<List<BaseMongoSubScanSpec>> minHeap = new PriorityQueue<>(
	numSlots, LIST_SIZE_COMPARATOR);
	PriorityQueue<List<BaseMongoSubScanSpec>> maxHeap = new PriorityQueue<>(
	numSlots, LIST_SIZE_COMPARATOR_REV);
	for (List<BaseMongoSubScanSpec> listOfScan : endpointFragmentMapping.values()) {
	if (listOfScan.size() < minPerEndpointSlot) {
	minHeap.offer(listOfScan);
	} else if (listOfScan.size() > minPerEndpointSlot) {
	maxHeap.offer(listOfScan);
	}
	}

	if (chunksToAssignSet.size() > 0) {
	for (Entry<String, List<ChunkInfo>> chunkEntry : chunksToAssignSet) {
	for (ChunkInfo chunkInfo : chunkEntry.getValue()) {
	List<BaseMongoSubScanSpec> smallestList = minHeap.poll();
	smallestList.add(buildSubScanSpecAndGet(chunkInfo));
	minHeap.offer(smallestList);
	}
	}
	}

	while (minHeap.peek() != null && minHeap.peek().size() < minPerEndpointSlot) {
	List<BaseMongoSubScanSpec> smallestList = minHeap.poll();
	List<BaseMongoSubScanSpec> largestList = maxHeap.poll();
	smallestList.add(largestList.remove(largestList.size() - 1));
	if (largestList.size() > minPerEndpointSlot) {
	maxHeap.offer(largestList);
	}
	if (smallestList.size() < minPerEndpointSlot) {
	minHeap.offer(smallestList);
	}
	}

	logger.debug(
	"Built assignment map in {} µs.\nEndpoints: {}.\nAssignment Map: {}",
	watch.elapsed(TimeUnit.NANOSECONDS) / 1000, endpoints,
	endpointFragmentMapping.toString());
	}

	private BaseMongoSubScanSpec buildSubScanSpecAndGet(ChunkInfo chunkInfo) {
	if (useAggregate) {
	return MongoSubScanSpec.builder()
	.operations(scanSpec.getOperations())
	.dbName(scanSpec.getDbName())
	.collectionName(scanSpec.getCollectionName())
	.hosts(chunkInfo.getChunkLocList())
	.build();
	}
	return ShardedMongoSubScanSpec.builder()
	.minFilters(chunkInfo.getMinFilters())
	.maxFilters(chunkInfo.getMaxFilters())
	.filter(scanSpec.getFilters())
	.dbName(scanSpec.getDbName())
	.collectionName(scanSpec.getCollectionName())
	.hosts(chunkInfo.getChunkLocList())
	.build();
	}

	@Override
	public MongoSubScan getSpecificScan(int minorFragmentId) {
	return new MongoSubScan(getUserName(), storagePlugin, storagePluginConfig,
	endpointFragmentMapping.get(minorFragmentId), columns);
	}

	@Override
	public int getMaxParallelizationWidth() {
	return chunksMapping.size();
	}

	@Override
	public String getDigest() {
	return toString();
	}

	@Override
	public ScanStats getScanStats() {
	try{
	MongoClient client = storagePlugin.getClient();
	MongoDatabase db = client.getDatabase(scanSpec.getDbName());
	MongoCollection<Document> collection = db.getCollection(scanSpec.getCollectionName());
	long recordCount = collection.estimatedDocumentCount();

	float approxDiskCost = 0;
	if (recordCount != 0) {
	//toJson should use client's codec, otherwise toJson could fail on
	// some types not known to DocumentCodec, e.g. DBRef.
	DocumentCodec codec = new DocumentCodec(db.getCodecRegistry(), new BsonTypeClassMap());
	String json = collection.find().first().toJson(codec);
	approxDiskCost = json.getBytes(StandardCharsets.UTF_8).length * recordCount;
	}
	return new ScanStats(GroupScanProperty.ESTIMATED_TOTAL_COST, recordCount, 1, approxDiskCost);
	} catch (Exception e) {
	throw new DrillRuntimeException(e.getMessage(), e);
	}
	}

	@Override
	public PhysicalOperator getNewWithChildren(List<PhysicalOperator> children) {
	Preconditions.checkArgument(children.isEmpty());
	return new MongoGroupScan(this);
	}

	@Override
	public List<EndpointAffinity> getOperatorAffinity() {
	watch.reset();
	watch.start();

	Map<String, DrillbitEndpoint> endpointMap = Maps.newHashMap();
	for (DrillbitEndpoint endpoint : storagePlugin.getContext().getBits()) {
	endpointMap.put(endpoint.getAddress(), endpoint);
	logger.debug("Endpoint address: {}", endpoint.getAddress());
	}

	Map<DrillbitEndpoint, EndpointAffinity> affinityMap = Maps.newHashMap();
	// As of now, considering only the first replica, though there may be
	// multiple replicas for each chunk.
	for (Set<ServerAddress> addressList : chunksMapping.values()) {
	// Each replica can be on multiple machines, take the first one, which
	// meets affinity.
	for (ServerAddress address : addressList) {
	DrillbitEndpoint ep = endpointMap.get(address.getHost());
	if (ep != null) {
	EndpointAffinity affinity = affinityMap.get(ep);
	if (affinity == null) {
	affinityMap.put(ep, new EndpointAffinity(ep, 1));
	} else {
	affinity.addAffinity(1);
	}
	break;
	}
	}
	}
	logger.debug("Took {} µs to get operator affinity",
	watch.elapsed(TimeUnit.NANOSECONDS) / 1000);
	logger.debug("Affined drillbits : " + affinityMap.values());
	return Lists.newArrayList(affinityMap.values());
	}

	@Override
	public boolean supportsLimitPushdown() {
	return true;
	}

	@Override
	public GroupScan applyLimit(int maxRecords) {
	return clone(maxRecords);
	}

	@Override
	@JsonProperty
	public List<SchemaPath> getColumns() {
	return columns;
	}

	@JsonProperty("mongoScanSpec")
	public MongoScanSpec getScanSpec() {
	return scanSpec;
	}

	@JsonProperty("storage")
	public MongoStoragePluginConfig getStorageConfig() {
	return storagePluginConfig;
	}

	@JsonIgnore
	public MongoStoragePlugin getStoragePlugin() {
	return storagePlugin;
	}

	@JsonProperty("useAggregate")
	public void setUseAggregate(boolean useAggregate) {
	this.useAggregate = useAggregate;
	}

	@JsonProperty("useAggregate")
	public boolean isUseAggregate() {
	return useAggregate;
	}

	@Override
	public String toString() {
	return new PlanStringBuilder(this)
	.field("MongoScanSpec", scanSpec)
	.field("columns", columns)
	.field("useAggregate", useAggregate)
	.toString();
	}

	@VisibleForTesting
	MongoGroupScan() {
	super((String) null);
	}

	@JsonIgnore
	@VisibleForTesting
	void setChunksMapping(Map<String, Set<ServerAddress>> chunksMapping) {
	this.chunksMapping = chunksMapping;
	}

	@JsonIgnore
	@VisibleForTesting
	void setScanSpec(MongoScanSpec scanSpec) {
	this.scanSpec = scanSpec;
	}

	@JsonIgnore
	@VisibleForTesting
	void setInverseChunksMapping(Map<String, List<ChunkInfo>> chunksInverseMapping) {
	this.chunksInverseMapping = chunksInverseMapping;
	}
	}