ldbc · szarnyasg · Jul 5, 2022 · Jul 5, 2022 · Jul 5, 2022 · Jul 5, 2022
diff --git a/src/main/scala/ldbc/snb/datagen/io/graphs.scala b/src/main/scala/ldbc/snb/datagen/io/graphs.scala
@@ -4,6 +4,7 @@ import ldbc.snb.datagen.io.dataframes.{DataFrameSink, DataFrameSource}
 import ldbc.snb.datagen.model.EntityType.{Attr, Edge, Node}
 import ldbc.snb.datagen.model.Mode.Raw
 import ldbc.snb.datagen.model._
+import ldbc.snb.datagen.model.raw.{ForumType, PersonType}
 import ldbc.snb.datagen.syntax._
 import ldbc.snb.datagen.util.{Logging, SparkUI}
 import org.apache.spark.sql.functions.col
@@ -36,10 +37,9 @@ object graphs {
         implicit val atNode = at[Node](n => Seq(n.name.hashCode, 0, n.hashCode()))
         implicit val atEdge = at[Edge](e => {
           val primary = e match {
-            case Edge("Likes", "Person", "Comment", _, _)   => "Comment"
-            case Edge("Likes", "Person", "Post", _, _)      => "Post"
-            case Edge("ContainerOf", "Forum", "Post", _, _) => "Post"
-            case Edge(_, source, _, _, _)                   => source
+            case Edge("Likes", PersonType, dest, _, _, _, _)   => dest.name
+            case Edge("ContainerOf", ForumType, dest, _, _, _, _) => dest.name
+            case Edge(_, source, _, _, _, _, _)                   => source.name
           }
           Seq(primary.hashCode, 2, e.hashCode)
         })

diff --git a/src/main/scala/ldbc/snb/datagen/model/package.scala b/src/main/scala/ldbc/snb/datagen/model/package.scala
@@ -48,36 +48,41 @@ package object model {
     private def s(isStatic: Boolean) = if (isStatic) "static" else "dynamic"
 
     final case class Node(name: String, isStatic: Boolean = false) extends EntityType {
-      override val entityPath: String      = s"${s(isStatic)}/${name}"
+      override val entityPath: String      = s"${s(isStatic)}/$name"
       override val primaryKey: Seq[String] = Seq("id")
       override def toString: String        = s"$name"
     }
 
     final case class Edge(
         `type`: String,
-        source: String,
-        destination: String,
+        source: Node,
+        destination: Node,
         cardinality: Cardinality,
-        isStatic: Boolean = false
+        isStatic: Boolean = false,
+        sourceNameOverride: Option[String] = None,
+        destinationNameOverride: Option[String] = None
     ) extends EntityType {
-      override val entityPath: String = s"${s(isStatic)}/${source}_${pascalToCamel(`type`)}_${destination}"
+      val sourceName: String = sourceNameOverride.getOrElse(source.name)
+      val destinationName: String = destinationNameOverride.getOrElse(destination.name)
 
-      override val primaryKey: Seq[String] = ((source, destination) match {
+      override val entityPath: String = s"${s(isStatic)}/${sourceName}_${pascalToCamel(`type`)}_${destinationName}"
+
+      override val primaryKey: Seq[String] = ((sourceName, destinationName) match {
         case (s, d) if s == d => Seq(s"${s}1", s"${d}2")
         case (s, d)           => Seq(s, d)
       }).map(name => s"${name}Id")
 
-      override def toString: String = s"$source -[${`type`}]-> $destination"
+      override def toString: String = s"${sourceName} -[${`type`}]-> ${destinationName}"
     }
 
-    final case class Attr(`type`: String, parent: String, attribute: String, isStatic: Boolean = false) extends EntityType {
-      override val entityPath: String = s"${s(isStatic)}/${parent}_${pascalToCamel(`type`)}_${attribute}"
+    final case class Attr(`type`: String, parent: Node, attribute: String, isStatic: Boolean = false) extends EntityType {
+      override val entityPath: String = s"${s(isStatic)}/${parent.name}_${pascalToCamel(`type`)}_${attribute}"
 
-      override val primaryKey: Seq[String] = ((parent, attribute) match {
+      override val primaryKey: Seq[String] = ((parent.name, attribute) match {
         case (s, d) if s == d => Seq(s"${s}1", s"${d}2")
         case (s, d)           => Seq(s, d)
       }).map(name => s"${name}Id")
-      override def toString: String = s"$parent ♢-[${`type`}]-> $attribute"
+      override def toString: String = s"${parent.name} ♢-[${`type`}]-> $attribute"
     }
 
   }

diff --git a/src/main/scala/ldbc/snb/datagen/model/raw.scala b/src/main/scala/ldbc/snb/datagen/model/raw.scala
@@ -174,19 +174,21 @@ object raw {
   val TagType                     = Node("Tag", isStatic = true)
   val TagClassType                = Node("TagClass", isStatic = true)
   val CommentType                 = Node("Comment")
-  val CommentHasTagType           = Edge("HasTag", "Comment", "Tag", NN)
-  val ForumType                   = Node("Forum")
-  val ForumHasMemberType          = Edge("HasMember", "Forum", "Person", NN)
-  val ForumHasTagType             = Edge("HasTag", "Forum", "Tag", NN)
   val PersonType                  = Node("Person")
-  val PersonHasInterestTagType    = Edge("HasInterest", "Person", "Tag", NN)
-  val PersonKnowsPersonType       = Edge("Knows", "Person", "Person", NN)
-  val PersonLikesCommentType      = Edge("Likes", "Person", "Comment", NN)
-  val PersonLikesPostType         = Edge("Likes", "Person", "Post", NN)
-  val PersonStudyAtUniversityType = Edge("StudyAt", "Person", "University", OneN)
-  val PersonWorkAtCompanyType     = Edge("WorkAt", "Person", "Company", NN)
+  val ForumType                   = Node("Forum")
   val PostType                    = Node("Post")
-  val PostHasTagType              = Edge("HasTag", "Post", "Tag", NN)
+
+  val CommentHasTagType           = Edge("HasTag", CommentType, TagType, NN)
+  val ForumHasMemberType          = Edge("HasMember", ForumType, PersonType, NN)
+  val ForumHasTagType             = Edge("HasTag", ForumType, TagType, NN)
+  val PersonHasInterestTagType    = Edge("HasInterest", PersonType, TagType, NN)
+  val PersonKnowsPersonType       = Edge("Knows", PersonType, PersonType, NN)
+  val PersonLikesCommentType      = Edge("Likes", PersonType, CommentType, NN)
+  val PersonLikesPostType         = Edge("Likes", PersonType, PostType, NN)
+  val PersonStudyAtUniversityType = Edge("StudyAt", PersonType, OrganisationType, OneN, destinationNameOverride = Some("University"))
+  val PersonWorkAtCompanyType     = Edge("WorkAt", PersonType, OrganisationType, NN, destinationNameOverride = Some("Company"))
+
+  val PostHasTagType              = Edge("HasTag", PostType, TagType, NN)
 
   trait EntityTraitsInstances {
     import EntityTraits._

diff --git a/src/main/scala/ldbc/snb/datagen/transformation/transform/ExplodeAttrs.scala b/src/main/scala/ldbc/snb/datagen/transformation/transform/ExplodeAttrs.scala
@@ -17,8 +17,8 @@ object ExplodeAttrs extends Transform[Mode.Raw.type, Mode.Raw.type] {
     val updatedEntities = input.entities
       .collect { case (k @ Node("Person", false), v) =>
         Map(
-          explodedAttr(Attr("Email", "Person", "EmailAddress"), v, $"email"),
-          explodedAttr(Attr("Speaks", "Person", "Language"), v, $"language"),
+          explodedAttr(Attr("Email", k, "EmailAddress"), v, $"email"),
+          explodedAttr(Attr("Speaks", k, "Language"), v, $"language"),
           k -> v.drop("email", "language")
         )
       }

diff --git a/src/main/scala/ldbc/snb/datagen/transformation/transform/ExplodeEdges.scala b/src/main/scala/ldbc/snb/datagen/transformation/transform/ExplodeEdges.scala
@@ -1,9 +1,10 @@
 package ldbc.snb.datagen.transformation.transform
 
-import ldbc.snb.datagen.model.Cardinality.{NN, NOne, OneN}
-import ldbc.snb.datagen.model.EntityType.{Edge, Node}
+import ldbc.snb.datagen.model.Cardinality._
+import ldbc.snb.datagen.model.EntityType._
 import ldbc.snb.datagen.model.Mode
 import ldbc.snb.datagen.model.Mode.Raw.withRawColumns
+import ldbc.snb.datagen.model.raw._
 import ldbc.snb.datagen.syntax._
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.{Column, DataFrame}
@@ -25,51 +26,51 @@ object ExplodeEdges extends Transform[Mode.Raw.type, Mode.Raw.type] {
 
     val updatedEntities = entities
       .collect {
-        case (k @ Node("Organisation", true), v) =>
+        case (k @ OrganisationType, v) =>
           Map(
-            explodedEdge(Edge("IsLocatedIn", "Organisation", "Place", OneN, isStatic = true), v, $"LocationPlaceId"),
+            explodedEdge(Edge("IsLocatedIn", k, PlaceType, OneN, isStatic = true), v, $"LocationPlaceId"),
             k -> v.drop("LocationPlaceId")
           )
-        case (k @ Node("Place", true), v) =>
+        case (k @ PlaceType, v) =>
           Map(
-            explodedEdge(Edge("IsPartOf", "Place", "Place", OneN, isStatic = true), v, $"PartOfPlaceId"),
+            explodedEdge(Edge("IsPartOf", k, k, OneN, isStatic = true), v, $"PartOfPlaceId"),
             k -> v.drop("PartOfPlaceId")
           )
-        case (k @ Node("Tag", true), v) =>
+        case (k @ TagType, v) =>
           Map(
-            explodedEdge(Edge("HasType", "Tag", "TagClass", OneN, isStatic = true), v, $"TypeTagClassId"),
+            explodedEdge(Edge("HasType", k, TagClassType, OneN, isStatic = true), v, $"TypeTagClassId"),
             k -> v.drop("TypeTagClassId")
           )
-        case (k @ Node("TagClass", true), v) =>
+        case (k @ TagClassType, v) =>
           Map(
-            explodedEdge(Edge("IsSubclassOf", "TagClass", "TagClass", OneN, isStatic = true), v, $"SubclassOfTagClassId"),
+            explodedEdge(Edge("IsSubclassOf", k, k, OneN, isStatic = true), v, $"SubclassOfTagClassId"),
             k -> v.drop("SubclassOfTagClassId")
           )
-        case (k @ Node("Comment", false), v) =>
+        case (k @ CommentType, v) =>
           Map(
-            explodedEdge(Edge("HasCreator", "Comment", "Person", OneN), v, $"CreatorPersonId"),
-            explodedEdge(Edge("IsLocatedIn", "Comment", "Country", OneN), v, $"LocationCountryId"),
-            explodedEdge(Edge("ReplyOf", "Comment", "Post", OneN), v, $"ParentPostId"),
-            explodedEdge(Edge("ReplyOf", "Comment", "Comment", OneN), v, $"ParentCommentId"),
+            explodedEdge(Edge("HasCreator", k, PersonType, OneN), v, $"CreatorPersonId"),
+            explodedEdge(Edge("IsLocatedIn", k, PlaceType, OneN, destinationNameOverride = Some("Country")), v, $"LocationCountryId"),
+            explodedEdge(Edge("ReplyOf", k, PostType, OneN), v, $"ParentPostId"),
+            explodedEdge(Edge("ReplyOf", k, k, OneN), v, $"ParentCommentId"),
             k -> v.drop("CreatorPersonId", "LocationCountryId", "ParentPostId", "ParentCommentId")
           )
-        case (k @ Node("Forum", false), v) =>
+        case (k @ ForumType, v) =>
           Map(
-            explodedEdge(Edge("HasModerator", "Forum", "Person", OneN), v, $"ModeratorPersonId"),
+            explodedEdge(Edge("HasModerator", k, PersonType, OneN), v, $"ModeratorPersonId"),
             k -> v.drop("ModeratorPersonId")
           )
 
-        case (k @ Node("Person", false), v) =>
+        case (k @ PersonType, v) =>
           Map(
-            explodedEdge(Edge("IsLocatedIn", "Person", "City", OneN), v, $"LocationCityId"),
+            explodedEdge(Edge("IsLocatedIn", k, PlaceType, OneN, destinationNameOverride = Some("City")), v, $"LocationCityId"),
             k -> v.drop("LocationCityId")
           )
 
-        case (k @ Node("Post", false), v) =>
+        case (k @ PostType, v) =>
           Map(
-            explodedEdge(Edge("HasCreator", "Post", "Person", OneN), v, $"CreatorPersonId"),
-            explodedEdge(Edge("ContainerOf", "Forum", "Post", NOne), v, $"ContainerForumId"),
-            explodedEdge(Edge("IsLocatedIn", "Post", "Country", OneN), v, $"LocationCountryId"),
+            explodedEdge(Edge("HasCreator", k, PersonType, OneN), v, $"CreatorPersonId"),
+            explodedEdge(Edge("ContainerOf", ForumType, k, NOne), v, $"ContainerForumId"),
+            explodedEdge(Edge("IsLocatedIn", k, PlaceType, OneN, destinationNameOverride = Some("Country")), v, $"LocationCountryId"),
             k -> v.drop("CreatorPersonId", "LocationCountryId", "ContainerForumId")
           )
       }

diff --git a/src/main/scala/ldbc/snb/datagen/transformation/transform/RawToBiTransform.scala b/src/main/scala/ldbc/snb/datagen/transformation/transform/RawToBiTransform.scala
@@ -1,5 +1,7 @@
 package ldbc.snb.datagen.transformation.transform
 
+import ldbc.snb.datagen.model.Cardinality._
+import ldbc.snb.datagen.model.EntityType._
 import ldbc.snb.datagen.model.Mode.BI
 import ldbc.snb.datagen.model._
 import ldbc.snb.datagen.syntax._
@@ -24,6 +26,13 @@ case class RawToBiTransform(mode: BI, simulationStart: Long, simulationEnd: Long
     case _        => throw new IllegalStateException("Unrecognized partition key")
   }
 
+  private def notDerived(entityType: EntityType): Boolean = entityType match {
+    case Edge(_, _, _, OneN, _, _, _) => false
+    case Edge(_, _, _, NOne, _, _, _) => false
+    case Attr(_, _, _, _) => false
+    case _ => true
+  }
+
   override def transform(input: In): Out = {
     val batch_id = (col: Column) => date_format(date_trunc(mode.batchPeriod, to_timestamp(col / lit(1000L))), batchPeriodFormat(mode.batchPeriod))
 
@@ -64,7 +73,7 @@ case class RawToBiTransform(mode: BI, simulationStart: Long, simulationEnd: Long
           tpe -> BatchedEntity(
             RawToInteractiveTransform.snapshotPart(tpe, v, bulkLoadThreshold, filterDeletion = false),
             Some(Batched(insertBatchPart(tpe, v, bulkLoadThreshold, simulationEnd), Seq("batch_id"), Seq($"creationDate"))),
-            if (keepImplicitDeletes || v.columns.contains("explicitlyDeleted"))
+            if (notDerived(tpe) && (keepImplicitDeletes || v.columns.contains("explicitlyDeleted")))
               Some(Batched(deleteBatchPart(tpe, v, bulkLoadThreshold, simulationEnd), Seq("batch_id"), Seq($"deletionDate")))
             else
               None

diff --git a/src/main/scala/ldbc/snb/datagen/transformation/transform/RawToInteractiveTransform.scala b/src/main/scala/ldbc/snb/datagen/transformation/transform/RawToInteractiveTransform.scala
@@ -1,13 +1,14 @@
 package ldbc.snb.datagen.transformation.transform
 
-import ldbc.snb.datagen.model.Cardinality.NN
-import ldbc.snb.datagen.model.EntityType.Edge
+import ldbc.snb.datagen.model.Cardinality._
+import ldbc.snb.datagen.model.EntityType._
+import ldbc.snb.datagen.model.raw._
 import ldbc.snb.datagen.model.{EntityType, Graph, Mode}
 import ldbc.snb.datagen.syntax._
 import ldbc.snb.datagen.util.Logging
 import ldbc.snb.datagen.util.sql._
 import org.apache.spark.sql.DataFrame
-import org.apache.spark.sql.functions.{col, lit, to_timestamp}
+import org.apache.spark.sql.functions.{col, lit}
 
 case class RawToInteractiveTransform(mode: Mode.Interactive, simulationStart: Long, simulationEnd: Long)
     extends Transform[Mode.Raw.type, Mode.Interactive]
@@ -29,7 +30,7 @@ object RawToInteractiveTransform {
 
   def columns(tpe: EntityType, cols: Seq[String]) = tpe match {
     case tpe if tpe.isStatic => cols
-    case Edge("Knows", "Person", "Person", NN, false) =>
+    case Edge("Knows", PersonType, PersonType, NN, false, _, _) =>
       val rawCols = Set("deletionDate", "explicitlyDeleted", "weight")
       cols.filter(!rawCols.contains(_))
     case _ =>