Skip to content

Commit 343e680

Browse files
xiangfu0xiangfu0
committed
Add a non-blocking and partitioned groupBy implementation
1 parent cd3e850 commit 343e680

13 files changed

Lines changed: 785 additions & 96 deletions

File tree

‎pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -403,6 +403,17 @@ public static Integer getMinInitialIndexedTableCapacity(Map<String, String> quer
403403
return checkedParseIntPositive(QueryOptionKey.MIN_INITIAL_INDEXED_TABLE_CAPACITY, minInitialIndexedTableCapacity);
404404
}
405405

406+
@Nullable
407+
public static String getGroupByAlgorithm(Map<String, String> queryOptions) {
408+
return queryOptions.get(QueryOptionKey.GROUP_BY_ALGORITHM);
409+
}
410+
411+
@Nullable
412+
public static Integer getNumGroupByPartitions(Map<String, String> queryOptions) {
413+
String numGroupByPartitionsString = queryOptions.get(QueryOptionKey.NUM_GROUP_BY_PARTITIONS);
414+
return checkedParseIntPositive(QueryOptionKey.NUM_GROUP_BY_PARTITIONS, numGroupByPartitionsString);
415+
}
416+
406417
public static boolean shouldDropResults(Map<String, String> queryOptions) {
407418
return Boolean.parseBoolean(queryOptions.get(CommonConstants.Broker.Request.QueryOptionKey.DROP_RESULTS));
408419
}

