com.facebook.presto.operator.window.AggregateWindowFunction.java Source code

Introduction

Here is the source code for com.facebook.presto.operator.window.AggregateWindowFunction.java

Source

/*
 * 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.window;

import com.facebook.presto.metadata.Signature;
import com.facebook.presto.operator.aggregation.Accumulator;
import com.facebook.presto.operator.aggregation.AccumulatorFactory;
import com.facebook.presto.operator.aggregation.InternalAggregationFunction;
import com.facebook.presto.spi.PageBuilder;
import com.facebook.presto.spi.block.BlockBuilder;
import com.google.common.collect.ImmutableList;
import com.google.common.primitives.Ints;

import java.util.List;
import java.util.Optional;

import static java.util.Objects.requireNonNull;

public class AggregateWindowFunction implements WindowFunction {
    private final InternalAggregationFunction function;
    private final int[] argumentChannels;
    private final AccumulatorFactory accumulatorFactory;

    private WindowIndex windowIndex;
    private Accumulator accumulator;
    private int currentStart;
    private int currentEnd;

    private AggregateWindowFunction(InternalAggregationFunction function, List<Integer> argumentChannels) {
        this.function = requireNonNull(function, "function is null");
        this.argumentChannels = Ints.toArray(argumentChannels);
        this.accumulatorFactory = function.bind(createArgs(function), Optional.empty(), Optional.empty(), 1.0);
    }

    @Override
    public void reset(WindowIndex windowIndex) {
        this.windowIndex = windowIndex;
        resetAccumulator();
    }

    @Override
    public void processRow(BlockBuilder output, int peerGroupStart, int peerGroupEnd, int frameStart,
            int frameEnd) {
        if (frameStart < 0) {
            // empty frame
            resetAccumulator();
        } else if ((frameStart == currentStart) && (frameEnd >= currentEnd)) {
            // same or expanding frame
            accumulate(currentEnd + 1, frameEnd);
            currentEnd = frameEnd;
        } else {
            // different frame
            resetAccumulator();
            accumulate(frameStart, frameEnd);
            currentStart = frameStart;
            currentEnd = frameEnd;
        }

        accumulator.evaluateFinal(output);
    }

    private void accumulate(int start, int end) {
        // TODO: add Accumulator method that does not require creating pages
        PageBuilder pageBuilder = new PageBuilder(function.getParameterTypes());
        for (int position = start; position <= end; position++) {
            for (int i = 0; i < function.getParameterTypes().size(); i++) {
                windowIndex.appendTo(argumentChannels[i], position, pageBuilder.getBlockBuilder(i));
            }
            pageBuilder.declarePosition();
        }
        accumulator.addInput(pageBuilder.build());
    }

    private void resetAccumulator() {
        if (currentStart >= 0) {
            accumulator = accumulatorFactory.createAccumulator();
            currentStart = -1;
            currentEnd = -1;
        }
    }

    public static WindowFunctionSupplier supplier(Signature signature, final InternalAggregationFunction function) {
        requireNonNull(function, "function is null");
        return new AbstractWindowFunctionSupplier(signature, null) {
            @Override
            protected WindowFunction newWindowFunction(List<Integer> inputs) {
                return new AggregateWindowFunction(function, inputs);
            }
        };
    }

    private static List<Integer> createArgs(InternalAggregationFunction function) {
        ImmutableList.Builder<Integer> list = ImmutableList.builder();
        for (int i = 0; i < function.getParameterTypes().size(); i++) {
            list.add(i);
        }
        return list.build();
    }
}