package org.apache.hadoop.hive.ql.exec.spark;

import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.hive.conf.HiveConf;
import org.apache.hadoop.hive.ql.DriverFactory;
import org.apache.hadoop.hive.ql.IDriver;
import org.apache.hadoop.hive.ql.exec.Utilities;
import org.apache.hadoop.hive.ql.io.HiveKey;
import org.apache.hadoop.hive.ql.security.authorization.plugin.sqlstd.SQLStdHiveAuthorizerFactory;
import org.apache.hadoop.hive.ql.session.SessionState;

import org.apache.hadoop.io.BytesWritable;
import org.apache.hadoop.mapred.JobConf;

import org.apache.spark.Dependency;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaSparkContext;

import org.apache.spark.rdd.HadoopRDD;
import org.apache.spark.rdd.MapPartitionsRDD;
import org.apache.spark.rdd.RDD;
import org.apache.spark.rdd.ShuffledRDD;

import org.junit.Assert;
import org.junit.Test;

import scala.Tuple2;
import scala.collection.JavaConversions;

import java.util.List;


public class TestSparkPlan {

  @Test
  public void testSetRDDCallSite() throws Exception {
    HiveConf conf = new HiveConf();
    conf.setVar(HiveConf.ConfVars.HIVE_AUTHORIZATION_MANAGER,
            SQLStdHiveAuthorizerFactory.class.getName());
    conf.setBoolVar(HiveConf.ConfVars.HIVE_SUPPORT_CONCURRENCY, false);
    conf.setVar(HiveConf.ConfVars.HIVE_EXECUTION_ENGINE, "spark");

    FileSystem fs = FileSystem.getLocal(conf);
    Path tmpDir = new Path("TestSparkPlan-tmp");

    SessionState.start(conf);

    IDriver driver = null;
    JavaSparkContext sc = null;

    try {
      driver = DriverFactory.newDriver(conf);
      Assert.assertEquals(0, driver.run("create table test (col int)").getResponseCode());

      driver.compile("select * from test order by col");
      List<SparkTask> sparkTasks = Utilities.getSparkTasks(driver.getPlan().getRootTasks());
      Assert.assertEquals(1, sparkTasks.size());

      SparkTask sparkTask = sparkTasks.get(0);

      JobConf jobConf = new JobConf(conf);

      SparkConf sparkConf = new SparkConf();
      sparkConf.setMaster("local");
      sparkConf.setAppName("TestSparkPlan-app");
      sc = new JavaSparkContext(sparkConf);

      SparkPlanGenerator sparkPlanGenerator = new SparkPlanGenerator(sc, null, jobConf, tmpDir,
              null);
      SparkPlan sparkPlan = sparkPlanGenerator.generate(sparkTask.getWork());
      RDD<Tuple2<HiveKey, BytesWritable>> reducerRdd = sparkPlan.generateGraph().rdd();

      Assert.assertTrue(reducerRdd.name().contains("Reducer 2"));
      Assert.assertTrue(reducerRdd instanceof MapPartitionsRDD);
      Assert.assertTrue(reducerRdd.creationSite().shortForm().contains("Reducer 2"));
      Assert.assertTrue(reducerRdd.creationSite().longForm().contains("Explain Plan"));
      Assert.assertTrue(reducerRdd.creationSite().longForm().contains("Reducer 2"));

      List<Dependency<?>> rdds = JavaConversions.seqAsJavaList(reducerRdd.dependencies());
      Assert.assertEquals(1, rdds.size());
      RDD shuffledRdd = rdds.get(0).rdd();

      Assert.assertTrue(shuffledRdd.name().contains("Reducer 2"));
      Assert.assertTrue(shuffledRdd.name().contains("SORT"));
      Assert.assertTrue(shuffledRdd instanceof ShuffledRDD);
      Assert.assertTrue(shuffledRdd.creationSite().shortForm().contains("Reducer 2"));
      Assert.assertTrue(shuffledRdd.creationSite().longForm().contains("Explain Plan"));
      Assert.assertTrue(shuffledRdd.creationSite().longForm().contains("Reducer 2"));

      rdds = JavaConversions.seqAsJavaList(shuffledRdd.dependencies());
      Assert.assertEquals(1, rdds.size());
      RDD mapRdd = rdds.get(0).rdd();

      Assert.assertTrue(mapRdd.name().contains("Map 1"));
      Assert.assertTrue(mapRdd instanceof MapPartitionsRDD);
      Assert.assertTrue(mapRdd.creationSite().shortForm().contains("Map 1"));
      Assert.assertTrue(mapRdd.creationSite().longForm().contains("Explain Plan"));
      Assert.assertTrue(mapRdd.creationSite().longForm().contains("Map 1"));

      rdds = JavaConversions.seqAsJavaList(mapRdd.dependencies());
      Assert.assertEquals(1, rdds.size());
      RDD hadoopRdd = rdds.get(0).rdd();

      Assert.assertTrue(hadoopRdd.name().contains("Map 1"));
      Assert.assertTrue(hadoopRdd.name().contains("test"));
      Assert.assertTrue(hadoopRdd instanceof HadoopRDD);
      Assert.assertTrue(hadoopRdd.creationSite().shortForm().contains("Map 1"));
    } finally {
      if (driver != null) {
        Assert.assertEquals(0, driver.run("drop table if exists test").getResponseCode());
        driver.destroy();
      }
      if (sc != null) {
        sc.close();
      }
      if (fs.exists(tmpDir)) {
        fs.delete(tmpDir, true);
      }
    }
  }
}

@@ -1,8 +1,815 @@
+/*%0A * Licensed to the Apache Software Foundation (ASF) under one%0A * or more contributor license agreements.  See the NOTICE file%0A * distributed with this work for additional information%0A * regarding copyright ownership.  The ASF licenses this file%0A * to you under the Apache License, Version 2.0 (the%0A * %22License%22); you may not use this file except in compliance%0A * with the License.  You may obtain a copy of the License at%0A *%0A *     http://www.apache.org/licenses/LICENSE-2.0%0A *%0A * Unless required by applicable law or agreed to in writing, software%0A * distributed under the License is distributed on an %22AS IS%22 BASIS,%0A * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.%0A * See the License for the specific language governing permissions and%0A * limitations under the License.%0A */%0A%0A
 package 

Add Apache license to TestSparkPlan.java (Sahil Takiar, reviewed by Vihang Karajgaonkar)