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

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

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.io.PrintWriter;
import java.io.StringWriter;
import java.util.Iterator;
import java.util.List;
import java.util.Properties;
import java.util.Random;
import java.util.regex.Pattern;

import org.apache.hadoop.mapreduce.Job;
import org.apache.pig.PigServer;
import org.apache.pig.ResourceSchema;
import org.apache.pig.backend.executionengine.ExecJob;
import org.apache.pig.backend.hadoop.executionengine.mapReduceLayer.InputSizeReducerEstimator;
import org.apache.pig.backend.hadoop.executionengine.tez.runtime.PigGraceShuffleVertexManager;
import org.apache.pig.builtin.PigStorage;
import org.apache.pig.data.Tuple;
import org.apache.pig.impl.plan.NodeIdGenerator;
import org.apache.pig.impl.util.ObjectSerializer;
import org.apache.pig.impl.util.UDFContext;
import org.apache.pig.test.Util;
import org.apache.tez.dag.library.vertexmanager.ShuffleVertexManager;
import org.junit.AfterClass;
import org.junit.Before;
import org.junit.BeforeClass;
import org.junit.Test;

public class TestTezGraceParallelism {
    private static PigServer pigServer;
    private static final String INPUT_FILE1 = TestTezGraceParallelism.class.getName() + "_1";
    private static final String INPUT_FILE2 = TestTezGraceParallelism.class.getName() + "_2";
    private static final String INPUT_DIR = Util.getTestDirectory(TestTezGraceParallelism.class);

    @Before
    public void setUp() throws Exception {
        pigServer = new PigServer(Util.getLocalTestMode());
    }

    private static void createFiles() throws IOException {
        new File(INPUT_DIR).mkdirs();

        PrintWriter w = new PrintWriter(new FileWriter(INPUT_DIR + "/" + INPUT_FILE1));

        String boyNames[] = {"Noah", "Liam", "Jacob", "Mason", "William",
                "Ethan", "Michael", "Alexander", "Jayden", "Daniel"};
        String girlNames[] = {"Sophia", "Emma", "Olivia", "Isabella", "Ava",
                "Mia", "Emily", "Abigail", "Madison", "Elizabeth"};

        String names[] = new String[boyNames.length + girlNames.length];
        for (int i=0;i<boyNames.length;i++) {
            names[i] = boyNames[i];
        }
        for (int i=0;i<girlNames.length;i++) {
            names[boyNames.length+i] = girlNames[i];
        }

        Random rand = new Random(1);
        for (int i=0;i<1000;i++) {
            w.println(names[rand.nextInt(names.length)] + "\t" + rand.nextInt(18));
        }
        w.close();

        w = new PrintWriter(new FileWriter(INPUT_DIR + "/" + INPUT_FILE2));
        for (String name : boyNames) {
            w.println(name + "\t" + "M");
        }
        for (String name : girlNames) {
            w.println(name + "\t" + "F");
        }
        w.close();
    }

    private static void deleteFiles() {
        Util.deleteDirectory(new File(INPUT_DIR));
    }

    @BeforeClass
    public static void oneTimeSetUp() throws Exception {
        createFiles();
    }

    @AfterClass
    public static void oneTimeTearDown() throws Exception {
        deleteFiles();
    }

