io.druid.query.groupby.epinephelinae.ConcurrentGrouper Maven / Gradle / Ivy
/*
* Licensed to Metamarkets Group Inc. (Metamarkets) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. Metamarkets 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 io.druid.query.groupby.epinephelinae;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.base.Preconditions;
import com.google.common.base.Supplier;
import com.google.common.base.Suppliers;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.ListeningExecutorService;
import io.druid.java.util.common.ISE;
import io.druid.query.AbstractPrioritizedCallable;
import io.druid.query.QueryInterruptedException;
import io.druid.query.aggregation.AggregatorFactory;
import io.druid.query.groupby.orderby.DefaultLimitSpec;
import io.druid.segment.ColumnSelectorFactory;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.CancellationException;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;
/**
* Grouper based around a set of underlying {@link SpillingGrouper} instances. Thread-safe.
*
* The passed-in buffer is cut up into concurrencyHint slices, and each slice is passed to a different underlying
* grouper. Access to each slice is separately synchronized. As long as the result set fits in memory, keys are
* partitioned between buffers based on their hash, and multiple threads can write into the same buffer. When
* it becomes clear that the result set does not fit in memory, the table switches to a mode where each thread
* gets its own buffer and its own spill files on disk.
*/
public class ConcurrentGrouper implements Grouper
{
private final List> groupers;
private final ThreadLocal> threadLocalGrouper;
private final AtomicInteger threadNumber = new AtomicInteger();
private volatile boolean spilling = false;
private volatile boolean closed = false;
private final Supplier bufferSupplier;
private final ColumnSelectorFactory columnSelectorFactory;
private final AggregatorFactory[] aggregatorFactories;
private final int bufferGrouperMaxSize;
private final float bufferGrouperMaxLoadFactor;
private final int bufferGrouperInitialBuckets;
private final LimitedTemporaryStorage temporaryStorage;
private final ObjectMapper spillMapper;
private final int concurrencyHint;
private final KeySerdeFactory keySerdeFactory;
private final DefaultLimitSpec limitSpec;
private final boolean sortHasNonGroupingFields;
private final Comparator> keyObjComparator;
private final ListeningExecutorService grouperSorter;
private final int priority;
private final boolean hasQueryTimeout;
private final long queryTimeoutAt;
private volatile boolean initialized = false;
public ConcurrentGrouper(
final Supplier bufferSupplier,
final KeySerdeFactory keySerdeFactory,
final ColumnSelectorFactory columnSelectorFactory,
final AggregatorFactory[] aggregatorFactories,
final int bufferGrouperMaxSize,
final float bufferGrouperMaxLoadFactor,
final int bufferGrouperInitialBuckets,
final LimitedTemporaryStorage temporaryStorage,
final ObjectMapper spillMapper,
final int concurrencyHint,
final DefaultLimitSpec limitSpec,
final boolean sortHasNonGroupingFields,
final ListeningExecutorService grouperSorter,
final int priority,
final boolean hasQueryTimeout,
final long queryTimeoutAt,
final int mergeBufferSize
)
{
Preconditions.checkArgument(concurrencyHint > 0, "concurrencyHint > 0");
this.groupers = new ArrayList<>(concurrencyHint);
this.threadLocalGrouper = new ThreadLocal>()
{
@Override
protected SpillingGrouper initialValue()
{
return groupers.get(threadNumber.getAndIncrement());
}
};
this.bufferSupplier = bufferSupplier;
this.columnSelectorFactory = columnSelectorFactory;
this.aggregatorFactories = aggregatorFactories;
this.bufferGrouperMaxSize = bufferGrouperMaxSize;
this.bufferGrouperMaxLoadFactor = bufferGrouperMaxLoadFactor;
this.bufferGrouperInitialBuckets = bufferGrouperInitialBuckets;
this.temporaryStorage = temporaryStorage;
this.spillMapper = spillMapper;
this.concurrencyHint = concurrencyHint;
this.keySerdeFactory = keySerdeFactory;
this.limitSpec = limitSpec;
this.sortHasNonGroupingFields = sortHasNonGroupingFields;
this.keyObjComparator = keySerdeFactory.objectComparator(sortHasNonGroupingFields);
this.grouperSorter = Preconditions.checkNotNull(grouperSorter);
this.priority = priority;
this.hasQueryTimeout = hasQueryTimeout;
this.queryTimeoutAt = queryTimeoutAt;
}
@Override
public void init()
{
if (!initialized) {
synchronized (bufferSupplier) {
if (!initialized) {
final ByteBuffer buffer = bufferSupplier.get();
final int sliceSize = (buffer.capacity() / concurrencyHint);
for (int i = 0; i < concurrencyHint; i++) {
final ByteBuffer slice = buffer.duplicate();
slice.position(sliceSize * i);
slice.limit(slice.position() + sliceSize);
final SpillingGrouper grouper = new SpillingGrouper<>(
Suppliers.ofInstance(slice.slice()),
keySerdeFactory,
columnSelectorFactory,
aggregatorFactories,
bufferGrouperMaxSize,
bufferGrouperMaxLoadFactor,
bufferGrouperInitialBuckets,
temporaryStorage,
spillMapper,
false,
limitSpec,
sortHasNonGroupingFields,
sliceSize
);
grouper.init();
groupers.add(grouper);
}
initialized = true;
}
}
}
}
@Override
public boolean isInitialized()
{
return initialized;
}
@Override
public AggregateResult aggregate(KeyType key, int keyHash)
{
if (!initialized) {
throw new ISE("Grouper is not initialized");
}
if (closed) {
throw new ISE("Grouper is closed");
}
if (!spilling) {
final SpillingGrouper hashBasedGrouper = groupers.get(grouperNumberForKeyHash(keyHash));
synchronized (hashBasedGrouper) {
if (!spilling) {
if (hashBasedGrouper.aggregate(key, keyHash).isOk()) {
return AggregateResult.ok();
} else {
spilling = true;
}
}
}
}
// At this point we know spilling = true
final SpillingGrouper tlGrouper = threadLocalGrouper.get();
synchronized (tlGrouper) {
tlGrouper.setSpillingAllowed(true);
return tlGrouper.aggregate(key, keyHash);
}
}
@Override
public void reset()
{
if (!initialized) {
throw new ISE("Grouper is not initialized");
}
if (closed) {
throw new ISE("Grouper is closed");
}
for (Grouper grouper : groupers) {
synchronized (grouper) {
grouper.reset();
}
}
}
@Override
public Iterator> iterator(final boolean sorted)
{
if (!initialized) {
throw new ISE("Grouper is not initialized");
}
if (closed) {
throw new ISE("Grouper is closed");
}
return Groupers.mergeIterators(
sorted && isParallelSortAvailable() ? parallelSortAndGetGroupersIterator() : getGroupersIterator(sorted),
sorted ? keyObjComparator : null
);
}
private boolean isParallelSortAvailable()
{
return concurrencyHint > 1;
}
private List>> parallelSortAndGetGroupersIterator()
{
// The number of groupers is same with the number of processing threads in grouperSorter
final ListenableFuture>>> future = Futures.allAsList(
groupers.stream()
.map(grouper ->
grouperSorter.submit(
new AbstractPrioritizedCallable>>(priority)
{
@Override
public Iterator> call() throws Exception
{
return grouper.iterator(true);
}
}
)
)
.collect(Collectors.toList())
);
try {
final long timeout = queryTimeoutAt - System.currentTimeMillis();
return hasQueryTimeout ? future.get(timeout, TimeUnit.MILLISECONDS) : future.get();
}
catch (InterruptedException | TimeoutException e) {
future.cancel(true);
throw new QueryInterruptedException(e);
}
catch (CancellationException e) {
throw new QueryInterruptedException(e);
}
catch (ExecutionException e) {
throw new RuntimeException(e.getCause());
}
}
private List>> getGroupersIterator(boolean sorted)
{
return groupers.stream()
.map(grouper -> grouper.iterator(sorted))
.collect(Collectors.toList());
}
@Override
public void close()
{
closed = true;
for (Grouper grouper : groupers) {
synchronized (grouper) {
grouper.close();
}
}
}
private int grouperNumberForKeyHash(int keyHash)
{
return keyHash % groupers.size();
}
}