add env switch

Author: dszakallas

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

@@ -123,12 +123,18 @@ object LdbcDatagen extends SparkApp {
 
   def run(args: Args): Unit = {
 
+    val irFormat = {
+      val _f = System.getenv("LDBC_DATAGEN_IR_FORMAT")
+      if (_f == null || _f == "") "parquet" else _f
+    }
+
     val generatorArgs = GenerationStage.Args(
       scaleFactor = args.scaleFactor,
       params = args.params,
       paramFile = args.paramFile,
       outputDir = args.outputDir,
       numThreads = args.numThreads,
+      format = irFormat,
       oversizeFactor = args.oversizeFactor
     )
 
@@ -155,6 +161,7 @@ object LdbcDatagen extends SparkApp {
         case "interactive" => Mode.Interactive(bulkLoadPortion = args.bulkloadPortion)
         case "raw"         => Mode.Raw
       },
+      irFormat,
       args.format,
       args.formatOptions
     )

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

@@ -13,13 +13,14 @@ object GenerationStage extends DatagenStage with Logging {
   val optimalPersonsPerFile = 500000
 
   case class Args(
-                   scaleFactor: String = "1",
-                   numThreads: Option[Int] = None,
-                   params: Map[String, String] = Map.empty,
-                   paramFile: Option[String] = None,
-                   outputDir: String = "out",
-                   oversizeFactor: Option[Double] = None
-                 )
+      scaleFactor: String = "1",
+      numThreads: Option[Int] = None,
+      params: Map[String, String] = Map.empty,
+      paramFile: Option[String] = None,
+      outputDir: String = "out",
+      format: String = "parquet",
+      oversizeFactor: Option[Double] = None
+  )
 
   def run(args: Args, config: GeneratorConfiguration)(implicit spark: SparkSession) = {
     val numPartitions   = config.getInt("hadoop.numThreads", spark.sparkContext.defaultParallelism)
@@ -44,9 +45,15 @@ object GenerationStage extends DatagenStage with Logging {
 
     val merged = SparkKnowsMerger(uniKnows, interestKnows, randomKnows).cache()
 
+    val format = args.format match {
+      case "csv"     => Csv
+      case "parquet" => Parquet
+      case a         => throw new IllegalArgumentException(s"Format `${a}` is not supported by the generator.")
+    }
+
     SparkUI.job(simpleNameOf[RawSerializer], "serialize persons") {
       val rawSerializer = new RawSerializer(randomRanker)
-      rawSerializer.write(merged, RawSink(Parquet, Some(numPartitions), config, oversizeFactor))
+      rawSerializer.write(merged, RawSink(format, Some(numPartitions), config, oversizeFactor))
     }
   }
 

src/main/scala/ldbc/snb/datagen/transformation/TransformationStage.scala

@@ -18,6 +18,7 @@ object TransformationStage extends DatagenStage with Logging {
       simulationStart: Long = 0,
       simulationEnd: Long = 0,
       mode: Mode = Mode.Raw,
+      irFormat: String = "parquet",
       format: String = "csv",
       formatOptions: Map[String, String] = Map.empty
   )
@@ -41,7 +42,7 @@ object TransformationStage extends DatagenStage with Logging {
       Graph[Mode.BI] :+:
       CNil
 
-    GraphSource(model.graphs.Raw.graphDef, args.outputDir, "parquet").read
+    GraphSource(model.graphs.Raw.graphDef, args.outputDir, args.irFormat).read
       .pipeFoldLeft(args.explodeAttrs.fork)((graph, _: Unit) => ExplodeAttrs.transform(graph))
       .pipeFoldLeft(args.explodeEdges.fork)((graph, _: Unit) => ExplodeEdges.transform(graph))
       .pipe(ConvertDates.transform)