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

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

import java.io.BufferedWriter;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.math.BigDecimal;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.List;

import org.apache.commons.lang.RandomStringUtils;
import org.apache.commons.lang.math.RandomUtils;
import org.apache.hadoop.hive.ql.io.sarg.SearchArgument;
import org.apache.log4j.FileAppender;
import org.apache.log4j.Level;
import org.apache.log4j.Logger;
import org.apache.log4j.SimpleLayout;
import org.apache.pig.Expression;
import org.apache.pig.Expression.OpType;
import org.apache.pig.PigServer;
import org.apache.pig.backend.executionengine.ExecJob;
import org.apache.pig.impl.PigImplConstants;
import org.apache.pig.impl.util.ObjectSerializer;
import org.apache.pig.newplan.Operator;
import org.apache.pig.newplan.PredicatePushDownFilterExtractor;
import org.apache.pig.newplan.logical.relational.LOFilter;
import org.apache.pig.newplan.logical.relational.LogicalPlan;
import org.apache.pig.newplan.logical.rules.ColumnPruneVisitor;
import org.apache.pig.test.MiniGenericCluster;
import org.apache.pig.test.Util;
import org.apache.pig.tools.pigstats.JobStats;
import org.joda.time.DateTime;
import org.joda.time.DateTimeZone;
import org.junit.After;
import org.junit.AfterClass;
import org.junit.Assert;
import org.junit.Before;
import org.junit.BeforeClass;
import org.junit.Test;

public class TestOrcStoragePushdown {

    private static List<OpType> supportedOpTypes;
    private static MiniGenericCluster cluster;
    private static PigServer pigServer;
    private String query = "a = load 'foo' as (srcid:int, mrkt:chararray, dstid:int, name:chararray, " +
            "age:int, browser:map[], location:tuple(country:chararray, zip:int));";
    private OrcStorage orcStorage;

    private static final String basedir = "test/org/apache/pig/builtin/orc/";
    private static final String inpbasedir = "build/test/TestOrcStorage_in/";
    private static final String outbasedir = "build/test/TestOrcStorage_out/";
    private static String INPUT = inpbasedir + "TestOrcStorage_1";
    private static String OUTPUT1 = outbasedir + "TestOrcStorage_1";
    private static String OUTPUT2 = outbasedir + "TestOrcStorage_2";
    private static String OUTPUT3 = outbasedir + "TestOrcStorage_3";
    private static String OUTPUT4 = outbasedir + "TestOrcStorage_4";

    private static File logFile;

    @BeforeClass
    public static void oneTimeSetup() throws Exception{
        cluster = MiniGenericCluster.buildCluster();
        Util.copyFromLocalToCluster(cluster, basedir + "orc-file-11-format.orc", basedir + "orc-file-11-format.orc");
        Util.copyFromLocalToCluster(cluster, basedir + "charvarchar.orc", basedir + "charvarchar.orc");
        createInputData();

        if(Util.WINDOWS){
            OUTPUT1 = OUTPUT1.replace("\\", "/");
            OUTPUT2 = OUTPUT2.replace("\\", "/");
            OUTPUT3 = OUTPUT3.replace("\\", "/");
            OUTPUT4 = OUTPUT4.replace("\\", "/");
        }

        supportedOpTypes = new ArrayList<OpType>();
        supportedOpTypes.add(OpType.OP_EQ);
        supportedOpTypes.add(OpType.OP_NE);
        supportedOpTypes.add(OpType.OP_GT);
        supportedOpTypes.add(OpType.OP_GE);
        supportedOpTypes.add(OpType.OP_LT);
        supportedOpTypes.add(OpType.OP_LE);
        supportedOpTypes.add(OpType.OP_IN);
        supportedOpTypes.add(OpType.OP_BETWEEN);
        supportedOpTypes.add(OpType.OP_NULL);
        supportedOpTypes.add(OpType.OP_NOT);
        supportedOpTypes.add(OpType.OP_AND);
        supportedOpTypes.add(OpType.OP_OR);

        Logger logger = Logger.getLogger(ColumnPruneVisitor.class);
        logger.removeAllAppenders();
        logger.setLevel(Level.INFO);
        SimpleLayout layout = new SimpleLayout();
        logFile = File.createTempFile("log", "");
        FileAppender appender = new FileAppender(layout, logFile.toString(), false, false, 0);
        logger.addAppender(appender);
    }

