Merge pull request #386 from ldbc/factors-q20

Add generation of 'sameUniversityKnows' factor table
ldbc · Apr 25, 2022 · 6e1c083 · 6e1c083
2 parents a5598cd + 986f9f0
commit 6e1c083
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/src/main/scala/ldbc/snb/datagen/factors/FactorGenerationStage.scala b/src/main/scala/ldbc/snb/datagen/factors/FactorGenerationStage.scala
@@ -421,6 +421,18 @@ object FactorGenerationStage extends DatagenStage with Logging {
         )
         .sort($"Person1Id", $"Person2Id")
         .limit(10000)
+    },
+    "sameUniversityKnows" -> Factor(PersonKnowsPersonType, PersonStudyAtUniversityType) {
+      case Seq(personKnowsPerson, studyAt) =>
+        undirectedKnows(personKnowsPerson)
+          .join(studyAt.as("studyAt1"), $"studyAt1.personId" === $"knows.person1Id")
+          .join(studyAt.as("studyAt2"), $"studyAt2.personId" === $"knows.person2Id")
+          .where($"studyAt1.universityId" === $"studyAt2.universityId")
+          .select(
+            $"knows.person1Id".as("person1Id"),
+            $"knows.person2Id".as("person2Id")
+          )
+          .distinct()
     }
   )
 }