/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF 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 org.apache.flink.test.operators;

import org.apache.flink.api.common.InvalidProgramException;
import org.apache.flink.api.common.functions.FilterFunction;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.functions.MapPartitionFunction;
import org.apache.flink.api.common.functions.ReduceFunction;
import org.apache.flink.api.common.functions.RichMapFunction;
import org.apache.flink.api.common.operators.Order;
import org.apache.flink.api.java.DataSet;
import org.apache.flink.api.java.ExecutionEnvironment;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.api.java.operators.AggregateOperator;
import org.apache.flink.api.java.operators.DataSource;
import org.apache.flink.api.java.operators.DeltaIteration;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.test.operators.util.CollectionDataSets;
import org.apache.flink.test.operators.util.CollectionDataSets.POJO;
import org.apache.flink.test.util.MultipleProgramsTestBase;
import org.apache.flink.util.Collector;

import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.Parameterized;

import java.io.Serializable;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;

/** Integration tests for {@link MapPartitionFunction}. */
@RunWith(Parameterized.class)
@SuppressWarnings("serial")
public class PartitionITCase extends MultipleProgramsTestBase {

    public PartitionITCase(TestExecutionMode mode) {
        super(mode);
    }

    @Test
    public void testHashPartitionByKeyField() throws Exception {
        /*
         * Test hash partition by key field
         */

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSet<Tuple3<Integer, Long, String>> ds = CollectionDataSets.get3TupleDataSet(env);
        DataSet<Long> uniqLongs = ds.partitionByHash(1).mapPartition(new UniqueTupleLongMapper());
        List<Long> result = uniqLongs.collect();

        String expected = "1\n" + "2\n" + "3\n" + "4\n" + "5\n" + "6\n";

        compareResultAsText(result, expected);
    }

    @Test
    public void testRangePartitionByKeyField() throws Exception {
        /*
         * Test range partition by key field
         */

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSet<Tuple3<Integer, Long, String>> ds = CollectionDataSets.get3TupleDataSet(env);
        DataSet<Long> uniqLongs = ds.partitionByRange(1).mapPartition(new UniqueTupleLongMapper());
        List<Long> result = uniqLongs.collect();

        String expected = "1\n" + "2\n" + "3\n" + "4\n" + "5\n" + "6\n";

        compareResultAsText(result, expected);
    }

    @Test
    public void testHashPartitionByKeyField2() throws Exception {
        /*
         * Test hash partition by key field
         */

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSet<Tuple3<Integer, Long, String>> ds = CollectionDataSets.get3TupleDataSet(env);
        AggregateOperator<Tuple3<Integer, Long, String>> sum =
                ds.map(new PrefixMapper()).partitionByHash(1, 2).groupBy(1, 2).sum(0);

        List<Tuple3<Integer, Long, String>> result = sum.collect();

        String expected =
                "(1,1,Hi)\n"
                        + "(5,2,Hello)\n"
                        + "(4,3,Hello)\n"
                        + "(5,3,I am )\n"
                        + "(6,3,Luke )\n"
                        + "(34,4,Comme)\n"
                        + "(65,5,Comme)\n"
                        + "(111,6,Comme)";

        compareResultAsText(result, expected);
    }

    @Test
    public void testRangePartitionByKeyField2() throws Exception {
        /*
         * Test range partition by key field
         */

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSet<Tuple3<Integer, Long, String>> ds = CollectionDataSets.get3TupleDataSet(env);
        AggregateOperator<Tuple3<Integer, Long, String>> sum =
                ds.map(new PrefixMapper()).partitionByRange(1, 2).groupBy(1, 2).sum(0);

        List<Tuple3<Integer, Long, String>> result = sum.collect();

        String expected =
                "(1,1,Hi)\n"
                        + "(5,2,Hello)\n"
                        + "(4,3,Hello)\n"
                        + "(5,3,I am )\n"
                        + "(6,3,Luke )\n"
                        + "(34,4,Comme)\n"
                        + "(65,5,Comme)\n"
                        + "(111,6,Comme)";

        compareResultAsText(result, expected);
    }

    @Test
    public void testHashPartitionOfAtomicType() throws Exception {
        /*
         * Test hash partition of atomic type
         */

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSet<Long> uniqLongs =
                env.generateSequence(1, 6)
                        .union(env.generateSequence(1, 6))
                        .rebalance()
                        .partitionByHash("*")
                        .mapPartition(new UniqueLongMapper());
        List<Long> result = uniqLongs.collect();

        String expected = "1\n" + "2\n" + "3\n" + "4\n" + "5\n" + "6\n";

        compareResultAsText(result, expected);
    }

