[dbnode] Fix AggregateQuery limits

src/dbnode/storage/index/aggregate_results.go

@@ -123,6 +123,7 @@ func (r *aggregatedResults) AggregateResultsOptions() AggregateResultsOptions {
 
 func (r *aggregatedResults) AddFields(batch []AggregateResultsEntry) (int, int) {
 	r.Lock()
+	maxInsertions := r.aggregateOpts.SizeLimit - r.totalDocsCount
 	valueInsertions := 0
 	for _, entry := range batch {
 		f := entry.Field
@@ -145,11 +146,17 @@ func (r *aggregatedResults) AddFields(batch []AggregateResultsEntry) (int, int)
 			if !valuesMap.Contains(t) {
 				// we can avoid the copy because we assume ownership of the passed ident.ID,
 				// but still need to finalize it.
-				valuesMap.SetUnsafe(t, struct{}{}, AggregateValuesMapSetUnsafeOptions{
-					NoCopyKey:     true,
-					NoFinalizeKey: false,
-				})
-				valueInsertions++
+				if maxInsertions > valueInsertions {
+					valuesMap.SetUnsafe(t, struct{}{}, AggregateValuesMapSetUnsafeOptions{
+						NoCopyKey:     true,
+						NoFinalizeKey: false,
+					})
+					valueInsertions++
+				} else {
+					// this value exceeds the limit, so should be released to the underling
+					// pool without adding to the map.
+					t.Finalize()
+				}
 			} else {
 				// because we already have a entry for this term, we release the ident back to
 				// the underlying pool.

src/dbnode/storage/index/block.go

@@ -641,15 +641,16 @@ func (b *block) aggregateWithSpan(
 	}
 
 	var (
-		source     = opts.Source
-		size       = results.Size()
-		docsCount  = results.TotalDocsCount()
-		batch      = b.opts.AggregateResultsEntryArrayPool().Get()
-		batchSize  = cap(batch)
-		iterClosed = false // tracking whether we need to free the iterator at the end.
+		source      = opts.Source
+		size        = results.Size()
+		resultCount = results.TotalDocsCount()
+		batch       = AggregateResultsEntries(b.opts.AggregateResultsEntryArrayPool().Get())
+		maxBatch    = cap(batch)
+		iterClosed  = false // tracking whether we need to free the iterator at the end.
+		exhaustive  = false
 	)
-	if batchSize == 0 {
-		batchSize = defaultAggregateResultsEntryBatchSize
+	if maxBatch == 0 {
+		maxBatch = defaultAggregateResultsEntryBatchSize
 	}
 
 	// cleanup at the end
@@ -675,8 +676,17 @@ func (b *block) aggregateWithSpan(
 		}))
 	}
 
+	if opts.SeriesLimit < maxBatch {
+		maxBatch = opts.SeriesLimit
+	}
+
+	if opts.DocsLimit < maxBatch {
+		maxBatch = opts.DocsLimit
+	}
+
 	for _, reader := range readers {
-		if opts.LimitsExceeded(size, docsCount) {
+		if opts.LimitsExceeded(size, resultCount) {
+			exhaustive = true
 			break
 		}
 
@@ -685,19 +695,19 @@ func (b *block) aggregateWithSpan(
 			return false, err
 		}
 		iterClosed = false // only once the iterator has been successfully Reset().
-
 		for iter.Next() {
-			if opts.LimitsExceeded(size, docsCount) {
+			if opts.LimitsExceeded(size, resultCount) {
+				exhaustive = true
 				break
 			}
 
 			field, term := iter.Current()
 			batch = b.appendFieldAndTermToBatch(batch, field, term, iterateTerms)
-			if len(batch) < batchSize {
+			if batch.Size() < maxBatch {
 				continue
 			}
 
-			batch, size, docsCount, err = b.addAggregateResults(cancellable, results, batch, source)
+			batch, size, resultCount, err = b.addAggregateResults(cancellable, results, batch, source)
 			if err != nil {
 				return false, err
 			}
@@ -714,14 +724,14 @@ func (b *block) aggregateWithSpan(
 	}
 
 	// Add last batch to results if remaining.
-	if len(batch) > 0 {
-		batch, size, docsCount, err = b.addAggregateResults(cancellable, results, batch, source)
+	for len(batch) > 0 {
+		batch, size, resultCount, err = b.addAggregateResults(cancellable, results, batch, source)
 		if err != nil {
 			return false, err
 		}
 	}
 
-	return opts.exhaustive(size, docsCount), nil
+	return exhaustive || opts.exhaustive(size, resultCount), nil
 }
 
 func (b *block) appendFieldAndTermToBatch(
@@ -783,6 +793,7 @@ func (b *block) appendFieldAndTermToBatch(
 	} else {
 		batch = append(batch, entry)
 	}
+
 	return batch
 }
 
@@ -797,12 +808,12 @@ func (b *block) pooledID(id []byte) ident.ID {
 func (b *block) addAggregateResults(
 	cancellable *xresource.CancellableLifetime,
 	results AggregateResults,
-	batch []AggregateResultsEntry,
+	batch AggregateResultsEntries,
 	source []byte,
-) ([]AggregateResultsEntry, int, int, error) {
+) (AggregateResultsEntries, int, int, error) {
 	// update recently queried docs to monitor memory.
 	if results.EnforceLimits() {
-		if err := b.docsLimit.Inc(len(batch), source); err != nil {
+		if err := b.docsLimit.Inc(batch.Size(), source); err != nil {
 			return batch, 0, 0, err
 		}
 	}
@@ -814,8 +825,8 @@ func (b *block) addAggregateResults(
 		return batch, 0, 0, errCancelledQuery
 	}
 
-	// try to add the docs to the xresource.
-	size, docsCount := results.AddFields(batch)
+	// try to add the docs to the resource.
+	size, resultCount := results.AddFields(batch)
 
 	// immediately release the checkout on the lifetime of query.
 	cancellable.ReleaseCheckout()
@@ -828,7 +839,7 @@ func (b *block) addAggregateResults(
 	batch = batch[:0]
 
 	// return results.
-	return batch, size, docsCount, nil
+	return batch, size, resultCount, nil
 }
 
 func (b *block) AddResults(