    private static void createInputData() throws Exception {
        pigServer = new PigServer(Util.getLocalTestMode());

        new File(inpbasedir).mkdirs();
        new File(outbasedir).mkdirs();
        String inputTxtFile = inpbasedir + File.separator + "input.txt";
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(new FileOutputStream(inputTxtFile), "UTF-8"));
        long[] lVal = new long[] {100L, 200L, 300L};
        float[] fVal = new float[] {50.0f, 100.0f, 200.0f, 300.0f};
        double[] dVal = new double[] {1000.11, 2000.22, 3000.33};
        StringBuilder sb = new StringBuilder();
        for (int i=1; i <= 10000; i++) {
            sb.append((i > 6500 && i <= 9000) ? true : false).append("\t"); //boolean
            sb.append((i > 1000 && i < 3000) ? 1 : 5).append("\t"); //byte
            sb.append((i > 2500 && i <= 4500) ? 100 : 200).append("\t"); //short
            sb.append(i).append("\t"); //int
            sb.append(lVal[i%3]).append("\t"); //long
            sb.append(fVal[i%4]).append("\t"); //float
            sb.append((i > 2500 && i < 3500) ? dVal[i%3] : dVal[i%1]).append("\t"); //double
            sb.append((i%2 == 1 ? "" : RandomStringUtils.random(100).replaceAll("\t", " ")
                    .replaceAll("\n", " ").replaceAll("\r", " "))).append("\t"); //bytearray
            sb.append((i%2 == 0 ? "" : RandomStringUtils.random(100).replaceAll("\t", " ")
                    .replaceAll("\n", " ").replaceAll("\r", " "))).append("\t"); //string
            int year;
            if (i > 5000 && i <= 8000) { //datetime
                year = RandomUtils.nextInt(4)+2010;
            } else {
                year = RandomUtils.nextInt(10)+2000;
            }
            sb.append(new DateTime(year, RandomUtils.nextInt(12)+1,
                    RandomUtils.nextInt(28)+1, RandomUtils.nextInt(24), RandomUtils.nextInt(60),
                    DateTimeZone.UTC).toString()).append("\t"); // datetime
            String bigString;
            if (i>7500) {
                bigString = RandomStringUtils.randomNumeric(9) + "." + RandomStringUtils.randomNumeric(5);
            } else {
                bigString = "1" + RandomStringUtils.randomNumeric(9) + "." + RandomStringUtils.randomNumeric(5);
            }
            sb.append(new BigDecimal(bigString)).append("\n"); //bigdecimal
            bw.write(sb.toString());
            sb.setLength(0);
        }
        bw.close();

