refactor io classes

Author: dszakallas

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

@@ -0,0 +1,34 @@
+package ldbc.snb.datagen.factors
+
+import ldbc.snb.datagen.model.Mode.Raw
+import ldbc.snb.datagen.model.{Graph, graphs}
+import ldbc.snb.datagen.syntax._
+import org.apache.spark.sql.functions.{broadcast, count}
+import org.apache.spark.sql.{Column, DataFrame}
+
+object Factors {
+
+  def frequency(df: DataFrame, value: Column, by: Seq[Column], sortBy: Seq[Column]) =
+    df
+      .groupBy(by: _*).agg(count(value).as("count"))
+      .select(by :+ $"count": _*)
+      .orderBy($"count".desc +: by: _*)
+
+
+  def countryNumPersons(graph: Graph[Raw.type]): FactorTable[Raw.type] = {
+    val places = graph.entities(graphs.Raw.entities.Place).cache()
+    val cities = places.where($"type" === "City")
+    val countries = places.where($"type" === "Country")
+
+    val persons = graph.entities(graphs.Raw.entities.Person)
+    val df = frequency(
+      persons.as("Person")
+        .join(broadcast(cities.as("City")), $"City.id" === $"Person.LocationCityId")
+        .join(broadcast(countries.as("Country")), $"Country.id" === $"City.PartOfPlaceId"),
+      value = $"Person.id",
+      by = Seq($"Country.id", $"Country.name"),
+      sortBy = Seq($"Country.id")
+    )
+    FactorTable[Raw.type](name="countryNumPersons", data=df, source=graph)
+  }
+}

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

@@ -0,0 +1,29 @@
+package ldbc.snb.datagen.factors
+
+import ldbc.snb.datagen.SparkApp
+import ldbc.snb.datagen.factors.Factors.countryNumPersons
+import ldbc.snb.datagen.factors.io.FactorTableSink
+import ldbc.snb.datagen.io.graphs.GraphSource
+import ldbc.snb.datagen.model
+import ldbc.snb.datagen.syntax._
+import ldbc.snb.datagen.util.Logging
+import org.apache.spark.sql.SparkSession
+
+object FactorGenerationStage extends SparkApp with Logging {
+  override def appName: String = "LDBC SNB Datagen for Spark: Factor Generation Stage"
+
+  case class Args(outputDir: String = "out")
+
+  def run(args: Args)(implicit spark: SparkSession): Unit = {
+    import ldbc.snb.datagen.io.instances._
+    import ldbc.snb.datagen.io.Reader.ops._
+    import ldbc.snb.datagen.io.Writer.ops._
+    import ldbc.snb.datagen.factors.io.instances._
+
+    GraphSource(model.graphs.Raw.graphDef, args.outputDir, "csv")
+      .read
+      .pipe(countryNumPersons)
+      .write(FactorTableSink(args.outputDir))
+  }
+
+}

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

@@ -0,0 +1,16 @@
+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]
+                            )
+
+case class FactorTable[M <: Mode](
+                        name: String,
+                        data: DataFrame,
+                        source: Graph[M]
+                      )

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

@@ -0,0 +1,28 @@
+package ldbc.snb.datagen.factors
+
+import ldbc.snb.datagen.io.{PathComponent, Writer}
+import ldbc.snb.datagen.util.Logging
+import ldbc.snb.datagen.model.{GraphLike, Mode}
+import better.files._
+import ldbc.snb.datagen.io.dataframes.DataFrameSink
+
+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._
+
+  private final class FactorTableWriter[M <: Mode] extends Writer[FactorTableSink] with Logging {
+    override type CoRet = FactorTable[M]
+
+    override def write(self: FactorTable[M], sink: FactorTableSink): Unit = {
+      val p = (sink.path / "factors" / sink.format / PathComponent[GraphLike[M]].path(self.source) / self.name).toString()
+      self.data.coalesce(1).write(DataFrameSink(p, sink.format))
+    }
+  }
+  trait WriterInstances {
+    implicit def factorTableWriter[M <: Mode]: Writer.Aux[FactorTableSink, FactorTable[M]] = new FactorTableWriter[M]
+  }
+
+  object instances extends WriterInstances
+}

src/main/scala/ldbc/snb/datagen/transformation/io/PathComponent.scala renamed to src/main/scala/ldbc/snb/datagen/io/PathComponent.scala

@@ -1,6 +1,6 @@
-package ldbc.snb.datagen.transformation.io
+package ldbc.snb.datagen.io
 
