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

import com.facebook.presto.common.Page;
import com.facebook.presto.operator.PartitionFunction;
import com.facebook.presto.operator.exchange.PageReference.PageReleasedListener;
import com.google.common.collect.ImmutableList;
import com.google.common.primitives.Ints;
import com.google.common.util.concurrent.ListenableFuture;
import it.unimi.dsi.fastutil.ints.IntArrayList;

import java.util.List;
import java.util.Optional;
import java.util.function.Consumer;

import static java.util.Objects.requireNonNull;

class PartitioningExchanger
        implements LocalExchanger
{
    private final List<Consumer<PageReference>> buffers;
    private final LocalExchangeMemoryManager memoryManager;
    private final PartitionFunction partitionFunction;
    private final int[] partitioningChannels;
    private final int hashChannel; // when >= 0, this is the precomputed raw hash column to partition on
    private final IntArrayList[] partitionAssignments;
    private final PageReleasedListener onPageReleased;

    public PartitioningExchanger(
            List<Consumer<PageReference>> partitions,
            LocalExchangeMemoryManager memoryManager,
            PartitionFunction partitionFunction,
            List<Integer> partitioningChannels,
            Optional<Integer> hashChannel)
    {
        this.buffers = ImmutableList.copyOf(requireNonNull(partitions, "partitions is null"));
        this.memoryManager = requireNonNull(memoryManager, "memoryManager is null");
        this.partitionFunction = requireNonNull(partitionFunction, "partitionFunction is null");
        this.partitioningChannels = Ints.toArray(requireNonNull(partitioningChannels, "partitioningChannels is null"));
        this.hashChannel = requireNonNull(hashChannel, "hashChannel is null").orElse(-1);
        this.onPageReleased = PageReleasedListener.forLocalExchangeMemoryManager(memoryManager);

        partitionAssignments = new IntArrayList[partitions.size()];
        for (int i = 0; i < partitionAssignments.length; i++) {
            partitionAssignments[i] = new IntArrayList();
        }
    }

    @Override
    public void accept(Page page)
    {
        // extract the partitioning channel before entering the critical section
        partitionPage(page, extractPartitioningChannels(page));
    }

    private synchronized void partitionPage(Page page, Page partitioningChannelsPage)
    {
        // assign each row to a partition. The assignments lists are all expected to cleared by the previous iterations
        for (int position = 0; position < partitioningChannelsPage.getPositionCount(); position++) {
            int partition = partitionFunction.getPartition(partitioningChannelsPage, position);
            partitionAssignments[partition].add(position);
        }

        // build a page for each partition
        for (int partition = 0; partition < partitionAssignments.length; partition++) {
            IntArrayList positions = partitionAssignments[partition];
            int partitionSize = positions.size();
            if (partitionSize == 0) {
                continue;
            }
            Page pageSplit;
            if (partitionSize == page.getPositionCount()) {
                pageSplit = page; // entire page will be sent to this partition, no copies necessary
            }
            else {
                pageSplit = page.copyPositions(positions.elements(), 0, partitionSize);
            }
            // clear the assigned positions list for the next iteration to start empty
            positions.clear();
            memoryManager.updateMemoryUsage(pageSplit.getRetainedSizeInBytes());
            buffers.get(partition).accept(new PageReference(pageSplit, 1, onPageReleased));
        }
    }

    private Page extractPartitioningChannels(Page inputPage)
    {
        // hash value is pre-computed, only needs to extract that channel
        if (hashChannel >= 0) {
            return inputPage.extractChannel(hashChannel);
        }

        // extract partitioning channels
        return inputPage.extractChannels(partitioningChannels);
    }

    @Override
    public ListenableFuture<?> waitForWriting()
    {
        return memoryManager.getNotFullFuture();
    }
}

@@ -975,16 +975,59 @@
 rayList;
+%0Aimport it.unimi.dsi.fastutil.ints.IntList;
 %0A%0Aimport
@@ -1460,94 +1460,38 @@
 nal 
-int hashChannel; // when %3E= 0, this is the precomputed raw hash column to partition on
+Optional%3CInteger%3E hashChannel;
 %0A   
@@ -2364,19 +2364,8 @@
 ll%22)
-.orElse(-1)
 ;%0A  
@@ -2692,16 +2692,29 @@
  public 
+synchronized 
 void acc
@@ -2749,360 +2749,270 @@
  // 
-extract the partitioning channel before entering the critical section%0A        partitionPage(page, extractPartitioningChannels(page));
+reset the assignment lists%0A        for (IntList partitionAssignment : partitionAssignments) %7B
 %0A    
-%7D%0A%0A
     
-private synchronized void partitionPage(Page page, Page partitioningChannelsPage)%0A    %7B%0A        // assign each row to a partition. The assignments lists are all expected to cleared by the previous iterations
+    partitionAssignment.clear();%0A        %7D%0A%0A        // assign each row to a partition%0A        Page partitioningChannelsPage = extractPartitioningChannels(page);
 %0A   
@@ -3366,35 +3366,22 @@
 n %3C 
-partitionAssignments.length
+buffers.size()
 ; pa
@@ -3480,565 +3480,135 @@
    i
-nt partitionSize = positions.size();%0A            if (partitionSize == 0) %7B%0A                continue;%0A            %7D%0A            Page pageSplit;%0A            if (partitionSize == page.getPositionCount()) %7B%0A                pageSplit = page; // entire page will be sent to this partition, no copies necessary%0A            %7D%0A            else %7B%0A                pageSplit = page.copyPositions(positions.elements(), 0, partitionSize);%0A            %7D%0A            // clear the assigned positions list for the next iteration to start empty%0A            positions.clear();%0A
+f (!positions.isEmpty()) %7B%0A                Page pageSplit = page.copyPositions(positions.elements(), 0, positions.size());%0A    
     
@@ -3680,24 +3680,28 @@
 InBytes());%0A
+    
             
@@ -3776,24 +3776,38 @@
 Released));%0A
+            %7D%0A
         %7D%0A  
@@ -3979,13 +3979,20 @@
 nnel
- %3E= 0
+.isPresent()
 ) %7B%0A
@@ -4002,32 +4002,41 @@
          return 
+new Page(
 inputPage.extrac
@@ -4029,30 +4029,24 @@
 putPage.
-extractChannel
+getBlock
 (hashCha
@@ -4049,16 +4049,23 @@
 hChannel
+.get())
 );%0A     

Revert "Improve PartitioningExchanger performance"