    @Test
    public void testRangePartitionOfAtomicType() throws Exception {
        /*
         * Test range partition of atomic type
         */

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSet<Long> uniqLongs =
                env.generateSequence(1, 6)
                        .union(env.generateSequence(1, 6))
                        .rebalance()
                        .partitionByRange("*")
                        .mapPartition(new UniqueLongMapper());
        List<Long> result = uniqLongs.collect();

        String expected = "1\n" + "2\n" + "3\n" + "4\n" + "5\n" + "6\n";

        compareResultAsText(result, expected);
    }

    @Test
    public void testHashPartitionByKeySelector() throws Exception {
        /*
         * Test hash partition by key selector
         */

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSet<Tuple3<Integer, Long, String>> ds = CollectionDataSets.get3TupleDataSet(env);
        DataSet<Long> uniqLongs =
                ds.partitionByHash(new KeySelector1()).mapPartition(new UniqueTupleLongMapper());
        List<Long> result = uniqLongs.collect();

        String expected = "1\n" + "2\n" + "3\n" + "4\n" + "5\n" + "6\n";

        compareResultAsText(result, expected);
    }

    private static class PrefixMapper
            implements MapFunction<Tuple3<Integer, Long, String>, Tuple3<Integer, Long, String>> {
        @Override
        public Tuple3<Integer, Long, String> map(Tuple3<Integer, Long, String> value)
                throws Exception {
            if (value.f2.length() > 5) {
                value.f2 = value.f2.substring(0, 5);
            }
            return value;
        }
    }

    @Test
    public void testRangePartitionByKeySelector() throws Exception {
        /*
         * Test range partition by key selector
         */

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSet<Tuple3<Integer, Long, String>> ds = CollectionDataSets.get3TupleDataSet(env);
        DataSet<Long> uniqLongs =
                ds.partitionByRange(new KeySelector1()).mapPartition(new UniqueTupleLongMapper());
        List<Long> result = uniqLongs.collect();

        String expected = "1\n" + "2\n" + "3\n" + "4\n" + "5\n" + "6\n";

        compareResultAsText(result, expected);
    }

    private static class KeySelector1 implements KeySelector<Tuple3<Integer, Long, String>, Long> {
        private static final long serialVersionUID = 1L;

        @Override
        public Long getKey(Tuple3<Integer, Long, String> value) throws Exception {
            return value.f1;
        }
    }

    @Test
    public void testForcedRebalancing() throws Exception {
        /*
         * Test forced rebalancing
         */

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        // generate some number in parallel
        DataSet<Long> ds = env.generateSequence(1, 3000);
        DataSet<Tuple2<Integer, Integer>> uniqLongs =
                ds
                        // introduce some partition skew by filtering
                        .filter(new Filter1())
                        // rebalance
                        .rebalance()
                        // count values in each partition
                        .map(new PartitionIndexMapper())
                        .groupBy(0)
                        .reduce(new Reducer1())
                        // round counts to mitigate runtime scheduling effects (lazy split
                        // assignment)
                        .map(new Mapper1());

        List<Tuple2<Integer, Integer>> result = uniqLongs.collect();

        StringBuilder expected = new StringBuilder();
        int numPerPartition = 2220 / env.getParallelism() / 10;
        for (int i = 0; i < env.getParallelism(); i++) {
            expected.append('(').append(i).append(',').append(numPerPartition).append(")\n");
        }

        compareResultAsText(result, expected.toString());
    }

    private static class Filter1 implements FilterFunction<Long> {
        private static final long serialVersionUID = 1L;

        @Override
        public boolean filter(Long value) throws Exception {
            return value > 780;
        }
    }

    private static class Reducer1 implements ReduceFunction<Tuple2<Integer, Integer>> {
        private static final long serialVersionUID = 1L;

        @Override
        public Tuple2<Integer, Integer> reduce(
                Tuple2<Integer, Integer> v1, Tuple2<Integer, Integer> v2) {
            return new Tuple2<>(v1.f0, v1.f1 + v2.f1);
        }
    }

    private static class Mapper1
            implements MapFunction<Tuple2<Integer, Integer>, Tuple2<Integer, Integer>> {
        private static final long serialVersionUID = 1L;

        @Override
        public Tuple2<Integer, Integer> map(Tuple2<Integer, Integer> value) throws Exception {
            value.f1 = (value.f1 / 10);
            return value;
        }
    }

