opt signIn

src/main/java/ldbc/finbench/datagen/entities/edges/CompanyOwnAccount.java

@@ -23,12 +23,12 @@ public CompanyOwnAccount(Company company, Account account, long creationDate, lo
     }
 
     public static void createCompanyOwnAccount(Company company, Account account, long creationDate) {
+        account.setOwnerType(PersonOrCompany.COMPANY);
+        account.setCompanyOwner(company);
         CompanyOwnAccount companyOwnAccount =
             new CompanyOwnAccount(company, account, creationDate, account.getDeletionDate(),
                                   account.isExplicitlyDeleted());
         company.getCompanyOwnAccounts().add(companyOwnAccount);
-        account.setOwnerType(PersonOrCompany.COMPANY);
-        account.setCompanyOwner(company);
     }
 
     public Company getCompany() {

src/main/java/ldbc/finbench/datagen/entities/edges/PersonApplyLoan.java

@@ -25,10 +25,10 @@ public PersonApplyLoan(Person person, Loan loan, long creationDate, long deletio
     }
 
     public static void createPersonApplyLoan(long creationDate, Person person, Loan loan, String organization) {
-        PersonApplyLoan personApplyLoan = new PersonApplyLoan(person, loan, creationDate, 0, false, organization);
-        person.getPersonApplyLoans().add(personApplyLoan);
         loan.setOwnerType(PersonOrCompany.PERSON);
         loan.setOwnerPerson(person);
+        PersonApplyLoan personApplyLoan = new PersonApplyLoan(person, loan, creationDate, 0, false, organization);
+        person.getPersonApplyLoans().add(personApplyLoan);
     }
 
     public Person getPerson() {

src/main/java/ldbc/finbench/datagen/entities/edges/PersonOwnAccount.java

@@ -26,12 +26,12 @@ public PersonOwnAccount(Person person, Account account, long creationDate, long
 
     public static void createPersonOwnAccount(Person person, Account account, long creationDate) {
         // Delete when account is deleted
+        account.setOwnerType(PersonOrCompany.PERSON);
+        account.setPersonOwner(person);
         PersonOwnAccount personOwnAccount = new PersonOwnAccount(person, account, creationDate,
                                                                  account.getDeletionDate(),
                                                                  account.isExplicitlyDeleted());
         person.getPersonOwnAccounts().add(personOwnAccount);
-        account.setOwnerType(PersonOrCompany.PERSON);
-        account.setPersonOwner(person);
     }
 
     public Person getPerson() {

src/main/java/ldbc/finbench/datagen/entities/edges/SignIn.java

@@ -1,7 +1,6 @@
 package ldbc.finbench.datagen.entities.edges;
 
 import java.io.Serializable;
-import java.util.Random;
 import ldbc.finbench.datagen.entities.DynamicActivity;
 import ldbc.finbench.datagen.entities.nodes.Account;
 import ldbc.finbench.datagen.entities.nodes.Medium;
@@ -15,7 +14,7 @@ public class SignIn implements DynamicActivity, Serializable {
     private long deletionDate;
     private int countryId;
     private int cityId;
-    private long multiplicityId;
+    private final long multiplicityId;
     private boolean isExplicitlyDeleted;
 
     public SignIn(Medium medium, Account account, int mid, long creationDate, long deletionDate,
@@ -37,7 +36,8 @@ public static SignIn createSignIn(int mid, RandomGeneratorFarm randomFarm, Mediu
         // Set country and city
         int countryId =
             Dictionaries.places.getCountryForPerson(randomFarm.get(RandomGeneratorFarm.Aspect.SIGNIN_COUNTRY));
-        signIn.setCountryId(countryId);
+        signIn.setCountryId(
+            Dictionaries.places.getCountryForPerson(randomFarm.get(RandomGeneratorFarm.Aspect.SIGNIN_COUNTRY)));
         signIn.setCityId(
             Dictionaries.places.getRandomCity(randomFarm.get(RandomGeneratorFarm.Aspect.SIGNIN_CITY), countryId));
 
@@ -85,10 +85,6 @@ public long getMultiplicityId() {
         return multiplicityId;
     }
 
-    public void setMultiplicityId(long multiplicityId) {
-        this.multiplicityId = multiplicityId;
-    }
-
     @Override
     public boolean isExplicitlyDeleted() {
         return isExplicitlyDeleted;
@@ -98,18 +94,10 @@ public void setExplicitlyDeleted(boolean explicitlyDeleted) {
         isExplicitlyDeleted = explicitlyDeleted;
     }
 
-    public int getCountryId() {
-        return countryId;
-    }
-
     public void setCountryId(int countryId) {
         this.countryId = countryId;
     }
 
-    public int getCityId() {
-        return cityId;
-    }
-
     public void setCityId(int cityId) {
         this.cityId = cityId;
     }

src/main/java/ldbc/finbench/datagen/generation/events/CompanyGuaranteeEvent.java

@@ -33,7 +33,7 @@ public List<Company> companyGuarantee(List<Company> companies, int blockId) {
         for (int i = 0; i < numCompaniesToTake; i++) {
             Company from = companies.get(pickCompanyRand.nextInt(companies.size()));
             int numGuarantees = numGuaranteesRand.nextInt(DatagenParams.maxTargetsToGuarantee);
-            for (int j = 0; j < numGuarantees; j++) {
+            for (int j = 0; j < Math.max(1,numGuarantees); j++) {
                 Company to = companies.get(randIndex.nextInt(companies.size()));
                 if (from.canGuarantee(to)) {
                     CompanyGuaranteeCompany.createCompanyGuaranteeCompany(dateRand, from, to);

src/main/java/ldbc/finbench/datagen/generation/events/CompanyRegisterEvent.java

@@ -25,13 +25,13 @@ public List<Company> companyRegister(List<Company> companies, AccountGenerator a
         resetState(blockId);
         accountGenerator.resetState(blockId);
         Random numAccRand = randomFarm.get(RandomGeneratorFarm.Aspect.NUM_ACCOUNTS_PER_COMPANY);
-        companies.forEach(company -> {
+        for (Company company : companies) {
             int numAccounts = numAccRand.nextInt(DatagenParams.maxAccountsPerOwner);
             for (int i = 0; i < Math.max(1, numAccounts); i++) {
                 Account account = accountGenerator.generateAccount(company.getCreationDate(), "company", blockId);
                 CompanyOwnAccount.createCompanyOwnAccount(company, account, account.getCreationDate());
             }
-        });
+        }
 
         return companies;
     }

src/main/java/ldbc/finbench/datagen/generation/events/PersonGuaranteeEvent.java

@@ -17,7 +17,6 @@ public PersonGuaranteeEvent() {
         randIndex = new Random(DatagenParams.defaultSeed);
     }
 
-
     private void resetState(int seed) {
         randomFarm.resetRandomGenerators(seed);
         randIndex.setSeed(seed);
@@ -28,12 +27,12 @@ public List<Person> personGuarantee(List<Person> persons, int blockId) {
 
         Random pickPersonRand = randomFarm.get(RandomGeneratorFarm.Aspect.PICK_PERSON_GUARANTEE);
         Random numGuaranteesRand = randomFarm.get(RandomGeneratorFarm.Aspect.NUM_GUARANTEES_PER_PERSON);
-        int numPersonsToTake = (int)(persons.size() * DatagenParams.personGuaranteeFraction);
+        int numPersonsToTake = (int) (persons.size() * DatagenParams.personGuaranteeFraction);
 
         for (int i = 0; i < numPersonsToTake; i++) {
             Person from = persons.get(pickPersonRand.nextInt(persons.size()));
-            int targetsToGuarantee = numGuaranteesRand.nextInt(DatagenParams.maxTargetsToGuarantee);
-            for (int j = 0; j < targetsToGuarantee; j++) {
+            int numGuarantees = numGuaranteesRand.nextInt(DatagenParams.maxTargetsToGuarantee);
+            for (int j = 0; j < Math.max(1, numGuarantees); j++) {
                 Person to = persons.get(randIndex.nextInt(persons.size()));
                 if (from.canGuarantee(to)) {
                     PersonGuaranteePerson.createPersonGuaranteePerson(randomFarm, from, to);

src/main/java/ldbc/finbench/datagen/generation/events/PersonRegisterEvent.java

@@ -2,6 +2,7 @@
 
 import java.io.Serializable;
 import java.util.List;
+import java.util.Random;
 import ldbc.finbench.datagen.entities.edges.PersonOwnAccount;
 import ldbc.finbench.datagen.entities.nodes.Account;
 import ldbc.finbench.datagen.entities.nodes.Person;
@@ -23,15 +24,15 @@ private void resetState(int seed) {
     public List<Person> personRegister(List<Person> persons, AccountGenerator accountGenerator, int blockId) {
         resetState(blockId);
         accountGenerator.resetState(blockId);
-        persons.forEach(person -> {
-            int numAccounts = randomFarm.get(RandomGeneratorFarm.Aspect.NUM_ACCOUNTS_PER_PERSON)
-                                        .nextInt(DatagenParams.maxAccountsPerOwner);
-            // Each person has at least one account
+
+        Random numAccRand = randomFarm.get(RandomGeneratorFarm.Aspect.NUM_ACCOUNTS_PER_PERSON);
+        for (Person person : persons) {
+            int numAccounts = numAccRand.nextInt(DatagenParams.maxAccountsPerOwner);
             for (int i = 0; i < Math.max(1, numAccounts); i++) {
                 Account account = accountGenerator.generateAccount(person.getCreationDate(), "person", blockId);
                 PersonOwnAccount.createPersonOwnAccount(person, account, account.getCreationDate());
             }
-        });
+        }
 
         return persons;
     }

src/main/java/ldbc/finbench/datagen/generation/events/SignInEvent.java

@@ -1,7 +1,7 @@
 package ldbc.finbench.datagen.generation.events;
 
 import java.io.Serializable;
-import java.util.ArrayList;
+import java.util.LinkedList;
 import java.util.List;
 import java.util.Random;
 import ldbc.finbench.datagen.entities.edges.SignIn;
@@ -13,42 +13,37 @@
 public class SignInEvent implements Serializable {
     private final RandomGeneratorFarm randomFarm;
     private final Random randIndex;
-    private final Random multiplicityRandom;
-    private final Random accountsToSignRandom;
 
     public SignInEvent() {
         randomFarm = new RandomGeneratorFarm();
         randIndex = new Random(DatagenParams.defaultSeed);
-        multiplicityRandom = new Random(DatagenParams.defaultSeed);
-        accountsToSignRandom = new Random(DatagenParams.defaultSeed);
     }
 
     private void resetState(int seed) {
         randomFarm.resetRandomGenerators(seed);
         randIndex.setSeed(seed);
-        multiplicityRandom.setSeed(seed);
-        accountsToSignRandom.setSeed(seed);
     }
 
-    public List<SignIn> signIn(List<Medium> media, List<Account> accounts, int blockId) {
+    public List<SignIn> signIn(Medium medium, List<Account> accounts, int blockId) {
         resetState(blockId);
-        List<SignIn> signIns = new ArrayList<>();
-        for (Medium medium : media) {
-            int numAccountsToSign = accountsToSignRandom.nextInt(DatagenParams.maxAccountToSignIn);
-            int signedCount = 0;
-            while (signedCount < numAccountsToSign) {
-                Account accountToSign = accounts.get(randIndex.nextInt(accounts.size()));
-                if (cannotSignIn(medium, accountToSign)) {
-                    continue;
-                }
-                int numSignIn = multiplicityRandom.nextInt(DatagenParams.maxSignInPerPair);
-                for (int mid = 0; mid < numSignIn; mid++) {
-                    SignIn signIn = SignIn.createSignIn(mid, randomFarm, medium, accountToSign);
-                    signIns.add(signIn);
-                }
-                signedCount++;
+
+        Random accountsToSignRand = randomFarm.get(RandomGeneratorFarm.Aspect.NUM_ACCOUNTS_SIGNIN_PER_MEDIUM);
+        Random multiplicityRandom = randomFarm.get(RandomGeneratorFarm.Aspect.MULTIPLICITY_SIGNIN);
+        int numAccountsToSign = accountsToSignRand.nextInt(DatagenParams.maxAccountToSignIn);
+
+        List<SignIn> signIns = new LinkedList<>();
+        for (int i = 0; i < Math.max(1, numAccountsToSign); i++) {
+            Account accountToSign = accounts.get(randIndex.nextInt(accounts.size()));
+            if (cannotSignIn(medium, accountToSign)) {
+                continue;
+            }
+            int numSignIn = multiplicityRandom.nextInt(DatagenParams.maxSignInPerPair);
+            for (int mid = 0; mid < Math.max(1, numSignIn); mid++) {
+                SignIn signIn = SignIn.createSignIn(mid, randomFarm, medium, accountToSign);
+                signIns.add(signIn);
             }
         }
+
         return signIns;
     }
 

src/main/java/ldbc/finbench/datagen/util/RandomGeneratorFarm.java

@@ -72,6 +72,8 @@ public enum Aspect {
         SIGNIN_DATE,
         SIGNIN_COUNTRY,
         SIGNIN_CITY,
+        NUM_ACCOUNTS_SIGNIN_PER_MEDIUM,
+        MULTIPLICITY_SIGNIN,
 
         // edge: person own account
         PERSON_OWN_ACCOUNT_DATE,

src/main/scala/ldbc/finbench/datagen/generation/generators/ActivityGenerator.scala

@@ -145,28 +145,30 @@ class ActivityGenerator()(implicit spark: SparkSession)
       mediumRDD: RDD[Medium],
       accountRDD: RDD[Account]
   ): RDD[SignIn] = {
-    val mediumParts = mediumRDD.partitions.length
-    val accountSampleList = new util.ArrayList[util.List[Account]](mediumParts)
-    for (i <- 1 to mediumParts) {
-      val sampleAccounts = accountRDD.sample(
+    val accountSampleList = accountRDD
+      .sample(
         withReplacement = false,
-        DatagenParams.accountSignedInFraction / mediumParts,
+        DatagenParams.accountSignedInFraction,
         sampleRandom.nextLong()
       )
-      accountSampleList.add(sampleAccounts.collect().toList.asJava)
-    }
+      .collect()
+      .grouped(accountRDD.partitions.length)
+      .map(_.toList.asJava)
+      .toList
+      .asJava
 
     val signInEvent = new SignInEvent
-    val signRels = mediumRDD.mapPartitions(mediums => {
-      val partitionId = TaskContext.getPartitionId()
-      val signInList = signInEvent.signIn(
-        mediums.toList.asJava,
-        accountSampleList.get(partitionId),
-        partitionId
-      )
-      for { signIn <- signInList.iterator().asScala } yield signIn
+    val signRels = mediumRDD.mapPartitionsWithIndex((index, mediums) => {
+      mediums.flatMap(medium => {
+        signInEvent
+          .signIn(
+            medium,
+            accountSampleList.get(index),
+            index
+          )
+          .asScala
+      })
     })
-
     signRels
   }