smaller refactors

dszakallas · dszakallas · commit 3ed8b536eed6 · 2021-12-13T11:12:21.000+01:00
diff --git a/src/main/scala/ldbc/snb/datagen/LdbcDatagen.scala b/src/main/scala/ldbc/snb/datagen/LdbcDatagen.scala
@@ -8,6 +8,8 @@ import ldbc.snb.datagen.transformation.TransformationStage
 import ldbc.snb.datagen.util.{SparkApp, lower}
 import shapeless.lens
 
+import scala.collection.JavaConverters._
+
 object LdbcDatagen extends SparkApp {
   val appName = "LDBC SNB Datagen for Spark"
 
@@ -123,10 +125,8 @@ 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 env      = System.getenv().asScala
+    val irFormat = env.getOrElse("LDBC_DATAGEN_IR_FORMAT", "parquet")
 
     val generatorArgs = GenerationStage.Args(
       scaleFactor = args.scaleFactor,
diff --git a/src/main/scala/ldbc/snb/datagen/factors/FactorGenerationStage.scala b/src/main/scala/ldbc/snb/datagen/factors/FactorGenerationStage.scala
@@ -2,8 +2,8 @@ package ldbc.snb.datagen.factors
 
 import ldbc.snb.datagen.factors.io.FactorTableSink
 import ldbc.snb.datagen.io.graphs.GraphSource
-import ldbc.snb.datagen.model.EntityType
 import ldbc.snb.datagen.model
+import ldbc.snb.datagen.model.EntityType
 import ldbc.snb.datagen.syntax._
 import ldbc.snb.datagen.util.{DatagenStage, Logging}
 import org.apache.spark.sql.functions.{broadcast, count, date_trunc, sum}
diff --git a/src/main/scala/ldbc/snb/datagen/factors/io/package.scala b/src/main/scala/ldbc/snb/datagen/factors/io/package.scala
@@ -1,16 +1,16 @@
 package ldbc.snb.datagen.factors
 
+import ldbc.snb.datagen.io.dataframes.DataFrameSink
 import ldbc.snb.datagen.io.{PathComponent, Writer}
-import ldbc.snb.datagen.util.Logging
 import ldbc.snb.datagen.model.{GraphLike, Mode}
 import ldbc.snb.datagen.syntax._
-import ldbc.snb.datagen.io.dataframes.DataFrameSink
+import ldbc.snb.datagen.util.Logging
 
 package object io {
   case class FactorTableSink(path: String, format: String = "csv")
 
-  import ldbc.snb.datagen.io.dataframes.instances._
   import ldbc.snb.datagen.io.Writer.ops._
+  import ldbc.snb.datagen.io.dataframes.instances._
 
   private final class FactorTableWriter[M <: Mode] extends Writer[FactorTableSink] with Logging {
     override type Data = FactorTable[M]
diff --git a/src/main/scala/ldbc/snb/datagen/generator/GenerationStage.scala b/src/main/scala/ldbc/snb/datagen/generator/GenerationStage.scala
@@ -4,7 +4,7 @@ import ldbc.snb.datagen.generator.generators.{SparkKnowsGenerator, SparkKnowsMer
 import ldbc.snb.datagen.generator.serializers.RawSerializer
 import ldbc.snb.datagen.io.raw.{Csv, Parquet, RawSink}
 import ldbc.snb.datagen.syntax._
-import ldbc.snb.datagen.util.{ConfigParser, DatagenStage, GeneratorConfiguration, Logging, SparkUI, simpleNameOf}
+import ldbc.snb.datagen.util._
 import org.apache.hadoop.fs.{FileSystem, Path}
 import org.apache.spark.sql.SparkSession
 
diff --git a/src/main/scala/ldbc/snb/datagen/generator/generators/SparkKnowsGenerator.scala b/src/main/scala/ldbc/snb/datagen/generator/generators/SparkKnowsGenerator.scala
@@ -1,17 +1,15 @@
 package ldbc.snb.datagen.generator.generators
 
 import ldbc.snb.datagen.entities.dynamic.person.Person
-
-import java.util
-import ldbc.snb.datagen.generator.{DatagenContext, DatagenParams}
 import ldbc.snb.datagen.generator.generators.knowsgenerators.KnowsGenerator
+import ldbc.snb.datagen.generator.{DatagenContext, DatagenParams}
 import ldbc.snb.datagen.util.GeneratorConfiguration
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.SparkSession
 
+import java.util
 import scala.collection.JavaConverters._
 import scala.collection.SortedMap
-import scala.reflect.ClassTag
 
 object SparkKnowsGenerator {
   def apply(
diff --git a/src/main/scala/ldbc/snb/datagen/generator/generators/SparkPersonGenerator.scala b/src/main/scala/ldbc/snb/datagen/generator/generators/SparkPersonGenerator.scala
@@ -1,11 +1,10 @@
 package ldbc.snb.datagen.generator.generators
 
 import ldbc.snb.datagen.entities.dynamic.person.Person
-import org.apache.spark.sql.SparkSession
 import ldbc.snb.datagen.generator.{DatagenContext, DatagenParams}
-import ldbc.snb.datagen.generator.generators.PersonGenerator
 import ldbc.snb.datagen.util.GeneratorConfiguration
 import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.SparkSession
 
 import scala.collection.JavaConverters._
 
diff --git a/src/main/scala/ldbc/snb/datagen/generator/serializers/ActivityOutputStream.scala b/src/main/scala/ldbc/snb/datagen/generator/serializers/ActivityOutputStream.scala
@@ -1,9 +1,9 @@
 package ldbc.snb.datagen.generator.serializers
 
-import ldbc.snb.datagen.generator.dictionary.Dictionaries
 import ldbc.snb.datagen.entities.dynamic.Forum
 import ldbc.snb.datagen.entities.dynamic.messages.{Comment, Photo, Post}
 import ldbc.snb.datagen.entities.dynamic.relations.{ForumMembership, Like}
+import ldbc.snb.datagen.generator.dictionary.Dictionaries
 import ldbc.snb.datagen.generator.generators.{GenActivity, GenWall}
 import ldbc.snb.datagen.io.raw.RecordOutputStream
 import ldbc.snb.datagen.model.raw
diff --git a/src/main/scala/ldbc/snb/datagen/generator/serializers/PersonOutputStream.scala b/src/main/scala/ldbc/snb/datagen/generator/serializers/PersonOutputStream.scala
@@ -1,7 +1,7 @@
 package ldbc.snb.datagen.generator.serializers
 
-import ldbc.snb.datagen.generator.dictionary.Dictionaries
 import ldbc.snb.datagen.entities.dynamic.person.Person
+import ldbc.snb.datagen.generator.dictionary.Dictionaries
 import ldbc.snb.datagen.io.raw.RecordOutputStream
 import ldbc.snb.datagen.model.raw
 
diff --git a/src/main/scala/ldbc/snb/datagen/generator/serializers/RawSerializer.scala b/src/main/scala/ldbc/snb/datagen/generator/serializers/RawSerializer.scala
@@ -1,42 +1,25 @@
 package ldbc.snb.datagen.generator.serializers
 
 import ldbc.snb.datagen.entities.dynamic.person.Person
+import ldbc.snb.datagen.generator.generators.{GenActivity, PersonActivityGenerator, SparkRanker}
 import ldbc.snb.datagen.generator.{DatagenContext, DatagenParams}
-import ldbc.snb.datagen.generator.generators.SparkRanker
-import ldbc.snb.datagen.generator.generators.{GenActivity, PersonActivityGenerator}
 import ldbc.snb.datagen.io.Writer
 import ldbc.snb.datagen.io.raw.csv.CsvRowEncoder
 import ldbc.snb.datagen.io.raw.parquet.ParquetRowEncoder
 import ldbc.snb.datagen.io.raw.{RawSink, WriteContext, createNewWriteContext, recordOutputStream}
-import ldbc.snb.datagen.model.raw.{
-  Comment,
-  CommentHasTag,
-  Forum,
-  ForumHasMember,
-  ForumHasTag,
-  Organisation,
-  PersonLikesComment,
-  PersonLikesPost,
-  Place,
-  Post,
-  PostHasTag,
-  Tag,
-  TagClass
-}
+import ldbc.snb.datagen.model.raw._
 import ldbc.snb.datagen.model.{EntityTraits, raw}
 import ldbc.snb.datagen.syntax._
 import ldbc.snb.datagen.util.SerializableConfiguration
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.{FileSystem, Path}
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.SparkSession
-import org.apache.spark.sql.internal.SQLConf
 
 import java.net.URI
 import java.util
 import java.util.Collections
 import java.util.function.Consumer
-import scala.collection.JavaConverters._
 
 class RawSerializer(ranker: SparkRanker)(implicit spark: SparkSession) extends Writer[RawSink] {
   override type Data = RDD[Person]
diff --git a/src/main/scala/ldbc/snb/datagen/generator/serializers/StaticOutputStream.scala b/src/main/scala/ldbc/snb/datagen/generator/serializers/StaticOutputStream.scala
@@ -1,12 +1,12 @@
 package ldbc.snb.datagen.generator.serializers
 
-import ldbc.snb.datagen.generator.dictionary.Dictionaries
-import ldbc.snb.datagen.entities.statictype.place.Place
 import ldbc.snb.datagen.entities.statictype.Organisation
+import ldbc.snb.datagen.entities.statictype.place.Place
+import ldbc.snb.datagen.generator.dictionary.Dictionaries
+import ldbc.snb.datagen.generator.vocabulary.{DBP, DBPOWL}
 import ldbc.snb.datagen.io.raw.RecordOutputStream
 import ldbc.snb.datagen.model.raw
 import ldbc.snb.datagen.util.StringUtils
-import ldbc.snb.datagen.generator.vocabulary.{DBP, DBPOWL}
 
 import scala.collection.JavaConverters._
 import scala.collection.mutable
diff --git a/src/main/scala/ldbc/snb/datagen/io/raw/package.scala b/src/main/scala/ldbc/snb/datagen/io/raw/package.scala
@@ -3,13 +3,13 @@ package ldbc.snb.datagen.io
 import ldbc.snb.datagen.io.raw.csv.{CsvRecordOutputStream, CsvRowEncoder}
 import ldbc.snb.datagen.io.raw.parquet.{ParquetRecordOutputStream, ParquetRowEncoder}
 import ldbc.snb.datagen.model.EntityTraits
-import ldbc.snb.datagen.syntax.fluentSyntaxOps
 import ldbc.snb.datagen.util.GeneratorConfiguration
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.{FSDataOutputStream, FileSystem, Path}
 import org.apache.hadoop.mapreduce.task.TaskAttemptContextImpl
-import org.apache.hadoop.mapreduce.{JobID, TaskAttemptContext, TaskAttemptID, TaskID, TaskType}
+import org.apache.hadoop.mapreduce._
 import org.apache.parquet.hadoop.ParquetOutputFormat
+import org.apache.parquet.hadoop.metadata.CompressionCodecName
 import org.apache.spark.TaskContext
 import org.apache.spark.sql.Encoder
 
@@ -51,6 +51,10 @@ package object raw {
       val fileSystem: FileSystem
   )
 
+  // See org.apache.parquet.hadoop.metadata.CompressionCodecName for other options. Note that not all are
+  // supported by Spark.
+  val DefaultParquetCompression = "snappy"
+
   private def fileOutputStream(fs: FileSystem, path: Path): FSDataOutputStream = fs.create(path, true, 131072)
 
   def csvRecordOutputStream[T <: Product: CsvRowEncoder](path: Path, writeContext: WriteContext) = {
@@ -61,30 +65,27 @@ package object raw {
     new ParquetRecordOutputStream[T](path, writeContext.taskAttemptContext, options)
   }
 
-  private def getExtension(format: String, compression: Option[String]) = format.pipeFoldLeft(compression) { (f, c) => s"$c.$f" }
-
-    def recordOutputStream[T <: Product: EntityTraits: CsvRowEncoder: ParquetRowEncoder](sink: RawSink, writeContext: WriteContext): RecordOutputStream[T] = {
+  def recordOutputStream[T <: Product: EntityTraits: CsvRowEncoder: ParquetRowEncoder](sink: RawSink, writeContext: WriteContext): RecordOutputStream[T] = {
     val et                           = EntityTraits[T]
     val numFiles                     = Math.ceil(et.sizeFactor / sink.oversizeFactor).toLong
     val partitionId                  = writeContext.taskContext.partitionId()
     implicit val encoder: Encoder[T] = ParquetRowEncoder[T].encoder
     val entityPath                   = et.`type`.entityPath
     val (formatOutputStream, extension) = sink.format match {
       case Parquet => {
-        val compression = Some("snappy")
-        val options = Map.empty[String, String]
-          .pipeFoldLeft(compression) { (m, v) => m + (ParquetOutputFormat.COMPRESSION -> v)}
+        val compression = CompressionCodecName.fromConf(DefaultParquetCompression)
+        val options     = Map { ParquetOutputFormat.COMPRESSION -> compression.toString }
         (
           parquetRecordOutputStream(_: Path, writeContext, options),
-          getExtension(sink.format.toString, compression)
+          s"${compression.getExtension}.${sink.format}"
         )
       }
-      case Csv     => (csvRecordOutputStream(_: Path, writeContext), sink.format.toString)
-      case x       => throw new UnsupportedOperationException(s"Raw serializer not implemented for format ${x}")
+      case Csv => (csvRecordOutputStream(_: Path, writeContext), sink.format.toString)
+      case x   => throw new UnsupportedOperationException(s"Raw serializer not implemented for format ${x}")
     }
     val files =
       for { i <- 0L until numFiles } yield new Path(
-        s"${sink.conf.getOutputDir}/graphs/${sink.format.toString}/raw/composite-merged-fk/${entityPath}/part_${partitionId}_${i}.${extension}"
+        s"${sink.conf.getOutputDir}/graphs/${sink.format.toString}/raw/composite-merged-fk/${entityPath}/part_${partitionId}_${i}${extension}"
       )
     new RoundRobinOutputStream(files.map(formatOutputStream))
   }
diff --git a/src/main/scala/ldbc/snb/datagen/io/raw/parquet.scala b/src/main/scala/ldbc/snb/datagen/io/raw/parquet.scala
@@ -1,17 +1,14 @@
 package ldbc.snb.datagen.io.raw
 
-import ldbc.snb.datagen.syntax.fluentSyntaxOps
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.Path
 import org.apache.hadoop.mapreduce.{RecordWriter, TaskAttemptContext}
 import org.apache.parquet.hadoop.ParquetOutputFormat
 import org.apache.parquet.hadoop.api.WriteSupport
 import org.apache.parquet.io.api.RecordConsumer
-import org.apache.spark.sql.{Encoder, SQLContext}
+import org.apache.spark.sql.Encoder
 import org.apache.spark.sql.catalyst.encoders.encoderFor
-import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap
-import org.apache.spark.sql.execution.datasources.parquet.{LdbcDatagenParquetWriteSupport, ParquetOptions}
-import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.execution.datasources.parquet.LdbcDatagenParquetWriteSupport
 
 object parquet {
 
@@ -39,14 +36,12 @@ object parquet {
   }
 
   class ParquetRecordOutputStream[T <: Product: ParquetRowEncoder](
-                                                                    path: Path,
-                                                                    taskAttemptContext: TaskAttemptContext,
-                                                                    options: Map[String, String]
-                                                                  ) extends RecordOutputStream[T] {
+      path: Path,
+      taskAttemptContext: TaskAttemptContext,
+      options: Map[String, String]
+  ) extends RecordOutputStream[T] {
     lazy val writer: RecordWriter[Void, T] = {
-      val compressionCodecClassName = options
-        .get(ParquetOutputFormat.COMPRESSION)
-        .map(ParquetOptions.getParquetCompressionCodecName)
+      val compressionCodecClassName = options.get(ParquetOutputFormat.COMPRESSION)
 
       implicit val encoder: Encoder[T] = implicitly[ParquetRowEncoder[T]].encoder
       val parquetWriteSupport          = parquetWriteSupportForEncodable[T](compressionCodecClassName)
diff --git a/src/main/scala/ldbc/snb/datagen/model/graphs.scala b/src/main/scala/ldbc/snb/datagen/model/graphs.scala
@@ -3,10 +3,11 @@ package ldbc.snb.datagen.model
 import ldbc.snb.datagen.model.raw.RawEntity
 
 object graphs {
+  import instances._
+
   object Raw {
 
-    import instances._
-    import ldbc.snb.datagen.model.raw.instances._
+    import raw.instances._
 
     val graphDef = GraphDef(
       isAttrExploded = false,
diff --git a/src/main/scala/ldbc/snb/datagen/model/package.scala b/src/main/scala/ldbc/snb/datagen/model/package.scala
@@ -169,6 +169,5 @@ package object model {
     case object Day   extends BatchPeriod
     case object Month extends BatchPeriod
   }
-
   object instances extends UntypedEntitiesInstances
 }
diff --git a/src/main/scala/ldbc/snb/datagen/model/raw.scala b/src/main/scala/ldbc/snb/datagen/model/raw.scala
@@ -18,7 +18,7 @@ object raw {
       `birthday`: Long,
       `locationIP`: String,
       `browserUsed`: String,
-      `LocationCityId`: Integer,
+      `LocationCityId`: Int,
       `language`: String,
       `email`: String
   ) extends RawEntity
diff --git a/src/main/scala/ldbc/snb/datagen/syntax/UseSyntax.scala b/src/main/scala/ldbc/snb/datagen/syntax/UseSyntax.scala
@@ -1,7 +1,5 @@
 package ldbc.snb.datagen.syntax
 
-import shapeless._
-
 import java.io.Closeable
 
 trait Usable[A] {
@@ -45,22 +43,6 @@ trait UsableInstances {
       }
     }
   }
-
-  implicit def usableForProduct[A <: Product, Repr](implicit gen: Generic.Aux[A, Repr], ev: Usable[Repr]): Usable[A] = new Usable[A] {
-    override def use[Ret](a: A, f: A => Ret): Ret = {
-      ev.use(gen.to(a), f.compose(gen.from))
-    }
-  }
-
-  implicit def usableForHCons[A, Rest <: HList](implicit ev1: Usable[A], ev2: Usable[Rest]): Usable[A :: Rest] = new Usable[A :: Rest] {
-    override def use[Ret](list: A :: Rest, f: (A :: Rest) => Ret): Ret = {
-      ev1.use(list.head, (a: A) => { ev2.use(list.tail, (rest: Rest) => f(a :: rest)) })
-    }
-  }
-
-  implicit val usableForHNil: Usable[HNil] = new Usable[HNil] {
-    override def use[Ret](a: HNil, f: HNil => Ret): Ret = f(a)
-  }
 }
 
 object UsableInstances extends UsableInstances
diff --git a/src/main/scala/ldbc/snb/datagen/transformation/TransformationStage.scala b/src/main/scala/ldbc/snb/datagen/transformation/TransformationStage.scala
@@ -4,7 +4,7 @@ import ldbc.snb.datagen.io.graphs.{GraphSink, GraphSource}
 import ldbc.snb.datagen.model
 import ldbc.snb.datagen.model.{BatchedEntity, Graph, Mode}
 import ldbc.snb.datagen.syntax._
-import ldbc.snb.datagen.transformation.transform.{ExplodeAttrs, ExplodeEdges, ConvertDates, RawToBiTransform, RawToInteractiveTransform}
+import ldbc.snb.datagen.transformation.transform._
 import ldbc.snb.datagen.util.{DatagenStage, Logging}
 import org.apache.spark.sql.{DataFrame, SparkSession}
 import shapeless._
@@ -23,12 +23,11 @@ object TransformationStage extends DatagenStage with Logging {
       formatOptions: Map[String, String] = Map.empty
   )
 
-  import ldbc.snb.datagen.io.instances._
-  import ldbc.snb.datagen.io.Writer.ops._
   import ldbc.snb.datagen.io.Reader.ops._
+  import ldbc.snb.datagen.io.Writer.ops._
+  import ldbc.snb.datagen.io.instances._
 
   def run(args: Args)(implicit spark: SparkSession) = {
-    import spark.implicits._
     object write extends Poly1 {
       implicit def caseSimple[M <: Mode](implicit ev: M#Layout =:= DataFrame) =
         at[Graph[M]](g => g.write(GraphSink(args.outputDir, args.format, args.formatOptions)))
diff --git a/src/main/scala/ldbc/snb/datagen/transformation/transform/ExplodeAttrs.scala b/src/main/scala/ldbc/snb/datagen/transformation/transform/ExplodeAttrs.scala
@@ -1,11 +1,11 @@
 package ldbc.snb.datagen.transformation.transform
 
-import ldbc.snb.datagen.syntax._
 import ldbc.snb.datagen.model.EntityType.{Attr, Node}
-import ldbc.snb.datagen.model.Mode.Raw.withRawColumns
 import ldbc.snb.datagen.model.Mode
-import org.apache.spark.sql.{Column, DataFrame}
+import ldbc.snb.datagen.model.Mode.Raw.withRawColumns
+import ldbc.snb.datagen.syntax._
 import org.apache.spark.sql.functions.{explode, split}
+import org.apache.spark.sql.{Column, DataFrame}
 import shapeless.lens
 
 object ExplodeAttrs extends Transform[Mode.Raw.type, Mode.Raw.type] {
diff --git a/src/main/scala/ldbc/snb/datagen/transformation/transform/ExplodeEdges.scala b/src/main/scala/ldbc/snb/datagen/transformation/transform/ExplodeEdges.scala
@@ -1,12 +1,12 @@
 package ldbc.snb.datagen.transformation.transform
 
-import ldbc.snb.datagen.syntax._
-import ldbc.snb.datagen.model.Mode
 import ldbc.snb.datagen.model.Cardinality.{NN, NOne, OneN}
 import ldbc.snb.datagen.model.EntityType.{Edge, Node}
+import ldbc.snb.datagen.model.Mode
 import ldbc.snb.datagen.model.Mode.Raw.withRawColumns
-import org.apache.spark.sql.{Column, DataFrame}
+import ldbc.snb.datagen.syntax._
 import org.apache.spark.sql.functions._
+import org.apache.spark.sql.{Column, DataFrame}
 import shapeless.lens
 
 object ExplodeEdges extends Transform[Mode.Raw.type, Mode.Raw.type] {
diff --git a/src/main/scala/ldbc/snb/datagen/transformation/transform/RawToBiTransform.scala b/src/main/scala/ldbc/snb/datagen/transformation/transform/RawToBiTransform.scala
diff --git a/src/main/scala/ldbc/snb/datagen/transformation/transform/RawToInteractiveTransform.scala b/src/main/scala/ldbc/snb/datagen/transformation/transform/RawToInteractiveTransform.scala
diff --git a/src/main/scala/ldbc/snb/datagen/transformation/transform/Transform.scala b/src/main/scala/ldbc/snb/datagen/transformation/transform/Transform.scala
diff --git a/src/main/scala/ldbc/snb/datagen/util/SerializableConfiguration.scala b/src/main/scala/ldbc/snb/datagen/util/SerializableConfiguration.scala
diff --git a/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/LdbcDatagenParquetWriteSupport.scala b/src/main/scala/org/apache/spark/sql/execution/datasources/parquet/LdbcDatagenParquetWriteSupport.scala
