Java Code Examples for com.mongodb.client.AggregateIterable#iterator()

/**
 * 根据统计字段计算统计结果（gte最小值）并排序
 *
 * @param collectionName 集合名
 * @param match          match条件
 * @param field          统计字段
 * @param minCount       最小值
 * @return
 */
public LinkedHashMap<String, Integer> sortMap(String collectionName, Document match, String field, int minCount) {
    AggregateIterable<Document> aggregate = getDB().getCollection(collectionName).aggregate(
            Arrays.asList(
                    match(match)
                    , group("$" + field, Accumulators.sum("_count", 1))
                    , match(new Document("_count", new Document("$gte", minCount)))
                    , sort(new Document("_count", -1))
            )
    );

    LinkedHashMap<String, Integer> map = new LinkedHashMap<String, Integer>();
    MongoCursor<Document> iterator = aggregate.iterator();
    while (iterator.hasNext()) {
        Document next = iterator.next();
        map.put(next.getString("_id"), next.getInteger("_count"));
    }
    return map;
}
 

/**
 * Get the MongoDB cursor.
 */
private MongoCursor<Document> getCursor(int skip) {
    if (cursor == null && cursorId == 0) {
        Document query = Document.parse(config.getMongo().getQuery());
        List<Bson> pipes = new ArrayList<>(3);
        pipes.add(match(query));
        pipes.add(skip(skip));

        Optional.ofNullable(config.getMongo().getProject()).ifPresent(p -> pipes.add(project(Document.parse(p))));

        AggregateIterable<Document> aggregate = collection.aggregate(pipes)
                .allowDiskUse(true)
                .useCursor(true);

        cursor = aggregate.iterator();

        // TODO: Persist cursor ID somewhere to allow restarts.
        Optional.ofNullable(cursor.getServerCursor()).ifPresent(serverCursor -> cursorId = serverCursor.getId());
    } else if (cursor == null && cursorId != 0) {
        // TODO: Lookup cursor ID for resume.
        // Open existing cursor in case of restart??
    }

    return cursor;
}
 

@Override
public <U> Iterator<U> aggregate(final String collectionName, final Class<U> target,
                                 final AggregationOptions options,
                                 final ReadPreference readPreference) {
    LOG.debug("stages = " + stages);


    AggregateIterable<U> cursor = collection.aggregate(stages, target);
    return cursor.iterator();
}
 

@SuppressWarnings({"unchecked", "deprecation"})
@Override
public Object executeQuery(QueryProvider queryProvider) {
  List<String> queries = queryProvider.getPipelines();
  Iterator<String> iterator = queries.iterator();
  int i = 0;

  String collectionName = queryProvider.getCollectionName();
  List<Bson> pipelineStages = new ArrayList<>();
  while (iterator.hasNext()) {
    String query = iterator.next();
    LOGGER.trace("Processing query string {} for pipeline stage {}", query, i++);
    Bson document = BsonDocument.parse(query);
    pipelineStages.add(document);
  }
  // run the pipeline and return a flux.
  MongoCollection<Document> collection = mongoOperations.getCollection(collectionName);

  // execute the query
  AggregateIterable<Document> aggregateIterable = collection.aggregate(pipelineStages)
                                                            .allowDiskUse(queryProvider.isAllowDiskUse())
                                                            .maxTime(queryProvider.getMaxTimeMS(), MILLISECONDS);
  if (isMongo360OrLater) {
    // after 3.6 CURSOR MODE is mandatory
    aggregateIterable.useCursor(true);
    LOGGER.debug("Mongo 3.6 detected - will use cursor mode for aggregate output");
  }
  try (MongoCursor<Document> cursor = aggregateIterable.iterator()) {
    // e.g. a pipeline with an @Out stage would not have any return value.
    if (isVoidReturnType(queryProvider)) {
      return null;
    }
    Class outputClass = queryProvider.getOutputClass();
    if (cursor.hasNext()) {
      if (!queryProvider.isPageable() || (queryProvider.isPageable() &&
                                          List.class.isAssignableFrom(queryProvider.getMethodReturnType()))) {
        return getNonPageResults(queryProvider, cursor);
      }
      else if (queryProvider.isPageable() && isPageReturnType(queryProvider)) {
        return getPageableResults(queryProvider, cursor);
      }
    }
  }

  return null;
}
 