    @Test
    public void testDecreaseParallelism() throws IOException{
        NodeIdGenerator.reset();
        PigServer.resetScope();
        StringWriter writer = new StringWriter();
        Util.createLogAppender("testDecreaseParallelism", writer, new Class[]{PigGraceShuffleVertexManager.class, ShuffleVertexManager.class});
        try {
            // DAG: 47 \
            //           -> 49(join) -> 52(distinct) -> 56(group)
            //      48 /
            // Parallelism at compile time:
            // DAG: 47(1) \
            //              -> 49(2) -> 52(20) -> 56(200)
            //      48(1) /
            // However, when 49 finishes, the actual output of 49 only justify parallelism 1.
            // We adjust the parallelism for 56 to 7 based on this.
            // At runtime, ShuffleVertexManager still kick in and further reduce parallelism from 7 to 1.
            //
            pigServer.registerQuery("A = load '" + INPUT_DIR + "/" + INPUT_FILE1 + "' as (name:chararray, age:int);");
            pigServer.registerQuery("B = load '" + INPUT_DIR + "/" + INPUT_FILE2 + "' as (name:chararray, gender:chararray);");
            pigServer.registerQuery("C = join A by name, B by name;");
            pigServer.registerQuery("D = foreach C generate A::name as name, A::age as age, gender;");
            pigServer.registerQuery("E = distinct D;");
            pigServer.registerQuery("F = group E by gender;");
            pigServer.registerQuery("G = foreach F generate group as gender, SUM(E.age);");
            Iterator<Tuple> iter = pigServer.openIterator("G");
            List<Tuple> expectedResults = Util
                    .getTuplesFromConstantTupleStrings(new String[] {
                            "('F',1349L)", "('M',1373L)"});
            Util.checkQueryOutputsAfterSort(iter, expectedResults);
            assertTrue(writer.toString().contains("Initialize parallelism for scope-52 to 18"));
            assertTrue(writer.toString().contains("Initialize parallelism for scope-56 to 7"));
            assertTrue(writer.toString().contains("Reduce auto parallelism for vertex: scope-49 to 1 from 2"));
            assertTrue(writer.toString().contains("Reduce auto parallelism for vertex: scope-52 to 1 from 18"));
            assertTrue(writer.toString().contains("Reduce auto parallelism for vertex: scope-56 to 1 from 7"));
        } finally {
            Util.removeLogAppender("testDecreaseParallelism", PigGraceShuffleVertexManager.class, ShuffleVertexManager.class);
        }
    }

    @Test
    public void testIncreaseParallelism() throws IOException{
        NodeIdGenerator.reset();
        PigServer.resetScope();
        StringWriter writer = new StringWriter();
        Util.createLogAppender("testIncreaseParallelism", writer, new Class[]{PigGraceShuffleVertexManager.class, ShuffleVertexManager.class});
        try {
            // DAG: 35 \             /  46(sample)   \
            //           -> 37(order)        ->    56(order) -> 58(order) -> 64(distinct)
            //      36 /
            // Parallelism at compile time:
            // DAG: 35(1) \          /  46(1)   \
            //              -> 37(2)     ->    56(-1) -> 58(-1) -> 64(20)
            //      36(1) /
            // However, when 56 finishes, the actual output of 56 need parallelism 5.
            // We adjust the parallelism for 64 to 50 based on this.
            // At runtime, ShuffleVertexManager will play and reduce parallelism from 50
            pigServer.getPigContext().getProperties().setProperty(InputSizeReducerEstimator.BYTES_PER_REDUCER_PARAM, "80000");
            pigServer.registerQuery("A = load '" + INPUT_DIR + "/" + INPUT_FILE1 + "' as (name:chararray, age:int);");
            pigServer.registerQuery("B = load '" + INPUT_DIR + "/" + INPUT_FILE2 + "' as (name:chararray, gender:chararray);");
            pigServer.registerQuery("C = join A by 1, B by 1;");
            pigServer.registerQuery("D = foreach C generate A::name as name, A::age as age, gender;");
            pigServer.registerQuery("E = order D by name;");
            pigServer.registerQuery("F = distinct E;");
            Iterator<Tuple> iter = pigServer.openIterator("F");
            int count = 0;
            while (iter.hasNext()) {
                iter.next();
                count++;
            }
            assertEquals(count, 644);
            System.out.println(writer.toString());
            assertTrue(writer.toString().contains("Initialize parallelism for scope-64 to 45"));
            // There are randomness in which task finishes first, so the auto parallelism could result different result
            assertTrue(Pattern.compile("Reduce auto parallelism for vertex: scope-64 to (\\d+)* from 45").matcher(writer.toString()).find());
        } finally {
            Util.removeLogAppender("testIncreaseParallelism", PigGraceShuffleVertexManager.class, ShuffleVertexManager.class);
        }
    }