        // Store only 1000 rows in each row block (MIN_ROW_INDEX_STRIDE is 1000. So can't use less than that)
        pigServer.registerQuery("A = load '" + Util.generateURI(inputTxtFile, pigServer.getPigContext()) + "' as (f1:boolean, f2:int, f3:int, f4:int, f5:long, f6:float, f7:double, f8:bytearray, f9:chararray, f10:datetime, f11:bigdecimal);");
        pigServer.registerQuery("store A into '" + Util.generateURI(INPUT, pigServer.getPigContext()) +"' using OrcStorage('-r 1000 -s 100000');");
        Util.copyFromLocalToCluster(cluster, INPUT, INPUT);
    }

    @AfterClass
    public static void tearDownAfterClass() throws Exception {
        Util.deleteDirectory(new File(inpbasedir));
        if (cluster != null) {
            Util.deleteFile(cluster, inpbasedir);
            cluster.shutDown();
        }
    }

    @Before
    public void setup() throws Exception{
        Util.resetStateForExecModeSwitch();
        pigServer = new PigServer(Util.getLocalTestMode());
        orcStorage = new OrcStorage();
    }

    @After
    public void teardown() throws IOException {
        if(pigServer != null) {
            pigServer.shutdown();
        }
        Util.deleteDirectory(new File(outbasedir));
        if (cluster != null) {
            Util.deleteFile(cluster, outbasedir);
        }
    }

    @Test
    public void testColumnPruning() throws Exception {
        Util.resetStateForExecModeSwitch();
        pigServer = new PigServer(cluster.getExecType(), cluster.getProperties());

        pigServer.registerQuery("A = load '" + basedir + "orc-file-11-format.orc' using OrcStorage();");
        ExecJob job = pigServer.store("A", OUTPUT1);
        JobStats stats = (JobStats) job.getStatistics().getJobGraph().getSources().get(0);
        long bytesWithoutPushdown = stats.getHdfsBytesRead();

        pigServer.registerQuery("PRUNE = load '" + basedir + "orc-file-11-format.orc' using OrcStorage();");
        pigServer.registerQuery("PRUNE = foreach PRUNE generate boolean1;");
        job = pigServer.store("PRUNE", OUTPUT2);
        Util.checkLogFileMessage(logFile, new String[]{"Columns pruned for PRUNE: $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13"}, true);
        stats = (JobStats) job.getStatistics().getJobGraph().getSources().get(0);
        long bytesWithPushdown = stats.getHdfsBytesRead();

        assertTrue((bytesWithoutPushdown - bytesWithPushdown) > 300000);
    }

    @Test
    public void testSimple() throws Exception {
        String q = query + "b = filter a by srcid == 10;" + "store b into 'out';";
        Expression expr = getExpressionForTest(q, Arrays.asList("srcid"));
        SearchArgument sarg = orcStorage.getSearchArgument(expr);
        assertEquals("leaf-0 = (EQUALS srcid 10)\n" +
                "expr = leaf-0", sarg.toString());
    }

    @Test
    public void testAndOr() throws Exception {
        String q = query + "b = filter a by (srcid > 10 or dstid <= 5) and name == 'foo' and mrkt is null;" + "store b into 'out';";
        Expression expr = getExpressionForTest(q, Arrays.asList("srcid", "dstid", "name", "mrkt"));
        SearchArgument sarg = orcStorage.getSearchArgument(expr);
        assertEquals("leaf-0 = (LESS_THAN_EQUALS srcid 10)\n" +
                "leaf-1 = (LESS_THAN_EQUALS dstid 5)\n" +
                "leaf-2 = (EQUALS name foo)\n" +
                "leaf-3 = (IS_NULL mrkt)\n" +
                "expr = (and (or (not leaf-0) leaf-1) leaf-2 leaf-3)", sarg.toString());
    }

    @Test
    public void testNot() throws Exception {
        String q = query + "b = filter a by srcid != 10 and mrkt is not null;" + "store b into 'out';";
        Expression expr = getExpressionForTest(q, Arrays.asList("srcid", "dstid", "name", "mrkt"));
        SearchArgument sarg = orcStorage.getSearchArgument(expr);
        assertEquals("leaf-0 = (EQUALS srcid 10)\n" +
                "leaf-1 = (IS_NULL mrkt)\n" +
                "expr = (and (not leaf-0) (not leaf-1))", sarg.toString());
    }

    @Test
    public void testBetweenExpression() throws Exception {
        // TODO: Add support for OP_BETWEEN expression type
        String q = query + "b = filter a by srcid > 10 or srcid < 20;" + "store b into 'out';";
        Expression expr = getExpressionForTest(q, Arrays.asList("srcid"));
        SearchArgument sarg = orcStorage.getSearchArgument(expr);
        assertEquals("leaf-0 = (LESS_THAN_EQUALS srcid 10)\n" +
                "leaf-1 = (LESS_THAN srcid 20)\n" +
                "expr = (or (not leaf-0) leaf-1)", sarg.toString());
    }

    @Test
    public void testInExpression() throws Exception {
        // TODO: Add support for OP_IN expression type
        String q = query + "b = filter a by srcid == 10 or srcid == 11;" + "store b into 'out';";
        Expression expr = getExpressionForTest(q, Arrays.asList("srcid"));
        SearchArgument sarg = orcStorage.getSearchArgument(expr);
        assertEquals("leaf-0 = (EQUALS srcid 10)\n" +
                "leaf-1 = (EQUALS srcid 11)\n" +
                "expr = (or leaf-0 leaf-1)", sarg.toString());
    }

    @Test
    public void testNegativeMatchesExpr() throws Exception {
        // matches operator is not a supported op type
        String q = query + "b = filter a by name matches 'foo*';" + "store b into 'out';";
        Expression expr = getExpressionForTest(q, Arrays.asList("name"));
        Assert.assertNull(expr);
        SearchArgument sarg = orcStorage.getSearchArgument(expr);
        Assert.assertNull(sarg);

        // AND in LHS/RHS
        q = query + "b = filter a by name matches 'foo*' and srcid == 10;" + "store b into 'out';";
        expr = getExpressionForTest(q, Arrays.asList("srcid", "name"));
        sarg = orcStorage.getSearchArgument(expr);
        assertEquals("leaf-0 = (EQUALS srcid 10)\n" +
                "expr = leaf-0", sarg.toString());

        q = query + "b = filter a by srcid == 10 and name matches 'foo*';" + "store b into 'out';";
        expr = getExpressionForTest(q, Arrays.asList("srcid", "name"));
        sarg = orcStorage.getSearchArgument(expr);
        assertEquals("leaf-0 = (EQUALS srcid 10)\n" +
                "expr = leaf-0", sarg.toString());

        // OR - Nothing should be pushed
        q = query + "b = filter a by name matches 'foo*' or srcid == 10;" + "store b into 'out';";
        expr = getExpressionForTest(q, Arrays.asList("srcid", "name"));
        Assert.assertNull(expr);
        sarg = orcStorage.getSearchArgument(expr);
        Assert.assertNull(sarg);
    }

    @Test
    public void testUnSupportedFields() throws Exception {
        //Struct, Map and Bag are not supported
        // TODO: Change the test to use ORCStorage to test OrcStorage.getPredicateFields()
        String q = query + "b = filter a by srcid == 10 and browser#'type' == 'IE';" +
                "store b into 'out';";
        Expression expr = getExpressionForTest(q, Arrays.asList("srcid"));
        SearchArgument sarg = orcStorage.getSearchArgument(expr);
        assertEquals("leaf-0 = (EQUALS srcid 10)\n" +
                "expr = leaf-0", sarg.toString());
    }

    @Test
    public void testPredicatePushdownBoolean() throws Exception {
        testPredicatePushdown(INPUT, "f1 == true", 2500, 1200000);
    }

    @Test
    public void testPredicatePushdownByteShort() throws Exception {
        testPredicatePushdown(INPUT, "f2 != 5 or f3 == 100", 3500, 1200000);
    }

    @Test
    public void testPredicatePushdownIntLongString() throws Exception {
        testPredicatePushdown(INPUT, "f4 >= 980 and f4 < 1010 and (f5 == 100 or f9 is not null)", 20, 1200000);
    }

    @Test
    public void testPredicatePushdownFloatDouble() throws Exception {
        testPredicatePushdown(INPUT, "f6 == 100.0 and f7 > 2000.00000001", 167, 1600000);
    }

    @Test
    public void testPredicatePushdownBigDecimal() throws Exception {
        testPredicatePushdown(INPUT, "f11 < (bigdecimal)'1000000000';", 2500, 1600000);
    }

    @Test
    public void testPredicatePushdownTimestamp() throws Exception {
        testPredicatePushdown(INPUT, "f10 >= ToDate('20100101', 'yyyyMMdd', 'UTC')", 3000, 400000);
    }

    private Expression getExpressionForTest(String query, List<String> predicateCols) throws Exception {
        LogicalPlan newLogicalPlan = Util.buildLp(pigServer, query);
        Operator op = newLogicalPlan.getSinks().get(0);
        LOFilter filter = (LOFilter) newLogicalPlan.getPredecessors(op).get(0);
        PredicatePushDownFilterExtractor filterExtractor = new PredicatePushDownFilterExtractor(filter.getFilterPlan(), predicateCols, supportedOpTypes);
        filterExtractor.visit();
        return filterExtractor.getPushDownExpression();
    }

    // For eclipse debugging
    private void testPredicatePushdownLocal(String filterStmt, int expectedRows) throws Exception {

        PigServer pigServer_disabledRule = new PigServer(Util.getLocalTestMode());
        // Test with PredicatePushdownOptimizer disabled.
        HashSet<String> disabledOptimizerRules = new HashSet<String>();
        disabledOptimizerRules.add("PredicatePushdownOptimizer");
        pigServer_disabledRule.getPigContext().getProperties().setProperty(PigImplConstants.PIG_OPTIMIZER_RULES_KEY,
                ObjectSerializer.serialize(disabledOptimizerRules));
        pigServer_disabledRule.registerQuery("B = load '" + INPUT + "' using OrcStorage();");
        pigServer_disabledRule.registerQuery("C = filter B by " + filterStmt + ";");

        // Test with PredicatePushdownOptimizer enabled.
        pigServer.registerQuery("D = load '" + INPUT + "' using OrcStorage();");
        pigServer.registerQuery("E = filter D by " + filterStmt + ";");

        //Verify that results are same
        Util.checkQueryOutputs(pigServer_disabledRule.openIterator("C"), pigServer.openIterator("E"), expectedRows);
    }

    private void testPredicatePushdown(String inputFile, String filterStmt, int expectedRows, int expectedBytesReadDiff) throws IOException {

        Util.resetStateForExecModeSwitch();
        // Minicluster is required to get hdfs bytes read counter value
        pigServer = new PigServer(cluster.getExecType(), cluster.getProperties());
        PigServer pigServer_disabledRule = new PigServer(cluster.getExecType(), cluster.getProperties());

        // Test with PredicatePushdownOptimizer disabled. All 3 blocks will be read
        HashSet<String> disabledOptimizerRules = new HashSet<String>();
        disabledOptimizerRules.add("PredicatePushdownOptimizer");
        pigServer_disabledRule.getPigContext().getProperties().setProperty(PigImplConstants.PIG_OPTIMIZER_RULES_KEY,
                ObjectSerializer.serialize(disabledOptimizerRules));
        pigServer_disabledRule.registerQuery("B = load '" + inputFile + "' using OrcStorage();");
        pigServer_disabledRule.registerQuery("C = filter B by " + filterStmt + ";");
        ExecJob job = pigServer_disabledRule.store("C", OUTPUT3);
        //Util.copyFromClusterToLocal(cluster, OUTPUT3 + "/part-m-00000", OUTPUT3);
        JobStats stats = (JobStats) job.getStatistics().getJobGraph().getSources().get(0);
        assertEquals(expectedRows, stats.getRecordWrittern());
        long bytesWithoutPushdown = stats.getHdfsBytesRead();

        // Test with PredicatePushdownOptimizer enabled. Only 2 blocks should be read
        pigServer.registerQuery("D = load '" + inputFile + "' using OrcStorage();");
        pigServer.registerQuery("E = filter D by " + filterStmt + ";");
        job = pigServer.store("E", OUTPUT4);
        //Util.copyFromClusterToLocal(cluster, OUTPUT4 + "/part-m-00000", OUTPUT4);
        stats = (JobStats) job.getStatistics().getJobGraph().getSources().get(0);
        assertEquals(expectedRows, stats.getRecordWrittern());
        long bytesWithPushdown = stats.getHdfsBytesRead();

        System.out.println("bytesWithoutPushdown was " + bytesWithoutPushdown +
                " and bytesWithPushdown was " + bytesWithPushdown);
        assertTrue("BytesWithoutPushdown was " + bytesWithoutPushdown +
                " and bytesWithPushdown was " + bytesWithPushdown,
                (bytesWithoutPushdown - bytesWithPushdown) > expectedBytesReadDiff);
        // Verify that results are same
        Util.checkQueryOutputs(pigServer_disabledRule.openIterator("C"), pigServer.openIterator("E"), expectedRows);
        pigServer_disabledRule.shutdown();

    }

    @Test
    public void testPredicatePushdownChar() throws Exception {
        testPredicatePushdown(basedir + "charvarchar.orc", "$0 == 'ulysses thompson'", 18, 18000);
    }

    @Test
    public void testPredicatePushdownVarchar() throws Exception {
        testPredicatePushdown(basedir + "charvarchar.orc", "$1 == 'alice allen         '", 19, 18000);
    }
}