    @Test
    public void testHashPartitionByKeyFieldAndDifferentParallelism() throws Exception {
        /*
         * Test hash partition by key field and different parallelism
         */

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(3);

        DataSet<Tuple3<Integer, Long, String>> ds = CollectionDataSets.get3TupleDataSet(env);
        DataSet<Long> uniqLongs =
                ds.partitionByHash(1).setParallelism(4).mapPartition(new UniqueTupleLongMapper());
        List<Long> result = uniqLongs.collect();

        String expected = "1\n" + "2\n" + "3\n" + "4\n" + "5\n" + "6\n";

        compareResultAsText(result, expected);
    }

    @Test
    public void testRangePartitionByKeyFieldAndDifferentParallelism() throws Exception {
        /*
         * Test range partition by key field and different parallelism
         */

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(3);

        DataSet<Tuple3<Integer, Long, String>> ds = CollectionDataSets.get3TupleDataSet(env);
        DataSet<Long> uniqLongs =
                ds.partitionByRange(1).setParallelism(4).mapPartition(new UniqueTupleLongMapper());
        List<Long> result = uniqLongs.collect();

        String expected = "1\n" + "2\n" + "3\n" + "4\n" + "5\n" + "6\n";

        compareResultAsText(result, expected);
    }

    @Test
    public void testHashPartitionWithKeyExpression() throws Exception {
        /*
         * Test hash partition with key expression
         */

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(3);

        DataSet<POJO> ds = CollectionDataSets.getDuplicatePojoDataSet(env);
        DataSet<Long> uniqLongs =
                ds.partitionByHash("nestedPojo.longNumber")
                        .setParallelism(4)
                        .mapPartition(new UniqueNestedPojoLongMapper());
        List<Long> result = uniqLongs.collect();

        String expected = "10000\n" + "20000\n" + "30000\n";

        compareResultAsText(result, expected);
    }

    @Test
    public void testRangePartitionWithKeyExpression() throws Exception {
        /*
         * Test range partition with key expression
         */

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(3);

        DataSet<POJO> ds = CollectionDataSets.getDuplicatePojoDataSet(env);
        DataSet<Long> uniqLongs =
                ds.partitionByRange("nestedPojo.longNumber")
                        .setParallelism(4)
                        .mapPartition(new UniqueNestedPojoLongMapper());
        List<Long> result = uniqLongs.collect();

        String expected = "10000\n" + "20000\n" + "30000\n";

        compareResultAsText(result, expected);
    }

    private static class UniqueTupleLongMapper
            implements MapPartitionFunction<Tuple3<Integer, Long, String>, Long> {
        private static final long serialVersionUID = 1L;

        @Override
        public void mapPartition(
                Iterable<Tuple3<Integer, Long, String>> records, Collector<Long> out)
                throws Exception {
            HashSet<Long> uniq = new HashSet<>();
            for (Tuple3<Integer, Long, String> t : records) {
                uniq.add(t.f1);
            }
            for (Long l : uniq) {
                out.collect(l);
            }
        }
    }

    private static class UniqueLongMapper implements MapPartitionFunction<Long, Long> {
        private static final long serialVersionUID = 1L;

        @Override
        public void mapPartition(Iterable<Long> longs, Collector<Long> out) throws Exception {
            HashSet<Long> uniq = new HashSet<>();
            for (Long l : longs) {
                uniq.add(l);
            }
            for (Long l : uniq) {
                out.collect(l);
            }
        }
    }

    private static class UniqueNestedPojoLongMapper implements MapPartitionFunction<POJO, Long> {
        private static final long serialVersionUID = 1L;

        @Override
        public void mapPartition(Iterable<POJO> records, Collector<Long> out) throws Exception {
            HashSet<Long> uniq = new HashSet<>();
            for (POJO t : records) {
                uniq.add(t.nestedPojo.longNumber);
            }
            for (Long l : uniq) {
                out.collect(l);
            }
        }
    }

    private static class PartitionIndexMapper
            extends RichMapFunction<Long, Tuple2<Integer, Integer>> {
        private static final long serialVersionUID = 1L;

        @Override
        public Tuple2<Integer, Integer> map(Long value) throws Exception {
            return new Tuple2<>(this.getRuntimeContext().getIndexOfThisSubtask(), 1);
        }
    }

