Skip to content

Commit 8b4c736

Browse files
jtuglu1jtuglu1
committed
perf: vectorize topN native engine
1 parent 8f3df74 commit 8b4c736

18 files changed

Lines changed: 2278 additions & 2 deletions

‎benchmarks/src/test/java/org/apache/druid/benchmark/query/TopNBenchmark.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121

2222
import com.fasterxml.jackson.core.JsonProcessingException;
2323
import com.fasterxml.jackson.databind.ObjectMapper;
24+
import com.google.common.collect.ImmutableMap;
2425
import org.apache.druid.collections.StupidPool;
2526
import org.apache.druid.jackson.DefaultObjectMapper;
2627
import org.apache.druid.java.util.common.FileUtils;
@@ -113,6 +114,9 @@ public class TopNBenchmark
113114
@Param({"all", "hour"})
114115
private String queryGranularity;
115116

117+
@Param({"false", "force"})
118+
private String vectorize;
119+
116120
private static final Logger log = new Logger(TopNBenchmark.class);
117121
private static final int RNG_SEED = 9999;
118122
private static final IndexMergerV9 INDEX_MERGER_V9;
@@ -161,6 +165,7 @@ private void setupQueries()
161165
.dimension("dimSequential")
162166
.metric("sumFloatNormal")
163167
.intervals(intervalSpec)
168+
.context(ImmutableMap.of("vectorize", vectorize))
164169
.aggregators(queryAggs);
165170

166171
basicQueries.put("A", queryBuilderA);
@@ -177,6 +182,7 @@ private void setupQueries()
177182
.dimension("dimUniform")
178183
.metric(new DimensionTopNMetricSpec(null, StringComparators.NUMERIC))
179184
.intervals(intervalSpec)
185+
.context(ImmutableMap.of("vectorize", vectorize))
180186
.aggregators(queryAggs);
181187

182188
basicQueries.put("numericSort", queryBuilderA);
@@ -193,6 +199,7 @@ private void setupQueries()
193199
.dimension("dimUniform")
194200
.metric(new DimensionTopNMetricSpec(null, StringComparators.ALPHANUMERIC))
195201
.intervals(intervalSpec)
202+
.context(ImmutableMap.of("vectorize", vectorize))
196203
.aggregators(queryAggs);
197204

198205
basicQueries.put("alphanumericSort", queryBuilderA);

‎processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/HeapVectorGrouper.java‎

