tmp

ldbc · Aug 22, 2024 · ad62be7 · ad62be7
1 parent b17eaf3
commit ad62be7
Showing 1 changed file with 10 additions and 6 deletions.
diff --git a/src/main/scala/ldbc/finbench/datagen/generation/generators/ActivityGenerator.scala b/src/main/scala/ldbc/finbench/datagen/generation/generators/ActivityGenerator.scala
@@ -94,17 +94,21 @@ class ActivityGenerator()(implicit spark: SparkSession)
         DatagenParams.companyInvestedFraction,
         sampleRandom.nextLong()
       )
-      .mapPartitionsWithIndex((index, targets) => {
+      .mapPartitionsWithIndex { (index, targets) =>
         personInvestEvent.resetState(index)
         companyInvestEvent.resetState(index)
         personInvestEvent
           .personInvestPartition(persons.value.asJava, targets.toList.asJava)
         companyInvestEvent
-          .companyInvestPartition(companies.value.asJava, targets.toList.asJava)
-        targets.map(target => target.scaleInvestmentRatios())
-      })
-
-    companyRDD
+          .companyInvestPartition(
+            companies.value.asJava,
+            targets.toList.asJava
+          )
+        targets.map { target =>
+          target.scaleInvestmentRatios()
+          target
+        }
+      }
   }
 
   def signInEvent(