remove Interactive object

Author: dszakallas

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

@@ -4,7 +4,7 @@ import ldbc.snb.datagen.model.{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, LongType, TimestampType}
+import org.apache.spark.sql.types.{DateType, TimestampType}
 import shapeless._
 
 object IrToRawTransform extends Transform[Mode.Raw.type, Mode.Raw.type] {

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

@@ -13,7 +13,7 @@ case class RawToBiTransform(mode: BI, simulationStart: Long, simulationEnd: Long
     with Logging {
   log.debug(s"BI Transformation parameters: $mode")
 
-  val bulkLoadThreshold = Interactive.calculateBulkLoadThreshold(mode.bulkloadPortion, simulationStart, simulationEnd)
+  val bulkLoadThreshold = RawToInteractiveTransform.calculateBulkLoadThreshold(mode.bulkloadPortion, simulationStart, simulationEnd)
 
   def batchPeriodFormat(batchPeriod: String) = batchPeriod match {
     case "year"   => "yyyy"
@@ -45,7 +45,7 @@ case class RawToBiTransform(mode: BI, simulationStart: Long, simulationEnd: Long
         .filter(inBatch($"creationDate", batchStart, batchEnd))
         .pipe(batched)
         .select(
-          Seq($"insert_batch_id".as("batch_id")) ++ Interactive.columns(tpe, df.columns).map(qcol): _*
+          Seq($"insert_batch_id".as("batch_id")) ++ RawToInteractiveTransform.columns(tpe, df.columns).map(qcol): _*
         )
         .repartitionByRange($"batch_id")
         .sortWithinPartitions($"creationDate")
@@ -70,7 +70,7 @@ case class RawToBiTransform(mode: BI, simulationStart: Long, simulationEnd: Long
         case (tpe, v) if tpe.isStatic => tpe -> BatchedEntity(v, None, None)
         case (tpe, v) =>
           tpe -> BatchedEntity(
-            Interactive.snapshotPart(tpe, v, bulkLoadThreshold, filterDeletion = false),
+            RawToInteractiveTransform.snapshotPart(tpe, v, bulkLoadThreshold, filterDeletion = false),
             Some(Batched(insertBatchPart(tpe, v, bulkLoadThreshold, simulationEnd), Seq("batch_id"))),
             if (keepImplicitDeletes || v.columns.contains("explicitlyDeleted"))
               Some(Batched(deleteBatchPart(tpe, v, bulkLoadThreshold, simulationEnd), Seq("batch_id")))

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

@@ -1,24 +1,60 @@
 package ldbc.snb.datagen.transformation.transform
 
-import ldbc.snb.datagen.model.{Graph, Mode}
+import ldbc.snb.datagen.model.Cardinality.NN
+import ldbc.snb.datagen.model.EntityType.Edge
+import ldbc.snb.datagen.model.{EntityType, Graph, Mode}
 import ldbc.snb.datagen.util.Logging
+import ldbc.snb.datagen.util.sql._
+import ldbc.snb.datagen.syntax._
 import org.apache.spark.sql.DataFrame
+import org.apache.spark.sql.functions.{col, lit, to_timestamp}
 
 case class RawToInteractiveTransform(mode: Mode.Interactive, simulationStart: Long, simulationEnd: Long)
     extends Transform[Mode.Raw.type, Mode.Interactive]
     with Logging {
   log.debug(s"Interactive Transformation parameters: $mode")
 
-  val bulkLoadThreshold = Interactive.calculateBulkLoadThreshold(mode.bulkLoadPortion, simulationStart, simulationEnd)
+  val bulkLoadThreshold = RawToInteractiveTransform.calculateBulkLoadThreshold(mode.bulkLoadPortion, simulationStart, simulationEnd)
 
   override def transform(input: In): Out = {
     val entities = input.entities
       .map { case (tpe, v) =>
         tpe -> IrToRawTransform.convertDates(tpe, v)
       }
       .map { case (tpe, v) =>
-        tpe -> Interactive.snapshotPart(tpe, v, bulkLoadThreshold, filterDeletion = true)
+        tpe -> RawToInteractiveTransform.snapshotPart(tpe, v, bulkLoadThreshold, filterDeletion = true)
       }
     Graph[Mode.Interactive](isAttrExploded = input.isAttrExploded, isEdgesExploded = input.isEdgesExploded, mode, entities)
   }
 }
+
+object RawToInteractiveTransform {
+
+  def columns(tpe: EntityType, cols: Seq[String]) = tpe match {
+    case tpe if tpe.isStatic => cols
+    case Edge("Knows", "Person", "Person", NN, false) =>
+      val rawCols = Set("deletionDate", "explicitlyDeleted", "weight")
+      cols.filter(!rawCols.contains(_))
+    case _ =>
+      val rawCols = Set("deletionDate", "explicitlyDeleted")
+      cols.filter(!rawCols.contains(_))
+  }
+
+  def calculateBulkLoadThreshold(bulkLoadPortion: Double, simulationStart: Long, simulationEnd: Long) = {
+    (simulationEnd - ((simulationEnd - simulationStart) * (1 - bulkLoadPortion)).toLong)
+  }
+
+  def snapshotPart(tpe: EntityType, df: DataFrame, bulkLoadThreshold: Long, filterDeletion: Boolean) = {
+    val filterBulkLoad = (ds: DataFrame) =>
+      ds
+        .filter(
+          $"creationDate" < to_timestamp(lit(bulkLoadThreshold / 1000)) &&
+            (!lit(filterDeletion) || $"deletionDate" >= to_timestamp(lit(bulkLoadThreshold / 1000)))
+        )
+
+    tpe match {
+      case tpe if tpe.isStatic => df
+      case tpe                 => filterBulkLoad(df).select(columns(tpe, df.columns).map(name => col(qualified(name))): _*)
+    }
+  }
+}