Merge pull request #401 from ldbc/dates-as-epoch-milli

Option to use epochMillis for datetime types

Author: szarnyasg

.circleci/config.yml

@@ -97,6 +97,11 @@ jobs:
           command: |
             tools/docker-run.sh --mode bi --scale-factor 0.003 --explode-edges --format-options header=false,quoteAll=true,compression=gzip
             mv out/ social-network-sf0.003-bi-composite-projected-fk-neo4j-compressed/
+      - run:
+          name: Generate SF0.003 / BI / compressed composite-projected CSVs for Neo4j with epoch milli timestamps
+          command: |
+            tools/docker-run.sh --mode bi --scale-factor 0.003 --explode-edges --epoch-millis --format-options header=false,quoteAll=true,compression=gzip
+            mv out/ social-network-sf0.003-bi-composite-projected-fk-neo4j-compressed-epoch-millis/
       # Interactive
       - run:
           name: Generate SF0.003 / Interactive / singular-projected CSVs

README.md

@@ -136,6 +136,12 @@ To get a complete list of the arguments, pass `--help` to the JAR file:
   ./tools/run.py ./target/ldbc_snb_datagen_${PLATFORM_VERSION}-${DATAGEN_VERSION}.jar -- --format parquet --scale-factor 0.003 --mode bi
   ```
 
+* Use epoch milliseconds encoded as longs (née `LongDateFormatter`) for serializing date and datetime values:
+
+  ```bash
+  ./tools/run.py ./target/ldbc_snb_datagen_${PLATFORM_VERSION}-${DATAGEN_VERSION}.jar -- --format csv --scale-factor 0.003 --mode bi --epoch-millis
+  ```
+
 * For the `interactive` and `bi` formats, the `--format-options` argument allows passing formatting options such as timestamp/date formats, the presence/abscence of headers (see the [Spark formatting options](https://spark.apache.org/docs/2.4.8/api/scala/index.html#org.apache.spark.sql.DataFrameWriter) for details), and whether quoting the fields in the CSV required:
 
   ```bash

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

@@ -28,7 +28,8 @@ object LdbcDatagen extends SparkApp {
       format: String = "csv",
       generateFactors: Boolean = false,
       formatOptions: Map[String, String] = Map.empty,
-      oversizeFactor: Option[Double] = None
+      oversizeFactor: Option[Double] = None,
+      epochMillis: Boolean = false
   )
 
   override type ArgsType = Args
@@ -118,6 +119,10 @@ object LdbcDatagen extends SparkApp {
         .text("Generate factor tables")
 
       help('h', "help").text("prints this usage text")
+
+      opt[Unit]("epoch-millis")
+        .action((x, c) => args.epochMillis.set(c)(true))
+        .text("Use longs with millis since Unix epoch instead of native dates")
     }
 
     val parsedArgs = parser.parse(args, Args()).getOrElse(throw new RuntimeException("Invalid arguments"))
@@ -159,7 +164,8 @@ object LdbcDatagen extends SparkApp {
       },
       irFormat,
       args.format,
-      args.formatOptions
+      args.formatOptions,
+      args.epochMillis
     )
 
     TransformationStage.run(transformArgs)

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

@@ -4,6 +4,7 @@ import ldbc.snb.datagen.factors.io.FactorTableSink
 import ldbc.snb.datagen.io.graphs.GraphSource
 import ldbc.snb.datagen.model
 import ldbc.snb.datagen.model.EntityType
+import ldbc.snb.datagen.model.Mode.Raw
 import ldbc.snb.datagen.syntax._
 import ldbc.snb.datagen.transformation.transform.ConvertDates
 import ldbc.snb.datagen.util.{DatagenStage, Logging}