    @Test
    public void testRangePartitionerOnSequenceData() throws Exception {
        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
        DataSource<Long> dataSource = env.generateSequence(0, 10000);
        KeySelector<Long, Long> keyExtractor = new ObjectSelfKeySelector();

        MapPartitionFunction<Long, Tuple2<Long, Long>> minMaxSelector =
                new MinMaxSelector<>(new LongComparator(true));

        Comparator<Tuple2<Long, Long>> tuple2Comparator =
                new Tuple2Comparator(new LongComparator(true));

        List<Tuple2<Long, Long>> collected =
                dataSource.partitionByRange(keyExtractor).mapPartition(minMaxSelector).collect();
        Collections.sort(collected, tuple2Comparator);

        long previousMax = -1;
        for (Tuple2<Long, Long> tuple2 : collected) {
            if (previousMax == -1) {
                previousMax = tuple2.f1;
            } else {
                long currentMin = tuple2.f0;
                assertTrue(tuple2.f0 < tuple2.f1);
                assertEquals(previousMax + 1, currentMin);
                previousMax = tuple2.f1;
            }
        }
    }

    @Test(expected = InvalidProgramException.class)
    public void testRangePartitionInIteration() throws Exception {

        // does not apply for collection execution
        if (super.mode == TestExecutionMode.COLLECTION) {
            throw new InvalidProgramException("Does not apply for collection execution");
        }

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
        DataSource<Long> source = env.generateSequence(0, 10000);

        DataSet<Tuple2<Long, String>> tuples =
                source.map(
                        new MapFunction<Long, Tuple2<Long, String>>() {
                            @Override
                            public Tuple2<Long, String> map(Long v) throws Exception {
                                return new Tuple2<>(v, Long.toString(v));
                            }
                        });

        DeltaIteration<Tuple2<Long, String>, Tuple2<Long, String>> it =
                tuples.iterateDelta(tuples, 10, 0);
        DataSet<Tuple2<Long, String>> body =
                it.getWorkset()
                        .partitionByRange(
                                1) // Verify that range partition is not allowed in iteration
                        .join(it.getSolutionSet())
                        .where(0)
                        .equalTo(0)
                        .projectFirst(0)
                        .projectSecond(1);
        DataSet<Tuple2<Long, String>> result = it.closeWith(body, body);

        result.collect(); // should fail
    }

    @Test
    public void testRangePartitionerOnSequenceDataWithOrders() throws Exception {
        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
        DataSet<Tuple2<Long, Long>> dataSet =
                env.generateSequence(0, 10000)
                        .map(
                                new MapFunction<Long, Tuple2<Long, Long>>() {
                                    @Override
                                    public Tuple2<Long, Long> map(Long value) throws Exception {
                                        return new Tuple2<>(value / 5000, value % 5000);
                                    }
                                });

        final Tuple2Comparator<Long> tuple2Comparator =
                new Tuple2Comparator<>(new LongComparator(true), new LongComparator(false));

        MinMaxSelector<Tuple2<Long, Long>> minMaxSelector = new MinMaxSelector<>(tuple2Comparator);

        final List<Tuple2<Tuple2<Long, Long>, Tuple2<Long, Long>>> collected =
                dataSet.partitionByRange(0, 1)
                        .withOrders(Order.ASCENDING, Order.DESCENDING)
                        .mapPartition(minMaxSelector)
                        .collect();

        Collections.sort(collected, new Tuple2Comparator<>(tuple2Comparator));

        Tuple2<Long, Long> previousMax = null;
        for (Tuple2<Tuple2<Long, Long>, Tuple2<Long, Long>> tuple2 : collected) {
            assertTrue(
                    "Min element in each partition should be smaller than max.",
                    tuple2Comparator.compare(tuple2.f0, tuple2.f1) <= 0);
            if (previousMax == null) {
                previousMax = tuple2.f1;
            } else {
                assertTrue(
                        "Partitions overlap. Previous max should be smaller than current min.",
                        tuple2Comparator.compare(previousMax, tuple2.f0) < 0);
                if (previousMax.f0.equals(tuple2.f0.f0)) {
                    // check that ordering on the second key is correct
                    assertEquals(
                            "Ordering on the second field should be continous.",
                            previousMax.f1 - 1,
                            tuple2.f0.f1.longValue());
                }
                previousMax = tuple2.f1;
            }
        }
    }

