/*
 * 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.distributions.DataDistribution;
import org.apache.flink.api.common.functions.CoGroupFunction;
import org.apache.flink.api.common.functions.RichCoGroupFunction;
import org.apache.flink.api.common.typeinfo.BasicTypeInfo;
import org.apache.flink.api.common.typeinfo.TypeInformation;
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.tuple.Tuple1;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.api.java.tuple.Tuple5;
import org.apache.flink.api.java.tuple.Tuple7;
import org.apache.flink.api.java.utils.DataSetUtils;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.core.memory.DataInputView;
import org.apache.flink.core.memory.DataOutputView;
import org.apache.flink.test.operators.util.CollectionDataSets;
import org.apache.flink.test.operators.util.CollectionDataSets.CustomType;
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.Assert;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.Parameterized;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;

/** Integration tests for {@link CoGroupFunction} and {@link RichCoGroupFunction}. */
@RunWith(Parameterized.class)
public class CoGroupITCase extends MultipleProgramsTestBase {

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

    /*
     * CoGroup on tuples with key field selector
     */
    @Test
    public void testCoGroupTuplesWithKeyFieldSelector() throws Exception {

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSet<Tuple5<Integer, Long, Integer, String, Long>> ds =
                CollectionDataSets.get5TupleDataSet(env);
        DataSet<Tuple5<Integer, Long, Integer, String, Long>> ds2 =
                CollectionDataSets.get5TupleDataSet(env);
        DataSet<Tuple2<Integer, Integer>> coGroupDs =
                ds.coGroup(ds2).where(0).equalTo(0).with(new Tuple5CoGroup());

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

        String expected = "1,0\n" + "2,6\n" + "3,24\n" + "4,60\n" + "5,120\n";

        compareResultAsTuples(result, expected);
    }

    @Test
    public void testCoGroupOnTwoCustomTypeInputsWithKeyExtractors() throws Exception {
        /*
         * CoGroup on two custom type inputs with key extractors
         */

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSet<CustomType> ds = CollectionDataSets.getCustomTypeDataSet(env);
        DataSet<CustomType> ds2 = CollectionDataSets.getCustomTypeDataSet(env);
        DataSet<CustomType> coGroupDs =
                ds.coGroup(ds2)
                        .where(new KeySelector4())
                        .equalTo(new KeySelector5())
                        .with(new CustomTypeCoGroup());

        List<CustomType> result = coGroupDs.collect();

        String expected =
                "1,0,test\n"
                        + "2,6,test\n"
                        + "3,24,test\n"
                        + "4,60,test\n"
                        + "5,120,test\n"
                        + "6,210,test\n";

        compareResultAsText(result, expected);
    }

    private static class KeySelector4 implements KeySelector<CustomType, Integer> {
        private static final long serialVersionUID = 1L;

        @Override
        public Integer getKey(CustomType in) {
            return in.myInt;
        }
    }

    private static class KeySelector5 implements KeySelector<CustomType, Integer> {
        private static final long serialVersionUID = 1L;

        @Override
        public Integer getKey(CustomType in) {
            return in.myInt;
        }
    }

    @Test
    public void testCorrectnessOfCoGroupIfUDFReturnsLeftInputObjects() throws Exception {
        /*
         * check correctness of cogroup if UDF returns left input objects
         */

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSet<Tuple3<Integer, Long, String>> ds = CollectionDataSets.get3TupleDataSet(env);
        DataSet<Tuple3<Integer, Long, String>> ds2 = CollectionDataSets.get3TupleDataSet(env);
        DataSet<Tuple3<Integer, Long, String>> coGroupDs =
                ds.coGroup(ds2).where(0).equalTo(0).with(new Tuple3ReturnLeft());

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

        String expected =
                "1,1,Hi\n"
                        + "2,2,Hello\n"
                        + "3,2,Hello world\n"
                        + "4,3,Hello world, how are you?\n"
                        + "5,3,I am fine.\n";

        compareResultAsTuples(result, expected);
    }

