Merge pull request #334 from ldbc/configurable-oversize-factor

Author: szarnyasg

src/main/scala/ldbc/snb/datagen/generation/GenerationStage.scala

@@ -18,14 +18,15 @@ object GenerationStage extends DatagenStage with Logging {
     numThreads: Option[Int] = None,
     params: Map[String, String] = Map.empty,
     paramFile: Option[String] = None,
-    outputDir: String = "out"
+    outputDir: String = "out",
+    oversizeFactor: Option[Double] = None
   )
 
-  def run(config: GeneratorConfiguration)(implicit spark: SparkSession) = {
+  def run(args: Args, config: GeneratorConfiguration)(implicit spark: SparkSession) = {
     val numPartitions = config.getInt("hadoop.numThreads", spark.sparkContext.defaultParallelism)
     val idealPartitions = DatagenParams.numPersons.toDouble / optimalPersonsPerFile
 
-    val oversizeFactor = Math.max(numPartitions / idealPartitions, 1.0)
+    val oversizeFactor = args.oversizeFactor.getOrElse(Math.max(numPartitions / idealPartitions, 1.0))
 
     val persons = SparkPersonGenerator(config)
 

src/main/scala/ldbc/snb/datagen/spark/LdbcDatagen.scala

@@ -27,7 +27,8 @@ object LdbcDatagen extends SparkApp {
     numThreads: Option[Int] = None,
     format: String = "csv",
     generateFactors: Boolean = false,
-    formatOptions: Map[String, String] = Map.empty
+    formatOptions: Map[String, String] = Map.empty,
+    oversizeFactor: Option[Double] = None
   )
 
   def main(args: Array[String]): Unit = {
@@ -65,6 +66,12 @@ object LdbcDatagen extends SparkApp {
         .action((x, c) => args.mode.set(c)(x))
         .text("Generation mode. Options: raw, bi, interactive. Default: raw")
 
+      opt[Double]("oversize-factor")
+        .action((x, c) => args.oversizeFactor.set(c)(Some(x)))
+        .text("Controls size of files relative to Persons. " +
+          "Values larger than 1 will result in less but larger files. " +
+          "Smaller values result in more, smaller files")
+
       opt[Double]("bulkload-portion")
         .action((x, c) => args.bulkloadPortion.set(c)(x))
         .text("Bulkload portion. Only applicable to BI and interactive modes")
@@ -113,14 +120,15 @@ object LdbcDatagen extends SparkApp {
       params = args.params,
       paramFile = args.paramFile,
       outputDir = args.outputDir,
-      numThreads = args.numThreads
+      numThreads = args.numThreads,
+      oversizeFactor = args.oversizeFactor
     )
 
     val generatorConfig = GenerationStage.buildConfig(generatorArgs)
 
     DatagenContext.initialize(generatorConfig)
 
-    GenerationStage.run(generatorConfig)
+    GenerationStage.run(generatorArgs, generatorConfig)
 
     if (args.generateFactors) {
       val factorArgs = FactorGenerationStage.Args()