[Backport] topn with granularity regression fixes

integration-tests/src/test/resources/queries/twitterstream_queries.json

@@ -94,7 +94,8 @@
             "context": {
                 "useCache": "true",
                 "populateCache": "true",
-                "timeout": 60000
+                "timeout": 60000,
+                "useTopNMultiPassPooledQueryGranularity": "true"
             }
         },
         "expectedResults": [
@@ -198,7 +199,8 @@
             "context": {
                 "useCache": "true",
                 "populateCache": "true",
-                "timeout": 60000
+                "timeout": 60000,
+                "useTopNMultiPassPooledQueryGranularity": "true"
             }
         },
         "expectedResults": [
@@ -322,7 +324,8 @@
             "context": {
                 "useCache": "true",
                 "populateCache": "true",
-                "timeout": 60000
+                "timeout": 60000,
+                "useTopNMultiPassPooledQueryGranularity": "true"
             }
         },
         "expectedResults": [
@@ -741,7 +744,8 @@
             "context": {
                 "useCache": "true",
                 "populateCache": "true",
-                "timeout": 60000
+                "timeout": 60000,
+                "useTopNMultiPassPooledQueryGranularity": "true"
             }
         },
         "expectedResults": [

processing/src/main/java/org/apache/druid/query/CursorGranularizer.java

@@ -24,6 +24,7 @@
 import com.google.common.collect.Lists;
 import org.apache.druid.error.DruidException;
 import org.apache.druid.java.util.common.DateTimes;
+import org.apache.druid.java.util.common.Intervals;
 import org.apache.druid.java.util.common.granularity.Granularities;
 import org.apache.druid.java.util.common.granularity.Granularity;
 import org.apache.druid.segment.ColumnValueSelector;
@@ -133,13 +134,18 @@ public DateTime getBucketStart()
     return DateTimes.utc(currentBucketStart);
   }
 
