gianm commented on code in PR #19353:
URL: https://github.com/apache/druid/pull/19353#discussion_r3264299994
##########
processing/src/main/java/org/apache/druid/query/topn/TopNQueryEngine.java:
##########
@@ -96,13 +98,34 @@ public Sequence<Result<TopNResultValue>> query(
if (cursorHolder.isPreAggregated()) {
query =
query.withAggregatorSpecs(Preconditions.checkNotNull(cursorHolder.getAggregatorsForPreAggregated()));
}
+
+ final TimeBoundaryInspector timeBoundaryInspector =
segment.as(TimeBoundaryInspector.class);
+
+ final boolean canVectorize = cursorHolder.canVectorize()
+ && VectorTopNEngine.canVectorize(query,
cursorFactory);
+ final boolean shouldVectorize =
query.context().getVectorize().shouldVectorize(canVectorize);
Review Comment:
I'd like to have a `vectorizeTopN` that is separate from `vectorize`, and
*both* must be set in order to vectorize topN. The idea is that if someone is
having issues with vectorized topNs, they can set `vectorizeTopN: false` across
the board rather than have to micro-manage the `vectorize` setting per-query.
Which isn't generally practical anyway.
##########
processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/HeapVectorGrouper.java:
##########
@@ -0,0 +1,218 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.query.groupby.epinephelinae;
+
+import it.unimi.dsi.fastutil.Hash;
+import it.unimi.dsi.fastutil.objects.Object2IntMap;
+import it.unimi.dsi.fastutil.objects.Object2IntOpenCustomHashMap;
+import org.apache.datasketches.memory.Memory;
+import org.apache.druid.java.util.common.ISE;
+import org.apache.druid.java.util.common.parsers.CloseableIterator;
+import org.apache.druid.query.aggregation.AggregatorAdapters;
+import org.apache.druid.query.groupby.epinephelinae.collection.MemoryPointer;
+
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.util.Arrays;
+import java.util.Iterator;
+
+/**
+ * On-heap {@link VectorGrouper} that grows aggregator state on demand, up to
maximum limit of 2GB.
+ *
+ * Vectorized analogue of {@link
org.apache.druid.query.topn.BaseTopNAlgorithm}'s
+ * {@code runWithCardinalityUnknown} path: used when dimension cardinality is
unknown (numeric columns,
+ * non-dict-encoded string virtual columns) or when a dict-encoded string
column's cardinality exceeds
+ * the processing buffer. Memory footprint is on-heap and grows with the
distinct-key count — matching
+ * the non-vectorized path's memory profile for the same queries.
+ */
+public class HeapVectorGrouper implements VectorGrouper
Review Comment:
This seems kind of similar to `HashVectorGrouper`, used by `groupBy`. The
difference I can see is that this one uses heap buffers rather than the
processing buffer. It'd be nice to share code and it'd be nice to use the
processing buffer here if possible.
Other than that, I'm also concerned that there are some aggregators which
tend to take up much more space for the `BufferAggregator` / `VectorAggregator`
form than for the `Aggregator` form for small numbers of aggregated data points
(certain sketches, especially theta sketches but also sometimes hll). With
`groupBy` we spill these and then the spill files tend to be very small. The
problem is described in #19357 and #19439.
So, overflowing into the heap when the processing buffer is full, but
continuing to use `VectorAggregator` in-heap, can cause large amounts of heap
to be used. I wonder if this would work well:
- Start `HashVectorGrouper` (the one from `groupBy`) in the processing buffer
- If aggregation completes with `HashVectorGrouper`, emit that
- If aggregation stops due to the processing buffer being full, "spill" to
heap and then continue aggregating in the processing buffer.
- If the total in-heap "spill" memory exceeds some threshold then fail the
query.
##########
processing/src/main/java/org/apache/druid/query/topn/vector/VectorTopNEngine.java:
##########
@@ -0,0 +1,373 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.query.topn.vector;
+
+import com.google.common.base.Suppliers;
+import org.apache.datasketches.memory.WritableMemory;
+import org.apache.druid.java.util.common.guava.BaseSequence;
+import org.apache.druid.java.util.common.guava.Sequence;
+import org.apache.druid.java.util.common.parsers.CloseableIterator;
+import org.apache.druid.query.Order;
+import org.apache.druid.query.Result;
+import org.apache.druid.query.aggregation.AggregatorAdapters;
+import org.apache.druid.query.aggregation.AggregatorFactory;
+import org.apache.druid.query.dimension.DimensionSpec;
+import org.apache.druid.query.groupby.epinephelinae.BufferArrayGrouper;
+import org.apache.druid.query.groupby.epinephelinae.Grouper;
+import org.apache.druid.query.groupby.epinephelinae.HeapVectorGrouper;
+import org.apache.druid.query.groupby.epinephelinae.VectorGrouper;
+import org.apache.druid.query.groupby.epinephelinae.collection.MemoryPointer;
+import org.apache.druid.query.topn.TopNQuery;
+import org.apache.druid.query.topn.TopNResultBuilder;
+import org.apache.druid.query.topn.TopNResultValue;
+import org.apache.druid.query.vector.VectorCursorGranularizer;
+import org.apache.druid.segment.ColumnInspector;
+import org.apache.druid.segment.ColumnProcessors;
+import org.apache.druid.segment.CursorHolder;
+import org.apache.druid.segment.TimeBoundaryInspector;
+import org.apache.druid.segment.column.ColumnCapabilities;
+import org.apache.druid.segment.column.Types;
+import org.apache.druid.segment.column.ValueType;
+import org.apache.druid.segment.vector.VectorColumnSelectorFactory;
+import org.apache.druid.segment.vector.VectorCursor;
+import org.apache.druid.segment.virtual.VirtualizedColumnInspector;
+import org.joda.time.DateTime;
+import org.joda.time.Interval;
+
+import javax.annotation.Nullable;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.Iterator;
+import java.util.List;
+import java.util.NoSuchElementException;
+
+/**
+ * Vectorized execution engine for {@link TopNQuery}, analogous to
+ * {@link
org.apache.druid.query.groupby.epinephelinae.vector.VectorGroupByEngine} for
groupBy.
+ *
+ * Uses a {@link VectorGrouper} for batch aggregation (with vectorized null
handling via
+ * {@link
org.apache.druid.segment.vector.VectorValueSelector#getNullVector()}) and then
applies top-N
+ * ordering via {@link TopNResultBuilder} after each time-bucket is fully
aggregated.
+ *
+ * @see org.apache.druid.query.topn.TopNQueryEngine for the entry point that
selects this path
+ */
+public class VectorTopNEngine
+{
+ private VectorTopNEngine()
+ {
+ // No instantiation.
+ }
+
+ public static Sequence<Result<TopNResultValue>> process(
+ final TopNQuery query,
+ @Nullable final TimeBoundaryInspector timeBoundaryInspector,
+ final CursorHolder cursorHolder,
+ final ByteBuffer processingBuffer
+ )
+ {
+ return new BaseSequence<>(
+ new BaseSequence.IteratorMaker<Result<TopNResultValue>,
CloseableIterator<Result<TopNResultValue>>>()
+ {
+ @Override
+ public CloseableIterator<Result<TopNResultValue>> make()
+ {
+ final VectorCursor cursor = cursorHolder.asVectorCursor();
+
+ if (cursor == null) {
+ return new CloseableIterator<>()
+ {
+ @Override
+ public boolean hasNext()
+ {
+ return false;
+ }
+
+ @Override
+ public Result<TopNResultValue> next()
+ {
+ throw new NoSuchElementException();
+ }
+
+ @Override
+ public void close()
+ {
+ // Nothing to do.
+ }
+ };
+ }
+
+ final VectorColumnSelectorFactory columnSelectorFactory =
cursor.getColumnSelectorFactory();
+ final TopNVectorColumnSelector selector =
ColumnProcessors.makeVectorProcessor(
+ query.getDimensionSpec(),
+ TopNVectorColumnProcessorFactory.instance(),
+ columnSelectorFactory
+ );
+
+ return new VectorTopNEngineIterator(
+ query,
+ timeBoundaryInspector,
+ cursor,
+ cursorHolder.getTimeOrder(),
+ selector,
+ processingBuffer
+ );
+ }
+
+ @Override
+ public void cleanup(final CloseableIterator<Result<TopNResultValue>>
iterFromMake)
+ {
+ try {
+ iterFromMake.close();
+ }
+ catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ }
+ }
+ );
+ }
+
+ /**
+ * Returns true if the given query is eligible for the vectorized topN path.
+ */
+ public static boolean canVectorize(final TopNQuery query, final
ColumnInspector inspector)
+ {
+ final DimensionSpec dimensionSpec = query.getDimensionSpec();
+
+ if (!dimensionSpec.canVectorize()) {
+ return false;
+ }
+
+ // Decorated specs (e.g. extraction functions that are not one-to-one)
change value semantics in ways that
Review Comment:
This comment doesn't sound right: `ExtractionDimensionSpec` does not in fact
decorate. (It returns `false` from `mustDecorate`.) There may be some other
reason why decoration isn't supported?
##########
processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/HeapVectorGrouper.java:
##########
@@ -0,0 +1,218 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.query.groupby.epinephelinae;
+
+import it.unimi.dsi.fastutil.Hash;
+import it.unimi.dsi.fastutil.objects.Object2IntMap;
+import it.unimi.dsi.fastutil.objects.Object2IntOpenCustomHashMap;
+import org.apache.datasketches.memory.Memory;
+import org.apache.druid.java.util.common.ISE;
+import org.apache.druid.java.util.common.parsers.CloseableIterator;
+import org.apache.druid.query.aggregation.AggregatorAdapters;
+import org.apache.druid.query.groupby.epinephelinae.collection.MemoryPointer;
+
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.util.Arrays;
+import java.util.Iterator;
+
+/**
+ * On-heap {@link VectorGrouper} that grows aggregator state on demand, up to
maximum limit of 2GB.
+ *
+ * Vectorized analogue of {@link
org.apache.druid.query.topn.BaseTopNAlgorithm}'s
+ * {@code runWithCardinalityUnknown} path: used when dimension cardinality is
unknown (numeric columns,
+ * non-dict-encoded string virtual columns) or when a dict-encoded string
column's cardinality exceeds
+ * the processing buffer. Memory footprint is on-heap and grows with the
distinct-key count — matching
+ * the non-vectorized path's memory profile for the same queries.
+ */
+public class HeapVectorGrouper implements VectorGrouper
+{
+ private static final Hash.Strategy<byte[]> BYTE_ARRAY_HASH_STRATEGY = new
Hash.Strategy<byte[]>()
+ {
+ @Override
+ public int hashCode(byte[] o)
+ {
+ return Arrays.hashCode(o);
+ }
+
+ @Override
+ public boolean equals(byte[] a, byte[] b)
+ {
+ return Arrays.equals(a, b);
+ }
+ };
+
+ private static final int MIN_INITIAL_STATE_BUFFER_SIZE = 4096;
+
+ private final AggregatorAdapters aggregators;
+ private final int keySize;
+ private final int aggStateSize;
+ private final Object2IntOpenCustomHashMap<byte[]> keyToOffset;
+
+ private boolean initialized;
+ private ByteBuffer aggStateBuffer;
+ private int aggStateEnd;
+
+ private int[] vAggregationPositions;
+ private int[] vAggregationRows;
+ private byte[] keyScratch;
+
+ public HeapVectorGrouper(final AggregatorAdapters aggregators, final int
keySize)
+ {
+ this.aggregators = aggregators;
+ this.keySize = keySize;
+ this.aggStateSize = aggregators.spaceNeeded();
+ this.keyToOffset = new
Object2IntOpenCustomHashMap<>(BYTE_ARRAY_HASH_STRATEGY);
+ this.keyToOffset.defaultReturnValue(-1);
+ }
+
+ @Override
+ public void initVectorized(final int maxVectorSize)
+ {
+ if (initialized) {
+ if (vAggregationPositions.length != maxVectorSize) {
+ throw new ISE(
+ "initVectorized called with different maxVectorSize (existing=%d,
new=%d)",
+ vAggregationPositions.length,
+ maxVectorSize
+ );
+ }
+ return;
+ }
+ this.aggStateBuffer = ByteBuffer.allocate(MIN_INITIAL_STATE_BUFFER_SIZE);
+ this.vAggregationPositions = new int[maxVectorSize];
+ this.vAggregationRows = new int[maxVectorSize];
+ this.keyScratch = new byte[keySize];
+ this.aggStateEnd = 0;
+ this.initialized = true;
+ }
+
+ /**
+ * Contract: keys for rows [startRow, endRow) must be packed contiguously at
{@code keySpace[0 ..
+ * numRows * keySize)}; {@code startRow}/{@code endRow} are source-vector
indices used to look up aggregator
+ * input values.
+ */
+ @Override
+ public AggregateResult aggregateVector(final Memory keySpace, final int
startRow, final int endRow)
+ {
+ final int numRows = endRow - startRow;
+
+ for (int i = 0; i < numRows; i++) {
+ keySpace.getByteArray((long) i * keySize, keyScratch, 0, keySize);
+ int offset = keyToOffset.getInt(keyScratch);
+ if (offset == -1) {
+ if ((long) aggStateEnd + aggStateSize > aggStateBuffer.capacity()) {
+ growBuffer((long) aggStateEnd + aggStateSize);
+ }
+ offset = aggStateEnd;
+ final byte[] keyCopy = Arrays.copyOf(keyScratch, keySize);
+ keyToOffset.put(keyCopy, offset);
+ aggregators.init(aggStateBuffer, offset);
+ aggStateEnd += aggStateSize;
+ }
+ vAggregationPositions[i] = offset;
+ }
+
+ aggregators.aggregateVector(
+ aggStateBuffer,
+ numRows,
+ vAggregationPositions,
+ Groupers.writeAggregationRows(vAggregationRows, startRow, endRow)
+ );
+
+ return AggregateResult.ok();
+ }
+
+ private void growBuffer(final long neededCapacity)
+ {
+ if (neededCapacity > Integer.MAX_VALUE) {
+ throw new ISE("Aggregator state exceeds 2 GB; cardinality too high for
HeapVectorGrouper");
Review Comment:
This should be a `DruidException` with `Persona.USER` and
`Category.RUNTIME_FAILURE`. It should include some advice about how to fix it,
such as setting `vectorizeTopN: false` in query context.
##########
processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/HeapVectorGrouper.java:
##########
@@ -0,0 +1,213 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.query.groupby.epinephelinae;
+
+import it.unimi.dsi.fastutil.Hash;
+import it.unimi.dsi.fastutil.objects.Object2IntMap;
+import it.unimi.dsi.fastutil.objects.Object2IntOpenCustomHashMap;
+import org.apache.datasketches.memory.Memory;
+import org.apache.druid.java.util.common.ISE;
+import org.apache.druid.java.util.common.parsers.CloseableIterator;
+import org.apache.druid.query.aggregation.AggregatorAdapters;
+import org.apache.druid.query.groupby.epinephelinae.collection.MemoryPointer;
+
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.util.Arrays;
+import java.util.Iterator;
+
+/**
+ * On-heap {@link VectorGrouper} that grows aggregator state on demand, up to
maximum limit of 2GB.
+ */
+public class HeapVectorGrouper implements VectorGrouper
+{
+ private static final Hash.Strategy<byte[]> BYTE_ARRAY_HASH_STRATEGY = new
Hash.Strategy<byte[]>()
+ {
+ @Override
+ public int hashCode(byte[] o)
+ {
+ return Arrays.hashCode(o);
+ }
+
+ @Override
+ public boolean equals(byte[] a, byte[] b)
+ {
+ return Arrays.equals(a, b);
+ }
+ };
+
+ private static final int MIN_INITIAL_STATE_BUFFER_SIZE = 4096;
+
+ private final AggregatorAdapters aggregators;
+ private final int keySize;
+ private final int aggStateSize;
+ private final Object2IntOpenCustomHashMap<byte[]> keyToOffset;
+
+ private boolean initialized;
+ private ByteBuffer aggStateBuffer;
+ private int aggStateEnd;
+
+ private int[] vAggregationPositions;
+ private int[] vAggregationRows;
+ private byte[] keyScratch;
+
+ public HeapVectorGrouper(final AggregatorAdapters aggregators, final int
keySize)
+ {
+ this.aggregators = aggregators;
+ this.keySize = keySize;
+ this.aggStateSize = aggregators.spaceNeeded();
+ this.keyToOffset = new
Object2IntOpenCustomHashMap<>(BYTE_ARRAY_HASH_STRATEGY);
+ this.keyToOffset.defaultReturnValue(-1);
+ }
+
+ @Override
+ public void initVectorized(final int maxVectorSize)
+ {
+ if (initialized) {
+ if (vAggregationPositions.length != maxVectorSize) {
+ throw new ISE(
+ "initVectorized called with different maxVectorSize (existing=%d,
new=%d)",
+ vAggregationPositions.length,
+ maxVectorSize
+ );
+ }
+ return;
+ }
+ this.aggStateBuffer = ByteBuffer.allocate(MIN_INITIAL_STATE_BUFFER_SIZE);
+ this.vAggregationPositions = new int[maxVectorSize];
+ this.vAggregationRows = new int[maxVectorSize];
+ this.keyScratch = new byte[keySize];
+ this.aggStateEnd = 0;
+ this.initialized = true;
+ }
+
+ /**
+ * Contract: keys for rows [startRow, endRow) must be packed contiguously at
{@code keySpace[0 ..
+ * numRows * keySize)}; {@code startRow}/{@code endRow} are source-vector
indices used to look up aggregator
+ * input values.
+ */
+ @Override
+ public AggregateResult aggregateVector(final Memory keySpace, final int
startRow, final int endRow)
+ {
+ final int numRows = endRow - startRow;
+
+ for (int i = 0; i < numRows; i++) {
+ keySpace.getByteArray((long) i * keySize, keyScratch, 0, keySize);
+ int offset = keyToOffset.getInt(keyScratch);
+ if (offset == -1) {
+ if ((long) aggStateEnd + aggStateSize > aggStateBuffer.capacity()) {
+ growBuffer((long) aggStateEnd + aggStateSize);
+ }
+ offset = aggStateEnd;
+ final byte[] keyCopy = Arrays.copyOf(keyScratch, keySize);
+ keyToOffset.put(keyCopy, offset);
+ aggregators.init(aggStateBuffer, offset);
+ aggStateEnd += aggStateSize;
+ }
+ vAggregationPositions[i] = offset;
+ }
+
+ aggregators.aggregateVector(
+ aggStateBuffer,
+ numRows,
+ vAggregationPositions,
+ Groupers.writeAggregationRows(vAggregationRows, startRow, endRow)
+ );
+
+ return AggregateResult.ok();
+ }
+
+ private void growBuffer(final long neededCapacity)
+ {
+ if (neededCapacity > Integer.MAX_VALUE) {
Review Comment:
Agreed, this should be configurable. It will help with testing, since we can
configure a low limit to test limit-hitting behavior without having to allocate
2 GB in the tests. It will also be useful in production since 2 GB per
processing thread can add up really quickly.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