    @Test
    public void testJoinWithDifferentDepth() throws IOException{
        NodeIdGenerator.reset();
        PigServer.resetScope();
        StringWriter writer = new StringWriter();
        Util.createLogAppender("testJoinWithDifferentDepth", writer, PigGraceShuffleVertexManager.class);
        try {
            // DAG:     /  51(sample) \
            //      42        ->        61(order) -> 63(order) -> 69(distinct) \
            //                                                                   -> 80(join)
            //                                             78  -> 79(group)    /
            // The join(80) has two inputs: 69 with deeper pipeline, 79 with narrower.
            // This test is to make sure 79 can start (by invoking 80.setParallelism) early,
            // don't need to wait for 63 complete
            pigServer.registerQuery("A = load '" + INPUT_DIR + "/" + INPUT_FILE2 + "' as (name:chararray, gender:chararray);");
            pigServer.registerQuery("B = order A by name;");
            pigServer.registerQuery("C = distinct B;");
            pigServer.registerQuery("D = load '" + INPUT_DIR + "/" + INPUT_FILE1 + "' as (name:chararray, age:int);");
            pigServer.registerQuery("E = group D by name;");
            pigServer.registerQuery("F = foreach E generate FLATTEN(group) as name, AVG(D.age) as avg_age;");
            pigServer.registerQuery("G = join C by name, F by name;");
            Iterator<Tuple> iter = pigServer.openIterator("G");
            int count = 0;
            while (iter.hasNext()) {
                iter.next();
                count++;
            }
            assertEquals(count, 20);
            assertTrue(writer.toString().contains("All predecessors for scope-79 are finished, time to set parallelism for scope-80"));
            assertTrue(writer.toString().contains("Initialize parallelism for scope-80 to 10"));
        } finally {
            Util.removeLogAppender("testJoinWithDifferentDepth", PigGraceShuffleVertexManager.class);
        }
    }

    @Test
    public void testJoinWithDifferentDepth2() throws IOException{
        NodeIdGenerator.reset();
        PigServer.resetScope();
        StringWriter writer = new StringWriter();
        Util.createLogAppender("testJoinWithDifferentDepth2", writer, PigGraceShuffleVertexManager.class);
        try {
            // DAG:     /  40(sample) \
            //      31        ->        50(order) -> 52(order) -> 58(distinct) \
            //                                                                   -> 68(join)
            //                                                           67    /
            // The join(68) should start immediately. We will not use PigGraceShuffleVertexManager in this case
            pigServer.registerQuery("A = load '" + INPUT_DIR + "/" + INPUT_FILE2 + "' as (name:chararray, gender:chararray);");
            pigServer.registerQuery("B = order A by name;");
            pigServer.registerQuery("C = distinct B;");
            pigServer.registerQuery("D = load '" + INPUT_DIR + "/" + INPUT_FILE1 + "' as (name:chararray, age:int);");
            pigServer.registerQuery("E = join C by name, D by name;");
            Iterator<Tuple> iter = pigServer.openIterator("E");
            int count = 0;
            while (iter.hasNext()) {
                iter.next();
                count++;
            }
            assertEquals(count, 1000);
            assertFalse(writer.toString().contains("scope-68"));
        } finally {
            Util.removeLogAppender("testJoinWithDifferentDepth2", PigGraceShuffleVertexManager.class);
        }
    }

    @Test
    // See PIG-4635 for a NPE in TezOperDependencyParallelismEstimator
    public void testJoinWithUnion() throws IOException{
        NodeIdGenerator.reset();
        PigServer.resetScope();
        StringWriter writer = new StringWriter();
        Util.createLogAppender("testJoinWithUnion", writer, PigGraceShuffleVertexManager.class);
        try {
            // DAG: 29 -> 32 -> 36 \
            //                       -> 55 (vertex group) -> 51
            //      37 -> 40 -> 44 /
            pigServer.registerQuery("A = load '" + INPUT_DIR + "/" + INPUT_FILE2 + "' as (name:chararray, gender:chararray);");
            pigServer.registerQuery("B = distinct A;");
            pigServer.registerQuery("C = group B by name;");
            pigServer.registerQuery("D = load '" + INPUT_DIR + "/" + INPUT_FILE2 + "' as (name:chararray, gender:chararray);");
            pigServer.registerQuery("E = distinct D;");
            pigServer.registerQuery("F = group E by name;");
            pigServer.registerQuery("G = union C, F;");
            pigServer.registerQuery("H = distinct G;");
            Iterator<Tuple> iter = pigServer.openIterator("H");
            int count = 0;
            while (iter.hasNext()) {
                iter.next();
                count++;
            }
            assertEquals(count, 20);
            assertTrue(writer.toString().contains("time to set parallelism for scope-36"));
            assertTrue(writer.toString().contains("time to set parallelism for scope-44"));
        } finally {
            Util.removeLogAppender("testJoinWithUnion", PigGraceShuffleVertexManager.class);
        }
    }

