/*
* Licensed 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 com.facebook.presto.operator.aggregation.builder;
import com.facebook.presto.array.IntBigArray;
import com.facebook.presto.memory.LocalMemoryContext;
import com.facebook.presto.operator.GroupByHash;
import com.facebook.presto.operator.GroupByIdBlock;
import com.facebook.presto.operator.HashCollisionsCounter;
import com.facebook.presto.operator.OperatorContext;
import com.facebook.presto.operator.aggregation.AccumulatorFactory;
import com.facebook.presto.operator.aggregation.GroupedAccumulator;
import com.facebook.presto.spi.Page;
import com.facebook.presto.spi.PageBuilder;
import com.facebook.presto.spi.block.BlockBuilder;
import com.facebook.presto.spi.type.Type;
import com.facebook.presto.sql.gen.JoinCompiler;
import com.facebook.presto.sql.planner.plan.AggregationNode;
import com.facebook.presto.sql.planner.plan.AggregationNode.Step;
import com.google.common.collect.AbstractIterator;
import com.google.common.collect.ImmutableList;
import com.google.common.primitives.Ints;
import com.google.common.util.concurrent.ListenableFuture;
import io.airlift.units.DataSize;
import it.unimi.dsi.fastutil.ints.AbstractIntIterator;
import it.unimi.dsi.fastutil.ints.IntIterator;
import it.unimi.dsi.fastutil.ints.IntIterators;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Optional;
import static com.facebook.presto.operator.GroupByHash.createGroupByHash;
import static com.facebook.presto.operator.Operator.NOT_BLOCKED;
import static com.facebook.presto.spi.type.BigintType.BIGINT;
import static com.google.common.base.Preconditions.checkArgument;
import static java.util.Objects.requireNonNull;
public class InMemoryHashAggregationBuilder
implements HashAggregationBuilder
{
private final GroupByHash groupByHash;
private final List<Aggregator> aggregators;
private final OperatorContext operatorContext;
private final boolean partial;
private final long maxPartialMemory;
private final LocalMemoryContext systemMemoryContext;
private boolean full;
public InMemoryHashAggregationBuilder(
List<AccumulatorFactory> accumulatorFactories,
Step step,
int expectedGroups,
List<Type> groupByTypes,
List<Integer> groupByChannels,
Optional<Integer> hashChannel,
OperatorContext operatorContext,
DataSize maxPartialMemory,
JoinCompiler joinCompiler)
{
this(accumulatorFactories,
step,
expectedGroups,
groupByTypes,
groupByChannels,
hashChannel,
operatorContext,
maxPartialMemory,
Optional.empty(),
joinCompiler);
}
public InMemoryHashAggregationBuilder(
List<AccumulatorFactory> accumulatorFactories,
AggregationNode.Step step,
int expectedGroups,
List<Type> groupByTypes,
List<Integer> groupByChannels,
Optional<Integer> hashChannel,
OperatorContext operatorContext,
DataSize maxPartialMemory,
Optional<Integer> overwriteIntermediateChannelOffset,
JoinCompiler joinCompiler)
{
this.groupByHash = createGroupByHash(operatorContext.getSession(), groupByTypes, Ints.toArray(groupByChannels), hashChannel, expectedGroups, joinCompiler);
this.operatorContext = operatorContext;
this.partial = step.isOutputPartial();
this.maxPartialMemory = maxPartialMemory.toBytes();
this.systemMemoryContext = operatorContext.getSystemMemoryContext().newLocalMemoryContext();
// wrapper each function with an aggregator
ImmutableList.Builder<Aggregator> builder = ImmutableList.builder();
requireNonNull(accumulatorFactories, "accumulatorFactories is null");
for (int i = 0; i < accumulatorFactories.size(); i++) {
AccumulatorFactory accumulatorFactory = accumulatorFactories.get(i);
Optional<Integer> overwriteIntermediateChannel = Optional.empty();
if (overwriteIntermediateChannelOffset.isPresent()) {
overwriteIntermediateChannel = Optional.of(overwriteIntermediateChannelOffset.get() + i);
}
builder.add(new Aggregator(accumulatorFactory, step, overwriteIntermediateChannel));
}
aggregators = builder.build();
}
@Override
public void close()
{
}
@Override
public void processPage(Page page)
{
if (aggregators.isEmpty()) {
groupByHash.addPage(page);
}
else {
GroupByIdBlock groupIds = groupByHash.getGroupIds(page);
for (Aggregator aggregator : aggregators) {
aggregator.processPage(groupIds, page);
}
}
}
@Override
public void updateMemory()
{
long memorySize = getSizeInMemory();
if (partial) {
systemMemoryContext.setBytes(memorySize);
full = (memorySize > maxPartialMemory);
}
else {
operatorContext.setMemoryReservation(memorySize);
}
}
@Override
public boolean isFull()
{
return full;
}
@Override
public void recordHashCollisions(HashCollisionsCounter hashCollisionsCounter)
{
hashCollisionsCounter.recordHashCollision(groupByHash.getHashCollisions(), groupByHash.getExpectedHashCollisions());
}
public long getHashCollisions()
{
return groupByHash.getHashCollisions();
}
public double getExpectedHashCollisions()
{
return groupByHash.getExpectedHashCollisions();
}
@Override
public ListenableFuture<?> isBlocked()
{
return NOT_BLOCKED;
}
public long getSizeInMemory()
{
long sizeInMemory = groupByHash.getEstimatedSize();
for (Aggregator aggregator : aggregators) {
sizeInMemory += aggregator.getEstimatedSize();
}
return sizeInMemory;
}
/**
* building hash sorted results requires memory for sorting group IDs.
* This method returns size of that memory requirement.
*/
public long getGroupIdsSortingSize()
{
return getGroupCount() * Integer.BYTES;
}
public void setOutputPartial()
{
for (Aggregator aggregator : aggregators) {
aggregator.setOutputPartial();
}
}
public int getKeyChannels()
{
return groupByHash.getTypes().size();
}
public long getGroupCount()
{
return groupByHash.getGroupCount();
}
@Override
public Iterator<Page> buildResult()
{
return buildResult(consecutiveGroupIds());
}
public Iterator<Page> buildHashSortedResult()
{
return buildResult(hashSortedGroupIds());
}
public List<Type> buildIntermediateTypes()
{
ArrayList<Type> types = new ArrayList<>(groupByHash.getTypes());
for (InMemoryHashAggregationBuilder.Aggregator aggregator : aggregators) {
types.add(aggregator.getIntermediateType());
}
return types;
}
private Iterator<Page> buildResult(IntIterator groupIds)
{
final PageBuilder pageBuilder = new PageBuilder(buildTypes());
return new AbstractIterator<Page>()
{
@Override
protected Page computeNext()
{
if (!groupIds.hasNext()) {
return endOfData();
}
pageBuilder.reset();
List<Type> types = groupByHash.getTypes();
while (!pageBuilder.isFull() && groupIds.hasNext()) {
int groupId = groupIds.nextInt();
groupByHash.appendValuesTo(groupId, pageBuilder, 0);
pageBuilder.declarePosition();
for (int i = 0; i < aggregators.size(); i++) {
Aggregator aggregator = aggregators.get(i);
BlockBuilder output = pageBuilder.getBlockBuilder(types.size() + i);
aggregator.evaluate(groupId, output);
}
}
return pageBuilder.build();
}
};
}
public List<Type> buildTypes()
{
ArrayList<Type> types = new ArrayList<>(groupByHash.getTypes());
for (Aggregator aggregator : aggregators) {
types.add(aggregator.getType());
}
return types;
}
private IntIterator consecutiveGroupIds()
{
return IntIterators.fromTo(0, groupByHash.getGroupCount());
}
private IntIterator hashSortedGroupIds()
{
IntBigArray groupIds = new IntBigArray();
groupIds.ensureCapacity(groupByHash.getGroupCount());
for (int i = 0; i < groupByHash.getGroupCount(); i++) {
groupIds.set(i, i);
}
groupIds.sort(0, groupByHash.getGroupCount(), (leftGroupId, rightGroupId) ->
Long.compare(groupByHash.getRawHash(leftGroupId), groupByHash.getRawHash(rightGroupId)));
return new AbstractIntIterator() {
private final int totalPositions = groupByHash.getGroupCount();
private int position = 0;
@Override
public boolean hasNext()
{
return position < totalPositions;
}
@Override
public int nextInt()
{
return groupIds.get(position++);
}
};
}
private static class Aggregator
{
private final GroupedAccumulator aggregation;
private AggregationNode.Step step;
private final int intermediateChannel;
private Aggregator(AccumulatorFactory accumulatorFactory, AggregationNode.Step step, Optional<Integer> overwriteIntermediateChannel)
{
if (step.isInputRaw()) {
this.intermediateChannel = -1;
this.aggregation = accumulatorFactory.createGroupedAccumulator();
}
else if (overwriteIntermediateChannel.isPresent()) {
this.intermediateChannel = overwriteIntermediateChannel.get();
this.aggregation = accumulatorFactory.createGroupedIntermediateAccumulator();
}
else {
checkArgument(accumulatorFactory.getInputChannels().size() == 1, "expected 1 input channel for intermediate aggregation");
this.intermediateChannel = accumulatorFactory.getInputChannels().get(0);
this.aggregation = accumulatorFactory.createGroupedIntermediateAccumulator();
}
this.step = step;
}
public long getEstimatedSize()
{
return aggregation.getEstimatedSize();
}
public Type getType()
{
if (step.isOutputPartial()) {
return aggregation.getIntermediateType();
}
else {
return aggregation.getFinalType();
}
}
public void processPage(GroupByIdBlock groupIds, Page page)
{
if (step.isInputRaw()) {
aggregation.addInput(groupIds, page);
}
else {
aggregation.addIntermediate(groupIds, page.getBlock(intermediateChannel));
}
}
public void evaluate(int groupId, BlockBuilder output)
{
if (step.isOutputPartial()) {
aggregation.evaluateIntermediate(groupId, output);
}
else {
aggregation.evaluateFinal(groupId, output);
}
}
public void setOutputPartial()
{
step = AggregationNode.Step.partialOutput(step);
}
public Type getIntermediateType()
{
return aggregation.getIntermediateType();
}
}
public static List<Type> toTypes(List<? extends Type> groupByType, Step step, List<AccumulatorFactory> factories, Optional<Integer> hashChannel)
{
ImmutableList.Builder<Type> types = ImmutableList.builder();
types.addAll(groupByType);
if (hashChannel.isPresent()) {
types.add(BIGINT);
}
for (AccumulatorFactory factory : factories) {
types.add(new Aggregator(factory, step, Optional.empty()).getType());
}
return types.build();
}
}