private Set<String> getMultiConsensusGroupIds(final MatchCollectionId collectionId,
                                              final boolean includeQGroupIds)
        throws IllegalArgumentException {

    final MongoCollection<Document> matchCollection = getExistingCollection(collectionId);

    final List<Document> pipeline = new ArrayList<>();
    pipeline.add(new Document("$match",
                              new Document("consensusSetData",
                                           new Document(QueryOperators.EXISTS, true))));

    if (includeQGroupIds) {

        // db.<matchCollection>.aggregate(
        //     [
        //         { "$match": { "consensusSetData": { "$exists": true } } },
        //         { "$group": { "_id": { "pGroupId": "$pGroupId", "qGroupId": "$qGroupId" } } }
        //     ]
        // )

        pipeline.add(new Document("$group",
                                  new Document("_id",
                                               new Document("pGroupId", "$pGroupId").
                                                       append("qGroupId", "$qGroupId"))));
    } else {

        // db.<matchCollection>.aggregate(
        //     [
        //         { "$match": { "consensusSetData": { "$exists": true } } },
        //         { "$group": { "_id": { "pGroupId": "$pGroupId" } } }
        //     ]
        // )

        pipeline.add(new Document("$group",
                                  new Document("_id",
                                               new Document("pGroupId", "$pGroupId"))));
    }

    if (LOG.isDebugEnabled()) {
        LOG.debug("getMultiConsensusGroupIds: running {}.aggregate({})",
                  MongoUtil.fullName(matchCollection),
                  MongoUtil.toJson(pipeline));
    }

    // sort and reduce to distinct set of group ids here instead of in pipeline
    final TreeSet<String> groupIdsWithMultiplePairs = new TreeSet<>();

    // mongodb java 3.0 driver notes:
    // -- need to set cursor batchSize to prevent NPE from cursor creation
    final AggregateIterable<Document> iterable = matchCollection.aggregate(pipeline).batchSize(1);
    try (final MongoCursor<Document> cursor = iterable.iterator()) {
        while (cursor.hasNext()) {
            final Document id = cursor.next().get("_id", Document.class);
            groupIdsWithMultiplePairs.add(id.getString("pGroupId"));
            if (includeQGroupIds) {
                groupIdsWithMultiplePairs.add(id.getString("qGroupId"));
            }
        }
    }

    return groupIdsWithMultiplePairs;
}
 

private <T> List<T> aggregate(String statement, Object parameter, ResultHandler handler, ReadPreference readPreference) {
  if (logger.isDebugEnabled()) {
    logger.debug("Execute 'aggregate' mongodb command. Statement '" + statement + "'.");
  }

  AggregateConfig config = (AggregateConfig) configuration.getStatement(statement);
  if (config == null) {
    throw new MongoDaoException(statement, "Aggregate statement id '" + statement + "' not found.");
  }

  String collection = config.getCollection();
  Map<String, NodeEntry> function = config.getFunction();
  NodeEntry field = config.getField();

  MongoDatabase db = getDatabase();
  MongoCollection<Document> coll = db.getCollection(collection).withReadPreference(readPreference);

  List<Document> operations = new ArrayList<Document>(function.size());
  for (Map.Entry<String, NodeEntry> entry : function.entrySet()) {
    NodeEntry ne = entry.getValue();
    Map<String, Object> op = (Map<String, Object>) ne.executorNode(config.getNamespace(), configuration, parameter);
    Document operation = new Document(op);
    operations.add(operation);
    if (logger.isDebugEnabled()) {
      logger.debug("Execute 'aggregate' mongodb command. Operation '" + operation + "'.");
    }
  }

  AggregateIterable<Document> iterable = coll.aggregate(operations);

  List<Document> list = new ArrayList<Document>();
  MongoCursor<Document> iterator = iterable.iterator();
  while (iterator.hasNext()) {
    list.add(iterator.next());
  }

  if (logger.isDebugEnabled()) {
    logger.debug("Execute 'aggregate' mongodb command. Result set '" + list + "'.");
  }

  if (handler != null) {
    handler.handleResult(new ResultContext() {
      @Override
      public Object getResultObject() {
        return list;
      }

      @Override
      public int getResultCount() {
        return list.size();
      }
    });
    return null;
  }

  List<T> result = new ArrayList<T>(list.size());
  for (Document doc : list) {
    T t = (T) helper.toResult(config.getNamespace(), field, doc);
    result.add(t);
  }
  return result;
}
 