Lines changed: 213 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,213 @@
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+
20+
package org.apache.druid.query.groupby.epinephelinae;
21+
22+
import it.unimi.dsi.fastutil.Hash;
23+
import it.unimi.dsi.fastutil.objects.Object2IntMap;
24+
import it.unimi.dsi.fastutil.objects.Object2IntOpenCustomHashMap;
25+
import org.apache.datasketches.memory.Memory;
26+
import org.apache.druid.java.util.common.ISE;
27+
import org.apache.druid.java.util.common.parsers.CloseableIterator;
28+
import org.apache.druid.query.aggregation.AggregatorAdapters;
29+
import org.apache.druid.query.groupby.epinephelinae.collection.MemoryPointer;
30+
31+
import java.nio.ByteBuffer;
32+
import java.nio.ByteOrder;
33+
import java.util.Arrays;
34+
import java.util.Iterator;
35+
36+
/**
37+
* On-heap {@link VectorGrouper} that grows aggregator state on demand, up to maximum limit of 2GB.
38+
*/
39+
public class HeapVectorGrouper implements VectorGrouper
40+
{
41+
private static final Hash.Strategy<byte[]> BYTE_ARRAY_HASH_STRATEGY = new Hash.Strategy<byte[]>()
42+
{
43+
@Override
44+
public int hashCode(byte[] o)
45+
{
46+
return Arrays.hashCode(o);
47+
}
48+
49+
@Override
50+
public boolean equals(byte[] a, byte[] b)
51+
{
52+
return Arrays.equals(a, b);
53+
}
54+
};
55+
56+
private static final int MIN_INITIAL_STATE_BUFFER_SIZE = 4096;
57+
58+
private final AggregatorAdapters aggregators;
59+
private final int keySize;
60+
private final int aggStateSize;
61+
private final Object2IntOpenCustomHashMap<byte[]> keyToOffset;
62+
63+
private boolean initialized;
64+
private ByteBuffer aggStateBuffer;
65+
private int aggStateEnd;
66+
67+
private int[] vAggregationPositions;
68+
private int[] vAggregationRows;
69+
private byte[] keyScratch;
70+
71+
public HeapVectorGrouper(final AggregatorAdapters aggregators, final int keySize)
72+
{
73+
this.aggregators = aggregators;
74+
this.keySize = keySize;
75+
this.aggStateSize = aggregators.spaceNeeded();
76+
this.keyToOffset = new Object2IntOpenCustomHashMap<>(BYTE_ARRAY_HASH_STRATEGY);
77+
this.keyToOffset.defaultReturnValue(-1);
78+
}
79+
80+
@Override
81+
public void initVectorized(final int maxVectorSize)
82+
{
83+
if (initialized) {
84+
if (vAggregationPositions.length != maxVectorSize) {
85+
throw new ISE(
86+
"initVectorized called with different maxVectorSize (existing=%d, new=%d)",
87+
vAggregationPositions.length,
88+
maxVectorSize
89+
);
90+
}
91+
return;
92+
}
93+
this.aggStateBuffer = ByteBuffer.allocate(MIN_INITIAL_STATE_BUFFER_SIZE);
94+
this.vAggregationPositions = new int[maxVectorSize];
95+
this.vAggregationRows = new int[maxVectorSize];
96+
this.keyScratch = new byte[keySize];
97+
this.aggStateEnd = 0;
98+
this.initialized = true;
99+
}
100+
101+
/**
102+
* Contract: keys for rows [startRow, endRow) must be packed contiguously at {@code keySpace[0 ..
103+
* numRows * keySize)}; {@code startRow}/{@code endRow} are source-vector indices used to look up aggregator
104+
* input values.
105+
*/
106+
@Override
107+
public AggregateResult aggregateVector(final Memory keySpace, final int startRow, final int endRow)
108+
{
109+
final int numRows = endRow - startRow;
110+
111+
for (int i = 0; i < numRows; i++) {
112+
keySpace.getByteArray((long) i * keySize, keyScratch, 0, keySize);
113+
int offset = keyToOffset.getInt(keyScratch);
114+
if (offset == -1) {
115+
if ((long) aggStateEnd + aggStateSize > aggStateBuffer.capacity()) {
116+
growBuffer((long) aggStateEnd + aggStateSize);
117+
}
118+
offset = aggStateEnd;
119+
final byte[] keyCopy = Arrays.copyOf(keyScratch, keySize);
120+
keyToOffset.put(keyCopy, offset);
121+
aggregators.init(aggStateBuffer, offset);
122+
aggStateEnd += aggStateSize;
123+
}
124+
vAggregationPositions[i] = offset;
125+
}
126+
127+
aggregators.aggregateVector(
128+
aggStateBuffer,
129+
numRows,
130+
vAggregationPositions,
131+
Groupers.writeAggregationRows(vAggregationRows, startRow, endRow)
132+
);
133+
134+
return AggregateResult.ok();
135+
}
136+
137+
private void growBuffer(final long neededCapacity)
138+
{
139+
if (neededCapacity > Integer.MAX_VALUE) {
140+
throw new ISE("Aggregator state exceeds 2 GB; cardinality too high for HeapVectorGrouper");
141+
}
142+
int newCapacity = aggStateBuffer.capacity();
143+
while (newCapacity < neededCapacity) {
144+
final long doubled = (long) newCapacity * 2;
145+
newCapacity = doubled > Integer.MAX_VALUE ? Integer.MAX_VALUE : (int) doubled;
146+
}
147+
148+
final ByteBuffer oldBuffer = aggStateBuffer;
149+
final ByteBuffer newBuffer = ByteBuffer.allocate(newCapacity);
150+
oldBuffer.position(0);
151+
oldBuffer.limit(aggStateEnd);
152+
newBuffer.put(oldBuffer);
153+
154+
for (int pos = 0; pos < aggStateEnd; pos += aggStateSize) {
155+
aggregators.relocate(pos, pos, oldBuffer, newBuffer);
156+
}
157+
158+
this.aggStateBuffer = newBuffer;
159+
}
160+
161+
@Override
162+
public void reset()
163+
{
164+
aggregators.reset();
165+
keyToOffset.clear();
166+
aggStateEnd = 0;
167+
}
168+
169+
@Override
170+
public void close()
171+
{
172+
aggregators.reset();
173+
keyToOffset.clear();
174+
}
175+
176+
@Override
177+
public CloseableIterator<Grouper.Entry<MemoryPointer>> iterator()
178+
{
179+
final Iterator<Object2IntMap.Entry<byte[]>> mapIter =
180+
keyToOffset.object2IntEntrySet().fastIterator();
181+
182+
return new CloseableIterator<>()
183+
{
184+
final ReusableEntry<MemoryPointer> reusableEntry =
185+
new ReusableEntry<>(new MemoryPointer(), new Object[aggregators.size()]);
186+
187+
@Override
188+
public boolean hasNext()
189+
{
190+
return mapIter.hasNext();
191+
}
192+
193+
@Override
194+
public Grouper.Entry<MemoryPointer> next()
195+
{
196+
final Object2IntMap.Entry<byte[]> mapEntry = mapIter.next();
197+
reusableEntry.getKey().set(Memory.wrap(mapEntry.getKey(), ByteOrder.nativeOrder()), 0);
198+
199+
final int position = mapEntry.getIntValue();
200+
for (int i = 0; i < aggregators.size(); i++) {
201+
reusableEntry.getValues()[i] = aggregators.get(aggStateBuffer, position, i);
202+
}
203+
return reusableEntry;
204+
}
205+
206+
@Override
207+
public void close()
208+
{
209+
// Nothing to close.
210+
}
211+
};
212+
}
213+
}

