org.apache.druid.query.groupby.epinephelinae.BufferHashGrouper Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of druid-processing Show documentation
Show all versions of druid-processing Show documentation
A module that is everything required to understands Druid Segments
/*
* 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 com.google.common.base.Supplier;
import org.apache.druid.java.util.common.CloseableIterators;
import org.apache.druid.java.util.common.IAE;
import org.apache.druid.java.util.common.parsers.CloseableIterator;
import org.apache.druid.query.aggregation.AggregatorAdapters;
import org.apache.druid.query.aggregation.AggregatorFactory;
import javax.annotation.Nullable;
import java.nio.ByteBuffer;
import java.util.AbstractList;
import java.util.Collections;
import java.util.List;
import java.util.NoSuchElementException;
import java.util.function.ToIntFunction;
public class BufferHashGrouper extends AbstractBufferHashGrouper
{
private static final int MIN_INITIAL_BUCKETS = 4;
private static final int DEFAULT_INITIAL_BUCKETS = 1024;
private static final float DEFAULT_MAX_LOAD_FACTOR = 0.7f;
private boolean initialized = false;
// The BufferHashGrouper normally sorts by all fields of the grouping key with lexicographic ascending order.
// However, when a query will have the limit push down optimization applied (see LimitedBufferHashGrouper),
// the optimization may not be applied on some nodes because of buffer capacity limits. In this case,
// those nodes will use BufferHashGrouper instead of LimitedBufferHashGrouper. In this mixed use case,
// nodes using BufferHashGrouper need to use the same sorting order as nodes using LimitedBufferHashGrouper, so that
// results are merged properly. When useDefaultSorting is false, we call keySerde.bufferComparatorWithAggregators()
// to get a comparator that uses the ordering defined by the OrderByColumnSpec of a query.
private final boolean useDefaultSorting;
@Nullable
private ByteBufferIntList offsetList;
public BufferHashGrouper(
final Supplier bufferSupplier,
final KeySerde keySerde,
final AggregatorAdapters aggregators,
final int bufferGrouperMaxSize,
final float maxLoadFactor,
final int initialBuckets,
final boolean useDefaultSorting
)
{
super(bufferSupplier, keySerde, aggregators, HASH_SIZE + keySerde.keySize(), bufferGrouperMaxSize);
this.maxLoadFactor = maxLoadFactor > 0 ? maxLoadFactor : DEFAULT_MAX_LOAD_FACTOR;
this.initialBuckets = initialBuckets > 0 ? Math.max(MIN_INITIAL_BUCKETS, initialBuckets) : DEFAULT_INITIAL_BUCKETS;
if (this.maxLoadFactor >= 1.0f) {
throw new IAE("Invalid maxLoadFactor[%f], must be < 1.0", maxLoadFactor);
}
this.bucketSize = HASH_SIZE + keySerde.keySize() + aggregators.spaceNeeded();
this.useDefaultSorting = useDefaultSorting;
}
@Override
public void init()
{
if (!initialized) {
ByteBuffer buffer = bufferSupplier.get();
int hashTableSize = ByteBufferHashTable.calculateTableArenaSizeWithPerBucketAdditionalSize(
buffer.capacity(),
bucketSize,
Integer.BYTES
);
hashTableBuffer = buffer.duplicate();
hashTableBuffer.position(0);
hashTableBuffer.limit(hashTableSize);
hashTableBuffer = hashTableBuffer.slice();
// Track the offsets of used buckets using this list.
// When a new bucket is initialized by initializeNewBucketKey(), an offset is added to this list.
// When expanding the table, the list is reset() and filled with the new offsets of the copied buckets.
ByteBuffer offsetListBuffer = buffer.duplicate();
offsetListBuffer.position(hashTableSize);
offsetListBuffer.limit(buffer.capacity());
offsetListBuffer = offsetListBuffer.slice();
this.offsetList = new ByteBufferIntList(
offsetListBuffer,
offsetListBuffer.capacity() / Integer.BYTES
);
this.hashTable = new ByteBufferHashTable(
maxLoadFactor,
initialBuckets,
bucketSize,
hashTableBuffer,
keySize,
bufferGrouperMaxSize,
new BufferGrouperBucketUpdateHandler()
);
reset();
initialized = true;
}
}
@Override
public boolean isInitialized()
{
return initialized;
}
@Override
public ToIntFunction hashFunction()
{
return Groupers::hashObject;
}
@Override
public void newBucketHook(int bucketOffset)
{
// Nothing needed.
}
@Override
public boolean canSkipAggregate(int bucketOffset)
{
return false;
}
@Override
public void afterAggregateHook(int bucketOffset)
{
// Nothing needed.
}
@Override
public void reset()
{
offsetList.reset();
hashTable.reset();
keySerde.reset();
}
@Override
public CloseableIterator> iterator(boolean sorted)
{
if (!initialized) {
// it's possible for iterator() to be called before initialization when
// a nested groupBy's subquery has an empty result set (see testEmptySubquery() in GroupByQueryRunnerTest)
return CloseableIterators.withEmptyBaggage(Collections.emptyIterator());
}
if (sorted) {
@SuppressWarnings("MismatchedQueryAndUpdateOfCollection")
final List wrappedOffsets = new AbstractList()
{
@Override
public Integer get(int index)
{
return offsetList.get(index);
}
@Override
public Integer set(int index, Integer element)
{
final Integer oldValue = get(index);
offsetList.set(index, element);
return oldValue;
}
@Override
public int size()
{
return hashTable.getSize();
}
};
final BufferComparator comparator;
if (useDefaultSorting) {
comparator = keySerde.bufferComparator();
} else {
comparator = keySerde.bufferComparatorWithAggregators(
aggregators.factories().toArray(new AggregatorFactory[0]),
aggregators.aggregatorPositions()
);
}
// Sort offsets in-place.
Collections.sort(
wrappedOffsets,
(lhs, rhs) -> {
final ByteBuffer tableBuffer = hashTable.getTableBuffer();
return comparator.compare(
tableBuffer,
tableBuffer,
lhs + HASH_SIZE,
rhs + HASH_SIZE
);
}
);
return new CloseableIterator>()
{
final ReusableEntry reusableEntry = ReusableEntry.create(keySerde, aggregators.size());
int curr = 0;
final int size = getSize();
@Override
public boolean hasNext()
{
return curr < size;
}
@Override
public Entry next()
{
if (curr >= size) {
throw new NoSuchElementException();
}
return populateBucketEntryForOffset(reusableEntry, wrappedOffsets.get(curr++));
}
@Override
public void remove()
{
throw new UnsupportedOperationException();
}
@Override
public void close()
{
// do nothing
}
};
} else {
// Unsorted iterator
return new CloseableIterator>()
{
final ReusableEntry reusableEntry = ReusableEntry.create(keySerde, aggregators.size());
int curr = 0;
final int size = getSize();
@Override
public boolean hasNext()
{
return curr < size;
}
@Override
public Entry next()
{
if (curr >= size) {
throw new NoSuchElementException();
}
final int offset = offsetList.get(curr);
final Entry entry = populateBucketEntryForOffset(reusableEntry, offset);
curr++;
return entry;
}
@Override
public void remove()
{
throw new UnsupportedOperationException();
}
@Override
public void close()
{
// do nothing
}
};
}
}
private class BufferGrouperBucketUpdateHandler implements ByteBufferHashTable.BucketUpdateHandler
{
@Override
public void handleNewBucket(int bucketOffset)
{
offsetList.add(bucketOffset);
}
@Override
public void handlePreTableSwap()
{
offsetList.reset();
}
@Override
public void handleBucketMove(int oldBucketOffset, int newBucketOffset, ByteBuffer oldBuffer, ByteBuffer newBuffer)
{
// relocate aggregators (see https://github.com/apache/druid/pull/4071)
aggregators.relocate(
oldBucketOffset + baseAggregatorOffset,
newBucketOffset + baseAggregatorOffset,
oldBuffer,
newBuffer
);
offsetList.add(newBucketOffset);
}
}
}