Skip to content

Commit

Permalink
topn with granularity regression fixes (apache#17565)
Browse files Browse the repository at this point in the history 
* 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 apache#16533
* fix issue where topN with query granularity other than ALL could incorrectly process values in the wrong time bucket, another regression from apache#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
  • Loading branch information
@clintropolis
clintropolis committed Dec 17, 2024
1 parent fe4d7f3 commit 5cd1870
Show file tree
Hide file tree
Showing 14 changed files with 1,213 additions and 239 deletions.
12 changes: 8 additions & 4 deletions integration-tests/src/test/resources/queries/twitterstream_queries.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,8 @@
"context": {
"useCache": "true",
"populateCache": "true",
"timeout": 60000
"timeout": 60000,
"useTopNMultiPassPooledQueryGranularity": "true"
}
},
"expectedResults": [
Expand Down Expand Up @@ -198,7 +199,8 @@
"context": {
"useCache": "true",
"populateCache": "true",
"timeout": 60000
"timeout": 60000,
"useTopNMultiPassPooledQueryGranularity": "true"
}
},
"expectedResults": [
Expand Down Expand Up @@ -322,7 +324,8 @@
"context": {
"useCache": "true",
"populateCache": "true",
"timeout": 60000
"timeout": 60000,
"useTopNMultiPassPooledQueryGranularity": "true"
}
},
"expectedResults": [
Expand Down Expand Up @@ -741,7 +744,8 @@
"context": {
"useCache": "true",
"populateCache": "true",
"timeout": 60000
"timeout": 60000,
"useTopNMultiPassPooledQueryGranularity": "true"
}
},
"expectedResults": [
Expand Down
16 changes: 13 additions & 3 deletions processing/src/main/java/org/apache/druid/query/CursorGranularizer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -94,11 +95,13 @@ public static CursorGranularizer create(
timeSelector = cursor.getColumnSelectorFactory().makeColumnValueSelector(ColumnHolder.TIME_COLUMN_NAME);
}

return new CursorGranularizer(cursor, bucketIterable, timeSelector, timeOrder == Order.DESCENDING);
return new CursorGranularizer(cursor, granularity, bucketIterable, timeSelector, timeOrder == Order.DESCENDING);
}

private final Cursor cursor;

private final Granularity granularity;

// Iterable that iterates over time buckets.
private final Iterable<Interval> bucketIterable;

Expand All @@ -112,12 +115,14 @@ public static CursorGranularizer create(

private CursorGranularizer(
Cursor cursor,
Granularity granularity,
Iterable<Interval> bucketIterable,
@Nullable ColumnValueSelector timeSelector,
boolean descending
)
{
this.cursor = cursor;
this.granularity = granularity;
this.bucketIterable = bucketIterable;
this.timeSelector = timeSelector;
this.descending = descending;
Expand All @@ -133,13 +138,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;
}
Expand Down
6 changes: 6 additions & 0 deletions processing/src/main/java/org/apache/druid/query/QueryContexts.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down
2 changes: 1 addition & 1 deletion ...essing/src/main/java/org/apache/druid/query/groupby/epinephelinae/GroupByQueryEngine.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
Expand Down
1 change: 1 addition & 0 deletions processing/src/main/java/org/apache/druid/query/topn/AggregateTopNMetricFirstAlgorithm.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
7 changes: 5 additions & 2 deletions processing/src/main/java/org/apache/druid/query/topn/BaseTopNAlgorithm.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand All @@ -125,6 +127,7 @@ private void runWithCardinalityKnown(
numProcessed += numToProcess;
if (numProcessed < cardinality) {
params.getCursor().reset();
params.getGranularizer().advanceToBucket(params.getGranularizer().getCurrentInterval());
}
}
if (queryMetrics != null) {
Expand Down
38 changes: 20 additions & 18 deletions ...sing/src/main/java/org/apache/druid/query/topn/Generic1AggPooledTopNScannerPrototype.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
46 changes: 24 additions & 22 deletions ...sing/src/main/java/org/apache/druid/query/topn/Generic2AggPooledTopNScannerPrototype.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Loading

0 comments on commit 5cd1870

Please sign in to comment.