    @Test
    public void testCorrectnessOfCoGroupIfUDFReturnsRightInputObjects() throws Exception {
        /*
         * check correctness of cogroup if UDF returns right input objects
         */

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSet<Tuple5<Integer, Long, Integer, String, Long>> ds =
                CollectionDataSets.get5TupleDataSet(env);
        DataSet<Tuple5<Integer, Long, Integer, String, Long>> ds2 =
                CollectionDataSets.get5TupleDataSet(env);
        DataSet<Tuple5<Integer, Long, Integer, String, Long>> coGroupDs =
                ds.coGroup(ds2).where(0).equalTo(0).with(new Tuple5ReturnRight());

        List<Tuple5<Integer, Long, Integer, String, Long>> result = coGroupDs.collect();

        String expected =
                "1,1,0,Hallo,1\n"
                        + "2,2,1,Hallo Welt,2\n"
                        + "2,3,2,Hallo Welt wie,1\n"
                        + "3,4,3,Hallo Welt wie gehts?,2\n"
                        + "3,5,4,ABC,2\n"
                        + "3,6,5,BCD,3\n";

        compareResultAsTuples(result, expected);
    }

    @Test
    public void testCoGroupWithBroadcastSet() throws Exception {
        /*
         * Reduce with broadcast set
         */

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSet<Integer> intDs = CollectionDataSets.getIntegerDataSet(env);

        DataSet<Tuple5<Integer, Long, Integer, String, Long>> ds =
                CollectionDataSets.get5TupleDataSet(env);
        DataSet<Tuple5<Integer, Long, Integer, String, Long>> ds2 =
                CollectionDataSets.get5TupleDataSet(env);
        DataSet<Tuple3<Integer, Integer, Integer>> coGroupDs =
                ds.coGroup(ds2)
                        .where(0)
                        .equalTo(0)
                        .with(new Tuple5CoGroupBC())
                        .withBroadcastSet(intDs, "ints");

        List<Tuple3<Integer, Integer, Integer>> result = coGroupDs.collect();

        String expected = "1,0,55\n" + "2,6,55\n" + "3,24,55\n" + "4,60,55\n" + "5,120,55\n";

        compareResultAsTuples(result, expected);
    }

    @Test
    public void testCoGroupOnATupleInputWithKeyFieldSelectorAndACustomTypeInputWithKeyExtractor()
            throws Exception {
        /*
         * CoGroup on a tuple input with key field selector and a custom type input with key extractor
         */

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSet<Tuple5<Integer, Long, Integer, String, Long>> ds =
                CollectionDataSets.get5TupleDataSet(env);
        DataSet<CustomType> ds2 = CollectionDataSets.getCustomTypeDataSet(env);
        DataSet<Tuple3<Integer, Long, String>> coGroupDs =
                ds.coGroup(ds2).where(2).equalTo(new KeySelector2()).with(new MixedCoGroup());

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

        String expected =
                "0,1,test\n"
                        + "1,2,test\n"
                        + "2,5,test\n"
                        + "3,15,test\n"
                        + "4,33,test\n"
                        + "5,63,test\n"
                        + "6,109,test\n"
                        + "7,4,test\n"
                        + "8,4,test\n"
                        + "9,4,test\n"
                        + "10,5,test\n"
                        + "11,5,test\n"
                        + "12,5,test\n"
                        + "13,5,test\n"
                        + "14,5,test\n";

        compareResultAsTuples(result, expected);
    }

    private static class KeySelector2 implements KeySelector<CustomType, Integer> {
        private static final long serialVersionUID = 1L;

        @Override
        public Integer getKey(CustomType in) {
            return in.myInt;
        }
    }

    @Test
    public void testCoGroupOnACustomTypeWithKeyExtractorAndATupleInputWithKeyFieldSelector()
            throws Exception {
        /*
         * CoGroup on a tuple input with key field selector and a custom type input with key extractor
         */

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSet<Tuple5<Integer, Long, Integer, String, Long>> ds =
                CollectionDataSets.get5TupleDataSet(env);
        DataSet<CustomType> ds2 = CollectionDataSets.getCustomTypeDataSet(env);
        DataSet<CustomType> coGroupDs =
                ds2.coGroup(ds).where(new KeySelector3()).equalTo(2).with(new MixedCoGroup2());

        List<CustomType> result = coGroupDs.collect();

        String expected =
                "0,1,test\n"
                        + "1,2,test\n"
                        + "2,5,test\n"
                        + "3,15,test\n"
                        + "4,33,test\n"
                        + "5,63,test\n"
                        + "6,109,test\n"
                        + "7,4,test\n"
                        + "8,4,test\n"
                        + "9,4,test\n"
                        + "10,5,test\n"
                        + "11,5,test\n"
                        + "12,5,test\n"
                        + "13,5,test\n"
                        + "14,5,test\n";

        compareResultAsText(result, expected);
    }

    private static class KeySelector3 implements KeySelector<CustomType, Integer> {
        private static final long serialVersionUID = 1L;

        @Override
        public Integer getKey(CustomType in) {
            return in.myInt;
        }
    }

