dry up ConvertDates

Author: dszakallas

src/main/scala/ldbc/snb/datagen/transformation/transform/ConvertDates.scala

@@ -1,13 +1,15 @@
 package ldbc.snb.datagen.transformation.transform
 
-import ldbc.snb.datagen.model.{Batched, BatchedEntity, EntityType, Mode}
+import ldbc.snb.datagen.model.{Batched, BatchedEntity, EntityType, Graph, Mode}
 import ldbc.snb.datagen.util.sql.qcol
 import org.apache.spark.sql.DataFrame
 import org.apache.spark.sql.functions.lit
 import org.apache.spark.sql.types.{DateType, TimestampType}
 import shapeless._
 
 trait ConvertDates[M <: Mode] extends Transform[M, M] {
+  def convertDatesInEntities(entities: Map[EntityType, M#Layout]): Map[EntityType, (Option[String], M#Layout)]
+
   def convertDates(tpe: EntityType, df: DataFrame): DataFrame = {
     tpe match {
       case tpe if !tpe.isStatic =>
@@ -19,48 +21,51 @@ trait ConvertDates[M <: Mode] extends Transform[M, M] {
       case _ => df
     }
   }
+
+  override def transform(input: Graph[M]): Graph[M] = {
+    if (input.definition.useTimestamp) {
+      throw new AssertionError("Already using timestamp for dates")
+    }
+
+    val updatedEntities = convertDatesInEntities(input.entities)
+
+    val modifiedEntities = updatedEntities
+      .map { case (k, (_, data)) => k -> data }
+
+    val modifiedEntityDefinitions = updatedEntities
+      .map { case (k, (schema, _)) => k -> schema }
+
+    val l = lens[In]
+    (l.definition.useTimestamp ~ l.definition.entities ~ l.entities).set(input)((true, modifiedEntityDefinitions, modifiedEntities))
+  }
 }
 
 object ConvertDates {
   def apply[T <: Mode : ConvertDates] = implicitly[ConvertDates[T]]
 
   object instances {
     implicit def batchedConvertDates[M <: Mode](implicit ev: BatchedEntity =:= M#Layout) = new ConvertDates[M] {
-      override def transform(input: In): Out = {
-        if (input.definition.useTimestamp) {
-          throw new AssertionError("Already using timestamp for dates")
-        }
-        val modifiedEntities = input.entities.map { case (tpe, layout) => tpe -> {
+      override def convertDatesInEntities(entities: Map[EntityType,M#Layout]): Map[EntityType, (Option[String], M#Layout)] = {
+        entities.map { case (tpe, layout) =>
           val be = layout.asInstanceOf[BatchedEntity]
-          ev(BatchedEntity(
-            convertDates(tpe, be.snapshot),
-            be.insertBatches.map(b => Batched(convertDates(tpe, b.entity), b.batchId, b.ordering)),
-            be.deleteBatches.map(b => Batched(convertDates(tpe, b.entity), b.batchId, b.ordering))
-          ))
-        }}
+          val convertedSnapshot = convertDates(tpe, be.snapshot)
 
-        val modifiedEntityDefinitions = modifiedEntities
-          .map { case (tpe, v) => tpe -> Some(v.asInstanceOf[BatchedEntity].snapshot.schema.toDDL) }
-
-        val l = lens[In]
-        (l.definition.useTimestamp ~ l.definition.entities ~ l.entities).set(input)((true, modifiedEntityDefinitions, modifiedEntities))
+          val convertedInserts = be.insertBatches.map(b => Batched(convertDates(tpe, b.entity), b.batchId, b.ordering))
+          val convertedDeletes = be.deleteBatches.map(b => Batched(convertDates(tpe, b.entity), b.batchId, b.ordering))
+          (tpe, (
+            Some(convertedSnapshot.schema.toDDL),
+            ev(BatchedEntity(convertedSnapshot, convertedInserts, convertedDeletes))
+          ))
+        }
       }
     }
 
     implicit def simpleConvertDates[M <: Mode](implicit ev: DataFrame =:= M#Layout) = new ConvertDates[M] {
-      override def transform(input: In): Out = {
-        if (input.definition.useTimestamp) {
-          throw new AssertionError("Already using timestamp for dates")
+      override def convertDatesInEntities(entities: Map[EntityType, M#Layout]): Map[EntityType, (Option[String], M#Layout)] = {
+        entities.map { case (tpe, v) =>
+          val convertedSnapshot = convertDates(tpe, v.asInstanceOf[DataFrame])
+          (tpe, (Some(convertedSnapshot.schema.toDDL), ev(convertedSnapshot)))
         }
-
-        val modifiedEntities = input.entities
-          .map { case (tpe, v) => tpe -> ev(convertDates(tpe, v.asInstanceOf[DataFrame])) }
-
-        val modifiedEntityDefinitions = modifiedEntities
-          .map { case (tpe, v) => tpe -> Some(v.asInstanceOf[DataFrame].schema.toDDL) }
-
-        val l = lens[In]
-        (l.definition.useTimestamp ~ l.definition.entities ~ l.entities).set(input)((true, modifiedEntityDefinitions, modifiedEntities))
       }
     }
   }