@Override
public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException {
    FlowFile flowFile = null;
    if (context.hasIncomingConnection()) {
        flowFile = session.get();

        if (flowFile == null && context.hasNonLoopConnection()) {
            return;
        }
    }

    final String query = context.getProperty(QUERY).evaluateAttributeExpressions(flowFile).getValue();
    final Boolean allowDiskUse = context.getProperty(ALLOW_DISK_USE).asBoolean();
    final String queryAttr = context.getProperty(QUERY_ATTRIBUTE).evaluateAttributeExpressions(flowFile).getValue();
    final Integer batchSize = context.getProperty(BATCH_SIZE).asInteger();
    final Integer resultsPerFlowfile = context.getProperty(RESULTS_PER_FLOWFILE).asInteger();
    final String jsonTypeSetting = context.getProperty(JSON_TYPE).getValue();
    final String dateFormat      = context.getProperty(DATE_FORMAT).evaluateAttributeExpressions(flowFile).getValue();

    configureMapper(jsonTypeSetting, dateFormat);

    Map<String, String> attrs = new HashMap<>();
    if (queryAttr != null && queryAttr.trim().length() > 0) {
        attrs.put(queryAttr, query);
    }

    MongoCursor<Document> iter = null;

    try {
        MongoCollection<Document> collection = getCollection(context, flowFile);
        List<Bson> aggQuery = buildAggregationQuery(query);
        AggregateIterable<Document> it = collection.aggregate(aggQuery).allowDiskUse(allowDiskUse);;
        it.batchSize(batchSize != null ? batchSize : 1);

        iter = it.iterator();
        List<Document> batch = new ArrayList<>();
        Boolean doneSomething = false;

        while (iter.hasNext()) {
            batch.add(iter.next());
            if (batch.size() == resultsPerFlowfile) {
                writeBatch(buildBatch(batch), flowFile, context, session, attrs, REL_RESULTS);
                batch = new ArrayList<>();
                doneSomething |= true;
            }
        }

        if (! batch.isEmpty()) {
            // Something remains in batch list, write it to RESULT
            writeBatch(buildBatch(batch), flowFile, context, session, attrs, REL_RESULTS);
        } else if (! doneSomething) {
            // The batch list is empty and no batch was written (empty result!), so write empty string to RESULT
            writeBatch("", flowFile, context, session, attrs, REL_RESULTS);
        }

        if (flowFile != null) {
            session.transfer(flowFile, REL_ORIGINAL);
        }
    } catch (Exception e) {
        getLogger().error("Error running MongoDB aggregation query.", e);
        if (flowFile != null) {
            session.transfer(flowFile, REL_FAILURE);
        }
    } finally {
        if (iter != null) {
            iter.close();
        }
    }
}
 

/**
 * Constructor.
 * @param aggIter Iterator of documents in AggregationPipelineQueryNode's
 *  intermediate solution representation.
 * @param varToOriginalName A mapping from field names in the pipeline
 *  result documents to equivalent variable names in the original query.
 *  Where an entry does not exist for a field, the field name and variable
 *  name are assumed to be the same.
 * @param bindings A partial solution. May be empty.
 */
public PipelineResultIteration(AggregateIterable<Document> aggIter,
        Map<String, String> varToOriginalName,
        BindingSet bindings) {
    this.varToOriginalName = Preconditions.checkNotNull(varToOriginalName);
    this.bindings = Preconditions.checkNotNull(bindings);
    Preconditions.checkNotNull(aggIter);
    aggIter.batchSize(BATCH_SIZE);
    this.cursor = aggIter.iterator();
}
 

private void submitBatchQuery() {
    int count = 0;
    executedRangeMap.clear();
    final List<Bson> pipeline = new ArrayList<>();
    final List<Bson> matches = new ArrayList<>();

    while (queryIterator.hasNext() && count < QUERY_BATCH_SIZE){
        count++;
        final RyaStatement query = queryIterator.next();
        executedRangeMap.putAll(query, rangeMap.get(query));
        final Document currentQuery = strategy.getQuery(query);
        matches.add(currentQuery);
    }

    final int numMatches = matches.size();
    if (numMatches > 1) {
        pipeline.add(Aggregates.match(Filters.or(matches)));
    } else if (numMatches == 1) {
        pipeline.add(Aggregates.match(matches.get(0)));
    } else {
        batchQueryResultsIterator = Iterators.emptyIterator();
        return;
    }

    // Executing redact aggregation to only return documents the user has access to.
    pipeline.addAll(AggregationUtil.createRedactPipeline(auths));
    log.trace(pipeline);

    final AggregateIterable<Document> aggIter = coll.aggregate(pipeline);
    aggIter.batchSize(1000);
    batchQueryResultsIterator = aggIter.iterator();
}