    @Test
    public void testRangePartitionerOnSequenceNestedDataWithOrders() throws Exception {
        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
        final DataSet<Tuple2<Tuple2<Long, Long>, Long>> dataSet =
                env.generateSequence(0, 10000)
                        .map(
                                new MapFunction<Long, Tuple2<Tuple2<Long, Long>, Long>>() {
                                    @Override
                                    public Tuple2<Tuple2<Long, Long>, Long> map(Long value)
                                            throws Exception {
                                        return new Tuple2<>(
                                                new Tuple2<>(value / 5000, value % 5000), value);
                                    }
                                });

        final Tuple2Comparator<Long> tuple2Comparator =
                new Tuple2Comparator<>(new LongComparator(true), new LongComparator(true));
        MinMaxSelector<Tuple2<Long, Long>> minMaxSelector = new MinMaxSelector<>(tuple2Comparator);

        final List<Tuple2<Tuple2<Long, Long>, Tuple2<Long, Long>>> collected =
                dataSet.partitionByRange(0)
                        .withOrders(Order.ASCENDING)
                        .mapPartition(
                                new MapPartitionFunction<
                                        Tuple2<Tuple2<Long, Long>, Long>, Tuple2<Long, Long>>() {
                                    @Override
                                    public void mapPartition(
                                            Iterable<Tuple2<Tuple2<Long, Long>, Long>> values,
                                            Collector<Tuple2<Long, Long>> out)
                                            throws Exception {
                                        for (Tuple2<Tuple2<Long, Long>, Long> value : values) {
                                            out.collect(value.f0);
                                        }
                                    }
                                })
                        .mapPartition(minMaxSelector)
                        .collect();

        Collections.sort(collected, new Tuple2Comparator<>(tuple2Comparator));

        Tuple2<Long, Long> previousMax = null;
        for (Tuple2<Tuple2<Long, Long>, Tuple2<Long, Long>> tuple2 : collected) {
            assertTrue(
                    "Min element in each partition should be smaller than max.",
                    tuple2Comparator.compare(tuple2.f0, tuple2.f1) <= 0);
            if (previousMax == null) {
                previousMax = tuple2.f1;
            } else {
                assertTrue(
                        "Partitions overlap. Previous max should be smaller than current min.",
                        tuple2Comparator.compare(previousMax, tuple2.f0) < 0);
                if (previousMax.f0.equals(tuple2.f0.f0)) {
                    assertEquals(
                            "Ordering on the second field should be continous.",
                            previousMax.f1 + 1,
                            tuple2.f0.f1.longValue());
                }
                previousMax = tuple2.f1;
            }
        }
    }

    @Test
    public void testRangePartitionerWithKeySelectorOnSequenceNestedDataWithOrders()
            throws Exception {
        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
        final DataSet<Tuple2<ComparablePojo, Long>> dataSet =
                env.generateSequence(0, 10000)
                        .map(
                                new MapFunction<Long, Tuple2<ComparablePojo, Long>>() {
                                    @Override
                                    public Tuple2<ComparablePojo, Long> map(Long value)
                                            throws Exception {
                                        return new Tuple2<>(
                                                new ComparablePojo(value / 5000, value % 5000),
                                                value);
                                    }
                                });

        final List<Tuple2<ComparablePojo, ComparablePojo>> collected =
                dataSet.partitionByRange(
                                new KeySelector<Tuple2<ComparablePojo, Long>, ComparablePojo>() {
                                    @Override
                                    public ComparablePojo getKey(Tuple2<ComparablePojo, Long> value)
                                            throws Exception {
                                        return value.f0;
                                    }
                                })
                        .withOrders(Order.ASCENDING)
                        .mapPartition(new MinMaxSelector<>(new ComparablePojoComparator()))
                        .mapPartition(new ExtractComparablePojo())
                        .collect();

        final Comparator<Tuple2<ComparablePojo, ComparablePojo>> pojoComparator =
                new Comparator<Tuple2<ComparablePojo, ComparablePojo>>() {
                    @Override
                    public int compare(
                            Tuple2<ComparablePojo, ComparablePojo> o1,
                            Tuple2<ComparablePojo, ComparablePojo> o2) {
                        return o1.f0.compareTo(o2.f1);
                    }
                };
        Collections.sort(collected, pojoComparator);

        ComparablePojo previousMax = null;
        for (Tuple2<ComparablePojo, ComparablePojo> element : collected) {
            assertTrue(
                    "Min element in each partition should be smaller than max.",
                    element.f0.compareTo(element.f1) <= 0);
            if (previousMax == null) {
                previousMax = element.f1;
            } else {
                assertTrue(
                        "Partitions overlap. Previous max should be smaller than current min.",
                        previousMax.compareTo(element.f0) < 0);
                if (previousMax.first.equals(element.f0.first)) {
                    assertEquals(
                            "Ordering on the second field should be continous.",
                            previousMax.second - 1,
                            element.f0.second.longValue());
                }
                previousMax = element.f1;
            }
        }
    }