‎processing/src/main/java/org/apache/druid/query/topn/TopNQueryEngine.java‎

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import org.apache.druid.java.util.common.granularity.Granularities;
2727
import org.apache.druid.java.util.common.guava.Sequence;
2828
import org.apache.druid.java.util.common.guava.Sequences;
29+
import org.apache.druid.java.util.common.io.Closer;
2930
import org.apache.druid.query.ColumnSelectorPlus;
3031
import org.apache.druid.query.CursorGranularizer;
3132
import org.apache.druid.query.QueryContexts;
@@ -35,6 +36,7 @@
3536
import org.apache.druid.query.extraction.ExtractionFn;
3637
import org.apache.druid.query.topn.types.TopNColumnAggregatesProcessor;
3738
import org.apache.druid.query.topn.types.TopNColumnAggregatesProcessorFactory;
39+
import org.apache.druid.query.topn.vector.VectorTopNEngine;
3840
import org.apache.druid.segment.ColumnSelectorFactory;
3941
import org.apache.druid.segment.Cursor;
4042
import org.apache.druid.segment.CursorBuildSpec;
@@ -96,13 +98,34 @@ public Sequence<Result<TopNResultValue>> query(
9698
if (cursorHolder.isPreAggregated()) {
9799
query = query.withAggregatorSpecs(Preconditions.checkNotNull(cursorHolder.getAggregatorsForPreAggregated()));
98100
}
101+
102+
final TimeBoundaryInspector timeBoundaryInspector = segment.as(TimeBoundaryInspector.class);
103+
104+
final boolean canVectorize = cursorHolder.canVectorize()
105+
&& VectorTopNEngine.canVectorize(query, cursorFactory);
106+
final boolean shouldVectorize = query.context().getVectorize().shouldVectorize(canVectorize);
107+
108+
if (shouldVectorize) {
109+
final ResourceHolder<ByteBuffer> bufHolder = bufferPool.take();
110+
try {
111+
final Closer resourceCloser = Closer.create();
112+
resourceCloser.register(bufHolder);
113+
resourceCloser.register(cursorHolder);
114+
return Sequences.filter(
115+
VectorTopNEngine.process(query, timeBoundaryInspector, cursorHolder, bufHolder.get()),
116+
Predicates.notNull()
117+
).withBaggage(resourceCloser);
118+
}
119+
catch (Throwable t) {
120+
throw CloseableUtils.closeAndWrapInCatch(t, bufHolder);
121+
}
122+
}
123+
99124
final Cursor cursor = cursorHolder.asCursor();
100125
if (cursor == null) {
101126
return Sequences.withBaggage(Sequences.empty(), cursorHolder);
102127
}
103128

104-
final TimeBoundaryInspector timeBoundaryInspector = segment.as(TimeBoundaryInspector.class);
105-
106129
final ColumnSelectorFactory factory = cursor.getColumnSelectorFactory();
107130

108131
final ColumnSelectorPlus<TopNColumnAggregatesProcessor<?>> selectorPlus =

0 commit comments

Comments
 (0)