Skip to content

Commit 5eee7a3

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

10 files changed

Lines changed: 488 additions & 8 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/data/table/IndexedTable.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -174,6 +174,9 @@ public int size() {
174174

175175
@Override
176176
public Iterator<Record> iterator() {
177+
if (_topRecords == null) {
178+
return _lookupMap.values().iterator();
179+
}
177180
return _topRecords.iterator();
178181
}
179182

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

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -57,15 +57,15 @@ public class GroupByCombineOperator extends BaseSingleBlockCombineOperator<Group
5757
private static final Logger LOGGER = LoggerFactory.getLogger(GroupByCombineOperator.class);
5858
private static final String EXPLAIN_NAME = "COMBINE_GROUP_BY";
5959

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

67-
private volatile IndexedTable _indexedTable;
68-
private volatile boolean _numGroupsLimitReached;
67+
protected volatile IndexedTable _indexedTable;
68+
protected volatile boolean _numGroupsLimitReached;
6969

7070
public GroupByCombineOperator(List<Operator> operators, QueryContext queryContext, ExecutorService executorService) {
7171
super(null, operators, overrideMaxExecutionThreads(queryContext, operators.size()), executorService);

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

Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,154 @@
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.ExecutorService;
26+
import org.apache.pinot.core.common.Operator;
27+
import org.apache.pinot.core.data.table.IndexedTable;
28+
import org.apache.pinot.core.data.table.IntermediateRecord;
29+
import org.apache.pinot.core.data.table.Key;
30+
import org.apache.pinot.core.data.table.Record;
31+
import org.apache.pinot.core.operator.AcquireReleaseColumnsSegmentOperator;
32+
import org.apache.pinot.core.operator.blocks.results.GroupByResultsBlock;
33+
import org.apache.pinot.core.query.aggregation.groupby.AggregationGroupByResult;
34+
import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGenerator;
35+
import org.apache.pinot.core.query.request.context.QueryContext;
36+
import org.apache.pinot.core.util.GroupByUtils;
37+
import org.apache.pinot.spi.trace.Tracing;
38+
import org.slf4j.Logger;
39+
import org.slf4j.LoggerFactory;
40+
41+
42+
/**
43+
* Combine operator for group-by queries.
44+
* TODO: Use CombineOperatorUtils.getNumThreadsForQuery() to get the parallelism of the query instead of using
45+
* all threads
46+
*/
47+
@SuppressWarnings("rawtypes")
48+
public class NonblockingGroupByCombineOperator extends GroupByCombineOperator {
49+
public static final String ALGORITHM = "NON-BLOCKING";
50+
51+
private static final Logger LOGGER = LoggerFactory.getLogger(NonblockingGroupByCombineOperator.class);
52+
private static final String EXPLAIN_NAME = "NON_BLOCKING_COMBINE_GROUP_BY";
53+
54+
public NonblockingGroupByCombineOperator(List<Operator> operators, QueryContext queryContext,
55+
ExecutorService executorService) {
56+
super(operators, queryContext, executorService);
57+
LOGGER.info("Using {} for group-by combine", ALGORITHM);
58+
}
59+
60+
@Override
61+
public String toExplainString() {
62+
return EXPLAIN_NAME;
63+
}
64+
65+
/**
66+
* Executes query on one segment in a worker thread and merges the results into the indexed table.
67+
*/
68+
@Override
69+
protected void processSegments() {
70+
int operatorId;
71+
while (_processingException.get() == null && (operatorId = _nextOperatorId.getAndIncrement()) < _numOperators) {
72+
Operator operator = _operators.get(operatorId);
73+
try {
74+
if (operator instanceof AcquireReleaseColumnsSegmentOperator) {
75+
((AcquireReleaseColumnsSegmentOperator) operator).acquire();
76+
}
77+
GroupByResultsBlock resultsBlock = (GroupByResultsBlock) operator.nextBlock();
78+
IndexedTable indexedTable = null;
79+
if (_indexedTable != null) {
80+
synchronized (this) {
81+
if (_indexedTable != null) {
82+
indexedTable = _indexedTable;
83+
_indexedTable = null;
84+
}
85+
}
86+
}
87+
if (indexedTable == null) {
88+
indexedTable = GroupByUtils.createIndexedTableForCombineOperator(resultsBlock, _queryContext, 1);
89+
}
90+
91+
// Set groups limit reached flag.
92+
if (resultsBlock.isNumGroupsLimitReached()) {
93+
_numGroupsLimitReached = true;
94+
}
95+
96+
// Merge aggregation group-by result.
97+
// Iterate over the group-by keys, for each key, update the group-by result in the indexedTable
98+
Collection<IntermediateRecord> intermediateRecords = resultsBlock.getIntermediateRecords();
99+
// Count the number of merged keys
100+
int mergedKeys = 0;
101+
// For now, only GroupBy OrderBy query has pre-constructed intermediate records
102+
if (intermediateRecords == null) {
103+
// Merge aggregation group-by result.
104+
AggregationGroupByResult aggregationGroupByResult = resultsBlock.getAggregationGroupByResult();
105+
if (aggregationGroupByResult != null) {
106+
// Iterate over the group-by keys, for each key, update the group-by result in the indexedTable
107+
Iterator<GroupKeyGenerator.GroupKey> dicGroupKeyIterator = aggregationGroupByResult.getGroupKeyIterator();
108+
while (dicGroupKeyIterator.hasNext()) {
109+
GroupKeyGenerator.GroupKey groupKey = dicGroupKeyIterator.next();
110+
Object[] keys = groupKey._keys;
111+
Object[] values = Arrays.copyOf(keys, _numColumns);
112+
int groupId = groupKey._groupId;
113+
for (int i = 0; i < _numAggregationFunctions; i++) {
114+
values[_numGroupByExpressions + i] = aggregationGroupByResult.getResultForGroupId(i, groupId);
115+
}
116+
indexedTable.upsert(new Key(keys), new Record(values));
117+
Tracing.ThreadAccountantOps.sampleAndCheckInterruptionPeriodically(mergedKeys);
118+
mergedKeys++;
119+
}
120+
}
121+
} else {
122+
for (IntermediateRecord intermediateResult : intermediateRecords) {
123+
//TODO: change upsert api so that it accepts intermediateRecord directly
124+
indexedTable.upsert(intermediateResult._key, intermediateResult._record);
125+
Tracing.ThreadAccountantOps.sampleAndCheckInterruptionPeriodically(mergedKeys);
126+
mergedKeys++;
127+
}
128+
}
129+
boolean setGroupByResult = false;
130+
while (!setGroupByResult) {
131+
IndexedTable indexedTableToMerge = null;
132+
synchronized (this) {
133+
if (_indexedTable == null) {
134+
_indexedTable = indexedTable;
135+
setGroupByResult = true;
136+
} else {
137+
indexedTableToMerge = _indexedTable;
138+
_indexedTable = null;
139+
}
140+
}
141+
if (indexedTableToMerge != null) {
142+
indexedTable.merge(indexedTableToMerge);
143+
}
144+
}
145+
} catch (RuntimeException e) {
146+
throw wrapOperatorException(operator, e);
147+
} finally {
148+
if (operator instanceof AcquireReleaseColumnsSegmentOperator) {
149+
((AcquireReleaseColumnsSegmentOperator) operator).release();
150+
}
151+
}
152+
}
153+
}
154+
}

‎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().toUpperCase()) {
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)