Skip to content

Commit ca13a30

Browse files
authored
[Backport] topn with granularity regression fixes (#17580)
* topn with granularity regression fixes (#17565) * topn with granularity regression fixes changes: * fix issue where topN with query granularity other than ALL would use the heap algorithm when it was actual able to use the pooled algorithm, and incorrectly used the pool algorithm in cases where it must use the heap algorithm, a regression from #16533 * fix issue where topN with query granularity other than ALL could incorrectly process values in the wrong time bucket, another regression from #16533 * move defensive check outside of loop * more test * extra layer of safety * move check outside of loop * fix spelling * add query context parameter to allow using pooled algorithm for topN when multi-passes is required even wihen query granularity is not all * add comment, revert IT context changes and add new context flag * remove unused
1 parent fe4d7f3 commit ca13a30

14 files changed

Lines changed: 1208 additions & 238 deletions

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

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,8 @@
9494
"context": {
9595
"useCache": "true",
9696
"populateCache": "true",
97-
"timeout": 60000
97+
"timeout": 60000,
98+
"useTopNMultiPassPooledQueryGranularity": "true"
9899
}
99100
},
100101
"expectedResults": [
@@ -198,7 +199,8 @@
198199
"context": {
199200
"useCache": "true",
200201
"populateCache": "true",
201-
"timeout": 60000
202+
"timeout": 60000,
203+
"useTopNMultiPassPooledQueryGranularity": "true"
202204
}
203205
},
204206
"expectedResults": [
@@ -322,7 +324,8 @@
322324
"context": {
323325
"useCache": "true",
324326
"populateCache": "true",
325-
"timeout": 60000
327+
"timeout": 60000,
328+
"useTopNMultiPassPooledQueryGranularity": "true"
326329
}
327330
},
328331
"expectedResults": [
@@ -741,7 +744,8 @@
741744
"context": {
742745
"useCache": "true",
743746
"populateCache": "true",
744-
"timeout": 60000
747+
"timeout": 60000,
748+
"useTopNMultiPassPooledQueryGranularity": "true"
745749
}
746750
},
747751
"expectedResults": [

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

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import com.google.common.collect.Lists;
2525
import org.apache.druid.error.DruidException;
2626
import org.apache.druid.java.util.common.DateTimes;
27+
import org.apache.druid.java.util.common.Intervals;
2728
import org.apache.druid.java.util.common.granularity.Granularities;
2829
import org.apache.druid.java.util.common.granularity.Granularity;
2930
import org.apache.druid.segment.ColumnValueSelector;
@@ -133,13 +134,18 @@ public DateTime getBucketStart()
133134
return DateTimes.utc(currentBucketStart);
134135
}
135136

137+
public Interval getCurrentInterval()
138+
{
139+
return Intervals.utc(currentBucketStart, currentBucketEnd);
140+
}
141+
136142
public boolean advanceToBucket(final Interval bucketInterval)
137143
{
144+
currentBucketStart = bucketInterval.getStartMillis();
145+
currentBucketEnd = bucketInterval.getEndMillis();
138146
if (cursor.isDone()) {
139147
return false;
140148
}
141-
currentBucketStart = bucketInterval.getStartMillis();
142-
currentBucketEnd = bucketInterval.getEndMillis();
143149
if (timeSelector == null) {
144150
return true;
145151
}

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,12 @@ public class QueryContexts
8888
public static final String UNCOVERED_INTERVALS_LIMIT_KEY = "uncoveredIntervalsLimit";
8989
public static final String MIN_TOP_N_THRESHOLD = "minTopNThreshold";
9090
public static final String CATALOG_VALIDATION_ENABLED = "catalogValidationEnabled";
91+
// this flag controls whether the topN engine can use the 'pooled' algorithm when query granularity is set to
92+
// anything other than 'ALL' and the cardinality + number of aggregators would require more size than is available
93+
// in the buffers and so must reset the cursor to use multiple passes. This is likely slower than the default
94+
// behavior of falling back to heap memory, but less dangerous since too large of a query can cause the heap to run
95+
// out of memory
96+
public static final String TOPN_USE_MULTI_PASS_POOLED_QUERY_GRANULARITY = "useTopNMultiPassPooledQueryGranularity";
9197

9298
// projection context keys
9399
public static final String NO_PROJECTIONS = "noProjections";

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -391,7 +391,7 @@ public boolean hasNext()
391391
if (delegate != null && delegate.hasNext()) {
392392
return true;
393393
} else {
394-
if (!cursor.isDone() && granularizer.currentOffsetWithinBucket()) {
394+
if (granularizer.currentOffsetWithinBucket()) {
395395
if (delegate != null) {
396396
delegate.close();
397397
}

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,7 @@ public void run(
113113
try {
114114
// reset cursor since we call run again
115115
params.getCursor().reset();
116+
params.getGranularizer().advanceToBucket(params.getGranularizer().getCurrentInterval());
116117
// Run topN for all metrics for top N dimension values
117118
allMetricsParam = allMetricAlgo.makeInitParams(params.getSelectorPlus(), params.getCursor(), params.getGranularizer());
118119
allMetricAlgo.run(

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

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -97,12 +97,14 @@ private void runWithCardinalityKnown(
9797
}
9898
boolean hasDimValSelector = (dimValSelector != null);
9999

100-
int cardinality = params.getCardinality();
100+
final int cardinality = params.getCardinality();
101+
final int numValuesPerPass = params.getNumValuesPerPass();
101102
int numProcessed = 0;
102103
long processedRows = 0;
103104
while (numProcessed < cardinality) {
104105
final int numToProcess;
105-
int maxNumToProcess = Math.min(params.getNumValuesPerPass(), cardinality - numProcessed);
106+
int maxNumToProcess = Math.min(numValuesPerPass, cardinality - numProcessed);
107+
106108

107109
DimValSelector theDimValSelector;
108110
if (!hasDimValSelector) {
@@ -125,6 +127,7 @@ private void runWithCardinalityKnown(
125127
numProcessed += numToProcess;
126128
if (numProcessed < cardinality) {
127129
params.getCursor().reset();
130+
params.getGranularizer().advanceToBucket(params.getGranularizer().getCurrentInterval());
128131
}
129132
}
130133
if (queryMetrics != null) {

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

Lines changed: 20 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -54,25 +54,27 @@ public long scanAndAggregate(
5454
{
5555
long processedRows = 0;
5656
int positionToAllocate = 0;
57-
while (!cursor.isDoneOrInterrupted()) {
58-
final IndexedInts dimValues = dimensionSelector.getRow();
59-
final int dimSize = dimValues.size();
60-
for (int i = 0; i < dimSize; i++) {
61-
int dimIndex = dimValues.get(i);
62-
int position = positions[dimIndex];
63-
if (position >= 0) {
64-
aggregator.aggregate(resultsBuffer, position);
65-
} else if (position == TopNAlgorithm.INIT_POSITION_VALUE) {
66-
positions[dimIndex] = positionToAllocate;
67-
position = positionToAllocate;
68-
aggregator.init(resultsBuffer, position);
69-
aggregator.aggregate(resultsBuffer, position);
70-
positionToAllocate += aggregatorSize;
57+
if (granularizer.currentOffsetWithinBucket()) {
58+
while (!cursor.isDoneOrInterrupted()) {
59+
final IndexedInts dimValues = dimensionSelector.getRow();
60+
final int dimSize = dimValues.size();
61+
for (int i = 0; i < dimSize; i++) {
62+
int dimIndex = dimValues.get(i);
63+
int position = positions[dimIndex];
64+
if (position >= 0) {
65+
aggregator.aggregate(resultsBuffer, position);
66+
} else if (position == TopNAlgorithm.INIT_POSITION_VALUE) {
67+
positions[dimIndex] = positionToAllocate;
68+
position = positionToAllocate;
69+
aggregator.init(resultsBuffer, position);
70+
aggregator.aggregate(resultsBuffer, position);
71+
positionToAllocate += aggregatorSize;
72+
}
73+
}
74+
processedRows++;
75+
if (!granularizer.advanceCursorWithinBucketUninterruptedly()) {
76+
break;
7177
}
72-
}
73-
processedRows++;
74-
if (!granularizer.advanceCursorWithinBucketUninterruptedly()) {
75-
break;
7678
}
7779
}
7880
return processedRows;

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

Lines changed: 24 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -57,29 +57,31 @@ public long scanAndAggregate(
5757
int totalAggregatorsSize = aggregator1Size + aggregator2Size;
5858
long processedRows = 0;
5959
int positionToAllocate = 0;
60-
while (!cursor.isDoneOrInterrupted()) {
61-
final IndexedInts dimValues = dimensionSelector.getRow();
62-
final int dimSize = dimValues.size();
63-
for (int i = 0; i < dimSize; i++) {
64-
int dimIndex = dimValues.get(i);
65-
int position = positions[dimIndex];
66-
if (position >= 0) {
67-
aggregator1.aggregate(resultsBuffer, position);
68-
aggregator2.aggregate(resultsBuffer, position + aggregator1Size);
69-
} else if (position == TopNAlgorithm.INIT_POSITION_VALUE) {
70-
positions[dimIndex] = positionToAllocate;
71-
position = positionToAllocate;
72-
aggregator1.init(resultsBuffer, position);
73-
aggregator1.aggregate(resultsBuffer, position);
74-
position += aggregator1Size;
75-
aggregator2.init(resultsBuffer, position);
76-
aggregator2.aggregate(resultsBuffer, position);
77-
positionToAllocate += totalAggregatorsSize;
60+
if (granularizer.currentOffsetWithinBucket()) {
61+
while (!cursor.isDoneOrInterrupted()) {
62+
final IndexedInts dimValues = dimensionSelector.getRow();
63+
final int dimSize = dimValues.size();
64+
for (int i = 0; i < dimSize; i++) {
65+
int dimIndex = dimValues.get(i);
66+
int position = positions[dimIndex];
67+
if (position >= 0) {
68+
aggregator1.aggregate(resultsBuffer, position);
69+
aggregator2.aggregate(resultsBuffer, position + aggregator1Size);
70+
} else if (position == TopNAlgorithm.INIT_POSITION_VALUE) {
71+
positions[dimIndex] = positionToAllocate;
72+
position = positionToAllocate;
73+
aggregator1.init(resultsBuffer, position);
74+
aggregator1.aggregate(resultsBuffer, position);
75+
position += aggregator1Size;
76+
aggregator2.init(resultsBuffer, position);
77+
aggregator2.aggregate(resultsBuffer, position);
78+
positionToAllocate += totalAggregatorsSize;
79+
}
80+
}
81+
processedRows++;
82+
if (!granularizer.advanceCursorWithinBucketUninterruptedly()) {
83+
break;
7884
}
79-
}
80-
processedRows++;
81-
if (!granularizer.advanceCursorWithinBucketUninterruptedly()) {
82-
break;
8385
}
8486
}
8587
return processedRows;

0 commit comments

Comments
 (0)