‎pinot-core/src/main/java/org/apache/pinot/core/data/table/IndexedTable.java‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -263,9 +263,29 @@ public int size() {
263263

264264
@Override
265265
public Iterator<Record> iterator() {
266+
if (_topRecords == null) {
267+
return _lookupMap.values().iterator();
268+
}
266269
return _topRecords.iterator();
267270
}
268271

272+
273+
public void mergePartitionTable(Table table) {
274+
if (table instanceof IndexedTable) {
275+
_lookupMap.putAll(((IndexedTable) table)._lookupMap);
276+
} else {
277+
Iterator<Record> iterator = table.iterator();
278+
while (iterator.hasNext()) {
279+
// NOTE: For performance concern, does not check the return value of the upsert(). Override this method if
280+
// upsert() can return false.
281+
upsert(iterator.next());
282+
}
283+
}
284+
if (_lookupMap.size() >= _trimThreshold) {
285+
resize();
286+
}
287+
}
288+
269289
public int getNumResizes() {
270290
return _numResizes;
271291
}

‎pinot-core/src/main/java/org/apache/pinot/core/operator/combine/GroupByCombineOperator.java‎

Lines changed: 88 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -53,20 +53,22 @@
5353
*/
5454
@SuppressWarnings("rawtypes")
5555
public class GroupByCombineOperator extends BaseSingleBlockCombineOperator<GroupByResultsBlock> {
56+
public static final String ALGORITHM = "DEFAULT";
57+
5658
private static final Logger LOGGER = LoggerFactory.getLogger(GroupByCombineOperator.class);
5759
private static final String EXPLAIN_NAME = "COMBINE_GROUP_BY";
5860

59-
private final int _numAggregationFunctions;
60-
private final int _numGroupByExpressions;
61-
private final int _numColumns;
61+
protected final int _numAggregationFunctions;
62+
protected final int _numGroupByExpressions;
63+
protected final int _numColumns;
6264
// We use a CountDownLatch to track if all Futures are finished by the query timeout, and cancel the unfinished
6365
// _futures (try to interrupt the execution if it already started).
64-
private final CountDownLatch _operatorLatch;
66+
protected final CountDownLatch _operatorLatch;
6567

66-
private volatile IndexedTable _indexedTable;
67-
private volatile boolean _groupsTrimmed;
68-
private volatile boolean _numGroupsLimitReached;
69-
private volatile boolean _numGroupsWarningLimitReached;
68+
protected volatile IndexedTable _indexedTable;
69+
protected volatile boolean _groupsTrimmed;
70+
protected volatile boolean _numGroupsLimitReached;
71+
protected volatile boolean _numGroupsWarningLimitReached;
7072

7173
public GroupByCombineOperator(List<Operator> operators, QueryContext queryContext, ExecutorService executorService) {
7274
super(null, operators, overrideMaxExecutionThreads(queryContext, operators.size()), executorService);
@@ -113,59 +115,11 @@ protected void processSegments() {
113115
if (_indexedTable == null) {
114116
synchronized (this) {
115117
if (_indexedTable == null) {
116-
_indexedTable = GroupByUtils.createIndexedTableForCombineOperator(resultsBlock, _queryContext, _numTasks,
117-
_executorService);
118-
}
119-
}
120-
}
121-
122-
if (resultsBlock.isGroupsTrimmed()) {
123-
_groupsTrimmed = true;
124-
}
125-
// Set groups limit reached flag.
126-
if (resultsBlock.isNumGroupsLimitReached()) {
127-
_numGroupsLimitReached = true;
128-
}
129-
if (resultsBlock.isNumGroupsWarningLimitReached()) {
130-
_numGroupsWarningLimitReached = true;
131-
}
132-
133-
// Merge aggregation group-by result.
134-
// Iterate over the group-by keys, for each key, update the group-by result in the indexedTable
135-
Collection<IntermediateRecord> intermediateRecords = resultsBlock.getIntermediateRecords();
136-
// Count the number of merged keys
137-
int mergedKeys = 0;
138-
// For now, only GroupBy OrderBy query has pre-constructed intermediate records
139-
if (intermediateRecords == null) {
140-
// Merge aggregation group-by result.
141-
AggregationGroupByResult aggregationGroupByResult = resultsBlock.getAggregationGroupByResult();
142-
if (aggregationGroupByResult != null) {
143-
// Iterate over the group-by keys, for each key, update the group-by result in the indexedTable
144-
try {
145-
Iterator<GroupKeyGenerator.GroupKey> dicGroupKeyIterator = aggregationGroupByResult.getGroupKeyIterator();
146-
while (dicGroupKeyIterator.hasNext()) {
147-
QueryThreadContext.checkTerminationAndSampleUsagePeriodically(mergedKeys++, EXPLAIN_NAME);
148-
GroupKeyGenerator.GroupKey groupKey = dicGroupKeyIterator.next();
149-
Object[] keys = groupKey._keys;
150-
Object[] values = Arrays.copyOf(keys, _numColumns);
151-
int groupId = groupKey._groupId;
152-
for (int i = 0; i < _numAggregationFunctions; i++) {
153-
values[_numGroupByExpressions + i] = aggregationGroupByResult.getResultForGroupId(i, groupId);
154-
}
155-
_indexedTable.upsert(new Key(keys), new Record(values));
156-
}
157-
} finally {
158-
// Release the resources used by the group key generator
159-
aggregationGroupByResult.closeGroupKeyGenerator();
118+
_indexedTable = createIndexedTable(resultsBlock, _numTasks);
160119
}
161120
}
162-
} else {
163-
for (IntermediateRecord intermediateResult : intermediateRecords) {
164-
QueryThreadContext.checkTerminationAndSampleUsagePeriodically(mergedKeys++, EXPLAIN_NAME);
165-
//TODO: change upsert api so that it accepts intermediateRecord directly
166-
_indexedTable.upsert(intermediateResult._key, intermediateResult._record);
167-
}
168121
}
122+
mergeGroupByResultsBlock(_indexedTable, resultsBlock, EXPLAIN_NAME);
169123
} catch (RuntimeException e) {
170124
throw wrapOperatorException(operator, e);
171125
} finally {
@@ -176,6 +130,56 @@ protected void processSegments() {
176130
}
177131
}
178132

133+
protected IndexedTable createIndexedTable(GroupByResultsBlock resultsBlock, int numThreads) {
134+
return GroupByUtils.createIndexedTableForCombineOperator(resultsBlock, _queryContext, numThreads, _executorService);
135+
}
136+
137+
protected void updateCombineResultsStats(GroupByResultsBlock resultsBlock) {
138+
if (resultsBlock.isGroupsTrimmed()) {
139+
_groupsTrimmed = true;
140+
}
141+
if (resultsBlock.isNumGroupsLimitReached()) {
142+
_numGroupsLimitReached = true;
143+
}
144+
if (resultsBlock.isNumGroupsWarningLimitReached()) {
145+
_numGroupsWarningLimitReached = true;
146+
}
147+
}
148+
149+
protected void mergeGroupByResultsBlock(IndexedTable indexedTable, GroupByResultsBlock resultsBlock,
150+
String explainName) {
151+
updateCombineResultsStats(resultsBlock);
152+
153+
Collection<IntermediateRecord> intermediateRecords = resultsBlock.getIntermediateRecords();
154+
int mergedKeys = 0;
155+
if (intermediateRecords == null) {
156+
AggregationGroupByResult aggregationGroupByResult = resultsBlock.getAggregationGroupByResult();
157+
if (aggregationGroupByResult != null) {
158+
try {
159+
Iterator<GroupKeyGenerator.GroupKey> dicGroupKeyIterator = aggregationGroupByResult.getGroupKeyIterator();
160+
while (dicGroupKeyIterator.hasNext()) {
161+
QueryThreadContext.checkTerminationAndSampleUsagePeriodically(mergedKeys++, explainName);
162+
GroupKeyGenerator.GroupKey groupKey = dicGroupKeyIterator.next();
163+
Object[] keys = groupKey._keys;
164+
Object[] values = Arrays.copyOf(keys, _numColumns);
165+
int groupId = groupKey._groupId;
166+
for (int i = 0; i < _numAggregationFunctions; i++) {
167+
values[_numGroupByExpressions + i] = aggregationGroupByResult.getResultForGroupId(i, groupId);
168+
}
169+
indexedTable.upsert(new Key(keys), new Record(values));
170+
}
171+
} finally {
172+
aggregationGroupByResult.closeGroupKeyGenerator();
173+
}
174+
}
175+
} else {
176+
for (IntermediateRecord intermediateResult : intermediateRecords) {
177+
QueryThreadContext.checkTerminationAndSampleUsagePeriodically(mergedKeys++, explainName);
178+
indexedTable.upsert(intermediateResult._key, intermediateResult._record);
179+
}
180+
}
181+
}
182+
179183
@Override
180184
public void onProcessSegmentsException(Throwable t) {
181185
_processingException.compareAndSet(null, t);
@@ -205,33 +209,41 @@ public BaseResultsBlock mergeResults()
205209
long timeoutMs = _queryContext.getEndTimeMs() - System.currentTimeMillis();
206210
boolean opCompleted = _operatorLatch.await(timeoutMs, TimeUnit.MILLISECONDS);
207211
if (!opCompleted) {
208-
// If this happens, the broker side should already timed out, just log the error and return
209-
String userError = "Timed out while combining group-by order-by results after " + timeoutMs + "ms";
210-
String logMsg = userError + ", queryContext = " + _queryContext;
211-
LOGGER.error(logMsg);
212-
return new ExceptionResultsBlock(new QueryErrorMessage(QueryErrorCode.EXECUTION_TIMEOUT, userError, logMsg));
212+
return getTimeoutResultsBlock(timeoutMs);
213213
}
214214

215215
Throwable ex = _processingException.get();
216216
if (ex != null) {
217-
String userError = "Caught exception while processing group-by order-by query";
218-
String devError = userError + ": " + ex.getMessage();
219-
QueryErrorMessage errMsg;
220-
if (ex instanceof QueryException) {
221-
// If the exception is a QueryException, use the error code from the exception and trust the error message
222-
errMsg = new QueryErrorMessage(((QueryException) ex).getErrorCode(), devError, devError);
223-
} else {
224-
// If the exception is not a QueryException, use the generic error code and don't expose the exception message
225-
errMsg = new QueryErrorMessage(QueryErrorCode.QUERY_EXECUTION, userError, devError);
226-
}
227-
return new ExceptionResultsBlock(errMsg);
217+
return getExceptionResultsBlock(ex);
228218
}
229219

230-
if (_indexedTable.isTrimmed() && _queryContext.isUnsafeTrim()) {
220+
return getMergedResultsBlock(_indexedTable);
221+
}
222+
223+
protected ExceptionResultsBlock getTimeoutResultsBlock(long timeoutMs) {
224+
String userError = "Timed out while combining group-by order-by results after " + timeoutMs + "ms";
225+
String logMsg = userError + ", queryContext = " + _queryContext;
226+
LOGGER.error(logMsg);
227+
return new ExceptionResultsBlock(new QueryErrorMessage(QueryErrorCode.EXECUTION_TIMEOUT, userError, logMsg));
228+
}
229+
230+
protected ExceptionResultsBlock getExceptionResultsBlock(Throwable ex) {
231+
String userError = "Caught exception while processing group-by order-by query";
232+
String devError = userError + ": " + ex.getMessage();
233+
QueryErrorMessage errMsg;
234+
if (ex instanceof QueryException) {
235+
errMsg = new QueryErrorMessage(((QueryException) ex).getErrorCode(), devError, devError);
236+
} else {
237+
errMsg = new QueryErrorMessage(QueryErrorCode.QUERY_EXECUTION, userError, devError);
238+
}
239+
return new ExceptionResultsBlock(errMsg);
240+
}
241+
242+
protected GroupByResultsBlock getMergedResultsBlock(IndexedTable indexedTable) {
243+
if (indexedTable.isTrimmed() && _queryContext.isUnsafeTrim()) {
231244
_groupsTrimmed = true;
232245
}
233246

234-
IndexedTable indexedTable = _indexedTable;
235247
if (_queryContext.isServerReturnFinalResult()) {
236248
indexedTable.finish(true, true);
237249
} else if (_queryContext.isServerReturnFinalResultKeyUnpartitioned()) {

‎pinot-core/src/main/java/org/apache/pinot/core/operator/combine/NonblockingGroupByCombineOperator.java‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
/**
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
package org.apache.pinot.core.operator.combine;
20+
21+
import java.util.List;
22+
import java.util.concurrent.ExecutorService;
23+
import org.apache.pinot.core.common.Operator;
24+
import org.apache.pinot.core.data.table.IndexedTable;
25+
import org.apache.pinot.core.operator.AcquireReleaseColumnsSegmentOperator;
26+
import org.apache.pinot.core.operator.blocks.results.GroupByResultsBlock;
27+
import org.apache.pinot.core.query.request.context.QueryContext;
28+
import org.slf4j.Logger;
29+
import org.slf4j.LoggerFactory;
30+
31+
32+
/**
33+
* Combine operator for group-by queries.
34+
* TODO: Use CombineOperatorUtils.getNumThreadsForQuery() to get the parallelism of the query instead of using
35+
* all threads
36+
*/
37+
@SuppressWarnings("rawtypes")
38+
public class NonblockingGroupByCombineOperator extends GroupByCombineOperator {
39+
public static final String ALGORITHM = "NON-BLOCKING";
40+
41+
private static final Logger LOGGER = LoggerFactory.getLogger(NonblockingGroupByCombineOperator.class);
42+
private static final String EXPLAIN_NAME = "NON_BLOCKING_COMBINE_GROUP_BY";
43+
44+
public NonblockingGroupByCombineOperator(List<Operator> operators, QueryContext queryContext,
45+
ExecutorService executorService) {
46+
super(operators, queryContext, executorService);
47+
LOGGER.info("Using {} for group-by combine with {} tasks", ALGORITHM, _numTasks);
48+
}
49+
50+
@Override
51+
public String toExplainString() {
52+
return EXPLAIN_NAME;
53+
}
54+
55+
/**
56+
* Executes query on one segment in a worker thread and merges the results into the indexed table.
57+
*/
58+
@Override
59+
protected void processSegments() {
60+
int operatorId;
61+
IndexedTable indexedTable = null;
62+
while (_processingException.get() == null && (operatorId = _nextOperatorId.getAndIncrement()) < _numOperators) {
63+
Operator operator = _operators.get(operatorId);
64+
try {
65+
if (operator instanceof AcquireReleaseColumnsSegmentOperator) {
66+
((AcquireReleaseColumnsSegmentOperator) operator).acquire();
67+
}
68+
GroupByResultsBlock resultsBlock = (GroupByResultsBlock) operator.nextBlock();
69+
if (indexedTable == null) {
70+
synchronized (this) {
71+
if (_indexedTable != null) {
72+
indexedTable = _indexedTable;
73+
_indexedTable = null;
74+
}
75+
}
76+
if (indexedTable == null) {
77+
indexedTable = createIndexedTable(resultsBlock, 1);
78+
}
79+
}
80+
mergeGroupByResultsBlock(indexedTable, resultsBlock, EXPLAIN_NAME);
81+
} catch (RuntimeException e) {
82+
throw wrapOperatorException(operator, e);
83+
} finally {
84+
if (operator instanceof AcquireReleaseColumnsSegmentOperator) {
85+
((AcquireReleaseColumnsSegmentOperator) operator).release();
86+
}
87+
}
88+
}
89+
90+
boolean setGroupByResult = false;
91+
while (indexedTable != null && !setGroupByResult) {
92+
IndexedTable indexedTableToMerge;
93+
synchronized (this) {
94+
if (_indexedTable == null) {
95+
_indexedTable = indexedTable;
96+
setGroupByResult = true;
97+
continue;
98+
} else {
99+
indexedTableToMerge = _indexedTable;
100+
_indexedTable = null;
101+
}
102+
}
103+
if (indexedTable.size() > indexedTableToMerge.size()) {
104+
indexedTable.merge(indexedTableToMerge);
105+
} else {
106+
indexedTableToMerge.merge(indexedTable);
107+
indexedTable = indexedTableToMerge;
108+
}
109+
}
110+
}
111+
}

0 commit comments

Comments
 (0)