Skip to content

Commit 7fe387a

Browse files
xiangfu0xiangfu0
committed
Add a non-blocking groupBy implementation
1 parent 616d691 commit 7fe387a

8 files changed

Lines changed: 569 additions & 2 deletions

File tree

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -281,6 +281,11 @@ public static Integer getMinInitialIndexedTableCapacity(Map<String, String> quer
281281
return checkedParseIntPositive(QueryOptionKey.MIN_INITIAL_INDEXED_TABLE_CAPACITY, minInitialIndexedTableCapacity);
282282
}
283283

284+
@Nullable
285+
public static String getGroupByAlgorithm(Map<String, String> queryOptions) {
286+
return queryOptions.get(QueryOptionKey.GROUP_BY_ALGORITHM);
287+
}
288+
284289
public static boolean shouldDropResults(Map<String, String> queryOptions) {
285290
return Boolean.parseBoolean(queryOptions.get(CommonConstants.Broker.Request.QueryOptionKey.DROP_RESULTS));
286291
}

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

Lines changed: 246 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,246 @@
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.Arrays;
22+
import java.util.Collection;
23+
import java.util.Iterator;
24+
import java.util.List;
25+
import java.util.concurrent.CountDownLatch;
26+
import java.util.concurrent.ExecutorService;
27+
import java.util.concurrent.TimeUnit;
28+
import java.util.concurrent.TimeoutException;
29+
import org.apache.pinot.core.common.Operator;
30+
import org.apache.pinot.core.data.table.IndexedTable;
31+
import org.apache.pinot.core.data.table.IntermediateRecord;
32+
import org.apache.pinot.core.data.table.Key;
33+
import org.apache.pinot.core.data.table.Record;
34+
import org.apache.pinot.core.operator.AcquireReleaseColumnsSegmentOperator;
35+
import org.apache.pinot.core.operator.blocks.results.BaseResultsBlock;
36+
import org.apache.pinot.core.operator.blocks.results.ExceptionResultsBlock;
37+
import org.apache.pinot.core.operator.blocks.results.GroupByResultsBlock;
38+
import org.apache.pinot.core.query.aggregation.function.AggregationFunction;
39+
import org.apache.pinot.core.query.aggregation.groupby.AggregationGroupByResult;
40+
import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGenerator;
41+
import org.apache.pinot.core.query.request.context.QueryContext;
42+
import org.apache.pinot.core.util.GroupByUtils;
43+
import org.apache.pinot.spi.trace.Tracing;
44+
import org.slf4j.Logger;
45+
import org.slf4j.LoggerFactory;
46+
47+
48+
/**
49+
* Combine operator for group-by queries.
50+
* TODO: Use CombineOperatorUtils.getNumThreadsForQuery() to get the parallelism of the query instead of using
51+
* all threads
52+
*/
53+
@SuppressWarnings("rawtypes")
54+
public class NonblockingGroupByCombineOperator extends BaseSingleBlockCombineOperator<GroupByResultsBlock> {
55+
public static final int MAX_TRIM_THRESHOLD = 1_000_000_000;
56+
public static final String ALGORITHM = "non-blocking";
57+
58+
private static final Logger LOGGER = LoggerFactory.getLogger(NonblockingGroupByCombineOperator.class);
59+
private static final String EXPLAIN_NAME = "NON_BLOCKING_COMBINE_GROUP_BY";
60+
61+
private final int _numAggregationFunctions;
62+
private final int _numGroupByExpressions;
63+
private final int _numColumns;
64+
// We use a CountDownLatch to track if all Futures are finished by the query timeout, and cancel the unfinished
65+
// _futures (try to interrupt the execution if it already started).
66+
private final CountDownLatch _operatorLatch;
67+
68+
private volatile IndexedTable _indexedTable;
69+
private volatile boolean _numGroupsLimitReached;
70+
71+
public NonblockingGroupByCombineOperator(List<Operator> operators, QueryContext queryContext,
72+
ExecutorService executorService) {
73+
super(null, operators, overrideMaxExecutionThreads(queryContext, operators.size()), executorService);
74+
75+
AggregationFunction[] aggregationFunctions = _queryContext.getAggregationFunctions();
76+
assert aggregationFunctions != null;
77+
_numAggregationFunctions = aggregationFunctions.length;
78+
assert _queryContext.getGroupByExpressions() != null;
79+
_numGroupByExpressions = _queryContext.getGroupByExpressions().size();
80+
_numColumns = _numGroupByExpressions + _numAggregationFunctions;
81+
_operatorLatch = new CountDownLatch(_numTasks);
82+
}
83+
84+
/**
85+
* For group-by queries, when maxExecutionThreads is not explicitly configured, create one task per operator.
86+
*/
87+
private static QueryContext overrideMaxExecutionThreads(QueryContext queryContext, int numOperators) {
88+
int maxExecutionThreads = queryContext.getMaxExecutionThreads();
89+
if (maxExecutionThreads <= 0) {
90+
queryContext.setMaxExecutionThreads(numOperators);
91+
}
92+
return queryContext;
93+
}
94+
95+
@Override
96+
public String toExplainString() {
97+
return EXPLAIN_NAME;
98+
}
99+
100+
/**
101+
* Executes query on one segment in a worker thread and merges the results into the indexed table.
102+
*/
103+
@Override
104+
protected void processSegments() {
105+
int operatorId;
106+
while (_processingException.get() == null && (operatorId = _nextOperatorId.getAndIncrement()) < _numOperators) {
107+
Operator operator = _operators.get(operatorId);
108+
try {
109+
if (operator instanceof AcquireReleaseColumnsSegmentOperator) {
110+
((AcquireReleaseColumnsSegmentOperator) operator).acquire();
111+
}
112+
GroupByResultsBlock resultsBlock = (GroupByResultsBlock) operator.nextBlock();
113+
IndexedTable indexedTable = null;
114+
if (_indexedTable != null) {
115+
synchronized (this) {
116+
if (_indexedTable != null) {
117+
indexedTable = _indexedTable;
118+
_indexedTable = null;
119+
}
120+
}
121+
}
122+
if (indexedTable == null) {
123+
indexedTable = GroupByUtils.createIndexedTableForCombineOperator(resultsBlock, _queryContext, 1);
124+
}
125+
126+
// Set groups limit reached flag.
127+
if (resultsBlock.isNumGroupsLimitReached()) {
128+
_numGroupsLimitReached = true;
129+
}
130+
131+
// Merge aggregation group-by result.
132+
// Iterate over the group-by keys, for each key, update the group-by result in the indexedTable
133+
Collection<IntermediateRecord> intermediateRecords = resultsBlock.getIntermediateRecords();
134+
// Count the number of merged keys
135+
int mergedKeys = 0;
136+
// For now, only GroupBy OrderBy query has pre-constructed intermediate records
137+
if (intermediateRecords == null) {
138+
// Merge aggregation group-by result.
139+
AggregationGroupByResult aggregationGroupByResult = resultsBlock.getAggregationGroupByResult();
140+
if (aggregationGroupByResult != null) {
141+
// Iterate over the group-by keys, for each key, update the group-by result in the indexedTable
142+
Iterator<GroupKeyGenerator.GroupKey> dicGroupKeyIterator = aggregationGroupByResult.getGroupKeyIterator();
143+
while (dicGroupKeyIterator.hasNext()) {
144+
GroupKeyGenerator.GroupKey groupKey = dicGroupKeyIterator.next();
145+
Object[] keys = groupKey._keys;
146+
Object[] values = Arrays.copyOf(keys, _numColumns);
147+
int groupId = groupKey._groupId;
148+
for (int i = 0; i < _numAggregationFunctions; i++) {
149+
values[_numGroupByExpressions + i] = aggregationGroupByResult.getResultForGroupId(i, groupId);
150+
}
151+
indexedTable.upsert(new Key(keys), new Record(values));
152+
Tracing.ThreadAccountantOps.sampleAndCheckInterruptionPeriodically(mergedKeys);
153+
mergedKeys++;
154+
}
155+
}
156+
} else {
157+
for (IntermediateRecord intermediateResult : intermediateRecords) {
158+
//TODO: change upsert api so that it accepts intermediateRecord directly
159+
indexedTable.upsert(intermediateResult._key, intermediateResult._record);
160+
Tracing.ThreadAccountantOps.sampleAndCheckInterruptionPeriodically(mergedKeys);
161+
mergedKeys++;
162+
}
163+
}
164+
boolean setGroupByResult = false;
165+
while (!setGroupByResult) {
166+
IndexedTable indexedTableToMerge = null;
167+
synchronized (this) {
168+
if (_indexedTable == null) {
169+
_indexedTable = indexedTable;
170+
setGroupByResult = true;
171+
} else {
172+
indexedTableToMerge = _indexedTable;
173+
_indexedTable = null;
174+
}
175+
}
176+
if (indexedTableToMerge != null) {
177+
indexedTable.merge(indexedTableToMerge);
178+
}
179+
}
180+
} catch (RuntimeException e) {
181+
throw wrapOperatorException(operator, e);
182+
} finally {
183+
if (operator instanceof AcquireReleaseColumnsSegmentOperator) {
184+
((AcquireReleaseColumnsSegmentOperator) operator).release();
185+
}
186+
}
187+
}
188+
}
189+
190+
@Override
191+
public void onProcessSegmentsException(Throwable t) {
192+
_processingException.compareAndSet(null, t);
193+
}
194+
195+
@Override
196+
public void onProcessSegmentsFinish() {
197+
_operatorLatch.countDown();
198+
}
199+
200+
/**
201+
* {@inheritDoc}
202+
*
203+
* <p>Combines intermediate selection result blocks from underlying operators and returns a merged one.
204+
* <ul>
205+
* <li>
206+
* Merges multiple intermediate selection result blocks as a merged one.
207+
* </li>
208+
* <li>
209+
* Set all exceptions encountered during execution into the merged result block
210+
* </li>
211+
* </ul>
212+
*/
213+
@Override
214+
public BaseResultsBlock mergeResults()
215+
throws Exception {
216+
long timeoutMs = _queryContext.getEndTimeMs() - System.currentTimeMillis();
217+
boolean opCompleted = _operatorLatch.await(timeoutMs, TimeUnit.MILLISECONDS);
218+
if (!opCompleted) {
219+
// If this happens, the broker side should already timed out, just log the error and return
220+
String errorMessage =
221+
String.format("Timed out while combining group-by order-by results after %dms, queryContext = %s", timeoutMs,
222+
_queryContext);
223+
LOGGER.error(errorMessage);
224+
return new ExceptionResultsBlock(new TimeoutException(errorMessage));
225+
}
226+
227+
Throwable processingException = _processingException.get();
228+
if (processingException != null) {
229+
return new ExceptionResultsBlock(processingException);
230+
}
231+
232+
IndexedTable indexedTable = _indexedTable;
233+
if (_queryContext.isServerReturnFinalResult()) {
234+
indexedTable.finish(true, true);
235+
} else if (_queryContext.isServerReturnFinalResultKeyUnpartitioned()) {
236+
indexedTable.finish(false, true);
237+
} else {
238+
indexedTable.finish(false);
239+
}
240+
GroupByResultsBlock mergedBlock = new GroupByResultsBlock(indexedTable, _queryContext);
241+
mergedBlock.setNumGroupsLimitReached(_numGroupsLimitReached);
242+
mergedBlock.setNumResizes(indexedTable.getNumResizes());
243+
mergedBlock.setResizeTimeMs(indexedTable.getResizeTimeMs());
244+
return mergedBlock;
245+
}
246+
}