    @Test
    public void testCoGroupWithMultipleKeyFieldsWithFieldSelector() throws Exception {
        /*
         * CoGroup with multiple key fields
         */

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSet<Tuple5<Integer, Long, Integer, String, Long>> ds1 =
                CollectionDataSets.get5TupleDataSet(env);
        DataSet<Tuple3<Integer, Long, String>> ds2 = CollectionDataSets.get3TupleDataSet(env);

        DataSet<Tuple3<Integer, Long, String>> coGrouped =
                ds1.coGroup(ds2).where(0, 4).equalTo(0, 1).with(new Tuple5Tuple3CoGroup());

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

        String expected =
                "1,1,Hallo\n"
                        + "2,2,Hallo Welt\n"
                        + "3,2,Hallo Welt wie gehts?\n"
                        + "3,2,ABC\n"
                        + "5,3,HIJ\n"
                        + "5,3,IJK\n";

        compareResultAsTuples(result, expected);
    }

    @Test
    public void testCoGroupWithMultipleKeyFieldsWithStaticClassKeyExtractor() throws Exception {
        /*
         * CoGroup with multiple key fields
         */

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSet<Tuple5<Integer, Long, Integer, String, Long>> ds1 =
                CollectionDataSets.get5TupleDataSet(env);
        DataSet<Tuple3<Integer, Long, String>> ds2 = CollectionDataSets.get3TupleDataSet(env);

        DataSet<Tuple3<Integer, Long, String>> coGrouped =
                ds1.coGroup(ds2)
                        .where(new KeySelector7())
                        .equalTo(new KeySelector8())
                        .with(new Tuple5Tuple3CoGroup());

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

        String expected =
                "1,1,Hallo\n"
                        + "2,2,Hallo Welt\n"
                        + "3,2,Hallo Welt wie gehts?\n"
                        + "3,2,ABC\n"
                        + "5,3,HIJ\n"
                        + "5,3,IJK\n";

        compareResultAsTuples(result, expected);
    }

    @Test
    public void testCoGroupWithMultipleKeyFieldsWithInnerClassKeyExtractorWithClosureCleaner()
            throws Exception {
        /*
         * CoGroup with multiple key fields, test working closure cleaner for inner classes
         */

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSet<Tuple5<Integer, Long, Integer, String, Long>> ds1 =
                CollectionDataSets.get5TupleDataSet(env);
        DataSet<Tuple3<Integer, Long, String>> ds2 = CollectionDataSets.get3TupleDataSet(env);

        DataSet<Tuple3<Integer, Long, String>> coGrouped =
                ds1.coGroup(ds2)
                        .where(
                                new KeySelector<
                                        Tuple5<Integer, Long, Integer, String, Long>,
                                        Tuple2<Integer, Long>>() {
                                    @Override
                                    public Tuple2<Integer, Long> getKey(
                                            Tuple5<Integer, Long, Integer, String, Long> t)
                                            throws Exception {
                                        return new Tuple2<Integer, Long>(t.f0, t.f4);
                                    }
                                })
                        .equalTo(
                                new KeySelector<
                                        Tuple3<Integer, Long, String>, Tuple2<Integer, Long>>() {

                                    @Override
                                    public Tuple2<Integer, Long> getKey(
                                            Tuple3<Integer, Long, String> t) {
                                        return new Tuple2<>(t.f0, t.f1);
                                    }
                                })
                        .with(
                                new CoGroupFunction<
                                        Tuple5<Integer, Long, Integer, String, Long>,
                                        Tuple3<Integer, Long, String>,
                                        Tuple3<Integer, Long, String>>() {
                                    @Override
                                    public void coGroup(
                                            Iterable<Tuple5<Integer, Long, Integer, String, Long>>
                                                    first,
                                            Iterable<Tuple3<Integer, Long, String>> second,
                                            Collector<Tuple3<Integer, Long, String>> out) {
                                        List<String> strs = new ArrayList<>();

                                        for (Tuple5<Integer, Long, Integer, String, Long> t :
                                                first) {
                                            strs.add(t.f3);
                                        }

                                        for (Tuple3<Integer, Long, String> t : second) {
                                            for (String s : strs) {
                                                out.collect(
                                                        new Tuple3<Integer, Long, String>(
                                                                t.f0, t.f1, s));
                                            }
                                        }
                                    }
                                });

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

        String expected =
                "1,1,Hallo\n"
                        + "2,2,Hallo Welt\n"
                        + "3,2,Hallo Welt wie gehts?\n"
                        + "3,2,ABC\n"
                        + "5,3,HIJ\n"
                        + "5,3,IJK\n";

        compareResultAsTuples(result, expected);
    }

