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

import com.facebook.presto.spi.block.BlockCursor;
import com.facebook.presto.spi.type.Type;
import com.google.common.collect.ImmutableList;
import com.google.common.primitives.Ints;
import com.google.common.util.concurrent.ListenableFuture;

import java.util.Arrays;
import java.util.List;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Preconditions.checkState;

public class DistinctLimitOperator
        implements Operator
{
    public static class DistinctLimitOperatorFactory
            implements OperatorFactory
    {
        private final int operatorId;
        private final List<Type> types;
        private final long limit;
        private boolean closed;

        public DistinctLimitOperatorFactory(int operatorId, List<? extends Type> types, long limit)
        {
            this.operatorId = operatorId;
            this.types = ImmutableList.copyOf(checkNotNull(types, "types is null"));
            checkArgument(limit >= 0, "limit must be at least zero");
            this.limit = limit;
        }

        @Override
        public List<Type> getTypes()
        {
            return types;
        }

        @Override
        public Operator createOperator(DriverContext driverContext)
        {
            checkState(!closed, "Factory is already closed");
            OperatorContext operatorContext = driverContext.addOperatorContext(operatorId, DistinctLimitOperator.class.getSimpleName());
            return new DistinctLimitOperator(operatorContext, types, limit);
        }

        @Override
        public void close()
        {
            closed = true;
        }
    }

    private final OperatorContext operatorContext;
    private final List<Type> types;
    private final BlockCursor[] cursors;

    private final PageBuilder pageBuilder;
    private Page outputPage;
    private long remainingLimit;

    private boolean finishing;

    private final GroupByHash groupByHash;
    private long nextDistinctId;

    public DistinctLimitOperator(OperatorContext operatorContext, List<Type> types, long limit)
    {
        this.operatorContext = checkNotNull(operatorContext, "operatorContext is null");
        this.types = ImmutableList.copyOf(checkNotNull(types, "types is null"));
        checkArgument(limit >= 0, "limit must be at least zero");

        ImmutableList.Builder<Type> distinctTypes = ImmutableList.builder();
        ImmutableList.Builder<Integer> distinctChannels = ImmutableList.builder();
        for (int i = 0; i < types.size(); i++) {
            distinctTypes.add(types.get(i));
            distinctChannels.add(i);
        }

        this.groupByHash = new GroupByHash(distinctTypes.build(), Ints.toArray(distinctChannels.build()), Math.min((int) limit, 10_000));

        this.cursors = new BlockCursor[types.size()];
        this.pageBuilder = new PageBuilder(getTypes());
        remainingLimit = limit;
    }

    @Override
    public OperatorContext getOperatorContext()
    {
        return operatorContext;
    }

    @Override
    public List<Type> getTypes()
    {
        return types;
    }

    @Override
    public void finish()
    {
        finishing = true;
        Arrays.fill(cursors, null);
        pageBuilder.reset();
    }

    @Override
    public boolean isFinished()
    {
        return (finishing && outputPage == null) || (remainingLimit == 0 && outputPage == null);
    }

    @Override
    public ListenableFuture<?> isBlocked()
    {
        return NOT_BLOCKED;
    }

    @Override
    public boolean needsInput()
    {
        operatorContext.setMemoryReservation(groupByHash.getEstimatedSize());
        return !finishing && remainingLimit > 0 && outputPage == null;
    }

    @Override
    public void addInput(Page page)
    {
        checkState(needsInput());
        operatorContext.setMemoryReservation(groupByHash.getEstimatedSize());

        // open cursors
        for (int i = 0; i < page.getChannelCount(); i++) {
            cursors[i] = page.getBlock(i).cursor();
        }
        pageBuilder.reset();

        GroupByIdBlock ids = groupByHash.getGroupIds(page);
        for (int i = 0; i < ids.getPositionCount(); i++) {
            checkState(advanceNextCursorPosition());
            if (ids.getGroupId(i) == nextDistinctId) {
                for (int j = 0; j < cursors.length; j++) {
                    cursors[j].appendTo(pageBuilder.getBlockBuilder(j));
                }
                remainingLimit--;
                nextDistinctId++;
                if (remainingLimit == 0) {
                    break;
                }
            }
        }
        if (!pageBuilder.isEmpty()) {
            outputPage = pageBuilder.build();
        }
    }

    private boolean advanceNextCursorPosition()
    {
        // advance all cursors
        boolean advanced = cursors[0].advanceNextPosition();
        for (int i = 1; i < cursors.length; i++) {
            checkState(advanced == cursors[i].advanceNextPosition());
        }

        if (!advanced) {
            Arrays.fill(cursors, null);
        }

        return advanced;
    }

    @Override
    public Page getOutput()
    {
        Page result = outputPage;
        outputPage = null;
        return result;
    }
}

@@ -598,58 +598,8 @@
 r;%0A%0A
-import com.facebook.presto.spi.block.BlockCursor;%0A
 impo
@@ -790,33 +790,8 @@
 e;%0A%0A
-import java.util.Arrays;%0A
 impo
@@ -2338,49 +2338,8 @@
 pes;
-%0A    private final BlockCursor%5B%5D cursors;
 %0A%0A  
@@ -3335,62 +3335,8 @@
 );%0A%0A
-        this.cursors = new BlockCursor%5Btypes.size()%5D;%0A
     
@@ -3374,26 +3374,21 @@
 Builder(
-getT
+t
 ypes
-()
 );%0A     
@@ -3685,44 +3685,8 @@
 ue;%0A
-        Arrays.fill(cursors, null);%0A
     
@@ -4352,153 +4352,8 @@
 );%0A%0A
-        // open cursors%0A        for (int i = 0; i %3C page.getChannelCount(); i++) %7B%0A            cursors%5Bi%5D = page.getBlock(i).cursor();%0A        %7D%0A
     
@@ -4455,24 +4455,38 @@
 or (int 
-i = 0; i
+position = 0; position
  %3C ids.g
@@ -4508,67 +4508,21 @@
 (); 
-i++) %7B%0A            checkState(advanceNextCursorP
+p
 osition
-());
+++) %7B
 %0A   
@@ -4549,17 +4549,24 @@
 GroupId(
-i
+position
 ) == nex
@@ -4600,141 +4600,45 @@
     
-for (int j = 0; j %3C cursors.length; j++) %7B%0A                    cursors%5Bj%5D.appendTo(pageBuilder.getBlockBuilder(j));%0A                %7D
+page.appendTo(position, pageBuilder);
 %0A   
@@ -4919,394 +4919,8 @@
  %7D%0A%0A
-    private boolean advanceNextCursorPosition()%0A    %7B%0A        // advance all cursors%0A        boolean advanced = cursors%5B0%5D.advanceNextPosition();%0A        for (int i = 1; i %3C cursors.length; i++) %7B%0A            checkState(advanced == cursors%5Bi%5D.advanceNextPosition());%0A        %7D%0A%0A        if (!advanced) %7B%0A            Arrays.fill(cursors, null);%0A        %7D%0A%0A        return advanced;%0A    %7D%0A%0A
     

Change DistinctLimitOperator to not use BlockCursor