    @Test
    // See PIG-4703
    public void testUDFContextSetInBackend() throws IOException{
        NodeIdGenerator.reset();
        PigServer.resetScope();
        File outputDir = File.createTempFile("intemediate", "txt");
        outputDir.delete();
        pigServer.registerQuery("A = load '" + INPUT_DIR + "/" + INPUT_FILE2 + "' as (name:chararray, gender:chararray);");
        pigServer.registerQuery("B = order A by name;");
        pigServer.registerQuery("C = distinct B;");
        pigServer.registerQuery("D = load '" + INPUT_DIR + "/" + INPUT_FILE1 + "' as (name:chararray, age:int);");
        pigServer.registerQuery("E = group D by name;");
        pigServer.registerQuery("F = foreach E generate group as name, AVG(D.age) as avg_age;");
        pigServer.registerQuery("G = join C by name left, F by name;");
        ExecJob job = pigServer.store("G", Util.removeColon(outputDir.getAbsolutePath()), StorerWithUDFContextCheck.class.getName());
        assertTrue(job.getStatus() == ExecJob.JOB_STATUS.COMPLETED);
    }

    static public class StorerWithUDFContextCheck extends PigStorage {
        @Override
        public void checkSchema(ResourceSchema resourceSchema) throws IOException {
            UDFContext.getUDFContext().getUDFProperties(this.getClass(), new String[]{signature})
            .setProperty("schema", ObjectSerializer.serialize(resourceSchema));
        }
        @Override
        public void setStoreLocation(String location, Job job) throws IOException {
            Properties udfProps = UDFContext.getUDFContext().getUDFProperties(
                    this.getClass(), new String[]{signature});
            ResourceSchema rs = (ResourceSchema)ObjectSerializer.deserialize(udfProps.getProperty("schema"));
            if (rs == null) {
                throw new IOException("Should not be null");
            }
            super.setStoreLocation(location, job);
        }
    }

    @Test
    // See PIG-4786
    public void testCross() throws IOException{
        // scope-90 is the cross vertex. It should not use PigGraceShuffleVertexManager
        NodeIdGenerator.reset();
        PigServer.resetScope();
        StringWriter writer = new StringWriter();
        Util.createLogAppender("testCross", writer, PigGraceShuffleVertexManager.class);
        File outputDir = File.createTempFile("intemediate", "txt");
        outputDir.delete();
        pigServer.getPigContext().getProperties().setProperty("mapreduce.input.fileinputformat.split.maxsize", "3000");
        pigServer.getPigContext().getProperties().setProperty("pig.noSplitCombination", "true");
        pigServer.registerQuery("A = load '" + INPUT_DIR + "/" + INPUT_FILE2 + "' as (name:chararray, gender:chararray);");
        pigServer.registerQuery("B = order A by name;");
        pigServer.registerQuery("C = distinct B;");
        pigServer.registerQuery("D = load '" + INPUT_DIR + "/" + INPUT_FILE1 + "' as (name:chararray, age:int);");
        pigServer.registerQuery("E = group D by name;");
        pigServer.registerQuery("F = foreach E generate group as name, AVG(D.age) as avg_age;");
        pigServer.registerQuery("G = cross C, F;");
        Iterator<Tuple> iter = pigServer.openIterator("G");
        int count = 0;
        while (iter.hasNext()) {
            iter.next();
            count++;
        }
        assertEquals(count, 400);
        assertFalse(writer.toString().contains("scope-90"));
    }
}