1 files changed, 57 insertions, 11 deletions

diff --git a/icing/icing-search-engine_benchmark.cc b/icing/icing-search-engine_benchmark.cc
index fb44595..18c6bb9 100644
--- a/icing/icing-search-engine_benchmark.cc
+++ b/icing/icing-search-engine_benchmark.cc
@@ -37,6 +37,7 @@
 #include "icing/join/join-processor.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/initialize.pb.h"
+#include "icing/proto/persist.pb.h"
 #include "icing/proto/reset.pb.h"
 #include "icing/proto/schema.pb.h"
 #include "icing/proto/scoring.pb.h"
@@ -1116,6 +1117,8 @@ void BM_JoinQueryQualifiedId(benchmark::State& state) {
   IcingSearchEngineOptions options;
   options.set_base_dir(test_dir);
   options.set_index_merge_size(kIcingFullIndexSize);
+  options.set_document_store_namespace_id_fingerprint(true);
+  options.set_use_new_qualified_id_join_index(true);
   std::unique_ptr<IcingSearchEngine> icing =
       std::make_unique<IcingSearchEngine>(options);
 
@@ -1139,7 +1142,7 @@ void BM_JoinQueryQualifiedId(benchmark::State& state) {
   }
 
   // Create Email documents (child)
-  static constexpr int kNumEmailDocuments = 10000;
+  static constexpr int kNumEmailDocuments = 1000;
   std::uniform_int_distribution<> distrib(0, kNumPersonDocuments - 1);
   std::default_random_engine e(/*seed=*/12345);
   for (int i = 0; i < kNumEmailDocuments; ++i) {
@@ -1199,20 +1202,63 @@ void BM_JoinQueryQualifiedId(benchmark::State& state) {
         std::reduce(results.results().begin(), results.results().end(), 0,
                     child_count_reduce_func);
 
-    // Get all pages.
-    while (results.next_page_token() != kInvalidNextPageToken) {
-      results = icing->GetNextPage(results.next_page_token());
-      total_parent_count += results.results_size();
-      total_child_count +=
-          std::reduce(results.results().begin(), results.results().end(), 0,
-                      child_count_reduce_func);
+    ASSERT_THAT(total_parent_count, Eq(kNumPerPage));
+    ASSERT_THAT(total_child_count, ::testing::Ge(0));
+  }
+}
+BENCHMARK(BM_JoinQueryQualifiedId);
+
+void BM_PersistToDisk(benchmark::State& state) {
+  // Initialize the filesystem
+  std::string test_dir = GetTestTempDir() + "/icing/benchmark";
+  Filesystem filesystem;
+  DestructibleDirectory ddir(filesystem, test_dir);
+
+  // Create the schema.
+  std::default_random_engine random;
+  int num_types = kAvgNumNamespaces * kAvgNumTypes;
+  ExactStringPropertyGenerator property_generator;
+  SchemaGenerator<ExactStringPropertyGenerator> schema_generator(
+      /*num_properties=*/state.range(1), &property_generator);
+  SchemaProto schema = schema_generator.GenerateSchema(num_types);
+  EvenDistributionTypeSelector type_selector(schema);
+
+  // Generate documents.
+  int num_docs = state.range(0);
+  std::vector<std::string> language = CreateLanguages(kLanguageSize, &random);
+  const std::vector<DocumentProto> random_docs =
+      GenerateRandomDocuments(&type_selector, num_docs, language);
+
+  for (auto _ : state) {
+    state.PauseTiming();
+    // Create the index.
+    IcingSearchEngineOptions options;
+    options.set_base_dir(test_dir);
+    options.set_index_merge_size(kIcingFullIndexSize);
+    options.set_use_persistent_hash_map(true);
+    std::unique_ptr<IcingSearchEngine> icing =
+        std::make_unique<IcingSearchEngine>(options);
+
+    ASSERT_THAT(icing->Reset().status(), ProtoIsOk());
+    ASSERT_THAT(icing->SetSchema(schema).status(), ProtoIsOk());
+
+    for (const DocumentProto& doc : random_docs) {
+      ASSERT_THAT(icing->Put(doc).status(), ProtoIsOk());
     }
 
-    ASSERT_THAT(total_parent_count, Eq(kNumPersonDocuments));
-    ASSERT_THAT(total_child_count, Eq(kNumEmailDocuments));
+    state.ResumeTiming();
+
+    ASSERT_THAT(icing->PersistToDisk(PersistType::FULL).status(), ProtoIsOk());
+
+    state.PauseTiming();
+    icing.reset();
+    ASSERT_TRUE(filesystem.DeleteDirectoryRecursively(test_dir.c_str()));
+    state.ResumeTiming();
   }
 }
-BENCHMARK(BM_JoinQueryQualifiedId);
+BENCHMARK(BM_PersistToDisk)
+    // Arguments: num_indexed_documents, num_sections
+    ->ArgPair(1024, 5);
 
 }  // namespace