Skip to content

Commit 40d0c26

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

13 files changed

Lines changed: 1493 additions & 2 deletions

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

Lines changed: 218 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,218 @@
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+
aggStateBuffer = null;
176+
vAggregationPositions = null;
177+
vAggregationRows = null;
178+
keyScratch = null;
179+
}
180+
181+
@Override
182+
public CloseableIterator<Grouper.Entry<MemoryPointer>> iterator()
183+
{
184+
final Iterator<Object2IntMap.Entry<byte[]>> mapIter =
185+
keyToOffset.object2IntEntrySet().fastIterator();
186+
187+
return new CloseableIterator<>()
188+
{
189+
final ReusableEntry<MemoryPointer> reusableEntry =
190+
new ReusableEntry<>(new MemoryPointer(), new Object[aggregators.size()]);
191+
192+
@Override
193+
public boolean hasNext()
194+
{
195+
return mapIter.hasNext();
196+
}
197+
198+
@Override
199+
public Grouper.Entry<MemoryPointer> next()
200+
{
201+
final Object2IntMap.Entry<byte[]> mapEntry = mapIter.next();
202+
reusableEntry.getKey().set(Memory.wrap(mapEntry.getKey(), ByteOrder.nativeOrder()), 0);
203+
204+
final int position = mapEntry.getIntValue();
205+
for (int i = 0; i < aggregators.size(); i++) {
206+
reusableEntry.getValues()[i] = aggregators.get(aggStateBuffer, position, i);
207+
}
208+
return reusableEntry;
209+
}
210+
211+
@Override
212+
public void close()
213+
{
214+
// Nothing to close.
215+
}
216+
};
217+
}
218+
}

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

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

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
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.topn.vector;
21+
22+
import it.unimi.dsi.fastutil.objects.Object2IntMap;
23+
import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap;
24+
import org.apache.datasketches.memory.WritableMemory;
25+
import org.apache.druid.query.groupby.epinephelinae.collection.MemoryPointer;
26+
import org.apache.druid.segment.DimensionHandlerUtils;
27+
import org.apache.druid.segment.vector.VectorObjectSelector;
28+
29+
import javax.annotation.Nullable;
30+
import java.util.ArrayList;
31+
import java.util.List;
32+
33+
/**
34+
* {@link TopNVectorColumnSelector} for single-valued STRING columns that are not natively dictionary-encoded,
35+
* such as expression virtual columns. Builds a local int dictionary on-the-fly and encodes keys as 4-byte
36+
* dictionary IDs, matching the key format of {@link SingleValueStringTopNVectorColumnSelector}.
37+
*/
38+
public class DictionaryBuildingSingleValueStringTopNVectorColumnSelector implements TopNVectorColumnSelector
39+
{
40+
private final VectorObjectSelector selector;
41+
private final List<String> dictionary = new ArrayList<>();
42+
private final Object2IntMap<String> reverseDictionary = new Object2IntOpenHashMap<>();
43+
44+
DictionaryBuildingSingleValueStringTopNVectorColumnSelector(final VectorObjectSelector selector)
45+
{
46+
this.selector = selector;
47+
reverseDictionary.defaultReturnValue(-1);
48+
}
49+
50+
@Override
51+
public int getGroupingKeySize()
52+
{
53+
return Integer.BYTES;
54+
}
55+
56+
@Override
57+
public void writeKeys(
58+
final WritableMemory keySpace,
59+
final int keySize,
60+
final int keyOffset,
61+
final int startRow,
62+
final int endRow
63+
)
64+
{
65+
final Object[] vector = selector.getObjectVector();
66+
67+
for (int i = startRow, j = keyOffset; i < endRow; i++, j += keySize) {
68+
final String value = DimensionHandlerUtils.convertObjectToString(vector[i]);
69+
int dictId = reverseDictionary.getInt(value);
70+
if (dictId < 0) {
71+
dictId = dictionary.size();
72+
dictionary.add(value);
73+
reverseDictionary.put(value, dictId);
74+
}
75+
keySpace.putInt(j, dictId);
76+
}
77+
}
78+
79+
@Override
80+
@Nullable
81+
public Object getDimensionValue(final MemoryPointer keyMemory, final int keyOffset)
82+
{
83+
return dictionary.get(keyMemory.memory().getInt(keyMemory.position() + keyOffset));
84+
}
85+
86+
@Override
87+
public void reset()
88+
{
89+
dictionary.clear();
90+
reverseDictionary.clear();
91+
}
92+
}

0 commit comments

Comments
 (0)