Minor updates to benchmarking

outr · Apr 6, 2024 · 7a5c41e · 7a5c41e
1 parent dd29029
commit 7a5c41e
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/benchmark/src/main/scala/benchmark/IMDBBenchmark.scala b/benchmark/src/main/scala/benchmark/IMDBBenchmark.scala
@@ -77,11 +77,11 @@ object IMDBBenchmark { // extends IOApp {
     val io = for {
       _ <- implementation.init()
       _ = scribe.info("--- Stage 1 ---")
-      akasFile <- downloadFile(new File(baseDirectory, "title.akas.tsv"), Limit.Unlimited).elapsed
+      akasFile <- downloadFile(new File(baseDirectory, "title.akas.tsv"), Limit.OneMillion).elapsed
       _ = scribe.info("--- Stage 2 ---")
       totalAka <- process(akasFile.value, implementation.map2TitleAka, implementation.persistTitleAka).elapsed
       _ = scribe.info("--- Stage 3 ---")
-      basicsFile <- downloadFile(new File(baseDirectory, "title.basics.tsv"), Limit.Unlimited).elapsed
+      basicsFile <- downloadFile(new File(baseDirectory, "title.basics.tsv"), Limit.OneMillion).elapsed
       _ = scribe.info("--- Stage 4 ---")
       totalBasics <- process(basicsFile.value, implementation.map2TitleBasics, implementation.persistTitleBasics).elapsed
       _ = scribe.info("--- Stage 5 ---")

diff --git a/benchmark/src/main/scala/benchmark/LightDBImplementation.scala b/benchmark/src/main/scala/benchmark/LightDBImplementation.scala
@@ -18,6 +18,8 @@ object LightDBImplementation extends BenchmarkImplementation {
 
   override def name: String = "LightDB"
 
+  override def init(): IO[Unit] = db.init(truncate = true)
+
   override def map2TitleAka(map: Map[String, String]): TitleAkaLDB = TitleAkaLDB(
     titleId = map.value("titleId"),
     ordering = map.int("ordering"),
@@ -55,7 +57,7 @@ object LightDBImplementation extends BenchmarkImplementation {
 
   override def get(id: String): IO[TitleAkaLDB] = db.titleAka.get(Id[TitleAkaLDB](id)).map(_.getOrElse(throw new RuntimeException(s"$id not found")))
 
-  override def findByTitleId(titleId: String): IO[List[TitleAkaLDB]] = db.titleAka.query.filter(TitleAkaLDB.titleId === titleId).search().compile.toList.flatMap(_.map(_.get()).sequence)
+  override def findByTitleId(titleId: String): IO[List[TitleAkaLDB]] = db.titleAka.query.filter(TitleAkaLDB.titleId === titleId).documentsStream().compile.toList
 
   override def flush(): IO[Unit] = for {
     _ <- db.titleAka.commit()