    @Test
    public void testCoGroupWithMultipleKeyFieldsWithInnerClassKeyExtractorWithoutClosureCleaner()
            throws Exception {
        /*
         * CoGroup with multiple key fields, test that disabling closure cleaner leads to an exception when using inner
         * classes.
         */

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
        env.getConfig().disableClosureCleaner();

        DataSet<Tuple5<Integer, Long, Integer, String, Long>> ds1 =
                CollectionDataSets.get5TupleDataSet(env);
        DataSet<Tuple3<Integer, Long, String>> ds2 = CollectionDataSets.get3TupleDataSet(env);
        boolean correctExceptionTriggered = false;
        try {
            DataSet<Tuple3<Integer, Long, String>> coGrouped =
                    ds1.coGroup(ds2)
                            .where(
                                    new KeySelector<
                                            Tuple5<Integer, Long, Integer, String, Long>,
                                            Tuple2<Integer, Long>>() {
                                        @Override
                                        public Tuple2<Integer, Long> getKey(
                                                Tuple5<Integer, Long, Integer, String, Long> t)
                                                throws Exception {
                                            return new Tuple2<Integer, Long>(t.f0, t.f4);
                                        }
                                    })
                            .equalTo(
                                    new KeySelector<
                                            Tuple3<Integer, Long, String>,
                                            Tuple2<Integer, Long>>() {

                                        @Override
                                        public Tuple2<Integer, Long> getKey(
                                                Tuple3<Integer, Long, String> t) {
                                            return new Tuple2<Integer, Long>(t.f0, t.f1);
                                        }
                                    })
                            .with(
                                    new CoGroupFunction<
                                            Tuple5<Integer, Long, Integer, String, Long>,
                                            Tuple3<Integer, Long, String>,
                                            Tuple3<Integer, Long, String>>() {
                                        @Override
                                        public void coGroup(
                                                Iterable<
                                                                Tuple5<
                                                                        Integer,
                                                                        Long,
                                                                        Integer,
                                                                        String,
                                                                        Long>>
                                                        first,
                                                Iterable<Tuple3<Integer, Long, String>> second,
                                                Collector<Tuple3<Integer, Long, String>> out) {
                                            List<String> strs = new ArrayList<String>();

                                            for (Tuple5<Integer, Long, Integer, String, Long> t :
                                                    first) {
                                                strs.add(t.f3);
                                            }

                                            for (Tuple3<Integer, Long, String> t : second) {
                                                for (String s : strs) {
                                                    out.collect(
                                                            new Tuple3<Integer, Long, String>(
                                                                    t.f0, t.f1, s));
                                                }
                                            }
                                        }
                                    });
        } catch (InvalidProgramException ex) {
            correctExceptionTriggered = (ex.getCause() instanceof java.io.NotSerializableException);
        }
        Assert.assertTrue(correctExceptionTriggered);
    }

    private static class KeySelector7
            implements KeySelector<
                    Tuple5<Integer, Long, Integer, String, Long>, Tuple2<Integer, Long>> {
        private static final long serialVersionUID = 1L;

        @Override
        public Tuple2<Integer, Long> getKey(Tuple5<Integer, Long, Integer, String, Long> t) {
            return new Tuple2<Integer, Long>(t.f0, t.f4);
        }
    }

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

