contrib/storage-mongo/src/main/java/org/apache/drill/exec/store/mongo/MongoRecordReader.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.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;

 import com.mongodb.client.MongoIterable;
 import com.mongodb.client.model.Aggregates;
 import org.apache.commons.collections.CollectionUtils;
 import org.apache.drill.common.exceptions.DrillRuntimeException;
 import org.apache.drill.common.exceptions.ExecutionSetupException;
 import org.apache.drill.common.expression.SchemaPath;
 import org.apache.drill.exec.ExecConstants;
 import org.apache.drill.exec.ops.FragmentContext;
 import org.apache.drill.exec.ops.OperatorContext;
 import org.apache.drill.exec.physical.impl.OutputMutator;
 import org.apache.drill.exec.store.AbstractRecordReader;
 import org.apache.drill.exec.store.bson.BsonRecordReader;
 import org.apache.drill.exec.vector.BaseValueVector;
 import org.apache.drill.exec.vector.complex.fn.JsonReader;
 import org.apache.drill.exec.vector.complex.impl.VectorContainerWriter;
 import org.bson.BsonDocument;
 import org.bson.BsonDocumentReader;
 import org.bson.Document;
 import org.bson.conversions.Bson;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 import com.google.common.base.Stopwatch;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Sets;
 import com.mongodb.client.MongoClient;
 import com.mongodb.ServerAddress;
 import com.mongodb.client.MongoCollection;
 import com.mongodb.client.MongoCursor;
 import com.mongodb.client.MongoDatabase;

 public class MongoRecordReader extends AbstractRecordReader {
   private static final Logger logger = LoggerFactory.getLogger(MongoRecordReader.class);

   private MongoCollection<BsonDocument> collection;
   private MongoCursor<BsonDocument> cursor;

   private JsonReader jsonReader;
   private BsonRecordReader bsonReader;
   private VectorContainerWriter writer;

   private Document filters;
   private List<Bson> operations;
   private final Document fields;

   private final FragmentContext fragmentContext;

   private final MongoStoragePlugin plugin;

   private final boolean enableAllTextMode;
   private final boolean enableNanInf;
   private final boolean readNumbersAsDouble;
   private boolean unionEnabled;
   private final boolean isBsonRecordReader;

   public MongoRecordReader(BaseMongoSubScanSpec subScanSpec, List<SchemaPath> projectedColumns,
       FragmentContext context, MongoStoragePlugin plugin) {

     fields = new Document();
     // exclude _id field, if not mentioned by user.
     fields.put(DrillMongoConstants.ID, 0);
     setColumns(projectedColumns);
     fragmentContext = context;
     this.plugin = plugin;
     filters = new Document();
     if (subScanSpec instanceof MongoSubScan.MongoSubScanSpec) {
       operations = ((MongoSubScan.MongoSubScanSpec) subScanSpec).getOperations().stream()
         .map(BsonDocument::parse)
         .collect(Collectors.toList());
     } else {
       MongoSubScan.ShardedMongoSubScanSpec shardedMongoSubScanSpec = (MongoSubScan.ShardedMongoSubScanSpec) subScanSpec;
       Map<String, List<Document>> mergedFilters = MongoUtils.mergeFilters(
           shardedMongoSubScanSpec.getMinFilters(), shardedMongoSubScanSpec.getMaxFilters());

       Document pushdownFilters = Optional.ofNullable(shardedMongoSubScanSpec.getFilter())
         .map(Document::parse)
         .orElse(null);
       buildFilters(pushdownFilters, mergedFilters);
     }
     enableAllTextMode = fragmentContext.getOptions().getOption(ExecConstants.MONGO_ALL_TEXT_MODE).bool_val;
     enableNanInf = fragmentContext.getOptions().getOption(ExecConstants.JSON_READER_NAN_INF_NUMBERS).bool_val;
     readNumbersAsDouble = fragmentContext.getOptions().getOption(ExecConstants.MONGO_READER_READ_NUMBERS_AS_DOUBLE).bool_val;
     isBsonRecordReader = fragmentContext.getOptions().getOption(ExecConstants.MONGO_BSON_RECORD_READER).bool_val;
     logger.debug("BsonRecordReader is enabled? " + isBsonRecordReader);
     init(subScanSpec);
   }

   @Override
   protected Collection<SchemaPath> transformColumns(Collection<SchemaPath> projectedColumns) {
     Set<SchemaPath> transformed = Sets.newLinkedHashSet();
     if (!isStarQuery()) {
       for (SchemaPath column : projectedColumns) {
         String fieldName = column.getRootSegment().getPath();
         transformed.add(column);
         this.fields.put(fieldName, 1);
       }
     } else {
       // Tale all the fields including the _id
       this.fields.remove(DrillMongoConstants.ID);
       transformed.add(SchemaPath.STAR_COLUMN);
     }
     return transformed;
   }

   private void buildFilters(Document pushdownFilters,
       Map<String, List<Document>> mergedFilters) {
     for (Entry<String, List<Document>> entry : mergedFilters.entrySet()) {
       List<Document> list = entry.getValue();
       if (list.size() == 1) {
         this.filters.putAll(list.get(0));
       } else {
         Document andQueryFilter = new Document();
         andQueryFilter.put("$and", list);
         this.filters.putAll(andQueryFilter);
       }
     }
     if (pushdownFilters != null && !pushdownFilters.isEmpty()) {
       if (!mergedFilters.isEmpty()) {
         this.filters = MongoUtils.andFilterAtIndex(this.filters, pushdownFilters);
       } else {
         this.filters = pushdownFilters;
       }
     }
   }

   private void init(BaseMongoSubScanSpec subScanSpec) {
     List<String> hosts = subScanSpec.getHosts();
     List<ServerAddress> addresses = Lists.newArrayList();
     for (String host : hosts) {
       addresses.add(new ServerAddress(host));
     }
     MongoClient client = plugin.getClient(addresses);
     MongoDatabase db = client.getDatabase(subScanSpec.getDbName());
     this.unionEnabled = fragmentContext.getOptions().getBoolean(ExecConstants.ENABLE_UNION_TYPE_KEY);
     collection = db.getCollection(subScanSpec.getCollectionName(), BsonDocument.class);
   }

   @Override
   public void setup(OperatorContext context, OutputMutator output) throws ExecutionSetupException {
     this.writer = new VectorContainerWriter(output, unionEnabled);
     // Default is BsonReader and all text mode will not be honored in
     // BsonRecordReader
     if (isBsonRecordReader) {
       this.bsonReader = new BsonRecordReader(fragmentContext.getManagedBuffer(), Lists.newArrayList(getColumns()),
           readNumbersAsDouble);
       logger.debug("Initialized BsonRecordReader. ");
     } else {
       this.jsonReader = new JsonReader.Builder(fragmentContext.getManagedBuffer())
           .schemaPathColumns(Lists.newArrayList(getColumns()))
           .allTextMode(enableAllTextMode)
           .readNumbersAsDouble(readNumbersAsDouble)
           .enableNanInf(enableNanInf)
           .build();
       logger.debug(" Intialized JsonRecordReader. ");
     }
   }

   @Override
   public int next() {
     if (cursor == null) {
       logger.debug("Filters Applied : " + filters);
       logger.debug("Fields Selected :" + fields);

       MongoIterable<BsonDocument> projection;
       if (CollectionUtils.isNotEmpty(operations)) {
         List<Bson> operations = new ArrayList<>(this.operations);
         if (!fields.isEmpty()) {
           operations.add(Aggregates.project(fields));
         }
         if (plugin.getConfig().allowDiskUse()) {
           projection = collection.aggregate(operations).allowDiskUse(true);
         } else {
           projection = collection.aggregate(operations);
         }
       } else {
         projection = collection.find(filters).projection(fields);
       }
       cursor = projection.batchSize(plugin.getConfig().getBatchSize()).iterator();
     }

     writer.allocate();
     writer.reset();

     int docCount = 0;
     Stopwatch watch = Stopwatch.createStarted();

     try {
       while (docCount < BaseValueVector.INITIAL_VALUE_ALLOCATION && cursor.hasNext()) {
         writer.setPosition(docCount);
         if (isBsonRecordReader) {
           BsonDocument bsonDocument = cursor.next();
           bsonReader.write(writer, new BsonDocumentReader(bsonDocument));
         } else {
           String doc = cursor.next().toJson();
           jsonReader.setSource(doc.getBytes(StandardCharsets.UTF_8));
           jsonReader.write(writer);
         }
         docCount++;
       }

       if (isBsonRecordReader) {
         bsonReader.ensureAtLeastOneField(writer);
       } else {
         jsonReader.ensureAtLeastOneField(writer);
       }

       writer.setValueCount(docCount);
       logger.debug("Took {} ms to get {} records", watch.elapsed(TimeUnit.MILLISECONDS), docCount);
       return docCount;
     } catch (IOException e) {
       String msg = "Failure while reading document. - Parser was at record: " + (docCount + 1);
       logger.error(msg, e);
       throw new DrillRuntimeException(msg, e);
     }
   }

   @Override
   public void close() {
   }

   @Override
   public String toString() {
     Object reader = isBsonRecordReader ? bsonReader : jsonReader;
     return "MongoRecordReader[reader=" + reader + "]";
   }
 }
	/*
	* 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.io.IOException;
	import java.nio.charset.StandardCharsets;
	import java.util.ArrayList;
	import java.util.Collection;
	import java.util.List;
	import java.util.Map;
	import java.util.Map.Entry;
	import java.util.Optional;
	import java.util.Set;
	import java.util.concurrent.TimeUnit;
	import java.util.stream.Collectors;

	import com.mongodb.client.MongoIterable;
	import com.mongodb.client.model.Aggregates;
	import org.apache.commons.collections.CollectionUtils;
	import org.apache.drill.common.exceptions.DrillRuntimeException;
	import org.apache.drill.common.exceptions.ExecutionSetupException;
	import org.apache.drill.common.expression.SchemaPath;
	import org.apache.drill.exec.ExecConstants;
	import org.apache.drill.exec.ops.FragmentContext;
	import org.apache.drill.exec.ops.OperatorContext;
	import org.apache.drill.exec.physical.impl.OutputMutator;
	import org.apache.drill.exec.store.AbstractRecordReader;
	import org.apache.drill.exec.store.bson.BsonRecordReader;
	import org.apache.drill.exec.vector.BaseValueVector;
	import org.apache.drill.exec.vector.complex.fn.JsonReader;
	import org.apache.drill.exec.vector.complex.impl.VectorContainerWriter;
	import org.bson.BsonDocument;
	import org.bson.BsonDocumentReader;
	import org.bson.Document;
	import org.bson.conversions.Bson;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	import com.google.common.base.Stopwatch;
	import com.google.common.collect.Lists;
	import com.google.common.collect.Sets;
	import com.mongodb.client.MongoClient;
	import com.mongodb.ServerAddress;
	import com.mongodb.client.MongoCollection;
	import com.mongodb.client.MongoCursor;
	import com.mongodb.client.MongoDatabase;

	public class MongoRecordReader extends AbstractRecordReader {
	private static final Logger logger = LoggerFactory.getLogger(MongoRecordReader.class);

	private MongoCollection<BsonDocument> collection;
	private MongoCursor<BsonDocument> cursor;

	private JsonReader jsonReader;
	private BsonRecordReader bsonReader;
	private VectorContainerWriter writer;

	private Document filters;
	private List<Bson> operations;
	private final Document fields;

	private final FragmentContext fragmentContext;

	private final MongoStoragePlugin plugin;

	private final boolean enableAllTextMode;
	private final boolean enableNanInf;
	private final boolean readNumbersAsDouble;
	private boolean unionEnabled;
	private final boolean isBsonRecordReader;

	public MongoRecordReader(BaseMongoSubScanSpec subScanSpec, List<SchemaPath> projectedColumns,
	FragmentContext context, MongoStoragePlugin plugin) {

	fields = new Document();
	// exclude _id field, if not mentioned by user.
	fields.put(DrillMongoConstants.ID, 0);
	setColumns(projectedColumns);
	fragmentContext = context;
	this.plugin = plugin;
	filters = new Document();
	if (subScanSpec instanceof MongoSubScan.MongoSubScanSpec) {
	operations = ((MongoSubScan.MongoSubScanSpec) subScanSpec).getOperations().stream()
	.map(BsonDocument::parse)
	.collect(Collectors.toList());
	} else {
	MongoSubScan.ShardedMongoSubScanSpec shardedMongoSubScanSpec = (MongoSubScan.ShardedMongoSubScanSpec) subScanSpec;
	Map<String, List<Document>> mergedFilters = MongoUtils.mergeFilters(
	shardedMongoSubScanSpec.getMinFilters(), shardedMongoSubScanSpec.getMaxFilters());

	Document pushdownFilters = Optional.ofNullable(shardedMongoSubScanSpec.getFilter())
	.map(Document::parse)
	.orElse(null);
	buildFilters(pushdownFilters, mergedFilters);
	}
	enableAllTextMode = fragmentContext.getOptions().getOption(ExecConstants.MONGO_ALL_TEXT_MODE).bool_val;
	enableNanInf = fragmentContext.getOptions().getOption(ExecConstants.JSON_READER_NAN_INF_NUMBERS).bool_val;
	readNumbersAsDouble = fragmentContext.getOptions().getOption(ExecConstants.MONGO_READER_READ_NUMBERS_AS_DOUBLE).bool_val;
	isBsonRecordReader = fragmentContext.getOptions().getOption(ExecConstants.MONGO_BSON_RECORD_READER).bool_val;
	logger.debug("BsonRecordReader is enabled? " + isBsonRecordReader);
	init(subScanSpec);
	}

	@Override
	protected Collection<SchemaPath> transformColumns(Collection<SchemaPath> projectedColumns) {
	Set<SchemaPath> transformed = Sets.newLinkedHashSet();
	if (!isStarQuery()) {
	for (SchemaPath column : projectedColumns) {
	String fieldName = column.getRootSegment().getPath();
	transformed.add(column);
	this.fields.put(fieldName, 1);
	}
	} else {
	// Tale all the fields including the _id
	this.fields.remove(DrillMongoConstants.ID);
	transformed.add(SchemaPath.STAR_COLUMN);
	}
	return transformed;
	}

	private void buildFilters(Document pushdownFilters,
	Map<String, List<Document>> mergedFilters) {
	for (Entry<String, List<Document>> entry : mergedFilters.entrySet()) {
	List<Document> list = entry.getValue();
	if (list.size() == 1) {
	this.filters.putAll(list.get(0));
	} else {
	Document andQueryFilter = new Document();
	andQueryFilter.put("$and", list);
	this.filters.putAll(andQueryFilter);
	}
	}
	if (pushdownFilters != null && !pushdownFilters.isEmpty()) {
	if (!mergedFilters.isEmpty()) {
	this.filters = MongoUtils.andFilterAtIndex(this.filters, pushdownFilters);
	} else {
	this.filters = pushdownFilters;
	}
	}
	}

	private void init(BaseMongoSubScanSpec subScanSpec) {
	List<String> hosts = subScanSpec.getHosts();
	List<ServerAddress> addresses = Lists.newArrayList();
	for (String host : hosts) {
	addresses.add(new ServerAddress(host));
	}
	MongoClient client = plugin.getClient(addresses);
	MongoDatabase db = client.getDatabase(subScanSpec.getDbName());
	this.unionEnabled = fragmentContext.getOptions().getBoolean(ExecConstants.ENABLE_UNION_TYPE_KEY);
	collection = db.getCollection(subScanSpec.getCollectionName(), BsonDocument.class);
	}

	@Override
	public void setup(OperatorContext context, OutputMutator output) throws ExecutionSetupException {
	this.writer = new VectorContainerWriter(output, unionEnabled);
	// Default is BsonReader and all text mode will not be honored in
	// BsonRecordReader
	if (isBsonRecordReader) {
	this.bsonReader = new BsonRecordReader(fragmentContext.getManagedBuffer(), Lists.newArrayList(getColumns()),
	readNumbersAsDouble);
	logger.debug("Initialized BsonRecordReader. ");
	} else {
	this.jsonReader = new JsonReader.Builder(fragmentContext.getManagedBuffer())
	.schemaPathColumns(Lists.newArrayList(getColumns()))
	.allTextMode(enableAllTextMode)
	.readNumbersAsDouble(readNumbersAsDouble)
	.enableNanInf(enableNanInf)
	.build();
	logger.debug(" Intialized JsonRecordReader. ");
	}
	}

	@Override
	public int next() {
	if (cursor == null) {
	logger.debug("Filters Applied : " + filters);
	logger.debug("Fields Selected :" + fields);

	MongoIterable<BsonDocument> projection;
	if (CollectionUtils.isNotEmpty(operations)) {
	List<Bson> operations = new ArrayList<>(this.operations);
	if (!fields.isEmpty()) {
	operations.add(Aggregates.project(fields));
	}
	if (plugin.getConfig().allowDiskUse()) {
	projection = collection.aggregate(operations).allowDiskUse(true);
	} else {
	projection = collection.aggregate(operations);
	}
	} else {
	projection = collection.find(filters).projection(fields);
	}
	cursor = projection.batchSize(plugin.getConfig().getBatchSize()).iterator();
	}

	writer.allocate();
	writer.reset();

	int docCount = 0;
	Stopwatch watch = Stopwatch.createStarted();

	try {
	while (docCount < BaseValueVector.INITIAL_VALUE_ALLOCATION && cursor.hasNext()) {
	writer.setPosition(docCount);
	if (isBsonRecordReader) {
	BsonDocument bsonDocument = cursor.next();
	bsonReader.write(writer, new BsonDocumentReader(bsonDocument));
	} else {
	String doc = cursor.next().toJson();
	jsonReader.setSource(doc.getBytes(StandardCharsets.UTF_8));
	jsonReader.write(writer);
	}
	docCount++;
	}

	if (isBsonRecordReader) {
	bsonReader.ensureAtLeastOneField(writer);
	} else {
	jsonReader.ensureAtLeastOneField(writer);
	}

	writer.setValueCount(docCount);
	logger.debug("Took {} ms to get {} records", watch.elapsed(TimeUnit.MILLISECONDS), docCount);
	return docCount;
	} catch (IOException e) {
	String msg = "Failure while reading document. - Parser was at record: " + (docCount + 1);
	logger.error(msg, e);
	throw new DrillRuntimeException(msg, e);
	}
	}

	@Override
	public void close() {
	}

	@Override
	public String toString() {
	Object reader = isBsonRecordReader ? bsonReader : jsonReader;
	return "MongoRecordReader[reader=" + reader + "]";
	}
	}