Merge pull request #342 from ldbc/apply-scalafmt

Author: szarnyasg

.scalafmt.conf

@@ -0,0 +1,4 @@
+version = 3.1.2
+runner.dialect = scala212
+align.preset = more
+maxColumn = 160

project/plugins.sbt

@@ -1,2 +1,3 @@
 addSbtPlugin("com.typesafe.sbt" % "sbt-pom-reader" % "2.1.0")
 addSbtPlugin("com.eed3si9n" % "sbt-assembly" % "0.15.0")
+addSbtPlugin("org.scalameta" % "sbt-scalafmt" % "2.4.3")

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

@@ -9,8 +9,7 @@ import ldbc.snb.datagen.util.{DatagenStage, Logging}
 import org.apache.spark.sql.functions.{broadcast, count, date_trunc, sum}
 import org.apache.spark.sql.{Column, DataFrame, SparkSession}
 
-
-case class Factor(requiredEntities: EntityType*)(f: Seq[DataFrame] => DataFrame)  extends (Seq[DataFrame] => DataFrame) {
+case class Factor(requiredEntities: EntityType*)(f: Seq[DataFrame] => DataFrame) extends (Seq[DataFrame] => DataFrame) {
   override def apply(v1: Seq[DataFrame]): DataFrame = f(v1)
 }
 
@@ -24,22 +23,23 @@ object FactorGenerationStage extends DatagenStage with Logging {
     import ldbc.snb.datagen.io.Writer.ops._
     import ldbc.snb.datagen.io.instances._
 
-    GraphSource(model.graphs.Raw.graphDef, args.outputDir, "csv")
-      .read
-      .pipe(g => rawFactors.map { case (name, calc) =>
-        val resolvedEntities = calc.requiredEntities.foldLeft(Seq.empty[DataFrame])((args, et) => args :+ g.entities(et))
-        FactorTable(name, calc(resolvedEntities), g)
-      })
+    GraphSource(model.graphs.Raw.graphDef, args.outputDir, "csv").read
+      .pipe(g =>
+        rawFactors.map { case (name, calc) =>
+          val resolvedEntities = calc.requiredEntities.foldLeft(Seq.empty[DataFrame])((args, et) => args :+ g.entities(et))
+          FactorTable(name, calc(resolvedEntities), g)
+        }
+      )
       .foreach(_.write(FactorTableSink(args.outputDir)))
   }
 
   private def frequency(df: DataFrame, value: Column, by: Seq[Column], agg: Column => Column = count) =
     df
-      .groupBy(by: _*).agg(agg(value).as("frequency"))
+      .groupBy(by: _*)
+      .agg(agg(value).as("frequency"))
       .select(by :+ $"frequency": _*)
       .orderBy($"frequency".desc +: by.map(_.asc): _*)
 
-
   private def messageTags(commentHasTag: DataFrame, postHasTag: DataFrame, tag: DataFrame) = {
     val messageHasTag = commentHasTag.select($"CommentId".as("id"), $"TagId") |+| postHasTag.select($"PostId".as("id"), $"TagId")
 
@@ -54,11 +54,12 @@ object FactorGenerationStage extends DatagenStage with Logging {
 
   private val rawFactors = Map(
     "countryNumPersons" -> Factor(Place, Person) { case Seq(places, persons) =>
-      val cities = places.where($"type" === "City").cache()
+      val cities    = places.where($"type" === "City").cache()
       val countries = places.where($"type" === "Country").cache()
 
       frequency(
-        persons.as("Person")
+        persons
+          .as("Person")
           .join(broadcast(cities.as("City")), $"City.id" === $"Person.LocationCityId")
           .join(broadcast(countries.as("Country")), $"Country.id" === $"City.PartOfPlaceId"),
         value = $"Person.id",
@@ -76,7 +77,8 @@ object FactorGenerationStage extends DatagenStage with Logging {
       val cities = places.where($"type" === "City").cache()
 
       frequency(
-        personKnowsPerson.alias("Knows")
+        personKnowsPerson
+          .alias("Knows")
           .join(persons.cache().as("Person1"), $"Person1.id" === $"Knows.Person1Id")
           .join(cities.cache().as("City1"), $"City1.id" === "Person1.LocationCityId")
           .join(persons.as("Person2"), $"Person2.id" === $"Knows.Person2Id")
@@ -93,11 +95,12 @@ object FactorGenerationStage extends DatagenStage with Logging {
       )
     },
     "countryPairsNumFriends" -> Factor(PersonKnowsPerson, Person, Place) { case Seq(personKnowsPerson, persons, places) =>
-      val cities = places.where($"type" === "City").cache()
+      val cities    = places.where($"type" === "City").cache()
       val countries = places.where($"type" === "Country").cache()
 
       frequency(
-        personKnowsPerson.alias("Knows")
+        personKnowsPerson
+          .alias("Knows")
           .join(persons.cache().as("Person1"), $"Person1.id" === $"Knows.Person1Id")
           .join(cities.cache().as("City1"), $"City1.id" === "Person1.LocationCityId")
           .join(countries.cache().as("Country1"), $"Country1.id" === "City1.PartOfPlaceId")
@@ -132,7 +135,8 @@ object FactorGenerationStage extends DatagenStage with Logging {
     },
     "messageTagClasses" -> Factor(CommentHasTag, PostHasTag, Tag, TagClass) { case Seq(commentHasTag, postHasTag, tag, tagClass) =>
       frequency(
-        messageTags(commentHasTag, postHasTag, tag).as("MessageTags")
+        messageTags(commentHasTag, postHasTag, tag)
+          .as("MessageTags")
           .join(tag.as("Tag"), $"MessageTags.tagId" === $"Tag.id")
           .join(tagClass.as("TagClass"), $"Tag.TypeTagClassId" === $"TagClass.id"),
         value = $"frequency",
@@ -141,10 +145,10 @@ object FactorGenerationStage extends DatagenStage with Logging {
       )
     },
     "personNumFriends" -> Factor(PersonKnowsPerson) { case Seq(knows) =>
-      frequency(knows, value=$"Person2Id", by=Seq($"Person1Id"))
+      frequency(knows, value = $"Person2Id", by = Seq($"Person1Id"))
     },
     "postLanguages" -> Factor(Post) { case Seq(post) =>
-      frequency(post.where($"language".isNotNull), value=$"id", by=Seq($"language"))
+      frequency(post.where($"language".isNotNull), value = $"id", by = Seq($"language"))
     },
     "tagClassNumTags" -> Factor(TagClass, Tag) { case Seq(tagClass, tag) =>
       frequency(
@@ -163,5 +167,3 @@ object FactorGenerationStage extends DatagenStage with Logging {
     }
   )
 }
-
-

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

@@ -3,14 +3,13 @@ package ldbc.snb.datagen.factors
 import ldbc.snb.datagen.model.{Graph, GraphDef, Mode}
 import org.apache.spark.sql.DataFrame
 
-
 case class FactorTableDef[M <: Mode](
-                                      name: String,
-                                      sourceDef: GraphDef[M]
-                                    )
+    name: String,
+    sourceDef: GraphDef[M]
+)
 
 case class FactorTable[M <: Mode](
-                                   name: String,
-                                   data: DataFrame,
-                                   source: Graph[M]
-                                 )
+    name: String,
+    data: DataFrame,
+    source: Graph[M]
+)

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

@@ -14,34 +14,34 @@ 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",
+      oversizeFactor: Option[Double] = None
   )
 
   def run(args: Args, config: GeneratorConfiguration)(implicit spark: SparkSession) = {
-    val numPartitions = config.getInt("hadoop.numThreads", spark.sparkContext.defaultParallelism)
+    val numPartitions   = config.getInt("hadoop.numThreads", spark.sparkContext.defaultParallelism)
     val idealPartitions = DatagenParams.numPersons.toDouble / optimalPersonsPerFile
 
     val oversizeFactor = args.oversizeFactor.getOrElse(Math.max(numPartitions / idealPartitions, 1.0))
 
     val persons = SparkPersonGenerator(config)
 
-    val percentages = Seq(0.45f, 0.45f, 0.1f)
+    val percentages             = Seq(0.45f, 0.45f, 0.1f)
     val knowsGeneratorClassName = DatagenParams.getKnowsGenerator
 
     import ldbc.snb.datagen.entities.Keys._
 
-    val uniRanker = SparkRanker.create(_.byUni)
+    val uniRanker      = SparkRanker.create(_.byUni)
     val interestRanker = SparkRanker.create(_.byInterest)
-    val randomRanker = SparkRanker.create(_.byRandomId)
+    val randomRanker   = SparkRanker.create(_.byRandomId)
 
-    val uniKnows = SparkKnowsGenerator(persons, uniRanker, config, percentages, 0, knowsGeneratorClassName)
+    val uniKnows      = SparkKnowsGenerator(persons, uniRanker, config, percentages, 0, knowsGeneratorClassName)
     val interestKnows = SparkKnowsGenerator(persons, interestRanker, config, percentages, 1, knowsGeneratorClassName)
-    val randomKnows = SparkKnowsGenerator(persons, randomRanker, config, percentages, 2, knowsGeneratorClassName)
+    val randomKnows   = SparkKnowsGenerator(persons, randomRanker, config, percentages, 2, knowsGeneratorClassName)
 
     val merged = SparkKnowsMerger(uniKnows, interestKnows, randomKnows).cache()
 

src/main/scala/ldbc/snb/datagen/generation/generator/SparkKnowsGenerator.scala

@@ -15,12 +15,12 @@ import scala.reflect.ClassTag
 
 object SparkKnowsGenerator {
   def apply(
-    persons: RDD[Person],
-    ranker: SparkRanker,
-    conf: GeneratorConfiguration,
-    percentages: Seq[Float],
-    stepIndex: Int,
-    knowsGeneratorClassName: String
+      persons: RDD[Person],
+      ranker: SparkRanker,
+      conf: GeneratorConfiguration,
+      percentages: Seq[Float],
+      stepIndex: Int,
+      knowsGeneratorClassName: String
   )(implicit spark: SparkSession) = {
     val blockSize = DatagenParams.blockSize
 
@@ -33,14 +33,14 @@ object SparkKnowsGenerator {
       // groupByKey wouldn't guarantee keeping the order inside groups
       // TODO check if it actually has better performance than sorting inside mapPartitions (probably not)
       .combineByKeyWithClassTag(
-          personByRank => SortedMap(personByRank),
-          (map: SortedMap[Long, Person], personByRank) => map + personByRank,
-          (a: SortedMap[Long, Person], b: SortedMap[Long, Person]) => a ++ b
-        )
+        personByRank => SortedMap(personByRank),
+        (map: SortedMap[Long, Person], personByRank) => map + personByRank,
+        (a: SortedMap[Long, Person], b: SortedMap[Long, Person]) => a ++ b
+      )
       .mapPartitions(groups => {
         DatagenContext.initialize(conf)
         val knowsGeneratorClass = Class.forName(knowsGeneratorClassName)
-        val knowsGenerator = knowsGeneratorClass.getConstructor().newInstance().asInstanceOf[KnowsGenerator]
+        val knowsGenerator      = knowsGeneratorClass.getConstructor().newInstance().asInstanceOf[KnowsGenerator]
         knowsGenerator.initialize(conf)
         val personSimilarity = DatagenParams.getPersonSimularity
 
@@ -53,9 +53,9 @@ object SparkKnowsGenerator {
           clonedPersons
         }
 
-        for { 
+        for {
           persons <- personGroups
-          person <- persons.iterator().asScala
+          person  <- persons.iterator().asScala
         } yield person
       })
   }

src/main/scala/ldbc/snb/datagen/generation/generator/SparkPersonGenerator.scala

@@ -21,7 +21,7 @@ object SparkPersonGenerator {
 
       for {
         i <- blocks
-        _ = println(s"Processing person block $i (${DatagenParams.blockSize})")
+        _    = println(s"Processing person block $i (${DatagenParams.blockSize})")
         size = Math.min(DatagenParams.numPersons - DatagenParams.blockSize * i, DatagenParams.blockSize).toInt
         person <- personGenerator.generatePersonBlock(i.toInt, DatagenParams.blockSize).asScala.take(size)
       } yield person

src/main/scala/ldbc/snb/datagen/generation/generator/SparkRanker.scala

@@ -24,9 +24,9 @@ object SparkRanker {
         .mapPartitionsWithIndex((i, ps) => Array((i, ps.size)).iterator, preservesPartitioning = true)
         .collectAsMap()
 
-      val aggregatedCounts = SortedMap(counts.toSeq : _*)
-        .foldLeft((0L, Map.empty[Int, Long])) {
-          case ((total, map), (i, c)) => (total + c, map + (i -> total))
+      val aggregatedCounts = SortedMap(counts.toSeq: _*)
+        .foldLeft((0L, Map.empty[Int, Long])) { case ((total, map), (i, c)) =>
+          (total + c, map + (i -> total))
         }
         ._2
 

src/main/scala/ldbc/snb/datagen/generation/serializer/SparkActivitySerializer.scala

@@ -20,7 +20,9 @@ import scala.collection.JavaConverters._
 
 object SparkActivitySerializer {
 
-  def apply(persons: RDD[Person], ranker: SparkRanker, conf: GeneratorConfiguration, partitions: Option[Int] = None, oversizeFactor: Double = 1.0)(implicit spark: SparkSession) = {
+  def apply(persons: RDD[Person], ranker: SparkRanker, conf: GeneratorConfiguration, partitions: Option[Int] = None, oversizeFactor: Double = 1.0)(implicit
+      spark: SparkSession
+  ) = {
 
     val blockSize = DatagenParams.blockSize
     val blocks = ranker(persons)
@@ -33,8 +35,8 @@ object SparkActivitySerializer {
     blocks.foreachPartition(groups => {
       DatagenContext.initialize(conf)
       val partitionId = TaskContext.getPartitionId()
-      val hadoopConf = serializableHadoopConf.value
-      val buildDir = conf.getOutputDir
+      val hadoopConf  = serializableHadoopConf.value
+      val buildDir    = conf.getOutputDir
 
       val fs = FileSystem.get(new URI(buildDir), hadoopConf)
       fs.mkdirs(new Path(buildDir))
@@ -44,10 +46,10 @@ object SparkActivitySerializer {
       dynamicActivitySerializer.initialize(fs, conf.getOutputDir, partitionId, oversizeFactor, false)
 
       val generator = new PersonActivityGenerator
-      val exporter = new PersonActivityExporter(dynamicActivitySerializer)
+      val exporter  = new PersonActivityExporter(dynamicActivitySerializer)
 
       try {
-        for {(blockId, persons) <- groups} {
+        for { (blockId, persons) <- groups } {
           val clonedPersons = new util.ArrayList[Person]
           for (p <- persons) {
             clonedPersons.add(new Person(p))

src/main/scala/ldbc/snb/datagen/generation/serializer/SparkPersonSerializer.scala

@@ -15,20 +15,20 @@ import java.net.URI
 object SparkPersonSerializer {
 
   def apply(
-    persons: RDD[Person],
-    conf: GeneratorConfiguration,
-    partitions: Option[Int] = None,
-    oversizeFactor: Double = 1.0
+      persons: RDD[Person],
+      conf: GeneratorConfiguration,
+      partitions: Option[Int] = None,
+      oversizeFactor: Double = 1.0
   )(implicit spark: SparkSession): Unit = {
     val serializableHadoopConf = new SerializableConfiguration(spark.sparkContext.hadoopConfiguration)
 
     persons
       .pipeFoldLeft(partitions)((rdd: RDD[Person], p: Int) => rdd.coalesce(p))
       .foreachPartition(persons => {
         val dynamicPersonSerializer = new DynamicPersonSerializer
-        val hadoopConf = serializableHadoopConf.value
-        val partitionId = TaskContext.getPartitionId()
-        val buildDir = conf.getOutputDir
+        val hadoopConf              = serializableHadoopConf.value
+        val partitionId             = TaskContext.getPartitionId()
+        val buildDir                = conf.getOutputDir
 
         val fs = FileSystem.get(new URI(buildDir), hadoopConf)
         fs.mkdirs(new Path(buildDir))
@@ -45,7 +45,7 @@ object SparkPersonSerializer {
 
         personExporter use { pe =>
           DatagenContext.initialize(conf)
-          for {p <- persons} {
+          for { p <- persons } {
             pe.export(p)
           }
         }