@@ -74,9 +75,10 @@ object FactorGenerationStage extends DatagenStage with Logging {
     import ldbc.snb.datagen.io.Reader.ops._
     import ldbc.snb.datagen.io.Writer.ops._
     import ldbc.snb.datagen.io.instances._
+    import ldbc.snb.datagen.transformation.transform.ConvertDates.instances._
 
     GraphSource(model.graphs.Raw.graphDef, args.outputDir, args.irFormat).read
-      .pipe(ConvertDates.transform)
+      .pipe(ConvertDates[Raw.type].transform)
       .pipe(g =>
         rawFactors
           .collect {

src/main/scala/ldbc/snb/datagen/transformation/TransformationStage.scala

@@ -6,7 +6,7 @@ import ldbc.snb.datagen.model.{BatchedEntity, Graph, Mode}
 import ldbc.snb.datagen.syntax._
 import ldbc.snb.datagen.transformation.transform._
 import ldbc.snb.datagen.util.{DatagenStage, Logging}
-import org.apache.spark.sql.{DataFrame, SparkSession}
+import org.apache.spark.sql.DataFrame
 import shapeless._
 
 object TransformationStage extends DatagenStage with Logging {
@@ -20,14 +20,16 @@ object TransformationStage extends DatagenStage with Logging {
       mode: Mode = Mode.Raw,
       irFormat: String = "parquet",
       format: String = "csv",
-      formatOptions: Map[String, String] = Map.empty
+      formatOptions: Map[String, String] = Map.empty,
+      epochMillis: Boolean = false
   )
 
   override type ArgsType = Args
 
   import ldbc.snb.datagen.io.Reader.ops._
   import ldbc.snb.datagen.io.Writer.ops._
   import ldbc.snb.datagen.io.instances._
+  import ldbc.snb.datagen.transformation.transform.ConvertDates.instances._
 
   def run(args: ArgsType) = {
     object write extends Poly1 {
@@ -38,12 +40,19 @@ object TransformationStage extends DatagenStage with Logging {
         at[Graph[M]](g => g.write(GraphSink(args.outputDir, args.format, args.formatOptions)))
     }
 
+    object convertDates extends Poly1 {
+      implicit def caseSimple[M <: Mode](implicit ev: DataFrame =:= M#Layout) =
+        at[Graph[M]](g => ConvertDates[M].transform(g))
+
+      implicit def caseBatched[M <: Mode](implicit ev: BatchedEntity =:= M#Layout) =
+        at[Graph[M]](g => ConvertDates[M].transform(g))
+    }
+
     type Out = Graph[Mode.Raw.type] :+: Graph[Mode.Interactive] :+: Graph[Mode.BI] :+: CNil
 
     GraphSource(model.graphs.Raw.graphDef, args.outputDir, args.irFormat).read
       .pipeFoldLeft(args.explodeAttrs.fork)((graph, _: Unit) => ExplodeAttrs.transform(graph))
       .pipeFoldLeft(args.explodeEdges.fork)((graph, _: Unit) => ExplodeEdges.transform(graph))
-      .pipe(ConvertDates.transform)
       .pipe[Out] { g =>
         args.mode match {
           case bi @ Mode.BI(_, _) => Coproduct[Out](RawToBiTransform(bi, args.simulationStart, args.simulationEnd, args.keepImplicitDeletes).transform(g))
@@ -52,6 +61,7 @@ object TransformationStage extends DatagenStage with Logging {
           case Mode.Raw => Coproduct[Out](g)
         }
       }
+      .pipeFoldLeft((!args.epochMillis).fork)((graph, _: Unit) => graph.map(convertDates))
       .map(write)
     ()
   }

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

@@ -1,29 +1,52 @@
 package ldbc.snb.datagen.transformation.transform
 
-import ldbc.snb.datagen.model.{EntityType, Mode}
+import ldbc.snb.datagen.model.{Batched, BatchedEntity, EntityType, 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._
 
-object ConvertDates extends Transform[Mode.Raw.type, Mode.Raw.type] {
-
+trait ConvertDates[M <: Mode] extends Transform[M, M] {
   def convertDates(tpe: EntityType, df: DataFrame): DataFrame = {
     tpe match {
       case tpe if !tpe.isStatic =>
         df.select(df.columns.map {
-          case col @ ("creationDate" | "deletionDate") => (qcol(col) / lit(1000L)).cast(TimestampType).as(col)
-          case col @ "birthday"                        => (qcol(col) / lit(1000L)).cast(TimestampType).cast(DateType).as(col)
-          case col                                     => qcol(col)
+          case col@("creationDate" | "deletionDate") => (qcol(col) / lit(1000L)).cast(TimestampType).as(col)
+          case col@"birthday" => (qcol(col) / lit(1000L)).cast(TimestampType).cast(DateType).as(col)
+          case col => qcol(col)
         }: _*)
       case _ => df
     }
   }
+}
+
+object ConvertDates {
+  def apply[T <: Mode : ConvertDates] = implicitly[ConvertDates[T]]
 
-  override def transform(input: In): Out = {
-    lens[In].entities.modify(input)(
-      _.map { case (tpe, v) => tpe -> convertDates(tpe, v) }
-    )
+  object instances {
+    implicit def batchedConvertDates[M <: Mode](implicit ev: BatchedEntity =:= M#Layout) = new ConvertDates[M] {
+      override def transform(input: In): Out = {
+        lens[In].entities.modify(input)(
+          _.map { case (tpe, layout) => tpe -> {
+            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))
+            ))
+          }
+          }
+        )
+      }
+    }
+
+    implicit def simpleConvertDates[M <: Mode](implicit ev: DataFrame =:= M#Layout) = new ConvertDates[M] {
+      override def transform(input: In): Out = {
+        lens[In].entities.modify(input)(
+          _.map { case (tpe, v) => tpe -> ev(convertDates(tpe, v.asInstanceOf[DataFrame])) }
+        )
+      }
+    }
   }
 }

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

@@ -25,11 +25,10 @@ case class RawToBiTransform(mode: BI, simulationStart: Long, simulationEnd: Long
   }
 
   override def transform(input: In): Out = {
-    val batch_id = (col: Column) => date_format(date_trunc(mode.batchPeriod, col), batchPeriodFormat(mode.batchPeriod))
+    val batch_id = (col: Column) => date_format(date_trunc(mode.batchPeriod, to_timestamp(col / lit(1000L))), batchPeriodFormat(mode.batchPeriod))
 
     def inBatch(col: Column, batchStart: Long, batchEnd: Long) =
-      col >= to_timestamp(lit(batchStart / 1000)) &&
-        col < to_timestamp(lit(batchEnd / 1000))
+      col >= lit(batchStart) && col < lit(batchEnd)
 
     val batched = (df: DataFrame) =>
       df

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

@@ -45,8 +45,8 @@ object RawToInteractiveTransform {
     val filterBulkLoad = (ds: DataFrame) =>
       ds
         .filter(
-          $"creationDate" < to_timestamp(lit(bulkLoadThreshold / 1000)) &&
-            (!lit(filterDeletion) || $"deletionDate" >= to_timestamp(lit(bulkLoadThreshold / 1000)))
+          $"creationDate" < lit(bulkLoadThreshold) &&
+            (!lit(filterDeletion) || $"deletionDate" >= lit(bulkLoadThreshold))
         )
 
     tpe match {