    private static class ExtractComparablePojo
            implements MapPartitionFunction<
                    Tuple2<Tuple2<ComparablePojo, Long>, Tuple2<ComparablePojo, Long>>,
                    Tuple2<ComparablePojo, ComparablePojo>> {

        @Override
        public void mapPartition(
                Iterable<Tuple2<Tuple2<ComparablePojo, Long>, Tuple2<ComparablePojo, Long>>> values,
                Collector<Tuple2<ComparablePojo, ComparablePojo>> out)
                throws Exception {
            for (Tuple2<Tuple2<ComparablePojo, Long>, Tuple2<ComparablePojo, Long>> value :
                    values) {
                out.collect(new Tuple2<>(value.f0.f0, value.f1.f0));
            }
        }
    }

    private static class ComparablePojoComparator
            implements Comparator<Tuple2<ComparablePojo, Long>>, Serializable {

        @Override
        public int compare(Tuple2<ComparablePojo, Long> o1, Tuple2<ComparablePojo, Long> o2) {
            return o1.f0.compareTo(o2.f0);
        }
    }

    private static class ComparablePojo implements Comparable<ComparablePojo> {
        private Long first;
        private Long second;

        public Long getFirst() {
            return first;
        }

        public void setFirst(Long first) {
            this.first = first;
        }

        public Long getSecond() {
            return second;
        }

        public void setSecond(Long second) {
            this.second = second;
        }

        public ComparablePojo(Long first, Long second) {
            this.first = first;
            this.second = second;
        }

        public ComparablePojo() {}

        @Override
        public int compareTo(ComparablePojo o) {
            final int firstResult = Long.compare(this.first, o.first);
            if (firstResult == 0) {
                return (-1) * Long.compare(this.second, o.second);
            }

            return firstResult;
        }
    }

    private static class ObjectSelfKeySelector implements KeySelector<Long, Long> {
        @Override
        public Long getKey(Long value) throws Exception {
            return value;
        }
    }

    private static class MinMaxSelector<T> implements MapPartitionFunction<T, Tuple2<T, T>> {

        private final Comparator<T> comparator;

        public MinMaxSelector(Comparator<T> comparator) {
            this.comparator = comparator;
        }

        @Override
        public void mapPartition(Iterable<T> values, Collector<Tuple2<T, T>> out) throws Exception {
            Iterator<T> itr = values.iterator();
            T min = itr.next();
            T max = min;
            T value;
            while (itr.hasNext()) {
                value = itr.next();
                if (comparator.compare(value, min) < 0) {
                    min = value;
                }
                if (comparator.compare(value, max) > 0) {
                    max = value;
                }
            }

            Tuple2<T, T> result = new Tuple2<>(min, max);
            out.collect(result);
        }
    }

    private static class Tuple2Comparator<T> implements Comparator<Tuple2<T, T>>, Serializable {

        private final Comparator<T> firstComparator;
        private final Comparator<T> secondComparator;

        public Tuple2Comparator(Comparator<T> comparator) {
            this(comparator, comparator);
        }

        public Tuple2Comparator(Comparator<T> firstComparator, Comparator<T> secondComparator) {
            this.firstComparator = firstComparator;
            this.secondComparator = secondComparator;
        }

        @Override
        public int compare(Tuple2<T, T> first, Tuple2<T, T> second) {
            long result = firstComparator.compare(first.f0, second.f0);
            if (result > 0) {
                return 1;
            } else if (result < 0) {
                return -1;
            }

            result = secondComparator.compare(first.f1, second.f1);
            if (result > 0) {
                return 1;
            } else if (result < 0) {
                return -1;
            }

            return 0;
        }
    }

    private static class LongComparator implements Comparator<Long>, Serializable {

        private final boolean ascending;

        public LongComparator(boolean ascending) {
            this.ascending = ascending;
        }

        @Override
        public int compare(Long o1, Long o2) {
            if (ascending) {
                return Long.compare(o1, o2);
            } else {
                return (-1) * Long.compare(o1, o2);
            }
        }
    }
}