        @Override
        public Tuple2<Integer, Long> getKey(Tuple3<Integer, Long, String> t) {
            return new Tuple2<Integer, Long>(t.f0, t.f1);
        }
    }

    @Test
    public void testCoGroupTwoCustomTypeInputsWithExpressionKeys() throws Exception {
        /*
         * CoGroup on two custom type inputs using expression keys
         */

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSet<CustomType> ds = CollectionDataSets.getCustomTypeDataSet(env);
        DataSet<CustomType> ds2 = CollectionDataSets.getCustomTypeDataSet(env);
        DataSet<CustomType> coGroupDs =
                ds.coGroup(ds2).where("myInt").equalTo("myInt").with(new CustomTypeCoGroup());

        List<CustomType> result = coGroupDs.collect();

        String expected =
                "1,0,test\n"
                        + "2,6,test\n"
                        + "3,24,test\n"
                        + "4,60,test\n"
                        + "5,120,test\n"
                        + "6,210,test\n";

        compareResultAsText(result, expected);
    }

    @Test
    public void testCoGroupOnTwoCustomTypeInputsWithExpressionKeyAndFieldSelector()
            throws Exception {
        /*
         * CoGroup on two custom type inputs using expression keys
         */

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSet<POJO> ds = CollectionDataSets.getSmallPojoDataSet(env);
        DataSet<Tuple7<Integer, String, Integer, Integer, Long, String, Long>> ds2 =
                CollectionDataSets.getSmallTuplebasedDataSet(env);
        DataSet<CustomType> coGroupDs =
                ds.coGroup(ds2).where("nestedPojo.longNumber").equalTo(6).with(new CoGroup1());

        List<CustomType> result = coGroupDs.collect();

        String expected = "-1,20000,Flink\n" + "-1,10000,Flink\n" + "-1,30000,Flink\n";

        compareResultAsText(result, expected);
    }

    private static class CoGroup1
            implements CoGroupFunction<
                    POJO,
                    Tuple7<Integer, String, Integer, Integer, Long, String, Long>,
                    CustomType> {
        private static final long serialVersionUID = 1L;

        @Override
        public void coGroup(
                Iterable<POJO> first,
                Iterable<Tuple7<Integer, String, Integer, Integer, Long, String, Long>> second,
                Collector<CustomType> out)
                throws Exception {
            for (POJO p : first) {
                for (Tuple7<Integer, String, Integer, Integer, Long, String, Long> t : second) {
                    Assert.assertTrue(p.nestedPojo.longNumber == t.f6);
                    out.collect(new CustomType(-1, p.nestedPojo.longNumber, "Flink"));
                }
            }
        }
    }

    @Test
    public void testCoGroupFieldSelectorAndComplicatedKeySelector() throws Exception {
        /*
         * CoGroup field-selector (expression keys) + key selector function
         * The key selector is unnecessary complicated (Tuple1) ;)
         */

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSet<POJO> ds = CollectionDataSets.getSmallPojoDataSet(env);
        DataSet<Tuple7<Integer, String, Integer, Integer, Long, String, Long>> ds2 =
                CollectionDataSets.getSmallTuplebasedDataSet(env);
        DataSet<CustomType> coGroupDs =
                ds.coGroup(ds2).where(new KeySelector6()).equalTo(6).with(new CoGroup3());

        List<CustomType> result = coGroupDs.collect();

        String expected = "-1,20000,Flink\n" + "-1,10000,Flink\n" + "-1,30000,Flink\n";

        compareResultAsText(result, expected);
    }

    private static class KeySelector6 implements KeySelector<POJO, Tuple1<Long>> {
        private static final long serialVersionUID = 1L;

        @Override
        public Tuple1<Long> getKey(POJO value) throws Exception {
            return new Tuple1<Long>(value.nestedPojo.longNumber);
        }
    }

    private static class CoGroup3
            implements CoGroupFunction<
                    POJO,
                    Tuple7<Integer, String, Integer, Integer, Long, String, Long>,
                    CustomType> {
        private static final long serialVersionUID = 1L;

        @Override
        public void coGroup(
                Iterable<POJO> first,
                Iterable<Tuple7<Integer, String, Integer, Integer, Long, String, Long>> second,
                Collector<CustomType> out)
                throws Exception {
            for (POJO p : first) {
                for (Tuple7<Integer, String, Integer, Integer, Long, String, Long> t : second) {
                    Assert.assertTrue(p.nestedPojo.longNumber == t.f6);
                    out.collect(new CustomType(-1, p.nestedPojo.longNumber, "Flink"));
                }
            }
        }
    }

    @Test
    public void testCoGroupFieldSelectorAndKeySelector() throws Exception {
        /*
         * CoGroup field-selector (expression keys) + key selector function
         * The key selector is simple here
         */

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSet<POJO> ds = CollectionDataSets.getSmallPojoDataSet(env);
        DataSet<Tuple7<Integer, String, Integer, Integer, Long, String, Long>> ds2 =
                CollectionDataSets.getSmallTuplebasedDataSet(env);
        DataSet<CustomType> coGroupDs =
                ds.coGroup(ds2).where(new KeySelector1()).equalTo(6).with(new CoGroup2());

        List<CustomType> result = coGroupDs.collect();

        String expected = "-1,20000,Flink\n" + "-1,10000,Flink\n" + "-1,30000,Flink\n";

        compareResultAsText(result, expected);
    }

    @Test
    public void testCoGroupWithAtomicType1() throws Exception {
        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
        DataSet<Tuple3<Integer, Long, String>> ds1 = CollectionDataSets.getSmall3TupleDataSet(env);
        DataSet<Integer> ds2 = env.fromElements(0, 1, 2);

        DataSet<Tuple3<Integer, Long, String>> coGroupDs =
                ds1.coGroup(ds2).where(0).equalTo("*").with(new CoGroupAtomic1());

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

        String expected = "(1,1,Hi)\n" + "(2,2,Hello)";

        compareResultAsText(result, expected);
    }

    @Test
    public void testCoGroupWithAtomicType2() throws Exception {
        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
        DataSet<Integer> ds1 = env.fromElements(0, 1, 2);
        DataSet<Tuple3<Integer, Long, String>> ds2 = CollectionDataSets.getSmall3TupleDataSet(env);

        DataSet<Tuple3<Integer, Long, String>> coGroupDs =
                ds1.coGroup(ds2).where("*").equalTo(0).with(new CoGroupAtomic2());

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

        String expected = "(1,1,Hi)\n" + "(2,2,Hello)";

        compareResultAsText(result, expected);
    }

    @Test
    public void testCoGroupWithRangePartitioning() throws Exception {
        /*
         * Test coGroup on tuples with multiple key field positions and same customized distribution
         */

        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSet<Tuple5<Integer, Long, Integer, String, Long>> ds1 =
                CollectionDataSets.get5TupleDataSet(env);
        DataSet<Tuple3<Integer, Long, String>> ds2 = CollectionDataSets.get3TupleDataSet(env);

        env.setParallelism(4);
        TestDistribution testDis = new TestDistribution();
        DataSet<Tuple3<Integer, Long, String>> coGrouped =
                DataSetUtils.partitionByRange(ds1, testDis, 0, 4)
                        .coGroup(DataSetUtils.partitionByRange(ds2, testDis, 0, 1))
                        .where(0, 4)
                        .equalTo(0, 1)
                        .with(new Tuple5Tuple3CoGroup());

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

        String expected =
                "1,1,Hallo\n"
                        + "2,2,Hallo Welt\n"
                        + "3,2,Hallo Welt wie gehts?\n"
                        + "3,2,ABC\n"
                        + "5,3,HIJ\n"
                        + "5,3,IJK\n";

        compareResultAsTuples(result, expected);
    }

    @Test
    public void testCoGroupLambda() throws Exception {
        final ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();

        DataSet<Tuple2<Integer, String>> left =
                env.fromElements(
                        new Tuple2<>(1, "hello"), new Tuple2<>(2, "what's"), new Tuple2<>(2, "up"));
        DataSet<Tuple2<Integer, String>> right =
                env.fromElements(
                        new Tuple2<>(1, "not"), new Tuple2<>(1, "much"), new Tuple2<>(2, "really"));
        DataSet<Integer> joined =
                left.coGroup(right)
                        .where(0)
                        .equalTo(0)
                        .with(
                                (Iterable<Tuple2<Integer, String>> values1,
                                        Iterable<Tuple2<Integer, String>> values2,
                                        Collector<Integer> out) -> {
                                    int sum = 0;
                                    for (Tuple2<Integer, String> next : values1) {
                                        sum += next.f0;
                                    }
                                    for (Tuple2<Integer, String> next : values2) {
                                        sum += next.f0;
                                    }
                                    out.collect(sum);
                                })
                        .returns(Integer.class);
        List<Integer> result = joined.collect();

        String expected = "6\n3\n";

        compareResultAsText(result, expected);
    }

    // --------------------------------------------------------------------------------------------
    //  UDF classes
    // --------------------------------------------------------------------------------------------

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

        @Override
        public Long getKey(POJO value) throws Exception {
            return value.nestedPojo.longNumber;
        }
    }

    private static class CoGroup2
            implements CoGroupFunction<
                    POJO,
                    Tuple7<Integer, String, Integer, Integer, Long, String, Long>,
                    CustomType> {
        private static final long serialVersionUID = 1L;

        @Override
        public void coGroup(
                Iterable<POJO> first,
                Iterable<Tuple7<Integer, String, Integer, Integer, Long, String, Long>> second,
                Collector<CustomType> out)
                throws Exception {
            for (POJO p : first) {
                for (Tuple7<Integer, String, Integer, Integer, Long, String, Long> t : second) {
                    Assert.assertTrue(p.nestedPojo.longNumber == t.f6);
                    out.collect(new CustomType(-1, p.nestedPojo.longNumber, "Flink"));
                }
            }
        }
    }

    private static class Tuple5CoGroup
            implements CoGroupFunction<
                    Tuple5<Integer, Long, Integer, String, Long>,
                    Tuple5<Integer, Long, Integer, String, Long>,
                    Tuple2<Integer, Integer>> {

        private static final long serialVersionUID = 1L;

        @Override
        public void coGroup(
                Iterable<Tuple5<Integer, Long, Integer, String, Long>> first,
                Iterable<Tuple5<Integer, Long, Integer, String, Long>> second,
                Collector<Tuple2<Integer, Integer>> out) {
            int sum = 0;
            int id = 0;

            for (Tuple5<Integer, Long, Integer, String, Long> element : first) {
                sum += element.f2;
                id = element.f0;
            }

            for (Tuple5<Integer, Long, Integer, String, Long> element : second) {
                sum += element.f2;
                id = element.f0;
            }

            out.collect(new Tuple2<Integer, Integer>(id, sum));
        }
    }

    private static class CustomTypeCoGroup
            implements CoGroupFunction<CustomType, CustomType, CustomType> {

        private static final long serialVersionUID = 1L;

        @Override
        public void coGroup(
                Iterable<CustomType> first,
                Iterable<CustomType> second,
                Collector<CustomType> out) {

            CustomType o = new CustomType(0, 0, "test");

            for (CustomType element : first) {
                o.myInt = element.myInt;
                o.myLong += element.myLong;
            }

            for (CustomType element : second) {
                o.myInt = element.myInt;
                o.myLong += element.myLong;
            }

            out.collect(o);
        }
    }

    private static class MixedCoGroup
            implements CoGroupFunction<
                    Tuple5<Integer, Long, Integer, String, Long>,
                    CustomType,
                    Tuple3<Integer, Long, String>> {

        private static final long serialVersionUID = 1L;

        @Override
        public void coGroup(
                Iterable<Tuple5<Integer, Long, Integer, String, Long>> first,
                Iterable<CustomType> second,
                Collector<Tuple3<Integer, Long, String>> out)
                throws Exception {

            long sum = 0;
            int id = 0;

            for (Tuple5<Integer, Long, Integer, String, Long> element : first) {
                sum += element.f0;
                id = element.f2;
            }

            for (CustomType element : second) {
                id = element.myInt;
                sum += element.myLong;
            }

            out.collect(new Tuple3<Integer, Long, String>(id, sum, "test"));
        }
    }

    private static class MixedCoGroup2
            implements CoGroupFunction<
                    CustomType, Tuple5<Integer, Long, Integer, String, Long>, CustomType> {

        private static final long serialVersionUID = 1L;

        @Override
        public void coGroup(
                Iterable<CustomType> first,
                Iterable<Tuple5<Integer, Long, Integer, String, Long>> second,
                Collector<CustomType> out) {
            CustomType o = new CustomType(0, 0, "test");

            for (CustomType element : first) {
                o.myInt = element.myInt;
                o.myLong += element.myLong;
            }

            for (Tuple5<Integer, Long, Integer, String, Long> element : second) {
                o.myInt = element.f2;
                o.myLong += element.f0;
            }

            out.collect(o);
        }
    }

    private static class Tuple3ReturnLeft
            implements CoGroupFunction<
                    Tuple3<Integer, Long, String>,
                    Tuple3<Integer, Long, String>,
                    Tuple3<Integer, Long, String>> {

        private static final long serialVersionUID = 1L;

        @Override
        public void coGroup(
                Iterable<Tuple3<Integer, Long, String>> first,
                Iterable<Tuple3<Integer, Long, String>> second,
                Collector<Tuple3<Integer, Long, String>> out) {
            for (Tuple3<Integer, Long, String> element : first) {
                if (element.f0 < 6) {
                    out.collect(element);
                }
            }
        }
    }

    private static class Tuple5ReturnRight
            implements CoGroupFunction<
                    Tuple5<Integer, Long, Integer, String, Long>,
                    Tuple5<Integer, Long, Integer, String, Long>,
                    Tuple5<Integer, Long, Integer, String, Long>> {

        private static final long serialVersionUID = 1L;

        @Override
        public void coGroup(
                Iterable<Tuple5<Integer, Long, Integer, String, Long>> first,
                Iterable<Tuple5<Integer, Long, Integer, String, Long>> second,
                Collector<Tuple5<Integer, Long, Integer, String, Long>> out) {
            for (Tuple5<Integer, Long, Integer, String, Long> element : second) {
                if (element.f0 < 4) {
                    out.collect(element);
                }
            }
        }
    }

    private static class Tuple5CoGroupBC
            extends RichCoGroupFunction<
                    Tuple5<Integer, Long, Integer, String, Long>,
                    Tuple5<Integer, Long, Integer, String, Long>,
                    Tuple3<Integer, Integer, Integer>> {

        private static final long serialVersionUID = 1L;

        private int broadcast = 42;

        @Override
        public void open(Configuration config) {

            Collection<Integer> ints = this.getRuntimeContext().getBroadcastVariable("ints");
            int sum = 0;
            for (Integer i : ints) {
                sum += i;
            }
            broadcast = sum;
        }

        @Override
        public void coGroup(
                Iterable<Tuple5<Integer, Long, Integer, String, Long>> first,
                Iterable<Tuple5<Integer, Long, Integer, String, Long>> second,
                Collector<Tuple3<Integer, Integer, Integer>> out) {
            int sum = 0;
            int id = 0;

            for (Tuple5<Integer, Long, Integer, String, Long> element : first) {
                sum += element.f2;
                id = element.f0;
            }

            for (Tuple5<Integer, Long, Integer, String, Long> element : second) {
                sum += element.f2;
                id = element.f0;
            }

            out.collect(new Tuple3<Integer, Integer, Integer>(id, sum, broadcast));
        }
    }

    private static class Tuple5Tuple3CoGroup
            implements CoGroupFunction<
                    Tuple5<Integer, Long, Integer, String, Long>,
                    Tuple3<Integer, Long, String>,
                    Tuple3<Integer, Long, String>> {

        private static final long serialVersionUID = 1L;

        @Override
        public void coGroup(
                Iterable<Tuple5<Integer, Long, Integer, String, Long>> first,
                Iterable<Tuple3<Integer, Long, String>> second,
                Collector<Tuple3<Integer, Long, String>> out) {
            List<String> strs = new ArrayList<String>();

            for (Tuple5<Integer, Long, Integer, String, Long> t : first) {
                strs.add(t.f3);
            }

            for (Tuple3<Integer, Long, String> t : second) {
                for (String s : strs) {
                    out.collect(new Tuple3<Integer, Long, String>(t.f0, t.f1, s));
                }
            }
        }
    }

    private static class CoGroupAtomic1
            implements CoGroupFunction<
                    Tuple3<Integer, Long, String>, Integer, Tuple3<Integer, Long, String>> {

        private static final long serialVersionUID = 1L;

        @Override
        public void coGroup(
                Iterable<Tuple3<Integer, Long, String>> first,
                Iterable<Integer> second,
                Collector<Tuple3<Integer, Long, String>> out)
                throws Exception {
            List<Integer> ints = new ArrayList<Integer>();

            for (Integer i : second) {
                ints.add(i);
            }

            for (Tuple3<Integer, Long, String> t : first) {
                for (Integer i : ints) {
                    if (t.f0.equals(i)) {
                        out.collect(t);
                    }
                }
            }
        }
    }

    private static class CoGroupAtomic2
            implements CoGroupFunction<
                    Integer, Tuple3<Integer, Long, String>, Tuple3<Integer, Long, String>> {

        private static final long serialVersionUID = 1L;

        @Override
        public void coGroup(
                Iterable<Integer> first,
                Iterable<Tuple3<Integer, Long, String>> second,
                Collector<Tuple3<Integer, Long, String>> out)
                throws Exception {
            List<Integer> ints = new ArrayList<Integer>();

            for (Integer i : first) {
                ints.add(i);
            }

            for (Tuple3<Integer, Long, String> t : second) {
                for (Integer i : ints) {
                    if (t.f0.equals(i)) {
                        out.collect(t);
                    }
                }
            }
        }
    }

    /** Test {@link DataDistribution}. */
    public static class TestDistribution implements DataDistribution {
        public Object[][] boundaries =
                new Object[][] {
                    new Object[] {2, 2L},
                    new Object[] {5, 4L},
                    new Object[] {10, 12L},
                    new Object[] {21, 6L}
                };

        public TestDistribution() {}

        @Override
        public Object[] getBucketBoundary(int bucketNum, int totalNumBuckets) {
            return boundaries[bucketNum];
        }

        @Override
        public int getNumberOfFields() {
            return 2;
        }

        @Override
        public TypeInformation[] getKeyTypes() {
            return new TypeInformation[] {
                BasicTypeInfo.INT_TYPE_INFO, BasicTypeInfo.LONG_TYPE_INFO
            };
        }

        @Override
        public void write(DataOutputView out) throws IOException {}

        @Override
        public void read(DataInputView in) throws IOException {}

        @Override
        public boolean equals(Object obj) {
            return obj instanceof TestDistribution;
        }
    }
}