Example usage for com.mongodb.client AggregateIterable batchSize

Introduction

In this page you can find the example usage for com.mongodb.client AggregateIterable batchSize.

Prototype

AggregateIterable<TResult> batchSize(int batchSize);

Source Link

Document

Sets the number of documents to return per batch.

Usage

From source file:org.apache.nifi.processors.mongodb.RunMongoAggregation.java

License:Apache License

@Override
public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException {
    FlowFile flowFile = null;//  w  w w  .ja  v  a 2  s .  c o  m
    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<>();

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

        if (batch.size() > 0) {
            writeBatch(buildBatch(batch), 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();
        }
    }
}

---

From source file:org.apache.rya.mongodb.aggregation.PipelineResultIteration.java

License:Apache License

/**
 * Constructor./*from w  w w . jav a  2 s  .c  om*/
 * @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();
}

---

From source file:org.apache.rya.mongodb.iter.RyaStatementBindingSetCursorIterator.java

License:Apache License

private void submitBatchQuery() {
    int count = 0;
    executedRangeMap.clear();/*  www .ja  va 2 s.  com*/
    final List<Document> pipeline = new ArrayList<>();
    final List<DBObject> match = new ArrayList<>();

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

    if (match.size() > 1) {
        pipeline.add(new Document("$match", new Document("$or", match)));
    } else if (match.size() == 1) {
        pipeline.add(new Document("$match", match.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.info(pipeline);

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

---