Merge branch 'tpch-benchmark' of github.com:octaviansima/opaque into tpch-benchmark

Author: octaviansima

src/main/scala/edu/berkeley/cs/rise/opaque/benchmark/Benchmark.scala

@@ -43,9 +43,7 @@ object Benchmark {
   val spark = SparkSession.builder()
       .appName("Benchmark")
       .getOrCreate()
-  var numPartitions = 2 * spark.sparkContext
-      .getConf
-      .getInt("spark.executor.instances", 2)
+  var numPartitions = spark.sparkContext.defaultParallelism
   var size = "sf_small"
 
   def dataDir: String = {

src/main/scala/edu/berkeley/cs/rise/opaque/benchmark/TPCH.scala

@@ -17,6 +17,7 @@
 
 package edu.berkeley.cs.rise.opaque.benchmark
 
+import java.io.File
 import scala.io.Source
 
 import org.apache.spark.sql.DataFrame
@@ -162,7 +163,7 @@ object TPCH {
       .option("delimiter", "|")
       .load(s"${Benchmark.dataDir}/tpch/$size/customer.tbl")
 
-  def generateMap(
+  def generateDFs(
       sqlContext: SQLContext, size: String)
       : Map[String, DataFrame] = {
     Map("part" -> part(sqlContext, size),
@@ -175,38 +176,70 @@ object TPCH {
     "customer" -> customer(sqlContext, size)
     ),
   }
-
-  def apply(sqlContext: SQLContext, size: String) : TPCH = {
-    val tpch = new TPCH(sqlContext, size)
-    tpch.tableNames = tableNames
-    tpch.nameToDF = generateMap(sqlContext, size)
-    tpch
-  }
 }
 
 class TPCH(val sqlContext: SQLContext, val size: String) {
 
-  var tableNames : Seq[String] = Seq()
-  var nameToDF : Map[String, DataFrame] = Map()
+  val tableNames = TPCH.tableNames
+  val nameToDF = TPCH.generateDFs(sqlContext, size)
 
-  def setupViews(securityLevel: SecurityLevel, numPartitions: Int) = {
-    for ((name, df) <- nameToDF) {
-      Utils.ensureCached(securityLevel.applyTo(df.repartition(numPartitions))).createOrReplaceTempView(name)
-    }
-  }
+  private var numPartitions: Int = -1
+  private var nameToPath = Map[String, File]()
+  private var nameToEncryptedPath = Map[String, File]()
 
   def getQuery(queryNumber: Int) : String = {
     val queryLocation = sys.env.getOrElse("OPAQUE_HOME", ".") + "/src/test/resources/tpch/"
     Source.fromFile(queryLocation + s"q$queryNumber.sql").getLines().mkString("\n")
   }
 
-  def performQuery(sqlContext: SQLContext, sqlStr: String) : DataFrame  = {
-    sqlContext.sparkSession.sql(sqlStr);
+  def generateFiles(numPartitions: Int) = {
+    if (numPartitions != this.numPartitions) {
+      this.numPartitions = numPartitions
+      for ((name, df) <- nameToDF) {
+        nameToPath.get(name).foreach{ path => Utils.deleteRecursively(path) }
+
+        nameToPath += (name -> createPath(df, Insecure, numPartitions))
+        nameToEncryptedPath += (name -> createPath(df, Encrypted, numPartitions))
+      }
+    }
+  }
+
+  private def createPath(df: DataFrame, securityLevel: SecurityLevel, numPartitions: Int): File = {
+    val partitionedDF = securityLevel.applyTo(df.repartition(numPartitions))
+    val path = Utils.createTempDir()
+    path.delete()
+    securityLevel match {
+      case Insecure => {
+        partitionedDF.write.format("com.databricks.spark.csv").save(path.toString)
+      }
+      case Encrypted => {
+        partitionedDF.write.format("edu.berkeley.cs.rise.opaque.EncryptedSource").save(path.toString)
+      }
+    }
+    path
+  }
+
+  private def loadViews(securityLevel: SecurityLevel) = {
+    val (map, formatStr) = if (securityLevel == Insecure) 
+        (nameToPath, "com.databricks.spark.csv") else 
+        (nameToEncryptedPath, "edu.berkeley.cs.rise.opaque.EncryptedSource")
+    for ((name, path) <- map) {
+        val df = sqlContext.sparkSession.read
+            .format(formatStr)
+            .schema(nameToDF.get(name).get.schema)
+            .load(path.toString)
+        df.createOrReplaceTempView(name)
+    }
+  }
+
+  def performQuery(sqlStr: String, securityLevel: SecurityLevel): DataFrame  = {
+    loadViews(securityLevel)
+    sqlContext.sparkSession.sql(sqlStr)
   }
 
-  def query(queryNumber: Int, securityLevel: SecurityLevel, sqlContext: SQLContext, numPartitions: Int) : DataFrame = {
-    setupViews(securityLevel, numPartitions)
+  def query(queryNumber: Int, securityLevel: SecurityLevel, numPartitions: Int): DataFrame = {
     val sqlStr = getQuery(queryNumber)
-    performQuery(sqlContext, sqlStr)
+    generateFiles(numPartitions)
+    performQuery(sqlStr, securityLevel)
   }
 }

src/main/scala/edu/berkeley/cs/rise/opaque/benchmark/TPCHBenchmark.scala

@@ -24,32 +24,27 @@ import org.apache.spark.sql.SQLContext
 object TPCHBenchmark {
   def query(queryNumber: Int, tpch: TPCH, sqlContext: SQLContext, numPartitions: Int) = {
     val sqlStr = tpch.getQuery(queryNumber)
+    tpch.generateFiles(numPartitions)
 
-    tpch.setupViews(Encrypted, numPartitions)
     Utils.timeBenchmark(
         "distributed" -> (numPartitions > 1),
         "query" -> s"TPC-H $queryNumber",
         "system" -> Encrypted.name) {
 
-      val df = tpch.performQuery(sqlContext, sqlStr)
-      Utils.force(df)
-      df
+      tpch.performQuery(sqlStr, Insecure).collect
     }
 
-    tpch.setupViews(Insecure, numPartitions)
     Utils.timeBenchmark(
         "distributed" -> (numPartitions > 1),
         "query" -> s"TPC-H $queryNumber",
         "system" -> Insecure.name) {
 
-      val df = tpch.performQuery(sqlContext, sqlStr)
-      Utils.force(df)
-      df
+      tpch.performQuery(sqlStr, Encrypted).collect
     }
   }
 
   def run(sqlContext: SQLContext, numPartitions: Int, size: String) = {
-    val tpch = TPCH(sqlContext, size)
+    val tpch = new TPCH(sqlContext, size)
 
     val supportedQueries = Seq(1, 3, 5, 6, 7, 8, 9, 10, 14, 17)
 

src/test/scala/edu/berkeley/cs/rise/opaque/TPCHTests.scala

@@ -26,94 +26,94 @@ import edu.berkeley.cs.rise.opaque.benchmark.TPCH
 trait TPCHTests extends OpaqueTestsBase { self => 
 
   def size = "sf_small"
-  def tpch = TPCH(spark.sqlContext, size)
+  def tpch = new TPCH(spark.sqlContext, size)
 
   testAgainstSpark("TPC-H 1") { securityLevel =>
-    tpch.query(1, securityLevel, spark.sqlContext, numPartitions).collect
+    tpch.query(1, securityLevel, numPartitions).collect
   }
 
   testAgainstSpark("TPC-H 2", ignore) { securityLevel =>
-    tpch.query(2, securityLevel, spark.sqlContext, numPartitions).collect
+    tpch.query(2, securityLevel, numPartitions).collect
   }
 
   testAgainstSpark("TPC-H 3") { securityLevel =>
-    tpch.query(3, securityLevel, spark.sqlContext, numPartitions).collect
+    tpch.query(3, securityLevel, numPartitions).collect
   }
 
   testAgainstSpark("TPC-H 4", ignore) { securityLevel =>
-    tpch.query(4, securityLevel, spark.sqlContext, numPartitions).collect
+    tpch.query(4, securityLevel, numPartitions).collect
   }
 
   testAgainstSpark("TPC-H 5") { securityLevel =>
-    tpch.query(5, securityLevel, spark.sqlContext, numPartitions).collect
+    tpch.query(5, securityLevel, numPartitions).collect
   }
 
   testAgainstSpark("TPC-H 6") { securityLevel =>
-    tpch.query(6, securityLevel, spark.sqlContext, numPartitions).collect.toSet
+    tpch.query(6, securityLevel, numPartitions).collect.toSet
   }
 
   testAgainstSpark("TPC-H 7") { securityLevel =>
-    tpch.query(7, securityLevel, spark.sqlContext, numPartitions).collect
+    tpch.query(7, securityLevel, numPartitions).collect
   }
 
   testAgainstSpark("TPC-H 8") { securityLevel =>
-    tpch.query(8, securityLevel, spark.sqlContext, numPartitions).collect
+    tpch.query(8, securityLevel, numPartitions).collect
   }
 
   testAgainstSpark("TPC-H 9") { securityLevel =>
-    tpch.query(9, securityLevel, spark.sqlContext, numPartitions).collect
+    tpch.query(9, securityLevel, numPartitions).collect
   }
 
   testAgainstSpark("TPC-H 10") { securityLevel =>
-    tpch.query(10, securityLevel, spark.sqlContext, numPartitions).collect
+    tpch.query(10, securityLevel, numPartitions).collect
   }
 
   testAgainstSpark("TPC-H 11", ignore) { securityLevel =>
-    tpch.query(11, securityLevel, spark.sqlContext, numPartitions).collect
+    tpch.query(11, securityLevel, numPartitions).collect
   }
 
   testAgainstSpark("TPC-H 12") { securityLevel =>
-    tpch.query(12, securityLevel, spark.sqlContext, numPartitions).collect
+    tpch.query(12, securityLevel, numPartitions).collect
   }
 
   testAgainstSpark("TPC-H 13", ignore) { securityLevel =>
-    tpch.query(13, securityLevel, spark.sqlContext, numPartitions).collect
+    tpch.query(13, securityLevel, numPartitions).collect
   }
 
   testAgainstSpark("TPC-H 14") { securityLevel =>
-    tpch.query(14, securityLevel, spark.sqlContext, numPartitions).collect.toSet
+    tpch.query(14, securityLevel, numPartitions).collect.toSet
   }
 
   testAgainstSpark("TPC-H 15", ignore) { securityLevel =>
-    tpch.query(15, securityLevel, spark.sqlContext, numPartitions).collect
+    tpch.query(15, securityLevel, numPartitions).collect
   }
 
   testAgainstSpark("TPC-H 16", ignore) { securityLevel =>
-    tpch.query(16, securityLevel, spark.sqlContext, numPartitions).collect
+    tpch.query(16, securityLevel, numPartitions).collect
   }
 
   testAgainstSpark("TPC-H 17") { securityLevel =>
-    tpch.query(17, securityLevel, spark.sqlContext, numPartitions).collect.toSet
+    tpch.query(17, securityLevel, numPartitions).collect.toSet
   }
 
   testAgainstSpark("TPC-H 18", ignore) { securityLevel =>
-    tpch.query(18, securityLevel, spark.sqlContext, numPartitions).collect
+    tpch.query(18, securityLevel, numPartitions).collect
   }
 
   testAgainstSpark("TPC-H 19") { securityLevel =>
-    tpch.query(19, securityLevel, spark.sqlContext, numPartitions).collect.toSet
+    tpch.query(19, securityLevel, numPartitions).collect.toSet
   }
 
   testAgainstSpark("TPC-H 20") { securityLevel =>
-    tpch.query(20, securityLevel, spark.sqlContext, numPartitions).collect.toSet
+    tpch.query(20, securityLevel, numPartitions).collect.toSet
   }
 
   testAgainstSpark("TPC-H 21", ignore) { securityLevel =>
-    tpch.query(21, securityLevel, spark.sqlContext, numPartitions).collect
+    tpch.query(21, securityLevel, numPartitions).collect
   }
 
   testAgainstSpark("TPC-H 22", ignore) { securityLevel =>
-    tpch.query(22, securityLevel, spark.sqlContext, numPartitions).collect
+    tpch.query(22, securityLevel, numPartitions).collect
   }
 }