Skip to content

Commit 16b8c6e

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

16 files changed

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

‎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)