Merge pull request #416 from ldbc/factor-format

Add factor table format option to executables

Author: szarnyasg

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

@@ -26,10 +26,11 @@ object LdbcDatagen extends SparkApp {
       batchPeriod: String = "day",
       numThreads: Option[Int] = None,
       format: String = "csv",
-      generateFactors: Boolean = false,
       formatOptions: Map[String, String] = Map.empty,
       oversizeFactor: Option[Double] = None,
-      epochMillis: Boolean = false
+      epochMillis: Boolean = false,
+      generateFactors: Boolean = false,
+      factorFormat: String = "parquet"
   )
 
   override type ArgsType = Args
@@ -118,6 +119,10 @@ object LdbcDatagen extends SparkApp {
         .action((x, c) => args.generateFactors.set(c)(true))
         .text("Generate factor tables")
 
+      opt[String]("factor-format")
+        .action((x, c) => args.factorFormat.set(c)(x))
+        .text("Output format of factor tables")
+
       help('h', "help").text("prints this usage text")
 
       opt[Unit]("epoch-millis")
@@ -146,7 +151,11 @@ object LdbcDatagen extends SparkApp {
     GenerationStage.run(generatorArgs)
 
     if (args.generateFactors) {
-      val factorArgs = FactorGenerationStage.Args(outputDir = args.outputDir, irFormat = irFormat)
+      val factorArgs = FactorGenerationStage.Args(
+        outputDir = args.outputDir,
+        irFormat = irFormat,
+        format = args.factorFormat
+      )
       FactorGenerationStage.run(factorArgs)
     }
 

src/main/scala/ldbc/snb/datagen/factors/FactorGenerationStage.scala

@@ -31,6 +31,7 @@ object FactorGenerationStage extends DatagenStage with Logging {
   case class Args(
       outputDir: String = "out",
       irFormat: String = "parquet",
+      format: String = "parquet",
       only: Option[Regex] = None,
       force: Boolean = false
   )
@@ -54,6 +55,10 @@ object FactorGenerationStage extends DatagenStage with Logging {
         .action((x, c) => args.irFormat.set(c)(x))
         .text("Format of the raw input")
 
+      opt[String]("format")
+        .action((x, c) => args.format.set(c)(x))
+        .text("Output format")
+
       opt[String]("only")
         .action((x, c) => args.only.set(c)(Some(x.r.anchored)))
         .text("Only generate factor tables whose name matches the supplied regex")
@@ -87,7 +92,7 @@ object FactorGenerationStage extends DatagenStage with Logging {
               FactorTable(name, calc(resolvedEntities), g)
           }
       )
-      .foreach(_.write(FactorTableSink(args.outputDir, overwrite = args.force)))
+      .foreach(_.write(FactorTableSink(args.outputDir, format = args.format, overwrite = args.force)))
   }
 
   private def frequency(df: DataFrame, value: Column, by: Seq[Column], agg: Column => Column = count) =

src/main/scala/ldbc/snb/datagen/factors/io/package.scala

@@ -8,7 +8,7 @@ import ldbc.snb.datagen.util.Logging
 import org.apache.spark.sql.SaveMode
 
 package object io {
-  case class FactorTableSink(path: String, format: String = "csv", overwrite: Boolean = false)
+  case class FactorTableSink(path: String, format: String = "parquet", overwrite: Boolean = false)
 
   import ldbc.snb.datagen.io.Writer.ops._
   import ldbc.snb.datagen.io.dataframes.instances._