Skip to content

Commit fc14886

Browse files
xiangfu0xiangfu0
committed
Optimize partitioned group-by combine
1 parent e762c0c commit fc14886

3 files changed

Lines changed: 466 additions & 47 deletions

File tree

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

Lines changed: 167 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,16 @@
1818
*/
1919
package org.apache.pinot.core.operator.combine;
2020

21+
import java.util.ArrayList;
2122
import java.util.Arrays;
2223
import java.util.Collection;
2324
import java.util.Iterator;
2425
import java.util.List;
26+
import java.util.concurrent.ExecutionException;
2527
import java.util.concurrent.ExecutorService;
28+
import java.util.concurrent.Future;
2629
import java.util.concurrent.TimeUnit;
30+
import java.util.concurrent.TimeoutException;
2731
import org.apache.pinot.core.common.Operator;
2832
import org.apache.pinot.core.data.table.IndexedTable;
2933
import org.apache.pinot.core.data.table.IntermediateRecord;
@@ -32,9 +36,11 @@
3236
import org.apache.pinot.core.operator.AcquireReleaseColumnsSegmentOperator;
3337
import org.apache.pinot.core.operator.blocks.results.BaseResultsBlock;
3438
import org.apache.pinot.core.operator.blocks.results.GroupByResultsBlock;
39+
import org.apache.pinot.core.operator.blocks.results.ResultsBlockUtils;
3540
import org.apache.pinot.core.query.aggregation.groupby.AggregationGroupByResult;
3641
import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGenerator;
3742
import org.apache.pinot.core.query.request.context.QueryContext;
43+
import org.apache.pinot.core.util.trace.TraceCallable;
3844
import org.apache.pinot.spi.query.QueryThreadContext;
3945
import org.slf4j.Logger;
4046
import org.slf4j.LoggerFactory;
@@ -53,13 +59,16 @@ public class PartitionedGroupByCombineOperator extends GroupByCombineOperator {
5359
private static final String EXPLAIN_NAME = "PARTITIONED_COMBINE_GROUP_BY";
5460

5561
protected final IndexedTable[] _indexedTables;
62+
private final Object[] _partitionLocks;
5663
private final int _partitionMask;
5764

5865
public PartitionedGroupByCombineOperator(List<Operator> operators, QueryContext queryContext,
5966
ExecutorService executorService) {
6067
super(operators, queryContext, executorService);
6168
int numGroupByPartitions = Math.max(1, _queryContext.getNumGroupByPartitions());
6269
_indexedTables = new IndexedTable[numGroupByPartitions];
70+
_partitionLocks = new Object[numGroupByPartitions];
71+
Arrays.setAll(_partitionLocks, ignored -> new Object());
6372
_partitionMask = Integer.bitCount(numGroupByPartitions) == 1 ? numGroupByPartitions - 1 : -1;
6473
LOGGER.info("Using {} for group-by combine, with {} partitions and {} numTasks", EXPLAIN_NAME, numGroupByPartitions,
6574
_numTasks);
@@ -76,54 +85,18 @@ public String toExplainString() {
7685
@Override
7786
protected void processSegments() {
7887
int operatorId;
88+
IndexedTable[] localIndexedTables = null;
7989
while (_processingException.get() == null && (operatorId = _nextOperatorId.getAndIncrement()) < _numOperators) {
8090
Operator operator = _operators.get(operatorId);
8191
try {
8292
if (operator instanceof AcquireReleaseColumnsSegmentOperator) {
8393
((AcquireReleaseColumnsSegmentOperator) operator).acquire();
8494
}
8595
GroupByResultsBlock resultsBlock = (GroupByResultsBlock) operator.nextBlock();
86-
if (_indexedTables[0] == null) {
87-
synchronized (this) {
88-
if (_indexedTables[0] == null) {
89-
for (int i = 0; i < _indexedTables.length; i++) {
90-
_indexedTables[i] = createIndexedTable(resultsBlock, _numTasks);
91-
}
92-
}
93-
}
94-
}
95-
96-
updateCombineResultsStats(resultsBlock);
97-
Collection<IntermediateRecord> intermediateRecords = resultsBlock.getIntermediateRecords();
98-
int mergedKeys = 0;
99-
if (intermediateRecords == null) {
100-
AggregationGroupByResult aggregationGroupByResult = resultsBlock.getAggregationGroupByResult();
101-
if (aggregationGroupByResult != null) {
102-
try {
103-
Iterator<GroupKeyGenerator.GroupKey> dicGroupKeyIterator = aggregationGroupByResult.getGroupKeyIterator();
104-
while (dicGroupKeyIterator.hasNext()) {
105-
QueryThreadContext.checkTerminationAndSampleUsagePeriodically(mergedKeys++, EXPLAIN_NAME);
106-
GroupKeyGenerator.GroupKey groupKey = dicGroupKeyIterator.next();
107-
Object[] keys = groupKey._keys;
108-
Object[] values = Arrays.copyOf(keys, _numColumns);
109-
int groupId = groupKey._groupId;
110-
for (int i = 0; i < _numAggregationFunctions; i++) {
111-
values[_numGroupByExpressions + i] = aggregationGroupByResult.getResultForGroupId(i, groupId);
112-
}
113-
Key key = new Key(keys);
114-
_indexedTables[getPartitionId(key)].upsert(key, new Record(values));
115-
}
116-
} finally {
117-
aggregationGroupByResult.closeGroupKeyGenerator();
118-
}
119-
}
120-
} else {
121-
for (IntermediateRecord intermediateResult : intermediateRecords) {
122-
QueryThreadContext.checkTerminationAndSampleUsagePeriodically(mergedKeys++, EXPLAIN_NAME);
123-
_indexedTables[getPartitionId(intermediateResult._key)].upsert(intermediateResult._key,
124-
intermediateResult._record);
125-
}
96+
if (localIndexedTables == null) {
97+
localIndexedTables = new IndexedTable[_indexedTables.length];
12698
}
99+
mergeResultsBlockIntoLocalPartitions(localIndexedTables, resultsBlock);
127100
} catch (RuntimeException e) {
128101
throw wrapOperatorException(operator, e);
129102
} finally {
@@ -132,6 +105,7 @@ protected void processSegments() {
132105
}
133106
}
134107
}
108+
publishLocalPartitions(localIndexedTables);
135109
}
136110

137111
/**
@@ -161,18 +135,164 @@ public BaseResultsBlock mergeResults()
161135
return getExceptionResultsBlock(processingException);
162136
}
163137

164-
IndexedTable indexedTable = _indexedTables[0];
165-
for (int i = 1; i < _indexedTables.length; i++) {
166-
if (_indexedTables[i].size() > indexedTable.size()) {
167-
indexedTable.mergePartitionTable(_indexedTables[i]);
168-
} else {
169-
_indexedTables[i].mergePartitionTable(indexedTable);
170-
indexedTable = _indexedTables[i];
138+
List<IndexedTable> partitionTables = new ArrayList<>(_indexedTables.length);
139+
for (IndexedTable partitionTable : _indexedTables) {
140+
if (partitionTable != null) {
141+
partitionTables.add(partitionTable);
171142
}
172143
}
144+
IndexedTable indexedTable;
145+
try {
146+
indexedTable = mergePartitionTables(partitionTables);
147+
} catch (TimeoutException e) {
148+
long remainingTimeMs = _queryContext.getEndTimeMs() - System.currentTimeMillis();
149+
return getTimeoutResultsBlock(Math.max(0L, remainingTimeMs));
150+
}
151+
if (indexedTable == null) {
152+
return ResultsBlockUtils.buildEmptyQueryResults(_queryContext);
153+
}
173154
return getMergedResultsBlock(indexedTable);
174155
}
175156

157+
private void mergeResultsBlockIntoLocalPartitions(IndexedTable[] localIndexedTables,
158+
GroupByResultsBlock resultsBlock) {
159+
updateCombineResultsStats(resultsBlock);
160+
161+
Collection<IntermediateRecord> intermediateRecords = resultsBlock.getIntermediateRecords();
162+
int mergedKeys = 0;
163+
if (intermediateRecords == null) {
164+
AggregationGroupByResult aggregationGroupByResult = resultsBlock.getAggregationGroupByResult();
165+
if (aggregationGroupByResult != null) {
166+
try {
167+
Iterator<GroupKeyGenerator.GroupKey> dicGroupKeyIterator = aggregationGroupByResult.getGroupKeyIterator();
168+
while (dicGroupKeyIterator.hasNext()) {
169+
QueryThreadContext.checkTerminationAndSampleUsagePeriodically(mergedKeys++, EXPLAIN_NAME);
170+
GroupKeyGenerator.GroupKey groupKey = dicGroupKeyIterator.next();
171+
Object[] keys = groupKey._keys;
172+
Object[] values = Arrays.copyOf(keys, _numColumns);
173+
int groupId = groupKey._groupId;
174+
for (int i = 0; i < _numAggregationFunctions; i++) {
175+
values[_numGroupByExpressions + i] = aggregationGroupByResult.getResultForGroupId(i, groupId);
176+
}
177+
Key key = new Key(keys);
178+
getOrCreateLocalPartitionTable(localIndexedTables, getPartitionId(key), resultsBlock).upsert(key,
179+
new Record(values));
180+
}
181+
} finally {
182+
aggregationGroupByResult.closeGroupKeyGenerator();
183+
}
184+
}
185+
} else {
186+
for (IntermediateRecord intermediateResult : intermediateRecords) {
187+
QueryThreadContext.checkTerminationAndSampleUsagePeriodically(mergedKeys++, EXPLAIN_NAME);
188+
getOrCreateLocalPartitionTable(localIndexedTables, getPartitionId(intermediateResult._key), resultsBlock)
189+
.upsert(intermediateResult._key, intermediateResult._record);
190+
}
191+
}
192+
}
193+
194+
private IndexedTable getOrCreateLocalPartitionTable(IndexedTable[] localIndexedTables, int partitionId,
195+
GroupByResultsBlock resultsBlock) {
196+
IndexedTable indexedTable = localIndexedTables[partitionId];
197+
if (indexedTable == null) {
198+
indexedTable = createIndexedTable(resultsBlock, 1);
199+
localIndexedTables[partitionId] = indexedTable;
200+
}
201+
return indexedTable;
202+
}
203+
204+
private void publishLocalPartitions(IndexedTable[] localIndexedTables) {
205+
if (localIndexedTables == null) {
206+
return;
207+
}
208+
for (int partitionId = 0; partitionId < localIndexedTables.length; partitionId++) {
209+
IndexedTable localIndexedTable = localIndexedTables[partitionId];
210+
if (localIndexedTable == null) {
211+
continue;
212+
}
213+
synchronized (_partitionLocks[partitionId]) {
214+
IndexedTable indexedTable = _indexedTables[partitionId];
215+
if (indexedTable == null) {
216+
_indexedTables[partitionId] = localIndexedTable;
217+
} else if (localIndexedTable.size() > indexedTable.size()) {
218+
localIndexedTable.merge(indexedTable);
219+
_indexedTables[partitionId] = localIndexedTable;
220+
} else {
221+
indexedTable.merge(localIndexedTable);
222+
}
223+
}
224+
}
225+
}
226+
227+
private IndexedTable mergePartitionTables(List<IndexedTable> partitionTables)
228+
throws Exception {
229+
int numPartitionTables = partitionTables.size();
230+
if (numPartitionTables == 0) {
231+
return null;
232+
}
233+
if (numPartitionTables == 1) {
234+
return partitionTables.get(0);
235+
}
236+
237+
List<IndexedTable> tablesToMerge = partitionTables;
238+
if (tablesToMerge.size() <= 2 || _numTasks <= 1) {
239+
return mergePartitionTablesSequentially(tablesToMerge);
240+
}
241+
242+
while (tablesToMerge.size() > 1) {
243+
if (_queryContext.getEndTimeMs() - System.currentTimeMillis() <= 0) {
244+
throw new TimeoutException("Timed out while reducing partitioned group-by results");
245+
}
246+
247+
int numPairs = tablesToMerge.size() / 2;
248+
List<Future<IndexedTable>> futures = new ArrayList<>(numPairs);
249+
List<IndexedTable> nextRoundTables = new ArrayList<>((tablesToMerge.size() + 1) / 2);
250+
for (int i = 0; i < numPairs; i++) {
251+
IndexedTable leftTable = tablesToMerge.get(i * 2);
252+
IndexedTable rightTable = tablesToMerge.get(i * 2 + 1);
253+
futures.add(_executorService.submit(new TraceCallable<IndexedTable>() {
254+
@Override
255+
public IndexedTable callJob() {
256+
return mergePartitionTables(leftTable, rightTable);
257+
}
258+
}));
259+
}
260+
if ((tablesToMerge.size() & 1) == 1) {
261+
nextRoundTables.add(tablesToMerge.get(tablesToMerge.size() - 1));
262+
}
263+
try {
264+
for (Future<IndexedTable> future : futures) {
265+
long timeoutMs = _queryContext.getEndTimeMs() - System.currentTimeMillis();
266+
if (timeoutMs <= 0) {
267+
throw new TimeoutException("Timed out while reducing partitioned group-by results");
268+
}
269+
nextRoundTables.add(future.get(timeoutMs, TimeUnit.MILLISECONDS));
270+
}
271+
} catch (ExecutionException e) {
272+
throw new RuntimeException("Error while reducing partitioned group-by results", e.getCause());
273+
}
274+
tablesToMerge = nextRoundTables;
275+
}
276+
return tablesToMerge.get(0);
277+
}
278+
279+
private IndexedTable mergePartitionTablesSequentially(List<IndexedTable> partitionTables) {
280+
IndexedTable indexedTable = partitionTables.get(0);
281+
for (int i = 1; i < partitionTables.size(); i++) {
282+
indexedTable = mergePartitionTables(indexedTable, partitionTables.get(i));
283+
}
284+
return indexedTable;
285+
}
286+
287+
private static IndexedTable mergePartitionTables(IndexedTable leftTable, IndexedTable rightTable) {
288+
if (leftTable.size() >= rightTable.size()) {
289+
leftTable.mergePartitionTable(rightTable);
290+
return leftTable;
291+
}
292+
rightTable.mergePartitionTable(leftTable);
293+
return rightTable;
294+
}
295+
176296
private int getPartitionId(Key key) {
177297
int hashCode = key.hashCode();
178298
return _partitionMask >= 0 ? (hashCode & _partitionMask) : Math.floorMod(hashCode, _indexedTables.length);

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

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
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.queries;
20+
21+
import java.util.Map;
22+
import org.apache.pinot.common.response.broker.ResultTable;
23+
import org.apache.pinot.core.plan.maker.InstancePlanMakerImplV2;
24+
import org.testng.annotations.Test;
25+
26+
27+
/**
28+
* Reuses the inter-segment group-by order-by coverage with the partitioned combine operator enabled.
29+
*/
30+
public class InterSegmentPartitionedGroupBySingleValueQueriesTest
31+
extends InterSegmentNonBlockingGroupBySingleValueQueriesTest {
32+
private static final InstancePlanMakerImplV2 TRIM_ENABLED_PLAN_MAKER = new InstancePlanMakerImplV2();
33+
private static final Map<String, String> QUERY_OPTIONS =
34+
Map.of("groupByAlgorithm", "partitioned", "numGroupByPartitions", "8");
35+
36+
static {
37+
TRIM_ENABLED_PLAN_MAKER.setMinSegmentGroupTrimSize(1);
38+
}
39+
40+
@Override
41+
@Test(dataProvider = "groupByOrderByDataProvider")
42+
public void testGroupByOrderBy(String query, long expectedNumEntriesScannedPostFilter,
43+
ResultTable expectedResultTable) {
44+
QueriesTestUtils.testInterSegmentsResult(getBrokerResponse(query, QUERY_OPTIONS), 120000L, 0L,
45+
expectedNumEntriesScannedPostFilter, 120000L, expectedResultTable);
46+
}
47+
48+
@Override
49+
@Test(dataProvider = "groupByOrderByDataProvider")
50+
public void testGroupByOrderByWithTrim(String query, long expectedNumEntriesScannedPostFilter,
51+
ResultTable expectedResultTable) {
52+
QueriesTestUtils.testInterSegmentsResult(getBrokerResponse(query, TRIM_ENABLED_PLAN_MAKER, QUERY_OPTIONS),
53+
120000L, 0L, expectedNumEntriesScannedPostFilter, 120000L, expectedResultTable);
54+
}
55+
}

0 commit comments

Comments
 (0)