Skip to content

Commit 530b877

Browse files
xiangfu0xiangfu0
committed
Add a non-blocking and partitioned groupBy implementation
1 parent 383bbce commit 530b877

14 files changed

Lines changed: 763 additions & 18 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
@@ -290,6 +290,17 @@ public static Integer getMinInitialIndexedTableCapacity(Map<String, String> quer
290290
return checkedParseIntPositive(QueryOptionKey.MIN_INITIAL_INDEXED_TABLE_CAPACITY, minInitialIndexedTableCapacity);
291291
}
292292

293+
@Nullable
294+
public static String getGroupByAlgorithm(Map<String, String> queryOptions) {
295+
return queryOptions.get(QueryOptionKey.GROUP_BY_ALGORITHM);
296+
}
297+
298+
@Nullable
299+
public static Integer getNumGroupByPartitions(Map<String, String> queryOptions) {
300+
String numGroupByPartitionsString = queryOptions.get(QueryOptionKey.NUM_GROUP_BY_PARTITIONS);
301+
return checkedParseIntPositive(QueryOptionKey.NUM_GROUP_BY_PARTITIONS, numGroupByPartitionsString);
302+
}
303+
293304
public static boolean shouldDropResults(Map<String, String> queryOptions) {
294305
return Boolean.parseBoolean(queryOptions.get(CommonConstants.Broker.Request.QueryOptionKey.DROP_RESULTS));
295306
}

‎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
@@ -174,9 +174,29 @@ 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

183+
184+
public void mergePartitionTable(Table table) {
185+
if (table instanceof IndexedTable) {
186+
_lookupMap.putAll(((IndexedTable) table)._lookupMap);
187+
} else {
188+
Iterator<Record> iterator = table.iterator();
189+
while (iterator.hasNext()) {
190+
// NOTE: For performance concern, does not check the return value of the upsert(). Override this method if
191+
// upsert() can return false.
192+
upsert(iterator.next());
193+
}
194+
}
195+
if (_lookupMap.size() >= _trimThreshold) {
196+
resize();
197+
}
198+
}
199+
180200
public int getNumResizes() {
181201
return _numResizes;
182202
}

‎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: 160 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,160 @@
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 with {} tasks", ALGORITHM, _numTasks);
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+
IndexedTable indexedTable = null;
72+
while (_processingException.get() == null && (operatorId = _nextOperatorId.getAndIncrement()) < _numOperators) {
73+
Operator operator = _operators.get(operatorId);
74+
try {
75+
if (operator instanceof AcquireReleaseColumnsSegmentOperator) {
76+
((AcquireReleaseColumnsSegmentOperator) operator).acquire();
77+
}
78+
GroupByResultsBlock resultsBlock = (GroupByResultsBlock) operator.nextBlock();
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+
} catch (RuntimeException e) {
130+
throw wrapOperatorException(operator, e);
131+
} finally {
132+
if (operator instanceof AcquireReleaseColumnsSegmentOperator) {
133+
((AcquireReleaseColumnsSegmentOperator) operator).release();
134+
}
135+
}
136+
}
137+
138+
boolean setGroupByResult = false;
139+
while (indexedTable != null && !setGroupByResult) {
140+
IndexedTable indexedTableToMerge = null;
141+
synchronized (this) {
142+
if (_indexedTable == null) {
143+
_indexedTable = indexedTable;
144+
setGroupByResult = true;
145+
} else {
146+
indexedTableToMerge = _indexedTable;
147+
_indexedTable = null;
148+
}
149+
}
150+
if (indexedTableToMerge != null) {
151+
if (indexedTable.size() > indexedTableToMerge.size()) {
152+
indexedTable.merge(indexedTableToMerge);
153+
} else {
154+
indexedTableToMerge.merge(indexedTable);
155+
indexedTable = indexedTableToMerge;
156+
}
157+
}
158+
}
159+
}
160+
}

0 commit comments

Comments
 (0)