-import ldbc.snb.datagen.transformation.model.{GraphLike, Mode}
+import ldbc.snb.datagen.model.{GraphLike, Mode}
 
 trait PathComponent[A] {
   def path(a: A): String

src/main/scala/ldbc/snb/datagen/io/Reader.scala

@@ -0,0 +1,35 @@
+package ldbc.snb.datagen.io
+
+trait Reader[T] {
+  type Ret
+
+  def read(self: T): Ret
+  def exists(self: T): Boolean
+}
+
+object Reader {
+  type Aux[T, R] = Reader[T] { type Ret = R }
+
+  def apply[T, R](implicit r: Reader.Aux[T, R]): Reader.Aux[T, R] = implicitly[Reader.Aux[T, R]]
+
+  trait ReaderOps[T] {
+    type Ret
+    def tcInstance: Reader.Aux[T, Ret]
+    def self: T
+    def read: Ret = tcInstance.read(self)
+  }
+
+  object ReaderOps {
+    type Aux[T, R] = ReaderOps[T] { type Ret = R }
+  }
+
+  object ops {
+    import scala.language.implicitConversions
+    implicit def toReaderOps[T, R](target: T)(implicit tc: Reader.Aux[T, R]): ReaderOps.Aux[T, R] = new ReaderOps[T] {
+      override type Ret = R
+      override def tcInstance: Aux[T, R] = tc
+      override def self: T = target
+    }
+  }
+}
+

src/main/scala/ldbc/snb/datagen/io/Writer.scala

@@ -0,0 +1,34 @@
+package ldbc.snb.datagen.io
+
+
+trait Writer[S] {
+  type CoRet
+  def write(self: CoRet, sink: S): Unit
+}
+
+object Writer {
+  type Aux[S, C] = Writer[S] { type CoRet = C }
+  def apply[S, C](implicit r: Writer.Aux[S, C]): Writer.Aux[S, C] = implicitly[Writer.Aux[S, C]]
+
+  trait WriterOps[CoRet] {
+    type Sink
+    def tcInstance: Writer.Aux[Sink, CoRet]
+    def self: CoRet
+    def write(sink: Sink): Unit = tcInstance.write(self, sink)
+  }
+
+  object WriterOps {
+    type Aux[CoRet, S] = WriterOps[CoRet] { type Sink = S }
+  }
+
+  object ops {
+    import scala.language.implicitConversions
+    implicit def toWriterOps[CoRet, S](target: CoRet)(implicit tc: Writer.Aux[S, CoRet]): WriterOps.Aux[CoRet, S] = new WriterOps[CoRet] {
+      override type Sink = S
+      override def tcInstance: Aux[S, CoRet] = tc
+      override def self: CoRet = target
+    }
+  }
+}
+
+

src/main/scala/ldbc/snb/datagen/io/dataframes.scala

@@ -0,0 +1,68 @@
+package ldbc.snb.datagen.io
+
+import ldbc.snb.datagen.syntax.fluentSyntaxOps
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.{DataFrame, SaveMode, SparkSession}
+
+
+object dataframes {
+
+  case class DataFrameSource(
+                            path: String,
+                            format: String,
+                            formatOptions: Map[String, String] = Map.empty,
+                            schema: Option[StructType] = None
+                            )
+
+  private class DataFrameReader(implicit spark: SparkSession) extends Reader[DataFrameSource] {
+    override type Ret = DataFrame
+
+    override def read(self: DataFrameSource): DataFrame = {
+      spark.read
+        .format(self.format)
+        .options(self.formatOptions)
+        .pipeFoldLeft(self.schema)(_ schema _)
+        .load(self.path)
+    }
+
+    override def exists(self: DataFrameSource): Boolean = utils.fileExists(self.path)
+  }
+
+  trait ReaderInstances {
+    implicit def dataFrameReader(implicit spark: SparkSession): Reader.Aux[DataFrameSource, DataFrame] = new DataFrameReader
+  }
+
+  case class DataFrameSink(path: String,
+                               format: String,
+                               formatOptions: Map[String, String] = Map.empty,
+                               mode: SaveMode = SaveMode.ErrorIfExists,
+                           partitionBy: Seq[String] = Seq.empty
+                              )
+
+  case class DataFrameWriterOptions(
+                                     format: String,
+                                     partitionBy: Seq[String] = Seq.empty,
+                                     formatOptions: Map[String, String] = Map.empty,
+                                     mode: SaveMode = SaveMode.ErrorIfExists
+                                   )
+
+  private object DataFrameWriter extends Writer[DataFrameSink] {
+    override type CoRet = DataFrame
+    override def write(self: DataFrame, sink: DataFrameSink): Unit = {
+      self.write
+        .partitionBy(sink.partitionBy: _*)
+        .format(sink.format)
+        .options(sink.formatOptions)
+        .mode(sink.mode)
+        .save(sink.path)
+    }
+  }
+
+  trait WriterInstances {
+    implicit val dataFrameWriter: Writer.Aux[DataFrameSink, DataFrame] = DataFrameWriter
+  }
+
+  trait Instances extends WriterInstances with ReaderInstances
+
+  object instances extends Instances
+}