‎pinot-core/src/main/java/org/apache/pinot/core/plan/CombinePlanNode.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
import org.apache.pinot.core.operator.combine.DistinctCombineOperator;
3131
import org.apache.pinot.core.operator.combine.GroupByCombineOperator;
3232
import org.apache.pinot.core.operator.combine.MinMaxValueBasedSelectionOrderByCombineOperator;
33+
import org.apache.pinot.core.operator.combine.NonblockingGroupByCombineOperator;
3334
import org.apache.pinot.core.operator.combine.SelectionOnlyCombineOperator;
3435
import org.apache.pinot.core.operator.combine.SelectionOrderByCombineOperator;
3536
import org.apache.pinot.core.operator.combine.TimeSeriesCombineOperator;
@@ -140,7 +141,12 @@ private BaseCombineOperator getCombineOperator() {
140141
return new AggregationCombineOperator(operators, _queryContext, _executorService);
141142
} else {
142143
// Aggregation group-by
143-
return new GroupByCombineOperator(operators, _queryContext, _executorService);
144+
switch (_queryContext.getGroupByAlgorithm()) {
145+
case NonblockingGroupByCombineOperator.ALGORITHM:
146+
return new NonblockingGroupByCombineOperator(operators, _queryContext, _executorService);
147+
default:
148+
return new GroupByCombineOperator(operators, _queryContext, _executorService);
149+
}
144150
}
145151
} else if (QueryContextUtils.isSelectionQuery(_queryContext)) {
146152
if (_queryContext.getLimit() == 0 || _queryContext.getOrderByExpressions() == null) {

‎pinot-core/src/main/java/org/apache/pinot/core/plan/maker/InstancePlanMakerImplV2.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,7 @@ public class InstancePlanMakerImplV2 implements PlanMaker {
8989
// set as pinot.server.query.executor.groupby.trim.threshold
9090
public static final String GROUPBY_TRIM_THRESHOLD_KEY = "groupby.trim.threshold";
9191
public static final int DEFAULT_GROUPBY_TRIM_THRESHOLD = 1_000_000;
92+
public static final String DEFAULT_GROUP_BY_ALGORITHM = "default";
9293

9394
private static final Logger LOGGER = LoggerFactory.getLogger(InstancePlanMakerImplV2.class);
9495

@@ -237,6 +238,13 @@ private void applyQueryOptions(QueryContext queryContext) {
237238
} else {
238239
queryContext.setMinInitialIndexedTableCapacity(_minInitialIndexedTableCapacity);
239240
}
241+
// Set groupByAlgorithm
242+
String groupByAlgorithm = QueryOptionsUtils.getGroupByAlgorithm(queryOptions);
243+
if (groupByAlgorithm != null) {
244+
queryContext.setGroupByAlgorithm(groupByAlgorithm);
245+
} else {
246+
queryContext.setGroupByAlgorithm(DEFAULT_GROUP_BY_ALGORITHM);
247+
}
240248
// Set numGroupsLimit
241249
Integer numGroupsLimit = QueryOptionsUtils.getNumGroupsLimit(queryOptions);
242250
if (numGroupsLimit != null) {

‎pinot-core/src/main/java/org/apache/pinot/core/query/request/context/QueryContext.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,8 @@ public class QueryContext {
116116
private int _maxInitialResultHolderCapacity = InstancePlanMakerImplV2.DEFAULT_MAX_INITIAL_RESULT_HOLDER_CAPACITY;
117117
// Initial capacity of the indexed table
118118
private int _minInitialIndexedTableCapacity = InstancePlanMakerImplV2.DEFAULT_MIN_INITIAL_INDEXED_TABLE_CAPACITY;
119+
// Algorithm to use for SQL GROUP BY
120+
private String _groupByAlgorithm = InstancePlanMakerImplV2.DEFAULT_GROUP_BY_ALGORITHM;
119121
// Limit of number of groups stored in each segment
120122
private int _numGroupsLimit = InstancePlanMakerImplV2.DEFAULT_NUM_GROUPS_LIMIT;
121123
// Minimum number of groups to keep per segment when trimming groups for SQL GROUP BY
@@ -378,6 +380,14 @@ public void setMinInitialIndexedTableCapacity(int minInitialIndexedTableCapacity
378380
_minInitialIndexedTableCapacity = minInitialIndexedTableCapacity;
379381
}
380382

383+
public String getGroupByAlgorithm() {
384+
return _groupByAlgorithm;
385+
}
386+
387+
public void setGroupByAlgorithm(String groupByAlgorithm) {
388+
_groupByAlgorithm = groupByAlgorithm;
389+
}
390+
381391
public int getNumGroupsLimit() {
382392
return _numGroupsLimit;
383393
}

‎pinot-core/src/test/java/org/apache/pinot/queries/BaseQueriesTest.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -180,7 +180,7 @@ protected BrokerResponseNative getBrokerResponse(
180180
* This can be particularly useful to test statistical aggregation functions.
181181
* @see StatisticalQueriesTest for an example use case.
182182
*/
183-
private BrokerResponseNative getBrokerResponse(@Language("sql") String query, PlanMaker planMaker,
183+
protected BrokerResponseNative getBrokerResponse(@Language("sql") String query, PlanMaker planMaker,
184184
@Nullable Map<String, String> extraQueryOptions) {
185185
PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query);
186186
if (extraQueryOptions != null) {

0 commit comments

Comments
 (0)