+  public Interval getCurrentInterval()
+  {
+    return Intervals.utc(currentBucketStart, currentBucketEnd);
+  }
+
   public boolean advanceToBucket(final Interval bucketInterval)
   {
+    currentBucketStart = bucketInterval.getStartMillis();
+    currentBucketEnd = bucketInterval.getEndMillis();
     if (cursor.isDone()) {
       return false;
     }
-    currentBucketStart = bucketInterval.getStartMillis();
-    currentBucketEnd = bucketInterval.getEndMillis();
     if (timeSelector == null) {
       return true;
     }

processing/src/main/java/org/apache/druid/query/QueryContexts.java

@@ -88,6 +88,12 @@ public class QueryContexts
   public static final String UNCOVERED_INTERVALS_LIMIT_KEY = "uncoveredIntervalsLimit";
   public static final String MIN_TOP_N_THRESHOLD = "minTopNThreshold";
   public static final String CATALOG_VALIDATION_ENABLED = "catalogValidationEnabled";
+  // this flag controls whether the topN engine can use the 'pooled' algorithm when query granularity is set to
+  // anything other than 'ALL' and the cardinality + number of aggregators would require more size than is available
+  // in the buffers and so must reset the cursor to use multiple passes. This is likely slower than the default
+  // behavior of falling back to heap memory, but less dangerous since too large of a query can cause the heap to run
+  // out of memory
+  public static final String TOPN_USE_MULTI_PASS_POOLED_QUERY_GRANULARITY = "useTopNMultiPassPooledQueryGranularity";
 
   // projection context keys
   public static final String NO_PROJECTIONS = "noProjections";

processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/GroupByQueryEngine.java

@@ -391,7 +391,7 @@ public boolean hasNext()
       if (delegate != null && delegate.hasNext()) {
         return true;
       } else {
-        if (!cursor.isDone() && granularizer.currentOffsetWithinBucket()) {
+        if (granularizer.currentOffsetWithinBucket()) {
           if (delegate != null) {
             delegate.close();
           }

processing/src/main/java/org/apache/druid/query/topn/AggregateTopNMetricFirstAlgorithm.java

@@ -113,6 +113,7 @@ public void run(
     try {
       // reset cursor since we call run again
       params.getCursor().reset();
+      params.getGranularizer().advanceToBucket(params.getGranularizer().getCurrentInterval());
       // Run topN for all metrics for top N dimension values
       allMetricsParam = allMetricAlgo.makeInitParams(params.getSelectorPlus(), params.getCursor(), params.getGranularizer());
       allMetricAlgo.run(

processing/src/main/java/org/apache/druid/query/topn/BaseTopNAlgorithm.java

@@ -97,12 +97,14 @@ private void runWithCardinalityKnown(
     }
     boolean hasDimValSelector = (dimValSelector != null);
 
-    int cardinality = params.getCardinality();
+    final int cardinality = params.getCardinality();
+    final int numValuesPerPass = params.getNumValuesPerPass();
     int numProcessed = 0;
     long processedRows = 0;
     while (numProcessed < cardinality) {
       final int numToProcess;
-      int maxNumToProcess = Math.min(params.getNumValuesPerPass(), cardinality - numProcessed);
+      int maxNumToProcess = Math.min(numValuesPerPass, cardinality - numProcessed);
+
 
       DimValSelector theDimValSelector;
       if (!hasDimValSelector) {
@@ -125,6 +127,7 @@ private void runWithCardinalityKnown(
       numProcessed += numToProcess;
       if (numProcessed < cardinality) {
         params.getCursor().reset();
+        params.getGranularizer().advanceToBucket(params.getGranularizer().getCurrentInterval());
       }
     }
     if (queryMetrics != null) {

processing/src/main/java/org/apache/druid/query/topn/Generic1AggPooledTopNScannerPrototype.java

@@ -54,25 +54,27 @@ public long scanAndAggregate(
   {
     long processedRows = 0;
     int positionToAllocate = 0;
-    while (!cursor.isDoneOrInterrupted()) {
-      final IndexedInts dimValues = dimensionSelector.getRow();
-      final int dimSize = dimValues.size();
-      for (int i = 0; i < dimSize; i++) {
-        int dimIndex = dimValues.get(i);
-        int position = positions[dimIndex];
-        if (position >= 0) {
-          aggregator.aggregate(resultsBuffer, position);
-        } else if (position == TopNAlgorithm.INIT_POSITION_VALUE) {
-          positions[dimIndex] = positionToAllocate;
-          position = positionToAllocate;
-          aggregator.init(resultsBuffer, position);
-          aggregator.aggregate(resultsBuffer, position);
-          positionToAllocate += aggregatorSize;
+    if (granularizer.currentOffsetWithinBucket()) {
+      while (!cursor.isDoneOrInterrupted()) {
+        final IndexedInts dimValues = dimensionSelector.getRow();
+        final int dimSize = dimValues.size();
+        for (int i = 0; i < dimSize; i++) {
+          int dimIndex = dimValues.get(i);
+          int position = positions[dimIndex];
+          if (position >= 0) {
+            aggregator.aggregate(resultsBuffer, position);
+          } else if (position == TopNAlgorithm.INIT_POSITION_VALUE) {
+            positions[dimIndex] = positionToAllocate;
+            position = positionToAllocate;
+            aggregator.init(resultsBuffer, position);
+            aggregator.aggregate(resultsBuffer, position);
+            positionToAllocate += aggregatorSize;
+          }
+        }
+        processedRows++;
+        if (!granularizer.advanceCursorWithinBucketUninterruptedly()) {
+          break;
         }
-      }
-      processedRows++;
-      if (!granularizer.advanceCursorWithinBucketUninterruptedly()) {
-        break;
       }
     }
     return processedRows;

processing/src/main/java/org/apache/druid/query/topn/Generic2AggPooledTopNScannerPrototype.java

@@ -57,29 +57,31 @@ public long scanAndAggregate(
     int totalAggregatorsSize = aggregator1Size + aggregator2Size;
     long processedRows = 0;
     int positionToAllocate = 0;
-    while (!cursor.isDoneOrInterrupted()) {
-      final IndexedInts dimValues = dimensionSelector.getRow();
-      final int dimSize = dimValues.size();
-      for (int i = 0; i < dimSize; i++) {
-        int dimIndex = dimValues.get(i);
-        int position = positions[dimIndex];
-        if (position >= 0) {
-          aggregator1.aggregate(resultsBuffer, position);
-          aggregator2.aggregate(resultsBuffer, position + aggregator1Size);
-        } else if (position == TopNAlgorithm.INIT_POSITION_VALUE) {
-          positions[dimIndex] = positionToAllocate;
-          position = positionToAllocate;
-          aggregator1.init(resultsBuffer, position);
-          aggregator1.aggregate(resultsBuffer, position);
-          position += aggregator1Size;
-          aggregator2.init(resultsBuffer, position);
-          aggregator2.aggregate(resultsBuffer, position);
-          positionToAllocate += totalAggregatorsSize;
+    if (granularizer.currentOffsetWithinBucket()) {
+      while (!cursor.isDoneOrInterrupted()) {
+        final IndexedInts dimValues = dimensionSelector.getRow();
+        final int dimSize = dimValues.size();
+        for (int i = 0; i < dimSize; i++) {
+          int dimIndex = dimValues.get(i);
+          int position = positions[dimIndex];
+          if (position >= 0) {
+            aggregator1.aggregate(resultsBuffer, position);
+            aggregator2.aggregate(resultsBuffer, position + aggregator1Size);
+          } else if (position == TopNAlgorithm.INIT_POSITION_VALUE) {
+            positions[dimIndex] = positionToAllocate;
+            position = positionToAllocate;
+            aggregator1.init(resultsBuffer, position);
+            aggregator1.aggregate(resultsBuffer, position);
+            position += aggregator1Size;
+            aggregator2.init(resultsBuffer, position);
+            aggregator2.aggregate(resultsBuffer, position);
+            positionToAllocate += totalAggregatorsSize;
+          }
+        }
+        processedRows++;
+        if (!granularizer.advanceCursorWithinBucketUninterruptedly()) {
+          break;
         }
-      }
-      processedRows++;
-      if (!granularizer.advanceCursorWithinBucketUninterruptedly()) {
-        break;
       }
     }
     return processedRows;