Snap for 15194150 from 031c80a6b6dbf285a2f9c65bd299d239032f11b9 to androidx-javascriptengine-release Change-Id: I0ddfbc979a72662aea3721e72045570bbef74b78

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 4b7c752..22afb9c 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt

@@ -16,7 +16,10 @@
 
 project(icing)
 
-add_definitions("-DICING_REVERSE_JNI_SEGMENTATION=1")
+add_definitions(
+        "-DICING_REVERSE_JNI_SEGMENTATION=1"
+        "-DICING_DISABLE_EIGEN"
+)
 set(VERSION_SCRIPT "${CMAKE_CURRENT_SOURCE_DIR}/icing/jni.lds")
 set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_SHARED_LINKER_FLAGS

diff --git a/icing/absl_ports/canonical_errors.cc b/icing/absl_ports/canonical_errors.cc
index 1f43505..bf8b1f5 100644
--- a/icing/absl_ports/canonical_errors.cc
+++ b/icing/absl_ports/canonical_errors.cc

@@ -23,6 +23,93 @@
 namespace lib {
 namespace absl_ports {
 
+libtextclassifier3::Status CancelledError(const char* error_message) {
+  return libtextclassifier3::Status(libtextclassifier3::StatusCode::CANCELLED,
+                                    std::string(error_message));
+}
+
+libtextclassifier3::Status UnknownError(const char* error_message) {
+  return libtextclassifier3::Status(libtextclassifier3::StatusCode::UNKNOWN,
+                                    std::string(error_message));
+}
+
+libtextclassifier3::Status InvalidArgumentError(const char* error_message) {
+  return libtextclassifier3::Status(
+      libtextclassifier3::StatusCode::INVALID_ARGUMENT,
+      std::string(error_message));
+}
+
+libtextclassifier3::Status DeadlineExceededError(const char* error_message) {
+  return libtextclassifier3::Status(
+      libtextclassifier3::StatusCode::DEADLINE_EXCEEDED,
+      std::string(error_message));
+}
+
+libtextclassifier3::Status NotFoundError(const char* error_message) {
+  return libtextclassifier3::Status(libtextclassifier3::StatusCode::NOT_FOUND,
+                                    std::string(error_message));
+}
+
+libtextclassifier3::Status AlreadyExistsError(const char* error_message) {
+  return libtextclassifier3::Status(
+      libtextclassifier3::StatusCode::ALREADY_EXISTS, error_message);
+}
+
+libtextclassifier3::Status PermissionDeniedError(const char* error_message) {
+  return libtextclassifier3::Status(
+      libtextclassifier3::StatusCode::PERMISSION_DENIED,
+      std::string(error_message));
+}
+
+libtextclassifier3::Status ResourceExhaustedError(const char* error_message) {
+  return libtextclassifier3::Status(
+      libtextclassifier3::StatusCode::RESOURCE_EXHAUSTED,
+      std::string(error_message));
+}
+
+libtextclassifier3::Status FailedPreconditionError(const char* error_message) {
+  return libtextclassifier3::Status(
+      libtextclassifier3::StatusCode::FAILED_PRECONDITION,
+      std::string(error_message));
+}
+
+libtextclassifier3::Status AbortedError(const char* error_message) {
+  return libtextclassifier3::Status(libtextclassifier3::StatusCode::ABORTED,
+                                    std::string(error_message));
+}
+
+libtextclassifier3::Status OutOfRangeError(const char* error_message) {
+  return libtextclassifier3::Status(
+      libtextclassifier3::StatusCode::OUT_OF_RANGE, std::string(error_message));
+}
+
+libtextclassifier3::Status UnimplementedError(const char* error_message) {
+  return libtextclassifier3::Status(
+      libtextclassifier3::StatusCode::UNIMPLEMENTED,
+      std::string(error_message));
+}
+
+libtextclassifier3::Status InternalError(const char* error_message) {
+  return libtextclassifier3::Status(libtextclassifier3::StatusCode::INTERNAL,
+                                    std::string(error_message));
+}
+
+libtextclassifier3::Status UnavailableError(const char* error_message) {
+  return libtextclassifier3::Status(libtextclassifier3::StatusCode::UNAVAILABLE,
+                                    std::string(error_message));
+}
+
+libtextclassifier3::Status DataLossError(const char* error_message) {
+  return libtextclassifier3::Status(libtextclassifier3::StatusCode::DATA_LOSS,
+                                    std::string(error_message));
+}
+
+libtextclassifier3::Status UnauthenticatedError(const char* error_message) {
+  return libtextclassifier3::Status(
+      libtextclassifier3::StatusCode::UNAUTHENTICATED,
+      std::string(error_message));
+}
+
 libtextclassifier3::Status CancelledError(std::string error_message) {
   return libtextclassifier3::Status(libtextclassifier3::StatusCode::CANCELLED,
                                     std::move(error_message));

diff --git a/icing/absl_ports/canonical_errors.h b/icing/absl_ports/canonical_errors.h
index 2232946..55bdcea 100644
--- a/icing/absl_ports/canonical_errors.h
+++ b/icing/absl_ports/canonical_errors.h

@@ -23,6 +23,38 @@
 namespace lib {
 namespace absl_ports {
 
+// Overload both const char* and std::string to support 2 different callers:
+// `FooError(StrCat("text", a_str, " and ", b_str));` and
+// `FooError("simple text");`
+//
+// - std::string is used for the first call to avoid additional copies (Note:
+//   if using std::string_view, then the caller cannot move the string).
+// - const char* is used for the second call.
+//   - If const char* is not overloaded, then it is still valid for the second
+//     call to use std::string.
+//   - However, in Android the compiler will implicitly compile some conversion
+//     code (from const char* to std::string) **on the caller side** and
+//     increase most of the classes' size by ~1 KB, and the size of libicing.so
+//     will be increased by ~100 KBs.
+//   - From the experiment, overloading both const char* and std::string can
+//     avoid this binary size increase.
+libtextclassifier3::Status CancelledError(const char* error_message);
+libtextclassifier3::Status UnknownError(const char* error_message);
+libtextclassifier3::Status InvalidArgumentError(const char* error_message);
+libtextclassifier3::Status DeadlineExceededError(const char* error_message);
+libtextclassifier3::Status NotFoundError(const char* error_message);
+libtextclassifier3::Status AlreadyExistsError(const char* error_message);
+libtextclassifier3::Status PermissionDeniedError(const char* error_message);
+libtextclassifier3::Status ResourceExhaustedError(const char* error_message);
+libtextclassifier3::Status FailedPreconditionError(const char* error_message);
+libtextclassifier3::Status AbortedError(const char* error_message);
+libtextclassifier3::Status OutOfRangeError(const char* error_message);
+libtextclassifier3::Status UnimplementedError(const char* error_message);
+libtextclassifier3::Status InternalError(const char* error_message);
+libtextclassifier3::Status UnavailableError(const char* error_message);
+libtextclassifier3::Status DataLossError(const char* error_message);
+libtextclassifier3::Status UnauthenticatedError(const char* error_message);
+
 libtextclassifier3::Status CancelledError(std::string error_message);
 libtextclassifier3::Status UnknownError(std::string error_message);
 libtextclassifier3::Status InvalidArgumentError(std::string error_message);

diff --git a/icing/absl_ports/mutex.h b/icing/absl_ports/mutex.h
index 382a100..aacebd5 100644
--- a/icing/absl_ports/mutex.h
+++ b/icing/absl_ports/mutex.h

@@ -52,6 +52,9 @@
       : lock_(*mu) {}
   ~unique_lock() ICING_UNLOCK_FUNCTION() = default;
 
+  // For fine-grained locking if necessary.
+  void unlock() ICING_UNLOCK_FUNCTION() { lock_.unlock(); }
+  void lock() ICING_EXCLUSIVE_LOCK_FUNCTION() { lock_.lock(); }
  private:
   std::unique_lock<shared_mutex> lock_;
 };

diff --git a/icing/feature-flags.h b/icing/feature-flags.h
index 0231fe5..1c3c977 100644
--- a/icing/feature-flags.h
+++ b/icing/feature-flags.h

@@ -15,6 +15,8 @@
 #ifndef ICING_FEATURE_FLAGS_H_
 #define ICING_FEATURE_FLAGS_H_
 
+#include <cstdint>
+
 namespace icing {
 namespace lib {
 
@@ -26,7 +28,21 @@
                         bool enable_repeated_field_joins,
                         bool enable_embedding_backup_generation,
                         bool enable_schema_database,
-                        bool release_backup_schema_file_if_overlay_present)
+                        bool release_backup_schema_file_if_overlay_present,
+                        bool enable_strict_page_byte_size_limit,
+                        bool enable_smaller_decompression_buffer_size,
+                        bool enable_eigen_embedding_scoring,
+                        bool enable_passing_filter_to_children,
+                        bool enable_proto_log_new_header_format,
+                        bool enable_embedding_iterator_v2,
+                        bool enable_reusable_decompression_buffer,
+                        bool enable_schema_type_id_optimization,
+                        bool enable_optimize_improvements,
+                        int64_t expired_document_purge_threshold_ms,
+                        bool enable_non_existent_qualified_id_join,
+                        bool enable_skip_set_schema_type_equality_check,
+                        bool enable_embed_query_optimization,
+                        bool enable_schema_definition_deduping)
       : allow_circular_schema_definitions_(allow_circular_schema_definitions),
         enable_scorable_properties_(enable_scorable_properties),
         enable_embedding_quantization_(enable_embedding_quantization),
@@ -34,7 +50,26 @@
         enable_embedding_backup_generation_(enable_embedding_backup_generation),
         enable_schema_database_(enable_schema_database),
         release_backup_schema_file_if_overlay_present_(
-            release_backup_schema_file_if_overlay_present) {}
+            release_backup_schema_file_if_overlay_present),
+        enable_strict_page_byte_size_limit_(enable_strict_page_byte_size_limit),
+        enable_smaller_decompression_buffer_size_(
+            enable_smaller_decompression_buffer_size),
+        enable_eigen_embedding_scoring_(enable_eigen_embedding_scoring),
+        enable_passing_filter_to_children_(enable_passing_filter_to_children),
+        enable_proto_log_new_header_format_(enable_proto_log_new_header_format),
+        enable_embedding_iterator_v2_(enable_embedding_iterator_v2),
+        enable_reusable_decompression_buffer_(
+            enable_reusable_decompression_buffer),
+        enable_schema_type_id_optimization_(enable_schema_type_id_optimization),
+        enable_optimize_improvements_(enable_optimize_improvements),
+        expired_document_purge_threshold_ms_(
+            expired_document_purge_threshold_ms),
+        enable_non_existent_qualified_id_join_(
+            enable_non_existent_qualified_id_join),
+        enable_skip_set_schema_type_equality_check_(
+            enable_skip_set_schema_type_equality_check),
+        enable_embed_query_optimization_(enable_embed_query_optimization),
+        enable_schema_definition_deduping_(enable_schema_definition_deduping) {}
 
   bool allow_circular_schema_definitions() const {
     return allow_circular_schema_definitions_;
@@ -62,6 +97,62 @@
     return release_backup_schema_file_if_overlay_present_;
   }
 
+  bool enable_strict_page_byte_size_limit() const {
+    return enable_strict_page_byte_size_limit_;
+  }
+
+  bool enable_smaller_decompression_buffer_size() const {
+    return enable_smaller_decompression_buffer_size_;
+  }
+
+  bool enable_eigen_embedding_scoring() const {
+    return enable_eigen_embedding_scoring_;
+  }
+
+  bool enable_passing_filter_to_children() const {
+    return enable_passing_filter_to_children_;
+  }
+
+  bool enable_proto_log_new_header_format() const {
+    return enable_proto_log_new_header_format_;
+  }
+
+  bool enable_embedding_iterator_v2() const {
+    return enable_embedding_iterator_v2_;
+  }
+
+  bool enable_reusable_decompression_buffer() const {
+    return enable_reusable_decompression_buffer_;
+  }
+
+  bool enable_schema_type_id_optimization() const {
+    return enable_schema_type_id_optimization_;
+  }
+
+  bool enable_optimize_improvements() const {
+    return enable_optimize_improvements_;
+  }
+
+  int64_t expired_document_purge_threshold_ms() const {
+    return expired_document_purge_threshold_ms_;
+  }
+
+  bool enable_non_existent_qualified_id_join() const {
+    return enable_non_existent_qualified_id_join_;
+  }
+
+  bool enable_skip_set_schema_type_equality_check() const {
+    return enable_skip_set_schema_type_equality_check_;
+  }
+
+  bool enable_embed_query_optimization() const {
+    return enable_embed_query_optimization_;
+  }
+
+  bool enable_schema_definition_deduping() const {
+    return enable_schema_definition_deduping_;
+  }
+
  private:
   // Whether to allow circular references in the schema definition. This was
   // added in the Android U timeline and is not a trunk-stable flag.
@@ -83,6 +174,79 @@
   bool enable_schema_database_;
 
   bool release_backup_schema_file_if_overlay_present_;
+
+  // Whether to enable strict page byte size limit enforcement in
+  // ResultRetrieverV2.
+  bool enable_strict_page_byte_size_limit_;
+
+  bool enable_smaller_decompression_buffer_size_;
+
+  // Whether to enable the Eigen library for embedding scoring.
+  // If set to true **and** Eigen is compiled in (when ICING_DISABLE_EIGEN is
+  // not defined), Eigen will be used for embedding scoring.
+  bool enable_eigen_embedding_scoring_;
+
+  bool enable_passing_filter_to_children_;
+
+  // Whether to enable the new header format (refactor legacy format and
+  // introduce unsynced tail checksum) related changes in
+  // PortableFileBackedProtoLog.
+  bool enable_proto_log_new_header_format_;
+
+  bool enable_embedding_iterator_v2_;
+
+  // Whether PortableFileBackedProtoLog should retain a decompression buffer
+  // that reads can reuse rather than allocating a new one for each read.
+  bool enable_reusable_decompression_buffer_;
+
+  bool enable_schema_type_id_optimization_;
+
+  // Whether to enable a few minor improvements to Optimize:
+  // 1. Avoid unnecessary Status allocs for deleted/expired docs
+  // 2. Remove an unnecessary persist to disk call
+  bool enable_optimize_improvements_;
+
+  // The time threshold for an expired document to be purged.
+  // - Since we schedule a task to purge expired documents according to the next
+  //   expiration time of the documents, it is possible that some documents
+  //   expire within a small time window and the task executes too frequently.
+  // - Therefore, we use this flag to purge more documents that also expire in a
+  //   short period of time after the current time.
+  //
+  // For example, if the value is 1000 ms and the current time is 10000 ms:
+  // - All documents that are expired before 10000 ms will be purged, since they
+  //   are already expired.
+  // - Additionally, we will also purge documents that expire in the next 1000
+  //   ms, i.e. (10000, 11000] ms.
+  int64_t expired_document_purge_threshold_ms_;
+
+  // Whether to allow a document to reference a join parent (by its qualified
+  // id) that does not yet exist. When enabled, the join index will handle cases
+  // where a child document is indexed before its parent. The join relationship
+  // will be established once the parent document is indexed. If disabled, the
+  // join relationship will be lost if the child is indexed before the parent.
+  bool enable_non_existent_qualified_id_join_;
+
+  // Whether to skip the schema type equality check during SetSchema. We
+  // serialize the schema proto to strings when doing this check, which is slow.
+  //
+  // This will be set to true in AppSearch.
+  // - AppSearch already checks if schema types are unchanged for a new
+  //   SetSchema request, and skips the Icing interaction entirely if that is
+  //   the case.
+  // - Therefore, the Icing-side equality check in SetSchema is redundant and
+  //   can be skipped if caller is AppSearch.
+  bool enable_skip_set_schema_type_equality_check_;
+
+  // Whether to enable a query optimization that will rewrite embedding query
+  // iterators that are being AND'ed with other iterators such that those other
+  // iterators can be pushed down into the embedding iterator as a delegate.
+  // This allows us to avoid reading and scoring embeddings for documents that
+  // don't match the other requirements of the query.
+  bool enable_embed_query_optimization_;
+
+  // Whether to enable deduping for the schema's type definitions.
+  bool enable_schema_definition_deduping_;
 };
 
 }  // namespace lib

diff --git a/icing/file/file-backed-proto-log.h b/icing/file/file-backed-proto-log.h
index 0c06f77..985f50d 100644
--- a/icing/file/file-backed-proto-log.h
+++ b/icing/file/file-backed-proto-log.h

@@ -573,7 +573,11 @@
   // Deserialize proto
   ProtoT proto;
   if (header_->compress) {
-    protobuf_ports::GzipInputStream decompress_stream(&proto_stream);
+    std::unique_ptr<uint8_t[]> buffer =
+        std::make_unique<uint8_t[]>(protobuf_ports::kDefaultBufferSize);
+    protobuf_ports::GzipInputStream decompress_stream(
+        &proto_stream, protobuf_ports::GzipInputStream::AUTO, buffer.get(),
+        protobuf_ports::kDefaultBufferSize);
     proto.ParseFromZeroCopyStream(&decompress_stream);
   } else {
     proto.ParseFromZeroCopyStream(&proto_stream);

diff --git a/icing/file/file-backed-proto.h b/icing/file/file-backed-proto.h
index 65566c6..3504f96 100644
--- a/icing/file/file-backed-proto.h
+++ b/icing/file/file-backed-proto.h

@@ -196,20 +196,24 @@
                 << " of size: " << file_size;
 
   Header header;
-  if (!filesystem_->PRead(fd.get(), &header, sizeof(Header), /*offset=*/0)) {
+  if (filesystem_->PRead(fd.get(), &header, sizeof(Header), /*offset=*/0) !=
+      sizeof(Header)) {
     return absl_ports::InternalError(
         absl_ports::StrCat("Unable to read header of: ", file_path_));
   }
 
   if (header.magic != Header::kMagic) {
-    return absl_ports::InternalError(
-        absl_ports::StrCat("Invalid header kMagic for: ", file_path_));
+    ICING_LOG(ERROR) << "Invalid header magic for FileBackedProto "
+                     << file_path_ << ". Expected: " << Header::kMagic
+                     << ", actual: " << header.magic;
+    return absl_ports::InternalError(absl_ports::StrCat(
+        "Invalid header magic for FileBackedProto: ", file_path_));
   }
 
   int proto_size = file_size - sizeof(Header);
   auto buffer = std::make_unique<uint8_t[]>(proto_size);
-  if (!filesystem_->PRead(fd.get(), buffer.get(), proto_size,
-                          /*offset=*/sizeof(Header))) {
+  if (filesystem_->PRead(fd.get(), buffer.get(), proto_size,
+                         /*offset=*/sizeof(Header)) != proto_size) {
     return absl_ports::InternalError(
         absl_ports::StrCat("File read failed: ", file_path_));
   }

diff --git a/icing/file/file-backed-vector.h b/icing/file/file-backed-vector.h
index 8bd7161..04ee7c3 100644
--- a/icing/file/file-backed-vector.h
+++ b/icing/file/file-backed-vector.h

@@ -701,9 +701,12 @@
   // Make sure the header is still valid before we use any of its values. This
   // should technically be included in the header_checksum check below, but this
   // is a quick/fast check that can save us from an extra crc computation.
-  if (header->kMagic != FileBackedVector<T>::Header::kMagic) {
-    return absl_ports::InternalError(
-        absl_ports::StrCat("Invalid header kMagic for ", file_path));
+  if (header->magic != Header::kMagic) {
+    ICING_LOG(ERROR) << "Invalid header magic for FileBackedVector "
+                     << file_path << ". Expected: " << Header::kMagic
+                     << ", actual: " << header->magic;
+    return absl_ports::InternalError(absl_ports::StrCat(
+        "Invalid header magic for FileBackedVector: ", file_path));
   }
 
   // Check header

diff --git a/icing/file/file-backed-vector_test.cc b/icing/file/file-backed-vector_test.cc
index 4561e39..c7fbedc 100644
--- a/icing/file/file-backed-vector_test.cc
+++ b/icing/file/file-backed-vector_test.cc

@@ -1364,7 +1364,7 @@
   // to corruption of the header.
   FileBackedVector<char>::Header header;
   ASSERT_THAT(filesystem_.PRead(fd_, &header, sizeof(header), /*offset=*/0),
-              IsTrue());
+              Eq(sizeof(header)));
   header.num_elements = 1;
   ASSERT_THAT(filesystem_.PWrite(fd_, /*offset=*/0, &header, sizeof(header)),
               IsTrue());
@@ -1396,7 +1396,7 @@
   // of the underlying file.
   FileBackedVector<char>::Header header;
   ASSERT_THAT(filesystem_.PRead(fd_, &header, sizeof(header), /*offset=*/0),
-              IsTrue());
+              Eq(sizeof(header)));
   int64_t file_size = filesystem_.GetFileSize(fd_);
   int64_t allocated_elements_size = file_size - sizeof(header);
   header.num_elements = (allocated_elements_size / sizeof(char)) + 1;

diff --git a/icing/file/filesystem.cc b/icing/file/filesystem.cc
index cd905e7..22d1984 100644
--- a/icing/file/filesystem.cc
+++ b/icing/file/filesystem.cc

@@ -26,11 +26,14 @@
 
 #include <algorithm>
 #include <cerrno>
+#include <cstddef>
 #include <cstdint>
+#include <cstring>
+#include <memory>
 #include <unordered_set>
+#include <vector>
 
 #include "icing/absl_ports/str_cat.h"
-#include "icing/legacy/core/icing-string-util.h"
 #include "icing/util/logging.h"
 
 using std::vector;
@@ -125,8 +128,14 @@
                            const std::unordered_set<std::string>& exclude,
                            bool recursive, const char* prefix,
                            std::vector<std::string>* entries) {
-  DIR* dir = opendir(dir_name);
-  if (!dir) {
+  auto closer = [](DIR* dir) {
+    if (closedir(dir) != 0) {
+      ICING_LOG(ERROR) << "Error closing dir (" << errno << ") "
+                       << strerror(errno);
+    }
+  };
+  std::unique_ptr<DIR, decltype(closer)> dir(opendir(dir_name), closer);
+  if (dir == nullptr) {
     LogOpenError("Unable to open directory ", dir_name, ": ", errno);
     return false;
   }
@@ -136,7 +145,7 @@
   // may be statically allocated, so don't free it.
   dirent* p;
   // readdir's implementation seems to be thread safe.
-  while ((p = readdir(dir)) != nullptr) {
+  while ((p = readdir(dir.get())) != nullptr) {
     std::string file_name(p->d_name);
     if (file_name == "." || file_name == ".." ||
         exclude.find(file_name) != exclude.end()) {
@@ -155,9 +164,6 @@
       }
     }
   }
-  if (closedir(dir) != 0) {
-    ICING_LOG(ERROR) << "Error closing " << dir_name << " " << strerror(errno);
-  }
   return true;
 }
 
@@ -182,7 +188,8 @@
   ICING_VLOG(1) << "Deleting file " << file_name;
   int ret = unlink(file_name);
   if (ret != 0 && errno != ENOENT) {
-    ICING_LOG(ERROR) << "Deleting file " << file_name << " failed: " << strerror(errno);
+    ICING_LOG(ERROR) << "Deleting file " << file_name << " failed: (" << errno
+                     << ") " << strerror(errno);
     return false;
   }
   return true;
@@ -191,7 +198,8 @@
 bool Filesystem::DeleteDirectory(const char* dir_name) const {
   int ret = rmdir(dir_name);
   if (ret != 0 && errno != ENOENT) {
-    ICING_LOG(ERROR) << "Deleting directory " << dir_name << " failed: " << strerror(errno);
+    ICING_LOG(ERROR) << "Deleting directory " << dir_name << " failed: ("
+                     << errno << ") " << strerror(errno);
     return false;
   }
   return true;
@@ -204,7 +212,8 @@
     if (errno == ENOENT) {
       return true;  // If directory didn't exist, this was successful.
     }
-    ICING_LOG(ERROR) << "Stat " << dir_name << " failed: " << strerror(errno);
+    ICING_LOG(ERROR) << "Stat " << dir_name << " failed: (" << errno << ") "
+                     << strerror(errno);
     return false;
   }
   vector<std::string> entries;
@@ -217,7 +226,8 @@
        ++i) {
     std::string filename = std::string(dir_name) + '/' + *i;
     if (stat(filename.c_str(), &st) < 0) {
-      ICING_LOG(ERROR) << "Stat " << filename << " failed: " << strerror(errno);
+      ICING_LOG(ERROR) << "Stat " << filename << " failed: (" << errno << ") "
+                       << strerror(errno);
       success = false;
     } else if (S_ISDIR(st.st_mode)) {
       success = DeleteDirectoryRecursively(filename.c_str()) && success;
@@ -240,7 +250,8 @@
     exists = S_ISREG(st.st_mode) != 0;
   } else {
     if (errno != ENOENT) {
-      ICING_LOG(ERROR) << "Unable to stat file " << file_name << ": " << strerror(errno);
+      ICING_LOG(ERROR) << "Unable to stat file " << file_name << ": (" << errno
+                       << ") " << strerror(errno);
     }
     exists = false;
   }
@@ -254,7 +265,8 @@
     exists = S_ISDIR(st.st_mode) != 0;
   } else {
     if (errno != ENOENT) {
-      ICING_LOG(ERROR) << "Unable to stat directory " << dir_name << ": " << strerror(errno);
+      ICING_LOG(ERROR) << "Unable to stat directory " << dir_name << ": ("
+                       << errno << ") " << strerror(errno);
     }
     exists = false;
   }
@@ -365,11 +377,8 @@
 int64_t Filesystem::GetFileSize(int fd) const {
   struct stat st;
   if (fstat(fd, &st) < 0) {
-    if (errno == ENOENT) {
-      ICING_VLOG(1) << "Unable to stat file: " << strerror(errno);
-    } else {
-      ICING_LOG(WARNING) << "Unable to stat file: " << strerror(errno);
-    }
+      ICING_LOG(WARNING) << "Unable to stat file: (" << errno << ") "
+                         << strerror(errno);
     return kBadFileSize;
   }
   return st.st_size;
@@ -379,9 +388,11 @@
   struct stat st;
   if (stat(filename, &st) < 0) {
     if (errno == ENOENT) {
-      ICING_VLOG(1) << "Unable to stat file " << filename << ": " << strerror(errno);
+      ICING_VLOG(1) << "Unable to stat file " << filename << ": "
+          << strerror(errno);
     } else {
-      ICING_LOG(WARNING) << "Unable to stat file " << filename << ": " << strerror(errno);
+      ICING_LOG(WARNING) << "Unable to stat file " << filename << ": (" << errno
+                         << ") " << strerror(errno);
     }
     return kBadFileSize;
   }
@@ -390,7 +401,8 @@
 
 bool Filesystem::Truncate(int fd, int64_t new_size) const {
   if (ftruncate(fd, new_size) != 0) {
-    ICING_LOG(ERROR) << "Unable to truncate file: " << strerror(errno);
+    ICING_LOG(ERROR) << "Unable to truncate file: (" << errno << ") "
+                     << strerror(errno);
     return false;
   }
   lseek(fd, new_size, SEEK_SET);
@@ -409,7 +421,8 @@
 
 bool Filesystem::Grow(int fd, int64_t new_size) const {
   if (ftruncate(fd, new_size) != 0) {
-    ICING_LOG(ERROR) << "Unable to grow file: " << strerror(errno);
+    ICING_LOG(ERROR) << "Unable to grow file: (" << errno << ") "
+                     << strerror(errno);
     return false;
   }
 
@@ -431,10 +444,14 @@
   size_t write_len = data_size;
   do {
     // Don't try to write too much at once.
-    size_t chunk_size = std::min<size_t>(write_len, 64u * 1024);
-    ssize_t wrote = write(fd, data, chunk_size);
+#ifdef __APPLE__
+    // TEMP_FAILURE_RETRY is not defined in unistd.h on iOS.
+    ssize_t wrote = write(fd, data, write_len);
+#else  // __APPLE__
+    ssize_t wrote = TEMP_FAILURE_RETRY(write(fd, data, write_len));
+#endif // __APPLE__
     if (wrote < 0) {
-      ICING_LOG(ERROR) << "Bad write: " << strerror(errno);
+      ICING_LOG(ERROR) << "Bad write: (" << errno << ") " << strerror(errno);
       return false;
     }
     data = static_cast<const uint8_t*>(data) + wrote;
@@ -469,7 +486,7 @@
   }
   uint64_t size = GetFileSize(*src_fd);
   std::unique_ptr<uint8_t[]> buf = std::make_unique<uint8_t[]>(size);
-  if (!Read(*src_fd, buf.get(), size)) {
+  if (Read(*src_fd, buf.get(), size) != size) {
     return false;
   }
   return Write(*dst_fd, buf.get(), size);
@@ -477,15 +494,21 @@
 
 bool Filesystem::CopyDirectory(const char* src_dir, const char* dst_dir,
                                bool recursive) const {
-  DIR* dir = opendir(src_dir);
-  if (!dir) {
+  auto closer = [](DIR* dir) {
+    if (closedir(dir) != 0) {
+      ICING_LOG(ERROR) << "Error closing dir (" << errno << ") "
+                       << strerror(errno);
+    }
+  };
+  std::unique_ptr<DIR, decltype(closer)> dir(opendir(src_dir), closer);
+  if (dir == nullptr) {
     LogOpenError("Unable to open directory ", src_dir, ": ", errno);
     return false;
   }
 
   dirent* p;
   // readdir's implementation seems to be thread safe.
-  while ((p = readdir(dir)) != nullptr) {
+  while ((p = readdir(dir.get())) != nullptr) {
     std::string file_name(p->d_name);
     if (file_name == "." || file_name == "..") {
       continue;
@@ -511,9 +534,6 @@
       }
     }
   }
-  if (closedir(dir) != 0) {
-    ICING_LOG(ERROR) << "Error closing " << src_dir << ": " << strerror(errno);
-  }
   return true;
 }
 
@@ -521,11 +541,14 @@
                         size_t data_size) const {
   size_t write_len = data_size;
   do {
-    // Don't try to write too much at once.
-    size_t chunk_size = std::min<size_t>(write_len, 64u * 1024);
-    ssize_t wrote = pwrite(fd, data, chunk_size, offset);
+#ifdef __APPLE__
+    // TEMP_FAILURE_RETRY is not defined in unistd.h on iOS.
+    ssize_t wrote = pwrite(fd, data, write_len, offset);
+#else  // __APPLE__
+    ssize_t wrote = TEMP_FAILURE_RETRY(pwrite(fd, data, write_len, offset));
+#endif // __APPLE__
     if (wrote < 0) {
-      ICING_LOG(ERROR) << "Bad write: " << strerror(errno);
+      ICING_LOG(ERROR) << "Bad write: (" << errno << ") " << strerror(errno);
       return false;
     }
     data = static_cast<const uint8_t*>(data) + wrote;
@@ -547,45 +570,78 @@
   return success;
 }
 
-bool Filesystem::Read(int fd, void* buf, size_t buf_size) const {
-  ssize_t read_status = read(fd, buf, buf_size);
-  if (read_status < 0) {
-    ICING_LOG(ERROR) << "Bad read: " << strerror(errno);
-    return false;
+ssize_t Filesystem::Read(int fd, void* buf, size_t buf_size) const {
+  // convenience for pointer arithmetic below.
+  char* buf_ptr = static_cast<char*>(buf);
+  ssize_t processed_size = 0;
+  while (processed_size < buf_size) {
+    ssize_t read_status =
+        read(fd, buf_ptr + processed_size, buf_size - processed_size);
+    if (read_status < 0) {
+      ICING_LOG(ERROR) << "Bad read: (" << errno << ") " << strerror(errno);
+      return read_status;
+    }
+    if (read_status < buf_size - processed_size) {
+      ICING_LOG(ERROR) << "Read less than requested: " << read_status << " < "
+                       << buf_size - processed_size;
+    }
+    if (read_status == 0) {
+      // EOF. Finish reading.
+      return processed_size;
+    }
+    processed_size += read_status;
   }
-  return true;
+  return processed_size;
 }
 
-bool Filesystem::Read(const char* filename, void* buf, size_t buf_size) const {
+ssize_t Filesystem::Read(const char* filename, void* buf,
+                         size_t buf_size) const {
   int fd = OpenForRead(filename);
   if (fd == -1) {
-    return false;
+    return -1;
   }
 
-  bool success = Read(fd, buf, buf_size);
+  ssize_t bytes_read = Read(fd, buf, buf_size);
   close(fd);
-  return success;
+  return bytes_read;
 }
 
-bool Filesystem::PRead(int fd, void* buf, size_t buf_size, off_t offset) const {
-  ssize_t read_status = pread(fd, buf, buf_size, offset);
-  if (read_status < 0) {
-    ICING_LOG(ERROR) << "Bad read: " << strerror(errno);
-    return false;
+ssize_t Filesystem::PRead(int fd, void* buf, size_t buf_size,
+                          off_t offset) const {
+  // convenience for pointer arithmetic below.
+  char* buf_ptr = static_cast<char*>(buf);
+  size_t processed_size = 0;
+  while (processed_size < buf_size) {
+    ssize_t read_status =
+        pread(fd, buf_ptr + processed_size, buf_size - processed_size,
+              offset + processed_size);
+    if (read_status < 0) {
+      ICING_LOG(ERROR) << "Bad read: (" << errno << ") " << strerror(errno);
+      return read_status;
+    }
+    if (read_status < buf_size - processed_size) {
+      ICING_LOG(ERROR) << "Read less than requested: " << read_status << " < "
+                       << buf_size - processed_size;
+    }
+    if (read_status == 0) {
+      // EOF. Finish reading.
+      return processed_size;
+    }
+    processed_size += read_status;
   }
-  return true;
+  return processed_size;
 }
 
-bool Filesystem::PRead(const char* filename, void* buf, size_t buf_size,
-                       off_t offset) const {
+ssize_t Filesystem::PRead(const char* filename, void* buf, size_t buf_size,
+                          off_t offset) const {
   int fd = OpenForRead(filename);
   if (fd == -1) {
-    return false;
+    return -1;
   }
 
-  bool success = PRead(fd, buf, buf_size, offset);
+  ssize_t bytes_read = PRead(fd, buf, buf_size, offset);
   close(fd);
-  return success;
+  return bytes_read;
 }
 
 bool Filesystem::DataSync(int fd) const {
@@ -596,7 +652,8 @@
 #endif
 
   if (result < 0) {
-    ICING_LOG(ERROR) << "Unable to sync data: " << strerror(errno);
+    ICING_LOG(ERROR) << "Unable to sync data: (" << errno << ") "
+                     << strerror(errno);
     return false;
   }
   return true;
@@ -604,7 +661,8 @@
 
 bool Filesystem::RenameFile(const char* old_name, const char* new_name) const {
   if (rename(old_name, new_name) < 0) {
-    ICING_LOG(ERROR) << "Unable to rename file " << old_name << " to " << new_name << ": " << strerror(errno);
+    ICING_LOG(ERROR) << "Unable to rename file " << old_name << " to "
+                     << new_name << ": (" << errno << ") " << strerror(errno);
     return false;
   }
   return true;
@@ -642,7 +700,8 @@
     if (mkdir(dir_name, S_IRUSR | S_IWUSR | S_IXUSR) == 0) {
       success = true;
     } else {
-      ICING_LOG(ERROR) << "Creating directory " << dir_name << " failed: " << strerror(errno);
+      ICING_LOG(ERROR) << "Creating directory " << dir_name << " failed: ("
+                       << errno << ") " << strerror(errno);
     }
   }
   return success;
@@ -662,7 +721,8 @@
 int64_t Filesystem::GetDiskUsage(int fd) const {
   struct stat st;
   if (fstat(fd, &st) < 0) {
-    ICING_LOG(ERROR) << "Unable to stat file: " << strerror(errno);
+    ICING_LOG(ERROR) << "Unable to stat file: (" << errno << ") "
+                     << strerror(errno);
     return kBadFileSize;
   }
   return st.st_blocks * kStatBlockSize;
@@ -671,7 +731,8 @@
 int64_t Filesystem::GetFileDiskUsage(const char* path) const {
   struct stat st;
   if (stat(path, &st) != 0) {
-    ICING_LOG(ERROR) << "Unable to stat " << path << ": " << strerror(errno);
+    ICING_LOG(ERROR) << "Unable to stat " << path << ": (" << errno << ") "
+                     << strerror(errno);
     return kBadFileSize;
   }
   return st.st_blocks * kStatBlockSize;
@@ -680,7 +741,8 @@
 int64_t Filesystem::GetDiskUsage(const char* path) const {
   struct stat st;
   if (stat(path, &st) != 0) {
-    ICING_LOG(ERROR) << "Unable to stat " << path << ": " << strerror(errno);
+    ICING_LOG(ERROR) << "Unable to stat " << path << ": (" << errno << ") "
+                     << strerror(errno);
     return kBadFileSize;
   }
   int64_t result = st.st_blocks * kStatBlockSize;

diff --git a/icing/file/filesystem.h b/icing/file/filesystem.h
index dd2c5d1..71eba0b 100644
--- a/icing/file/filesystem.h
+++ b/icing/file/filesystem.h

@@ -183,11 +183,11 @@
   // Reads from a file. Returns true if data was successfully read out. If the
   // file is seekable, read starts at the file offset, and the file offset is
   // incremented by number of bytes read.
-  virtual bool Read(int fd, void* buf, size_t buf_size) const;
-  virtual bool Read(const char* filename, void* buf, size_t buf_size) const;
-  virtual bool PRead(int fd, void* buf, size_t buf_size, off_t offset) const;
-  virtual bool PRead(const char* filename, void* buf, size_t buf_size,
-                     off_t offset) const;
+  virtual ssize_t Read(int fd, void* buf, size_t buf_size) const;
+  virtual ssize_t Read(const char* filename, void* buf, size_t buf_size) const;
+  virtual ssize_t PRead(int fd, void* buf, size_t buf_size, off_t offset) const;
+  virtual ssize_t PRead(const char* filename, void* buf, size_t buf_size,
+                        off_t offset) const;
 
   // Syncs the file to disk (fdatasync). Returns true on success.
   virtual bool DataSync(int fd) const;

diff --git a/icing/file/filesystem_test.cc b/icing/file/filesystem_test.cc
index 214180e..378ef98 100644
--- a/icing/file/filesystem_test.cc
+++ b/icing/file/filesystem_test.cc

@@ -17,8 +17,11 @@
 #include "icing/file/filesystem.h"
 
 #include <algorithm>
+#include <chrono>
 #include <cstdint>
+#include <cstring>
 #include <string>
+#include <thread>
 #include <unordered_set>
 #include <vector>
 
@@ -419,38 +422,89 @@
   EXPECT_TRUE(filesystem.Write(fd, data.c_str(), strlen(data.c_str())));
 
   std::string hello;
-  hello.resize(strlen("hello"));
-  EXPECT_TRUE(filesystem.Read(foo_file.c_str(), &hello[0], strlen("hello")));
+  size_t hello_len = strlen("hello");
+  hello.resize(hello_len);
+  EXPECT_THAT(filesystem.Read(foo_file.c_str(), &hello[0], hello_len),
+              Eq(hello_len));
   EXPECT_THAT(hello, Eq("hello"));
 
   // Read starts from wherever file offset is at the moment.
   filesystem.SetPosition(fd, 0);
   hello.clear();
-  hello.resize(strlen("hello"));
-  EXPECT_TRUE(filesystem.Read(fd, &hello[0], strlen("hello")));
+  hello.resize(hello_len);
+  EXPECT_THAT(filesystem.Read(fd, &hello[0], hello_len), Eq(hello_len));
   EXPECT_THAT(hello, Eq("hello"));
 
   // Shouldn't need to move file offset anymore since file offset gets updated
   // after the read.
   std::string world;
-  world.resize(strlen(" world"));
-  EXPECT_TRUE(filesystem.Read(fd, &world[0], strlen(" world")));
+  size_t world_len = strlen(" world");
+  world.resize(world_len);
+  EXPECT_THAT(filesystem.Read(fd, &world[0], world_len), Eq(world_len));
   EXPECT_THAT(world, Eq(" world"));
 
   // PRead should not be dependent on the file offset
   world.clear();
-  world.resize(strlen(" world"));
-  EXPECT_TRUE(
-      filesystem.PRead(fd, &world[0], strlen(" world"), strlen("hello")));
+  world.resize(world_len);
+  EXPECT_THAT(filesystem.PRead(fd, &world[0], world_len, hello_len),
+              Eq(world_len));
   EXPECT_THAT(world, Eq(" world"));
 
   hello.clear();
-  hello.resize(strlen("hello"));
-  EXPECT_TRUE(
-      filesystem.PRead(foo_file.c_str(), &hello[0], strlen("hello"), 0));
+  hello.resize(hello_len);
+  EXPECT_THAT(filesystem.PRead(foo_file.c_str(), &hello[0], hello_len, 0),
+              Eq(hello_len));
   EXPECT_THAT(hello, Eq("hello"));
 }
 
+TEST_F(FilesystemTest, ReadInChunks) {
+  int pipe_fd[2];
+  ASSERT_THAT(pipe(pipe_fd), Eq(0));
+  ScopedFd read_fd(pipe_fd[0]);
+  ScopedFd write_fd(pipe_fd[1]);
+
+  std::string read_data(200, 'a');
+  bool read_success = false;
+  Filesystem filesystem;
+  auto read_callable = [&]() {
+    read_success = filesystem.Read(read_fd.get(), &read_data[0],
+                                   read_data.size()) == read_data.size();
+  };
+
+  std::string write_chunks(50, 'b');
+  auto write_callable = [&]() {
+    for (int i = 0; i < 4; ++i) {
+      filesystem.Write(write_fd.get(), &write_chunks[0], write_chunks.size());
+      std::this_thread::sleep_for(std::chrono::milliseconds(1));
+    }
+  };
+
+  std::thread read_thread(read_callable);
+  std::thread write_thread(write_callable);
+  read_thread.join();
+  write_thread.join();
+  EXPECT_TRUE(read_success);
+  std::string expected_data(200, 'b');
+  EXPECT_THAT(read_data, Eq(expected_data));
+}
+
+TEST_F(FilesystemTest, ReadFileSmallerThanBufferSize) {
+  Filesystem filesystem;
+  const std::string foo_file = temp_dir_ + "/foo_file";
+  ScopedFd fd(filesystem.OpenForWrite(foo_file.c_str()));;
+  ASSERT_TRUE(fd.is_valid());
+
+  std::string write_buf(100, 'b');
+  ASSERT_TRUE(filesystem.Write(fd.get(), &write_buf[0], write_buf.length()));
+
+  std::string read_buf(200, 'a');
+  EXPECT_THAT(filesystem.PRead(fd.get(), &read_buf[0], read_buf.length(), 0),
+              Eq(write_buf.length()));
+
+  std::string expected_read_buf = std::string(100, 'b') + std::string(100, 'a');
+  EXPECT_THAT(read_buf, Eq(expected_read_buf));
+}
+
 TEST_F(FilesystemTest, CopyDirectory) {
   Filesystem filesystem;
 

diff --git a/icing/file/marker-file_test.cc b/icing/file/marker-file_test.cc
index cd7d655..ba51769 100644
--- a/icing/file/marker-file_test.cc
+++ b/icing/file/marker-file_test.cc

@@ -123,8 +123,7 @@
   EXPECT_THAT(filesystem_.FileExists(file_path.c_str()), IsTrue());
   EXPECT_THAT(filesystem_.GetFileSize(file_path.c_str()), Eq(4));
   char buf[4];
-  EXPECT_THAT(filesystem_.Read(file_path.c_str(), buf, /*buf_size=*/4),
-              IsTrue());
+  EXPECT_THAT(filesystem_.Read(file_path.c_str(), buf, /*buf_size=*/4), Eq(4));
   EXPECT_THAT(std::string(buf, 4), Eq("test"));
 }
 

diff --git a/icing/file/memory-mapped-file-backed-proto-log.h b/icing/file/memory-mapped-file-backed-proto-log.h
index cc8f35f..ab0ddcc 100644
--- a/icing/file/memory-mapped-file-backed-proto-log.h
+++ b/icing/file/memory-mapped-file-backed-proto-log.h

@@ -31,6 +31,7 @@
 #include "icing/file/memory-mapped-file.h"
 #include "icing/legacy/core/icing-string-util.h"
 #include "icing/util/crc32.h"
+#include "icing/util/logging.h"
 #include "icing/util/status-macros.h"
 
 namespace icing {
@@ -159,6 +160,9 @@
     ProtoMetadata proto_metadata) {
   uint8_t magic_number = proto_metadata >> 24;
   if (magic_number != kProtoMagic) {
+    ICING_LOG(ERROR)
+        << "Invalid header magic for MemoryMappedFileBackedProtoLog. Expected: "
+        << kProtoMagic << ", actual: " << magic_number;
     return absl_ports::InvalidArgumentError(
         "Proto metadata has invalid magic number");
   }

diff --git a/icing/file/memory-mapped-file_test.cc b/icing/file/memory-mapped-file_test.cc
index 16f76e6..8c3891d 100644
--- a/icing/file/memory-mapped-file_test.cc
+++ b/icing/file/memory-mapped-file_test.cc

@@ -161,8 +161,9 @@
     ASSERT_TRUE(sfd.is_valid());
     int buf_size = 10;
     auto buf = std::make_unique<char[]>(buf_size);
-    ASSERT_TRUE(filesystem_.PRead(sfd.get(), buf.get(), buf_size,
-                                  pre_mapping_file_offset));
+    ASSERT_THAT(filesystem_.PRead(sfd.get(), buf.get(), buf_size,
+                                  pre_mapping_file_offset),
+                Eq(buf_size));
     EXPECT_THAT(buf.get()[0], Eq('a'));
   }
 }
@@ -300,8 +301,9 @@
   ASSERT_TRUE(sfd.is_valid());
   int buf_size = 1;
   auto buf = std::make_unique<char[]>(buf_size);
-  ASSERT_TRUE(filesystem_.PRead(sfd.get(), buf.get(), buf_size,
-                                pre_mapping_file_offset));
+  ASSERT_THAT(filesystem_.PRead(sfd.get(), buf.get(), buf_size,
+                                pre_mapping_file_offset),
+              Eq(buf_size));
   EXPECT_THAT(buf.get()[0], Eq('a'));
 }
 

diff --git a/icing/file/mock-filesystem.h b/icing/file/mock-filesystem.h
index 32817d4..02740f6 100644
--- a/icing/file/mock-filesystem.h
+++ b/icing/file/mock-filesystem.h

@@ -299,15 +299,15 @@
                size_t data_size),
               (const));
 
-  MOCK_METHOD(bool, Read, (int fd, void* buf, size_t buf_size), (const));
+  MOCK_METHOD(ssize_t, Read, (int fd, void* buf, size_t buf_size), (const));
 
-  MOCK_METHOD(bool, Read, (const char* filename, void* buf, size_t buf_size),
+  MOCK_METHOD(ssize_t, Read, (const char* filename, void* buf, size_t buf_size),
               (const));
 
-  MOCK_METHOD(bool, PRead, (int fd, void* buf, size_t buf_size, off_t offset),
-              (const));
+  MOCK_METHOD(ssize_t, PRead,
+              (int fd, void* buf, size_t buf_size, off_t offset), (const));
 
-  MOCK_METHOD(bool, PRead,
+  MOCK_METHOD(ssize_t, PRead,
               (const char* filename, void* buf, size_t buf_size, off_t offset),
               (const));
 

diff --git a/icing/file/persistent-hash-map.cc b/icing/file/persistent-hash-map.cc
index 148c7ea..3daeadd 100644
--- a/icing/file/persistent-hash-map.cc
+++ b/icing/file/persistent-hash-map.cc

@@ -29,6 +29,7 @@
 #include "icing/file/memory-mapped-file.h"
 #include "icing/file/persistent-storage.h"
 #include "icing/util/crc32.h"
+#include "icing/util/logging.h"
 #include "icing/util/status-macros.h"
 
 namespace icing {
@@ -504,8 +505,13 @@
 
   // Magic should be the same.
   if (persistent_hash_map->info().magic != Info::kMagic) {
+    ICING_LOG(ERROR) << "Invalid header magic for PersistentHashMap "
+                     << persistent_hash_map->working_path_
+                     << ". Expected: " << Info::kMagic
+                     << ", actual: " << persistent_hash_map->info().magic;
     return absl_ports::FailedPreconditionError(
-        "PersistentHashMap header magic mismatch");
+        absl_ports::StrCat("Invalid header magic for PersistentHashMap: ",
+                           persistent_hash_map->working_path_));
   }
 
   // Value type size should be consistent.

diff --git a/icing/file/persistent-hash-map_test.cc b/icing/file/persistent-hash-map_test.cc
index d2e211d..6accfbb 100644
--- a/icing/file/persistent-hash-map_test.cc
+++ b/icing/file/persistent-hash-map_test.cc

@@ -235,8 +235,9 @@
 
   // Check info section
   Info info;
-  ASSERT_TRUE(filesystem_.PRead(metadata_sfd.get(), &info, sizeof(Info),
-                                PersistentHashMap::kInfoMetadataFileOffset));
+  ASSERT_THAT(filesystem_.PRead(metadata_sfd.get(), &info, sizeof(Info),
+                                PersistentHashMap::kInfoMetadataFileOffset),
+              Eq(sizeof(Info)));
   EXPECT_THAT(info.magic, Eq(Info::kMagic));
   EXPECT_THAT(info.value_type_size, Eq(sizeof(int)));
   EXPECT_THAT(info.max_load_factor_percent,
@@ -246,8 +247,9 @@
 
   // Check crcs section
   Crcs crcs;
-  ASSERT_TRUE(filesystem_.PRead(metadata_sfd.get(), &crcs, sizeof(Crcs),
-                                PersistentHashMap::kCrcsMetadataFileOffset));
+  ASSERT_THAT(filesystem_.PRead(metadata_sfd.get(), &crcs, sizeof(Crcs),
+                                PersistentHashMap::kCrcsMetadataFileOffset),
+              Eq(sizeof(Crcs)));
   // # of elements in bucket_storage should be 1, so it should have non-zero
   // all storages crc value.
   EXPECT_THAT(crcs.component_crcs.storages_crc, Ne(0));
@@ -487,12 +489,14 @@
     ASSERT_TRUE(metadata_sfd.is_valid());
 
     Crcs crcs;
-    ASSERT_TRUE(filesystem_.PRead(metadata_sfd.get(), &crcs, sizeof(Crcs),
-                                  PersistentHashMap::kCrcsMetadataFileOffset));
+    ASSERT_THAT(filesystem_.PRead(metadata_sfd.get(), &crcs, sizeof(Crcs),
+                                  PersistentHashMap::kCrcsMetadataFileOffset),
+                Eq(sizeof(Crcs)));
 
     Info info;
-    ASSERT_TRUE(filesystem_.PRead(metadata_sfd.get(), &info, sizeof(Info),
-                                  PersistentHashMap::kInfoMetadataFileOffset));
+    ASSERT_THAT(filesystem_.PRead(metadata_sfd.get(), &info, sizeof(Info),
+                                  PersistentHashMap::kInfoMetadataFileOffset),
+                Eq(sizeof(Info)));
 
     // Manually change magic and update checksums.
     info.magic += kCorruptedValueOffset;
@@ -515,7 +519,7 @@
     EXPECT_THAT(persistent_hash_map_or,
                 StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
     EXPECT_THAT(persistent_hash_map_or.status().error_message(),
-                HasSubstr("PersistentHashMap header magic mismatch"));
+                HasSubstr("Invalid header magic for PersistentHashMap"));
   }
 }
 
@@ -612,8 +616,9 @@
   ASSERT_TRUE(metadata_sfd.is_valid());
 
   Crcs crcs;
-  ASSERT_TRUE(filesystem_.PRead(metadata_sfd.get(), &crcs, sizeof(Crcs),
-                                PersistentHashMap::kCrcsMetadataFileOffset));
+  ASSERT_THAT(filesystem_.PRead(metadata_sfd.get(), &crcs, sizeof(Crcs),
+                                PersistentHashMap::kCrcsMetadataFileOffset),
+              Eq(sizeof(Crcs)));
 
   // Manually corrupt all_crc
   crcs.all_crc += kCorruptedValueOffset;
@@ -656,8 +661,9 @@
   ASSERT_TRUE(metadata_sfd.is_valid());
 
   Info info;
-  ASSERT_TRUE(filesystem_.PRead(metadata_sfd.get(), &info, sizeof(Info),
-                                PersistentHashMap::kInfoMetadataFileOffset));
+  ASSERT_THAT(filesystem_.PRead(metadata_sfd.get(), &info, sizeof(Info),
+                                PersistentHashMap::kInfoMetadataFileOffset),
+              Eq(sizeof(Info)));
 
   // Modify info, but don't update the checksum. This would be similar to
   // corruption of info.
@@ -863,8 +869,9 @@
   ASSERT_TRUE(metadata_sfd.is_valid());
 
   Info info;
-  ASSERT_TRUE(filesystem_.PRead(metadata_sfd.get(), &info, sizeof(Info),
-                                PersistentHashMap::kInfoMetadataFileOffset));
+  ASSERT_THAT(filesystem_.PRead(metadata_sfd.get(), &info, sizeof(Info),
+                                PersistentHashMap::kInfoMetadataFileOffset),
+              Eq(sizeof(Info)));
   EXPECT_THAT(info.max_load_factor_percent,
               Eq(options.max_load_factor_percent));
 

diff --git a/icing/file/portable-file-backed-proto-log.h b/icing/file/portable-file-backed-proto-log.h
index b498563..6188171 100644
--- a/icing/file/portable-file-backed-proto-log.h
+++ b/icing/file/portable-file-backed-proto-log.h

@@ -53,6 +53,7 @@
 #ifndef ICING_FILE_PORTABLE_FILE_BACKED_PROTO_LOG_H_
 #define ICING_FILE_PORTABLE_FILE_BACKED_PROTO_LOG_H_
 
+#include <algorithm>
 #include <cstddef>
 #include <cstdint>
 #include <cstring>
@@ -60,11 +61,11 @@
 #include <string>
 #include <string_view>
 #include <utility>
-#include <vector>
 
 #include "icing/text_classifier/lib3/utils/base/status.h"
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "icing/absl_ports/canonical_errors.h"
+#include "icing/absl_ports/mutex.h"
 #include "icing/absl_ports/str_cat.h"
 #include "icing/file/constants.h"
 #include "icing/file/filesystem.h"
@@ -73,8 +74,9 @@
 #include "icing/portable/endian.h"
 #include "icing/portable/gzip_stream.h"
 #include "icing/portable/platform.h"
-#include "icing/portable/zlib.h"
+#include "icing/proto/persist.pb.h"
 #include "icing/util/bit-util.h"
+#include "icing/util/clock.h"
 #include "icing/util/crc32.h"
 #include "icing/util/data-loss.h"
 #include "icing/util/logging.h"
@@ -111,20 +113,56 @@
     // BEST_COMPRESSION = 9
     const int32_t compression_level;
 
+    // The threshold in bytes for compression if enabled. If the proto is larger
+    // than or equal to this threshold, it will be compressed.
+    const uint32_t compression_threshold_bytes;
+
+    // Level of memory usage for compression if enabled, BEST_MEMORY = 1,
+    // BEST_COMPRESSION and SPEED = 9
+    const int32_t compression_mem_level;
+
+    // Whether to use a smaller decompression buffer size. If false, the
+    // decompression buffer size will be the default size of 64MiB.
+    const bool enable_smaller_decompression_buffer_size;
+
+    // Whether to enable the new header format and unsynced tail checksum
+    // related changes.
+    const bool enable_new_header_format;
+
+    // Whether to retain a decompression buffer that reads can reuse rather than
+    // allocating a new one for each read.
+    const bool enable_reusable_decompression_buffer;
+
     // Must specify values for options.
     Options() = delete;
-    explicit Options(
-        bool compress_in,
-        const int32_t max_proto_size_in = constants::kMaxProtoSize,
-        const int32_t compression_level_in = kDefaultCompressionLevel)
+    explicit Options(bool compress_in, const int32_t max_proto_size_in,
+                     const int32_t compression_level_in,
+                     const uint32_t compression_threshold_bytes_in,
+                     const int32_t compression_mem_level_in,
+                     const bool enable_smaller_decompression_buffer_size_in,
+                     const bool enable_new_header_format_in,
+                     const bool enable_reusable_decompression_buffer_in)
         : compress(compress_in),
           max_proto_size(max_proto_size_in),
-          compression_level(compression_level_in) {}
+          compression_level(compression_level_in),
+          compression_threshold_bytes(compression_threshold_bytes_in),
+          compression_mem_level(compression_mem_level_in),
+          enable_smaller_decompression_buffer_size(
+              enable_smaller_decompression_buffer_size_in),
+          enable_new_header_format(enable_new_header_format_in),
+          enable_reusable_decompression_buffer(
+              enable_reusable_decompression_buffer_in) {}
   };
 
   // Level of compression, BEST_SPEED = 1, BEST_COMPRESSION = 9
   static constexpr int kDefaultCompressionLevel = 3;
 
+  // The default compression threshold is 0, which means always compress.
+  static constexpr uint32_t kDefaultCompressionThresholdBytes = 0;
+
+  // The compression ratio to use for decompression buffer size.
+  static constexpr int kProtoCompressionRatio = 3;
+
   // Number of bytes we reserve for the heading at the beginning of the proto
   // log. We reserve this so the header can grow without running into the
   // contents of the proto log, triggering an unnecessary migration of the data.
@@ -140,21 +178,58 @@
     // format.
     static constexpr int32_t kFileFormatVersion = 0;
 
-    uint32_t CalculateHeaderChecksum() const {
+    // Legacy header section padding offset.
+    // - In the legacy header definition, we never declared the padding bytes.
+    //   Instead, it was implicitly added by C++ compiler.
+    // - In order to make the legacy header checksum calculation compatible, we
+    //   need to declare the padding bytes explicitly before the new fields.
+    static constexpr int32_t kLegacyHeaderSectionPaddingOffset =
+        offsetof(Header, reserved_padding1);
+
+    // Legacy header section size, including reserved_padding1. This size must
+    // not be changed due to compatibility issues.
+    static constexpr int32_t kLegacyHeaderSectionSize = 32;
+
+    uint32_t CalculateLegacyHeaderChecksum() const {
+      static_assert(offsetof(Header, flags_) == 28,
+                    "Incompatible header layout for flags and paddings!");
+      static_assert(kLegacyHeaderSectionPaddingOffset +
+                        sizeof(reserved_padding1) ==
+                    kLegacyHeaderSectionSize);
+
       Crc32 crc;
 
-      // Get a string_view of all the fields of the Header, excluding the
-      // magic_nbytes_ and header_checksum_nbytes_
-      std::string_view header_str(
+      // Get a string_view of all the fields of the legacy header section,
+      // excluding the magic_nbytes_ and legacy_header_checksum_nbytes_.
+      std::string_view legacy_header_str(
           reinterpret_cast<const char*>(this) +
-              offsetof(Header, header_checksum_nbytes_) +
-              sizeof(header_checksum_nbytes_),
-          sizeof(Header) - sizeof(magic_nbytes_) -
-              sizeof(header_checksum_nbytes_));
-      crc.Append(header_str);
+              offsetof(Header, legacy_header_checksum_nbytes_) +
+              sizeof(legacy_header_checksum_nbytes_),
+          kLegacyHeaderSectionSize - sizeof(magic_nbytes_) -
+              sizeof(legacy_header_checksum_nbytes_));
+      crc.Append(legacy_header_str);
       return crc.Get();
     }
 
+    Crc32 CalculateHeaderChecksum() const {
+      // The header_checksum field must be at offset 252.
+      static_assert(offsetof(Header, header_checksum_nbytes_) == 252);
+
+      // Get a string_view of all the fields of the header, excluding
+      // header_checksum_nbytes_ itself.
+      std::string_view header_str(
+          reinterpret_cast<const char*>(this),
+          sizeof(Header) - sizeof(header_checksum_nbytes_));
+      return Crc32(header_str);
+    }
+
+    void UpdateHeaderChecksums(bool enable_new_header_format) {
+      SetLegacyHeaderChecksum(CalculateLegacyHeaderChecksum());
+      if (enable_new_header_format) {
+        SetHeaderChecksum(CalculateHeaderChecksum().Get());
+      }
+    }
+
     int32_t GetMagic() const { return GNetworkToHostL(magic_nbytes_); }
 
     void SetMagic(int32_t magic_in) {
@@ -193,12 +268,13 @@
       rewind_offset_nbytes_ = GHostToNetworkLL(rewind_offset_in);
     }
 
-    int32_t GetHeaderChecksum() const {
-      return GNetworkToHostL(header_checksum_nbytes_);
+    uint32_t GetLegacyHeaderChecksum() const {
+      return GNetworkToHostL(legacy_header_checksum_nbytes_);
     }
 
-    void SetHeaderChecksum(int32_t header_checksum_in) {
-      header_checksum_nbytes_ = GHostToNetworkL(header_checksum_in);
+    void SetLegacyHeaderChecksum(uint32_t legacy_header_checksum_in) {
+      legacy_header_checksum_nbytes_ =
+          GHostToNetworkL(legacy_header_checksum_in);
     }
 
     bool GetCompressFlag() const { return GetFlag(kCompressBit); }
@@ -209,6 +285,35 @@
 
     void SetDirtyFlag(bool dirty) { SetFlag(kDirtyBit, dirty); }
 
+    uint32_t GetUnsyncedTailChecksum() const {
+      return GNetworkToHostL(unsynced_tail_checksum_nbytes_);
+    }
+
+    void SetUnsyncedTailChecksum(uint32_t unsynced_tail_checksum_in) {
+      unsynced_tail_checksum_nbytes_ =
+          GHostToNetworkL(unsynced_tail_checksum_in);
+    }
+
+    // Resets paddings to 0.
+    //
+    // Normally a new header will have default 0 bytes for all paddings since we
+    // declared "= {0}" for them, but if the header was read from an existing
+    // file generated by older versions, the paddings may contain some random
+    // values. Therefore, InitializeExistingFile will reset the paddings to 0
+    // just in case.
+    void ResetPaddings() {
+      memset(reserved_padding1, 0, sizeof(reserved_padding1));
+      memset(available_padding, 0, sizeof(available_padding));
+    }
+
+    uint32_t GetHeaderChecksum() const {
+      return GNetworkToHostL(header_checksum_nbytes_);
+    }
+
+    void SetHeaderChecksum(uint32_t header_checksum_in) {
+      header_checksum_nbytes_ = GHostToNetworkL(header_checksum_in);
+    }
+
    private:
     // The least-significant bit offset at which the compress flag is stored in
     // 'flags_nbytes_'. Represents whether the protos in the log are compressed
@@ -228,16 +333,55 @@
       bit_util::BitfieldSet(value, offset, /*len=*/1, &flags_);
     }
 
+    // Header bytes layout:
+    //             +--------- 4 bytes ---------+--------- 4 bytes ---------+
+    // offset 0    |           magic           |  legacy_header_checksum   |
+    // offset 8    |                     rewind_offset                     |
+    // offset 16   |    file_format_version    |      max_proto_size       |
+    // offset 24   |       log_checksum        | flags | reserved_padding1 |
+    //
+    // <legacy header section: offset 0 to 31>
+    //
+    // offset 32   |  unsynced_tail_checksum   |    available_padding[]    |
+    // offset 40   |                  available_padding[]                  |
+    // ...
+    // offset 240  |                  available_padding[]                  |
+    // offset 248  |    available_padding[]    |      header_checksum      |
+    //             +--------- 4 bytes ---------+--------- 4 bytes ---------+
+    //
+    // Note:
+    // - Due to compatibility issues, legacy header checksum must always be
+    //   computed by the fixed legacy header section, and the legacy header
+    //   section must not be changed in the future.
+    // - The new "header_checksum" field is added to the end of the header
+    //   section, and it is computed by the entire header section (excluding
+    //   itself), i.e. offset 0 to 251.
+    // - Any new fields can be added into the section of available_padding, and
+    //   the developer must:
+    //   - Maintain the fields and essential paddings according to the C++
+    //     struct bytes layout and padding rules
+    //   - Recompute the available_padding size to make sizeof(Header) == 256
+    //     and header_checksum_nbytes_ remain at the end of the header (i.e.
+    //     offset 252-255).
+
+    //// START OF LEGACY HEADER SECTION ////
+
     // Holds the magic as a quick sanity check against file corruption.
     //
     // Field is in network-byte order.
     int32_t magic_nbytes_ = GHostToNetworkL(kMagic);
 
     // Must be at the beginning after kMagic. Contains the crc checksum of
-    // the following fields.
+    // the entire legacy header section excluding magic and itself (i.e. offset
+    // 8 to 31).
+    //
+    // Need to keep this field and it should be computed by the same way
+    // forever, due to Android mainline rollback policy. We have to maintain
+    // "forward compatibility", i.e. data generated by the new code should be
+    // able to be interpreted by the old code.
     //
     // Field is in network-byte order.
-    uint32_t header_checksum_nbytes_ = 0;
+    uint32_t legacy_header_checksum_nbytes_ = 0;
 
     // Last known good offset at which the log and its checksum were updated.
     // If we crash between writing to the log and updating the checksum, we can
@@ -272,12 +416,30 @@
     // Field is only 1 byte, so is byte-order agnostic.
     uint8_t flags_ = 0;
 
-    // NOTE: New fields should *almost always* be added to the end here. Since
-    // this class may have already been written to disk, appending fields
-    // increases the chances that changes are backwards-compatible.
+    // Reserved padding after flags_ for the C++ struct bytes layout.
+    // Due to compatibility with legacy_header_checksum_nbytes_, we have to pad
+    // the legacy header section to 32 bytes.
+    uint8_t reserved_padding1[3] = {0};
+
+    //// END OF LEGACY HEADER SECTION ////
+
+    // Checksum of the unsynced tail elements.
+    //
+    // Field is in network-byte order.
+    uint32_t unsynced_tail_checksum_nbytes_ = 0;
+
+    // Available bytes for future use.
+    uint8_t available_padding[216] = {0};
+
+    // Must be at the end of the header, i.e. offset 252-255. Contains the crc
+    // checksum of the entire header section excluding itself (i.e. offset 0 to
+    // 251).
+    //
+    // Field is in network-byte order.
+    uint32_t header_checksum_nbytes_ = 0;
   };
-  static_assert(sizeof(Header) <= kHeaderReservedBytes,
-                "Header has grown past our reserved bytes!");
+  static_assert(sizeof(Header) == kHeaderReservedBytes,
+                "Header size should be the same as the reserved bytes!");
 
   struct CreateResult {
     // A successfully initialized log.
@@ -381,6 +543,9 @@
   //   INTERNAL_ERROR on IO error
   libtextclassifier3::StatusOr<int64_t> GetElementsFileSize() const;
 
+  // For testing only.
+  Header* header() { return header_.get(); }
+
   // An iterator helping to find offsets of all the protos in file.
   // Example usage:
   //
@@ -480,7 +645,9 @@
   // Returns:
   //   OK on success
   //   INTERNAL_ERROR on IO error
-  libtextclassifier3::Status PersistToDisk();
+  libtextclassifier3::Status PersistToDisk(
+      PersistToDiskStatsProto* persist_stats = nullptr,
+      const Clock* clock = nullptr);
 
   // Calculates the checksum of the log contents (excluding the header) and
   // updates the header.
@@ -500,10 +667,13 @@
 
  private:
   // Object can only be instantiated via the ::Create factory.
-  PortableFileBackedProtoLog(const Filesystem* filesystem,
-                             const std::string& file_path,
-                             std::unique_ptr<Header> header, int64_t file_size,
-                             int32_t compression_level);
+  explicit PortableFileBackedProtoLog(
+      const Filesystem* filesystem, const std::string& file_path,
+      std::unique_ptr<Header> header, int64_t file_size,
+      int32_t compression_level, uint32_t compression_threshold_bytes,
+      int32_t compression_mem_level,
+      bool enable_smaller_decompression_buffer_size,
+      bool enable_new_header_format, bool enable_reusable_decompression_buffer);
 
   // Initializes a new proto log.
   //
@@ -553,9 +723,9 @@
   // metadata and converts it into a portable metadata before writing.
   //
   // Returns:
-  //   OK on success
+  //   Crc32 of the written metadata bytes on success
   //   INTERNAL_ERROR on any IO errors
-  static libtextclassifier3::Status WriteProtoMetadata(
+  static libtextclassifier3::StatusOr<Crc32> WriteProtoMetadata(
       const Filesystem* filesystem, int fd, int32_t host_order_metadata);
 
   static bool IsEmptyBuffer(const char* buffer, int size) {
@@ -571,6 +741,30 @@
   // Metadata format: 8 bits magic + 24 bits size
   static uint8_t GetProtoMagic(int metadata) { return metadata >> 24; }
 
+  class BufferHolder {
+   public:
+    BufferHolder(const PortableFileBackedProtoLog* log,
+                 std::unique_ptr<uint8_t[]> buffer, size_t size)
+        : log_(log), buffer_(std::move(buffer)), size_(size) {}
+
+    ~BufferHolder() { log_->ReturnBuffer(std::move(buffer_), size_); }
+
+    uint8_t* data() { return buffer_.get(); }
+    size_t size() const { return size_; }
+
+   private:
+    const PortableFileBackedProtoLog* log_;
+    std::unique_ptr<uint8_t[]> buffer_;
+    size_t size_;
+  };
+
+  // Returns the buffer to the log. It will either be cached or destroyed.
+  void ReturnBuffer(std::unique_ptr<uint8_t[]> buffer, size_t size) const;
+
+  // Returns a buffer holder to the caller - either a newly constructed one or
+  // the cached one if it is big enough.
+  BufferHolder PossiblyBorrowBuffer(size_t size) const;
+
   // Magic number added in front of every proto. Used when reading out protos
   // as a first check for corruption in each entry in the file. Even if there is
   // a corruption, the best we can do is roll back to our last recovery point
@@ -589,18 +783,38 @@
   std::unique_ptr<Header> header_;
   int64_t file_size_;
   const int32_t compression_level_;
+  const uint32_t compression_threshold_bytes_;
+  const int32_t compression_mem_level_;
+  const bool enable_smaller_decompression_buffer_size_;
+  const bool enable_new_header_format_;
+  const bool enable_reusable_decompression_buffer_;
+
+  mutable std::unique_ptr<uint8_t[]> read_buffer_ ICING_GUARDED_BY(mutex_);
+  mutable size_t read_buffer_size_ ICING_GUARDED_BY(mutex_);
+  mutable absl_ports::shared_mutex mutex_;
 };
 
 template <typename ProtoT>
 PortableFileBackedProtoLog<ProtoT>::PortableFileBackedProtoLog(
     const Filesystem* filesystem, const std::string& file_path,
     std::unique_ptr<Header> header, int64_t file_size,
-    int32_t compression_level)
+    int32_t compression_level, uint32_t compression_threshold_bytes,
+    int32_t compression_mem_level,
+    bool enable_smaller_decompression_buffer_size,
+    bool enable_new_header_format, bool enable_reusable_decompression_buffer)
     : filesystem_(filesystem),
       file_path_(file_path),
       header_(std::move(header)),
       file_size_(file_size),
-      compression_level_(compression_level) {
+      compression_level_(compression_level),
+      compression_threshold_bytes_(compression_threshold_bytes),
+      compression_mem_level_(compression_mem_level),
+      enable_smaller_decompression_buffer_size_(
+          enable_smaller_decompression_buffer_size),
+      enable_new_header_format_(enable_new_header_format),
+      enable_reusable_decompression_buffer_(
+          enable_reusable_decompression_buffer),
+      read_buffer_size_(0) {
   fd_.reset(filesystem_->OpenForAppend(file_path.c_str()));
 }
 
@@ -639,6 +853,13 @@
         options.compression_level));
   }
 
+  if (options.compression_mem_level < 1 || options.compression_mem_level > 9) {
+    return absl_ports::InvalidArgumentError(
+        IcingStringUtil::StringPrintf("options.compression_mem_level must be "
+                                      "between 1 and 9 inclusive, was %d",
+                                      options.compression_mem_level));
+  }
+
   if (!filesystem->FileExists(file_path.c_str())) {
     return InitializeNewFile(filesystem, file_path, options);
   }
@@ -668,22 +889,50 @@
         absl_ports::StrCat("Failed to initialize file size: ", file_path));
   }
 
-  // Create the header
+  // Create the header.
+  //
+  // No need to memset padding bytes to 0 since we declared = {0} in the struct
+  // definition and the new object created by std::make_unique will initialize
+  // them to 0.
   std::unique_ptr<Header> header = std::make_unique<Header>();
   header->SetCompressFlag(options.compress);
   header->SetMaxProtoSize(options.max_proto_size);
-  header->SetHeaderChecksum(header->CalculateHeaderChecksum());
+  header->UpdateHeaderChecksums(options.enable_new_header_format);
 
-  if (!filesystem->Write(file_path.c_str(), header.get(), sizeof(Header))) {
-    return absl_ports::InternalError(
-        absl_ports::StrCat("Failed to write header for file: ", file_path));
+  {
+    ScopedFd fd(filesystem->OpenForWrite(file_path.c_str()));
+    if (!fd.is_valid()) {
+      return absl_ports::InternalError(
+          absl_ports::StrCat("Failed to open file for write: ", file_path));
+    }
+
+    if (!filesystem->Write(fd.get(), header.get(), sizeof(Header))) {
+      return absl_ports::InternalError(
+          absl_ports::StrCat("Failed to write header for file: ", file_path));
+    }
+
+    // Sync the file to disk to ensure that the header is flushed to disk.
+    // - Otherwise, if the app crashes before the header is flushed, the next
+    //   initialize may fail.
+    // - This is especially important for this class, since it is used to store
+    //   ground truth data. If magic or checksum is wrong, then Icing cannot
+    //   recover it from this state and therefore end up with entire data loss.
+    if (!filesystem->DataSync(fd.get())) {
+      return absl_ports::InternalError(
+          absl_ports::StrCat("Failed to sync file: ", file_path));
+    }
   }
 
   CreateResult create_result = {
       std::unique_ptr<PortableFileBackedProtoLog<ProtoT>>(
           new PortableFileBackedProtoLog<ProtoT>(
               filesystem, file_path, std::move(header),
-              /*file_size=*/kHeaderReservedBytes, options.compression_level)),
+              /*file_size=*/kHeaderReservedBytes, options.compression_level,
+              options.compression_threshold_bytes,
+              options.compression_mem_level,
+              options.enable_smaller_decompression_buffer_size,
+              options.enable_new_header_format,
+              options.enable_reusable_decompression_buffer)),
       /*data_loss=*/DataLoss::NONE, /*recalculated_checksum=*/false};
 
   return create_result;
@@ -695,30 +944,62 @@
 PortableFileBackedProtoLog<ProtoT>::InitializeExistingFile(
     const Filesystem* filesystem, const std::string& file_path,
     const Options& options, int64_t file_size) {
+  bool legacy_header_section_changed = false;
   bool header_changed = false;
   if (file_size < kHeaderReservedBytes) {
+    ICING_LOG(ERROR) << "Invalid file size for PortableFileBackedProtoLog "
+                     << file_path
+                     << " for header reserved bytes. File size: " << file_size
+                     << ", header reserved bytes: " << kHeaderReservedBytes;
     return absl_ports::InternalError(
         absl_ports::StrCat("File header too short for: ", file_path));
   }
 
   std::unique_ptr<Header> header = std::make_unique<Header>();
-  if (!filesystem->PRead(file_path.c_str(), header.get(), sizeof(Header),
-                         /*offset=*/0)) {
+  if (filesystem->PRead(file_path.c_str(), header.get(), sizeof(Header),
+                        /*offset=*/0) != sizeof(Header)) {
     return absl_ports::InternalError(
         absl_ports::StrCat("Failed to read header for file: ", file_path));
   }
 
   // Make sure the header is still valid before we use any of its values. This
-  // is covered by the header_checksum check below, but this is a quick check
-  // that can save us from an extra crc computation.
+  // is covered by the legacy_header_checksum check below, but this is a quick
+  // check that can save us from an extra crc computation.
   if (header->GetMagic() != Header::kMagic) {
-    return absl_ports::InternalError(
-        absl_ports::StrCat("Invalid header kMagic for file: ", file_path));
+    ICING_LOG(ERROR) << "Invalid header magic for PortableFileBackedProtoLog "
+                     << file_path << ". Expected: " << Header::kMagic
+                     << ", actual: " << header->GetMagic();
+    return absl_ports::InternalError(absl_ports::StrCat(
+        "Invalid header magic for PortableFileBackedProtoLog: ", file_path));
   }
 
-  if (header->GetHeaderChecksum() != header->CalculateHeaderChecksum()) {
+  // Check the new header checksum.
+  bool header_checksum_mismatch = false;
+  if (options.enable_new_header_format) {
+    header_checksum_mismatch =
+        header->GetHeaderChecksum() != header->CalculateHeaderChecksum().Get();
+    if (header_checksum_mismatch) {
+      // If upgrade or rollforward happens, then existing data's header_checksum
+      // may not be written by the old code, so we need to fallback to legacy
+      // checksum.
+      ICING_LOG(WARNING)
+          << "Invalid header checksum for PortableFileBackedProtoLog "
+          << file_path
+          << ". Fallback to legacy header checksum validation, and need to "
+             "rewind unsynced tail.";
+    }
+  }
+
+  // Check the legacy header checksum.
+  if (header->GetLegacyHeaderChecksum() !=
+      header->CalculateLegacyHeaderChecksum()) {
+    // Old code and new code should write legacy header checksum correctly, so
+    // if it doesn't match, then it's likely that the file is corrupted.
+    ICING_LOG(ERROR)
+        << "Invalid legacy header checksum for PortableFileBackedProtoLog "
+        << file_path;
     return absl_ports::InternalError(
-        absl_ports::StrCat("Invalid header checksum for: ", file_path));
+        absl_ports::StrCat("Invalid legacy header checksum for: ", file_path));
   }
 
   if (header->GetRewindOffset() < kHeaderReservedBytes) {
@@ -749,21 +1030,55 @@
     // It's fine if our new max size is greater than our previous one. Existing
     // data is still valid.
     header->SetMaxProtoSize(options.max_proto_size);
-    header_changed = true;
+    legacy_header_section_changed = true;
   }
 
   DataLoss data_loss = DataLoss::NONE;
 
-  // If we have any documents in our tail, get rid of them since they're not in
-  // our checksum. Our checksum reflects content up to the rewind offset.
+  // If we have any documents in our tail, validate the unsynced tail checksum
+  // and get rid of the unsynced tail if they don't match.
   if (file_size > header->GetRewindOffset()) {
-    if (!filesystem->Truncate(file_path.c_str(), header->GetRewindOffset())) {
-      return absl_ports::InternalError(IcingStringUtil::StringPrintf(
-          "Failed to truncate '%s' to size %lld", file_path.data(),
-          static_cast<long long>(header->GetRewindOffset())));
+    // - If enable_new_header_format is false, then we need to throw away all
+    //   data in the unsynced tail.
+    // - Otherwise we use the new header format logic:
+    //   - If the header checksum doesn't match, then we need to throw away all
+    //     data in the unsynced tail, regardless of the unsynced tail checksum.
+    //   - Otherwise, validate the unsynced tail checksum to determine if we
+    //     need to throw away all data in the unsynced tail.
+    bool need_rewind_unsynced_tail =
+        !options.enable_new_header_format || header_checksum_mismatch;
+    if (!need_rewind_unsynced_tail) {
+      // Recompute the unsynced tail's checksum only if the header checksum
+      // matches. This saves us a crc computation if the header checksum already
+      // doesn't match.
+      ICING_ASSIGN_OR_RETURN(
+          Crc32 calculated_unsynced_tail_checksum,
+          GetPartialChecksum(filesystem, file_path, Crc32(),
+                             /*start=*/header->GetRewindOffset(),
+                             /*end=*/file_size, file_size));
+      if (header->GetUnsyncedTailChecksum() !=
+          calculated_unsynced_tail_checksum.Get()) {
+        // Unsynced tail checksum doesn't match, so we need to throw it out.
+        ICING_LOG(WARNING)
+            << "Mismatch unsynced tail checksum. Expected (from data): "
+            << calculated_unsynced_tail_checksum.Get()
+            << ", actual (stored in the header): "
+            << header->GetUnsyncedTailChecksum();
+        need_rewind_unsynced_tail = true;
+      }
     }
-    file_size = header->GetRewindOffset();
-    data_loss = DataLoss::PARTIAL;
+
+    if (need_rewind_unsynced_tail) {
+      if (!filesystem->Truncate(file_path.c_str(), header->GetRewindOffset())) {
+        return absl_ports::InternalError(IcingStringUtil::StringPrintf(
+            "Failed to truncate '%s' to size %lld", file_path.data(),
+            static_cast<long long>(header->GetRewindOffset())));
+      }
+      header->SetUnsyncedTailChecksum(0);
+      file_size = header->GetRewindOffset();
+      data_loss = DataLoss::PARTIAL;
+      header_changed = true;
+    }
   }
 
   bool recalculated_checksum = false;
@@ -801,11 +1116,15 @@
 
     // Update our header.
     header->SetDirtyFlag(false);
-    header_changed = true;
+    legacy_header_section_changed = true;
   }
 
-  if (header_changed) {
-    header->SetHeaderChecksum(header->CalculateHeaderChecksum());
+  if (legacy_header_section_changed || header_changed ||
+      header_checksum_mismatch) {
+    // Reset all padding bytes to 0 before recompute the checksums, just in
+    // case we had some random values generated by the old code.
+    header->ResetPaddings();
+    header->UpdateHeaderChecksums(options.enable_new_header_format);
 
     if (!filesystem->PWrite(file_path.c_str(), /*offset=*/0, header.get(),
                             sizeof(Header))) {
@@ -816,9 +1135,13 @@
 
   CreateResult create_result = {
       std::unique_ptr<PortableFileBackedProtoLog<ProtoT>>(
-          new PortableFileBackedProtoLog<ProtoT>(filesystem, file_path,
-                                                 std::move(header), file_size,
-                                                 options.compression_level)),
+          new PortableFileBackedProtoLog<ProtoT>(
+              filesystem, file_path, std::move(header), file_size,
+              options.compression_level, options.compression_threshold_bytes,
+              options.compression_mem_level,
+              options.enable_smaller_decompression_buffer_size,
+              options.enable_new_header_format,
+              options.enable_reusable_decompression_buffer)),
       data_loss, recalculated_checksum};
 
   return create_result;
@@ -829,12 +1152,6 @@
 PortableFileBackedProtoLog<ProtoT>::GetPartialChecksum(
     const Filesystem* filesystem, const std::string& file_path,
     Crc32 initial_crc, int64_t start, int64_t end, int64_t file_size) {
-  ICING_ASSIGN_OR_RETURN(
-      MemoryMappedFile mmapped_file,
-      MemoryMappedFile::Create(*filesystem, file_path,
-                               MemoryMappedFile::Strategy::READ_ONLY));
-  Crc32 new_crc(initial_crc.Get());
-
   if (start < 0) {
     return absl_ports::InvalidArgumentError(IcingStringUtil::StringPrintf(
         "Starting checksum offset of file '%s' must be greater than 0, was "
@@ -858,6 +1175,16 @@
         static_cast<long long>(end)));
   }
 
+  if (start == end) {
+    return initial_crc;
+  }
+
+  ICING_ASSIGN_OR_RETURN(
+      MemoryMappedFile mmapped_file,
+      MemoryMappedFile::Create(*filesystem, file_path,
+                               MemoryMappedFile::Strategy::READ_ONLY));
+  Crc32 new_crc(initial_crc.Get());
+
   Architecture architecture = GetArchitecture();
   switch (architecture) {
     case Architecture::BIT_64: {
@@ -922,7 +1249,17 @@
   if (header_->GetCompressFlag()) {
     protobuf_ports::GzipOutputStream::Options options;
     options.format = protobuf_ports::GzipOutputStream::ZLIB;
-    options.compression_level = compression_level_;
+
+    if (proto_size >= compression_threshold_bytes_) {
+      options.compression_level = compression_level_;
+    } else {
+      options.compression_level = 0;
+    }
+    options.buffer_size = proto_size;
+    if (proto_size < 0 ||
+        static_cast<size_t>(proto_size) > protobuf_ports::kDefaultBufferSize) {
+      options.buffer_size = protobuf_ports::kDefaultBufferSize;
+    }
 
     protobuf_ports::GzipOutputStream compressing_stream(&proto_stream, options);
 
@@ -955,7 +1292,8 @@
 
   // Actually write metadata, has to be done after we know the possibly
   // compressed proto size
-  ICING_RETURN_IF_ERROR(
+  ICING_ASSIGN_OR_RETURN(
+      Crc32 metadata_crc,
       WriteProtoMetadata(filesystem_, fd_.get(), host_order_metadata));
 
   // Write the serialized proto
@@ -967,6 +1305,23 @@
   // Update file size. The file should have grown by sizeof(Metadata) + size of
   // the serialized proto.
   file_size_ += sizeof(host_order_metadata) + final_size;
+
+  if (enable_new_header_format_) {
+    // Compute the new unsynced tail checksum.
+    Crc32 new_unsynced_tail_crc(header_->GetUnsyncedTailChecksum());
+    new_unsynced_tail_crc.Combine(metadata_crc, sizeof(host_order_metadata));
+    new_unsynced_tail_crc.Append(proto_str);
+
+    // Set the new unsynced tail checksum and header checksum. Write the header.
+    header_->SetUnsyncedTailChecksum(new_unsynced_tail_crc.Get());
+    header_->SetHeaderChecksum(header_->CalculateHeaderChecksum().Get());
+    if (!filesystem_->PWrite(fd_.get(), /*offset=*/0, header_.get(),
+                             sizeof(Header))) {
+      return absl_ports::InternalError(absl_ports::StrCat(
+          "Failed to update header after write proto: ", file_path_));
+    }
+  }
+
   return current_position;
 }
 
@@ -990,7 +1345,8 @@
                                       static_cast<long long>(file_size_ - 1)));
   }
   auto buf = std::make_unique<char[]>(stored_size);
-  if (!filesystem_->PRead(fd_.get(), buf.get(), stored_size, file_offset)) {
+  if (filesystem_->PRead(fd_.get(), buf.get(), stored_size, file_offset) !=
+      stored_size) {
     return absl_ports::InternalError("");
   }
 
@@ -1003,7 +1359,15 @@
   // Deserialize proto
   ProtoT proto;
   if (header_->GetCompressFlag()) {
-    protobuf_ports::GzipInputStream decompress_stream(&proto_stream);
+    size_t buffer_size = protobuf_ports::kDefaultBufferSize;
+    if (enable_smaller_decompression_buffer_size_ && stored_size >= 0) {
+      buffer_size = std::min(buffer_size, kProtoCompressionRatio *
+                                              static_cast<size_t>(stored_size));
+    }
+    BufferHolder buffer_holder = PossiblyBorrowBuffer(buffer_size);
+    protobuf_ports::GzipInputStream decompress_stream(
+        &proto_stream, protobuf_ports::GzipInputStream::AUTO,
+        buffer_holder.data(), buffer_holder.size());
     proto.ParseFromZeroCopyStream(&decompress_stream);
   } else {
     proto.ParseFromZeroCopyStream(&proto_stream);
@@ -1030,28 +1394,44 @@
   }
   auto buf = std::make_unique<char[]>(stored_size);
 
-  // We need to update the crc checksum if the erased area is before the
-  // rewind position.
-  int32_t new_crc;
+  // Update the checksum depending on whether the erased area is before or after
+  // rewind offset.
+  uint32_t new_crc = 0;  // Reused for both cases.
   if (file_offset < header_->GetRewindOffset()) {
+    // Case 1: erasing bytes from log checksummed area (before rewind offset).
+    //
+    // kHeaderReservedBytes                     RewindOffset           file_size
+    //          v                                    v                     v
+    // +--------+------------------------------------+---------------------+
+    // | Header |        Log checksummed area        |    Unsynced tail    |
+    // +--------+------------------------------------+---------------------+
+    //          |<-          log_checksum          ->|<-   ut_checksum   ->|
+    //                 xxxxxxxxxxxx
+    //                 ^           ^
+    //            file_offset file_offset
+    //                          + stored_size
+
     // Set to "dirty" before we start writing anything.
     header_->SetDirtyFlag(true);
-    header_->SetHeaderChecksum(header_->CalculateHeaderChecksum());
+    header_->UpdateHeaderChecksums(enable_new_header_format_);
+
     if (!filesystem_->PWrite(fd_.get(), /*offset=*/0, header_.get(),
                              sizeof(Header))) {
       return absl_ports::InternalError(absl_ports::StrCat(
           "Failed to update dirty bit of header to: ", file_path_));
     }
 
-    // We need to calculate [original string xor 0s].
-    // The xored string is the same as the original string because 0 xor 0 =
-    // 0, 1 xor 0 = 1.
     // Read the compressed proto out.
-    if (!filesystem_->PRead(fd_.get(), buf.get(), stored_size, file_offset)) {
-      return absl_ports::InternalError("");
+    if (filesystem_->PRead(fd_.get(), buf.get(), stored_size, file_offset) !=
+        stored_size) {
+      return absl_ports::InternalError("Cannot read proto from file");
     }
+    // We're going to erase the bytes. For example, "ABCDEF" -> "A\0\0DEF". The
+    // checksum can be easily updated by ("BC" XOR "\0\0"), which remains "BC",
+    // so the XORed string is the original string.
     const std::string_view xored_str(buf.get(), stored_size);
 
+    // Recompute the log checksum by the XORed data.
     Crc32 crc(header_->GetLogChecksum());
     ICING_ASSIGN_OR_RETURN(
         new_crc,
@@ -1059,6 +1439,42 @@
                           /*full_data_size=*/header_->GetRewindOffset() -
                               kHeaderReservedBytes,
                           /*position=*/file_offset - kHeaderReservedBytes));
+  } else {
+    // Case 2: erasing bytes from unsynced tail checksum area (after rewind
+    // offset).
+    //
+    // kHeaderReservedBytes                     RewindOffset           file_size
+    //          v                                    v                     v
+    // +--------+------------------------------------+---------------------+
+    // | Header |        Log checksummed area        |    Unsynced tail    |
+    // +--------+------------------------------------+---------------------+
+    //          |<-          log_checksum          ->|<-   ut_checksum   ->|
+    //                                                    xxxxxxxxxxxx
+    //                                                    ^           ^
+    //                                               file_offset file_offset
+    //                                                             + stored_size
+
+    // Read the compressed proto out and re-calculate the unsynced tail checksum
+    // only if the flag is enabled.
+    if (enable_new_header_format_) {
+      if (filesystem_->PRead(fd_.get(), buf.get(), stored_size, file_offset) !=
+          stored_size) {
+        return absl_ports::InternalError("Cannot read proto from file");
+      }
+      // We're going to erase the bytes. For example, "ABCDEF" -> "A\0\0DEF".
+      // The checksum can be easily updated by ("BC" XOR "\0\0"), which remains
+      // "BC", so the XORed string is the original string.
+      const std::string_view xored_str(buf.get(), stored_size);
+
+      // Recompute the unsynced tail checksum by the XORed data.
+      Crc32 crc(header_->GetUnsyncedTailChecksum());
+      ICING_ASSIGN_OR_RETURN(
+          new_crc,
+          crc.UpdateWithXor(
+              xored_str,
+              /*full_data_size=*/file_size_ - header_->GetRewindOffset(),
+              /*position=*/file_offset - header_->GetRewindOffset()));
+    }
   }
 
   // Clear the region.
@@ -1067,18 +1483,30 @@
     return absl_ports::InternalError("");
   }
 
-  // If we cleared something in our checksummed area, we should update our
-  // checksum and reset our dirty bit.
+  // Update header.
+  bool write_header = false;
   if (file_offset < header_->GetRewindOffset()) {
+    // If we cleared something in the log checksummed area, we should reset the
+    // dirty bit and update related all checksums.
     header_->SetDirtyFlag(false);
     header_->SetLogChecksum(new_crc);
-    header_->SetHeaderChecksum(header_->CalculateHeaderChecksum());
+    header_->UpdateHeaderChecksums(enable_new_header_format_);
 
-    if (!filesystem_->PWrite(fd_.get(), /*offset=*/0, header_.get(),
-                             sizeof(Header))) {
-      return absl_ports::InternalError(
-          absl_ports::StrCat("Failed to update header to: ", file_path_));
-    }
+    write_header = true;
+  } else if (enable_new_header_format_) {
+    // If we cleared something in the unsynced tail section, we should update
+    // the unsynced tail checksum and header checksum only if the flag is
+    // enabled.
+    header_->SetUnsyncedTailChecksum(new_crc);
+    header_->SetHeaderChecksum(header_->CalculateHeaderChecksum().Get());
+
+    write_header = true;
+  }
+
+  if (write_header && !filesystem_->PWrite(fd_.get(), /*offset=*/0,
+                                           header_.get(), sizeof(Header))) {
+    return absl_ports::InternalError(
+        absl_ports::StrCat("Failed to update header to: ", file_path_));
   }
 
   return libtextclassifier3::Status::OK;
@@ -1168,7 +1596,8 @@
         static_cast<long long>(file_size)));
   }
 
-  if (!filesystem->PRead(fd, &portable_metadata, metadata_size, file_offset)) {
+  if (filesystem->PRead(fd, &portable_metadata, metadata_size, file_offset) !=
+      metadata_size) {
     return absl_ports::InternalError("");
   }
 
@@ -1187,7 +1616,7 @@
 }
 
 template <typename ProtoT>
-libtextclassifier3::Status
+libtextclassifier3::StatusOr<Crc32>
 PortableFileBackedProtoLog<ProtoT>::WriteProtoMetadata(
     const Filesystem* filesystem, int fd, int32_t host_order_metadata) {
   // Convert it into portable endian format before writing to disk
@@ -1200,18 +1629,37 @@
         absl_ports::StrCat("Failed to write proto metadata."));
   }
 
-  return libtextclassifier3::Status::OK;
+  std::string_view metadata_bytes(
+      reinterpret_cast<const char*>(&portable_metadata),
+      sizeof(portable_metadata));
+  return Crc32(metadata_bytes);
 }
 
 template <typename ProtoT>
-libtextclassifier3::Status PortableFileBackedProtoLog<ProtoT>::PersistToDisk() {
+libtextclassifier3::Status PortableFileBackedProtoLog<ProtoT>::PersistToDisk(
+    PersistToDiskStatsProto* persist_stats, const Clock* clock) {
   if (file_size_ == header_->GetRewindOffset()) {
     // No new protos appended, don't need to update the checksum.
     return libtextclassifier3::Status::OK;
   }
 
+  std::unique_ptr<Timer> persist_timer;
+  if (persist_stats != nullptr && clock != nullptr) {
+    persist_timer = clock->GetNewTimer();
+  }
   ICING_RETURN_IF_ERROR(UpdateChecksum());
-  if (!filesystem_->DataSync(fd_.get())) {
+  if (persist_stats != nullptr && clock != nullptr) {
+    persist_stats->set_document_log_checksum_update_latency_ms(
+        persist_timer->GetElapsedMilliseconds());
+    persist_timer = clock->GetNewTimer();
+  }
+
+  bool datasync_succeeded = filesystem_->DataSync(fd_.get());
+  if (persist_stats != nullptr && clock != nullptr) {
+    persist_stats->set_document_log_data_sync_latency_ms(
+        persist_timer->GetElapsedMilliseconds());
+  }
+  if (!datasync_succeeded) {
     return absl_ports::InternalError(
         absl_ports::StrCat("Failed to sync data to disk: ", file_path_));
   }
@@ -1227,8 +1675,9 @@
   }
   ICING_ASSIGN_OR_RETURN(Crc32 crc, GetChecksum());
   header_->SetLogChecksum(crc.Get());
+  header_->SetUnsyncedTailChecksum(0);
   header_->SetRewindOffset(file_size_);
-  header_->SetHeaderChecksum(header_->CalculateHeaderChecksum());
+  header_->UpdateHeaderChecksums(enable_new_header_format_);
 
   if (!filesystem_->PWrite(fd_.get(), /*offset=*/0, header_.get(),
                            sizeof(Header))) {
@@ -1251,6 +1700,13 @@
                               /*start=*/kHeaderReservedBytes,
                               /*end=*/file_size_, file_size_);
   } else {
+    if (enable_new_header_format_) {
+      // Combine unsynced tail checksum with log checksum.
+      Crc32 log_checksum(header_->GetLogChecksum());
+      Crc32 unsynced_tail_checksum(header_->GetUnsyncedTailChecksum());
+      log_checksum.Combine(unsynced_tail_checksum, new_content_size);
+      return log_checksum;
+    }
     // Append new changes to the existing checksum.
     return GetPartialChecksum(
         filesystem_, file_path_, Crc32(header_->GetLogChecksum()),
@@ -1258,6 +1714,36 @@
   }
 }
 
+template <typename ProtoT>
+void PortableFileBackedProtoLog<ProtoT>::ReturnBuffer(
+    std::unique_ptr<uint8_t[]> buffer, size_t size) const {
+  if (!enable_reusable_decompression_buffer_) {
+    return;
+  }
+  absl_ports::unique_lock lock(&mutex_);
+  if (read_buffer_ == nullptr || read_buffer_size_ < size) {
+    read_buffer_.swap(buffer);
+    read_buffer_size_ = size;
+  }
+}
+
+template <typename ProtoT>
+typename PortableFileBackedProtoLog<ProtoT>::BufferHolder
+PortableFileBackedProtoLog<ProtoT>::PossiblyBorrowBuffer(size_t size) const {
+  if (!enable_reusable_decompression_buffer_) {
+    return BufferHolder(this, std::make_unique<uint8_t[]>(size), size);
+  }
+  absl_ports::unique_lock lock(&mutex_);
+  if (read_buffer_ == nullptr || read_buffer_size_ < size) {
+    return BufferHolder(this, std::make_unique<uint8_t[]>(size), size);
+  }
+  std::unique_ptr<uint8_t[]> temp;
+  temp.swap(read_buffer_);
+  size_t temp_size = read_buffer_size_;
+  read_buffer_size_ = 0;
+  return BufferHolder(this, std::move(temp), temp_size);
+}
+
 }  // namespace lib
 }  // namespace icing
 

diff --git a/icing/file/portable-file-backed-proto-log_benchmark.cc b/icing/file/portable-file-backed-proto-log_benchmark.cc
index e6935ab..36ab812 100644
--- a/icing/file/portable-file-backed-proto-log_benchmark.cc
+++ b/icing/file/portable-file-backed-proto-log_benchmark.cc

@@ -13,14 +13,16 @@
 // limitations under the License.
 
 #include <cstdint>
+#include <memory>
 #include <random>
+#include <string>
 
 #include "testing/base/public/benchmark.h"
-#include "gmock/gmock.h"
 #include "icing/document-builder.h"
 #include "icing/file/filesystem.h"
 #include "icing/file/portable-file-backed-proto-log.h"
 #include "icing/legacy/core/icing-string-util.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/proto/document.pb.h"
 #include "icing/testing/common-matchers.h"
 #include "icing/testing/random-string.h"
@@ -55,6 +57,25 @@
 
 namespace {
 
+std::unique_ptr<PortableFileBackedProtoLog<DocumentProto>> CreateProtoLog(
+    const Filesystem& filesystem, const std::string& file_path,
+    int max_proto_size, bool compress) {
+  return PortableFileBackedProtoLog<DocumentProto>::Create(
+             &filesystem, file_path,
+             PortableFileBackedProtoLog<DocumentProto>::Options(
+                 compress, max_proto_size,
+                 PortableFileBackedProtoLog<
+                     DocumentProto>::kDefaultCompressionLevel,
+                 PortableFileBackedProtoLog<
+                     DocumentProto>::kDefaultCompressionThresholdBytes,
+                 protobuf_ports::kDefaultMemLevel,
+                 /*enable_smaller_decompression_buffer_size=*/true,
+                 /*enable_new_header_format=*/true,
+                 /*enable_reusable_decompression_buffer=*/true))
+      .ValueOrDie()
+      .proto_log;
+}
+
 void BM_Write(benchmark::State& state) {
   const Filesystem filesystem;
   int string_length = state.range(0);
@@ -66,12 +87,8 @@
   // Make sure it doesn't already exist.
   filesystem.DeleteFile(file_path.c_str());
 
-  auto proto_log = PortableFileBackedProtoLog<DocumentProto>::Create(
-                       &filesystem, file_path,
-                       PortableFileBackedProtoLog<DocumentProto>::Options(
-                           compress, max_proto_size))
-                       .ValueOrDie()
-                       .proto_log;
+  auto proto_log =
+      CreateProtoLog(filesystem, file_path, max_proto_size, compress);
 
   DocumentProto document = DocumentBuilder().SetKey("namespace", "uri").Build();
 
@@ -119,12 +136,8 @@
   // Make sure it doesn't already exist.
   filesystem.DeleteFile(file_path.c_str());
 
-  auto proto_log = PortableFileBackedProtoLog<DocumentProto>::Create(
-                       &filesystem, file_path,
-                       PortableFileBackedProtoLog<DocumentProto>::Options(
-                           compress, max_proto_size))
-                       .ValueOrDie()
-                       .proto_log;
+  auto proto_log =
+      CreateProtoLog(filesystem, file_path, max_proto_size, compress);
 
   DocumentProto document = DocumentBuilder().SetKey("namespace", "uri").Build();
 
@@ -174,12 +187,8 @@
   // Make sure it doesn't already exist.
   filesystem.DeleteFile(file_path.c_str());
 
-  auto proto_log = PortableFileBackedProtoLog<DocumentProto>::Create(
-                       &filesystem, file_path,
-                       PortableFileBackedProtoLog<DocumentProto>::Options(
-                           compress, max_proto_size))
-                       .ValueOrDie()
-                       .proto_log;
+  auto proto_log =
+      CreateProtoLog(filesystem, file_path, max_proto_size, compress);
 
   DocumentProto document = DocumentBuilder().SetKey("namespace", "uri").Build();
 
@@ -213,12 +222,8 @@
   // Make sure it doesn't already exist.
   filesystem.DeleteFile(file_path.c_str());
 
-  auto proto_log = PortableFileBackedProtoLog<DocumentProto>::Create(
-                       &filesystem, file_path,
-                       PortableFileBackedProtoLog<DocumentProto>::Options(
-                           compress, max_proto_size))
-                       .ValueOrDie()
-                       .proto_log;
+  auto proto_log =
+      CreateProtoLog(filesystem, file_path, max_proto_size, compress);
 
   DocumentProto document = DocumentBuilder().SetKey("namespace", "uri").Build();
 
@@ -255,12 +260,8 @@
   // Make sure it doesn't already exist.
   filesystem.DeleteFile(file_path.c_str());
 
-  auto proto_log = PortableFileBackedProtoLog<DocumentProto>::Create(
-                       &filesystem, file_path,
-                       PortableFileBackedProtoLog<DocumentProto>::Options(
-                           compress, max_proto_size))
-                       .ValueOrDie()
-                       .proto_log;
+  auto proto_log =
+      CreateProtoLog(filesystem, file_path, max_proto_size, compress);
 
   DocumentProto document = DocumentBuilder().SetKey("namespace", "uri").Build();
 
@@ -299,12 +300,8 @@
   // Make sure it doesn't already exist.
   filesystem.DeleteFile(file_path.c_str());
 
-  auto proto_log = PortableFileBackedProtoLog<DocumentProto>::Create(
-                       &filesystem, file_path,
-                       PortableFileBackedProtoLog<DocumentProto>::Options(
-                           compress, max_proto_size))
-                       .ValueOrDie()
-                       .proto_log;
+  auto proto_log =
+      CreateProtoLog(filesystem, file_path, max_proto_size, compress);
 
   DocumentProto document = DocumentBuilder().SetKey("namespace", "uri").Build();
 

diff --git a/icing/file/portable-file-backed-proto-log_test.cc b/icing/file/portable-file-backed-proto-log_test.cc
index 587bfdc..97dbbe6 100644
--- a/icing/file/portable-file-backed-proto-log_test.cc
+++ b/icing/file/portable-file-backed-proto-log_test.cc

@@ -16,16 +16,27 @@
 
 #include <cstdint>
 #include <cstdlib>
+#include <cstring>
+#include <memory>
+#include <string>
+#include <utility>
 
+#include "icing/text_classifier/lib3/utils/base/status.h"
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include "icing/document-builder.h"
 #include "icing/file/filesystem.h"
+#include "icing/file/memory-mapped-file.h"
 #include "icing/file/mock-filesystem.h"
 #include "icing/portable/equals-proto.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/proto/document.pb.h"
 #include "icing/testing/common-matchers.h"
 #include "icing/testing/tmp-directory.h"
+#include "icing/util/crc32.h"
+#include "icing/util/data-loss.h"
+#include "icing/util/status-macros.h"
 
 namespace icing {
 namespace lib {
@@ -33,14 +44,12 @@
 namespace {
 
 using ::icing::lib::portable_equals_proto::EqualsProto;
-using ::testing::A;
 using ::testing::Eq;
 using ::testing::Gt;
 using ::testing::HasSubstr;
+using ::testing::Ne;
 using ::testing::Not;
 using ::testing::NotNull;
-using ::testing::Pair;
-using ::testing::Return;
 
 using Header = PortableFileBackedProtoLog<DocumentProto>::Header;
 
@@ -56,7 +65,34 @@
   filesystem.Write(file_path.c_str(), &header, sizeof(Header));
 }
 
-class PortableFileBackedProtoLogTest : public ::testing::Test {
+libtextclassifier3::StatusOr<std::pair<int64_t, int64_t>>
+WriteProtoAndReturnOffsetAndSize(
+    PortableFileBackedProtoLog<DocumentProto>* proto_log,
+    const DocumentProto& document) {
+  ICING_ASSIGN_OR_RETURN(int64_t document_log_size_before,
+                         proto_log->GetElementsFileSize());
+  ICING_ASSIGN_OR_RETURN(int64_t offset, proto_log->WriteProto(document));
+  ICING_ASSIGN_OR_RETURN(int64_t document_log_size_after,
+                         proto_log->GetElementsFileSize());
+  int64_t size_written = document_log_size_after - document_log_size_before;
+  return std::make_pair(offset, size_written);
+}
+
+void CheckNewChecksums(const Header& header,
+                       uint32_t expected_unsynced_tail_checksum,
+                       uint32_t expected_header_checksum,
+                       bool enable_new_header_format) {
+  if (enable_new_header_format) {
+    EXPECT_THAT(header.GetUnsyncedTailChecksum(),
+                Eq(expected_unsynced_tail_checksum));
+    EXPECT_THAT(header.GetHeaderChecksum(), Eq(expected_header_checksum));
+  } else {
+    EXPECT_THAT(header.GetUnsyncedTailChecksum(), Eq(0u));
+    EXPECT_THAT(header.GetHeaderChecksum(), Eq(0u));
+  }
+}
+
+class PortableFileBackedProtoLogTest : public ::testing::TestWithParam<bool> {
  protected:
   // Adds a user-defined default construct because a const member variable may
   // make the compiler accidentally delete the default constructor.
@@ -75,16 +111,24 @@
   bool compress_ = true;
   int32_t compression_level_ =
       PortableFileBackedProtoLog<DocumentProto>::kDefaultCompressionLevel;
+  uint32_t compression_threshold_bytes_ = PortableFileBackedProtoLog<
+      DocumentProto>::kDefaultCompressionThresholdBytes;
+  int32_t compression_mem_level_ = protobuf_ports::kDefaultMemLevel;
   int64_t max_proto_size_ = 256 * 1024;  // 256 KiB
+  bool enable_smaller_decompression_buffer_size_ = true;
 };
 
-TEST_F(PortableFileBackedProtoLogTest, Initialize) {
+TEST_P(PortableFileBackedProtoLogTest, Initialize) {
   ICING_ASSERT_OK_AND_ASSIGN(
       PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
       PortableFileBackedProtoLog<DocumentProto>::Create(
           &filesystem_, file_path_,
           PortableFileBackedProtoLog<DocumentProto>::Options(
-              compress_, max_proto_size_, compression_level_)));
+              compress_, max_proto_size_, compression_level_,
+              compression_threshold_bytes_, compression_mem_level_,
+              enable_smaller_decompression_buffer_size_,
+              /*enable_new_header_format_in=*/GetParam(),
+              /*enable_reusable_decompression_buffer=*/true)));
   EXPECT_THAT(create_result.proto_log, NotNull());
   EXPECT_FALSE(create_result.has_data_loss());
   EXPECT_FALSE(create_result.recalculated_checksum);
@@ -93,17 +137,47 @@
   ASSERT_THAT(PortableFileBackedProtoLog<DocumentProto>::Create(
                   &filesystem_, file_path_,
                   PortableFileBackedProtoLog<DocumentProto>::Options(
-                      !compress_, max_proto_size_, compression_level_)),
+                      !compress_, max_proto_size_, compression_level_,
+                      compression_threshold_bytes_, compression_mem_level_,
+                      enable_smaller_decompression_buffer_size_,
+                      /*enable_new_header_format_in=*/GetParam(),
+                      /*enable_reusable_decompression_buffer=*/true)),
               StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
 }
 
-TEST_F(PortableFileBackedProtoLogTest, InitializeValidatesOptions) {
+TEST_P(PortableFileBackedProtoLogTest, NewAndEmptyFileShouldFlushHeader) {
+  // Mock the filesystem to verify that DataSync is called.
+  auto mock_filesystem = std::make_unique<MockFilesystem>();
+  EXPECT_CALL(*mock_filesystem, DataSync(_)).Times(1);
+
+  {
+    ICING_ASSERT_OK_AND_ASSIGN(
+        PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
+        PortableFileBackedProtoLog<DocumentProto>::Create(
+            mock_filesystem.get(), file_path_,
+            PortableFileBackedProtoLog<DocumentProto>::Options(
+                compress_, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
+    EXPECT_THAT(create_result.proto_log, NotNull());
+    EXPECT_FALSE(create_result.has_data_loss());
+    EXPECT_FALSE(create_result.recalculated_checksum);
+  }
+}
+
+TEST_P(PortableFileBackedProtoLogTest, InitializeValidatesOptions) {
   // max_proto_size must be greater than 0
   int invalid_max_proto_size = 0;
   ASSERT_THAT(PortableFileBackedProtoLog<DocumentProto>::Create(
                   &filesystem_, file_path_,
                   PortableFileBackedProtoLog<DocumentProto>::Options(
-                      compress_, invalid_max_proto_size, compression_level_)),
+                      compress_, invalid_max_proto_size, compression_level_,
+                      compression_threshold_bytes_, compression_mem_level_,
+                      enable_smaller_decompression_buffer_size_,
+                      /*enable_new_header_format_in=*/GetParam(),
+                      /*enable_reusable_decompression_buffer=*/true)),
               StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
 
   // max_proto_size must be under 16 MiB
@@ -111,7 +185,11 @@
   ASSERT_THAT(PortableFileBackedProtoLog<DocumentProto>::Create(
                   &filesystem_, file_path_,
                   PortableFileBackedProtoLog<DocumentProto>::Options(
-                      compress_, invalid_max_proto_size, compression_level_)),
+                      compress_, invalid_max_proto_size, compression_level_,
+                      compression_threshold_bytes_, compression_mem_level_,
+                      enable_smaller_decompression_buffer_size_,
+                      /*enable_new_header_format_in=*/GetParam(),
+                      /*enable_reusable_decompression_buffer=*/true)),
               StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
 
   // compression_level must be between 0 and 9 inclusive
@@ -119,7 +197,11 @@
   ASSERT_THAT(PortableFileBackedProtoLog<DocumentProto>::Create(
                   &filesystem_, file_path_,
                   PortableFileBackedProtoLog<DocumentProto>::Options(
-                      compress_, max_proto_size_, invalid_compression_level)),
+                      compress_, max_proto_size_, invalid_compression_level,
+                      compression_threshold_bytes_, compression_mem_level_,
+                      enable_smaller_decompression_buffer_size_,
+                      /*enable_new_header_format_in=*/GetParam(),
+                      /*enable_reusable_decompression_buffer=*/true)),
               StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
 
   // compression_level must be between 0 and 9 inclusive
@@ -127,17 +209,51 @@
   ASSERT_THAT(PortableFileBackedProtoLog<DocumentProto>::Create(
                   &filesystem_, file_path_,
                   PortableFileBackedProtoLog<DocumentProto>::Options(
-                      compress_, max_proto_size_, invalid_compression_level)),
+                      compress_, max_proto_size_, invalid_compression_level,
+                      compression_threshold_bytes_, compression_mem_level_,
+                      enable_smaller_decompression_buffer_size_,
+                      /*enable_new_header_format_in=*/GetParam(),
+                      /*enable_reusable_decompression_buffer=*/true)),
               StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
+
+  // compression_mem_level must be between 1 and 9 inclusive
+  int invalid_compression_mem_level = 0;
+  ASSERT_THAT(
+      PortableFileBackedProtoLog<DocumentProto>::Create(
+          &filesystem_, file_path_,
+          PortableFileBackedProtoLog<DocumentProto>::Options(
+              compress_, max_proto_size_, compression_level_,
+              compression_threshold_bytes_, invalid_compression_mem_level,
+              enable_smaller_decompression_buffer_size_,
+              /*enable_new_header_format_in=*/GetParam(),
+              /*enable_reusable_decompression_buffer=*/true)),
+      StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
+
+  // compression_mem_level must be between 1 and 9 inclusive
+  invalid_compression_mem_level = 10;
+  ASSERT_THAT(
+      PortableFileBackedProtoLog<DocumentProto>::Create(
+          &filesystem_, file_path_,
+          PortableFileBackedProtoLog<DocumentProto>::Options(
+              compress_, max_proto_size_, compression_level_,
+              compression_threshold_bytes_, invalid_compression_mem_level,
+              enable_smaller_decompression_buffer_size_,
+              /*enable_new_header_format_in=*/GetParam(),
+              /*enable_reusable_decompression_buffer=*/true)),
+      StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
 }
 
-TEST_F(PortableFileBackedProtoLogTest, ReservedSpaceForHeader) {
+TEST_P(PortableFileBackedProtoLogTest, ReservedSpaceForHeader) {
   ICING_ASSERT_OK_AND_ASSIGN(
       PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
       PortableFileBackedProtoLog<DocumentProto>::Create(
           &filesystem_, file_path_,
           PortableFileBackedProtoLog<DocumentProto>::Options(
-              compress_, max_proto_size_, compression_level_)));
+              compress_, max_proto_size_, compression_level_,
+              compression_threshold_bytes_, compression_mem_level_,
+              enable_smaller_decompression_buffer_size_,
+              /*enable_new_header_format_in=*/GetParam(),
+              /*enable_reusable_decompression_buffer=*/true)));
 
   // With no protos written yet, the log should be minimum the size of the
   // reserved header space.
@@ -145,14 +261,18 @@
             PortableFileBackedProtoLog<DocumentProto>::kHeaderReservedBytes);
 }
 
-TEST_F(PortableFileBackedProtoLogTest, WriteProtoTooLarge) {
+TEST_P(PortableFileBackedProtoLogTest, WriteProtoTooLarge) {
   int max_proto_size = 1;
   ICING_ASSERT_OK_AND_ASSIGN(
       PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
       PortableFileBackedProtoLog<DocumentProto>::Create(
           &filesystem_, file_path_,
           PortableFileBackedProtoLog<DocumentProto>::Options(
-              compress_, max_proto_size, compression_level_)));
+              compress_, max_proto_size, compression_level_,
+              compression_threshold_bytes_, compression_mem_level_,
+              enable_smaller_decompression_buffer_size_,
+              /*enable_new_header_format_in=*/GetParam(),
+              /*enable_reusable_decompression_buffer=*/true)));
   auto proto_log = std::move(create_result.proto_log);
   ASSERT_FALSE(create_result.has_data_loss());
 
@@ -163,13 +283,17 @@
               StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
 }
 
-TEST_F(PortableFileBackedProtoLogTest, ReadProtoWrongKProtoMagic) {
+TEST_P(PortableFileBackedProtoLogTest, ReadProtoWrongKProtoMagic) {
   ICING_ASSERT_OK_AND_ASSIGN(
       PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
       PortableFileBackedProtoLog<DocumentProto>::Create(
           &filesystem_, file_path_,
           PortableFileBackedProtoLog<DocumentProto>::Options(
-              compress_, max_proto_size_, compression_level_)));
+              compress_, max_proto_size_, compression_level_,
+              compression_threshold_bytes_, compression_mem_level_,
+              enable_smaller_decompression_buffer_size_,
+              /*enable_new_header_format_in=*/GetParam(),
+              /*enable_reusable_decompression_buffer=*/true)));
   auto proto_log = std::move(create_result.proto_log);
   ASSERT_FALSE(create_result.has_data_loss());
 
@@ -195,7 +319,7 @@
               StatusIs(libtextclassifier3::StatusCode::INTERNAL));
 }
 
-TEST_F(PortableFileBackedProtoLogTest, ReadWriteUncompressedProto) {
+TEST_P(PortableFileBackedProtoLogTest, ReadWriteUncompressedProto) {
   int last_offset;
   {
     ICING_ASSERT_OK_AND_ASSIGN(
@@ -203,7 +327,11 @@
         PortableFileBackedProtoLog<DocumentProto>::Create(
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
-                /*compress_in=*/false, max_proto_size_, compression_level_)));
+                /*compress_in=*/false, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
     auto proto_log = std::move(create_result.proto_log);
     ASSERT_FALSE(create_result.has_data_loss());
 
@@ -250,7 +378,11 @@
         PortableFileBackedProtoLog<DocumentProto>::Create(
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
-                /*compress_in=*/false, max_proto_size_, compression_level_)));
+                /*compress_in=*/false, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
     auto recreated_proto_log = std::move(create_result.proto_log);
     ASSERT_FALSE(create_result.has_data_loss());
 
@@ -263,7 +395,7 @@
   }
 }
 
-TEST_F(PortableFileBackedProtoLogTest, ReadWriteCompressedProto) {
+TEST_P(PortableFileBackedProtoLogTest, ReadWriteCompressedProto) {
   int last_offset;
 
   {
@@ -272,7 +404,11 @@
         PortableFileBackedProtoLog<DocumentProto>::Create(
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
-                /*compress_in=*/true, max_proto_size_, compression_level_)));
+                /*compress_in=*/true, max_proto_size_, compression_level_,
+                /*compression_threshold_bytes_in=*/0, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
     auto proto_log = std::move(create_result.proto_log);
     ASSERT_FALSE(create_result.has_data_loss());
 
@@ -319,7 +455,11 @@
         PortableFileBackedProtoLog<DocumentProto>::Create(
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
-                /*compress_in=*/true, max_proto_size_, compression_level_)));
+                /*compress_in=*/true, max_proto_size_, compression_level_,
+                /*compression_threshold_bytes_in=*/0, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
     auto recreated_proto_log = std::move(create_result.proto_log);
     ASSERT_FALSE(create_result.has_data_loss());
 
@@ -332,7 +472,7 @@
   }
 }
 
-TEST_F(PortableFileBackedProtoLogTest, ReadWriteDifferentCompressionLevel) {
+TEST_P(PortableFileBackedProtoLogTest, ReadWriteDifferentCompressionLevel) {
   int document1_offset;
   int document2_offset;
   int document3_offset;
@@ -356,7 +496,11 @@
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
                 /*compress_in=*/true, max_proto_size_,
-                /*compression_level_in=*/3)));
+                /*compression_level_in=*/3,
+                /*compression_threshold_bytes_in=*/0, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
     auto proto_log = std::move(create_result.proto_log);
     ASSERT_FALSE(create_result.has_data_loss());
 
@@ -381,7 +525,11 @@
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
                 /*compress_in=*/true, max_proto_size_,
-                /*compression_level_in=*/9)));
+                /*compression_level_in=*/9,
+                /*compression_threshold_bytes_in=*/0, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
     auto recreated_proto_log = std::move(create_result.proto_log);
     ASSERT_FALSE(create_result.has_data_loss());
 
@@ -410,7 +558,11 @@
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
                 /*compress_in=*/true, max_proto_size_,
-                /*compression_level=*/0)));
+                /*compression_level_in=*/0,
+                /*compression_threshold_bytes_in=*/0, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
     auto recreated_proto_log = std::move(create_result.proto_log);
     ASSERT_FALSE(create_result.has_data_loss());
 
@@ -434,7 +586,233 @@
   }
 }
 
-TEST_F(PortableFileBackedProtoLogTest,
+TEST_P(PortableFileBackedProtoLogTest, ReadWriteDifferentCompressionMemLevel) {
+  int document1_offset;
+  int document2_offset;
+  int document3_offset;
+
+  // The first proto to write that's close to the max size. Leave some room for
+  // the rest of the proto properties.
+  std::string long_str(max_proto_size_ - 1024, 'a');
+  DocumentProto document1 = DocumentBuilder()
+                                .SetKey("namespace1", "uri1")
+                                .AddStringProperty("long_str", long_str)
+                                .Build();
+  DocumentProto document2 =
+      DocumentBuilder().SetKey("namespace2", "uri2").Build();
+  DocumentProto document3 =
+      DocumentBuilder().SetKey("namespace3", "uri3").Build();
+
+  {
+    ICING_ASSERT_OK_AND_ASSIGN(
+        PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
+        PortableFileBackedProtoLog<DocumentProto>::Create(
+            &filesystem_, file_path_,
+            PortableFileBackedProtoLog<DocumentProto>::Options(
+                /*compress_in=*/true, max_proto_size_, compression_level_,
+                compression_threshold_bytes_,
+                /*compression_mem_level_in=*/8,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
+    auto proto_log = std::move(create_result.proto_log);
+    ASSERT_FALSE(create_result.has_data_loss());
+
+    // Write the first proto
+    ICING_ASSERT_OK_AND_ASSIGN(document1_offset,
+                               proto_log->WriteProto(document1));
+
+    // Check that what we read is what we wrote
+    ASSERT_THAT(proto_log->ReadProto(document1_offset),
+                IsOkAndHolds(EqualsProto(document1)));
+
+    ICING_ASSERT_OK(proto_log->PersistToDisk());
+  }
+
+  // Make a new proto_log with the same file_path but different compression mem
+  // level, and make sure we can still read from and write to the same
+  // underlying file.
+  {
+    ICING_ASSERT_OK_AND_ASSIGN(
+        PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
+        PortableFileBackedProtoLog<DocumentProto>::Create(
+            &filesystem_, file_path_,
+            PortableFileBackedProtoLog<DocumentProto>::Options(
+                /*compress_in=*/true, max_proto_size_, compression_level_,
+                compression_threshold_bytes_,
+                /*compression_mem_level_in=*/1,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
+    auto recreated_proto_log = std::move(create_result.proto_log);
+    ASSERT_FALSE(create_result.has_data_loss());
+
+    // Check the first proto
+    ASSERT_THAT(recreated_proto_log->ReadProto(document1_offset),
+                IsOkAndHolds(EqualsProto(document1)));
+
+    // Write a second proto
+    ICING_ASSERT_OK_AND_ASSIGN(document2_offset,
+                               recreated_proto_log->WriteProto(document2));
+
+    ASSERT_GT(document2_offset, document1_offset);
+
+    // Check the second proto
+    ASSERT_THAT(recreated_proto_log->ReadProto(document2_offset),
+                IsOkAndHolds(EqualsProto(document2)));
+
+    ICING_ASSERT_OK(recreated_proto_log->PersistToDisk());
+  }
+
+  // One more time but with 9 compression mem level
+  {
+    ICING_ASSERT_OK_AND_ASSIGN(
+        PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
+        PortableFileBackedProtoLog<DocumentProto>::Create(
+            &filesystem_, file_path_,
+            PortableFileBackedProtoLog<DocumentProto>::Options(
+                /*compress_in=*/true, max_proto_size_, compression_level_,
+                compression_threshold_bytes_,
+                /*compression_mem_level_in=*/9,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
+    auto recreated_proto_log = std::move(create_result.proto_log);
+    ASSERT_FALSE(create_result.has_data_loss());
+
+    // Check the first proto
+    ASSERT_THAT(recreated_proto_log->ReadProto(document1_offset),
+                IsOkAndHolds(EqualsProto(document1)));
+
+    // Check the second proto
+    ASSERT_THAT(recreated_proto_log->ReadProto(document2_offset),
+                IsOkAndHolds(EqualsProto(document2)));
+
+    // Write a third proto
+    ICING_ASSERT_OK_AND_ASSIGN(document3_offset,
+                               recreated_proto_log->WriteProto(document3));
+
+    ASSERT_GT(document3_offset, document2_offset);
+
+    // Check the third proto
+    ASSERT_THAT(recreated_proto_log->ReadProto(document3_offset),
+                IsOkAndHolds(EqualsProto(document3)));
+  }
+}
+
+TEST_P(PortableFileBackedProtoLogTest,
+       ReadWriteEnableAndDisableSmallerDecompressionBufferSize) {
+  int document1_offset;
+  int document2_offset;
+  int document3_offset;
+
+  // The first proto to write that's close to the max size. Leave some room for
+  // the rest of the proto properties.
+  std::string long_str(max_proto_size_ - 1024, 'a');
+  DocumentProto document1 = DocumentBuilder()
+                                .SetKey("namespace1", "uri1")
+                                .AddStringProperty("long_str", long_str)
+                                .Build();
+  DocumentProto document2 =
+      DocumentBuilder().SetKey("namespace2", "uri2").Build();
+  DocumentProto document3 =
+      DocumentBuilder().SetKey("namespace3", "uri3").Build();
+
+  {
+    ICING_ASSERT_OK_AND_ASSIGN(
+        PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
+        PortableFileBackedProtoLog<DocumentProto>::Create(
+            &filesystem_, file_path_,
+            PortableFileBackedProtoLog<DocumentProto>::Options(
+                /*compress_in=*/true, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                /*enable_smaller_decompression_buffer_size_in=*/false,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
+    auto proto_log = std::move(create_result.proto_log);
+    ASSERT_FALSE(create_result.has_data_loss());
+
+    // Write the first proto
+    ICING_ASSERT_OK_AND_ASSIGN(document1_offset,
+                               proto_log->WriteProto(document1));
+
+    // Check that what we read is what we wrote
+    ASSERT_THAT(proto_log->ReadProto(document1_offset),
+                IsOkAndHolds(EqualsProto(document1)));
+
+    ICING_ASSERT_OK(proto_log->PersistToDisk());
+  }
+
+  // Make a new proto_log with the same file_path with smaller decompression
+  // buffer size enabled, and make sure we can still read from and write to the
+  // same underlying file.
+  {
+    ICING_ASSERT_OK_AND_ASSIGN(
+        PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
+        PortableFileBackedProtoLog<DocumentProto>::Create(
+            &filesystem_, file_path_,
+            PortableFileBackedProtoLog<DocumentProto>::Options(
+                /*compress_in=*/true, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                /*enable_smaller_decompression_buffer_size_in=*/true,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
+    auto recreated_proto_log = std::move(create_result.proto_log);
+    ASSERT_FALSE(create_result.has_data_loss());
+
+    // Check the first proto
+    ASSERT_THAT(recreated_proto_log->ReadProto(document1_offset),
+                IsOkAndHolds(EqualsProto(document1)));
+
+    // Write a second proto
+    ICING_ASSERT_OK_AND_ASSIGN(document2_offset,
+                               recreated_proto_log->WriteProto(document2));
+
+    ASSERT_GT(document2_offset, document1_offset);
+
+    // Check the second proto
+    ASSERT_THAT(recreated_proto_log->ReadProto(document2_offset),
+                IsOkAndHolds(EqualsProto(document2)));
+
+    ICING_ASSERT_OK(recreated_proto_log->PersistToDisk());
+  }
+
+  // One more time but with smaller decompression buffer size disabled
+  {
+    ICING_ASSERT_OK_AND_ASSIGN(
+        PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
+        PortableFileBackedProtoLog<DocumentProto>::Create(
+            &filesystem_, file_path_,
+            PortableFileBackedProtoLog<DocumentProto>::Options(
+                /*compress_in=*/true, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                /*enable_smaller_decompression_buffer_size_in=*/false,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
+    auto recreated_proto_log = std::move(create_result.proto_log);
+    ASSERT_FALSE(create_result.has_data_loss());
+
+    // Check the first proto
+    ASSERT_THAT(recreated_proto_log->ReadProto(document1_offset),
+                IsOkAndHolds(EqualsProto(document1)));
+
+    // Check the second proto
+    ASSERT_THAT(recreated_proto_log->ReadProto(document2_offset),
+                IsOkAndHolds(EqualsProto(document2)));
+
+    // Write a third proto
+    ICING_ASSERT_OK_AND_ASSIGN(document3_offset,
+                               recreated_proto_log->WriteProto(document3));
+
+    ASSERT_GT(document3_offset, document2_offset);
+
+    // Check the third proto
+    ASSERT_THAT(recreated_proto_log->ReadProto(document3_offset),
+                IsOkAndHolds(EqualsProto(document3)));
+  }
+}
+
+TEST_P(PortableFileBackedProtoLogTest,
        WriteDifferentCompressionLevelDifferentSizes) {
   int document_log_size_with_compression_3;
   int document_log_size_with_no_compression;
@@ -454,7 +832,11 @@
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
                 /*compress_in=*/true, max_proto_size_,
-                /*compression_level_in=*/3)));
+                /*compression_level_in=*/3, compression_threshold_bytes_,
+                compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
     auto proto_log = std::move(create_result.proto_log);
     ASSERT_FALSE(create_result.has_data_loss());
 
@@ -476,7 +858,11 @@
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
                 /*compress_in=*/true, max_proto_size_,
-                /*compression_level_in=*/0)));
+                /*compression_level_in=*/0, compression_threshold_bytes_,
+                compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
     auto proto_log = std::move(create_result.proto_log);
     ASSERT_FALSE(create_result.has_data_loss());
 
@@ -494,44 +880,505 @@
   }
 }
 
-TEST_F(PortableFileBackedProtoLogTest, CorruptHeader) {
+TEST_P(PortableFileBackedProtoLogTest, CompressionThreshold) {
+  // Create document1 with size of 1024 bytes.
+  int64_t document1_size = 1024;
+  int64_t document1_offset;
+  DocumentProto document1 =
+      DocumentBuilder()
+          .SetKey("namespace1", "uri1")
+          .AddStringProperty("long_str", std::string(990, 'a'))
+          .Build();
+  ASSERT_THAT(document1.ByteSizeLong(), Eq(document1_size));
+
+  // Create document2 with size of 900 bytes.
+  int64_t document2_size = 900;
+  int64_t document2_offset;
+  DocumentProto document2 =
+      DocumentBuilder()
+          .SetKey("namespace1", "uri1")
+          .AddStringProperty("long_str", std::string(866, 'a'))
+          .Build();
+  ASSERT_THAT(document2.ByteSizeLong(), Eq(document2_size));
+
+  {
+    // Create a proto_log with compression threshold of 1000 bytes
+    ICING_ASSERT_OK_AND_ASSIGN(
+        PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
+        PortableFileBackedProtoLog<DocumentProto>::Create(
+            &filesystem_, file_path_,
+            PortableFileBackedProtoLog<DocumentProto>::Options(
+                /*compress_in=*/true, max_proto_size_,
+                /*compression_level_in=*/3,
+                /*compression_threshold_bytes_in=*/1000, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
+    auto proto_log = std::move(create_result.proto_log);
+    ASSERT_FALSE(create_result.has_data_loss());
+
+    ICING_ASSERT_OK_AND_ASSIGN(
+        auto document1_write_result,
+        WriteProtoAndReturnOffsetAndSize(proto_log.get(), document1));
+    document1_offset = document1_write_result.first;
+    int64_t document1_size_written = document1_write_result.second;
+    // Check that what we read is what we wrote
+    ASSERT_THAT(proto_log->ReadProto(document1_offset),
+                IsOkAndHolds(EqualsProto(document1)));
+    // The document should be compressed since it's size is larger than the
+    // compression threshold.
+    ASSERT_LT(document1_size_written, document1_size);
+
+    ICING_ASSERT_OK_AND_ASSIGN(
+        auto document2_write_result,
+        WriteProtoAndReturnOffsetAndSize(proto_log.get(), document2));
+    document2_offset = document2_write_result.first;
+    int64_t document2_size_written = document2_write_result.second;
+    // Check that what we read is what we wrote
+    ASSERT_THAT(proto_log->ReadProto(document2_offset),
+                IsOkAndHolds(EqualsProto(document2)));
+    // The document should be uncompressed since it's size is smaller than the
+    // compression threshold.
+    ASSERT_GE(document2_size_written, document2_size);
+
+    ICING_ASSERT_OK(proto_log->PersistToDisk());
+  }
+
+  {
+    // Make a new proto_log with the same file_path, and make sure we can still
+    // read and write to the same underlying file.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
+        PortableFileBackedProtoLog<DocumentProto>::Create(
+            &filesystem_, file_path_,
+            PortableFileBackedProtoLog<DocumentProto>::Options(
+                /*compress_in=*/true, max_proto_size_,
+                /*compression_level_in=*/3,
+                /*compression_threshold_bytes_in=*/1000, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
+    auto recreated_proto_log = std::move(create_result.proto_log);
+    ASSERT_FALSE(create_result.has_data_loss());
+
+    // Check the first proto
+    ASSERT_THAT(recreated_proto_log->ReadProto(document1_offset),
+                IsOkAndHolds(EqualsProto(document1)));
+
+    // Check the second proto
+    ASSERT_THAT(recreated_proto_log->ReadProto(document2_offset),
+                IsOkAndHolds(EqualsProto(document2)));
+
+    // Write a third proto
+    DocumentProto document3 =
+        DocumentBuilder().SetKey("namespace3", "uri3").Build();
+    ICING_ASSERT_OK_AND_ASSIGN(int64_t document3_offset,
+                               recreated_proto_log->WriteProto(document3));
+    ASSERT_GT(document3_offset, document2_offset);
+    // Check the third proto
+    ASSERT_THAT(recreated_proto_log->ReadProto(document3_offset),
+                IsOkAndHolds(EqualsProto(document3)));
+  }
+}
+
+TEST_P(PortableFileBackedProtoLogTest, ChangingCompressionThresholdIsOk) {
+  // Create document1 with size of 1024 bytes.
+  int64_t document1_size = 1024;
+  int64_t document1_offset;
+  int64_t document1_size_written;
+  DocumentProto document1 =
+      DocumentBuilder()
+          .SetKey("namespace1", "uri1")
+          .AddStringProperty("long_str", std::string(990, 'a'))
+          .Build();
+  ASSERT_THAT(document1.ByteSizeLong(), Eq(document1_size));
+
+  // Create document2 with size of 900 bytes.
+  int64_t document2_size = 900;
+  int64_t document2_offset;
+  int64_t document2_size_written;
+  DocumentProto document2 =
+      DocumentBuilder()
+          .SetKey("namespace1", "uri1")
+          .AddStringProperty("long_str", std::string(866, 'a'))
+          .Build();
+  ASSERT_THAT(document2.ByteSizeLong(), Eq(document2_size));
+
+  {
+    // Create a proto_log with compression threshold of 1000 bytes
+    ICING_ASSERT_OK_AND_ASSIGN(
+        PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
+        PortableFileBackedProtoLog<DocumentProto>::Create(
+            &filesystem_, file_path_,
+            PortableFileBackedProtoLog<DocumentProto>::Options(
+                /*compress_in=*/true, max_proto_size_,
+                /*compression_level_in=*/3,
+                /*compression_threshold_bytes_in=*/1000, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
+    auto proto_log = std::move(create_result.proto_log);
+    ASSERT_FALSE(create_result.has_data_loss());
+
+    ICING_ASSERT_OK_AND_ASSIGN(
+        auto document1_write_result,
+        WriteProtoAndReturnOffsetAndSize(proto_log.get(), document1));
+    document1_offset = document1_write_result.first;
+    document1_size_written = document1_write_result.second;
+    // Check that what we read is what we wrote
+    ASSERT_THAT(proto_log->ReadProto(document1_offset),
+                IsOkAndHolds(EqualsProto(document1)));
+    // The document should be compressed since it's size is larger than the
+    // compression threshold.
+    ASSERT_LT(document1_size_written, document1_size);
+
+    ICING_ASSERT_OK_AND_ASSIGN(
+        auto document2_write_result,
+        WriteProtoAndReturnOffsetAndSize(proto_log.get(), document2));
+    document2_offset = document2_write_result.first;
+    document2_size_written = document2_write_result.second;
+    // Check that what we read is what we wrote
+    ASSERT_THAT(proto_log->ReadProto(document2_offset),
+                IsOkAndHolds(EqualsProto(document2)));
+    // The document should be uncompressed since it's size is smaller than the
+    // compression threshold.
+    ASSERT_GE(document2_size_written, document2_size);
+
+    ICING_ASSERT_OK(proto_log->PersistToDisk());
+  }
+
+  {
+    // Make a new proto_log with the same file_path with a different compression
+    // threshold, and make sure we can still read and write to the same
+    // underlying file.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
+        PortableFileBackedProtoLog<DocumentProto>::Create(
+            &filesystem_, file_path_,
+            PortableFileBackedProtoLog<DocumentProto>::Options(
+                /*compress_in=*/true, max_proto_size_,
+                /*compression_level_in=*/3,
+                /*compression_threshold_bytes_in=*/100, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
+    auto recreated_proto_log = std::move(create_result.proto_log);
+    ASSERT_FALSE(create_result.has_data_loss());
+
+    // Check the first proto
+    ASSERT_THAT(recreated_proto_log->ReadProto(document1_offset),
+                IsOkAndHolds(EqualsProto(document1)));
+
+    // Check the second proto
+    ASSERT_THAT(recreated_proto_log->ReadProto(document2_offset),
+                IsOkAndHolds(EqualsProto(document2)));
+
+    // Write document1 again as the third proto;
+    ICING_ASSERT_OK_AND_ASSIGN(
+        auto document3_write_result,
+        WriteProtoAndReturnOffsetAndSize(recreated_proto_log.get(), document1));
+    int64_t document3_offset = document3_write_result.first;
+    int64_t document3_size_written = document3_write_result.second;
+    // With this new compression threshold, document1 should still be
+    // compressed.
+    ASSERT_EQ(document3_size_written, document1_size_written);
+    ASSERT_GT(document3_offset, document2_offset);
+    // Check the third proto
+    ASSERT_THAT(recreated_proto_log->ReadProto(document3_offset),
+                IsOkAndHolds(EqualsProto(document1)));
+
+    // Write document2 again as the fourth proto;
+    ICING_ASSERT_OK_AND_ASSIGN(
+        auto document4_write_result,
+        WriteProtoAndReturnOffsetAndSize(recreated_proto_log.get(), document2));
+    int64_t document4_offset = document4_write_result.first;
+    int64_t document4_size_written = document4_write_result.second;
+    // With this new compression threshold, document2 should now be compressed.
+    ASSERT_LT(document4_size_written, document2_size_written);
+    ASSERT_LT(document4_size_written, document2_size);
+    ASSERT_GT(document4_offset, document3_offset);
+    // Check the fourth proto
+    ASSERT_THAT(recreated_proto_log->ReadProto(document4_offset),
+                IsOkAndHolds(EqualsProto(document2)));
+  }
+}
+
+TEST_P(PortableFileBackedProtoLogTest, CorruptedLegacyHeaderSection) {
+  // This test simulates the following scenario:
+  // - Log checksum matches.
+  // - Unsynced tail checksum matches.
+  // - Legacy header checksum doesn't match.
+  // - Header checksum matches.
+  //
+  // Expected behavior: full data loss.
   {
     ICING_ASSERT_OK_AND_ASSIGN(
         PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
         PortableFileBackedProtoLog<DocumentProto>::Create(
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
-                compress_, max_proto_size_, compression_level_)));
+                compress_, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
     auto recreated_proto_log = std::move(create_result.proto_log);
     EXPECT_FALSE(create_result.has_data_loss());
   }
 
-  int corrupt_checksum = 24;
+  int corrupt_legacy_header_checksum = 24;
 
-  // Write the corrupted header
-  Header header = ReadHeader(filesystem_, file_path_);
-  header.SetHeaderChecksum(corrupt_checksum);
-  WriteHeader(filesystem_, file_path_, header);
+  // Manually overwrite the legacy header checksum to simulate legacy header
+  // section is corrupted.
+  Header corrupted_header = ReadHeader(filesystem_, file_path_);
+  corrupted_header.SetLegacyHeaderChecksum(corrupt_legacy_header_checksum);
+  if (GetParam()) {
+    corrupted_header.SetHeaderChecksum(
+        corrupted_header.CalculateHeaderChecksum().Get());
+  }
+  WriteHeader(filesystem_, file_path_, corrupted_header);
 
   {
     // Reinitialize the same proto_log
     ASSERT_THAT(PortableFileBackedProtoLog<DocumentProto>::Create(
                     &filesystem_, file_path_,
                     PortableFileBackedProtoLog<DocumentProto>::Options(
-                        compress_, max_proto_size_, compression_level_)),
+                        compress_, max_proto_size_, compression_level_,
+                        compression_threshold_bytes_, compression_mem_level_,
+                        enable_smaller_decompression_buffer_size_,
+                        /*enable_new_header_format_in=*/GetParam(),
+                        /*enable_reusable_decompression_buffer=*/true)),
                 StatusIs(libtextclassifier3::StatusCode::INTERNAL,
-                         HasSubstr("Invalid header checksum")));
+                         HasSubstr("Invalid legacy header checksum")));
   }
 }
 
-TEST_F(PortableFileBackedProtoLogTest, DifferentMagic) {
+TEST_P(PortableFileBackedProtoLogTest, CorruptedHeaderChecksum) {
+  if (!GetParam()) {
+    GTEST_SKIP() << "This test is only relevant to the new header format.";
+  }
+
+  // This test simulates the following scenario:
+  // - Log checksum matches.
+  // - Unsynced tail checksum matches.
+  // - Legacy header checksum matches.
+  // - Header checksum doesn't match.
+  //
+  // Expected behavior:
+  // - Due to compatibility issue, PortableFileBackedProtoLog should assume that
+  //   the existing dataset is generated by an older version which does not have
+  //   the new header checksum logic.
+  // - Should only throw away the data in the unsynced tail section.
+  // - Related checksums should be recomputed and rewritten during
+  //   initialization.
+
+  DocumentProto document1 =
+      DocumentBuilder().SetKey("namespace", "uri1").Build();
+  DocumentProto document2 =
+      DocumentBuilder().SetKey("namespace", "uri2").Build();
+  DocumentProto document3 =
+      DocumentBuilder().SetKey("namespace", "uri3").Build();
+  DocumentProto document4 =
+      DocumentBuilder().SetKey("namespace", "uri4").Build();
+  DocumentProto document5 =
+      DocumentBuilder().SetKey("namespace", "uri5").Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
+      PortableFileBackedProtoLog<DocumentProto>::Create(
+          &filesystem_, file_path_,
+          PortableFileBackedProtoLog<DocumentProto>::Options(
+              compress_, max_proto_size_, compression_level_,
+              compression_threshold_bytes_, compression_mem_level_,
+              enable_smaller_decompression_buffer_size_,
+              /*enable_new_header_format_in=*/GetParam(),
+              /*enable_reusable_decompression_buffer=*/true)));
+  auto proto_log = std::move(create_result.proto_log);
+  EXPECT_FALSE(create_result.has_data_loss());
+
+  // Write 3 protos to the log and call PersistToDisk to make them into log
+  // checksummed section. Record the header at this state.
+  ICING_ASSERT_OK(proto_log->WriteProto(document1));
+  ICING_ASSERT_OK(proto_log->WriteProto(document2));
+  ICING_ASSERT_OK(proto_log->WriteProto(document3));
+  ICING_ASSERT_OK(proto_log->PersistToDisk());
+  EXPECT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Crc32(3230460271u)));
+
+  // Read the header for the state of doc1-to-doc3.
+  Header header_doc1_to_doc3 = ReadHeader(filesystem_, file_path_);
+  // Sanity check that the unsynced tail checksum is 0.
+  ASSERT_THAT(header_doc1_to_doc3.GetUnsyncedTailChecksum(), Eq(0));
+
+  // Write 2 more protos and keep them in the unsynced tail section (don't call
+  // PersistToDisk, UpdateChecksum, or destructor).
+  ICING_ASSERT_OK(proto_log->WriteProto(document4));
+  ICING_ASSERT_OK(proto_log->WriteProto(document5));
+  EXPECT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Crc32(705827520u)));
+
+  // Read the header for the state of doc1-to-doc5.
+  Header header_doc1_to_doc5 = ReadHeader(filesystem_, file_path_);
+  // Sanity check:
+  // - The legacy and log checksums are same as doc1-to-doc3.
+  // - The unsynced tail checksum and header checksum are different from state
+  //   doc1-to-doc3.
+  ASSERT_THAT(header_doc1_to_doc5.GetLegacyHeaderChecksum(),
+              Eq(header_doc1_to_doc3.GetLegacyHeaderChecksum()));
+  ASSERT_THAT(header_doc1_to_doc5.GetLogChecksum(),
+              Eq(header_doc1_to_doc3.GetLogChecksum()));
+  ASSERT_THAT(header_doc1_to_doc5.GetUnsyncedTailChecksum(),
+              Ne(header_doc1_to_doc3.GetUnsyncedTailChecksum()));
+  ASSERT_THAT(header_doc1_to_doc5.GetHeaderChecksum(),
+              Ne(header_doc1_to_doc3.GetHeaderChecksum()));
+
+  // Manually overwrite the header checksum.
+  Header corrupted_header = header_doc1_to_doc5;
+  corrupted_header.SetHeaderChecksum(corrupted_header.GetHeaderChecksum() +
+                                     12345);
+  WriteHeader(filesystem_, file_path_, corrupted_header);
+
+  // Initialize another instance with the same proto file. This should only
+  // trigger partial data loss. IOW, we only throw away the data in the unsynced
+  // tail section, so the data state is reverted back to doc1-to-doc3.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result2,
+      PortableFileBackedProtoLog<DocumentProto>::Create(
+          &filesystem_, file_path_,
+          PortableFileBackedProtoLog<DocumentProto>::Options(
+              compress_, max_proto_size_, compression_level_,
+              compression_threshold_bytes_, compression_mem_level_,
+              enable_smaller_decompression_buffer_size_,
+              /*enable_new_header_format_in=*/GetParam(),
+              /*enable_reusable_decompression_buffer=*/true)));
+  EXPECT_THAT(create_result2.data_loss, Eq(DataLoss::PARTIAL));
+
+  // Check the overall checksum.
+  EXPECT_THAT(create_result2.proto_log->GetChecksum(),
+              IsOkAndHolds(Crc32(3230460271u)));
+
+  // Manually check the header.
+  Header header_after = ReadHeader(filesystem_, file_path_);
+  EXPECT_THAT(header_after.GetUnsyncedTailChecksum(),
+              Eq(header_doc1_to_doc3.GetUnsyncedTailChecksum()));
+  EXPECT_THAT(header_after.GetHeaderChecksum(),
+              Eq(header_doc1_to_doc3.GetHeaderChecksum()));
+}
+
+TEST_P(PortableFileBackedProtoLogTest, CorruptedUnsyncedTailChecksum) {
+  if (!GetParam()) {
+    GTEST_SKIP() << "This test is only relevant to the new header format.";
+  }
+
+  // This test simulates the following scenario:
+  // - Log checksum matches.
+  // - Unsynced tail checksum doesn't match.
+  // - Legacy header checksum matches.
+  // - Header checksum matches.
+  //
+  // Expected behavior:
+  // - Should only throw away the data in the unsynced tail section.
+  // - Related checksums should be recomputed and rewritten during
+  //   initialization.
+
+  DocumentProto document1 =
+      DocumentBuilder().SetKey("namespace", "uri1").Build();
+  DocumentProto document2 =
+      DocumentBuilder().SetKey("namespace", "uri2").Build();
+  DocumentProto document3 =
+      DocumentBuilder().SetKey("namespace", "uri3").Build();
+  DocumentProto document4 =
+      DocumentBuilder().SetKey("namespace", "uri4").Build();
+  DocumentProto document5 =
+      DocumentBuilder().SetKey("namespace", "uri5").Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
+      PortableFileBackedProtoLog<DocumentProto>::Create(
+          &filesystem_, file_path_,
+          PortableFileBackedProtoLog<DocumentProto>::Options(
+              compress_, max_proto_size_, compression_level_,
+              compression_threshold_bytes_, compression_mem_level_,
+              enable_smaller_decompression_buffer_size_,
+              /*enable_new_header_format_in=*/GetParam(),
+              /*enable_reusable_decompression_buffer=*/true)));
+  auto proto_log = std::move(create_result.proto_log);
+  EXPECT_FALSE(create_result.has_data_loss());
+
+  // Write 3 protos to the log and call PersistToDisk to make them into log
+  // checksummed section. Record the header at this state.
+  ICING_ASSERT_OK(proto_log->WriteProto(document1));
+  ICING_ASSERT_OK(proto_log->WriteProto(document2));
+  ICING_ASSERT_OK(proto_log->WriteProto(document3));
+  ICING_ASSERT_OK(proto_log->PersistToDisk());
+  EXPECT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Crc32(3230460271u)));
+
+  // Read the header for the state of doc1-to-doc3.
+  Header header_doc1_to_doc3 = ReadHeader(filesystem_, file_path_);
+  // Sanity check that the unsynced tail checksum is 0.
+  ASSERT_THAT(header_doc1_to_doc3.GetUnsyncedTailChecksum(), Eq(0));
+
+  // Write 2 more protos and keep them in the unsynced tail section (don't call
+  // PersistToDisk, UpdateChecksum, or destructor).
+  ICING_ASSERT_OK(proto_log->WriteProto(document4));
+  ICING_ASSERT_OK(proto_log->WriteProto(document5));
+  EXPECT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Crc32(705827520u)));
+
+  // Read the header for the state of doc1-to-doc5.
+  Header header_doc1_to_doc5 = ReadHeader(filesystem_, file_path_);
+  // Sanity check that the unsynced tail checksum and header checksum are
+  // different from state doc1-to-doc3.
+  ASSERT_THAT(header_doc1_to_doc5.GetUnsyncedTailChecksum(),
+              Ne(header_doc1_to_doc3.GetUnsyncedTailChecksum()));
+  ASSERT_THAT(header_doc1_to_doc5.GetHeaderChecksum(),
+              Ne(header_doc1_to_doc3.GetHeaderChecksum()));
+
+  // Manually overwrite the unsynced tail checksum and recompute header
+  // checksum.
+  Header corrupted_header = header_doc1_to_doc5;
+  corrupted_header.SetUnsyncedTailChecksum(
+      corrupted_header.GetUnsyncedTailChecksum() + 12345);
+  corrupted_header.SetHeaderChecksum(
+      corrupted_header.CalculateHeaderChecksum().Get());
+  WriteHeader(filesystem_, file_path_, corrupted_header);
+
+  // Initialize another instance with the same proto file. This should only
+  // trigger partial data loss. IOW, we only throw away the data in the unsynced
+  // tail section, so the data state is reverted back to doc1-to-doc3.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result2,
+      PortableFileBackedProtoLog<DocumentProto>::Create(
+          &filesystem_, file_path_,
+          PortableFileBackedProtoLog<DocumentProto>::Options(
+              compress_, max_proto_size_, compression_level_,
+              compression_threshold_bytes_, compression_mem_level_,
+              enable_smaller_decompression_buffer_size_,
+              /*enable_new_header_format_in=*/GetParam(),
+              /*enable_reusable_decompression_buffer=*/true)));
+  EXPECT_THAT(create_result2.data_loss, Eq(DataLoss::PARTIAL));
+
+  // Check the overall checksum.
+  EXPECT_THAT(create_result2.proto_log->GetChecksum(),
+              IsOkAndHolds(Crc32(3230460271u)));
+
+  // Manually check the header.
+  Header header_after = ReadHeader(filesystem_, file_path_);
+  EXPECT_THAT(header_after.GetUnsyncedTailChecksum(),
+              Eq(header_doc1_to_doc3.GetUnsyncedTailChecksum()));
+  EXPECT_THAT(header_after.GetHeaderChecksum(),
+              Eq(header_doc1_to_doc3.GetHeaderChecksum()));
+}
+
+TEST_P(PortableFileBackedProtoLogTest, DifferentMagic) {
   {
     ICING_ASSERT_OK_AND_ASSIGN(
         PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
         PortableFileBackedProtoLog<DocumentProto>::Create(
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
-                compress_, max_proto_size_, compression_level_)));
+                compress_, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
     auto recreated_proto_log = std::move(create_result.proto_log);
     EXPECT_FALSE(create_result.has_data_loss());
 
@@ -547,16 +1394,22 @@
 
   {
     // Reinitialize the same proto_log
-    ASSERT_THAT(PortableFileBackedProtoLog<DocumentProto>::Create(
-                    &filesystem_, file_path_,
-                    PortableFileBackedProtoLog<DocumentProto>::Options(
-                        compress_, max_proto_size_, compression_level_)),
-                StatusIs(libtextclassifier3::StatusCode::INTERNAL,
-                         HasSubstr("Invalid header kMagic")));
+    ASSERT_THAT(
+        PortableFileBackedProtoLog<DocumentProto>::Create(
+            &filesystem_, file_path_,
+            PortableFileBackedProtoLog<DocumentProto>::Options(
+                compress_, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)),
+        StatusIs(
+            libtextclassifier3::StatusCode::INTERNAL,
+            HasSubstr("Invalid header magic for PortableFileBackedProtoLog")));
   }
 }
 
-TEST_F(PortableFileBackedProtoLogTest,
+TEST_P(PortableFileBackedProtoLogTest,
        UnableToDetectCorruptContentWithoutDirtyBit) {
   // This is intentional that we can't detect corruption. We're trading off
   // earlier corruption detection for lower initialization latency. By not
@@ -573,7 +1426,11 @@
         PortableFileBackedProtoLog<DocumentProto>::Create(
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
-                compress_, max_proto_size_, compression_level_)));
+                compress_, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
     auto proto_log = std::move(create_result.proto_log);
     EXPECT_FALSE(create_result.has_data_loss());
 
@@ -600,7 +1457,11 @@
         PortableFileBackedProtoLog<DocumentProto>::Create(
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
-                compress_, max_proto_size_, compression_level_)));
+                compress_, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
     auto proto_log = std::move(create_result.proto_log);
     EXPECT_FALSE(create_result.has_data_loss());
     EXPECT_THAT(create_result.data_loss, Eq(DataLoss::NONE));
@@ -614,7 +1475,7 @@
   }
 }
 
-TEST_F(PortableFileBackedProtoLogTest,
+TEST_P(PortableFileBackedProtoLogTest,
        DetectAndThrowOutCorruptContentWithDirtyBit) {
   {
     ICING_ASSERT_OK_AND_ASSIGN(
@@ -622,7 +1483,11 @@
         PortableFileBackedProtoLog<DocumentProto>::Create(
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
-                compress_, max_proto_size_, compression_level_)));
+                compress_, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
     auto proto_log = std::move(create_result.proto_log);
     ASSERT_FALSE(create_result.has_data_loss());
 
@@ -657,7 +1522,7 @@
 
     // Set dirty bit to true to reflect that something changed in the log.
     header.SetDirtyFlag(true);
-    header.SetHeaderChecksum(header.CalculateHeaderChecksum());
+    header.UpdateHeaderChecksums(/*enable_new_header_format=*/GetParam());
 
     WriteHeader(filesystem_, file_path_, header);
   }
@@ -668,7 +1533,11 @@
         PortableFileBackedProtoLog<DocumentProto>::Create(
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
-                compress_, max_proto_size_, compression_level_)));
+                compress_, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
     auto proto_log = std::move(create_result.proto_log);
     EXPECT_TRUE(create_result.has_data_loss());
     EXPECT_THAT(create_result.data_loss, Eq(DataLoss::COMPLETE));
@@ -687,7 +1556,7 @@
   }
 }
 
-TEST_F(PortableFileBackedProtoLogTest, DirtyBitFalseAlarmKeepsData) {
+TEST_P(PortableFileBackedProtoLogTest, DirtyBitFalseAlarmKeepsData) {
   DocumentProto document =
       DocumentBuilder().SetKey("namespace1", "uri1").Build();
   int64_t document_offset;
@@ -697,7 +1566,11 @@
         PortableFileBackedProtoLog<DocumentProto>::Create(
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
-                compress_, max_proto_size_, compression_level_)));
+                compress_, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
     auto proto_log = std::move(create_result.proto_log);
     ASSERT_FALSE(create_result.has_data_loss());
 
@@ -716,7 +1589,7 @@
     // Simulate the dirty flag set as true, but no data has been changed yet.
     // Maybe we crashed between writing the dirty flag and erasing a proto.
     header.SetDirtyFlag(true);
-    header.SetHeaderChecksum(header.CalculateHeaderChecksum());
+    header.UpdateHeaderChecksums(/*enable_new_header_format=*/GetParam());
 
     WriteHeader(filesystem_, file_path_, header);
   }
@@ -727,7 +1600,11 @@
         PortableFileBackedProtoLog<DocumentProto>::Create(
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
-                compress_, max_proto_size_, compression_level_)));
+                compress_, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
     auto proto_log = std::move(create_result.proto_log);
     EXPECT_FALSE(create_result.has_data_loss());
 
@@ -744,7 +1621,7 @@
   }
 }
 
-TEST_F(PortableFileBackedProtoLogTest,
+TEST_P(PortableFileBackedProtoLogTest,
        PersistToDiskKeepsPersistedDataAndTruncatesExtraData) {
   DocumentProto document1 =
       DocumentBuilder().SetKey("namespace1", "uri1").Build();
@@ -759,7 +1636,11 @@
         PortableFileBackedProtoLog<DocumentProto>::Create(
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
-                compress_, max_proto_size_, compression_level_)));
+                compress_, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
     auto proto_log = std::move(create_result.proto_log);
     ASSERT_FALSE(create_result.has_data_loss());
 
@@ -805,7 +1686,11 @@
         PortableFileBackedProtoLog<DocumentProto>::Create(
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
-                compress_, max_proto_size_, compression_level_)));
+                compress_, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
     auto proto_log = std::move(create_result.proto_log);
     ASSERT_TRUE(create_result.has_data_loss());
     ASSERT_THAT(create_result.data_loss, Eq(DataLoss::PARTIAL));
@@ -822,7 +1707,83 @@
   }
 }
 
-TEST_F(PortableFileBackedProtoLogTest,
+TEST_P(PortableFileBackedProtoLogTest, UnsyncedTailChecksumKeepsData) {
+  if (!GetParam()) {
+    GTEST_SKIP() << "This test is only relevant to the new header format.";
+  }
+
+  DocumentProto document1 =
+      DocumentBuilder().SetKey("namespace", "uri1").Build();
+  DocumentProto document2 =
+      DocumentBuilder().SetKey("namespace", "uri2").Build();
+  DocumentProto document3 =
+      DocumentBuilder().SetKey("namespace", "uri3").Build();
+  DocumentProto document4 =
+      DocumentBuilder().SetKey("namespace", "uri4").Build();
+  DocumentProto document5 =
+      DocumentBuilder().SetKey("namespace", "uri5").Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
+      PortableFileBackedProtoLog<DocumentProto>::Create(
+          &filesystem_, file_path_,
+          PortableFileBackedProtoLog<DocumentProto>::Options(
+              compress_, max_proto_size_, compression_level_,
+              compression_threshold_bytes_, compression_mem_level_,
+              enable_smaller_decompression_buffer_size_,
+              /*enable_new_header_format_in=*/GetParam(),
+              /*enable_reusable_decompression_buffer=*/true)));
+  auto proto_log = std::move(create_result.proto_log);
+  EXPECT_FALSE(create_result.has_data_loss());
+
+  // Write 3 protos to the log and call PersistToDisk to make them into log
+  // checksummed section. Record the header at this state.
+  ICING_ASSERT_OK(proto_log->WriteProto(document1));
+  ICING_ASSERT_OK(proto_log->WriteProto(document2));
+  ICING_ASSERT_OK(proto_log->WriteProto(document3));
+  ICING_ASSERT_OK(proto_log->PersistToDisk());
+  EXPECT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Crc32(3230460271u)));
+
+  // Write 2 more protos and keep them in the unsynced tail section (don't call
+  // PersistToDisk, UpdateChecksum, or destructor).
+  ICING_ASSERT_OK(proto_log->WriteProto(document4));
+  ICING_ASSERT_OK(proto_log->WriteProto(document5));
+  EXPECT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Crc32(705827520u)));
+
+  // Read the header for the state of doc1-to-doc5.
+  Header header_doc1_to_doc5 = ReadHeader(filesystem_, file_path_);
+
+  // Initialize another instance with the same proto file. Thanks to the
+  // unsynced tail checksum, document4 and document5 should be kept.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result2,
+      PortableFileBackedProtoLog<DocumentProto>::Create(
+          &filesystem_, file_path_,
+          PortableFileBackedProtoLog<DocumentProto>::Options(
+              compress_, max_proto_size_, compression_level_,
+              compression_threshold_bytes_, compression_mem_level_,
+              enable_smaller_decompression_buffer_size_,
+              /*enable_new_header_format_in=*/GetParam(),
+              /*enable_reusable_decompression_buffer=*/true)));
+  EXPECT_FALSE(create_result2.has_data_loss());
+
+  // Check the overall checksum. Should remain the same.
+  EXPECT_THAT(create_result2.proto_log->GetChecksum(),
+              IsOkAndHolds(Crc32(705827520u)));
+
+  // Manually check the header. Should remain the same.
+  Header header_after = ReadHeader(filesystem_, file_path_);
+  EXPECT_THAT(header_after.GetLegacyHeaderChecksum(),
+              Eq(header_doc1_to_doc5.GetLegacyHeaderChecksum()));
+  EXPECT_THAT(header_after.GetRewindOffset(),
+              Eq(header_doc1_to_doc5.GetRewindOffset()));
+  EXPECT_THAT(header_after.GetUnsyncedTailChecksum(),
+              Eq(header_doc1_to_doc5.GetUnsyncedTailChecksum()));
+  EXPECT_THAT(header_after.GetHeaderChecksum(),
+              Eq(header_doc1_to_doc5.GetHeaderChecksum()));
+}
+
+TEST_P(PortableFileBackedProtoLogTest,
        DirtyBitIsFalseAfterPutAndPersistToDisk) {
   {
     ICING_ASSERT_OK_AND_ASSIGN(
@@ -830,7 +1791,11 @@
         PortableFileBackedProtoLog<DocumentProto>::Create(
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
-                compress_, max_proto_size_, compression_level_)));
+                compress_, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
     auto proto_log = std::move(create_result.proto_log);
     ASSERT_FALSE(create_result.has_data_loss());
 
@@ -853,7 +1818,11 @@
         PortableFileBackedProtoLog<DocumentProto>::Create(
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
-                compress_, max_proto_size_, compression_level_)));
+                compress_, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
 
     // We previously persisted to disk so everything should be in a perfect
     // state.
@@ -865,7 +1834,7 @@
   }
 }
 
-TEST_F(PortableFileBackedProtoLogTest,
+TEST_P(PortableFileBackedProtoLogTest,
        DirtyBitIsFalseAfterDeleteAndPersistToDisk) {
   {
     ICING_ASSERT_OK_AND_ASSIGN(
@@ -873,7 +1842,11 @@
         PortableFileBackedProtoLog<DocumentProto>::Create(
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
-                compress_, max_proto_size_, compression_level_)));
+                compress_, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
     auto proto_log = std::move(create_result.proto_log);
     ASSERT_FALSE(create_result.has_data_loss());
 
@@ -897,7 +1870,11 @@
         PortableFileBackedProtoLog<DocumentProto>::Create(
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
-                compress_, max_proto_size_, compression_level_)));
+                compress_, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
 
     // We previously persisted to disk so everything should be in a perfect
     // state.
@@ -909,14 +1886,18 @@
   }
 }
 
-TEST_F(PortableFileBackedProtoLogTest, DirtyBitIsFalseAfterPutAndDestructor) {
+TEST_P(PortableFileBackedProtoLogTest, DirtyBitIsFalseAfterPutAndDestructor) {
   {
     ICING_ASSERT_OK_AND_ASSIGN(
         PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
         PortableFileBackedProtoLog<DocumentProto>::Create(
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
-                compress_, max_proto_size_, compression_level_)));
+                compress_, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
     auto proto_log = std::move(create_result.proto_log);
     ASSERT_FALSE(create_result.has_data_loss());
 
@@ -941,7 +1922,11 @@
         PortableFileBackedProtoLog<DocumentProto>::Create(
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
-                compress_, max_proto_size_, compression_level_)));
+                compress_, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
 
     // We previously persisted to disk so everything should be in a perfect
     // state.
@@ -953,7 +1938,7 @@
   }
 }
 
-TEST_F(PortableFileBackedProtoLogTest,
+TEST_P(PortableFileBackedProtoLogTest,
        DirtyBitIsFalseAfterDeleteAndDestructor) {
   {
     ICING_ASSERT_OK_AND_ASSIGN(
@@ -961,7 +1946,11 @@
         PortableFileBackedProtoLog<DocumentProto>::Create(
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
-                compress_, max_proto_size_, compression_level_)));
+                compress_, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
     auto proto_log = std::move(create_result.proto_log);
     ASSERT_FALSE(create_result.has_data_loss());
 
@@ -987,7 +1976,11 @@
         PortableFileBackedProtoLog<DocumentProto>::Create(
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
-                compress_, max_proto_size_, compression_level_)));
+                compress_, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
 
     // We previously persisted to disk so everything should be in a perfect
     // state.
@@ -999,7 +1992,7 @@
   }
 }
 
-TEST_F(PortableFileBackedProtoLogTest, Iterator) {
+TEST_P(PortableFileBackedProtoLogTest, Iterator) {
   DocumentProto document1 =
       DocumentBuilder().SetKey("namespace", "uri1").Build();
   DocumentProto document2 =
@@ -1010,7 +2003,11 @@
       PortableFileBackedProtoLog<DocumentProto>::Create(
           &filesystem_, file_path_,
           PortableFileBackedProtoLog<DocumentProto>::Options(
-              compress_, max_proto_size_, compression_level_)));
+              compress_, max_proto_size_, compression_level_,
+              compression_threshold_bytes_, compression_mem_level_,
+              enable_smaller_decompression_buffer_size_,
+              /*enable_new_header_format_in=*/GetParam(),
+              /*enable_reusable_decompression_buffer=*/true)));
   auto proto_log = std::move(create_result.proto_log);
   ASSERT_FALSE(create_result.has_data_loss());
 
@@ -1040,7 +2037,7 @@
   }
 }
 
-TEST_F(PortableFileBackedProtoLogTest, UpdateChecksum) {
+TEST_P(PortableFileBackedProtoLogTest, UpdateChecksum) {
   DocumentProto document = DocumentBuilder().SetKey("namespace", "uri").Build();
   Crc32 checksum;
 
@@ -1050,7 +2047,11 @@
         PortableFileBackedProtoLog<DocumentProto>::Create(
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
-                compress_, max_proto_size_, compression_level_)));
+                compress_, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
     auto proto_log = std::move(create_result.proto_log);
     ASSERT_FALSE(create_result.has_data_loss());
 
@@ -1070,7 +2071,11 @@
         PortableFileBackedProtoLog<DocumentProto>::Create(
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
-                compress_, max_proto_size_, compression_level_)));
+                compress_, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
     auto proto_log = std::move(create_result.proto_log);
     ASSERT_FALSE(create_result.has_data_loss());
 
@@ -1091,7 +2096,7 @@
   }
 }
 
-TEST_F(PortableFileBackedProtoLogTest, EraseProtoShouldSetZero) {
+TEST_P(PortableFileBackedProtoLogTest, EraseProtoShouldSetZero) {
   DocumentProto document1 =
       DocumentBuilder().SetKey("namespace", "uri1").Build();
 
@@ -1100,7 +2105,11 @@
       PortableFileBackedProtoLog<DocumentProto>::Create(
           &filesystem_, file_path_,
           PortableFileBackedProtoLog<DocumentProto>::Options(
-              compress_, max_proto_size_, compression_level_)));
+              compress_, max_proto_size_, compression_level_,
+              compression_threshold_bytes_, compression_mem_level_,
+              enable_smaller_decompression_buffer_size_,
+              /*enable_new_header_format_in=*/GetParam(),
+              /*enable_reusable_decompression_buffer=*/true)));
   auto proto_log = std::move(create_result.proto_log);
   ASSERT_FALSE(create_result.has_data_loss());
 
@@ -1125,7 +2134,7 @@
   }
 }
 
-TEST_F(PortableFileBackedProtoLogTest, EraseProtoShouldReturnNotFound) {
+TEST_P(PortableFileBackedProtoLogTest, EraseProtoShouldReturnNotFound) {
   DocumentProto document1 =
       DocumentBuilder().SetKey("namespace", "uri1").Build();
   DocumentProto document2 =
@@ -1136,7 +2145,11 @@
       PortableFileBackedProtoLog<DocumentProto>::Create(
           &filesystem_, file_path_,
           PortableFileBackedProtoLog<DocumentProto>::Options(
-              compress_, max_proto_size_, compression_level_)));
+              compress_, max_proto_size_, compression_level_,
+              compression_threshold_bytes_, compression_mem_level_,
+              enable_smaller_decompression_buffer_size_,
+              /*enable_new_header_format_in=*/GetParam(),
+              /*enable_reusable_decompression_buffer=*/true)));
   auto proto_log = std::move(create_result.proto_log);
   ASSERT_FALSE(create_result.has_data_loss());
 
@@ -1157,7 +2170,201 @@
               IsOkAndHolds(EqualsProto(document2)));
 }
 
-TEST_F(PortableFileBackedProtoLogTest, ChecksumShouldBeCorrectWithErasedProto) {
+TEST_P(PortableFileBackedProtoLogTest, GetChecksumShouldNotUpdateHeader) {
+  // This test verifies that GetChecksum() returns the correct checksum after
+  // each change, but the header checksum is not updated if UpdateChecksum() or
+  // PersistToDisk() is not called.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
+      PortableFileBackedProtoLog<DocumentProto>::Create(
+          &filesystem_, file_path_,
+          PortableFileBackedProtoLog<DocumentProto>::Options(
+              compress_, max_proto_size_, compression_level_,
+              compression_threshold_bytes_, compression_mem_level_,
+              enable_smaller_decompression_buffer_size_,
+              /*enable_new_header_format_in=*/GetParam(),
+              /*enable_reusable_decompression_buffer=*/true)));
+  auto proto_log = std::move(create_result.proto_log);
+  ASSERT_FALSE(create_result.has_data_loss());
+
+  // Sanity check: empty checksum.
+  ASSERT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(0u))));
+
+  // Put document1 and persist to disk.
+  DocumentProto document1 =
+      DocumentBuilder().SetKey("namespace", "uri1").Build();
+  ICING_ASSERT_OK(proto_log->WriteProto(document1));
+  EXPECT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(1883553267u))));
+  ICING_ASSERT_OK(proto_log->PersistToDisk());
+  EXPECT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(1883553267u))));
+
+  // Put document2 without calling UpdateChecksum() or PersistToDisk().
+  DocumentProto document2 =
+      DocumentBuilder().SetKey("namespace", "uri2").Build();
+  ICING_ASSERT_OK(proto_log->WriteProto(document2));
+
+  // Read the header.
+  Header header_before_get_checksum = ReadHeader(filesystem_, file_path_);
+  EXPECT_THAT(header_before_get_checksum.GetLegacyHeaderChecksum(),
+              Eq(1460433589u));
+  EXPECT_THAT(header_before_get_checksum.GetLogChecksum(), Eq(1883553267u));
+  CheckNewChecksums(header_before_get_checksum,
+                    /*expected_unsynced_tail_checksum=*/3503817807u,
+                    /*expected_header_checksum=*/1080094584u,
+                    /*enable_new_header_format=*/GetParam());
+
+  // Since there is a new document (document2) in the unsynced tail, the file is
+  // dirty. GetChecksum():
+  // - Should return the checksum reflecting the new written proto (document2).
+  // - Should not update anything in the header.
+  EXPECT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(1632103647u))));
+
+  // Manually check the header: although GetChecksum() reflects the latest
+  // changes, the header file should remain unchanged.
+  Header header_after_get_checksum = ReadHeader(filesystem_, file_path_);
+  EXPECT_THAT(header_after_get_checksum.GetLegacyHeaderChecksum(),
+              Eq(1460433589u));
+  EXPECT_THAT(header_after_get_checksum.GetLogChecksum(), Eq(1883553267u));
+  CheckNewChecksums(header_after_get_checksum,
+                    /*expected_unsynced_tail_checksum=*/3503817807u,
+                    /*expected_header_checksum=*/1080094584u,
+                    /*enable_new_header_format=*/GetParam());
+  EXPECT_THAT(header_after_get_checksum.GetRewindOffset(),
+              Eq(header_before_get_checksum.GetRewindOffset()));
+
+  // Call UpdateChecksum(). Should update the header.
+  // - Legacy header checksum should be updated.
+  // - Log checksum should be updated.
+  // - Rewind offset should be updated.
+  ICING_ASSERT_OK(proto_log->UpdateChecksum());
+  Header header_after_update_checksum = ReadHeader(filesystem_, file_path_);
+  EXPECT_THAT(header_after_update_checksum.GetLegacyHeaderChecksum(),
+              Eq(455544443u));
+  EXPECT_THAT(header_after_update_checksum.GetLogChecksum(), Eq(1632103647u));
+  // If the flag is enabled, the unsynced tail checksum and header checksum
+  // should be updated. Since we merged unsynced tail section into log
+  // checksummed section, the unsynced tail checksum should be updated to 0.
+  CheckNewChecksums(header_after_update_checksum,
+                    /*expected_unsynced_tail_checksum=*/0u,
+                    /*expected_header_checksum=*/2472053009u,
+                    /*enable_new_header_format=*/GetParam());
+  EXPECT_THAT(header_after_update_checksum.GetRewindOffset(),
+              Ne(header_after_get_checksum.GetRewindOffset()));
+}
+
+TEST_P(PortableFileBackedProtoLogTest,
+       GetChecksumShouldBeCorrectWithWriteProto) {
+  // This test verifies that GetChecksum() returns the correct checksum after
+  // each change, even if UpdateChecksum() or PersistToDisk() is not called.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
+      PortableFileBackedProtoLog<DocumentProto>::Create(
+          &filesystem_, file_path_,
+          PortableFileBackedProtoLog<DocumentProto>::Options(
+              compress_, max_proto_size_, compression_level_,
+              compression_threshold_bytes_, compression_mem_level_,
+              enable_smaller_decompression_buffer_size_,
+              /*enable_new_header_format_in=*/GetParam(),
+              /*enable_reusable_decompression_buffer=*/true)));
+  auto proto_log = std::move(create_result.proto_log);
+  ASSERT_FALSE(create_result.has_data_loss());
+
+  DocumentProto document1 =
+      DocumentBuilder().SetKey("namespace", "uri1").Build();
+  ICING_ASSERT_OK(proto_log->WriteProto(document1));
+  // Although we didn't call UpdateChecksum() or PersistToDisk(), GetChecksum()
+  // should return the value reflecting the new written proto (document1).
+  EXPECT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(1883553267u))));
+
+  DocumentProto document2 =
+      DocumentBuilder().SetKey("namespace", "uri2").Build();
+  ICING_ASSERT_OK(proto_log->WriteProto(document2));
+
+  const Crc32 kChecksumWithDocument2(1632103647u);
+
+  // Although we didn't call UpdateChecksum() or PersistToDisk(), GetChecksum()
+  // should return the value reflecting the new written proto (document1,
+  // document2).
+  EXPECT_THAT(proto_log->GetChecksum(), IsOkAndHolds(kChecksumWithDocument2));
+
+  // Call PersistToDisk.
+  ICING_ASSERT_OK(proto_log->PersistToDisk());
+  EXPECT_THAT(proto_log->GetChecksum(), IsOkAndHolds(kChecksumWithDocument2));
+
+  // Read the header.
+  Header header_before_write_proto = ReadHeader(filesystem_, file_path_);
+
+  // Put document3 without calling UpdateChecksum() or PersistToDisk().
+  DocumentProto document3 =
+      DocumentBuilder().SetKey("namespace", "uri3").Build();
+  ICING_ASSERT_OK(proto_log->WriteProto(document3));
+
+  const Crc32 kChecksumWithDocument3(3230460271u);
+  // Although we didn't call UpdateChecksum() or PersistToDisk(), GetChecksum()
+  // should return the value reflecting the new written proto (document3).
+  EXPECT_THAT(proto_log->GetChecksum(), IsOkAndHolds(kChecksumWithDocument3));
+
+  // Manually check the header.
+  Header header_after_write_proto = ReadHeader(filesystem_, file_path_);
+  EXPECT_THAT(header_after_write_proto.GetLegacyHeaderChecksum(),
+              Eq(header_before_write_proto.GetLegacyHeaderChecksum()));
+  EXPECT_THAT(header_after_write_proto.GetRewindOffset(),
+              Eq(header_before_write_proto.GetRewindOffset()));
+  EXPECT_THAT(header_after_write_proto.GetLogChecksum(),
+              Eq(header_before_write_proto.GetLogChecksum()));
+  if (GetParam()) {
+    // If the flag is enabled, when writing document3, the unsynced tail
+    // checksum and header checksum should be updated.
+    EXPECT_THAT(header_after_write_proto.GetUnsyncedTailChecksum(),
+                Ne(header_before_write_proto.GetUnsyncedTailChecksum()));
+    EXPECT_THAT(header_after_write_proto.GetHeaderChecksum(),
+                Ne(header_before_write_proto.GetHeaderChecksum()));
+
+    // Create another instance of the same proto log.
+    // - Nothing will be discarded thanks to the unsynced tail checksum.
+    // - GetChecksum() should return the same checksum after writing document3.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result2,
+        PortableFileBackedProtoLog<DocumentProto>::Create(
+            &filesystem_, file_path_,
+            PortableFileBackedProtoLog<DocumentProto>::Options(
+                compress_, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
+    auto proto_log2 = std::move(create_result2.proto_log);
+    EXPECT_FALSE(create_result2.has_data_loss());
+    EXPECT_THAT(proto_log2->GetChecksum(),
+                IsOkAndHolds(kChecksumWithDocument3));
+  } else {
+    // If the flag is disabled, the unsynced tail checksum and header checksum
+    // should remain 0.
+    EXPECT_THAT(header_after_write_proto.GetUnsyncedTailChecksum(), Eq(0u));
+    EXPECT_THAT(header_after_write_proto.GetHeaderChecksum(), Eq(0u));
+
+    // Create another instance of the same proto log.
+    // - document3 will be discarded since it is after the rewind offset.
+    // - GetChecksum() will return the old one before writing document3.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result2,
+        PortableFileBackedProtoLog<DocumentProto>::Create(
+            &filesystem_, file_path_,
+            PortableFileBackedProtoLog<DocumentProto>::Options(
+                compress_, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
+    auto proto_log2 = std::move(create_result2.proto_log);
+    EXPECT_THAT(create_result2.data_loss, Eq(DataLoss::PARTIAL));
+    EXPECT_THAT(proto_log2->GetChecksum(),
+                IsOkAndHolds(kChecksumWithDocument2));
+  }
+}
+
+TEST_P(PortableFileBackedProtoLogTest,
+       GetChecksumShouldBeCorrectWithErasedProtoAfterRewindOffset) {
   DocumentProto document1 =
       DocumentBuilder().SetKey("namespace", "uri1").Build();
   DocumentProto document2 =
@@ -1166,80 +2373,69 @@
       DocumentBuilder().SetKey("namespace", "uri3").Build();
   DocumentProto document4 =
       DocumentBuilder().SetKey("namespace", "uri4").Build();
-
-  int64_t document2_offset;
-  int64_t document3_offset;
+  DocumentProto document5 =
+      DocumentBuilder().SetKey("namespace", "uri5").Build();
+  DocumentProto document6 =
+      DocumentBuilder().SetKey("namespace", "uri6").Build();
 
   {
-    // Erase data after the rewind position. This won't update the checksum
-    // immediately.
+    // Write 3 protos.
     ICING_ASSERT_OK_AND_ASSIGN(
         PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
         PortableFileBackedProtoLog<DocumentProto>::Create(
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
-                compress_, max_proto_size_, compression_level_)));
+                compress_, max_proto_size_, compression_level_,
+                /*compression_threshold_bytes_in=*/0, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
     auto proto_log = std::move(create_result.proto_log);
     ASSERT_FALSE(create_result.has_data_loss());
 
+    // Sanity check: empty checksum.
+    ASSERT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(0u))));
+
     // Writes 3 protos
-    ICING_ASSERT_OK_AND_ASSIGN(int64_t document1_offset,
-                               proto_log->WriteProto(document1));
-    ICING_ASSERT_OK_AND_ASSIGN(document2_offset,
-                               proto_log->WriteProto(document2));
-    ICING_ASSERT_OK_AND_ASSIGN(document3_offset,
-                               proto_log->WriteProto(document3));
-
-    // Erases the 1st proto, checksum won't be updated immediately because the
-    // rewind position is 0.
-    ICING_ASSERT_OK(proto_log->EraseProto(document1_offset));
-
-    EXPECT_THAT(proto_log->UpdateChecksum(),
-                IsOkAndHolds(Eq(Crc32(2175574628))));
-  }  // New checksum is updated in destructor.
+    ICING_ASSERT_OK(proto_log->WriteProto(document1));
+    ICING_ASSERT_OK(proto_log->WriteProto(document2));
+    ICING_ASSERT_OK(proto_log->WriteProto(document3));
+    // Although we didn't call UpdateChecksum() or PersistToDisk(),
+    // GetChecksum() should return the value reflecting the new written proto.
+    EXPECT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(3230460271u))));
+  }  // New checksum is written to the header and flushed in destructor.
 
   {
-    // Erase data before the rewind position. This will update the checksum
-    // immediately.
+    // Erase data after the rewind position.
     ICING_ASSERT_OK_AND_ASSIGN(
         PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
         PortableFileBackedProtoLog<DocumentProto>::Create(
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
-                compress_, max_proto_size_, compression_level_)));
+                compress_, max_proto_size_, compression_level_,
+                /*compression_threshold_bytes_in=*/0, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
     auto proto_log = std::move(create_result.proto_log);
     ASSERT_FALSE(create_result.has_data_loss());
 
-    // Erases the 2nd proto that is now before the rewind position. Checksum
-    // is updated.
-    ICING_ASSERT_OK(proto_log->EraseProto(document2_offset));
+    // Sanity check that the checksum is identical to the previous one.
+    ASSERT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(3230460271u))));
 
-    EXPECT_THAT(proto_log->UpdateChecksum(),
-                IsOkAndHolds(Eq(Crc32(790877774))));
-  }
+    // Write 2 more documents to the unsynced tail.
+    ICING_ASSERT_OK_AND_ASSIGN(int64_t document4_offset,
+                               proto_log->WriteProto(document4));
+    ICING_ASSERT_OK(proto_log->WriteProto(document5));
+    // Although we didn't call UpdateChecksum() or PersistToDisk(),
+    // GetChecksum() should return the value reflecting the new written proto.
+    EXPECT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(705827520u))));
 
-  {
-    // Append data and erase data before the rewind position. This will update
-    // the checksum twice: in EraseProto() and destructor.
-    ICING_ASSERT_OK_AND_ASSIGN(
-        PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
-        PortableFileBackedProtoLog<DocumentProto>::Create(
-            &filesystem_, file_path_,
-            PortableFileBackedProtoLog<DocumentProto>::Options(
-                compress_, max_proto_size_, compression_level_)));
-    auto proto_log = std::move(create_result.proto_log);
-    ASSERT_FALSE(create_result.has_data_loss());
-
-    // Append a new document which is after the rewind position.
-    ICING_ASSERT_OK(proto_log->WriteProto(document4));
-
-    // Erases the 3rd proto that is now before the rewind position. Checksum
-    // is updated.
-    ICING_ASSERT_OK(proto_log->EraseProto(document3_offset));
-
-    EXPECT_THAT(proto_log->UpdateChecksum(),
-                IsOkAndHolds(Eq(Crc32(2344803210))));
-  }  // Checksum is updated with the newly appended document.
+    // Erases the 4th proto that is after the rewind position. GetChecksum()
+    // should return the new checksum reflecting the erased proto.
+    ICING_ASSERT_OK(proto_log->EraseProto(document4_offset));
+    EXPECT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(1052805596u))));
+  }  // New checksum is written to the header and flushed in destructor.
 
   {
     // A successful creation means that the checksum matches.
@@ -1248,12 +2444,419 @@
         PortableFileBackedProtoLog<DocumentProto>::Create(
             &filesystem_, file_path_,
             PortableFileBackedProtoLog<DocumentProto>::Options(
-                compress_, max_proto_size_, compression_level_)));
+                compress_, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
     auto proto_log = std::move(create_result.proto_log);
     EXPECT_FALSE(create_result.has_data_loss());
+
+    // The checksum should match the one from the previous case.
+    EXPECT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(1052805596u))));
   }
 }
 
+TEST_P(PortableFileBackedProtoLogTest,
+       GetChecksumShouldBeCorrectWithErasedProtoBeforeRewindOffset) {
+  DocumentProto document1 =
+      DocumentBuilder().SetKey("namespace", "uri1").Build();
+  DocumentProto document2 =
+      DocumentBuilder().SetKey("namespace", "uri2").Build();
+  DocumentProto document3 =
+      DocumentBuilder().SetKey("namespace", "uri3").Build();
+  DocumentProto document4 =
+      DocumentBuilder().SetKey("namespace", "uri4").Build();
+  DocumentProto document5 =
+      DocumentBuilder().SetKey("namespace", "uri5").Build();
+  DocumentProto document6 =
+      DocumentBuilder().SetKey("namespace", "uri6").Build();
+
+  int64_t document2_offset;
+  int64_t document3_offset;
+  {
+    // Write 3 protos.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
+        PortableFileBackedProtoLog<DocumentProto>::Create(
+            &filesystem_, file_path_,
+            PortableFileBackedProtoLog<DocumentProto>::Options(
+                compress_, max_proto_size_, compression_level_,
+                /*compression_threshold_bytes_in=*/0, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
+    auto proto_log = std::move(create_result.proto_log);
+    ASSERT_FALSE(create_result.has_data_loss());
+
+    // Sanity check: empty checksum.
+    ASSERT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(0u))));
+
+    // Writes 3 protos
+    ICING_ASSERT_OK(proto_log->WriteProto(document1));
+    ICING_ASSERT_OK_AND_ASSIGN(document2_offset,
+                               proto_log->WriteProto(document2));
+    ICING_ASSERT_OK_AND_ASSIGN(document3_offset,
+                               proto_log->WriteProto(document3));
+    // Although we didn't call UpdateChecksum() or PersistToDisk(),
+    // GetChecksum() should return the value reflecting the new written proto.
+    EXPECT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(3230460271u))));
+
+    // Call PersistToDisk() so that the unsynced tail is flushed and becomes
+    // synced section.
+    ICING_ASSERT_OK(proto_log->PersistToDisk());
+  }
+
+  {
+    // Erase data before the rewind position.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
+        PortableFileBackedProtoLog<DocumentProto>::Create(
+            &filesystem_, file_path_,
+            PortableFileBackedProtoLog<DocumentProto>::Options(
+                compress_, max_proto_size_, compression_level_,
+                /*compression_threshold_bytes_in=*/0, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
+    auto proto_log = std::move(create_result.proto_log);
+    ASSERT_FALSE(create_result.has_data_loss());
+
+    // Sanity check that the checksum is identical to the previous one.
+    ASSERT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(3230460271u))));
+
+    // Erases the 2nd proto that is before the rewind position.
+    // GetChecksum() should return the new checksum reflecting the erased proto.
+    ICING_ASSERT_OK(proto_log->EraseProto(document2_offset));
+    EXPECT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(1845730629u))));
+  }  // New checksum is written to the header and flushed in destructor.
+
+  {
+    // Append data to the unsynced tail and erase data before the rewind
+    // position.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
+        PortableFileBackedProtoLog<DocumentProto>::Create(
+            &filesystem_, file_path_,
+            PortableFileBackedProtoLog<DocumentProto>::Options(
+                compress_, max_proto_size_, compression_level_,
+                /*compression_threshold_bytes_in=*/0, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
+    auto proto_log = std::move(create_result.proto_log);
+    ASSERT_FALSE(create_result.has_data_loss());
+
+    // Sanity check that the checksum is identical to the previous one.
+    ASSERT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(1845730629u))));
+
+    // Append a new document which is after the rewind position.
+    ICING_ASSERT_OK(proto_log->WriteProto(document6));
+
+    // Erases the 3rd proto that is before the rewind position.
+    // GetChecksum() should return the new checksum reflecting document6 and the
+    // erased proto.
+    ICING_ASSERT_OK(proto_log->EraseProto(document3_offset));
+    EXPECT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(3659111045u))));
+  }  // New checksum is written to the header and flushed in destructor.
+
+  {
+    // A successful creation means that the checksum matches.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
+        PortableFileBackedProtoLog<DocumentProto>::Create(
+            &filesystem_, file_path_,
+            PortableFileBackedProtoLog<DocumentProto>::Options(
+                compress_, max_proto_size_, compression_level_,
+                compression_threshold_bytes_, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
+    auto proto_log = std::move(create_result.proto_log);
+    EXPECT_FALSE(create_result.has_data_loss());
+
+    // The checksum should match the one from the previous case.
+    EXPECT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(3659111045u))));
+  }
+}
+
+TEST_P(PortableFileBackedProtoLogTest, ErasedProtoAfterRewindOffset) {
+  DocumentProto document1 =
+      DocumentBuilder().SetKey("namespace", "uri1").Build();
+  DocumentProto document2 =
+      DocumentBuilder().SetKey("namespace", "uri2").Build();
+  DocumentProto document3 =
+      DocumentBuilder().SetKey("namespace", "uri3").Build();
+  DocumentProto document4 =
+      DocumentBuilder().SetKey("namespace", "uri4").Build();
+  DocumentProto document5 =
+      DocumentBuilder().SetKey("namespace", "uri5").Build();
+
+  {
+    // Write 3 protos.
+    ICING_ASSERT_OK_AND_ASSIGN(
+        PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
+        PortableFileBackedProtoLog<DocumentProto>::Create(
+            &filesystem_, file_path_,
+            PortableFileBackedProtoLog<DocumentProto>::Options(
+                compress_, max_proto_size_, compression_level_,
+                /*compression_threshold_bytes_in=*/0, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
+    auto proto_log = std::move(create_result.proto_log);
+    ASSERT_FALSE(create_result.has_data_loss());
+
+    // Sanity check: empty checksum.
+    ASSERT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(0u))));
+
+    // Writes 3 protos
+    ICING_ASSERT_OK(proto_log->WriteProto(document1));
+    ICING_ASSERT_OK(proto_log->WriteProto(document2));
+    ICING_ASSERT_OK(proto_log->WriteProto(document3));
+    // Although we didn't call UpdateChecksum() or PersistToDisk(),
+    // GetChecksum() should return the value reflecting the new written proto.
+    EXPECT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(3230460271u))));
+  }  // New checksum is written to the header and flushed in destructor.
+
+  // Erase data after the rewind position.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
+      PortableFileBackedProtoLog<DocumentProto>::Create(
+          &filesystem_, file_path_,
+          PortableFileBackedProtoLog<DocumentProto>::Options(
+              compress_, max_proto_size_, compression_level_,
+              /*compression_threshold_bytes_in=*/0, compression_mem_level_,
+              enable_smaller_decompression_buffer_size_,
+              /*enable_new_header_format_in=*/GetParam(),
+              /*enable_reusable_decompression_buffer=*/true)));
+  auto proto_log = std::move(create_result.proto_log);
+  ASSERT_FALSE(create_result.has_data_loss());
+
+  // Sanity check that the checksum is identical to the previous one.
+  ASSERT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(3230460271u))));
+
+  // Write 2 more documents to the unsynced tail.
+  ICING_ASSERT_OK_AND_ASSIGN(int64_t document4_offset,
+                             proto_log->WriteProto(document4));
+  ICING_ASSERT_OK(proto_log->WriteProto(document5));
+  // Although we didn't call UpdateChecksum() or PersistToDisk(),
+  // GetChecksum() should return the value reflecting the new written proto.
+  EXPECT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(705827520u))));
+
+  // Read the header.
+  Header header_before_erase = ReadHeader(filesystem_, file_path_);
+
+  // Erases the 4th proto that is after the rewind position. GetChecksum()
+  // should return the new checksum reflecting the erased proto.
+  ICING_ASSERT_OK(proto_log->EraseProto(document4_offset));
+  EXPECT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(1052805596u))));
+
+  // Manually check the header.
+  Header header_after_erase = ReadHeader(filesystem_, file_path_);
+  EXPECT_THAT(header_after_erase.GetLegacyHeaderChecksum(),
+              Eq(header_before_erase.GetLegacyHeaderChecksum()));
+  EXPECT_THAT(header_after_erase.GetRewindOffset(),
+              Eq(header_before_erase.GetRewindOffset()));
+  EXPECT_THAT(header_after_erase.GetLogChecksum(),
+              Eq(header_before_erase.GetLogChecksum()));
+  if (GetParam()) {
+    // If the flag is enabled, when erasing document4, the unsynced tail
+    // checksum and header checksum should be updated.
+    EXPECT_THAT(header_after_erase.GetUnsyncedTailChecksum(),
+                Ne(header_before_erase.GetUnsyncedTailChecksum()));
+    EXPECT_THAT(header_after_erase.GetHeaderChecksum(),
+                Ne(header_before_erase.GetHeaderChecksum()));
+  } else {
+    // If the flag is disabled, the unsynced tail checksum and header checksum
+    // should remain 0.
+    EXPECT_THAT(header_after_erase.GetUnsyncedTailChecksum(), Eq(0u));
+    EXPECT_THAT(header_after_erase.GetHeaderChecksum(), Eq(0u));
+  }
+}
+
+TEST_P(PortableFileBackedProtoLogTest,
+       ErasedProtoBeforeRewindOffsetShouldUpdateHeader) {
+  DocumentProto document1 =
+      DocumentBuilder().SetKey("namespace", "uri1").Build();
+  DocumentProto document2 =
+      DocumentBuilder().SetKey("namespace", "uri2").Build();
+  DocumentProto document3 =
+      DocumentBuilder().SetKey("namespace", "uri3").Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
+      PortableFileBackedProtoLog<DocumentProto>::Create(
+          &filesystem_, file_path_,
+          PortableFileBackedProtoLog<DocumentProto>::Options(
+              compress_, max_proto_size_, compression_level_,
+              /*compression_threshold_bytes_in=*/0, compression_mem_level_,
+              enable_smaller_decompression_buffer_size_,
+              /*enable_new_header_format_in=*/GetParam(),
+              /*enable_reusable_decompression_buffer=*/true)));
+  auto proto_log = std::move(create_result.proto_log);
+  ASSERT_FALSE(create_result.has_data_loss());
+
+  // Sanity check: empty checksum.
+  ASSERT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(0u))));
+
+  // Writes 3 protos
+  ICING_ASSERT_OK(proto_log->WriteProto(document1));
+  ICING_ASSERT_OK_AND_ASSIGN(int64_t document2_offset,
+                             proto_log->WriteProto(document2));
+  ICING_ASSERT_OK(proto_log->WriteProto(document3));
+  EXPECT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(3230460271u))));
+
+  // Persist to disk in order to make document1 to document3 in the log
+  // checksummed section.
+  ICING_ASSERT_OK(proto_log->PersistToDisk());
+
+  // Read the header.
+  Header header_before_erase = ReadHeader(filesystem_, file_path_);
+
+  // Erases the 2nd proto that is before the rewind position. GetChecksum()
+  // should return the new checksum reflecting the erased proto.
+  ICING_ASSERT_OK(proto_log->EraseProto(document2_offset));
+  EXPECT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(1845730629u))));
+
+  // Manually check the header:
+  // - Log checksum is updated because the log checksummed section was updated.
+  // - Legacy header checksum is updated because log checksum was updated.
+  // - RewindOffset should remain the same.
+  // - (Only if enable_new_header_format flag is true) Header checksum is
+  //   updated because the header was updated.
+  Header header_after_erase = ReadHeader(filesystem_, file_path_);
+  EXPECT_THAT(header_after_erase.GetLegacyHeaderChecksum(),
+              Ne(header_before_erase.GetLegacyHeaderChecksum()));
+  EXPECT_THAT(header_after_erase.GetRewindOffset(),
+              Eq(header_before_erase.GetRewindOffset()));
+  EXPECT_THAT(header_after_erase.GetLogChecksum(),
+              Ne(header_before_erase.GetLogChecksum()));
+  if (GetParam()) {
+    EXPECT_THAT(header_after_erase.GetHeaderChecksum(),
+                Ne(header_before_erase.GetHeaderChecksum()));
+  }
+}
+
+TEST_P(PortableFileBackedProtoLogTest,
+       ReadWriteHeaderShouldIncludePaddingBytes) {
+  // The header should have some padding bytes to make it 8-byte aligned.
+  static_assert(Header::kLegacyHeaderSectionSize -
+                    Header::kLegacyHeaderSectionPaddingOffset >
+                0);
+  std::string bytes = "abc";
+
+  DocumentProto document1 =
+      DocumentBuilder().SetKey("namespace", "uri1").Build();
+  DocumentProto document2 =
+      DocumentBuilder().SetKey("namespace", "uri2").Build();
+  DocumentProto document3 =
+      DocumentBuilder().SetKey("namespace", "uri3").Build();
+
+  {
+    ICING_ASSERT_OK_AND_ASSIGN(
+        PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
+        PortableFileBackedProtoLog<DocumentProto>::Create(
+            &filesystem_, file_path_,
+            PortableFileBackedProtoLog<DocumentProto>::Options(
+                compress_, max_proto_size_, compression_level_,
+                /*compression_threshold_bytes_in=*/0, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
+    auto proto_log = std::move(create_result.proto_log);
+    ASSERT_FALSE(create_result.has_data_loss());
+
+    // Sanity check: empty checksum.
+    ASSERT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(0u))));
+
+    // Writes 2 protos
+    ICING_ASSERT_OK(proto_log->WriteProto(document1));
+    ICING_ASSERT_OK(proto_log->WriteProto(document2));
+    EXPECT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(1632103647u))));
+    // Persist to disk.
+    ICING_ASSERT_OK(proto_log->PersistToDisk());
+
+    // Overwrite the padding bytes with some random bytes.
+    Header* header = proto_log->header();
+    memcpy(reinterpret_cast<char*>(header) +
+               Header::kLegacyHeaderSectionPaddingOffset,
+           bytes.data(),
+           Header::kLegacyHeaderSectionSize -
+               Header::kLegacyHeaderSectionPaddingOffset);
+
+    // Add 1 more proto. It is a hack to make the proto log dirty and make the
+    // next PersistToDisk take effect.
+    ICING_ASSERT_OK(proto_log->WriteProto(document3));
+    EXPECT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(3230460271u))));
+    ICING_ASSERT_OK(proto_log->PersistToDisk());
+
+    // Manually check the header: the padding bytes contain the overwritten
+    // values.
+    // Ensures that write header operation also includes the padding bytes.
+    Header header_after = ReadHeader(filesystem_, file_path_);
+    EXPECT_THAT(std::string(reinterpret_cast<const char*>(&header_after) +
+                                Header::kLegacyHeaderSectionPaddingOffset,
+                            Header::kLegacyHeaderSectionSize -
+                                Header::kLegacyHeaderSectionPaddingOffset),
+                Eq(bytes));
+  }
+
+  {
+    ICING_ASSERT_OK_AND_ASSIGN(
+        PortableFileBackedProtoLog<DocumentProto>::CreateResult create_result,
+        PortableFileBackedProtoLog<DocumentProto>::Create(
+            &filesystem_, file_path_,
+            PortableFileBackedProtoLog<DocumentProto>::Options(
+                compress_, max_proto_size_, compression_level_,
+                /*compression_threshold_bytes_in=*/0, compression_mem_level_,
+                enable_smaller_decompression_buffer_size_,
+                /*enable_new_header_format_in=*/GetParam(),
+                /*enable_reusable_decompression_buffer=*/true)));
+    auto proto_log = std::move(create_result.proto_log);
+    ASSERT_FALSE(create_result.has_data_loss());
+
+    // Sanity check that the checksum is identical to the previous one.
+    ASSERT_THAT(proto_log->GetChecksum(), IsOkAndHolds(Eq(Crc32(3230460271u))));
+
+    // Verify that read header operation includes the padding bytes.
+    Header* header = proto_log->header();
+    EXPECT_THAT(std::string(reinterpret_cast<const char*>(header) +
+                                Header::kLegacyHeaderSectionPaddingOffset,
+                            Header::kLegacyHeaderSectionSize -
+                                Header::kLegacyHeaderSectionPaddingOffset),
+                Eq(bytes));
+  }
+}
+
+TEST_P(PortableFileBackedProtoLogTest,
+       CalculateLegacyHeaderChecksumShouldIncludePadding) {
+  // The legacy header should have some padding bytes to make it 8-byte aligned.
+  static_assert(Header::kLegacyHeaderSectionSize -
+                    Header::kLegacyHeaderSectionPaddingOffset >
+                0);
+  std::string bytes = "abc";
+
+  // Create a header instance and set all bytes to 1 for the header.
+  // Note that -1 means 0b111111...111.
+  Header header;
+  memset(&header, -1, sizeof(Header));
+  EXPECT_THAT(header.CalculateLegacyHeaderChecksum(), Eq(2132597986u));
+
+  // Set different bytes to the padding section. Verify that the legacy header
+  // checksum reflects the padding byte changes.
+  memcpy(reinterpret_cast<char*>(&header) +
+             Header::kLegacyHeaderSectionPaddingOffset,
+         bytes.data(),
+         Header::kLegacyHeaderSectionSize -
+             Header::kLegacyHeaderSectionPaddingOffset);
+  EXPECT_THAT(header.CalculateLegacyHeaderChecksum(), Eq(3049742880u));
+}
+
+INSTANTIATE_TEST_SUITE_P(PortableFileBackedProtoLogTest,
+                         PortableFileBackedProtoLogTest,
+                         testing::Values(true, false));
+
 }  // namespace
 }  // namespace lib
 }  // namespace icing

diff --git a/icing/file/portable-file-backed-proto-log_thread-safety_test.cc b/icing/file/portable-file-backed-proto-log_thread-safety_test.cc
new file mode 100644
index 0000000..1797e98
--- /dev/null
+++ b/icing/file/portable-file-backed-proto-log_thread-safety_test.cc

@@ -0,0 +1,117 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <memory>
+#include <string>
+#include <thread>
+#include <vector>
+
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+#include "icing/document-builder.h"
+#include "icing/file/filesystem.h"
+#include "icing/file/portable-file-backed-proto-log.h"
+#include "icing/portable/equals-proto.h"
+#include "icing/proto/document.pb.h"
+#include "icing/testing/common-matchers.h"
+#include "icing/testing/tmp-directory.h"
+
+namespace icing {
+namespace lib {
+namespace {
+
+using ::icing::lib::portable_equals_proto::EqualsProto;
+using ::testing::NotNull;
+
+class PortableFileBackedProtoLogThreadSafetyTest
+    : public ::testing::TestWithParam<bool> {
+ protected:
+  void SetUp() override {
+    file_path_ = GetTestTempDir() + "/proto_log_thread_safety";
+    filesystem_.DeleteFile(file_path_.c_str());
+  }
+
+  void TearDown() override { filesystem_.DeleteFile(file_path_.c_str()); }
+
+  const Filesystem filesystem_;
+  std::string file_path_;
+};
+
+TEST_P(PortableFileBackedProtoLogThreadSafetyTest, ConcurrentReadProto) {
+  // Write some documents to the log.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      auto create_result,
+      PortableFileBackedProtoLog<DocumentProto>::Create(
+          &filesystem_, file_path_,
+          PortableFileBackedProtoLog<DocumentProto>::Options(
+              /*compress_in=*/true,
+              /*max_proto_size_in=*/1024 * 1024,
+              PortableFileBackedProtoLog<
+                  DocumentProto>::kDefaultCompressionLevel,
+              PortableFileBackedProtoLog<
+                  DocumentProto>::kDefaultCompressionThresholdBytes,
+              /*compression_mem_level_in=*/1,
+              /*enable_smaller_decompression_buffer_size_in=*/true,
+              /*enable_new_header_format_in=*/true,
+              /*enable_reusable_decompression_buffer_in=*/GetParam())));
+  std::unique_ptr<PortableFileBackedProtoLog<DocumentProto>> proto_log =
+      std::move(create_result.proto_log);
+  ASSERT_THAT(proto_log, NotNull());
+
+  constexpr int kNumDocuments = 20;
+  std::vector<DocumentProto> documents;
+  std::vector<int64_t> offsets;
+  for (int i = 0; i < kNumDocuments; ++i) {
+    documents.push_back(
+        DocumentBuilder()
+            .SetKey("namespace", "uri" + std::to_string(i))
+            .AddStringProperty("prop", std::string(100 * i, 'a'))
+            .Build());
+    ICING_ASSERT_OK_AND_ASSIGN(int64_t offset,
+                               proto_log->WriteProto(documents.back()));
+    offsets.push_back(offset);
+  }
+  ICING_ASSERT_OK(proto_log->PersistToDisk());
+
+  // Create kNumThreads to call ReadProto concurrently.
+  constexpr int kNumThreads = 50;
+  constexpr int kNumReadsPerThread = 100;
+
+  auto reader_task = [&]() {
+    for (int i = 0; i < kNumReadsPerThread; ++i) {
+      int doc_index = i % kNumDocuments;
+      ICING_ASSERT_OK_AND_ASSIGN(DocumentProto read_doc,
+                                 proto_log->ReadProto(offsets[doc_index]));
+      EXPECT_THAT(read_doc, EqualsProto(documents[doc_index]));
+    }
+  };
+
+  std::vector<std::thread> threads;
+  threads.reserve(kNumThreads);
+  for (int i = 0; i < kNumThreads; ++i) {
+    threads.emplace_back(reader_task);
+  }
+
+  for (auto& thread : threads) {
+    thread.join();
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(PortableFileBackedProtoLogThreadSafetyTest,
+                         PortableFileBackedProtoLogThreadSafetyTest,
+                         testing::Values(true, false));
+
+}  // namespace
+}  // namespace lib
+}  // namespace icing
\ No newline at end of file

diff --git a/icing/file/posting_list/flash-index-storage-header.h b/icing/file/posting_list/flash-index-storage-header.h
index f7b331c..67fc4a6 100644
--- a/icing/file/posting_list/flash-index-storage-header.h
+++ b/icing/file/posting_list/flash-index-storage-header.h

@@ -62,7 +62,7 @@
   static libtextclassifier3::StatusOr<HeaderBlock> Read(
       const Filesystem* filesystem, int fd, int block_size) {
     std::unique_ptr<uint8_t[]> buffer = std::make_unique<uint8_t[]>(block_size);
-    if (!filesystem->PRead(fd, buffer.get(), block_size, 0)) {
+    if (filesystem->PRead(fd, buffer.get(), block_size, 0) != block_size) {
       return absl_ports::InternalError("Unable to reader header block!");
     }
     return HeaderBlock(filesystem, std::move(buffer), block_size);

diff --git a/icing/file/posting_list/flash-index-storage.cc b/icing/file/posting_list/flash-index-storage.cc
index 2198d2c..68f8e4c 100644
--- a/icing/file/posting_list/flash-index-storage.cc
+++ b/icing/file/posting_list/flash-index-storage.cc

@@ -21,15 +21,23 @@
 #include <cerrno>
 #include <cinttypes>
 #include <cstdint>
+#include <cstring>
 #include <memory>
+#include <string>
+#include <utility>
 
 #include "icing/text_classifier/lib3/utils/base/status.h"
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "icing/absl_ports/canonical_errors.h"
 #include "icing/absl_ports/str_cat.h"
+#include "icing/file/filesystem.h"
+#include "icing/file/posting_list/flash-index-storage-header.h"
 #include "icing/file/posting_list/index-block.h"
 #include "icing/file/posting_list/posting-list-common.h"
+#include "icing/file/posting_list/posting-list-identifier.h"
+#include "icing/file/posting_list/posting-list-used.h"
 #include "icing/legacy/core/icing-string-util.h"
+#include "icing/store/document-id.h"
 #include "icing/util/logging.h"
 #include "icing/util/math-util.h"
 #include "icing/util/status-macros.h"
@@ -110,20 +118,24 @@
   // Look for an existing file size.
   int64_t file_size = filesystem_->GetFileSize(storage_sfd_.get());
   if (file_size == Filesystem::kBadFileSize) {
-    ICING_LOG(ERROR) << "Could not initialize main index. Bad file size.";
+    ICING_LOG(ERROR)
+        << "Could not initialize flash index storage. Bad file size. Path: "
+        << index_filename_;
     return false;
   }
 
   if (file_size == 0) {
     if (!CreateHeader()) {
-      ICING_LOG(ERROR)
-          << "Could not initialize main index. Unable to create header.";
+      ICING_LOG(ERROR) << "Could not initialize flash index storage. Unable to "
+                          "create header. Path: "
+                       << index_filename_;
       return false;
     }
   } else {
     if (!OpenHeader(file_size)) {
-      ICING_LOG(ERROR)
-          << "Could not initialize main index. Unable to open header.";
+      ICING_LOG(ERROR) << "Could not initialize flash index storage. Unable to "
+                          "open header. Path: "
+                       << index_filename_;
       return false;
     }
   }
@@ -181,25 +193,32 @@
       HeaderBlock read_header,
       HeaderBlock::Read(filesystem_, storage_sfd_.get(), block_size), false);
   if (read_header.header()->magic != HeaderBlock::Header::kMagic) {
-    ICING_LOG(ERROR) << "Index header block wrong magic";
+    ICING_LOG(ERROR) << "Invalid FlashIndexStorage header for "
+                     << index_filename_ << ": wrong magic. Expected: "
+                     << HeaderBlock::Header::kMagic
+                     << ", actual: " << read_header.header()->magic;
     return false;
   }
   if (file_size % read_header.header()->block_size != 0) {
-    ICING_LOG(ERROR) << "Index size " << file_size
-                     << " not a multiple of block size "
+    ICING_LOG(ERROR) << "Invalid FlashIndexStorage header for "
+                     << index_filename_ << ": file size " << file_size
+                     << " is not a multiple of block size "
                      << read_header.header()->block_size;
     return false;
   }
 
   if (file_size < static_cast<int64_t>(read_header.header()->block_size)) {
-    ICING_LOG(ERROR) << "Index size " << file_size
-                     << " shorter than block size "
+    ICING_LOG(ERROR) << "Invalid FlashIndexStorage header for "
+                     << index_filename_ << ": file size " << file_size
+                     << " is shorter than block size "
                      << read_header.header()->block_size;
     return false;
   }
 
   if (read_header.header()->block_size % getpagesize() != 0) {
-    ICING_LOG(ERROR) << "Block size " << read_header.header()->block_size
+    ICING_LOG(ERROR) << "Invalid FlashIndexStorage header for "
+                     << index_filename_ << ": block size "
+                     << read_header.header()->block_size
                      << " is not a multiple of page size " << getpagesize();
     return false;
   }
@@ -211,7 +230,8 @@
     // still use the main index, but reads/writes won't be as efficient in terms
     // of flash IO because the 'blocks' that we're reading are actually multiple
     // pages long.
-    ICING_LOG(ERROR) << "Block size of existing header ("
+    ICING_LOG(ERROR) << "Invalid FlashIndexStorage header for "
+                     << index_filename_ << ": block size of existing header ("
                      << read_header.header()->block_size
                      << ") does not match the requested block size ("
                      << block_size << "). Defaulting to existing block size "
@@ -244,7 +264,8 @@
 bool FlashIndexStorage::PersistToDisk() {
   // First, write header.
   if (!header_block_->Write(storage_sfd_.get())) {
-    ICING_LOG(ERROR) << "Write index header failed: " << strerror(errno);
+    ICING_LOG(ERROR) << "Write index header failed: (" << errno << ") "
+                     << strerror(errno);
     return false;
   }
 
@@ -535,7 +556,8 @@
   if (!filesystem_->Grow(
           storage_sfd_.get(),
           static_cast<uint64_t>(num_blocks_ + 1) * block_size())) {
-    ICING_VLOG(1) << "Error growing index file: " << strerror(errno);
+    ICING_LOG(ERROR) << "Error growing index file: (" << errno << ") "
+                     << strerror(errno);
     return kInvalidBlockIndex;
   }
 

diff --git a/icing/file/posting_list/index-block.cc b/icing/file/posting_list/index-block.cc
index 3fa397c..1f47da2 100644
--- a/icing/file/posting_list/index-block.cc
+++ b/icing/file/posting_list/index-block.cc

@@ -67,7 +67,8 @@
   }
 
   BlockHeader header;
-  if (!filesystem->PRead(fd, &header, sizeof(BlockHeader), block_file_offset)) {
+  if (filesystem->PRead(fd, &header, sizeof(BlockHeader), block_file_offset) !=
+      sizeof(BlockHeader)) {
     return absl_ports::InternalError("PRead block header error");
   }
 
@@ -284,8 +285,8 @@
 libtextclassifier3::StatusOr<IndexBlock::BlockHeader> IndexBlock::ReadHeader()
     const {
   BlockHeader header;
-  if (!filesystem_->PRead(fd_, &header, sizeof(BlockHeader),
-                          block_file_offset_)) {
+  if (filesystem_->PRead(fd_, &header, sizeof(BlockHeader),
+                         block_file_offset_) != sizeof(BlockHeader)) {
     return absl_ports::InternalError(
         absl_ports::StrCat("PRead block header error: ", strerror(errno)));
   }
@@ -301,8 +302,9 @@
 libtextclassifier3::StatusOr<std::unique_ptr<uint8_t[]>>
 IndexBlock::ReadPostingList(PostingListIndex posting_list_index) const {
   auto posting_list_buffer = std::make_unique<uint8_t[]>(posting_list_bytes_);
-  if (!filesystem_->PRead(fd_, posting_list_buffer.get(), posting_list_bytes_,
-                          get_posting_list_file_offset(posting_list_index))) {
+  if (filesystem_->PRead(fd_, posting_list_buffer.get(), posting_list_bytes_,
+                         get_posting_list_file_offset(posting_list_index)) !=
+      posting_list_bytes_) {
     return absl_ports::InternalError(
         absl_ports::StrCat("PRead posting list error: ", strerror(errno)));
   }

diff --git a/icing/file/version-util.cc b/icing/file/version-util.cc
index c44ac69..2bc2122 100644
--- a/icing/file/version-util.cc
+++ b/icing/file/version-util.cc

@@ -30,6 +30,7 @@
 #include "icing/file/filesystem.h"
 #include "icing/index/index.h"
 #include "icing/proto/initialize.pb.h"
+#include "icing/util/logging.h"
 #include "icing/util/status-macros.h"
 
 namespace icing {
@@ -45,10 +46,15 @@
   // 1. Read the version info.
   const std::string v1_version_filepath =
       MakeVersionFilePath(version_file_dir, kVersionFilenameV1);
+  bool file_exists = filesystem.FileExists(v1_version_filepath.c_str());
+  ICING_LOG(INFO) << "v1 version file exists: " << file_exists;
+
   VersionInfo existing_version_info(-1, -1);
-  if (filesystem.FileExists(v1_version_filepath.c_str()) &&
-      !filesystem.PRead(v1_version_filepath.c_str(), &existing_version_info,
-                        sizeof(VersionInfo), /*offset=*/0)) {
+  if (file_exists &&
+      filesystem.PRead(v1_version_filepath.c_str(), &existing_version_info,
+                       sizeof(VersionInfo),
+                       /*offset=*/0) != sizeof(VersionInfo)) {
+    ICING_LOG(ERROR) << "Failed to read v1 version file";
     return absl_ports::InternalError("Failed to read v1 version file");
   }
 
@@ -56,6 +62,12 @@
   libtextclassifier3::StatusOr<int> existing_flash_index_magic =
       Index::ReadFlashIndexMagic(&filesystem, index_base_dir);
   if (!existing_flash_index_magic.ok()) {
+    ICING_LOG(WARNING) << "Cannot read index flash index magic for v1 version "
+                          "check. Error code: "
+                       << existing_flash_index_magic.status().error_code()
+                       << ", message: "
+                       << existing_flash_index_magic.status().error_message();
+
     if (absl_ports::IsNotFound(existing_flash_index_magic.status())) {
       // Flash index magic doesn't exist. In this case, we're unable to
       // determine the version change state correctly (regardless of the
@@ -67,7 +79,10 @@
     // Real error.
     return std::move(existing_flash_index_magic).status();
   }
-  if (existing_flash_index_magic.ValueOrDie() == kVersionZeroFlashIndexMagic) {
+
+  int flash_index_magic = existing_flash_index_magic.ValueOrDie();
+  ICING_LOG(INFO) << "Index flash index magic: " << flash_index_magic;
+  if (flash_index_magic == kVersionZeroFlashIndexMagic) {
     existing_version_info.version = 0;
     if (existing_version_info.max_version == -1) {
       existing_version_info.max_version = 0;
@@ -83,6 +98,9 @@
   // IcingSearchEngineVersionProto as a file-backed proto.
   const std::string v2_version_filepath =
       MakeVersionFilePath(version_file_dir, kVersionFilenameV2);
+  ICING_LOG(INFO) << "v2 version file exists: "
+                  << filesystem.FileExists(v2_version_filepath.c_str());
+
   FileBackedProto<IcingSearchEngineVersionProto> v2_version_file(
       filesystem, v2_version_filepath);
   ICING_ASSIGN_OR_RETURN(const IcingSearchEngineVersionProto* v2_version_proto,
@@ -112,6 +130,11 @@
   // 2. Read the v2 version file
   auto v2_version_proto = ReadV2VersionInfo(filesystem, version_file_dir);
   if (!v2_version_proto.ok()) {
+    ICING_LOG(WARNING) << "Cannot read v2 version file. Error code: "
+                       << v2_version_proto.status().error_code()
+                       << ", message: "
+                       << v2_version_proto.status().error_message();
+
     if (!absl_ports::IsNotFound(v2_version_proto.status())) {
       // Real error.
       return std::move(v2_version_proto).status();
@@ -141,6 +164,11 @@
       // should have been written.
       // Return an invalid version number in this case and trigger rebuilding
       // everything.
+      ICING_LOG(WARNING)
+          << "v1 version file has a version number greater or equal to "
+             "kFirstV2Version, but failed to read v2 version file. Return "
+             "version -1 and rebuild everything.";
+
       version_proto.set_version(-1);
       version_proto.set_max_version(v1_version_info.max_version);
     }
@@ -153,6 +181,10 @@
   IcingSearchEngineVersionProto v2_version_proto_value =
       std::move(v2_version_proto).ValueOrDie();
   if (v1_version_info.version != v2_version_proto_value.version()) {
+    ICING_LOG(WARNING)
+        << "v1 and v2 version files have different version numbers, probably "
+           "due to roll forward. Return version -1 and rebuild everything.";
+
     v2_version_proto_value.set_version(-1);
     v2_version_proto_value.mutable_enabled_features()->Clear();
   }
@@ -165,26 +197,51 @@
                                           const VersionInfo& version_info) {
   ScopedFd scoped_fd(filesystem.OpenForWrite(
       MakeVersionFilePath(version_file_dir, kVersionFilenameV1).c_str()));
-  if (!scoped_fd.is_valid() ||
-      !filesystem.PWrite(scoped_fd.get(), /*offset=*/0, &version_info,
-                         sizeof(VersionInfo)) ||
-      !filesystem.DataSync(scoped_fd.get())) {
+  if (!scoped_fd.is_valid()) {
+    ICING_LOG(ERROR) << "Failed to open v1 version file for write";
+    return absl_ports::InternalError(
+        "Failed to open v1 version file for write");
+  }
+
+  if (!filesystem.PWrite(scoped_fd.get(), /*offset=*/0, &version_info,
+                         sizeof(VersionInfo))) {
+    ICING_LOG(ERROR) << "Failed to write v1 version file";
     return absl_ports::InternalError("Failed to write v1 version file");
   }
+
+  if (!filesystem.DataSync(scoped_fd.get())) {
+    ICING_LOG(ERROR) << "Failed to sync v1 version file";
+    return absl_ports::InternalError("Failed to sync v1 version file");
+  }
+
+  ICING_LOG(INFO)
+      << "Successfully write and sync v1 version file with version = "
+      << version_info.version;
   return libtextclassifier3::Status::OK;
 }
 
 libtextclassifier3::Status WriteV2Version(
     const Filesystem& filesystem, const std::string& version_file_dir,
     std::unique_ptr<IcingSearchEngineVersionProto> version_proto) {
+  int32_t version = version_proto->version();
+
+  // FileBackedProto::Write also syncs the file, so we don't need to call
+  // DataSync explicitly.
   FileBackedProto<IcingSearchEngineVersionProto> v2_version_file(
       filesystem, MakeVersionFilePath(version_file_dir, kVersionFilenameV2));
   libtextclassifier3::Status v2_write_status =
       v2_version_file.Write(std::move(version_proto));
   if (!v2_write_status.ok()) {
+    ICING_LOG(ERROR) << "Failed to write v2 version file. Error code: "
+                     << v2_write_status.error_code()
+                     << ", message: " << v2_write_status.error_message();
     return absl_ports::InternalError(absl_ports::StrCat(
         "Failed to write v2 version file: ", v2_write_status.error_message()));
   }
+
+  ICING_LOG(INFO)
+      << "Successfully write and sync v2 version file with version = "
+      << version;
   return libtextclassifier3::Status::OK;
 }
 
@@ -329,6 +386,18 @@
         // version 5 -> version 6 upgrade, no need to rebuild
         break;
       }
+      case 6: {
+        // version 6 -> version 7 upgrade, no need to rebuild
+        break;
+      }
+      case 7: {
+        // version 7 -> version 8 upgrade, no need to rebuild
+        break;
+      }
+      case 8: {
+        // version 8 -> version 9 upgrade, no need to rebuild
+        break;
+      }
       default:
         // This should not happen. Rebuild anyway if unsure.
         should_rebuild |= true;
@@ -389,8 +458,37 @@
           /*needs_qualified_id_join_index_rebuild=*/false,
           /*needs_embedding_index_rebuild=*/false);
     }
+    case IcingSearchEngineFeatureInfoProto::FEATURE_QUALIFIED_ID_JOIN_INDEX_V3:
     case IcingSearchEngineFeatureInfoProto::
-        FEATURE_QUALIFIED_ID_JOIN_INDEX_V3: {
+        FEATURE_NON_EXISTENT_QUALIFIED_ID_JOIN: {
+      return derived_file_util::DerivedFilesRebuildInfo(
+          /*needs_document_store_derived_files_rebuild=*/false,
+          /*needs_schema_store_derived_files_rebuild=*/false,
+          /*needs_term_index_rebuild=*/false,
+          /*needs_integer_index_rebuild=*/false,
+          /*needs_qualified_id_join_index_rebuild=*/true,
+          /*needs_embedding_index_rebuild=*/false);
+    }
+    case IcingSearchEngineFeatureInfoProto::FEATURE_DELETE_PROPAGATION_FROM: {
+      // By rebuilding qualified id join index, we will re-validate document
+      // dependency: if a schema enables delete propagation but a document of
+      // this schema references a deleted/expired/non-existent parent document,
+      // then Icing will remove this child document.
+      //
+      // Note: document dependency re-validation can be done by rebuilding any
+      //   index since the process is done in TokenizedDocument generation, but
+      //   we choose to rebuild the join index because:
+      //   - Rebuilding the join index guarantees to "replay" document protos
+      //     from doc id 0. All alive documents will be iterated and
+      //     re-validated.
+      //   - Delete propagation is related to the join index.
+      //   - We could've only generated TokenizedDocument and re-validated the
+      //     document dependency without rebuilding any indices, but:
+      //     - This would require more changes to Icing's derived files recovery
+      //       logic (i.e. RestoreIndexIfNeeded()).
+      //     - Tt is rare to switch this flag on and off too frequently.
+      //     - Therefore, let's just simply reuse the existing derived files
+      //       recovery logic to trigger document dependency re-validation.
       return derived_file_util::DerivedFilesRebuildInfo(
           /*needs_document_store_derived_files_rebuild=*/false,
           /*needs_schema_store_derived_files_rebuild=*/false,
@@ -478,6 +576,17 @@
     enabled_features->Add(GetFeatureInfoProto(
         IcingSearchEngineFeatureInfoProto::FEATURE_QUALIFIED_ID_JOIN_INDEX_V3));
   }
+  // DeletePropagation PROPAGATE_FROM feature
+  if (options.enable_delete_propagation_from()) {
+    enabled_features->Add(GetFeatureInfoProto(
+        IcingSearchEngineFeatureInfoProto::FEATURE_DELETE_PROPAGATION_FROM));
+  }
+  // NonExistentQualifiedIdJoin feature
+  if (options.enable_non_existent_qualified_id_join()) {
+    enabled_features->Add(
+        GetFeatureInfoProto(IcingSearchEngineFeatureInfoProto::
+                                FEATURE_NON_EXISTENT_QUALIFIED_ID_JOIN));
+  }
 }
 
 }  // namespace version_util

diff --git a/icing/file/version-util.h b/icing/file/version-util.h
index 4b32e0c..3ee187c 100644
--- a/icing/file/version-util.h
+++ b/icing/file/version-util.h

@@ -39,7 +39,7 @@
 // - Version 3: M-2024-02. Schema is compatible with v1 and v2.
 // - Version 4: Android V base. Schema is compatible with v1, v2 and v3.
 // - Version 5: M-2025-02. Schema is compatible with v1, v2, v3 and v4.
-inline static constexpr int32_t kVersion = 6;
+inline static constexpr int32_t kVersion = 9;
 inline static constexpr int32_t kVersionOne = 1;
 inline static constexpr int32_t kVersionTwo = 2;
 inline static constexpr int32_t kVersionThree = 3;

diff --git a/icing/file/version-util_test.cc b/icing/file/version-util_test.cc
index deec9d1..ca54af9 100644
--- a/icing/file/version-util_test.cc
+++ b/icing/file/version-util_test.cc

@@ -1076,6 +1076,92 @@
                 /*needs_term_index_rebuild_in=*/false,
                 /*needs_integer_index_rebuild_in=*/false,
                 /*needs_qualified_id_join_index_rebuild_in=*/true,
+                /*needs_embedding_index_rebuild_in=*/false)),
+
+        // - Existing version 8, max_version 8
+        // - Existing enabled features = {}
+        // - Current version = 8
+        // - Current enabled features = {FEATURE_DELETE_PROPAGATION_FROM}
+        //
+        // - Result: rebuild qualified id join index
+        VersionUtilDerivedFilesRebuildTestParam(
+            /*existing_version_in=*/8, /*max_version_in=*/8,
+            /*existing_enabled_features_in=*/{},
+            /*curr_version_in=*/8, /*curr_enabled_features_in=*/
+            {IcingSearchEngineFeatureInfoProto::
+                 FEATURE_DELETE_PROPAGATION_FROM},
+            /*expected_derived_files_rebuild_info_in=*/
+            DerivedFilesRebuildInfo(
+                /*needs_document_store_derived_files_rebuild_in=*/false,
+                /*needs_schema_store_derived_files_rebuild_in=*/false,
+                /*needs_term_index_rebuild_in=*/false,
+                /*needs_integer_index_rebuild_in=*/false,
+                /*needs_qualified_id_join_index_rebuild_in=*/true,
+                /*needs_embedding_index_rebuild_in=*/false)),
+
+        // - Existing version 8, max_version 8
+        // - Existing enabled features =
+        //   {FEATURE_DELETE_PROPAGATION_FROM}
+        // - Current version = 8
+        // - Current enabled features = {}
+        //
+        // - Result: rebuild qualified id join index
+        VersionUtilDerivedFilesRebuildTestParam(
+            /*existing_version_in=*/8, /*max_version_in=*/8,
+            /*existing_enabled_features_in=*/
+            {IcingSearchEngineFeatureInfoProto::
+                 FEATURE_DELETE_PROPAGATION_FROM},
+            /*curr_version_in=*/8, /*curr_enabled_features_in=*/{},
+            /*expected_derived_files_rebuild_info_in=*/
+            DerivedFilesRebuildInfo(
+                /*needs_document_store_derived_files_rebuild_in=*/false,
+                /*needs_schema_store_derived_files_rebuild_in=*/false,
+                /*needs_term_index_rebuild_in=*/false,
+                /*needs_integer_index_rebuild_in=*/false,
+                /*needs_qualified_id_join_index_rebuild_in=*/true,
+                /*needs_embedding_index_rebuild_in=*/false)),
+
+        // - Existing version 5, max_version 5
+        // - Existing enabled features = {}
+        // - Current version = 5
+        // - Current enabled features = {FEATURE_NON_EXISTENT_QUALIFIED_ID_JOIN}
+        //
+        // - Result: rebuild qualified id join index
+        VersionUtilDerivedFilesRebuildTestParam(
+            /*existing_version_in=*/5, /*max_version_in=*/5,
+            /*existing_enabled_features_in=*/{},
+            /*curr_version_in=*/5, /*curr_enabled_features_in=*/
+            {IcingSearchEngineFeatureInfoProto::
+                 FEATURE_NON_EXISTENT_QUALIFIED_ID_JOIN},
+            /*expected_derived_files_rebuild_info_in=*/
+            DerivedFilesRebuildInfo(
+                /*needs_document_store_derived_files_rebuild_in=*/false,
+                /*needs_schema_store_derived_files_rebuild_in=*/false,
+                /*needs_term_index_rebuild_in=*/false,
+                /*needs_integer_index_rebuild_in=*/false,
+                /*needs_qualified_id_join_index_rebuild_in=*/true,
+                /*needs_embedding_index_rebuild_in=*/false)),
+
+        // - Existing version 5, max_version 5
+        // - Existing enabled features =
+        //   {FEATURE_NON_EXISTENT_QUALIFIED_ID_JOIN}
+        // - Current version = 5
+        // - Current enabled features = {}
+        //
+        // - Result: rebuild qualified id join index
+        VersionUtilDerivedFilesRebuildTestParam(
+            /*existing_version_in=*/5, /*max_version_in=*/5,
+            /*existing_enabled_features_in=*/
+            {IcingSearchEngineFeatureInfoProto::
+                 FEATURE_NON_EXISTENT_QUALIFIED_ID_JOIN},
+            /*curr_version_in=*/5, /*curr_enabled_features_in=*/{},
+            /*expected_derived_files_rebuild_info_in=*/
+            DerivedFilesRebuildInfo(
+                /*needs_document_store_derived_files_rebuild_in=*/false,
+                /*needs_schema_store_derived_files_rebuild_in=*/false,
+                /*needs_term_index_rebuild_in=*/false,
+                /*needs_integer_index_rebuild_in=*/false,
+                /*needs_qualified_id_join_index_rebuild_in=*/true,
                 /*needs_embedding_index_rebuild_in=*/false))));
 
 TEST(VersionUtilTest, ShouldRebuildDerivedFilesUndeterminedVersion) {
@@ -1293,6 +1379,34 @@
           /*needs_embedding_index_rebuild_in=*/false)));
 }
 
+TEST(VersionUtilTest,
+     GetFeatureDerivedFilesRebuildInfo_featureDeletePropagationFrom) {
+  EXPECT_THAT(
+      GetFeatureDerivedFilesRebuildInfo(
+          IcingSearchEngineFeatureInfoProto::FEATURE_DELETE_PROPAGATION_FROM),
+      Eq(DerivedFilesRebuildInfo(
+          /*needs_document_store_derived_files_rebuild_in=*/false,
+          /*needs_schema_store_derived_files_rebuild_in=*/false,
+          /*needs_term_index_rebuild_in=*/false,
+          /*needs_integer_index_rebuild_in=*/false,
+          /*needs_qualified_id_join_index_rebuild_in=*/true,
+          /*needs_embedding_index_rebuild_in=*/false)));
+}
+
+TEST(VersionUtilTest,
+     GetFeatureDerivedFilesRebuildInfo_featureNonExistentQualifiedIdJoin) {
+  EXPECT_THAT(GetFeatureDerivedFilesRebuildInfo(
+                  IcingSearchEngineFeatureInfoProto::
+                      FEATURE_NON_EXISTENT_QUALIFIED_ID_JOIN),
+              Eq(DerivedFilesRebuildInfo(
+                  /*needs_document_store_derived_files_rebuild_in=*/false,
+                  /*needs_schema_store_derived_files_rebuild_in=*/false,
+                  /*needs_term_index_rebuild_in=*/false,
+                  /*needs_integer_index_rebuild_in=*/false,
+                  /*needs_qualified_id_join_index_rebuild_in=*/true,
+                  /*needs_embedding_index_rebuild_in=*/false)));
+}
+
 TEST(VersionUtilTest, SchemaDatabaseMigrationRequired) {
   // Migration is required if the previous version is less than the version at
   // which the database field is introduced.
@@ -1365,6 +1479,34 @@
                                 FEATURE_QUALIFIED_ID_JOIN_INDEX_V3))));
 }
 
+TEST(VersionUtilTest,
+     IcingSearchEngineOptionsToVersionProto_deletePropagationFrom_enabled) {
+  IcingSearchEngineOptions options;
+  options.set_enable_delete_propagation_from(true);
+
+  IcingSearchEngineVersionProto version_proto;
+  AddEnabledFeatures(options, &version_proto);
+  EXPECT_THAT(
+      version_proto.enabled_features(),
+      Contains(Property(
+          &IcingSearchEngineFeatureInfoProto::feature_type,
+          IcingSearchEngineFeatureInfoProto::FEATURE_DELETE_PROPAGATION_FROM)));
+}
+
+TEST(VersionUtilTest,
+     IcingSearchEngineOptionsToVersionProto_deletePropagationFrom_disabled) {
+  IcingSearchEngineOptions options;
+  options.set_enable_delete_propagation_from(false);
+
+  IcingSearchEngineVersionProto version_proto;
+  AddEnabledFeatures(options, &version_proto);
+  EXPECT_THAT(
+      version_proto.enabled_features(),
+      Not(Contains(Property(&IcingSearchEngineFeatureInfoProto::feature_type,
+                            IcingSearchEngineFeatureInfoProto::
+                                FEATURE_DELETE_PROPAGATION_FROM))));
+}
+
 class VersionUtilFeatureProtoTest
     : public ::testing::TestWithParam<
           IcingSearchEngineFeatureInfoProto::FlaggedFeatureType> {};
@@ -1402,7 +1544,10 @@
         IcingSearchEngineFeatureInfoProto::FEATURE_EMBEDDING_INDEX,
         IcingSearchEngineFeatureInfoProto::FEATURE_EMBEDDING_QUANTIZATION,
         IcingSearchEngineFeatureInfoProto::FEATURE_SCHEMA_DATABASE,
-        IcingSearchEngineFeatureInfoProto::FEATURE_QUALIFIED_ID_JOIN_INDEX_V3));
+        IcingSearchEngineFeatureInfoProto::FEATURE_QUALIFIED_ID_JOIN_INDEX_V3,
+        IcingSearchEngineFeatureInfoProto::FEATURE_DELETE_PROPAGATION_FROM,
+        IcingSearchEngineFeatureInfoProto::
+            FEATURE_NON_EXISTENT_QUALIFIED_ID_JOIN));
 
 }  // namespace
 

diff --git a/icing/graph/graph-interface.h b/icing/graph/graph-interface.h
new file mode 100644
index 0000000..6605f28
--- /dev/null
+++ b/icing/graph/graph-interface.h

@@ -0,0 +1,74 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef ICING_GRAPH_GRAPH_INTERFACE_H_
+#define ICING_GRAPH_GRAPH_INTERFACE_H_
+
+#include <memory>
+
+#include "icing/text_classifier/lib3/utils/base/status.h"
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
+
+namespace icing {
+namespace lib {
+
+namespace graph {
+
+// Graph interface with integer node id and given EdgeType.
+template <typename EdgeType>
+class GraphInterface {
+ public:
+  // Edge iterator interface.
+  class EdgeIteratorIf {
+   public:
+    virtual ~EdgeIteratorIf() = default;
+
+    // Advances to the next edge.
+    //
+    // Returns:
+    // - OK if successfully advanced to the next edge.
+    // - RESOURCE_EXHAUSTED_ERROR if there is no more edge to advance to.
+    // - Any other errors from the underlying implementation.
+    virtual libtextclassifier3::Status Advance() = 0;
+
+    // Gets the current edge.
+    //
+    // REQUIRES: preceding Advance() succeeded.
+    virtual const EdgeType& Get() const = 0;
+  };
+
+  virtual ~GraphInterface() = default;
+
+  virtual int GetNumNodes() const = 0;
+
+  // Returns an iterator to the (out) edges of the given node.
+  //
+  // Returns:
+  // - On success, a non-null unique pointer of EdgeIteratorIf. If there is no
+  //   edge for a valid node_id, then it should still return a valid iterator
+  //   with no edge to advance to.
+  // - INVALID_ARGUMENT_ERROR if node_id is invalid.
+  // - Any other errors from the underlying implementation.
+  virtual libtextclassifier3::StatusOr<std::unique_ptr<EdgeIteratorIf>>
+  GetEdgesIterator(int node_id) const = 0;
+
+  // Add GetInEdgesIterator if needed.
+};
+
+}  // namespace graph
+
+}  // namespace lib
+}  // namespace icing
+
+#endif  // ICING_GRAPH_GRAPH_INTERFACE_H_

diff --git a/icing/graph/simple-graph.h b/icing/graph/simple-graph.h
new file mode 100644
index 0000000..13088b3
--- /dev/null
+++ b/icing/graph/simple-graph.h

@@ -0,0 +1,107 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef ICING_GRAPH_SIMPLE_GRAPH_H_
+#define ICING_GRAPH_SIMPLE_GRAPH_H_
+
+#include <memory>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "icing/text_classifier/lib3/utils/base/status.h"
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
+#include "icing/absl_ports/canonical_errors.h"
+#include "icing/graph/graph-interface.h"
+
+namespace icing {
+namespace lib {
+
+namespace graph {
+
+// A simple in-memory graph data structure without edge weights.
+// - Node ids are from 0 to GetNumNodes() - 1.
+// - Edges only contain the connected node ids, and are stored in adjacent
+//   lists in memory.
+class SimpleGraph : public GraphInterface<int> {
+ public:
+  // Builder class for SimpleGraph.
+  class Builder {
+   public:
+    explicit Builder(int num_nodes) : out_(num_nodes) {}
+
+    // Builds the graph. It is undefined behavior to use the builder after this
+    // call.
+    SimpleGraph Build() { return SimpleGraph(std::move(out_)); }
+
+    // Adds an edge from node u to node v.
+    //
+    // REQUIRES: 0 <= u, v < num_nodes.
+    Builder& AddEdge(int u, int v) {
+      out_[u].insert(v);
+      return *this;
+    }
+
+   private:
+    std::vector<std::unordered_set<int>> out_;
+  };
+
+  int GetNumNodes() const override { return out_edges_.size(); }
+
+  libtextclassifier3::StatusOr<std::unique_ptr<EdgeIteratorIf>>
+  GetEdgesIterator(int node_id) const override {
+    if (node_id < 0 || node_id >= GetNumNodes()) {
+      return absl_ports::InvalidArgumentError("Invalid node id.");
+    }
+    return std::make_unique<EdgeIterator>(out_edges_[node_id].cbegin(),
+                                          out_edges_[node_id].size());
+  }
+
+ private:
+  class EdgeIterator : public EdgeIteratorIf {
+   public:
+    explicit EdgeIterator(std::unordered_set<int>::const_iterator it, int len)
+        : it_(std::move(it)), len_(len), num_advanced_(0) {}
+
+    libtextclassifier3::Status Advance() override {
+      if (num_advanced_ >= len_) {
+        return absl_ports::OutOfRangeError("No more edges to advance to.");
+      }
+      if (num_advanced_ != 0) {
+        ++it_;
+      }
+      ++num_advanced_;
+      return libtextclassifier3::Status::OK;
+    }
+
+    const int& Get() const override { return *it_; }
+
+   private:
+    std::unordered_set<int>::const_iterator it_;
+    int len_;
+    int num_advanced_;
+  };
+
+  explicit SimpleGraph(std::vector<std::unordered_set<int>>&& out_edges)
+      : out_edges_(std::move(out_edges)) {}
+
+  std::vector<std::unordered_set<int>> out_edges_;
+};
+
+}  // namespace graph
+
+}  // namespace lib
+}  // namespace icing
+
+#endif  // ICING_GRAPH_SIMPLE_GRAPH_H_

diff --git a/icing/icing-search-engine.cc b/icing/icing-search-engine.cc
index 0f85590..6b6d5db 100644
--- a/icing/icing-search-engine.cc
+++ b/icing/icing-search-engine.cc

@@ -18,7 +18,9 @@
 #include <cstddef>
 #include <cstdint>
 #include <functional>
+#include <limits>
 #include <memory>
+#include <optional>
 #include <string>
 #include <string_view>
 #include <unordered_map>
@@ -51,6 +53,8 @@
 #include "icing/index/term-indexing-handler.h"
 #include "icing/index/term-metadata.h"
 #include "icing/jni/jni-cache.h"
+#include "icing/join/delete-propagation-handler.h"
+#include "icing/join/document-dependency-processor.h"
 #include "icing/join/join-children-fetcher.h"
 #include "icing/join/join-processor.h"
 #include "icing/join/qualified-id-join-index-impl-v2.h"
@@ -105,6 +109,7 @@
 #include "icing/util/data-loss.h"
 #include "icing/util/icu-data-file-helper.h"
 #include "icing/util/logging.h"
+#include "icing/util/simple-task-scheduler.h"
 #include "icing/util/status-macros.h"
 #include "icing/util/status-util.h"
 #include "icing/util/tokenized-document.h"
@@ -118,6 +123,8 @@
 
 using ::icing::lib::status_util::TransformStatus;
 
+constexpr SimpleTaskScheduler::TaskId kHandleExpiredDocumentsTaskId = 1;
+
 constexpr std::string_view kDocumentSubfolderName = "document_dir";
 constexpr std::string_view kBlobSubfolderName = "blob_dir";
 constexpr std::string_view kIndexSubfolderName = "index_dir";
@@ -192,17 +199,17 @@
          result_grouping.entry_groupings()) {
       const std::string& name_space = entry.namespace_();
       const std::string& schema = entry.schema();
-      auto entry_id_or = document_store->GetResultGroupingEntryId(
-          result_grouping_type, name_space, schema);
-      if (!entry_id_or.ok()) {
+      std::optional<int32_t> entry_id =
+          document_store->GetResultGroupingEntryId(result_grouping_type,
+                                                   name_space, schema);
+      if (!entry_id.has_value()) {
         continue;
       }
-      int32_t entry_id = entry_id_or.ValueOrDie();
-      if (unique_entry_ids.find(entry_id) != unique_entry_ids.end()) {
+      if (unique_entry_ids.find(*entry_id) != unique_entry_ids.end()) {
         return absl_ports::InvalidArgumentError(
             "Entry Ids must be unique across result groups.");
       }
-      unique_entry_ids.insert(entry_id);
+      unique_entry_ids.insert(*entry_id);
     }
   }
   return libtextclassifier3::Status::OK;
@@ -384,21 +391,6 @@
   }
 }
 
-// Prepares the document for indexing. This includes tokenization and dependency
-// enforcement.
-libtextclassifier3::StatusOr<TokenizedDocument> PrepareDocumentForIndexing(
-    const SchemaStore* schema_store,
-    const LanguageSegmenter* language_segmenter, DocumentProto&& document) {
-  ICING_ASSIGN_OR_RETURN(
-      TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store, language_segmenter,
-                                std::move(document)));
-
-  // TODO(b/384947619): apply dependency enforcement.
-
-  return tokenized_document;
-}
-
 libtextclassifier3::Status RetrieveAndAddDocumentInfo(
     const DocumentStore* document_store, DeleteByQueryResultProto& result_proto,
     std::unordered_map<NamespaceTypePair,
@@ -508,7 +500,21 @@
                      options_.enable_repeated_field_joins(),
                      options_.enable_embedding_backup_generation(),
                      options_.enable_schema_database(),
-                     options_.release_backup_schema_file_if_overlay_present()),
+                     options_.release_backup_schema_file_if_overlay_present(),
+                     options_.enable_strict_page_byte_size_limit(),
+                     options_.enable_smaller_decompression_buffer_size(),
+                     options_.enable_eigen_embedding_scoring(),
+                     options_.enable_passing_filter_to_children(),
+                     options_.enable_proto_log_new_header_format(),
+                     options_.enable_embedding_iterator_v2(),
+                     options_.enable_reusable_decompression_buffer(),
+                     options_.enable_schema_type_id_optimization(),
+                     options_.enable_optimize_improvements(),
+                     options_.expired_document_purge_threshold_ms(),
+                     options_.enable_non_existent_qualified_id_join(),
+                     options_.enable_skip_set_schema_type_equality_check(),
+                     options_.enable_embed_query_optimization(),
+                     options_.enable_schema_definition_deduping()),
       filesystem_(std::move(filesystem)),
       icing_filesystem_(std::move(icing_filesystem)),
       clock_(std::move(clock)),
@@ -517,10 +523,16 @@
 }
 
 IcingSearchEngine::~IcingSearchEngine() {
+  // Stop all scheduled background tasks before persisting to disk.
+  DestroyTaskScheduler();
+
   if (initialized_) {
     if (PersistToDisk(PersistType::FULL).status().code() != StatusProto::OK) {
       ICING_LOG(ERROR)
           << "Error persisting to disk in IcingSearchEngine destructor";
+    } else {
+      ICING_LOG(INFO)
+          << "Done persisting to disk in IcingSearchEngine destructor";
     }
   }
 }
@@ -530,10 +542,32 @@
   // locks (reader and writer) has the chance to be interrupted during
   // switching.
   absl_ports::unique_lock l(&mutex_);
-  return InternalInitialize();
+
+  InitializeResultProto result = InitializeLocked();
+  ICING_LOG(INFO)
+      << "Initialize result: document store recovery cause = "
+      << static_cast<int>(
+             result.initialize_stats().document_store_recovery_cause())
+      << ", schema store recovery cause = "
+      << static_cast<int>(
+             result.initialize_stats().schema_store_recovery_cause())
+      << ", index restoration cause = "
+      << static_cast<int>(result.initialize_stats().index_restoration_cause())
+      << ", integer index restoration cause = "
+      << static_cast<int>(
+             result.initialize_stats().integer_index_restoration_cause())
+      << ", qualified id join index restoration cause = "
+      << static_cast<int>(result.initialize_stats()
+                              .qualified_id_join_index_restoration_cause())
+      << ", embedding index restoration cause = "
+      << static_cast<int>(
+             result.initialize_stats().embedding_index_restoration_cause());
+  result.set_vm_binder_transaction_latency_start_time_ms(
+      clock_->GetSystemTimeMilliseconds());
+  return result;
 }
 
-void IcingSearchEngine::ResetMembers() {
+void IcingSearchEngine::ResetMembersLocked() {
   // Reset all members in the reverse order of their initialization to ensure
   // the dependencies are not violated.
   embedding_index_.reset();
@@ -566,12 +600,13 @@
   libtextclassifier3::Status status;
   if (file_exists &&
       filesystem_->PRead(marker_file_fd->get(), &network_init_attempts,
-                         sizeof(network_init_attempts), /*offset=*/0)) {
+                         sizeof(network_init_attempts),
+                         /*offset=*/0) == sizeof(network_init_attempts)) {
     host_init_attempts = GNetworkToHostL(network_init_attempts);
     if (host_init_attempts > kMaxUnsuccessfulInitAttempts) {
       // We're tried and failed to init too many times. We need to throw
       // everything out and start from scratch.
-      ResetMembers();
+      ResetMembersLocked();
       marker_file_fd.reset();
 
       // Delete the entire base directory.
@@ -614,7 +649,7 @@
   return status;
 }
 
-InitializeResultProto IcingSearchEngine::InternalInitialize() {
+InitializeResultProto IcingSearchEngine::InitializeLocked() {
   ICING_VLOG(1) << "Initializing IcingSearchEngine in dir: "
                 << options_.base_dir();
 
@@ -635,11 +670,14 @@
   }
 
   if (options_.enable_delete_propagation_from() &&
-      !options_.enable_qualified_id_join_index_v3()) {
+      (!options_.enable_qualified_id_join_index_v3() ||
+       !options_.enable_soft_index_restoration())) {
     result_status->set_code(StatusProto::INVALID_ARGUMENT);
     result_status->set_message(
-        "Delete propagation is enabled but qualified id join index v3 is not "
-        "enabled.");
+        "Delete propagation is enabled but qualified id join index v3 or soft "
+        "index restoration is not enabled.");
+    initialize_stats->set_failure_stage(
+        InitializeStatsProto::FailureStage::OPTIONS_VALIDATION);
     return result_proto;
   }
 
@@ -653,6 +691,24 @@
       status = absl_ports::InternalError("Failed to delete init marker file!");
     } else {
       initialized_ = true;
+
+      // Set needs_persist_type if anything was rebuilt.
+      if (initialize_stats->schema_store_recovery_cause() !=
+              InitializeStatsProto::NONE ||
+          initialize_stats->document_store_data_status() !=
+              InitializeStatsProto::NO_DATA_LOSS ||
+          initialize_stats->document_store_recovery_cause() !=
+              InitializeStatsProto::NONE ||
+          initialize_stats->index_restoration_cause() !=
+              InitializeStatsProto::NONE ||
+          initialize_stats->integer_index_restoration_cause() !=
+              InitializeStatsProto::NONE ||
+          initialize_stats->qualified_id_join_index_restoration_cause() !=
+              InitializeStatsProto::NONE ||
+          initialize_stats->embedding_index_restoration_cause() !=
+              InitializeStatsProto::NONE) {
+        result_proto.set_needs_persist_type(PersistType::RECOVERY_PROOF);
+      }
     }
   }
   TransformStatus(status, result_status);
@@ -665,6 +721,8 @@
   ICING_RETURN_ERROR_IF_NULL(initialize_stats);
   // Make sure the base directory exists
   if (!filesystem_->CreateDirectoryRecursively(options_.base_dir().c_str())) {
+    initialize_stats->set_failure_stage(
+        InitializeStatsProto::FailureStage::BASE_DIRECTORY_CREATION);
     return absl_ports::InternalError(absl_ports::StrCat(
         "Could not create directory: ", options_.base_dir()));
   }
@@ -673,6 +731,8 @@
   // max number of init attempts.
   libtextclassifier3::Status status = CheckInitMarkerFile(initialize_stats);
   if (!status.ok() && !absl_ports::IsDataLoss(status)) {
+    initialize_stats->set_failure_stage(
+        InitializeStatsProto::FailureStage::INIT_MARKER_FILE);
     return status;
   }
 
@@ -686,10 +746,24 @@
   // Read version file, determine the state change and rebuild derived files if
   // needed.
   const std::string index_dir = MakeIndexDirectoryPath(options_.base_dir());
-  ICING_ASSIGN_OR_RETURN(
-      IcingSearchEngineVersionProto stored_version_proto,
-      version_util::ReadVersion(
-          *filesystem_, /*version_file_dir=*/options_.base_dir(), index_dir));
+  auto stored_version_proto_or = version_util::ReadVersion(
+      *filesystem_, /*version_file_dir=*/options_.base_dir(), index_dir);
+  if (!stored_version_proto_or.ok()) {
+    initialize_stats->set_failure_stage(
+        InitializeStatsProto::FailureStage::READ_VERSION_FILE);
+    ICING_LOG(ERROR) << "Failed to read version file. Error: "
+                     << stored_version_proto_or.status().error_code()
+                     << ", message: "
+                     << stored_version_proto_or.status().error_message();
+    return std::move(stored_version_proto_or).status();
+  }
+  IcingSearchEngineVersionProto stored_version_proto =
+      std::move(stored_version_proto_or).ValueOrDie();
+  ICING_LOG(INFO)
+      << "Successfully read version proto from existing dataset. Version: "
+      << stored_version_proto.version()
+      << ", max_version: " << stored_version_proto.max_version();
+
   version_util::VersionInfo stored_version_info =
       version_util::GetVersionInfoFromProto(stored_version_proto);
   version_util::StateChange version_state_change =
@@ -707,10 +781,15 @@
   bool perform_schema_database_migration =
       version_util::SchemaDatabaseMigrationRequired(stored_version_proto) &&
       options_.enable_schema_database();
-  ICING_RETURN_IF_ERROR(SchemaStore::MigrateSchema(
+  auto migrate_status = SchemaStore::MigrateSchema(
       filesystem_.get(), MakeSchemaDirectoryPath(options_.base_dir()),
       version_state_change, version_util::kVersion,
-      perform_schema_database_migration));
+      perform_schema_database_migration);
+  if (!migrate_status.ok()) {
+    initialize_stats->set_failure_stage(
+        InitializeStatsProto::FailureStage::MIGRATE_SCHEMA);
+    return migrate_status;
+  }
 
   // Step 2: Discard derived files that need to be rebuilt
   derived_file_util::DerivedFilesRebuildInfo required_derived_files_rebuild =
@@ -719,17 +798,32 @@
   if (existing_marker_proto != nullptr) {
     required_derived_files_rebuild.RebuildAll();
   }
-  ICING_RETURN_IF_ERROR(DiscardDerivedFiles(required_derived_files_rebuild));
+  auto discard_status = DiscardDerivedFiles(required_derived_files_rebuild);
+  if (!discard_status.ok()) {
+    initialize_stats->set_failure_stage(
+        InitializeStatsProto::FailureStage::DISCARD_DERIVED_FILES);
+    return discard_status;
+  }
 
   // Step 3: update version files. We need to update both the V1 and V2
   // version files.
-  ICING_RETURN_IF_ERROR(version_util::WriteV1Version(
+  auto write_version_status = version_util::WriteV1Version(
       *filesystem_, /*version_file_dir=*/options_.base_dir(),
-      version_util::GetVersionInfoFromProto(current_version_proto)));
-  ICING_RETURN_IF_ERROR(version_util::WriteV2Version(
+      version_util::GetVersionInfoFromProto(current_version_proto));
+  if (!write_version_status.ok()) {
+    initialize_stats->set_failure_stage(
+        InitializeStatsProto::FailureStage::WRITE_VERSION_FILE);
+    return write_version_status;
+  }
+  write_version_status = version_util::WriteV2Version(
       *filesystem_, /*version_file_dir=*/options_.base_dir(),
       std::make_unique<IcingSearchEngineVersionProto>(
-          std::move(current_version_proto))));
+          std::move(current_version_proto)));
+  if (!write_version_status.ok()) {
+    initialize_stats->set_failure_stage(
+        InitializeStatsProto::FailureStage::WRITE_VERSION_FILE);
+    return write_version_status;
+  }
 
   ICING_RETURN_IF_ERROR(InitializeSchemaStore(initialize_stats));
 
@@ -748,13 +842,35 @@
   // TODO(b/156383798) : Resolve how to specify the locale.
   language_segmenter_factory::SegmenterOptions segmenter_options(
       ULOC_US, jni_cache_.get(), enable_icu);
-  TC3_ASSIGN_OR_RETURN(language_segmenter_, language_segmenter_factory::Create(
-                                                std::move(segmenter_options)));
+  StatusProto* icu_segmenter_creation_status = nullptr;
+  if (enable_icu) {
+    icu_segmenter_creation_status =
+        initialize_stats->mutable_icu_segmenter_creation_status();
+  }
+  auto language_segmenter_or = language_segmenter_factory::Create(
+      std::move(segmenter_options), icu_segmenter_creation_status);
+  if (!language_segmenter_or.ok()) {
+    initialize_stats->set_failure_stage(
+        InitializeStatsProto::FailureStage::LANGUAGE_SEGMENTER_CREATION);
+    return std::move(language_segmenter_or).status();
+  }
+  language_segmenter_ = std::move(language_segmenter_or).ValueOrDie();
 
   NormalizerOptions normalizer_options(
       /*max_term_byte_size=*/options_.max_token_length(), enable_icu);
-  TC3_ASSIGN_OR_RETURN(normalizer_,
-                       normalizer_factory::Create(normalizer_options));
+  StatusProto* icu_normalizer_creation_status = nullptr;
+  if (enable_icu) {
+    icu_normalizer_creation_status =
+        initialize_stats->mutable_icu_normalizer_creation_status();
+  }
+  auto normalizer_or = normalizer_factory::Create(
+      normalizer_options, icu_normalizer_creation_status);
+  if (!normalizer_or.ok()) {
+    initialize_stats->set_failure_stage(
+        InitializeStatsProto::FailureStage::NORMALIZER_CREATION);
+    return std::move(normalizer_or).status();
+  }
+  normalizer_ = std::move(normalizer_or).ValueOrDie();
 
   libtextclassifier3::Status index_init_status;
   if (absl_ports::IsNotFound(schema_store_->GetSchema().status())) {
@@ -782,6 +898,8 @@
              .ok() ||
         !EmbeddingIndex::Discard(*filesystem_, embedding_index_dir).ok() ||
         !filesystem_->DeleteDirectoryRecursively(blob_store_dir.c_str())) {
+      initialize_stats->set_failure_stage(
+          InitializeStatsProto::FailureStage::DISCARD_COMPONENTS);
       return absl_ports::InternalError(absl_ports::StrCat(
           "Could not delete directories: ", index_dir, ", ", integer_index_dir,
           ", ", qualified_id_join_index_dir, ", ", embedding_index_dir, ", ",
@@ -790,9 +908,15 @@
 
     // Initialize (empty) blob store.
     if (options_.enable_blob_store()) {
-      ICING_RETURN_IF_ERROR(
+      auto blob_store_init_status =
           InitializeBlobStore(options_.orphan_blob_time_to_live_ms(),
-                              options_.blob_store_compression_level()));
+                              options_.blob_store_compression_level(),
+                              options_.blob_store_compression_mem_level());
+      if (!blob_store_init_status.ok()) {
+        initialize_stats->set_failure_stage(
+            InitializeStatsProto::FailureStage::BLOB_STORE_INSTANTIATION);
+        return blob_store_init_status;
+      }
     }
 
     // Initialize (empty) document store.
@@ -825,9 +949,15 @@
     // We just need to re-initialize each component here.
 
     if (options_.enable_blob_store()) {
-      ICING_RETURN_IF_ERROR(
+      auto blob_store_init_status =
           InitializeBlobStore(options_.orphan_blob_time_to_live_ms(),
-                              options_.blob_store_compression_level()));
+                              options_.blob_store_compression_level(),
+                              options_.blob_store_compression_mem_level());
+      if (!blob_store_init_status.ok()) {
+        initialize_stats->set_failure_stage(
+            InitializeStatsProto::FailureStage::BLOB_STORE_INSTANTIATION);
+        return blob_store_init_status;
+      }
     }
 
     // Initialize document store. This also rebuilds all derived files in the
@@ -868,9 +998,15 @@
 
     // Initialize blob store.
     if (options_.enable_blob_store()) {
-      ICING_RETURN_IF_ERROR(
+      auto blob_store_init_status =
           InitializeBlobStore(options_.orphan_blob_time_to_live_ms(),
-                              options_.blob_store_compression_level()));
+                              options_.blob_store_compression_level(),
+                              options_.blob_store_compression_mem_level());
+      if (!blob_store_init_status.ok()) {
+        initialize_stats->set_failure_stage(
+            InitializeStatsProto::FailureStage::BLOB_STORE_INSTANTIATION);
+        return blob_store_init_status;
+      }
     }
 
     // Initialize document store. This also rebuilds all derived files in the
@@ -918,10 +1054,40 @@
     }
   }
 
+  if (options_.enable_background_task_scheduler() &&
+      task_scheduler_ == nullptr) {
+    // Initialize the task scheduler.
+    //
+    // Note: InitializeMembers may be called by ResetLocked() when put API
+    //   fails. In that case, we can still use the existing task_scheduler_
+    //   and no need to re-initialize it.
+    task_scheduler_ = SimpleTaskScheduler::Create(*clock_);
+  }
+
   if (status.ok()) {
     status = index_init_status;
   }
 
+  // Call HandleExpiredDocumentsLocked() to handle documents that expire during
+  // Icing was off. This function will reschedule another task with the next
+  // expiration timestamp if task_scheduler_ is not null (i.e.
+  // options_.enable_background_task_scheduler() is true).
+  if ((status.ok() || absl_ports::IsDataLoss(status)) &&
+      options_.enable_delete_propagation_from()) {
+    // Call HandleExpiredDocumentsLocked() to handle documents that expire
+    // during Icing was off. This function will reschedule another task with the
+    // next expiration timestamp and activate the task.
+    HandleExpiredDocumentsResultProto result = HandleExpiredDocumentsLocked();
+    if (result.status().code() != StatusProto::OK) {
+      ICING_LOG(ERROR) << "Failed to handle expired documents during "
+                          "initialization: "
+                       << result.status().message();
+      return absl_ports::InternalError(result.status().message());
+    }
+    initialize_stats->set_next_expiration_timestamp_ms(
+        result.next_expiration_timestamp_ms());
+  }
+
   result_state_manager_ = std::make_unique<ResultStateManager>(
       performance_configuration_.max_num_total_hits, *document_store_);
 
@@ -936,13 +1102,21 @@
       MakeSchemaDirectoryPath(options_.base_dir());
   // Make sure the sub-directory exists
   if (!filesystem_->CreateDirectoryRecursively(schema_store_dir.c_str())) {
+    initialize_stats->set_failure_stage(
+        InitializeStatsProto::FailureStage::SCHEMA_STORE_DIRECTORY_CREATION);
     return absl_ports::InternalError(
         absl_ports::StrCat("Could not create directory: ", schema_store_dir));
   }
-  ICING_ASSIGN_OR_RETURN(
-      schema_store_,
+
+  auto schema_store_or =
       SchemaStore::Create(filesystem_.get(), schema_store_dir, clock_.get(),
-                          &feature_flags_, initialize_stats));
+                          &feature_flags_, initialize_stats);
+  if (!schema_store_or.ok()) {
+    initialize_stats->set_failure_stage(
+        InitializeStatsProto::FailureStage::SCHEMA_STORE_INSTANTIATION);
+    return schema_store_or.status();
+  }
+  schema_store_ = std::move(schema_store_or).ValueOrDie();
 
   return libtextclassifier3::Status::OK;
 }
@@ -956,22 +1130,33 @@
       MakeDocumentDirectoryPath(options_.base_dir());
   // Make sure the sub-directory exists
   if (!filesystem_->CreateDirectoryRecursively(document_dir.c_str())) {
+    initialize_stats->set_failure_stage(
+        InitializeStatsProto::FailureStage::DOCUMENT_STORE_DIRECTORY_CREATION);
     return absl_ports::InternalError(
         absl_ports::StrCat("Could not create directory: ", document_dir));
   }
-  ICING_ASSIGN_OR_RETURN(
-      DocumentStore::CreateResult create_result,
-      DocumentStore::Create(
-          filesystem_.get(), document_dir, clock_.get(), schema_store_.get(),
-          &feature_flags_, force_recovery_and_revalidate_documents,
-          /*pre_mapping_fbv=*/false, /*use_persistent_hash_map=*/true,
-          options_.compression_level(), initialize_stats));
+
+  auto create_result_or = DocumentStore::Create(
+      filesystem_.get(), document_dir, clock_.get(), schema_store_.get(),
+      &feature_flags_, force_recovery_and_revalidate_documents,
+      /*pre_mapping_fbv=*/false, /*use_persistent_hash_map=*/true,
+      options_.compression_level(), options_.compression_threshold_bytes(),
+      options_.compression_mem_level(), initialize_stats);
+  if (!create_result_or.ok()) {
+    initialize_stats->set_failure_stage(
+        InitializeStatsProto::FailureStage::DOCUMENT_STORE_INSTANTIATION);
+    return std::move(create_result_or).status();
+  }
+  DocumentStore::CreateResult create_result =
+      std::move(create_result_or).ValueOrDie();
+
   document_store_ = std::move(create_result.document_store);
   return create_result.derived_files_regenerated;
 }
 
 libtextclassifier3::Status IcingSearchEngine::InitializeBlobStore(
-    int32_t orphan_blob_time_to_live_ms, int32_t blob_store_compression_level) {
+    int32_t orphan_blob_time_to_live_ms, int32_t blob_store_compression_level,
+    int32_t blob_store_compression_mem_level) {
   std::string blob_dir = MakeBlobDirectoryPath(options_.base_dir());
   // Make sure the sub-directory exists
   if (!filesystem_->CreateDirectoryRecursively(blob_dir.c_str())) {
@@ -984,7 +1169,8 @@
       BlobStore::Create(filesystem_.get(), blob_dir, clock_.get(),
                         orphan_blob_time_to_live_ms,
                         blob_store_compression_level,
-                        options_.manage_blob_files()));
+                        blob_store_compression_mem_level,
+                        options_.manage_blob_files(), &feature_flags_));
   blob_store_ = std::make_unique<BlobStore>(std::move(blob_store_or));
   return libtextclassifier3::Status::OK;
 }
@@ -997,6 +1183,8 @@
   const std::string index_dir = MakeIndexDirectoryPath(options_.base_dir());
   // Make sure the sub-directory exists
   if (!filesystem_->CreateDirectoryRecursively(index_dir.c_str())) {
+    initialize_stats->set_failure_stage(
+        InitializeStatsProto::FailureStage::TERM_INDEX_DIRECTORY);
     return absl_ports::InternalError(
         absl_ports::StrCat("Could not create directory: ", index_dir));
   }
@@ -1006,11 +1194,13 @@
 
   // Term index
   InitializeStatsProto::RecoveryCause index_recovery_cause;
-  auto index_or =
-      Index::Create(index_options, filesystem_.get(), icing_filesystem_.get());
+  auto index_or = Index::Create(index_options, filesystem_.get(),
+                                icing_filesystem_.get(), &feature_flags_);
   if (!index_or.ok()) {
     if (!filesystem_->DeleteDirectoryRecursively(index_dir.c_str()) ||
         !filesystem_->CreateDirectoryRecursively(index_dir.c_str())) {
+      initialize_stats->set_failure_stage(
+          InitializeStatsProto::FailureStage::TERM_INDEX_DIRECTORY);
       return absl_ports::InternalError(
           absl_ports::StrCat("Could not recreate directory: ", index_dir));
     }
@@ -1018,9 +1208,14 @@
     index_recovery_cause = InitializeStatsProto::IO_ERROR;
 
     // Try recreating it from scratch and re-indexing everything.
-    ICING_ASSIGN_OR_RETURN(index_,
-                           Index::Create(index_options, filesystem_.get(),
-                                         icing_filesystem_.get()));
+    index_or = Index::Create(index_options, filesystem_.get(),
+                             icing_filesystem_.get(), &feature_flags_);
+    if (!index_or.ok()) {
+      initialize_stats->set_failure_stage(
+          InitializeStatsProto::FailureStage::TERM_INDEX_INSTANTIATION);
+      return std::move(index_or).status();
+    }
+    index_ = std::move(index_or).ValueOrDie();
   } else {
     // Index was created fine.
     index_ = std::move(index_or).ValueOrDie();
@@ -1038,17 +1233,27 @@
                            options_.integer_index_bucket_split_threshold(),
                            options_.pre_mapping_fbv());
   if (!integer_index_or.ok()) {
-    ICING_RETURN_IF_ERROR(
-        IntegerIndex::Discard(*filesystem_, integer_index_dir));
+    auto discard_status =
+        IntegerIndex::Discard(*filesystem_, integer_index_dir);
+    if (!discard_status.ok()) {
+      initialize_stats->set_failure_stage(
+          InitializeStatsProto::FailureStage::INTEGER_INDEX_DIRECTORY);
+      return discard_status;
+    }
 
     integer_index_recovery_cause = InitializeStatsProto::IO_ERROR;
 
     // Try recreating it from scratch and re-indexing everything.
-    ICING_ASSIGN_OR_RETURN(
-        integer_index_,
+    integer_index_or =
         IntegerIndex::Create(*filesystem_, std::move(integer_index_dir),
                              options_.integer_index_bucket_split_threshold(),
-                             options_.pre_mapping_fbv()));
+                             options_.pre_mapping_fbv());
+    if (!integer_index_or.ok()) {
+      initialize_stats->set_failure_stage(
+          InitializeStatsProto::FailureStage::INTEGER_INDEX_INSTANTIATION);
+      return std::move(integer_index_or).status();
+    }
+    integer_index_ = std::move(integer_index_or).ValueOrDie();
   } else {
     // Integer index was created fine.
     integer_index_ = std::move(integer_index_or).ValueOrDie();
@@ -1067,14 +1272,25 @@
     // V2 qualified id join index depends on document store derived files, so we
     // have to rebuild it from scratch if
     // document_store_derived_files_regenerated is true.
-    ICING_RETURN_IF_ERROR(QualifiedIdJoinIndex::Discard(
-        *filesystem_, qualified_id_join_index_dir));
+    auto discard_status = QualifiedIdJoinIndex::Discard(
+        *filesystem_, qualified_id_join_index_dir);
+    if (!discard_status.ok()) {
+      initialize_stats->set_failure_stage(
+          InitializeStatsProto::FailureStage::
+              QUALIFIED_ID_JOIN_INDEX_DIRECTORY);
+      return discard_status;
+    }
 
-    ICING_ASSIGN_OR_RETURN(
-        qualified_id_join_index_,
-        CreateQualifiedIdJoinIndex(*filesystem_,
-                                   std::move(qualified_id_join_index_dir),
-                                   options_, feature_flags_));
+    auto qualified_id_join_index_or = CreateQualifiedIdJoinIndex(
+        *filesystem_, qualified_id_join_index_dir, options_, feature_flags_);
+    if (!qualified_id_join_index_or.ok()) {
+      initialize_stats->set_failure_stage(
+          InitializeStatsProto::FailureStage::
+              QUALIFIED_ID_JOIN_INDEX_INSTANTIATION);
+      return std::move(qualified_id_join_index_or).status();
+    }
+    qualified_id_join_index_ =
+        std::move(qualified_id_join_index_or).ValueOrDie();
 
     qualified_id_join_index_recovery_cause =
         InitializeStatsProto::DEPENDENCIES_CHANGED;
@@ -1082,17 +1298,29 @@
     auto qualified_id_join_index_or = CreateQualifiedIdJoinIndex(
         *filesystem_, qualified_id_join_index_dir, options_, feature_flags_);
     if (!qualified_id_join_index_or.ok()) {
-      ICING_RETURN_IF_ERROR(QualifiedIdJoinIndex::Discard(
-          *filesystem_, qualified_id_join_index_dir));
+      auto discard_status = QualifiedIdJoinIndex::Discard(
+          *filesystem_, qualified_id_join_index_dir);
+      if (!discard_status.ok()) {
+        initialize_stats->set_failure_stage(
+            InitializeStatsProto::FailureStage::
+                QUALIFIED_ID_JOIN_INDEX_DIRECTORY);
+        return discard_status;
+      }
 
       qualified_id_join_index_recovery_cause = InitializeStatsProto::IO_ERROR;
 
       // Try recreating it from scratch and rebuild everything.
-      ICING_ASSIGN_OR_RETURN(
-          qualified_id_join_index_,
-          CreateQualifiedIdJoinIndex(*filesystem_,
-                                     std::move(qualified_id_join_index_dir),
-                                     options_, feature_flags_));
+      qualified_id_join_index_or = CreateQualifiedIdJoinIndex(
+          *filesystem_, std::move(qualified_id_join_index_dir), options_,
+          feature_flags_);
+      if (!qualified_id_join_index_or.ok()) {
+        initialize_stats->set_failure_stage(
+            InitializeStatsProto::FailureStage::
+                QUALIFIED_ID_JOIN_INDEX_INSTANTIATION);
+        return std::move(qualified_id_join_index_or).status();
+      }
+      qualified_id_join_index_ =
+          std::move(qualified_id_join_index_or).ValueOrDie();
     } else {
       // Qualified id join index was created fine.
       qualified_id_join_index_ =
@@ -1109,17 +1337,28 @@
       MakeEmbeddingIndexWorkingPath(options_.base_dir());
   InitializeStatsProto::RecoveryCause embedding_index_recovery_cause;
   auto embedding_index_or = EmbeddingIndex::Create(
-      filesystem_.get(), embedding_dir, clock_.get(), &feature_flags_);
+      filesystem_.get(), embedding_dir, clock_.get(), &feature_flags_,
+      options_.embedding_index_num_shards());
   if (!embedding_index_or.ok()) {
-    ICING_RETURN_IF_ERROR(EmbeddingIndex::Discard(*filesystem_, embedding_dir));
+    auto discard_status = EmbeddingIndex::Discard(*filesystem_, embedding_dir);
+    if (!discard_status.ok()) {
+      initialize_stats->set_failure_stage(
+          InitializeStatsProto::FailureStage::EMBEDDING_INDEX_DIRECTORY);
+      return discard_status;
+    }
 
     embedding_index_recovery_cause = InitializeStatsProto::IO_ERROR;
 
     // Try recreating it from scratch and re-indexing everything.
-    ICING_ASSIGN_OR_RETURN(
-        embedding_index_,
-        EmbeddingIndex::Create(filesystem_.get(), embedding_dir, clock_.get(),
-                               &feature_flags_));
+    embedding_index_or = EmbeddingIndex::Create(
+        filesystem_.get(), embedding_dir, clock_.get(), &feature_flags_,
+        options_.embedding_index_num_shards());
+    if (!embedding_index_or.ok()) {
+      initialize_stats->set_failure_stage(
+          InitializeStatsProto::FailureStage::EMBEDDING_INDEX_INSTANTIATION);
+      return std::move(embedding_index_or).status();
+    }
+    embedding_index_ = std::move(embedding_index_or).ValueOrDie();
   } else {
     // Embedding index was created fine.
     embedding_index_ = std::move(embedding_index_or).ValueOrDie();
@@ -1156,6 +1395,11 @@
           embedding_index_recovery_cause);
     }
   }
+  if (!restore_result.status.ok() &&
+      !absl_ports::IsDataLoss(restore_result.status)) {
+    initialize_stats->set_failure_stage(
+        InitializeStatsProto::FailureStage::RESTORE_INDEX);
+  }
   return restore_result.status;
 }
 
@@ -1170,8 +1414,10 @@
   *set_schema_request.mutable_schema() = std::move(new_schema);
   set_schema_request.set_ignore_errors_and_delete_documents(
       ignore_errors_and_delete_documents);
-
-  return SetSchema(std::move(set_schema_request));
+  SetSchemaResultProto result_proto = SetSchema(std::move(set_schema_request));
+  result_proto.set_vm_binder_transaction_latency_start_time_ms(
+      clock_->GetSystemTimeMilliseconds());
+  return result_proto;
 }
 
 SetSchemaResultProto IcingSearchEngine::SetSchema(
@@ -1179,12 +1425,15 @@
   ICING_VLOG(1) << "Setting new Schema";
 
   SetSchemaResultProto result_proto;
+  SetSchemaStatsProto* set_schema_stats =
+      result_proto.mutable_set_schema_stats();
   StatusProto* result_status = result_proto.mutable_status();
 
   absl_ports::unique_lock l(&mutex_);
-  ScopedTimer timer(clock_->GetNewTimer(), [&result_proto](int64_t t) {
-    result_proto.set_latency_ms(t);
-  });
+  ScopedTimer overall_timer(clock_->GetNewTimer(),
+                            [&set_schema_stats](int64_t t) {
+                              set_schema_stats->set_overall_latency_ms(t);
+                            });
   if (!initialized_) {
     result_status->set_code(StatusProto::FAILED_PRECONDITION);
     result_status->set_message("IcingSearchEngine has not been initialized!");
@@ -1261,68 +1510,134 @@
   //   stores schema type id (+ joinable property path) as a key to group join
   //   data.
   bool join_incompatible =
-      !set_schema_result.schema_types_join_incompatible_by_name.empty() ||
-      !set_schema_result.old_schema_type_ids_changed.empty();
+      !set_schema_result.schema_types_join_incompatible_by_name.empty();
+  if (!options_.enable_qualified_id_join_index_v3()) {
+    join_incompatible |= !set_schema_result.old_schema_type_ids_changed.empty();
+  }
   for (const std::string& join_incompatible_type :
        set_schema_result.schema_types_join_incompatible_by_name) {
     result_proto.add_join_incompatible_changed_schema_types(
         std::move(join_incompatible_type));
   }
 
+  set_schema_stats->set_schema_store_set_schema_latency_ms(
+      overall_timer.timer().GetElapsedMilliseconds());
+
   libtextclassifier3::Status status;
   if (set_schema_result.success) {
+    // - No need to manually persist the schema file since
+    //   SchemaStore::SetSchema already handles schema store file persistence
+    //   internally.
+    // - We just need to detect ground truth and derived files changes from the
+    //   document store and indices.
+    bool needs_flush_ground_truth = false;
+    bool needs_flush_derived_files = false;
+
+    // Update document store if necessary.
+    std::optional<DocumentStore::UpdateSchemaStoreResult> update_result;
     if (lost_previous_schema) {
+      ScopedTimer update_schema_store_timer(
+          clock_->GetNewTimer(), [&set_schema_stats](int64_t t) {
+            set_schema_stats->set_document_store_update_schema_latency_ms(t);
+          });
       // No previous schema to calculate a diff against. We have to go through
       // and revalidate all the Documents in the DocumentStore
-      status = document_store_->UpdateSchemaStore(schema_store_.get());
-      if (!status.ok()) {
-        TransformStatus(status, result_status);
+      auto update_result_or =
+          document_store_->UpdateSchemaStore(schema_store_.get());
+      if (!update_result_or.ok()) {
+        TransformStatus(update_result_or.status(), result_status);
         return result_proto;
       }
+      update_result =
+          std::make_optional(std::move(update_result_or).ValueOrDie());
     } else if (!set_schema_result.old_schema_type_ids_changed.empty() ||
                !set_schema_result.schema_types_incompatible_by_id.empty() ||
                !set_schema_result.schema_types_deleted_by_id.empty()) {
-      status = document_store_->OptimizedUpdateSchemaStore(schema_store_.get(),
-                                                           set_schema_result);
-      if (!status.ok()) {
-        TransformStatus(status, result_status);
+      ScopedTimer optimized_update_schema_store_timer(
+          clock_->GetNewTimer(), [&set_schema_stats](int64_t t) {
+            set_schema_stats
+                ->set_document_store_optimized_update_schema_latency_ms(t);
+          });
+      auto update_result_or = document_store_->OptimizedUpdateSchemaStore(
+          schema_store_.get(), set_schema_result);
+      if (!update_result_or.ok()) {
+        TransformStatus(update_result_or.status(), result_status);
         return result_proto;
       }
+      update_result =
+          std::make_optional(std::move(update_result_or).ValueOrDie());
+    }
+    if (update_result.has_value()) {
+      result_proto.set_deleted_document_count(
+          update_result->deleted_document_count);
+      needs_flush_ground_truth |= (update_result->deleted_document_count > 0);
+      needs_flush_derived_files |= update_result->derived_files_changed;
     }
 
-    if (lost_previous_schema || index_incompatible) {
-      // Clears search indices
-      status = ClearSearchIndices();
-      if (!status.ok()) {
-        TransformStatus(status, result_status);
-        return result_proto;
+    {
+      // Restore indices if needed.
+      ScopedTimer index_restoration_timer(
+          clock_->GetNewTimer(), [&set_schema_stats](int64_t t) {
+            set_schema_stats->set_index_restoration_latency_ms(t);
+          });
+
+      if (lost_previous_schema || index_incompatible) {
+        // Clears search indices
+        status = ClearSearchIndices();
+        if (!status.ok()) {
+          TransformStatus(status, result_status);
+          return result_proto;
+        }
       }
-    }
 
-    if (lost_previous_schema || join_incompatible) {
-      // Clears join indices
-      status = ClearJoinIndices();
-      if (!status.ok()) {
-        TransformStatus(status, result_status);
-        return result_proto;
+      if (lost_previous_schema || join_incompatible) {
+        // Clears join indices
+        status = ClearJoinIndices();
+        if (!status.ok()) {
+          TransformStatus(status, result_status);
+          return result_proto;
+        }
       }
-    }
 
-    if (lost_previous_schema || index_incompatible || join_incompatible) {
-      IndexRestorationResult restore_result = RestoreIndexIfNeeded();
-      // DATA_LOSS means that we have successfully re-added content to the
-      // index. Some indexed content was lost, but otherwise the index is in a
-      // valid state and can be queried.
-      if (!restore_result.status.ok() &&
-          !absl_ports::IsDataLoss(restore_result.status)) {
-        TransformStatus(status, result_status);
-        return result_proto;
+      if (lost_previous_schema || index_incompatible || join_incompatible) {
+        needs_flush_derived_files = true;
+
+        IndexRestorationResult restore_result = RestoreIndexIfNeeded();
+        result_proto.set_has_term_index_restored(
+            restore_result.has_index_restored);
+        result_proto.set_has_integer_index_restored(
+            restore_result.has_integer_index_restored);
+        result_proto.set_has_qualified_id_join_index_restored(
+            restore_result.has_qualified_id_join_index_restored);
+        result_proto.set_has_embedding_index_restored(
+            restore_result.has_embedding_index_restored);
+        // DATA_LOSS means that we have successfully re-added content to the
+        // index. Some indexed content was lost, but otherwise the index is in a
+        // valid state and can be queried.
+        if (!restore_result.status.ok() &&
+            !absl_ports::IsDataLoss(restore_result.status)) {
+          TransformStatus(status, result_status);
+          return result_proto;
+        }
       }
     }
 
     if (feature_flags_.enable_scorable_properties()) {
       if (!set_schema_result.schema_types_scorable_property_inconsistent_by_id
                .empty()) {
+        needs_flush_derived_files = true;
+
+        ScopedTimer scorable_property_cache_regeneration_timer(
+            clock_->GetNewTimer(), [&set_schema_stats](int64_t t) {
+              set_schema_stats
+                  ->set_scorable_property_cache_regeneration_latency_ms(t);
+            });
+        for (const std::string& scorable_property_incompatible_type :
+             set_schema_result
+                 .schema_types_scorable_property_inconsistent_by_name) {
+          result_proto.add_scorable_property_incompatible_changed_schema_types(
+              scorable_property_incompatible_type);
+        }
         status = document_store_->RegenerateScorablePropertyCache(
             set_schema_result
                 .schema_types_scorable_property_inconsistent_by_id);
@@ -1332,8 +1647,28 @@
         }
       }
     }
-
     result_status->set_code(StatusProto::OK);
+    if (needs_flush_derived_files) {
+      // If derived files need to be flushed, then we need RECOVERY_PROOF which:
+      // - Updates all checksums of derived files.
+      // - Flushes ground truth data.
+      //
+      // Here, it is ok to use RECOVERY_PROOF even if ground truth does not need
+      // to be flushed.
+      result_proto.set_needs_persist_type(PersistType::RECOVERY_PROOF);
+    } else if (needs_flush_ground_truth) {
+      // If derived files are unchanged but ground truth needs to be flushed,
+      // then we need LITE which only flushes ground truth data.
+      //
+      // Theoretically, it is impossible to have this case since:
+      // - The only possibility of ground truth change is incompatible document
+      //   deletion.
+      // - When deleting incompatible documents, derived files are always
+      //   changed, so it should belong to the RECOVERY_PROOF case.
+      result_proto.set_needs_persist_type(PersistType::LITE);
+    } else {
+      result_proto.set_needs_persist_type(PersistType::UNKNOWN);
+    }
   } else {
     result_status->set_code(StatusProto::FAILED_PRECONDITION);
     result_status->set_message("Schema is incompatible.");
@@ -1396,17 +1731,23 @@
   if (!initialized_) {
     result_status->set_code(StatusProto::FAILED_PRECONDITION);
     result_status->set_message("IcingSearchEngine has not been initialized!");
+    result_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
     return result_proto;
   }
 
   auto type_config_or = schema_store_->GetSchemaTypeConfig(schema_type);
   if (!type_config_or.ok()) {
     TransformStatus(type_config_or.status(), result_status);
+    result_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
     return result_proto;
   }
 
   result_status->set_code(StatusProto::OK);
   *result_proto.mutable_schema_type_config() = *(type_config_or.ValueOrDie());
+  result_proto.set_vm_binder_transaction_latency_start_time_ms(
+      clock_->GetSystemTimeMilliseconds());
   return result_proto;
 }
 
@@ -1414,19 +1755,52 @@
     PutDocumentRequest&& put_document_request) {
   BatchPutResultProto batch_put_result_proto;
 
-  // TODO(b/394875109): right now we lock in the Put(DocumentProto&&) for each
-  // document. We should considering just locking once for the whole batch here.
+  absl_ports::unique_lock l(&mutex_);  // Acquire lock once for the batch
+
+  if (!initialized_) {
+    // Handle not initialized case for all documents
+    for (const DocumentProto& document_proto :
+         put_document_request.documents()) {
+      PutResultProto* put_result =
+          batch_put_result_proto.mutable_put_result_protos()->Add();
+      put_result->set_uri(document_proto.uri());
+      put_result->mutable_status()->set_code(StatusProto::FAILED_PRECONDITION);
+    }
+    batch_put_result_proto.mutable_status()->set_message(
+        "IcingSearchEngine has not been initialized!");
+    batch_put_result_proto.mutable_status()->set_code(
+        StatusProto::FAILED_PRECONDITION);
+    batch_put_result_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
+    return batch_put_result_proto;
+  }
+
+  int64_t current_time_ms = clock_->GetSystemTimeMilliseconds();
+
   for (DocumentProto& document_proto :
        *(put_document_request.mutable_documents())) {
     batch_put_result_proto.mutable_put_result_protos()->Add(
-        Put(std::move(document_proto)));
+        PutLocked(std::move(document_proto),
+                  current_time_ms));  // Call the locked version
   }
 
   if (put_document_request.persist_type() != PersistType::UNKNOWN) {
-    *batch_put_result_proto.mutable_persist_to_disk_result_proto() =
-        PersistToDisk(put_document_request.persist_type());
+    // Measure the latency of the persist process.
+    std::unique_ptr<Timer> persist_timer = clock_->GetNewTimer();
+    PersistToDiskStatsProto* persist_stats =
+        batch_put_result_proto.mutable_persist_to_disk_result_proto()
+            ->mutable_persist_stats();
+    auto status =
+        PersistToDiskLocked(put_document_request.persist_type(), persist_stats);
+    TransformStatus(
+        status, batch_put_result_proto.mutable_persist_to_disk_result_proto()
+                    ->mutable_status());
+    persist_stats->set_latency_ms(persist_timer->GetElapsedMilliseconds());
   }
 
+  batch_put_result_proto.mutable_status()->set_code(StatusProto::OK);
+  batch_put_result_proto.set_vm_binder_transaction_latency_start_time_ms(
+      clock_->GetSystemTimeMilliseconds());
   return batch_put_result_proto;
 }
 
@@ -1435,8 +1809,31 @@
 }
 
 PutResultProto IcingSearchEngine::Put(DocumentProto&& document) {
+  absl_ports::unique_lock l(&mutex_);
   ICING_VLOG(1) << "Writing document to document store";
 
+  if (!initialized_) {
+    PutResultProto result_proto;
+    result_proto.set_uri(document.uri());
+    StatusProto* result_status = result_proto.mutable_status();
+    result_status->set_code(StatusProto::FAILED_PRECONDITION);
+    result_status->set_message("IcingSearchEngine has not been initialized!");
+    result_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
+    return result_proto;
+  }
+
+  int64_t current_time_ms = clock_->GetSystemTimeMilliseconds();
+
+  PutResultProto result_proto = PutLocked(std::move(document), current_time_ms);
+  result_proto.set_vm_binder_transaction_latency_start_time_ms(
+      clock_->GetSystemTimeMilliseconds());
+  return result_proto;
+}
+
+// PutLocked to be called when mutex_ is already held.
+PutResultProto IcingSearchEngine::PutLocked(DocumentProto&& document,
+                                            int64_t current_time_ms) {
   PutResultProto result_proto;
   result_proto.set_uri(document.uri());
 
@@ -1450,37 +1847,41 @@
   // Lock must be acquired before validation because the DocumentStore uses
   // the schema file to validate, and the schema could be changed in
   // SetSchema() which is protected by the same mutex.
-  absl_ports::unique_lock l(&mutex_);
-  if (!initialized_) {
-    result_status->set_code(StatusProto::FAILED_PRECONDITION);
-    result_status->set_message("IcingSearchEngine has not been initialized!");
-    return result_proto;
-  }
+  // NO LOCK ACQUISITION HERE - mutex_ is already held.
 
-  auto tokenized_document_or = PrepareDocumentForIndexing(
-      schema_store_.get(), language_segmenter_.get(), std::move(document));
+  auto tokenized_document_or =
+      PrepareDocumentsForIndexing(std::move(document), current_time_ms);
   if (!tokenized_document_or.ok()) {
     TransformStatus(tokenized_document_or.status(), result_status);
+    result_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
     return result_proto;
   }
   TokenizedDocument tokenized_document(
       std::move(tokenized_document_or).ValueOrDie());
 
   auto put_result_or = document_store_->Put(
-      tokenized_document.document(), tokenized_document.num_string_tokens(),
-      put_document_stats);
+      tokenized_document.document_wrapper(), put_document_stats);
   if (!put_result_or.ok()) {
+    // TODO(b/384947619): revisit here if it is document replacement. Determine
+    //   which stage of DocumentStore::Put failed and decide whether to delete
+    //   the document and run delete propagation.
     TransformStatus(put_result_or.status(), result_status);
+    result_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
     return result_proto;
   }
   DocumentId old_document_id = put_result_or.ValueOrDie().old_document_id;
   DocumentId document_id = put_result_or.ValueOrDie().new_document_id;
-  result_proto.set_was_replacement(
-      put_result_or.ValueOrDie().was_replacement());
+  int64_t expiration_timestamp_ms =
+      put_result_or.ValueOrDie().expiration_timestamp_ms;
+  result_proto.set_was_replacement(put_result_or.ValueOrDie().was_replacement);
 
   auto data_indexing_handlers_or = CreateDataIndexingHandlers();
   if (!data_indexing_handlers_or.ok()) {
     TransformStatus(data_indexing_handlers_or.status(), result_status);
+    result_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
     return result_proto;
   }
   IndexProcessor index_processor(
@@ -1493,9 +1894,35 @@
   //   threshold)
   auto index_status = index_processor.IndexDocument(
       tokenized_document, document_id, old_document_id, put_document_stats);
-  // Getting an internal error from the index could possibly mean that the index
-  // is broken. Try to rebuild them to recover.
-  if (absl_ports::IsInternal(index_status)) {
+  if (index_status.ok()) {
+    result_proto.set_document_expiration_timestamp_ms(expiration_timestamp_ms);
+    result_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
+    if (options_.enable_delete_propagation_from() &&
+        task_scheduler_ != nullptr) {
+      // Reschedule purging expired document task if:
+      // - expiration_timestamp_ms is not INT64_MAX (i.e. the new document never
+      //   expires). Note: it is fine to schedule a task with at t = INT64_MAX
+      //   since it will never be executed, but let's avoid this edge case.
+      // - AND one of the following is true:
+      //   - existing_scheduled_time_ms < 0: there was no task scheduled, so we
+      //     need to schedule one.
+      //   - expiration_timestamp_ms < existing_scheduled_time_ms: the new
+      //     document has an expiration time that is earlier than the existing
+      //     scheduled.
+      int64_t existing_scheduled_time_ms =
+          task_scheduler_->GetScheduledTimeMs(kHandleExpiredDocumentsTaskId);
+      if (expiration_timestamp_ms < std::numeric_limits<int64_t>::max() &&
+          (existing_scheduled_time_ms < 0 ||
+           expiration_timestamp_ms < existing_scheduled_time_ms)) {
+        task_scheduler_->ScheduleAt(kHandleExpiredDocumentsTaskId,
+                                    CreateHandleExpiredDocumentsTask(),
+                                    expiration_timestamp_ms);
+      }
+    }
+  } else if (absl_ports::IsInternal(index_status)) {
+    // Getting an internal error from the index could possibly mean that the
+    // index is broken. Try to rebuild them to recover.
     ICING_LOG(ERROR) << "Got an internal error from the index. Trying to "
                         "rebuild the index!\n"
                      << index_status.error_message();
@@ -1515,29 +1942,41 @@
   if (!index_status.ok()) {
     // If we encountered a failure or cannot resolve an internal error while
     // indexing this document, then mark it as deleted.
-    int64_t current_time_ms = clock_->GetSystemTimeMilliseconds();
     libtextclassifier3::Status delete_status =
         document_store_->Delete(document_id, current_time_ms);
     if (!delete_status.ok()) {
       // This is pretty dire (and, hopefully, unlikely). We can't roll back the
-      // document that we just added. Wipeout the whole index.
+      // document that we just added. Wipeout the whole database.
       ICING_LOG(ERROR) << "Cannot delete the document that is failed to index. "
                           "Wiping out the whole Icing search engine.";
-      ResetInternal();
+      ResetLocked();
     }
+    // TODO(b/384947619): revisit here and decide whether to propagate delete to
+    //   all dependents if it is a replacement. We might not be able to rely on
+    //   normal delete propagation to remove the dependents here, since the join
+    //   index may be broken at this point.
   }
 
   TransformStatus(index_status, result_status);
+  result_proto.set_vm_binder_transaction_latency_start_time_ms(
+      clock_->GetSystemTimeMilliseconds());
   return result_proto;
 }
 
 GetResultProto IcingSearchEngine::Get(const std::string_view name_space,
                                       const std::string_view uri,
                                       const GetResultSpecProto& result_spec) {
+  absl_ports::shared_lock l(&mutex_);
+  return GetLocked(name_space, uri, result_spec);
+}
+
+// GetLocked to be called when mutex_ is already held.
+GetResultProto IcingSearchEngine::GetLocked(
+    const std::string_view name_space, const std::string_view uri,
+    const GetResultSpecProto& result_spec) {
   GetResultProto result_proto;
   StatusProto* result_status = result_proto.mutable_status();
 
-  absl_ports::shared_lock l(&mutex_);
   if (!initialized_) {
     result_status->set_code(StatusProto::FAILED_PRECONDITION);
     result_status->set_message("IcingSearchEngine has not been initialized!");
@@ -1577,6 +2016,105 @@
   return result_proto;
 }
 
+BatchGetResultProto IcingSearchEngine::BatchGet(
+    GetResultSpecProto&& get_result_spec) {
+  const std::string& name_space = get_result_spec.namespace_requested();
+  BatchGetResultProto batch_get_result_proto;
+
+  if (get_result_spec.num_total_document_bytes_to_return() <= 0) {
+    StatusProto* result_status = batch_get_result_proto.mutable_status();
+    result_status->set_code(StatusProto::INVALID_ARGUMENT);
+    result_status->set_message(
+        "num_total_document_bytes_to_return must be greater than 0.");
+    return batch_get_result_proto;
+  }
+
+  absl_ports::shared_lock l(&mutex_);  // Acquire lock once for the batch
+
+  if (!initialized_) {
+    // Handle not initialized case for all documents
+    for (const std::string& id : get_result_spec.ids()) {
+      GetResultProto* result_proto =
+          batch_get_result_proto.mutable_get_result_protos()->Add();
+      result_proto->set_uri(id);
+      result_proto->mutable_status()->set_code(
+          StatusProto::FAILED_PRECONDITION);
+    }
+    batch_get_result_proto.mutable_status()->set_message(
+        "IcingSearchEngine has not been initialized!");
+    batch_get_result_proto.mutable_status()->set_code(
+        StatusProto::FAILED_PRECONDITION);
+    return batch_get_result_proto;
+  }
+
+  int32_t total_docs_bytes_so_far = 0;
+  bool skip_remaining_docs = false;
+  int32_t num_doc_returned = 0;
+  for (std::string& id : *get_result_spec.mutable_ids()) {
+    if (skip_remaining_docs) {
+      // We simply set the status to OUT_OF_SPACE for the remaining docs, even
+      // if some docs might be small enough to fit in the result. We are doing
+      // this to avoid the worst case that all remaining docs are too big, and
+      // we read all of them out but fail to put them in the result.
+      GetResultProto result_proto;
+      // Id is redundant for a single Get call, so we only set it in BatchGet.
+      result_proto.set_uri(std::move(id));
+      StatusProto* result_status = result_proto.mutable_status();
+      result_status->set_code(StatusProto::ABORTED);
+      batch_get_result_proto.mutable_get_result_protos()->Add(
+          std::move(result_proto));
+      continue;
+    }
+
+    // We try to check each doc so smaller docs at the end of the list can
+    // still be put into the result.
+    GetResultProto result_proto =
+        GetLocked(name_space, id, get_result_spec);  // Call locked version
+    // Id is redundant for a single Get call, so we only set it in BatchGet.
+    result_proto.set_uri(std::move(id));
+    if (result_proto.status().code() == StatusProto::OK) {
+      // We get the doc successfully. Check if we can add it to the result.
+      size_t document_bytes = result_proto.document().ByteSizeLong();
+      if (num_doc_returned == 0 ||
+          document_bytes <=
+              get_result_spec.num_total_document_bytes_to_return() -
+                  total_docs_bytes_so_far) {
+        total_docs_bytes_so_far += document_bytes;
+        ++num_doc_returned;
+        // We skip the remaining docs if this 1st doc is already too big.
+        // But we always return 1st one no matter how big it is.
+        if (document_bytes >
+            get_result_spec.num_total_document_bytes_to_return()) {
+          skip_remaining_docs = true;
+        }
+      } else {
+        ICING_LOG(INFO) << "Skipping document due to byte size threshold. "
+                           "Current num docs: "
+                        << num_doc_returned
+                        << ", total byte size: " << total_docs_bytes_so_far
+                        << ", next doc byte size: " << document_bytes
+                        << ", threshold: "
+                        << get_result_spec.num_total_document_bytes_to_return();
+
+        result_proto.clear_document();
+        StatusProto* result_status = result_proto.mutable_status();
+        result_status->set_code(StatusProto::ABORTED);
+        // We skip the remaining docs even if there might be smaller docs
+        // afterwards. This is to avoid the worst case that all remaining docs
+        // are too big, and we read all of them out but fail to put them in the
+        // result.
+        skip_remaining_docs = true;
+      }
+    }
+    batch_get_result_proto.mutable_get_result_protos()->Add(
+        std::move(result_proto));
+  }
+  batch_get_result_proto.mutable_status()->set_code(
+      icing::lib::StatusProto::OK);
+
+  return batch_get_result_proto;
+}
+
 ReportUsageResultProto IcingSearchEngine::ReportUsage(
     const UsageReport& usage_report) {
   ReportUsageResultProto result_proto;
@@ -1586,12 +2124,16 @@
   if (!initialized_) {
     result_status->set_code(StatusProto::FAILED_PRECONDITION);
     result_status->set_message("IcingSearchEngine has not been initialized!");
+    result_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
     return result_proto;
   }
 
   libtextclassifier3::Status status =
       document_store_->ReportUsage(usage_report);
   TransformStatus(status, result_status);
+  result_proto.set_vm_binder_transaction_latency_start_time_ms(
+      clock_->GetSystemTimeMilliseconds());
   return result_proto;
 }
 
@@ -1656,11 +2198,14 @@
     // It is possible that the document has expired and the delete operation
     // fails with NOT_FOUND_ERROR. In this case, we should still propagate the
     // delete operation, regardless of the outcome of the delete operation.
-    libtextclassifier3::StatusOr<int> propagated_child_docs_deleted_or =
-        PropagateDelete(/*deleted_document_ids=*/{document_id},
-                        current_time_ms);
+    // TODO(b/384947619): add metadata of propagated documents to
+    //   DeleteResultProto for observer.
+    libtextclassifier3::StatusOr<std::vector<DocumentStore::DocumentMetadata>>
+        propagated_child_docs_deleted_or = PropagateDelete(
+            /*deleted_document_ids=*/{document_id}, current_time_ms);
     if (propagated_child_docs_deleted_or.ok()) {
-      num_documents_deleted += propagated_child_docs_deleted_or.ValueOrDie();
+      num_documents_deleted += static_cast<int>(
+          propagated_child_docs_deleted_or.ValueOrDie().size());
     } else {
       propagate_delete_status =
           std::move(propagated_child_docs_deleted_or).status();
@@ -1874,15 +2419,19 @@
   }
 
   // Propagate deletion.
-  libtextclassifier3::StatusOr<int> propagated_child_docs_deleted_or =
-      PropagateDelete(deleted_document_ids, current_time_ms);
+  // TODO(b/384947619): add metadata of propagated documents to
+  //   DeleteResultProto for observer.
+  libtextclassifier3::StatusOr<std::vector<DocumentStore::DocumentMetadata>>
+      propagated_child_docs_deleted_or =
+          PropagateDelete(deleted_document_ids, current_time_ms);
   if (!propagated_child_docs_deleted_or.ok()) {
     TransformStatus(propagated_child_docs_deleted_or.status(), result_status);
     delete_stats->set_document_removal_latency_ms(
         component_timer->GetElapsedMilliseconds());
     return result_proto;
   }
-  num_deleted += propagated_child_docs_deleted_or.ValueOrDie();
+  num_deleted +=
+      static_cast<int>(propagated_child_docs_deleted_or.ValueOrDie().size());
 
   delete_stats->set_document_removal_latency_ms(
       component_timer->GetElapsedMilliseconds());
@@ -1903,68 +2452,63 @@
   return result_proto;
 }
 
-libtextclassifier3::StatusOr<int> IcingSearchEngine::PropagateDelete(
+// TODO(b/384947619): remove this function once we fully ramp
+// enable_delete_propagation_from.
+libtextclassifier3::StatusOr<std::vector<DocumentStore::DocumentMetadata>>
+IcingSearchEngine::PropagateDelete(
     const std::unordered_set<DocumentId>& deleted_document_ids,
     int64_t current_time_ms) {
-  int propagated_child_docs_deleted = 0;
-
-  if (!options_.enable_delete_propagation_from()) {
-    // No-op if delete propagation is disabled.
-    return propagated_child_docs_deleted;
+  if (!options_.enable_delete_propagation_from() ||
+      deleted_document_ids.empty()) {
+    // No-op if delete propagation is disabled or no deleted document ids, so
+    // return an empty vector.
+    return std::vector<DocumentStore::DocumentMetadata>();
   }
 
-  if (qualified_id_join_index_->version() !=
-      QualifiedIdJoinIndex::Version::kV3) {
-    // This should not happen since Icing should've failed initialization with
-    // delete propagation enabled and join index v3 disabled.
-    // But let's check it here again just in case.
-    return absl_ports::FailedPreconditionError(
-        "Delete propagation is enabled but qualified id join index v3 is not "
-        "used.");
-  }
-
-  // Create join processor to get propagated child documents to delete.
-  JoinProcessor join_processor(document_store_.get(), schema_store_.get(),
-                               qualified_id_join_index_.get(), current_time_ms);
   ICING_ASSIGN_OR_RETURN(
-      std::unordered_set<DocumentId> child_docs_to_delete,
-      join_processor.GetPropagatedChildDocumentsToDelete(deleted_document_ids));
-
-  // Delete all propagated child documents.
-  for (DocumentId child_doc_id : child_docs_to_delete) {
-    auto status = document_store_->Delete(child_doc_id, current_time_ms);
-    if (!status.ok()) {
-      if (absl_ports::IsNotFound(status)) {
-        // The child document has already been deleted or expired, so skip the
-        // error.
-        continue;
-      }
-
-      // Real error.
-      return status;
-    }
-    ++propagated_child_docs_deleted;
-  }
-
-  return propagated_child_docs_deleted;
+      DeletePropagationHandler delete_propagation_handler,
+      DeletePropagationHandler::Create(schema_store_.get(),
+                                       qualified_id_join_index_.get(),
+                                       document_store_.get(), current_time_ms));
+  return delete_propagation_handler.Handle(deleted_document_ids);
 }
 
 PersistToDiskResultProto IcingSearchEngine::PersistToDisk(
     PersistType::Code persist_type) {
-  ICING_VLOG(1) << "Persisting data to disk";
+  // Measure the latency of the persist process.
+  std::unique_ptr<Timer> persist_timer = clock_->GetNewTimer();
 
   PersistToDiskResultProto result_proto;
   StatusProto* result_status = result_proto.mutable_status();
+  PersistToDiskStatsProto* persist_stats = result_proto.mutable_persist_stats();
 
   absl_ports::unique_lock l(&mutex_);
+
+  ICING_LOG(INFO) << "Persisting data to disk with mode "
+                  << static_cast<int>(persist_type);
+
   if (!initialized_) {
+    ICING_LOG(WARNING) << "Attempt to persist data to disk for an "
+                          "uninitialized IcingSearchEngine.";
     result_status->set_code(StatusProto::FAILED_PRECONDITION);
     result_status->set_message("IcingSearchEngine has not been initialized!");
+    result_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
     return result_proto;
   }
 
-  auto status = InternalPersistToDisk(persist_type);
+  auto status = PersistToDiskLocked(persist_type, persist_stats);
+  if (status.ok()) {
+    ICING_LOG(INFO) << "PersistToDisk completed.";
+  } else {
+    ICING_LOG(ERROR) << "PersistToDisk failed. Error code: "
+                     << status.error_code()
+                     << ", message: " << status.error_message();
+  }
   TransformStatus(status, result_status);
+  persist_stats->set_latency_ms(persist_timer->GetElapsedMilliseconds());
+  result_proto.set_vm_binder_transaction_latency_start_time_ms(
+      clock_->GetSystemTimeMilliseconds());
   return result_proto;
 }
 
@@ -1975,7 +2519,7 @@
 // 2. Copy data needed to a tmp directory.
 // 3. Swap current directory and tmp directory.
 OptimizeResultProto IcingSearchEngine::Optimize() {
-  ICING_VLOG(1) << "Optimizing icing storage";
+  ICING_LOG(INFO) << "Optimizing icing storage";
 
   OptimizeResultProto result_proto;
   StatusProto* result_status = result_proto.mutable_status();
@@ -1984,6 +2528,8 @@
   if (!initialized_) {
     result_status->set_code(StatusProto::FAILED_PRECONDITION);
     result_status->set_message("IcingSearchEngine has not been initialized!");
+    result_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
     return result_proto;
   }
 
@@ -1992,6 +2538,8 @@
   ScopedTimer optimize_timer(
       clock_->GetNewTimer(),
       [optimize_stats](int64_t t) { optimize_stats->set_latency_ms(t); });
+  result_proto.set_vm_binder_transaction_latency_start_time_ms(
+      clock_->GetSystemTimeMilliseconds());
 
   // Read the optimize status and assign previous_optimize_status. This is the
   // time that we last ran optimize.
@@ -2040,11 +2588,26 @@
     }
   }
 
-  // Flushes data to disk before doing optimization
-  auto status = InternalPersistToDisk(PersistType::FULL);
-  if (!status.ok()) {
-    TransformStatus(status, result_status);
-    return result_proto;
+  libtextclassifier3::Status status;
+  std::unique_ptr<Timer> persist_timer;
+  if (!feature_flags_.enable_optimize_improvements()) {
+    // Flushes data to disk before doing optimization.
+    // This really is not necessary. Therefore, if the improvements flag is
+    // enabled, just skip this step.
+    PersistToDiskStatsProto* before_optimize_persist_stats =
+        optimize_stats->mutable_before_optimize_persist_stats();
+
+    persist_timer = clock_->GetNewTimer();
+    status =
+        PersistToDiskLocked(PersistType::FULL, before_optimize_persist_stats);
+    before_optimize_persist_stats->set_latency_ms(
+        persist_timer->GetElapsedMilliseconds());
+    if (!status.ok()) {
+      TransformStatus(status, result_status);
+      result_proto.set_vm_binder_transaction_latency_start_time_ms(
+          clock_->GetSystemTimeMilliseconds());
+      return result_proto;
+    }
   }
 
   // Get all expired blob handles. This can be done before the marker file since
@@ -2072,6 +2635,8 @@
             IcingSearchEngineMarkerProto::OperationType::OPTIMIZE);
     if (!marker_file_or.ok()) {
       TransformStatus(marker_file_or.status(), result_status);
+      result_proto.set_vm_binder_transaction_latency_start_time_ms(
+          clock_->GetSystemTimeMilliseconds());
       return result_proto;
     }
     marker_file = std::move(marker_file_or).ValueOrDie();
@@ -2097,6 +2662,8 @@
     // If INTERNAL_ERROR, we're having IO errors or other errors that we can't
     // recover from.
     TransformStatus(optimize_result_or.status(), result_status);
+    result_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
     return result_proto;
   }
 
@@ -2106,7 +2673,7 @@
     // optimize blob store
     libtextclassifier3::StatusOr<std::vector<std::string>>
         blob_file_names_to_remove_or = blob_store_->Optimize(
-            optimize_result_or.ValueOrDie().dead_blob_handles);
+            optimize_result_or.ValueOrDie().dead_blob_handles, &feature_flags_);
     if (!blob_file_names_to_remove_or.ok()) {
       TransformStatus(blob_file_names_to_remove_or.status(), result_status);
       return result_proto;
@@ -2157,7 +2724,7 @@
 
     libtextclassifier3::Status qualified_id_join_index_optimize_status =
         qualified_id_join_index_->Optimize(
-            optimize_result.document_id_old_to_new,
+            document_store_.get(), optimize_result.document_id_old_to_new,
             optimize_result.namespace_id_old_to_new,
             document_store_->last_added_document_id());
     if (!qualified_id_join_index_optimize_status.ok()) {
@@ -2198,6 +2765,8 @@
       TransformStatus(status, result_status);
       optimize_stats->set_index_restoration_latency_ms(
           optimize_index_timer->GetElapsedMilliseconds());
+      result_proto.set_vm_binder_transaction_latency_start_time_ms(
+          clock_->GetSystemTimeMilliseconds());
       return result_proto;
     }
 
@@ -2215,6 +2784,8 @@
       TransformStatus(status, result_status);
       optimize_stats->set_index_restoration_latency_ms(
           optimize_index_timer->GetElapsedMilliseconds());
+      result_proto.set_vm_binder_transaction_latency_start_time_ms(
+          clock_->GetSystemTimeMilliseconds());
       return result_proto;
     }
   }
@@ -2246,9 +2817,16 @@
   }
 
   // Flushes data to disk after doing optimization
-  status = InternalPersistToDisk(PersistType::FULL);
+  PersistToDiskStatsProto* after_optimize_persist_stats =
+      optimize_stats->mutable_after_optimize_persist_stats();
+  persist_timer = clock_->GetNewTimer();
+  status = PersistToDiskLocked(PersistType::FULL, after_optimize_persist_stats);
+  after_optimize_persist_stats->set_latency_ms(
+      persist_timer->GetElapsedMilliseconds());
   if (!status.ok()) {
     TransformStatus(status, result_status);
+    result_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
     return result_proto;
   }
 
@@ -2257,6 +2835,9 @@
       Filesystem::SanitizeFileSize(after_size));
 
   TransformStatus(doc_store_optimize_result_status, result_status);
+  result_proto.set_vm_binder_transaction_latency_start_time_ms(
+      clock_->GetSystemTimeMilliseconds());
+  ICING_LOG(INFO) << "Finished optimizing icing storage";
   return result_proto;
 }
 
@@ -2270,36 +2851,52 @@
   if (!initialized_) {
     result_status->set_code(StatusProto::FAILED_PRECONDITION);
     result_status->set_message("IcingSearchEngine has not been initialized!");
+    result_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
     return result_proto;
   }
 
+  int64_t current_time_ms = clock_->GetSystemTimeMilliseconds();
+
   // Read the optimize status to get the time that we last ran.
   std::string optimize_status_filename =
       absl_ports::StrCat(options_.base_dir(), "/", kOptimizeStatusFilename);
   FileBackedProto<OptimizeStatusProto> optimize_status_file(
       *filesystem_, optimize_status_filename);
   auto optimize_status_or = optimize_status_file.Read();
-  int64_t current_time = clock_->GetSystemTimeMilliseconds();
 
-  if (optimize_status_or.ok()) {
-    // If we have trouble reading the status or this is the first time that
-    // we've ever run, don't set this field.
+  if (!optimize_status_or.ok()) {
+    // We have trouble reading the status; or we've never run optimize before.
+    result_proto.set_no_previous_optimize_info(true);
+  } else {
     int64_t time_since_last_optimize_ms;
     if (options_.calculate_time_since_last_attempted_optimize()) {
       time_since_last_optimize_ms = GetTimeSinceLastOptimizeMs(
-          current_time, *optimize_status_or.ValueOrDie());
+          current_time_ms, *optimize_status_or.ValueOrDie());
     } else {
       time_since_last_optimize_ms =
-          current_time - optimize_status_or.ValueOrDie()
-                             ->last_successful_optimize_run_time_ms();
+          current_time_ms - optimize_status_or.ValueOrDie()
+                                ->last_successful_optimize_run_time_ms();
+    }
+    if (time_since_last_optimize_ms < 0) {
+      ICING_LOG(WARNING) << "Time since last optimize is negative: "
+                         << time_since_last_optimize_ms
+                         << ". Setting no_previous_optimize_info to true.";
+      result_proto.set_no_previous_optimize_info(true);
     }
     result_proto.set_time_since_last_optimize_ms(time_since_last_optimize_ms);
   }
 
+  // Get stats from ResultStateManager.
+  result_proto.set_num_active_result_states(
+      result_state_manager_->GetNumActiveResultStates(current_time_ms));
+
   // Get stats from DocumentStore
   auto doc_store_optimize_info_or = document_store_->GetOptimizeInfo();
   if (!doc_store_optimize_info_or.ok()) {
     TransformStatus(doc_store_optimize_info_or.status(), result_status);
+    result_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
     return result_proto;
   }
   DocumentStore::OptimizeInfo doc_store_optimize_info =
@@ -2310,6 +2907,8 @@
     // Can return early since there's nothing to calculate on the index side
     result_proto.set_estimated_optimizable_bytes(0);
     result_status->set_code(StatusProto::OK);
+    result_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
     return result_proto;
   }
 
@@ -2317,6 +2916,8 @@
   auto index_elements_size_or = index_->GetElementsSize();
   if (!index_elements_size_or.ok()) {
     TransformStatus(index_elements_size_or.status(), result_status);
+    result_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
     return result_proto;
   }
   int64_t index_elements_size = index_elements_size_or.ValueOrDie();
@@ -2330,6 +2931,8 @@
       doc_store_optimize_info.estimated_optimizable_bytes);
 
   result_status->set_code(StatusProto::OK);
+  result_proto.set_vm_binder_transaction_latency_start_time_ms(
+      clock_->GetSystemTimeMilliseconds());
   return result_proto;
 }
 
@@ -2358,6 +2961,8 @@
       result.mutable_status()->set_code(StatusProto::INTERNAL);
       result.mutable_status()->set_message(
           namespace_blob_storage_infos_or.status().error_message());
+      result.set_vm_binder_transaction_latency_start_time_ms(
+          clock_->GetSystemTimeMilliseconds());
       return result;
     }
     std::vector<NamespaceBlobStorageInfoProto> namespace_blob_storage_infos =
@@ -2372,6 +2977,8 @@
   }
   // TODO(b/259744228): add stats for integer index
   result.mutable_status()->set_code(StatusProto::OK);
+  result.set_vm_binder_transaction_latency_start_time_ms(
+      clock_->GetSystemTimeMilliseconds());
   return result;
 }
 
@@ -2417,30 +3024,79 @@
   return debug_info;
 }
 
-libtextclassifier3::Status IcingSearchEngine::InternalPersistToDisk(
-    PersistType::Code persist_type) {
+libtextclassifier3::Status IcingSearchEngine::PersistToDiskLocked(
+    PersistType::Code persist_type, PersistToDiskStatsProto* persist_stats) {
+  std::unique_ptr<Timer> overall_timer = clock_->GetNewTimer();
+  persist_stats->set_persist_type(persist_type);
+
   if (blob_store_ != nullptr) {
     // For all valid PersistTypes, we persist the ground truth. The ground truth
     // in the blob_store is a proto log file, which is need to be called when
     // persist_type is LITE.
     ICING_RETURN_IF_ERROR(blob_store_->PersistToDisk());
+    persist_stats->set_blob_store_persist_latency_ms(
+        overall_timer->GetElapsedMilliseconds());
   }
-  ICING_RETURN_IF_ERROR(document_store_->PersistToDisk(persist_type));
+
+  overall_timer = clock_->GetNewTimer();
+  ICING_RETURN_IF_ERROR(
+      document_store_->PersistToDisk(persist_type, persist_stats));
+  persist_stats->set_document_store_total_persist_latency_ms(
+      overall_timer->GetElapsedMilliseconds());
+
   if (persist_type == PersistType::RECOVERY_PROOF) {
     // Persist RECOVERY_PROOF will persist the ground truth and then update all
     // checksums. There is no need to call document_store_->UpdateChecksum()
     // because PersistToDisk(RECOVERY_PROOF) will update the checksum anyways.
+    overall_timer = clock_->GetNewTimer();
     ICING_RETURN_IF_ERROR(schema_store_->UpdateChecksum());
+    persist_stats->set_schema_store_persist_latency_ms(
+        overall_timer->GetElapsedMilliseconds());
+
+    overall_timer = clock_->GetNewTimer();
     index_->UpdateChecksum();
+    persist_stats->set_index_persist_latency_ms(
+        overall_timer->GetElapsedMilliseconds());
+
+    overall_timer = clock_->GetNewTimer();
     ICING_RETURN_IF_ERROR(integer_index_->UpdateChecksums());
+    persist_stats->set_integer_index_persist_latency_ms(
+        overall_timer->GetElapsedMilliseconds());
+
+    overall_timer = clock_->GetNewTimer();
     ICING_RETURN_IF_ERROR(qualified_id_join_index_->UpdateChecksums());
+    persist_stats->set_qualified_id_join_index_persist_latency_ms(
+        overall_timer->GetElapsedMilliseconds());
+
+    overall_timer = clock_->GetNewTimer();
     ICING_RETURN_IF_ERROR(embedding_index_->UpdateChecksums());
+    persist_stats->set_embedding_index_persist_latency_ms(
+        overall_timer->GetElapsedMilliseconds());
   } else if (persist_type == PersistType::FULL) {
+    overall_timer = clock_->GetNewTimer();
     ICING_RETURN_IF_ERROR(schema_store_->PersistToDisk());
+    persist_stats->set_schema_store_persist_latency_ms(
+        overall_timer->GetElapsedMilliseconds());
+
+    overall_timer = clock_->GetNewTimer();
     ICING_RETURN_IF_ERROR(index_->PersistToDisk());
+    persist_stats->set_index_persist_latency_ms(
+        overall_timer->GetElapsedMilliseconds());
+
+    overall_timer = clock_->GetNewTimer();
     ICING_RETURN_IF_ERROR(integer_index_->PersistToDisk());
+    persist_stats->set_integer_index_persist_latency_ms(
+        overall_timer->GetElapsedMilliseconds());
+
+    overall_timer = clock_->GetNewTimer();
     ICING_RETURN_IF_ERROR(qualified_id_join_index_->PersistToDisk());
+    persist_stats->set_qualified_id_join_index_persist_latency_ms(
+        overall_timer->GetElapsedMilliseconds());
+
+    overall_timer = clock_->GetNewTimer();
     ICING_RETURN_IF_ERROR(embedding_index_->PersistToDisk());
+    persist_stats->set_embedding_index_persist_latency_ms(
+        overall_timer->GetElapsedMilliseconds());
   }
 
   return libtextclassifier3::Status::OK;
@@ -2468,12 +3124,14 @@
   int64_t lock_acquisition_latency = overall_timer->GetElapsedMilliseconds();
 
   SearchResultProto result_proto =
-      InternalSearch(search_spec, scoring_spec, result_spec);
+      SearchLocked(search_spec, scoring_spec, result_spec);
 
   result_proto.mutable_query_stats()->set_lock_acquisition_latency_ms(
       lock_acquisition_latency);
   result_proto.mutable_query_stats()->set_latency_ms(
       overall_timer->GetElapsedMilliseconds());
+  result_proto.set_vm_binder_transaction_latency_start_time_ms(
+      clock_->GetSystemTimeMilliseconds());
   return result_proto;
 }
 
@@ -2487,16 +3145,18 @@
   int64_t lock_acquisition_latency = overall_timer->GetElapsedMilliseconds();
 
   SearchResultProto result_proto =
-      InternalSearch(search_spec, scoring_spec, result_spec);
+      SearchLocked(search_spec, scoring_spec, result_spec);
 
   result_proto.mutable_query_stats()->set_lock_acquisition_latency_ms(
       lock_acquisition_latency);
   result_proto.mutable_query_stats()->set_latency_ms(
       overall_timer->GetElapsedMilliseconds());
+  result_proto.set_vm_binder_transaction_latency_start_time_ms(
+      clock_->GetSystemTimeMilliseconds());
   return result_proto;
 }
 
-SearchResultProto IcingSearchEngine::InternalSearch(
+SearchResultProto IcingSearchEngine::SearchLocked(
     const SearchSpecProto& search_spec, const ScoringSpecProto& scoring_spec,
     const ResultSpecProto& result_spec) {
   SearchResultProto result_proto;
@@ -2696,9 +3356,9 @@
   std::unique_ptr<Timer> component_timer = clock_->GetNewTimer();
   // CacheAndRetrieveFirstPage and retrieves the document protos and snippets if
   // requested
-  auto result_retriever_or =
-      ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get());
+  auto result_retriever_or = ResultRetrieverV2::Create(
+      document_store_.get(), schema_store_.get(), language_segmenter_.get(),
+      normalizer_.get(), &feature_flags_);
   if (!result_retriever_or.ok()) {
     TransformStatus(result_retriever_or.status(), result_status);
     query_stats->set_document_retrieval_latency_ms(
@@ -2712,7 +3372,7 @@
       page_result_info_or = result_state_manager_->CacheAndRetrieveFirstPage(
           std::move(ranker), std::move(parent_result_adjustment_info),
           std::move(child_result_adjustment_info), result_spec,
-          *document_store_, *result_retriever, current_time_ms);
+          *document_store_, *result_retriever, current_time_ms, query_stats);
   if (!page_result_info_or.ok()) {
     TransformStatus(page_result_info_or.status(), result_status);
     query_stats->set_document_retrieval_latency_ms(
@@ -2750,6 +3410,32 @@
   return result_proto;
 }
 
+libtextclassifier3::StatusOr<TokenizedDocument>
+IcingSearchEngine::PrepareDocumentsForIndexing(DocumentProto&& document,
+                                               int64_t current_time_ms) {
+  ICING_ASSIGN_OR_RETURN(
+      TokenizedDocument tokenized_document,
+      TokenizedDocument::Create(schema_store_.get(), language_segmenter_.get(),
+                                current_time_ms, std::move(document)));
+
+  if (!options_.enable_delete_propagation_from()) {
+    return tokenized_document;
+  }
+
+  // Make a temporary vector to hold the single tokenized document for
+  // DocumentDependencyProcessor.
+  std::vector<TokenizedDocument> tmp_tokenized_documents;
+  tmp_tokenized_documents.push_back(std::move(tokenized_document));
+
+  ICING_ASSIGN_OR_RETURN(
+      DocumentDependencyProcessor dependency_processor,
+      DocumentDependencyProcessor::Create(
+          document_store_.get(), tmp_tokenized_documents, current_time_ms));
+  ICING_RETURN_IF_ERROR(dependency_processor.Evaluate());
+
+  return std::move(tmp_tokenized_documents[0]);
+}
+
 IcingSearchEngine::QueryScoringResults IcingSearchEngine::ProcessQueryAndScore(
     const SearchSpecProto& search_spec, const ScoringSpecProto& scoring_spec,
     const ResultSpecProto& result_spec,
@@ -2842,6 +3528,14 @@
 }
 
 SearchResultProto IcingSearchEngine::GetNextPage(uint64_t next_page_token) {
+  GetNextPageRequestProto request_proto;
+  request_proto.set_next_page_token(next_page_token);
+
+  return GetNextPage(std::move(request_proto));
+}
+
+SearchResultProto IcingSearchEngine::GetNextPage(
+    GetNextPageRequestProto&& get_next_page_request) {
   SearchResultProto result_proto;
   StatusProto* result_status = result_proto.mutable_status();
 
@@ -2859,9 +3553,9 @@
     return result_proto;
   }
 
-  auto result_retriever_or =
-      ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get());
+  auto result_retriever_or = ResultRetrieverV2::Create(
+      document_store_.get(), schema_store_.get(), language_segmenter_.get(),
+      normalizer_.get(), &feature_flags_);
   if (!result_retriever_or.ok()) {
     TransformStatus(result_retriever_or.status(), result_status);
     return result_proto;
@@ -2869,14 +3563,32 @@
   std::unique_ptr<ResultRetrieverV2> result_retriever =
       std::move(result_retriever_or).ValueOrDie();
 
+  uint64_t next_page_token = get_next_page_request.next_page_token();
   int64_t current_time_ms = clock_->GetSystemTimeMilliseconds();
   libtextclassifier3::StatusOr<std::pair<uint64_t, PageResult>>
       page_result_info_or = result_state_manager_->GetNextPage(
-          next_page_token, *result_retriever, current_time_ms);
+          next_page_token,
+          get_next_page_request.max_results_to_retrieve_from_page(),
+          *result_retriever, current_time_ms);
   if (!page_result_info_or.ok()) {
     if (absl_ports::IsNotFound(page_result_info_or.status())) {
-      // NOT_FOUND means an empty result.
+      // - If calling GetNextPage with an invalid page token, return OK with an
+      //   empty result.
+      // - If the token is valid but getting NOT_FOUND error, then it is likely
+      //   that the corresponding ResultState has been removed due to cache
+      //   eviction (caused by cache budget limit exceeded or
+      //   SetSchema/Optimize). In this case, we should return an additional
+      //   (warning) field to the client indicating that the pagination is
+      //   incomplete.
       result_status->set_code(StatusProto::OK);
+
+      if (next_page_token != kInvalidNextPageToken) {
+        result_proto.set_page_token_not_found(true);
+        query_stats->set_page_token_type(
+            QueryStatsProto::PageTokenType::NOT_FOUND);
+      } else {
+        query_stats->set_page_token_type(QueryStatsProto::PageTokenType::EMPTY);
+      }
     } else {
       // Real error, pass up.
       TransformStatus(page_result_info_or.status(), result_status);
@@ -2916,6 +3628,7 @@
   query_stats->set_num_results_with_snippets(
       page_result_info.second.num_results_with_snippets);
   query_stats->set_num_joined_results_returned_current_page(child_count);
+  query_stats->set_page_token_type(QueryStatsProto::PageTokenType::VALID);
 
   return result_proto;
 }
@@ -2929,6 +3642,119 @@
   result_state_manager_->InvalidateResultState(next_page_token);
 }
 
+HandleExpiredDocumentsResultProto IcingSearchEngine::HandleExpiredDocuments() {
+  absl_ports::unique_lock l(&mutex_);
+
+  if (!initialized_) {
+    HandleExpiredDocumentsResultProto result_proto;
+    StatusProto* result_status = result_proto.mutable_status();
+    result_status->set_code(StatusProto::FAILED_PRECONDITION);
+    result_status->set_message("IcingSearchEngine has not been initialized!");
+    return result_proto;
+  }
+
+  return HandleExpiredDocumentsLocked();
+}
+
+HandleExpiredDocumentsResultProto
+IcingSearchEngine::HandleExpiredDocumentsLocked() {
+  HandleExpiredDocumentsResultProto result_proto;
+  StatusProto* result_status = result_proto.mutable_status();
+
+  // We don't need to purge expired documents before delete propagation is
+  // supported, since in search API it will automatically filter out expired
+  // documents. However, when delete propagation is enabled, we need to purge
+  // expired documents explicitly in order to propagate deletion to their
+  // children. Therefore, this API is implemented and can be released together
+  // with the delete propagation feature.
+  if (!options_.enable_delete_propagation_from()) {
+    result_status->set_code(StatusProto::FAILED_PRECONDITION);
+    result_status->set_message(
+        "Delete propagation is not enabled in this Icing instance!");
+    return result_proto;
+  }
+
+  int64_t current_time_ms = clock_->GetSystemTimeMilliseconds();
+  // Step 1: purge expired documents.
+  auto expired_docs_or =
+      document_store_->PurgeExpiredDocuments(current_time_ms);
+  if (!expired_docs_or.ok()) {
+    TransformStatus(expired_docs_or.status(), result_status);
+    return result_proto;
+  }
+  std::vector<DocumentStore::DocumentMetadata> expired_docs =
+      std::move(expired_docs_or).ValueOrDie();
+
+  // Step 2: propagate deletion to their children.
+  std::unordered_set<DocumentId> expired_doc_ids;
+  expired_doc_ids.reserve(expired_docs.size());
+  for (const DocumentStore::DocumentMetadata& metadata : expired_docs) {
+    expired_doc_ids.insert(metadata.document_id);
+  }
+  auto propagated_deleted_docs_or =
+      PropagateDelete(expired_doc_ids, current_time_ms);
+  if (!propagated_deleted_docs_or.ok()) {
+    TransformStatus(propagated_deleted_docs_or.status(), result_status);
+    return result_proto;
+  }
+  std::vector<DocumentStore::DocumentMetadata> propagated_deleted_docs =
+      std::move(propagated_deleted_docs_or).ValueOrDie();
+
+  // Get the next expiration timestamp.
+  int64_t next_expired_doc_ts_ms =
+      document_store_->GetNextExpiredDocumentTimestampMs(current_time_ms);
+  result_proto.set_next_expiration_timestamp_ms(next_expired_doc_ts_ms);
+
+  // Step 3: reschedule the task to handle next expired document(s).
+  if (task_scheduler_ != nullptr) {
+    if (next_expired_doc_ts_ms < 0) {
+      // No more expired documents. Cancel the scheduled task.
+      task_scheduler_->Cancel(kHandleExpiredDocumentsTaskId);
+    } else {
+      task_scheduler_->ScheduleAt(kHandleExpiredDocumentsTaskId,
+                                  CreateHandleExpiredDocumentsTask(),
+                                  next_expired_doc_ts_ms);
+    }
+  }
+
+  result_proto.set_num_expired_documents(
+      static_cast<int32_t>(expired_docs.size()));
+  result_proto.set_num_propagated_deleted_documents(
+      static_cast<int32_t>(propagated_deleted_docs.size()));
+
+  // Add all deleted documents to the result proto, grouped by NamespaceTypePair
+  // (namespace, schema).
+  std::unordered_map<NamespaceTypePair,
+                     HandleExpiredDocumentsResultProto::DocumentGroupInfo*,
+                     NamespaceTypePairHasher>
+      group_map;
+  const auto add_fn =
+      [&result_proto,
+       &group_map](std::vector<DocumentStore::DocumentMetadata>&& metadata_list)
+      -> void {
+    for (DocumentStore::DocumentMetadata& metadata : metadata_list) {
+      NamespaceTypePair group_key = {std::move(metadata.name_space),
+                                     std::move(metadata.schema_type_name)};
+      auto itr = group_map.find(group_key);
+      if (itr == group_map.end()) {
+        HandleExpiredDocumentsResultProto::DocumentGroupInfo* entry =
+            result_proto.add_deleted_documents();
+        entry->set_name_space(group_key.namespace_);
+        entry->set_schema(group_key.type);
+        entry->add_uris(std::move(metadata.uri));
+        group_map.insert({std::move(group_key), entry});
+      } else {
+        itr->second->add_uris(std::move(metadata.uri));
+      }
+    }
+  };
+  add_fn(std::move(expired_docs));
+  add_fn(std::move(propagated_deleted_docs));
+
+  result_status->set_code(StatusProto::OK);
+  return result_proto;
+}
+
 BlobProto IcingSearchEngine::OpenWriteBlob(
     const PropertyProto::BlobHandleProto& blob_handle) {
   BlobProto blob_proto;
@@ -2945,10 +3771,14 @@
   if (!initialized_) {
     status->set_code(StatusProto::FAILED_PRECONDITION);
     status->set_message("IcingSearchEngine has not been initialized!");
+    blob_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
     return blob_proto;
   }
-
-  return blob_store_->OpenWrite(blob_handle);
+  blob_proto = blob_store_->OpenWrite(blob_handle);
+  blob_proto.set_vm_binder_transaction_latency_start_time_ms(
+      clock_->GetSystemTimeMilliseconds());
+  return blob_proto;
 }
 
 BlobProto IcingSearchEngine::RemoveBlob(
@@ -2960,16 +3790,22 @@
   if (blob_store_ == nullptr) {
     status->set_code(StatusProto::FAILED_PRECONDITION);
     status->set_message("Remove blob is not supported in this Icing instance!");
+    blob_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
     return blob_proto;
   }
 
   if (!initialized_) {
     status->set_code(StatusProto::FAILED_PRECONDITION);
     status->set_message("IcingSearchEngine has not been initialized!");
+    blob_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
     return blob_proto;
   }
-
-  return blob_store_->RemoveBlob(blob_handle);
+  blob_proto = blob_store_->RemoveBlob(blob_handle);
+  blob_proto.set_vm_binder_transaction_latency_start_time_ms(
+      clock_->GetSystemTimeMilliseconds());
+  return blob_proto;
 }
 
 BlobProto IcingSearchEngine::OpenReadBlob(
@@ -2981,17 +3817,23 @@
     status->set_code(StatusProto::FAILED_PRECONDITION);
     status->set_message(
         "Open read blob is not supported in this Icing instance!");
+    blob_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
     return blob_proto;
   }
 
   if (!initialized_) {
     status->set_code(StatusProto::FAILED_PRECONDITION);
     status->set_message("IcingSearchEngine has not been initialized!");
+    blob_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
     ICING_LOG(ERROR) << status->message();
     return blob_proto;
   }
-
-  return blob_store_->OpenRead(blob_handle);
+  blob_proto = blob_store_->OpenRead(blob_handle);
+  blob_proto.set_vm_binder_transaction_latency_start_time_ms(
+      clock_->GetSystemTimeMilliseconds());
+  return blob_proto;
 }
 
 BlobProto IcingSearchEngine::CommitBlob(
@@ -3002,6 +3844,8 @@
   if (blob_store_ == nullptr) {
     status->set_code(StatusProto::FAILED_PRECONDITION);
     status->set_message("Commit blob is not supported in this Icing instance!");
+    blob_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
     return blob_proto;
   }
 
@@ -3009,10 +3853,69 @@
     status->set_code(StatusProto::FAILED_PRECONDITION);
     status->set_message("IcingSearchEngine has not been initialized!");
     ICING_LOG(ERROR) << status->message();
+    blob_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
+    return blob_proto;
+  }
+  blob_proto = blob_store_->CommitBlob(blob_handle);
+  blob_proto.set_vm_binder_transaction_latency_start_time_ms(
+      clock_->GetSystemTimeMilliseconds());
+  return blob_proto;
+}
+
+BlobProto IcingSearchEngine::GetAllBlobInfos() {
+  BlobProto blob_proto;
+  StatusProto* status = blob_proto.mutable_status();
+  absl_ports::unique_lock l(&mutex_);
+  if (blob_store_ == nullptr) {
+    status->set_code(StatusProto::FAILED_PRECONDITION);
+    status->set_message(
+        "Get all blob info is not supported in this Icing instance!");
+    blob_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
     return blob_proto;
   }
 
-  return blob_store_->CommitBlob(blob_handle);
+  if (!initialized_) {
+    status->set_code(StatusProto::FAILED_PRECONDITION);
+    status->set_message("IcingSearchEngine has not been initialized!");
+    ICING_LOG(ERROR) << status->message();
+    blob_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
+    return blob_proto;
+  }
+  blob_proto = blob_store_->GetAllBlobInfos();
+  blob_proto.set_vm_binder_transaction_latency_start_time_ms(
+      clock_->GetSystemTimeMilliseconds());
+  return blob_proto;
+}
+
+BlobProto IcingSearchEngine::PutBlobInfos(const BlobProto& blob_info_protos) {
+  BlobProto result_blob_proto;
+  StatusProto* status = result_blob_proto.mutable_status();
+  absl_ports::unique_lock l(&mutex_);
+  if (blob_store_ == nullptr) {
+    status->set_code(StatusProto::FAILED_PRECONDITION);
+    status->set_message(
+        "Put blob info is not supported in this Icing instance!");
+    result_blob_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
+    return result_blob_proto;
+  }
+
+  if (!initialized_) {
+    status->set_code(StatusProto::FAILED_PRECONDITION);
+    status->set_message("IcingSearchEngine has not been initialized!");
+    ICING_LOG(ERROR) << status->message();
+    result_blob_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
+    return result_blob_proto;
+  }
+
+  result_blob_proto = blob_store_->PutBlobInfos(blob_info_protos);
+  result_blob_proto.set_vm_binder_transaction_latency_start_time_ms(
+      clock_->GetSystemTimeMilliseconds());
+  return result_blob_proto;
 }
 
 libtextclassifier3::StatusOr<DocumentStore::OptimizeResult>
@@ -3049,6 +3952,8 @@
 
   // result_state_manager_ depends on document_store_. So we need to reset it at
   // the same time that we reset the document_store_.
+  ICING_LOG(INFO) << "Resetting result state manager due to optimize. All "
+                     "existing result states will be invalidated.";
   result_state_manager_.reset();
   document_store_.reset();
 
@@ -3076,7 +3981,9 @@
         schema_store_.get(), &feature_flags_,
         /*force_recovery_and_revalidate_documents=*/false,
         /*pre_mapping_fbv=*/false, /*use_persistent_hash_map=*/true,
-        options_.compression_level(), /*initialize_stats=*/nullptr);
+        options_.compression_level(), options_.compression_threshold_bytes(),
+        options_.compression_mem_level(),
+        /*initialize_stats=*/nullptr);
     // TODO(b/144458732): Implement a more robust version of
     // TC_ASSIGN_OR_RETURN that can support error logging.
     if (!create_result_or.ok()) {
@@ -3104,7 +4011,9 @@
       schema_store_.get(), &feature_flags_,
       /*force_recovery_and_revalidate_documents=*/false,
       /*pre_mapping_fbv=*/false, /*use_persistent_hash_map=*/true,
-      options_.compression_level(), /*initialize_stats=*/nullptr);
+      options_.compression_level(), options_.compression_threshold_bytes(),
+      options_.compression_mem_level(),
+      /*initialize_stats=*/nullptr);
   if (!create_result_or.ok()) {
     // Unable to create DocumentStore from the new file. Mark as uninitialized
     // and return INTERNAL.
@@ -3213,9 +4122,7 @@
 
     libtextclassifier3::Status status;
     libtextclassifier3::StatusOr<TokenizedDocument> tokenized_document_or =
-        PrepareDocumentForIndexing(schema_store_.get(),
-                                   language_segmenter_.get(),
-                                   std::move(document));
+        PrepareDocumentsForIndexing(std::move(document), current_time_ms);
     if (!tokenized_document_or.ok()) {
       status = std::move(tokenized_document_or).status();
     } else {
@@ -3252,20 +4159,39 @@
     }
   }
 
-  // Finally, delete all failed documents.
+  // Delete all failed documents and propagate deletion to their child
+  // documents. If there is any error, log it without failing index restoration.
+  // TODO(b/384947619): add metadata of deleted documents into
+  //   InitializeResultProto.
   if (options_.enable_soft_index_restoration()) {
     for (DocumentId document_id : failed_document_ids) {
-      libtextclassifier3::Status delete_status =
-          document_store_->Delete(document_id, current_time_ms);
-      if (!delete_status.ok()) {
+      libtextclassifier3::StatusOr<DocumentStore::DocumentMetadata>
+          deleted_metadata_or = document_store_->ForceDelete(document_id);
+      if (!deleted_metadata_or.ok()) {
         // This is pretty dire (and, hopefully, unlikely). Log the error and
         // skip it.
         ICING_LOG(WARNING) << "Cannot delete document " << document_id
                            << " that which failed to index: "
-                           << delete_status.error_message();
+                           << deleted_metadata_or.status().error_message();
       }
     }
-    // TODO(b/384947619): apply delete propagation on these failed documents.
+
+    // Propagate deletion to child documents. Call PropagateDelete directly here
+    // since the delete propagation flag will be checked there.
+    auto propagated_deleted_child_docs_or =
+        PropagateDelete(failed_document_ids, current_time_ms);
+    if (!propagated_deleted_child_docs_or.ok()) {
+      ICING_LOG(WARNING)
+          << "Cannot propagate deletion for child documents of failed "
+             "documents during index restoration: "
+          << propagated_deleted_child_docs_or.status().error_message();
+    } else {
+      ICING_LOG(INFO)
+          << "Successfully deleted " << failed_document_ids.size()
+          << " documents that failed to index, and propagated deletion to "
+          << propagated_deleted_child_docs_or.ValueOrDie().size()
+          << " child documents during index restoration.";
+    }
   }
 
   return IndexRestorationResult(std::move(overall_status),
@@ -3317,11 +4243,11 @@
   handlers.push_back(std::move(integer_section_indexing_handler));
 
   // Qualified id join index handler
-  ICING_ASSIGN_OR_RETURN(
-      std::unique_ptr<QualifiedIdJoinIndexingHandler>
-          qualified_id_join_indexing_handler,
-      QualifiedIdJoinIndexingHandler::Create(
-          clock_.get(), document_store_.get(), qualified_id_join_index_.get()));
+  ICING_ASSIGN_OR_RETURN(std::unique_ptr<QualifiedIdJoinIndexingHandler>
+                             qualified_id_join_indexing_handler,
+                         QualifiedIdJoinIndexingHandler::Create(
+                             clock_.get(), document_store_.get(),
+                             qualified_id_join_index_.get(), &feature_flags_));
   handlers.push_back(std::move(qualified_id_join_indexing_handler));
 
   // Embedding index handler
@@ -3533,25 +4459,62 @@
   return libtextclassifier3::Status::OK;
 }
 
-ResetResultProto IcingSearchEngine::Reset() {
-  absl_ports::unique_lock l(&mutex_);
-  return ResetInternal();
+ResetResultProto IcingSearchEngine::ClearAndDestroy() {
+  // Destroy the task scheduler first to make sure the background tasks finish
+  // before we delete the IcingSearchEngine object.
+  //
+  // Note: it is fine to not destroy the task scheduler since all the tasks will
+  //   fail silently due to initialized_ being false, but let's do it for
+  //   cleanup.
+  DestroyTaskScheduler();
+
+  {
+    absl_ports::unique_lock l(&mutex_);
+    ResetResultProto result_proto = ClearAndDestroyLocked();
+    result_proto.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
+    return result_proto;
+  }
 }
 
-ResetResultProto IcingSearchEngine::ResetInternal() {
-  ICING_VLOG(1) << "Resetting IcingSearchEngine";
+ResetResultProto IcingSearchEngine::ClearAndDestroyLocked() {
+  ICING_LOG(INFO) << "Removing Icing Search Engine directory: "
+                  << options_.base_dir() << ".";
 
   ResetResultProto result_proto;
   StatusProto* result_status = result_proto.mutable_status();
 
   initialized_ = false;
-  ResetMembers();
+  ResetMembersLocked();
   if (!filesystem_->DeleteDirectoryRecursively(options_.base_dir().c_str())) {
     result_status->set_code(StatusProto::INTERNAL);
     return result_proto;
   }
 
-  if (InternalInitialize().status().code() != StatusProto::OK) {
+  result_status->set_code(StatusProto::OK);
+  return result_proto;
+}
+
+ResetResultProto IcingSearchEngine::Reset() {
+  absl_ports::unique_lock l(&mutex_);
+  ResetResultProto result_proto = ResetLocked();
+  result_proto.set_vm_binder_transaction_latency_start_time_ms(
+      clock_->GetSystemTimeMilliseconds());
+  return result_proto;
+}
+
+ResetResultProto IcingSearchEngine::ResetLocked() {
+  ICING_VLOG(1) << "Resetting IcingSearchEngine";
+
+  ResetResultProto result_proto;
+  StatusProto* result_status = result_proto.mutable_status();
+
+  if (ClearAndDestroyLocked().status().code() != StatusProto::OK) {
+    result_status->set_code(StatusProto::INTERNAL);
+    return result_proto;
+  }
+
+  if (InitializeLocked().status().code() != StatusProto::OK) {
     // We shouldn't hit the following Initialize errors:
     //   NOT_FOUND: all data was cleared, we aren't expecting anything
     //   DATA_LOSS: all data was cleared, we aren't expecting anything
@@ -3580,6 +4543,8 @@
   if (!initialized_) {
     response_status->set_code(StatusProto::FAILED_PRECONDITION);
     response_status->set_message("IcingSearchEngine has not been initialized!");
+    response.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
     return response;
   }
 
@@ -3587,6 +4552,8 @@
       ValidateSuggestionSpec(suggestion_spec, performance_configuration_);
   if (!status.ok()) {
     TransformStatus(status, response_status);
+    response.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
     return response;
   }
 
@@ -3597,6 +4564,8 @@
       schema_store_.get(), clock_.get(), &feature_flags_);
   if (!suggestion_processor_or.ok()) {
     TransformStatus(suggestion_processor_or.status(), response_status);
+    response.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
     return response;
   }
   std::unique_ptr<SuggestionProcessor> suggestion_processor =
@@ -3608,6 +4577,8 @@
       suggestion_processor->QuerySuggestions(suggestion_spec, current_time_ms);
   if (!terms_or.ok()) {
     TransformStatus(terms_or.status(), response_status);
+    response.set_vm_binder_transaction_latency_start_time_ms(
+        clock_->GetSystemTimeMilliseconds());
     return response;
   }
 
@@ -3618,8 +4589,43 @@
     response.mutable_suggestions()->Add(std::move(suggestion));
   }
   response_status->set_code(StatusProto::OK);
+  response.set_vm_binder_transaction_latency_start_time_ms(
+      clock_->GetSystemTimeMilliseconds());
   return response;
 }
 
+void IcingSearchEngine::DestroyTaskScheduler() {
+  // We have to acquire the lock before accessing and resetting task_scheduler_.
+  // Otherwise, if the background task is still running in the critical section
+  // and accessing task_scheduler_, then it will cause data race issues on the
+  // pointer.
+  SimpleTaskScheduler* task_scheduler_local = nullptr;
+  {
+    absl_ports::unique_lock l(&mutex_);
+
+    if (task_scheduler_ != nullptr) {
+      task_scheduler_local = task_scheduler_.release();
+    }
+  }
+
+  // Delete the task scheduler OUTSIDE of the critical section. Otherwise, it is
+  // possible that a background task starts to run and blocks at acquiring the
+  // global lock, which causes deadlock.
+  if (task_scheduler_local != nullptr) {
+    delete task_scheduler_local;
+  }
+}
+
+std::function<void()> IcingSearchEngine::CreateHandleExpiredDocumentsTask() {
+  return [this]() -> void {
+    HandleExpiredDocumentsResultProto result = HandleExpiredDocuments();
+    if (result.status().code() != StatusProto::OK) {
+      ICING_LOG(ERROR)
+          << "Failed to handle expired documents in the background: "
+          << result.status().message();
+    }
+  };
+}
+
 }  // namespace lib
 }  // namespace icing

diff --git a/icing/icing-search-engine.h b/icing/icing-search-engine.h
index 1fc8e23..f1bb1af 100644
--- a/icing/icing-search-engine.h
+++ b/icing/icing-search-engine.h

@@ -16,6 +16,7 @@
 #define ICING_ICING_SEARCH_ENGINE_H_
 
 #include <cstdint>
+#include <functional>
 #include <memory>
 #include <string>
 #include <string_view>
@@ -63,7 +64,8 @@
 #include "icing/tokenization/language-segmenter.h"
 #include "icing/transform/normalizer.h"
 #include "icing/util/clock.h"
-#include "icing/util/icu-data-file-helper.h"
+#include "icing/util/simple-task-scheduler.h"
+#include "icing/util/tokenized-document.h"
 
 namespace icing {
 namespace lib {
@@ -321,6 +323,11 @@
   GetResultProto Get(std::string_view name_space, std::string_view uri,
                      const GetResultSpecProto& result_spec);
 
+  // Finds and returns the documents identified by the given GetResultSpecProto.
+  // Returns:
+  //   A BatchGetResultProto with a list of GetResultProto.
+  BatchGetResultProto BatchGet(GetResultSpecProto&& get_result_spec);
+
   // Reports usage. The corresponding usage scores of the specified document in
   // the report will be updated.
   //
@@ -440,6 +447,11 @@
   //   ABORTED if failed to get results but existing data is not affected
   //   FAILED_PRECONDITION IcingSearchEngine has not been initialized yet
   //   INTERNAL_ERROR on any other errors
+  SearchResultProto GetNextPage(GetNextPageRequestProto&& get_next_page_request)
+      ICING_LOCKS_EXCLUDED(mutex_);
+
+  // TODO: b/417644758 - Remove this method once all old callers are migrated to
+  // the new GetNextPage API. Internally, this should just be used in tests.
   SearchResultProto GetNextPage(uint64_t next_page_token)
       ICING_LOCKS_EXCLUDED(mutex_);
 
@@ -448,6 +460,28 @@
   void InvalidateNextPageToken(uint64_t next_page_token)
       ICING_LOCKS_EXCLUDED(mutex_);
 
+  // Handles expired documents.
+  // - Purges expired documents from the document store. Note: it also purges
+  //   documents that expire before
+  //   current_time_ms + options_.expired_document_purge_threshold_ms().
+  // - Propagates the deletion to their child documents.
+  // - Reschedules another (background) task to handle the next expired
+  //   document(s).
+  //
+  // This method is called by:
+  // - (Icing lib direct client only) Icing's internal task scheduler.
+  // - AppSearch task scheduler.
+  // - In the last step of Initialize(), to purge documents that expire during
+  //   Icing was off. Note: Initialize() calls the locked version of this
+  //   method.
+  //
+  // Returns a HandleExpiredDocumentsResultProto with status:
+  // - OK with results on success
+  // - FAILED_PRECONDITION_ERROR IcingSearchEngine has not been initialized yet
+  // - INTERNAL_ERROR on any other errors
+  HandleExpiredDocumentsResultProto HandleExpiredDocuments()
+      ICING_LOCKS_EXCLUDED(mutex_);
+
   // Gets or creates a file for write only purpose for the given blob handle.
   // To mark the blob is completed written, CommitBlob must be called. Once
   // CommitBlob is called, the blob is sealed and rewrite is not allowed.
@@ -524,6 +558,22 @@
   BlobProto CommitBlob(const PropertyProto::BlobHandleProto& blob_handle)
       ICING_LOCKS_EXCLUDED(mutex_);
 
+  // Gets all the blob info from the blob store.
+  //
+  // Returns:
+  //   BlobProto with all the blob info on success
+  //   InternalError on IO error
+  BlobProto GetAllBlobInfos() ICING_LOCKS_EXCLUDED(mutex_);
+
+  // Puts the blob info protos from the blob proto to the blob info proto log
+  // file.
+  //
+  // Returns:
+  //   BlobProto with all the blob info on success
+  //   InternalError on IO error
+  BlobProto PutBlobInfos(const BlobProto& blob_info_protos)
+      ICING_LOCKS_EXCLUDED(mutex_);
+
   // Makes sure that every update/delete received till this point is flushed
   // to disk. If the app crashes after a call to PersistToDisk(), Icing
   // would be able to fully recover all data written up to this point.
@@ -605,16 +655,24 @@
   //   INTERNAL_ERROR if internal state is no longer consistent
   ResetResultProto Reset() ICING_LOCKS_EXCLUDED(mutex_);
 
+  // Clears all data from Icing. Clients DO need to call Initialize again.
+  //
+  // Returns:
+  //   OK on success
+  //   INTERNAL_ERROR if failed to delete underlying files
+  ResetResultProto ClearAndDestroy() ICING_LOCKS_EXCLUDED(mutex_);
+
   // Disallow copy and move.
   IcingSearchEngine(const IcingSearchEngine&) = delete;
   IcingSearchEngine& operator=(const IcingSearchEngine&) = delete;
 
  protected:
-  IcingSearchEngine(IcingSearchEngineOptions options,
-                    std::unique_ptr<const Filesystem> filesystem,
-                    std::unique_ptr<const IcingFilesystem> icing_filesystem,
-                    std::unique_ptr<Clock> clock,
-                    std::unique_ptr<const JniCache> jni_cache = nullptr);
+  explicit IcingSearchEngine(
+      IcingSearchEngineOptions options,
+      std::unique_ptr<const Filesystem> filesystem,
+      std::unique_ptr<const IcingFilesystem> icing_filesystem,
+      std::unique_ptr<Clock> clock,
+      std::unique_ptr<const JniCache> jni_cache = nullptr);
 
  private:
   const IcingSearchEngineOptions options_;
@@ -673,15 +731,52 @@
   // Storage for all hits of embedding contents from the document store.
   std::unique_ptr<EmbeddingIndex> embedding_index_ ICING_GUARDED_BY(mutex_);
 
+  // Simple task scheduler for background tasks. Currently, it is used for:
+  // - Handle (purge) expired documents.
+  //
+  // Note: after acquiring the global lock, checking task_scheduler_ is nullptr
+  //   or not before accessing it **is ESSENTIAL**.
+  // - task_scheduler_ may be always nullptr if
+  //   options_.enable_background_task_scheduler() is false.
+  // - Race condition during destruction. Consider an example:
+  //   - Task scheduler thread: wakes up and ready to run
+  //     HandleExpiredDocuments()
+  //   - Main thread (which holds IcingSearchEngine object): destructs
+  //     IcingSearchEngine instance. IcingSearchEngine destructor destructs
+  //     task_scheduler_ via DestroyTaskScheduler() and wait to join the task
+  //     scheduler thread. At this point, task_scheduler_ is nullptr.
+  //   - Task scheduler thread: start to run HandleExpiredDocuments() and
+  //     reaches the point to reschedule the task.
+  //   - If task_scheduler_ is not checked, then the code will crash at runtime
+  //     because task_scheduler_ is already set to nullptr.
+  std::unique_ptr<SimpleTaskScheduler> task_scheduler_ ICING_GUARDED_BY(mutex_);
+
   // Pointer to JNI class references
   const std::unique_ptr<const JniCache> jni_cache_;
 
   // Resets all members that are created during Initialize.
-  void ResetMembers() ICING_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
+  //
+  // Note: this method DOES NOT reset task_scheduler_. task_scheduler_ should be
+  //   reset only if we want to destroy IcingSearchEngine, and in this case,
+  //   DestroyTaskScheduler should be called.
+  void ResetMembersLocked() ICING_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
 
   // Resets all members that are created during Initialize, deletes all
   // underlying files and initializes a fresh index.
-  ResetResultProto ResetInternal() ICING_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
+  ResetResultProto ResetLocked() ICING_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
+
+  // Resets all members and clears all data from Icing.
+  ResetResultProto ClearAndDestroyLocked()
+      ICING_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
+
+  // Assumes mutex_ is already held.
+  PutResultProto PutLocked(DocumentProto&& document, int64_t current_time_ms)
+      ICING_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
+
+  // Assumes mutex_ is already held.
+  GetResultProto GetLocked(std::string_view name_space, std::string_view uri,
+    const GetResultSpecProto& result_spec)
+      ICING_SHARED_LOCKS_REQUIRED(mutex_);
 
   // Checks for the existence of the init marker file. If the failed init count
   // exceeds kMaxUnsuccessfulInitAttempts, all data is deleted and the index is
@@ -699,13 +794,22 @@
   // separate method so that other public methods don't need to call
   // PersistToDisk(). Public methods calling each other may cause deadlock
   // issues.
-  libtextclassifier3::Status InternalPersistToDisk(
-      PersistType::Code persist_type) ICING_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
+  //
+  // @param persist_type: The type of persistence guarantee that PersistToDisk
+  // should provide.
+  // @param persist_stats: The NON-null stats about the PersistToDisk call.
+  libtextclassifier3::Status PersistToDiskLocked(
+      PersistType::Code persist_type, PersistToDiskStatsProto* persist_stats)
+      ICING_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
 
   // Helper method to the actual work to Initialize. We need this separate
   // method so that other public methods don't need to call Initialize(). Public
   // methods calling each other may cause deadlock issues.
-  InitializeResultProto InternalInitialize()
+  InitializeResultProto InitializeLocked()
+      ICING_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
+
+  // Helper method to the actual work to handle expired documents.
+  HandleExpiredDocumentsResultProto HandleExpiredDocumentsLocked()
       ICING_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
 
   // Helper method to initialize member variables.
@@ -753,8 +857,8 @@
   //   OK on success
   //   FAILED_PRECONDITION if initialize_stats is null
   libtextclassifier3::Status InitializeBlobStore(
-      int32_t orphan_blob_time_to_live_ms, int32_t compression_level)
-      ICING_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
+      int32_t orphan_blob_time_to_live_ms, int32_t compression_level,
+      int32_t compression_mem_level) ICING_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
 
   // Do any initialization/recovery necessary to create term index, integer
   // index, and qualified id join index instances.
@@ -792,11 +896,24 @@
 
   // Helper method for the actual work to Search. We need this separate
   // method to manage locking for Search.
-  SearchResultProto InternalSearch(const SearchSpecProto& search_spec,
+  SearchResultProto SearchLocked(const SearchSpecProto& search_spec,
                                    const ScoringSpecProto& scoring_spec,
                                    const ResultSpecProto& result_spec)
       ICING_SHARED_LOCKS_REQUIRED(mutex_);
 
+  // Helper method to prepare a document for indexing. This includes
+  // tokenization and dependency evaluation.
+  //
+  // TODO(b/397769319): change the input and output to a vector of documents to
+  //   support batch indexing.
+  //
+  // Returns:
+  //   On success, TokenizedDocument.
+  //   Any errors from tokenization or dependency evaluation.
+  libtextclassifier3::StatusOr<TokenizedDocument> PrepareDocumentsForIndexing(
+      DocumentProto&& document, int64_t current_time_ms)
+      ICING_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
+
   // Processes query and scores according to the specs. It is a helper function
   // (called by Search) to process and score normal query and the nested child
   // query for join search.
@@ -836,11 +953,12 @@
   // joinable properties with delete propagation enabled.
   //
   // Returns:
-  //   Number of propagated documents deleted on success
+  //   A list of metadata of the propagated documents deleted on success
   //   INTERNAL_ERROR on any I/O errors
-  libtextclassifier3::StatusOr<int> PropagateDelete(
-      const std::unordered_set<DocumentId>& deleted_document_ids,
-      int64_t current_time_ms) ICING_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
+  libtextclassifier3::StatusOr<std::vector<DocumentStore::DocumentMetadata>>
+  PropagateDelete(const std::unordered_set<DocumentId>& deleted_document_ids,
+                  int64_t current_time_ms)
+      ICING_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
 
   // Discards derived data that requires rebuild based on rebuild_info.
   //
@@ -902,6 +1020,12 @@
       std::vector<std::unique_ptr<DataIndexingHandler>>>
   CreateDataIndexingHandlers() ICING_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
 
+  // Destroys the task scheduler.
+  //
+  // Task scheduler should be destroyed ONLY IF we want to destroy
+  // IcingSearchEngine (e.g. destructor, ClearAndDestroy API, etc).
+  void DestroyTaskScheduler() ICING_LOCKS_EXCLUDED(mutex_);
+
   // Helper method to discard parts of (term, integer, qualified id join,
   // embedding) indices if they contain data for document ids greater than
   // last_stored_document_id.
@@ -1030,6 +1154,11 @@
   //   INTERNAL_ERROR on any I/O errors
   libtextclassifier3::Status ClearAllIndices()
       ICING_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
+
+  // Helper method to create a lambda function for handling expired documents
+  // task. This is used to schedule the task with the task scheduler.
+  std::function<void()> CreateHandleExpiredDocumentsTask()
+      ICING_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
 };
 
 }  // namespace lib

diff --git a/icing/icing-search-engine_benchmark.cc b/icing/icing-search-engine_benchmark.cc
index 3061300..82fc9a6 100644
--- a/icing/icing-search-engine_benchmark.cc
+++ b/icing/icing-search-engine_benchmark.cc

@@ -14,18 +14,19 @@
 
 #include <unistd.h>
 
-#include <fstream>
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
 #include <iostream>
 #include <limits>
 #include <memory>
 #include <numeric>
 #include <ostream>
 #include <random>
-#include <sstream>
-#include <stdexcept>
 #include <string>
 #include <string_view>
 #include <unordered_set>
+#include <utility>
 #include <vector>
 
 #include "testing/base/public/benchmark.h"
@@ -45,7 +46,9 @@
 #include "icing/proto/status.pb.h"
 #include "icing/proto/term.pb.h"
 #include "icing/query/query-features.h"
+#include "icing/result/result-state-manager.h"
 #include "icing/schema-builder.h"
+#include "icing/store/document-id.h"
 #include "icing/testing/common-matchers.h"
 #include "icing/testing/document-generator.h"
 #include "icing/testing/numeric/number-generator.h"
@@ -150,6 +153,7 @@
   explicit DestructibleDirectory(const Filesystem& filesystem,
                                  const std::string& dir)
       : filesystem_(filesystem), dir_(dir) {
+    filesystem_.DeleteDirectoryRecursively(dir_.c_str());
     filesystem_.CreateDirectoryRecursively(dir_.c_str());
   }
   ~DestructibleDirectory() {
@@ -204,6 +208,18 @@
       /*range_upper=*/static_cast<int64_t>(num_documents) * 10 - 1);
 }
 
+PropertyProto::VectorProto GenerateRandomVector(
+    std::default_random_engine& random, const std::string& model_signature,
+    int dimension) {
+  std::uniform_real_distribution<float> dis(-1.0, 1.0);
+  PropertyProto::VectorProto vector;
+  vector.set_model_signature(model_signature);
+  for (int i = 0; i < dimension; ++i) {
+    vector.add_values(dis(random));
+  }
+  return vector;
+}
+
 void BM_IndexLatency(benchmark::State& state) {
   // Initialize the filesystem
   std::string test_dir = GetTestTempDir() + "/icing/benchmark";
@@ -222,7 +238,6 @@
   // Create the index.
   IcingSearchEngineOptions options;
   options.set_base_dir(test_dir);
-  options.set_index_merge_size(kIcingFullIndexSize);
   std::unique_ptr<IcingSearchEngine> icing =
       std::make_unique<IcingSearchEngine>(options);
 
@@ -291,6 +306,135 @@
     // Arguments: num_indexed_documents, num_sections
     ->ArgPair(1000000, 2);
 
+void BM_EmbeddingQueryLatency(benchmark::State& state) {
+  // Initialize the filesystem
+  std::string test_dir = GetTestTempDir() + "/icing/benchmark";
+  Filesystem filesystem;
+  DestructibleDirectory ddir(filesystem, test_dir);
+
+  std::default_random_engine random;
+
+  // Create the schema.
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Type1").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("embedding")
+                  .SetDataTypeVector(EMBEDDING_INDEXING_LINEAR_SEARCH)
+                  .SetCardinality(CARDINALITY_REPEATED)))
+          .AddType(SchemaTypeConfigBuilder().SetType("Type2").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("embedding")
+                  .SetDataTypeVector(EMBEDDING_INDEXING_LINEAR_SEARCH)
+                  .SetCardinality(CARDINALITY_REPEATED)))
+          .Build();
+
+  // Create the index.
+  IcingSearchEngineOptions options;
+  options.set_base_dir(test_dir);
+  options.set_index_merge_size(kIcingFullIndexSize);
+  options.set_enable_passing_filter_to_children(true);
+  std::unique_ptr<IcingSearchEngine> icing =
+      std::make_unique<IcingSearchEngine>(options);
+
+  ASSERT_THAT(icing->Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing->SetSchema(schema).status(), ProtoIsOk());
+
+  int dimension = state.range(0);
+  int num_vectors_per_doc = state.range(1);
+  int num_type1_docs = state.range(2);
+  int num_type2_docs = state.range(3);
+  int filter_type = state.range(4);
+
+  // Type 1
+  for (int i = 0; i < num_type1_docs; ++i) {
+    DocumentProto doc = DocumentBuilder()
+                            .SetKey("icing", "type1" + std::to_string(i))
+                            .SetSchema("Type1")
+                            .SetCreationTimestampMs(1)
+                            .Build();
+    auto embedding_property = doc.add_properties();
+    embedding_property->set_name("embedding");
+    for (int j = 0; j < num_vectors_per_doc; ++j) {
+      embedding_property->mutable_vector_values()->Add(
+          GenerateRandomVector(random, /*model_signature=*/"model", dimension));
+    }
+    ASSERT_THAT(icing->Put(doc).status(), ProtoIsOk());
+  }
+
+  // Type 2
+  for (int i = 0; i < num_type2_docs; ++i) {
+    DocumentProto doc = DocumentBuilder()
+                            .SetKey("icing", "type2" + std::to_string(i))
+                            .SetSchema("Type2")
+                            .SetCreationTimestampMs(1)
+                            .Build();
+    auto embedding_property = doc.add_properties();
+    embedding_property->set_name("embedding");
+    for (int j = 0; j < num_vectors_per_doc; ++j) {
+      embedding_property->mutable_vector_values()->Add(
+          GenerateRandomVector(random, /*model_signature=*/"model", dimension));
+    }
+    ASSERT_THAT(icing->Put(doc).status(), ProtoIsOk());
+  }
+
+  SearchSpecProto search_spec;
+  search_spec.set_term_match_type(TermMatchType::EXACT_ONLY);
+  search_spec.set_embedding_query_metric_type(
+      SearchSpecProto::EmbeddingQueryMetricType::DOT_PRODUCT);
+  search_spec.add_enabled_features(
+      std::string(kListFilterQueryLanguageFeature));
+  *search_spec.add_embedding_query_vectors() =
+      GenerateRandomVector(random, /*model_signature=*/"model", dimension);
+  search_spec.set_query("semanticSearch(getEmbeddingParameter(0))");
+
+  int expected_num_docs_scored = 0;
+  if (filter_type == 1) {
+    search_spec.add_schema_type_filters("Type1");
+    expected_num_docs_scored = num_type1_docs;
+  } else if (filter_type == 2) {
+    search_spec.add_schema_type_filters("Type2");
+    expected_num_docs_scored = num_type2_docs;
+  } else {
+    expected_num_docs_scored = num_type1_docs + num_type2_docs;
+  }
+
+  ScoringSpecProto scoring_spec = CreateScoringSpec(
+      ScoringSpecProto::RankingStrategy::ADVANCED_SCORING_EXPRESSION);
+  scoring_spec.set_advanced_scoring_expression(
+      "len(this.matchedSemanticScores(getEmbeddingParameter(0)))");
+  ResultSpecProto result_spec = ResultSpecProto::default_instance();
+  result_spec.set_num_to_score(10000000);
+
+  std::string scoring_latency_ms_str;
+  for (auto _ : state) {
+    SearchResultProto results =
+        icing->Search(search_spec, scoring_spec, result_spec);
+
+    ASSERT_THAT(results.status(), ProtoIsOk());
+    for (const SearchResultProto::ResultProto& result : results.results()) {
+      ASSERT_THAT(result.score(), num_vectors_per_doc);
+    }
+    ASSERT_THAT(
+        results.query_stats().parent_search_stats().num_documents_scored(),
+        Eq(expected_num_docs_scored));
+    int scoring_latency_ms =
+        results.query_stats().parent_search_stats().scoring_latency_ms();
+    scoring_latency_ms_str += std::to_string(scoring_latency_ms) + ", ";
+  }
+  state.SetLabel("Scoring latency ms: " + scoring_latency_ms_str);
+}
+BENCHMARK(BM_EmbeddingQueryLatency)
+    // Arguments:
+    // - dimension
+    // - num_vectors_per_doc
+    // - num_type1_docs,
+    // - num_type2_docs
+    // - filter_type (1: filter by type1, 2: filter by type2, Other: no filter)
+    ->Args({768, 2, 4000, 25000, 0})
+    ->Args({768, 2, 4000, 25000, 1})
+    ->Args({768, 2, 4000, 25000, 2});
+
 void BM_IndexThroughput(benchmark::State& state) {
   // Initialize the filesystem
   std::string test_dir = GetTestTempDir() + "/icing/benchmark";
@@ -1205,6 +1349,83 @@
 }
 BENCHMARK(BM_JoinQueryQualifiedId);
 
+void BM_Optimize(benchmark::State& state) {
+  int num_docs = state.range(0);
+  int num_deleted = state.range(1);
+
+  // Initialize the filesystem
+  std::string test_dir = GetTestTempDir() + "/icing/benchmark";
+  Filesystem filesystem;
+
+  // Create the schema.
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Message").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("body")
+                  .SetDataTypeString(TermMatchType::PREFIX,
+                                     StringIndexingConfig::TokenizerType::PLAIN)
+                  .SetCardinality(PropertyConfigProto::Cardinality::OPTIONAL)))
+          .Build();
+
+  // Create documents.
+  DocumentProto base_document = DocumentBuilder()
+                                    .SetSchema("Message")
+                                    .SetNamespace("namespace")
+                                    .AddStringProperty("body", "foo")
+                                    .Build();
+  std::vector<std::string> uris;
+  uris.reserve(num_docs);
+  std::vector<DocumentProto> documents;
+  documents.reserve(num_docs);
+  for (int i = 0; i < num_docs; ++i) {
+    std::string uri = std::to_string(i);
+    uris.push_back(uri);
+    documents.push_back(DocumentBuilder(base_document).SetUri(uri).Build());
+  }
+
+  std::vector<std::string> uris_to_delete = uris;
+  std::shuffle(uris_to_delete.begin(), uris_to_delete.end(),
+               std::default_random_engine(12345));
+  uris_to_delete.resize(num_deleted);
+
+  for (auto s : state) {
+    state.PauseTiming();
+    DestructibleDirectory ddir(filesystem, test_dir);
+    // Create the index.
+    IcingSearchEngineOptions options;
+    options.set_enable_optimize_improvements(true);
+    options.set_base_dir(test_dir);
+    options.set_index_merge_size(kIcingFullIndexSize);
+    std::unique_ptr<IcingSearchEngine> icing =
+        std::make_unique<IcingSearchEngine>(options);
+
+    ASSERT_THAT(icing->Initialize().status(), ProtoIsOk());
+    ASSERT_THAT(icing->SetSchema(schema).status(), ProtoIsOk());
+
+    for (const auto& doc : documents) {
+      ASSERT_THAT(icing->Put(doc).status(), ProtoIsOk());
+    }
+
+    for (int i = 0; i < num_deleted; ++i) {
+      ASSERT_THAT(icing->Delete("namespace", uris_to_delete[i]).status(),
+                  ProtoIsOk());
+    }
+    state.ResumeTiming();
+
+    ASSERT_THAT(icing->Optimize().status(), ProtoIsOk());
+  }
+}
+
+BENCHMARK(BM_Optimize)
+    // Arguments: num_documents, num_deleted
+    ->ArgPair(1024, 256)
+    ->ArgPair(1024, 512)
+    ->ArgPair(1024, 1024)
+    ->ArgPair(8192, 2048)
+    ->ArgPair(8192, 4096)
+    ->ArgPair(8192, 8192);
+
 void BM_PersistToDisk(benchmark::State& state) {
   // Initialize the filesystem
   std::string test_dir = GetTestTempDir() + "/icing/benchmark";

diff --git a/icing/icing-search-engine_blob_test.cc b/icing/icing-search-engine_blob_test.cc
index b13d258..655c5c6 100644
--- a/icing/icing-search-engine_blob_test.cc
+++ b/icing/icing-search-engine_blob_test.cc

@@ -47,6 +47,7 @@
 namespace {
 
 using ::icing::lib::portable_equals_proto::EqualsProto;
+using ::testing::Eq;
 using ::testing::IsEmpty;
 using ::testing::SizeIs;
 using ::testing::UnorderedElementsAre;
@@ -234,13 +235,96 @@
     uint64_t size = filesystem()->GetFileSize(*read_fd);
     std::unique_ptr<unsigned char[]> buf =
         std::make_unique<unsigned char[]>(size);
-    EXPECT_TRUE(filesystem()->Read(read_fd.get(), buf.get(), size));
+    EXPECT_THAT(filesystem()->Read(read_fd.get(), buf.get(), size), Eq(size));
     std::string expected_data = std::string(data.begin(), data.end());
     std::string actual_data = std::string(buf.get(), buf.get() + size);
     EXPECT_EQ(expected_data, actual_data);
   }
 }
 
+TEST_P(IcingSearchEngineBlobTest, GetAllBlobInfo) {
+  IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+
+  PropertyProto::BlobHandleProto blob_handle;
+  std::vector<unsigned char> data = GenerateRandomBytes(24);
+  std::array<uint8_t, 32> digest = CalculateDigest(data);
+  blob_handle.set_digest((void*)digest.data(), digest.size());
+  blob_handle.set_namespace_("namespaceA");
+
+  BlobProto write_blob_proto = icing.OpenWriteBlob(blob_handle);
+  ASSERT_THAT(write_blob_proto.status(), ProtoIsOk());
+  {
+    ScopedFd write_fd(GetScopedFdFromBlobProto(write_blob_proto));
+    ASSERT_TRUE(filesystem()->Write(write_fd.get(), data.data(), data.size()));
+  }
+
+  BlobProto commit_blob_proto = icing.CommitBlob(blob_handle);
+  ASSERT_THAT(commit_blob_proto.status(), ProtoIsOk());
+
+  BlobProto blob_proto = icing.GetAllBlobInfos();
+  EXPECT_THAT(blob_proto.status(), ProtoIsOk());
+  EXPECT_THAT(blob_proto.blob_info_protos_size(), Eq(1));
+  EXPECT_THAT(blob_proto.blob_info_protos(0).blob_handle().namespace_(),
+              Eq("namespaceA"));
+}
+
+TEST_P(IcingSearchEngineBlobTest, PutBlobInfos) {
+  bool manage_blob_files = GetParam();
+
+  IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+
+  PropertyProto::BlobHandleProto blob_handle;
+  std::vector<unsigned char> data = GenerateRandomBytes(24);
+  std::array<uint8_t, 32> digest = CalculateDigest(data);
+  blob_handle.set_digest((void*)digest.data(), digest.size());
+  blob_handle.set_namespace_("namespaceA");
+
+  BlobProto write_blob_proto = icing.OpenWriteBlob(blob_handle);
+  ASSERT_THAT(write_blob_proto.status(), ProtoIsOk());
+  {
+    ScopedFd write_fd(GetScopedFdFromBlobProto(write_blob_proto));
+    ASSERT_TRUE(filesystem()->Write(write_fd.get(), data.data(), data.size()));
+  }
+
+  BlobProto commit_blob_proto = icing.CommitBlob(blob_handle);
+  ASSERT_THAT(commit_blob_proto.status(), ProtoIsOk());
+
+  BlobProto blob_proto = icing.GetAllBlobInfos();
+  EXPECT_THAT(blob_proto.status(), ProtoIsOk());
+  EXPECT_THAT(blob_proto.blob_info_protos_size(), Eq(1));
+  EXPECT_THAT(blob_proto.blob_info_protos(0).blob_handle().namespace_(),
+              Eq("namespaceA"));
+
+  // Delete the blob info proto log file.
+  filesystem()->DeleteFile(
+      absl_ports::StrCat(GetTestBlobDir(), "/blob_info_proto_file").c_str());
+
+  if (manage_blob_files) {
+    ASSERT_THAT(icing.PutBlobInfos(std::move(blob_proto)).status(),
+                ProtoStatusIs(StatusProto::FAILED_PRECONDITION));
+  } else {
+    // Put the blob info proto log file.
+    BlobProto result_blob_proto = icing.PutBlobInfos(std::move(blob_proto));
+    EXPECT_THAT(result_blob_proto.status(), ProtoIsOk());
+
+    BlobProto read_blob_proto = icing.OpenReadBlob(blob_handle);
+    ASSERT_THAT(read_blob_proto.status(), ProtoIsOk());
+    {
+      ScopedFd read_fd(GetScopedFdFromBlobProto(read_blob_proto));
+
+      uint64_t size = filesystem()->GetFileSize(*read_fd);
+      std::unique_ptr<unsigned char[]> buf =
+          std::make_unique<unsigned char[]>(size);
+      EXPECT_THAT(filesystem()->Read(read_fd.get(), buf.get(), size), Eq(size));
+      std::string expected_data = std::string(data.begin(), data.end());
+      std::string actual_data = std::string(buf.get(), buf.get() + size);
+      EXPECT_EQ(expected_data, actual_data);
+    }
+  }
+}
+
 TEST_P(IcingSearchEngineBlobTest, RemovePendingBlob) {
   bool manage_blob_files = GetParam();
 
@@ -372,7 +456,7 @@
 
     uint64_t size = filesystem()->GetFileSize(*read_fd);
     std::unique_ptr<uint8_t[]> buf = std::make_unique<uint8_t[]>(size);
-    EXPECT_TRUE(filesystem()->Read(read_fd.get(), buf.get(), size));
+    EXPECT_THAT(filesystem()->Read(read_fd.get(), buf.get(), size), Eq(size));
 
     std::string expected_data = std::string(data.begin(), data.end());
     std::string actual_data = std::string(buf.get(), buf.get() + size);
@@ -478,7 +562,7 @@
     ScopedFd read_fd(GetScopedFdFromBlobProto(read_blob_proto));
     uint64_t size = filesystem()->GetFileSize(*read_fd);
     std::unique_ptr<uint8_t[]> buf = std::make_unique<uint8_t[]>(size);
-    EXPECT_TRUE(filesystem()->Read(read_fd.get(), buf.get(), size));
+    EXPECT_THAT(filesystem()->Read(read_fd.get(), buf.get(), size), Eq(size));
     std::string expected_data = std::string(data.begin(), data.end());
     std::string actual_data = std::string(buf.get(), buf.get() + size);
     EXPECT_EQ(expected_data, actual_data);
@@ -522,7 +606,7 @@
     ScopedFd read_fd(GetScopedFdFromBlobProto(read_blob_proto));
     uint64_t size = filesystem()->GetFileSize(*read_fd);
     std::unique_ptr<uint8_t[]> buf = std::make_unique<uint8_t[]>(size);
-    EXPECT_TRUE(filesystem()->Read(read_fd.get(), buf.get(), size));
+    EXPECT_THAT(filesystem()->Read(read_fd.get(), buf.get(), size), Eq(size));
     std::string expected_data = std::string(data.begin(), data.end());
     std::string actual_data = std::string(buf.get(), buf.get() + size);
     EXPECT_EQ(expected_data, actual_data);
@@ -585,8 +669,7 @@
 
   uint64_t size = filesystem()->GetFileSize(*read_fd);
   std::unique_ptr<uint8_t[]> buf = std::make_unique<uint8_t[]>(size);
-  filesystem()->Read(read_fd.get(), buf.get(), size);
-  close(read_fd.get());
+  EXPECT_THAT(filesystem()->Read(read_fd.get(), buf.get(), size), Eq(size));
 
   std::string expected_data = std::string(data.begin(), data.end());
   std::string actual_data = std::string(buf.get(), buf.get() + size);
@@ -699,7 +782,6 @@
 
   ScopedFd write_fd(GetScopedFdFromBlobProto(writeBlobProto));
   ASSERT_TRUE(filesystem()->Write(write_fd.get(), data.data(), data.size()));
-  close(write_fd.get());
 
   BlobProto commitBlobProto = icing.CommitBlob(blob_handle);
   ASSERT_THAT(commitBlobProto.status(), ProtoIsOk());
@@ -731,7 +813,7 @@
     ScopedFd read_fd(GetScopedFdFromBlobProto(readBlobProto));
     uint64_t size = filesystem()->GetFileSize(*read_fd);
     std::unique_ptr<uint8_t[]> buf = std::make_unique<uint8_t[]>(size);
-    ASSERT_TRUE(filesystem()->Read(read_fd.get(), buf.get(), size));
+    ASSERT_THAT(filesystem()->Read(read_fd.get(), buf.get(), size), Eq(size));
 
     std::string expected_data = std::string(data.begin(), data.end());
     std::string actual_data = std::string(buf.get(), buf.get() + size);
@@ -748,7 +830,7 @@
 
     uint64_t size = filesystem()->GetFileSize(*read_fd2);
     std::unique_ptr<uint8_t[]> buf = std::make_unique<uint8_t[]>(size);
-    ASSERT_TRUE(filesystem()->Read(read_fd2.get(), buf.get(), size));
+    ASSERT_THAT(filesystem()->Read(read_fd2.get(), buf.get(), size), Eq(size));
 
     std::string expected_data = std::string(data.begin(), data.end());
     std::string actual_data = std::string(buf.get(), buf.get() + size);
@@ -781,7 +863,6 @@
 
   ScopedFd write_fd(GetScopedFdFromBlobProto(writeBlobProto));
   ASSERT_TRUE(filesystem()->Write(write_fd.get(), data.data(), data.size()));
-  close(write_fd.get());
 
   BlobProto commitBlobProto = icing.CommitBlob(blob_handle);
   ASSERT_THAT(commitBlobProto.status(), ProtoIsOk());
@@ -842,7 +923,7 @@
     ScopedFd read_fd(GetScopedFdFromBlobProto(readBlobProto));
     uint64_t size = filesystem()->GetFileSize(*read_fd);
     std::unique_ptr<uint8_t[]> buf = std::make_unique<uint8_t[]>(size);
-    ASSERT_TRUE(filesystem()->Read(read_fd.get(), buf.get(), size));
+    ASSERT_THAT(filesystem()->Read(read_fd.get(), buf.get(), size), Eq(size));
 
     std::string expected_data = std::string(data.begin(), data.end());
     std::string actual_data = std::string(buf.get(), buf.get() + size);
@@ -859,7 +940,7 @@
 
     uint64_t size = filesystem()->GetFileSize(*read_fd2);
     std::unique_ptr<uint8_t[]> buf = std::make_unique<uint8_t[]>(size);
-    ASSERT_TRUE(filesystem()->Read(read_fd2.get(), buf.get(), size));
+    ASSERT_THAT(filesystem()->Read(read_fd2.get(), buf.get(), size), Eq(size));
 
     std::string expected_data = std::string(data.begin(), data.end());
     std::string actual_data = std::string(buf.get(), buf.get() + size);
@@ -931,8 +1012,7 @@
 
     uint64_t size = filesystem()->GetFileSize(*read_fd);
     std::unique_ptr<uint8_t[]> buf = std::make_unique<uint8_t[]>(size);
-    filesystem()->Read(read_fd.get(), buf.get(), size);
-    close(read_fd.get());
+    EXPECT_THAT(filesystem()->Read(read_fd.get(), buf.get(), size), Eq(size));
 
     std::string expected_data = std::string(data.begin(), data.end());
     std::string actual_data = std::string(buf.get(), buf.get() + size);
@@ -951,8 +1031,7 @@
 
     uint64_t size = filesystem()->GetFileSize(*read_fd2);
     std::unique_ptr<uint8_t[]> buf = std::make_unique<uint8_t[]>(size);
-    filesystem()->Read(read_fd2.get(), buf.get(), size);
-    close(read_fd2.get());
+    EXPECT_THAT(filesystem()->Read(read_fd2.get(), buf.get(), size), Eq(size));
 
     std::string expected_data = std::string(data.begin(), data.end());
     std::string actual_data = std::string(buf.get(), buf.get() + size);
@@ -1160,8 +1239,7 @@
 
   uint64_t size = filesystem()->GetFileSize(*read_fd);
   std::unique_ptr<uint8_t[]> buf = std::make_unique<uint8_t[]>(size);
-  filesystem()->Read(read_fd.get(), buf.get(), size);
-  close(read_fd.get());
+  EXPECT_THAT(filesystem()->Read(read_fd.get(), buf.get(), size), Eq(size));
 
   std::string expected_data = std::string(data.begin(), data.end());
   std::string actual_data = std::string(buf.get(), buf.get() + size);

diff --git a/icing/icing-search-engine_delete_test.cc b/icing/icing-search-engine_delete_test.cc
index a7e370c..13b6a2f 100644
--- a/icing/icing-search-engine_delete_test.cc
+++ b/icing/icing-search-engine_delete_test.cc

@@ -12,10 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <chrono>  // NOLINT
 #include <cstdint>
 #include <limits>
 #include <memory>
 #include <string>
+#include <thread>  // NOLINT
 #include <utility>
 
 #include "icing/text_classifier/lib3/utils/base/status.h"
@@ -58,11 +60,13 @@
 namespace {
 
 using ::icing::lib::portable_equals_proto::EqualsProto;
+using ::testing::AllOf;
 using ::testing::Eq;
 using ::testing::Ge;
 using ::testing::Gt;
 using ::testing::HasSubstr;
 using ::testing::IsEmpty;
+using ::testing::Property;
 using ::testing::Return;
 using ::testing::SizeIs;
 using ::testing::StrEq;
@@ -119,6 +123,7 @@
   IcingSearchEngineOptions icing_options;
   icing_options.set_base_dir(GetTestBaseDir());
   icing_options.set_enable_qualified_id_join_index_v3(true);
+  icing_options.set_enable_soft_index_restoration(true);
   icing_options.set_enable_delete_propagation_from(false);
   return icing_options;
 }
@@ -1155,6 +1160,689 @@
               EqualsProto(expected_get_result_proto6));
 }
 
+TEST_F(IcingSearchEngineDeleteTest,
+       HandleExpiredDocuments_taskSchedulerDisabled) {
+  SchemaProto schema =
+      SchemaBuilder().AddType(CreateMessageSchemaTypeConfig()).Build();
+
+  DocumentProto document = DocumentBuilder()
+                               .SetKey("namespace", "uri")
+                               .SetSchema("Message")
+                               .SetCreationTimestampMs(10)
+                               .SetTtlMs(1000)  // Expired at 1010 ms.
+                               .AddStringProperty("body", "message body1")
+                               .Build();
+
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  options.set_enable_background_task_scheduler(false);
+  options.set_enable_delete_propagation_from(true);
+  options.set_expired_document_purge_threshold_ms(0);
+
+  auto fake_clock = std::make_unique<FakeClock>();
+  FakeClock* fake_clock_ptr = fake_clock.get();
+  TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                              std::make_unique<IcingFilesystem>(),
+                              std::move(fake_clock), GetTestJniCache());
+  InitializeResultProto initialize_result = icing.Initialize();
+  ASSERT_THAT(initialize_result.status(), ProtoIsOk());
+  EXPECT_THAT(
+      initialize_result.initialize_stats().next_expiration_timestamp_ms(),
+      Eq(-1));  // No next expiration timestamp since the database is empty.
+
+  ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(document).status(), ProtoIsOk());
+
+  // Sanity check that the document is present.
+  GetResultProto expected_get_result_proto1;
+  expected_get_result_proto1.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_proto1.mutable_document() = document;
+  ASSERT_THAT(
+      icing.Get("namespace", "uri", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto1));
+
+  // Adjust the clock to 500 ms and call HandleExpiredDocuments. The document
+  // should still be present.
+  fake_clock_ptr->SetSystemTimeMilliseconds(500);
+  HandleExpiredDocumentsResultProto result_proto1 =
+      icing.HandleExpiredDocuments();
+  EXPECT_THAT(result_proto1.status(), ProtoIsOk());
+  EXPECT_THAT(result_proto1.num_expired_documents(), Eq(0));
+  EXPECT_THAT(result_proto1.num_propagated_deleted_documents(), Eq(0));
+  EXPECT_THAT(result_proto1.deleted_documents(), IsEmpty());
+  EXPECT_THAT(result_proto1.next_expiration_timestamp_ms(), Eq(1010));
+  EXPECT_THAT(
+      icing.Get("namespace", "uri", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto1));
+
+  // Adjust the clock to 1010 ms and call HandleExpiredDocuments. The document
+  // should be deleted.
+  fake_clock_ptr->SetSystemTimeMilliseconds(1010);
+  HandleExpiredDocumentsResultProto result_google::protobuf =
+      icing.HandleExpiredDocuments();
+  EXPECT_THAT(result_google::protobuf.status(), ProtoIsOk());
+  EXPECT_THAT(result_google::protobuf.num_expired_documents(), Eq(1));
+  EXPECT_THAT(result_google::protobuf.num_propagated_deleted_documents(), Eq(0));
+  EXPECT_THAT(
+      result_google::protobuf.deleted_documents(),
+      UnorderedElementsAre(AllOf(
+          Property(
+              &HandleExpiredDocumentsResultProto::DocumentGroupInfo::name_space,
+              Eq("namespace")),
+          Property(
+              &HandleExpiredDocumentsResultProto::DocumentGroupInfo::schema,
+              Eq("Message")),
+          Property(&HandleExpiredDocumentsResultProto::DocumentGroupInfo::uris,
+                   UnorderedElementsAre("uri")))));
+  EXPECT_THAT(result_google::protobuf.next_expiration_timestamp_ms(), Eq(-1));
+
+  GetResultProto expected_get_result_google::protobuf;
+  expected_get_result_google::protobuf.mutable_status()->set_code(StatusProto::NOT_FOUND);
+  expected_get_result_google::protobuf.mutable_status()->set_message(
+      "Document (namespace, uri) not found.");
+  EXPECT_THAT(
+      icing.Get("namespace", "uri", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_google::protobuf));
+
+  // Adjust the clock back to 500 ms and get the document again. Should get
+  // NOT_FOUND error.
+  // This is a hack to make sure that the document is "purged" when we handle
+  // expired documents instead of just marking deleted.
+  fake_clock_ptr->SetSystemTimeMilliseconds(500);
+  EXPECT_THAT(
+      icing.Get("namespace", "uri", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_google::protobuf));
+}
+
+TEST_F(
+    IcingSearchEngineDeleteTest,
+    HandleExpiredDocuments_taskSchedulerDisabled_propagateToChildrenWithDeletePropagationEnabled) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("name")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .AddType(
+              SchemaTypeConfigBuilder()
+                  .SetType("Email")
+                  .AddProperty(
+                      PropertyConfigBuilder()
+                          .SetName("subject")
+                          .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                          .SetCardinality(CARDINALITY_OPTIONAL))
+                  .AddProperty(PropertyConfigBuilder()
+                                   .SetName("sender")
+                                   .SetDataTypeJoinableString(
+                                       JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                       DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                                   .SetCardinality(CARDINALITY_OPTIONAL)))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("Message")
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("body")
+                                        .SetDataTypeString(TERM_MATCH_PREFIX,
+                                                           TOKENIZER_PLAIN)
+                                        .SetCardinality(CARDINALITY_REQUIRED))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("sender")
+                                        .SetDataTypeJoinableString(
+                                            JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                            DELETE_PROPAGATION_TYPE_NONE)
+                                        .SetCardinality(CARDINALITY_OPTIONAL)))
+          .AddType(
+              SchemaTypeConfigBuilder()
+                  .SetType("Label")
+                  .AddProperty(
+                      PropertyConfigBuilder()
+                          .SetName("name")
+                          .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                          .SetCardinality(CARDINALITY_REQUIRED))
+                  .AddProperty(PropertyConfigBuilder()
+                                   .SetName("target")
+                                   .SetDataTypeJoinableString(
+                                       JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                       DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                                   .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  DocumentProto person1 = DocumentBuilder()
+                              .SetKey("namespace", "person1")
+                              .SetSchema("Person")
+                              .SetCreationTimestampMs(10)
+                              .SetTtlMs(1000)  // Expired at 1010 ms.
+                              .AddStringProperty("name", "Alice")
+                              .Build();
+  DocumentProto person2 = DocumentBuilder()
+                              .SetKey("namespace", "person2")
+                              .SetSchema("Person")
+                              .SetCreationTimestampMs(10)
+                              .SetTtlMs(3000)  // Expired at 3010 ms.
+                              .AddStringProperty("name", "Bob")
+                              .Build();
+  DocumentProto email1 = DocumentBuilder()
+                             .SetKey("namespace", "email1")
+                             .SetSchema("Email")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(0)  // Never expire.
+                             .AddStringProperty("subject", "test")
+                             .AddStringProperty("sender", "namespace#person1")
+                             .Build();
+  DocumentProto email2 = DocumentBuilder()
+                             .SetKey("namespace", "email2")
+                             .SetSchema("Email")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(0)  // Never expire.
+                             .AddStringProperty("subject", "test")
+                             .AddStringProperty("sender", "namespace#person2")
+                             .Build();
+  DocumentProto message1 = DocumentBuilder()
+                               .SetKey("namespace", "message1")
+                               .SetSchema("Message")
+                               .SetCreationTimestampMs(10)
+                               .SetTtlMs(0)  // Never expire.
+                               .AddStringProperty("body", "test")
+                               .AddStringProperty("sender", "namespace#person1")
+                               .Build();
+  DocumentProto message2 = DocumentBuilder()
+                               .SetKey("namespace", "message2")
+                               .SetSchema("Message")
+                               .SetCreationTimestampMs(10)
+                               .SetTtlMs(0)  // Never expire.
+                               .AddStringProperty("body", "test")
+                               .AddStringProperty("sender", "namespace#person2")
+                               .Build();
+  DocumentProto label1 = DocumentBuilder()
+                             .SetKey("namespace", "label1")
+                             .SetSchema("Label")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(0)  // Never expire.
+                             .AddStringProperty("name", "label1")
+                             .AddStringProperty("target", "namespace#email1")
+                             .Build();
+  DocumentProto label2 = DocumentBuilder()
+                             .SetKey("namespace", "label2")
+                             .SetSchema("Label")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(0)  // Never expire.
+                             .AddStringProperty("name", "label2")
+                             .AddStringProperty("target", "namespace#email2")
+                             .Build();
+  DocumentProto label3 = DocumentBuilder()
+                             .SetKey("namespace", "label3")
+                             .SetSchema("Label")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(0)  // Never expire.
+                             .AddStringProperty("name", "label3")
+                             .AddStringProperty("target", "namespace#message1")
+                             .Build();
+  DocumentProto label4 = DocumentBuilder()
+                             .SetKey("namespace", "label4")
+                             .SetSchema("Label")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(0)  // Never expire.
+                             .AddStringProperty("name", "label4")
+                             .AddStringProperty("target", "namespace#message2")
+                             .Build();
+
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  options.set_enable_background_task_scheduler(false);
+  options.set_enable_delete_propagation_from(true);
+  options.set_expired_document_purge_threshold_ms(0);
+
+  auto fake_clock = std::make_unique<FakeClock>();
+  FakeClock* fake_clock_ptr = fake_clock.get();
+  TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                              std::make_unique<IcingFilesystem>(),
+                              std::move(fake_clock), GetTestJniCache());
+  InitializeResultProto initialize_result = icing.Initialize();
+  ASSERT_THAT(initialize_result.status(), ProtoIsOk());
+  EXPECT_THAT(
+      initialize_result.initialize_stats().next_expiration_timestamp_ms(),
+      Eq(-1));  // No next expiration timestamp since the database is empty.
+
+  ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(person1).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(person2).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(email1).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(email2).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(message1).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(message2).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(label1).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(label2).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(label3).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(label4).status(), ProtoIsOk());
+
+  // Adjust the clock to 1010 ms and call HandleExpiredDocuments. person1 should
+  // be deleted and delete propagation should be triggered for email1 (child)
+  // and label1 (grandchild).
+  fake_clock_ptr->SetSystemTimeMilliseconds(1010);
+  HandleExpiredDocumentsResultProto result_proto =
+      icing.HandleExpiredDocuments();
+  EXPECT_THAT(result_proto.status(), ProtoIsOk());
+  EXPECT_THAT(result_proto.num_expired_documents(), Eq(1));
+  EXPECT_THAT(result_proto.num_propagated_deleted_documents(), Eq(2));
+  EXPECT_THAT(
+      result_proto.deleted_documents(),
+      UnorderedElementsAre(
+          AllOf(
+              Property(&HandleExpiredDocumentsResultProto::DocumentGroupInfo::
+                           name_space,
+                       Eq("namespace")),
+              Property(
+                  &HandleExpiredDocumentsResultProto::DocumentGroupInfo::schema,
+                  Eq("Person")),
+              Property(
+                  &HandleExpiredDocumentsResultProto::DocumentGroupInfo::uris,
+                  UnorderedElementsAre("person1"))),
+          AllOf(
+              Property(&HandleExpiredDocumentsResultProto::DocumentGroupInfo::
+                           name_space,
+                       Eq("namespace")),
+              Property(
+                  &HandleExpiredDocumentsResultProto::DocumentGroupInfo::schema,
+                  Eq("Email")),
+              Property(
+                  &HandleExpiredDocumentsResultProto::DocumentGroupInfo::uris,
+                  UnorderedElementsAre("email1"))),
+          AllOf(
+              Property(&HandleExpiredDocumentsResultProto::DocumentGroupInfo::
+                           name_space,
+                       Eq("namespace")),
+              Property(
+                  &HandleExpiredDocumentsResultProto::DocumentGroupInfo::schema,
+                  Eq("Label")),
+              Property(
+                  &HandleExpiredDocumentsResultProto::DocumentGroupInfo::uris,
+                  UnorderedElementsAre("label1")))));
+  EXPECT_THAT(result_proto.next_expiration_timestamp_ms(), Eq(3010));
+
+  // Adjust the clock back to 500 ms and verify Get API for email, message and
+  // label documents.
+  // This is a hack to make sure that the document is "purged" when we handle
+  // expired documents instead of just marking deleted.
+  fake_clock_ptr->SetSystemTimeMilliseconds(500);
+
+  // Email1 should be deleted. The joinable property "sender" in schema type
+  // "Email" has delete propagation type PROPAGATE_FROM and the referenced
+  // document "person1" is deleted.
+  GetResultProto expected_get_result_proto1;
+  expected_get_result_proto1.mutable_status()->set_code(StatusProto::NOT_FOUND);
+  expected_get_result_proto1.mutable_status()->set_message(
+      "Document (namespace, email1) not found.");
+  EXPECT_THAT(
+      icing.Get("namespace", "email1", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto1));
+
+  // Email2 should still exist. The joinable property "sender" in schema type
+  // "Email" has delete propagation type PROPAGATE_FROM but the referenced
+  // document "person2" is not deleted.
+  GetResultProto expected_get_result_google::protobuf;
+  expected_get_result_google::protobuf.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_google::protobuf.mutable_document() = email2;
+  EXPECT_THAT(
+      icing.Get("namespace", "email2", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_google::protobuf));
+
+  // Message1 should still exist. The joinable property "sender" in schema type
+  // "Message" has delete propagation type NONE.
+  GetResultProto expected_get_result_proto3;
+  expected_get_result_proto3.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_proto3.mutable_document() = message1;
+  EXPECT_THAT(icing.Get("namespace", "message1",
+                        GetResultSpecProto::default_instance()),
+              EqualsProto(expected_get_result_proto3));
+
+  // Message2 should still exist. The joinable property "sender" in schema type
+  // "Message" has delete propagation type NONE, and the referenced document
+  // "person2" is not deleted.
+  GetResultProto expected_get_result_proto4;
+  expected_get_result_proto4.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_proto4.mutable_document() = message2;
+  EXPECT_THAT(icing.Get("namespace", "message2",
+                        GetResultSpecProto::default_instance()),
+              EqualsProto(expected_get_result_proto4));
+
+  // Label1 should be deleted (propagated from email1).
+  GetResultProto expected_get_result_proto5;
+  expected_get_result_proto5.mutable_status()->set_code(StatusProto::NOT_FOUND);
+  expected_get_result_proto5.mutable_status()->set_message(
+      "Document (namespace, label1) not found.");
+  EXPECT_THAT(
+      icing.Get("namespace", "label1", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto5));
+
+  // Label2 should still exist.
+  GetResultProto expected_get_result_proto6;
+  expected_get_result_proto6.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_proto6.mutable_document() = label2;
+  EXPECT_THAT(
+      icing.Get("namespace", "label2", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto6));
+
+  // Label3 should still exist.
+  GetResultProto expected_get_result_proto7;
+  expected_get_result_proto7.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_proto7.mutable_document() = label3;
+  EXPECT_THAT(
+      icing.Get("namespace", "label3", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto7));
+
+  // Label4 should still exist.
+  GetResultProto expected_get_result_proto8;
+  expected_get_result_proto8.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_proto8.mutable_document() = label4;
+  EXPECT_THAT(
+      icing.Get("namespace", "label4", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto8));
+}
+
+TEST_F(
+    IcingSearchEngineDeleteTest,
+    HandleExpiredDocuments_taskSchedulerDisabled_shouldPurgeDocumentsThatExpireWithinThreshold) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("name")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .AddType(
+              SchemaTypeConfigBuilder()
+                  .SetType("Email")
+                  .AddProperty(
+                      PropertyConfigBuilder()
+                          .SetName("subject")
+                          .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                          .SetCardinality(CARDINALITY_OPTIONAL))
+                  .AddProperty(PropertyConfigBuilder()
+                                   .SetName("sender")
+                                   .SetDataTypeJoinableString(
+                                       JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                       DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                                   .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  DocumentProto person1 = DocumentBuilder()
+                              .SetKey("namespace", "person1")
+                              .SetSchema("Person")
+                              .SetCreationTimestampMs(10)
+                              .SetTtlMs(1000)  // Expired at 1010 ms.
+                              .AddStringProperty("name", "Alice")
+                              .Build();
+  DocumentProto person2 = DocumentBuilder()
+                              .SetKey("namespace", "person2")
+                              .SetSchema("Person")
+                              .SetCreationTimestampMs(10)
+                              .SetTtlMs(1101)  // Expired at 1111 ms.
+                              .AddStringProperty("name", "Bob")
+                              .Build();
+  DocumentProto person3 = DocumentBuilder()
+                              .SetKey("namespace", "person3")
+                              .SetSchema("Person")
+                              .SetCreationTimestampMs(10)
+                              .SetTtlMs(1100)  // Expired at 1110 ms.
+                              .AddStringProperty("name", "Bob")
+                              .Build();
+  DocumentProto email1 = DocumentBuilder()
+                             .SetKey("namespace", "email1")
+                             .SetSchema("Email")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(0)  // Never expire.
+                             .AddStringProperty("subject", "test")
+                             .AddStringProperty("sender", "namespace#person1")
+                             .Build();
+  DocumentProto email2 = DocumentBuilder()
+                             .SetKey("namespace", "email2")
+                             .SetSchema("Email")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(0)  // Never expire.
+                             .AddStringProperty("subject", "test")
+                             .AddStringProperty("sender", "namespace#person2")
+                             .Build();
+  DocumentProto email3 = DocumentBuilder()
+                             .SetKey("namespace", "email3")
+                             .SetSchema("Email")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(0)  // Never expire.
+                             .AddStringProperty("subject", "test")
+                             .AddStringProperty("sender", "namespace#person3")
+                             .Build();
+
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  options.set_enable_background_task_scheduler(false);
+  options.set_enable_delete_propagation_from(true);
+  options.set_expired_document_purge_threshold_ms(100);
+
+  auto fake_clock = std::make_unique<FakeClock>();
+  FakeClock* fake_clock_ptr = fake_clock.get();
+  TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                              std::make_unique<IcingFilesystem>(),
+                              std::move(fake_clock), GetTestJniCache());
+  InitializeResultProto initialize_result = icing.Initialize();
+  ASSERT_THAT(initialize_result.status(), ProtoIsOk());
+  EXPECT_THAT(
+      initialize_result.initialize_stats().next_expiration_timestamp_ms(),
+      Eq(-1));  // No next expiration timestamp since the database is empty.
+
+  ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(person1).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(person2).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(person3).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(email1).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(email2).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(email3).status(), ProtoIsOk());
+
+  // Adjust the clock to 1010 ms and call HandleExpiredDocuments. person1 and
+  // person3 should be deleted, and delete propagation should be triggered for
+  // email1 and email3.
+  fake_clock_ptr->SetSystemTimeMilliseconds(1010);
+  HandleExpiredDocumentsResultProto result_proto =
+      icing.HandleExpiredDocuments();
+  EXPECT_THAT(result_proto.status(), ProtoIsOk());
+  EXPECT_THAT(result_proto.num_expired_documents(), Eq(2));
+  EXPECT_THAT(result_proto.num_propagated_deleted_documents(), Eq(2));
+  EXPECT_THAT(
+      result_proto.deleted_documents(),
+      UnorderedElementsAre(
+          AllOf(
+              Property(&HandleExpiredDocumentsResultProto::DocumentGroupInfo::
+                           name_space,
+                       Eq("namespace")),
+              Property(
+                  &HandleExpiredDocumentsResultProto::DocumentGroupInfo::schema,
+                  Eq("Person")),
+              Property(
+                  &HandleExpiredDocumentsResultProto::DocumentGroupInfo::uris,
+                  UnorderedElementsAre("person1", "person3"))),
+          AllOf(
+              Property(&HandleExpiredDocumentsResultProto::DocumentGroupInfo::
+                           name_space,
+                       Eq("namespace")),
+              Property(
+                  &HandleExpiredDocumentsResultProto::DocumentGroupInfo::schema,
+                  Eq("Email")),
+              Property(
+                  &HandleExpiredDocumentsResultProto::DocumentGroupInfo::uris,
+                  UnorderedElementsAre("email1", "email3")))));
+  EXPECT_THAT(result_proto.next_expiration_timestamp_ms(), Eq(1111));
+}
+
+TEST_F(
+    IcingSearchEngineDeleteTest,
+    HandleExpiredDocuments_taskSchedulerEnabled_shouldScheduleNextPurgingExpirationTask) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("name")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .AddType(
+              SchemaTypeConfigBuilder()
+                  .SetType("Email")
+                  .AddProperty(
+                      PropertyConfigBuilder()
+                          .SetName("subject")
+                          .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                          .SetCardinality(CARDINALITY_OPTIONAL))
+                  .AddProperty(PropertyConfigBuilder()
+                                   .SetName("sender")
+                                   .SetDataTypeJoinableString(
+                                       JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                       DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                                   .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  DocumentProto person1 = DocumentBuilder()
+                              .SetKey("namespace", "person1")
+                              .SetSchema("Person")
+                              .SetCreationTimestampMs(10)
+                              .SetTtlMs(5000)  // Expired at 5010 ms.
+                              .AddStringProperty("name", "Alice")
+                              .Build();
+  DocumentProto person2 = DocumentBuilder()
+                              .SetKey("namespace", "person2")
+                              .SetSchema("Person")
+                              .SetCreationTimestampMs(10)
+                              .SetTtlMs(5400)  // Expired at 5410 ms.
+                              .AddStringProperty("name", "Bob")
+                              .Build();
+  DocumentProto email1 = DocumentBuilder()
+                             .SetKey("namespace", "email1")
+                             .SetSchema("Email")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(0)  // Never expire.
+                             .AddStringProperty("subject", "test")
+                             .AddStringProperty("sender", "namespace#person1")
+                             .Build();
+  DocumentProto email2 = DocumentBuilder()
+                             .SetKey("namespace", "email2")
+                             .SetSchema("Email")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(0)  // Never expire.
+                             .AddStringProperty("subject", "test")
+                             .AddStringProperty("sender", "namespace#person2")
+                             .Build();
+
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  options.set_enable_background_task_scheduler(true);
+  options.set_enable_delete_propagation_from(true);
+  options.set_expired_document_purge_threshold_ms(100);
+
+  {
+    // Initialize Icing and put all documents. Destruct Icing.
+    TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                                std::make_unique<IcingFilesystem>(),
+                                std::make_unique<FakeClock>(),
+                                GetTestJniCache());
+    InitializeResultProto initialize_result = icing.Initialize();
+    ASSERT_THAT(initialize_result.status(), ProtoIsOk());
+    EXPECT_THAT(
+        initialize_result.initialize_stats().next_expiration_timestamp_ms(),
+        Eq(-1));  // No next expiration timestamp since the database is empty.
+
+    ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(person1).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(person2).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(email1).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(email2).status(), ProtoIsOk());
+  }
+
+  // Initialize Icing again with a fake clock and t = 3500 ms. Initialization
+  // should schedule the purging expiration task at t = 5010 ms.
+  auto fake_clock = std::make_unique<FakeClock>();
+  FakeClock* fake_clock_ptr = fake_clock.get();
+  fake_clock->SetSystemTimeMilliseconds(3500);
+  TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                              std::make_unique<IcingFilesystem>(),
+                              std::move(fake_clock), GetTestJniCache());
+  InitializeResultProto initialize_result = icing.Initialize();
+  ASSERT_THAT(initialize_result.status(), ProtoIsOk());
+  EXPECT_THAT(
+      initialize_result.initialize_stats().next_expiration_timestamp_ms(),
+      Eq(5010));
+
+  // Sanity check that person1, person2, email1 and email2 are present.
+  GetResultProto expected_get_result_proto1;
+  expected_get_result_proto1.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_proto1.mutable_document() = person1;
+  EXPECT_THAT(
+      icing.Get("namespace", "person1", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto1));
+
+  GetResultProto expected_get_result_google::protobuf;
+  expected_get_result_google::protobuf.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_google::protobuf.mutable_document() = person2;
+  EXPECT_THAT(
+      icing.Get("namespace", "person2", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_google::protobuf));
+
+  GetResultProto expected_get_result_proto3;
+  expected_get_result_proto3.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_proto3.mutable_document() = email1;
+  EXPECT_THAT(
+      icing.Get("namespace", "email1", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto3));
+
+  GetResultProto expected_get_result_proto4;
+  expected_get_result_proto4.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_proto4.mutable_document() = email2;
+  EXPECT_THAT(
+      icing.Get("namespace", "email2", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto4));
+
+  // 1. Adjust the clock from 3500 to 5010 ms and sleep for 1550 ms. person1 and
+  //    email1 should be purged by the scheduled task, but person2 and email2
+  //    should still be alive.
+  fake_clock_ptr->SetSystemTimeMilliseconds(5010);
+  std::this_thread::sleep_for(std::chrono::milliseconds(1550));
+
+  GetResultProto expected_get_result_proto5;
+  expected_get_result_proto5.mutable_status()->set_code(StatusProto::NOT_FOUND);
+  expected_get_result_proto5.mutable_status()->set_message(
+      "Document (namespace, person1) not found.");
+  EXPECT_THAT(
+      icing.Get("namespace", "person1", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto5));
+
+  GetResultProto expected_get_result_proto6;
+  expected_get_result_proto6.mutable_status()->set_code(StatusProto::NOT_FOUND);
+  expected_get_result_proto6.mutable_status()->set_message(
+      "Document (namespace, email1) not found.");
+  EXPECT_THAT(
+      icing.Get("namespace", "email1", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto6));
+
+  EXPECT_THAT(
+      icing.Get("namespace", "person2", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_google::protobuf));
+  EXPECT_THAT(
+      icing.Get("namespace", "email2", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto4));
+
+  // 2. Adjust the clock from 5010 to 5410 ms and sleep for 400 ms. person2 and
+  //    email2 should be purged by the scheduled task.
+  fake_clock_ptr->SetSystemTimeMilliseconds(5410);
+  std::this_thread::sleep_for(std::chrono::milliseconds(400));
+
+  GetResultProto expected_get_result_proto7;
+  expected_get_result_proto7.mutable_status()->set_code(StatusProto::NOT_FOUND);
+  expected_get_result_proto7.mutable_status()->set_message(
+      "Document (namespace, person2) not found.");
+  EXPECT_THAT(
+      icing.Get("namespace", "person2", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto7));
+
+  GetResultProto expected_get_result_proto8;
+  expected_get_result_proto8.mutable_status()->set_code(StatusProto::NOT_FOUND);
+  expected_get_result_proto8.mutable_status()->set_message(
+      "Document (namespace, email2) not found.");
+  EXPECT_THAT(
+      icing.Get("namespace", "email2", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto8));
+}
+
 }  // namespace
 }  // namespace lib
 }  // namespace icing

diff --git a/icing/icing-search-engine_fuzz_test.cc b/icing/icing-search-engine_fuzz_test.cc
index 6f1cf69..72ac01d 100644
--- a/icing/icing-search-engine_fuzz_test.cc
+++ b/icing/icing-search-engine_fuzz_test.cc

@@ -12,20 +12,23 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include <cstddef>
-#include <cstdint>
 
 #include "icing/text_classifier/lib3/utils/base/status.h"
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
+#include "gmock/gmock.h"
+#include "testing/fuzzing/fuzztest.h"
 #include "icing/document-builder.h"
+#include "icing/file/filesystem.h"
 #include "icing/icing-search-engine.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/initialize.pb.h"
 #include "icing/proto/schema.pb.h"
 #include "icing/proto/scoring.pb.h"
 #include "icing/proto/search.pb.h"
+#include "icing/proto/status.pb.h"
 #include "icing/proto/term.pb.h"
 #include "icing/schema-builder.h"
+#include "icing/testing/common-matchers.h"
 #include "icing/testing/test-data.h"
 #include "icing/testing/tmp-directory.h"
 #include "icing/util/icu-data-file-helper.h"
@@ -40,32 +43,28 @@
   return icing_options;
 }
 
-DocumentProto MakeDocument(const uint8_t* data, size_t size) {
-  // TODO (sidchhabra): Added more optimized fuzzing techniques.
-  DocumentProto document;
-  std::string string_prop(reinterpret_cast<const char*>(data), size);
+DocumentProto MakeDocument(const std::string& data) {
   return DocumentBuilder()
       .SetKey("namespace", "uri1")
       .SetSchema("Message")
-      .AddStringProperty("body", string_prop)
+      .AddStringProperty("body", data)
+      .SetCreationTimestampMs(0L)
       .Build();
 }
 
-SearchSpecProto SetSearchSpec(const uint8_t* data, size_t size) {
+SearchSpecProto SetSearchSpec(const std::string& data) {
   SearchSpecProto search_spec;
-  search_spec.set_term_match_type(TermMatchType::PREFIX);
-  // TODO (sidchhabra): Added more optimized fuzzing techniques.
-  std::string query_string(reinterpret_cast<const char*>(data), size);
-  search_spec.set_query(query_string);
+  search_spec.set_term_match_type(TermMatchType::EXACT_ONLY);
+  search_spec.set_query(data);
   return search_spec;
 }
 
-extern "C" int LLVMFuzzerTestOneInput(const uint8_t* data, size_t size) {
+void StringFuzzTest(const std::string& data) {
   // Initialize
   IcingSearchEngineOptions icing_options = Setup();
   std::string icu_data_file_path = GetTestFilePath("icing/icu.dat");
   if (!icu_data_file_helper::SetUpIcuDataFile(icu_data_file_path).ok()) {
-    return 1;
+    return;
   }
   IcingSearchEngine icing(icing_options);
   const Filesystem filesystem_;
@@ -84,19 +83,36 @@
   icing.SetSchema(schema_proto);
 
   // Index
-  DocumentProto document = MakeDocument(data, size);
+  DocumentProto document = MakeDocument(data);
   icing.Put(document);
 
   // Query
-  SearchSpecProto search_spec = SetSearchSpec(data, size);
+  SearchSpecProto search_spec = SetSearchSpec(data);
   ScoringSpecProto scoring_spec;
   scoring_spec.set_rank_by(ScoringSpecProto::RankingStrategy::DOCUMENT_SCORE);
-  ResultSpecProto result_spec;
-  libtextclassifier3::StatusOr<SearchResultProto> result =
-      icing.Search(search_spec, scoring_spec, result_spec);
-  return 0;
+  SearchResultProto result = icing.Search(search_spec, scoring_spec,
+                                          ResultSpecProto::default_instance());
+  EXPECT_THAT(result.results(0).document().uri(), document.uri());
 }
 
+// vector of emojis to be used in the fuzz test
+const std::vector<std::string> emojis = {"😀", "😂", "😊", "😘",
+  "😠", "😢", "😱", "🍻"};
+
+// add a random emoji to a string containing a decomposed character
+std::string addEmoji(const std::string& str, const std::string& emoji) {
+  return str + emoji;
+}
+
+// dictionary including Spanish words with accents (lowercase and uppercase)
+std::vector<std::string> data = {"dìas", "Sí", "Más", "Tú", "Él",
+  "Dónde", "Qué", "miércoles", "café",
+  "sábado"};
+
+FUZZ_TEST(TestSuite, StringFuzzTest)
+    .WithDomains(fuzztest::Map(addEmoji, fuzztest::ElementOf(data),
+                               fuzztest::ElementOf(emojis)));
+
 }  // namespace
 }  // namespace lib
 }  // namespace icing

diff --git a/icing/icing-search-engine_initialization_test.cc b/icing/icing-search-engine_initialization_test.cc
index 5fc862b..836f849 100644
--- a/icing/icing-search-engine_initialization_test.cc
+++ b/icing/icing-search-engine_initialization_test.cc

@@ -13,11 +13,13 @@
 // limitations under the License.
 
 #include <algorithm>
+#include <chrono>  // NOLINT
 #include <cstdint>
 #include <limits>
 #include <memory>
 #include <string>
 #include <string_view>
+#include <thread>  // NOLINT
 #include <tuple>
 #include <unordered_set>
 #include <utility>
@@ -57,6 +59,7 @@
 #include "icing/legacy/index/icing-mock-filesystem.h"
 #include "icing/portable/endian.h"
 #include "icing/portable/equals-proto.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/portable/platform.h"
 #include "icing/proto/debug.pb.h"
 #include "icing/proto/document.pb.h"
@@ -96,6 +99,7 @@
 #include "icing/transform/normalizer-options.h"
 #include "icing/transform/normalizer.h"
 #include "icing/util/clock.h"
+#include "icing/util/document-util.h"
 #include "icing/util/icu-data-file-helper.h"
 #include "icing/util/tokenized-document.h"
 #include "unicode/uloc.h"
@@ -260,8 +264,12 @@
   icing_options.set_enable_embedding_quantization(true);
   icing_options.set_enable_blob_store(true);
   icing_options.set_enable_qualified_id_join_index_v3(true);
+  icing_options.set_enable_soft_index_restoration(true);
   icing_options.set_enable_delete_propagation_from(false);
   icing_options.set_enable_marker_file_for_optimize(true);
+  icing_options.set_enable_proto_log_new_header_format(true);
+  icing_options.set_embedding_index_num_shards(32);
+  icing_options.set_enable_non_existent_qualified_id_join(true);
   return icing_options;
 }
 
@@ -500,6 +508,7 @@
        DeletePropagationEnabledAndJoinIndexV3DisabledReturnsInvalidArgument) {
   IcingSearchEngineOptions icing_options = GetDefaultIcingOptions();
   icing_options.set_enable_qualified_id_join_index_v3(false);
+  icing_options.set_enable_soft_index_restoration(true);
   icing_options.set_enable_delete_propagation_from(true);
 
   IcingSearchEngine icing(icing_options, GetTestJniCache());
@@ -508,7 +517,24 @@
               ProtoStatusIs(StatusProto::INVALID_ARGUMENT));
   EXPECT_THAT(initialize_result_proto.status().message(),
               HasSubstr("Delete propagation is enabled but qualified id join "
-                        "index v3 is not enabled."));
+                        "index v3 or soft index restoration is not enabled."));
+}
+
+TEST_F(
+    IcingSearchEngineInitializationTest,
+    DeletePropagationEnabledAndSoftIndexRestorationDisabledReturnsInvalidArgument) {
+  IcingSearchEngineOptions icing_options = GetDefaultIcingOptions();
+  icing_options.set_enable_qualified_id_join_index_v3(true);
+  icing_options.set_enable_soft_index_restoration(false);
+  icing_options.set_enable_delete_propagation_from(true);
+
+  IcingSearchEngine icing(icing_options, GetTestJniCache());
+  InitializeResultProto initialize_result_proto = icing.Initialize();
+  EXPECT_THAT(initialize_result_proto.status(),
+              ProtoStatusIs(StatusProto::INVALID_ARGUMENT));
+  EXPECT_THAT(initialize_result_proto.status().message(),
+              HasSubstr("Delete propagation is enabled but qualified id join "
+                        "index v3 or soft index restoration is not enabled."));
 }
 
 TEST_F(IcingSearchEngineInitializationTest, GoodCompressionLevelReturnsOk) {
@@ -543,6 +569,99 @@
   }
 }
 
+TEST_F(IcingSearchEngineInitializationTest,
+       OutOfRangeCompressionMemLevelReturnsInvalidArgument) {
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+
+  // Mem level must be between 1 and 9 inclusive.
+  options.set_compression_mem_level(-1);
+  {
+    IcingSearchEngine icing(options, GetTestJniCache());
+    EXPECT_THAT(icing.Initialize().status(),
+                ProtoStatusIs(StatusProto::INVALID_ARGUMENT));
+  }
+
+  options.set_compression_mem_level(10);
+  {
+    IcingSearchEngine icing(options, GetTestJniCache());
+    EXPECT_THAT(icing.Initialize().status(),
+                ProtoStatusIs(StatusProto::INVALID_ARGUMENT));
+  }
+
+  options.set_compression_mem_level(0);
+  {
+    IcingSearchEngine icing(options, GetTestJniCache());
+    EXPECT_THAT(icing.Initialize().status(),
+                ProtoStatusIs(StatusProto::INVALID_ARGUMENT));
+  }
+}
+
+TEST_F(IcingSearchEngineInitializationTest,
+       OutOfRangeBlobStoreCompressionMemLevelReturnsInvalidArgument) {
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+
+  // Mem level must be between 1 and 9 inclusive.
+  options.set_blob_store_compression_mem_level(-1);
+  {
+    IcingSearchEngine icing(options, GetTestJniCache());
+    EXPECT_THAT(icing.Initialize().status(),
+                ProtoStatusIs(StatusProto::INVALID_ARGUMENT));
+  }
+
+  options.set_blob_store_compression_mem_level(10);
+  {
+    IcingSearchEngine icing(options, GetTestJniCache());
+    EXPECT_THAT(icing.Initialize().status(),
+                ProtoStatusIs(StatusProto::INVALID_ARGUMENT));
+  }
+
+  options.set_blob_store_compression_mem_level(0);
+  {
+    IcingSearchEngine icing(options, GetTestJniCache());
+    EXPECT_THAT(icing.Initialize().status(),
+                ProtoStatusIs(StatusProto::INVALID_ARGUMENT));
+  }
+}
+
+TEST_F(IcingSearchEngineInitializationTest,
+       ReinitializingWithDifferentCompressionMemLevelOk) {
+  DocumentProto document =
+      DocumentBuilder()
+          .SetKey("icing", "fake_type/0")
+          .SetSchema("Message")
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .AddStringProperty("body", "message body")
+          .AddInt64Property("indexableInteger", 123)
+          .Build();
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  {
+    // Initialize and put document
+    IcingSearchEngine icing(options, GetTestJniCache());
+    ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+    ASSERT_THAT(icing.SetSchema(CreateMessageSchema()).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(document).status(), ProtoIsOk());
+    ASSERT_THAT(icing.PersistToDisk(PersistType::FULL).status(), ProtoIsOk());
+  }
+
+  // Reinitialize with different compression mem level and check that the
+  // document is still searchable.
+  options.set_compression_mem_level(1);
+  IcingSearchEngine icing(options, GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+
+  SearchSpecProto search_spec;
+  search_spec.set_query("message");
+  search_spec.set_term_match_type(TermMatchType::EXACT_ONLY);
+  SearchResultProto expected_search_result_proto;
+  expected_search_result_proto.mutable_status()->set_code(StatusProto::OK);
+  *expected_search_result_proto.mutable_results()->Add()->mutable_document() =
+      document;
+  EXPECT_THAT(
+      icing.Search(search_spec, GetDefaultScoringSpec(),
+                   ResultSpecProto::default_instance()),
+      EqualsSearchResultIgnoreStatsAndScores(expected_search_result_proto));
+}
+
 TEST_F(IcingSearchEngineInitializationTest, FailToCreateDocStore) {
   auto mock_filesystem = std::make_unique<MockFilesystem>();
   // This fails DocumentStore::Create()
@@ -934,6 +1053,9 @@
   EXPECT_THAT(
       initialize_result.initialize_stats().embedding_index_restoration_cause(),
       Eq(InitializeStatsProto::NONE));
+  // needs_persist_type should be set to RECOVERY_PROOF.
+  EXPECT_THAT(initialize_result.needs_persist_type(),
+              Eq(PersistType::RECOVERY_PROOF));
 
   // ("namespace", "uri1") should be found.
   GetResultProto expected_get_result_proto;
@@ -1264,7 +1386,8 @@
                            /*lite_index_sort_size=*/1024 * 8);
     ICING_ASSERT_OK_AND_ASSIGN(
         std::unique_ptr<Index> index,
-        Index::Create(options, filesystem(), icing_filesystem()));
+        Index::Create(options, filesystem(), icing_filesystem(),
+                      feature_flags_.get()));
     ICING_ASSERT_OK(index->Merge());
     ICING_ASSERT_OK(index->PersistToDisk());
   }
@@ -1294,14 +1417,18 @@
     std::string doc_store_dir = GetDocumentDir();
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
-        DocumentStore::Create(filesystem(), doc_store_dir, &fake_clock,
-                              schema_store.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            filesystem(), doc_store_dir, &fake_clock, schema_store.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     std::unique_ptr<DocumentStore> document_store =
         std::move(create_result.document_store);
 
@@ -1473,23 +1600,29 @@
             filesystem(), GetBlobDir(), &fake_clock,
             /*orphan_blob_time_to_live_ms=*/0,
             PortableFileBackedProtoLog<BlobInfoProto>::kDefaultCompressionLevel,
-            /*manage_blob_files=*/true));
+            protobuf_ports::kDefaultMemLevel,
+            /*manage_blob_files=*/true, feature_flags_.get()));
 
     // Puts message2 into DocumentStore but doesn't index it.
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
-        DocumentStore::Create(filesystem(), GetDocumentDir(), &fake_clock,
-                              schema_store.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            filesystem(), GetDocumentDir(), &fake_clock, schema_store.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     std::unique_ptr<DocumentStore> document_store =
         std::move(create_result.document_store);
 
-    ICING_EXPECT_OK(document_store->Put(message2));
+    ICING_EXPECT_OK(
+        document_store->Put(document_util::CreateDocumentWrapper(message2)));
   }
 
   // Mock filesystem to observe and check the behavior of all indices.
@@ -1789,6 +1922,9 @@
   EXPECT_THAT(initialize_result.initialize_stats()
                   .qualified_id_join_index_restoration_cause(),
               Eq(InitializeStatsProto::NONE));
+  // needs_persist_type should be set to RECOVERY_PROOF.
+  EXPECT_THAT(initialize_result.needs_persist_type(),
+              Eq(PersistType::RECOVERY_PROOF));
 
   // Verify join search: join a query for `name:person` with a child query for
   // `body:message` based on the child's `senderQualifiedId` field. message2
@@ -2485,7 +2621,7 @@
                                      /*index_merge_size=*/100,
                                      /*lite_index_sort_at_indexing=*/true,
                                      /*lite_index_sort_size=*/50),
-                      filesystem(), icing_filesystem()));
+                      filesystem(), icing_filesystem(), feature_flags_.get()));
     ICING_ASSERT_OK(index->PersistToDisk());
   }
 
@@ -3091,7 +3227,7 @@
                            /*index_merge_size=*/message.ByteSizeLong(),
                            /*lite_index_sort_at_indexing=*/true,
                            /*lite_index_sort_size=*/8),
-            filesystem(), icing_filesystem()));
+            filesystem(), icing_filesystem(), feature_flags_.get()));
     DocumentId original_last_added_doc_id = index->last_added_document_id();
     index->set_last_added_document_id(original_last_added_doc_id + 1);
     Index::Editor editor =
@@ -3223,7 +3359,7 @@
                            /*index_merge_size=*/message.ByteSizeLong(),
                            /*lite_index_sort_at_indexing=*/true,
                            /*lite_index_sort_size=*/8),
-            filesystem(), icing_filesystem()));
+            filesystem(), icing_filesystem(), feature_flags_.get()));
     ICING_ASSERT_OK_AND_ASSIGN(
         std::unique_ptr<DocHitInfoIterator> doc_hit_info_iter,
         index->GetIterator("foo", /*term_start_index=*/0,
@@ -3337,7 +3473,7 @@
                            /*index_merge_size=*/message.ByteSizeLong(),
                            /*lite_index_sort_at_indexing=*/true,
                            /*lite_index_sort_size=*/8),
-            filesystem(), icing_filesystem()));
+            filesystem(), icing_filesystem(), feature_flags_.get()));
     DocumentId original_last_added_doc_id = index->last_added_document_id();
     index->set_last_added_document_id(original_last_added_doc_id + 1);
     Index::Editor editor =
@@ -3474,7 +3610,7 @@
                            /*index_merge_size=*/message.ByteSizeLong(),
                            /*lite_index_sort_at_indexing=*/true,
                            /*lite_index_sort_size=*/8),
-            filesystem(), icing_filesystem()));
+            filesystem(), icing_filesystem(), feature_flags_.get()));
     ICING_ASSERT_OK_AND_ASSIGN(
         std::unique_ptr<DocHitInfoIterator> doc_hit_info_iter,
         index->GetIterator("foo", /*term_start_index=*/0,
@@ -3536,7 +3672,7 @@
             Index::Options(GetIndexDir(), /*index_merge_size=*/100,
                            /*lite_index_sort_at_indexing=*/true,
                            /*lite_index_sort_size=*/50),
-            filesystem(), icing_filesystem()));
+            filesystem(), icing_filesystem(), feature_flags_.get()));
     // Add hits for document 0 and merge.
     ASSERT_THAT(index->last_added_document_id(), kInvalidDocumentId);
     index->set_last_added_document_id(0);
@@ -3613,7 +3749,7 @@
         Index::Create(Index::Options(GetIndexDir(), /*index_merge_size=*/100,
                                      /*lite_index_sort_at_indexing=*/true,
                                      /*lite_index_sort_size=*/50),
-                      filesystem(), icing_filesystem()));
+                      filesystem(), icing_filesystem(), feature_flags_.get()));
     ICING_ASSERT_OK_AND_ASSIGN(
         std::unique_ptr<DocHitInfoIterator> doc_hit_info_iter,
         index->GetIterator("foo", /*term_start_index=*/0,
@@ -3731,7 +3867,7 @@
                            /*index_merge_size=*/message.ByteSizeLong(),
                            /*lite_index_sort_at_indexing=*/true,
                            /*lite_index_sort_size=*/8),
-            filesystem(), icing_filesystem()));
+            filesystem(), icing_filesystem(), feature_flags_.get()));
     // Add hits for document 4 and merge.
     DocumentId original_last_added_doc_id = index->last_added_document_id();
     index->set_last_added_document_id(original_last_added_doc_id + 1);
@@ -3874,7 +4010,7 @@
         Index::Create(Index::Options(GetIndexDir(), /*index_merge_size=*/100,
                                      /*lite_index_sort_at_indexing=*/true,
                                      /*lite_index_sort_size=*/50),
-                      filesystem(), icing_filesystem()));
+                      filesystem(), icing_filesystem(), feature_flags_.get()));
     ICING_ASSERT_OK_AND_ASSIGN(
         std::unique_ptr<DocHitInfoIterator> doc_hit_info_iter,
         index->GetIterator("foo", /*term_start_index=*/0,
@@ -4378,7 +4514,8 @@
         QualifiedIdJoinIndexImplV3::Create(
             filesystem, GetQualifiedIdJoinIndexDir(), *feature_flags_));
     EXPECT_THAT(qualified_id_join_index, Pointee(IsEmpty()));
-    EXPECT_THAT(qualified_id_join_index->Get(/*parent_document_id=*/0),
+    EXPECT_THAT(qualified_id_join_index->GetDocumentJoinIdPairArrayView(
+                    /*parent_document_id=*/0),
                 IsOkAndHolds(IsEmpty()));
   }
 }
@@ -4669,6 +4806,11 @@
     EXPECT_THAT(init_result.initialize_stats()
                     .qualified_id_join_index_restoration_cause(),
                 Eq(InitializeStatsProto::NONE));
+    EXPECT_THAT(
+        init_result.initialize_stats().embedding_index_restoration_cause(),
+        Eq(InitializeStatsProto::NONE));
+    // PersistToDisk is not needed.
+    EXPECT_THAT(init_result.needs_persist_type(), Eq(PersistType::UNKNOWN));
   }
 }
 
@@ -4735,10 +4877,505 @@
     EXPECT_THAT(init_result.initialize_stats()
                     .qualified_id_join_index_restoration_cause(),
                 Eq(InitializeStatsProto::NONE));
+    EXPECT_THAT(
+        init_result.initialize_stats().embedding_index_restoration_cause(),
+        Eq(InitializeStatsProto::NONE));
+    // PersistToDisk is not needed.
+    EXPECT_THAT(init_result.needs_persist_type(), Eq(PersistType::UNKNOWN));
   }
 }
 
 TEST_F(IcingSearchEngineInitializationTest,
+       InitializeShouldPurgeDocumentsThatExpiredDuringOffline) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("name")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_REQUIRED)))
+          .AddType(
+              SchemaTypeConfigBuilder()
+                  .SetType("Message")
+                  .AddProperty(
+                      PropertyConfigBuilder()
+                          .SetName("body")
+                          .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                          .SetCardinality(CARDINALITY_REQUIRED))
+                  .AddProperty(PropertyConfigBuilder()
+                                   .SetName("sender")
+                                   .SetDataTypeJoinableString(
+                                       JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                       DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                                   .SetCardinality(CARDINALITY_REQUIRED)))
+          .Build();
+
+  DocumentProto person = DocumentBuilder()
+                             .SetKey("namespace", "person")
+                             .SetSchema("Person")
+                             .AddStringProperty("name", "person")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(10000)  // Expire at 10010 ms.
+                             .Build();
+  DocumentProto message = DocumentBuilder()
+                              .SetKey("namespace", "message/1")
+                              .SetSchema("Message")
+                              .AddStringProperty("body", "message body")
+                              .AddStringProperty("sender", "namespace#person")
+                              .SetCreationTimestampMs(10)
+                              .SetTtlMs(0)  // Never expire.
+                              .Build();
+
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  options.set_enable_qualified_id_join_index_v3(true);
+  options.set_enable_soft_index_restoration(true);
+  options.set_enable_delete_propagation_from(true);
+  options.set_expired_document_purge_threshold_ms(0);
+
+  {
+    // Initialize Icing and put person and message documents. Destruct Icing.
+    TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                                std::make_unique<IcingFilesystem>(),
+                                std::make_unique<FakeClock>(),
+                                GetTestJniCache());
+    ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+    ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(person).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(message).status(), ProtoIsOk());
+  }
+
+  // Initialize Icing again with a fake clock and t = 20000 ms. Initialization
+  // should purge the person document (since it expired at t = 10010 ms) and
+  // apply delete propagation to the message document.
+  auto fake_clock = std::make_unique<FakeClock>();
+  fake_clock->SetSystemTimeMilliseconds(20000);
+  TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                              std::make_unique<IcingFilesystem>(),
+                              std::move(fake_clock), GetTestJniCache());
+  InitializeResultProto initialize_result = icing.Initialize();
+  EXPECT_THAT(initialize_result.status(), ProtoIsOk());
+
+  // Verify that person and message documents are purged.
+  // - At t = 20000 ms, the person document is expired, but DocumentFilterData
+  //   has already filtered it out, so we're unsure whether it is purged or not.
+  // - But the message document never expires and DocumentFilterData doesn't
+  //   filter it out. If Initialize() had not handled expired documents
+  //   correctly, then expiration propagation would have not been triggered and
+  //   the message document would've been alive. Therefore, we can verify it by
+  //   checking the message document.
+  GetResultProto expected_get_result_proto1;
+  expected_get_result_proto1.mutable_status()->set_code(StatusProto::NOT_FOUND);
+  expected_get_result_proto1.mutable_status()->set_message(
+      "Document (namespace, person) not found.");
+  EXPECT_THAT(
+      icing.Get("namespace", "person", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto1));
+
+  GetResultProto expected_get_result_google::protobuf;
+  expected_get_result_google::protobuf.mutable_status()->set_code(StatusProto::NOT_FOUND);
+  expected_get_result_google::protobuf.mutable_status()->set_message(
+      "Document (namespace, message/1) not found.");
+  EXPECT_THAT(icing.Get("namespace", "message/1",
+                        GetResultSpecProto::default_instance()),
+              EqualsProto(expected_get_result_google::protobuf));
+}
+
+TEST_F(
+    IcingSearchEngineInitializationTest,
+    Initialize_taskSchedulerEnabled_shouldSchedulePurgingExpiredDocumentsTask) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("name")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_REQUIRED)))
+          .AddType(
+              SchemaTypeConfigBuilder()
+                  .SetType("Message")
+                  .AddProperty(
+                      PropertyConfigBuilder()
+                          .SetName("body")
+                          .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                          .SetCardinality(CARDINALITY_REQUIRED))
+                  .AddProperty(PropertyConfigBuilder()
+                                   .SetName("sender")
+                                   .SetDataTypeJoinableString(
+                                       JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                       DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                                   .SetCardinality(CARDINALITY_REQUIRED)))
+          .Build();
+
+  DocumentProto person1 = DocumentBuilder()
+                              .SetKey("namespace", "person1")
+                              .SetSchema("Person")
+                              .AddStringProperty("name", "person")
+                              .SetCreationTimestampMs(10)
+                              .SetTtlMs(10000)  // Expire at 10010 ms.
+                              .Build();
+  DocumentProto person2 = DocumentBuilder()
+                              .SetKey("namespace", "person2")
+                              .SetSchema("Person")
+                              .AddStringProperty("name", "person")
+                              .SetCreationTimestampMs(10)
+                              .SetTtlMs(21000)  // Expire at 21010 ms.
+                              .Build();
+  DocumentProto message1 = DocumentBuilder()
+                               .SetKey("namespace", "message/1")
+                               .SetSchema("Message")
+                               .AddStringProperty("body", "message body")
+                               .AddStringProperty("sender", "namespace#person1")
+                               .SetCreationTimestampMs(10)
+                               .SetTtlMs(0)  // Never expire.
+                               .Build();
+  DocumentProto message2 = DocumentBuilder()
+                               .SetKey("namespace", "message/2")
+                               .SetSchema("Message")
+                               .AddStringProperty("body", "message body")
+                               .AddStringProperty("sender", "namespace#person2")
+                               .SetCreationTimestampMs(10)
+                               .SetTtlMs(0)  // Never expire.
+                               .Build();
+
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  options.set_enable_background_task_scheduler(true);
+  options.set_enable_qualified_id_join_index_v3(true);
+  options.set_enable_soft_index_restoration(true);
+  options.set_enable_delete_propagation_from(true);
+  options.set_expired_document_purge_threshold_ms(0);
+
+  {
+    // Initialize Icing and put all documents. Destruct Icing.
+    TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                                std::make_unique<IcingFilesystem>(),
+                                std::make_unique<FakeClock>(),
+                                GetTestJniCache());
+    ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+    ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(person1).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(person2).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(message1).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(message2).status(), ProtoIsOk());
+  }
+
+  // Initialize Icing again with a fake clock and t = 20000 ms. Initialization
+  // should purge person1 (since it expired at t = 10010 ms) and apply delete
+  // propagation to message1. Also it should schedule the puring expired
+  // documents task for the next expiration event.
+  auto fake_clock = std::make_unique<FakeClock>();
+  FakeClock* fake_clock_ptr = fake_clock.get();
+  fake_clock->SetSystemTimeMilliseconds(20000);
+  TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                              std::make_unique<IcingFilesystem>(),
+                              std::move(fake_clock), GetTestJniCache());
+  InitializeResultProto initialize_result = icing.Initialize();
+  EXPECT_THAT(initialize_result.status(), ProtoIsOk());
+  // Next expiration timestamp should be set correctly.
+  EXPECT_THAT(
+      initialize_result.initialize_stats().next_expiration_timestamp_ms(),
+      Eq(21010));
+
+  // Sanity check: person2 and message2 are still alive.
+  GetResultProto expected_get_result_proto1;
+  expected_get_result_proto1.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_proto1.mutable_document() = person2;
+  EXPECT_THAT(
+      icing.Get("namespace", "person2", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto1));
+
+  GetResultProto expected_get_result_google::protobuf;
+  expected_get_result_google::protobuf.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_google::protobuf.mutable_document() = message2;
+  EXPECT_THAT(icing.Get("namespace", "message/2",
+                        GetResultSpecProto::default_instance()),
+              EqualsProto(expected_get_result_google::protobuf));
+
+  // Adjust the clock to 21010 ms and sleep for 1010 ms. The purging expiration
+  // task should execute and purge person2 and message2.
+  // - At t = 21010 ms, person2 is expired, but DocumentFilterData has already
+  //   filtered it out, so we're unsure whether it is purged or not.
+  // - But message2 never expires and DocumentFilterData doesn't filter it out.
+  //   If Initialize() had not scheduled the purging expired documents task
+  //   correctly for the next expiration event (i.e. for person2), then
+  //   expiration propagation would have not been triggered and message2
+  //   would've been alive. Therefore, we can verify it by checking message2
+  fake_clock_ptr->SetSystemTimeMilliseconds(21010);
+  std::this_thread::sleep_for(std::chrono::milliseconds(1010));
+
+  GetResultProto expected_get_result_proto3;
+  expected_get_result_proto3.mutable_status()->set_code(StatusProto::NOT_FOUND);
+  expected_get_result_proto3.mutable_status()->set_message(
+      "Document (namespace, person2) not found.");
+  EXPECT_THAT(
+      icing.Get("namespace", "person2", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto3));
+
+  GetResultProto expected_get_result_proto4;
+  expected_get_result_proto4.mutable_status()->set_code(StatusProto::NOT_FOUND);
+  expected_get_result_proto4.mutable_status()->set_message(
+      "Document (namespace, message/2) not found.");
+  EXPECT_THAT(icing.Get("namespace", "message/2",
+                        GetResultSpecProto::default_instance()),
+              EqualsProto(expected_get_result_proto4));
+}
+
+TEST_F(
+    IcingSearchEngineInitializationTest,
+    Initialize_taskSchedulerDisabled_shouldPurgeExpiredDocumentsButNoRescheduling) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("name")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_REQUIRED)))
+          .AddType(
+              SchemaTypeConfigBuilder()
+                  .SetType("Message")
+                  .AddProperty(
+                      PropertyConfigBuilder()
+                          .SetName("body")
+                          .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                          .SetCardinality(CARDINALITY_REQUIRED))
+                  .AddProperty(PropertyConfigBuilder()
+                                   .SetName("sender")
+                                   .SetDataTypeJoinableString(
+                                       JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                       DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                                   .SetCardinality(CARDINALITY_REQUIRED)))
+          .Build();
+
+  DocumentProto person1 = DocumentBuilder()
+                              .SetKey("namespace", "person1")
+                              .SetSchema("Person")
+                              .AddStringProperty("name", "person")
+                              .SetCreationTimestampMs(10)
+                              .SetTtlMs(10000)  // Expire at 10010 ms.
+                              .Build();
+  DocumentProto person2 = DocumentBuilder()
+                              .SetKey("namespace", "person2")
+                              .SetSchema("Person")
+                              .AddStringProperty("name", "person")
+                              .SetCreationTimestampMs(10)
+                              .SetTtlMs(21000)  // Expire at 21010 ms.
+                              .Build();
+  DocumentProto message1 = DocumentBuilder()
+                               .SetKey("namespace", "message/1")
+                               .SetSchema("Message")
+                               .AddStringProperty("body", "message body")
+                               .AddStringProperty("sender", "namespace#person1")
+                               .SetCreationTimestampMs(10)
+                               .SetTtlMs(0)  // Never expire.
+                               .Build();
+  DocumentProto message2 = DocumentBuilder()
+                               .SetKey("namespace", "message/2")
+                               .SetSchema("Message")
+                               .AddStringProperty("body", "message body")
+                               .AddStringProperty("sender", "namespace#person2")
+                               .SetCreationTimestampMs(10)
+                               .SetTtlMs(0)  // Never expire.
+                               .Build();
+
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  options.set_enable_background_task_scheduler(false);
+  options.set_enable_qualified_id_join_index_v3(true);
+  options.set_enable_soft_index_restoration(true);
+  options.set_enable_delete_propagation_from(true);
+  options.set_expired_document_purge_threshold_ms(0);
+
+  {
+    // Initialize Icing and put all documents. Destruct Icing.
+    TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                                std::make_unique<IcingFilesystem>(),
+                                std::make_unique<FakeClock>(),
+                                GetTestJniCache());
+    ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+    ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(person1).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(person2).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(message1).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(message2).status(), ProtoIsOk());
+  }
+
+  // Initialize Icing again with a fake clock and t = 20000 ms.
+  // - Initialization should purge person1 (since it expired at t = 10010 ms)
+  //   and apply delete propagation to message1.
+  // - Since the task scheduler is disabled, it should NOT schedule the puring
+  //   expired documents task for the next expiration event.
+  auto fake_clock = std::make_unique<FakeClock>();
+  FakeClock* fake_clock_ptr = fake_clock.get();
+  fake_clock->SetSystemTimeMilliseconds(20000);
+  TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                              std::make_unique<IcingFilesystem>(),
+                              std::move(fake_clock), GetTestJniCache());
+  InitializeResultProto initialize_result = icing.Initialize();
+  EXPECT_THAT(initialize_result.status(), ProtoIsOk());
+  // Next expiration timestamp should be set correctly.
+  EXPECT_THAT(
+      initialize_result.initialize_stats().next_expiration_timestamp_ms(),
+      Eq(21010));
+
+  // Sanity check: person2 and message2 are still alive.
+  GetResultProto expected_get_result_proto1;
+  expected_get_result_proto1.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_proto1.mutable_document() = person2;
+  EXPECT_THAT(
+      icing.Get("namespace", "person2", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto1));
+
+  GetResultProto expected_get_result_google::protobuf;
+  expected_get_result_google::protobuf.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_google::protobuf.mutable_document() = message2;
+  EXPECT_THAT(icing.Get("namespace", "message/2",
+                        GetResultSpecProto::default_instance()),
+              EqualsProto(expected_get_result_google::protobuf));
+
+  // Adjust the clock to 21010 ms and sleep for 1010 ms. message2 should still
+  // be present.
+  fake_clock_ptr->SetSystemTimeMilliseconds(21010);
+  std::this_thread::sleep_for(std::chrono::milliseconds(1010));
+  EXPECT_THAT(icing.Get("namespace", "message/2",
+                        GetResultSpecProto::default_instance()),
+              EqualsProto(expected_get_result_google::protobuf));
+}
+
+TEST_F(IcingSearchEngineInitializationTest,
+       IndexRecovery_failingDocumentShouldPropagateDeletionToChildDocuments) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("Person")
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("name")
+                                        .SetDataTypeString(TERM_MATCH_PREFIX,
+                                                           TOKENIZER_PLAIN)
+                                        .SetCardinality(CARDINALITY_REQUIRED))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("timestamp")
+                                        .SetDataTypeInt64(NUMERIC_MATCH_RANGE)
+                                        .SetCardinality(CARDINALITY_REQUIRED)))
+          .AddType(
+              SchemaTypeConfigBuilder()
+                  .SetType("Message")
+                  .AddProperty(
+                      PropertyConfigBuilder()
+                          .SetName("body")
+                          .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                          .SetCardinality(CARDINALITY_REQUIRED))
+                  .AddProperty(PropertyConfigBuilder()
+                                   .SetName("sender")
+                                   .SetDataTypeJoinableString(
+                                       JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                       DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                                   .SetCardinality(CARDINALITY_REQUIRED)))
+          .Build();
+
+  DocumentProto person = DocumentBuilder()
+                             .SetKey("namespace", "person1")
+                             .SetSchema("Person")
+                             .AddStringProperty("name", "person")
+                             .AddInt64Property("timestamp", 1024)
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(10000)  // Expire at 10010 ms.
+                             .Build();
+  DocumentProto message = DocumentBuilder()
+                              .SetKey("namespace", "message/1")
+                              .SetSchema("Message")
+                              .AddStringProperty("body", "message body")
+                              .AddStringProperty("sender", "namespace#person1")
+                              .SetCreationTimestampMs(10)
+                              .SetTtlMs(0)  // Never expire.
+                              .Build();
+
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  options.set_enable_qualified_id_join_index_v3(true);
+  options.set_enable_soft_index_restoration(true);
+  options.set_enable_delete_propagation_from(true);
+  options.set_expired_document_purge_threshold_ms(0);
+
+  {
+    // Initialize Icing and put all documents. Destruct Icing.
+    TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                                std::make_unique<IcingFilesystem>(),
+                                std::make_unique<FakeClock>(),
+                                GetTestJniCache());
+    ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+    ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(person).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(message).status(), ProtoIsOk());
+  }
+
+  // 2. Delete integer index to trigger index restoration.
+  ASSERT_TRUE(
+      filesystem()->DeleteDirectoryRecursively(GetIntegerIndexDir().c_str()));
+
+  // 3. Mock filesystem to fail creating "timestamp" integer index storage.
+  auto mock_filesystem = std::make_unique<MockFilesystem>();
+  ON_CALL(*mock_filesystem,
+          CreateDirectory(HasSubstr(GetIntegerIndexDir() + "/timestamp")))
+      .WillByDefault(Return(false));
+
+  // 4. Initialize IcingSearchEngine again with the mock filesystem. When
+  //    indexing document "person1", it will fail to create "timestamp" integer
+  //    index storage, but soft index restoration mechanism should skip the
+  //    error and delete the document without failing initialization.
+  auto fake_clock = std::make_unique<FakeClock>();
+  fake_clock->SetTimerElapsedMilliseconds(10);
+  TestIcingSearchEngine icing(options, std::move(mock_filesystem),
+                              std::make_unique<IcingFilesystem>(),
+                              std::move(fake_clock), GetTestJniCache());
+
+  InitializeResultProto initialize_result = icing.Initialize();
+  EXPECT_THAT(initialize_result.status(),
+              ProtoStatusIs(StatusProto::WARNING_DATA_LOSS));
+
+  EXPECT_THAT(
+      initialize_result.initialize_stats().document_store_recovery_cause(),
+      Eq(InitializeStatsProto::NONE));
+  // Indices should be restored. "person1" should fail to be reindexed.
+  EXPECT_THAT(
+      initialize_result.initialize_stats().index_restoration_latency_ms(),
+      Eq(10));
+  EXPECT_THAT(
+      initialize_result.initialize_stats().num_failed_reindexed_documents(),
+      Eq(1));
+  EXPECT_THAT(initialize_result.initialize_stats().index_restoration_cause(),
+              Eq(InitializeStatsProto::NONE));
+  EXPECT_THAT(
+      initialize_result.initialize_stats().integer_index_restoration_cause(),
+      Eq(InitializeStatsProto::INCONSISTENT_WITH_GROUND_TRUTH));
+  EXPECT_THAT(initialize_result.initialize_stats()
+                  .qualified_id_join_index_restoration_cause(),
+              Eq(InitializeStatsProto::NONE));
+  EXPECT_THAT(
+      initialize_result.initialize_stats().embedding_index_restoration_cause(),
+      Eq(InitializeStatsProto::NONE));
+  // needs_persist_type should be set to RECOVERY_PROOF.
+  EXPECT_THAT(initialize_result.needs_persist_type(),
+              Eq(PersistType::RECOVERY_PROOF));
+
+  // ("namespace", "person1") should be deleted.
+  GetResultProto expected_get_result_proto1;
+  expected_get_result_proto1.mutable_status()->set_code(StatusProto::NOT_FOUND);
+  expected_get_result_proto1.mutable_status()->set_message(
+      "Document (namespace, person1) not found.");
+  EXPECT_THAT(
+      icing.Get("namespace", "person1", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto1));
+
+  // ("namespace", "message/1") should be deleted. Although this document should
+  // successfully be reindexed, when "person1" fails to be reindexed, all
+  // of its child documents should be deleted.
+  GetResultProto expected_get_result_google::protobuf;
+  expected_get_result_google::protobuf.mutable_status()->set_code(StatusProto::NOT_FOUND);
+  expected_get_result_google::protobuf.mutable_status()->set_message(
+      "Document (namespace, message/1) not found.");
+  EXPECT_THAT(icing.Get("namespace", "message/1",
+                        GetResultSpecProto::default_instance()),
+              EqualsProto(expected_get_result_google::protobuf));
+}
+
+TEST_F(IcingSearchEngineInitializationTest,
        InitializeShouldLogFunctionLatency) {
   auto fake_clock = std::make_unique<FakeClock>();
   fake_clock->SetTimerElapsedMilliseconds(10);
@@ -4810,6 +5447,12 @@
                               std::move(fake_clock), GetTestJniCache());
   InitializeResultProto initialize_result_proto = icing.Initialize();
   EXPECT_THAT(initialize_result_proto.status(), ProtoIsOk());
+  EXPECT_THAT(
+      initialize_result_proto.initialize_stats().schema_store_recovery_cause(),
+      Eq(InitializeStatsProto::NONE));
+  EXPECT_THAT(initialize_result_proto.initialize_stats()
+                  .schema_store_recovery_latency_ms(),
+              Eq(0));
   EXPECT_THAT(initialize_result_proto.initialize_stats()
                   .document_store_recovery_cause(),
               Eq(InitializeStatsProto::NONE));
@@ -4828,15 +5471,59 @@
   EXPECT_THAT(initialize_result_proto.initialize_stats()
                   .qualified_id_join_index_restoration_cause(),
               Eq(InitializeStatsProto::NONE));
+  EXPECT_THAT(initialize_result_proto.initialize_stats()
+                  .embedding_index_restoration_cause(),
+              Eq(InitializeStatsProto::NONE));
   EXPECT_THAT(
       initialize_result_proto.initialize_stats().index_restoration_latency_ms(),
       Eq(0));
+  // PersistToDisk is not needed.
+  EXPECT_THAT(initialize_result_proto.needs_persist_type(),
+              Eq(PersistType::UNKNOWN));
+
+  // Initialize another Icing instance with the "existing empty database". Even
+  // though we never call PersistToDisk(), there should be no recovery.
+  auto fake_clock2 = std::make_unique<FakeClock>();
+  fake_clock2->SetTimerElapsedMilliseconds(10);
+  TestIcingSearchEngine another_icing(
+      GetDefaultIcingOptions(), std::make_unique<Filesystem>(),
+      std::make_unique<IcingFilesystem>(), std::move(fake_clock2),
+      GetTestJniCache());
+  InitializeResultProto initialize_result_google::protobuf = another_icing.Initialize();
+  EXPECT_THAT(initialize_result_google::protobuf.status(), ProtoIsOk());
   EXPECT_THAT(
-      initialize_result_proto.initialize_stats().schema_store_recovery_cause(),
+      initialize_result_google::protobuf.initialize_stats().schema_store_recovery_cause(),
       Eq(InitializeStatsProto::NONE));
-  EXPECT_THAT(initialize_result_proto.initialize_stats()
+  EXPECT_THAT(initialize_result_google::protobuf.initialize_stats()
                   .schema_store_recovery_latency_ms(),
               Eq(0));
+  EXPECT_THAT(initialize_result_google::protobuf.initialize_stats()
+                  .document_store_recovery_cause(),
+              Eq(InitializeStatsProto::NONE));
+  EXPECT_THAT(initialize_result_google::protobuf.initialize_stats()
+                  .document_store_recovery_latency_ms(),
+              Eq(0));
+  EXPECT_THAT(
+      initialize_result_google::protobuf.initialize_stats().document_store_data_status(),
+      Eq(InitializeStatsProto::NO_DATA_LOSS));
+  EXPECT_THAT(
+      initialize_result_google::protobuf.initialize_stats().index_restoration_cause(),
+      Eq(InitializeStatsProto::NONE));
+  EXPECT_THAT(initialize_result_google::protobuf.initialize_stats()
+                  .integer_index_restoration_cause(),
+              Eq(InitializeStatsProto::NONE));
+  EXPECT_THAT(initialize_result_google::protobuf.initialize_stats()
+                  .qualified_id_join_index_restoration_cause(),
+              Eq(InitializeStatsProto::NONE));
+  EXPECT_THAT(initialize_result_google::protobuf.initialize_stats()
+                  .embedding_index_restoration_cause(),
+              Eq(InitializeStatsProto::NONE));
+  EXPECT_THAT(initialize_result_google::protobuf.initialize_stats()
+                  .index_restoration_latency_ms(),
+              Eq(0));
+  // PersistToDisk is not needed.
+  EXPECT_THAT(initialize_result_google::protobuf.needs_persist_type(),
+              Eq(PersistType::UNKNOWN));
 }
 
 TEST_F(IcingSearchEngineInitializationTest,
@@ -4903,6 +5590,9 @@
                     .qualified_id_join_index_restoration_cause(),
                 Eq(InitializeStatsProto::NONE));
     EXPECT_THAT(initialize_result_proto.initialize_stats()
+                    .embedding_index_restoration_cause(),
+                Eq(InitializeStatsProto::NONE));
+    EXPECT_THAT(initialize_result_proto.initialize_stats()
                     .index_restoration_latency_ms(),
                 Eq(0));
     EXPECT_THAT(initialize_result_proto.initialize_stats()
@@ -4911,6 +5601,10 @@
     EXPECT_THAT(initialize_result_proto.initialize_stats()
                     .schema_store_recovery_latency_ms(),
                 Eq(0));
+
+    // needs_persist_type should be set to RECOVERY_PROOF.
+    EXPECT_THAT(initialize_result_proto.needs_persist_type(),
+                Eq(PersistType::RECOVERY_PROOF));
   }
 }
 
@@ -4960,7 +5654,7 @@
 
     // Set dirty bit to true to reflect that something changed in the log.
     header.SetDirtyFlag(true);
-    header.SetHeaderChecksum(header.CalculateHeaderChecksum());
+    header.UpdateHeaderChecksums(/*enable_new_header_format=*/true);
 
     WriteDocumentLogHeader(*filesystem(), document_log_file, header);
   }
@@ -4997,6 +5691,9 @@
                     .qualified_id_join_index_restoration_cause(),
                 Eq(InitializeStatsProto::NONE));
     EXPECT_THAT(initialize_result_proto.initialize_stats()
+                    .embedding_index_restoration_cause(),
+                Eq(InitializeStatsProto::NONE));
+    EXPECT_THAT(initialize_result_proto.initialize_stats()
                     .index_restoration_latency_ms(),
                 Eq(0));
     EXPECT_THAT(initialize_result_proto.initialize_stats()
@@ -5005,6 +5702,10 @@
     EXPECT_THAT(initialize_result_proto.initialize_stats()
                     .schema_store_recovery_latency_ms(),
                 Eq(0));
+
+    // needs_persist_type should be set to RECOVERY_PROOF.
+    EXPECT_THAT(initialize_result_proto.needs_persist_type(),
+                Eq(PersistType::RECOVERY_PROOF));
   }
 }
 
@@ -5035,7 +5736,7 @@
                                      /*index_merge_size=*/100,
                                      /*lite_index_sort_at_indexing=*/true,
                                      /*lite_index_sort_size=*/50),
-                      filesystem(), icing_filesystem()));
+                      filesystem(), icing_filesystem(), feature_flags_.get()));
     ICING_ASSERT_OK(index->PersistToDisk());
   }
 
@@ -5060,6 +5761,9 @@
                     .qualified_id_join_index_restoration_cause(),
                 Eq(InitializeStatsProto::NONE));
     EXPECT_THAT(initialize_result_proto.initialize_stats()
+                    .embedding_index_restoration_cause(),
+                Eq(InitializeStatsProto::NONE));
+    EXPECT_THAT(initialize_result_proto.initialize_stats()
                     .index_restoration_latency_ms(),
                 Eq(10));
     EXPECT_THAT(initialize_result_proto.initialize_stats()
@@ -5077,6 +5781,10 @@
     EXPECT_THAT(initialize_result_proto.initialize_stats()
                     .schema_store_recovery_latency_ms(),
                 Eq(0));
+
+    // needs_persist_type should be set to RECOVERY_PROOF.
+    EXPECT_THAT(initialize_result_proto.needs_persist_type(),
+                Eq(PersistType::RECOVERY_PROOF));
   }
 }
 
@@ -5124,6 +5832,9 @@
                     .qualified_id_join_index_restoration_cause(),
                 Eq(InitializeStatsProto::NONE));
     EXPECT_THAT(initialize_result_proto.initialize_stats()
+                    .embedding_index_restoration_cause(),
+                Eq(InitializeStatsProto::NONE));
+    EXPECT_THAT(initialize_result_proto.initialize_stats()
                     .index_restoration_latency_ms(),
                 Eq(10));
     EXPECT_THAT(initialize_result_proto.initialize_stats()
@@ -5141,6 +5852,10 @@
     EXPECT_THAT(initialize_result_proto.initialize_stats()
                     .schema_store_recovery_latency_ms(),
                 Eq(0));
+
+    // needs_persist_type should be set to RECOVERY_PROOF.
+    EXPECT_THAT(initialize_result_proto.needs_persist_type(),
+                Eq(PersistType::RECOVERY_PROOF));
   }
 }
 
@@ -5226,6 +5941,9 @@
                     .qualified_id_join_index_restoration_cause(),
                 Eq(InitializeStatsProto::INCONSISTENT_WITH_GROUND_TRUTH));
     EXPECT_THAT(initialize_result_proto.initialize_stats()
+                    .embedding_index_restoration_cause(),
+                Eq(InitializeStatsProto::NONE));
+    EXPECT_THAT(initialize_result_proto.initialize_stats()
                     .index_restoration_latency_ms(),
                 Eq(10));
     EXPECT_THAT(initialize_result_proto.initialize_stats()
@@ -5243,6 +5961,10 @@
     EXPECT_THAT(initialize_result_proto.initialize_stats()
                     .schema_store_recovery_latency_ms(),
                 Eq(0));
+
+    // needs_persist_type should be set to RECOVERY_PROOF.
+    EXPECT_THAT(initialize_result_proto.needs_persist_type(),
+                Eq(PersistType::RECOVERY_PROOF));
   }
 }
 
@@ -5296,7 +6018,7 @@
   }
 
   {
-    // Both document store and index should be recovered from checksum mismatch.
+    // All derived files should be rebuilt.
     auto fake_clock = std::make_unique<FakeClock>();
     fake_clock->SetTimerElapsedMilliseconds(10);
     TestIcingSearchEngine icing(GetDefaultIcingOptions(),
@@ -5347,10 +6069,14 @@
     EXPECT_THAT(initialize_result_proto.initialize_stats()
                     .embedding_index_restoration_cause(),
                 Eq(InitializeStatsProto::SCHEMA_CHANGES_OUT_OF_SYNC));
+
+    // needs_persist_type should be set to RECOVERY_PROOF.
+    EXPECT_THAT(initialize_result_proto.needs_persist_type(),
+                Eq(PersistType::RECOVERY_PROOF));
   }
 
   {
-    // No recovery should be needed.
+    // Initialize again. No recovery should be needed.
     auto fake_clock = std::make_unique<FakeClock>();
     fake_clock->SetTimerElapsedMilliseconds(10);
     TestIcingSearchEngine icing(GetDefaultIcingOptions(),
@@ -5401,6 +6127,10 @@
     EXPECT_THAT(initialize_result_proto.initialize_stats()
                     .embedding_index_restoration_cause(),
                 Eq(InitializeStatsProto::NONE));
+
+    // PersistToDisk is not needed.
+    EXPECT_THAT(initialize_result_proto.needs_persist_type(),
+                Eq(PersistType::UNKNOWN));
   }
 }
 
@@ -5443,14 +6173,18 @@
     std::string doc_store_dir = GetDocumentDir();
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
-        DocumentStore::Create(filesystem(), doc_store_dir, fake_clock.get(),
-                              schema_store.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            filesystem(), doc_store_dir, fake_clock.get(), schema_store.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     std::unique_ptr<DocumentStore> document_store =
         std::move(create_result.document_store);
 
@@ -5529,6 +6263,10 @@
     EXPECT_THAT(initialize_result_proto.initialize_stats()
                     .embedding_index_restoration_cause(),
                 Eq(InitializeStatsProto::OPTIMIZE_OUT_OF_SYNC));
+
+    // needs_persist_type should be set to RECOVERY_PROOF.
+    EXPECT_THAT(initialize_result_proto.needs_persist_type(),
+                Eq(PersistType::RECOVERY_PROOF));
   }
 
   {
@@ -5583,6 +6321,10 @@
     EXPECT_THAT(initialize_result_proto.initialize_stats()
                     .embedding_index_restoration_cause(),
                 Eq(InitializeStatsProto::NONE));
+
+    // PersistToDisk is not needed.
+    EXPECT_THAT(initialize_result_proto.needs_persist_type(),
+                Eq(PersistType::UNKNOWN));
   }
 }
 
@@ -5631,6 +6373,9 @@
   EXPECT_THAT(initialize_result_proto.initialize_stats()
                   .qualified_id_join_index_restoration_cause(),
               Eq(InitializeStatsProto::NONE));
+  EXPECT_THAT(initialize_result_proto.initialize_stats()
+                  .embedding_index_restoration_cause(),
+              Eq(InitializeStatsProto::NONE));
   EXPECT_THAT(
       initialize_result_proto.initialize_stats().index_restoration_latency_ms(),
       Eq(10));
@@ -5649,6 +6394,10 @@
   EXPECT_THAT(initialize_result_proto.initialize_stats()
                   .schema_store_recovery_latency_ms(),
               Eq(0));
+
+  // needs_persist_type should be set to RECOVERY_PROOF.
+  EXPECT_THAT(initialize_result_proto.needs_persist_type(),
+              Eq(PersistType::RECOVERY_PROOF));
 }
 
 TEST_F(IcingSearchEngineInitializationTest,
@@ -5694,6 +6443,9 @@
   EXPECT_THAT(initialize_result_proto.initialize_stats()
                   .qualified_id_join_index_restoration_cause(),
               Eq(InitializeStatsProto::NONE));
+  EXPECT_THAT(initialize_result_proto.initialize_stats()
+                  .embedding_index_restoration_cause(),
+              Eq(InitializeStatsProto::NONE));
   EXPECT_THAT(
       initialize_result_proto.initialize_stats().index_restoration_latency_ms(),
       Eq(10));
@@ -5712,6 +6464,10 @@
   EXPECT_THAT(initialize_result_proto.initialize_stats()
                   .schema_store_recovery_latency_ms(),
               Eq(0));
+
+  // needs_persist_type should be set to RECOVERY_PROOF.
+  EXPECT_THAT(initialize_result_proto.needs_persist_type(),
+              Eq(PersistType::RECOVERY_PROOF));
 }
 
 TEST_F(IcingSearchEngineInitializationTest,
@@ -5797,6 +6553,9 @@
   EXPECT_THAT(initialize_result_proto.initialize_stats()
                   .qualified_id_join_index_restoration_cause(),
               Eq(InitializeStatsProto::IO_ERROR));
+  EXPECT_THAT(initialize_result_proto.initialize_stats()
+                  .embedding_index_restoration_cause(),
+              Eq(InitializeStatsProto::NONE));
   EXPECT_THAT(
       initialize_result_proto.initialize_stats().index_restoration_latency_ms(),
       Eq(10));
@@ -5815,6 +6574,103 @@
   EXPECT_THAT(initialize_result_proto.initialize_stats()
                   .schema_store_recovery_latency_ms(),
               Eq(0));
+
+  // needs_persist_type should be set to RECOVERY_PROOF.
+  EXPECT_THAT(initialize_result_proto.needs_persist_type(),
+              Eq(PersistType::RECOVERY_PROOF));
+}
+
+TEST_F(IcingSearchEngineInitializationTest,
+       InitializeShouldLogRecoveryCauseEmbeddingIndexIOError) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("Email")
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("subject")
+                                        .SetDataTypeString(TERM_MATCH_EXACT,
+                                                           TOKENIZER_PLAIN)
+                                        .SetCardinality(CARDINALITY_REPEATED))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("embedding")
+                                        .SetDataTypeVector(
+                                            EMBEDDING_INDEXING_LINEAR_SEARCH)
+                                        .SetCardinality(CARDINALITY_REPEATED)))
+          .Build();
+
+  DocumentProto email =
+      DocumentBuilder()
+          .SetKey("namespace", "email")
+          .SetSchema("Email")
+          .AddStringProperty("subject", "test subject")
+          .AddVectorProperty(
+              "embedding",
+              CreateVector("my_model_v1", {0.1, 0.2, 0.3, 0.4, 0.5}),
+              CreateVector("my_model_v1", {1, 2, 3, 4, 5}),
+              CreateVector("my_model_v1", {0.6, 0.7, 0.8, 0.9, -1}))
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .Build();
+
+  {
+    // Initialize and put documents.
+    IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
+    ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+    ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(email).status(), ProtoIsOk());
+  }
+
+  std::string embedding_index_metadata_file =
+      absl_ports::StrCat(GetEmbeddingIndexDir(), "/metadata");
+  auto mock_filesystem = std::make_unique<MockFilesystem>();
+  EXPECT_CALL(*mock_filesystem, OpenForWrite(_)).WillRepeatedly(DoDefault());
+  // This fails EmbeddingIndex::Create() once.
+  EXPECT_CALL(*mock_filesystem, OpenForWrite(Eq(embedding_index_metadata_file)))
+      .WillOnce(Return(-1))
+      .WillRepeatedly(DoDefault());
+
+  auto fake_clock = std::make_unique<FakeClock>();
+  fake_clock->SetTimerElapsedMilliseconds(10);
+  TestIcingSearchEngine icing(GetDefaultIcingOptions(),
+                              std::move(mock_filesystem),
+                              std::make_unique<IcingFilesystem>(),
+                              std::move(fake_clock), GetTestJniCache());
+
+  InitializeResultProto initialize_result_proto = icing.Initialize();
+  EXPECT_THAT(initialize_result_proto.status(), ProtoIsOk());
+  EXPECT_THAT(
+      initialize_result_proto.initialize_stats().index_restoration_cause(),
+      Eq(InitializeStatsProto::NONE));
+  EXPECT_THAT(initialize_result_proto.initialize_stats()
+                  .integer_index_restoration_cause(),
+              Eq(InitializeStatsProto::NONE));
+  EXPECT_THAT(initialize_result_proto.initialize_stats()
+                  .qualified_id_join_index_restoration_cause(),
+              Eq(InitializeStatsProto::NONE));
+  EXPECT_THAT(initialize_result_proto.initialize_stats()
+                  .embedding_index_restoration_cause(),
+              Eq(InitializeStatsProto::IO_ERROR));
+  EXPECT_THAT(
+      initialize_result_proto.initialize_stats().index_restoration_latency_ms(),
+      Eq(10));
+  EXPECT_THAT(initialize_result_proto.initialize_stats()
+                  .document_store_recovery_cause(),
+              Eq(InitializeStatsProto::NONE));
+  EXPECT_THAT(initialize_result_proto.initialize_stats()
+                  .document_store_recovery_latency_ms(),
+              Eq(0));
+  EXPECT_THAT(
+      initialize_result_proto.initialize_stats().document_store_data_status(),
+      Eq(InitializeStatsProto::NO_DATA_LOSS));
+  EXPECT_THAT(
+      initialize_result_proto.initialize_stats().schema_store_recovery_cause(),
+      Eq(InitializeStatsProto::NONE));
+  EXPECT_THAT(initialize_result_proto.initialize_stats()
+                  .schema_store_recovery_latency_ms(),
+              Eq(0));
+
+  // needs_persist_type should be set to RECOVERY_PROOF.
+  EXPECT_THAT(initialize_result_proto.needs_persist_type(),
+              Eq(PersistType::RECOVERY_PROOF));
 }
 
 TEST_F(IcingSearchEngineInitializationTest,
@@ -5872,6 +6728,9 @@
   EXPECT_THAT(initialize_result_proto.initialize_stats()
                   .qualified_id_join_index_restoration_cause(),
               Eq(InitializeStatsProto::NONE));
+  EXPECT_THAT(initialize_result_proto.initialize_stats()
+                  .embedding_index_restoration_cause(),
+              Eq(InitializeStatsProto::NONE));
   EXPECT_THAT(
       initialize_result_proto.initialize_stats().index_restoration_latency_ms(),
       Eq(0));
@@ -5881,6 +6740,10 @@
   EXPECT_THAT(initialize_result_proto.initialize_stats()
                   .schema_store_recovery_latency_ms(),
               Eq(0));
+
+  // needs_persist_type should be set to RECOVERY_PROOF.
+  EXPECT_THAT(initialize_result_proto.needs_persist_type(),
+              Eq(PersistType::RECOVERY_PROOF));
 }
 
 TEST_F(IcingSearchEngineInitializationTest,
@@ -5933,8 +6796,15 @@
                     .qualified_id_join_index_restoration_cause(),
                 Eq(InitializeStatsProto::NONE));
     EXPECT_THAT(initialize_result_proto.initialize_stats()
+                    .embedding_index_restoration_cause(),
+                Eq(InitializeStatsProto::NONE));
+    EXPECT_THAT(initialize_result_proto.initialize_stats()
                     .index_restoration_latency_ms(),
                 Eq(0));
+
+    // needs_persist_type should be set to RECOVERY_PROOF.
+    EXPECT_THAT(initialize_result_proto.needs_persist_type(),
+                Eq(PersistType::RECOVERY_PROOF));
   }
 }
 
@@ -5977,6 +6847,379 @@
   }
 }
 
+TEST_F(IcingSearchEngineInitializationTest,
+       ProtoLogNewHeaderFormat_switchFlagFromOnToOffShouldDiscardUnsyncedTail) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("name")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_REQUIRED)))
+          .Build();
+
+  DocumentProto document0 =
+      DocumentBuilder()
+          .SetKey("namespace", "uri/0")
+          .SetSchema("Person")
+          .AddStringProperty("name", "person")
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .Build();
+  DocumentProto document1 =
+      DocumentBuilder()
+          .SetKey("namespace", "uri/1")
+          .SetSchema("Person")
+          .AddStringProperty("name", "person")
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .Build();
+
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  options.set_enable_proto_log_new_header_format(true);
+  IcingSearchEngine icing(options, GetTestJniCache());
+
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+
+  // Put document0 and persist to disk.
+  ASSERT_THAT(icing.Put(document0).status(), ProtoIsOk());
+  ASSERT_THAT(icing.PersistToDisk(PersistType::FULL).status(), ProtoIsOk());
+
+  // Put document1 and don't persist to disk.
+  ASSERT_THAT(icing.Put(document1).status(), ProtoIsOk());
+
+  // Initialize another Icing instance with the flag off.
+  IcingSearchEngineOptions another_options = GetDefaultIcingOptions();
+  another_options.set_enable_proto_log_new_header_format(false);
+  IcingSearchEngine another_icing(another_options, GetTestJniCache());
+
+  InitializeResultProto initialize_result = another_icing.Initialize();
+  // Partial data loss should be reported.
+  EXPECT_THAT(initialize_result.initialize_stats().document_store_data_status(),
+              Eq(InitializeStatsProto::PARTIAL_LOSS));
+
+  // Verify that the persisted document can be retrieved.
+  GetResultProto expected_get_result_proto0;
+  expected_get_result_proto0.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_proto0.mutable_document() = document0;
+  EXPECT_THAT(another_icing.Get("namespace", "uri/0",
+                                GetResultSpecProto::default_instance()),
+              EqualsProto(expected_get_result_proto0));
+  // Verify that document in the unsynced tail cannot be retrieved.
+  EXPECT_THAT(
+      another_icing
+          .Get("namespace", "uri/1", GetResultSpecProto::default_instance())
+          .status(),
+      ProtoStatusIs(StatusProto::NOT_FOUND));
+}
+
+TEST_F(IcingSearchEngineInitializationTest,
+       ProtoLogNewHeaderFormat_switchFlagOffToOnShouldDiscardUnsyncedTail) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("name")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_REQUIRED)))
+          .Build();
+
+  DocumentProto document0 =
+      DocumentBuilder()
+          .SetKey("namespace", "uri/0")
+          .SetSchema("Person")
+          .AddStringProperty("name", "person")
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .Build();
+  DocumentProto document1 =
+      DocumentBuilder()
+          .SetKey("namespace", "uri/1")
+          .SetSchema("Person")
+          .AddStringProperty("name", "person")
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .Build();
+
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  options.set_enable_proto_log_new_header_format(false);
+  IcingSearchEngine icing(options, GetTestJniCache());
+
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+
+  // Put document0 and persist to disk.
+  ASSERT_THAT(icing.Put(document0).status(), ProtoIsOk());
+  ASSERT_THAT(icing.PersistToDisk(PersistType::FULL).status(), ProtoIsOk());
+
+  // Put document1 and don't persist to disk.
+  ASSERT_THAT(icing.Put(document1).status(), ProtoIsOk());
+
+  // Initialize another Icing instance with the flag off.
+  IcingSearchEngineOptions another_options = GetDefaultIcingOptions();
+  another_options.set_enable_proto_log_new_header_format(true);
+  IcingSearchEngine another_icing(another_options, GetTestJniCache());
+
+  InitializeResultProto initialize_result = another_icing.Initialize();
+  // Partial data loss should be reported.
+  EXPECT_THAT(initialize_result.initialize_stats().document_store_data_status(),
+              Eq(InitializeStatsProto::PARTIAL_LOSS));
+
+  // Verify that the persisted document can be retrieved.
+  GetResultProto expected_get_result_proto0;
+  expected_get_result_proto0.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_proto0.mutable_document() = document0;
+  EXPECT_THAT(another_icing.Get("namespace", "uri/0",
+                                GetResultSpecProto::default_instance()),
+              EqualsProto(expected_get_result_proto0));
+  // Verify that document in the unsynced tail cannot be retrieved. Although we
+  // enable the flag, the existing data generated by the flag off instance
+  // cannot be recovered since the unsynced tail checksum doesn't match.
+  EXPECT_THAT(
+      another_icing
+          .Get("namespace", "uri/1", GetResultSpecProto::default_instance())
+          .status(),
+      ProtoStatusIs(StatusProto::NOT_FOUND));
+}
+
+TEST_F(IcingSearchEngineInitializationTest,
+       DeletePropagationFrom_switchFlagOffToOnShouldRevalidateDependency) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Label").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("target")
+                  .SetDataTypeJoinableString(
+                      JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                      DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  // Create 3 label documents with the following relationship:
+  // label1 -> label2 -> label3
+  DocumentProto label1 = DocumentBuilder()
+                             .SetKey("namespace", "label/1")
+                             .SetSchema("Label")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(10000)  // Expire at t = 10010 ms
+                             .Build();
+  DocumentProto label2 = DocumentBuilder()
+                             .SetKey("namespace", "label/2")
+                             .SetSchema("Label")
+                             .AddStringProperty("target", "namespace#label/1")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(0)  // Never expire.
+                             .Build();
+  DocumentProto label3 = DocumentBuilder()
+                             .SetKey("namespace", "label/3")
+                             .SetSchema("Label")
+                             .AddStringProperty("target", "namespace#label/2")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(0)  // Never expire.
+                             .Build();
+
+  {
+    IcingSearchEngineOptions options = GetDefaultIcingOptions();
+    options.set_enable_qualified_id_join_index_v3(true);
+    options.set_enable_soft_index_restoration(true);
+    options.set_enable_delete_propagation_from(false);
+    options.set_expired_document_purge_threshold_ms(0);
+
+    // Initialize Icing instance with the flag off and put all label documents.
+    auto fake_clock = std::make_unique<FakeClock>();
+    fake_clock->SetSystemTimeMilliseconds(10);
+
+    TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                                std::make_unique<IcingFilesystem>(),
+                                std::move(fake_clock), GetTestJniCache());
+    ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+    ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(label1).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(label2).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(label3).status(), ProtoIsOk());
+  }
+
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  options.set_enable_qualified_id_join_index_v3(true);
+  options.set_enable_soft_index_restoration(true);
+  options.set_enable_delete_propagation_from(true);
+  options.set_expired_document_purge_threshold_ms(0);
+
+  auto fake_clock = std::make_unique<FakeClock>();
+  fake_clock->SetSystemTimeMilliseconds(20000);
+
+  // Initialize Icing instance again with the flag on, at t = 20000 ms.
+  TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                              std::make_unique<IcingFilesystem>(),
+                              std::move(fake_clock), GetTestJniCache());
+  InitializeResultProto initialize_result = icing.Initialize();
+
+  // Join index should be rebuilt due to the flag change.
+  // - Label1 is expired at this moment.
+  // - Label2 is not expired, but delete propagation flag is enabled now and its
+  //   parent (label1) is expired, so when re-indexing it, it should fail and be
+  //   deleted.
+  // - Label3 is not expired and its parent (label2) is alive originally, so
+  //   when re-indexing it, it should succeed. But since label2 failed to be
+  //   re-indexed, Icing should delete label2 and propagate the deletion to
+  //   label3.
+  EXPECT_THAT(initialize_result.initialize_stats().document_store_data_status(),
+              Eq(InitializeStatsProto::NONE));
+  EXPECT_THAT(initialize_result.initialize_stats().index_restoration_cause(),
+              Eq(InitializeStatsProto::NONE));
+  EXPECT_THAT(
+      initialize_result.initialize_stats().integer_index_restoration_cause(),
+      Eq(InitializeStatsProto::NONE));
+  EXPECT_THAT(initialize_result.initialize_stats()
+                  .qualified_id_join_index_restoration_cause(),
+              Eq(InitializeStatsProto::FEATURE_FLAG_CHANGED));
+  EXPECT_THAT(
+      initialize_result.initialize_stats().embedding_index_restoration_cause(),
+      Eq(InitializeStatsProto::NONE));
+  EXPECT_THAT(
+      initialize_result.initialize_stats().num_failed_reindexed_documents(),
+      Eq(1));  // Label2 failed to be re-indexed.
+
+  GetResultProto expected_get_result_proto1;
+  expected_get_result_proto1.mutable_status()->set_code(StatusProto::NOT_FOUND);
+  expected_get_result_proto1.mutable_status()->set_message(
+      "Document (namespace, label/1) not found.");
+  EXPECT_THAT(
+      icing.Get("namespace", "label/1", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto1));
+
+  GetResultProto expected_get_result_google::protobuf;
+  expected_get_result_google::protobuf.mutable_status()->set_code(StatusProto::NOT_FOUND);
+  expected_get_result_google::protobuf.mutable_status()->set_message(
+      "Document (namespace, label/2) not found.");
+  EXPECT_THAT(
+      icing.Get("namespace", "label/2", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_google::protobuf));
+
+  GetResultProto expected_get_result_proto3;
+  expected_get_result_proto3.mutable_status()->set_code(StatusProto::NOT_FOUND);
+  expected_get_result_proto3.mutable_status()->set_message(
+      "Document (namespace, label/3) not found.");
+  EXPECT_THAT(
+      icing.Get("namespace", "label/3", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto3));
+}
+
+TEST_F(IcingSearchEngineInitializationTest,
+       DeletePropagationFrom_switchflagOnToOffShouldNotRevalidateDependency) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Label").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("target")
+                  .SetDataTypeJoinableString(
+                      JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                      DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  // Create 3 label documents with the following relationship:
+  // label1 -> label2 -> label3
+  DocumentProto label1 = DocumentBuilder()
+                             .SetKey("namespace", "label/1")
+                             .SetSchema("Label")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(10000)  // Expire at t = 10010 ms
+                             .Build();
+  DocumentProto label2 = DocumentBuilder()
+                             .SetKey("namespace", "label/2")
+                             .SetSchema("Label")
+                             .AddStringProperty("target", "namespace#label/1")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(0)  // Never expire.
+                             .Build();
+  DocumentProto label3 = DocumentBuilder()
+                             .SetKey("namespace", "label/3")
+                             .SetSchema("Label")
+                             .AddStringProperty("target", "namespace#label/2")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(0)  // Never expire.
+                             .Build();
+
+  {
+    IcingSearchEngineOptions options = GetDefaultIcingOptions();
+    options.set_enable_qualified_id_join_index_v3(true);
+    options.set_enable_soft_index_restoration(true);
+    options.set_enable_delete_propagation_from(true);
+    options.set_expired_document_purge_threshold_ms(0);
+
+    // Initialize Icing instance with the flag on and put all label documents.
+    auto fake_clock = std::make_unique<FakeClock>();
+    fake_clock->SetSystemTimeMilliseconds(10);
+
+    TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                                std::make_unique<IcingFilesystem>(),
+                                std::move(fake_clock), GetTestJniCache());
+    ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+    ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(label1).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(label2).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(label3).status(), ProtoIsOk());
+  }
+
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  options.set_enable_qualified_id_join_index_v3(true);
+  options.set_enable_soft_index_restoration(true);
+  options.set_enable_delete_propagation_from(false);
+  options.set_expired_document_purge_threshold_ms(0);
+
+  auto fake_clock = std::make_unique<FakeClock>();
+  fake_clock->SetSystemTimeMilliseconds(20000);
+
+  // Initialize Icing instance again with the flag off, at t = 20000 ms.
+  TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                              std::make_unique<IcingFilesystem>(),
+                              std::move(fake_clock), GetTestJniCache());
+  InitializeResultProto initialize_result = icing.Initialize();
+
+  // Join index should be rebuilt due to the flag change.
+  // - Label1 is expired at this moment.
+  // - Label2 is not expired. Although the schema sets delete propagation for
+  //   the joinable property, since the flag is off, we won't check the
+  //   dependency and it will be re-indexed successfully.
+  // - Label3 is re-indexed successfully.
+  EXPECT_THAT(initialize_result.initialize_stats().document_store_data_status(),
+              Eq(InitializeStatsProto::NONE));
+  EXPECT_THAT(initialize_result.initialize_stats().index_restoration_cause(),
+              Eq(InitializeStatsProto::NONE));
+  EXPECT_THAT(
+      initialize_result.initialize_stats().integer_index_restoration_cause(),
+      Eq(InitializeStatsProto::NONE));
+  EXPECT_THAT(initialize_result.initialize_stats()
+                  .qualified_id_join_index_restoration_cause(),
+              Eq(InitializeStatsProto::FEATURE_FLAG_CHANGED));
+  EXPECT_THAT(
+      initialize_result.initialize_stats().embedding_index_restoration_cause(),
+      Eq(InitializeStatsProto::NONE));
+  EXPECT_THAT(
+      initialize_result.initialize_stats().num_failed_reindexed_documents(),
+      Eq(0));
+
+  GetResultProto expected_get_result_proto1;
+  expected_get_result_proto1.mutable_status()->set_code(StatusProto::NOT_FOUND);
+  expected_get_result_proto1.mutable_status()->set_message(
+      "Document (namespace, label/1) not found.");
+  EXPECT_THAT(
+      icing.Get("namespace", "label/1", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto1));
+
+  GetResultProto expected_get_result_google::protobuf;
+  expected_get_result_google::protobuf.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_google::protobuf.mutable_document() = label2;
+  EXPECT_THAT(
+      icing.Get("namespace", "label/2", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_google::protobuf));
+
+  GetResultProto expected_get_result_proto3;
+  expected_get_result_proto3.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_proto3.mutable_document() = label3;
+  EXPECT_THAT(
+      icing.Get("namespace", "label/3", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto3));
+}
+
 struct IcingSearchEngineInitializationVersionChangeTestParam {
   version_util::VersionInfo existing_version_info;
   std::unordered_set<IcingSearchEngineFeatureInfoProto::FlaggedFeatureType>
@@ -6092,23 +7335,29 @@
             filesystem(), GetBlobDir(), &fake_clock,
             /*orphan_blob_time_to_live_ms=*/0,
             PortableFileBackedProtoLog<BlobInfoProto>::kDefaultCompressionLevel,
-            /*manage_blob_files=*/true));
+            protobuf_ports::kDefaultMemLevel,
+            /*manage_blob_files=*/true, feature_flags_.get()));
 
     // Put message into DocumentStore
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
-        DocumentStore::Create(filesystem(), GetDocumentDir(), &fake_clock,
-                              schema_store.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            filesystem(), GetDocumentDir(), &fake_clock, schema_store.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     std::unique_ptr<DocumentStore> document_store =
         std::move(create_result.document_store);
-    ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                               document_store->Put(message));
+    ICING_ASSERT_OK_AND_ASSIGN(
+        DocumentStore::PutResult put_result,
+        document_store->Put(document_util::CreateDocumentWrapper(message)));
     DocumentId doc_id = put_result.new_document_id;
 
     // Index doc_id with incorrect data
@@ -6117,7 +7366,8 @@
                            /*lite_index_sort_size=*/1024 * 8);
     ICING_ASSERT_OK_AND_ASSIGN(
         std::unique_ptr<Index> index,
-        Index::Create(options, filesystem(), icing_filesystem()));
+        Index::Create(options, filesystem(), icing_filesystem(),
+                      feature_flags_.get()));
 
     ICING_ASSERT_OK_AND_ASSIGN(
         std::unique_ptr<IntegerIndex> integer_index,
@@ -6133,7 +7383,8 @@
     ICING_ASSERT_OK_AND_ASSIGN(
         std::unique_ptr<EmbeddingIndex> embedding_index,
         EmbeddingIndex::Create(filesystem(), GetEmbeddingIndexDir(),
-                               &fake_clock, feature_flags_.get()));
+                               &fake_clock, feature_flags_.get(),
+                               /*num_shards=*/32));
 
     ICING_ASSERT_OK_AND_ASSIGN(
         std::unique_ptr<TermIndexingHandler> term_indexing_handler,
@@ -6148,7 +7399,8 @@
         std::unique_ptr<QualifiedIdJoinIndexingHandler>
             qualified_id_join_indexing_handler,
         QualifiedIdJoinIndexingHandler::Create(
-            &fake_clock, document_store.get(), qualified_id_join_index.get()));
+            &fake_clock, document_store.get(), qualified_id_join_index.get(),
+            feature_flags_.get()));
     ICING_ASSERT_OK_AND_ASSIGN(
         std::unique_ptr<EmbeddingIndexingHandler> embedding_indexing_handler,
         EmbeddingIndexingHandler::Create(&fake_clock, embedding_index.get(),
@@ -6174,8 +7426,10 @@
             .Build();
     ICING_ASSERT_OK_AND_ASSIGN(
         TokenizedDocument tokenized_document,
-        TokenizedDocument::Create(schema_store.get(), lang_segmenter_.get(),
-                                  std::move(incorrect_message)));
+        TokenizedDocument::Create(
+            schema_store.get(), lang_segmenter_.get(),
+            /*current_time_ms=*/fake_clock.GetSystemTimeMilliseconds(),
+            std::move(incorrect_message)));
     ICING_ASSERT_OK(index_processor.IndexDocument(tokenized_document, doc_id,
                                                   put_result.old_document_id));
 
@@ -6246,6 +7500,10 @@
       initialize_result.initialize_stats().embedding_index_restoration_cause(),
       Eq(expected_recovery_cause));
 
+  // needs_persist_type should be set to RECOVERY_PROOF.
+  EXPECT_THAT(initialize_result.needs_persist_type(),
+              Eq(PersistType::RECOVERY_PROOF));
+
   // Manually check version file
   ICING_ASSERT_OK_AND_ASSIGN(
       IcingSearchEngineVersionProto version_proto_after_init,
@@ -6814,6 +8072,146 @@
                     std::vector<bool>{true, true, true, true},
                     std::vector<bool>{false, false, false, false}));
 
+class IcingSearchEngineInitializationChangeEmbeddingIndexNumShardsTest
+    : public IcingSearchEngineInitializationTest,
+      public ::testing::WithParamInterface<std::vector<uint32_t>> {};
+
+TEST_P(IcingSearchEngineInitializationChangeEmbeddingIndexNumShardsTest,
+       ChangeEmbeddingIndexNumShardsTest) {
+  constexpr float eps = 0.0001f;
+  std::vector<uint32_t> num_shards = GetParam();
+
+  // Create icing and add 128 documents in different schemas, so that the
+  // embeddings will be stored in different shards.
+  SchemaBuilder schema = SchemaBuilder();
+  std::vector<DocumentProto> documents;
+  documents.reserve(128);
+  for (int i = 0; i < 128; ++i) {
+    schema.AddType(
+        SchemaTypeConfigBuilder()
+            .SetType("Type" + std::to_string(i))
+            .AddProperty(
+                PropertyConfigBuilder()
+                    .SetName("embeddingUnquantized")
+                    .SetDataTypeVector(
+                        EMBEDDING_INDEXING_LINEAR_SEARCH,
+                        EmbeddingIndexingConfig::QuantizationType::NONE)
+                    .SetCardinality(CARDINALITY_OPTIONAL))
+            .AddProperty(
+                PropertyConfigBuilder()
+                    .SetName("embeddingQuantized")
+                    .SetDataTypeVector(EMBEDDING_INDEXING_LINEAR_SEARCH,
+                                       EmbeddingIndexingConfig::
+                                           QuantizationType::QUANTIZE_8_BIT)
+                    .SetCardinality(CARDINALITY_OPTIONAL)));
+    PropertyProto::VectorProto vector =
+        CreateVector("my_model", {0., 0., (float)i});
+    documents.push_back(DocumentBuilder()
+                            .SetKey("icing", "uri" + std::to_string(i))
+                            .SetSchema("Type" + std::to_string(i))
+                            .SetCreationTimestampMs(1)
+                            .AddVectorProperty("embeddingQuantized", vector)
+                            .AddVectorProperty("embeddingUnquantized", vector)
+                            .Build());
+  }
+  {
+    IcingSearchEngineOptions options = GetDefaultIcingOptions();
+    options.set_enable_embedding_index(true);
+    options.set_enable_embedding_quantization(true);
+    options.set_embedding_index_num_shards(num_shards[0]);
+    TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                                std::make_unique<IcingFilesystem>(),
+                                std::make_unique<FakeClock>(),
+                                GetTestJniCache());
+    ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+    ASSERT_THAT(icing.SetSchema(schema.Build()).status(), ProtoIsOk());
+    for (const auto& document : documents) {
+      ASSERT_THAT(icing.Put(document).status(), ProtoIsOk());
+    }
+  }
+
+  // Create icing multiple times with different num_shards.
+  for (int i = 1; i < num_shards.size(); ++i) {
+    bool num_shards_changed = num_shards[i] != num_shards[i - 1];
+
+    // Ensure that the embedding index is rebuilt if num_shards is changed.
+    auto mock_filesystem = std::make_unique<MockFilesystem>();
+    EXPECT_CALL(*mock_filesystem, DeleteDirectoryRecursively(_))
+        .WillRepeatedly(DoDefault());
+    EXPECT_CALL(*mock_filesystem,
+                DeleteDirectoryRecursively(EndsWith("/embedding_index_dir")))
+        .Times(num_shards_changed ? 1 : 0);
+
+    IcingSearchEngineOptions options = GetDefaultIcingOptions();
+    options.set_enable_embedding_index(true);
+    options.set_enable_embedding_quantization(true);
+    options.set_embedding_index_num_shards(num_shards[i]);
+    TestIcingSearchEngine icing(options, std::move(mock_filesystem),
+                                std::make_unique<IcingFilesystem>(),
+                                std::make_unique<FakeClock>(),
+                                GetTestJniCache());
+    InitializeResultProto initialize_result = icing.Initialize();
+    ASSERT_THAT(initialize_result.status(), ProtoIsOk());
+    // Ensure that the embedding index is rebuilt if num_shards is changed.
+    EXPECT_THAT(initialize_result.initialize_stats()
+                    .embedding_index_restoration_cause(),
+                Eq(num_shards_changed ? InitializeStatsProto::IO_ERROR
+                                      : InitializeStatsProto::NONE));
+    EXPECT_THAT(initialize_result.initialize_stats().index_restoration_cause(),
+                Eq(InitializeStatsProto::NONE));
+    EXPECT_THAT(
+        initialize_result.initialize_stats().integer_index_restoration_cause(),
+        Eq(InitializeStatsProto::NONE));
+    EXPECT_THAT(initialize_result.initialize_stats()
+                    .qualified_id_join_index_restoration_cause(),
+                Eq(InitializeStatsProto::NONE));
+
+    // Write an embedding query that always matches all the documents.
+    SearchSpecProto search_spec;
+    search_spec.set_query("semanticSearch(getEmbeddingParameter(0))");
+    *search_spec.add_embedding_query_vectors() =
+        CreateVector("my_model", {1, 1, 1});
+    search_spec.set_embedding_query_metric_type(
+        SearchSpecProto::EmbeddingQueryMetricType::DOT_PRODUCT);
+    search_spec.add_enabled_features(
+        std::string(kListFilterQueryLanguageFeature));
+    ScoringSpecProto scoring_spec = GetDefaultScoringSpec();
+    scoring_spec.set_rank_by(
+        ScoringSpecProto::RankingStrategy::ADVANCED_SCORING_EXPRESSION);
+    // Set the advanced scoring expression to a constant so that the documents
+    // will be returned according to their order they were added.
+    scoring_spec.set_advanced_scoring_expression("0");
+    // Set the additional advanced scoring expression to check the embedding
+    // scores.
+    // Each embedding in document[i] should have a score of i, and every
+    // document has two embeddings, one quantized and one unquantized. So
+    // document[i] should have a total score of i * 2.
+    scoring_spec.add_additional_advanced_scoring_expressions(
+        "sum(this.matchedSemanticScores(getEmbeddingParameter(0)))");
+    ResultSpecProto result_spec;
+    result_spec.set_num_per_page(10000);
+    SearchResultProto results =
+        icing.Search(search_spec, scoring_spec, result_spec);
+    ASSERT_THAT(results.status(), ProtoIsOk());
+    ASSERT_THAT(results.results(), SizeIs(128));
+    for (int i = 0; i < 128; ++i) {
+      EXPECT_THAT(results.results(127 - i).document(),
+                  EqualsProto(documents[i]));
+      EXPECT_NEAR(results.results(127 - i).additional_scores(0), i * 2.0f, eps);
+    }
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    IcingSearchEngineInitializationChangeEmbeddingIndexNumShardsTest,
+    IcingSearchEngineInitializationChangeEmbeddingIndexNumShardsTest,
+    testing::Values(std::vector<uint32_t>{1, 32, 1, 32, 32, 32, 1, 1, 32},
+                    std::vector<uint32_t>{32, 1, 32, 1, 1, 1, 32, 32, 1},
+                    std::vector<uint32_t>{1, 2, 4, 8, 16, 32, 1, 1, 32},
+                    std::vector<uint32_t>{32, 16, 8, 4, 2, 1, 32, 32, 1},
+                    std::vector<uint32_t>{1, 1, 1, 1},
+                    std::vector<uint32_t>{32, 32, 32, 32}));
+
 class IcingSearchEngineInitializationChangeEnableScorablePropertiesFlagTest
     : public IcingSearchEngineInitializationTest,
       public ::testing::WithParamInterface<std::vector<bool>> {};
@@ -6885,6 +8283,9 @@
         initialize_result.initialize_stats().document_store_recovery_cause(),
         Eq(flag_changed ? InitializeStatsProto::FEATURE_FLAG_CHANGED
                         : InitializeStatsProto::NONE));
+    EXPECT_THAT(
+        initialize_result.needs_persist_type(),
+        Eq(flag_changed ? PersistType::RECOVERY_PROOF : PersistType::UNKNOWN));
 
     // Schema store and all indices should be unaffected.
     EXPECT_THAT(
@@ -7005,9 +8406,9 @@
                            .SetCardinality(CARDINALITY_OPTIONAL))
           .Build();
   SchemaTypeConfigProto db1_email_type_with_db =
-      SchemaTypeConfigBuilder(db1_email_type).SetDatabase("db1").Build();
+      SchemaTypeConfigBuilder(db1_email_type).SetDatabase("db1/").Build();
   SchemaTypeConfigProto db2_email_type_with_db =
-      SchemaTypeConfigBuilder(db2_email_type).SetDatabase("db2").Build();
+      SchemaTypeConfigBuilder(db2_email_type).SetDatabase("db2/").Build();
 
   SchemaProto previous_version_db1_schema;
   SchemaProto previous_version_db2_schema;
@@ -7075,18 +8476,12 @@
     EXPECT_THAT(icing.Initialize().status(), ProtoIsOk());
     // 1. Set schema.
     if (options.enable_schema_database()) {
-      // Can only set schema for a single database at a time.
+      // Use the SetSchemaRequestProto with empty database field to populate
+      // both databases at once.
       ASSERT_THAT(icing
                       .SetSchema(CreateSetSchemaRequestProto(
-                          previous_version_db1_schema,
-                          previous_version_db1_schema.types(0).database(),
-                          /*ignore_errors_and_delete_documents=*/false))
-                      .status(),
-                  ProtoIsOk());
-      ASSERT_THAT(icing
-                      .SetSchema(CreateSetSchemaRequestProto(
-                          previous_version_db2_schema,
-                          previous_version_db2_schema.types(0).database(),
+                          previous_version_schema,
+                          /*database=*/"",
                           /*ignore_errors_and_delete_documents=*/false))
                       .status(),
                   ProtoIsOk());
@@ -7166,9 +8561,9 @@
   // Verify GetSchema
   if (enable_schema_database) {
     SchemaTypeConfigProto db1_email_type_with_db =
-        SchemaTypeConfigBuilder(db1_email_type).SetDatabase("db1").Build();
+        SchemaTypeConfigBuilder(db1_email_type).SetDatabase("db1/").Build();
     SchemaTypeConfigProto db2_email_type_with_db =
-        SchemaTypeConfigBuilder(db2_email_type).SetDatabase("db2").Build();
+        SchemaTypeConfigBuilder(db2_email_type).SetDatabase("db2/").Build();
     SchemaProto full_schema_with_database = SchemaBuilder()
                                                 .AddType(db1_email_type_with_db)
                                                 .AddType(db2_email_type_with_db)
@@ -7190,14 +8585,14 @@
     expected_get_schema_result_proto_db1.mutable_status()->set_code(
         StatusProto::OK);
     *expected_get_schema_result_proto_db1.mutable_schema() = db1_schema;
-    EXPECT_THAT(icing.GetSchema("db1"),
+    EXPECT_THAT(icing.GetSchema("db1/"),
                 EqualsProto(expected_get_schema_result_proto_db1));
 
     GetSchemaResultProto expected_get_schema_result_proto_db2;
     expected_get_schema_result_proto_db2.mutable_status()->set_code(
         StatusProto::OK);
     *expected_get_schema_result_proto_db2.mutable_schema() = db2_schema;
-    EXPECT_THAT(icing.GetSchema("db2"),
+    EXPECT_THAT(icing.GetSchema("db2/"),
                 EqualsProto(expected_get_schema_result_proto_db2));
   } else {
     GetSchemaResultProto expected_get_schema_result_proto;
@@ -7337,6 +8732,9 @@
                     .qualified_id_join_index_restoration_cause(),
                 Eq(flag_changed ? InitializeStatsProto::FEATURE_FLAG_CHANGED
                                 : InitializeStatsProto::NONE));
+    EXPECT_THAT(
+        initialize_result.needs_persist_type(),
+        Eq(flag_changed ? PersistType::RECOVERY_PROOF : PersistType::UNKNOWN));
 
     // Schema store, document store and all other indices should be unaffected.
     EXPECT_THAT(
@@ -7405,6 +8803,261 @@
                     std::vector<bool>{true, true, true, true},
                     std::vector<bool>{false, false, false, false}));
 
+class IcingSearchEngineInitializationChangeNonExistentQualifiedIdJoinFlagTest
+    : public IcingSearchEngineInitializationTest,
+      public ::testing::WithParamInterface<std::vector<bool>> {};
+TEST_P(IcingSearchEngineInitializationChangeNonExistentQualifiedIdJoinFlagTest,
+       ChangeNonExistentQualifiedIdJoinFlagTest) {
+  std::vector<bool> enable_non_existent_qualified_id_join_flags = GetParam();
+
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("name")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_REQUIRED)))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("Message")
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("body")
+                                        .SetDataTypeString(TERM_MATCH_PREFIX,
+                                                           TOKENIZER_PLAIN)
+                                        .SetCardinality(CARDINALITY_REQUIRED))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("senderQualifiedId")
+                                        .SetDataTypeJoinableString(
+                                            JOINABLE_VALUE_TYPE_QUALIFIED_ID)
+                                        .SetCardinality(CARDINALITY_REQUIRED)))
+          .Build();
+
+  DocumentProto person =
+      DocumentBuilder()
+          .SetKey("namespace", "person")
+          .SetSchema("Person")
+          .AddStringProperty("name", "person")
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .Build();
+  DocumentProto message =
+      DocumentBuilder()
+          .SetKey("namespace", "message/1")
+          .SetSchema("Message")
+          .AddStringProperty("body", "message body")
+          .AddStringProperty("senderQualifiedId", "namespace#person")
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .Build();
+
+  {
+    IcingSearchEngineOptions options = GetDefaultIcingOptions();
+    options.set_enable_non_existent_qualified_id_join(
+        enable_non_existent_qualified_id_join_flags.at(0));
+    TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                                std::make_unique<IcingFilesystem>(),
+                                std::make_unique<FakeClock>(),
+                                GetTestJniCache());
+
+    ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+    ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(person).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(message).status(), ProtoIsOk());
+  }
+
+  // Create icing multiple times with different
+  // enable_non_existent_qualified_id_join flags.
+  for (int i = 1; i < enable_non_existent_qualified_id_join_flags.size(); ++i) {
+    bool flag_changed = enable_non_existent_qualified_id_join_flags[i] !=
+                        enable_non_existent_qualified_id_join_flags[i - 1];
+
+    // Ensure that the qualified id join index is rebuilt if the flag is
+    // changed.
+    IcingSearchEngineOptions options = GetDefaultIcingOptions();
+    options.set_enable_non_existent_qualified_id_join(
+        enable_non_existent_qualified_id_join_flags[i]);
+    TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                                std::make_unique<IcingFilesystem>(),
+                                std::make_unique<FakeClock>(),
+                                GetTestJniCache());
+    InitializeResultProto initialize_result = icing.Initialize();
+    ASSERT_THAT(initialize_result.status(), ProtoIsOk());
+
+    // Qualified id join index recovery cause should be FEATURE_FLAG_CHANGED if
+    // flag is changed.
+    EXPECT_THAT(initialize_result.initialize_stats()
+                    .qualified_id_join_index_restoration_cause(),
+                Eq(flag_changed ? InitializeStatsProto::FEATURE_FLAG_CHANGED
+                                : InitializeStatsProto::NONE));
+    EXPECT_THAT(
+        initialize_result.needs_persist_type(),
+        Eq(flag_changed ? PersistType::RECOVERY_PROOF : PersistType::UNKNOWN));
+
+    // Schema store, document store and all other indices should be unaffected.
+    EXPECT_THAT(
+        initialize_result.initialize_stats().schema_store_recovery_cause(),
+        Eq(InitializeStatsProto::NONE));
+    EXPECT_THAT(
+        initialize_result.initialize_stats().document_store_recovery_cause(),
+        Eq(InitializeStatsProto::NONE));
+    EXPECT_THAT(initialize_result.initialize_stats().index_restoration_cause(),
+                Eq(InitializeStatsProto::NONE));
+    EXPECT_THAT(
+        initialize_result.initialize_stats().integer_index_restoration_cause(),
+        Eq(InitializeStatsProto::NONE));
+    EXPECT_THAT(initialize_result.initialize_stats()
+                    .embedding_index_restoration_cause(),
+                Eq(InitializeStatsProto::NONE));
+
+    // Prepare join search spec to join a query for `name:person` with a child
+    // query for `body:message` based on the child's `senderQualifiedId` field.
+    //
+    // No matter what the flag value is, the join API should always return the
+    // expected result.
+    SearchSpecProto search_spec;
+    search_spec.set_term_match_type(TermMatchType::EXACT_ONLY);
+    search_spec.set_query("name:person");
+    JoinSpecProto* join_spec = search_spec.mutable_join_spec();
+    join_spec->set_parent_property_expression(
+        std::string(JoinProcessor::kQualifiedIdExpr));
+    join_spec->set_child_property_expression("senderQualifiedId");
+    join_spec->set_aggregation_scoring_strategy(
+        JoinSpecProto::AggregationScoringStrategy::COUNT);
+    JoinSpecProto::NestedSpecProto* nested_spec =
+        join_spec->mutable_nested_spec();
+    SearchSpecProto* nested_search_spec = nested_spec->mutable_search_spec();
+    nested_search_spec->set_term_match_type(TermMatchType::EXACT_ONLY);
+    nested_search_spec->set_query("body:message");
+    *nested_spec->mutable_scoring_spec() = GetDefaultScoringSpec();
+    *nested_spec->mutable_result_spec() = ResultSpecProto::default_instance();
+
+    ResultSpecProto result_spec = ResultSpecProto::default_instance();
+    result_spec.set_max_joined_children_per_parent_to_return(
+        std::numeric_limits<int32_t>::max());
+
+    SearchResultProto expected_search_result_proto;
+    expected_search_result_proto.mutable_status()->set_code(StatusProto::OK);
+    SearchResultProto::ResultProto* result_proto =
+        expected_search_result_proto.mutable_results()->Add();
+    *result_proto->mutable_document() = person;
+    *result_proto->mutable_joined_results()->Add()->mutable_document() =
+        message;
+
+    SearchResultProto search_result_proto =
+        icing.Search(search_spec, GetDefaultScoringSpec(), result_spec);
+    EXPECT_THAT(search_result_proto, EqualsSearchResultIgnoreStatsAndScores(
+                                         expected_search_result_proto));
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    IcingSearchEngineInitializationChangeNonExistentQualifiedIdJoinFlagTest,
+    IcingSearchEngineInitializationChangeNonExistentQualifiedIdJoinFlagTest,
+    testing::Values(std::vector<bool>{false, true, false, true, false, true},
+                    std::vector<bool>{true, false, true, false, true, false},
+                    std::vector<bool>{false, true, true, true, false, true},
+                    std::vector<bool>{true, false, false, false, true, false},
+                    std::vector<bool>{true, true, true, true},
+                    std::vector<bool>{false, false, false, false}));
+
+class IcingSearchEngineInitializationChangeEnableProtoLogNewHeaderFormatFlagTest
+    : public IcingSearchEngineInitializationTest,
+      public ::testing::WithParamInterface<std::vector<bool>> {};
+
+TEST_P(
+    IcingSearchEngineInitializationChangeEnableProtoLogNewHeaderFormatFlagTest,
+    ChangeFlag) {
+  std::vector<bool> enable_proto_log_new_header_format_flags = GetParam();
+
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("name")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_REQUIRED)))
+          .Build();
+
+  DocumentProto document =
+      DocumentBuilder()
+          .SetKey("namespace", "uri")
+          .SetSchema("Person")
+          .AddStringProperty("name", "person")
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .Build();
+
+  {
+    IcingSearchEngineOptions options = GetDefaultIcingOptions();
+    options.set_enable_proto_log_new_header_format(
+        enable_proto_log_new_header_format_flags[0]);
+    IcingSearchEngine icing(options, GetTestJniCache());
+
+    ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+    ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(document).status(), ProtoIsOk());
+    ASSERT_THAT(icing.PersistToDisk(PersistType::FULL).status(), ProtoIsOk());
+  }
+
+  // Create icing multiple times with different
+  // enable_proto_log_new_header_format flags.
+  for (int i = 1; i < enable_proto_log_new_header_format_flags.size(); ++i) {
+    // Ensure that persisted data is kept.
+    IcingSearchEngineOptions options = GetDefaultIcingOptions();
+    options.set_enable_proto_log_new_header_format(
+        enable_proto_log_new_header_format_flags[i]);
+    IcingSearchEngine icing(options, GetTestJniCache());
+
+    InitializeResultProto initialize_result = icing.Initialize();
+    ASSERT_THAT(initialize_result.status(), ProtoIsOk());
+
+    // No data loss should be reported.
+    EXPECT_THAT(
+        initialize_result.initialize_stats().document_store_data_status(),
+        Eq(InitializeStatsProto::NONE));
+
+    // We don't really care about the derived files recovery cause, but normally
+    // this flag change should not affect derived files recovery, so let's check
+    // it here.
+    EXPECT_THAT(
+        initialize_result.initialize_stats().schema_store_recovery_cause(),
+        Eq(InitializeStatsProto::NONE));
+    EXPECT_THAT(
+        initialize_result.initialize_stats().document_store_recovery_cause(),
+        Eq(InitializeStatsProto::NONE));
+    EXPECT_THAT(initialize_result.initialize_stats().index_restoration_cause(),
+                Eq(InitializeStatsProto::NONE));
+    EXPECT_THAT(
+        initialize_result.initialize_stats().integer_index_restoration_cause(),
+        Eq(InitializeStatsProto::NONE));
+    EXPECT_THAT(initialize_result.initialize_stats()
+                    .qualified_id_join_index_restoration_cause(),
+                Eq(InitializeStatsProto::NONE));
+    EXPECT_THAT(initialize_result.initialize_stats()
+                    .embedding_index_restoration_cause(),
+                Eq(InitializeStatsProto::NONE));
+
+    // PersistToDisk is not needed.
+    EXPECT_THAT(initialize_result.needs_persist_type(),
+                Eq(PersistType::UNKNOWN));
+
+    // Verify that the persisted document can be retrieved.
+    GetResultProto expected_get_result_proto;
+    expected_get_result_proto.mutable_status()->set_code(StatusProto::OK);
+    *expected_get_result_proto.mutable_document() = document;
+
+    EXPECT_THAT(
+        icing.Get("namespace", "uri", GetResultSpecProto::default_instance()),
+        EqualsProto(expected_get_result_proto));
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    IcingSearchEngineInitializationChangeEnableProtoLogNewHeaderFormatFlagTest,
+    IcingSearchEngineInitializationChangeEnableProtoLogNewHeaderFormatFlagTest,
+    testing::Values(std::vector<bool>{false, true, false, true, false, true},
+                    std::vector<bool>{true, false, true, false, true, false},
+                    std::vector<bool>{false, true, true, true, false, true},
+                    std::vector<bool>{true, false, false, false, true, false},
+                    std::vector<bool>{true, true, true, true},
+                    std::vector<bool>{false, false, false, false}));
+
 }  // namespace
 }  // namespace lib
 }  // namespace icing

diff --git a/icing/icing-search-engine_optimize_test.cc b/icing/icing-search-engine_optimize_test.cc
index cd36074..d8018db 100644
--- a/icing/icing-search-engine_optimize_test.cc
+++ b/icing/icing-search-engine_optimize_test.cc

@@ -75,6 +75,7 @@
 using ::testing::Gt;
 using ::testing::HasSubstr;
 using ::testing::Lt;
+using ::testing::Ne;
 using ::testing::Return;
 using ::testing::SizeIs;
 
@@ -230,25 +231,28 @@
       document2;
   SearchResultProto search_result_proto =
       icing.Search(search_spec, GetDefaultScoringSpec(), result_spec);
-  EXPECT_THAT(search_result_proto.next_page_token(), Gt(kInvalidNextPageToken));
   uint64_t next_page_token = search_result_proto.next_page_token();
+
   // Since the token is a random number, we don't need to verify
   expected_search_result_proto.set_next_page_token(next_page_token);
+  EXPECT_THAT(search_result_proto.next_page_token(), Ne(kInvalidNextPageToken));
   EXPECT_THAT(search_result_proto, EqualsSearchResultIgnoreStatsAndScores(
                                        expected_search_result_proto));
-  // Now document1 is still to be fetched.
 
+  // Now there are more pages to be fetched (document1). Call Optimize.
   OptimizeResultProto optimize_result_proto;
   optimize_result_proto.mutable_status()->set_code(StatusProto::OK);
   optimize_result_proto.mutable_status()->set_message("");
   OptimizeResultProto actual_result = icing.Optimize();
   actual_result.clear_optimize_stats();
+  actual_result.clear_vm_binder_transaction_latency_start_time_ms();
   ASSERT_THAT(actual_result, EqualsProto(optimize_result_proto));
 
   // Tries to fetch the second page, no results since all tokens have been
-  // invalidated during Optimize()
+  // invalidated during Optimize().
   expected_search_result_proto.clear_results();
   expected_search_result_proto.clear_next_page_token();
+  expected_search_result_proto.set_page_token_not_found(true);
   search_result_proto = icing.GetNextPage(next_page_token);
   EXPECT_THAT(search_result_proto, EqualsSearchResultIgnoreStatsAndScores(
                                        expected_search_result_proto));
@@ -361,10 +365,17 @@
           .AddStringProperty("body", "message body one")
           .SetCreationTimestampMs(kDefaultCreationTimestampMs)
           .Build();
-  DocumentProto document2 = DocumentBuilder()
-                                .SetKey("namespace", "uri2")
+  DocumentProto document2 =
+      DocumentBuilder()
+          .SetKey("namespace", "uri2")
+          .SetSchema("Message")
+          .AddStringProperty("body", "message body two")
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .Build();
+  DocumentProto document3 = DocumentBuilder()
+                                .SetKey("namespace", "uri3")
                                 .SetSchema("Message")
-                                .AddStringProperty("body", "message body two")
+                                .AddStringProperty("body", "message body three")
                                 .SetCreationTimestampMs(100)
                                 .SetTtlMs(500)
                                 .Build();
@@ -385,9 +396,13 @@
     EXPECT_THAT(optimize_info.optimizable_docs(), Eq(0));
     EXPECT_THAT(optimize_info.estimated_optimizable_bytes(), Eq(0));
     EXPECT_THAT(optimize_info.time_since_last_optimize_ms(), Eq(0));
+    EXPECT_TRUE(optimize_info.no_previous_optimize_info());
+    EXPECT_THAT(optimize_info.num_active_result_states(), Eq(0));
 
+    // Set schema and add 2 documents.
     ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
     ASSERT_THAT(icing.Put(document1).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(document2).status(), ProtoIsOk());
 
     // Only have active documents, nothing is optimizable yet.
     optimize_info = icing.GetOptimizeInfo();
@@ -395,32 +410,64 @@
     EXPECT_THAT(optimize_info.optimizable_docs(), Eq(0));
     EXPECT_THAT(optimize_info.estimated_optimizable_bytes(), Eq(0));
     EXPECT_THAT(optimize_info.time_since_last_optimize_ms(), Eq(0));
+    EXPECT_TRUE(optimize_info.no_previous_optimize_info());
+    EXPECT_THAT(optimize_info.num_active_result_states(), Eq(0));
 
-    // Deletes document1
-    ASSERT_THAT(icing.Delete("namespace", "uri1").status(), ProtoIsOk());
+    // Send a search request to create a result state.
+    SearchSpecProto search_spec;
+    search_spec.set_query("body:message");
+    search_spec.set_term_match_type(TermMatchType::EXACT_ONLY);
+    ResultSpecProto result_spec;
+    result_spec.set_num_per_page(1);
+    SearchResultProto search_result =
+        icing.Search(search_spec, GetDefaultScoringSpec(), result_spec);
+    ASSERT_THAT(search_result.status(), ProtoIsOk());
+    ASSERT_THAT(search_result.results_size(), Eq(1));
+    ASSERT_THAT(search_result.next_page_token(), Ne(kInvalidNextPageToken));
 
     optimize_info = icing.GetOptimizeInfo();
     EXPECT_THAT(optimize_info.status(), ProtoIsOk());
-    EXPECT_THAT(optimize_info.optimizable_docs(), Eq(1));
-    EXPECT_THAT(optimize_info.estimated_optimizable_bytes(), Gt(0));
+    EXPECT_THAT(optimize_info.optimizable_docs(), Eq(0));
+    EXPECT_THAT(optimize_info.estimated_optimizable_bytes(), Eq(0));
     EXPECT_THAT(optimize_info.time_since_last_optimize_ms(), Eq(0));
-    int64_t first_estimated_optimizable_bytes =
-        optimize_info.estimated_optimizable_bytes();
+    EXPECT_TRUE(optimize_info.no_previous_optimize_info());
+    EXPECT_THAT(optimize_info.num_active_result_states(), Eq(1));
 
-    // Add a second document, but it'll be expired since the time (1000) is
-    // greater than the document's creation timestamp (100) + the document's ttl
-    // (500)
-    ASSERT_THAT(icing.Put(document2).status(), ProtoIsOk());
+    // Deletes document1 and document2.
+    ASSERT_THAT(icing.Delete("namespace", "uri1").status(), ProtoIsOk());
+    ASSERT_THAT(icing.Delete("namespace", "uri2").status(), ProtoIsOk());
 
     optimize_info = icing.GetOptimizeInfo();
     EXPECT_THAT(optimize_info.status(), ProtoIsOk());
     EXPECT_THAT(optimize_info.optimizable_docs(), Eq(2));
+    EXPECT_THAT(optimize_info.estimated_optimizable_bytes(), Gt(0));
+    EXPECT_THAT(optimize_info.time_since_last_optimize_ms(), Eq(0));
+    EXPECT_TRUE(optimize_info.no_previous_optimize_info());
+    int64_t first_estimated_optimizable_bytes =
+        optimize_info.estimated_optimizable_bytes();
+    EXPECT_THAT(optimize_info.num_active_result_states(), Eq(1));
+
+    // Add third document, but it'll be expired since the time (1000) is
+    // greater than the document's creation timestamp (100) + the document's ttl
+    // (500)
+    ASSERT_THAT(icing.Put(document3).status(), ProtoIsOk());
+
+    optimize_info = icing.GetOptimizeInfo();
+    EXPECT_THAT(optimize_info.status(), ProtoIsOk());
+    EXPECT_THAT(optimize_info.optimizable_docs(), Eq(3));
     EXPECT_THAT(optimize_info.estimated_optimizable_bytes(),
                 Gt(first_estimated_optimizable_bytes));
     EXPECT_THAT(optimize_info.time_since_last_optimize_ms(), Eq(0));
+    EXPECT_TRUE(optimize_info.no_previous_optimize_info());
+    EXPECT_THAT(optimize_info.num_active_result_states(), Eq(1));
 
     // Optimize
     ASSERT_THAT(icing.Optimize().status(), ProtoIsOk());
+
+    // Result state manager is reset after optimize.
+    optimize_info = icing.GetOptimizeInfo();
+    EXPECT_THAT(optimize_info.status(), ProtoIsOk());
+    EXPECT_THAT(optimize_info.num_active_result_states(), Eq(0));
   }
 
   {
@@ -440,6 +487,125 @@
     EXPECT_THAT(optimize_info.optimizable_docs(), Eq(0));
     EXPECT_THAT(optimize_info.estimated_optimizable_bytes(), Eq(0));
     EXPECT_THAT(optimize_info.time_since_last_optimize_ms(), Eq(4000));
+    EXPECT_FALSE(optimize_info.no_previous_optimize_info());
+    EXPECT_THAT(optimize_info.num_active_result_states(), Eq(0));
+  }
+}
+
+TEST_F(IcingSearchEngineOptimizeTest,
+       NegativeTimeSinceLastOptimizeResetsToZero) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Message").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("body")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_REQUIRED)))
+          .Build();
+
+  DocumentProto document1 =
+      DocumentBuilder()
+          .SetKey("namespace", "uri1")
+          .SetSchema("Message")
+          .AddStringProperty("body", "message body one")
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .Build();
+  DocumentProto document2 = DocumentBuilder()
+                                .SetKey("namespace", "uri2")
+                                .SetSchema("Message")
+                                .AddStringProperty("body", "message body two")
+                                .SetCreationTimestampMs(100)
+                                .SetTtlMs(500)
+                                .Build();
+
+  IcingSearchEngineOptions icing_options = GetDefaultIcingOptions();
+  icing_options.set_calculate_time_since_last_attempted_optimize(true);
+  {
+    auto fake_clock = std::make_unique<FakeClock>();
+    fake_clock->SetSystemTimeMilliseconds(1000);
+    TestIcingSearchEngine icing(icing_options, std::make_unique<Filesystem>(),
+                                std::make_unique<IcingFilesystem>(),
+                                std::move(fake_clock), GetTestJniCache());
+    ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+
+    // Just initialized, nothing is optimizable yet.
+    GetOptimizeInfoResultProto optimize_info = icing.GetOptimizeInfo();
+    EXPECT_THAT(optimize_info.status(), ProtoIsOk());
+    EXPECT_THAT(optimize_info.optimizable_docs(), Eq(0));
+    EXPECT_THAT(optimize_info.time_since_last_optimize_ms(), Eq(0));
+    EXPECT_TRUE(optimize_info.no_previous_optimize_info());
+    EXPECT_THAT(optimize_info.num_active_result_states(), Eq(0));
+
+    // Call some APIs
+    ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(document1).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Delete("namespace", "uri1").status(), ProtoIsOk());
+    // Add a second document, but it'll be expired since the time (1000) is
+    // greater than the document's creation timestamp (100) + the document's ttl
+    // (500)
+    ASSERT_THAT(icing.Put(document2).status(), ProtoIsOk());
+
+    optimize_info = icing.GetOptimizeInfo();
+    EXPECT_THAT(optimize_info.status(), ProtoIsOk());
+    EXPECT_THAT(optimize_info.optimizable_docs(), Eq(2));
+    EXPECT_THAT(optimize_info.time_since_last_optimize_ms(), Eq(0));
+    EXPECT_TRUE(optimize_info.no_previous_optimize_info());
+    EXPECT_THAT(optimize_info.num_active_result_states(), Eq(0));
+
+    // Optimize
+    OptimizeResultProto optimize_result = icing.Optimize();
+    EXPECT_THAT(optimize_result.status(), ProtoIsOk());
+    EXPECT_THAT(optimize_result.optimize_stats().time_since_last_optimize_ms(),
+                Eq(0));
+    EXPECT_THAT(optimize_result.optimize_stats()
+                    .time_since_last_successful_optimize_ms(),
+                Eq(0));
+  }
+
+  {
+    // Recreate with new time that's earlier than the last successful optimize
+    // run time. no_previous_optimize_info should be true and time since last
+    // optimize values should be negative.
+    auto fake_clock = std::make_unique<FakeClock>();
+    fake_clock->SetSystemTimeMilliseconds(500);
+    TestIcingSearchEngine icing(icing_options, std::make_unique<Filesystem>(),
+                                std::make_unique<IcingFilesystem>(),
+                                std::move(fake_clock), GetTestJniCache());
+    ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+
+    GetOptimizeInfoResultProto optimize_info = icing.GetOptimizeInfo();
+    EXPECT_THAT(optimize_info.status(), ProtoIsOk());
+    EXPECT_THAT(optimize_info.optimizable_docs(), Eq(0));
+    EXPECT_THAT(optimize_info.time_since_last_optimize_ms(), Eq(-500));
+    EXPECT_TRUE(optimize_info.no_previous_optimize_info());
+    EXPECT_THAT(optimize_info.num_active_result_states(), Eq(0));
+
+    // Optimize.
+    OptimizeResultProto optimize_result = icing.Optimize();
+    EXPECT_THAT(optimize_result.status(), ProtoIsOk());
+    EXPECT_THAT(optimize_result.optimize_stats().time_since_last_optimize_ms(),
+                Eq(-500));
+    EXPECT_THAT(optimize_result.optimize_stats()
+                    .time_since_last_successful_optimize_ms(),
+                Eq(-500));
+  }
+
+  {
+    // Recreate with new timer and check that time_since_last_optimize_ms is
+    // populated correctly.
+    auto fake_clock = std::make_unique<FakeClock>();
+    fake_clock->SetSystemTimeMilliseconds(800);
+    TestIcingSearchEngine icing(icing_options, std::make_unique<Filesystem>(),
+                                std::make_unique<IcingFilesystem>(),
+                                std::move(fake_clock), GetTestJniCache());
+    ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+
+    GetOptimizeInfoResultProto optimize_info = icing.GetOptimizeInfo();
+    EXPECT_THAT(optimize_info.status(), ProtoIsOk());
+    EXPECT_THAT(optimize_info.optimizable_docs(), Eq(0));
+    EXPECT_THAT(optimize_info.time_since_last_optimize_ms(), Eq(300));
+    EXPECT_FALSE(optimize_info.no_previous_optimize_info());
+    EXPECT_THAT(optimize_info.num_active_result_states(), Eq(0));
   }
 }
 
@@ -487,6 +653,8 @@
     EXPECT_THAT(optimize_info.status(), ProtoIsOk());
     EXPECT_THAT(optimize_info.optimizable_docs(), Eq(0));
     EXPECT_THAT(optimize_info.time_since_last_optimize_ms(), Eq(0));
+    EXPECT_TRUE(optimize_info.no_previous_optimize_info());
+    EXPECT_THAT(optimize_info.num_active_result_states(), Eq(0));
 
     // Call some APIs
     ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
@@ -501,6 +669,8 @@
     EXPECT_THAT(optimize_info.status(), ProtoIsOk());
     EXPECT_THAT(optimize_info.optimizable_docs(), Eq(2));
     EXPECT_THAT(optimize_info.time_since_last_optimize_ms(), Eq(0));
+    EXPECT_TRUE(optimize_info.no_previous_optimize_info());
+    EXPECT_THAT(optimize_info.num_active_result_states(), Eq(0));
 
     // Optimize
     OptimizeResultProto optimize_result = icing.Optimize();
@@ -536,6 +706,8 @@
     EXPECT_THAT(optimize_info.status(), ProtoIsOk());
     EXPECT_THAT(optimize_info.optimizable_docs(), Eq(0));
     EXPECT_THAT(optimize_info.time_since_last_optimize_ms(), Eq(500));
+    EXPECT_FALSE(optimize_info.no_previous_optimize_info());
+    EXPECT_THAT(optimize_info.num_active_result_states(), Eq(0));
 
     // Optimize again -- this should fail because of the mock filesystem.
     OptimizeResultProto optimize_result = icing.Optimize();
@@ -567,6 +739,8 @@
     EXPECT_THAT(optimize_info.status(), ProtoIsOk());
     EXPECT_THAT(optimize_info.optimizable_docs(), Eq(0));
     EXPECT_THAT(optimize_info.time_since_last_optimize_ms(), Eq(1300));
+    EXPECT_FALSE(optimize_info.no_previous_optimize_info());
+    EXPECT_THAT(optimize_info.num_active_result_states(), Eq(0));
 
     // Optimize again -- this should fail because of the mock filesystem, but
     // the time since last optimize should be populated.
@@ -598,6 +772,8 @@
     EXPECT_THAT(optimize_info.status(), ProtoIsOk());
     EXPECT_THAT(optimize_info.optimizable_docs(), Eq(0));
     EXPECT_THAT(optimize_info.time_since_last_optimize_ms(), Eq(2700));
+    EXPECT_FALSE(optimize_info.no_previous_optimize_info());
+    EXPECT_THAT(optimize_info.num_active_result_states(), Eq(0));
 
     // Optimize
     OptimizeResultProto optimize_result = icing.Optimize();
@@ -655,6 +831,8 @@
     EXPECT_THAT(optimize_info.optimizable_docs(), Eq(0));
     EXPECT_THAT(optimize_info.estimated_optimizable_bytes(), Eq(0));
     EXPECT_THAT(optimize_info.time_since_last_optimize_ms(), Eq(0));
+    EXPECT_TRUE(optimize_info.no_previous_optimize_info());
+    EXPECT_THAT(optimize_info.num_active_result_states(), Eq(0));
 
     // Call some APIs
     ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
@@ -669,6 +847,8 @@
     EXPECT_THAT(optimize_info.status(), ProtoIsOk());
     EXPECT_THAT(optimize_info.optimizable_docs(), Eq(2));
     EXPECT_THAT(optimize_info.time_since_last_optimize_ms(), Eq(0));
+    EXPECT_TRUE(optimize_info.no_previous_optimize_info());
+    EXPECT_THAT(optimize_info.num_active_result_states(), Eq(0));
 
     // Optimize
     OptimizeResultProto optimize_result = icing.Optimize();
@@ -704,6 +884,8 @@
     EXPECT_THAT(optimize_info.status(), ProtoIsOk());
     EXPECT_THAT(optimize_info.optimizable_docs(), Eq(0));
     EXPECT_THAT(optimize_info.time_since_last_optimize_ms(), Eq(500));
+    EXPECT_FALSE(optimize_info.no_previous_optimize_info());
+    EXPECT_THAT(optimize_info.num_active_result_states(), Eq(0));
 
     // Optimize again -- this should fail because of the mock filesystem.
     OptimizeResultProto optimize_result = icing.Optimize();
@@ -740,6 +922,8 @@
     EXPECT_THAT(optimize_info.status(), ProtoIsOk());
     EXPECT_THAT(optimize_info.optimizable_docs(), Eq(0));
     EXPECT_THAT(optimize_info.time_since_last_optimize_ms(), Eq(1300));
+    EXPECT_FALSE(optimize_info.no_previous_optimize_info());
+    EXPECT_THAT(optimize_info.num_active_result_states(), Eq(0));
 
     // Optimize again -- this should fail because of the mock filesystem.
     OptimizeResultProto optimize_result = icing.Optimize();
@@ -764,6 +948,7 @@
     GetOptimizeInfoResultProto optimize_info = icing.GetOptimizeInfo();
     EXPECT_THAT(optimize_info.status(), ProtoIsOk());
     EXPECT_THAT(optimize_info.time_since_last_optimize_ms(), Eq(4000));
+    EXPECT_FALSE(optimize_info.no_previous_optimize_info());
 
     // Optimize
     OptimizeResultProto optimize_result = icing.Optimize();
@@ -1939,6 +2124,8 @@
               Ge(result.optimize_stats().storage_size_after() - page_size));
   result.mutable_optimize_stats()->clear_storage_size_before();
   result.mutable_optimize_stats()->clear_storage_size_after();
+  result.mutable_optimize_stats()->clear_before_optimize_persist_stats();
+  result.mutable_optimize_stats()->clear_after_optimize_persist_stats();
   EXPECT_THAT(result.optimize_stats(), EqualsProto(expected));
 
   fake_clock = std::make_unique<FakeClock>();
@@ -1969,6 +2156,8 @@
               Eq(result.optimize_stats().storage_size_after()));
   result.mutable_optimize_stats()->clear_storage_size_before();
   result.mutable_optimize_stats()->clear_storage_size_after();
+  result.mutable_optimize_stats()->clear_before_optimize_persist_stats();
+  result.mutable_optimize_stats()->clear_after_optimize_persist_stats();
   EXPECT_THAT(result.optimize_stats(), EqualsProto(expected));
 
   // Delete the last document.
@@ -1995,6 +2184,8 @@
               Ge(result.optimize_stats().storage_size_after()));
   result.mutable_optimize_stats()->clear_storage_size_before();
   result.mutable_optimize_stats()->clear_storage_size_after();
+  result.mutable_optimize_stats()->clear_before_optimize_persist_stats();
+  result.mutable_optimize_stats()->clear_after_optimize_persist_stats();
   EXPECT_THAT(result.optimize_stats(), EqualsProto(expected));
 }
 
@@ -2059,18 +2250,76 @@
   // Delete the first document.
   ASSERT_THAT(icing->Delete(document1.namespace_(), document1.uri()).status(),
               ProtoIsOk());
-  ASSERT_THAT(icing->PersistToDisk(PersistType::FULL).status(), ProtoIsOk());
+  PersistToDiskResultProto persist_result =
+      icing->PersistToDisk(PersistType::FULL);
+  ASSERT_THAT(persist_result.status(), ProtoIsOk());
 
-  OptimizeStatsProto expected;
-  expected.set_latency_ms(5);
-  expected.set_document_store_optimize_latency_ms(5);
-  expected.set_index_restoration_latency_ms(5);
-  expected.set_num_original_documents(3);
-  expected.set_num_deleted_documents(1);
-  expected.set_num_expired_documents(1);
-  expected.set_num_original_namespaces(1);
-  expected.set_num_deleted_namespaces(0);
-  expected.set_index_restoration_mode(OptimizeStatsProto::FULL_INDEX_REBUILD);
+  PersistToDiskStatsProto expected_persist_stats;
+  expected_persist_stats.set_persist_type(PersistType::FULL);
+  expected_persist_stats.set_latency_ms(5);
+  expected_persist_stats.set_schema_store_persist_latency_ms(5);
+  expected_persist_stats.set_document_store_total_persist_latency_ms(5);
+  expected_persist_stats.set_document_store_components_persist_latency_ms(5);
+  expected_persist_stats.set_document_store_checksum_update_latency_ms(5);
+  expected_persist_stats.set_document_log_checksum_update_latency_ms(5);
+  expected_persist_stats.set_document_log_data_sync_latency_ms(5);
+  expected_persist_stats.set_index_persist_latency_ms(5);
+  expected_persist_stats.set_integer_index_persist_latency_ms(5);
+  expected_persist_stats.set_qualified_id_join_index_persist_latency_ms(5);
+  expected_persist_stats.set_embedding_index_persist_latency_ms(5);
+  EXPECT_THAT(persist_result.persist_stats(),
+              EqualsProto(expected_persist_stats));
+
+  OptimizeStatsProto expected_optimize_stats;
+  expected_optimize_stats.set_latency_ms(5);
+  expected_optimize_stats.set_document_store_optimize_latency_ms(5);
+  expected_optimize_stats.set_index_restoration_latency_ms(5);
+  expected_optimize_stats.set_num_original_documents(3);
+  expected_optimize_stats.set_num_deleted_documents(1);
+  expected_optimize_stats.set_num_expired_documents(1);
+  expected_optimize_stats.set_num_original_namespaces(1);
+  expected_optimize_stats.set_num_deleted_namespaces(0);
+  expected_optimize_stats.set_index_restoration_mode(
+      OptimizeStatsProto::FULL_INDEX_REBUILD);
+
+  PersistToDiskStatsProto* expected_persist_stats_before_optimize =
+      expected_optimize_stats.mutable_before_optimize_persist_stats();
+  expected_persist_stats_before_optimize->set_persist_type(PersistType::FULL);
+  expected_persist_stats_before_optimize->set_latency_ms(5);
+  expected_persist_stats_before_optimize->set_schema_store_persist_latency_ms(
+      5);
+  expected_persist_stats_before_optimize
+      ->set_document_store_total_persist_latency_ms(5);
+  expected_persist_stats_before_optimize
+      ->set_document_store_components_persist_latency_ms(5);
+  expected_persist_stats_before_optimize
+      ->set_document_store_checksum_update_latency_ms(5);
+  expected_persist_stats_before_optimize->set_index_persist_latency_ms(5);
+  expected_persist_stats_before_optimize->set_integer_index_persist_latency_ms(
+      5);
+  expected_persist_stats_before_optimize
+      ->set_qualified_id_join_index_persist_latency_ms(5);
+  expected_persist_stats_before_optimize
+      ->set_embedding_index_persist_latency_ms(5);
+
+  PersistToDiskStatsProto* expected_persist_stats_after_optimize =
+      expected_optimize_stats.mutable_after_optimize_persist_stats();
+  expected_persist_stats_after_optimize->set_persist_type(PersistType::FULL);
+  expected_persist_stats_after_optimize->set_latency_ms(5);
+  expected_persist_stats_after_optimize->set_schema_store_persist_latency_ms(5);
+  expected_persist_stats_after_optimize
+      ->set_document_store_total_persist_latency_ms(5);
+  expected_persist_stats_after_optimize
+      ->set_document_store_components_persist_latency_ms(5);
+  expected_persist_stats_after_optimize
+      ->set_document_store_checksum_update_latency_ms(5);
+  expected_persist_stats_after_optimize->set_index_persist_latency_ms(5);
+  expected_persist_stats_after_optimize->set_integer_index_persist_latency_ms(
+      5);
+  expected_persist_stats_after_optimize
+      ->set_qualified_id_join_index_persist_latency_ms(5);
+  expected_persist_stats_after_optimize->set_embedding_index_persist_latency_ms(
+      5);
 
   // Run Optimize
   OptimizeResultProto result = icing->Optimize();
@@ -2084,7 +2333,7 @@
               Ge(result.optimize_stats().storage_size_after() - page_size));
   result.mutable_optimize_stats()->clear_storage_size_before();
   result.mutable_optimize_stats()->clear_storage_size_after();
-  EXPECT_THAT(result.optimize_stats(), EqualsProto(expected));
+  EXPECT_THAT(result.optimize_stats(), EqualsProto(expected_optimize_stats));
 
   fake_clock = std::make_unique<FakeClock>();
   fake_clock->SetTimerElapsedMilliseconds(5);
@@ -2097,18 +2346,58 @@
       GetTestJniCache());
   ASSERT_THAT(icing->Initialize().status(), ProtoIsOk());
 
-  expected = OptimizeStatsProto();
-  expected.set_latency_ms(5);
-  expected.set_document_store_optimize_latency_ms(5);
-  expected.set_index_restoration_latency_ms(5);
-  expected.set_num_original_documents(1);
-  expected.set_num_deleted_documents(0);
-  expected.set_num_expired_documents(0);
-  expected.set_num_original_namespaces(1);
-  expected.set_num_deleted_namespaces(0);
-  expected.set_time_since_last_optimize_ms(10000);
-  expected.set_time_since_last_successful_optimize_ms(10000);
-  expected.set_index_restoration_mode(OptimizeStatsProto::FULL_INDEX_REBUILD);
+  expected_optimize_stats = OptimizeStatsProto();
+  expected_optimize_stats.set_latency_ms(5);
+  expected_optimize_stats.set_document_store_optimize_latency_ms(5);
+  expected_optimize_stats.set_index_restoration_latency_ms(5);
+  expected_optimize_stats.set_num_original_documents(1);
+  expected_optimize_stats.set_num_deleted_documents(0);
+  expected_optimize_stats.set_num_expired_documents(0);
+  expected_optimize_stats.set_num_original_namespaces(1);
+  expected_optimize_stats.set_num_deleted_namespaces(0);
+  expected_optimize_stats.set_time_since_last_optimize_ms(10000);
+  expected_optimize_stats.set_time_since_last_successful_optimize_ms(10000);
+  expected_optimize_stats.set_index_restoration_mode(
+      OptimizeStatsProto::FULL_INDEX_REBUILD);
+
+  expected_persist_stats_before_optimize =
+      expected_optimize_stats.mutable_before_optimize_persist_stats();
+  expected_persist_stats_before_optimize->set_persist_type(PersistType::FULL);
+  expected_persist_stats_before_optimize->set_latency_ms(5);
+  expected_persist_stats_before_optimize->set_schema_store_persist_latency_ms(
+      5);
+  expected_persist_stats_before_optimize
+      ->set_document_store_total_persist_latency_ms(5);
+  expected_persist_stats_before_optimize
+      ->set_document_store_components_persist_latency_ms(5);
+  expected_persist_stats_before_optimize
+      ->set_document_store_checksum_update_latency_ms(5);
+  expected_persist_stats_before_optimize->set_index_persist_latency_ms(5);
+  expected_persist_stats_before_optimize->set_integer_index_persist_latency_ms(
+      5);
+  expected_persist_stats_before_optimize
+      ->set_qualified_id_join_index_persist_latency_ms(5);
+  expected_persist_stats_before_optimize
+      ->set_embedding_index_persist_latency_ms(5);
+
+  expected_persist_stats_after_optimize =
+      expected_optimize_stats.mutable_after_optimize_persist_stats();
+  expected_persist_stats_after_optimize->set_persist_type(PersistType::FULL);
+  expected_persist_stats_after_optimize->set_latency_ms(5);
+  expected_persist_stats_after_optimize->set_schema_store_persist_latency_ms(5);
+  expected_persist_stats_after_optimize
+      ->set_document_store_total_persist_latency_ms(5);
+  expected_persist_stats_after_optimize
+      ->set_document_store_components_persist_latency_ms(5);
+  expected_persist_stats_after_optimize
+      ->set_document_store_checksum_update_latency_ms(5);
+  expected_persist_stats_after_optimize->set_index_persist_latency_ms(5);
+  expected_persist_stats_after_optimize->set_integer_index_persist_latency_ms(
+      5);
+  expected_persist_stats_after_optimize
+      ->set_qualified_id_join_index_persist_latency_ms(5);
+  expected_persist_stats_after_optimize->set_embedding_index_persist_latency_ms(
+      5);
 
   // Run Optimize
   result = icing->Optimize();
@@ -2116,24 +2405,65 @@
               Eq(result.optimize_stats().storage_size_after()));
   result.mutable_optimize_stats()->clear_storage_size_before();
   result.mutable_optimize_stats()->clear_storage_size_after();
-  EXPECT_THAT(result.optimize_stats(), EqualsProto(expected));
+  EXPECT_THAT(result.optimize_stats(), EqualsProto(expected_optimize_stats));
 
   // Delete the last document.
   ASSERT_THAT(icing->Delete(document3.namespace_(), document3.uri()).status(),
               ProtoIsOk());
 
-  expected = OptimizeStatsProto();
-  expected.set_latency_ms(5);
-  expected.set_document_store_optimize_latency_ms(5);
-  expected.set_index_restoration_latency_ms(5);
-  expected.set_num_original_documents(1);
-  expected.set_num_deleted_documents(1);
-  expected.set_num_expired_documents(0);
-  expected.set_num_original_namespaces(1);
-  expected.set_num_deleted_namespaces(1);
-  expected.set_time_since_last_optimize_ms(0);
-  expected.set_time_since_last_successful_optimize_ms(0);
-  expected.set_index_restoration_mode(OptimizeStatsProto::FULL_INDEX_REBUILD);
+  expected_optimize_stats = OptimizeStatsProto();
+  expected_optimize_stats.set_latency_ms(5);
+  expected_optimize_stats.set_document_store_optimize_latency_ms(5);
+  expected_optimize_stats.set_index_restoration_latency_ms(5);
+  expected_optimize_stats.set_num_original_documents(1);
+  expected_optimize_stats.set_num_deleted_documents(1);
+  expected_optimize_stats.set_num_expired_documents(0);
+  expected_optimize_stats.set_num_original_namespaces(1);
+  expected_optimize_stats.set_num_deleted_namespaces(1);
+  expected_optimize_stats.set_time_since_last_optimize_ms(0);
+  expected_optimize_stats.set_time_since_last_successful_optimize_ms(0);
+  expected_optimize_stats.set_index_restoration_mode(
+      OptimizeStatsProto::FULL_INDEX_REBUILD);
+
+  expected_persist_stats_before_optimize =
+      expected_optimize_stats.mutable_before_optimize_persist_stats();
+  expected_persist_stats_before_optimize->set_persist_type(PersistType::FULL);
+  expected_persist_stats_before_optimize->set_latency_ms(5);
+  expected_persist_stats_before_optimize->set_schema_store_persist_latency_ms(
+      5);
+  expected_persist_stats_before_optimize
+      ->set_document_store_total_persist_latency_ms(5);
+
+  expected_persist_stats_before_optimize
+      ->set_document_store_components_persist_latency_ms(5);
+  expected_persist_stats_before_optimize
+      ->set_document_store_checksum_update_latency_ms(5);
+  expected_persist_stats_before_optimize->set_index_persist_latency_ms(5);
+  expected_persist_stats_before_optimize->set_integer_index_persist_latency_ms(
+      5);
+  expected_persist_stats_before_optimize
+      ->set_qualified_id_join_index_persist_latency_ms(5);
+  expected_persist_stats_before_optimize
+      ->set_embedding_index_persist_latency_ms(5);
+
+  expected_persist_stats_after_optimize =
+      expected_optimize_stats.mutable_after_optimize_persist_stats();
+  expected_persist_stats_after_optimize->set_persist_type(PersistType::FULL);
+  expected_persist_stats_after_optimize->set_latency_ms(5);
+  expected_persist_stats_after_optimize->set_schema_store_persist_latency_ms(5);
+  expected_persist_stats_after_optimize
+      ->set_document_store_total_persist_latency_ms(5);
+  expected_persist_stats_after_optimize
+      ->set_document_store_components_persist_latency_ms(5);
+  expected_persist_stats_after_optimize
+      ->set_document_store_checksum_update_latency_ms(5);
+  expected_persist_stats_after_optimize->set_index_persist_latency_ms(5);
+  expected_persist_stats_after_optimize->set_integer_index_persist_latency_ms(
+      5);
+  expected_persist_stats_after_optimize
+      ->set_qualified_id_join_index_persist_latency_ms(5);
+  expected_persist_stats_after_optimize->set_embedding_index_persist_latency_ms(
+      5);
 
   // Run Optimize
   result = icing->Optimize();
@@ -2141,7 +2471,7 @@
               Ge(result.optimize_stats().storage_size_after()));
   result.mutable_optimize_stats()->clear_storage_size_before();
   result.mutable_optimize_stats()->clear_storage_size_after();
-  EXPECT_THAT(result.optimize_stats(), EqualsProto(expected));
+  EXPECT_THAT(result.optimize_stats(), EqualsProto(expected_optimize_stats));
 }
 
 TEST_F(IcingSearchEngineOptimizeTest,
@@ -2170,6 +2500,7 @@
           .Build();
   IcingSearchEngineOptions icing_options = GetDefaultIcingOptions();
   icing_options.set_compression_level(3);
+  icing_options.set_compression_threshold_bytes(0);
   int64_t document_log_size_compression_3;
   int64_t document_log_size_after_opti_no_compression;
   int64_t document_log_size_after_opti_compression_3;

diff --git a/icing/icing-search-engine_put_test.cc b/icing/icing-search-engine_put_test.cc
index 37df1a1..ad76b98 100644
--- a/icing/icing-search-engine_put_test.cc
+++ b/icing/icing-search-engine_put_test.cc

@@ -12,10 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <chrono>  // NOLINT
 #include <cstdint>
 #include <limits>
 #include <memory>
 #include <string>
+#include <thread>  // NOLINT
 #include <utility>
 
 #include "icing/text_classifier/lib3/utils/base/status.h"
@@ -45,6 +47,7 @@
 #include "icing/proto/term.pb.h"
 #include "icing/proto/usage.pb.h"
 #include "icing/schema-builder.h"
+#include "icing/store/document-log-creator.h"
 #include "icing/testing/common-matchers.h"
 #include "icing/testing/fake-clock.h"
 #include "icing/testing/jni-test-helpers.h"
@@ -136,6 +139,10 @@
 IcingSearchEngineOptions GetDefaultIcingOptions() {
   IcingSearchEngineOptions icing_options;
   icing_options.set_base_dir(GetTestBaseDir());
+  icing_options.set_enable_repeated_field_joins(true);
+  icing_options.set_enable_soft_index_restoration(true);
+  icing_options.set_enable_qualified_id_join_index_v3(true);
+  icing_options.set_enable_delete_propagation_from(true);
   return icing_options;
 }
 
@@ -602,6 +609,1108 @@
               ProtoStatusIs(StatusProto::INVALID_ARGUMENT));
 }
 
+TEST_F(IcingSearchEnginePutTest, DocumentCompressionThreshold) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Message").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("body")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_REQUIRED)))
+          .Build();
+
+  // Create document1 with size of 1024 bytes.
+  int64_t document1_size = 1024;
+  DocumentProto document1 =
+      DocumentBuilder()
+          .SetKey("namespace", "uri1")
+          .SetSchema("Message")
+          .AddStringProperty("body", std::string(979, 'a'))
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .Build();
+  ASSERT_THAT(document1.ByteSizeLong(), Eq(document1_size));
+
+  // Create document2 with size of 900 bytes.
+  int64_t document2_size = 900;
+  DocumentProto document2 =
+      DocumentBuilder()
+          .SetKey("namespace", "uri2")
+          .SetSchema("Message")
+          .AddStringProperty("body", std::string(855, 'a'))
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .Build();
+  ASSERT_THAT(document2.ByteSizeLong(), Eq(document2_size));
+
+  // Create icing instance with compression threshold of 1000 bytes.
+  IcingSearchEngineOptions icing_options = GetDefaultIcingOptions();
+  icing_options.set_compression_level(3);
+
+  int64_t document_log_size_full_compression;
+  int64_t document_log_size_part_compression;
+  int64_t document_log_size_no_compression;
+  const std::string document_log_path =
+      icing_options.base_dir() + "/document_dir/" +
+      DocumentLogCreator::GetDocumentLogFilename();
+
+  // Set the compression threshold to 1000 bytes to only compress document1.
+  {
+    icing_options.set_compression_threshold_bytes(1000);
+    IcingSearchEngine icing(icing_options, GetTestJniCache());
+    ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+    ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(document1).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(document2).status(), ProtoIsOk());
+    ASSERT_THAT(icing.PersistToDisk(PersistType::FULL).status(), ProtoIsOk());
+
+    document_log_size_part_compression =
+        filesystem()->GetFileSize(document_log_path.c_str());
+
+    // Calling Optimize() will not change the size of the document log since
+    // the compression setting is the same.
+    ASSERT_THAT(icing.Optimize().status(), ProtoIsOk());
+    ASSERT_EQ(filesystem()->GetFileSize(document_log_path.c_str()),
+              document_log_size_part_compression);
+  }
+
+  // Set the compression threshold to 900 bytes to compress both documents.
+  {
+    icing_options.set_compression_threshold_bytes(900);
+    IcingSearchEngine icing(icing_options, GetTestJniCache());
+    ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+
+    // Call Optimize() to refresh the document log with the new compression
+    // threshold.
+    ASSERT_THAT(icing.Optimize().status(), ProtoIsOk());
+    document_log_size_full_compression =
+        filesystem()->GetFileSize(document_log_path.c_str());
+
+    // Calling Optimize() again will not change the size of the document log
+    // since the compression setting is the same.
+    ASSERT_THAT(icing.Optimize().status(), ProtoIsOk());
+    ASSERT_EQ(filesystem()->GetFileSize(document_log_path.c_str()),
+              document_log_size_full_compression);
+  }
+
+  // Set the compression threshold to 10000 bytes will turn off compression.
+  {
+    icing_options.set_compression_threshold_bytes(10000);
+    IcingSearchEngine icing(icing_options, GetTestJniCache());
+    ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+
+    // Call Optimize() to refresh the document log with the new compression
+    // threshold.
+    ASSERT_THAT(icing.Optimize().status(), ProtoIsOk());
+    document_log_size_no_compression =
+        filesystem()->GetFileSize(document_log_path.c_str());
+
+    // Calling Optimize() again will not change the size of the document log
+    // since the compression setting is the same.
+    ASSERT_THAT(icing.Optimize().status(), ProtoIsOk());
+    ASSERT_EQ(filesystem()->GetFileSize(document_log_path.c_str()),
+              document_log_size_no_compression);
+  }
+
+  // Check that no_compression > part_compression > full_compression
+  ASSERT_GT(document_log_size_part_compression,
+            document_log_size_full_compression);
+  ASSERT_GT(document_log_size_no_compression,
+            document_log_size_part_compression);
+}
+
+TEST_F(
+    IcingSearchEnginePutTest,
+    PutDocument_taskSchedulerDisabled_shouldNotReschedulePurgingExpirationTask) {
+  // This test verifies that when enable_delete_propagation_from is true but
+  // enable_background_task_scheduler is false, then:
+  // - No background tasks are scheduled.
+  // - Since task_scheduler_ is null, APIs attempt to (re)schedule a task should
+  //   check the nullness of task_scheduler_ before using it.
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("name")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .AddType(
+              SchemaTypeConfigBuilder()
+                  .SetType("Email")
+                  .AddProperty(
+                      PropertyConfigBuilder()
+                          .SetName("subject")
+                          .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                          .SetCardinality(CARDINALITY_OPTIONAL))
+                  .AddProperty(PropertyConfigBuilder()
+                                   .SetName("sender")
+                                   .SetDataTypeJoinableString(
+                                       JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                       DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                                   .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  DocumentProto person = DocumentBuilder()
+                             .SetKey("namespace", "person")
+                             .SetSchema("Person")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(1000)  // Expire at 1010 ms.
+                             .AddStringProperty("name", "Alice")
+                             .Build();
+  DocumentProto email = DocumentBuilder()
+                            .SetKey("namespace", "email")
+                            .SetSchema("Email")
+                            .SetCreationTimestampMs(10)
+                            .SetTtlMs(0)  // Never expire.
+                            .AddStringProperty("subject", "test")
+                            .AddStringProperty("sender", "namespace#person")
+                            .Build();
+
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  options.set_enable_background_task_scheduler(false);
+  options.set_enable_qualified_id_join_index_v3(true);
+  options.set_enable_soft_index_restoration(true);
+  options.set_enable_delete_propagation_from(true);
+  options.set_expired_document_purge_threshold_ms(0);
+
+  // Initialize Icing with a fake clock and t = 10 ms.
+  auto fake_clock = std::make_unique<FakeClock>();
+  FakeClock* fake_clock_ptr = fake_clock.get();
+  fake_clock->SetSystemTimeMilliseconds(10);
+  TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                              std::make_unique<IcingFilesystem>(),
+                              std::move(fake_clock), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+
+  // Put person and email. Since enable_background_task_scheduler is false,
+  // there should be no background tasks scheduled.
+  ASSERT_THAT(icing.Put(person).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(email).status(), ProtoIsOk());
+
+  // Sanity check that person and email are present.
+  GetResultProto expected_get_result_proto1;
+  expected_get_result_proto1.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_proto1.mutable_document() = person;
+  EXPECT_THAT(
+      icing.Get("namespace", "person", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto1));
+
+  GetResultProto expected_get_result_google::protobuf;
+  expected_get_result_google::protobuf.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_google::protobuf.mutable_document() = email;
+  EXPECT_THAT(
+      icing.Get("namespace", "email", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_google::protobuf));
+
+  // Adjust the clock to 1010 ms and sleep for 1010 ms. Email document should
+  // still be present since no background tasks were and therefore, expiration
+  // propagation and purging did not run for the child document.
+  fake_clock_ptr->SetSystemTimeMilliseconds(1010);
+  std::this_thread::sleep_for(std::chrono::milliseconds(1010));
+
+  GetResultProto expected_get_result_proto4;
+  expected_get_result_proto4.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_proto4.mutable_document() = email;
+  EXPECT_THAT(
+      icing.Get("namespace", "email", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto4));
+}
+
+TEST_F(IcingSearchEnginePutTest,
+       PutDocument_shouldReschedulePurgingExpirationTaskForNewExpTs) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("name")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .AddType(
+              SchemaTypeConfigBuilder()
+                  .SetType("Email")
+                  .AddProperty(
+                      PropertyConfigBuilder()
+                          .SetName("subject")
+                          .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                          .SetCardinality(CARDINALITY_OPTIONAL))
+                  .AddProperty(PropertyConfigBuilder()
+                                   .SetName("sender")
+                                   .SetDataTypeJoinableString(
+                                       JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                       DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                                   .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  DocumentProto person = DocumentBuilder()
+                             .SetKey("namespace", "person")
+                             .SetSchema("Person")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(1000)  // Expire at 1010 ms.
+                             .AddStringProperty("name", "Alice")
+                             .Build();
+  DocumentProto email = DocumentBuilder()
+                            .SetKey("namespace", "email")
+                            .SetSchema("Email")
+                            .SetCreationTimestampMs(10)
+                            .SetTtlMs(0)  // Never expire.
+                            .AddStringProperty("subject", "test")
+                            .AddStringProperty("sender", "namespace#person")
+                            .Build();
+
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  options.set_enable_background_task_scheduler(true);
+  options.set_enable_qualified_id_join_index_v3(true);
+  options.set_enable_soft_index_restoration(true);
+  options.set_enable_delete_propagation_from(true);
+  options.set_expired_document_purge_threshold_ms(0);
+
+  // Initialize Icing with a fake clock and t = 10 ms. At this moment, there is
+  // no expiration document so the purging expired document task is not
+  // scheduled.
+  auto fake_clock = std::make_unique<FakeClock>();
+  FakeClock* fake_clock_ptr = fake_clock.get();
+  fake_clock->SetSystemTimeMilliseconds(10);
+  TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                              std::make_unique<IcingFilesystem>(),
+                              std::move(fake_clock), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+
+  // Put person and email. The scheduler should schedule the purging expiration
+  // task at t = 1010 ms.
+  ASSERT_THAT(icing.Put(person).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(email).status(), ProtoIsOk());
+
+  // Sanity check that person and email are present.
+  GetResultProto expected_get_result_proto1;
+  expected_get_result_proto1.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_proto1.mutable_document() = person;
+  EXPECT_THAT(
+      icing.Get("namespace", "person", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto1));
+
+  GetResultProto expected_get_result_google::protobuf;
+  expected_get_result_google::protobuf.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_google::protobuf.mutable_document() = email;
+  EXPECT_THAT(
+      icing.Get("namespace", "email", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_google::protobuf));
+
+  // Adjust the clock to 1010 ms and sleep for 1010 ms. The purging expiration
+  // task should execute and purge person and email.
+  // - At t = 1010 ms, person is expired, but DocumentFilterData has already
+  //   filtered it out, so we're unsure whether it is purged or not.
+  // - But the child document email never expires and DocumentFilterData
+  //   doesn't filter it out. If the scheduled task had not worked correctly,
+  //   then expiration propagation would have not been triggered and email
+  //   would've been alive. Therefore, we can verify it by checking email.
+  fake_clock_ptr->SetSystemTimeMilliseconds(1010);
+  std::this_thread::sleep_for(std::chrono::milliseconds(1010));
+
+  GetResultProto expected_get_result_proto3;
+  expected_get_result_proto3.mutable_status()->set_code(StatusProto::NOT_FOUND);
+  expected_get_result_proto3.mutable_status()->set_message(
+      "Document (namespace, person) not found.");
+  EXPECT_THAT(
+      icing.Get("namespace", "person", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto3));
+
+  GetResultProto expected_get_result_proto4;
+  expected_get_result_proto4.mutable_status()->set_code(StatusProto::NOT_FOUND);
+  expected_get_result_proto4.mutable_status()->set_message(
+      "Document (namespace, email) not found.");
+  EXPECT_THAT(
+      icing.Get("namespace", "email", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto4));
+}
+
+TEST_F(IcingSearchEnginePutTest,
+       PutDocument_shouldReschedulePurgingExpirationTaskForSmallerExpTs) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("name")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .AddType(
+              SchemaTypeConfigBuilder()
+                  .SetType("Email")
+                  .AddProperty(
+                      PropertyConfigBuilder()
+                          .SetName("subject")
+                          .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                          .SetCardinality(CARDINALITY_OPTIONAL))
+                  .AddProperty(PropertyConfigBuilder()
+                                   .SetName("sender")
+                                   .SetDataTypeJoinableString(
+                                       JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                       DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                                   .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  DocumentProto person1 = DocumentBuilder()
+                              .SetKey("namespace", "person1")
+                              .SetSchema("Person")
+                              .SetCreationTimestampMs(10)
+                              .SetTtlMs(2000000)  // Expire at 2000010 ms.
+                              .AddStringProperty("name", "Alice")
+                              .Build();
+  DocumentProto person2 = DocumentBuilder()
+                              .SetKey("namespace", "person2")
+                              .SetSchema("Person")
+                              .SetCreationTimestampMs(10)
+                              .SetTtlMs(1000)  // Expire at 1010 ms.
+                              .AddStringProperty("name", "Bob")
+                              .Build();
+  DocumentProto email1 = DocumentBuilder()
+                             .SetKey("namespace", "email1")
+                             .SetSchema("Email")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(0)  // Never expire.
+                             .AddStringProperty("subject", "test")
+                             .AddStringProperty("sender", "namespace#person1")
+                             .Build();
+  DocumentProto email2 = DocumentBuilder()
+                             .SetKey("namespace", "email2")
+                             .SetSchema("Email")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(0)  // Never expire.
+                             .AddStringProperty("subject", "test")
+                             .AddStringProperty("sender", "namespace#person2")
+                             .Build();
+
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  options.set_enable_background_task_scheduler(true);
+  options.set_enable_qualified_id_join_index_v3(true);
+  options.set_enable_soft_index_restoration(true);
+  options.set_enable_delete_propagation_from(true);
+  options.set_expired_document_purge_threshold_ms(0);
+
+  {
+    // Initialize Icing and put person1 and email1. Destruct Icing.
+    TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                                std::make_unique<IcingFilesystem>(),
+                                std::make_unique<FakeClock>(),
+                                GetTestJniCache());
+    ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+    ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(person1).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(email1).status(), ProtoIsOk());
+  }
+
+  // Initialize Icing again with a fake clock and t = 10 ms. Initialization
+  // should schedule the purging expiration task at t = 2000010 ms.
+  auto fake_clock = std::make_unique<FakeClock>();
+  FakeClock* fake_clock_ptr = fake_clock.get();
+  fake_clock->SetSystemTimeMilliseconds(10);
+  TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                              std::make_unique<IcingFilesystem>(),
+                              std::move(fake_clock), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+
+  // Put person2 and email2. It should reschedule the purging expiration task at
+  // t = 1010 ms.
+  ASSERT_THAT(icing.Put(person2).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(email2).status(), ProtoIsOk());
+
+  // Sanity check that person2 and email2 are present.
+  GetResultProto expected_get_result_proto1;
+  expected_get_result_proto1.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_proto1.mutable_document() = person2;
+  EXPECT_THAT(
+      icing.Get("namespace", "person2", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto1));
+
+  GetResultProto expected_get_result_google::protobuf;
+  expected_get_result_google::protobuf.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_google::protobuf.mutable_document() = email2;
+  EXPECT_THAT(
+      icing.Get("namespace", "email2", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_google::protobuf));
+
+  // Adjust the clock to 1010 ms and sleep for 1010 ms. The purging expiration
+  // task should execute and purge person2 and email2.
+  // - At t = 1010 ms, person2 is expired, but DocumentFilterData has already
+  //   filtered it out, so we're unsure whether it is purged or not.
+  // - But the child document email2 never expires and DocumentFilterData
+  //   doesn't filter it out. If the scheduled task had not worked correctly,
+  //   then expiration propagation would have not been triggered and email2
+  //   would've been alive. Therefore, we can verify it by checking email2.
+  fake_clock_ptr->SetSystemTimeMilliseconds(1010);
+  std::this_thread::sleep_for(std::chrono::milliseconds(1010));
+
+  GetResultProto expected_get_result_proto3;
+  expected_get_result_proto3.mutable_status()->set_code(StatusProto::NOT_FOUND);
+  expected_get_result_proto3.mutable_status()->set_message(
+      "Document (namespace, person2) not found.");
+  EXPECT_THAT(
+      icing.Get("namespace", "person2", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto3));
+
+  GetResultProto expected_get_result_proto4;
+  expected_get_result_proto4.mutable_status()->set_code(StatusProto::NOT_FOUND);
+  expected_get_result_proto4.mutable_status()->set_message(
+      "Document (namespace, email2) not found.");
+  EXPECT_THAT(
+      icing.Get("namespace", "email2", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto4));
+}
+
+TEST_F(IcingSearchEnginePutTest,
+       PutDocument_shouldNotReschedulePurgingExpirationTaskForGreaterExpTs) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("name")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .AddType(
+              SchemaTypeConfigBuilder()
+                  .SetType("Email")
+                  .AddProperty(
+                      PropertyConfigBuilder()
+                          .SetName("subject")
+                          .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                          .SetCardinality(CARDINALITY_OPTIONAL))
+                  .AddProperty(PropertyConfigBuilder()
+                                   .SetName("sender")
+                                   .SetDataTypeJoinableString(
+                                       JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                       DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                                   .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  DocumentProto person1 = DocumentBuilder()
+                              .SetKey("namespace", "person1")
+                              .SetSchema("Person")
+                              .SetCreationTimestampMs(10)
+                              .SetTtlMs(10000)  // Expire at 10010 ms.
+                              .AddStringProperty("name", "Alice")
+                              .Build();
+  DocumentProto person2 = DocumentBuilder()
+                              .SetKey("namespace", "person2")
+                              .SetSchema("Person")
+                              .SetCreationTimestampMs(10)
+                              .SetTtlMs(2000000)  // Expire at 2000010 ms.
+                              .AddStringProperty("name", "Bob")
+                              .Build();
+  DocumentProto email1 = DocumentBuilder()
+                             .SetKey("namespace", "email1")
+                             .SetSchema("Email")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(0)  // Never expire.
+                             .AddStringProperty("subject", "test")
+                             .AddStringProperty("sender", "namespace#person1")
+                             .Build();
+  DocumentProto email2 = DocumentBuilder()
+                             .SetKey("namespace", "email2")
+                             .SetSchema("Email")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(0)  // Never expire.
+                             .AddStringProperty("subject", "test")
+                             .AddStringProperty("sender", "namespace#person2")
+                             .Build();
+
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  options.set_enable_background_task_scheduler(true);
+  options.set_enable_qualified_id_join_index_v3(true);
+  options.set_enable_soft_index_restoration(true);
+  options.set_enable_delete_propagation_from(true);
+  options.set_expired_document_purge_threshold_ms(0);
+
+  {
+    // Initialize Icing and put person1 and email1. Destruct Icing.
+    TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                                std::make_unique<IcingFilesystem>(),
+                                std::make_unique<FakeClock>(),
+                                GetTestJniCache());
+    ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+    ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(person1).status(), ProtoIsOk());
+    ASSERT_THAT(icing.Put(email1).status(), ProtoIsOk());
+  }
+
+  // Initialize Icing again with a fake clock and t = 9010 ms. Initialization
+  // should schedule the purging expiration task at t = 10010 ms.
+  auto fake_clock = std::make_unique<FakeClock>();
+  FakeClock* fake_clock_ptr = fake_clock.get();
+  fake_clock->SetSystemTimeMilliseconds(9010);
+  TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                              std::make_unique<IcingFilesystem>(),
+                              std::move(fake_clock), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+
+  // Put person2 and email2. It should NOT reschedule the purging expiration
+  // task since the new expiration timestamp is greater than the current
+  // scheduled one.
+  ASSERT_THAT(icing.Put(person2).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(email2).status(), ProtoIsOk());
+
+  // Sanity check that person1 and email1 are present.
+  GetResultProto expected_get_result_proto1;
+  expected_get_result_proto1.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_proto1.mutable_document() = person1;
+  EXPECT_THAT(
+      icing.Get("namespace", "person1", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto1));
+
+  GetResultProto expected_get_result_google::protobuf;
+  expected_get_result_google::protobuf.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_google::protobuf.mutable_document() = email1;
+  EXPECT_THAT(
+      icing.Get("namespace", "email1", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_google::protobuf));
+
+  // Adjust the clock to 10010 ms and sleep for 1100 ms. The purging expiration
+  // task should execute and purge person1 and email1.
+  // - If Put API had falsely rescheduled the task, then at t = 10010 ms person1
+  //   and email1 would've been alive. We need to verify that the task was not
+  //   rescheduled with the greater expiration timestamp.
+  // - At t = 10010 ms, person1 is expired, but DocumentFilterData has already
+  //   filtered it out, so we're unsure whether it is purged or not.
+  // - But the child document email1 never expires and DocumentFilterData
+  //   doesn't filter it out. If the scheduled task had not worked correctly,
+  //   then expiration propagation would have not been triggered and email1
+  //   would've been alive. Therefore, we can verify it by checking email1.
+  fake_clock_ptr->SetSystemTimeMilliseconds(10010);
+  std::this_thread::sleep_for(std::chrono::milliseconds(1100));
+
+  GetResultProto expected_get_result_proto3;
+  expected_get_result_proto3.mutable_status()->set_code(StatusProto::NOT_FOUND);
+  expected_get_result_proto3.mutable_status()->set_message(
+      "Document (namespace, person1) not found.");
+  EXPECT_THAT(
+      icing.Get("namespace", "person1", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto3));
+
+  GetResultProto expected_get_result_proto4;
+  expected_get_result_proto4.mutable_status()->set_code(StatusProto::NOT_FOUND);
+  expected_get_result_proto4.mutable_status()->set_message(
+      "Document (namespace, email1) not found.");
+  EXPECT_THAT(
+      icing.Get("namespace", "email1", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto4));
+}
+
+TEST_F(IcingSearchEnginePutTest,
+       PutDocument_shouldSucceedWithDeletePropagationEnabledAndValidParent) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("name")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .AddType(
+              SchemaTypeConfigBuilder()
+                  .SetType("Email")
+                  .AddProperty(
+                      PropertyConfigBuilder()
+                          .SetName("subject")
+                          .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                          .SetCardinality(CARDINALITY_OPTIONAL))
+                  .AddProperty(PropertyConfigBuilder()
+                                   .SetName("sender")
+                                   .SetDataTypeJoinableString(
+                                       JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                       DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                                   .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  DocumentProto person = DocumentBuilder()
+                             .SetKey("namespace", "person")
+                             .SetSchema("Person")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(10000)  // Expire at 10010 ms.
+                             .AddStringProperty("name", "Alice")
+                             .Build();
+  DocumentProto email = DocumentBuilder()
+                            .SetKey("namespace", "email")
+                            .SetSchema("Email")
+                            .SetCreationTimestampMs(10)
+                            .SetTtlMs(0)  // Never expire.
+                            .AddStringProperty("subject", "test")
+                            .AddStringProperty("sender", "namespace#person")
+                            .Build();
+
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  options.set_enable_qualified_id_join_index_v3(true);
+  options.set_enable_soft_index_restoration(true);
+  options.set_enable_delete_propagation_from(true);
+  options.set_expired_document_purge_threshold_ms(0);
+
+  // Initialize Icing.
+  auto fake_clock = std::make_unique<FakeClock>();
+  fake_clock->SetSystemTimeMilliseconds(10);
+  TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                              std::make_unique<IcingFilesystem>(),
+                              std::move(fake_clock), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(person).status(), ProtoIsOk());
+
+  // Put email document should succeed.
+  EXPECT_THAT(icing.Put(email).status(), ProtoIsOk());
+
+  // Can get email document.
+  GetResultProto expected_get_result_proto;
+  expected_get_result_proto.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_proto.mutable_document() = email;
+  EXPECT_THAT(
+      icing.Get("namespace", "email", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto));
+}
+
+TEST_F(
+    IcingSearchEnginePutTest,
+    PutDocument_shouldSucceedWithDeletePropagationEnabledAndEmptyQualifiedId) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("name")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .AddType(
+              SchemaTypeConfigBuilder()
+                  .SetType("Email")
+                  .AddProperty(
+                      PropertyConfigBuilder()
+                          .SetName("subject")
+                          .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                          .SetCardinality(CARDINALITY_OPTIONAL))
+                  .AddProperty(PropertyConfigBuilder()
+                                   .SetName("sender")
+                                   .SetDataTypeJoinableString(
+                                       JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                       DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                                   .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  DocumentProto email =
+      DocumentBuilder()
+          .SetKey("namespace", "email")
+          .SetSchema("Email")
+          .SetCreationTimestampMs(10)
+          .SetTtlMs(0)  // Never expire.
+          .AddStringProperty("subject", "test")
+          .AddStringProperty("sender", "")  // Empty qualified id.
+          .Build();
+
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  options.set_enable_qualified_id_join_index_v3(true);
+  options.set_enable_soft_index_restoration(true);
+  options.set_enable_delete_propagation_from(true);
+  options.set_expired_document_purge_threshold_ms(0);
+
+  // Initialize Icing.
+  auto fake_clock = std::make_unique<FakeClock>();
+  fake_clock->SetSystemTimeMilliseconds(10);
+  TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                              std::make_unique<IcingFilesystem>(),
+                              std::move(fake_clock), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+
+  // Put email document should succeed.
+  EXPECT_THAT(icing.Put(email).status(), ProtoIsOk());
+
+  // Can get email document.
+  GetResultProto expected_get_result_proto;
+  expected_get_result_proto.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_proto.mutable_document() = email;
+  EXPECT_THAT(
+      icing.Get("namespace", "email", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto));
+}
+
+TEST_F(
+    IcingSearchEnginePutTest,
+    PutDocument_shouldFailWithDeletePropagationEnabledAndInvalidQualifiedId) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("name")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .AddType(
+              SchemaTypeConfigBuilder()
+                  .SetType("Email")
+                  .AddProperty(
+                      PropertyConfigBuilder()
+                          .SetName("subject")
+                          .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                          .SetCardinality(CARDINALITY_OPTIONAL))
+                  .AddProperty(PropertyConfigBuilder()
+                                   .SetName("sender")
+                                   .SetDataTypeJoinableString(
+                                       JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                       DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                                   .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  DocumentProto email = DocumentBuilder()
+                            .SetKey("namespace", "email")
+                            .SetSchema("Email")
+                            .SetCreationTimestampMs(10)
+                            .SetTtlMs(0)  // Never expire.
+                            .AddStringProperty("subject", "test")
+                            .AddStringProperty("sender", "InvalidQualifiedId")
+                            .Build();
+
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  options.set_enable_qualified_id_join_index_v3(true);
+  options.set_enable_soft_index_restoration(true);
+  options.set_enable_delete_propagation_from(true);
+  options.set_expired_document_purge_threshold_ms(0);
+
+  // Initialize Icing.
+  auto fake_clock = std::make_unique<FakeClock>();
+  fake_clock->SetSystemTimeMilliseconds(10);
+  TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                              std::make_unique<IcingFilesystem>(),
+                              std::move(fake_clock), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+
+  // Put email document should fail.
+  EXPECT_THAT(icing.Put(email).status(),
+              ProtoStatusIs(StatusProto::INVALID_ARGUMENT));
+}
+
+TEST_F(IcingSearchEnginePutTest,
+       PutDocument_shouldFailWithDeletePropagationEnabledAndNonExistentParent) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("name")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .AddType(
+              SchemaTypeConfigBuilder()
+                  .SetType("Email")
+                  .AddProperty(
+                      PropertyConfigBuilder()
+                          .SetName("subject")
+                          .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                          .SetCardinality(CARDINALITY_OPTIONAL))
+                  .AddProperty(PropertyConfigBuilder()
+                                   .SetName("sender")
+                                   .SetDataTypeJoinableString(
+                                       JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                       DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                                   .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  DocumentProto email =
+      DocumentBuilder()
+          .SetKey("namespace", "email")
+          .SetSchema("Email")
+          .SetCreationTimestampMs(10)
+          .SetTtlMs(0)  // Never expire.
+          .AddStringProperty("subject", "test")
+          .AddStringProperty("sender",
+                             "namespace#person")  // Non-existent parent.
+          .Build();
+
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  options.set_enable_qualified_id_join_index_v3(true);
+  options.set_enable_soft_index_restoration(true);
+  options.set_enable_delete_propagation_from(true);
+  options.set_expired_document_purge_threshold_ms(0);
+
+  // Initialize Icing.
+  auto fake_clock = std::make_unique<FakeClock>();
+  fake_clock->SetSystemTimeMilliseconds(10);
+  TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                              std::make_unique<IcingFilesystem>(),
+                              std::move(fake_clock), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+
+  // Put email document should fail.
+  EXPECT_THAT(icing.Put(email).status(),
+              ProtoStatusIs(StatusProto::INVALID_ARGUMENT));
+}
+
+TEST_F(IcingSearchEnginePutTest,
+       PutDocument_shouldFailWithDeletePropagationEnabledAndExpiredParent) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("name")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .AddType(
+              SchemaTypeConfigBuilder()
+                  .SetType("Email")
+                  .AddProperty(
+                      PropertyConfigBuilder()
+                          .SetName("subject")
+                          .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                          .SetCardinality(CARDINALITY_OPTIONAL))
+                  .AddProperty(PropertyConfigBuilder()
+                                   .SetName("sender")
+                                   .SetDataTypeJoinableString(
+                                       JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                       DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                                   .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  DocumentProto person = DocumentBuilder()
+                             .SetKey("namespace", "person")
+                             .SetSchema("Person")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(10000)  // Expire at 10010 ms.
+                             .AddStringProperty("name", "Alice")
+                             .Build();
+  DocumentProto email = DocumentBuilder()
+                            .SetKey("namespace", "email")
+                            .SetSchema("Email")
+                            .SetCreationTimestampMs(10)
+                            .SetTtlMs(0)  // Never expire.
+                            .AddStringProperty("subject", "test")
+                            .AddStringProperty("sender", "namespace#person")
+                            .Build();
+
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  options.set_enable_qualified_id_join_index_v3(true);
+  options.set_enable_soft_index_restoration(true);
+  options.set_enable_delete_propagation_from(true);
+  options.set_expired_document_purge_threshold_ms(0);
+
+  // Initialize Icing.
+  auto fake_clock = std::make_unique<FakeClock>();
+  FakeClock* fake_clock_ptr = fake_clock.get();
+  fake_clock->SetSystemTimeMilliseconds(10);
+  TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                              std::make_unique<IcingFilesystem>(),
+                              std::move(fake_clock), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(person).status(), ProtoIsOk());
+
+  // Adjust the clock to 20000 ms and put email document. Since person document
+  // is expired, put email document should fail.
+  fake_clock_ptr->SetSystemTimeMilliseconds(20000);
+  EXPECT_THAT(icing.Put(email).status(),
+              ProtoStatusIs(StatusProto::INVALID_ARGUMENT));
+}
+
+TEST_F(IcingSearchEnginePutTest,
+       PutDocument_shouldFailWithDeletePropagationEnabledAndDeletedParent) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("name")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .AddType(
+              SchemaTypeConfigBuilder()
+                  .SetType("Email")
+                  .AddProperty(
+                      PropertyConfigBuilder()
+                          .SetName("subject")
+                          .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                          .SetCardinality(CARDINALITY_OPTIONAL))
+                  .AddProperty(PropertyConfigBuilder()
+                                   .SetName("sender")
+                                   .SetDataTypeJoinableString(
+                                       JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                       DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                                   .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  DocumentProto person = DocumentBuilder()
+                             .SetKey("namespace", "person")
+                             .SetSchema("Person")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(0)  // Never expire.
+                             .AddStringProperty("name", "Alice")
+                             .Build();
+  DocumentProto email = DocumentBuilder()
+                            .SetKey("namespace", "email")
+                            .SetSchema("Email")
+                            .SetCreationTimestampMs(10)
+                            .SetTtlMs(0)  // Never expire.
+                            .AddStringProperty("subject", "test")
+                            .AddStringProperty("sender", "namespace#person")
+                            .Build();
+
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  options.set_enable_qualified_id_join_index_v3(true);
+  options.set_enable_soft_index_restoration(true);
+  options.set_enable_delete_propagation_from(true);
+  options.set_expired_document_purge_threshold_ms(0);
+
+  // Initialize Icing.
+  auto fake_clock = std::make_unique<FakeClock>();
+  fake_clock->SetSystemTimeMilliseconds(10);
+  TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                              std::make_unique<IcingFilesystem>(),
+                              std::move(fake_clock), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(person).status(), ProtoIsOk());
+
+  // Delete person document. Put email document should fail afterwards.
+  ASSERT_THAT(icing.Delete("namespace", "person").status(), ProtoIsOk());
+  EXPECT_THAT(icing.Put(email).status(),
+              ProtoStatusIs(StatusProto::INVALID_ARGUMENT));
+}
+
+TEST_F(
+    IcingSearchEnginePutTest,
+    PutDocument_shouldSucceedWithDeletePropagationDisabledAndNonAliveParents) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("name")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .AddType(
+              SchemaTypeConfigBuilder()
+                  .SetType("Email")
+                  .AddProperty(PropertyConfigBuilder()
+                                   .SetName("subject")
+                                   .SetDataTypeString(TERM_MATCH_PREFIX,
+                                                      TOKENIZER_PLAIN)
+                                   .SetCardinality(CARDINALITY_OPTIONAL))
+                  .AddProperty(
+                      PropertyConfigBuilder()
+                          .SetName("sender")
+                          .SetDataTypeJoinableString(
+                              JOINABLE_VALUE_TYPE_QUALIFIED_ID)  // No delete
+                                                                 // propagation.
+                          .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  DocumentProto person1 = DocumentBuilder()
+                              .SetKey("namespace", "person/1")
+                              .SetSchema("Person")
+                              .SetCreationTimestampMs(10)
+                              .SetTtlMs(10000)  // Expire at 10010 ms.
+                              .AddStringProperty("name", "Alice")
+                              .Build();
+  DocumentProto person2 = DocumentBuilder()
+                              .SetKey("namespace", "person/2")
+                              .SetSchema("Person")
+                              .SetCreationTimestampMs(10)
+                              .SetTtlMs(0)  // Never expire.
+                              .AddStringProperty("name", "Bob")
+                              .Build();
+  DocumentProto email1 = DocumentBuilder()
+                             .SetKey("namespace", "email/1")
+                             .SetSchema("Email")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(0)  // Never expire.
+                             .AddStringProperty("subject", "test")
+                             .AddStringProperty("sender", "InvalidQualifiedId")
+                             .Build();
+  DocumentProto email2 =
+      DocumentBuilder()
+          .SetKey("namespace", "email/2")
+          .SetSchema("Email")
+          .SetCreationTimestampMs(10)
+          .SetTtlMs(0)  // Never expire.
+          .AddStringProperty("subject", "test")
+          .AddStringProperty(
+              "sender", "namespace#NonExistentParent")  // Non-existent parent.
+          .Build();
+  DocumentProto email3 = DocumentBuilder()
+                             .SetKey("namespace", "email/3")
+                             .SetSchema("Email")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(0)  // Never expire.
+                             .AddStringProperty("subject", "test")
+                             .AddStringProperty("sender", "namespace#person/1")
+                             .Build();
+  DocumentProto email4 = DocumentBuilder()
+                             .SetKey("namespace", "email/4")
+                             .SetSchema("Email")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(0)  // Never expire.
+                             .AddStringProperty("subject", "test")
+                             .AddStringProperty("sender", "namespace#person/2")
+                             .Build();
+
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  options.set_enable_qualified_id_join_index_v3(true);
+  options.set_enable_soft_index_restoration(true);
+  options.set_enable_delete_propagation_from(true);
+  options.set_expired_document_purge_threshold_ms(0);
+
+  // Initialize Icing.
+  auto fake_clock = std::make_unique<FakeClock>();
+  FakeClock* fake_clock_ptr = fake_clock.get();
+  fake_clock->SetSystemTimeMilliseconds(10);
+  TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                              std::make_unique<IcingFilesystem>(),
+                              std::move(fake_clock), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(person1).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(person2).status(), ProtoIsOk());
+
+  // Put email1, email2 should succeed, even though they contain invalid
+  // qualified id and non-existent parent respectively.
+  EXPECT_THAT(icing.Put(email1).status(), ProtoIsOk());
+  EXPECT_THAT(icing.Put(email2).status(), ProtoIsOk());
+
+  // Adjust the clock to 20000 ms and put email3. Although person1 is expired
+  // and email3 contains a joinable qualified id to person1, putting email3
+  // should succeed since delete propagation is disabled.
+  fake_clock_ptr->SetSystemTimeMilliseconds(20000);
+  EXPECT_THAT(icing.Put(email3).status(), ProtoIsOk());
+
+  // Delete person2. Although email4 contains a joinable qualified id to
+  // person2, putting email4 should succeed since delete propagation is
+  // disabled.
+  ASSERT_THAT(icing.Delete("namespace", "person/2").status(), ProtoIsOk());
+  EXPECT_THAT(icing.Put(email4).status(), ProtoIsOk());
+}
+
+TEST_F(IcingSearchEnginePutTest,
+       PutDocument_shouldSetExpirationTimestampInPutResultProto) {
+  auto fake_clock = std::make_unique<FakeClock>();
+  fake_clock->SetSystemTimeMilliseconds(10);
+  TestIcingSearchEngine icing(GetDefaultIcingOptions(),
+                              std::make_unique<Filesystem>(),
+                              std::make_unique<IcingFilesystem>(),
+                              std::move(fake_clock), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(CreateMessageSchema()).status(), ProtoIsOk());
+
+  DocumentProto document1 = DocumentBuilder()
+                                .SetKey("icing", "message/1")
+                                .SetSchema("Message")
+                                .SetCreationTimestampMs(10)
+                                .SetTtlMs(10000)  // Expire at 10010 ms.
+                                .AddStringProperty("body", "message body")
+                                .Build();
+  PutResultProto put_result_proto1 = icing.Put(document1);
+  ASSERT_THAT(put_result_proto1.status(), ProtoIsOk());
+  EXPECT_THAT(put_result_proto1.document_expiration_timestamp_ms(), Eq(10010));
+
+  DocumentProto document2 = DocumentBuilder()
+                                .SetKey("icing", "message/2")
+                                .SetSchema("Message")
+                                .SetCreationTimestampMs(10)
+                                .SetTtlMs(0)  // Never expire.
+                                .AddStringProperty("body", "message body")
+                                .Build();
+  PutResultProto put_result_google::protobuf = icing.Put(document2);
+  ASSERT_THAT(put_result_google::protobuf.status(), ProtoIsOk());
+  EXPECT_THAT(put_result_google::protobuf.document_expiration_timestamp_ms(),
+              Eq(std::numeric_limits<int64_t>::max()));
+}
+
 }  // namespace
 }  // namespace lib
 }  // namespace icing

diff --git a/icing/icing-search-engine_schema_test.cc b/icing/icing-search-engine_schema_test.cc
index 269daaa..735d8d6 100644
--- a/icing/icing-search-engine_schema_test.cc
+++ b/icing/icing-search-engine_schema_test.cc

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <algorithm>
 #include <cstdint>
 #include <limits>
 #include <memory>
@@ -132,6 +133,7 @@
   icing_options.set_enable_schema_database(true);
   icing_options.set_enable_qualified_id_join_index_v3(true);
   icing_options.set_enable_delete_propagation_from(false);
+  icing_options.set_enable_scorable_properties(true);
   return icing_options;
 }
 
@@ -411,10 +413,13 @@
 
     SetSchemaResultProto set_schema_result = icing.SetSchema(schema);
     // Ignore latency numbers. They're covered elsewhere.
-    set_schema_result.clear_latency_ms();
+    set_schema_result.clear_set_schema_stats();
+    set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
     SetSchemaResultProto expected_set_schema_result;
     expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
     expected_set_schema_result.mutable_new_schema_types()->Add("Email");
+    expected_set_schema_result.set_needs_persist_type(
+        PersistType::UNKNOWN);  // Flush is not needed for the first SetSchema.
     EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
     EXPECT_THAT(icing.GetSchema().schema().types(0).version(), Eq(1));
@@ -441,11 +446,14 @@
     // 3. SetSchema should succeed and the version number should be updated.
     SetSchemaResultProto set_schema_result = icing.SetSchema(schema, true);
     // Ignore latency numbers. They're covered elsewhere.
-    set_schema_result.clear_latency_ms();
+    set_schema_result.clear_set_schema_stats();
+    set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
     SetSchemaResultProto expected_set_schema_result;
     expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
     expected_set_schema_result.mutable_fully_compatible_changed_schema_types()
         ->Add("Email");
+    expected_set_schema_result.set_needs_persist_type(
+        PersistType::UNKNOWN);  // Flush is not needed for the first SetSchema.
     EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
     EXPECT_THAT(icing.GetSchema().schema().types(0).version(), Eq(2));
@@ -665,12 +673,20 @@
   SetSchemaResultProto set_schema_result = icing.SetSchema(invalid_schema);
   EXPECT_THAT(set_schema_result.status(),
               ProtoStatusIs(StatusProto::INVALID_ARGUMENT));
-  EXPECT_THAT(set_schema_result.latency_ms(), Eq(1000));
+  EXPECT_THAT(set_schema_result.set_schema_stats().overall_latency_ms(),
+              Eq(1000));
 
   // Can add an document of a set schema
   set_schema_result = icing.SetSchema(schema_with_message);
   EXPECT_THAT(set_schema_result.status(), ProtoStatusIs(StatusProto::OK));
-  EXPECT_THAT(set_schema_result.latency_ms(), Eq(1000));
+  EXPECT_THAT(set_schema_result.set_schema_stats().overall_latency_ms(),
+              Eq(1000));
+  EXPECT_THAT(
+      set_schema_result.set_schema_stats().schema_store_set_schema_latency_ms(),
+      Eq(1000));
+  EXPECT_THAT(set_schema_result.needs_persist_type(),
+              Eq(PersistType::UNKNOWN));  // Flush is not needed for the first
+                                          // SetSchema.
   EXPECT_THAT(icing.Put(message_document).status(), ProtoIsOk());
 
   // Schema with Email doesn't have Message, so would result incompatible
@@ -678,13 +694,21 @@
   set_schema_result = icing.SetSchema(schema_with_email);
   EXPECT_THAT(set_schema_result.status(),
               ProtoStatusIs(StatusProto::FAILED_PRECONDITION));
-  EXPECT_THAT(set_schema_result.latency_ms(), Eq(1000));
+  EXPECT_THAT(set_schema_result.set_schema_stats().overall_latency_ms(),
+              Eq(1000));
+  EXPECT_THAT(
+      set_schema_result.set_schema_stats().schema_store_set_schema_latency_ms(),
+      Eq(1000));
 
-  // Can expand the set of schema types and add an document of a new
-  // schema type
+  // Can expand the set of schema types and add an document of a new schema
+  // type.
   set_schema_result = icing.SetSchema(schema_with_email_and_message);
   EXPECT_THAT(set_schema_result.status(), ProtoStatusIs(StatusProto::OK));
-  EXPECT_THAT(set_schema_result.latency_ms(), Eq(1000));
+  EXPECT_THAT(set_schema_result.set_schema_stats().overall_latency_ms(),
+              Eq(1000));
+  EXPECT_THAT(
+      set_schema_result.set_schema_stats().schema_store_set_schema_latency_ms(),
+      Eq(1000));
 
   EXPECT_THAT(icing.Put(message_document).status(), ProtoIsOk());
   // Can't add an document whose schema isn't set
@@ -702,14 +726,84 @@
 TEST_F(IcingSearchEngineSchemaTest, SetSchema_schemaTypeIdChanged) {
   auto fake_clock = std::make_unique<FakeClock>();
   fake_clock->SetTimerElapsedMilliseconds(1000);
-  TestIcingSearchEngine icing(GetDefaultIcingOptions(),
-                              std::make_unique<Filesystem>(),
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
                               std::make_unique<IcingFilesystem>(),
                               std::move(fake_clock), GetTestJniCache());
   ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
 
+  // We need to add 2 dummy types to the schema so that once deleted, both
+  // Person and Email types' schema type ids will change.
   SchemaProto schema_one =
       SchemaBuilder()
+          .AddType(
+              SchemaTypeConfigBuilder().SetType("DummyA"))  // schema type id 0
+          .AddType(
+              SchemaTypeConfigBuilder().SetType("DummyB"))  // schema type id 1
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("Person")  // schema type id 2
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("name")
+                                        .SetDataTypeString(TERM_MATCH_PREFIX,
+                                                           TOKENIZER_PLAIN)
+                                        .SetCardinality(CARDINALITY_REQUIRED))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("age")
+                                        .SetDataTypeInt64(NUMERIC_MATCH_RANGE)
+                                        .SetCardinality(CARDINALITY_OPTIONAL)))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("Email")  // schema type id 3
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("subject")
+                                        .SetDataTypeString(TERM_MATCH_PREFIX,
+                                                           TOKENIZER_PLAIN)
+                                        .SetCardinality(CARDINALITY_REQUIRED))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("sender")
+                                        .SetDataTypeJoinableString(
+                                            JOINABLE_VALUE_TYPE_QUALIFIED_ID)
+                                        .SetCardinality(CARDINALITY_REQUIRED)))
+          .Build();
+
+  // Set schema with Person and Email types.
+  SetSchemaResultProto set_schema_result1 = icing.SetSchema(schema_one);
+  EXPECT_THAT(set_schema_result1.status(), ProtoStatusIs(StatusProto::OK));
+  EXPECT_THAT(set_schema_result1.set_schema_stats().overall_latency_ms(),
+              Eq(1000));
+  EXPECT_THAT(set_schema_result1.set_schema_stats()
+                  .schema_store_set_schema_latency_ms(),
+              Eq(1000));
+  EXPECT_THAT(set_schema_result1.needs_persist_type(),
+              Eq(PersistType::UNKNOWN));  // Flush is not needed for the first
+                                          // SetSchema.
+
+  // Put a person document and an email document.
+  DocumentProto person_document = DocumentBuilder()
+                                      .SetKey("namespace", "person/1")
+                                      .SetSchema("Person")
+                                      .SetCreationTimestampMs(1000)
+                                      .AddStringProperty("name", "John")
+                                      .AddInt64Property("age", 20)
+                                      .Build();
+  DocumentProto email_document =
+      DocumentBuilder()
+          .SetKey("namespace", "email/1")
+          .SetSchema("Email")
+          .SetCreationTimestampMs(1000)
+          .AddStringProperty("subject", "subject")
+          .AddStringProperty("sender", "namespace#person/1")
+          .Build();
+  EXPECT_THAT(icing.Put(person_document).status(), ProtoIsOk());
+  EXPECT_THAT(icing.Put(email_document).status(), ProtoIsOk());
+
+  // SetSchema again with both the dummy schema types deleted.
+  // - This will cause the schema type ids of both Person and Email types to be
+  //   shifted up.
+  //
+  // This should succeed and all search features should still work after the
+  // schema id change.
+  SchemaProto schema_two =
+      SchemaBuilder()
           .AddType(SchemaTypeConfigBuilder()
                        .SetType("Person")  // schema type id 0
                        .AddProperty(PropertyConfigBuilder()
@@ -734,64 +828,24 @@
                                             JOINABLE_VALUE_TYPE_QUALIFIED_ID)
                                         .SetCardinality(CARDINALITY_REQUIRED)))
           .Build();
-
-  // Set schema with Person and Email types.
-  SetSchemaResultProto set_schema_result1 = icing.SetSchema(schema_one);
-  EXPECT_THAT(set_schema_result1.status(), ProtoStatusIs(StatusProto::OK));
-  EXPECT_THAT(set_schema_result1.latency_ms(), Eq(1000));
-
-  // Put a person document and an email document.
-  DocumentProto person_document = DocumentBuilder()
-                                      .SetKey("namespace", "person/1")
-                                      .SetSchema("Person")
-                                      .SetCreationTimestampMs(1000)
-                                      .AddStringProperty("name", "John")
-                                      .AddInt64Property("age", 20)
-                                      .Build();
-  DocumentProto email_document =
-      DocumentBuilder()
-          .SetKey("namespace", "email/1")
-          .SetSchema("Email")
-          .SetCreationTimestampMs(1000)
-          .AddStringProperty("subject", "subject")
-          .AddStringProperty("sender", "namespace#person/1")
-          .Build();
-  EXPECT_THAT(icing.Put(person_document).status(), ProtoIsOk());
-  EXPECT_THAT(icing.Put(email_document).status(), ProtoIsOk());
-
-  // SetSchema again with the schema types reordered. Schema type id will
-  // change. This should succeed and all search features should still work after
-  // schema id changed.
-  SchemaProto schema_two =
-      SchemaBuilder()
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("Email")  // schema type id 0
-                       .AddProperty(PropertyConfigBuilder()
-                                        .SetName("subject")
-                                        .SetDataTypeString(TERM_MATCH_PREFIX,
-                                                           TOKENIZER_PLAIN)
-                                        .SetCardinality(CARDINALITY_REQUIRED))
-                       .AddProperty(PropertyConfigBuilder()
-                                        .SetName("sender")
-                                        .SetDataTypeJoinableString(
-                                            JOINABLE_VALUE_TYPE_QUALIFIED_ID)
-                                        .SetCardinality(CARDINALITY_REQUIRED)))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("Person")  // schema type id 1
-                       .AddProperty(PropertyConfigBuilder()
-                                        .SetName("name")
-                                        .SetDataTypeString(TERM_MATCH_PREFIX,
-                                                           TOKENIZER_PLAIN)
-                                        .SetCardinality(CARDINALITY_REQUIRED))
-                       .AddProperty(PropertyConfigBuilder()
-                                        .SetName("age")
-                                        .SetDataTypeInt64(NUMERIC_MATCH_RANGE)
-                                        .SetCardinality(CARDINALITY_OPTIONAL)))
-
-          .Build();
-  SetSchemaResultProto set_schema_result2 = icing.SetSchema(schema_two);
+  SetSchemaResultProto set_schema_result2 =
+      icing.SetSchema(schema_two, /*ignore_errors_and_delete_documents=*/true);
   EXPECT_THAT(set_schema_result2.status(), ProtoStatusIs(StatusProto::OK));
-  EXPECT_THAT(set_schema_result2.latency_ms(), Eq(1000));
+  EXPECT_THAT(set_schema_result2.set_schema_stats().overall_latency_ms(),
+              Eq(1000));
+  EXPECT_THAT(set_schema_result2.set_schema_stats()
+                  .schema_store_set_schema_latency_ms(),
+              Eq(1000));
+  // Since the schema types have been changed, derived files were changed. So
+  // we need to do a recovery proof flush.
+  EXPECT_THAT(set_schema_result2.needs_persist_type(),
+              Eq(PersistType::RECOVERY_PROOF));
+  // This will trigger join index restoration for JoinIndex V1 and V2.
+  if (options.enable_qualified_id_join_index_v3()) {
+    EXPECT_FALSE(set_schema_result2.has_qualified_id_join_index_restored());
+  } else {
+    EXPECT_TRUE(set_schema_result2.has_qualified_id_join_index_restored());
+  }
 
   ResultSpecProto result_spec = ResultSpecProto::default_instance();
   result_spec.set_max_joined_children_per_parent_to_return(
@@ -866,8 +920,8 @@
   // - 'b': int64 type, indexed.
   SchemaTypeConfigProto db1_type =
       SchemaTypeConfigBuilder()
-          .SetType("db1_type")
-          .SetDatabase("db1")
+          .SetType("db1/type")
+          .SetDatabase("db1/")
           .AddProperty(PropertyConfigBuilder()
                            .SetName("a")
                            .SetDataTypeString(TERM_MATCH_EXACT, TOKENIZER_PLAIN)
@@ -880,18 +934,20 @@
   SchemaProto db1_schema = SchemaBuilder().AddType(db1_type).Build();
   SetSchemaResultProto set_schema_result =
       icing.SetSchema(CreateSetSchemaRequestProto(
-          db1_schema, "db1", /*ignore_errors_and_delete_documents=*/false));
+          db1_schema, "db1/", /*ignore_errors_and_delete_documents=*/false));
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
   SetSchemaResultProto expected_set_schema_result;
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
-  expected_set_schema_result.mutable_new_schema_types()->Add("db1_type");
+  expected_set_schema_result.mutable_new_schema_types()->Add("db1/type");
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::UNKNOWN);  // Flush is not needed for the first SetSchema.
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   DocumentProto db1_document =
       DocumentBuilder()
           .SetKey("namespace", "uri1")
-          .SetSchema("db1_type")
+          .SetSchema("db1/type")
           .AddStringProperty("a", "message body")
           .AddInt64Property("b", 123)
           .SetCreationTimestampMs(kDefaultCreationTimestampMs)
@@ -930,8 +986,8 @@
   // - 'c': int64 type, indexed.
   SchemaTypeConfigProto db2_type =
       SchemaTypeConfigBuilder()
-          .SetType("db2_type")
-          .SetDatabase("db2")
+          .SetType("db2/type")
+          .SetDatabase("db2/")
           .AddProperty(
               PropertyConfigBuilder()
                   .SetName("a")
@@ -944,18 +1000,22 @@
           .Build();
   SchemaProto db2_schema = SchemaBuilder().AddType(db2_type).Build();
   set_schema_result = icing.SetSchema(CreateSetSchemaRequestProto(
-      db2_schema, "db2", /*ignore_errors_and_delete_documents=*/false));
+      db2_schema, "db2/", /*ignore_errors_and_delete_documents=*/false));
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
   expected_set_schema_result = SetSchemaResultProto();
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
-  expected_set_schema_result.mutable_new_schema_types()->Add("db2_type");
+  expected_set_schema_result.mutable_new_schema_types()->Add("db2/type");
+  // SetSchema db operation will preserve the existing schema type ids as much
+  // as possible. In this case, since "db1/type"'s id is unchanged, the document
+  // store and index are not updated. Therefore, no need to flush.
+  expected_set_schema_result.set_needs_persist_type(PersistType::UNKNOWN);
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   DocumentProto db2_document =
       DocumentBuilder()
           .SetKey("namespace", "uri2")
-          .SetSchema("db2_type")
+          .SetSchema("db2/type")
           .AddStringProperty("a", "message body")
           .AddInt64Property("c", 123)
           .SetCreationTimestampMs(kDefaultCreationTimestampMs)
@@ -999,7 +1059,7 @@
   expected_get_schema_result_proto_db1_full.mutable_status()->set_code(
       StatusProto::OK);
   *expected_get_schema_result_proto_db1_full.mutable_schema() = db1_schema;
-  EXPECT_THAT(icing.GetSchema("db1"),
+  EXPECT_THAT(icing.GetSchema("db1/"),
               EqualsProto(expected_get_schema_result_proto_db1_full));
 
   // Get db2 schema
@@ -1007,7 +1067,188 @@
   expected_get_schema_result_proto_db2_full.mutable_status()->set_code(
       StatusProto::OK);
   *expected_get_schema_result_proto_db2_full.mutable_schema() = db2_schema;
-  EXPECT_THAT(icing.GetSchema("db2"),
+  EXPECT_THAT(icing.GetSchema("db2/"),
+              EqualsProto(expected_get_schema_result_proto_db2_full));
+}
+
+TEST_F(IcingSearchEngineSchemaTest,
+       SetSchemaWithEmptyDatabaseResetsFullSchema) {
+  IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+
+  // Create and set schema in db1 with 2 properties:
+  // - 'a': string type, indexed.
+  // - 'b': int64 type, indexed.
+  SchemaTypeConfigProto db1_type =
+      SchemaTypeConfigBuilder()
+          .SetType("db1/type")
+          .SetDatabase("db1/")
+          .AddProperty(PropertyConfigBuilder()
+                           .SetName("a")
+                           .SetDataTypeString(TERM_MATCH_EXACT, TOKENIZER_PLAIN)
+                           .SetCardinality(CARDINALITY_REQUIRED))
+          .AddProperty(PropertyConfigBuilder()
+                           .SetName("b")
+                           .SetDataTypeInt64(NUMERIC_MATCH_RANGE)
+                           .SetCardinality(CARDINALITY_REQUIRED))
+          .Build();
+  SchemaProto db1_schema = SchemaBuilder().AddType(db1_type).Build();
+  SetSchemaResultProto set_schema_result =
+      icing.SetSchema(CreateSetSchemaRequestProto(
+          db1_schema, "db1/", /*ignore_errors_and_delete_documents=*/false));
+  // Ignore latency numbers. They're covered elsewhere.
+  set_schema_result.clear_set_schema_stats();
+  SetSchemaResultProto expected_set_schema_result;
+  expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
+  expected_set_schema_result.mutable_new_schema_types()->Add("db1/type");
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::UNKNOWN);  // Flush is not needed for the first SetSchema.
+  EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
+
+  DocumentProto db1_document =
+      DocumentBuilder()
+          .SetKey("namespace", "uri1")
+          .SetSchema("db1/type")
+          .AddStringProperty("a", "message body")
+          .AddInt64Property("b", 123)
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .Build();
+  ASSERT_THAT(icing.Put(db1_document).status(), ProtoIsOk());
+
+  SearchResultProto expected_search_result_proto_db1;
+  expected_search_result_proto_db1.mutable_status()->set_code(StatusProto::OK);
+  *expected_search_result_proto_db1.mutable_results()
+       ->Add()
+       ->mutable_document() = db1_document;
+
+  // Verify term search
+  SearchSpecProto search_spec1;
+  search_spec1.set_query("message");
+  search_spec1.set_term_match_type(TermMatchType::EXACT_ONLY);
+
+  SearchResultProto actual_results =
+      icing.Search(search_spec1, GetDefaultScoringSpec(),
+                   ResultSpecProto::default_instance());
+  EXPECT_THAT(actual_results, EqualsSearchResultIgnoreStatsAndScores(
+                                  expected_search_result_proto_db1));
+
+  // Verify numeric (integer) search
+  SearchSpecProto search_spec2;
+  search_spec2.set_query("b == 123");
+  search_spec2.add_enabled_features(std::string(kNumericSearchFeature));
+
+  actual_results = icing.Search(search_spec2, GetDefaultScoringSpec(),
+                                ResultSpecProto::default_instance());
+  EXPECT_THAT(actual_results, EqualsSearchResultIgnoreStatsAndScores(
+                                  expected_search_result_proto_db1));
+
+  // Reset the full schema. Add a type each for db1 and db2.
+  SchemaTypeConfigProto db1_type_2 =
+      SchemaTypeConfigBuilder()
+          .SetType("db1/type_2")
+          .SetDatabase("db1/")
+          .AddProperty(PropertyConfigBuilder()
+                           .SetName("d")
+                           .SetDataTypeString(TERM_MATCH_EXACT, TOKENIZER_PLAIN)
+                           .SetCardinality(CARDINALITY_REQUIRED))
+          .Build();
+  SchemaTypeConfigProto db2_type =
+      SchemaTypeConfigBuilder()
+          .SetType("db2/type")
+          .SetDatabase("db2/")
+          .AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("a")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_REQUIRED))
+          .AddProperty(PropertyConfigBuilder()
+                           .SetName("c")
+                           .SetDataTypeInt64(NUMERIC_MATCH_RANGE)
+                           .SetCardinality(CARDINALITY_REQUIRED))
+          .Build();
+  SchemaProto full_schema = SchemaBuilder()
+                                .AddType(db1_type)
+                                .AddType(db1_type_2)
+                                .AddType(db2_type)
+                                .Build();
+
+  // Set the full schema using an empty database.
+  SetSchemaRequestProto set_schema_request;
+  *set_schema_request.mutable_schema() = full_schema;
+  set_schema_request.set_ignore_errors_and_delete_documents(false);
+  set_schema_request.clear_database();
+  set_schema_result = icing.SetSchema(std::move(set_schema_request));
+  std::sort(set_schema_result.mutable_new_schema_types()->begin(),
+            set_schema_result.mutable_new_schema_types()->end());
+  // Ignore latency numbers. They're covered elsewhere.
+  set_schema_result.clear_set_schema_stats();
+  expected_set_schema_result = SetSchemaResultProto();
+  expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
+  expected_set_schema_result.mutable_new_schema_types()->Add("db1/type_2");
+  expected_set_schema_result.mutable_new_schema_types()->Add("db2/type");
+  // SetSchema db operation will preserve the existing schema type ids as much
+  // as possible. In this case, since "db1/type"'s id is unchanged, the document
+  // store and index are not updated. Therefore, no need to flush.
+  expected_set_schema_result.set_needs_persist_type(PersistType::UNKNOWN);
+  EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
+
+  DocumentProto db2_document =
+      DocumentBuilder()
+          .SetKey("namespace", "uri2")
+          .SetSchema("db2/type")
+          .AddStringProperty("a", "message body")
+          .AddInt64Property("c", 123)
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .Build();
+  EXPECT_THAT(icing.Put(db2_document).status(), ProtoIsOk());
+
+  SearchResultProto expected_search_result_proto_both_docs;
+  expected_search_result_proto_both_docs.mutable_status()->set_code(
+      StatusProto::OK);
+  *expected_search_result_proto_both_docs.mutable_results()
+       ->Add()
+       ->mutable_document() = db2_document;
+  *expected_search_result_proto_both_docs.mutable_results()
+       ->Add()
+       ->mutable_document() = db1_document;
+
+  // Verify term search: will get both documents for query "a:message".
+  actual_results = icing.Search(search_spec1, GetDefaultScoringSpec(),
+                                ResultSpecProto::default_instance());
+  EXPECT_THAT(actual_results, EqualsSearchResultIgnoreStatsAndScores(
+                                  expected_search_result_proto_both_docs));
+
+  // Verify numeric (integer) search: will only get first document for query
+  // "b == 123".
+  actual_results = icing.Search(search_spec2, GetDefaultScoringSpec(),
+                                ResultSpecProto::default_instance());
+  EXPECT_THAT(actual_results, EqualsSearchResultIgnoreStatsAndScores(
+                                  expected_search_result_proto_db1));
+
+  // Get full schema
+  GetSchemaResultProto expected_get_schema_result_proto_full;
+  expected_get_schema_result_proto_full.mutable_status()->set_code(
+      StatusProto::OK);
+  *expected_get_schema_result_proto_full.mutable_schema() = full_schema;
+  EXPECT_THAT(icing.GetSchema(),
+              EqualsProto(expected_get_schema_result_proto_full));
+
+  // Get db1 schema
+  GetSchemaResultProto expected_get_schema_result_proto_db1_full;
+  expected_get_schema_result_proto_db1_full.mutable_status()->set_code(
+      StatusProto::OK);
+  *expected_get_schema_result_proto_db1_full.mutable_schema() =
+      SchemaBuilder().AddType(db1_type).AddType(db1_type_2).Build();
+  EXPECT_THAT(icing.GetSchema("db1/"),
+              EqualsProto(expected_get_schema_result_proto_db1_full));
+
+  // Get db2 schema
+  GetSchemaResultProto expected_get_schema_result_proto_db2_full;
+  expected_get_schema_result_proto_db2_full.mutable_status()->set_code(
+      StatusProto::OK);
+  *expected_get_schema_result_proto_db2_full.mutable_schema() =
+      SchemaBuilder().AddType(db2_type).Build();
+  EXPECT_THAT(icing.GetSchema("db2/"),
               EqualsProto(expected_get_schema_result_proto_db2_full));
 }
 
@@ -1020,8 +1261,8 @@
   // - 'c': int64 type, indexed.
   SchemaTypeConfigProto db1_type =
       SchemaTypeConfigBuilder()
-          .SetType("db1_type")
-          .SetDatabase("db1")
+          .SetType("db1/type")
+          .SetDatabase("db1/")
           .AddProperty(PropertyConfigBuilder()
                            .SetName("b")
                            .SetDataTypeString(TERM_MATCH_EXACT, TOKENIZER_PLAIN)
@@ -1034,12 +1275,14 @@
   SchemaProto db1_schema = SchemaBuilder().AddType(db1_type).Build();
   SetSchemaResultProto set_schema_result =
       icing.SetSchema(CreateSetSchemaRequestProto(
-          db1_schema, "db1", /*ignore_errors_and_delete_documents=*/false));
+          db1_schema, "db1/", /*ignore_errors_and_delete_documents=*/false));
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
   SetSchemaResultProto expected_set_schema_result;
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
-  expected_set_schema_result.mutable_new_schema_types()->Add("db1_type");
+  expected_set_schema_result.mutable_new_schema_types()->Add("db1/type");
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::UNKNOWN);  // Flush is not needed for the first SetSchema.
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   // Add a schema for db2:
@@ -1047,8 +1290,8 @@
   // - 'd': int64 type, indexed.
   SchemaTypeConfigProto db2_type =
       SchemaTypeConfigBuilder()
-          .SetType("db2_type")
-          .SetDatabase("db2")
+          .SetType("db2/type")
+          .SetDatabase("db2/")
           .AddProperty(
               PropertyConfigBuilder()
                   .SetName("b")
@@ -1061,19 +1304,23 @@
           .Build();
   SchemaProto db2_schema = SchemaBuilder().AddType(db2_type).Build();
   set_schema_result = icing.SetSchema(CreateSetSchemaRequestProto(
-      db2_schema, "db2", /*ignore_errors_and_delete_documents=*/false));
+      db2_schema, "db2/", /*ignore_errors_and_delete_documents=*/false));
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
   expected_set_schema_result = SetSchemaResultProto();
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
-  expected_set_schema_result.mutable_new_schema_types()->Add("db2_type");
+  expected_set_schema_result.mutable_new_schema_types()->Add("db2/type");
+  // SetSchema db operation will preserve the existing schema type ids as much
+  // as possible. In this case, since "db1/type"'s id is unchanged, the document
+  // store and index are not updated. Therefore, no need to flush.
+  expected_set_schema_result.set_needs_persist_type(PersistType::UNKNOWN);
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   // Add documents
   DocumentProto db1_document1 =
       DocumentBuilder()
           .SetKey("namespace", "uri1")
-          .SetSchema("db1_type")
+          .SetSchema("db1/type")
           .AddStringProperty("b", "message body")
           .AddInt64Property("c", 123)
           .SetCreationTimestampMs(kDefaultCreationTimestampMs)
@@ -1081,7 +1328,7 @@
   DocumentProto db2_document =
       DocumentBuilder()
           .SetKey("namespace", "uri2")
-          .SetSchema("db2_type")
+          .SetSchema("db2/type")
           .AddStringProperty("b", "message body")
           .AddInt64Property("d", 123)
           .SetCreationTimestampMs(kDefaultCreationTimestampMs)
@@ -1138,20 +1385,28 @@
           .Build());
   db1_schema = SchemaBuilder().AddType(db1_type).Build();
   set_schema_result = icing.SetSchema(CreateSetSchemaRequestProto(
-      db1_schema, "db1", /*ignore_errors_and_delete_documents=*/false));
+      db1_schema, "db1/", /*ignore_errors_and_delete_documents=*/false));
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
   expected_set_schema_result = SetSchemaResultProto();
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
   expected_set_schema_result.mutable_index_incompatible_changed_schema_types()
-      ->Add("db1_type");
+      ->Add("db1/type");
+  expected_set_schema_result.set_has_term_index_restored(true);
+  expected_set_schema_result.set_has_integer_index_restored(true);
+  expected_set_schema_result.set_has_embedding_index_restored(true);
+  expected_set_schema_result.set_has_qualified_id_join_index_restored(false);
+  // Since it is an index incompatible change, we need to do a recovery proof
+  // flush.
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::RECOVERY_PROOF);
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   // Add new document
   DocumentProto db1_document2 =
       DocumentBuilder()
           .SetKey("namespace", "uri3")
-          .SetSchema("db1_type")
+          .SetSchema("db1/type")
           .AddStringProperty("a", "message body")
           .AddStringProperty("b", "string value")
           .AddInt64Property("c", 123)
@@ -1214,7 +1469,7 @@
   expected_get_schema_result_proto_db1_full.mutable_status()->set_code(
       StatusProto::OK);
   *expected_get_schema_result_proto_db1_full.mutable_schema() = db1_schema;
-  EXPECT_THAT(icing.GetSchema("db1"),
+  EXPECT_THAT(icing.GetSchema("db1/"),
               EqualsProto(expected_get_schema_result_proto_db1_full));
 
   // Get db2 schema
@@ -1222,12 +1477,13 @@
   expected_get_schema_result_proto_db2_full.mutable_status()->set_code(
       StatusProto::OK);
   *expected_get_schema_result_proto_db2_full.mutable_schema() = db2_schema;
-  EXPECT_THAT(icing.GetSchema("db2"),
+  EXPECT_THAT(icing.GetSchema("db2/"),
               EqualsProto(expected_get_schema_result_proto_db2_full));
 }
 
 TEST_F(IcingSearchEngineSchemaTest, SetSchemaEmptySchemaClearsDatabase) {
-  IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  IcingSearchEngine icing(options, GetTestJniCache());
   ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
 
   // Create and set schema in db1 with 2 properties:
@@ -1235,8 +1491,8 @@
   // - 'c': int64 type, indexed.
   SchemaTypeConfigProto db1_type =
       SchemaTypeConfigBuilder()
-          .SetType("db1_type")
-          .SetDatabase("db1")
+          .SetType("db1/type")
+          .SetDatabase("db1/")
           .AddProperty(PropertyConfigBuilder()
                            .SetName("b")
                            .SetDataTypeString(TERM_MATCH_EXACT, TOKENIZER_PLAIN)
@@ -1249,12 +1505,14 @@
   SchemaProto db1_schema = SchemaBuilder().AddType(db1_type).Build();
   SetSchemaResultProto set_schema_result =
       icing.SetSchema(CreateSetSchemaRequestProto(
-          db1_schema, "db1", /*ignore_errors_and_delete_documents=*/true));
+          db1_schema, "db1/", /*ignore_errors_and_delete_documents=*/true));
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
   SetSchemaResultProto expected_set_schema_result;
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
-  expected_set_schema_result.mutable_new_schema_types()->Add("db1_type");
+  expected_set_schema_result.mutable_new_schema_types()->Add("db1/type");
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::UNKNOWN);  // Flush is not needed for the first SetSchema.
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   // Add a schema for db2:
@@ -1262,8 +1520,8 @@
   // - 'd': int64 type, indexed.
   SchemaTypeConfigProto db2_type =
       SchemaTypeConfigBuilder()
-          .SetType("db2_type")
-          .SetDatabase("db2")
+          .SetType("db2/type")
+          .SetDatabase("db2/")
           .AddProperty(
               PropertyConfigBuilder()
                   .SetName("b")
@@ -1276,19 +1534,23 @@
           .Build();
   SchemaProto db2_schema = SchemaBuilder().AddType(db2_type).Build();
   set_schema_result = icing.SetSchema(CreateSetSchemaRequestProto(
-      db2_schema, "db2", /*ignore_errors_and_delete_documents=*/true));
+      db2_schema, "db2/", /*ignore_errors_and_delete_documents=*/true));
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
   expected_set_schema_result = SetSchemaResultProto();
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
-  expected_set_schema_result.mutable_new_schema_types()->Add("db2_type");
+  expected_set_schema_result.mutable_new_schema_types()->Add("db2/type");
+  // SetSchema db operation will preserve the existing schema type ids as much
+  // as possible. In this case, since "db1/type"'s id is unchanged, the document
+  // store and index are not updated. Therefore, no need to flush.
+  expected_set_schema_result.set_needs_persist_type(PersistType::UNKNOWN);
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   // Add documents
   DocumentProto db1_document1 =
       DocumentBuilder()
           .SetKey("namespace", "uri1")
-          .SetSchema("db1_type")
+          .SetSchema("db1/type")
           .AddStringProperty("b", "message body")
           .AddInt64Property("c", 123)
           .SetCreationTimestampMs(kDefaultCreationTimestampMs)
@@ -1296,7 +1558,7 @@
   DocumentProto db2_document =
       DocumentBuilder()
           .SetKey("namespace", "uri2")
-          .SetSchema("db2_type")
+          .SetSchema("db2/type")
           .AddStringProperty("b", "message body")
           .AddInt64Property("d", 123)
           .SetCreationTimestampMs(kDefaultCreationTimestampMs)
@@ -1344,19 +1606,31 @@
   //   - 'd': int64 type, indexed.
   db1_schema = SchemaProto();
   set_schema_result = icing.SetSchema(CreateSetSchemaRequestProto(
-      db1_schema, "db1", /*ignore_errors_and_delete_documents=*/true));
+      db1_schema, "db1/", /*ignore_errors_and_delete_documents=*/true));
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
   expected_set_schema_result = SetSchemaResultProto();
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
-  expected_set_schema_result.mutable_deleted_schema_types()->Add("db1_type");
+  expected_set_schema_result.mutable_deleted_schema_types()->Add("db1/type");
+  expected_set_schema_result.set_deleted_document_count(1);
+  // Since a document was deleted, we need to do a recovery proof flush.
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::RECOVERY_PROOF);
+  // Deleting db1_schema results in a schema type id reassignment, which would
+  // trigger join index restoration if using join index v1 and v2.
+  if (!options.enable_qualified_id_join_index_v3()) {
+    expected_set_schema_result.set_has_qualified_id_join_index_restored(true);
+    expected_set_schema_result.set_has_term_index_restored(false);
+    expected_set_schema_result.set_has_integer_index_restored(false);
+    expected_set_schema_result.set_has_embedding_index_restored(false);
+  }
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   // Adding new document fails because db1_type is deleted.
   DocumentProto db1_document2 =
       DocumentBuilder()
           .SetKey("namespace", "uri3")
-          .SetSchema("db1_type")
+          .SetSchema("db1/type")
           .AddStringProperty("a", "message body")
           .AddStringProperty("b", "string value")
           .AddInt64Property("c", 123)
@@ -1405,7 +1679,7 @@
               EqualsProto(expected_get_schema_result_proto_full));
 
   // Get db1 schema should return NOT_FOUND.
-  EXPECT_THAT(icing.GetSchema("db1").status(),
+  EXPECT_THAT(icing.GetSchema("db1/").status(),
               ProtoStatusIs(StatusProto::NOT_FOUND));
 
   // Get db2 schema
@@ -1413,7 +1687,7 @@
   expected_get_schema_result_proto_db2_full.mutable_status()->set_code(
       StatusProto::OK);
   *expected_get_schema_result_proto_db2_full.mutable_schema() = db2_schema;
-  EXPECT_THAT(icing.GetSchema("db2"),
+  EXPECT_THAT(icing.GetSchema("db2/"),
               EqualsProto(expected_get_schema_result_proto_db2_full));
 }
 
@@ -1442,10 +1716,13 @@
 
   SetSchemaResultProto set_schema_result = icing.SetSchema(schema_one);
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
   SetSchemaResultProto expected_set_schema_result;
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
   expected_set_schema_result.mutable_new_schema_types()->Add("Schema");
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::UNKNOWN);  // Flush is not needed for the first SetSchema.
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   DocumentProto document =
@@ -1505,11 +1782,20 @@
   // restoration should use new section ids to rebuild.
   set_schema_result = icing.SetSchema(schema_two);
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
   expected_set_schema_result = SetSchemaResultProto();
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
   expected_set_schema_result.mutable_index_incompatible_changed_schema_types()
       ->Add("Schema");
+  expected_set_schema_result.set_has_term_index_restored(true);
+  expected_set_schema_result.set_has_integer_index_restored(true);
+  expected_set_schema_result.set_has_qualified_id_join_index_restored(false);
+  expected_set_schema_result.set_has_embedding_index_restored(true);
+  // Since it is an index incompatible change, we need to do a recovery proof
+  // flush.
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::RECOVERY_PROOF);
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   // Verify term search: will get document now.
@@ -1551,10 +1837,13 @@
 
   SetSchemaResultProto set_schema_result = icing.SetSchema(schema_one);
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
   SetSchemaResultProto expected_set_schema_result;
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
   expected_set_schema_result.mutable_new_schema_types()->Add("Schema");
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::UNKNOWN);  // Flush is not needed for the first SetSchema.
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   DocumentProto document =
@@ -1610,11 +1899,20 @@
   // restoration should use new section ids to rebuild.
   set_schema_result = icing.SetSchema(schema_two);
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
   expected_set_schema_result = SetSchemaResultProto();
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
   expected_set_schema_result.mutable_index_incompatible_changed_schema_types()
       ->Add("Schema");
+  expected_set_schema_result.set_has_term_index_restored(true);
+  expected_set_schema_result.set_has_integer_index_restored(true);
+  expected_set_schema_result.set_has_qualified_id_join_index_restored(false);
+  expected_set_schema_result.set_has_embedding_index_restored(true);
+  // Since it is an index incompatible change, we need to do a recovery proof
+  // flush.
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::RECOVERY_PROOF);
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   // Verify term search: will still get document.
@@ -1839,13 +2137,22 @@
   // restoration should use new section ids to rebuild.
   SetSchemaResultProto set_schema_result = icing.SetSchema(schema_two);
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
   SetSchemaResultProto expected_set_schema_result = SetSchemaResultProto();
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
   expected_set_schema_result.mutable_index_incompatible_changed_schema_types()
       ->Add("Person");
   expected_set_schema_result.mutable_join_incompatible_changed_schema_types()
       ->Add("Person");
+  expected_set_schema_result.set_has_term_index_restored(true);
+  expected_set_schema_result.set_has_integer_index_restored(true);
+  expected_set_schema_result.set_has_qualified_id_join_index_restored(true);
+  expected_set_schema_result.set_has_embedding_index_restored(true);
+  // Since it is an index incompatible change, we need to do a recovery proof
+  // flush.
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::RECOVERY_PROOF);
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   // Verify term search:
@@ -1925,11 +2232,14 @@
 
   SetSchemaResultProto set_schema_result = icing.SetSchema(nested_schema);
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
   SetSchemaResultProto expected_set_schema_result;
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
   expected_set_schema_result.mutable_new_schema_types()->Add("Email");
   expected_set_schema_result.mutable_new_schema_types()->Add("Person");
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::UNKNOWN);  // Flush is not needed for the first SetSchema.
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   DocumentProto document =
@@ -2027,11 +2337,20 @@
 
   set_schema_result = icing.SetSchema(no_nested_schema);
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
   expected_set_schema_result = SetSchemaResultProto();
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
   expected_set_schema_result.mutable_index_incompatible_changed_schema_types()
       ->Add("Email");
+  expected_set_schema_result.set_has_term_index_restored(true);
+  expected_set_schema_result.set_has_integer_index_restored(true);
+  expected_set_schema_result.set_has_qualified_id_join_index_restored(false);
+  expected_set_schema_result.set_has_embedding_index_restored(true);
+  // Since it is an index incompatible change, we need to do a recovery proof
+  // flush.
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::RECOVERY_PROOF);
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   // Verify term search
@@ -2145,11 +2464,14 @@
 
   SetSchemaResultProto set_schema_result = icing.SetSchema(nested_schema);
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
   SetSchemaResultProto expected_set_schema_result;
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
   expected_set_schema_result.mutable_new_schema_types()->Add("Email");
   expected_set_schema_result.mutable_new_schema_types()->Add("Person");
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::UNKNOWN);  // Flush is not needed for the first SetSchema.
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   DocumentProto document =
@@ -2317,11 +2639,20 @@
 
   set_schema_result = icing.SetSchema(nested_schema_with_less_props);
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
   expected_set_schema_result = SetSchemaResultProto();
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
   expected_set_schema_result.mutable_index_incompatible_changed_schema_types()
       ->Add("Email");
+  expected_set_schema_result.set_has_term_index_restored(true);
+  expected_set_schema_result.set_has_integer_index_restored(true);
+  expected_set_schema_result.set_has_qualified_id_join_index_restored(false);
+  expected_set_schema_result.set_has_embedding_index_restored(true);
+  // Since it is an index incompatible change, we need to do a recovery proof
+  // flush.
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::RECOVERY_PROOF);
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   // Verify term search
@@ -2380,11 +2711,14 @@
 
   SetSchemaResultProto set_schema_result = icing.SetSchema(schema_one);
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
   SetSchemaResultProto expected_set_schema_result;
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
   expected_set_schema_result.mutable_new_schema_types()->Add("Message");
   expected_set_schema_result.mutable_new_schema_types()->Add("Person");
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::UNKNOWN);  // Flush is not needed for the first SetSchema.
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   DocumentProto person1 =
@@ -2496,11 +2830,20 @@
   // index restoration should use new joinable property ids to rebuild.
   set_schema_result = icing.SetSchema(schema_two);
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
   expected_set_schema_result = SetSchemaResultProto();
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
   expected_set_schema_result.mutable_join_incompatible_changed_schema_types()
       ->Add("Message");
+  expected_set_schema_result.set_has_term_index_restored(false);
+  expected_set_schema_result.set_has_integer_index_restored(false);
+  expected_set_schema_result.set_has_qualified_id_join_index_restored(true);
+  expected_set_schema_result.set_has_embedding_index_restored(false);
+  // Since it is a join incompatible change, we need to do a recovery proof
+  // flush.
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::RECOVERY_PROOF);
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   // Verify join search: join a query for `name:person` with a child query for
@@ -2658,10 +3001,13 @@
   SetSchemaResultProto set_schema_result =
       icing.SetSchema(email_with_body_schema);
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
   SetSchemaResultProto expected_set_schema_result;
   expected_set_schema_result.mutable_new_schema_types()->Add("Email");
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::UNKNOWN);  // Flush is not needed for the first SetSchema.
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   // Create a document with only subject and timestamp2 property.
@@ -2733,12 +3079,22 @@
   set_schema_result = icing.SetSchema(
       email_no_body_schema, /*ignore_errors_and_delete_documents=*/true);
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
   expected_set_schema_result = SetSchemaResultProto();
   expected_set_schema_result.mutable_incompatible_schema_types()->Add("Email");
   expected_set_schema_result.mutable_index_incompatible_changed_schema_types()
       ->Add("Email");
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
+  expected_set_schema_result.set_deleted_document_count(0);
+  expected_set_schema_result.set_has_term_index_restored(true);
+  expected_set_schema_result.set_has_integer_index_restored(true);
+  expected_set_schema_result.set_has_qualified_id_join_index_restored(false);
+  expected_set_schema_result.set_has_embedding_index_restored(true);
+  // Since it is an index incompatible change, we need to do a recovery proof
+  // flush.
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::RECOVERY_PROOF);
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   // Verify term search
@@ -2799,11 +3155,14 @@
   SetSchemaResultProto set_schema_result =
       icing.SetSchema(email_with_receiver_schema);
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
   SetSchemaResultProto expected_set_schema_result;
   expected_set_schema_result.mutable_new_schema_types()->Add("Email");
   expected_set_schema_result.mutable_new_schema_types()->Add("Person");
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::UNKNOWN);  // Flush is not needed for the first SetSchema.
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   DocumentProto person = DocumentBuilder()
@@ -2899,12 +3258,22 @@
       icing.SetSchema(email_without_receiver_schema,
                       /*ignore_errors_and_delete_documents=*/true);
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
   expected_set_schema_result = SetSchemaResultProto();
   expected_set_schema_result.mutable_incompatible_schema_types()->Add("Email");
   expected_set_schema_result.mutable_join_incompatible_changed_schema_types()
       ->Add("Email");
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
+  expected_set_schema_result.set_deleted_document_count(0);
+  expected_set_schema_result.set_has_term_index_restored(false);
+  expected_set_schema_result.set_has_integer_index_restored(false);
+  expected_set_schema_result.set_has_qualified_id_join_index_restored(true);
+  expected_set_schema_result.set_has_embedding_index_restored(false);
+  // Since it is a join incompatible change, we need to do a recovery proof
+  // flush.
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::RECOVERY_PROOF);
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   // Verify join search: join a query for `name:person` with a child query for
@@ -2949,10 +3318,13 @@
   SetSchemaResultProto set_schema_result =
       icing.SetSchema(email_with_body_schema);
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
   SetSchemaResultProto expected_set_schema_result;
   expected_set_schema_result.mutable_new_schema_types()->Add("Email");
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::UNKNOWN);  // Flush is not needed for the first SetSchema.
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   // Create a document with only subject and timestamp property.
@@ -3030,12 +3402,22 @@
   set_schema_result = icing.SetSchema(
       email_no_body_schema, /*ignore_errors_and_delete_documents=*/true);
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
   expected_set_schema_result = SetSchemaResultProto();
   expected_set_schema_result.mutable_incompatible_schema_types()->Add("Email");
   expected_set_schema_result.mutable_index_incompatible_changed_schema_types()
       ->Add("Email");
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
+  expected_set_schema_result.set_deleted_document_count(0);
+  expected_set_schema_result.set_has_term_index_restored(true);
+  expected_set_schema_result.set_has_integer_index_restored(true);
+  expected_set_schema_result.set_has_qualified_id_join_index_restored(false);
+  expected_set_schema_result.set_has_embedding_index_restored(true);
+  // Since it is an index incompatible change, we need to do a recovery proof
+  // flush.
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::RECOVERY_PROOF);
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   // Verify term search
@@ -3097,11 +3479,14 @@
   SetSchemaResultProto set_schema_result =
       icing.SetSchema(email_with_body_schema);
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
   SetSchemaResultProto expected_set_schema_result;
   expected_set_schema_result.mutable_new_schema_types()->Add("Email");
   expected_set_schema_result.mutable_new_schema_types()->Add("Person");
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::UNKNOWN);  // Flush is not needed for the first SetSchema.
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   DocumentProto person = DocumentBuilder()
@@ -3199,12 +3584,22 @@
   set_schema_result = icing.SetSchema(
       email_no_body_schema, /*ignore_errors_and_delete_documents=*/true);
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
   expected_set_schema_result = SetSchemaResultProto();
   expected_set_schema_result.mutable_incompatible_schema_types()->Add("Email");
   expected_set_schema_result.mutable_join_incompatible_changed_schema_types()
       ->Add("Email");
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
+  expected_set_schema_result.set_deleted_document_count(0);
+  expected_set_schema_result.set_has_term_index_restored(false);
+  expected_set_schema_result.set_has_integer_index_restored(false);
+  expected_set_schema_result.set_has_qualified_id_join_index_restored(true);
+  expected_set_schema_result.set_has_embedding_index_restored(false);
+  // Since it is a join incompatible change, we need to do a recovery proof
+  // flush.
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::RECOVERY_PROOF);
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   // Verify join search: join a query for `name:person` with a child query for
@@ -3255,11 +3650,14 @@
 
   SetSchemaResultProto set_schema_result = icing.SetSchema(nested_schema);
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
   SetSchemaResultProto expected_set_schema_result;
   expected_set_schema_result.mutable_new_schema_types()->Add("Email");
   expected_set_schema_result.mutable_new_schema_types()->Add("Person");
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::UNKNOWN);  // Flush is not needed for the first SetSchema.
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   // Create two documents - a person document and an email document - both docs
@@ -3312,7 +3710,8 @@
   set_schema_result = icing.SetSchema(
       nested_schema, /*ignore_errors_and_delete_documents=*/true);
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
   expected_set_schema_result = SetSchemaResultProto();
   expected_set_schema_result.mutable_incompatible_schema_types()->Add("Person");
   expected_set_schema_result.mutable_incompatible_schema_types()->Add("Email");
@@ -3321,6 +3720,15 @@
   expected_set_schema_result.mutable_index_incompatible_changed_schema_types()
       ->Add("Person");
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
+  expected_set_schema_result.set_deleted_document_count(2);
+  expected_set_schema_result.set_has_term_index_restored(true);
+  expected_set_schema_result.set_has_integer_index_restored(true);
+  expected_set_schema_result.set_has_qualified_id_join_index_restored(false);
+  expected_set_schema_result.set_has_embedding_index_restored(true);
+  // Since it is an index incompatible change, we need to do a recovery proof
+  // flush.
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::RECOVERY_PROOF);
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   // Both documents should be deleted now.
@@ -3333,6 +3741,141 @@
   EXPECT_THAT(get_result.status(), ProtoStatusIs(StatusProto::NOT_FOUND));
 }
 
+TEST_F(
+    IcingSearchEngineSchemaTest,
+    SetSchemaAddDeletePropagationTriggersIndexRestorationAndRevalidatesDependency) {
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  options.set_enable_qualified_id_join_index_v3(true);
+  options.set_enable_soft_index_restoration(true);
+  options.set_enable_delete_propagation_from(true);
+  options.set_expired_document_purge_threshold_ms(0);
+
+  auto fake_clock = std::make_unique<FakeClock>();
+  FakeClock* fake_clock_ptr = fake_clock.get();
+
+  TestIcingSearchEngine icing(options, std::make_unique<Filesystem>(),
+                              std::make_unique<IcingFilesystem>(),
+                              std::move(fake_clock), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+
+  // Create "Email" schema with a joinable property "senderQualifiedId" and
+  // delete propagation disabled.
+  SchemaProto schema1 =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("name")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_REQUIRED)))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("Email")
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("subject")
+                                        .SetDataTypeString(TERM_MATCH_PREFIX,
+                                                           TOKENIZER_PLAIN)
+                                        .SetCardinality(CARDINALITY_OPTIONAL))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("senderQualifiedId")
+                                        .SetDataTypeJoinableString(
+                                            JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                            DELETE_PROPAGATION_TYPE_NONE)
+                                        .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  SetSchemaResultProto set_schema_result = icing.SetSchema(schema1);
+  // Ignore latency numbers. They're covered elsewhere.
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
+  SetSchemaResultProto expected_set_schema_result;
+  expected_set_schema_result.mutable_new_schema_types()->Add("Email");
+  expected_set_schema_result.mutable_new_schema_types()->Add("Person");
+  expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::UNKNOWN);  // Flush is not needed for the first SetSchema.
+  EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
+
+  DocumentProto person = DocumentBuilder()
+                             .SetKey("namespace", "person")
+                             .SetSchema("Person")
+                             .SetCreationTimestampMs(10)
+                             .SetTtlMs(10000)  // Expire at t = 10010.
+                             .AddStringProperty("name", "person")
+                             .Build();
+  // Create an email document.
+  DocumentProto email =
+      DocumentBuilder()
+          .SetKey("namespace", "email")
+          .SetSchema("Email")
+          .SetCreationTimestampMs(10)
+          .SetTtlMs(0)  // Never expire.
+          .AddStringProperty("subject", "test")
+          .AddStringProperty("senderQualifiedId", "namespace#person")
+          .Build();
+
+  ASSERT_THAT(icing.Put(person).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(email).status(), ProtoIsOk());
+
+  // Now update the schema to enable delete propagation for "senderQualifiedId".
+  // It is:
+  // - Schema compatible.
+  // - Join incompatible. This will trigger index restoration, and revalidation
+  //   of dependency. If any child documents are referring to a non-alive
+  //   parent, then they will be deleted.
+  SchemaProto schema2 =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("name")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_REQUIRED)))
+          .AddType(
+              SchemaTypeConfigBuilder()
+                  .SetType("Email")
+                  .AddProperty(
+                      PropertyConfigBuilder()
+                          .SetName("subject")
+                          .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                          .SetCardinality(CARDINALITY_OPTIONAL))
+                  .AddProperty(PropertyConfigBuilder()
+                                   .SetName("senderQualifiedId")
+                                   .SetDataTypeJoinableString(
+                                       JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                       DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                                   .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  // Adjust the timestamp to t = 20000 so that the person document is expired
+  // during the schema update.
+  fake_clock_ptr->SetSystemTimeMilliseconds(20000);
+  set_schema_result = icing.SetSchema(schema2);
+  // Ignore latency numbers. They're covered elsewhere.
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
+  expected_set_schema_result = SetSchemaResultProto();
+  expected_set_schema_result.mutable_join_incompatible_changed_schema_types()
+      ->Add("Email");
+  expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
+  expected_set_schema_result.set_has_term_index_restored(false);
+  expected_set_schema_result.set_has_integer_index_restored(false);
+  expected_set_schema_result.set_has_qualified_id_join_index_restored(true);
+  expected_set_schema_result.set_has_embedding_index_restored(false);
+  // Since it is a join incompatible change, we need to do a recovery proof
+  // flush.
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::RECOVERY_PROOF);
+  EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
+
+  // Attempt to get the email document. It should be deleted due to the
+  // dependency on the expired person document.
+  GetResultProto expected_get_result_proto;
+  expected_get_result_proto.mutable_status()->set_code(StatusProto::NOT_FOUND);
+  expected_get_result_proto.mutable_status()->set_message(
+      "Document (namespace, email) not found.");
+  EXPECT_THAT(
+      icing.Get("namespace", "email", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto));
+}
+
 TEST_F(IcingSearchEngineSchemaTest, SetSchemaRevalidatesDocumentsAndReturnsOk) {
   IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
   ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
@@ -3381,14 +3924,14 @@
   SetSchemaResultProto set_schema_result =
       icing.SetSchema(schema_with_required_subject);
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
   SetSchemaResultProto expected_set_schema_result_proto;
   expected_set_schema_result_proto.mutable_status()->set_code(
       StatusProto::FAILED_PRECONDITION);
   expected_set_schema_result_proto.mutable_status()->set_message(
       "Schema is incompatible.");
   expected_set_schema_result_proto.add_incompatible_schema_types("email");
-
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result_proto));
 
   // Force set it
@@ -3396,9 +3939,14 @@
       icing.SetSchema(schema_with_required_subject,
                       /*ignore_errors_and_delete_documents=*/true);
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
   expected_set_schema_result_proto.mutable_status()->set_code(StatusProto::OK);
   expected_set_schema_result_proto.mutable_status()->clear_message();
+  expected_set_schema_result_proto.set_deleted_document_count(1);
+  // Since a document was deleted, we need to do a recovery proof flush.
+  expected_set_schema_result_proto.set_needs_persist_type(
+      PersistType::RECOVERY_PROOF);
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result_proto));
 
   GetResultProto expected_get_result_proto;
@@ -3458,7 +4006,8 @@
   // Can't set the schema since it's incompatible
   SetSchemaResultProto set_schema_result = icing.SetSchema(new_schema);
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
   SetSchemaResultProto expected_result;
   expected_result.mutable_status()->set_code(StatusProto::FAILED_PRECONDITION);
   expected_result.mutable_status()->set_message("Schema is incompatible.");
@@ -3471,9 +4020,13 @@
       icing.SetSchema(new_schema,
                       /*ignore_errors_and_delete_documents=*/true);
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
   expected_result.mutable_status()->set_code(StatusProto::OK);
   expected_result.mutable_status()->clear_message();
+  expected_result.set_deleted_document_count(1);
+  // Since a document was deleted, we need to do a recovery proof flush.
+  expected_result.set_needs_persist_type(PersistType::RECOVERY_PROOF);
   EXPECT_THAT(set_schema_result, EqualsProto(expected_result));
 
   // "email" document is still there
@@ -3523,8 +4076,8 @@
   // Add a schema for db1:
   SchemaTypeConfigProto db1_type =
       SchemaTypeConfigBuilder()
-          .SetType("db1_type")
-          .SetDatabase("db1")
+          .SetType("db1/type")
+          .SetDatabase("db1/")
           .AddProperty(PropertyConfigBuilder()
                            .SetName("a")
                            .SetDataTypeString(TERM_MATCH_EXACT, TOKENIZER_PLAIN)
@@ -3538,19 +4091,21 @@
   SchemaProto db1_schema = SchemaBuilder().AddType(db1_type).Build();
   SetSchemaResultProto set_schema_result =
       icing.SetSchema(CreateSetSchemaRequestProto(
-          db1_schema, "db1", /*ignore_errors_and_delete_documents=*/false));
+          db1_schema, "db1/", /*ignore_errors_and_delete_documents=*/false));
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
   SetSchemaResultProto expected_set_schema_result;
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
-  expected_set_schema_result.mutable_new_schema_types()->Add("db1_type");
+  expected_set_schema_result.mutable_new_schema_types()->Add("db1/type");
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::UNKNOWN);  // Flush is not needed for the first SetSchema.
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   // Add a schema for db2:
   SchemaTypeConfigProto db2_type =
       SchemaTypeConfigBuilder()
-          .SetType("db2_type")
-          .SetDatabase("db2")
+          .SetType("db2/type")
+          .SetDatabase("db2/")
           .AddProperty(
               PropertyConfigBuilder()
                   .SetName("a")
@@ -3564,12 +4119,17 @@
 
   SchemaProto db2_schema = SchemaBuilder().AddType(db2_type).Build();
   set_schema_result = icing.SetSchema(CreateSetSchemaRequestProto(
-      db2_schema, "db2", /*ignore_errors_and_delete_documents=*/false));
+      db2_schema, "db2/", /*ignore_errors_and_delete_documents=*/false));
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
   expected_set_schema_result = SetSchemaResultProto();
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
-  expected_set_schema_result.mutable_new_schema_types()->Add("db2_type");
+  expected_set_schema_result.mutable_new_schema_types()->Add("db2/type");
+  // No need to flush since:
+  // - Schema type id is not changed.
+  // - No document is deleted.
+  // - Only a new schema type is added and indices are not rebuilt.
+  expected_set_schema_result.set_needs_persist_type(PersistType::UNKNOWN);
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   // GetSchema per database
@@ -3577,13 +4137,13 @@
   expected_get_schema_result_proto_db1.mutable_status()->set_code(
       StatusProto::OK);
   *expected_get_schema_result_proto_db1.mutable_schema() = db1_schema;
-  EXPECT_THAT(icing.GetSchema("db1"),
+  EXPECT_THAT(icing.GetSchema("db1/"),
               EqualsProto(expected_get_schema_result_proto_db1));
   GetSchemaResultProto expected_get_schema_result_proto_db2;
   expected_get_schema_result_proto_db2.mutable_status()->set_code(
       StatusProto::OK);
   *expected_get_schema_result_proto_db2.mutable_schema() = db2_schema;
-  EXPECT_THAT(icing.GetSchema("db2"),
+  EXPECT_THAT(icing.GetSchema("db2/"),
               EqualsProto(expected_get_schema_result_proto_db2));
 
   // Get full schema
@@ -3611,8 +4171,8 @@
   SchemaProto db1_schema =
       SchemaBuilder()
           .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db1_type")
-                       .SetDatabase("db1")
+                       .SetType("db1/type")
+                       .SetDatabase("db1/")
                        .AddProperty(PropertyConfigBuilder()
                                         .SetName("a")
                                         .SetDataTypeString(TERM_MATCH_EXACT,
@@ -3625,12 +4185,15 @@
           .Build();
   SetSchemaResultProto set_schema_result =
       icing.SetSchema(CreateSetSchemaRequestProto(
-          db1_schema, "db1", /*ignore_errors_and_delete_documents=*/false));
+          db1_schema, "db1/", /*ignore_errors_and_delete_documents=*/false));
   // Ignore latency numbers. They're covered elsewhere.
-  set_schema_result.clear_latency_ms();
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
   SetSchemaResultProto expected_set_schema_result;
   expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
-  expected_set_schema_result.mutable_new_schema_types()->Add("db1_type");
+  expected_set_schema_result.mutable_new_schema_types()->Add("db1/type");
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::UNKNOWN);  // Flush is not needed for the first SetSchema.
   EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   get_schema_result_proto = icing.GetSchema("nonexistent_db");
@@ -3703,7 +4266,11 @@
       StatusProto::OK);
   *expected_get_schema_type_result_proto.mutable_schema_type_config() =
       schema_type_config;
-  EXPECT_THAT(icing.GetSchemaType("SchemaType"),
+  GetSchemaTypeResultProto actual_get_schema_type_result_proto =
+      icing.GetSchemaType("SchemaType");
+  actual_get_schema_type_result_proto
+      .clear_vm_binder_transaction_latency_start_time_ms();
+  EXPECT_THAT(actual_get_schema_type_result_proto,
               EqualsProto(expected_get_schema_type_result_proto));
 }
 
@@ -3790,7 +4357,6 @@
   // Setting the new, different schema will remove incompatible documents
   IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
   EXPECT_THAT(icing.Initialize().status(), ProtoIsOk());
-  EXPECT_THAT(icing.SetSchema(incompatible_schema).status(), ProtoIsOk());
 
   // Can't retrieve by namespace/uri
   GetResultProto expected_get_result_proto;
@@ -4152,6 +4718,113 @@
               Eq(IcingSearchEngineMarkerProto::OperationType::SET_SCHEMA));
 }
 
+TEST_F(IcingSearchEngineSchemaTest, SetSchemaStatsArePopulated) {
+  auto fake_clock = std::make_unique<FakeClock>();
+  fake_clock->SetTimerElapsedMilliseconds(1000);
+  TestIcingSearchEngine icing(GetDefaultIcingOptions(),
+                              std::make_unique<Filesystem>(),
+                              std::make_unique<IcingFilesystem>(),
+                              std::move(fake_clock), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+
+  SchemaProto message_schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("Message")
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("body")
+                                        .SetDataTypeString(TERM_MATCH_UNKNOWN,
+                                                           TOKENIZER_NONE)
+                                        .SetCardinality(CARDINALITY_REQUIRED))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("indexableInteger")
+                                        .SetDataTypeInt64(NUMERIC_MATCH_RANGE)
+                                        .SetCardinality(CARDINALITY_REQUIRED)))
+          .Build();
+
+  SetSchemaRequestProto set_schema_request =
+      CreateSetSchemaRequestProto(message_schema, /*database=*/"",
+                                  /*ignore_errors_and_delete_documents=*/false);
+  SetSchemaResultProto set_schema_result =
+      icing.SetSchema(std::move(set_schema_request));
+  EXPECT_THAT(set_schema_result.status(), ProtoStatusIs(StatusProto::OK));
+  EXPECT_THAT(set_schema_result.set_schema_stats().overall_latency_ms(),
+              Eq(1000));
+  EXPECT_THAT(
+      set_schema_result.set_schema_stats().schema_store_set_schema_latency_ms(),
+      Eq(1000));
+
+  // Put a document to trigger index restoration later.
+  DocumentProto message_document =
+      DocumentBuilder()
+          .SetKey("namespace", "message")
+          .SetSchema("Message")
+          .AddStringProperty("body", "test body")
+          .AddInt64Property("indexableInteger", 123)
+          .Build();
+  EXPECT_THAT(icing.Put(message_document).status(),
+              ProtoStatusIs(StatusProto::OK));
+
+  // This schema update will trigger and set latencies for index restoration and
+  // scorable property cache regeneration.
+  SchemaProto updated_schema =
+      SchemaBuilder()
+          .AddType(
+              SchemaTypeConfigBuilder()
+                  .SetType("Message")
+                  .AddProperty(
+                      PropertyConfigBuilder()
+                          .SetName("body")
+                          .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                          .SetCardinality(CARDINALITY_REQUIRED))
+                  .AddProperty(PropertyConfigBuilder()
+                                   .SetName("indexableInteger")
+                                   .SetDataTypeInt64(NUMERIC_MATCH_UNKNOWN)
+                                   .SetCardinality(CARDINALITY_REQUIRED))
+                  .AddProperty(
+                      PropertyConfigBuilder()
+                          .SetName("score")
+                          .SetDataType(PropertyConfigProto::DataType::DOUBLE)
+                          .SetScorableType(SCORABLE_TYPE_ENABLED)
+                          .SetCardinality(CARDINALITY_REPEATED))
+                  .AddProperty(PropertyConfigBuilder()
+                                   .SetName("senderQualifiedId")
+                                   .SetDataTypeJoinableString(
+                                       JOINABLE_VALUE_TYPE_QUALIFIED_ID)
+                                   .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+  SetSchemaRequestProto set_schema_request2 =
+      CreateSetSchemaRequestProto(updated_schema, /*database=*/"",
+                                  /*ignore_errors_and_delete_documents=*/false);
+  set_schema_result = icing.SetSchema(std::move(set_schema_request2));
+  EXPECT_THAT(set_schema_result.status(), ProtoStatusIs(StatusProto::OK));
+  EXPECT_THAT(set_schema_result.set_schema_stats().overall_latency_ms(),
+              Eq(1000));
+  EXPECT_THAT(
+      set_schema_result.set_schema_stats().index_restoration_latency_ms(),
+      Eq(1000));
+  EXPECT_THAT(set_schema_result.set_schema_stats()
+                  .scorable_property_cache_regeneration_latency_ms(),
+              Eq(1000));
+
+  // Clear schema -- this should trigger and update document store and set
+  // latencies for document store update schema
+  SchemaProto empty_schema = SchemaBuilder().Build();
+  SetSchemaRequestProto set_schema_request3 =
+      CreateSetSchemaRequestProto(empty_schema, /*database=*/"",
+                                  /*ignore_errors_and_delete_documents=*/true);
+  set_schema_result = icing.SetSchema(std::move(set_schema_request3));
+  EXPECT_THAT(set_schema_result.status(), ProtoStatusIs(StatusProto::OK));
+  EXPECT_THAT(set_schema_result.set_schema_stats().overall_latency_ms(),
+              Eq(1000));
+  EXPECT_THAT(
+      set_schema_result.set_schema_stats().schema_store_set_schema_latency_ms(),
+      Eq(1000));
+  EXPECT_THAT(set_schema_result.set_schema_stats()
+                  .document_store_optimized_update_schema_latency_ms(),
+              Eq(1000));
+}
+
 }  // namespace
 }  // namespace lib
 }  // namespace icing

diff --git a/icing/icing-search-engine_search_test.cc b/icing/icing-search-engine_search_test.cc
index 46040a4..0835439 100644
--- a/icing/icing-search-engine_search_test.cc
+++ b/icing/icing-search-engine_search_test.cc

@@ -19,6 +19,7 @@
 #include <memory>
 #include <string>
 #include <string_view>
+#include <tuple>
 #include <utility>
 #include <vector>
 
@@ -67,13 +68,15 @@
 namespace {
 
 using ::icing::lib::portable_equals_proto::EqualsProto;
+using ::testing::AnyOf;
 using ::testing::DoubleEq;
 using ::testing::DoubleNear;
 using ::testing::ElementsAre;
 using ::testing::Eq;
-using ::testing::Gt;
 using ::testing::HasSubstr;
 using ::testing::IsEmpty;
+using ::testing::IsFalse;
+using ::testing::IsTrue;
 using ::testing::Lt;
 using ::testing::Ne;
 using ::testing::SizeIs;
@@ -137,8 +140,13 @@
   icing_options.set_enable_embedding_index(true);
   icing_options.set_enable_scorable_properties(true);
   icing_options.set_enable_embedding_quantization(true);
+  icing_options.set_enable_repeated_field_joins(true);
+  icing_options.set_enable_soft_index_restoration(true);
   icing_options.set_enable_qualified_id_join_index_v3(true);
-  icing_options.set_enable_delete_propagation_from(false);
+  icing_options.set_enable_delete_propagation_from(true);
+  icing_options.set_enable_passing_filter_to_children(true);
+  icing_options.set_enable_embedding_iterator_v2(true);
+  icing_options.set_enable_non_existent_qualified_id_join(true);
   return icing_options;
 }
 
@@ -936,10 +944,11 @@
       document4;
   SearchResultProto search_result_proto =
       icing.Search(search_spec, GetDefaultScoringSpec(), result_spec);
-  EXPECT_THAT(search_result_proto.next_page_token(), Gt(kInvalidNextPageToken));
   uint64_t next_page_token = search_result_proto.next_page_token();
+
   // Since the token is a random number, we don't need to verify
   expected_search_result_proto.set_next_page_token(next_page_token);
+  EXPECT_THAT(search_result_proto.next_page_token(), Ne(kInvalidNextPageToken));
   EXPECT_THAT(search_result_proto, EqualsSearchResultIgnoreStatsAndScores(
                                        expected_search_result_proto));
 
@@ -950,6 +959,7 @@
   *expected_search_result_proto.mutable_results()->Add()->mutable_document() =
       document2;
   search_result_proto = icing.GetNextPage(next_page_token);
+  EXPECT_THAT(search_result_proto.next_page_token(), Ne(kInvalidNextPageToken));
   EXPECT_THAT(search_result_proto, EqualsSearchResultIgnoreStatsAndScores(
                                        expected_search_result_proto));
 
@@ -961,12 +971,153 @@
   // token.
   expected_search_result_proto.clear_next_page_token();
   search_result_proto = icing.GetNextPage(next_page_token);
+  EXPECT_THAT(search_result_proto.next_page_token(),
+              Eq(kInvalidNextPageToken));  // No more results.
+  EXPECT_THAT(search_result_proto, EqualsSearchResultIgnoreStatsAndScores(
+                                       expected_search_result_proto));
+}
+
+TEST_F(IcingSearchEngineSearchTest,
+       SearchShouldReturnMultiplePages_withMaxResultsLimit) {
+  IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(CreateMessageSchema()).status(), ProtoIsOk());
+
+  // Creates and inserts 5 documents
+  DocumentProto document1 = CreateMessageDocument("namespace", "uri1");
+  DocumentProto document2 = CreateMessageDocument("namespace", "uri2");
+  DocumentProto document3 = CreateMessageDocument("namespace", "uri3");
+  DocumentProto document4 = CreateMessageDocument("namespace", "uri4");
+  DocumentProto document5 = CreateMessageDocument("namespace", "uri5");
+  ASSERT_THAT(icing.Put(document1).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(document2).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(document3).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(document4).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(document5).status(), ProtoIsOk());
+
+  SearchSpecProto search_spec;
+  search_spec.set_term_match_type(TermMatchType::PREFIX);
+  search_spec.set_query("message");
+
+  ResultSpecProto result_spec;
+  result_spec.set_num_per_page(2);
+
+  // Searches and gets the first page, 2 results
+  SearchResultProto expected_search_result_proto;
+  expected_search_result_proto.mutable_status()->set_code(StatusProto::OK);
+  *expected_search_result_proto.mutable_results()->Add()->mutable_document() =
+      document5;
+  *expected_search_result_proto.mutable_results()->Add()->mutable_document() =
+      document4;
+  SearchResultProto search_result_proto =
+      icing.Search(search_spec, GetDefaultScoringSpec(), result_spec);
+  uint64_t next_page_token = search_result_proto.next_page_token();
+
+  // Since the token is a random number, we don't need to verify
+  expected_search_result_proto.set_next_page_token(next_page_token);
+  EXPECT_THAT(search_result_proto.next_page_token(), Ne(kInvalidNextPageToken));
   EXPECT_THAT(search_result_proto, EqualsSearchResultIgnoreStatsAndScores(
                                        expected_search_result_proto));
 
-  // No more results
+  // Second page, with a max limit of 1 result
   expected_search_result_proto.clear_results();
-  search_result_proto = icing.GetNextPage(next_page_token);
+  *expected_search_result_proto.mutable_results()->Add()->mutable_document() =
+      document3;
+  GetNextPageRequestProto get_next_page_request;
+  get_next_page_request.set_next_page_token(next_page_token);
+  get_next_page_request.set_max_results_to_retrieve_from_page(1);
+  search_result_proto = icing.GetNextPage(std::move(get_next_page_request));
+  EXPECT_THAT(search_result_proto.next_page_token(), Ne(kInvalidNextPageToken));
+  EXPECT_THAT(search_result_proto, EqualsSearchResultIgnoreStatsAndScores(
+                                       expected_search_result_proto));
+
+  // Third page, 2 results
+  expected_search_result_proto.clear_results();
+  *expected_search_result_proto.mutable_results()->Add()->mutable_document() =
+      document2;
+  *expected_search_result_proto.mutable_results()->Add()->mutable_document() =
+      document1;
+  get_next_page_request = GetNextPageRequestProto();
+  get_next_page_request.set_next_page_token(next_page_token);
+  search_result_proto = icing.GetNextPage(std::move(get_next_page_request));
+  // Because there are no more results, we should not return the next page
+  // token.
+  expected_search_result_proto.clear_next_page_token();
+  EXPECT_THAT(search_result_proto.next_page_token(),
+              Eq(kInvalidNextPageToken));  // No more results.
+  EXPECT_THAT(search_result_proto, EqualsSearchResultIgnoreStatsAndScores(
+                                       expected_search_result_proto));
+}
+
+TEST_F(IcingSearchEngineSearchTest,
+       SearchShouldReturnMultiplePages_maxResultsLimitLargerThanPageSize) {
+  IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(CreateMessageSchema()).status(), ProtoIsOk());
+
+  // Creates and inserts 5 documents
+  DocumentProto document1 = CreateMessageDocument("namespace", "uri1");
+  DocumentProto document2 = CreateMessageDocument("namespace", "uri2");
+  DocumentProto document3 = CreateMessageDocument("namespace", "uri3");
+  DocumentProto document4 = CreateMessageDocument("namespace", "uri4");
+  DocumentProto document5 = CreateMessageDocument("namespace", "uri5");
+  ASSERT_THAT(icing.Put(document1).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(document2).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(document3).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(document4).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(document5).status(), ProtoIsOk());
+
+  SearchSpecProto search_spec;
+  search_spec.set_term_match_type(TermMatchType::PREFIX);
+  search_spec.set_query("message");
+
+  ResultSpecProto result_spec;
+  result_spec.set_num_per_page(2);
+
+  // Searches and gets the first page -- this should return 2 results
+  SearchResultProto expected_search_result_proto;
+  expected_search_result_proto.mutable_status()->set_code(StatusProto::OK);
+  *expected_search_result_proto.mutable_results()->Add()->mutable_document() =
+      document5;
+  *expected_search_result_proto.mutable_results()->Add()->mutable_document() =
+      document4;
+  SearchResultProto search_result_proto =
+      icing.Search(search_spec, GetDefaultScoringSpec(), result_spec);
+  uint64_t next_page_token = search_result_proto.next_page_token();
+
+  // Since the token is a random number, we don't need to verify
+  expected_search_result_proto.set_next_page_token(next_page_token);
+  EXPECT_THAT(search_result_proto.next_page_token(), Ne(kInvalidNextPageToken));
+  EXPECT_THAT(search_result_proto, EqualsSearchResultIgnoreStatsAndScores(
+                                       expected_search_result_proto));
+
+  // Second page with a max limit of 5. This should still only return 2
+  // results.
+  expected_search_result_proto.clear_results();
+  *expected_search_result_proto.mutable_results()->Add()->mutable_document() =
+      document3;
+  *expected_search_result_proto.mutable_results()->Add()->mutable_document() =
+      document2;
+  GetNextPageRequestProto get_next_page_request;
+  get_next_page_request.set_next_page_token(next_page_token);
+  get_next_page_request.set_max_results_to_retrieve_from_page(5);
+  search_result_proto = icing.GetNextPage(std::move(get_next_page_request));
+  EXPECT_THAT(search_result_proto.next_page_token(), Ne(kInvalidNextPageToken));
+  EXPECT_THAT(search_result_proto, EqualsSearchResultIgnoreStatsAndScores(
+                                       expected_search_result_proto));
+
+  // Third page, only 1 result left.
+  expected_search_result_proto.clear_results();
+  *expected_search_result_proto.mutable_results()->Add()->mutable_document() =
+      document1;
+  get_next_page_request = GetNextPageRequestProto();
+  get_next_page_request.set_next_page_token(next_page_token);
+  search_result_proto = icing.GetNextPage(std::move(get_next_page_request));
+  // Because there are no more results, we should not return the next page
+  // token.
+  expected_search_result_proto.clear_next_page_token();
+  EXPECT_THAT(search_result_proto.next_page_token(),
+              Eq(kInvalidNextPageToken));  // No more results.
   EXPECT_THAT(search_result_proto, EqualsSearchResultIgnoreStatsAndScores(
                                        expected_search_result_proto));
 }
@@ -1008,10 +1159,11 @@
       document4;
   SearchResultProto search_result_proto =
       icing.Search(search_spec, scoring_spec, result_spec);
-  EXPECT_THAT(search_result_proto.next_page_token(), Gt(kInvalidNextPageToken));
   uint64_t next_page_token = search_result_proto.next_page_token();
+
   // Since the token is a random number, we don't need to verify
   expected_search_result_proto.set_next_page_token(next_page_token);
+  EXPECT_THAT(search_result_proto.next_page_token(), Ne(kInvalidNextPageToken));
   EXPECT_THAT(search_result_proto, EqualsSearchResultIgnoreStatsAndScores(
                                        expected_search_result_proto));
 
@@ -1022,6 +1174,7 @@
   *expected_search_result_proto.mutable_results()->Add()->mutable_document() =
       document2;
   search_result_proto = icing.GetNextPage(next_page_token);
+  EXPECT_THAT(search_result_proto.next_page_token(), Ne(kInvalidNextPageToken));
   EXPECT_THAT(search_result_proto, EqualsSearchResultIgnoreStatsAndScores(
                                        expected_search_result_proto));
 
@@ -1033,12 +1186,8 @@
   // token.
   expected_search_result_proto.clear_next_page_token();
   search_result_proto = icing.GetNextPage(next_page_token);
-  EXPECT_THAT(search_result_proto, EqualsSearchResultIgnoreStatsAndScores(
-                                       expected_search_result_proto));
-
-  // No more results
-  expected_search_result_proto.clear_results();
-  search_result_proto = icing.GetNextPage(next_page_token);
+  EXPECT_THAT(search_result_proto.next_page_token(),
+              Eq(kInvalidNextPageToken));  // No more results.
   EXPECT_THAT(search_result_proto, EqualsSearchResultIgnoreStatsAndScores(
                                        expected_search_result_proto));
 }
@@ -1094,7 +1243,7 @@
       icing.Search(search_spec, GetDefaultScoringSpec(), result_spec);
   ASSERT_THAT(search_result.status(), ProtoIsOk());
   ASSERT_THAT(search_result.results(), SizeIs(2));
-  ASSERT_THAT(search_result.next_page_token(), Gt(kInvalidNextPageToken));
+  ASSERT_THAT(search_result.next_page_token(), Ne(kInvalidNextPageToken));
 
   const DocumentProto& document_result_1 = search_result.results(0).document();
   EXPECT_THAT(document_result_1, EqualsProto(document5));
@@ -1123,7 +1272,7 @@
   search_result = icing.GetNextPage(search_result.next_page_token());
   ASSERT_THAT(search_result.status(), ProtoIsOk());
   ASSERT_THAT(search_result.results(), SizeIs(2));
-  ASSERT_THAT(search_result.next_page_token(), Gt(kInvalidNextPageToken));
+  ASSERT_THAT(search_result.next_page_token(), Ne(kInvalidNextPageToken));
 
   const DocumentProto& document_result_3 = search_result.results(0).document();
   EXPECT_THAT(document_result_3, EqualsProto(document3));
@@ -1173,7 +1322,7 @@
       document2;
   SearchResultProto search_result_proto =
       icing.Search(search_spec, GetDefaultScoringSpec(), result_spec);
-  EXPECT_THAT(search_result_proto.next_page_token(), Gt(kInvalidNextPageToken));
+  EXPECT_THAT(search_result_proto.next_page_token(), Ne(kInvalidNextPageToken));
   uint64_t next_page_token = search_result_proto.next_page_token();
   // Since the token is a random number, we don't need to verify
   expected_search_result_proto.set_next_page_token(next_page_token);
@@ -1184,9 +1333,12 @@
   // Invalidates token
   icing.InvalidateNextPageToken(next_page_token);
 
-  // Tries to fetch the second page, no result since it's invalidated
+  // Tries to fetch the second page, no result since it's invalidated. Also
+  // page_token_not_found should be set to true in order to hint the client that
+  // the pagination is incomplete.
   expected_search_result_proto.clear_results();
   expected_search_result_proto.clear_next_page_token();
+  expected_search_result_proto.set_page_token_not_found(true);
   search_result_proto = icing.GetNextPage(next_page_token);
   EXPECT_THAT(search_result_proto, EqualsSearchResultIgnoreStatsAndScores(
                                        expected_search_result_proto));
@@ -1277,7 +1429,7 @@
   // Time just has to be greater than the document's creation timestamp (100) +
   // the document's ttl (500)
   auto fake_clock = std::make_unique<FakeClock>();
-  fake_clock->SetSystemTimeMilliseconds(700);
+  FakeClock* fake_clock_ptr = fake_clock.get();
 
   TestIcingSearchEngine icing(GetDefaultIcingOptions(),
                               std::make_unique<Filesystem>(),
@@ -1288,6 +1440,10 @@
   EXPECT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
   EXPECT_THAT(icing.Put(document).status(), ProtoIsOk());
 
+  // Adjust the clock to be greater than the document's creation timestamp (100)
+  // + the document's ttl (500).
+  fake_clock_ptr->SetSystemTimeMilliseconds(700);
+
   // Check that the document is not returned as part of search results
   SearchResultProto search_result_proto =
       icing.Search(search_spec, GetDefaultScoringSpec(),
@@ -3910,71 +4066,6 @@
   EXPECT_THAT(results.results(), IsEmpty());
 }
 
-TEST_F(IcingSearchEngineSearchTest, EmptySearchWithPropertyFilter) {
-  IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
-  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
-  ASSERT_THAT(icing.SetSchema(CreatePersonAndEmailSchema()).status(),
-              ProtoIsOk());
-
-  // 1. Add two email documents
-  DocumentProto document_one =
-      DocumentBuilder()
-          .SetKey("namespace", "uri1")
-          .SetCreationTimestampMs(1000)
-          .SetSchema("Email")
-          .AddDocumentProperty(
-              "sender",
-              DocumentBuilder()
-                  .SetKey("namespace", "uri1")
-                  .SetSchema("Person")
-                  .AddStringProperty("name", "Meg Ryan")
-                  .AddStringProperty("emailAddress", "hellogirl@aol.com")
-                  .Build())
-          .AddStringProperty("subject", "Hello World!")
-          .AddStringProperty(
-              "body", "Oh what a beautiful morning! Oh what a beautiful day!")
-          .Build();
-  ASSERT_THAT(icing.Put(document_one).status(), ProtoIsOk());
-
-  DocumentProto document_two =
-      DocumentBuilder()
-          .SetKey("namespace", "uri2")
-          .SetCreationTimestampMs(1000)
-          .SetSchema("Email")
-          .AddDocumentProperty(
-              "sender", DocumentBuilder()
-                            .SetKey("namespace", "uri2")
-                            .SetSchema("Person")
-                            .AddStringProperty("name", "Tom Hanks")
-                            .AddStringProperty("emailAddress", "ny152@aol.com")
-                            .Build())
-          .AddStringProperty("subject", "Goodnight Moon!")
-          .AddStringProperty("body",
-                             "Count all the sheep and tell them 'Hello'.")
-          .Build();
-  ASSERT_THAT(icing.Put(document_two).status(), ProtoIsOk());
-
-  // 2. Issue a query with a property filter
-  auto search_spec = std::make_unique<SearchSpecProto>();
-  search_spec->set_term_match_type(TermMatchType::PREFIX);
-  search_spec->set_query("");
-
-  TypePropertyMask* email_property_filters =
-      search_spec->add_type_property_filters();
-  email_property_filters->set_schema_type("Email");
-  email_property_filters->add_paths("subject");
-
-  auto result_spec = std::make_unique<ResultSpecProto>();
-
-  // 3. Verify that both documents are returned.
-  auto scoring_spec = std::make_unique<ScoringSpecProto>();
-  *scoring_spec = GetDefaultScoringSpec();
-  SearchResultProto results =
-      icing.Search(*search_spec, *scoring_spec, *result_spec);
-  EXPECT_THAT(results.status(), ProtoIsOk());
-  EXPECT_THAT(results.results(), SizeIs(2));
-}
-
 TEST_F(IcingSearchEngineSearchTest, EmptySearchWithEmptyPropertyFilter) {
   IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
   ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
@@ -4713,6 +4804,10 @@
   exp_parent_search_stats->set_query_processor_parser_consume_query_latency_ms(
       5);
   exp_parent_search_stats->set_query_processor_query_visitor_latency_ms(5);
+  exp_parent_search_stats->set_num_unquantized_embeddings_scored(0);
+  exp_parent_search_stats->set_num_quantized_embeddings_scored(0);
+  exp_parent_search_stats->set_num_embedding_shards_read(0);
+  exp_parent_search_stats->set_num_embedding_bytes_read(0);
 
   EXPECT_THAT(search_result.query_stats(), EqualsProto(exp_stats));
 
@@ -4720,7 +4815,7 @@
   search_result = icing.GetNextPage(search_result.next_page_token());
   ASSERT_THAT(search_result.status(), ProtoIsOk());
   ASSERT_THAT(search_result.results(), SizeIs(2));
-  ASSERT_THAT(search_result.next_page_token(), Gt(kInvalidNextPageToken));
+  ASSERT_THAT(search_result.next_page_token(), Ne(kInvalidNextPageToken));
 
   exp_stats = QueryStatsProto();
   exp_stats.set_is_first_page(false);
@@ -4731,6 +4826,7 @@
   exp_stats.set_document_retrieval_latency_ms(5);
   exp_stats.set_lock_acquisition_latency_ms(5);
   exp_stats.set_num_joined_results_returned_current_page(0);
+  exp_stats.set_page_token_type(QueryStatsProto::PageTokenType::VALID);
   EXPECT_THAT(search_result.query_stats(), EqualsProto(exp_stats));
 
   // Third page, 1 result with 0 snippets
@@ -4748,6 +4844,75 @@
   exp_stats.set_document_retrieval_latency_ms(5);
   exp_stats.set_lock_acquisition_latency_ms(5);
   exp_stats.set_num_joined_results_returned_current_page(0);
+  exp_stats.set_page_token_type(QueryStatsProto::PageTokenType::VALID);
+  EXPECT_THAT(search_result.query_stats(), EqualsProto(exp_stats));
+
+  // Fetch the next page with kInvalidNextPageToken.
+  search_result = icing.GetNextPage(kInvalidNextPageToken);
+  ASSERT_THAT(search_result.status(), ProtoIsOk());
+  ASSERT_THAT(search_result.results(), IsEmpty());
+  ASSERT_THAT(search_result.next_page_token(), Eq(kInvalidNextPageToken));
+
+  exp_stats = QueryStatsProto();
+  exp_stats.set_is_first_page(false);
+  exp_stats.set_lock_acquisition_latency_ms(5);
+  exp_stats.set_page_token_type(QueryStatsProto::PageTokenType::EMPTY);
+  EXPECT_THAT(search_result.query_stats(), EqualsProto(exp_stats));
+}
+
+TEST_F(IcingSearchEngineSearchTest, GetNextPage_withNotFoundPageToken) {
+  auto fake_clock = std::make_unique<FakeClock>();
+  fake_clock->SetTimerElapsedMilliseconds(5);
+
+  TestIcingSearchEngine icing(GetDefaultIcingOptions(),
+                              std::make_unique<Filesystem>(),
+                              std::make_unique<IcingFilesystem>(),
+                              std::move(fake_clock), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+
+  // Call GetNextPage with a not found page token.
+  SearchResultProto search_result = icing.GetNextPage(/*next_page_token=*/123);
+  ASSERT_THAT(search_result.status(), ProtoIsOk());
+  ASSERT_THAT(search_result.results(), IsEmpty());
+  ASSERT_THAT(search_result.next_page_token(), Eq(kInvalidNextPageToken));
+  // page_token_not_found is set to true.
+  // - Usually the client uses a valid page token to fetch the next page. If the
+  //   page token is not found, it is "likely" that the corresponding
+  //   ResultState has been removed due to cache eviction.
+  // - If the client really tries to fetch the next page with an non-existing
+  //   page token, then we cannot tell whether the search is "invalid" or
+  //   "incomplete". Still, it is just a hint warning field for the 1st case.
+  EXPECT_THAT(search_result.page_token_not_found(), IsTrue());
+
+  QueryStatsProto exp_stats = QueryStatsProto();
+  exp_stats.set_is_first_page(false);
+  exp_stats.set_lock_acquisition_latency_ms(5);
+  exp_stats.set_page_token_type(QueryStatsProto::PageTokenType::NOT_FOUND);
+  EXPECT_THAT(search_result.query_stats(), EqualsProto(exp_stats));
+}
+
+TEST_F(IcingSearchEngineSearchTest, GetNextPage_withInvalidPageToken) {
+  auto fake_clock = std::make_unique<FakeClock>();
+  fake_clock->SetTimerElapsedMilliseconds(5);
+
+  TestIcingSearchEngine icing(GetDefaultIcingOptions(),
+                              std::make_unique<Filesystem>(),
+                              std::make_unique<IcingFilesystem>(),
+                              std::move(fake_clock), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+
+  // Call GetNextPage with kInvalidNextPageToken.
+  SearchResultProto search_result = icing.GetNextPage(kInvalidNextPageToken);
+  ASSERT_THAT(search_result.status(), ProtoIsOk());
+  ASSERT_THAT(search_result.results(), IsEmpty());
+  ASSERT_THAT(search_result.next_page_token(), Eq(kInvalidNextPageToken));
+  // page_token_not_found should be false.
+  EXPECT_THAT(search_result.page_token_not_found(), IsFalse());
+
+  QueryStatsProto exp_stats = QueryStatsProto();
+  exp_stats.set_is_first_page(false);
+  exp_stats.set_lock_acquisition_latency_ms(5);
+  exp_stats.set_page_token_type(QueryStatsProto::PageTokenType::EMPTY);
   EXPECT_THAT(search_result.query_stats(), EqualsProto(exp_stats));
 }
 
@@ -4995,6 +5160,10 @@
   exp_parent_search_stats->set_query_processor_parser_consume_query_latency_ms(
       5);
   exp_parent_search_stats->set_query_processor_query_visitor_latency_ms(5);
+  exp_parent_search_stats->set_num_unquantized_embeddings_scored(0);
+  exp_parent_search_stats->set_num_quantized_embeddings_scored(0);
+  exp_parent_search_stats->set_num_embedding_shards_read(0);
+  exp_parent_search_stats->set_num_embedding_bytes_read(0);
 
   QueryStatsProto::SearchStats* exp_child_search_stats =
       exp_stats.mutable_child_search_stats();
@@ -5014,6 +5183,10 @@
   exp_child_search_stats->set_query_processor_parser_consume_query_latency_ms(
       5);
   exp_child_search_stats->set_query_processor_query_visitor_latency_ms(5);
+  exp_child_search_stats->set_num_unquantized_embeddings_scored(0);
+  exp_child_search_stats->set_num_quantized_embeddings_scored(0);
+  exp_child_search_stats->set_num_embedding_shards_read(0);
+  exp_child_search_stats->set_num_embedding_bytes_read(0);
 
   EXPECT_THAT(search_result.query_stats(), EqualsProto(exp_stats));
 
@@ -5034,6 +5207,7 @@
   exp_stats.set_document_retrieval_latency_ms(5);
   exp_stats.set_lock_acquisition_latency_ms(5);
   exp_stats.set_num_joined_results_returned_current_page(1);
+  exp_stats.set_page_token_type(QueryStatsProto::PageTokenType::VALID);
   EXPECT_THAT(search_result.query_stats(), EqualsProto(exp_stats));
 
   // Third page, 0 child docs.
@@ -5052,6 +5226,7 @@
   exp_stats.set_document_retrieval_latency_ms(5);
   exp_stats.set_lock_acquisition_latency_ms(5);
   exp_stats.set_num_results_with_snippets(0);
+  exp_stats.set_page_token_type(QueryStatsProto::PageTokenType::VALID);
   ASSERT_THAT(search_result,
               EqualsSearchResultIgnoreStatsAndScores(expected_result3));
   EXPECT_THAT(search_result.query_stats(), EqualsProto(exp_stats));
@@ -5066,6 +5241,7 @@
   exp_stats = QueryStatsProto();
   exp_stats.set_is_first_page(false);
   exp_stats.set_lock_acquisition_latency_ms(5);
+  exp_stats.set_page_token_type(QueryStatsProto::PageTokenType::EMPTY);
   EXPECT_THAT(search_result.query_stats(), EqualsProto(exp_stats));
 }
 
@@ -5553,6 +5729,361 @@
               EqualsSearchResultIgnoreStatsAndScores(expected_result3));
 }
 
+TEST_F(IcingSearchEngineSearchTest,
+       JoinByQualifiedId_ChildPutBeforeParent_FixDisabled) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("firstName")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("Email")
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("subject")
+                                        .SetDataTypeString(TERM_MATCH_PREFIX,
+                                                           TOKENIZER_PLAIN)
+                                        .SetCardinality(CARDINALITY_OPTIONAL))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("personQualifiedId")
+                                        .SetDataTypeJoinableString(
+                                            JOINABLE_VALUE_TYPE_QUALIFIED_ID)
+                                        .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  DocumentProto person =
+      DocumentBuilder()
+          .SetKey("pkg$db/namespace", "person")
+          .SetSchema("Person")
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .AddStringProperty("firstName", "first")
+          .Build();
+
+  DocumentProto email =
+      DocumentBuilder()
+          .SetKey("namespace", "email")
+          .SetSchema("Email")
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .AddStringProperty("subject", "test subject")
+          .AddStringProperty("personQualifiedId", "pkg$db/namespace#person")
+          .Build();
+
+  IcingSearchEngineOptions icing_options = GetDefaultIcingOptions();
+  icing_options.set_enable_non_existent_qualified_id_join(false);
+  IcingSearchEngine icing(icing_options, GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(email).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(person).status(), ProtoIsOk());
+
+  // Parent SearchSpec
+  SearchSpecProto search_spec;
+  search_spec.set_term_match_type(TermMatchType::PREFIX);
+  search_spec.set_query("firstName:first");
+
+  // JoinSpec
+  JoinSpecProto* join_spec = search_spec.mutable_join_spec();
+  join_spec->set_parent_property_expression(
+      std::string(JoinProcessor::kQualifiedIdExpr));
+  join_spec->set_child_property_expression("personQualifiedId");
+  JoinSpecProto::NestedSpecProto* nested_spec =
+      join_spec->mutable_nested_spec();
+  SearchSpecProto* nested_search_spec = nested_spec->mutable_search_spec();
+  nested_search_spec->set_term_match_type(TermMatchType::PREFIX);
+  nested_search_spec->set_query("subject:test");
+
+  *nested_spec->mutable_scoring_spec() = GetDefaultScoringSpec();
+  *nested_spec->mutable_result_spec() = ResultSpecProto::default_instance();
+
+  // Parent ScoringSpec
+  ScoringSpecProto scoring_spec = GetDefaultScoringSpec();
+
+  // Parent ResultSpec
+  ResultSpecProto result_spec;
+  result_spec.set_max_joined_children_per_parent_to_return(
+      std::numeric_limits<int32_t>::max());
+
+  // The join relationship should not be established.
+  SearchResultProto expected_result;
+  expected_result.mutable_status()->set_code(StatusProto::OK);
+  SearchResultProto::ResultProto* result_proto =
+      expected_result.mutable_results()->Add();
+  *result_proto->mutable_document() = person;
+
+  EXPECT_THAT(icing.Search(search_spec, scoring_spec, result_spec),
+              EqualsSearchResultIgnoreStatsAndScores(expected_result));
+}
+
+TEST_F(IcingSearchEngineSearchTest, JoinByQualifiedId_ChildPutBeforeParent) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("firstName")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("Email")
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("subject")
+                                        .SetDataTypeString(TERM_MATCH_PREFIX,
+                                                           TOKENIZER_PLAIN)
+                                        .SetCardinality(CARDINALITY_OPTIONAL))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("personQualifiedId")
+                                        .SetDataTypeJoinableString(
+                                            JOINABLE_VALUE_TYPE_QUALIFIED_ID)
+                                        .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  DocumentProto person =
+      DocumentBuilder()
+          .SetKey("pkg$db/namespace", "person")
+          .SetSchema("Person")
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .AddStringProperty("firstName", "first")
+          .Build();
+
+  DocumentProto email =
+      DocumentBuilder()
+          .SetKey("namespace", "email")
+          .SetSchema("Email")
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .AddStringProperty("subject", "test subject")
+          .AddStringProperty("personQualifiedId", "pkg$db/namespace#person")
+          .Build();
+
+  IcingSearchEngineOptions icing_options = GetDefaultIcingOptions();
+  icing_options.set_enable_non_existent_qualified_id_join(true);
+  IcingSearchEngine icing(icing_options, GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(email).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(person).status(), ProtoIsOk());
+
+  // Parent SearchSpec
+  SearchSpecProto search_spec;
+  search_spec.set_term_match_type(TermMatchType::PREFIX);
+  search_spec.set_query("firstName:first");
+
+  // JoinSpec
+  JoinSpecProto* join_spec = search_spec.mutable_join_spec();
+  join_spec->set_parent_property_expression(
+      std::string(JoinProcessor::kQualifiedIdExpr));
+  join_spec->set_child_property_expression("personQualifiedId");
+  JoinSpecProto::NestedSpecProto* nested_spec =
+      join_spec->mutable_nested_spec();
+  SearchSpecProto* nested_search_spec = nested_spec->mutable_search_spec();
+  nested_search_spec->set_term_match_type(TermMatchType::PREFIX);
+  nested_search_spec->set_query("subject:test");
+
+  *nested_spec->mutable_scoring_spec() = GetDefaultScoringSpec();
+  *nested_spec->mutable_result_spec() = ResultSpecProto::default_instance();
+
+  // Parent ScoringSpec
+  ScoringSpecProto scoring_spec = GetDefaultScoringSpec();
+
+  // Parent ResultSpec
+  ResultSpecProto result_spec;
+  result_spec.set_max_joined_children_per_parent_to_return(
+      std::numeric_limits<int32_t>::max());
+
+  // The join relationship should be established.
+  SearchResultProto expected_result;
+  expected_result.mutable_status()->set_code(StatusProto::OK);
+  SearchResultProto::ResultProto* result_proto =
+      expected_result.mutable_results()->Add();
+  *result_proto->mutable_document() = person;
+  *result_proto->mutable_joined_results()->Add()->mutable_document() = email;
+
+  EXPECT_THAT(icing.Search(search_spec, scoring_spec, result_spec),
+              EqualsSearchResultIgnoreStatsAndScores(expected_result));
+}
+
+TEST_F(IcingSearchEngineSearchTest,
+       JoinByQualifiedId_ChildPutBeforeParent_DeleteAndPutParentAgain) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("firstName")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("Email")
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("subject")
+                                        .SetDataTypeString(TERM_MATCH_PREFIX,
+                                                           TOKENIZER_PLAIN)
+                                        .SetCardinality(CARDINALITY_OPTIONAL))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("personQualifiedId")
+                                        .SetDataTypeJoinableString(
+                                            JOINABLE_VALUE_TYPE_QUALIFIED_ID)
+                                        .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  DocumentProto person =
+      DocumentBuilder()
+          .SetKey("pkg$db/namespace", "person")
+          .SetSchema("Person")
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .AddStringProperty("firstName", "first")
+          .Build();
+
+  DocumentProto email =
+      DocumentBuilder()
+          .SetKey("namespace", "email")
+          .SetSchema("Email")
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .AddStringProperty("subject", "test subject")
+          .AddStringProperty("personQualifiedId", "pkg$db/namespace#person")
+          .Build();
+
+  IcingSearchEngineOptions icing_options = GetDefaultIcingOptions();
+  IcingSearchEngine icing(icing_options, GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(email).status(), ProtoIsOk());
+  // Put, delete and put again the parent document.
+  ASSERT_THAT(icing.Put(person).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Delete(person.namespace_(), person.uri()).status(),
+              ProtoIsOk());
+  ASSERT_THAT(icing.Put(person).status(), ProtoIsOk());
+
+  // Parent SearchSpec
+  SearchSpecProto search_spec;
+  search_spec.set_term_match_type(TermMatchType::PREFIX);
+  search_spec.set_query("firstName:first");
+
+  // JoinSpec
+  JoinSpecProto* join_spec = search_spec.mutable_join_spec();
+  join_spec->set_parent_property_expression(
+      std::string(JoinProcessor::kQualifiedIdExpr));
+  join_spec->set_child_property_expression("personQualifiedId");
+  JoinSpecProto::NestedSpecProto* nested_spec =
+      join_spec->mutable_nested_spec();
+  SearchSpecProto* nested_search_spec = nested_spec->mutable_search_spec();
+  nested_search_spec->set_term_match_type(TermMatchType::PREFIX);
+  nested_search_spec->set_query("subject:test");
+
+  *nested_spec->mutable_scoring_spec() = GetDefaultScoringSpec();
+  *nested_spec->mutable_result_spec() = ResultSpecProto::default_instance();
+
+  // Parent ScoringSpec
+  ScoringSpecProto scoring_spec = GetDefaultScoringSpec();
+
+  // Parent ResultSpec
+  ResultSpecProto result_spec;
+  result_spec.set_max_joined_children_per_parent_to_return(
+      std::numeric_limits<int32_t>::max());
+
+  // The join relationship should be established.
+  SearchResultProto expected_result;
+  expected_result.mutable_status()->set_code(StatusProto::OK);
+  SearchResultProto::ResultProto* result_proto =
+      expected_result.mutable_results()->Add();
+  *result_proto->mutable_document() = person;
+  *result_proto->mutable_joined_results()->Add()->mutable_document() = email;
+
+  EXPECT_THAT(icing.Search(search_spec, scoring_spec, result_spec),
+              EqualsSearchResultIgnoreStatsAndScores(expected_result));
+}
+
+TEST_F(
+    IcingSearchEngineSearchTest,
+    JoinByQualifiedId_ChildPutBeforeParent_DeleteAndPutParentAgain_Optimize) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("firstName")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("Email")
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("subject")
+                                        .SetDataTypeString(TERM_MATCH_PREFIX,
+                                                           TOKENIZER_PLAIN)
+                                        .SetCardinality(CARDINALITY_OPTIONAL))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("personQualifiedId")
+                                        .SetDataTypeJoinableString(
+                                            JOINABLE_VALUE_TYPE_QUALIFIED_ID)
+                                        .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  DocumentProto person =
+      DocumentBuilder()
+          .SetKey("pkg$db/namespace", "person")
+          .SetSchema("Person")
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .AddStringProperty("firstName", "first")
+          .Build();
+
+  DocumentProto email =
+      DocumentBuilder()
+          .SetKey("namespace", "email")
+          .SetSchema("Email")
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .AddStringProperty("subject", "test subject")
+          .AddStringProperty("personQualifiedId", "pkg$db/namespace#person")
+          .Build();
+
+  IcingSearchEngineOptions icing_options = GetDefaultIcingOptions();
+  icing_options.set_optimize_rebuild_index_threshold(1.0);
+  IcingSearchEngine icing(icing_options, GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(email).status(), ProtoIsOk());
+  // Put, delete, optimize and put again the parent document.
+  ASSERT_THAT(icing.Put(person).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Delete(person.namespace_(), person.uri()).status(),
+              ProtoIsOk());
+  ASSERT_THAT(icing.Optimize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(person).status(), ProtoIsOk());
+
+  // Parent SearchSpec
+  SearchSpecProto search_spec;
+  search_spec.set_term_match_type(TermMatchType::PREFIX);
+  search_spec.set_query("firstName:first");
+
+  // JoinSpec
+  JoinSpecProto* join_spec = search_spec.mutable_join_spec();
+  join_spec->set_parent_property_expression(
+      std::string(JoinProcessor::kQualifiedIdExpr));
+  join_spec->set_child_property_expression("personQualifiedId");
+  JoinSpecProto::NestedSpecProto* nested_spec =
+      join_spec->mutable_nested_spec();
+  SearchSpecProto* nested_search_spec = nested_spec->mutable_search_spec();
+  nested_search_spec->set_term_match_type(TermMatchType::PREFIX);
+  nested_search_spec->set_query("subject:test");
+
+  *nested_spec->mutable_scoring_spec() = GetDefaultScoringSpec();
+  *nested_spec->mutable_result_spec() = ResultSpecProto::default_instance();
+
+  // Parent ScoringSpec
+  ScoringSpecProto scoring_spec = GetDefaultScoringSpec();
+
+  // Parent ResultSpec
+  ResultSpecProto result_spec;
+  result_spec.set_max_joined_children_per_parent_to_return(
+      std::numeric_limits<int32_t>::max());
+
+  // The join relationship should be established.
+  SearchResultProto expected_result;
+  expected_result.mutable_status()->set_code(StatusProto::OK);
+  SearchResultProto::ResultProto* result_proto =
+      expected_result.mutable_results()->Add();
+  *result_proto->mutable_document() = person;
+  *result_proto->mutable_joined_results()->Add()->mutable_document() = email;
+
+  EXPECT_THAT(icing.Search(search_spec, scoring_spec, result_spec),
+              EqualsSearchResultIgnoreStatsAndScores(expected_result));
+}
+
 TEST_F(IcingSearchEngineSearchTest, JoinByQualifiedIdMultipleNamespaces) {
   SchemaProto schema =
       SchemaBuilder()
@@ -7255,6 +7786,10 @@
   exp_parent_search_stats->set_query_processor_parser_consume_query_latency_ms(
       5);
   exp_parent_search_stats->set_query_processor_query_visitor_latency_ms(5);
+  exp_parent_search_stats->set_num_unquantized_embeddings_scored(0);
+  exp_parent_search_stats->set_num_quantized_embeddings_scored(0);
+  exp_parent_search_stats->set_num_embedding_shards_read(0);
+  exp_parent_search_stats->set_num_embedding_bytes_read(0);
 
   EXPECT_THAT(results.query_stats(), EqualsProto(exp_stats));
 }
@@ -7745,7 +8280,7 @@
 }
 
 class IcingSearchEngineEmbeddingSearchTest
-    : public ::testing::TestWithParam<bool> {
+    : public ::testing::TestWithParam<std::tuple<bool, bool>> {
  protected:
   void SetUp() override {
     if (!IsCfStringTokenization() && !IsReverseJniTokenization()) {
@@ -7774,6 +8309,9 @@
 };
 
 TEST_P(IcingSearchEngineEmbeddingSearchTest, EmbeddingSearch) {
+  bool get_embedding_match_info = std::get<0>(GetParam());
+  bool enable_eigen_embedding_scoring = std::get<1>(GetParam());
+
   SchemaProto schema =
       SchemaBuilder()
           .AddType(SchemaTypeConfigBuilder()
@@ -7822,7 +8360,9 @@
                              CreateVector("my_model_v2", {0.6, 0.7, -0.8}))
           .Build();
 
-  IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  options.set_enable_eigen_embedding_scoring(enable_eigen_embedding_scoring);
+  IcingSearchEngine icing(options, GetTestJniCache());
   ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
   ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
   ASSERT_THAT(icing.Put(document0).status(), ProtoIsOk());
@@ -7850,7 +8390,6 @@
   scoring_spec.set_rank_by(
       ScoringSpecProto::RankingStrategy::ADVANCED_SCORING_EXPRESSION);
 
-  bool get_embedding_match_info = GetParam();
   ResultSpecProto result_spec = ResultSpecProto::default_instance();
   result_spec.mutable_snippet_spec()->set_num_to_snippet(3);
   result_spec.mutable_snippet_spec()->set_num_matches_per_property(5);
@@ -8168,10 +8707,12 @@
   }
 }
 
-INSTANTIATE_TEST_SUITE_P(IcingSearchEngineEmbeddingSearchTest,
-                         IcingSearchEngineEmbeddingSearchTest,
-                         testing::Values(/*enable_embedding_quantization=*/true,
-                                         false));
+INSTANTIATE_TEST_SUITE_P(
+    IcingSearchEngineEmbeddingSearchTest, IcingSearchEngineEmbeddingSearchTest,
+    testing::Values(std::make_tuple(/*get_embedding_match_info=*/false,
+                                    /*enable_eigen_embedding_scoring=*/false),
+                    std::make_tuple(false, true), std::make_tuple(true, false),
+                    std::make_tuple(true, true)));
 
 TEST_F(IcingSearchEngineSearchTest, CannotScoreUnqueriedEmbedding) {
   SchemaProto schema =
@@ -8339,7 +8880,37 @@
               DoubleNear(0 + 0.5, kEps));
 }
 
-TEST_F(IcingSearchEngineSearchTest, EmbeddingSearchWithQuantizedProperty) {
+class IcingSearchEngineEmbeddingSearchQuantizationTest
+    : public ::testing::TestWithParam<bool> {
+ protected:
+  void SetUp() override {
+    if (!IsCfStringTokenization() && !IsReverseJniTokenization()) {
+      // If we've specified using the reverse-JNI method for segmentation (i.e.
+      // not ICU), then we won't have the ICU data file included to set up.
+      // Technically, we could choose to use reverse-JNI for segmentation AND
+      // include an ICU data file, but that seems unlikely and our current BUILD
+      // setup doesn't do this.
+      // File generated via icu_data_file rule in //icing/BUILD.
+      std::string icu_data_file_path =
+          GetTestFilePath("icing/icu.dat");
+      ICING_ASSERT_OK(
+          icu_data_file_helper::SetUpIcuDataFile(icu_data_file_path));
+    }
+    filesystem_.CreateDirectoryRecursively(GetTestBaseDir().c_str());
+  }
+
+  void TearDown() override {
+    filesystem_.DeleteDirectoryRecursively(GetTestBaseDir().c_str());
+  }
+
+  const Filesystem* filesystem() const { return &filesystem_; }
+
+ private:
+  Filesystem filesystem_;
+};
+
+TEST_P(IcingSearchEngineEmbeddingSearchQuantizationTest,
+       EmbeddingSearchWithQuantizedProperty) {
   constexpr float eps = 0.0001f;
 
   SchemaProto schema =
@@ -8385,7 +8956,9 @@
           .AddVectorProperty("embeddingQuantized", vector1, vector2)
           .Build();
 
-  IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  options.set_enable_eigen_embedding_scoring(GetParam());
+  IcingSearchEngine icing(options, GetTestJniCache());
   ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
   ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
   ASSERT_THAT(icing.Put(document_with_original_embedding).status(),
@@ -8457,6 +9030,11 @@
   EXPECT_THAT(results.results(1).additional_scores(0), DoubleNear(256.45, eps));
 }
 
+INSTANTIATE_TEST_SUITE_P(
+    IcingSearchEngineEmbeddingSearchQuantizationTest,
+    IcingSearchEngineEmbeddingSearchQuantizationTest,
+    testing::Values(/*enable_eigen_embedding_scoring=*/true, false));
+
 TEST_F(IcingSearchEngineSearchTest,
        AdditionalScoresOnlyAllowedInAdvancedScoring) {
   SchemaProto schema =
@@ -8492,6 +9070,216 @@
   EXPECT_THAT(results.status(), ProtoStatusIs(StatusProto::INVALID_ARGUMENT));
 }
 
+TEST_F(IcingSearchEngineSearchTest, EmbeddingSearchStats) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("Email1")
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("body")
+                                        .SetDataTypeString(TERM_MATCH_EXACT,
+                                                           TOKENIZER_PLAIN)
+                                        .SetCardinality(CARDINALITY_REPEATED))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("embedding")
+                                        .SetDataTypeVector(
+                                            EMBEDDING_INDEXING_LINEAR_SEARCH)
+                                        .SetCardinality(CARDINALITY_REPEATED))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("embeddingQuantized")
+                                        .SetDataTypeVector(
+                                            EMBEDDING_INDEXING_LINEAR_SEARCH,
+                                            QUANTIZATION_TYPE_QUANTIZE_8_BIT)
+                                        .SetCardinality(CARDINALITY_REPEATED)))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("Email2")
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("body")
+                                        .SetDataTypeString(TERM_MATCH_EXACT,
+                                                           TOKENIZER_PLAIN)
+                                        .SetCardinality(CARDINALITY_REPEATED))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("embedding")
+                                        .SetDataTypeVector(
+                                            EMBEDDING_INDEXING_LINEAR_SEARCH)
+                                        .SetCardinality(CARDINALITY_REPEATED))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("embeddingQuantized")
+                                        .SetDataTypeVector(
+                                            EMBEDDING_INDEXING_LINEAR_SEARCH,
+                                            QUANTIZATION_TYPE_QUANTIZE_8_BIT)
+                                        .SetCardinality(CARDINALITY_REPEATED)))
+          .Build();
+  // Unquantized size: 12 bytes
+  // Quantized size: 3 + sizeof(Quantizer) = 11 bytes.
+  PropertyProto::VectorProto model1_vector =
+      CreateVector("my_model1", {-1, -1, -1});
+  PropertyProto::VectorProto model2_vector =
+      CreateVector("my_model2", {-1, -1, -1});
+  DocumentProto email1_document =
+      DocumentBuilder()
+          .SetKey("icing", "uri0")
+          .SetSchema("Email1")
+          .SetCreationTimestampMs(1)
+          .AddStringProperty("body", "foo")
+          .AddVectorProperty("embedding", model1_vector, model2_vector)
+          .AddVectorProperty("embeddingQuantized", model1_vector, model2_vector)
+          .Build();
+  DocumentProto email2_document =
+      DocumentBuilder()
+          .SetKey("icing", "uri1")
+          .SetSchema("Email2")
+          .SetCreationTimestampMs(1)
+          .AddStringProperty("body", "foo")
+          .AddVectorProperty("embedding", model1_vector, model2_vector)
+          .AddVectorProperty("embeddingQuantized", model1_vector, model2_vector)
+          .Build();
+
+  IcingSearchEngineOptions icing_options = GetDefaultIcingOptions();
+  // Set the number of shards to a large number to avoid hash collisions.
+  icing_options.set_embedding_index_num_shards(8192);
+  IcingSearchEngine icing(icing_options, GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(email1_document).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(email2_document).status(), ProtoIsOk());
+
+  SearchSpecProto search_spec;
+  search_spec.set_term_match_type(TermMatchType::EXACT_ONLY);
+  search_spec.set_embedding_query_metric_type(
+      SearchSpecProto::EmbeddingQueryMetricType::DOT_PRODUCT);
+  search_spec.add_enabled_features(
+      std::string(kListFilterQueryLanguageFeature));
+  *search_spec.add_embedding_query_vectors() = model1_vector;
+  *search_spec.add_embedding_query_vectors() = model2_vector;
+  ScoringSpecProto scoring_spec = GetDefaultScoringSpec();
+
+  // Perform an embedding query that matches all embeddings.
+  search_spec.set_query(
+      "semanticSearch(getEmbeddingParameter(0)) OR "
+      "semanticSearch(getEmbeddingParameter(1))");
+  SearchResultProto results = icing.Search(search_spec, scoring_spec,
+                                           ResultSpecProto::default_instance());
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  EXPECT_THAT(results.results(), SizeIs(2));
+  EXPECT_THAT(results.results(0).document(), EqualsProto(email2_document));
+  EXPECT_THAT(results.results(1).document(), EqualsProto(email1_document));
+  // Check embedding search stats.
+  QueryStatsProto::SearchStats search_stats =
+      results.query_stats().parent_search_stats();
+  // We have 8 embeddings in total in the two documents. 4 of them are
+  // unquantized and 4 of them are quantized.
+  EXPECT_THAT(search_stats.num_unquantized_embeddings_scored(), Eq(4));
+  EXPECT_THAT(search_stats.num_quantized_embeddings_scored(), Eq(4));
+  // We have 2 types and 2 models, which corresponds to 4 shards for unquantized
+  // embeddings and 4 shards for quantized embeddings.
+  EXPECT_THAT(search_stats.num_embedding_shards_read(), Eq(8));
+  // 4 unquantized embeddings * 12 bytes + 4 quantized embeddings * 11 bytes.
+  EXPECT_THAT(search_stats.num_embedding_bytes_read(), Eq(92));
+
+  // Perform an embedding query that only matches model1 embeddings.
+  search_spec.set_query("semanticSearch(getEmbeddingParameter(0))");
+  results = icing.Search(search_spec, scoring_spec,
+                         ResultSpecProto::default_instance());
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  EXPECT_THAT(results.results(), SizeIs(2));
+  EXPECT_THAT(results.results(0).document(), EqualsProto(email2_document));
+  EXPECT_THAT(results.results(1).document(), EqualsProto(email1_document));
+  // Check embedding search stats.
+  search_stats = results.query_stats().parent_search_stats();
+  // We have 4 embeddings in total in the two documents. 2 of them are
+  // unquantized and 2 of them are quantized.
+  EXPECT_THAT(search_stats.num_unquantized_embeddings_scored(), Eq(2));
+  EXPECT_THAT(search_stats.num_quantized_embeddings_scored(), Eq(2));
+  // We have 2 types and 1 models matched, which corresponds to 2 shards for
+  // unquantized embeddings and 2 shards for quantized embeddings.
+  EXPECT_THAT(search_stats.num_embedding_shards_read(), Eq(4));
+  // 2 unquantized embeddings * 12 bytes + 2 quantized embeddings * 11 bytes.
+  EXPECT_THAT(search_stats.num_embedding_bytes_read(), Eq(46));
+
+  // Perform an embedding query that only matches model1 embeddings and filter
+  // on Email1.
+  search_spec.set_query("semanticSearch(getEmbeddingParameter(0))");
+  search_spec.add_schema_type_filters("Email1");
+  results = icing.Search(search_spec, scoring_spec,
+                         ResultSpecProto::default_instance());
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  EXPECT_THAT(results.results(), SizeIs(1));
+  EXPECT_THAT(results.results(0).document(), EqualsProto(email1_document));
+  // Check embedding search stats.
+  search_stats = results.query_stats().parent_search_stats();
+  // We have 2 embeddings in total in the Email1 document. 1 of them are
+  // unquantized and 1 of them are quantized.
+  EXPECT_THAT(search_stats.num_unquantized_embeddings_scored(), Eq(1));
+  EXPECT_THAT(search_stats.num_quantized_embeddings_scored(), Eq(1));
+  // We have 1 types and 1 models matched, which corresponds to 1 shards for
+  // unquantized embeddings and 1 shards for quantized embeddings.
+  EXPECT_THAT(search_stats.num_embedding_shards_read(), Eq(2));
+  // 1 unquantized embeddings * 12 bytes + 1 quantized embeddings * 11 bytes.
+  EXPECT_THAT(search_stats.num_embedding_bytes_read(), Eq(23));
+
+  // Now add a property filter on Email1.embeddingQuantized, which will only
+  // score quantized embeddings in Email1.
+  TypePropertyMask* property_filters = search_spec.add_type_property_filters();
+  property_filters->set_schema_type("Email1");
+  property_filters->add_paths("embeddingQuantized");
+  results = icing.Search(search_spec, scoring_spec,
+                         ResultSpecProto::default_instance());
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  EXPECT_THAT(results.results(), SizeIs(1));
+  EXPECT_THAT(results.results(0).document(), EqualsProto(email1_document));
+  // Check embedding search stats.
+  search_stats = results.query_stats().parent_search_stats();
+  // We have 1 quantized embeddings in the Email1 document.
+  EXPECT_THAT(search_stats.num_unquantized_embeddings_scored(), Eq(0));
+  EXPECT_THAT(search_stats.num_quantized_embeddings_scored(), Eq(1));
+  // We have 1 types and 1 models matched, which corresponds to 0 shard for
+  // unquantized embeddings and 1 shard for quantized embeddings.
+  EXPECT_THAT(search_stats.num_embedding_shards_read(), Eq(1));
+  // 1 quantized embeddings * 11 bytes.
+  EXPECT_THAT(search_stats.num_embedding_bytes_read(), Eq(11));
+}
+
+TEST_F(IcingSearchEngineSearchTest, EmbeddingSearchStats_EmptyIndex) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Email").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("embedding")
+                  .SetDataTypeVector(EMBEDDING_INDEXING_LINEAR_SEARCH)
+                  .SetCardinality(CARDINALITY_REPEATED)))
+          .Build();
+
+  IcingSearchEngineOptions icing_options = GetDefaultIcingOptions();
+  IcingSearchEngine icing(icing_options, GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+
+  SearchSpecProto search_spec;
+  search_spec.set_term_match_type(TermMatchType::EXACT_ONLY);
+  search_spec.set_embedding_query_metric_type(
+      SearchSpecProto::EmbeddingQueryMetricType::DOT_PRODUCT);
+  search_spec.add_enabled_features(
+      std::string(kListFilterQueryLanguageFeature));
+  *search_spec.add_embedding_query_vectors() =
+      CreateVector("my_model", {-1, -1, -1});
+  ScoringSpecProto scoring_spec = GetDefaultScoringSpec();
+
+  // Perform an embedding query.
+  search_spec.set_query("semanticSearch(getEmbeddingParameter(0))");
+  SearchResultProto results = icing.Search(search_spec, scoring_spec,
+                                           ResultSpecProto::default_instance());
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  EXPECT_THAT(results.results(), IsEmpty());
+  // Check embedding search stats.
+  QueryStatsProto::SearchStats search_stats =
+      results.query_stats().parent_search_stats();
+  EXPECT_THAT(search_stats.num_unquantized_embeddings_scored(), Eq(0));
+  EXPECT_THAT(search_stats.num_quantized_embeddings_scored(), Eq(0));
+  EXPECT_THAT(search_stats.num_embedding_shards_read(), Eq(0));
+  EXPECT_THAT(search_stats.num_embedding_bytes_read(), Eq(0));
+}
+
 TEST_F(IcingSearchEngineSearchTest,
        EmbeddingSearchWithManyFilteredOutDocuments) {
   SchemaProto schema =
@@ -8659,6 +9447,213 @@
   EXPECT_THAT(results.results(), IsEmpty());
 }
 
+TEST_F(IcingSearchEngineSearchTest, SearchWithTypeFiltersEmbedding) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("TypeA").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("embedding")
+                  .SetDataTypeVector(EMBEDDING_INDEXING_LINEAR_SEARCH)
+                  .SetCardinality(CARDINALITY_REPEATED)))
+          .AddType(SchemaTypeConfigBuilder().SetType("TypeB").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("embedding")
+                  .SetDataTypeVector(EMBEDDING_INDEXING_LINEAR_SEARCH)
+                  .SetCardinality(CARDINALITY_REPEATED)))
+          .Build();
+
+  DocumentProto doc_a =
+      DocumentBuilder()
+          .SetKey("icing", "uriA")
+          .SetSchema("TypeA")
+          .SetCreationTimestampMs(1)
+          .AddVectorProperty("embedding",
+                             CreateVector("my_model", {0.1, 0.2, 0.3}))
+          .Build();
+  DocumentProto doc_b =
+      DocumentBuilder()
+          .SetKey("icing", "uriB")
+          .SetSchema("TypeB")
+          .SetCreationTimestampMs(1)
+          .AddVectorProperty("embedding",
+                             CreateVector("my_model", {0.4, 0.5, 0.6}))
+          .Build();
+
+  IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(doc_a).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(doc_b).status(), ProtoIsOk());
+
+  SearchSpecProto search_spec;
+  search_spec.set_embedding_query_metric_type(
+      SearchSpecProto::EmbeddingQueryMetricType::DOT_PRODUCT);
+  search_spec.add_enabled_features(
+      std::string(kListFilterQueryLanguageFeature));
+  *search_spec.add_embedding_query_vectors() =
+      CreateVector("my_model", {1, 1, 1});
+  search_spec.set_query("semanticSearch(getEmbeddingParameter(0))");
+  ScoringSpecProto scoring_spec = GetDefaultScoringSpec();
+  scoring_spec.set_rank_by(
+      ScoringSpecProto::RankingStrategy::ADVANCED_SCORING_EXPRESSION);
+  scoring_spec.set_advanced_scoring_expression(
+      "sum(this.matchedSemanticScores(getEmbeddingParameter(0)))");
+
+  // Filter for TypeA only
+  search_spec.add_schema_type_filters("TypeA");
+  SearchResultProto results = icing.Search(search_spec, scoring_spec,
+                                           ResultSpecProto::default_instance());
+  // Verify that only doc_a is returned due to the schema type filter.
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  ASSERT_THAT(results.results(), SizeIs(1));
+  EXPECT_THAT(results.results(0).document(), EqualsProto(doc_a));
+  // Score should be 0.1 + 0.2 + 0.3 = 0.6
+  EXPECT_THAT(results.results(0).score(), DoubleNear(0.6, kEps));
+
+  // Filter for TypeB only
+  search_spec.clear_schema_type_filters();
+  search_spec.add_schema_type_filters("TypeB");
+  results = icing.Search(search_spec, scoring_spec,
+                         ResultSpecProto::default_instance());
+  // Verify that only doc_b is returned due to the schema type filter.
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  ASSERT_THAT(results.results(), SizeIs(1));
+  EXPECT_THAT(results.results(0).document(), EqualsProto(doc_b));
+  // Score should be 0.4 + 0.5 + 0.6 = 1.5
+  EXPECT_THAT(results.results(0).score(), DoubleNear(1.5, kEps));
+}
+
+TEST_F(IcingSearchEngineSearchTest, SearchWithTypeFiltersEmbedding_manyHits) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("TypeA").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("embedding")
+                  .SetDataTypeVector(EMBEDDING_INDEXING_LINEAR_SEARCH)
+                  .SetCardinality(CARDINALITY_REPEATED)))
+          .AddType(SchemaTypeConfigBuilder().SetType("TypeB").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("embedding")
+                  .SetDataTypeVector(EMBEDDING_INDEXING_LINEAR_SEARCH)
+                  .SetCardinality(CARDINALITY_REPEATED)))
+          .AddType(SchemaTypeConfigBuilder().SetType("TypeC").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("embedding")
+                  .SetDataTypeVector(EMBEDDING_INDEXING_LINEAR_SEARCH)
+                  .SetCardinality(CARDINALITY_REPEATED)))
+          .Build();
+  IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+
+  // Add 600 documents of TypeA, each of which has two embedding hits.
+  for (int i = 0; i < 600; ++i) {
+    DocumentProto doc =
+        DocumentBuilder()
+            .SetKey("icing", "uriA" + std::to_string(i))
+            .SetSchema("TypeA")
+            .SetCreationTimestampMs(1)
+            .AddVectorProperty("embedding",
+                               CreateVector("my_model", {0.1, 0.2, 0.3}),
+                               CreateVector("my_model", {0.1, 0.2, 0.3}))
+            .Build();
+    ASSERT_THAT(icing.Put(doc).status(), ProtoIsOk());
+  }
+  // Add 600 documents of TypeB, each of which has two embedding hits.
+  for (int i = 0; i < 600; ++i) {
+    DocumentProto doc =
+        DocumentBuilder()
+            .SetKey("icing", "uriB" + std::to_string(i))
+            .SetSchema("TypeB")
+            .SetCreationTimestampMs(1)
+            .AddVectorProperty("embedding",
+                               CreateVector("my_model", {0.4, 0.5, 0.6}),
+                               CreateVector("my_model", {0.4, 0.5, 0.6}))
+            .Build();
+    ASSERT_THAT(icing.Put(doc).status(), ProtoIsOk());
+  }
+  // Add 600 documents of TypeC, each of which has two embedding hits.
+  for (int i = 0; i < 600; ++i) {
+    DocumentProto doc =
+        DocumentBuilder()
+            .SetKey("icing", "uriC" + std::to_string(i))
+            .SetSchema("TypeC")
+            .SetCreationTimestampMs(1)
+            .AddVectorProperty("embedding",
+                               CreateVector("my_model", {0.7, 0.8, 0.9}),
+                               CreateVector("my_model", {0.7, 0.8, 0.9}))
+            .Build();
+    ASSERT_THAT(icing.Put(doc).status(), ProtoIsOk());
+  }
+
+  SearchSpecProto search_spec;
+  search_spec.set_embedding_query_metric_type(
+      SearchSpecProto::EmbeddingQueryMetricType::DOT_PRODUCT);
+  search_spec.add_enabled_features(
+      std::string(kListFilterQueryLanguageFeature));
+  *search_spec.add_embedding_query_vectors() =
+      CreateVector("my_model", {1, 1, 1});
+  search_spec.set_query("semanticSearch(getEmbeddingParameter(0))");
+  ScoringSpecProto scoring_spec = GetDefaultScoringSpec();
+  scoring_spec.set_rank_by(
+      ScoringSpecProto::RankingStrategy::ADVANCED_SCORING_EXPRESSION);
+  // Let the score be the number of embedding hits.
+  scoring_spec.set_advanced_scoring_expression(
+      "len(this.matchedSemanticScores(getEmbeddingParameter(0)))");
+  ResultSpecProto result_spec;
+  result_spec.set_num_per_page(10000);
+
+  // Filter for TypeA only
+  search_spec.add_schema_type_filters("TypeA");
+  SearchResultProto results =
+      icing.Search(search_spec, scoring_spec, result_spec);
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  ASSERT_THAT(results.results(), SizeIs(600));
+  for (int i = 0; i < 600; ++i) {
+    EXPECT_THAT(results.results(i).document().schema(), Eq("TypeA"));
+    // Each document has 2 embedding hits.
+    EXPECT_THAT(results.results(i).score(), Eq(2));
+  }
+
+  // Filter for TypeB only
+  search_spec.clear_schema_type_filters();
+  search_spec.add_schema_type_filters("TypeB");
+  results = icing.Search(search_spec, scoring_spec, result_spec);
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  ASSERT_THAT(results.results(), SizeIs(600));
+  for (int i = 0; i < 600; ++i) {
+    EXPECT_THAT(results.results(i).document().schema(), Eq("TypeB"));
+    // Each document has 2 embedding hits.
+    EXPECT_THAT(results.results(i).score(), Eq(2));
+  }
+
+  // Filter for TypeC only
+  search_spec.clear_schema_type_filters();
+  search_spec.add_schema_type_filters("TypeC");
+  results = icing.Search(search_spec, scoring_spec, result_spec);
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  ASSERT_THAT(results.results(), SizeIs(600));
+  for (int i = 0; i < 600; ++i) {
+    EXPECT_THAT(results.results(i).document().schema(), Eq("TypeC"));
+    // Each document has 2 embedding hits.
+    EXPECT_THAT(results.results(i).score(), Eq(2));
+  }
+
+  // Filter for TypeA and TypeC
+  search_spec.clear_schema_type_filters();
+  search_spec.add_schema_type_filters("TypeA");
+  search_spec.add_schema_type_filters("TypeC");
+  results = icing.Search(search_spec, scoring_spec, result_spec);
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  ASSERT_THAT(results.results(), SizeIs(1200));
+  for (int i = 0; i < 1200; ++i) {
+    EXPECT_THAT(results.results(i).document().schema(),
+                AnyOf(Eq("TypeA"), Eq("TypeC")));
+    // Each document has 2 embedding hits.
+    EXPECT_THAT(results.results(i).score(), Eq(2));
+  }
+}
+
 TEST_F(IcingSearchEngineSearchTest, SearchWithUriFilters) {
   IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
   ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
@@ -8775,62 +9770,6 @@
               HasSubstr("does not appear in the namespace filter"));
 }
 
-TEST_F(IcingSearchEngineSearchTest, UriFiltersWorkWithPropertyFilters) {
-  IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
-  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
-  ASSERT_THAT(icing.SetSchema(CreateEmailSchema()).status(), ProtoIsOk());
-
-  // Add three documents with different URIs.
-  DocumentProto document_one =
-      DocumentBuilder()
-          .SetKey("namespace", "uri1")
-          .SetSchema("Email")
-          .AddStringProperty("subject", "foo")
-          .AddStringProperty("body", "foo")
-          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
-          .Build();
-  ASSERT_THAT(icing.Put(document_one).status(), ProtoIsOk());
-  DocumentProto document_two =
-      DocumentBuilder()
-          .SetKey("namespace", "uri2")
-          .SetSchema("Email")
-          .AddStringProperty("subject", "foo")
-          .AddStringProperty("body", "foo")
-          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
-          .Build();
-  ASSERT_THAT(icing.Put(document_two).status(), ProtoIsOk());
-  DocumentProto document_three =
-      DocumentBuilder()
-          .SetKey("namespace", "uri3")
-          .SetSchema("Email")
-          .AddStringProperty("subject", "baz")
-          .AddStringProperty("body", "foo")
-          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
-          .Build();
-  ASSERT_THAT(icing.Put(document_three).status(), ProtoIsOk());
-
-  // Create a search spec with uri filters that should only match document1 and
-  // document3, and a property filter that only searches for "subject". As a
-  // result, only document1 should be returned.
-  SearchSpecProto search_spec;
-  search_spec.set_term_match_type(TermMatchType::PREFIX);
-  search_spec.set_query("foo");
-  NamespaceDocumentUriGroup* uris = search_spec.add_document_uri_filters();
-  uris->set_namespace_("namespace");
-  uris->add_document_uris("uri1");
-  uris->add_document_uris("uri3");
-  TypePropertyMask* property_filters = search_spec.add_type_property_filters();
-  property_filters->set_schema_type("Email");
-  property_filters->add_paths("subject");
-
-  // Check results
-  SearchResultProto results = icing.Search(search_spec, GetDefaultScoringSpec(),
-                                           ResultSpecProto::default_instance());
-  EXPECT_THAT(results.status(), ProtoIsOk());
-  EXPECT_THAT(results.results(), SizeIs(1));
-  EXPECT_THAT(results.results(0).document(), EqualsProto(document_one));
-}
-
 TEST_F(IcingSearchEngineSearchTest, SearchWithRankingByScorableProperty) {
   SchemaProto schema =
       SchemaBuilder()
@@ -9692,7 +10631,20 @@
                           .SetScorableType(SCORABLE_TYPE_ENABLED)
                           .SetCardinality(CARDINALITY_REPEATED)))
           .Build();
-  EXPECT_THAT(icing.SetSchema(new_schema).status(), ProtoIsOk());
+  SetSchemaResultProto set_schema_result = icing.SetSchema(new_schema);
+  // Ignore latency numbers as they're covered elsewhere
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
+  SetSchemaResultProto expected_set_schema_result = SetSchemaResultProto();
+  expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
+  expected_set_schema_result
+      .mutable_scorable_property_incompatible_changed_schema_types()
+      ->Add("Person");
+  // Since it is a scorable property incompatible change, we need to do a
+  // recovery proof flush.
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::RECOVERY_PROOF);
+  EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   actual_search_result_proto = icing.Search(
       search_spec, scoring_spec, ResultSpecProto::default_instance());
@@ -9767,7 +10719,20 @@
                           .SetScorableType(SCORABLE_TYPE_DISABLED)
                           .SetCardinality(CARDINALITY_REPEATED)))
           .Build();
-  EXPECT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+  SetSchemaResultProto set_schema_result = icing.SetSchema(schema);
+  // Ignore latency numbers as they're covered elsewhere
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
+  SetSchemaResultProto expected_set_schema_result = SetSchemaResultProto();
+  expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
+  expected_set_schema_result
+      .mutable_scorable_property_incompatible_changed_schema_types()
+      ->Add("Person");
+  // Since it is a scorable property incompatible change, we need to do a
+  // recovery proof flush.
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::RECOVERY_PROOF);
+  EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   SearchResultProto search_result_proto = icing.Search(
       search_spec, scoring_spec, ResultSpecProto::default_instance());
@@ -9859,7 +10824,20 @@
                           .SetScorableType(SCORABLE_TYPE_ENABLED)
                           .SetCardinality(CARDINALITY_REPEATED)))
           .Build();
-  EXPECT_THAT(icing.SetSchema(new_schema).status(), ProtoIsOk());
+  SetSchemaResultProto set_schema_result = icing.SetSchema(new_schema);
+  // Ignore latency numbers as they're covered elsewhere
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
+  SetSchemaResultProto expected_set_schema_result = SetSchemaResultProto();
+  expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
+  expected_set_schema_result
+      .mutable_scorable_property_incompatible_changed_schema_types()
+      ->Add("Person");
+  // Since it is a scorable property incompatible change, we need to do a
+  // recovery proof flush.
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::RECOVERY_PROOF);
+  EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   SearchSpecProto search_spec;
   ScoringSpecProto scoring_spec = GetDefaultScoringSpec();
@@ -9956,7 +10934,14 @@
                           .SetScorableType(SCORABLE_TYPE_ENABLED)
                           .SetCardinality(CARDINALITY_REPEATED)))
           .Build();
-  EXPECT_THAT(icing.SetSchema(new_schema).status(), ProtoIsOk());
+  SetSchemaResultProto set_schema_result = icing.SetSchema(new_schema);
+  EXPECT_THAT(set_schema_result.status(), ProtoIsOk());
+  EXPECT_THAT(
+      set_schema_result.scorable_property_incompatible_changed_schema_types(),
+      IsEmpty());
+  // Flush is not needed since the schema change is (scorable property)
+  // compatible.
+  EXPECT_THAT(set_schema_result.needs_persist_type(), PersistType::UNKNOWN);
 
   SearchSpecProto search_spec;
   ScoringSpecProto scoring_spec = GetDefaultScoringSpec();
@@ -10086,7 +11071,23 @@
                                    "Person", /*index_nested_properties=*/true)
                                .SetCardinality(CARDINALITY_REPEATED)))
           .Build();
-  EXPECT_THAT(icing.SetSchema(schema_proto).status(), ProtoIsOk());
+  SetSchemaResultProto set_schema_result = icing.SetSchema(schema_proto);
+  // Ignore latency numbers as they're covered elsewhere
+  set_schema_result.clear_set_schema_stats();
+  set_schema_result.clear_vm_binder_transaction_latency_start_time_ms();
+  SetSchemaResultProto expected_set_schema_result = SetSchemaResultProto();
+  expected_set_schema_result.mutable_status()->set_code(StatusProto::OK);
+  expected_set_schema_result
+      .mutable_scorable_property_incompatible_changed_schema_types()
+      ->Add("Email");
+  expected_set_schema_result
+      .mutable_scorable_property_incompatible_changed_schema_types()
+      ->Add("Person");
+  // Since it is a scorable property incompatible change, we need to do a
+  // recovery proof flush.
+  expected_set_schema_result.set_needs_persist_type(
+      PersistType::RECOVERY_PROOF);
+  EXPECT_THAT(set_schema_result, EqualsProto(expected_set_schema_result));
 
   SearchSpecProto search_spec;
   ScoringSpecProto scoring_spec = GetDefaultScoringSpec();
@@ -10570,6 +11571,494 @@
   EXPECT_THAT(results.results(1).joined_results(), IsEmpty());
 }
 
+// Section restriction should not be applicable to
+// DocHitInfoIteratorAllDocumentId. Otherwise, no document will be returned,
+// since the iterator does not have any section information.
+TEST_F(IcingSearchEngineSearchTest, PropertyRestrictionWorksWithEmptySearch) {
+  IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(CreatePersonAndEmailSchema()).status(),
+              ProtoIsOk());
+
+  // 1. Add two email documents
+  DocumentProto document_one =
+      DocumentBuilder()
+          .SetKey("namespace", "uri1")
+          .SetCreationTimestampMs(1000)
+          .SetSchema("Email")
+          .AddDocumentProperty(
+              "sender",
+              DocumentBuilder()
+                  .SetKey("namespace", "uri1")
+                  .SetSchema("Person")
+                  .AddStringProperty("name", "Meg Ryan")
+                  .AddStringProperty("emailAddress", "hellogirl@aol.com")
+                  .Build())
+          .AddStringProperty("subject", "Hello World!")
+          .AddStringProperty(
+              "body", "Oh what a beautiful morning! Oh what a beautiful day!")
+          .Build();
+  ASSERT_THAT(icing.Put(document_one).status(), ProtoIsOk());
+
+  DocumentProto document_two =
+      DocumentBuilder()
+          .SetKey("namespace", "uri2")
+          .SetCreationTimestampMs(1000)
+          .SetSchema("Email")
+          .AddDocumentProperty(
+              "sender", DocumentBuilder()
+                            .SetKey("namespace", "uri2")
+                            .SetSchema("Person")
+                            .AddStringProperty("name", "Tom Hanks")
+                            .AddStringProperty("emailAddress", "ny152@aol.com")
+                            .Build())
+          .AddStringProperty("subject", "Goodnight Moon!")
+          .AddStringProperty("body",
+                             "Count all the sheep and tell them 'Hello'.")
+          .Build();
+  ASSERT_THAT(icing.Put(document_two).status(), ProtoIsOk());
+
+  // 2. Issue a query with a property filter
+  auto search_spec = std::make_unique<SearchSpecProto>();
+  search_spec->set_term_match_type(TermMatchType::PREFIX);
+  search_spec->set_query("");
+
+  TypePropertyMask* email_property_filters =
+      search_spec->add_type_property_filters();
+  email_property_filters->set_schema_type("Email");
+  email_property_filters->add_paths("subject");
+
+  auto result_spec = std::make_unique<ResultSpecProto>();
+
+  // 3. Verify that both documents are returned.
+  auto scoring_spec = std::make_unique<ScoringSpecProto>();
+  *scoring_spec = GetDefaultScoringSpec();
+  SearchResultProto results =
+      icing.Search(*search_spec, *scoring_spec, *result_spec);
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  EXPECT_THAT(results.results(), SizeIs(2));
+}
+
+// Section restriction should not be applicable to DocHitInfoIteratorByUri.
+// Otherwise, no document will be returned, since the iterator does not have any
+// section information.
+TEST_F(IcingSearchEngineSearchTest, PropertyRestrictionWorksWithUriFilters) {
+  IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(CreateEmailSchema()).status(), ProtoIsOk());
+
+  // Add three documents with different URIs.
+  DocumentProto document_one =
+      DocumentBuilder()
+          .SetKey("namespace", "uri1")
+          .SetSchema("Email")
+          .AddStringProperty("subject", "foo")
+          .AddStringProperty("body", "foo")
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .Build();
+  ASSERT_THAT(icing.Put(document_one).status(), ProtoIsOk());
+  DocumentProto document_two =
+      DocumentBuilder()
+          .SetKey("namespace", "uri2")
+          .SetSchema("Email")
+          .AddStringProperty("subject", "foo")
+          .AddStringProperty("body", "foo")
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .Build();
+  ASSERT_THAT(icing.Put(document_two).status(), ProtoIsOk());
+  DocumentProto document_three =
+      DocumentBuilder()
+          .SetKey("namespace", "uri3")
+          .SetSchema("Email")
+          .AddStringProperty("subject", "baz")
+          .AddStringProperty("body", "foo")
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .Build();
+  ASSERT_THAT(icing.Put(document_three).status(), ProtoIsOk());
+
+  // Create a search spec with uri filters that should only match document1 and
+  // document3, and a property filter that only searches for "foo" in "subject".
+  // As a result, only document1 should be returned.
+  SearchSpecProto search_spec1;
+  search_spec1.set_term_match_type(TermMatchType::PREFIX);
+  search_spec1.set_query("foo");
+  NamespaceDocumentUriGroup* uris = search_spec1.add_document_uri_filters();
+  uris->set_namespace_("namespace");
+  uris->add_document_uris("uri1");
+  uris->add_document_uris("uri3");
+  TypePropertyMask* property_filters = search_spec1.add_type_property_filters();
+  property_filters->set_schema_type("Email");
+  property_filters->add_paths("subject");
+  SearchResultProto results =
+      icing.Search(search_spec1, GetDefaultScoringSpec(),
+                   ResultSpecProto::default_instance());
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  EXPECT_THAT(results.results(), SizeIs(1));
+  EXPECT_THAT(results.results(0).document(), EqualsProto(document_one));
+
+  // Same as above with property filter in the query expression.
+  SearchSpecProto search_spec2;
+  search_spec2.set_term_match_type(TermMatchType::PREFIX);
+  search_spec2.set_query("subject:foo");
+  uris = search_spec2.add_document_uri_filters();
+  uris->set_namespace_("namespace");
+  uris->add_document_uris("uri1");
+  uris->add_document_uris("uri3");
+  results = icing.Search(search_spec2, GetDefaultScoringSpec(),
+                         ResultSpecProto::default_instance());
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  EXPECT_THAT(results.results(), SizeIs(1));
+  EXPECT_THAT(results.results(0).document(), EqualsProto(document_one));
+
+  // For empty query, property filter is not applicable, since there is no
+  // terms. For this case, only uri filters are applied.
+  SearchSpecProto search_spec3;
+  search_spec3.set_term_match_type(TermMatchType::PREFIX);
+  search_spec3.set_query("");
+  uris = search_spec3.add_document_uri_filters();
+  uris->set_namespace_("namespace");
+  uris->add_document_uris("uri1");
+  uris->add_document_uris("uri3");
+  property_filters = search_spec3.add_type_property_filters();
+  property_filters->set_schema_type("Email");
+  // This property filter should be ignored.
+  property_filters->add_paths("non_existent_property");
+  results = icing.Search(search_spec3, GetDefaultScoringSpec(),
+                         ResultSpecProto::default_instance());
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  EXPECT_THAT(results.results(), SizeIs(2));
+  EXPECT_THAT(results.results(0).document(), EqualsProto(document_three));
+  EXPECT_THAT(results.results(1).document(), EqualsProto(document_one));
+}
+
+// Section restriction should not be applicable to
+// DocHitInfoIteratorMatchScoreExpression. Otherwise, no document will be
+// returned, since the iterator does not have any section information.
+TEST_F(IcingSearchEngineSearchTest,
+       PropertyRestrictionWorksWithMatchScoreExpression) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("Message")
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("title")
+                                        .SetDataTypeString(TERM_MATCH_EXACT,
+                                                           TOKENIZER_PLAIN)
+                                        .SetCardinality(CARDINALITY_REPEATED))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("body")
+                                        .SetDataTypeString(TERM_MATCH_EXACT,
+                                                           TOKENIZER_PLAIN)
+                                        .SetCardinality(CARDINALITY_REPEATED)))
+          .Build();
+
+  IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+
+  // Add three documents with different document scores.
+  DocumentProto document1 =
+      DocumentBuilder()
+          .SetKey("namespace", "uri1")
+          .SetSchema("Message")
+          .AddStringProperty("body", "foo")
+          .SetScore(2)
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .Build();
+  DocumentProto document2 =
+      DocumentBuilder()
+          .SetKey("namespace", "uri2")
+          .SetSchema("Message")
+          .AddStringProperty("title", "foo")
+          .SetScore(3)
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .Build();
+  DocumentProto document3 =
+      DocumentBuilder()
+          .SetKey("namespace", "uri3")
+          .SetSchema("Message")
+          .AddStringProperty("body", "foo")
+          .SetScore(4)
+          .SetCreationTimestampMs(kDefaultCreationTimestampMs)
+          .Build();
+
+  ASSERT_THAT(icing.Put(document1).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(document2).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(document3).status(), ProtoIsOk());
+
+  // Get documents with a document score in [3, 4], which matches document 2
+  // and 3.
+  SearchSpecProto search_spec;
+  search_spec.set_term_match_type(TermMatchType::EXACT_ONLY);
+  search_spec.add_enabled_features(
+      std::string(kMatchScoreExpressionFunctionFeature));
+  search_spec.add_enabled_features(
+      std::string(kListFilterQueryLanguageFeature));
+  search_spec.set_query("matchScoreExpression(\"this.documentScore()\", 3, 4)");
+  SearchResultProto results = icing.Search(search_spec, GetDefaultScoringSpec(),
+                                           ResultSpecProto::default_instance());
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  EXPECT_THAT(results.results(), SizeIs(2));
+  EXPECT_THAT(results.results(0).document(), EqualsProto(document3));
+  EXPECT_THAT(results.results(1).document(), EqualsProto(document2));
+
+  // Check that the property restriction does not apply to matchScoreExpression.
+  search_spec.set_query(
+      "body:matchScoreExpression(\"this.documentScore()\", 3, 4)");
+  results = icing.Search(search_spec, GetDefaultScoringSpec(),
+                         ResultSpecProto::default_instance());
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  EXPECT_THAT(results.results(), SizeIs(2));
+  EXPECT_THAT(results.results(0).document(), EqualsProto(document3));
+  EXPECT_THAT(results.results(1).document(), EqualsProto(document2));
+
+  // Same as above.
+  search_spec.set_query(
+      "non_existent_property:"
+      "matchScoreExpression(\"this.documentScore()\", 3, 4)");
+  results = icing.Search(search_spec, GetDefaultScoringSpec(),
+                         ResultSpecProto::default_instance());
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  EXPECT_THAT(results.results(), SizeIs(2));
+  EXPECT_THAT(results.results(0).document(), EqualsProto(document3));
+  EXPECT_THAT(results.results(1).document(), EqualsProto(document2));
+
+  // Property restriction works as normal for "foo" in "body", so document 2
+  // will not be returned.
+  search_spec.set_query(
+      "body:(foo AND matchScoreExpression(\"this.documentScore()\", 3, 4))");
+  results = icing.Search(search_spec, GetDefaultScoringSpec(),
+                         ResultSpecProto::default_instance());
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  EXPECT_THAT(results.results(), SizeIs(1));
+  EXPECT_THAT(results.results(0).document(), EqualsProto(document3));
+
+  // Property restriction works as normal for "foo" in "non_existent_property",
+  // which should return no results.
+  search_spec.set_query(
+      "non_existent_property:(foo AND "
+      "matchScoreExpression(\"this.documentScore()\", 3, 4))");
+  results = icing.Search(search_spec, GetDefaultScoringSpec(),
+                         ResultSpecProto::default_instance());
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  EXPECT_THAT(results.results(), IsEmpty());
+}
+
+// Section restriction should not be applicable to
+// DocHitInfoIteratorPropertyInDocument. Otherwise, no document will be
+// returned, since the iterator does not have any section information.
+TEST_F(IcingSearchEngineSearchTest, PropertyRestrictionWorksWithHasProperty) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("Value")
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("title")
+                                        .SetDataTypeString(TERM_MATCH_EXACT,
+                                                           TOKENIZER_PLAIN)
+                                        .SetCardinality(CARDINALITY_REPEATED))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("body")
+                                        .SetDataTypeString(TERM_MATCH_EXACT,
+                                                           TOKENIZER_PLAIN)
+                                        .SetCardinality(CARDINALITY_REPEATED)))
+          .Build();
+
+  // A document with every property.
+  DocumentProto document0 = DocumentBuilder()
+                                .SetKey("icing", "uri0")
+                                .SetSchema("Value")
+                                .SetCreationTimestampMs(1)
+                                .AddStringProperty("title", "foo")
+                                .AddStringProperty("body", "foo")
+                                .Build();
+  // A document with missing body.
+  DocumentProto document1 = DocumentBuilder()
+                                .SetKey("icing", "uri1")
+                                .SetSchema("Value")
+                                .SetCreationTimestampMs(1)
+                                .AddStringProperty("title", "bar")
+                                .Build();
+  // A document with missing title.
+  DocumentProto document2 = DocumentBuilder()
+                                .SetKey("icing", "uri2")
+                                .SetSchema("Value")
+                                .SetCreationTimestampMs(1)
+                                .AddStringProperty("body", "bar")
+                                .Build();
+
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  options.set_build_property_existence_metadata_hits(true);
+  IcingSearchEngine icing(options, GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(document0).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(document1).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(document2).status(), ProtoIsOk());
+
+  // Get all documents that have "body".
+  SearchSpecProto search_spec;
+  search_spec.set_term_match_type(TermMatchType::EXACT_ONLY);
+
+  search_spec.add_enabled_features(std::string(kHasPropertyFunctionFeature));
+  search_spec.add_enabled_features(
+      std::string(kListFilterQueryLanguageFeature));
+  search_spec.set_query("hasProperty(\"body\")");
+  SearchResultProto results = icing.Search(search_spec, GetDefaultScoringSpec(),
+                                           ResultSpecProto::default_instance());
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  EXPECT_THAT(results.results(), SizeIs(2));
+  EXPECT_THAT(results.results(0).document(), EqualsProto(document2));
+  EXPECT_THAT(results.results(1).document(), EqualsProto(document0));
+
+  // Check that the property restriction does not apply to hasProperty.
+  search_spec.set_query("body:hasProperty(\"body\")");
+  results = icing.Search(search_spec, GetDefaultScoringSpec(),
+                         ResultSpecProto::default_instance());
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  EXPECT_THAT(results.results(), SizeIs(2));
+  EXPECT_THAT(results.results(0).document(), EqualsProto(document2));
+  EXPECT_THAT(results.results(1).document(), EqualsProto(document0));
+
+  // Same as above.
+  search_spec.set_query("title:hasProperty(\"body\")");
+  results = icing.Search(search_spec, GetDefaultScoringSpec(),
+                         ResultSpecProto::default_instance());
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  EXPECT_THAT(results.results(), SizeIs(2));
+  EXPECT_THAT(results.results(0).document(), EqualsProto(document2));
+  EXPECT_THAT(results.results(1).document(), EqualsProto(document0));
+
+  // Same as above.
+  search_spec.set_query("non_existent_property:hasProperty(\"body\")");
+  results = icing.Search(search_spec, GetDefaultScoringSpec(),
+                         ResultSpecProto::default_instance());
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  EXPECT_THAT(results.results(), SizeIs(2));
+  EXPECT_THAT(results.results(0).document(), EqualsProto(document2));
+  EXPECT_THAT(results.results(1).document(), EqualsProto(document0));
+
+  // Property restriction works as normal for "foo" in "body".
+  search_spec.set_query("body:(foo AND hasProperty(\"body\"))");
+  results = icing.Search(search_spec, GetDefaultScoringSpec(),
+                         ResultSpecProto::default_instance());
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  EXPECT_THAT(results.results(), SizeIs(1));
+  EXPECT_THAT(results.results(0).document(), EqualsProto(document0));
+
+  // Property restriction works as normal for "foo" in "non_existent_property",
+  // which should return no results.
+  search_spec.set_query(
+      "non_existent_property:(foo AND hasProperty(\"body\"))");
+  results = icing.Search(search_spec, GetDefaultScoringSpec(),
+                         ResultSpecProto::default_instance());
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  EXPECT_THAT(results.results(), IsEmpty());
+}
+
+// Section restriction should not be applicable to
+// DocHitInfoIteratorPropertyInSchema. Otherwise, no document will be returned,
+// since the iterator does not have any section information.
+TEST_F(IcingSearchEngineSearchTest,
+       PropertyRestrictionWorksWithPropertyDefined) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("Value")
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("title")
+                                        .SetDataTypeString(TERM_MATCH_EXACT,
+                                                           TOKENIZER_PLAIN)
+                                        .SetCardinality(CARDINALITY_REPEATED))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("body")
+                                        .SetDataTypeString(TERM_MATCH_EXACT,
+                                                           TOKENIZER_PLAIN)
+                                        .SetCardinality(CARDINALITY_REPEATED)))
+          .Build();
+
+  // Create documents.
+  DocumentProto document0 = DocumentBuilder()
+                                .SetKey("icing", "uri0")
+                                .SetSchema("Value")
+                                .SetCreationTimestampMs(1)
+                                .AddStringProperty("title", "foo")
+                                .AddStringProperty("body", "foo")
+                                .Build();
+  DocumentProto document1 = DocumentBuilder()
+                                .SetKey("icing", "uri1")
+                                .SetSchema("Value")
+                                .SetCreationTimestampMs(1)
+                                .AddStringProperty("title", "bar")
+                                .AddStringProperty("body", "bar")
+                                .Build();
+
+  IcingSearchEngineOptions options = GetDefaultIcingOptions();
+  options.set_build_property_existence_metadata_hits(true);
+  IcingSearchEngine icing(options, GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(document0).status(), ProtoIsOk());
+  ASSERT_THAT(icing.Put(document1).status(), ProtoIsOk());
+
+  // Check that propertyDefined works as normal.
+  SearchSpecProto search_spec;
+  search_spec.set_term_match_type(TermMatchType::EXACT_ONLY);
+  search_spec.add_enabled_features(
+      std::string(kListFilterQueryLanguageFeature));
+  search_spec.set_query("propertyDefined(\"body\")");
+  SearchResultProto results = icing.Search(search_spec, GetDefaultScoringSpec(),
+                                           ResultSpecProto::default_instance());
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  EXPECT_THAT(results.results(), SizeIs(2));
+  EXPECT_THAT(results.results(0).document(), EqualsProto(document1));
+  EXPECT_THAT(results.results(1).document(), EqualsProto(document0));
+
+  // Check that the property restriction does not apply to propertyDefined.
+  search_spec.set_query("body:propertyDefined(\"body\")");
+  results = icing.Search(search_spec, GetDefaultScoringSpec(),
+                         ResultSpecProto::default_instance());
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  EXPECT_THAT(results.results(), SizeIs(2));
+  EXPECT_THAT(results.results(0).document(), EqualsProto(document1));
+  EXPECT_THAT(results.results(1).document(), EqualsProto(document0));
+
+  // Same as above.
+  search_spec.set_query("title:propertyDefined(\"body\")");
+  results = icing.Search(search_spec, GetDefaultScoringSpec(),
+                         ResultSpecProto::default_instance());
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  EXPECT_THAT(results.results(), SizeIs(2));
+  EXPECT_THAT(results.results(0).document(), EqualsProto(document1));
+  EXPECT_THAT(results.results(1).document(), EqualsProto(document0));
+
+  // Same as above.
+  search_spec.set_query("non_existent_property:propertyDefined(\"body\")");
+  results = icing.Search(search_spec, GetDefaultScoringSpec(),
+                         ResultSpecProto::default_instance());
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  EXPECT_THAT(results.results(), SizeIs(2));
+  EXPECT_THAT(results.results(0).document(), EqualsProto(document1));
+  EXPECT_THAT(results.results(1).document(), EqualsProto(document0));
+
+  // Property restriction works as normal for "foo" in "body".
+  search_spec.set_query("body:(foo AND propertyDefined(\"body\"))");
+  results = icing.Search(search_spec, GetDefaultScoringSpec(),
+                         ResultSpecProto::default_instance());
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  EXPECT_THAT(results.results(), SizeIs(1));
+  EXPECT_THAT(results.results(0).document(), EqualsProto(document0));
+
+  // Property restriction works as normal for "foo" in "non_existent_property",
+  // which should return no results.
+  search_spec.set_query(
+      "non_existent_property:(foo AND propertyDefined(\"body\"))");
+  results = icing.Search(search_spec, GetDefaultScoringSpec(),
+                         ResultSpecProto::default_instance());
+  EXPECT_THAT(results.status(), ProtoIsOk());
+  EXPECT_THAT(results.results(), IsEmpty());
+}
+
 }  // namespace
 }  // namespace lib
 }  // namespace icing

diff --git a/icing/icing-search-engine_suggest_test.cc b/icing/icing-search-engine_suggest_test.cc
index 1f814aa..0fa586f 100644
--- a/icing/icing-search-engine_suggest_test.cc
+++ b/icing/icing-search-engine_suggest_test.cc

@@ -144,7 +144,14 @@
                                .SetName("sender")
                                .SetDataTypeDocument(
                                    "Person", /*index_nested_properties=*/true)
-                               .SetCardinality(CARDINALITY_OPTIONAL)))
+                               .SetCardinality(CARDINALITY_OPTIONAL))
+              .AddProperty(
+                  PropertyConfigBuilder()
+                      .SetName("embed")
+                      .SetDataTypeVector(
+                          EmbeddingIndexingConfig_EmbeddingIndexingType::
+                              LINEAR_SEARCH)
+                      .SetCardinality(CARDINALITY_OPTIONAL)))
       .Build();
 }
 
@@ -464,7 +471,7 @@
   SuggestionResponse response = icing.SearchSuggestions(suggestion_spec);
   ASSERT_THAT(response.status(), ProtoIsOk());
   ASSERT_THAT(response.suggestions(),
-              UnorderedElementsAre(EqualsProto(suggestionFool)));
+              ElementsAre(EqualsProto(suggestionFool)));
 
   // namespace2 has this suggestion
   suggestion_spec.clear_namespace_filters();
@@ -472,7 +479,7 @@
   response = icing.SearchSuggestions(suggestion_spec);
   ASSERT_THAT(response.status(), ProtoIsOk());
   ASSERT_THAT(response.suggestions(),
-              UnorderedElementsAre(EqualsProto(suggestionFool)));
+              ElementsAre(EqualsProto(suggestionFool)));
 
   // delete document from namespace 1
   EXPECT_THAT(icing.Delete("namespace1", "uri1").status(), ProtoIsOk());
@@ -490,7 +497,7 @@
   response = icing.SearchSuggestions(suggestion_spec);
   ASSERT_THAT(response.status(), ProtoIsOk());
   ASSERT_THAT(response.suggestions(),
-              UnorderedElementsAre(EqualsProto(suggestionFool)));
+              ElementsAre(EqualsProto(suggestionFool)));
 }
 
 TEST_F(IcingSearchEngineSuggestTest,
@@ -536,7 +543,7 @@
   SuggestionResponse response = icing.SearchSuggestions(suggestion_spec);
   ASSERT_THAT(response.status(), ProtoIsOk());
   ASSERT_THAT(response.suggestions(),
-              UnorderedElementsAre(EqualsProto(suggestionFool)));
+              ElementsAre(EqualsProto(suggestionFool)));
 
   // Only search in namespace1,uri2
   suggestion_spec.clear_document_uri_filters();
@@ -548,7 +555,7 @@
   response = icing.SearchSuggestions(suggestion_spec);
   ASSERT_THAT(response.status(), ProtoIsOk());
   ASSERT_THAT(response.suggestions(),
-              UnorderedElementsAre(EqualsProto(suggestionFoo)));
+              ElementsAre(EqualsProto(suggestionFoo)));
 }
 
 TEST_F(IcingSearchEngineSuggestTest,
@@ -928,7 +935,7 @@
   SuggestionResponse response = icing.SearchSuggestions(suggestion_spec);
   ASSERT_THAT(response.status(), ProtoIsOk());
   ASSERT_THAT(response.suggestions(),
-              UnorderedElementsAre(EqualsProto(suggestionFool)));
+              ElementsAre(EqualsProto(suggestionFool)));
 
   // Search in subject and sender.name
   suggestion_spec.clear_type_property_filters();
@@ -1214,7 +1221,7 @@
     SuggestionResponse response = icing.SearchSuggestions(suggestion_spec);
     ASSERT_THAT(response.status(), ProtoIsOk());
     ASSERT_THAT(response.suggestions(),
-                UnorderedElementsAre(EqualsProto(suggestionFool)));
+                ElementsAre(EqualsProto(suggestionFool)));
 
     // namespace2 has this suggestion
     suggestion_spec.clear_namespace_filters();
@@ -1222,7 +1229,7 @@
     response = icing.SearchSuggestions(suggestion_spec);
     ASSERT_THAT(response.status(), ProtoIsOk());
     ASSERT_THAT(response.suggestions(),
-                UnorderedElementsAre(EqualsProto(suggestionFool)));
+                ElementsAre(EqualsProto(suggestionFool)));
   }
   // We reinitialize here so we can feed in a fake clock this time
   {
@@ -1261,7 +1268,7 @@
     response = icing.SearchSuggestions(suggestion_spec);
     ASSERT_THAT(response.status(), ProtoIsOk());
     ASSERT_THAT(response.suggestions(),
-                UnorderedElementsAre(EqualsProto(suggestionFool)));
+                ElementsAre(EqualsProto(suggestionFool)));
   }
 }
 
@@ -1334,7 +1341,7 @@
   SuggestionResponse response = icing.SearchSuggestions(suggestion_spec);
   ASSERT_THAT(response.status(), ProtoIsOk());
   ASSERT_THAT(response.suggestions(),
-              UnorderedElementsAre(EqualsProto(suggestionBarFo)));
+              ElementsAre(EqualsProto(suggestionBarFo)));
 }
 
 TEST_F(IcingSearchEngineSuggestTest, SearchSuggestionsTest_MultipleTerms_Or) {
@@ -1423,7 +1430,7 @@
   SuggestionResponse response = icing.SearchSuggestions(suggestion_spec);
   ASSERT_THAT(response.status(), ProtoIsOk());
   ASSERT_THAT(response.suggestions(),
-              UnorderedElementsAre(EqualsProto(suggestionSubjectFool)));
+              ElementsAre(EqualsProto(suggestionSubjectFool)));
 
   // Add property restriction, only search for nested sender.name
   suggestion_spec.set_prefix("sender.name:f");
@@ -1433,7 +1440,7 @@
   response = icing.SearchSuggestions(suggestion_spec);
   ASSERT_THAT(response.status(), ProtoIsOk());
   ASSERT_THAT(response.suggestions(),
-              UnorderedElementsAre(EqualsProto(suggestionSenderNameFoo)));
+              ElementsAre(EqualsProto(suggestionSenderNameFoo)));
 
   // Add property restriction, only search for nonExist section
   suggestion_spec.set_prefix("none:f");
@@ -1504,7 +1511,7 @@
   response = icing.SearchSuggestions(suggestion_spec);
   ASSERT_THAT(response.status(), ProtoIsOk());
   ASSERT_THAT(response.suggestions(),
-              UnorderedElementsAre(EqualsProto(suggestionBarCatSubjectFoo)));
+              ElementsAre(EqualsProto(suggestionBarCatSubjectFoo)));
 }
 
 TEST_F(IcingSearchEngineSuggestTest, SearchSuggestionsTest_InvalidPrefixTest) {

diff --git a/icing/icing-search-engine_task_scheduler_thread_safety_test.cc b/icing/icing-search-engine_task_scheduler_thread_safety_test.cc
new file mode 100644
index 0000000..21cd66c
--- /dev/null
+++ b/icing/icing-search-engine_task_scheduler_thread_safety_test.cc

@@ -0,0 +1,274 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <chrono>  // NOLINT
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <thread>  // NOLINT
+
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+#include "icing/document-builder.h"
+#include "icing/file/filesystem.h"
+#include "icing/icing-search-engine.h"
+#include "icing/portable/equals-proto.h"
+#include "icing/portable/platform.h"
+#include "icing/proto/debug.pb.h"
+#include "icing/proto/document.pb.h"
+#include "icing/proto/document_wrapper.pb.h"
+#include "icing/proto/initialize.pb.h"
+#include "icing/proto/logging.pb.h"
+#include "icing/proto/optimize.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"
+#include "icing/proto/search.pb.h"
+#include "icing/proto/status.pb.h"
+#include "icing/proto/storage.pb.h"
+#include "icing/proto/term.pb.h"
+#include "icing/proto/usage.pb.h"
+#include "icing/schema-builder.h"
+#include "icing/testing/common-matchers.h"
+#include "icing/testing/jni-test-helpers.h"
+#include "icing/testing/test-data.h"
+#include "icing/testing/tmp-directory.h"
+#include "icing/util/clock.h"
+#include "icing/util/icu-data-file-helper.h"
+
+namespace icing {
+namespace lib {
+
+namespace {
+
+using ::testing::Eq;
+
+static constexpr int kNumIterations = 10;
+
+std::string GetTestBaseDir() { return GetTestTempDir() + "/icing"; }
+
+// Thread safety test for IcingSearchEngine's task scheduler.
+//
+// WARNING: if this test timeouts, then it's likely there is a deadlock between
+//   IcingSearchEngine and the scheduled background task. Please revisit all
+//   logics of task scheduling and task implementation to fix the issue.
+class IcingSearchEngineTaskSchedulerThreadSafetyTest : public testing::Test {
+ protected:
+  void SetUp() override {
+    if (!IsCfStringTokenization() && !IsReverseJniTokenization()) {
+      // If we've specified using the reverse-JNI method for segmentation (i.e.
+      // not ICU), then we won't have the ICU data file included to set up.
+      // Technically, we could choose to use reverse-JNI for segmentation AND
+      // include an ICU data file, but that seems unlikely and our current BUILD
+      // setup doesn't do this.
+      // File generated via icu_data_file rule in //icing/BUILD.
+      std::string icu_data_file_path =
+          GetTestFilePath("icing/icu.dat");
+      ICING_ASSERT_OK(
+          icu_data_file_helper::SetUpIcuDataFile(icu_data_file_path));
+    }
+    filesystem_.CreateDirectoryRecursively(GetTestBaseDir().c_str());
+  }
+
+  void TearDown() override {
+    filesystem_.DeleteDirectoryRecursively(GetTestBaseDir().c_str());
+  }
+
+  Filesystem filesystem_;
+  Clock clock_;
+};
+
+IcingSearchEngineOptions GetDefaultIcingOptions() {
+  IcingSearchEngineOptions icing_options;
+  icing_options.set_base_dir(GetTestBaseDir());
+  icing_options.set_enable_repeated_field_joins(true);
+  icing_options.set_enable_soft_index_restoration(true);
+  icing_options.set_enable_qualified_id_join_index_v3(true);
+  icing_options.set_enable_delete_propagation_from(true);
+  icing_options.set_enable_background_task_scheduler(true);
+  icing_options.set_expired_document_purge_threshold_ms(0);
+  return icing_options;
+}
+
+TEST_F(IcingSearchEngineTaskSchedulerThreadSafetyTest,
+       Destructor_TaskSchedulerShouldNotCauseDeadlockOrCrash) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("name")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  for (int i = 0; i < kNumIterations; ++i) {
+    // Initialize Icing and set schema.
+    auto icing = std::make_unique<IcingSearchEngine>(GetDefaultIcingOptions(),
+                                                     GetTestJniCache());
+    InitializeResultProto initialize_result_proto = icing->Initialize();
+
+    // There should be no recovery.
+    ASSERT_THAT(initialize_result_proto.status(), ProtoIsOk());
+    EXPECT_THAT(
+        initialize_result_proto.initialize_stats().document_store_data_status(),
+        Eq(InitializeStatsProto::NO_DATA_LOSS));
+    EXPECT_THAT(initialize_result_proto.initialize_stats()
+                    .schema_store_recovery_cause(),
+                Eq(InitializeStatsProto::NONE));
+    EXPECT_THAT(initialize_result_proto.initialize_stats()
+                    .document_store_recovery_cause(),
+                Eq(InitializeStatsProto::NONE));
+    EXPECT_THAT(
+        initialize_result_proto.initialize_stats().index_restoration_cause(),
+        Eq(InitializeStatsProto::NONE));
+    EXPECT_THAT(initialize_result_proto.initialize_stats()
+                    .integer_index_restoration_cause(),
+                Eq(InitializeStatsProto::NONE));
+    EXPECT_THAT(initialize_result_proto.initialize_stats()
+                    .qualified_id_join_index_restoration_cause(),
+                Eq(InitializeStatsProto::NONE));
+    EXPECT_THAT(initialize_result_proto.initialize_stats()
+                    .embedding_index_restoration_cause(),
+                Eq(InitializeStatsProto::NONE));
+
+    ASSERT_THAT(icing->SetSchema(schema).status(), ProtoIsOk());
+
+    // Generate 5 person documents that expire in 2 seconds.
+    int64_t current_time_ms = clock_.GetSystemTimeMilliseconds();
+    for (int j = 0; j < 5; ++j) {
+      DocumentProto person =
+          DocumentBuilder()
+              .SetKey("namespace", "person/" + std::to_string(j))
+              .SetSchema("Person")
+              .SetCreationTimestampMs(current_time_ms)
+              .SetTtlMs(2000)
+              .AddStringProperty("name", "person name")
+              .Build();
+      ASSERT_THAT(icing->Put(person).status(), ProtoIsOk());
+    }
+
+    int64_t expected_first_expiration_time_ms = current_time_ms + 2000;
+    // Sleep until the expiration time.
+    std::this_thread::sleep_for(
+        std::chrono::milliseconds(expected_first_expiration_time_ms -
+                                  clock_.GetSystemTimeMilliseconds()));
+
+    // Destroy Icing at the same time when HandleExpiredDocuments task is about
+    // to fire. Destructor should join the task scheduler thread and destruct
+    // the task scheduler.
+    // - There should be no crash or deadlock.
+    // - PersistToDisk should be called after terminating/wating the last task
+    //   if present. The next initialization should not have data recovery.
+    icing.reset();
+  }
+}
+
+TEST_F(IcingSearchEngineTaskSchedulerThreadSafetyTest,
+       ClearAndDestroy_TaskSchedulerShouldNotCauseDeadlockOrCrash) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("name")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  for (int i = 0; i < kNumIterations; ++i) {
+    // Initialize Icing and set schema.
+    auto icing = std::make_unique<IcingSearchEngine>(GetDefaultIcingOptions(),
+                                                     GetTestJniCache());
+    ASSERT_THAT(icing->Initialize().status(), ProtoIsOk());
+    ASSERT_THAT(icing->SetSchema(schema).status(), ProtoIsOk());
+
+    // Generate 5 person documents that expire in 2 seconds.
+    int64_t current_time_ms = clock_.GetSystemTimeMilliseconds();
+    for (int j = 0; j < 5; ++j) {
+      DocumentProto person =
+          DocumentBuilder()
+              .SetKey("namespace", "person/" + std::to_string(j))
+              .SetSchema("Person")
+              .SetCreationTimestampMs(current_time_ms)
+              .SetTtlMs(2000)
+              .AddStringProperty("name", "person name")
+              .Build();
+      ASSERT_THAT(icing->Put(person).status(), ProtoIsOk());
+    }
+
+    int64_t expected_first_expiration_time_ms = current_time_ms + 2000;
+    // Sleep until the expiration time.
+    std::this_thread::sleep_for(
+        std::chrono::milliseconds(expected_first_expiration_time_ms -
+                                  clock_.GetSystemTimeMilliseconds()));
+
+    // Call ClearAndDestroy() at the same time when HandleExpiredDocuments task.
+    // is about to fire.
+    // - It should either cancel or wait for the ongoing task to finish; and
+    //   destruct the task scheduler before erasing the database.
+    // - There should be no crash or deadlock.
+    ResetResultProto clear_and_destroy_result = icing->ClearAndDestroy();
+    EXPECT_THAT(clear_and_destroy_result.status(), ProtoIsOk());
+  }
+}
+
+TEST_F(IcingSearchEngineTaskSchedulerThreadSafetyTest,
+       Reset_TaskSchedulerShouldNotCauseDeadlockOrCrash) {
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("name")
+                  .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  for (int i = 0; i < kNumIterations; ++i) {
+    // Initialize Icing and set schema.
+    auto icing = std::make_unique<IcingSearchEngine>(GetDefaultIcingOptions(),
+                                                     GetTestJniCache());
+    ASSERT_THAT(icing->Initialize().status(), ProtoIsOk());
+    ASSERT_THAT(icing->SetSchema(schema).status(), ProtoIsOk());
+
+    // Generate 5 person documents that expire in 2 seconds.
+    int64_t current_time_ms = clock_.GetSystemTimeMilliseconds();
+    for (int j = 0; j < 5; ++j) {
+      DocumentProto person =
+          DocumentBuilder()
+              .SetKey("namespace", "person/" + std::to_string(j))
+              .SetSchema("Person")
+              .SetCreationTimestampMs(current_time_ms)
+              .SetTtlMs(2000)
+              .AddStringProperty("name", "person name")
+              .Build();
+      ASSERT_THAT(icing->Put(person).status(), ProtoIsOk());
+    }
+
+    int64_t expected_first_expiration_time_ms = current_time_ms + 2000;
+    // Sleep until the expiration time.
+    std::this_thread::sleep_for(
+        std::chrono::milliseconds(expected_first_expiration_time_ms -
+                                  clock_.GetSystemTimeMilliseconds()));
+
+    // Call Reset() at the same time when HandleExpiredDocuments task is about
+    // to fire. There should be no crash or deadlock.
+    ResetResultProto reset_result = icing->Reset();
+    EXPECT_THAT(reset_result.status(), ProtoIsOk());
+  }
+}
+
+}  // namespace
+
+}  // namespace lib
+}  // namespace icing

diff --git a/icing/icing-search-engine_test.cc b/icing/icing-search-engine_test.cc
index 95146ed..5bc6b22 100644
--- a/icing/icing-search-engine_test.cc
+++ b/icing/icing-search-engine_test.cc

@@ -14,20 +14,21 @@
 
 #include "icing/icing-search-engine.h"
 
+#include <cstddef>
 #include <cstdint>
-#include <limits>
+#include <cstring>
 #include <memory>
 #include <string>
 #include <utility>
 
-#include "icing/text_classifier/lib3/utils/base/status.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include "icing/document-builder.h"
 #include "icing/file/filesystem.h"
 #include "icing/file/mock-filesystem.h"
+#include "icing/file/portable-file-backed-proto-log.h"
 #include "icing/jni/jni-cache.h"
-#include "icing/portable/endian.h"
+#include "icing/legacy/index/icing-filesystem.h"
 #include "icing/portable/equals-proto.h"
 #include "icing/portable/platform.h"
 #include "icing/proto/debug.pb.h"
@@ -64,6 +65,7 @@
 using ::testing::Gt;
 using ::testing::HasSubstr;
 using ::testing::IsEmpty;
+using ::testing::Ne;
 using ::testing::Return;
 using ::testing::SizeIs;
 using ::testing::StrEq;
@@ -237,6 +239,396 @@
               EqualsProto(expected_get_result_proto));
 }
 
+TEST_F(IcingSearchEngineTest, BatchGetDocumentResultSizeLimitOver1) {
+  IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(CreateMessageSchema()).status(), ProtoIsOk());
+
+  DocumentProto document1 = CreateMessageDocument("namespace", "uri1");
+  DocumentProto biggerDocument2 =
+      CreateMessageDocument("namespace", "uri2ForBiggerDocument");
+  DocumentProto document3 = CreateMessageDocument("namespace", "uri3");
+  size_t doc1_size = document1.ByteSizeLong();
+  size_t doc3_size = document3.ByteSizeLong();
+
+  //
+  // Expected result
+  //
+  BatchGetResultProto expected_batch_get_result_proto;
+  expected_batch_get_result_proto.mutable_status()->set_code(StatusProto::OK);
+
+  // doc1 should be OK
+  GetResultProto expected_get_result_proto1;
+  expected_get_result_proto1.mutable_status()->set_code(StatusProto::OK);
+  expected_get_result_proto1.set_uri("uri1");
+  *expected_get_result_proto1.mutable_document() = document1;
+  expected_batch_get_result_proto.mutable_get_result_protos()->Add(
+      std::move(expected_get_result_proto1));
+
+  // result for doc2 should be ABORTED
+  GetResultProto expected_get_result_google::protobuf;
+  expected_get_result_google::protobuf.mutable_status()->set_code(
+      StatusProto::ABORTED);
+  expected_get_result_google::protobuf.set_uri("uri2ForBiggerDocument");
+  expected_batch_get_result_proto.mutable_get_result_protos()->Add(
+      std::move(expected_get_result_google::protobuf));
+
+  // doc3 should be ABORTED
+  GetResultProto expected_get_result_proto3;
+  expected_get_result_proto3.set_uri("uri3");
+  expected_get_result_proto3.mutable_status()->set_code(
+      StatusProto::ABORTED);
+  expected_batch_get_result_proto.mutable_get_result_protos()->Add(
+      std::move(expected_get_result_proto3));
+
+  PutDocumentRequest put_document_request;
+  put_document_request.mutable_documents()->Add(std::move(document1));
+  put_document_request.mutable_documents()->Add(std::move(biggerDocument2));
+  put_document_request.mutable_documents()->Add(std::move(document3));
+
+  ASSERT_THAT(icing.BatchPut(std::move(put_document_request)).status(),
+              ProtoIsOk());
+
+  GetResultSpecProto get_result_spec;
+  get_result_spec.set_namespace_requested("namespace");
+  get_result_spec.add_ids("uri1");
+  get_result_spec.add_ids("uri2ForBiggerDocument");
+  get_result_spec.add_ids("uri3");
+  get_result_spec.set_num_total_document_bytes_to_return(doc1_size + doc3_size);
+  ASSERT_THAT(icing.BatchGet(std::move(get_result_spec)),
+              EqualsProto(expected_batch_get_result_proto));
+}
+
+TEST_F(IcingSearchEngineTest, BatchGetDocumentResultSizeLimitOver2) {
+  IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(CreateMessageSchema()).status(), ProtoIsOk());
+
+  DocumentProto document1 = CreateMessageDocument("namespace", "uri1");
+  DocumentProto biggerDocument2 =
+      CreateMessageDocument("namespace", "uri2ForBiggerDocument");
+  DocumentProto document3 = CreateMessageDocument("namespace", "uri3");
+  size_t doc1_size = document1.ByteSizeLong();
+  size_t doc2_size = biggerDocument2.ByteSizeLong();
+  size_t doc3_size = document3.ByteSizeLong();
+
+  //
+  // Expected result
+  //
+  BatchGetResultProto expected_batch_get_result_proto;
+  expected_batch_get_result_proto.mutable_status()->set_code(StatusProto::OK);
+
+  // doc1 should be OK
+  GetResultProto expected_get_result_proto1;
+  expected_get_result_proto1.mutable_status()->set_code(StatusProto::OK);
+  expected_get_result_proto1.set_uri("uri1");
+  *expected_get_result_proto1.mutable_document() = document1;
+  expected_batch_get_result_proto.mutable_get_result_protos()->Add(
+      std::move(expected_get_result_proto1));
+
+  // doc2 should be OK.
+  GetResultProto expected_get_result_google::protobuf;
+  expected_get_result_google::protobuf.mutable_status()->set_code(StatusProto::OK);
+  expected_get_result_google::protobuf.set_uri("uri2ForBiggerDocument");
+  *expected_get_result_google::protobuf.mutable_document() = biggerDocument2;
+  expected_batch_get_result_proto.mutable_get_result_protos()->Add(
+      std::move(expected_get_result_google::protobuf));
+
+  // doc3 should be ABORTED
+  GetResultProto expected_get_result_proto3;
+  expected_get_result_proto3.mutable_status()->set_code(
+      StatusProto::ABORTED);
+  expected_get_result_proto3.set_uri("uri3");
+  expected_batch_get_result_proto.mutable_get_result_protos()->Add(
+      std::move(expected_get_result_proto3));
+
+  PutDocumentRequest put_document_request;
+  put_document_request.mutable_documents()->Add(std::move(document1));
+  put_document_request.mutable_documents()->Add(std::move(biggerDocument2));
+  put_document_request.mutable_documents()->Add(std::move(document3));
+
+  ASSERT_THAT(icing.BatchPut(std::move(put_document_request)).status(),
+              ProtoIsOk());
+
+  GetResultSpecProto get_result_spec;
+  get_result_spec.set_namespace_requested("namespace");
+  get_result_spec.add_ids("uri1");
+  get_result_spec.add_ids("uri2ForBiggerDocument");
+  get_result_spec.add_ids("uri3");
+  get_result_spec.set_num_total_document_bytes_to_return(doc1_size + doc2_size +
+                                                         doc3_size - 1);
+  ASSERT_THAT(icing.BatchGet(std::move(get_result_spec)),
+              EqualsProto(expected_batch_get_result_proto));
+}
+
+TEST_F(IcingSearchEngineTest,
+       BatchGetDocumentResultSizeLimitOriginalErrorKept) {
+  IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(CreateMessageSchema()).status(), ProtoIsOk());
+
+  DocumentProto document1 = CreateMessageDocument("namespace", "uri1");
+  DocumentProto document2 = CreateMessageDocument("namespace", "uri2");
+  DocumentProto biggerDocument3 =
+      CreateMessageDocument("namespace", "uri3ForBiggerDocument");
+  size_t doc1_size = document1.ByteSizeLong();
+
+  //
+  // Expected result
+  //
+  BatchGetResultProto expected_batch_get_result_proto;
+  expected_batch_get_result_proto.mutable_status()->set_code(StatusProto::OK);
+
+  // doc1 should be OK
+  GetResultProto expected_get_result_proto1;
+  expected_get_result_proto1.mutable_status()->set_code(StatusProto::OK);
+  expected_get_result_proto1.set_uri("uri1");
+  *expected_get_result_proto1.mutable_document() = document1;
+  expected_batch_get_result_proto.mutable_get_result_protos()->Add(
+      std::move(expected_get_result_proto1));
+
+  // uriNotExist should not be found.
+  GetResultProto expected_get_result_google::protobuf;
+  expected_get_result_google::protobuf.set_uri("uriNotExist");
+  expected_get_result_google::protobuf.mutable_status()->set_code(StatusProto::NOT_FOUND);
+  expected_get_result_google::protobuf.mutable_status()->set_message(
+      "Document (namespace, uriNotExist) not found.");
+  expected_batch_get_result_proto.mutable_get_result_protos()->Add(
+      std::move(expected_get_result_google::protobuf));
+
+  // doc3 should be ABORTED
+  GetResultProto expected_get_result_proto3;
+  expected_get_result_proto3.mutable_status()->set_code(
+      StatusProto::ABORTED);
+  expected_get_result_proto3.set_uri("uri3ForBiggerDocument");
+  expected_batch_get_result_proto.mutable_get_result_protos()->Add(
+      std::move(expected_get_result_proto3));
+
+  // uriNotExist should be ABORTED as we have reached the limit.
+  // Even though the doc doesn't exist.
+  GetResultProto expected_get_result_proto4;
+  expected_get_result_proto4.mutable_status()->set_code(
+      StatusProto::ABORTED);
+  expected_get_result_proto4.set_uri("uriNotExist");
+  expected_batch_get_result_proto.mutable_get_result_protos()->Add(
+      std::move(expected_get_result_proto4));
+
+  PutDocumentRequest put_document_request;
+  put_document_request.mutable_documents()->Add(std::move(document1));
+  put_document_request.mutable_documents()->Add(std::move(document2));
+  put_document_request.mutable_documents()->Add(std::move(biggerDocument3));
+
+  ASSERT_THAT(icing.BatchPut(std::move(put_document_request)).status(),
+              ProtoIsOk());
+
+  GetResultSpecProto get_result_spec;
+  get_result_spec.set_namespace_requested("namespace");
+  get_result_spec.add_ids("uri1");
+  get_result_spec.add_ids("uriNotExist");
+  get_result_spec.add_ids("uri3ForBiggerDocument");
+  get_result_spec.add_ids("uriNotExist");
+  get_result_spec.set_num_total_document_bytes_to_return(doc1_size);
+  ASSERT_THAT(icing.BatchGet(std::move(get_result_spec)),
+              EqualsProto(expected_batch_get_result_proto));
+}
+
+TEST_F(IcingSearchEngineTest, BatchGetDocumentAlwaysReturnOneDoc) {
+  IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(CreateMessageSchema()).status(), ProtoIsOk());
+
+  DocumentProto document1 = CreateMessageDocument("namespace", "uri1");
+  DocumentProto document2 = CreateMessageDocument("namespace", "uri2");
+  DocumentProto document3 = CreateMessageDocument("namespace", "uri3");
+
+  //
+  // Expected result
+  //
+  BatchGetResultProto expected_batch_get_result_proto;
+  expected_batch_get_result_proto.mutable_status()->set_code(StatusProto::OK);
+
+  // doc1 should be OK
+  GetResultProto expected_get_result_proto1;
+  expected_get_result_proto1.mutable_status()->set_code(StatusProto::OK);
+  expected_get_result_proto1.set_uri("uri1");
+  *expected_get_result_proto1.mutable_document() = document1;
+  expected_batch_get_result_proto.mutable_get_result_protos()->Add(
+      std::move(expected_get_result_proto1));
+
+  // result for doc2 should be ABORTED
+  GetResultProto expected_get_result_google::protobuf;
+  expected_get_result_google::protobuf.mutable_status()->set_code(StatusProto::ABORTED);
+  expected_get_result_google::protobuf.set_uri("uri2");
+  expected_batch_get_result_proto.mutable_get_result_protos()->Add(
+      std::move(expected_get_result_google::protobuf));
+
+  // doc3 should be ABORTED
+  GetResultProto expected_get_result_proto3;
+  expected_get_result_proto3.set_uri("uri3");
+  expected_get_result_proto3.mutable_status()->set_code(StatusProto::ABORTED);
+  expected_batch_get_result_proto.mutable_get_result_protos()->Add(
+      std::move(expected_get_result_proto3));
+
+  PutDocumentRequest put_document_request;
+  put_document_request.mutable_documents()->Add(std::move(document1));
+  put_document_request.mutable_documents()->Add(std::move(document2));
+  put_document_request.mutable_documents()->Add(std::move(document3));
+
+  ASSERT_THAT(icing.BatchPut(std::move(put_document_request)).status(),
+              ProtoIsOk());
+
+  GetResultSpecProto get_result_spec;
+  get_result_spec.set_namespace_requested("namespace");
+  get_result_spec.add_ids("uri1");
+  get_result_spec.add_ids("uri2");
+  get_result_spec.add_ids("uri3");
+  // very small limit. We should always return at least one doc.
+  get_result_spec.set_num_total_document_bytes_to_return(1);
+  ASSERT_THAT(icing.BatchGet(std::move(get_result_spec)),
+              EqualsProto(expected_batch_get_result_proto));
+}
+
+TEST_F(IcingSearchEngineTest, BatchGetDocumentResultSizeLimitNotOver) {
+  IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(CreateMessageSchema()).status(), ProtoIsOk());
+
+  DocumentProto document1 = CreateMessageDocument("namespace", "uri1");
+  DocumentProto biggerDocument2 =
+      CreateMessageDocument("namespace", "uri2ForBiggerDocument");
+  DocumentProto document3 = CreateMessageDocument("namespace", "uri3");
+  size_t doc1_size = document1.ByteSizeLong();
+  size_t doc2_size = biggerDocument2.ByteSizeLong();
+  size_t doc3_size = document3.ByteSizeLong();
+
+  //
+  // Expected result
+  //
+  BatchGetResultProto expected_batch_get_result_proto;
+  expected_batch_get_result_proto.mutable_status()->set_code(StatusProto::OK);
+
+  // doc1 should be OK
+  GetResultProto expected_get_result_proto1;
+  expected_get_result_proto1.mutable_status()->set_code(StatusProto::OK);
+  expected_get_result_proto1.set_uri("uri1");
+  *expected_get_result_proto1.mutable_document() = document1;
+  expected_batch_get_result_proto.mutable_get_result_protos()->Add(
+      std::move(expected_get_result_proto1));
+
+  // doc2 should be OK.
+  GetResultProto expected_get_result_google::protobuf;
+  expected_get_result_google::protobuf.mutable_status()->set_code(StatusProto::OK);
+  expected_get_result_google::protobuf.set_uri("uri2ForBiggerDocument");
+  *expected_get_result_google::protobuf.mutable_document() = biggerDocument2;
+  expected_batch_get_result_proto.mutable_get_result_protos()->Add(
+      std::move(expected_get_result_google::protobuf));
+
+  // doc3 should be OK
+  GetResultProto expected_get_result_proto3;
+  expected_get_result_proto3.mutable_status()->set_code(StatusProto::OK);
+  expected_get_result_proto3.set_uri("uri3");
+  *expected_get_result_proto3.mutable_document() = document3;
+  expected_batch_get_result_proto.mutable_get_result_protos()->Add(
+      std::move(expected_get_result_proto3));
+
+  PutDocumentRequest put_document_request;
+  put_document_request.mutable_documents()->Add(std::move(document1));
+  put_document_request.mutable_documents()->Add(std::move(biggerDocument2));
+  put_document_request.mutable_documents()->Add(std::move(document3));
+
+  ASSERT_THAT(icing.BatchPut(std::move(put_document_request)).status(),
+              ProtoIsOk());
+
+  GetResultSpecProto get_result_spec1;
+  get_result_spec1.set_namespace_requested("namespace");
+  get_result_spec1.add_ids("uri1");
+  get_result_spec1.add_ids("uri2ForBiggerDocument");
+  get_result_spec1.add_ids("uri3");
+  get_result_spec1.set_num_total_document_bytes_to_return(
+      doc1_size + doc2_size + doc3_size);
+  ASSERT_THAT(icing.BatchGet(std::move(get_result_spec1)),
+              EqualsProto(expected_batch_get_result_proto));
+
+  // Return bytes limit with default value(INT_MAX).
+  GetResultSpecProto get_result_spec2;
+  get_result_spec2.set_namespace_requested("namespace");
+  get_result_spec2.add_ids("uri1");
+  get_result_spec2.add_ids("uri2ForBiggerDocument");
+  get_result_spec2.add_ids("uri3");
+
+  ASSERT_THAT(icing.BatchGet(std::move(get_result_spec2)),
+              EqualsProto(expected_batch_get_result_proto));
+}
+
+TEST_F(IcingSearchEngineTest, BatchGetDocumentResultSizeLimitInvalidValue) {
+  IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(CreateMessageSchema()).status(), ProtoIsOk());
+
+  GetResultSpecProto get_result_spec;
+  get_result_spec.set_namespace_requested("namespace");
+
+  BatchGetResultProto expected_batch_get_result_proto;
+  expected_batch_get_result_proto.mutable_status()->set_code(
+      StatusProto::INVALID_ARGUMENT);
+  expected_batch_get_result_proto.mutable_status()->set_message(
+      "num_total_document_bytes_to_return must be greater than 0.");
+
+  get_result_spec.set_num_total_document_bytes_to_return(0);
+  ASSERT_THAT(icing.BatchGet(GetResultSpecProto(get_result_spec)),
+              EqualsProto(expected_batch_get_result_proto));
+
+  get_result_spec.set_num_total_document_bytes_to_return(-1);
+  ASSERT_THAT(icing.BatchGet(std::move(get_result_spec)),
+              EqualsProto(expected_batch_get_result_proto));
+}
+
+TEST_F(IcingSearchEngineTest, BatchGetNotInitialized) {
+  // Declare but don't initialize icing
+  IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
+  GetResultSpecProto get_result_spec;
+  get_result_spec.set_namespace_requested("namespace");
+  get_result_spec.add_ids("uri1");
+  get_result_spec.add_ids("uri2");
+  get_result_spec.set_num_total_document_bytes_to_return(1000);
+
+  BatchGetResultProto result = icing.BatchGet(std::move(get_result_spec));
+  EXPECT_THAT(result.status(), ProtoStatusIs(StatusProto::FAILED_PRECONDITION));
+  EXPECT_THAT(result.get_result_protos(), SizeIs(2));
+  for (const auto& get_result : result.get_result_protos()) {
+    EXPECT_THAT(get_result.status(),
+                ProtoStatusIs(StatusProto::FAILED_PRECONDITION));
+  }
+}
+
+TEST_F(IcingSearchEngineTest, BatchPutNotInitialized) {
+  // Declare but don't initialize icing
+  IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
+  PutDocumentRequest put_document_request;
+  *put_document_request.mutable_documents()->Add() =
+      DocumentBuilder()
+          .SetKey("namespace", "uri1")
+          .SetSchema("Email")
+          .SetCreationTimestampMs(1000)
+          .AddStringProperty("subject", "subject1")
+          .Build();
+  *put_document_request.mutable_documents()->Add() =
+      DocumentBuilder()
+          .SetKey("namespace", "uri2")
+          .SetSchema("Email")
+          .SetCreationTimestampMs(1000)
+          .AddStringProperty("subject", "subject2")
+          .Build();
+
+  BatchPutResultProto result = icing.BatchPut(std::move(put_document_request));
+  EXPECT_THAT(result.status(), ProtoStatusIs(StatusProto::FAILED_PRECONDITION));
+  EXPECT_THAT(result.put_result_protos(), SizeIs(2));
+  for (const auto& put_result : result.put_result_protos()) {
+    EXPECT_THAT(put_result.status(),
+                ProtoStatusIs(StatusProto::FAILED_PRECONDITION));
+  }
+}
+
 TEST_F(IcingSearchEngineTest, GetDocumentWithBadString) {
   IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
   ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
@@ -1062,8 +1454,53 @@
                                   expected_search_result_proto));
 }
 
-TEST_F(IcingSearchEngineTest, NoPersistToDiskLosesAllDocumentsAndIndex) {
-  IcingSearchEngine icing1(GetDefaultIcingOptions(), GetTestJniCache());
+TEST_F(IcingSearchEngineTest, PersistToDiskLogging) {
+  DocumentProto document = CreateMessageDocument("namespace", "uri");
+  auto fake_clock = std::make_unique<FakeClock>();
+  fake_clock->SetTimerElapsedMilliseconds(10);
+  // Add schema and documents to our first icing1 instance.
+
+  IcingSearchEngineOptions icing_options;
+  icing_options.set_base_dir(GetTestBaseDir());
+  icing_options.set_enable_blob_store(true);
+  icing_options.set_enable_marker_file_for_optimize(true);
+
+  TestIcingSearchEngine icing(icing_options, std::make_unique<Filesystem>(),
+                              std::make_unique<IcingFilesystem>(),
+                              std::move(fake_clock), GetTestJniCache());
+  EXPECT_THAT(icing.Initialize().status(), ProtoIsOk());
+  EXPECT_THAT(icing.SetSchema(CreateMessageSchema()).status(), ProtoIsOk());
+  EXPECT_THAT(icing.Put(document).status(), ProtoIsOk());
+
+  PersistToDiskResultProto persist_result =
+      icing.PersistToDisk(PersistType::FULL);
+  EXPECT_THAT(persist_result.status(), ProtoIsOk());
+  PersistToDiskStatsProto persist_stats = persist_result.persist_stats();
+  EXPECT_THAT(persist_stats.latency_ms(), Eq(10));
+  EXPECT_THAT(persist_stats.persist_type(), Eq(PersistType::FULL));
+  EXPECT_THAT(persist_stats.blob_store_persist_latency_ms(), Eq(10));
+  EXPECT_THAT(persist_stats.document_store_total_persist_latency_ms(), Eq(10));
+  EXPECT_THAT(persist_stats.document_store_components_persist_latency_ms(),
+              Eq(10));
+  EXPECT_THAT(persist_stats.document_store_checksum_update_latency_ms(),
+              Eq(10));
+  EXPECT_THAT(persist_stats.document_log_checksum_update_latency_ms(), Eq(10));
+  EXPECT_THAT(persist_stats.document_log_data_sync_latency_ms(), Eq(10));
+  EXPECT_THAT(persist_stats.schema_store_persist_latency_ms(), Eq(10));
+  EXPECT_THAT(persist_stats.index_persist_latency_ms(), Eq(10));
+  EXPECT_THAT(persist_stats.integer_index_persist_latency_ms(), Eq(10));
+  EXPECT_THAT(persist_stats.qualified_id_join_index_persist_latency_ms(),
+              Eq(10));
+  EXPECT_THAT(persist_stats.embedding_index_persist_latency_ms(), Eq(10));
+}
+
+TEST_F(
+    IcingSearchEngineTest,
+    NoPersistToDisk_protoLogNewHeaderFormatFlagOffLosesAllDocumentsAndIndex) {
+  IcingSearchEngineOptions icing_options = GetDefaultIcingOptions();
+  icing_options.set_enable_proto_log_new_header_format(false);
+  IcingSearchEngine icing1(icing_options, GetTestJniCache());
+
   EXPECT_THAT(icing1.Initialize().status(), ProtoIsOk());
   EXPECT_THAT(icing1.SetSchema(CreateMessageSchema()).status(), ProtoIsOk());
   DocumentProto document = CreateMessageDocument("namespace", "uri");
@@ -1111,6 +1548,261 @@
                                   expected_search_result_proto));
 }
 
+TEST_F(
+    IcingSearchEngineTest,
+    NoPersistToDisk_protoLogNewHeaderFormatFlagOnWithImplicitFlushKeepsDocumentsAndRebuildIndex) {
+  IcingSearchEngineOptions icing_options = GetDefaultIcingOptions();
+  icing_options.set_enable_proto_log_new_header_format(true);
+  IcingSearchEngine icing1(icing_options, GetTestJniCache());
+
+  EXPECT_THAT(icing1.Initialize().status(), ProtoIsOk());
+  EXPECT_THAT(icing1.SetSchema(CreateMessageSchema()).status(), ProtoIsOk());
+  DocumentProto document = CreateMessageDocument("namespace", "uri");
+  EXPECT_THAT(icing1.Put(document).status(), ProtoIsOk());
+  EXPECT_THAT(
+      icing1.Get("namespace", "uri", GetResultSpecProto::default_instance())
+          .document(),
+      EqualsProto(document));
+
+  // It's intentional that no PersistToDisk call is made before initializing a
+  // second instance of icing. We won't lose ground truth, but we will get IO
+  // errors (due to mismatch checksums) for derived files and rebuild them.
+
+  IcingSearchEngine icing2(icing_options, GetTestJniCache());
+  InitializeResultProto init_result = icing2.Initialize();
+  EXPECT_THAT(init_result.status(), ProtoIsOk());
+  EXPECT_THAT(init_result.initialize_stats().schema_store_recovery_cause(),
+              Eq(InitializeStatsProto::NONE));
+  EXPECT_THAT(init_result.initialize_stats().document_store_data_status(),
+              Eq(InitializeStatsProto::NO_DATA_LOSS));
+  EXPECT_THAT(init_result.initialize_stats().document_store_recovery_cause(),
+              Eq(InitializeStatsProto::IO_ERROR));
+  EXPECT_THAT(init_result.initialize_stats().index_restoration_cause(),
+              Eq(InitializeStatsProto::IO_ERROR));
+
+  // The document can still be retrieved.
+  EXPECT_THAT(
+      icing2.Get("namespace", "uri", GetResultSpecProto::default_instance())
+          .document(),
+      EqualsProto(document));
+
+  // Searching also should get the document since index is rebuilt.
+  SearchSpecProto search_spec;
+  search_spec.set_term_match_type(TermMatchType::PREFIX);
+  search_spec.set_query("message");  // Content in the Message document.
+
+  SearchResultProto expected_search_result_proto;
+  expected_search_result_proto.mutable_status()->set_code(StatusProto::OK);
+  *expected_search_result_proto.mutable_results()->Add()->mutable_document() =
+      document;
+
+  SearchResultProto actual_results =
+      icing2.Search(search_spec, GetDefaultScoringSpec(),
+                    ResultSpecProto::default_instance());
+  EXPECT_THAT(actual_results, EqualsSearchResultIgnoreStatsAndScores(
+                                  expected_search_result_proto));
+}
+
+TEST_F(
+    IcingSearchEngineTest,
+    NoPersistToDisk_protoLogNewHeaderFormatFlagOnWithoutImplicitFlushLosesNewDocumentsAndIndex) {
+  // Mock the filesystem for document_log to avoid implicit flush when writing,
+  // and only DataSync will do the actual write and flush.
+  std::string mock_file_buffer;
+  off_t document_log_fd_offset = -1;
+  int document_log_fd = -1;
+
+  auto mock_filesystem = std::make_unique<MockFilesystem>();
+  MockFilesystem* mock_filesystem_ptr = mock_filesystem.get();
+  ON_CALL(*mock_filesystem, OpenForAppend(HasSubstr("document_log")))
+      .WillByDefault([mock_filesystem_ptr, &mock_file_buffer, &document_log_fd,
+                      &document_log_fd_offset](const char* file_name) -> int {
+        // Call the real OpenForAppend() get the file descriptor, and capture it
+        // in the test fixture.
+        document_log_fd =
+            mock_filesystem_ptr->Filesystem::OpenForAppend(file_name);
+        int64_t file_size =
+            mock_filesystem_ptr->Filesystem::GetFileSize(document_log_fd);
+        document_log_fd_offset = file_size;
+
+        // Load data from disk to the mock file buffer.
+        mock_file_buffer.resize(file_size);
+        mock_filesystem_ptr->Filesystem::PRead(
+            document_log_fd, mock_file_buffer.data(), file_size,
+            /*offset=*/0);
+        return document_log_fd;
+      });
+  ON_CALL(*mock_filesystem, GetCurrentPosition(_))
+      .WillByDefault([mock_filesystem_ptr, &document_log_fd,
+                      &document_log_fd_offset](int fd) -> int64_t {
+        if (fd != document_log_fd) {
+          return mock_filesystem_ptr->Filesystem::GetCurrentPosition(fd);
+        }
+
+        return document_log_fd_offset;
+      });
+  ON_CALL(*mock_filesystem, Read(A<int>(), _, _))
+      .WillByDefault([mock_filesystem_ptr, &mock_file_buffer, &document_log_fd,
+                      &document_log_fd_offset](int fd, void* buf,
+                                               size_t buf_size) -> ssize_t {
+        if (fd != document_log_fd) {
+          return mock_filesystem_ptr->Filesystem::Read(fd, buf, buf_size);
+        }
+
+        memcpy(buf, mock_file_buffer.data() + document_log_fd_offset, buf_size);
+        document_log_fd_offset += buf_size;
+        return buf_size;
+      });
+  ON_CALL(*mock_filesystem, Write(A<int>(), _, _))
+      .WillByDefault([mock_filesystem_ptr, &mock_file_buffer, &document_log_fd,
+                      &document_log_fd_offset](int fd, const void* data,
+                                               size_t size) -> bool {
+        if (fd != document_log_fd) {
+          return mock_filesystem_ptr->Filesystem::Write(fd, data, size);
+        }
+
+        // Extend the buffer if needed.
+        if (document_log_fd_offset + size > mock_file_buffer.size()) {
+          mock_file_buffer.resize(document_log_fd_offset + size);
+        }
+
+        // Write to the mock file buffer and adjust the fd offset.
+        memcpy(mock_file_buffer.data() + document_log_fd_offset, data, size);
+        document_log_fd_offset += size;
+        return true;
+      });
+  ON_CALL(*mock_filesystem, PRead(A<int>(), _, _, _))
+      .WillByDefault(
+          [mock_filesystem_ptr, &mock_file_buffer, &document_log_fd](
+              int fd, void* buf, size_t buf_size, off_t offset) -> ssize_t {
+            if (fd != document_log_fd) {
+              return mock_filesystem_ptr->Filesystem::PRead(fd, buf, buf_size,
+                                                            offset);
+            }
+
+            memcpy(buf, mock_file_buffer.data() + offset, buf_size);
+            return buf_size;
+          });
+  ON_CALL(*mock_filesystem, PWrite(A<int>(), _, _, _))
+      .WillByDefault([mock_filesystem_ptr, &mock_file_buffer, &document_log_fd](
+                         int fd, off_t offset, const void* data,
+                         size_t data_size) -> bool {
+        if (fd != document_log_fd) {
+          return mock_filesystem_ptr->Filesystem::PWrite(fd, offset, data,
+                                                         data_size);
+        }
+
+        // Extend the buffer if needed.
+        if (offset + data_size > mock_file_buffer.size()) {
+          mock_file_buffer.resize(offset + data_size);
+        }
+        // Write to the mock file buffer.
+        memcpy(mock_file_buffer.data() + offset, data, data_size);
+        return true;
+      });
+  ON_CALL(*mock_filesystem, DataSync(_))
+      .WillByDefault([mock_filesystem_ptr, &mock_file_buffer,
+                      &document_log_fd](int fd) -> bool {
+        if (fd != document_log_fd) {
+          return mock_filesystem_ptr->Filesystem::DataSync(fd);
+        }
+
+        mock_filesystem_ptr->Filesystem::PWrite(
+            fd, /*offset=*/0, mock_file_buffer.data(), mock_file_buffer.size());
+        return mock_filesystem_ptr->Filesystem::DataSync(fd);
+      });
+
+  IcingSearchEngineOptions icing_options = GetDefaultIcingOptions();
+  icing_options.set_enable_proto_log_new_header_format(true);
+  TestIcingSearchEngine icing1(icing_options, std::move(mock_filesystem),
+                               std::make_unique<IcingFilesystem>(),
+                               std::make_unique<FakeClock>(),
+                               GetTestJniCache());
+
+  EXPECT_THAT(icing1.Initialize().status(), ProtoIsOk());
+  // After initialization, document log fd should be captured.
+  ASSERT_THAT(document_log_fd, Ne(-1));
+
+  EXPECT_THAT(icing1.SetSchema(CreateMessageSchema()).status(), ProtoIsOk());
+
+  DocumentProto document1 = CreateMessageDocument("namespace", "uri1");
+  DocumentProto document2 = CreateMessageDocument("namespace", "uri2");
+  DocumentProto document3 = CreateMessageDocument("namespace", "uri3");
+
+  // Put document1 and PersistToDisk to flush it to disk.
+  EXPECT_THAT(icing1.Put(document1).status(), ProtoIsOk());
+  EXPECT_THAT(icing1.PersistToDisk(PersistType::LITE).status(), ProtoIsOk());
+
+  // Put document2 and document3 without flushing.
+  EXPECT_THAT(icing1.Put(document2).status(), ProtoIsOk());
+  EXPECT_THAT(icing1.Put(document3).status(), ProtoIsOk());
+
+  // Sanity check that document1, document2 and document3 are present.
+  EXPECT_THAT(
+      icing1.Get("namespace", "uri1", GetResultSpecProto::default_instance())
+          .document(),
+      EqualsProto(document1));
+  EXPECT_THAT(
+      icing1.Get("namespace", "uri2", GetResultSpecProto::default_instance())
+          .document(),
+      EqualsProto(document2));
+  EXPECT_THAT(
+      icing1.Get("namespace", "uri3", GetResultSpecProto::default_instance())
+          .document(),
+      EqualsProto(document3));
+
+  // When writing a document, we separate it into several steps:
+  // - Write document content to document_log (unsynced tail).
+  // - Compute unsynced tail checksum and header checksum.
+  // - PWrite the header.
+  //
+  // Let's simulate the following scenario by only flushing contents to disk:
+  // - document2 and document3 contents are flushed into disk.
+  // - But before the checksums are computed, there is a power loss and the new
+  //   checksums are never written.
+
+  int64_t content_offset =
+      PortableFileBackedProtoLog<DocumentProto>::kHeaderReservedBytes;
+  ASSERT_TRUE(filesystem()->PWrite(document_log_fd, content_offset,
+                                   mock_file_buffer.data() + content_offset,
+                                   mock_file_buffer.size() - content_offset));
+  ASSERT_TRUE(filesystem()->DataSync(document_log_fd));
+
+  // Create a second instance of icing.
+  // - Document2 and document3 contents are in the file.
+  // - But unsynced tail checksum doesn't match.
+  //
+  // Therefore, document log will rewound to document1 and discard document2 +
+  // document3.
+  IcingSearchEngine icing2(icing_options, GetTestJniCache());
+  InitializeResultProto init_result = icing2.Initialize();
+  EXPECT_THAT(init_result.status(), ProtoIsOk());
+  EXPECT_THAT(init_result.initialize_stats().schema_store_recovery_cause(),
+              Eq(InitializeStatsProto::NONE));
+  EXPECT_THAT(init_result.initialize_stats().document_store_data_status(),
+              Eq(InitializeStatsProto::PARTIAL_LOSS));
+  EXPECT_THAT(init_result.initialize_stats().document_store_recovery_cause(),
+              Eq(InitializeStatsProto::DATA_LOSS));
+  EXPECT_THAT(init_result.initialize_stats().index_restoration_cause(),
+              Eq(InitializeStatsProto::IO_ERROR));
+
+  // document1 can still be retrieved.
+  EXPECT_THAT(
+      icing2.Get("namespace", "uri1", GetResultSpecProto::default_instance())
+          .document(),
+      EqualsProto(document1));
+
+  // document2 and document3 are lost.
+  EXPECT_THAT(
+      icing2.Get("namespace", "uri2", GetResultSpecProto::default_instance())
+          .status(),
+      ProtoStatusIs(StatusProto::NOT_FOUND));
+  EXPECT_THAT(
+      icing2.Get("namespace", "uri3", GetResultSpecProto::default_instance())
+          .status(),
+      ProtoStatusIs(StatusProto::NOT_FOUND));
+}
+
 TEST_F(IcingSearchEngineTest, PersistToDiskLiteSavesGroundTruth) {
   DocumentProto document = CreateMessageDocument("namespace", "uri");
 
@@ -1494,6 +2186,49 @@
   ASSERT_THAT(result.status(), ProtoIsOk());
 }
 
+TEST_F(IcingSearchEngineTest, ClearAndDestroySucceeds) {
+  IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache());
+
+  // Obtain empty size of the base directory.
+  int64_t empty_state_size =
+      filesystem()->GetFileDiskUsage(GetTestBaseDir().c_str());
+
+  ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
+  ASSERT_THAT(icing.SetSchema(CreateMessageSchema()).status(), ProtoIsOk());
+
+  // Simple put and get
+  ASSERT_THAT(icing.Put(CreateMessageDocument("namespace", "uri")).status(),
+              ProtoIsOk());
+
+  GetResultProto expected_get_result_proto;
+  expected_get_result_proto.mutable_status()->set_code(StatusProto::OK);
+  *expected_get_result_proto.mutable_document() =
+      CreateMessageDocument("namespace", "uri");
+  ASSERT_THAT(
+      icing.Get("namespace", "uri", GetResultSpecProto::default_instance()),
+      EqualsProto(expected_get_result_proto));
+
+  // Sanity check that things have been added
+  EXPECT_THAT(filesystem()->GetDiskUsage(GetTestBaseDir().c_str()),
+              Gt(empty_state_size));
+  EXPECT_EQ(filesystem()->DirectoryExists(GetTestBaseDir().c_str()), true);
+
+  // Clear all data for the icing instance
+  ResetResultProto clear_result = icing.ClearAndDestroy();
+  EXPECT_THAT(clear_result.status(), ProtoIsOk());
+
+  // Check that the directory no longer exists after deletion.
+  EXPECT_EQ(filesystem()->DirectoryExists(GetTestBaseDir().c_str()), false);
+
+  // Try to put and get again, but it should fail since the instance is
+  // uninitialized.
+  EXPECT_THAT(icing.Put(CreateMessageDocument("namespace", "uri")).status(),
+              ProtoStatusIs(StatusProto::FAILED_PRECONDITION));
+  EXPECT_THAT(
+      icing.Get("namespace", "uri", GetResultSpecProto::default_instance())
+          .status(),
+      ProtoStatusIs(StatusProto::FAILED_PRECONDITION));
+}
 }  // namespace
 }  // namespace lib
 }  // namespace icing

diff --git a/icing/index/embed/doc-hit-info-iterator-embedding.cc b/icing/index/embed/doc-hit-info-iterator-embedding-v1.cc
similarity index 63%
rename from icing/index/embed/doc-hit-info-iterator-embedding.cc
rename to icing/index/embed/doc-hit-info-iterator-embedding-v1.cc
index 292fb03..b991e4a 100644
--- a/icing/index/embed/doc-hit-info-iterator-embedding.cc
+++ b/icing/index/embed/doc-hit-info-iterator-embedding-v1.cc

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "icing/index/embed/doc-hit-info-iterator-embedding.h"
+#include "icing/index/embed/doc-hit-info-iterator-embedding-v1.h"
 
 #include <cstdint>
 #include <memory>
@@ -26,7 +26,6 @@
 #include "icing/index/embed/embedding-index.h"
 #include "icing/index/embed/embedding-query-results.h"
 #include "icing/index/embed/embedding-scorer.h"
-#include "icing/index/embed/posting-list-embedding-hit-accessor.h"
 #include "icing/index/hit/doc-hit-info.h"
 #include "icing/index/hit/hit.h"
 #include "icing/proto/search.pb.h"
@@ -40,50 +39,55 @@
 namespace icing {
 namespace lib {
 
-libtextclassifier3::StatusOr<std::unique_ptr<DocHitInfoIteratorEmbedding>>
-DocHitInfoIteratorEmbedding::Create(
+libtextclassifier3::StatusOr<std::unique_ptr<DocHitInfoIteratorEmbeddingV1>>
+DocHitInfoIteratorEmbeddingV1::Create(
     const PropertyProto::VectorProto* query,
     SearchSpecProto::EmbeddingQueryMetricType::Code metric_type,
-    double score_low, double score_high, bool get_embedding_match_info,
+    double score_low, double score_high,
     EmbeddingQueryResults::EmbeddingQueryMatchInfoMap* info_map,
+    std::vector<double>* global_scores,
+    std::vector<EmbeddingMatchInfos::EmbeddingMatchSectionInfo>*
+        global_section_infos,
     const EmbeddingIndex* embedding_index, const DocumentStore* document_store,
     const SchemaStore* schema_store, int64_t current_time_ms) {
   ICING_RETURN_ERROR_IF_NULL(query);
   ICING_RETURN_ERROR_IF_NULL(embedding_index);
   ICING_RETURN_ERROR_IF_NULL(info_map);
+  ICING_RETURN_ERROR_IF_NULL(global_scores);
   ICING_RETURN_ERROR_IF_NULL(document_store);
   ICING_RETURN_ERROR_IF_NULL(schema_store);
 
-  libtextclassifier3::StatusOr<std::unique_ptr<PostingListEmbeddingHitAccessor>>
-      pl_accessor_or = embedding_index->GetAccessorForVector(*query);
-  std::unique_ptr<PostingListEmbeddingHitAccessor> pl_accessor;
-  if (pl_accessor_or.ok()) {
-    pl_accessor = std::move(pl_accessor_or).ValueOrDie();
-  } else if (absl_ports::IsNotFound(pl_accessor_or.status())) {
+  libtextclassifier3::StatusOr<
+      std::unique_ptr<EmbeddingIndex::EmbeddingHitAccessor>>
+      embedding_hit_accessor_or = embedding_index->GetAccessorForVector(*query);
+  std::unique_ptr<EmbeddingIndex::EmbeddingHitAccessor> embedding_hit_accessor;
+  if (embedding_hit_accessor_or.ok()) {
+    embedding_hit_accessor = std::move(embedding_hit_accessor_or).ValueOrDie();
+  } else if (absl_ports::IsNotFound(embedding_hit_accessor_or.status())) {
     // A not-found error should be fine, since that means there is no matching
     // embedding hits in the index.
-    pl_accessor = nullptr;
+    embedding_hit_accessor = nullptr;
   } else {
     // Otherwise, return the error as is.
-    return pl_accessor_or.status();
+    return embedding_hit_accessor_or.status();
   }
 
   ICING_ASSIGN_OR_RETURN(std::unique_ptr<EmbeddingScorer> embedding_scorer,
                          EmbeddingScorer::Create(metric_type));
 
-  return std::unique_ptr<DocHitInfoIteratorEmbedding>(
-      new DocHitInfoIteratorEmbedding(
+  return std::unique_ptr<DocHitInfoIteratorEmbeddingV1>(
+      new DocHitInfoIteratorEmbeddingV1(
           query, metric_type, std::move(embedding_scorer), score_low,
-          score_high, get_embedding_match_info, info_map, embedding_index,
-          std::move(pl_accessor), document_store, schema_store,
-          current_time_ms));
+          score_high, info_map, global_scores, global_section_infos,
+          embedding_index, std::move(embedding_hit_accessor), document_store,
+          schema_store, current_time_ms));
 }
 
 libtextclassifier3::StatusOr<const EmbeddingHit*>
-DocHitInfoIteratorEmbedding::AdvanceToNextEmbeddingHit() {
+DocHitInfoIteratorEmbeddingV1::AdvanceToNextEmbeddingHit() {
   if (cached_embedding_hits_idx_ == cached_embedding_hits_.size()) {
     ICING_ASSIGN_OR_RETURN(cached_embedding_hits_,
-                           posting_list_accessor_->GetNextHitsBatch());
+                           embedding_hit_accessor_->GetNextHitsBatch());
     cached_embedding_hits_idx_ = 0;
     if (cached_embedding_hits_.empty()) {
       no_more_hit_ = true;
@@ -94,14 +98,25 @@
       cached_embedding_hits_[cached_embedding_hits_idx_];
   if (doc_hit_info_.document_id() == kInvalidDocumentId) {
     doc_hit_info_.set_document_id(embedding_hit.basic_hit().document_id());
-    current_allowed_sections_mask_ =
-        ComputeAllowedSectionsMask(doc_hit_info_.document_id());
+    if (DoesDocumentPassAllFilters(doc_hit_info_.document_id())) {
+      current_allowed_sections_mask_ =
+          ComputeAllowedSectionsMask(doc_hit_info_.document_id());
 
-    schema_type_id_ = document_store_.GetSchemaTypeId(
-        doc_hit_info_.document_id(), current_time_ms_);
-    if (schema_type_id_ == kInvalidSchemaTypeId) {
-      // This means that the document is deleted or expired, so update
-      // current_allowed_sections_mask_ to skip the document.
+      schema_type_id_ = document_store_.GetSchemaTypeId(
+          doc_hit_info_.document_id(), current_time_ms_);
+      if (schema_type_id_ == kInvalidSchemaTypeId) {
+        // This means that the document is deleted or expired, so update
+        // current_allowed_sections_mask_ to skip the document.
+        current_allowed_sections_mask_ = kSectionIdMaskNone;
+      } else {
+        ICING_ASSIGN_OR_RETURN(
+            schema_name_hash_,
+            schema_store_.GetSchemaNameHash(schema_type_id_));
+      }
+    } else {
+      // This means that the document is filtered out by the document filter
+      // predicate, so update current_allowed_sections_mask_ to skip the
+      // document.
       current_allowed_sections_mask_ = kSectionIdMaskNone;
     }
   } else if (doc_hit_info_.document_id() !=
@@ -113,14 +128,15 @@
 }
 
 libtextclassifier3::Status
-DocHitInfoIteratorEmbedding::AdvanceToNextUnfilteredDocument() {
-  if (no_more_hit_ || posting_list_accessor_ == nullptr) {
+DocHitInfoIteratorEmbeddingV1::AdvanceToNextUnfilteredDocument() {
+  if (no_more_hit_ || embedding_hit_accessor_ == nullptr) {
     return absl_ports::ResourceExhaustedError(
         "No more DocHitInfos in iterator");
   }
 
   doc_hit_info_ = DocHitInfo(kInvalidDocumentId, kSectionIdMaskNone);
   schema_type_id_ = kInvalidSchemaTypeId;
+  schema_name_hash_ = 0;
   EmbeddingMatchInfos* matched_infos = nullptr;
   current_allowed_sections_mask_ = kSectionIdMaskAll;
   SectionId current_section_id = kInvalidSectionId;
@@ -142,6 +158,12 @@
       continue;
     }
 
+    if (schema_type_id_ == kInvalidSchemaTypeId) {
+      // This should never happen, because current_allowed_sections_mask_ would
+      // have been updated to kSectionIdMaskNone for this case.
+      continue;
+    }
+
     // We've reached a new section. Reset the match count and retrieve the
     // quantization type for the new section.
     if (current_section_id != embedding_hit->basic_hit().section_id()) {
@@ -151,17 +173,16 @@
       // current_allowed_sections_mask_ should be assigned to kSectionIdMaskNone
       // by AdvanceToNextEmbeddingHit, and the embedding hit should have been
       // skipped above.
-      ICING_ASSIGN_OR_RETURN(
-          quantization_type,
-          schema_store_.GetQuantizationType(
-              schema_type_id_, current_section_id));
+      ICING_ASSIGN_OR_RETURN(quantization_type,
+                             schema_store_.GetQuantizationType(
+                                 schema_type_id_, current_section_id));
     }
 
     // Calculate the semantic score.
-    ICING_ASSIGN_OR_RETURN(
-        float semantic_score,
-        embedding_index_.ScoreEmbeddingHit(*embedding_scorer_, query_,
-                                           *embedding_hit, quantization_type));
+    ICING_ASSIGN_OR_RETURN(float semantic_score,
+                           embedding_hit_accessor_->ScoreEmbeddingHit(
+                               *embedding_scorer_, query_, *embedding_hit,
+                               quantization_type, schema_name_hash_));
 
     // If the semantic score is within the desired score range, update
     // doc_hit_info_ and info_map_.
@@ -170,11 +191,13 @@
       if (matched_infos == nullptr) {
         matched_infos = &(info_map_[doc_hit_info_.document_id()]);
       }
-      matched_infos->AppendScore(semantic_score);
-      if (get_embedding_match_info_) {
+      ICING_RETURN_IF_ERROR(
+          matched_infos->AppendScore(global_scores_, semantic_score));
+      if (global_section_infos_ != nullptr) {
         // Add the section info for this embedding match.
-        matched_infos->AppendSectionInfo(current_section_id,
-                                         current_section_match_count);
+        ICING_RETURN_IF_ERROR(matched_infos->AppendSectionInfo(
+            *global_section_infos_, current_section_id,
+            current_section_match_count));
       }
     }
     ++current_section_match_count;
@@ -187,7 +210,7 @@
   return libtextclassifier3::Status::OK;
 }
 
-libtextclassifier3::Status DocHitInfoIteratorEmbedding::Advance() {
+libtextclassifier3::Status DocHitInfoIteratorEmbeddingV1::Advance() {
   do {
     ICING_RETURN_IF_ERROR(AdvanceToNextUnfilteredDocument());
   } while (doc_hit_info_.hit_section_ids_mask() == kSectionIdMaskNone);

diff --git a/icing/index/embed/doc-hit-info-iterator-embedding.h b/icing/index/embed/doc-hit-info-iterator-embedding-v1.h
similarity index 76%
rename from icing/index/embed/doc-hit-info-iterator-embedding.h
rename to icing/index/embed/doc-hit-info-iterator-embedding-v1.h
index ed2f972..fc7268a 100644
--- a/icing/index/embed/doc-hit-info-iterator-embedding.h
+++ b/icing/index/embed/doc-hit-info-iterator-embedding-v1.h

@@ -12,13 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#ifndef ICING_INDEX_EMBED_DOC_HIT_INFO_ITERATOR_EMBEDDING_H_
-#define ICING_INDEX_EMBED_DOC_HIT_INFO_ITERATOR_EMBEDDING_H_
+#ifndef ICING_INDEX_EMBED_DOC_HIT_INFO_ITERATOR_EMBEDDING_V1_H_
+#define ICING_INDEX_EMBED_DOC_HIT_INFO_ITERATOR_EMBEDDING_V1_H_
 
 #include <cstdint>
 #include <memory>
 #include <string>
-#include <string_view>
 #include <utility>
 #include <vector>
 
@@ -29,8 +28,8 @@
 #include "icing/index/embed/embedding-index.h"
 #include "icing/index/embed/embedding-query-results.h"
 #include "icing/index/embed/embedding-scorer.h"
-#include "icing/index/embed/posting-list-embedding-hit-accessor.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
+#include "icing/index/iterator/document-filter-predicate.h"
 #include "icing/index/iterator/section-restrict-data.h"
 #include "icing/proto/search.pb.h"
 #include "icing/schema/schema-store.h"
@@ -41,8 +40,9 @@
 namespace icing {
 namespace lib {
 
-class DocHitInfoIteratorEmbedding
-    : public DocHitInfoIteratorHandlingSectionRestrict {
+class DocHitInfoIteratorEmbeddingV1
+    : public DocHitInfoIteratorHandlingSectionRestrict,
+      public DocHitInfoIteratorHandlingFilter {
  public:
   // Create a DocHitInfoIterator for iterating through all docs which have an
   // embedding matched with the provided query with a score in the range of
@@ -55,14 +55,17 @@
   // help of DocHitInfoIteratorHandlingSectionRestrict.
   //
   // Returns:
-  //   - a DocHitInfoIteratorEmbedding instance on success.
+  //   - a DocHitInfoIteratorEmbeddingV1 instance on success.
   //   - Any error from posting lists.
   static libtextclassifier3::StatusOr<
-      std::unique_ptr<DocHitInfoIteratorEmbedding>>
+      std::unique_ptr<DocHitInfoIteratorEmbeddingV1>>
   Create(const PropertyProto::VectorProto* query,
          SearchSpecProto::EmbeddingQueryMetricType::Code metric_type,
-         double score_low, double score_high, bool get_embedding_match_info,
+         double score_low, double score_high,
          EmbeddingQueryResults::EmbeddingQueryMatchInfoMap* info_map,
+         std::vector<double>* global_scores,
+         std::vector<EmbeddingMatchInfos::EmbeddingMatchSectionInfo>*
+             global_section_infos,
          const EmbeddingIndex* embedding_index,
          const DocumentStore* document_store, const SchemaStore* schema_store,
          int64_t current_time_ms);
@@ -74,13 +77,18 @@
         "Query suggestions for the semanticSearch function are not supported");
   }
 
+  std::vector<std::unique_ptr<DocHitInfoIterator>*> GetChildren() override {
+    return {};
+  }
+
   CallStats GetCallStats() const override {
     return CallStats(
         /*num_leaf_advance_calls_lite_index_in=*/num_advance_calls_,
         /*num_leaf_advance_calls_main_index_in=*/0,
         /*num_leaf_advance_calls_integer_index_in=*/0,
         /*num_leaf_advance_calls_no_index_in=*/0,
-        /*num_blocks_inspected_in=*/0);
+        /*num_blocks_inspected_in=*/0,
+        /*embedding_stats_in=*/{});
   }
 
   std::string ToString() const override { return "embedding_iterator"; }
@@ -91,14 +99,18 @@
       SectionIdMask filtering_section_mask) const override {}
 
  private:
-  explicit DocHitInfoIteratorEmbedding(
+  explicit DocHitInfoIteratorEmbeddingV1(
       const PropertyProto::VectorProto* query,
       SearchSpecProto::EmbeddingQueryMetricType::Code metric_type,
       std::unique_ptr<EmbeddingScorer> embedding_scorer, double score_low,
-      double score_high, bool get_embedding_match_info,
+      double score_high,
       EmbeddingQueryResults::EmbeddingQueryMatchInfoMap* info_map,
+      std::vector<double>* global_scores,
+      std::vector<EmbeddingMatchInfos::EmbeddingMatchSectionInfo>*
+          global_section_infos,
       const EmbeddingIndex* embedding_index,
-      std::unique_ptr<PostingListEmbeddingHitAccessor> posting_list_accessor,
+      std::unique_ptr<EmbeddingIndex::EmbeddingHitAccessor>
+          embedding_hit_accessor,
       const DocumentStore* document_store, const SchemaStore* schema_store,
       int64_t current_time_ms)
       : query_(*query),
@@ -106,14 +118,16 @@
         embedding_scorer_(std::move(embedding_scorer)),
         score_low_(score_low),
         score_high_(score_high),
-        get_embedding_match_info_(get_embedding_match_info),
         info_map_(*info_map),
+        global_scores_(*global_scores),
+        global_section_infos_(global_section_infos),
         embedding_index_(*embedding_index),
-        posting_list_accessor_(std::move(posting_list_accessor)),
+        embedding_hit_accessor_(std::move(embedding_hit_accessor)),
         cached_embedding_hits_idx_(0),
         current_allowed_sections_mask_(kSectionIdMaskAll),
         no_more_hit_(false),
         schema_type_id_(kInvalidSchemaTypeId),
+        schema_name_hash_(0),
         document_store_(*document_store),
         schema_store_(*schema_store),
         current_time_ms_(current_time_ms),
@@ -153,15 +167,16 @@
   double score_low_;
   double score_high_;
 
-  // Snippet arguments
-  bool get_embedding_match_info_;
-
   // MatchInfo map
   EmbeddingQueryResults::EmbeddingQueryMatchInfoMap& info_map_;  // Does not own
+  std::vector<double>& global_scores_;                           // Does not own
+  // Nullable, and does not own. If null, section info will not be populated.
+  std::vector<EmbeddingMatchInfos::EmbeddingMatchSectionInfo>*
+      global_section_infos_;
 
   // Access to embeddings index data
   const EmbeddingIndex& embedding_index_;
-  std::unique_ptr<PostingListEmbeddingHitAccessor> posting_list_accessor_;
+  std::unique_ptr<EmbeddingIndex::EmbeddingHitAccessor> embedding_hit_accessor_;
 
   // Cached data from the embeddings index
   std::vector<EmbeddingHit> cached_embedding_hits_;
@@ -169,6 +184,10 @@
   SectionIdMask current_allowed_sections_mask_;
   bool no_more_hit_;
   SchemaTypeId schema_type_id_;  // The schema type id for the current document.
+  // The schema name hash for the current document, which will be updated
+  // together with schema_type_id_ when advancing to the next document. This
+  // value is valid only if schema_type_id_ is valid.
+  uint32_t schema_name_hash_;
 
   const DocumentStore& document_store_;
   const SchemaStore& schema_store_;
@@ -179,4 +198,4 @@
 }  // namespace lib
 }  // namespace icing
 
-#endif  // ICING_INDEX_EMBED_DOC_HIT_INFO_ITERATOR_EMBEDDING_H_
+#endif  // ICING_INDEX_EMBED_DOC_HIT_INFO_ITERATOR_EMBEDDING_V1_H_

diff --git a/icing/index/embed/doc-hit-info-iterator-embedding-v2.cc b/icing/index/embed/doc-hit-info-iterator-embedding-v2.cc
new file mode 100644
index 0000000..de16bf8
--- /dev/null
+++ b/icing/index/embed/doc-hit-info-iterator-embedding-v2.cc

@@ -0,0 +1,313 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "icing/index/embed/doc-hit-info-iterator-embedding-v2.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <memory>
+#include <numeric>
+#include <utility>
+#include <vector>
+
+#include "icing/text_classifier/lib3/utils/base/status.h"
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
+#include "icing/absl_ports/canonical_errors.h"
+#include "icing/index/embed/embedding-hit.h"
+#include "icing/index/embed/embedding-index.h"
+#include "icing/index/embed/embedding-query-results.h"
+#include "icing/index/embed/embedding-scorer.h"
+#include "icing/index/hit/doc-hit-info.h"
+#include "icing/index/hit/hit.h"
+#include "icing/proto/search.pb.h"
+#include "icing/schema/schema-store.h"
+#include "icing/schema/section.h"
+#include "icing/store/document-filter-data.h"
+#include "icing/store/document-id.h"
+#include "icing/store/document-store.h"
+#include "icing/util/status-macros.h"
+
+namespace icing {
+namespace lib {
+
+libtextclassifier3::StatusOr<std::unique_ptr<DocHitInfoIteratorEmbeddingV2>>
+DocHitInfoIteratorEmbeddingV2::Create(
+    const PropertyProto::VectorProto* query,
+    SearchSpecProto::EmbeddingQueryMetricType::Code metric_type,
+    double score_low, double score_high,
+    EmbeddingQueryResults::EmbeddingQueryMatchInfoMap* info_map,
+    std::vector<double>* global_scores,
+    std::vector<EmbeddingMatchInfos::EmbeddingMatchSectionInfo>*
+        global_section_infos,
+    const EmbeddingIndex* embedding_index, const DocumentStore* document_store,
+    const SchemaStore* schema_store, int64_t current_time_ms) {
+  ICING_RETURN_ERROR_IF_NULL(query);
+  ICING_RETURN_ERROR_IF_NULL(embedding_index);
+  ICING_RETURN_ERROR_IF_NULL(info_map);
+  ICING_RETURN_ERROR_IF_NULL(global_scores);
+  ICING_RETURN_ERROR_IF_NULL(document_store);
+  ICING_RETURN_ERROR_IF_NULL(schema_store);
+
+  libtextclassifier3::StatusOr<
+      std::unique_ptr<EmbeddingIndex::EmbeddingHitAccessor>>
+      embedding_hit_accessor_or = embedding_index->GetAccessorForVector(*query);
+  std::unique_ptr<EmbeddingIndex::EmbeddingHitAccessor> embedding_hit_accessor;
+  if (embedding_hit_accessor_or.ok()) {
+    embedding_hit_accessor = std::move(embedding_hit_accessor_or).ValueOrDie();
+  } else if (absl_ports::IsNotFound(embedding_hit_accessor_or.status())) {
+    // A not-found error should be fine, since that means there is no matching
+    // embedding hits in the index.
+    embedding_hit_accessor = nullptr;
+  } else {
+    // Otherwise, return the error as is.
+    return embedding_hit_accessor_or.status();
+  }
+
+  ICING_ASSIGN_OR_RETURN(std::unique_ptr<EmbeddingScorer> embedding_scorer,
+                         EmbeddingScorer::Create(metric_type));
+
+  return std::unique_ptr<DocHitInfoIteratorEmbeddingV2>(
+      new DocHitInfoIteratorEmbeddingV2(
+          query, metric_type, std::move(embedding_scorer), score_low,
+          score_high, info_map, global_scores, global_section_infos,
+          embedding_index, std::move(embedding_hit_accessor), document_store,
+          schema_store, current_time_ms));
+}
+
+libtextclassifier3::Status
+DocHitInfoIteratorEmbeddingV2::RetrieveNextHitsBatch() {
+  ICING_ASSIGN_OR_RETURN(std::vector<EmbeddingHit> embedding_hits,
+                         embedding_hit_accessor_->GetNextHitsBatch());
+  if (embedding_hits.empty()) {
+    no_more_hit_ = true;
+  }
+
+  // Create a sequential access order of the embedding hits for optimized
+  // memory access for embedding data.
+  cached_hit_scores_idx_ = 0;
+  cached_hit_scores_.clear();
+  cached_hit_scores_.resize(embedding_hits.size());
+  cached_section_id_masks_.clear();
+  if (delegate_ != nullptr) {
+    cached_section_id_masks_.resize(embedding_hits.size());
+  }
+
+  std::vector<int> access_order(embedding_hits.size());
+  std::iota(access_order.begin(), access_order.end(), 0);
+  std::sort(access_order.begin(), access_order.end(),
+            [&embedding_hits](int i, int j) {
+              return embedding_hits[i].location() <
+                     embedding_hits[j].location();
+            });
+  if (delegate_ != nullptr) {
+    // Add to access_order and filter any docids that don't match the delegate.
+    for (int i = 0; i < embedding_hits.size(); ++i) {
+      while (delegate_->doc_hit_info().document_id() == kInvalidDocumentId ||
+             delegate_->doc_hit_info().document_id() >
+                 embedding_hits[i].basic_hit().document_id()) {
+        if (!delegate_->Advance().ok()) {
+          break;
+        }
+      }
+      SectionIdMask section_id_mask = kSectionIdMaskNone;
+      if (delegate_->doc_hit_info().document_id() ==
+          embedding_hits[i].basic_hit().document_id()) {
+        section_id_mask = delegate_->doc_hit_info().hit_section_ids_mask();
+      }
+      cached_section_id_masks_[i] = section_id_mask;
+    }
+  }
+
+  DocumentId document_id = kInvalidDocumentId;
+  SchemaTypeId schema_type_id = kInvalidSchemaTypeId;
+  uint32_t schema_name_hash = 0;
+  SectionIdMask allowed_sections_mask = kSectionIdMaskNone;
+  for (int i = 0; i < access_order.size(); ++i) {
+    const EmbeddingHit& embedding_hit = embedding_hits[access_order[i]];
+    // Update document id, schema type id, and allowed sections mask for the
+    // new document.
+    if (embedding_hit.basic_hit().document_id() != document_id) {
+      document_id = embedding_hit.basic_hit().document_id();
+      schema_type_id = kInvalidSchemaTypeId;
+      schema_name_hash = 0;
+      allowed_sections_mask = kSectionIdMaskNone;
+
+      if (!DoesDocumentPassAllFilters(document_id)) {
+        // This means that the document is filtered out by the document filter
+        // predicate. Just skip the document.
+        continue;
+      }
+      if (delegate_ != nullptr &&
+          cached_section_id_masks_[access_order[i]] == kSectionIdMaskNone) {
+        // The document has been filtered out by the delegate. Therefore, there
+        // is no need to calculate the embedding score for this document.
+        continue;
+      }
+      schema_type_id =
+          document_store_.GetSchemaTypeId(document_id, current_time_ms_);
+      if (schema_type_id == kInvalidSchemaTypeId) {
+        // This means that the document is deleted or expired.
+        // No need to update allowed_sections_mask back to kSectionIdMaskNone,
+        // since we will always check schema_type_id itself.
+        continue;
+      }
+      ICING_ASSIGN_OR_RETURN(schema_name_hash,
+                             schema_store_.GetSchemaNameHash(schema_type_id));
+      allowed_sections_mask = ComputeAllowedSectionsMask(document_id);
+    }
+
+    SectionId section_id = embedding_hit.basic_hit().section_id();
+
+    // Filter out the embedding hit according to the section restriction.
+    if (((UINT64_C(1) << section_id) & allowed_sections_mask) == 0) {
+      continue;
+    }
+
+    // Filter out the embedding hit if the schema type id is invalid.
+    if (schema_type_id == kInvalidSchemaTypeId) {
+      continue;
+    }
+
+    // The schema type id is guaranteed to be valid here. Otherwise,
+    // allowed_sections_mask should be assigned to kSectionIdMaskNone, and the
+    // embedding hit should have been skipped above.
+    ICING_ASSIGN_OR_RETURN(
+        EmbeddingIndexingConfig::QuantizationType::Code quantization_type,
+        schema_store_.GetQuantizationType(schema_type_id, section_id));
+
+    // Calculate the semantic score.
+    ICING_ASSIGN_OR_RETURN(float semantic_score,
+                           embedding_hit_accessor_->ScoreEmbeddingHit(
+                               *embedding_scorer_, query_, embedding_hit,
+                               quantization_type, schema_name_hash));
+    cached_hit_scores_[access_order[i]] = {embedding_hit.basic_hit(),
+                                           semantic_score};
+  }
+
+  // Remove the embedding hits that are filtered out.
+  int hit_cnt = 0;
+  for (int i = 0; i < cached_hit_scores_.size(); ++i) {
+    const HitWithScore& hit_with_score = cached_hit_scores_[i];
+    if (hit_with_score.hit.is_valid()) {
+      cached_hit_scores_[hit_cnt] = hit_with_score;
+      if (delegate_ != nullptr && i < cached_section_id_masks_.size()) {
+        cached_section_id_masks_[hit_cnt] = cached_section_id_masks_[i];
+      }
+      ++hit_cnt;
+    }
+  }
+  cached_hit_scores_.resize(hit_cnt);
+  if (!cached_section_id_masks_.empty()) {
+    cached_section_id_masks_.resize(hit_cnt);
+  }
+  return libtextclassifier3::Status::OK;
+}
+
+libtextclassifier3::StatusOr<const DocHitInfoIteratorEmbeddingV2::HitWithScore*>
+DocHitInfoIteratorEmbeddingV2::AdvanceToNextEmbeddingHit() {
+  while (cached_hit_scores_idx_ == cached_hit_scores_.size()) {
+    ICING_RETURN_IF_ERROR(RetrieveNextHitsBatch());
+    if (no_more_hit_) {
+      return nullptr;
+    }
+  }
+  const HitWithScore& embedding_hit_score =
+      cached_hit_scores_[cached_hit_scores_idx_];
+  DocumentId document_id = embedding_hit_score.hit.document_id();
+  if (doc_hit_info_.document_id() == kInvalidDocumentId) {
+    doc_hit_info_.set_document_id(document_id);
+  } else if (doc_hit_info_.document_id() != document_id) {
+    return nullptr;
+  }
+  ++cached_hit_scores_idx_;
+  return &embedding_hit_score;
+}
+
+libtextclassifier3::Status
+DocHitInfoIteratorEmbeddingV2::AdvanceToNextUnfilteredDocument() {
+  if (no_more_hit_ || embedding_hit_accessor_ == nullptr) {
+    return absl_ports::ResourceExhaustedError(
+        "No more DocHitInfos in iterator");
+  }
+
+  doc_hit_info_ = DocHitInfo(kInvalidDocumentId, kSectionIdMaskNone);
+  EmbeddingMatchInfos* matched_infos = nullptr;
+  SectionId current_section_id = kInvalidSectionId;
+  int current_section_match_count = 0;
+
+  SectionIdMask delegate_section_id_mask = kSectionIdMaskNone;
+  while (true) {
+    ICING_ASSIGN_OR_RETURN(const HitWithScore* embedding_hit_score,
+                           AdvanceToNextEmbeddingHit());
+    if (embedding_hit_score == nullptr) {
+      // No more hits for the current document.
+      break;
+    }
+    if (delegate_ != nullptr) {
+      delegate_section_id_mask |=
+          cached_section_id_masks_[cached_hit_scores_idx_ - 1];
+    }
+
+    // We've reached a new section. Reset the match count and retrieve the
+    // quantization type for the new section.
+    if (current_section_id != embedding_hit_score->hit.section_id()) {
+      current_section_match_count = 0;
+      current_section_id = embedding_hit_score->hit.section_id();
+    }
+
+    float semantic_score = embedding_hit_score->score;
+    // If the semantic score is within the desired score range, update
+    // doc_hit_info_ and info_map_.
+    if (score_low_ <= semantic_score && semantic_score <= score_high_) {
+      doc_hit_info_.UpdateSection(embedding_hit_score->hit.section_id());
+      if (matched_infos == nullptr) {
+        matched_infos = &(info_map_[doc_hit_info_.document_id()]);
+      }
+      ICING_RETURN_IF_ERROR(
+          matched_infos->AppendScore(global_scores_, semantic_score));
+      if (global_section_infos_ != nullptr) {
+        // Add the section info for this embedding match.
+        ICING_RETURN_IF_ERROR(matched_infos->AppendSectionInfo(
+            *global_section_infos_, current_section_id,
+            current_section_match_count));
+      }
+    }
+    ++current_section_match_count;
+  }
+  if (doc_hit_info_.hit_section_ids_mask() != kSectionIdMaskNone) {
+    // The hit sections id mask is used by the caller of this function
+    // (DocHitInfoIteratorEmbeddingV2::Advance) to determine whether
+    // the current document has any hits meeting our score thresholds.
+    // Therefore, we can only merge here so long as at least one embedding hit
+    // has passed the score threshold.
+    doc_hit_info_.MergeSectionsFrom(delegate_section_id_mask);
+  }
+
+  if (doc_hit_info_.document_id() == kInvalidDocumentId) {
+    return absl_ports::ResourceExhaustedError(
+        "No more DocHitInfos in iterator");
+  }
+  return libtextclassifier3::Status::OK;
+}
+
+libtextclassifier3::Status DocHitInfoIteratorEmbeddingV2::Advance() {
+  do {
+    ICING_RETURN_IF_ERROR(AdvanceToNextUnfilteredDocument());
+  } while (doc_hit_info_.hit_section_ids_mask() == kSectionIdMaskNone);
+  ++num_advance_calls_;
+  return libtextclassifier3::Status::OK;
+}
+
+}  // namespace lib
+}  // namespace icing

diff --git a/icing/index/embed/doc-hit-info-iterator-embedding-v2.h b/icing/index/embed/doc-hit-info-iterator-embedding-v2.h
new file mode 100644
index 0000000..366674c
--- /dev/null
+++ b/icing/index/embed/doc-hit-info-iterator-embedding-v2.h

@@ -0,0 +1,257 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef ICING_INDEX_EMBED_DOC_HIT_INFO_ITERATOR_EMBEDDING_V2_H_
+#define ICING_INDEX_EMBED_DOC_HIT_INFO_ITERATOR_EMBEDDING_V2_H_
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "icing/text_classifier/lib3/utils/base/status.h"
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
+#include "icing/absl_ports/canonical_errors.h"
+#include "icing/absl_ports/str_cat.h"
+#include "icing/index/embed/embedding-index.h"
+#include "icing/index/embed/embedding-query-results.h"
+#include "icing/index/embed/embedding-scorer.h"
+#include "icing/index/hit/hit.h"
+#include "icing/index/iterator/doc-hit-info-iterator-section-restrict.h"
+#include "icing/index/iterator/doc-hit-info-iterator.h"
+#include "icing/index/iterator/document-filter-predicate.h"
+#include "icing/index/iterator/section-restrict-data.h"
+#include "icing/proto/search.pb.h"
+#include "icing/schema/schema-store.h"
+#include "icing/schema/section.h"
+#include "icing/store/document-store.h"
+
+namespace icing {
+namespace lib {
+
+class DocHitInfoIteratorEmbeddingV2
+    : public DocHitInfoIteratorHandlingSectionRestrict,
+      public DocHitInfoIteratorHandlingFilter {
+ public:
+  // Create a DocHitInfoIterator for iterating through all docs which have an
+  // embedding matched with the provided query with a score in the range of
+  // [score_low, score_high], using the provided metric_type.
+  //
+  // The iterator will store the matched embedding scores in info_map to
+  // prepare for scoring and snippeting.
+  //
+  // The iterator will handle the section restriction logic internally with the
+  // help of DocHitInfoIteratorHandlingSectionRestrict.
+  //
+  // Returns:
+  //   - a DocHitInfoIteratorEmbeddingV2 instance on success.
+  //   - Any error from posting lists.
+  static libtextclassifier3::StatusOr<
+      std::unique_ptr<DocHitInfoIteratorEmbeddingV2>>
+  Create(const PropertyProto::VectorProto* query,
+         SearchSpecProto::EmbeddingQueryMetricType::Code metric_type,
+         double score_low, double score_high,
+         EmbeddingQueryResults::EmbeddingQueryMatchInfoMap* info_map,
+         std::vector<double>* global_scores,
+         std::vector<EmbeddingMatchInfos::EmbeddingMatchSectionInfo>*
+             global_section_infos,
+         const EmbeddingIndex* embedding_index,
+         const DocumentStore* document_store, const SchemaStore* schema_store,
+         int64_t current_time_ms);
+
+  bool CanAdoptDelegate() const override { return true; }
+
+  void AdoptDelegate(std::unique_ptr<DocHitInfoIterator> delegate,
+                     bool delegate_node_is_right_most) override {
+    delegate_ = std::move(delegate);
+    delegate_node_is_right_most_ = delegate_node_is_right_most;
+  }
+
+  bool HandleSectionRestriction(SectionRestrictData* other_data) override {
+    // Apply section restriction to delegate if we have one.
+    if (delegate_ != nullptr) {
+      delegate_ = DocHitInfoIteratorSectionRestrict::ApplyRestrictions(
+          std::move(delegate_), other_data);
+    }
+    return DocHitInfoIteratorHandlingSectionRestrict::HandleSectionRestriction(
+        other_data);
+  }
+
+  libtextclassifier3::Status Advance() override;
+
+  libtextclassifier3::StatusOr<TrimmedNode> TrimRightMostNode() && override {
+    if (delegate_ != nullptr && !delegate_node_is_right_most_) {
+      return std::move(*delegate_).TrimRightMostNode();
+    }
+    return absl_ports::InvalidArgumentError(
+        "Query suggestions for the semanticSearch function are not supported");
+  }
+
+  std::vector<std::unique_ptr<DocHitInfoIterator>*> GetChildren() override {
+    if (delegate_ != nullptr) {
+      return {&delegate_};
+    }
+    return {};
+  }
+
+  CallStats GetCallStats() const override {
+    CallStats call_stats(
+        /*num_leaf_advance_calls_lite_index_in=*/num_advance_calls_,
+        /*num_leaf_advance_calls_main_index_in=*/0,
+        /*num_leaf_advance_calls_integer_index_in=*/0,
+        /*num_leaf_advance_calls_no_index_in=*/0,
+        /*num_blocks_inspected_in=*/0,
+        embedding_hit_accessor_ != nullptr
+            ? embedding_hit_accessor_->GetEmbeddingStats()
+            : CallStats::EmbeddingStats{});
+    if (delegate_ != nullptr) {
+      call_stats += delegate_->GetCallStats();
+    }
+    return call_stats;
+  }
+
+  std::string ToString() const override {
+    if (delegate_ != nullptr) {
+      return absl_ports::StrCat("embedding_iterator with delegate (",
+                                delegate_->ToString(), ")");
+    }
+    return "embedding_iterator";
+  }
+
+  // PopulateMatchedTermsStats is not applicable to embedding search.
+  void PopulateMatchedTermsStats(
+      std::vector<TermMatchInfo>* matched_terms_stats,
+      SectionIdMask filtering_section_mask) const override {
+    if (delegate_ != nullptr) {
+      delegate_->PopulateMatchedTermsStats(matched_terms_stats,
+                                           filtering_section_mask);
+    }
+  }
+
+ private:
+  struct HitWithScore {
+    BasicHit hit;
+    float score;
+  };
+
+  explicit DocHitInfoIteratorEmbeddingV2(
+      const PropertyProto::VectorProto* query,
+      SearchSpecProto::EmbeddingQueryMetricType::Code metric_type,
+      std::unique_ptr<EmbeddingScorer> embedding_scorer, double score_low,
+      double score_high,
+      EmbeddingQueryResults::EmbeddingQueryMatchInfoMap* info_map,
+      std::vector<double>* global_scores,
+      std::vector<EmbeddingMatchInfos::EmbeddingMatchSectionInfo>*
+          global_section_infos,
+      const EmbeddingIndex* embedding_index,
+      std::unique_ptr<EmbeddingIndex::EmbeddingHitAccessor>
+          embedding_hit_accessor,
+      const DocumentStore* document_store, const SchemaStore* schema_store,
+      int64_t current_time_ms)
+      : query_(*query),
+        metric_type_(metric_type),
+        embedding_scorer_(std::move(embedding_scorer)),
+        score_low_(score_low),
+        score_high_(score_high),
+        info_map_(*info_map),
+        global_scores_(*global_scores),
+        global_section_infos_(global_section_infos),
+        embedding_index_(*embedding_index),
+        embedding_hit_accessor_(std::move(embedding_hit_accessor)),
+        cached_hit_scores_idx_(0),
+        no_more_hit_(false),
+        document_store_(*document_store),
+        schema_store_(*schema_store),
+        current_time_ms_(current_time_ms),
+        num_advance_calls_(0) {}
+
+  // Retrieve the next batch of embedding hits from the posting list.
+  //
+  // Hits that do not pass section restriction or document filter will be
+  // filtered out. Otherwise, the hits will be scored and added to
+  // cached_hit_scores_.
+  //
+  // Returns:
+  //   - OK, if it is able to retrieve the next batch of embedding hits.
+  //   - Any error from posting lists.
+  libtextclassifier3::Status RetrieveNextHitsBatch();
+
+  // Advance to the next embedding hit of the current document. If the current
+  // document id is kInvalidDocumentId, the method will advance to the first
+  // embedding hit of the next document and update doc_hit_info_.
+  //
+  // This method also properly updates cached_hit_scores_,
+  // cached_hit_scores_idx_, and no_more_hit_ to reflect the current
+  // state.
+  //
+  // Returns:
+  //   - a const pointer to the next embedding hit on success.
+  //   - nullptr, if there is no more hit for the current document, or no more
+  //     hit in general if the current document id is kInvalidDocumentId.
+  //   - Any error from posting lists.
+  libtextclassifier3::StatusOr<const HitWithScore*> AdvanceToNextEmbeddingHit();
+
+  // Similar to Advance(), this method advances the iterator to the next
+  // document, but it does not guarantee that the next document will have
+  // a matched embedding hit within the score range.
+  //
+  // Returns:
+  //   - OK, if it is able to advance to a new document_id.
+  //   - RESOUCE_EXHAUSTED, if we have run out of document_ids to iterate over.
+  //   - Any error from posting lists.
+  libtextclassifier3::Status AdvanceToNextUnfilteredDocument();
+
+  std::unique_ptr<DocHitInfoIterator> delegate_;
+  // Whether the delegate is a node to the right of the current node. This
+  // affects the behavior of TrimRightMostNode.
+  bool delegate_node_is_right_most_;
+
+  // Query information
+  const PropertyProto::VectorProto& query_;  // Does not own
+
+  // Scoring arguments
+  SearchSpecProto::EmbeddingQueryMetricType::Code metric_type_;
+  std::unique_ptr<EmbeddingScorer> embedding_scorer_;
+  double score_low_;
+  double score_high_;
+
+  // MatchInfo map
+  EmbeddingQueryResults::EmbeddingQueryMatchInfoMap& info_map_;  // Does not own
+  std::vector<double>& global_scores_;                           // Does not own
+  // Nullable, and does not own. If null, section info will not be populated.
+  std::vector<EmbeddingMatchInfos::EmbeddingMatchSectionInfo>*
+      global_section_infos_;
+
+  // Access to embeddings index data
+  const EmbeddingIndex& embedding_index_;
+  std::unique_ptr<EmbeddingIndex::EmbeddingHitAccessor>
+      embedding_hit_accessor_;  // Nullable.
+
+  // Cached data from the embeddings index
+  std::vector<HitWithScore> cached_hit_scores_;
+  std::vector<SectionIdMask> cached_section_id_masks_;
+  int cached_hit_scores_idx_;
+  bool no_more_hit_;
+
+  const DocumentStore& document_store_;
+  const SchemaStore& schema_store_;
+  int64_t current_time_ms_;
+  int num_advance_calls_;
+};
+
+}  // namespace lib
+}  // namespace icing
+
+#endif  // ICING_INDEX_EMBED_DOC_HIT_INFO_ITERATOR_EMBEDDING_V2_H_

diff --git a/icing/index/embed/embedding-index.cc b/icing/index/embed/embedding-index.cc
index 028d85a..f05b9d1 100644
--- a/icing/index/embed/embedding-index.cc
+++ b/icing/index/embed/embedding-index.cc

@@ -17,6 +17,7 @@
 #include <algorithm>
 #include <cstdint>
 #include <cstring>
+#include <limits>
 #include <memory>
 #include <string>
 #include <string_view>
@@ -59,15 +60,17 @@
 // The maximum size of the embedding hit list mmapper.
 // We use 64MiB for 32-bit platforms and 128MiB for 64-bit platforms.
 #ifdef ICING_ARCH_BIT_64
-  constexpr uint32_t kEmbeddingHitListMapperMaxSize = 128 * 1024 * 1024;
+constexpr uint32_t kEmbeddingHitListMapperMaxSize = 128 * 1024 * 1024;
 #else
-  constexpr uint32_t kEmbeddingHitListMapperMaxSize = 64 * 1024 * 1024;
+constexpr uint32_t kEmbeddingHitListMapperMaxSize = 64 * 1024 * 1024;
 #endif
 
 // The maximum length returned by encode_util::EncodeIntToCString is 5 for
 // uint32_t.
 constexpr uint32_t kEncodedDimensionLength = 5;
 
+constexpr uint32_t kInvalidShardId = std::numeric_limits<uint32_t>::max();
+
 std::string GetMetadataFilePath(std::string_view working_path) {
   return absl_ports::StrCat(working_path, "/metadata");
 }
@@ -80,13 +83,28 @@
   return absl_ports::StrCat(working_path, "/embedding_hit_list_mapper");
 }
 
-std::string GetEmbeddingVectorsFilePath(std::string_view working_path) {
-  return absl_ports::StrCat(working_path, "/embedding_vectors");
+std::string AppendShardId(std::string&& file_path, uint32_t shard_id,
+                          uint32_t num_shards) {
+  if (num_shards > 1) {
+    return absl_ports::StrCat(std::move(file_path), "_",
+                              std::to_string(shard_id));
+  }
+  return file_path;
 }
 
-std::string GetQuantizedEmbeddingVectorsFilePath(
-    std::string_view working_path) {
-  return absl_ports::StrCat(working_path, "/quantized_embedding_vectors");
+std::string GetEmbeddingVectorsFilePath(std::string_view working_path,
+                                        uint32_t shard_id,
+                                        uint32_t num_shards) {
+  return AppendShardId(absl_ports::StrCat(working_path, "/embedding_vectors"),
+                       shard_id, num_shards);
+}
+
+std::string GetQuantizedEmbeddingVectorsFilePath(std::string_view working_path,
+                                                 uint32_t shard_id,
+                                                 uint32_t num_shards) {
+  return AppendShardId(
+      absl_ports::StrCat(working_path, "/quantized_embedding_vectors"),
+      shard_id, num_shards);
 }
 
 // An injective function that maps the ordered pair (dimension, model_signature)
@@ -126,24 +144,33 @@
 
 }  // namespace
 
+uint32_t EmbeddingIndex::GetShardId(uint32_t dimension,
+                                    std::string_view model_signature,
+                                    std::string_view schema_name) const {
+  return GetShardId(
+      GetPostingListKeyHash(GetPostingListKey(dimension, model_signature)),
+      SchemaStore::GetSchemaNameHash(schema_name));
+}
+
 libtextclassifier3::StatusOr<std::unique_ptr<EmbeddingIndex>>
 EmbeddingIndex::Create(const Filesystem* filesystem, std::string working_path,
-                       const Clock* clock, const FeatureFlags* feature_flags) {
+                       const Clock* clock, const FeatureFlags* feature_flags,
+                       uint32_t num_shards) {
   ICING_RETURN_ERROR_IF_NULL(filesystem);
   ICING_RETURN_ERROR_IF_NULL(clock);
 
-  std::unique_ptr<EmbeddingIndex> index =
-      std::unique_ptr<EmbeddingIndex>(new EmbeddingIndex(
-          *filesystem, std::move(working_path), clock, feature_flags));
+  if (num_shards == 0) {
+    return absl_ports::InvalidArgumentError("Number of shards cannot be 0");
+  }
+
+  std::unique_ptr<EmbeddingIndex> index = std::unique_ptr<EmbeddingIndex>(
+      new EmbeddingIndex(*filesystem, std::move(working_path), clock,
+                         feature_flags, num_shards));
   ICING_RETURN_IF_ERROR(index->Initialize());
   return index;
 }
 
-libtextclassifier3::Status EmbeddingIndex::CreateStorageDataIfNonEmpty() {
-  if (is_empty()) {
-    return libtextclassifier3::Status::OK;
-  }
-
+libtextclassifier3::Status EmbeddingIndex::CreateStorageData() {
   ICING_ASSIGN_OR_RETURN(FlashIndexStorage flash_index_storage,
                          FlashIndexStorage::Create(
                              GetFlashIndexStorageFilePath(working_path_),
@@ -157,17 +184,21 @@
           filesystem_, GetEmbeddingHitListMapperPath(working_path_),
           kEmbeddingHitListMapperMaxSize));
 
-  ICING_ASSIGN_OR_RETURN(
-      embedding_vectors_,
-      FileBackedVector<float>::Create(
-          filesystem_, GetEmbeddingVectorsFilePath(working_path_),
-          MemoryMappedFile::READ_WRITE_AUTO_SYNC));
-
-  ICING_ASSIGN_OR_RETURN(
-      quantized_embedding_vectors_,
-      FileBackedVector<char>::Create(
-          filesystem_, GetQuantizedEmbeddingVectorsFilePath(working_path_),
-          MemoryMappedFile::READ_WRITE_AUTO_SYNC));
+  // Scan for existing sharded vector files and load them.
+  for (int i = 0; i < num_shards_; ++i) {
+    // Non-quantized vectors
+    if (filesystem_.FileExists(
+            GetEmbeddingVectorsFilePath(working_path_, i, num_shards_)
+                .c_str())) {
+      ICING_RETURN_IF_ERROR(GetOrCreateEmbeddingVector(i));
+    }
+    // Quantized vectors
+    if (filesystem_.FileExists(
+            GetQuantizedEmbeddingVectorsFilePath(working_path_, i, num_shards_)
+                .c_str())) {
+      ICING_RETURN_IF_ERROR(GetOrCreateQuantizedEmbeddingVector(i));
+    }
+  }
 
   return libtextclassifier3::Status::OK;
 }
@@ -176,8 +207,43 @@
   if (!is_empty()) {
     return libtextclassifier3::Status::OK;
   }
+  ICING_RETURN_IF_ERROR(CreateStorageData());
   info().is_empty = false;
-  return CreateStorageDataIfNonEmpty();
+  return libtextclassifier3::Status::OK;
+}
+
+libtextclassifier3::StatusOr<FileBackedVector<float>*>
+EmbeddingIndex::GetOrCreateEmbeddingVector(uint32_t shard_id) {
+  if (shard_id >= num_shards_) {
+    return absl_ports::InvalidArgumentError("Shard id is out of range.");
+  }
+  auto& fbv_ptr = embedding_vectors_[shard_id];
+  if (fbv_ptr == nullptr) {
+    ICING_ASSIGN_OR_RETURN(
+        fbv_ptr,
+        FileBackedVector<float>::Create(
+            filesystem_,
+            GetEmbeddingVectorsFilePath(working_path_, shard_id, num_shards_),
+            MemoryMappedFile::READ_WRITE_AUTO_SYNC));
+  }
+  return fbv_ptr.get();
+}
+
+libtextclassifier3::StatusOr<FileBackedVector<char>*>
+EmbeddingIndex::GetOrCreateQuantizedEmbeddingVector(uint32_t shard_id) {
+  if (shard_id >= num_shards_) {
+    return absl_ports::InvalidArgumentError("Shard id is out of range.");
+  }
+  auto& fbv_ptr = quantized_embedding_vectors_[shard_id];
+  if (fbv_ptr == nullptr) {
+    ICING_ASSIGN_OR_RETURN(fbv_ptr,
+                           FileBackedVector<char>::Create(
+                               filesystem_,
+                               GetQuantizedEmbeddingVectorsFilePath(
+                                   working_path_, shard_id, num_shards_),
+                               MemoryMappedFile::READ_WRITE_AUTO_SYNC));
+  }
+  return fbv_ptr.get();
 }
 
 libtextclassifier3::Status EmbeddingIndex::Initialize() {
@@ -207,7 +273,8 @@
     info().magic = Info::kMagic;
     info().last_added_document_id = kInvalidDocumentId;
     info().is_empty = true;
-    memset(Info().padding_, 0, Info::kPaddingSize);
+    info().num_shards = num_shards_;
+    memset(info().padding_, 0, Info::kPaddingSize);
     ICING_RETURN_IF_ERROR(InitializeNewStorage());
   } else {
     if (metadata_mmapped_file_->available_size() != kMetadataFileSize) {
@@ -215,10 +282,39 @@
           "Incorrect metadata file size");
     }
     if (info().magic != Info::kMagic) {
-      return absl_ports::FailedPreconditionError("Incorrect magic value");
+      ICING_LOG(ERROR) << "Invalid header magic for EmbeddingIndex "
+                       << working_path_ << ". Expected: " << Info::kMagic
+                       << ", actual: " << info().magic;
+      return absl_ports::FailedPreconditionError(absl_ports::StrCat(
+          "Invalid header magic for EmbeddingIndex: ", working_path_));
     }
-    ICING_RETURN_IF_ERROR(CreateStorageDataIfNonEmpty());
+    uint32_t num_shards_read = info().num_shards;
+    // This happens for old versions of embedding index that did not have
+    // num_shards set. Just treat it as 1.
+    if (num_shards_read == 0) {
+      ICING_LOG(INFO)
+          << "Number of shards not set in metadata. Defaulting to 1.";
+      num_shards_read = 1;
+    }
+    // If num_shards doesn't match, the index need to be rebuilt.
+    if (num_shards_read != num_shards_) {
+      ICING_LOG(ERROR) << "Mismatched number of shards. Expected "
+                       << num_shards_ << ", actual " << num_shards_read;
+      return absl_ports::FailedPreconditionError(absl_ports::StrCat(
+          "Mismatched number of shards in metadata for EmbeddingIndex: ",
+          working_path_));
+    }
+    if (!info().is_empty) {
+      ICING_RETURN_IF_ERROR(CreateStorageData());
+    }
     ICING_RETURN_IF_ERROR(InitializeExistingStorage());
+    if (info().num_shards == 0) {
+      // This means that num_shards isn't set in the header, but we are still
+      // considering it a match for num_shards_. The only possibility is that
+      // num_shards_ == 1 and we have treated info().num_shards as 1.
+      // Now, update the header to record the correct num_shards_.
+      info().num_shards = num_shards_;
+    }
   }
   return libtextclassifier3::Status::OK;
 }
@@ -228,8 +324,10 @@
   metadata_mmapped_file_.reset();
   flash_index_storage_.reset();
   embedding_posting_list_mapper_.reset();
-  embedding_vectors_.reset();
-  quantized_embedding_vectors_.reset();
+  for (int i = 0; i < num_shards_; ++i) {
+    embedding_vectors_[i].reset();
+    quantized_embedding_vectors_[i].reset();
+  }
   if (filesystem_.DirectoryExists(working_path_.c_str())) {
     ICING_RETURN_IF_ERROR(Discard(filesystem_, working_path_));
   }
@@ -237,7 +335,8 @@
   return Initialize();
 }
 
-libtextclassifier3::StatusOr<std::unique_ptr<PostingListEmbeddingHitAccessor>>
+libtextclassifier3::StatusOr<
+    std::unique_ptr<EmbeddingIndex::EmbeddingHitAccessor>>
 EmbeddingIndex::GetAccessor(uint32_t dimension,
                             std::string_view model_signature) const {
   if (dimension == 0) {
@@ -250,24 +349,33 @@
   std::string key = GetPostingListKey(dimension, model_signature);
   ICING_ASSIGN_OR_RETURN(PostingListIdentifier posting_list_id,
                          embedding_posting_list_mapper_->Get(key));
-  return PostingListEmbeddingHitAccessor::CreateFromExisting(
-      flash_index_storage_.get(), posting_list_hit_serializer_.get(),
-      posting_list_id);
+  ICING_ASSIGN_OR_RETURN(
+      std::unique_ptr<PostingListEmbeddingHitAccessor> pl_accessor,
+      PostingListEmbeddingHitAccessor::CreateFromExisting(
+          flash_index_storage_.get(), posting_list_hit_serializer_.get(),
+          posting_list_id));
+  return std::make_unique<EmbeddingHitAccessor>(this, std::move(pl_accessor),
+                                                key);
 }
 
 libtextclassifier3::StatusOr<uint32_t> EmbeddingIndex::AppendEmbeddingVector(
     const PropertyProto::VectorProto& vector,
-    EmbeddingIndexingConfig::QuantizationType::Code quantization_type) {
+    EmbeddingIndexingConfig::QuantizationType::Code quantization_type,
+    uint32_t shard_id) {
   uint32_t dimension = vector.values().size();
   uint32_t location;
   if (!feature_flags_->enable_embedding_quantization() ||
       quantization_type == EmbeddingIndexingConfig::QuantizationType::NONE) {
-    location = embedding_vectors_->num_elements();
+    ICING_ASSIGN_OR_RETURN(FileBackedVector<float> * fbv_ptr,
+                           GetOrCreateEmbeddingVector(shard_id));
+    location = fbv_ptr->num_elements();
     ICING_ASSIGN_OR_RETURN(
         FileBackedVector<float>::MutableArrayView mutable_arr,
-        embedding_vectors_->Allocate(dimension));
+        fbv_ptr->Allocate(dimension));
     mutable_arr.SetArray(/*idx=*/0, vector.values().data(), dimension);
   } else {
+    ICING_ASSIGN_OR_RETURN(FileBackedVector<char> * fbv_ptr,
+                           GetOrCreateQuantizedEmbeddingVector(shard_id));
     ICING_ASSIGN_OR_RETURN(Quantizer quantizer, CreateQuantizer(vector));
     // Quantize the vector
     std::vector<uint8_t> quantized_values;
@@ -277,10 +385,9 @@
     }
 
     // Store the quantizer and the quantized vector
-    location = quantized_embedding_vectors_->num_elements();
-    ICING_ASSIGN_OR_RETURN(
-        FileBackedVector<char>::MutableArrayView mutable_arr,
-        quantized_embedding_vectors_->Allocate(sizeof(Quantizer) + dimension));
+    location = fbv_ptr->num_elements();
+    ICING_ASSIGN_OR_RETURN(FileBackedVector<char>::MutableArrayView mutable_arr,
+                           fbv_ptr->Allocate(sizeof(Quantizer) + dimension));
     mutable_arr.SetArray(/*idx=*/0, reinterpret_cast<char*>(&quantizer),
                          sizeof(Quantizer));
     mutable_arr.SetArray(/*idx=*/sizeof(Quantizer),
@@ -292,15 +399,19 @@
 
 libtextclassifier3::Status EmbeddingIndex::BufferEmbedding(
     const BasicHit& basic_hit, const PropertyProto::VectorProto& vector,
-    EmbeddingIndexingConfig::QuantizationType::Code quantization_type) {
+    EmbeddingIndexingConfig::QuantizationType::Code quantization_type,
+    std::string_view schema_name) {
   if (vector.values().empty()) {
     return absl_ports::InvalidArgumentError("Vector dimension is 0");
   }
   ICING_RETURN_IF_ERROR(MarkIndexNonEmpty());
 
   std::string key = GetPostingListKey(vector);
-  ICING_ASSIGN_OR_RETURN(uint32_t location,
-                         AppendEmbeddingVector(vector, quantization_type));
+  uint32_t shard_id = GetShardId(GetPostingListKeyHash(key),
+                                 SchemaStore::GetSchemaNameHash(schema_name));
+  ICING_ASSIGN_OR_RETURN(
+      uint32_t location,
+      AppendEmbeddingVector(vector, quantization_type, shard_id));
 
   // Buffer the embedding hit.
   pending_embedding_hits_.push_back(
@@ -373,28 +484,33 @@
 libtextclassifier3::StatusOr<uint32_t> EmbeddingIndex::TransferEmbeddingVector(
     const EmbeddingHit& old_hit, uint32_t dimension,
     EmbeddingIndexingConfig::QuantizationType::Code quantization_type,
-    EmbeddingIndex* new_index) const {
+    uint32_t shard_id, EmbeddingIndex* new_index) const {
   uint32_t new_location;
   if (!feature_flags_->enable_embedding_quantization() ||
       quantization_type == EmbeddingIndexingConfig::QuantizationType::NONE) {
     ICING_ASSIGN_OR_RETURN(const float* old_vector,
-                           GetEmbeddingVector(old_hit, dimension));
-    new_location = new_index->embedding_vectors_->num_elements();
+                           GetEmbeddingVector(old_hit, dimension, shard_id));
+
+    // Copy the embedding vector of the hit to the new index.
+    ICING_ASSIGN_OR_RETURN(FileBackedVector<float> * fbv_ptr,
+                           new_index->GetOrCreateEmbeddingVector(shard_id));
+    new_location = fbv_ptr->num_elements();
+    ICING_ASSIGN_OR_RETURN(
+        FileBackedVector<float>::MutableArrayView mutable_arr,
+        fbv_ptr->Allocate(dimension));
+    mutable_arr.SetArray(/*idx=*/0, old_vector, dimension);
+  } else {
+    ICING_ASSIGN_OR_RETURN(
+        const char* old_data,
+        GetQuantizedEmbeddingVector(old_hit, dimension, shard_id));
 
     // Copy the embedding vector of the hit to the new index.
     ICING_ASSIGN_OR_RETURN(
-        FileBackedVector<float>::MutableArrayView mutable_arr,
-        new_index->embedding_vectors_->Allocate(dimension));
-    mutable_arr.SetArray(/*idx=*/0, old_vector, dimension);
-  } else {
-    ICING_ASSIGN_OR_RETURN(const char* old_data,
-                           GetQuantizedEmbeddingVector(old_hit, dimension));
-    new_location = new_index->quantized_embedding_vectors_->num_elements();
-
-    // Copy the embedding vector of the hit to the new index.
+        FileBackedVector<char> * fbv_ptr,
+        new_index->GetOrCreateQuantizedEmbeddingVector(shard_id));
+    new_location = fbv_ptr->num_elements();
     ICING_ASSIGN_OR_RETURN(FileBackedVector<char>::MutableArrayView mutable_arr,
-                           new_index->quantized_embedding_vectors_->Allocate(
-                               sizeof(Quantizer) + dimension));
+                           fbv_ptr->Allocate(sizeof(Quantizer) + dimension));
     mutable_arr.SetArray(/*idx=*/0, old_data, sizeof(Quantizer) + dimension);
   }
   return new_location;
@@ -431,6 +547,7 @@
             /*existing_posting_list_id=*/itr->GetValue()));
     DocumentId last_new_document_id = kInvalidDocumentId;
     SchemaTypeId schema_type_id = kInvalidSchemaTypeId;
+    uint32_t shard_id = kInvalidShardId;
     while (true) {
       ICING_ASSIGN_OR_RETURN(std::vector<EmbeddingHit> batch,
                              old_pl_accessor->GetNextHitsBatch());
@@ -459,9 +576,22 @@
         if (new_document_id != last_new_document_id) {
           schema_type_id =
               document_store.GetSchemaTypeId(new_document_id, current_time_ms);
+          libtextclassifier3::StatusOr<uint32_t> schema_name_hash_or =
+              schema_store.GetSchemaNameHash(schema_type_id);
+          // Shard id only depends on posting list key and schema name hash.
+          // Posting list key will not change in this scope, so we only update
+          // the shard id here for a new document, since it can have a
+          // different schema type.
+          if (schema_name_hash_or.ok()) {
+            shard_id = GetShardId(GetPostingListKeyHash(key),
+                                  schema_name_hash_or.ValueOrDie());
+          } else {
+            shard_id = kInvalidShardId;
+          }
         }
         last_new_document_id = new_document_id;
-        if (schema_type_id == kInvalidSchemaTypeId) {
+        if (schema_type_id == kInvalidSchemaTypeId ||
+            shard_id == kInvalidShardId) {
           // This should not happen, since document store is optimized first,
           // so that new_document_id here should be alive.
           continue;
@@ -475,7 +605,7 @@
         ICING_ASSIGN_OR_RETURN(
             uint32_t new_location,
             TransferEmbeddingVector(old_hit, dimension, quantization_type,
-                                    new_index));
+                                    shard_id, new_index));
         new_hits.push_back(EmbeddingHit(
             BasicHit(old_hit.basic_hit().section_id(), new_document_id),
             new_location));
@@ -540,7 +670,7 @@
     ICING_ASSIGN_OR_RETURN(
         std::unique_ptr<EmbeddingIndex> new_index,
         EmbeddingIndex::Create(&filesystem_, temporary_index_dir.dir(), &clock_,
-                               feature_flags_));
+                               feature_flags_, num_shards_));
     ICING_RETURN_IF_ERROR(TransferIndex(*document_store, *schema_store,
                                         document_id_old_to_new,
                                         new_index.get()));
@@ -552,8 +682,10 @@
   metadata_mmapped_file_.reset();
   flash_index_storage_.reset();
   embedding_posting_list_mapper_.reset();
-  embedding_vectors_.reset();
-  quantized_embedding_vectors_.reset();
+  for (int i = 0; i < num_shards_; ++i) {
+    embedding_vectors_[i].reset();
+    quantized_embedding_vectors_[i].reset();
+  }
 
   if (!filesystem_.SwapFiles(temporary_index_dir.dir().c_str(),
                              working_path_.c_str())) {
@@ -566,28 +698,55 @@
   return Initialize();
 }
 
-libtextclassifier3::StatusOr<float> EmbeddingIndex::ScoreEmbeddingHit(
+libtextclassifier3::StatusOr<float>
+EmbeddingIndex::EmbeddingHitAccessor::ScoreEmbeddingHit(
     const EmbeddingScorer& scorer, const PropertyProto::VectorProto& query,
     const EmbeddingHit& hit,
-    EmbeddingIndexingConfig::QuantizationType::Code quantization_type) const {
+    EmbeddingIndexingConfig::QuantizationType::Code quantization_type,
+    uint32_t schema_name_hash) {
+  uint32_t shard_id =
+      embedding_index_.GetShardId(posting_list_key_hash_, schema_name_hash);
   int dimension = query.values().size();
   float semantic_score;
-  if (!feature_flags_->enable_embedding_quantization() ||
+  if (!embedding_index_.feature_flags_->enable_embedding_quantization() ||
       quantization_type == EmbeddingIndexingConfig::QuantizationType::NONE) {
-    ICING_ASSIGN_OR_RETURN(const float* vector,
-                           GetEmbeddingVector(hit, dimension));
-    semantic_score = scorer.Score(dimension,
-                                  /*v1=*/query.values().data(),
-                                  /*v2=*/vector);
+    ICING_ASSIGN_OR_RETURN(
+        const float* vector,
+        embedding_index_.GetEmbeddingVector(hit, dimension, shard_id));
+    if (embedding_index_.feature_flags_->enable_eigen_embedding_scoring()) {
+      semantic_score = scorer.EigenScore(dimension,
+                                         /*v1=*/query.values().data(),
+                                         /*v2=*/vector);
+    } else {
+      semantic_score = scorer.Score(dimension,
+                                    /*v1=*/query.values().data(),
+                                    /*v2=*/vector);
+    }
+    ++embedding_stats_.num_unquantized_embeddings_scored;
+    embedding_stats_.unquantized_shards_read.insert(shard_id);
+    embedding_stats_.num_embedding_bytes_read +=
+        static_cast<int64_t>(sizeof(float)) * dimension;
   } else {
-    ICING_ASSIGN_OR_RETURN(const char* data,
-                           GetQuantizedEmbeddingVector(hit, dimension));
+    ICING_ASSIGN_OR_RETURN(
+        const char* data,
+        embedding_index_.GetQuantizedEmbeddingVector(hit, dimension, shard_id));
     Quantizer quantizer(data);
     const uint8_t* quantized_vector =
         reinterpret_cast<const uint8_t*>(data + sizeof(Quantizer));
-    semantic_score = scorer.Score(dimension,
-                                  /*v1=*/query.values().data(),
-                                  /*v2=*/quantized_vector, quantizer);
+    if (embedding_index_.feature_flags_->enable_eigen_embedding_scoring()) {
+      semantic_score = scorer.EigenScore(dimension,
+                                         /*v1=*/query.values().data(),
+                                         /*v2=*/quantized_vector, quantizer);
+    } else {
+      semantic_score = scorer.Score(dimension,
+                                    /*v1=*/query.values().data(),
+                                    /*v2=*/quantized_vector, quantizer);
+    }
+    ++embedding_stats_.num_quantized_embeddings_scored;
+    embedding_stats_.quantized_shards_read.insert(shard_id);
+    embedding_stats_.num_embedding_bytes_read +=
+        static_cast<int64_t>(sizeof(Quantizer)) +
+        static_cast<int64_t>(sizeof(uint8_t)) * dimension;
   }
   return semantic_score;
 }
@@ -604,8 +763,14 @@
     return absl_ports::InternalError("Fail to persist flash index to disk");
   }
   ICING_RETURN_IF_ERROR(embedding_posting_list_mapper_->PersistToDisk());
-  ICING_RETURN_IF_ERROR(embedding_vectors_->PersistToDisk());
-  ICING_RETURN_IF_ERROR(quantized_embedding_vectors_->PersistToDisk());
+  for (int i = 0; i < num_shards_; ++i) {
+    if (embedding_vectors_[i] != nullptr) {
+      ICING_RETURN_IF_ERROR(embedding_vectors_[i]->PersistToDisk());
+    }
+    if (quantized_embedding_vectors_[i] != nullptr) {
+      ICING_RETURN_IF_ERROR(quantized_embedding_vectors_[i]->PersistToDisk());
+    }
+  }
   return libtextclassifier3::Status::OK;
 }
 
@@ -615,13 +780,21 @@
   }
   ICING_ASSIGN_OR_RETURN(Crc32 embedding_posting_list_mapper_crc,
                          embedding_posting_list_mapper_->UpdateChecksum());
-  ICING_ASSIGN_OR_RETURN(Crc32 embedding_vectors_crc,
-                         embedding_vectors_->UpdateChecksum());
-  ICING_ASSIGN_OR_RETURN(Crc32 quantized_embedding_vectors_crc,
-                         quantized_embedding_vectors_->UpdateChecksum());
-  return Crc32(embedding_posting_list_mapper_crc.Get() ^
-               embedding_vectors_crc.Get() ^
-               quantized_embedding_vectors_crc.Get());
+  uint32_t checksum = embedding_posting_list_mapper_crc.Get();
+
+  for (int i = 0; i < num_shards_; ++i) {
+    if (embedding_vectors_[i] != nullptr) {
+      ICING_ASSIGN_OR_RETURN(Crc32 crc,
+                             embedding_vectors_[i]->UpdateChecksum());
+      checksum ^= crc.Get();
+    }
+    if (quantized_embedding_vectors_[i] != nullptr) {
+      ICING_ASSIGN_OR_RETURN(Crc32 crc,
+                             quantized_embedding_vectors_[i]->UpdateChecksum());
+      checksum ^= crc.Get();
+    }
+  }
+  return Crc32(checksum);
 }
 
 libtextclassifier3::StatusOr<Crc32> EmbeddingIndex::GetStoragesChecksum()
@@ -631,12 +804,17 @@
   }
   ICING_ASSIGN_OR_RETURN(Crc32 embedding_posting_list_mapper_crc,
                          embedding_posting_list_mapper_->GetChecksum());
-  Crc32 embedding_vectors_crc = embedding_vectors_->GetChecksum();
-  Crc32 quantized_embedding_vectors_crc =
-      quantized_embedding_vectors_->GetChecksum();
-  return Crc32(embedding_posting_list_mapper_crc.Get() ^
-               embedding_vectors_crc.Get() ^
-               quantized_embedding_vectors_crc.Get());
+  uint32_t checksum = embedding_posting_list_mapper_crc.Get();
+
+  for (int i = 0; i < num_shards_; ++i) {
+    if (embedding_vectors_[i] != nullptr) {
+      checksum ^= embedding_vectors_[i]->GetChecksum().Get();
+    }
+    if (quantized_embedding_vectors_[i] != nullptr) {
+      checksum ^= quantized_embedding_vectors_[i]->GetChecksum().Get();
+    }
+  }
+  return Crc32(checksum);
 }
 
 }  // namespace lib

diff --git a/icing/index/embed/embedding-index.h b/icing/index/embed/embedding-index.h
index a2dfb7a..680e2f7 100644
--- a/icing/index/embed/embedding-index.h
+++ b/icing/index/embed/embedding-index.h

@@ -25,6 +25,7 @@
 #include "icing/text_classifier/lib3/utils/base/status.h"
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "icing/absl_ports/canonical_errors.h"
+#include "icing/absl_ports/str_cat.h"
 #include "icing/feature-flags.h"
 #include "icing/file/file-backed-vector.h"
 #include "icing/file/filesystem.h"
@@ -38,6 +39,7 @@
 #include "icing/index/embed/posting-list-embedding-hit-serializer.h"
 #include "icing/index/embed/quantizer.h"
 #include "icing/index/hit/hit.h"
+#include "icing/index/iterator/doc-hit-info-iterator.h"
 #include "icing/schema/schema-store.h"
 #include "icing/store/document-id.h"
 #include "icing/store/document-store.h"
@@ -57,8 +59,9 @@
     int32_t magic;
     DocumentId last_added_document_id;
     bool is_empty;
+    uint32_t num_shards;
 
-    static constexpr int kPaddingSize = 1000;
+    static constexpr int kPaddingSize = 996;
     // Padding exists just to reserve space for additional values.
     uint8_t padding_[kPaddingSize];
 
@@ -92,7 +95,8 @@
   //     DynamicTrieKeyMapper, or FileBackedVector.
   static libtextclassifier3::StatusOr<std::unique_ptr<EmbeddingIndex>> Create(
       const Filesystem* filesystem, std::string working_path,
-      const Clock* clock, const FeatureFlags* feature_flags);
+      const Clock* clock, const FeatureFlags* feature_flags,
+      uint32_t num_shards);
 
   static libtextclassifier3::Status Discard(const Filesystem& filesystem,
                                             const std::string& working_path) {
@@ -121,7 +125,8 @@
   //   - INTERNAL_ERROR on I/O error
   libtextclassifier3::Status BufferEmbedding(
       const BasicHit& basic_hit, const PropertyProto::VectorProto& vector,
-      EmbeddingIndexingConfig::QuantizationType::Code quantization_type);
+      EmbeddingIndexingConfig::QuantizationType::Code quantization_type,
+      std::string_view schema_name);
 
   // Commit the embedding hits in the buffer to the index.
   //
@@ -131,26 +136,66 @@
   //   - Any error from posting lists
   libtextclassifier3::Status CommitBufferToIndex();
 
-  // Returns a PostingListEmbeddingHitAccessor for all embedding hits that match
+  class EmbeddingHitAccessor {
+   public:
+    EmbeddingHitAccessor(
+        const EmbeddingIndex* embedding_index,
+        std::unique_ptr<PostingListEmbeddingHitAccessor> pl_accessor,
+        std::string_view posting_list_key)
+        : embedding_index_(*embedding_index),
+          pl_accessor_(std::move(pl_accessor)),
+          posting_list_key_hash_(
+              EmbeddingIndex::GetPostingListKeyHash(posting_list_key)) {}
+
+    libtextclassifier3::StatusOr<std::vector<EmbeddingHit>> GetNextHitsBatch() {
+      return pl_accessor_->GetNextHitsBatch();
+    }
+
+    // Calculates the score for the given embedding hit with the given query.
+    //
+    // Returns:
+    //   - The score on success.
+    //   - OUT_OF_RANGE_ERROR if the referred vector is out of range based on
+    //     the location and dimension.
+    //   - Any error from schema store when getting the quantization type.
+    libtextclassifier3::StatusOr<float> ScoreEmbeddingHit(
+        const EmbeddingScorer& scorer, const PropertyProto::VectorProto& query,
+        const EmbeddingHit& hit,
+        EmbeddingIndexingConfig::QuantizationType::Code quantization_type,
+        uint32_t schema_name_hash);
+
+    const DocHitInfoIterator::CallStats::EmbeddingStats& GetEmbeddingStats()
+        const {
+      return embedding_stats_;
+    }
+
+   private:
+    const EmbeddingIndex& embedding_index_;
+    std::unique_ptr<PostingListEmbeddingHitAccessor> pl_accessor_;
+    const uint32_t posting_list_key_hash_;
+    DocHitInfoIterator::CallStats::EmbeddingStats embedding_stats_;
+  };
+
+  // Returns a EmbeddingHitAccessor for all embedding hits that match
   // with the provided dimension and signature.
   //
   // Returns:
-  //   - a PostingListEmbeddingHitAccessor instance on success.
+  //   - a EmbeddingHitAccessor instance on success.
   //   - INVALID_ARGUMENT error if the dimension is 0.
   //   - NOT_FOUND error if there is no matching embedding hit.
   //   - Any error from posting lists.
-  libtextclassifier3::StatusOr<std::unique_ptr<PostingListEmbeddingHitAccessor>>
+  libtextclassifier3::StatusOr<std::unique_ptr<EmbeddingHitAccessor>>
   GetAccessor(uint32_t dimension, std::string_view model_signature) const;
 
-  // Returns a PostingListEmbeddingHitAccessor for all embedding hits that match
+  // Returns a EmbeddingHitAccessor for all embedding hits that match
   // with the provided vector's dimension and signature.
   //
   // Returns:
-  //   - a PostingListEmbeddingHitAccessor instance on success.
+  //   - a EmbeddingHitAccessor instance on success.
   //   - INVALID_ARGUMENT error if the dimension is 0.
   //   - NOT_FOUND error if there is no matching embedding hit.
   //   - Any error from posting lists.
-  libtextclassifier3::StatusOr<std::unique_ptr<PostingListEmbeddingHitAccessor>>
+  libtextclassifier3::StatusOr<std::unique_ptr<EmbeddingHitAccessor>>
   GetAccessorForVector(const PropertyProto::VectorProto& vector) const {
     return GetAccessor(vector.values_size(), vector.model_signature());
   }
@@ -172,62 +217,69 @@
   //
   // Returns:
   //   - a pointer to the embedding vector on success.
+  //   - INVALID_ARGUMENT if the shard does not exist.
   //   - OUT_OF_RANGE error if the referred vector is out of range based on the
   //     location and dimension.
   libtextclassifier3::StatusOr<const float*> GetEmbeddingVector(
-      const EmbeddingHit& hit, uint32_t dimension) const {
+      const EmbeddingHit& hit, uint32_t dimension, uint32_t shard_id) const {
+    if (shard_id >= num_shards_ || embedding_vectors_[shard_id] == nullptr) {
+      return absl_ports::InvalidArgumentError(
+          "Attempting to query a non-existent storage shard.");
+    }
+    const auto& fbv = embedding_vectors_[shard_id];
     if (static_cast<int64_t>(hit.location()) + dimension >
-        GetTotalVectorSize()) {
+        fbv->num_elements()) {
       return absl_ports::OutOfRangeError(
           "Got an embedding hit that refers to a vector out of range.");
     }
-    return embedding_vectors_->array() + hit.location();
+    return fbv->array() + hit.location();
   }
   libtextclassifier3::StatusOr<const char*> GetQuantizedEmbeddingVector(
-      const EmbeddingHit& hit, uint32_t dimension) const {
+      const EmbeddingHit& hit, uint32_t dimension, uint32_t shard_id) const {
+    if (shard_id >= num_shards_ ||
+        quantized_embedding_vectors_[shard_id] == nullptr) {
+      return absl_ports::InvalidArgumentError(
+          "Attempting to query a non-existent storage shard.");
+    }
+    const auto& fbv = quantized_embedding_vectors_[shard_id];
     // quantized_embedding_vectors_ stores data in char format. Every quantized
     // embedding vector contains a Quantizer header followed by the actual
     // vector, and every value in the vector is stored in uint8_t.
     if (static_cast<int64_t>(hit.location()) + sizeof(Quantizer) +
             sizeof(uint8_t) * dimension >
-        GetTotalQuantizedVectorSize()) {
+        fbv->num_elements()) {
       return absl_ports::OutOfRangeError(
           "Got an embedding hit that refers to a vector out of range.");
     }
-    return quantized_embedding_vectors_->array() + hit.location();
+    return fbv->array() + hit.location();
   }
 
-  // Calculates the score for the given embedding hit with the given query.
-  //
-  // Returns:
-  //   - The score on success.
-  //   - OUT_OF_RANGE_ERROR if the referred vector is out of range based on the
-  //     location and dimension.
-  //   - Any error from schema store when getting the quantization type.
-  libtextclassifier3::StatusOr<float> ScoreEmbeddingHit(
-      const EmbeddingScorer& scorer, const PropertyProto::VectorProto& query,
-      const EmbeddingHit& hit,
-      EmbeddingIndexingConfig::QuantizationType::Code quantization_type) const;
-
-  libtextclassifier3::StatusOr<const float*> GetRawEmbeddingData() const {
+  libtextclassifier3::StatusOr<const float*> GetRawEmbeddingData(
+      uint32_t shard_id) const {
     if (is_empty()) {
       return absl_ports::NotFoundError("EmbeddingIndex is empty");
     }
-    return embedding_vectors_->array();
+    if (shard_id >= num_shards_ || embedding_vectors_[shard_id] == nullptr) {
+      return absl_ports::NotFoundError(absl_ports::StrCat(
+          "Embedding storage shard ", std::to_string(shard_id), " not found"));
+    }
+    return embedding_vectors_[shard_id]->array();
   }
 
-  int32_t GetTotalVectorSize() const {
-    if (is_empty()) {
+  int32_t GetTotalVectorSize(uint32_t shard_id) const {
+    if (is_empty() || shard_id >= num_shards_ ||
+        embedding_vectors_[shard_id] == nullptr) {
       return 0;
     }
-    return embedding_vectors_->num_elements();
+    return embedding_vectors_[shard_id]->num_elements();
   }
 
-  int32_t GetTotalQuantizedVectorSize() const {
-    if (is_empty()) {
+  int32_t GetTotalQuantizedVectorSize(uint32_t shard_id) const {
+    if (is_empty() || shard_id >= num_shards_ ||
+        quantized_embedding_vectors_[shard_id] == nullptr) {
       return 0;
     }
-    return quantized_embedding_vectors_->num_elements();
+    return quantized_embedding_vectors_[shard_id]->num_elements();
   }
 
   DocumentId last_added_document_id() const {
@@ -244,32 +296,63 @@
 
   bool is_empty() const { return info().is_empty; }
 
+  uint32_t GetShardId(uint32_t dimension, std::string_view model_signature,
+                      std::string_view schema_name) const;
+
+  uint32_t num_shards() const { return num_shards_; }
+
+  Info& info() {
+    return *reinterpret_cast<Info*>(metadata_mmapped_file_->mutable_region() +
+                                    kInfoMetadataBufferOffset);
+  }
+
+  const Info& info() const {
+    return *reinterpret_cast<const Info*>(metadata_mmapped_file_->region() +
+                                          kInfoMetadataBufferOffset);
+  }
+
  private:
   explicit EmbeddingIndex(const Filesystem& filesystem,
                           std::string working_path, const Clock* clock,
-                          const FeatureFlags* feature_flags)
+                          const FeatureFlags* feature_flags,
+                          uint32_t num_shards)
       : PersistentStorage(filesystem, std::move(working_path),
                           kWorkingPathType),
         clock_(*clock),
-        feature_flags_(feature_flags) {}
+        feature_flags_(feature_flags),
+        num_shards_(num_shards),
+        embedding_vectors_(num_shards),
+        quantized_embedding_vectors_(num_shards) {}
 
-  // Creates the storage data if the index is not empty. This will initialize
-  // flash_index_storage_, embedding_posting_list_mapper_, embedding_vectors_.
+  friend class EmbeddingHitAccessor;
+
+  static uint32_t GetPostingListKeyHash(std::string_view posting_list_key) {
+    return Crc32(posting_list_key).Get();
+  }
+
+  uint32_t GetShardId(uint32_t posting_list_key_hash,
+                      uint32_t schema_name_hash) const {
+    return (posting_list_key_hash * 31 + schema_name_hash) % num_shards_;
+  }
+
+  // Creates the storage data. This will initialize flash_index_storage_,
+  // embedding_posting_list_mapper_, and scan and initialize for existing vector
+  // storage files.
   //
   // Returns:
   //   - OK on success
   //   - Any error from FlashIndexStorage, DynamicTrieKeyMapper, or
   //     FileBackedVector.
-  libtextclassifier3::Status CreateStorageDataIfNonEmpty();
+  libtextclassifier3::Status CreateStorageData();
 
   // Marks the index's header to indicate that the index is non-empty.
   //
   // If the index is already marked as non-empty, this is a no-op. Otherwise,
-  // CreateStorageDataIfNonEmpty will be called to create the storage data.
+  // CreateStorageData will be called to create the storage data.
   //
   // Returns:
   //   - OK on success
-  //   - Any error when calling CreateStorageDataIfNonEmpty.
+  //   - Any error when calling CreateStorageData.
   libtextclassifier3::Status MarkIndexNonEmpty();
 
   libtextclassifier3::Status Initialize();
@@ -283,7 +366,7 @@
   libtextclassifier3::StatusOr<uint32_t> TransferEmbeddingVector(
       const EmbeddingHit& old_hit, uint32_t dimension,
       EmbeddingIndexingConfig::QuantizationType::Code quantization_type,
-      EmbeddingIndex* new_index) const;
+      uint32_t shard_id, EmbeddingIndex* new_index) const;
 
   // Transfers embedding data and hits from the current index to new_index.
   //
@@ -317,16 +400,17 @@
   libtextclassifier3::StatusOr<Crc32> GetStoragesChecksum() const override;
 
   // Appends the given embedding vector to the appropriate vector storage
-  // (embedding_vectors_ or quantized_embedding_vectors_) based on the
-  // quantization type.
+  // shard based on the quantization type and shard_id. If the storage shard
+  // does not exist, it will be created.
   //
   // Returns:
   //   - The location of the appended vector (i.e., the starting index within
-  //     the vector storage).
+  //     the vector storage shard).
   //   - Any error when allocating the vector storage.
   libtextclassifier3::StatusOr<uint32_t> AppendEmbeddingVector(
       const PropertyProto::VectorProto& vector,
-      EmbeddingIndexingConfig::QuantizationType::Code quantization_type);
+      EmbeddingIndexingConfig::QuantizationType::Code quantization_type,
+      uint32_t shard_id);
 
   Crcs& crcs() override {
     return *reinterpret_cast<Crcs*>(metadata_mmapped_file_->mutable_region() +
@@ -338,18 +422,15 @@
                                           kCrcsMetadataBufferOffset);
   }
 
-  Info& info() {
-    return *reinterpret_cast<Info*>(metadata_mmapped_file_->mutable_region() +
-                                    kInfoMetadataBufferOffset);
-  }
+  libtextclassifier3::StatusOr<FileBackedVector<float>*>
+  GetOrCreateEmbeddingVector(uint32_t shard_id);
 
-  const Info& info() const {
-    return *reinterpret_cast<const Info*>(metadata_mmapped_file_->region() +
-                                          kInfoMetadataBufferOffset);
-  }
+  libtextclassifier3::StatusOr<FileBackedVector<char>*>
+  GetOrCreateQuantizedEmbeddingVector(uint32_t shard_id);
 
   const Clock& clock_;
   const FeatureFlags* feature_flags_;  // Does not own.
+  const uint32_t num_shards_;
 
   // In memory data:
   // Pending embedding hits with their embedding keys used for
@@ -377,11 +458,13 @@
   std::unique_ptr<KeyMapper<PostingListIdentifier>>
       embedding_posting_list_mapper_;
 
-  // A single FileBackedVector that holds all embedding vectors.
+  // An array of FileBackedVectors that hold all embedding vectors, sharded by
+  // a hash.
   //
-  // null if the index is empty.
-  std::unique_ptr<FileBackedVector<float>> embedding_vectors_;
-  std::unique_ptr<FileBackedVector<char>> quantized_embedding_vectors_;
+  // An element is null if its corresponding file does not exist.
+  std::vector<std::unique_ptr<FileBackedVector<float>>> embedding_vectors_;
+  std::vector<std::unique_ptr<FileBackedVector<char>>>
+      quantized_embedding_vectors_;
 };
 
 }  // namespace lib

diff --git a/icing/index/embed/embedding-index_test.cc b/icing/index/embed/embedding-index_test.cc
index 12ba95e..80031ec 100644
--- a/icing/index/embed/embedding-index_test.cc
+++ b/icing/index/embed/embedding-index_test.cc

@@ -20,6 +20,7 @@
 #include <memory>
 #include <string>
 #include <string_view>
+#include <unordered_set>
 #include <utility>
 #include <vector>
 
@@ -28,6 +29,7 @@
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include "icing/absl_ports/canonical_errors.h"
+#include "icing/absl_ports/str_cat.h"
 #include "icing/document-builder.h"
 #include "icing/feature-flags.h"
 #include "icing/file/filesystem.h"
@@ -36,6 +38,7 @@
 #include "icing/index/embed/quantizer.h"
 #include "icing/index/hit/hit.h"
 #include "icing/legacy/index/icing-filesystem.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/proto/document.pb.h"
 #include "icing/schema-builder.h"
 #include "icing/schema/schema-store.h"
@@ -48,6 +51,7 @@
 #include "icing/testing/tmp-directory.h"
 #include "icing/util/clock.h"
 #include "icing/util/crc32.h"
+#include "icing/util/document-util.h"
 
 namespace icing {
 namespace lib {
@@ -59,12 +63,17 @@
 using ::testing::FloatNear;
 using ::testing::HasSubstr;
 using ::testing::IsEmpty;
+using ::testing::Not;
 using ::testing::Pointwise;
 using ::testing::Test;
 
 static constexpr SectionId kSectionIdQuantizedEmbedding = 2;
 static constexpr float kEpsQuantized = 0.01f;
 
+static constexpr uint32_t kDefaultDimension = 3;
+static const char kDefaultModelSignature[] = "model";
+static constexpr std::string_view kDefaultSchemaName = "type";
+
 class EmbeddingIndexTest : public Test {
  protected:
   void SetUp() override {
@@ -77,26 +86,35 @@
     filesystem_.CreateDirectoryRecursively(document_store_dir_.c_str());
     filesystem_.CreateDirectoryRecursively(schema_store_dir_.c_str());
 
+    test_vector1_ = CreateVector(kDefaultModelSignature, {0.1, 0.2, 0.3});
+    test_vector2_ = CreateVector(kDefaultModelSignature, {-0.1, -0.2, -0.3});
+    test_vector3_ = CreateVector(kDefaultModelSignature, {0.4, 0.5, 0.6});
+
     ICING_ASSERT_OK_AND_ASSIGN(
         schema_store_, SchemaStore::Create(&filesystem_, schema_store_dir_,
                                            &clock_, feature_flags_.get()));
 
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
-        DocumentStore::Create(&filesystem_, document_store_dir_, &clock_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, document_store_dir_, &clock_, schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     document_store_ = std::move(create_result.document_store);
 
     ICING_ASSERT_OK_AND_ASSIGN(
         embedding_index_,
         EmbeddingIndex::Create(&filesystem_, embedding_index_dir_, &clock_,
-                               feature_flags_.get()));
+                               feature_flags_.get(),
+                               /*num_shards=*/32));
 
     ICING_ASSERT_OK(schema_store_->SetSchema(
         SchemaBuilder()
@@ -122,12 +140,15 @@
                             .SetCardinality(CARDINALITY_OPTIONAL)))
             .Build(),
         /*ignore_errors_and_delete_documents=*/false));
-    ICING_ASSERT_OK(document_store_->Put(
-        DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build()));
-    ICING_ASSERT_OK(document_store_->Put(
-        DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build()));
-    ICING_ASSERT_OK(document_store_->Put(
-        DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build()));
+    ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+        DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
+    ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+        DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build())));
+    ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+        DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build())));
+
+    default_shard_id_ = embedding_index_->GetShardId(
+        kDefaultDimension, kDefaultModelSignature, kDefaultSchemaName);
   }
 
   void TearDown() override {
@@ -157,8 +178,35 @@
   std::unique_ptr<SchemaStore> schema_store_;
   std::unique_ptr<DocumentStore> document_store_;
   std::unique_ptr<EmbeddingIndex> embedding_index_;
+
+  PropertyProto::VectorProto test_vector1_;
+  PropertyProto::VectorProto test_vector2_;
+  PropertyProto::VectorProto test_vector3_;
+
+  uint32_t default_shard_id_;
 };
 
+TEST_F(EmbeddingIndexTest, GetShardId) {
+  // Hardcode some inputs to the GetShardId function, so that we can be aware of
+  // any changes to the hashing function.
+  EXPECT_EQ(embedding_index_->GetShardId(768, "model1", "schema1"), 10);
+  EXPECT_EQ(embedding_index_->GetShardId(768, "model1", "schema2"), 4);
+  EXPECT_EQ(embedding_index_->GetShardId(768, "model2", "schema1"), 20);
+  EXPECT_EQ(embedding_index_->GetShardId(768, "model2", "schema2"), 14);
+  EXPECT_EQ(embedding_index_->GetShardId(1024, "model1", "schema1"), 27);
+  EXPECT_EQ(embedding_index_->GetShardId(1024, "model1", "schema2"), 21);
+  EXPECT_EQ(embedding_index_->GetShardId(1024, "model2", "schema1"), 1);
+  EXPECT_EQ(embedding_index_->GetShardId(1024, "model2", "schema2"), 27);
+  EXPECT_EQ(embedding_index_->GetShardId(100, "aa", "bb"), 4);
+  EXPECT_EQ(embedding_index_->GetShardId(100, "bb", "aa"), 20);
+  EXPECT_EQ(embedding_index_->GetShardId(100, "aa", "aa"), 27);
+  EXPECT_EQ(embedding_index_->GetShardId(100, "bb", "bb"), 29);
+  EXPECT_EQ(embedding_index_->GetShardId(100, "aa", "aaa"), 18);
+  EXPECT_EQ(embedding_index_->GetShardId(100, "bb", "bbb"), 11);
+  EXPECT_EQ(embedding_index_->GetShardId(100, "aaa", "aa"), 4);
+  EXPECT_EQ(embedding_index_->GetShardId(100, "bbb", "bb"), 13);
+}
+
 TEST_F(EmbeddingIndexTest, EmptyIndexContainsMetadataOnly) {
   EXPECT_THAT(IndexContainsMetadataOnly(), IsOkAndHolds(true));
 }
@@ -168,11 +216,13 @@
       GetTestTempDir() + "/embedding_index_test_local";
 
   EXPECT_THAT(EmbeddingIndex::Create(nullptr, embedding_index_dir, &clock_,
-                                     feature_flags_.get()),
+                                     feature_flags_.get(),
+                                     /*num_shards=*/32),
               StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
 
   EXPECT_THAT(EmbeddingIndex::Create(&filesystem_, embedding_index_dir, nullptr,
-                                     feature_flags_.get()),
+                                     feature_flags_.get(),
+                                     /*num_shards=*/32),
               StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
 }
 
@@ -184,22 +234,23 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<EmbeddingIndex> embedding_index,
       EmbeddingIndex::Create(&filesystem_, embedding_index_dir, &clock_,
-                             feature_flags_.get()));
+                             feature_flags_.get(),
+                             /*num_shards=*/32));
 
-  PropertyProto::VectorProto vector = CreateVector("model", {0.1, 0.2, 0.3});
   ICING_ASSERT_OK(embedding_index->BufferEmbedding(
-      BasicHit(/*section_id=*/0, /*document_id=*/0), vector,
-      QUANTIZATION_TYPE_NONE));
+      BasicHit(/*section_id=*/0, /*document_id=*/0), test_vector1_,
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index->CommitBufferToIndex());
   embedding_index->set_last_added_document_id(0);
 
   EXPECT_THAT(
       GetEmbeddingHitsFromIndex(embedding_index.get(), /*dimension=*/3,
-                                /*model_signature=*/"model"),
+                                kDefaultModelSignature),
       IsOkAndHolds(ElementsAre(EmbeddingHit(
           BasicHit(/*section_id=*/0, /*document_id=*/0), /*location=*/0))));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index.get()),
-              ElementsAre(0.1, 0.2, 0.3));
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index.get(), default_shard_id_),
+      ElementsAre(0.1, 0.2, 0.3));
   EXPECT_EQ(embedding_index->last_added_document_id(), 0);
   // GetChecksum should succeed without updating the checksum.
   ICING_EXPECT_OK(embedding_index->GetChecksum());
@@ -207,13 +258,101 @@
   // 2. Try to create another index with the same directory. This should fail
   // due to checksum mismatch.
   EXPECT_THAT(EmbeddingIndex::Create(&filesystem_, embedding_index_dir, &clock_,
-                                     feature_flags_.get()),
+                                     feature_flags_.get(),
+                                     /*num_shards=*/32),
               StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
 
   embedding_index.reset();
   filesystem_.DeleteDirectoryRecursively(embedding_index_dir.c_str());
 }
 
+TEST_F(EmbeddingIndexTest, InitializationShouldFailWithZeroShards) {
+  std::string embedding_index_dir =
+      GetTestTempDir() + "/embedding_index_test_local";
+  EXPECT_THAT(EmbeddingIndex::Create(&filesystem_, embedding_index_dir, &clock_,
+                                     feature_flags_.get(),
+                                     /*num_shards=*/0),
+              StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
+}
+
+TEST_F(EmbeddingIndexTest, InitializationShouldFailWithMismatchedNumShards) {
+  std::string embedding_index_dir =
+      GetTestTempDir() + "/embedding_index_test_local";
+  // 1. Create an index with num_shards = 1.
+  {
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<EmbeddingIndex> embedding_index,
+        EmbeddingIndex::Create(&filesystem_, embedding_index_dir, &clock_,
+                               feature_flags_.get(), /*num_shards=*/1));
+    ICING_ASSERT_OK(embedding_index->PersistToDisk());
+  }
+
+  // 2. Try to create another index with a different num_shards. This should
+  // fail.
+  EXPECT_THAT(EmbeddingIndex::Create(&filesystem_, embedding_index_dir, &clock_,
+                                     feature_flags_.get(),
+                                     /*num_shards=*/32),
+              StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION,
+                       HasSubstr("Mismatched number of shards")));
+
+  filesystem_.DeleteDirectoryRecursively(embedding_index_dir.c_str());
+}
+
+TEST_F(EmbeddingIndexTest,
+       InitializationShouldSucceedWithNumShardsUpgradeFromZeroToOne) {
+  std::string embedding_index_dir =
+      GetTestTempDir() + "/embedding_index_test_local";
+
+  // 1. Create an index with num_shards = 1, and manually set num_shards to 0 in
+  // the header.
+  {
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<EmbeddingIndex> embedding_index,
+        EmbeddingIndex::Create(&filesystem_, embedding_index_dir, &clock_,
+                               feature_flags_.get(), /*num_shards=*/1));
+    embedding_index->info().num_shards = 0;
+    ICING_ASSERT_OK(embedding_index->PersistToDisk());
+  }
+
+  // 2. Re-initialize with num_shards = 1. It should succeed.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<EmbeddingIndex> embedding_index,
+      EmbeddingIndex::Create(&filesystem_, embedding_index_dir, &clock_,
+                             feature_flags_.get(), /*num_shards=*/1));
+
+  // 3. Check that num_shards in the header is now 1.
+  EXPECT_EQ(embedding_index->num_shards(), 1);
+  EXPECT_EQ(embedding_index->info().num_shards, 1);
+
+  filesystem_.DeleteDirectoryRecursively(embedding_index_dir.c_str());
+}
+
+TEST_F(EmbeddingIndexTest,
+       InitializationShouldFailWithNumShardsUpgradeFromZeroToThirtyTwo) {
+  std::string embedding_index_dir =
+      GetTestTempDir() + "/embedding_index_test_local";
+
+  // 1. Create an index with num_shards = 1, and manually set num_shards to 0 in
+  // the header.
+  {
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<EmbeddingIndex> embedding_index,
+        EmbeddingIndex::Create(&filesystem_, embedding_index_dir, &clock_,
+                               feature_flags_.get(), /*num_shards=*/1));
+    embedding_index->info().num_shards = 0;
+    ICING_ASSERT_OK(embedding_index->PersistToDisk());
+  }
+
+  // 2. Re-initialize with num_shards = 32. It should fail.
+  EXPECT_THAT(EmbeddingIndex::Create(&filesystem_, embedding_index_dir, &clock_,
+                                     feature_flags_.get(),
+                                     /*num_shards=*/32),
+              StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION,
+                       HasSubstr("Mismatched number of shards")));
+
+  filesystem_.DeleteDirectoryRecursively(embedding_index_dir.c_str());
+}
+
 TEST_F(EmbeddingIndexTest, InitializationShouldSucceedWithUpdateChecksums) {
   // 1. Create index and confirm that data was properly added.
   std::string embedding_index_dir =
@@ -221,22 +360,23 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<EmbeddingIndex> embedding_index,
       EmbeddingIndex::Create(&filesystem_, embedding_index_dir, &clock_,
-                             feature_flags_.get()));
+                             feature_flags_.get(),
+                             /*num_shards=*/32));
 
-  PropertyProto::VectorProto vector = CreateVector("model", {0.1, 0.2, 0.3});
   ICING_ASSERT_OK(embedding_index->BufferEmbedding(
-      BasicHit(/*section_id=*/0, /*document_id=*/0), vector,
-      QUANTIZATION_TYPE_NONE));
+      BasicHit(/*section_id=*/0, /*document_id=*/0), test_vector1_,
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index->CommitBufferToIndex());
   embedding_index->set_last_added_document_id(0);
 
   EXPECT_THAT(
       GetEmbeddingHitsFromIndex(embedding_index.get(), /*dimension=*/3,
-                                /*model_signature=*/"model"),
+                                kDefaultModelSignature),
       IsOkAndHolds(ElementsAre(EmbeddingHit(
           BasicHit(/*section_id=*/0, /*document_id=*/0), /*location=*/0))));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index.get()),
-              ElementsAre(0.1, 0.2, 0.3));
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index.get(), default_shard_id_),
+      ElementsAre(0.1, 0.2, 0.3));
   EXPECT_EQ(embedding_index->last_added_document_id(), 0);
 
   // 2. Update checksums to reflect the new content.
@@ -248,14 +388,16 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<EmbeddingIndex> embedding_index_two,
       EmbeddingIndex::Create(&filesystem_, embedding_index_dir, &clock_,
-                             feature_flags_.get()));
+                             feature_flags_.get(),
+                             /*num_shards=*/32));
 
   EXPECT_THAT(
       GetEmbeddingHitsFromIndex(embedding_index_two.get(), /*dimension=*/3,
-                                /*model_signature=*/"model"),
+                                kDefaultModelSignature),
       IsOkAndHolds(ElementsAre(EmbeddingHit(
           BasicHit(/*section_id=*/0, /*document_id=*/0), /*location=*/0))));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_two.get()),
+  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_two.get(),
+                                           default_shard_id_),
               ElementsAre(0.1, 0.2, 0.3));
   EXPECT_EQ(embedding_index_two->last_added_document_id(), 0);
 
@@ -271,22 +413,23 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<EmbeddingIndex> embedding_index,
       EmbeddingIndex::Create(&filesystem_, embedding_index_dir, &clock_,
-                             feature_flags_.get()));
+                             feature_flags_.get(),
+                             /*num_shards=*/32));
 
-  PropertyProto::VectorProto vector = CreateVector("model", {0.1, 0.2, 0.3});
   ICING_ASSERT_OK(embedding_index->BufferEmbedding(
-      BasicHit(/*section_id=*/0, /*document_id=*/0), vector,
-      QUANTIZATION_TYPE_NONE));
+      BasicHit(/*section_id=*/0, /*document_id=*/0), test_vector1_,
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index->CommitBufferToIndex());
   embedding_index->set_last_added_document_id(0);
 
   EXPECT_THAT(
       GetEmbeddingHitsFromIndex(embedding_index.get(), /*dimension=*/3,
-                                /*model_signature=*/"model"),
+                                kDefaultModelSignature),
       IsOkAndHolds(ElementsAre(EmbeddingHit(
           BasicHit(/*section_id=*/0, /*document_id=*/0), /*location=*/0))));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index.get()),
-              ElementsAre(0.1, 0.2, 0.3));
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index.get(), default_shard_id_),
+      ElementsAre(0.1, 0.2, 0.3));
   EXPECT_EQ(embedding_index->last_added_document_id(), 0);
 
   // 2. Update checksums to reflect the new content.
@@ -296,14 +439,16 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<EmbeddingIndex> embedding_index_two,
       EmbeddingIndex::Create(&filesystem_, embedding_index_dir, &clock_,
-                             feature_flags_.get()));
+                             feature_flags_.get(),
+                             /*num_shards=*/32));
 
   EXPECT_THAT(
       GetEmbeddingHitsFromIndex(embedding_index_two.get(), /*dimension=*/3,
-                                /*model_signature=*/"model"),
+                                kDefaultModelSignature),
       IsOkAndHolds(ElementsAre(EmbeddingHit(
           BasicHit(/*section_id=*/0, /*document_id=*/0), /*location=*/0))));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_two.get()),
+  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_two.get(),
+                                           default_shard_id_),
               ElementsAre(0.1, 0.2, 0.3));
   EXPECT_EQ(embedding_index_two->last_added_document_id(), 0);
 
@@ -314,114 +459,109 @@
 
 TEST_F(EmbeddingIndexTest, GetEmbeddingVectorShouldFailWhenOutOfRange) {
   BasicHit basic_hit(/*section_id=*/0, /*document_id=*/0);
-  PropertyProto::VectorProto vector = CreateVector("model", {0.1, 0.2, 0.3});
-  ICING_ASSERT_OK(embedding_index_->BufferEmbedding(basic_hit, vector,
-                                                    QUANTIZATION_TYPE_NONE));
+  ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
+      basic_hit, test_vector1_, QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
 
   EmbeddingHit embedding_hit(basic_hit, /*location=*/0);
   uint32_t dimension = 3;
-  ICING_ASSERT_OK(
-      embedding_index_->GetEmbeddingVector(embedding_hit, dimension));
-  EXPECT_THAT(
-      embedding_index_->GetEmbeddingVector(embedding_hit, dimension + 1),
-      StatusIs(libtextclassifier3::StatusCode::OUT_OF_RANGE));
+  ICING_ASSERT_OK(embedding_index_->GetEmbeddingVector(embedding_hit, dimension,
+                                                       default_shard_id_));
+  EXPECT_THAT(embedding_index_->GetEmbeddingVector(embedding_hit, dimension + 1,
+                                                   default_shard_id_),
+              StatusIs(libtextclassifier3::StatusCode::OUT_OF_RANGE));
 }
 
 TEST_F(EmbeddingIndexTest,
        GetQuantizedEmbeddingVectorShouldFailWhenOutOfRange) {
   BasicHit basic_hit(kSectionIdQuantizedEmbedding, /*document_id=*/0);
-  PropertyProto::VectorProto vector = CreateVector("model", {0.1, 0.2, 0.3});
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-      basic_hit, vector, QUANTIZATION_TYPE_QUANTIZE_8_BIT));
+      basic_hit, test_vector1_, QUANTIZATION_TYPE_QUANTIZE_8_BIT,
+      kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
 
   EmbeddingHit embedding_hit(basic_hit, /*location=*/0);
   uint32_t dimension = 3;
-  ICING_ASSERT_OK(
-      embedding_index_->GetQuantizedEmbeddingVector(embedding_hit, dimension));
-  EXPECT_THAT(embedding_index_->GetQuantizedEmbeddingVector(embedding_hit,
-                                                            dimension + 1),
+  ICING_ASSERT_OK(embedding_index_->GetQuantizedEmbeddingVector(
+      embedding_hit, dimension, default_shard_id_));
+  EXPECT_THAT(embedding_index_->GetQuantizedEmbeddingVector(
+                  embedding_hit, dimension + 1, default_shard_id_),
               StatusIs(libtextclassifier3::StatusCode::OUT_OF_RANGE));
 }
 
 TEST_F(EmbeddingIndexTest, AddSingleEmbedding) {
-  PropertyProto::VectorProto vector = CreateVector("model", {0.1, 0.2, 0.3});
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-      BasicHit(/*section_id=*/0, /*document_id=*/0), vector,
-      QUANTIZATION_TYPE_NONE));
+      BasicHit(/*section_id=*/0, /*document_id=*/0), test_vector1_,
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
   embedding_index_->set_last_added_document_id(0);
 
   EXPECT_THAT(
       GetEmbeddingHitsFromIndex(embedding_index_.get(), /*dimension=*/3,
-                                /*model_signature=*/"model"),
+                                kDefaultModelSignature),
       IsOkAndHolds(ElementsAre(EmbeddingHit(
           BasicHit(/*section_id=*/0, /*document_id=*/0), /*location=*/0))));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
-              ElementsAre(0.1, 0.2, 0.3));
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), default_shard_id_),
+      ElementsAre(0.1, 0.2, 0.3));
   EXPECT_EQ(embedding_index_->last_added_document_id(), 0);
 }
 
 TEST_F(EmbeddingIndexTest, AddSingleQuantizedEmbedding) {
-  PropertyProto::VectorProto vector = CreateVector("model", {0.1, 0.2, 0.3});
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-      BasicHit(kSectionIdQuantizedEmbedding, /*document_id=*/0), vector,
-      QUANTIZATION_TYPE_QUANTIZE_8_BIT));
+      BasicHit(kSectionIdQuantizedEmbedding, /*document_id=*/0), test_vector1_,
+      QUANTIZATION_TYPE_QUANTIZE_8_BIT, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
   embedding_index_->set_last_added_document_id(0);
 
   EmbeddingHit hit(BasicHit(kSectionIdQuantizedEmbedding, /*document_id=*/0),
                    /*location=*/0);
   EXPECT_THAT(GetEmbeddingHitsFromIndex(embedding_index_.get(), /*dimension=*/3,
-                                        /*model_signature=*/"model"),
+                                        kDefaultModelSignature),
               IsOkAndHolds(ElementsAre(hit)));
-  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(),
+  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(default_shard_id_),
               Eq(3 + sizeof(Quantizer)));
   EXPECT_THAT(
-      GetAndRestoreQuantizedEmbeddingVectorFromIndex(embedding_index_.get(),
-                                                     hit,
-                                                     /*dimension=*/3),
+      GetAndRestoreQuantizedEmbeddingVectorFromIndex(
+          embedding_index_.get(), hit,
+          /*dimension=*/3, kDefaultModelSignature, kDefaultSchemaName),
       IsOkAndHolds(Pointwise(FloatNear(kEpsQuantized), {0.1, 0.2, 0.3})));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()), IsEmpty());
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), default_shard_id_),
+      IsEmpty());
   EXPECT_EQ(embedding_index_->last_added_document_id(), 0);
 }
 
 TEST_F(EmbeddingIndexTest, AddMultipleEmbeddingsInTheSameSection) {
-  PropertyProto::VectorProto vector1 = CreateVector("model", {0.1, 0.2, 0.3});
-  PropertyProto::VectorProto vector2 =
-      CreateVector("model", {-0.1, -0.2, -0.3});
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-      BasicHit(/*section_id=*/0, /*document_id=*/0), vector1,
-      QUANTIZATION_TYPE_NONE));
+      BasicHit(/*section_id=*/0, /*document_id=*/0), test_vector1_,
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-      BasicHit(/*section_id=*/0, /*document_id=*/0), vector2,
-      QUANTIZATION_TYPE_NONE));
+      BasicHit(/*section_id=*/0, /*document_id=*/0), test_vector2_,
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
   embedding_index_->set_last_added_document_id(0);
 
   EXPECT_THAT(GetEmbeddingHitsFromIndex(embedding_index_.get(), /*dimension=*/3,
-                                        /*model_signature=*/"model"),
+                                        kDefaultModelSignature),
               IsOkAndHolds(ElementsAre(
                   EmbeddingHit(BasicHit(/*section_id=*/0, /*document_id=*/0),
                                /*location=*/0),
                   EmbeddingHit(BasicHit(/*section_id=*/0, /*document_id=*/0),
                                /*location=*/3))));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
-              ElementsAre(0.1, 0.2, 0.3, -0.1, -0.2, -0.3));
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), default_shard_id_),
+      ElementsAre(0.1, 0.2, 0.3, -0.1, -0.2, -0.3));
   EXPECT_EQ(embedding_index_->last_added_document_id(), 0);
 }
 
 TEST_F(EmbeddingIndexTest, AddMultipleQuantizedEmbeddingsInTheSameSection) {
-  PropertyProto::VectorProto vector1 = CreateVector("model", {0.1, 0.2, 0.3});
-  PropertyProto::VectorProto vector2 =
-      CreateVector("model", {-0.1, -0.2, -0.3});
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-      BasicHit(kSectionIdQuantizedEmbedding, /*document_id=*/0), vector1,
-      QUANTIZATION_TYPE_QUANTIZE_8_BIT));
+      BasicHit(kSectionIdQuantizedEmbedding, /*document_id=*/0), test_vector1_,
+      QUANTIZATION_TYPE_QUANTIZE_8_BIT, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-      BasicHit(kSectionIdQuantizedEmbedding, /*document_id=*/0), vector2,
-      QUANTIZATION_TYPE_QUANTIZE_8_BIT));
+      BasicHit(kSectionIdQuantizedEmbedding, /*document_id=*/0), test_vector2_,
+      QUANTIZATION_TYPE_QUANTIZE_8_BIT, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
   embedding_index_->set_last_added_document_id(0);
 
@@ -430,69 +570,69 @@
   EmbeddingHit hit2(BasicHit(kSectionIdQuantizedEmbedding, /*document_id=*/0),
                     /*location=*/3 + sizeof(Quantizer));
   EXPECT_THAT(GetEmbeddingHitsFromIndex(embedding_index_.get(), /*dimension=*/3,
-                                        /*model_signature=*/"model"),
+                                        kDefaultModelSignature),
               IsOkAndHolds(ElementsAre(hit1, hit2)));
-  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(),
+  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(default_shard_id_),
               Eq(2 * (3 + sizeof(Quantizer))));  // Two quantized vectors
   EXPECT_THAT(
-      GetAndRestoreQuantizedEmbeddingVectorFromIndex(embedding_index_.get(),
-                                                     hit1, /*dimension=*/3),
+      GetAndRestoreQuantizedEmbeddingVectorFromIndex(
+          embedding_index_.get(), hit1, /*dimension=*/3, kDefaultModelSignature,
+          kDefaultSchemaName),
       IsOkAndHolds(Pointwise(FloatNear(kEpsQuantized), {0.1, 0.2, 0.3})));
   EXPECT_THAT(
-      GetAndRestoreQuantizedEmbeddingVectorFromIndex(embedding_index_.get(),
-                                                     hit2, /*dimension=*/3),
+      GetAndRestoreQuantizedEmbeddingVectorFromIndex(
+          embedding_index_.get(), hit2, /*dimension=*/3, kDefaultModelSignature,
+          kDefaultSchemaName),
       IsOkAndHolds(Pointwise(FloatNear(kEpsQuantized), {-0.1, -0.2, -0.3})));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()), IsEmpty());
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), default_shard_id_),
+      IsEmpty());
   EXPECT_EQ(embedding_index_->last_added_document_id(), 0);
 }
 
 TEST_F(EmbeddingIndexTest, HitsWithLowerSectionIdReturnedFirst) {
-  PropertyProto::VectorProto vector1 = CreateVector("model", {0.1, 0.2, 0.3});
-  PropertyProto::VectorProto vector2 =
-      CreateVector("model", {-0.1, -0.2, -0.3});
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-      BasicHit(/*section_id=*/5, /*document_id=*/0), vector1,
-      QUANTIZATION_TYPE_NONE));
+      BasicHit(/*section_id=*/5, /*document_id=*/0), test_vector1_,
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-      BasicHit(/*section_id=*/2, /*document_id=*/0), vector2,
-      QUANTIZATION_TYPE_NONE));
+      BasicHit(/*section_id=*/2, /*document_id=*/0), test_vector2_,
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
   embedding_index_->set_last_added_document_id(0);
 
   EXPECT_THAT(GetEmbeddingHitsFromIndex(embedding_index_.get(), /*dimension=*/3,
-                                        /*model_signature=*/"model"),
+                                        kDefaultModelSignature),
               IsOkAndHolds(ElementsAre(
                   EmbeddingHit(BasicHit(/*section_id=*/2, /*document_id=*/0),
                                /*location=*/3),
                   EmbeddingHit(BasicHit(/*section_id=*/5, /*document_id=*/0),
                                /*location=*/0))));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
-              ElementsAre(0.1, 0.2, 0.3, -0.1, -0.2, -0.3));
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), default_shard_id_),
+      ElementsAre(0.1, 0.2, 0.3, -0.1, -0.2, -0.3));
   EXPECT_EQ(embedding_index_->last_added_document_id(), 0);
 }
 
 TEST_F(EmbeddingIndexTest, HitsWithHigherDocumentIdReturnedFirst) {
-  PropertyProto::VectorProto vector1 = CreateVector("model", {0.1, 0.2, 0.3});
-  PropertyProto::VectorProto vector2 =
-      CreateVector("model", {-0.1, -0.2, -0.3});
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-      BasicHit(/*section_id=*/0, /*document_id=*/0), vector1,
-      QUANTIZATION_TYPE_NONE));
+      BasicHit(/*section_id=*/0, /*document_id=*/0), test_vector1_,
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-      BasicHit(/*section_id=*/0, /*document_id=*/1), vector2,
-      QUANTIZATION_TYPE_NONE));
+      BasicHit(/*section_id=*/0, /*document_id=*/1), test_vector2_,
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
   embedding_index_->set_last_added_document_id(1);
 
   EXPECT_THAT(GetEmbeddingHitsFromIndex(embedding_index_.get(), /*dimension=*/3,
-                                        /*model_signature=*/"model"),
+                                        kDefaultModelSignature),
               IsOkAndHolds(ElementsAre(
                   EmbeddingHit(BasicHit(/*section_id=*/0, /*document_id=*/1),
                                /*location=*/3),
                   EmbeddingHit(BasicHit(/*section_id=*/0, /*document_id=*/0),
                                /*location=*/0))));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
-              ElementsAre(0.1, 0.2, 0.3, -0.1, -0.2, -0.3));
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), default_shard_id_),
+      ElementsAre(0.1, 0.2, 0.3, -0.1, -0.2, -0.3));
   EXPECT_EQ(embedding_index_->last_added_document_id(), 1);
 }
 
@@ -502,10 +642,10 @@
       CreateVector("model2", {-0.1, -0.2, -0.3});
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(/*section_id=*/0, /*document_id=*/0), vector1,
-      QUANTIZATION_TYPE_NONE));
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(/*section_id=*/0, /*document_id=*/0), vector2,
-      QUANTIZATION_TYPE_NONE));
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
   embedding_index_->set_last_added_document_id(0);
 
@@ -518,42 +658,69 @@
                                         /*model_signature=*/"model2"),
               IsOkAndHolds(ElementsAre(
                   EmbeddingHit(BasicHit(/*section_id=*/0, /*document_id=*/0),
-                               /*location=*/2))));
+                               /*location=*/0))));
   EXPECT_THAT(GetEmbeddingHitsFromIndex(
                   embedding_index_.get(),
                   /*dimension=*/5, /*model_signature=*/"non-existent-model"),
               IsOkAndHolds(IsEmpty()));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
-              ElementsAre(0.1, 0.2, -0.1, -0.2, -0.3));
+  // Check the shard for vector1.
+  EXPECT_THAT(GetRawEmbeddingDataFromIndex(
+                  embedding_index_.get(),
+                  embedding_index_->GetShardId(/*dimension=*/2,
+                                               /*model_signature=*/"model1",
+                                               kDefaultSchemaName)),
+              ElementsAre(0.1, 0.2));
+  // Check the shard for vector2.
+  EXPECT_THAT(GetRawEmbeddingDataFromIndex(
+                  embedding_index_.get(),
+                  embedding_index_->GetShardId(/*dimension=*/3,
+                                               /*model_signature=*/"model2",
+                                               kDefaultSchemaName)),
+              ElementsAre(-0.1, -0.2, -0.3));
   EXPECT_EQ(embedding_index_->last_added_document_id(), 0);
 }
 
 TEST_F(EmbeddingIndexTest,
        AddEmbeddingsWithSameSignatureButDifferentDimension) {
-  PropertyProto::VectorProto vector1 = CreateVector("model", {0.1, 0.2});
+  PropertyProto::VectorProto vector1 =
+      CreateVector(kDefaultModelSignature, {0.1, 0.2});
   PropertyProto::VectorProto vector2 =
-      CreateVector("model", {-0.1, -0.2, -0.3});
+      CreateVector(kDefaultModelSignature, {-0.1, -0.2, -0.3});
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(/*section_id=*/0, /*document_id=*/0), vector1,
-      QUANTIZATION_TYPE_NONE));
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(/*section_id=*/0, /*document_id=*/0), vector2,
-      QUANTIZATION_TYPE_NONE));
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
   embedding_index_->set_last_added_document_id(0);
 
   EXPECT_THAT(GetEmbeddingHitsFromIndex(embedding_index_.get(), /*dimension=*/2,
-                                        /*model_signature=*/"model"),
+                                        kDefaultModelSignature),
               IsOkAndHolds(ElementsAre(
                   EmbeddingHit(BasicHit(/*section_id=*/0, /*document_id=*/0),
                                /*location=*/0))));
   EXPECT_THAT(GetEmbeddingHitsFromIndex(embedding_index_.get(), /*dimension=*/3,
-                                        /*model_signature=*/"model"),
+                                        kDefaultModelSignature),
               IsOkAndHolds(ElementsAre(
                   EmbeddingHit(BasicHit(/*section_id=*/0, /*document_id=*/0),
-                               /*location=*/2))));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
-              ElementsAre(0.1, 0.2, -0.1, -0.2, -0.3));
+                               /*location=*/0))));
+  // Check the shard for vector1.
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(
+          embedding_index_.get(),
+          embedding_index_->GetShardId(
+              /*dimension=*/2,
+              /*model_signature=*/kDefaultModelSignature, kDefaultSchemaName)),
+      ElementsAre(0.1, 0.2));
+  // Check the shard for vector2.
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(
+          embedding_index_.get(),
+          embedding_index_->GetShardId(
+              /*dimension=*/3,
+              /*model_signature=*/kDefaultModelSignature, kDefaultSchemaName)),
+      ElementsAre(-0.1, -0.2, -0.3));
   EXPECT_EQ(embedding_index_->last_added_document_id(), 0);
 }
 
@@ -561,20 +728,15 @@
   // Loop the same logic twice to make sure that clear works as expected, and
   // the index is still valid after clearing.
   for (int i = 0; i < 2; i++) {
-    PropertyProto::VectorProto vector1 = CreateVector("model", {0.1, 0.2, 0.3});
-    PropertyProto::VectorProto vector2 =
-        CreateVector("model", {-0.1, -0.2, -0.3});
-    PropertyProto::VectorProto vector3 = CreateVector("model", {0.4, 0.5, 0.6});
-
     ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-        BasicHit(/*section_id=*/1, /*document_id=*/0), vector1,
-        QUANTIZATION_TYPE_NONE));
+        BasicHit(/*section_id=*/1, /*document_id=*/0), test_vector1_,
+        QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
     ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-        BasicHit(/*section_id=*/2, /*document_id=*/1), vector2,
-        QUANTIZATION_TYPE_NONE));
+        BasicHit(/*section_id=*/2, /*document_id=*/1), test_vector2_,
+        QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
     ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-        BasicHit(kSectionIdQuantizedEmbedding, /*document_id=*/2), vector3,
-        QUANTIZATION_TYPE_QUANTIZE_8_BIT));
+        BasicHit(kSectionIdQuantizedEmbedding, /*document_id=*/2),
+        test_vector3_, QUANTIZATION_TYPE_QUANTIZE_8_BIT, kDefaultSchemaName));
     ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
     embedding_index_->set_last_added_document_id(2);
 
@@ -587,17 +749,20 @@
 
     EXPECT_THAT(
         GetEmbeddingHitsFromIndex(embedding_index_.get(), /*dimension=*/3,
-                                  /*model_signature=*/"model"),
+                                  kDefaultModelSignature),
         IsOkAndHolds(ElementsAre(hit1, hit2, hit3)));
-    EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
-                ElementsAre(0.1, 0.2, 0.3, -0.1, -0.2, -0.3));
-    EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(),
-                Eq(3 + sizeof(Quantizer)));
     EXPECT_THAT(
-        GetAndRestoreQuantizedEmbeddingVectorFromIndex(embedding_index_.get(),
-                                                       hit1,
-                                                       /*dimension=*/3),
-        IsOkAndHolds(Pointwise(FloatNear(kEpsQuantized), vector3.values())));
+        GetRawEmbeddingDataFromIndex(embedding_index_.get(), default_shard_id_),
+        ElementsAre(0.1, 0.2, 0.3, -0.1, -0.2, -0.3));
+    EXPECT_THAT(
+        embedding_index_->GetTotalQuantizedVectorSize(default_shard_id_),
+        Eq(3 + sizeof(Quantizer)));
+    EXPECT_THAT(
+        GetAndRestoreQuantizedEmbeddingVectorFromIndex(
+            embedding_index_.get(), hit1,
+            /*dimension=*/3, kDefaultModelSignature, kDefaultSchemaName),
+        IsOkAndHolds(
+            Pointwise(FloatNear(kEpsQuantized), test_vector3_.values())));
     EXPECT_EQ(embedding_index_->last_added_document_id(), 2);
     EXPECT_FALSE(embedding_index_->is_empty());
     EXPECT_THAT(IndexContainsMetadataOnly(), IsOkAndHolds(false));
@@ -606,9 +771,12 @@
     ICING_ASSERT_OK(embedding_index_->Clear());
     EXPECT_TRUE(embedding_index_->is_empty());
     EXPECT_THAT(IndexContainsMetadataOnly(), IsOkAndHolds(true));
-    EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
-                IsEmpty());
-    EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(), Eq(0));
+    EXPECT_THAT(
+        GetRawEmbeddingDataFromIndex(embedding_index_.get(), default_shard_id_),
+        IsEmpty());
+    EXPECT_THAT(
+        embedding_index_->GetTotalQuantizedVectorSize(default_shard_id_),
+        Eq(0));
     EXPECT_EQ(embedding_index_->last_added_document_id(), kInvalidDocumentId);
   }
 }
@@ -617,20 +785,15 @@
   // Loop the same logic twice to make sure that Discard works as expected, and
   // the index is still valid after discarding.
   for (int i = 0; i < 2; i++) {
-    PropertyProto::VectorProto vector1 = CreateVector("model", {0.1, 0.2, 0.3});
-    PropertyProto::VectorProto vector2 =
-        CreateVector("model", {-0.1, -0.2, -0.3});
-    PropertyProto::VectorProto vector3 = CreateVector("model", {0.4, 0.5, 0.6});
-
     ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-        BasicHit(/*section_id=*/1, /*document_id=*/0), vector1,
-        QUANTIZATION_TYPE_NONE));
+        BasicHit(/*section_id=*/1, /*document_id=*/0), test_vector1_,
+        QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
     ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-        BasicHit(/*section_id=*/2, /*document_id=*/1), vector2,
-        QUANTIZATION_TYPE_NONE));
+        BasicHit(/*section_id=*/2, /*document_id=*/1), test_vector2_,
+        QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
     ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-        BasicHit(kSectionIdQuantizedEmbedding, /*document_id=*/2), vector3,
-        QUANTIZATION_TYPE_QUANTIZE_8_BIT));
+        BasicHit(kSectionIdQuantizedEmbedding, /*document_id=*/2),
+        test_vector3_, QUANTIZATION_TYPE_QUANTIZE_8_BIT, kDefaultSchemaName));
     ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
     embedding_index_->set_last_added_document_id(2);
 
@@ -642,17 +805,20 @@
                       /*location=*/0);
     EXPECT_THAT(
         GetEmbeddingHitsFromIndex(embedding_index_.get(), /*dimension=*/3,
-                                  /*model_signature=*/"model"),
+                                  kDefaultModelSignature),
         IsOkAndHolds(ElementsAre(hit1, hit2, hit3)));
-    EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
-                ElementsAre(0.1, 0.2, 0.3, -0.1, -0.2, -0.3));
-    EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(),
-                Eq(3 + sizeof(Quantizer)));
     EXPECT_THAT(
-        GetAndRestoreQuantizedEmbeddingVectorFromIndex(embedding_index_.get(),
-                                                       hit1,
-                                                       /*dimension=*/3),
-        IsOkAndHolds(Pointwise(FloatNear(kEpsQuantized), vector3.values())));
+        GetRawEmbeddingDataFromIndex(embedding_index_.get(), default_shard_id_),
+        ElementsAre(0.1, 0.2, 0.3, -0.1, -0.2, -0.3));
+    EXPECT_THAT(
+        embedding_index_->GetTotalQuantizedVectorSize(default_shard_id_),
+        Eq(3 + sizeof(Quantizer)));
+    EXPECT_THAT(
+        GetAndRestoreQuantizedEmbeddingVectorFromIndex(
+            embedding_index_.get(), hit1,
+            /*dimension=*/3, kDefaultModelSignature, kDefaultSchemaName),
+        IsOkAndHolds(
+            Pointwise(FloatNear(kEpsQuantized), test_vector3_.values())));
     EXPECT_EQ(embedding_index_->last_added_document_id(), 2);
     EXPECT_FALSE(embedding_index_->is_empty());
     EXPECT_THAT(IndexContainsMetadataOnly(), IsOkAndHolds(false));
@@ -663,12 +829,16 @@
     ICING_ASSERT_OK_AND_ASSIGN(
         embedding_index_,
         EmbeddingIndex::Create(&filesystem_, embedding_index_dir_, &clock_,
-                               feature_flags_.get()));
+                               feature_flags_.get(),
+                               /*num_shards=*/32));
     EXPECT_TRUE(embedding_index_->is_empty());
     EXPECT_THAT(IndexContainsMetadataOnly(), IsOkAndHolds(true));
-    EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
-                IsEmpty());
-    EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(), Eq(0));
+    EXPECT_THAT(
+        GetRawEmbeddingDataFromIndex(embedding_index_.get(), default_shard_id_),
+        IsEmpty());
+    EXPECT_THAT(
+        embedding_index_->GetTotalQuantizedVectorSize(default_shard_id_),
+        Eq(0));
     EXPECT_EQ(embedding_index_->last_added_document_id(), kInvalidDocumentId);
   }
 }
@@ -677,68 +847,60 @@
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
   EXPECT_TRUE(embedding_index_->is_empty());
   EXPECT_THAT(IndexContainsMetadataOnly(), IsOkAndHolds(true));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()), IsEmpty());
-  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(), Eq(0));
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), default_shard_id_),
+      IsEmpty());
+  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(default_shard_id_),
+              Eq(0));
 }
 
 TEST_F(EmbeddingIndexTest, MultipleCommits) {
-  PropertyProto::VectorProto vector1 = CreateVector("model", {0.1, 0.2, 0.3});
-  PropertyProto::VectorProto vector2 =
-      CreateVector("model", {-0.1, -0.2, -0.3});
-
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-      BasicHit(/*section_id=*/1, /*document_id=*/0), vector1,
-      QUANTIZATION_TYPE_NONE));
+      BasicHit(/*section_id=*/1, /*document_id=*/0), test_vector1_,
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
 
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-      BasicHit(/*section_id=*/0, /*document_id=*/0), vector2,
-      QUANTIZATION_TYPE_NONE));
+      BasicHit(/*section_id=*/0, /*document_id=*/0), test_vector2_,
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
 
   EXPECT_THAT(GetEmbeddingHitsFromIndex(embedding_index_.get(), /*dimension=*/3,
-                                        /*model_signature=*/"model"),
+                                        kDefaultModelSignature),
               IsOkAndHolds(ElementsAre(
                   EmbeddingHit(BasicHit(/*section_id=*/0, /*document_id=*/0),
                                /*location=*/3),
                   EmbeddingHit(BasicHit(/*section_id=*/1, /*document_id=*/0),
                                /*location=*/0))));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
-              ElementsAre(0.1, 0.2, 0.3, -0.1, -0.2, -0.3));
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), default_shard_id_),
+      ElementsAre(0.1, 0.2, 0.3, -0.1, -0.2, -0.3));
 }
 
 TEST_F(EmbeddingIndexTest,
        InvalidCommit_SectionIdCanOnlyDecreaseForSingleDocument) {
-  PropertyProto::VectorProto vector1 = CreateVector("model", {0.1, 0.2, 0.3});
-  PropertyProto::VectorProto vector2 =
-      CreateVector("model", {-0.1, -0.2, -0.3});
-
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-      BasicHit(/*section_id=*/0, /*document_id=*/0), vector1,
-      QUANTIZATION_TYPE_NONE));
+      BasicHit(/*section_id=*/0, /*document_id=*/0), test_vector1_,
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
 
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-      BasicHit(/*section_id=*/1, /*document_id=*/0), vector2,
-      QUANTIZATION_TYPE_NONE));
+      BasicHit(/*section_id=*/1, /*document_id=*/0), test_vector2_,
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   // Posting list with delta encoding can only allow decreasing values.
   EXPECT_THAT(embedding_index_->CommitBufferToIndex(),
               StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
 }
 
 TEST_F(EmbeddingIndexTest, InvalidCommit_DocumentIdCanOnlyIncrease) {
-  PropertyProto::VectorProto vector1 = CreateVector("model", {0.1, 0.2, 0.3});
-  PropertyProto::VectorProto vector2 =
-      CreateVector("model", {-0.1, -0.2, -0.3});
-
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-      BasicHit(/*section_id=*/0, /*document_id=*/1), vector1,
-      QUANTIZATION_TYPE_NONE));
+      BasicHit(/*section_id=*/0, /*document_id=*/1), test_vector1_,
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
 
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-      BasicHit(/*section_id=*/0, /*document_id=*/0), vector2,
-      QUANTIZATION_TYPE_NONE));
+      BasicHit(/*section_id=*/0, /*document_id=*/0), test_vector2_,
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   // Posting list with delta encoding can only allow decreasing values, which
   // means document ids must be committed increasingly, since document ids are
   // inverted in hit values.
@@ -761,10 +923,9 @@
 }
 
 TEST_F(EmbeddingIndexTest, OptimizeShouldFailWhenDocumentIdMapIsTooSmall) {
-  PropertyProto::VectorProto vector = CreateVector("model", {0.1, 0.2, 0.3});
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-      BasicHit(/*section_id=*/0, /*document_id=*/2), vector,
-      QUANTIZATION_TYPE_NONE));
+      BasicHit(/*section_id=*/0, /*document_id=*/2), test_vector1_,
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
   embedding_index_->set_last_added_document_id(2);
 
@@ -785,26 +946,29 @@
       /*new_last_added_document_id=*/kInvalidDocumentId));
   EXPECT_TRUE(embedding_index_->is_empty());
   EXPECT_THAT(IndexContainsMetadataOnly(), IsOkAndHolds(true));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()), IsEmpty());
-  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(), Eq(0));
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), default_shard_id_),
+      IsEmpty());
+  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(default_shard_id_),
+              Eq(0));
 }
 
 TEST_F(EmbeddingIndexTest, OptimizeSingleEmbeddingSingleDocument) {
-  PropertyProto::VectorProto vector = CreateVector("model", {0.1, 0.2, 0.3});
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-      BasicHit(/*section_id=*/0, /*document_id=*/2), vector,
-      QUANTIZATION_TYPE_NONE));
+      BasicHit(/*section_id=*/0, /*document_id=*/2), test_vector1_,
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
   embedding_index_->set_last_added_document_id(2);
 
   // Before optimize
   EXPECT_THAT(
       GetEmbeddingHitsFromIndex(embedding_index_.get(), /*dimension=*/3,
-                                /*model_signature=*/"model"),
+                                kDefaultModelSignature),
       IsOkAndHolds(ElementsAre(EmbeddingHit(
           BasicHit(/*section_id=*/0, /*document_id=*/2), /*location=*/0))));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
-              ElementsAre(0.1, 0.2, 0.3));
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), default_shard_id_),
+      ElementsAre(0.1, 0.2, 0.3));
   EXPECT_EQ(embedding_index_->last_added_document_id(), 2);
 
   // Run optimize without deleting any documents, and check that the index is
@@ -815,11 +979,12 @@
                                  /*new_last_added_document_id=*/2));
   EXPECT_THAT(
       GetEmbeddingHitsFromIndex(embedding_index_.get(), /*dimension=*/3,
-                                /*model_signature=*/"model"),
+                                kDefaultModelSignature),
       IsOkAndHolds(ElementsAre(EmbeddingHit(
           BasicHit(/*section_id=*/0, /*document_id=*/2), /*location=*/0))));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
-              ElementsAre(0.1, 0.2, 0.3));
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), default_shard_id_),
+      ElementsAre(0.1, 0.2, 0.3));
   EXPECT_EQ(embedding_index_->last_added_document_id(), 2);
 
   // Run optimize to map document id 2 to 1, and check that the index is
@@ -830,11 +995,12 @@
       /*new_last_added_document_id=*/1));
   EXPECT_THAT(
       GetEmbeddingHitsFromIndex(embedding_index_.get(), /*dimension=*/3,
-                                /*model_signature=*/"model"),
+                                kDefaultModelSignature),
       IsOkAndHolds(ElementsAre(EmbeddingHit(
           BasicHit(/*section_id=*/0, /*document_id=*/1), /*location=*/0))));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
-              ElementsAre(0.1, 0.2, 0.3));
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), default_shard_id_),
+      ElementsAre(0.1, 0.2, 0.3));
   EXPECT_EQ(embedding_index_->last_added_document_id(), 1);
 
   // Run optimize to delete the document.
@@ -843,19 +1009,20 @@
       /*document_id_old_to_new=*/{0, kInvalidDocumentId},
       /*new_last_added_document_id=*/0));
   EXPECT_THAT(GetEmbeddingHitsFromIndex(embedding_index_.get(), /*dimension=*/3,
-                                        /*model_signature=*/"model"),
+                                        kDefaultModelSignature),
               IsOkAndHolds(IsEmpty()));
   EXPECT_TRUE(embedding_index_->is_empty());
   EXPECT_THAT(IndexContainsMetadataOnly(), IsOkAndHolds(true));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()), IsEmpty());
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), default_shard_id_),
+      IsEmpty());
   EXPECT_EQ(embedding_index_->last_added_document_id(), 0);
 }
 
 TEST_F(EmbeddingIndexTest, OptimizeSingleQuantizedEmbeddingSingleDocument) {
-  PropertyProto::VectorProto vector = CreateVector("model", {0.1, 0.2, 0.3});
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-      BasicHit(kSectionIdQuantizedEmbedding, /*document_id=*/2), vector,
-      QUANTIZATION_TYPE_QUANTIZE_8_BIT));
+      BasicHit(kSectionIdQuantizedEmbedding, /*document_id=*/2), test_vector1_,
+      QUANTIZATION_TYPE_QUANTIZE_8_BIT, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
   embedding_index_->set_last_added_document_id(2);
 
@@ -863,15 +1030,18 @@
   EmbeddingHit hit(BasicHit(kSectionIdQuantizedEmbedding, /*document_id=*/2),
                    /*location=*/0);
   EXPECT_THAT(GetEmbeddingHitsFromIndex(embedding_index_.get(), /*dimension=*/3,
-                                        /*model_signature=*/"model"),
+                                        kDefaultModelSignature),
               IsOkAndHolds(ElementsAre(hit)));
-  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(),
+  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(default_shard_id_),
               Eq(3 + sizeof(Quantizer)));
   EXPECT_THAT(
-      GetAndRestoreQuantizedEmbeddingVectorFromIndex(embedding_index_.get(),
-                                                     hit, /*dimension=*/3),
+      GetAndRestoreQuantizedEmbeddingVectorFromIndex(
+          embedding_index_.get(), hit, /*dimension=*/3, kDefaultModelSignature,
+          kDefaultSchemaName),
       IsOkAndHolds(Pointwise(FloatNear(kEpsQuantized), {0.1, 0.2, 0.3})));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()), IsEmpty());
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), default_shard_id_),
+      IsEmpty());
   EXPECT_EQ(embedding_index_->last_added_document_id(), 2);
 
   // Run optimize without deleting any documents, and check that the index is
@@ -881,15 +1051,18 @@
                                  /*document_id_old_to_new=*/{0, 1, 2},
                                  /*new_last_added_document_id=*/2));
   EXPECT_THAT(GetEmbeddingHitsFromIndex(embedding_index_.get(), /*dimension=*/3,
-                                        /*model_signature=*/"model"),
+                                        kDefaultModelSignature),
               IsOkAndHolds(ElementsAre(hit)));
-  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(),
+  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(default_shard_id_),
               Eq(3 + sizeof(Quantizer)));
   EXPECT_THAT(
-      GetAndRestoreQuantizedEmbeddingVectorFromIndex(embedding_index_.get(),
-                                                     hit, /*dimension=*/3),
+      GetAndRestoreQuantizedEmbeddingVectorFromIndex(
+          embedding_index_.get(), hit, /*dimension=*/3, kDefaultModelSignature,
+          kDefaultSchemaName),
       IsOkAndHolds(Pointwise(FloatNear(kEpsQuantized), {0.1, 0.2, 0.3})));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()), IsEmpty());
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), default_shard_id_),
+      IsEmpty());
   EXPECT_EQ(embedding_index_->last_added_document_id(), 2);
 
   // Run optimize to map document id 2 to 1, and check that the index is
@@ -901,15 +1074,18 @@
   hit = EmbeddingHit(BasicHit(kSectionIdQuantizedEmbedding, /*document_id=*/1),
                      /*location=*/0);
   EXPECT_THAT(GetEmbeddingHitsFromIndex(embedding_index_.get(), /*dimension=*/3,
-                                        /*model_signature=*/"model"),
+                                        kDefaultModelSignature),
               IsOkAndHolds(ElementsAre(hit)));
-  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(),
+  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(default_shard_id_),
               Eq(3 + sizeof(Quantizer)));
   EXPECT_THAT(
-      GetAndRestoreQuantizedEmbeddingVectorFromIndex(embedding_index_.get(),
-                                                     hit, /*dimension=*/3),
+      GetAndRestoreQuantizedEmbeddingVectorFromIndex(
+          embedding_index_.get(), hit, /*dimension=*/3, kDefaultModelSignature,
+          kDefaultSchemaName),
       IsOkAndHolds(Pointwise(FloatNear(kEpsQuantized), {0.1, 0.2, 0.3})));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()), IsEmpty());
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), default_shard_id_),
+      IsEmpty());
   EXPECT_EQ(embedding_index_->last_added_document_id(), 1);
 
   // Run optimize to delete the document
@@ -918,30 +1094,28 @@
       /*document_id_old_to_new=*/{0, kInvalidDocumentId},
       /*new_last_added_document_id=*/0));
   EXPECT_THAT(GetEmbeddingHitsFromIndex(embedding_index_.get(), /*dimension=*/3,
-                                        /*model_signature=*/"model"),
+                                        kDefaultModelSignature),
               IsOkAndHolds(IsEmpty()));
   EXPECT_TRUE(embedding_index_->is_empty());
   EXPECT_THAT(IndexContainsMetadataOnly(), IsOkAndHolds(true));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()), IsEmpty());
-  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(), Eq(0));
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), default_shard_id_),
+      IsEmpty());
+  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(default_shard_id_),
+              Eq(0));
   EXPECT_EQ(embedding_index_->last_added_document_id(), 0);
 }
 
 TEST_F(EmbeddingIndexTest, OptimizeMultipleEmbeddingsSingleDocument) {
-  PropertyProto::VectorProto vector1 = CreateVector("model", {0.1, 0.2, 0.3});
-  PropertyProto::VectorProto vector2 =
-      CreateVector("model", {-0.1, -0.2, -0.3});
-  PropertyProto::VectorProto vector3 = CreateVector("model", {0.4, 0.5, 0.6});
-
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-      BasicHit(/*section_id=*/0, /*document_id=*/2), vector1,
-      QUANTIZATION_TYPE_NONE));
+      BasicHit(/*section_id=*/0, /*document_id=*/2), test_vector1_,
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-      BasicHit(/*section_id=*/0, /*document_id=*/2), vector2,
-      QUANTIZATION_TYPE_NONE));
+      BasicHit(/*section_id=*/0, /*document_id=*/2), test_vector2_,
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-      BasicHit(kSectionIdQuantizedEmbedding, /*document_id=*/2), vector3,
-      QUANTIZATION_TYPE_QUANTIZE_8_BIT));
+      BasicHit(kSectionIdQuantizedEmbedding, /*document_id=*/2), test_vector3_,
+      QUANTIZATION_TYPE_QUANTIZE_8_BIT, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
   embedding_index_->set_last_added_document_id(2);
 
@@ -950,22 +1124,23 @@
       BasicHit(kSectionIdQuantizedEmbedding, /*document_id=*/2),
       /*location=*/0);
   EXPECT_THAT(GetEmbeddingHitsFromIndex(embedding_index_.get(), /*dimension=*/3,
-                                        /*model_signature=*/"model"),
+                                        kDefaultModelSignature),
               IsOkAndHolds(ElementsAre(
                   EmbeddingHit(BasicHit(/*section_id=*/0, /*document_id=*/2),
                                /*location=*/0),
                   EmbeddingHit(BasicHit(/*section_id=*/0, /*document_id=*/2),
                                /*location=*/3),
                   quantized_hit)));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
-              ElementsAre(0.1, 0.2, 0.3, -0.1, -0.2, -0.3));
-  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(),
-              Eq(3 + sizeof(Quantizer)));
   EXPECT_THAT(
-      GetAndRestoreQuantizedEmbeddingVectorFromIndex(embedding_index_.get(),
-                                                     quantized_hit,
-                                                     /*dimension=*/3),
-      IsOkAndHolds(Pointwise(FloatNear(kEpsQuantized), vector3.values())));
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), default_shard_id_),
+      ElementsAre(0.1, 0.2, 0.3, -0.1, -0.2, -0.3));
+  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(default_shard_id_),
+              Eq(3 + sizeof(Quantizer)));
+  EXPECT_THAT(GetAndRestoreQuantizedEmbeddingVectorFromIndex(
+                  embedding_index_.get(), quantized_hit,
+                  /*dimension=*/3, kDefaultModelSignature, kDefaultSchemaName),
+              IsOkAndHolds(
+                  Pointwise(FloatNear(kEpsQuantized), test_vector3_.values())));
   EXPECT_EQ(embedding_index_->last_added_document_id(), 2);
 
   // Run optimize without deleting any documents, and check that the index is
@@ -975,22 +1150,23 @@
                                  /*document_id_old_to_new=*/{0, 1, 2},
                                  /*new_last_added_document_id=*/2));
   EXPECT_THAT(GetEmbeddingHitsFromIndex(embedding_index_.get(), /*dimension=*/3,
-                                        /*model_signature=*/"model"),
+                                        kDefaultModelSignature),
               IsOkAndHolds(ElementsAre(
                   EmbeddingHit(BasicHit(/*section_id=*/0, /*document_id=*/2),
                                /*location=*/0),
                   EmbeddingHit(BasicHit(/*section_id=*/0, /*document_id=*/2),
                                /*location=*/3),
                   quantized_hit)));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
-              ElementsAre(0.1, 0.2, 0.3, -0.1, -0.2, -0.3));
-  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(),
-              Eq(3 + sizeof(Quantizer)));
   EXPECT_THAT(
-      GetAndRestoreQuantizedEmbeddingVectorFromIndex(embedding_index_.get(),
-                                                     quantized_hit,
-                                                     /*dimension=*/3),
-      IsOkAndHolds(Pointwise(FloatNear(kEpsQuantized), vector3.values())));
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), default_shard_id_),
+      ElementsAre(0.1, 0.2, 0.3, -0.1, -0.2, -0.3));
+  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(default_shard_id_),
+              Eq(3 + sizeof(Quantizer)));
+  EXPECT_THAT(GetAndRestoreQuantizedEmbeddingVectorFromIndex(
+                  embedding_index_.get(), quantized_hit,
+                  /*dimension=*/3, kDefaultModelSignature, kDefaultSchemaName),
+              IsOkAndHolds(
+                  Pointwise(FloatNear(kEpsQuantized), test_vector3_.values())));
   EXPECT_EQ(embedding_index_->last_added_document_id(), 2);
 
   // Run optimize to map document id 2 to 1, and check that the index is
@@ -1003,22 +1179,23 @@
       EmbeddingHit(BasicHit(kSectionIdQuantizedEmbedding, /*document_id=*/1),
                    /*location=*/0);
   EXPECT_THAT(GetEmbeddingHitsFromIndex(embedding_index_.get(), /*dimension=*/3,
-                                        /*model_signature=*/"model"),
+                                        kDefaultModelSignature),
               IsOkAndHolds(ElementsAre(
                   EmbeddingHit(BasicHit(/*section_id=*/0, /*document_id=*/1),
                                /*location=*/0),
                   EmbeddingHit(BasicHit(/*section_id=*/0, /*document_id=*/1),
                                /*location=*/3),
                   quantized_hit)));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
-              ElementsAre(0.1, 0.2, 0.3, -0.1, -0.2, -0.3));
-  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(),
-              Eq(3 + sizeof(Quantizer)));
   EXPECT_THAT(
-      GetAndRestoreQuantizedEmbeddingVectorFromIndex(embedding_index_.get(),
-                                                     quantized_hit,
-                                                     /*dimension=*/3),
-      IsOkAndHolds(Pointwise(FloatNear(kEpsQuantized), vector3.values())));
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), default_shard_id_),
+      ElementsAre(0.1, 0.2, 0.3, -0.1, -0.2, -0.3));
+  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(default_shard_id_),
+              Eq(3 + sizeof(Quantizer)));
+  EXPECT_THAT(GetAndRestoreQuantizedEmbeddingVectorFromIndex(
+                  embedding_index_.get(), quantized_hit,
+                  /*dimension=*/3, kDefaultModelSignature, kDefaultSchemaName),
+              IsOkAndHolds(
+                  Pointwise(FloatNear(kEpsQuantized), test_vector3_.values())));
   EXPECT_EQ(embedding_index_->last_added_document_id(), 1);
 
   // Run optimize to delete the document.
@@ -1027,34 +1204,40 @@
       /*document_id_old_to_new=*/{0, kInvalidDocumentId},
       /*new_last_added_document_id=*/0));
   EXPECT_THAT(GetEmbeddingHitsFromIndex(embedding_index_.get(), /*dimension=*/3,
-                                        /*model_signature=*/"model"),
+                                        kDefaultModelSignature),
               IsOkAndHolds(IsEmpty()));
   EXPECT_TRUE(embedding_index_->is_empty());
   EXPECT_THAT(IndexContainsMetadataOnly(), IsOkAndHolds(true));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()), IsEmpty());
-  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(), Eq(0));
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), default_shard_id_),
+      IsEmpty());
+  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(default_shard_id_),
+              Eq(0));
   EXPECT_EQ(embedding_index_->last_added_document_id(), 0);
 }
 
 TEST_F(EmbeddingIndexTest, OptimizeMultipleEmbeddingsMultipleDocument) {
-  PropertyProto::VectorProto vector1 = CreateVector("model", {0.1, 0.2, 0.3});
-  PropertyProto::VectorProto vector2 = CreateVector("model", {1, 2, 3});
+  PropertyProto::VectorProto vector1 =
+      CreateVector(kDefaultModelSignature, {0.1, 0.2, 0.3});
+  PropertyProto::VectorProto vector2 =
+      CreateVector(kDefaultModelSignature, {1, 2, 3});
   PropertyProto::VectorProto vector3 =
-      CreateVector("model", {-0.1, -0.2, -0.3});
-  PropertyProto::VectorProto vector4 = CreateVector("model", {0.4, 0.5, 0.6});
+      CreateVector(kDefaultModelSignature, {-0.1, -0.2, -0.3});
+  PropertyProto::VectorProto vector4 =
+      CreateVector(kDefaultModelSignature, {0.4, 0.5, 0.6});
 
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(/*section_id=*/0, /*document_id=*/0), vector1,
-      QUANTIZATION_TYPE_NONE));
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(/*section_id=*/1, /*document_id=*/0), vector2,
-      QUANTIZATION_TYPE_NONE));
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(/*section_id=*/0, /*document_id=*/1), vector3,
-      QUANTIZATION_TYPE_NONE));
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(kSectionIdQuantizedEmbedding, /*document_id=*/1), vector4,
-      QUANTIZATION_TYPE_QUANTIZE_8_BIT));
+      QUANTIZATION_TYPE_QUANTIZE_8_BIT, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
   embedding_index_->set_last_added_document_id(1);
 
@@ -1063,7 +1246,7 @@
       BasicHit(kSectionIdQuantizedEmbedding, /*document_id=*/1),
       /*location=*/0);
   EXPECT_THAT(GetEmbeddingHitsFromIndex(embedding_index_.get(), /*dimension=*/3,
-                                        /*model_signature=*/"model"),
+                                        kDefaultModelSignature),
               IsOkAndHolds(ElementsAre(
                   EmbeddingHit(BasicHit(/*section_id=*/0, /*document_id=*/1),
                                /*location=*/6),
@@ -1072,14 +1255,15 @@
                                /*location=*/0),
                   EmbeddingHit(BasicHit(/*section_id=*/1, /*document_id=*/0),
                                /*location=*/3))));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
-              ElementsAre(0.1, 0.2, 0.3, 1, 2, 3, -0.1, -0.2, -0.3));
-  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(),
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), default_shard_id_),
+      ElementsAre(0.1, 0.2, 0.3, 1, 2, 3, -0.1, -0.2, -0.3));
+  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(default_shard_id_),
               Eq(3 + sizeof(Quantizer)));
   EXPECT_THAT(
-      GetAndRestoreQuantizedEmbeddingVectorFromIndex(embedding_index_.get(),
-                                                     quantized_hit,
-                                                     /*dimension=*/3),
+      GetAndRestoreQuantizedEmbeddingVectorFromIndex(
+          embedding_index_.get(), quantized_hit,
+          /*dimension=*/3, kDefaultModelSignature, kDefaultSchemaName),
       IsOkAndHolds(Pointwise(FloatNear(kEpsQuantized), vector4.values())));
   EXPECT_EQ(embedding_index_->last_added_document_id(), 1);
 
@@ -1096,7 +1280,7 @@
                                    /*new_last_added_document_id=*/1));
     EXPECT_THAT(
         GetEmbeddingHitsFromIndex(embedding_index_.get(), /*dimension=*/3,
-                                  /*model_signature=*/"model"),
+                                  kDefaultModelSignature),
         IsOkAndHolds(ElementsAre(
             EmbeddingHit(BasicHit(/*section_id=*/0, /*document_id=*/1),
                          /*location=*/0),
@@ -1105,14 +1289,16 @@
                          /*location=*/3),
             EmbeddingHit(BasicHit(/*section_id=*/1, /*document_id=*/0),
                          /*location=*/6))));
-    EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
-                ElementsAre(-0.1, -0.2, -0.3, 0.1, 0.2, 0.3, 1, 2, 3));
-    EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(),
-                Eq(3 + sizeof(Quantizer)));
     EXPECT_THAT(
-        GetAndRestoreQuantizedEmbeddingVectorFromIndex(embedding_index_.get(),
-                                                       quantized_hit,
-                                                       /*dimension=*/3),
+        GetRawEmbeddingDataFromIndex(embedding_index_.get(), default_shard_id_),
+        ElementsAre(-0.1, -0.2, -0.3, 0.1, 0.2, 0.3, 1, 2, 3));
+    EXPECT_THAT(
+        embedding_index_->GetTotalQuantizedVectorSize(default_shard_id_),
+        Eq(3 + sizeof(Quantizer)));
+    EXPECT_THAT(
+        GetAndRestoreQuantizedEmbeddingVectorFromIndex(
+            embedding_index_.get(), quantized_hit,
+            /*dimension=*/3, kDefaultModelSignature, kDefaultSchemaName),
         IsOkAndHolds(Pointwise(FloatNear(kEpsQuantized), vector4.values())));
     EXPECT_EQ(embedding_index_->last_added_document_id(), 1);
   }
@@ -1127,19 +1313,20 @@
       EmbeddingHit(BasicHit(kSectionIdQuantizedEmbedding, /*document_id=*/0),
                    /*location=*/0);
   EXPECT_THAT(GetEmbeddingHitsFromIndex(embedding_index_.get(), /*dimension=*/3,
-                                        /*model_signature=*/"model"),
+                                        kDefaultModelSignature),
               IsOkAndHolds(ElementsAre(
                   EmbeddingHit(BasicHit(/*section_id=*/0, /*document_id=*/0),
                                /*location=*/0),
                   quantized_hit)));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
-              ElementsAre(-0.1, -0.2, -0.3));
-  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(),
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), default_shard_id_),
+      ElementsAre(-0.1, -0.2, -0.3));
+  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(default_shard_id_),
               Eq(3 + sizeof(Quantizer)));
   EXPECT_THAT(
-      GetAndRestoreQuantizedEmbeddingVectorFromIndex(embedding_index_.get(),
-                                                     quantized_hit,
-                                                     /*dimension=*/3),
+      GetAndRestoreQuantizedEmbeddingVectorFromIndex(
+          embedding_index_.get(), quantized_hit,
+          /*dimension=*/3, kDefaultModelSignature, kDefaultSchemaName),
       IsOkAndHolds(Pointwise(FloatNear(kEpsQuantized), vector4.values())));
   EXPECT_EQ(embedding_index_->last_added_document_id(), 0);
 }
@@ -1149,15 +1336,19 @@
   PropertyProto::VectorProto vector2 = CreateVector("model1", {1, 2});
   PropertyProto::VectorProto vector3 =
       CreateVector("model2", {-0.1, -0.2, -0.3});
+  uint32_t vector1_and_vector2_shard_id = embedding_index_->GetShardId(
+      /*dimension=*/2, "model1", kDefaultSchemaName);
+  uint32_t vector3_shard_id = embedding_index_->GetShardId(
+      /*dimension=*/3, "model2", kDefaultSchemaName);
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(/*section_id=*/0, /*document_id=*/0), vector1,
-      QUANTIZATION_TYPE_NONE));
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(/*section_id=*/0, /*document_id=*/1), vector2,
-      QUANTIZATION_TYPE_NONE));
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(/*section_id=*/1, /*document_id=*/1), vector3,
-      QUANTIZATION_TYPE_NONE));
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
   embedding_index_->set_last_added_document_id(1);
 
@@ -1173,19 +1364,17 @@
                                         /*model_signature=*/"model2"),
               IsOkAndHolds(ElementsAre(
                   EmbeddingHit(BasicHit(/*section_id=*/1, /*document_id=*/1),
-                               /*location=*/4))));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
-              ElementsAre(0.1, 0.2, 1, 2, -0.1, -0.2, -0.3));
+                               /*location=*/0))));
+  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get(),
+                                           vector1_and_vector2_shard_id),
+              ElementsAre(0.1, 0.2, 1, 2));
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), vector3_shard_id),
+      ElementsAre(-0.1, -0.2, -0.3));
   EXPECT_EQ(embedding_index_->last_added_document_id(), 1);
 
-  // Run optimize without deleting any documents. It is expected to see that the
-  // raw embedding data is rearranged, since during index transfer:
-  // - Embedding vectors with lower keys, which are the string encoded ordered
-  //   pairs (dimension, model_signature), are iterated first.
-  // - Embedding vectors from higher document ids are added first.
-  //
-  // Also keep in mind that once the raw data is rearranged, calling another
-  // Optimize subsequently will not change the raw data again.
+  // Run optimize without deleting any documents. We should see the same data in
+  // all the shards, since each shard only contains one vector.
   for (int i = 0; i < 2; i++) {
     ICING_ASSERT_OK(
         embedding_index_->Optimize(document_store_.get(), schema_store_.get(),
@@ -1204,9 +1393,13 @@
                                   /*model_signature=*/"model2"),
         IsOkAndHolds(ElementsAre(
             EmbeddingHit(BasicHit(/*section_id=*/1, /*document_id=*/1),
-                         /*location=*/4))));
-    EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
-                ElementsAre(1, 2, 0.1, 0.2, -0.1, -0.2, -0.3));
+                         /*location=*/0))));
+    EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get(),
+                                             vector1_and_vector2_shard_id),
+                ElementsAre(1, 2, 0.1, 0.2));
+    EXPECT_THAT(
+        GetRawEmbeddingDataFromIndex(embedding_index_.get(), vector3_shard_id),
+        ElementsAre(-0.1, -0.2, -0.3));
     EXPECT_EQ(embedding_index_->last_added_document_id(), 1);
   }
 
@@ -1224,8 +1417,12 @@
   EXPECT_THAT(GetEmbeddingHitsFromIndex(embedding_index_.get(), /*dimension=*/3,
                                         /*model_signature=*/"model2"),
               IsOkAndHolds(IsEmpty()));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
+  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get(),
+                                           vector1_and_vector2_shard_id),
               ElementsAre(0.1, 0.2));
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), vector3_shard_id),
+      IsEmpty());
   EXPECT_EQ(embedding_index_->last_added_document_id(), 0);
 }
 
@@ -1234,12 +1431,16 @@
   PropertyProto::VectorProto vector1 = CreateVector("model1", {0.1, 0.2});
   PropertyProto::VectorProto vector2 =
       CreateVector("model2", {-0.1, -0.2, -0.3});
+  uint32_t vector1_shard_id = embedding_index_->GetShardId(
+      /*dimension=*/2, "model1", kDefaultSchemaName);
+  uint32_t vector2_shard_id = embedding_index_->GetShardId(
+      /*dimension=*/3, "model2", kDefaultSchemaName);
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(/*section_id=*/0, /*document_id=*/0), vector1,
-      QUANTIZATION_TYPE_NONE));
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(/*section_id=*/1, /*document_id=*/1), vector2,
-      QUANTIZATION_TYPE_NONE));
+      QUANTIZATION_TYPE_NONE, kDefaultSchemaName));
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
   embedding_index_->set_last_added_document_id(1);
 
@@ -1253,9 +1454,13 @@
                                         /*model_signature=*/"model2"),
               IsOkAndHolds(ElementsAre(
                   EmbeddingHit(BasicHit(/*section_id=*/1, /*document_id=*/1),
-                               /*location=*/2))));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
-              ElementsAre(0.1, 0.2, -0.1, -0.2, -0.3));
+                               /*location=*/0))));
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), vector1_shard_id),
+      ElementsAre(0.1, 0.2));
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), vector2_shard_id),
+      ElementsAre(-0.1, -0.2, -0.3));
   EXPECT_EQ(embedding_index_->last_added_document_id(), 1);
 
   // Run optimize to delete document 0, and check that the index is
@@ -1272,11 +1477,59 @@
               IsOkAndHolds(ElementsAre(
                   EmbeddingHit(BasicHit(/*section_id=*/1, /*document_id=*/0),
                                /*location=*/0))));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
-              ElementsAre(-0.1, -0.2, -0.3));
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), vector1_shard_id),
+      IsEmpty());
+  EXPECT_THAT(
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), vector2_shard_id),
+      ElementsAre(-0.1, -0.2, -0.3));
   EXPECT_EQ(embedding_index_->last_added_document_id(), 0);
 }
 
 }  // namespace
+
+TEST_F(EmbeddingIndexTest, ShardFileShouldBeLazilyCreated) {
+  // Check no shard file exists after initialization.
+  std::vector<std::string> files;
+  ASSERT_TRUE(filesystem_.ListDirectory(embedding_index_dir_.c_str(),
+                                        /*exclude=*/{},
+                                        /*recursive=*/true, &files));
+  for (const std::string& file : files) {
+    EXPECT_THAT(file, Not(HasSubstr("embedding_vectors")));
+  }
+
+  // Add 5 embeddings, each with a different schema name, and should (possibly)
+  // have a different shard id.
+  std::unordered_set<uint32_t> shard_ids;
+  for (int i = 0; i < 5; i++) {
+    std::string schema_name = "schema" + std::to_string(i);
+    uint32_t shard_id = embedding_index_->GetShardId(
+        kDefaultDimension, kDefaultModelSignature, schema_name);
+    shard_ids.insert(shard_id);
+
+    // Add a non-quantized embedding, and check if its corresponding shard
+    // file is created.
+    std::string vector_file = absl_ports::StrCat(
+        embedding_index_dir_, "/embedding_vectors_", std::to_string(shard_id));
+    EXPECT_FALSE(filesystem_.FileExists(vector_file.c_str()));
+    ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
+        BasicHit(/*section_id=*/0, /*document_id=*/0), test_vector1_,
+        QUANTIZATION_TYPE_NONE, schema_name));
+    EXPECT_TRUE(filesystem_.FileExists(vector_file.c_str()));
+
+    // Add a quantized embedding, and check if its corresponding shard file
+    // is created.
+    std::string quantized_vector_file = absl_ports::StrCat(
+        embedding_index_dir_, "/quantized_embedding_vectors_",
+        std::to_string(shard_id));
+    EXPECT_FALSE(filesystem_.FileExists(quantized_vector_file.c_str()));
+    ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
+        BasicHit(kSectionIdQuantizedEmbedding, /*document_id=*/0),
+        test_vector1_, QUANTIZATION_TYPE_QUANTIZE_8_BIT, schema_name));
+    EXPECT_TRUE(filesystem_.FileExists(quantized_vector_file.c_str()));
+  }
+  EXPECT_EQ(shard_ids.size(), 5);
+}
+
 }  // namespace lib
 }  // namespace icing

diff --git a/icing/index/embed/embedding-query-results.h b/icing/index/embed/embedding-query-results.h
index a90ad10..15c597d 100644
--- a/icing/index/embed/embedding-query-results.h
+++ b/icing/index/embed/embedding-query-results.h

@@ -15,21 +15,33 @@
 #ifndef ICING_INDEX_EMBED_EMBEDDING_QUERY_RESULTS_H_
 #define ICING_INDEX_EMBED_EMBEDDING_QUERY_RESULTS_H_
 
+#include <cstdint>
 #include <memory>
+#include <optional>
 #include <unordered_map>
 #include <vector>
 
+#include "icing/text_classifier/lib3/utils/base/status.h"
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
+#include "icing/absl_ports/canonical_errors.h"
 #include "icing/legacy/core/icing-packed-pod.h"
 #include "icing/proto/search.pb.h"
 #include "icing/schema/section.h"
+#include "icing/scoring/advanced_scoring/double-list.h"
 #include "icing/store/document-id.h"
+#include "icing/util/embedding-util.h"
+#include "icing/util/status-macros.h"
 
 namespace icing {
 namespace lib {
 
+// Stores the matched embedding infos for a single document.
 struct EmbeddingMatchInfos {
-  // A vector of semantic scores of matched embeddings.
-  std::vector<double> scores;
+  // [score_start_index, score_end_index) is the range of score indexes in the
+  // global_scores vector. If embedding section info is enabled, the same range
+  // will be used for section infos in global_section_infos as well.
+  int32_t score_start_index = 0;
+  int32_t score_end_index = 0;
 
   struct EmbeddingMatchSectionInfo {
     // The position of the matched embedding vector in a section relative to
@@ -56,33 +68,72 @@
   static_assert(icing_is_packed_pod<EmbeddingMatchSectionInfo>::value,
                 "go/icing-ubsan");
 
-  // A vector of section infos on the matched embeddings. This will be nullptr
-  // if embedding match info is not enabled for this query.
-  //
-  // When non-null, section_infos must have a 1:1 mapping with the scores
-  // vector.
-  std::unique_ptr<std::vector<EmbeddingMatchSectionInfo>> section_infos;
-
   EmbeddingMatchInfos() = default;
-
   EmbeddingMatchInfos(const EmbeddingMatchInfos& other) = delete;
   EmbeddingMatchInfos& operator=(const EmbeddingMatchInfos& other) = delete;
 
-  // Appends a score to the scores vector.
-  void AppendScore(double score) { scores.push_back(score); }
-
-  // Appends a section info to the section_infos vector, allocating if needed.
-  void AppendSectionInfo(SectionId section_id, int position) {
-    if (!section_infos) {
-      section_infos =
-          std::make_unique<std::vector<EmbeddingMatchSectionInfo>>();
+  // Appends a score to the scores vector, which is stored in the global_scores
+  // vector. score_start_index and score_end_index will be updated accordingly.
+  //
+  // Returns:
+  //   - OK, if the score is appended successfully.
+  //   - FailedPreconditionError, if the score is not contiguous with the
+  //     existing scores.
+  libtextclassifier3::Status AppendScore(std::vector<double>& global_scores,
+                                         double score) {
+    if (score_end_index == 0) {
+      score_start_index = score_end_index = global_scores.size();
     }
-    section_infos->push_back({.position = position, .section_id = section_id});
+    if (score_end_index != global_scores.size()) {
+      return absl_ports::FailedPreconditionError(
+          "Scores for the same document should be contiguous.");
+    }
+    global_scores.push_back(score);
+    score_end_index += 1;
+    return libtextclassifier3::Status::OK;
+  }
+
+  // Appends a section info to the section info vector, which is stored in the
+  // global_section_infos vector.
+  //
+  // Returns:
+  //   - OK, if the section info is appended successfully.
+  //   - FailedPreconditionError, if the section info is not appended
+  //     immediately after AppendScore.
+  libtextclassifier3::Status AppendSectionInfo(
+      std::vector<EmbeddingMatchInfos::EmbeddingMatchSectionInfo>&
+          global_section_infos,
+      SectionId section_id, int position) {
+    if (score_end_index != global_section_infos.size() + 1) {
+      return absl_ports::FailedPreconditionError(
+          "Section infos must be appended immediately after AppendScore.");
+    }
+    global_section_infos.push_back(
+        {.position = position, .section_id = section_id});
+    return libtextclassifier3::Status::OK;
   }
 };
 
 // A class to store results generated from embedding queries.
-struct EmbeddingQueryResults {
+class EmbeddingQueryResults {
+ public:
+  // Creates an empty EmbeddingQueryResults instance.
+  EmbeddingQueryResults() : EmbeddingQueryResults(/*num_query_vectors=*/0) {}
+
+  // Creates an EmbeddingQueryResults instance with the given number of query
+  // vectors.
+  EmbeddingQueryResults(int num_query_vectors)
+      : result_infos_size_(embedding_util::kEmbeddingQueryMetricTypes.size() *
+                           num_query_vectors),
+        result_infos_(
+            std::make_unique<std::optional<EmbeddingQueryMatchInfoMap>[]>(
+                result_infos_size_)) {}
+
+  int GetNumQueryVectors() const {
+    return result_infos_size_ /
+           embedding_util::kEmbeddingQueryMetricTypes.size();
+  }
+
   // Maps from DocumentId to matched embedding infos for that document.
   // For each document, its embedding match info consists of two vectors:
   // - The scores vector, which will be used in the advanced scoring language
@@ -93,11 +144,35 @@
   using EmbeddingQueryMatchInfoMap =
       std::unordered_map<DocumentId, EmbeddingMatchInfos>;
 
-  // Maps from (query_vector_index, metric_type) to EmbeddingQueryMatchInfoMap.
-  std::unordered_map<
-      int, std::unordered_map<SearchSpecProto::EmbeddingQueryMetricType::Code,
-                              EmbeddingQueryMatchInfoMap>>
-      result_infos;
+  // A centralized vector of scores for all documents. This is used to store the
+  // scores for the "this.matchedSemanticScores(...)" function.
+  std::unique_ptr<std::vector<double>> global_scores =
+      std::make_unique<std::vector<double>>();
+
+  // A centralized vector of EmbeddingMatchSectionInfo for all documents. This
+  // is used to store the section infos for the embedding query.
+  std::unique_ptr<std::vector<EmbeddingMatchInfos::EmbeddingMatchSectionInfo>>
+      global_section_infos = std::make_unique<
+          std::vector<EmbeddingMatchInfos::EmbeddingMatchSectionInfo>>();
+
+  // Get or create the MatchedInfo map for the given query_vector_index and
+  // metric_type.
+  //
+  // Returns:
+  //   - The pointer to the EmbeddingQueryMatchInfoMap map, if the map is found
+  //     or created.
+  //   - InvalidArgumentError, if the index is out of bounds.
+  libtextclassifier3::StatusOr<EmbeddingQueryMatchInfoMap*>
+  GetOrCreateMatchInfoMap(
+      int query_vector_index,
+      SearchSpecProto::EmbeddingQueryMetricType::Code metric_type) const {
+    ICING_ASSIGN_OR_RETURN(int index,
+                           GetResultInfoIndex(query_vector_index, metric_type));
+    if (!result_infos_[index].has_value()) {
+      result_infos_[index] = EmbeddingQueryMatchInfoMap();
+    }
+    return &result_infos_[index].value();
+  }
 
   // Get the MatchedInfo map for the given query_vector_index and metric_type.
   // Returns nullptr if (query_vector_index, metric_type) does not exist in the
@@ -105,17 +180,12 @@
   const EmbeddingQueryMatchInfoMap* GetMatchInfoMap(
       int query_vector_index,
       SearchSpecProto::EmbeddingQueryMetricType::Code metric_type) const {
-    // Check if a mapping exists for the query_vector_index
-    auto outer_it = result_infos.find(query_vector_index);
-    if (outer_it == result_infos.end()) {
+    libtextclassifier3::StatusOr<int> index =
+        GetResultInfoIndex(query_vector_index, metric_type);
+    if (!index.ok() || !result_infos_[index.ValueOrDie()].has_value()) {
       return nullptr;
     }
-    // Check if a mapping exists for the metric_type
-    auto inner_it = outer_it->second.find(metric_type);
-    if (inner_it == outer_it->second.end()) {
-      return nullptr;
-    }
-    return &inner_it->second;
+    return &result_infos_[index.ValueOrDie()].value();
   }
 
   // Returns the matched infos for the given query_vector_index, metric_type,
@@ -139,19 +209,62 @@
   }
 
   // Returns the matched scores for the given query_vector_index, metric_type,
-  // and doc_id. Returns nullptr if (query_vector_index, metric_type, doc_id)
-  // does not exist in the result_scores map.
-  const std::vector<double>* GetMatchedScoresForDocument(
+  // and doc_id. Returns an empty DoubleList if (query_vector_index,
+  // metric_type, doc_id) does not exist in the result_scores map.
+  //
+  // The returned DoubleList is a non-owning view of the scores vector stored
+  // within the EmbeddingQueryResults instance. The caller must ensure the
+  // lifetime of the EmbeddingQueryResults exceeds the lifetime of the returned
+  // DoubleList.
+  DoubleList GetMatchedScoresForDocument(
       int query_vector_index,
       SearchSpecProto::EmbeddingQueryMetricType::Code metric_type,
       DocumentId doc_id) const {
     const EmbeddingMatchInfos* match_infos =
         GetMatchedInfosForDocument(query_vector_index, metric_type, doc_id);
     if (match_infos == nullptr) {
-      return nullptr;
+      return DoubleList();
     }
-    return &match_infos->scores;
+    return GetMatchedScoresFromEmbeddingMatchInfos(*match_infos);
   };
+
+  // Returns the matched scores for the given EmbeddingMatchInfos, which stores
+  // the match infos for a single document.
+  //
+  // The returned DoubleList is a non-owning view of the scores vector stored
+  // within the EmbeddingQueryResults instance. The caller must ensure the
+  // lifetime of the EmbeddingQueryResults exceeds the lifetime of the returned
+  // DoubleList.
+  DoubleList GetMatchedScoresFromEmbeddingMatchInfos(
+      const EmbeddingMatchInfos& match_infos) const {
+    return DoubleList(
+        global_scores->data() + match_infos.score_start_index,
+        match_infos.score_end_index - match_infos.score_start_index);
+  }
+
+ private:
+  // Maps from (query_vector_index, metric_type) to EmbeddingQueryMatchInfoMap.
+  int result_infos_size_;
+  std::unique_ptr<std::optional<EmbeddingQueryMatchInfoMap>[]> result_infos_;
+
+  // Returns the index of the result info for the given query_vector_index and
+  // metric_type.
+  //
+  // Returns:
+  //   - The index of the result info, if the index is valid.
+  //   - InvalidArgumentError, if the index is out of bounds.
+  libtextclassifier3::StatusOr<int> GetResultInfoIndex(
+      int query_vector_index,
+      SearchSpecProto::EmbeddingQueryMetricType::Code metric_type) const {
+    int index =
+        query_vector_index * embedding_util::kEmbeddingQueryMetricTypes.size() +
+        (metric_type - embedding_util::kEmbeddingQueryMetricTypes[0]);
+    if (result_infos_ == nullptr || index < 0 || index >= result_infos_size_) {
+      return absl_ports::InvalidArgumentError(
+          "result_infos_ index out of bounds.");
+    }
+    return index;
+  }
 };
 
 }  // namespace lib

diff --git a/icing/index/embed/embedding-scorer.cc b/icing/index/embed/embedding-scorer.cc
index 0a3ed6c..64007d4 100644
--- a/icing/index/embed/embedding-scorer.cc
+++ b/icing/index/embed/embedding-scorer.cc

@@ -26,6 +26,11 @@
 #include "icing/index/embed/quantizer.h"
 #include "icing/proto/search.pb.h"
 
+#ifndef ICING_DISABLE_EIGEN
+#include <Eigen/Core>
+#include <Eigen/Dense>
+#endif  // ICING_DISABLE_EIGEN
+
 namespace icing {
 namespace lib {
 
@@ -85,6 +90,30 @@
   return std::sqrt(result);
 }
 
+#ifndef ICING_DISABLE_EIGEN
+// Returns a lazily-evaluated (due to the return type "auto") expression that
+// dequantizes the quantized vector.
+//
+// It requires that the scale factor is not 0. If it is 0, the caller should
+// handle the case specifically.
+inline auto DequantizeEigenVector(
+    const Quantizer& quantizer,
+    Eigen::Ref<const Eigen::VectorX<uint8_t>> quantized_vec) {
+  return (quantized_vec.cast<float>().array() / quantizer.scale_factor() +
+          quantizer.float_min())
+      .matrix();
+}
+
+float EigenCosine(Eigen::Ref<const Eigen::VectorXf> v1,
+                  Eigen::Ref<const Eigen::VectorXf> v2) {
+  float divisor = v1.norm() * v2.norm();
+  if (divisor == 0.0) {
+    return 0.0;
+  }
+  return v1.dot(v2) / divisor;
+}
+#endif  // ICING_DISABLE_EIGEN
+
 }  // namespace
 
 libtextclassifier3::StatusOr<std::unique_ptr<EmbeddingScorer>>
@@ -136,5 +165,103 @@
   return CalculateEuclideanDistance(dimension, v1, v2, quantizer);
 }
 
+#ifndef ICING_DISABLE_EIGEN
+float CosineEmbeddingScorer::EigenScore(int dimension, const float* v1,
+                                        const float* v2) const {
+  Eigen::Map<const Eigen::VectorXf> vec1(v1, dimension);
+  Eigen::Map<const Eigen::VectorXf> vec2(v2, dimension);
+  return EigenCosine(vec1, vec2);
+}
+
+float DotProductEmbeddingScorer::EigenScore(int dimension, const float* v1,
+                                            const float* v2) const {
+  Eigen::Map<const Eigen::VectorXf> vec1(v1, dimension);
+  Eigen::Map<const Eigen::VectorXf> vec2(v2, dimension);
+  return vec1.dot(vec2);
+}
+
+float EuclideanDistanceEmbeddingScorer::EigenScore(int dimension,
+                                                   const float* v1,
+                                                   const float* v2) const {
+  Eigen::Map<const Eigen::VectorXf> vec1(v1, dimension);
+  Eigen::Map<const Eigen::VectorXf> vec2(v2, dimension);
+  return (vec1 - vec2).norm();
+}
+
+float CosineEmbeddingScorer::EigenScore(int dimension, const float* v1,
+                                        const uint8_t* v2,
+                                        const Quantizer& quantizer) const {
+  Eigen::Map<const Eigen::VectorXf> vec1(v1, dimension);
+  Eigen::Map<const Eigen::Matrix<uint8_t, Eigen::Dynamic, 1>> vec2(v2,
+                                                                   dimension);
+  if (quantizer.scale_factor() == 0.0) {
+    return EigenCosine(
+        vec1, Eigen::VectorXf::Constant(dimension, quantizer.float_min()));
+  }
+  return EigenCosine(vec1, DequantizeEigenVector(quantizer, vec2));
+}
+
+float DotProductEmbeddingScorer::EigenScore(int dimension, const float* v1,
+                                            const uint8_t* v2,
+                                            const Quantizer& quantizer) const {
+  Eigen::Map<const Eigen::VectorXf> vec1(v1, dimension);
+  Eigen::Map<const Eigen::Matrix<uint8_t, Eigen::Dynamic, 1>> vec2(v2,
+                                                                   dimension);
+  if (quantizer.scale_factor() == 0.0) {
+    return vec1.dot(
+        Eigen::VectorXf::Constant(dimension, quantizer.float_min()));
+  }
+  return vec1.dot(DequantizeEigenVector(quantizer, vec2));
+}
+
+float EuclideanDistanceEmbeddingScorer::EigenScore(
+    int dimension, const float* v1, const uint8_t* v2,
+    const Quantizer& quantizer) const {
+  Eigen::Map<const Eigen::VectorXf> vec1(v1, dimension);
+  Eigen::Map<const Eigen::Matrix<uint8_t, Eigen::Dynamic, 1>> vec2(v2,
+                                                                   dimension);
+  if (quantizer.scale_factor() == 0.0) {
+    return (vec1.array() - quantizer.float_min()).matrix().norm();
+  }
+  return (vec1 - DequantizeEigenVector(quantizer, vec2)).norm();
+}
+#else   // ICING_DISABLE_EIGEN
+// If Eigen is disabled, just fall back to the regular Score() function.
+
+float CosineEmbeddingScorer::EigenScore(int dimension, const float* v1,
+                                        const float* v2) const {
+  return Score(dimension, v1, v2);
+}
+
+float DotProductEmbeddingScorer::EigenScore(int dimension, const float* v1,
+                                            const float* v2) const {
+  return Score(dimension, v1, v2);
+}
+
+float EuclideanDistanceEmbeddingScorer::EigenScore(int dimension,
+                                                   const float* v1,
+                                                   const float* v2) const {
+  return Score(dimension, v1, v2);
+}
+
+float CosineEmbeddingScorer::EigenScore(int dimension, const float* v1,
+                                        const uint8_t* v2,
+                                        const Quantizer& quantizer) const {
+  return Score(dimension, v1, v2, quantizer);
+}
+
+float DotProductEmbeddingScorer::EigenScore(int dimension, const float* v1,
+                                            const uint8_t* v2,
+                                            const Quantizer& quantizer) const {
+  return Score(dimension, v1, v2, quantizer);
+}
+
+float EuclideanDistanceEmbeddingScorer::EigenScore(
+    int dimension, const float* v1, const uint8_t* v2,
+    const Quantizer& quantizer) const {
+  return Score(dimension, v1, v2, quantizer);
+}
+#endif  // ICING_DISABLE_EIGEN
+
 }  // namespace lib
 }  // namespace icing

diff --git a/icing/index/embed/embedding-scorer.h b/icing/index/embed/embedding-scorer.h
index b0e0f75..403616a 100644
--- a/icing/index/embed/embedding-scorer.h
+++ b/icing/index/embed/embedding-scorer.h

@@ -35,6 +35,11 @@
   virtual float Score(int dimension, const float* v1, const uint8_t* v2,
                       const Quantizer& quantizer) const = 0;
 
+  virtual float EigenScore(int dimension, const float* v1,
+                           const float* v2) const = 0;
+  virtual float EigenScore(int dimension, const float* v1, const uint8_t* v2,
+                           const Quantizer& quantizer) const = 0;
+
   virtual ~EmbeddingScorer() = default;
 };
 
@@ -43,6 +48,11 @@
   float Score(int dimension, const float* v1, const float* v2) const override;
   float Score(int dimension, const float* v1, const uint8_t* v2,
               const Quantizer& quantizer) const override;
+
+  float EigenScore(int dimension, const float* v1,
+                   const float* v2) const override;
+  float EigenScore(int dimension, const float* v1, const uint8_t* v2,
+                   const Quantizer& quantizer) const override;
 };
 
 class DotProductEmbeddingScorer : public EmbeddingScorer {
@@ -50,6 +60,11 @@
   float Score(int dimension, const float* v1, const float* v2) const override;
   float Score(int dimension, const float* v1, const uint8_t* v2,
               const Quantizer& quantizer) const override;
+
+  float EigenScore(int dimension, const float* v1,
+                   const float* v2) const override;
+  float EigenScore(int dimension, const float* v1, const uint8_t* v2,
+                   const Quantizer& quantizer) const override;
 };
 
 class EuclideanDistanceEmbeddingScorer : public EmbeddingScorer {
@@ -57,6 +72,11 @@
   float Score(int dimension, const float* v1, const float* v2) const override;
   float Score(int dimension, const float* v1, const uint8_t* v2,
               const Quantizer& quantizer) const override;
+
+  float EigenScore(int dimension, const float* v1,
+                   const float* v2) const override;
+  float EigenScore(int dimension, const float* v1, const uint8_t* v2,
+                   const Quantizer& quantizer) const override;
 };
 
 }  // namespace lib

diff --git a/icing/index/embed/embedding-scorer_test.cc b/icing/index/embed/embedding-scorer_test.cc
index 162b969..035526d 100644
--- a/icing/index/embed/embedding-scorer_test.cc
+++ b/icing/index/embed/embedding-scorer_test.cc

@@ -14,8 +14,11 @@
 
 #include "icing/index/embed/embedding-scorer.h"
 
+#include <algorithm>
 #include <cstdint>
 #include <memory>
+#include <random>
+#include <tuple>
 #include <vector>
 
 #include "gtest/gtest.h"
@@ -120,6 +123,120 @@
               expected_euclidean, eps_quantized);
 }
 
+class EmbeddingScorerEigenTest
+    : public testing::TestWithParam<
+          std::tuple<SearchSpecProto::EmbeddingQueryMetricType::Code, int>> {
+ protected:
+  void SetUp() override {
+    metric_ = std::get<0>(GetParam());
+    dimension_ = std::get<1>(GetParam());
+    ICING_ASSERT_OK_AND_ASSIGN(embedding_scorer_,
+                               EmbeddingScorer::Create(metric_));
+
+    // Initialize random number generator
+    random_ = std::default_random_engine(std::random_device()());
+    dist_ = std::uniform_real_distribution<float>(-3.0f, 3.0f);
+  }
+
+  // Generates a random vector of the specified dimension.
+  std::vector<float> GenerateRandomVector() {
+    std::vector<float> vec(dimension_);
+    for (int i = 0; i < dimension_; ++i) {
+      vec[i] = dist_(random_);
+    }
+    return vec;
+  }
+
+  std::vector<float> GenerateRandomConstantVector() {
+    float value = dist_(random_);
+    std::vector<float> vec(dimension_, value);
+    return vec;
+  }
+
+  const int kNumRandomPairs = 1000;
+  const float kEps = 0.001f;
+
+  SearchSpecProto::EmbeddingQueryMetricType::Code metric_;
+  int dimension_;
+  std::unique_ptr<EmbeddingScorer> embedding_scorer_;
+  std::default_random_engine random_;
+  std::uniform_real_distribution<float> dist_;
+};
+
+// Test that the EigenScore function matches the Score function for a variety
+// of random vectors.
+TEST_P(EmbeddingScorerEigenTest, EigenScoreMatchesScore) {
+  for (int i = 0; i < kNumRandomPairs; ++i) {
+    std::vector<float> v1 = GenerateRandomVector();
+    std::vector<float> v2 = GenerateRandomVector();
+
+    // Compare scores
+    float score_val =
+        embedding_scorer_->Score(dimension_, v1.data(), v2.data());
+    float eigen_score_val =
+        embedding_scorer_->EigenScore(dimension_, v1.data(), v2.data());
+    ASSERT_NEAR(score_val, eigen_score_val, kEps);
+  }
+}
+
+// Test that the EigenScore function matches the Score function for a variety
+// of random quantized vectors.
+TEST_P(EmbeddingScorerEigenTest, EigenScoreMatchesScoreForQuantizedVectors) {
+  for (int i = 0; i < kNumRandomPairs; ++i) {
+    std::vector<float> v1 = GenerateRandomVector();
+    std::vector<float> v2 = GenerateRandomVector();
+
+    // Quantize v2
+    auto v2_minmax_pair = std::minmax_element(v2.begin(), v2.end());
+    ICING_ASSERT_OK_AND_ASSIGN(
+        Quantizer quantizer,
+        Quantizer::Create(*v2_minmax_pair.first, *v2_minmax_pair.second));
+    std::vector<uint8_t> v2_quantized = QuantizeVector(v2, quantizer);
+
+    // Compare scores
+    float score_val = embedding_scorer_->Score(dimension_, v1.data(),
+                                               v2_quantized.data(), quantizer);
+    float eigen_score_val = embedding_scorer_->EigenScore(
+        dimension_, v1.data(), v2_quantized.data(), quantizer);
+    ASSERT_NEAR(score_val, eigen_score_val, kEps);
+  }
+}
+
+// Test that the EigenScore function matches the Score function for constant
+// vectors (i.e. all values are the same) to be quantized.
+TEST_P(EmbeddingScorerEigenTest,
+       EigenScoreMatchesScoreForQuantizedConstantVectors) {
+  for (int i = 0; i < kNumRandomPairs; ++i) {
+    std::vector<float> v1 = GenerateRandomVector();
+    std::vector<float> v2 = GenerateRandomConstantVector();
+
+    // Check that v2 is constant.
+    auto v2_minmax_pair = std::minmax_element(v2.begin(), v2.end());
+    ASSERT_TRUE(*v2_minmax_pair.first == *v2_minmax_pair.second);
+    // Quantize v2
+    ICING_ASSERT_OK_AND_ASSIGN(
+        Quantizer quantizer,
+        Quantizer::Create(*v2_minmax_pair.first, *v2_minmax_pair.second));
+    ASSERT_EQ(quantizer.scale_factor(), 0.f);
+    std::vector<uint8_t> v2_quantized = QuantizeVector(v2, quantizer);
+
+    // Compare scores
+    float score_val = embedding_scorer_->Score(dimension_, v1.data(),
+                                               v2_quantized.data(), quantizer);
+    float eigen_score_val = embedding_scorer_->EigenScore(
+        dimension_, v1.data(), v2_quantized.data(), quantizer);
+    ASSERT_NEAR(score_val, eigen_score_val, kEps);
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    EigenVsScoreComparison, EmbeddingScorerEigenTest,
+    testing::Combine(
+        testing::Values(SearchSpecProto::EmbeddingQueryMetricType::DOT_PRODUCT,
+                        SearchSpecProto::EmbeddingQueryMetricType::COSINE,
+                        SearchSpecProto::EmbeddingQueryMetricType::EUCLIDEAN),
+        testing::Values(128, 512, 768, 1024)));
+
 }  // namespace
 
 }  // namespace lib

diff --git a/icing/index/embed/quantizer.h b/icing/index/embed/quantizer.h
index 099c96d..7095263 100644
--- a/icing/index/embed/quantizer.h
+++ b/icing/index/embed/quantizer.h

@@ -74,6 +74,9 @@
     return (quantized / scale_factor_) + float_min_;
   }
 
+  float float_min() const { return float_min_; }
+  float scale_factor() const { return scale_factor_; }
+
  private:
   static constexpr uint8_t kMaxQuantizedValue =
       std::numeric_limits<uint8_t>::max();

diff --git a/icing/index/embedding-indexing-handler.cc b/icing/index/embedding-indexing-handler.cc
index 28d6568..ec95534 100644
--- a/icing/index/embedding-indexing-handler.cc
+++ b/icing/index/embedding-indexing-handler.cc

@@ -72,7 +72,9 @@
       BasicHit hit(/*section_id=*/vector_section.metadata.id, document_id);
       for (const PropertyProto::VectorProto& vector : vector_section.content) {
         ICING_RETURN_IF_ERROR(embedding_index_.BufferEmbedding(
-            hit, vector, vector_section.metadata.quantization_type));
+            hit, vector, vector_section.metadata.quantization_type,
+            /*schema_name=*/
+            tokenized_document.document_wrapper().document().schema()));
       }
     }
     ICING_RETURN_IF_ERROR(embedding_index_.CommitBufferToIndex());

diff --git a/icing/index/embedding-indexing-handler_test.cc b/icing/index/embedding-indexing-handler_test.cc
index f2352a5..3435951 100644
--- a/icing/index/embedding-indexing-handler_test.cc
+++ b/icing/index/embedding-indexing-handler_test.cc

@@ -14,6 +14,7 @@
 
 #include "icing/index/embedding-indexing-handler.h"
 
+#include <cstdint>
 #include <initializer_list>
 #include <memory>
 #include <string>
@@ -31,6 +32,7 @@
 #include "icing/index/embed/embedding-index.h"
 #include "icing/index/embed/quantizer.h"
 #include "icing/index/hit/hit.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/portable/platform.h"
 #include "icing/proto/document_wrapper.pb.h"
 #include "icing/proto/schema.pb.h"
@@ -198,20 +200,25 @@
         filesystem_.CreateDirectoryRecursively(document_store_dir_.c_str()));
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult doc_store_create_result,
-        DocumentStore::Create(&filesystem_, document_store_dir_, &fake_clock_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, document_store_dir_, &fake_clock_,
+            schema_store_.get(), feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     document_store_ = std::move(doc_store_create_result.document_store);
 
     ICING_ASSERT_OK_AND_ASSIGN(
         embedding_index_,
         EmbeddingIndex::Create(&filesystem_, embedding_index_working_path_,
-                               &fake_clock_, feature_flags_.get()));
+                               &fake_clock_, feature_flags_.get(),
+                               /*num_shards=*/32));
   }
 
   void TearDown() override {
@@ -269,13 +276,18 @@
           .AddVectorProperty(std::string(kPropertyNonIndexableEmbedding),
                              CreateVector("model", {1.1, 1.2, 1.3}))
           .Build();
+  uint32_t shard_id = embedding_index_->GetShardId(
+      /*dimension=*/3, /*model_signature=*/"model",
+      /*schema_name=*/kFakeType);
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document)));
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result,
-      document_store_->Put(tokenized_document.document()));
+      document_store_->Put(tokenized_document.document_wrapper()));
   DocumentId document_id = put_result.new_document_id;
 
   ASSERT_THAT(embedding_index_->last_added_document_id(),
@@ -310,20 +322,22 @@
               IsOkAndHolds(ElementsAre(hit1, hit2, quantized_hit1,
                                        quantized_hit2, hit3)));
   // Check unquantized embedding data
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
+  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get(), shard_id),
               ElementsAre(0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.1, 0.2, 0.3));
   // Check quantized embedding data
-  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(),
+  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(shard_id),
               Eq(6 + 2 * sizeof(Quantizer)));
   EXPECT_THAT(
-      GetAndRestoreQuantizedEmbeddingVectorFromIndex(embedding_index_.get(),
-                                                     quantized_hit1,
-                                                     /*dimension=*/3),
+      GetAndRestoreQuantizedEmbeddingVectorFromIndex(
+          embedding_index_.get(), quantized_hit1,
+          /*dimension=*/3, /*model_signature=*/"model",
+          /*schema_name=*/kFakeType),
       IsOkAndHolds(Pointwise(FloatNear(kEpsQuantized), {0.1, 0.2, 0.3})));
   EXPECT_THAT(
-      GetAndRestoreQuantizedEmbeddingVectorFromIndex(embedding_index_.get(),
-                                                     quantized_hit2,
-                                                     /*dimension=*/3),
+      GetAndRestoreQuantizedEmbeddingVectorFromIndex(
+          embedding_index_.get(), quantized_hit2,
+          /*dimension=*/3, /*model_signature=*/"model",
+          /*schema_name=*/kFakeType),
       IsOkAndHolds(Pointwise(FloatNear(kEpsQuantized), {0.4, 0.5, 0.6})));
 
   EXPECT_THAT(embedding_index_->last_added_document_id(), Eq(document_id));
@@ -344,13 +358,18 @@
           .AddVectorProperty(std::string(kPropertyNonIndexableEmbedding),
                              CreateVector("model", {1.1, 1.2, 1.3}))
           .Build();
+  uint32_t shard_id = embedding_index_->GetShardId(
+      /*dimension=*/3, /*model_signature=*/"model",
+      /*schema_name=*/kFakeType);
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document)));
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result,
-      document_store_->Put(tokenized_document.document()));
+      document_store_->Put(tokenized_document.document_wrapper()));
   DocumentId document_id = put_result.new_document_id;
 
   ASSERT_THAT(embedding_index_->last_added_document_id(),
@@ -370,7 +389,8 @@
   EXPECT_THAT(GetEmbeddingHitsFromIndex(embedding_index_.get(), /*dimension=*/3,
                                         /*model_signature=*/"model"),
               IsOkAndHolds(IsEmpty()));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()), IsEmpty());
+  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get(), shard_id),
+              IsEmpty());
   EXPECT_THAT(embedding_index_->last_added_document_id(), Eq(document_id));
 }
 
@@ -401,13 +421,18 @@
           .AddVectorProperty(std::string(kPropertyFullDocEmbedding),
                              CreateVector("model", {2.1, 2.2, 2.3}))
           .Build();
+  uint32_t shard_id = embedding_index_->GetShardId(
+      /*dimension=*/3, /*model_signature=*/"model",
+      /*schema_name=*/kFakeCollectionType);
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document)));
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result,
-      document_store_->Put(tokenized_document.document()));
+      document_store_->Put(tokenized_document.document_wrapper()));
   DocumentId document_id = put_result.new_document_id;
 
   ASSERT_THAT(embedding_index_->last_added_document_id(),
@@ -445,20 +470,22 @@
                                        quantized_hit2, hit3, hit4)));
   // Check unquantized embedding data
   EXPECT_THAT(
-      GetRawEmbeddingDataFromIndex(embedding_index_.get()),
+      GetRawEmbeddingDataFromIndex(embedding_index_.get(), shard_id),
       ElementsAre(0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.1, 0.2, 0.3, 2.1, 2.2, 2.3));
   // Check quantized embedding data
-  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(),
+  EXPECT_THAT(embedding_index_->GetTotalQuantizedVectorSize(shard_id),
               Eq(6 + 2 * sizeof(Quantizer)));
   EXPECT_THAT(
-      GetAndRestoreQuantizedEmbeddingVectorFromIndex(embedding_index_.get(),
-                                                     quantized_hit1,
-                                                     /*dimension=*/3),
+      GetAndRestoreQuantizedEmbeddingVectorFromIndex(
+          embedding_index_.get(), quantized_hit1,
+          /*dimension=*/3, /*model_signature=*/"model",
+          /*schema_name=*/kFakeCollectionType),
       IsOkAndHolds(Pointwise(FloatNear(kEpsQuantized), {0.1, 0.2, 0.3})));
   EXPECT_THAT(
-      GetAndRestoreQuantizedEmbeddingVectorFromIndex(embedding_index_.get(),
-                                                     quantized_hit2,
-                                                     /*dimension=*/3),
+      GetAndRestoreQuantizedEmbeddingVectorFromIndex(
+          embedding_index_.get(), quantized_hit2,
+          /*dimension=*/3, /*model_signature=*/"model",
+          /*schema_name=*/kFakeCollectionType),
       IsOkAndHolds(Pointwise(FloatNear(kEpsQuantized), {0.4, 0.5, 0.6})));
 
   EXPECT_THAT(embedding_index_->last_added_document_id(), Eq(document_id));
@@ -480,11 +507,16 @@
           .AddVectorProperty(std::string(kPropertyNonIndexableEmbedding),
                              CreateVector("model", {1.1, 1.2, 1.3}))
           .Build();
+  uint32_t shard_id = embedding_index_->GetShardId(
+      /*dimension=*/3, /*model_signature=*/"model",
+      /*schema_name=*/kFakeType);
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document)));
-  ICING_ASSERT_OK(document_store_->Put(tokenized_document.document()));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document)));
+  ICING_ASSERT_OK(document_store_->Put(tokenized_document.document_wrapper()));
 
   static constexpr DocumentId kCurrentDocumentId = 3;
   embedding_index_->set_last_added_document_id(kCurrentDocumentId);
@@ -510,7 +542,8 @@
                                         /*model_signature=*/"model"),
               IsOkAndHolds(IsEmpty()));
   EXPECT_TRUE(embedding_index_->is_empty());
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()), IsEmpty());
+  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get(), shard_id),
+              IsEmpty());
 
   // Recovery mode should get the same result.
   EXPECT_THAT(
@@ -525,7 +558,8 @@
                                         /*model_signature=*/"model"),
               IsOkAndHolds(IsEmpty()));
   EXPECT_TRUE(embedding_index_->is_empty());
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()), IsEmpty());
+  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get(), shard_id),
+              IsEmpty());
 }
 
 TEST_F(EmbeddingIndexingHandlerTest,
@@ -544,13 +578,18 @@
           .AddVectorProperty(std::string(kPropertyNonIndexableEmbedding),
                              CreateVector("model", {1.1, 1.2, 1.3}))
           .Build();
+  uint32_t shard_id = embedding_index_->GetShardId(
+      /*dimension=*/3, /*model_signature=*/"model",
+      /*schema_name=*/kFakeType);
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document)));
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result,
-      document_store_->Put(tokenized_document.document()));
+      document_store_->Put(tokenized_document.document_wrapper()));
   DocumentId document_id = put_result.new_document_id;
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -574,7 +613,8 @@
                                         /*model_signature=*/"model"),
               IsOkAndHolds(IsEmpty()));
   EXPECT_TRUE(embedding_index_->is_empty());
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()), IsEmpty());
+  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get(), shard_id),
+              IsEmpty());
 
   // Handling document with document_id < last_added_document_id should cause a
   // failure, and both index data and last_added_document_id should remain
@@ -592,7 +632,8 @@
                                         /*model_signature=*/"model"),
               IsOkAndHolds(IsEmpty()));
   EXPECT_TRUE(embedding_index_->is_empty());
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()), IsEmpty());
+  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get(), shard_id),
+              IsEmpty());
 }
 
 TEST_F(EmbeddingIndexingHandlerTest,
@@ -625,21 +666,28 @@
           .AddVectorProperty(std::string(kPropertyNonIndexableEmbedding),
                              CreateVector("model", {11.1, 11.2, 11.3}))
           .Build();
+  uint32_t shard_id = embedding_index_->GetShardId(
+      /*dimension=*/3, /*model_signature=*/"model",
+      /*schema_name=*/kFakeType);
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document1,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document1)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document1)));
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document2,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document2)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document2)));
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result1,
-      document_store_->Put(tokenized_document1.document()));
+      document_store_->Put(tokenized_document1.document_wrapper()));
   DocumentId document_id1 = put_result1.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result2,
-      document_store_->Put(tokenized_document2.document()));
+      document_store_->Put(tokenized_document2.document_wrapper()));
   DocumentId document_id2 = put_result2.new_document_id;
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -667,7 +715,7 @@
                        /*location=*/3),
           EmbeddingHit(BasicHit(kSectionIdTitleEmbedding, /*document_id=*/0),
                        /*location=*/6))));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
+  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get(), shard_id),
               ElementsAre(0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.1, 0.2, 0.3));
 
   // Handle document with document_id == last_added_document_id in recovery
@@ -694,7 +742,7 @@
                        /*location=*/3),
           EmbeddingHit(BasicHit(kSectionIdTitleEmbedding, /*document_id=*/0),
                        /*location=*/6))));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
+  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get(), shard_id),
               ElementsAre(0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.1, 0.2, 0.3));
 
   // Handle document with document_id < last_added_document_id in recovery mode.
@@ -720,7 +768,7 @@
                        /*location=*/3),
           EmbeddingHit(BasicHit(kSectionIdTitleEmbedding, /*document_id=*/0),
                        /*location=*/6))));
-  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get()),
+  EXPECT_THAT(GetRawEmbeddingDataFromIndex(embedding_index_.get(), shard_id),
               ElementsAre(0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.1, 0.2, 0.3));
 }
 

diff --git a/icing/index/index-processor_benchmark.cc b/icing/index/index-processor_benchmark.cc
index 7adaac4..ff41adc 100644
--- a/icing/index/index-processor_benchmark.cc
+++ b/icing/index/index-processor_benchmark.cc

@@ -158,11 +158,13 @@
 
 std::unique_ptr<Index> CreateIndex(const IcingFilesystem& icing_filesystem,
                                    const Filesystem& filesystem,
-                                   const std::string& index_dir) {
+                                   const std::string& index_dir,
+                                   const FeatureFlags& feature_flags) {
   Index::Options options(index_dir, /*index_merge_size=*/1024 * 1024 * 10,
                          /*lite_index_sort_at_indexing=*/true,
                          /*lite_index_sort_size=*/1024 * 8);
-  return Index::Create(options, &filesystem, &icing_filesystem).ValueOrDie();
+  return Index::Create(options, &filesystem, &icing_filesystem, &feature_flags)
+      .ValueOrDie();
 }
 
 std::unique_ptr<Normalizer> CreateNormalizer() {
@@ -235,7 +237,7 @@
               IsTrue());
 
   std::unique_ptr<Index> index =
-      CreateIndex(icing_filesystem, filesystem, index_dir);
+      CreateIndex(icing_filesystem, filesystem, index_dir, feature_flags);
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<NumericIndex<int64_t>> integer_index,
       IntegerIndex::Create(filesystem, integer_index_dir,
@@ -258,8 +260,9 @@
 
   DocumentProto input_document = CreateDocumentWithOneProperty(state.range(0));
   TokenizedDocument tokenized_document(std::move(
-      TokenizedDocument::Create(schema_store.get(), language_segmenter.get(),
-                                input_document)
+      TokenizedDocument::Create(
+          schema_store.get(), language_segmenter.get(),
+          /*current_time_ms=*/clock.GetSystemTimeMilliseconds(), input_document)
           .ValueOrDie()));
 
   DocumentId old_document_id = kInvalidDocumentId;
@@ -314,7 +317,7 @@
               IsTrue());
 
   std::unique_ptr<Index> index =
-      CreateIndex(icing_filesystem, filesystem, index_dir);
+      CreateIndex(icing_filesystem, filesystem, index_dir, feature_flags);
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<NumericIndex<int64_t>> integer_index,
       IntegerIndex::Create(filesystem, integer_index_dir,
@@ -338,8 +341,9 @@
   DocumentProto input_document =
       CreateDocumentWithTenProperties(state.range(0));
   TokenizedDocument tokenized_document(std::move(
-      TokenizedDocument::Create(schema_store.get(), language_segmenter.get(),
-                                input_document)
+      TokenizedDocument::Create(
+          schema_store.get(), language_segmenter.get(),
+          /*current_time_ms=*/clock.GetSystemTimeMilliseconds(), input_document)
           .ValueOrDie()));
 
   DocumentId old_document_id = kInvalidDocumentId;
@@ -394,7 +398,7 @@
               IsTrue());
 
   std::unique_ptr<Index> index =
-      CreateIndex(icing_filesystem, filesystem, index_dir);
+      CreateIndex(icing_filesystem, filesystem, index_dir, feature_flags);
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<NumericIndex<int64_t>> integer_index,
       IntegerIndex::Create(filesystem, integer_index_dir,
@@ -418,8 +422,9 @@
   DocumentProto input_document =
       CreateDocumentWithDiacriticLetters(state.range(0));
   TokenizedDocument tokenized_document(std::move(
-      TokenizedDocument::Create(schema_store.get(), language_segmenter.get(),
-                                input_document)
+      TokenizedDocument::Create(
+          schema_store.get(), language_segmenter.get(),
+          /*current_time_ms=*/clock.GetSystemTimeMilliseconds(), input_document)
           .ValueOrDie()));
 
   DocumentId old_document_id = kInvalidDocumentId;
@@ -474,7 +479,7 @@
               IsTrue());
 
   std::unique_ptr<Index> index =
-      CreateIndex(icing_filesystem, filesystem, index_dir);
+      CreateIndex(icing_filesystem, filesystem, index_dir, feature_flags);
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<NumericIndex<int64_t>> integer_index,
       IntegerIndex::Create(filesystem, integer_index_dir,
@@ -497,8 +502,9 @@
 
   DocumentProto input_document = CreateDocumentWithHiragana(state.range(0));
   TokenizedDocument tokenized_document(std::move(
-      TokenizedDocument::Create(schema_store.get(), language_segmenter.get(),
-                                input_document)
+      TokenizedDocument::Create(
+          schema_store.get(), language_segmenter.get(),
+          /*current_time_ms=*/clock.GetSystemTimeMilliseconds(), input_document)
           .ValueOrDie()));
 
   DocumentId old_document_id = kInvalidDocumentId;

diff --git a/icing/index/index-processor_test.cc b/icing/index/index-processor_test.cc
index cba75f6..2b2f4e7 100644
--- a/icing/index/index-processor_test.cc
+++ b/icing/index/index-processor_test.cc

@@ -48,6 +48,7 @@
 #include "icing/join/qualified-id-join-indexing-handler.h"
 #include "icing/legacy/index/icing-filesystem.h"
 #include "icing/legacy/index/icing-mock-filesystem.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/portable/platform.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/schema.pb.h"
@@ -179,7 +180,8 @@
                            /*lite_index_sort_at_indexing=*/true,
                            /*lite_index_sort_size=*/1024 * 8);
     ICING_ASSERT_OK_AND_ASSIGN(
-        index_, Index::Create(options, &filesystem_, &icing_filesystem_));
+        index_, Index::Create(options, &filesystem_, &icing_filesystem_,
+                              feature_flags_.get()));
 
     ICING_ASSERT_OK_AND_ASSIGN(
         integer_index_,
@@ -288,14 +290,18 @@
     ASSERT_TRUE(filesystem_.CreateDirectoryRecursively(doc_store_dir_.c_str()));
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
-        DocumentStore::Create(&filesystem_, doc_store_dir_, &fake_clock_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, doc_store_dir_, &fake_clock_, schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     doc_store_ = std::move(create_result.document_store);
 
     ICING_ASSERT_OK_AND_ASSIGN(
@@ -311,7 +317,8 @@
         std::unique_ptr<QualifiedIdJoinIndexingHandler>
             qualified_id_join_indexing_handler,
         QualifiedIdJoinIndexingHandler::Create(&fake_clock_, doc_store_.get(),
-                                               qualified_id_join_index_.get()));
+                                               qualified_id_join_index_.get(),
+                                               feature_flags_.get()));
     std::vector<std::unique_ptr<DataIndexingHandler>> handlers;
     handlers.push_back(std::move(term_indexing_handler));
     handlers.push_back(std::move(integer_section_indexing_handler));
@@ -393,8 +400,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(
       index_processor_->IndexDocument(tokenized_document, kDocumentId0,
                                       /*old_document_id=*/kInvalidDocumentId),
@@ -411,8 +420,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(
       index_processor_->IndexDocument(tokenized_document, kDocumentId0,
                                       /*old_document_id=*/kInvalidDocumentId),
@@ -429,8 +440,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(
       index_processor_->IndexDocument(tokenized_document, kDocumentId0,
                                       /*old_document_id=*/kInvalidDocumentId),
@@ -466,8 +479,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(
       index_processor_->IndexDocument(tokenized_document, kDocumentId0,
                                       /*old_document_id=*/kInvalidDocumentId),
@@ -489,8 +504,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(
       index_processor_->IndexDocument(tokenized_document, kDocumentId1,
                                       /*old_document_id=*/kInvalidDocumentId),
@@ -551,8 +568,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(
       index_processor_->IndexDocument(tokenized_document, kDocumentId0,
                                       /*old_document_id=*/kInvalidDocumentId),
@@ -580,8 +599,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(
       index_processor_->IndexDocument(tokenized_document, kDocumentId0,
                                       /*old_document_id=*/kInvalidDocumentId),
@@ -618,8 +639,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(
       index_processor_->IndexDocument(tokenized_document, kDocumentId0,
                                       /*old_document_id=*/kInvalidDocumentId),
@@ -642,8 +665,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(
       index_processor_->IndexDocument(tokenized_document, kDocumentId0,
                                       /*old_document_id=*/kInvalidDocumentId),
@@ -680,8 +705,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(
       index_processor_->IndexDocument(tokenized_document, kDocumentId0,
                                       /*old_document_id=*/kInvalidDocumentId),
@@ -727,8 +754,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(tokenized_document.num_string_tokens(), Eq(8));
 
   EXPECT_THAT(
@@ -773,8 +802,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(
       index_processor_->IndexDocument(tokenized_document, kDocumentId0,
                                       /*old_document_id=*/kInvalidDocumentId),
@@ -789,8 +820,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(
       index_processor_->IndexDocument(tokenized_document, kDocumentId1,
                                       /*old_document_id=*/kInvalidDocumentId),
@@ -817,8 +850,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(
       index_processor_->IndexDocument(tokenized_document, kDocumentId0,
                                       /*old_document_id=*/kInvalidDocumentId),
@@ -833,8 +868,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(
       index_processor_->IndexDocument(tokenized_document, kDocumentId1,
                                       /*old_document_id=*/kInvalidDocumentId),
@@ -864,8 +901,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(
       index_processor_->IndexDocument(tokenized_document, kDocumentId1,
                                       /*old_document_id=*/kInvalidDocumentId),
@@ -888,8 +927,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(
       index_processor_->IndexDocument(tokenized_document, kDocumentId0,
                                       /*old_document_id=*/kInvalidDocumentId),
@@ -928,7 +969,8 @@
       std::unique_ptr<QualifiedIdJoinIndexingHandler>
           qualified_id_join_indexing_handler,
       QualifiedIdJoinIndexingHandler::Create(&fake_clock_, doc_store_.get(),
-                                             qualified_id_join_index_.get()));
+                                             qualified_id_join_index_.get(),
+                                             feature_flags_.get()));
   std::vector<std::unique_ptr<DataIndexingHandler>> handlers;
   handlers.push_back(std::move(term_indexing_handler));
   handlers.push_back(std::move(integer_section_indexing_handler));
@@ -946,8 +988,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(
       index_processor.IndexDocument(tokenized_document, kDocumentId1,
                                     /*old_document_id=*/kInvalidDocumentId),
@@ -971,8 +1015,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(
       index_processor.IndexDocument(tokenized_document, kDocumentId0,
                                     /*old_document_id=*/kInvalidDocumentId),
@@ -1013,8 +1059,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(
       index_processor_->IndexDocument(tokenized_document, kDocumentId0,
                                       /*old_document_id=*/kInvalidDocumentId),
@@ -1049,8 +1097,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document_one));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document_one));
   EXPECT_THAT(
       index_processor_->IndexDocument(tokenized_document, kDocumentId0,
                                       /*old_document_id=*/kInvalidDocumentId),
@@ -1072,14 +1122,17 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   Index::Options options(index_dir_,
                          /*index_merge_size=*/document.ByteSizeLong() * 100,
                          /*lite_index_sort_at_indexing=*/true,
                          /*lite_index_sort_size=*/64);
   ICING_ASSERT_OK_AND_ASSIGN(
-      index_, Index::Create(options, &filesystem_, &icing_filesystem_));
+      index_, Index::Create(options, &filesystem_, &icing_filesystem_,
+                            feature_flags_.get()));
 
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<TermIndexingHandler> term_indexing_handler,
@@ -1140,8 +1193,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
 
   // 2. Recreate the index with the mock filesystem and a merge size that will
   // only allow one document to be added before requiring a merge.
@@ -1150,8 +1205,8 @@
                          /*lite_index_sort_at_indexing=*/true,
                          /*lite_index_sort_size=*/16);
   ICING_ASSERT_OK_AND_ASSIGN(
-      index_,
-      Index::Create(options, &filesystem_, mock_icing_filesystem_.get()));
+      index_, Index::Create(options, &filesystem_, mock_icing_filesystem_.get(),
+                            feature_flags_.get()));
 
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<TermIndexingHandler> term_indexing_handler,
@@ -1198,8 +1253,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(tokenized_document.num_string_tokens(), Eq(1));
 
   EXPECT_THAT(
@@ -1232,8 +1289,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(tokenized_document.num_string_tokens(), Eq(1));
 
   EXPECT_THAT(
@@ -1268,8 +1327,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(tokenized_document.num_string_tokens(), Eq(1));
 
   EXPECT_THAT(
@@ -1301,8 +1362,10 @@
                                .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(tokenized_document.num_string_tokens(), Eq(7));
 
   EXPECT_THAT(
@@ -1352,8 +1415,10 @@
                                .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(tokenized_document.num_string_tokens(), Eq(7));
 
   EXPECT_THAT(
@@ -1384,8 +1449,10 @@
                                .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(tokenized_document.num_string_tokens(), Eq(7));
 
   EXPECT_THAT(
@@ -1437,8 +1504,10 @@
                                .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(tokenized_document.num_string_tokens(), Eq(7));
 
   EXPECT_THAT(
@@ -1470,8 +1539,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(tokenized_document.num_string_tokens(), Eq(7));
 
   EXPECT_THAT(
@@ -1530,8 +1601,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(tokenized_document.num_string_tokens(), Eq(8));
 
   EXPECT_THAT(
@@ -1574,8 +1647,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(tokenized_document.num_string_tokens(), Eq(7));
 
   EXPECT_THAT(
@@ -1628,8 +1703,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   EXPECT_THAT(tokenized_document.num_string_tokens(), Eq(8));
 
   EXPECT_THAT(
@@ -1682,8 +1759,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   // Expected to have 1 integer section.
   EXPECT_THAT(tokenized_document.integer_sections(), SizeIs(1));
 
@@ -1714,8 +1793,10 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          document));
   // Expected to have 1 integer section.
   EXPECT_THAT(tokenized_document.integer_sections(), SizeIs(1));
 

diff --git a/icing/index/index.cc b/icing/index/index.cc
index a03f07a..2da5cc8 100644
--- a/icing/index/index.cc
+++ b/icing/index/index.cc

@@ -27,6 +27,7 @@
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "icing/absl_ports/canonical_errors.h"
 #include "icing/absl_ports/str_cat.h"
+#include "icing/feature-flags.h"
 #include "icing/file/filesystem.h"
 #include "icing/index/hit/hit.h"
 #include "icing/index/iterator/doc-hit-info-iterator-or.h"
@@ -148,9 +149,11 @@
 
 libtextclassifier3::StatusOr<std::unique_ptr<Index>> Index::Create(
     const Options& options, const Filesystem* filesystem,
-    const IcingFilesystem* icing_filesystem) {
+    const IcingFilesystem* icing_filesystem,
+    const FeatureFlags* feature_flags) {
   ICING_RETURN_ERROR_IF_NULL(filesystem);
   ICING_RETURN_ERROR_IF_NULL(icing_filesystem);
+  ICING_RETURN_ERROR_IF_NULL(feature_flags);
 
   ICING_ASSIGN_OR_RETURN(LiteIndex::Options lite_index_options,
                          CreateLiteIndexOptions(options));
@@ -175,9 +178,9 @@
       std::unique_ptr<MainIndex> main_index,
       MainIndex::Create(MakeMainIndexFilepath(options.base_dir), filesystem,
                         icing_filesystem));
-  return std::unique_ptr<Index>(new Index(options, std::move(term_id_codec),
-                                          std::move(lite_index),
-                                          std::move(main_index), filesystem));
+  return std::unique_ptr<Index>(
+      new Index(options, std::move(term_id_codec), std::move(lite_index),
+                std::move(main_index), filesystem, feature_flags));
 }
 
 /* static */ libtextclassifier3::StatusOr<int> Index::ReadFlashIndexMagic(

diff --git a/icing/index/index.h b/icing/index/index.h
index fe3415c..774815c 100644
--- a/icing/index/index.h
+++ b/icing/index/index.h

@@ -25,6 +25,7 @@
 
 #include "icing/text_classifier/lib3/utils/base/status.h"
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
+#include "icing/feature-flags.h"
 #include "icing/file/filesystem.h"
 #include "icing/index/hit/hit.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
@@ -95,7 +96,8 @@
   //   INTERNAL on I/O error
   static libtextclassifier3::StatusOr<std::unique_ptr<Index>> Create(
       const Options& options, const Filesystem* filesystem,
-      const IcingFilesystem* icing_filesystem);
+      const IcingFilesystem* icing_filesystem,
+      const FeatureFlags* feature_flags);
 
   // Reads magic from existing flash (main) index file header. We need this
   // during Icing initialization phase to determine the version.
@@ -312,7 +314,9 @@
     ICING_RETURN_IF_ERROR(main_index_->AddHits(
         *term_id_codec_, std::move(outputs.backfill_map),
         std::move(term_id_hit_pairs), lite_index_->last_added_document_id()));
-    ICING_RETURN_IF_ERROR(main_index_->PersistToDisk());
+    if (!feature_flags_.enable_optimize_improvements()) {
+      ICING_RETURN_IF_ERROR(main_index_->PersistToDisk());
+    }
     return lite_index_->Reset();
   }
 
@@ -340,14 +344,18 @@
       DocumentId new_last_added_document_id);
 
  private:
-  Index(const Options& options, std::unique_ptr<TermIdCodec> term_id_codec,
-        std::unique_ptr<LiteIndex> lite_index,
-        std::unique_ptr<MainIndex> main_index, const Filesystem* filesystem)
+  explicit Index(const Options& options,
+                 std::unique_ptr<TermIdCodec> term_id_codec,
+                 std::unique_ptr<LiteIndex> lite_index,
+                 std::unique_ptr<MainIndex> main_index,
+                 const Filesystem* filesystem,
+                 const FeatureFlags* feature_flags)
       : lite_index_(std::move(lite_index)),
         main_index_(std::move(main_index)),
         options_(options),
         term_id_codec_(std::move(term_id_codec)),
-        filesystem_(filesystem) {}
+        filesystem_(filesystem),
+        feature_flags_(*feature_flags) {}
 
   libtextclassifier3::StatusOr<std::vector<TermMetadata>> FindLiteTermsByPrefix(
       const std::string& prefix,
@@ -358,7 +366,8 @@
   std::unique_ptr<MainIndex> main_index_;
   const Options options_;
   std::unique_ptr<TermIdCodec> term_id_codec_;
-  const Filesystem* filesystem_;
+  const Filesystem* filesystem_;       // Does not own.
+  const FeatureFlags& feature_flags_;  // Does not own.
 };
 
 }  // namespace lib

diff --git a/icing/index/index_test.cc b/icing/index/index_test.cc
index 7f73590..4925f38 100644
--- a/icing/index/index_test.cc
+++ b/icing/index/index_test.cc

@@ -30,6 +30,7 @@
 #include "icing/text_classifier/lib3/utils/base/status.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
+#include "icing/feature-flags.h"
 #include "icing/file/filesystem.h"
 #include "icing/index/hit/doc-hit-info.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
@@ -45,9 +46,11 @@
 #include "icing/testing/always-true-suggestion-result-checker-impl.h"
 #include "icing/testing/common-matchers.h"
 #include "icing/testing/random-string.h"
+#include "icing/testing/test-feature-flags.h"
 #include "icing/testing/tmp-directory.h"
 #include "icing/util/crc32.h"
 #include "icing/util/logging.h"
+#include "icing/util/status-macros.h"
 
 namespace icing {
 namespace lib {
@@ -77,12 +80,15 @@
 class IndexTest : public Test {
  protected:
   void SetUp() override {
+    feature_flags_ = std::make_unique<FeatureFlags>(GetTestFeatureFlags());
+
     index_dir_ = GetTestTempDir() + "/index_test/";
     Index::Options options(index_dir_, /*index_merge_size=*/1024 * 1024,
                            /*lite_index_sort_at_indexing=*/true,
                            /*lite_index_sort_size=*/1024 * 8);
     ICING_ASSERT_OK_AND_ASSIGN(
-        index_, Index::Create(options, &filesystem_, &icing_filesystem_));
+        index_, Index::Create(options, &filesystem_, &icing_filesystem_,
+                              feature_flags_.get()));
   }
 
   void TearDown() override {
@@ -109,6 +115,7 @@
     return GetHits(std::move(itr));
   }
 
+  std::unique_ptr<FeatureFlags> feature_flags_;
   Filesystem filesystem_;
   IcingFilesystem icing_filesystem_;
   std::string index_dir_;
@@ -154,12 +161,15 @@
   Index::Options options(index_dir_, /*index_merge_size=*/1024 * 1024,
                          /*lite_index_sort_at_indexing=*/true,
                          /*lite_index_sort_size=*/1024 * 8);
-  EXPECT_THAT(
-      Index::Create(options, &filesystem_, /*icing_filesystem=*/nullptr),
-      StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
-  EXPECT_THAT(
-      Index::Create(options, /*filesystem=*/nullptr, &icing_filesystem_),
-      StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
+  EXPECT_THAT(Index::Create(options, &filesystem_, /*icing_filesystem=*/nullptr,
+                            feature_flags_.get()),
+              StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
+  EXPECT_THAT(Index::Create(options, /*filesystem=*/nullptr, &icing_filesystem_,
+                            feature_flags_.get()),
+              StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
+  EXPECT_THAT(Index::Create(options, &filesystem_, &icing_filesystem_,
+                            /*feature_flags=*/nullptr),
+              StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
 }
 
 TEST_F(IndexTest, EmptyIndex) {
@@ -206,7 +216,8 @@
                          /*lite_index_sort_at_indexing=*/false,
                          /*lite_index_sort_size=*/16);
   ICING_ASSERT_OK_AND_ASSIGN(
-      index_, Index::Create(options, &filesystem_, &icing_filesystem_));
+      index_, Index::Create(options, &filesystem_, &icing_filesystem_,
+                            feature_flags_.get()));
 
   Index::Editor edit =
       index_->Edit(kDocumentId0, kSectionId2, /*namespace_id=*/0);
@@ -221,7 +232,8 @@
                            /*lite_index_sort_at_indexing=*/true,
                            /*lite_index_sort_size=*/16);
   ICING_ASSERT_OK_AND_ASSIGN(
-      index_, Index::Create(options, &filesystem_, &icing_filesystem_));
+      index_, Index::Create(options, &filesystem_, &icing_filesystem_,
+                            feature_flags_.get()));
 
   // Check that the index is sorted after recreating with
   // lite_index_sort_at_indexing, with the unsorted HitBuffer exceeding the sort
@@ -307,44 +319,48 @@
                           TermMatchType::EXACT_ONLY));
 
   // Before Advance().
-  EXPECT_THAT(
-      itr->GetCallStats(),
-      EqualsDocHitInfoIteratorCallStats(
-          /*num_leaf_advance_calls_lite_index=*/0,
-          /*num_leaf_advance_calls_main_index=*/0,
-          /*num_leaf_advance_calls_integer_index=*/0,
-          /*num_leaf_advance_calls_no_index=*/0, /*num_blocks_inspected=*/0));
+  EXPECT_THAT(itr->GetCallStats(),
+              EqualsDocHitInfoIteratorCallStats(
+                  /*num_leaf_advance_calls_lite_index=*/0,
+                  /*num_leaf_advance_calls_main_index=*/0,
+                  /*num_leaf_advance_calls_integer_index=*/0,
+                  /*num_leaf_advance_calls_no_index=*/0,
+                  /*num_blocks_inspected=*/0,
+                  DocHitInfoIterator::CallStats::EmbeddingStats()));
 
   // 1st Advance().
   ICING_ASSERT_OK(itr->Advance());
-  EXPECT_THAT(
-      itr->GetCallStats(),
-      EqualsDocHitInfoIteratorCallStats(
-          /*num_leaf_advance_calls_lite_index=*/0,
-          /*num_leaf_advance_calls_main_index=*/1,
-          /*num_leaf_advance_calls_integer_index=*/0,
-          /*num_leaf_advance_calls_no_index=*/0, /*num_blocks_inspected=*/1));
+  EXPECT_THAT(itr->GetCallStats(),
+              EqualsDocHitInfoIteratorCallStats(
+                  /*num_leaf_advance_calls_lite_index=*/0,
+                  /*num_leaf_advance_calls_main_index=*/1,
+                  /*num_leaf_advance_calls_integer_index=*/0,
+                  /*num_leaf_advance_calls_no_index=*/0,
+                  /*num_blocks_inspected=*/1,
+                  DocHitInfoIterator::CallStats::EmbeddingStats()));
 
   // 2nd Advance().
   ICING_ASSERT_OK(itr->Advance());
-  EXPECT_THAT(
-      itr->GetCallStats(),
-      EqualsDocHitInfoIteratorCallStats(
-          /*num_leaf_advance_calls_lite_index=*/0,
-          /*num_leaf_advance_calls_main_index=*/2,
-          /*num_leaf_advance_calls_integer_index=*/0,
-          /*num_leaf_advance_calls_no_index=*/0, /*num_blocks_inspected=*/1));
+  EXPECT_THAT(itr->GetCallStats(),
+              EqualsDocHitInfoIteratorCallStats(
+                  /*num_leaf_advance_calls_lite_index=*/0,
+                  /*num_leaf_advance_calls_main_index=*/2,
+                  /*num_leaf_advance_calls_integer_index=*/0,
+                  /*num_leaf_advance_calls_no_index=*/0,
+                  /*num_blocks_inspected=*/1,
+                  DocHitInfoIterator::CallStats::EmbeddingStats()));
 
   // 3rd Advance().
   ASSERT_THAT(itr->Advance(),
               StatusIs(libtextclassifier3::StatusCode::RESOURCE_EXHAUSTED));
-  EXPECT_THAT(
-      itr->GetCallStats(),
-      EqualsDocHitInfoIteratorCallStats(
-          /*num_leaf_advance_calls_lite_index=*/0,
-          /*num_leaf_advance_calls_main_index=*/2,
-          /*num_leaf_advance_calls_integer_index=*/0,
-          /*num_leaf_advance_calls_no_index=*/0, /*num_blocks_inspected=*/1));
+  EXPECT_THAT(itr->GetCallStats(),
+              EqualsDocHitInfoIteratorCallStats(
+                  /*num_leaf_advance_calls_lite_index=*/0,
+                  /*num_leaf_advance_calls_main_index=*/2,
+                  /*num_leaf_advance_calls_integer_index=*/0,
+                  /*num_leaf_advance_calls_no_index=*/0,
+                  /*num_blocks_inspected=*/1,
+                  DocHitInfoIterator::CallStats::EmbeddingStats()));
 }
 
 TEST_F(IndexTest, IteratorGetCallStats_liteIndexOnly) {
@@ -366,44 +382,48 @@
                           TermMatchType::EXACT_ONLY));
 
   // Before Advance().
-  EXPECT_THAT(
-      itr->GetCallStats(),
-      EqualsDocHitInfoIteratorCallStats(
-          /*num_leaf_advance_calls_lite_index=*/0,
-          /*num_leaf_advance_calls_main_index=*/0,
-          /*num_leaf_advance_calls_integer_index=*/0,
-          /*num_leaf_advance_calls_no_index=*/0, /*num_blocks_inspected=*/0));
+  EXPECT_THAT(itr->GetCallStats(),
+              EqualsDocHitInfoIteratorCallStats(
+                  /*num_leaf_advance_calls_lite_index=*/0,
+                  /*num_leaf_advance_calls_main_index=*/0,
+                  /*num_leaf_advance_calls_integer_index=*/0,
+                  /*num_leaf_advance_calls_no_index=*/0,
+                  /*num_blocks_inspected=*/0,
+                  DocHitInfoIterator::CallStats::EmbeddingStats()));
 
   // 1st Advance().
   ICING_ASSERT_OK(itr->Advance());
-  EXPECT_THAT(
-      itr->GetCallStats(),
-      EqualsDocHitInfoIteratorCallStats(
-          /*num_leaf_advance_calls_lite_index=*/1,
-          /*num_leaf_advance_calls_main_index=*/0,
-          /*num_leaf_advance_calls_integer_index=*/0,
-          /*num_leaf_advance_calls_no_index=*/0, /*num_blocks_inspected=*/0));
+  EXPECT_THAT(itr->GetCallStats(),
+              EqualsDocHitInfoIteratorCallStats(
+                  /*num_leaf_advance_calls_lite_index=*/1,
+                  /*num_leaf_advance_calls_main_index=*/0,
+                  /*num_leaf_advance_calls_integer_index=*/0,
+                  /*num_leaf_advance_calls_no_index=*/0,
+                  /*num_blocks_inspected=*/0,
+                  DocHitInfoIterator::CallStats::EmbeddingStats()));
 
   // 2nd Advance().
   ICING_ASSERT_OK(itr->Advance());
-  EXPECT_THAT(
-      itr->GetCallStats(),
-      EqualsDocHitInfoIteratorCallStats(
-          /*num_leaf_advance_calls_lite_index=*/2,
-          /*num_leaf_advance_calls_main_index=*/0,
-          /*num_leaf_advance_calls_integer_index=*/0,
-          /*num_leaf_advance_calls_no_index=*/0, /*num_blocks_inspected=*/0));
+  EXPECT_THAT(itr->GetCallStats(),
+              EqualsDocHitInfoIteratorCallStats(
+                  /*num_leaf_advance_calls_lite_index=*/2,
+                  /*num_leaf_advance_calls_main_index=*/0,
+                  /*num_leaf_advance_calls_integer_index=*/0,
+                  /*num_leaf_advance_calls_no_index=*/0,
+                  /*num_blocks_inspected=*/0,
+                  DocHitInfoIterator::CallStats::EmbeddingStats()));
 
   // 3rd Advance().
   ASSERT_THAT(itr->Advance(),
               StatusIs(libtextclassifier3::StatusCode::RESOURCE_EXHAUSTED));
-  EXPECT_THAT(
-      itr->GetCallStats(),
-      EqualsDocHitInfoIteratorCallStats(
-          /*num_leaf_advance_calls_lite_index=*/2,
-          /*num_leaf_advance_calls_main_index=*/0,
-          /*num_leaf_advance_calls_integer_index=*/0,
-          /*num_leaf_advance_calls_no_index=*/0, /*num_blocks_inspected=*/0));
+  EXPECT_THAT(itr->GetCallStats(),
+              EqualsDocHitInfoIteratorCallStats(
+                  /*num_leaf_advance_calls_lite_index=*/2,
+                  /*num_leaf_advance_calls_main_index=*/0,
+                  /*num_leaf_advance_calls_integer_index=*/0,
+                  /*num_leaf_advance_calls_no_index=*/0,
+                  /*num_blocks_inspected=*/0,
+                  DocHitInfoIterator::CallStats::EmbeddingStats()));
 }
 
 TEST_F(IndexTest, IteratorGetCallStats) {
@@ -439,72 +459,78 @@
                           TermMatchType::EXACT_ONLY));
 
   // Before Advance().
-  EXPECT_THAT(
-      itr->GetCallStats(),
-      EqualsDocHitInfoIteratorCallStats(
-          /*num_leaf_advance_calls_lite_index=*/0,
-          /*num_leaf_advance_calls_main_index=*/0,
-          /*num_leaf_advance_calls_integer_index=*/0,
-          /*num_leaf_advance_calls_no_index=*/0, /*num_blocks_inspected=*/0));
+  EXPECT_THAT(itr->GetCallStats(),
+              EqualsDocHitInfoIteratorCallStats(
+                  /*num_leaf_advance_calls_lite_index=*/0,
+                  /*num_leaf_advance_calls_main_index=*/0,
+                  /*num_leaf_advance_calls_integer_index=*/0,
+                  /*num_leaf_advance_calls_no_index=*/0,
+                  /*num_blocks_inspected=*/0,
+                  DocHitInfoIterator::CallStats::EmbeddingStats()));
 
   // 1st Advance(). DocHitInfoIteratorOr will advance both left and right
   // iterator (i.e. lite and main index iterator) once, compare document ids,
   // and return the hit with larger document id. In this case, hit from lite
   // index will be chosen and returned.
   ICING_ASSERT_OK(itr->Advance());
-  EXPECT_THAT(
-      itr->GetCallStats(),
-      EqualsDocHitInfoIteratorCallStats(
-          /*num_leaf_advance_calls_lite_index=*/1,
-          /*num_leaf_advance_calls_main_index=*/1,
-          /*num_leaf_advance_calls_integer_index=*/0,
-          /*num_leaf_advance_calls_no_index=*/0, /*num_blocks_inspected=*/1));
+  EXPECT_THAT(itr->GetCallStats(),
+              EqualsDocHitInfoIteratorCallStats(
+                  /*num_leaf_advance_calls_lite_index=*/1,
+                  /*num_leaf_advance_calls_main_index=*/1,
+                  /*num_leaf_advance_calls_integer_index=*/0,
+                  /*num_leaf_advance_calls_no_index=*/0,
+                  /*num_blocks_inspected=*/1,
+                  DocHitInfoIterator::CallStats::EmbeddingStats()));
 
   // 2nd Advance(). Since lite index iterator has larger document id in the
   // previous round, we advance lite index iterator in this round. We still
   // choose and return hit from lite index.
   ICING_ASSERT_OK(itr->Advance());
-  EXPECT_THAT(
-      itr->GetCallStats(),
-      EqualsDocHitInfoIteratorCallStats(
-          /*num_leaf_advance_calls_lite_index=*/2,
-          /*num_leaf_advance_calls_main_index=*/1,
-          /*num_leaf_advance_calls_integer_index=*/0,
-          /*num_leaf_advance_calls_no_index=*/0, /*num_blocks_inspected=*/1));
+  EXPECT_THAT(itr->GetCallStats(),
+              EqualsDocHitInfoIteratorCallStats(
+                  /*num_leaf_advance_calls_lite_index=*/2,
+                  /*num_leaf_advance_calls_main_index=*/1,
+                  /*num_leaf_advance_calls_integer_index=*/0,
+                  /*num_leaf_advance_calls_no_index=*/0,
+                  /*num_blocks_inspected=*/1,
+                  DocHitInfoIterator::CallStats::EmbeddingStats()));
 
   // 3rd Advance(). Since lite index iterator has larger document id in the
   // previous round, we advance lite index iterator in this round. However,
   // there is no hit from lite index anymore, so we choose and return hit from
   // main index.
   ICING_ASSERT_OK(itr->Advance());
-  EXPECT_THAT(
-      itr->GetCallStats(),
-      EqualsDocHitInfoIteratorCallStats(
-          /*num_leaf_advance_calls_lite_index=*/2,
-          /*num_leaf_advance_calls_main_index=*/1,
-          /*num_leaf_advance_calls_integer_index=*/0,
-          /*num_leaf_advance_calls_no_index=*/0, /*num_blocks_inspected=*/1));
+  EXPECT_THAT(itr->GetCallStats(),
+              EqualsDocHitInfoIteratorCallStats(
+                  /*num_leaf_advance_calls_lite_index=*/2,
+                  /*num_leaf_advance_calls_main_index=*/1,
+                  /*num_leaf_advance_calls_integer_index=*/0,
+                  /*num_leaf_advance_calls_no_index=*/0,
+                  /*num_blocks_inspected=*/1,
+                  DocHitInfoIterator::CallStats::EmbeddingStats()));
 
   // 4th Advance(). Advance main index.
   ICING_ASSERT_OK(itr->Advance());
-  EXPECT_THAT(
-      itr->GetCallStats(),
-      EqualsDocHitInfoIteratorCallStats(
-          /*num_leaf_advance_calls_lite_index=*/2,
-          /*num_leaf_advance_calls_main_index=*/2,
-          /*num_leaf_advance_calls_integer_index=*/0,
-          /*num_leaf_advance_calls_no_index=*/0, /*num_blocks_inspected=*/1));
+  EXPECT_THAT(itr->GetCallStats(),
+              EqualsDocHitInfoIteratorCallStats(
+                  /*num_leaf_advance_calls_lite_index=*/2,
+                  /*num_leaf_advance_calls_main_index=*/2,
+                  /*num_leaf_advance_calls_integer_index=*/0,
+                  /*num_leaf_advance_calls_no_index=*/0,
+                  /*num_blocks_inspected=*/1,
+                  DocHitInfoIterator::CallStats::EmbeddingStats()));
 
   // 5th Advance(). Reach the end.
   ASSERT_THAT(itr->Advance(),
               StatusIs(libtextclassifier3::StatusCode::RESOURCE_EXHAUSTED));
-  EXPECT_THAT(
-      itr->GetCallStats(),
-      EqualsDocHitInfoIteratorCallStats(
-          /*num_leaf_advance_calls_lite_index=*/2,
-          /*num_leaf_advance_calls_main_index=*/2,
-          /*num_leaf_advance_calls_integer_index=*/0,
-          /*num_leaf_advance_calls_no_index=*/0, /*num_blocks_inspected=*/1));
+  EXPECT_THAT(itr->GetCallStats(),
+              EqualsDocHitInfoIteratorCallStats(
+                  /*num_leaf_advance_calls_lite_index=*/2,
+                  /*num_leaf_advance_calls_main_index=*/2,
+                  /*num_leaf_advance_calls_integer_index=*/0,
+                  /*num_leaf_advance_calls_no_index=*/0,
+                  /*num_blocks_inspected=*/1,
+                  DocHitInfoIterator::CallStats::EmbeddingStats()));
 }
 
 TEST_F(IndexTest, SingleHitSingleTermIndex) {
@@ -1230,7 +1256,8 @@
                          /*lite_index_sort_at_indexing=*/true,
                          /*lite_index_sort_size=*/64);
   ICING_ASSERT_OK_AND_ASSIGN(
-      index_, Index::Create(options, &filesystem_, &icing_filesystem_));
+      index_, Index::Create(options, &filesystem_, &icing_filesystem_,
+                            feature_flags_.get()));
 
   std::default_random_engine random;
   std::vector<std::string> query_terms;
@@ -1299,7 +1326,8 @@
                          /*lite_index_sort_at_indexing=*/true,
                          /*lite_index_sort_size=*/64);
   ICING_ASSERT_OK_AND_ASSIGN(
-      index_, Index::Create(options, &filesystem_, &icing_filesystem_));
+      index_, Index::Create(options, &filesystem_, &icing_filesystem_,
+                            feature_flags_.get()));
 
   std::default_random_engine random;
   std::vector<std::string> query_terms;
@@ -1634,7 +1662,8 @@
   Index::Options options(index_dir_, /*index_merge_size=*/1024 * 1024,
                          /*lite_index_sort_at_indexing=*/true,
                          /*lite_index_sort_size=*/1024 * 8);
-  EXPECT_THAT(Index::Create(options, &filesystem_, &mock_icing_filesystem),
+  EXPECT_THAT(Index::Create(options, &filesystem_, &mock_icing_filesystem,
+                            feature_flags_.get()),
               StatusIs(libtextclassifier3::StatusCode::INTERNAL));
 }
 
@@ -1667,7 +1696,8 @@
   Index::Options options(index_dir_, /*index_merge_size=*/1024 * 1024,
                          /*lite_index_sort_at_indexing=*/true,
                          /*lite_index_sort_size=*/1024 * 8);
-  EXPECT_THAT(Index::Create(options, &filesystem_, &icing_filesystem_),
+  EXPECT_THAT(Index::Create(options, &filesystem_, &icing_filesystem_,
+                            feature_flags_.get()),
               StatusIs(libtextclassifier3::StatusCode::DATA_LOSS));
 }
 
@@ -1719,7 +1749,8 @@
                          /*lite_index_sort_at_indexing=*/true,
                          /*lite_index_sort_size=*/1024 * 8);
   ICING_ASSERT_OK_AND_ASSIGN(
-      index_, Index::Create(options, &filesystem_, &icing_filesystem_));
+      index_, Index::Create(options, &filesystem_, &icing_filesystem_,
+                            feature_flags_.get()));
 
   // Check that the hits are present.
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -1750,7 +1781,8 @@
                          /*lite_index_sort_at_indexing=*/true,
                          /*lite_index_sort_size=*/1024 * 8);
   ICING_ASSERT_OK_AND_ASSIGN(
-      index_, Index::Create(options, &filesystem_, &icing_filesystem_));
+      index_, Index::Create(options, &filesystem_, &icing_filesystem_,
+                            feature_flags_.get()));
 
   // Check that the hits are present.
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -1767,7 +1799,8 @@
   Index::Options options(
       index_dir_, /*index_merge_size=*/std::numeric_limits<uint32_t>::max(),
       /*lite_index_sort_at_indexing=*/true, /*lite_index_sort_size=*/1024 * 8);
-  EXPECT_THAT(Index::Create(options, &filesystem_, &icing_filesystem_),
+  EXPECT_THAT(Index::Create(options, &filesystem_, &icing_filesystem_,
+                            feature_flags_.get()),
               StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
 }
 

diff --git a/icing/index/integer-section-indexing-handler_test.cc b/icing/index/integer-section-indexing-handler_test.cc
index 0f70387..b1c7841 100644
--- a/icing/index/integer-section-indexing-handler_test.cc
+++ b/icing/index/integer-section-indexing-handler_test.cc

@@ -28,10 +28,12 @@
 #include "icing/document-builder.h"
 #include "icing/feature-flags.h"
 #include "icing/file/filesystem.h"
+#include "icing/file/portable-file-backed-proto-log.h"
 #include "icing/index/hit/doc-hit-info.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
 #include "icing/index/numeric/integer-index.h"
 #include "icing/index/numeric/numeric-index.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/portable/platform.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/schema.pb.h"
@@ -169,14 +171,18 @@
         filesystem_.CreateDirectoryRecursively(document_store_dir_.c_str()));
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult doc_store_create_result,
-        DocumentStore::Create(&filesystem_, document_store_dir_, &fake_clock_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, document_store_dir_, &fake_clock_,
+            schema_store_.get(), feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     document_store_ = std::move(doc_store_create_result.document_store);
   }
 
@@ -232,11 +238,13 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document)));
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result,
-      document_store_->Put(tokenized_document.document()));
+      document_store_->Put(tokenized_document.document_wrapper()));
 
   ASSERT_THAT(integer_index_->last_added_document_id(), Eq(kInvalidDocumentId));
   // Handle document.
@@ -284,11 +292,13 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(nested_document)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(nested_document)));
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result,
-      document_store_->Put(tokenized_document.document()));
+      document_store_->Put(tokenized_document.document_wrapper()));
 
   ASSERT_THAT(integer_index_->last_added_document_id(), Eq(kInvalidDocumentId));
   // Handle nested_document.
@@ -349,11 +359,13 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document)));
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result,
-      document_store_->Put(tokenized_document.document()));
+      document_store_->Put(tokenized_document.document_wrapper()));
 
   ASSERT_THAT(integer_index_->last_added_document_id(), Eq(kInvalidDocumentId));
   // Handle document. Index data should remain unchanged since there is no
@@ -392,9 +404,11 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document)));
-  ICING_ASSERT_OK(document_store_->Put(tokenized_document.document()));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document)));
+  ICING_ASSERT_OK(document_store_->Put(tokenized_document.document_wrapper()));
 
   static constexpr DocumentId kCurrentDocumentId = 3;
   integer_index_->set_last_added_document_id(kCurrentDocumentId);
@@ -453,11 +467,13 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document)));
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result,
-      document_store_->Put(tokenized_document.document()));
+      document_store_->Put(tokenized_document.document_wrapper()));
 
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<IntegerSectionIndexingHandler> handler,
@@ -531,18 +547,22 @@
           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document1,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document1)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document1)));
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document2,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document2)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document2)));
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result1,
-      document_store_->Put(tokenized_document1.document()));
+      document_store_->Put(tokenized_document1.document_wrapper()));
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result2,
-      document_store_->Put(tokenized_document2.document()));
+      document_store_->Put(tokenized_document2.document_wrapper()));
 
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<IntegerSectionIndexingHandler> handler,

diff --git a/icing/index/iterator/doc-hit-info-iterator-all-document-id.h b/icing/index/iterator/doc-hit-info-iterator-all-document-id.h
index 60c5e0c..53ecc38 100644
--- a/icing/index/iterator/doc-hit-info-iterator-all-document-id.h
+++ b/icing/index/iterator/doc-hit-info-iterator-all-document-id.h

@@ -15,11 +15,12 @@
 #ifndef ICING_INDEX_ITERATOR_DOC_HIT_INFO_ITERATOR_ALL_DOCUMENT_ID_H_
 #define ICING_INDEX_ITERATOR_DOC_HIT_INFO_ITERATOR_ALL_DOCUMENT_ID_H_
 
-#include <cstdint>
+#include <memory>
 #include <string>
+#include <vector>
 
 #include "icing/text_classifier/lib3/utils/base/status.h"
-#include "icing/absl_ports/str_cat.h"
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
 #include "icing/legacy/core/icing-string-util.h"
 #include "icing/store/document-id.h"
@@ -29,7 +30,8 @@
 
 // Iterator for all DocumentIds in range [0, document_id_limit_]: 0 inclusive,
 // document_id_limit_ inclusive. Returns DocumentIds in descending order.
-class DocHitInfoIteratorAllDocumentId : public DocHitInfoIterator {
+class DocHitInfoIteratorAllDocumentId
+    : public DocHitInfoIteratorSectionRestrictionNotApplicable {
  public:
   explicit DocHitInfoIteratorAllDocumentId(DocumentId document_id_limit);
 
@@ -37,7 +39,9 @@
 
   libtextclassifier3::StatusOr<TrimmedNode> TrimRightMostNode() && override;
 
-  void MapChildren(const ChildrenMapper& mapper) override {}
+  std::vector<std::unique_ptr<DocHitInfoIterator>*> GetChildren() override {
+    return {};
+  }
 
   CallStats GetCallStats() const override {
     return CallStats(
@@ -46,7 +50,8 @@
         /*num_leaf_advance_calls_integer_index_in=*/0,
         /*num_leaf_advance_calls_no_index_in=*/document_id_limit_ -
             current_document_id_,
-        /*num_blocks_inspected_in=*/0);
+        /*num_blocks_inspected_in=*/0,
+        /*embedding_stats_in=*/{});
   }
 
   std::string ToString() const override {

diff --git a/icing/index/iterator/doc-hit-info-iterator-all-document-id_test.cc b/icing/index/iterator/doc-hit-info-iterator-all-document-id_test.cc
index 379cb4d..fc61efe 100644
--- a/icing/index/iterator/doc-hit-info-iterator-all-document-id_test.cc
+++ b/icing/index/iterator/doc-hit-info-iterator-all-document-id_test.cc

@@ -55,23 +55,25 @@
 
 TEST(DocHitInfoIteratorAllDocumentIdTest, GetCallStats) {
   DocHitInfoIteratorAllDocumentId all_it(100);
-  EXPECT_THAT(
-      all_it.GetCallStats(),
-      EqualsDocHitInfoIteratorCallStats(
-          /*num_leaf_advance_calls_lite_index=*/0,
-          /*num_leaf_advance_calls_main_index=*/0,
-          /*num_leaf_advance_calls_integer_index=*/0,
-          /*num_leaf_advance_calls_no_index=*/0, /*num_blocks_inspected=*/0));
+  EXPECT_THAT(all_it.GetCallStats(),
+              EqualsDocHitInfoIteratorCallStats(
+                  /*num_leaf_advance_calls_lite_index=*/0,
+                  /*num_leaf_advance_calls_main_index=*/0,
+                  /*num_leaf_advance_calls_integer_index=*/0,
+                  /*num_leaf_advance_calls_no_index=*/0,
+                  /*num_blocks_inspected=*/0,
+                  DocHitInfoIterator::CallStats::EmbeddingStats()));
 
   for (int i = 1; i <= 5; ++i) {
     EXPECT_THAT(all_it.Advance(), IsOk());
-    EXPECT_THAT(
-        all_it.GetCallStats(),
-        EqualsDocHitInfoIteratorCallStats(
-            /*num_leaf_advance_calls_lite_index=*/0,
-            /*num_leaf_advance_calls_main_index=*/0,
-            /*num_leaf_advance_calls_integer_index=*/0,
-            /*num_leaf_advance_calls_no_index=*/i, /*num_blocks_inspected=*/0));
+    EXPECT_THAT(all_it.GetCallStats(),
+                EqualsDocHitInfoIteratorCallStats(
+                    /*num_leaf_advance_calls_lite_index=*/0,
+                    /*num_leaf_advance_calls_main_index=*/0,
+                    /*num_leaf_advance_calls_integer_index=*/0,
+                    /*num_leaf_advance_calls_no_index=*/i,
+                    /*num_blocks_inspected=*/0,
+                    DocHitInfoIterator::CallStats::EmbeddingStats()));
   }
 }
 

diff --git a/icing/index/iterator/doc-hit-info-iterator-and.h b/icing/index/iterator/doc-hit-info-iterator-and.h
index 8c52ac9..963b380 100644
--- a/icing/index/iterator/doc-hit-info-iterator-and.h
+++ b/icing/index/iterator/doc-hit-info-iterator-and.h

@@ -15,14 +15,16 @@
 #ifndef ICING_INDEX_ITERATOR_DOC_HIT_INFO_ITERATOR_AND_H_
 #define ICING_INDEX_ITERATOR_DOC_HIT_INFO_ITERATOR_AND_H_
 
-#include <cstdint>
+#include <cstddef>
 #include <memory>
 #include <string>
-#include <utility>
 #include <vector>
 
 #include "icing/text_classifier/lib3/utils/base/status.h"
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
+#include "icing/schema/section.h"
+#include "icing/store/document-id.h"
 
 namespace icing {
 namespace lib {
@@ -33,7 +35,8 @@
     std::vector<std::unique_ptr<DocHitInfoIterator>> iterators);
 
 // Iterate over a logical AND of two child iterators.
-class DocHitInfoIteratorAnd : public DocHitInfoIterator {
+class DocHitInfoIteratorAnd
+    : public DocHitInfoIteratorSectionRestrictionApplyToChildren {
  public:
   // Set the shorter iterator to short_it to get performance benefits
   // for when an underlying iterator has a more efficient AdvanceTo.
@@ -49,9 +52,8 @@
 
   std::string ToString() const override;
 
-  void MapChildren(const ChildrenMapper& mapper) override {
-    short_ = mapper(std::move(short_));
-    long_ = mapper(std::move(long_));
+  std::vector<std::unique_ptr<DocHitInfoIterator>*> GetChildren() override {
+    return {&short_, &long_};
   }
 
   void PopulateMatchedTermsStats(
@@ -75,7 +77,8 @@
 // Iterate over a logical AND of multiple child iterators.
 // NOTE: DocHitInfoIteratorAnd is a faster alternative to AND exactly 2
 // iterators.
-class DocHitInfoIteratorAndNary : public DocHitInfoIterator {
+class DocHitInfoIteratorAndNary
+    : public DocHitInfoIteratorSectionRestrictionApplyToChildren {
  public:
   explicit DocHitInfoIteratorAndNary(
       std::vector<std::unique_ptr<DocHitInfoIterator>> iterators);
@@ -88,10 +91,13 @@
 
   std::string ToString() const override;
 
-  void MapChildren(const ChildrenMapper& mapper) override {
+  std::vector<std::unique_ptr<DocHitInfoIterator>*> GetChildren() override {
+    std::vector<std::unique_ptr<DocHitInfoIterator>*> children;
+    children.reserve(iterators_.size());
     for (int i = 0; i < iterators_.size(); ++i) {
-      iterators_[i] = mapper(std::move(iterators_[i]));
+      children.push_back(&iterators_[i]);
     }
+    return children;
   }
 
   void PopulateMatchedTermsStats(

diff --git a/icing/index/iterator/doc-hit-info-iterator-and_test.cc b/icing/index/iterator/doc-hit-info-iterator-and_test.cc
index f204ada..f4faaf8 100644
--- a/icing/index/iterator/doc-hit-info-iterator-and_test.cc
+++ b/icing/index/iterator/doc-hit-info-iterator-and_test.cc

@@ -14,12 +14,16 @@
 
 #include "icing/index/iterator/doc-hit-info-iterator-and.h"
 
-#include <string>
+#include <memory>
+#include <unordered_map>
+#include <utility>
 #include <vector>
 
+#include "icing/text_classifier/lib3/utils/base/status.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include "icing/index/hit/doc-hit-info.h"
+#include "icing/index/hit/hit.h"
 #include "icing/index/iterator/doc-hit-info-iterator-test-util.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
 #include "icing/schema/section.h"
@@ -34,6 +38,8 @@
 using ::testing::ElementsAre;
 using ::testing::Eq;
 using ::testing::IsEmpty;
+using ::testing::Pointee;
+using ::testing::Pointer;
 
 TEST(CreateAndIteratorTest, And) {
   // Basic test that we can create a working And iterator. Further testing of
@@ -84,7 +90,12 @@
       /*num_leaf_advance_calls_main_index_in=*/5,
       /*num_leaf_advance_calls_integer_index_in=*/3,
       /*num_leaf_advance_calls_no_index_in=*/1,
-      /*num_blocks_inspected_in=*/4);  // arbitrary value
+      /*num_blocks_inspected_in=*/4,
+      /*embedding_stats_in=*/
+      {.num_unquantized_embeddings_scored = 2,
+       .num_quantized_embeddings_scored = 3,
+       .unquantized_shards_read = {1, 2},
+       .quantized_shards_read{3, 4}});  // arbitrary value
   auto first_iter = std::make_unique<DocHitInfoIteratorDummy>();
   first_iter->SetCallStats(first_iter_call_stats);
 
@@ -93,7 +104,12 @@
       /*num_leaf_advance_calls_main_index_in=*/2,
       /*num_leaf_advance_calls_integer_index_in=*/10,
       /*num_leaf_advance_calls_no_index_in=*/3,
-      /*num_blocks_inspected_in=*/7);  // arbitrary value
+      /*num_blocks_inspected_in=*/7,
+      /*embedding_stats_in=*/
+      {.num_unquantized_embeddings_scored = 4,
+       .num_quantized_embeddings_scored = 5,
+       .unquantized_shards_read = {5, 6},
+       .quantized_shards_read{7}});  // arbitrary value
   auto second_iter = std::make_unique<DocHitInfoIteratorDummy>();
   second_iter->SetCallStats(second_iter_call_stats);
 
@@ -447,7 +463,12 @@
       /*num_leaf_advance_calls_main_index_in=*/5,
       /*num_leaf_advance_calls_integer_index_in=*/3,
       /*num_leaf_advance_calls_no_index_in=*/1,
-      /*num_blocks_inspected_in=*/4);  // arbitrary value
+      /*num_blocks_inspected_in=*/4,
+      /*embedding_stats_in=*/
+      {.num_unquantized_embeddings_scored = 2,
+       .num_quantized_embeddings_scored = 3,
+       .unquantized_shards_read = {1, 2},
+       .quantized_shards_read{3, 4}});  // arbitrary value
   auto first_iter = std::make_unique<DocHitInfoIteratorDummy>();
   first_iter->SetCallStats(first_iter_call_stats);
 
@@ -456,7 +477,12 @@
       /*num_leaf_advance_calls_main_index_in=*/2,
       /*num_leaf_advance_calls_integer_index_in=*/10,
       /*num_leaf_advance_calls_no_index_in=*/3,
-      /*num_blocks_inspected_in=*/7);  // arbitrary value
+      /*num_blocks_inspected_in=*/7,
+      /*embedding_stats_in=*/
+      {.num_unquantized_embeddings_scored = 4,
+       .num_quantized_embeddings_scored = 5,
+       .unquantized_shards_read = {5, 6},
+       .quantized_shards_read{7}});  // arbitrary value
   auto second_iter = std::make_unique<DocHitInfoIteratorDummy>();
   second_iter->SetCallStats(second_iter_call_stats);
 
@@ -465,7 +491,12 @@
       /*num_leaf_advance_calls_main_index_in=*/2000,
       /*num_leaf_advance_calls_integer_index_in=*/3000,
       /*num_leaf_advance_calls_no_index_in=*/0,
-      /*num_blocks_inspected_in=*/200);  // arbitrary value
+      /*num_blocks_inspected_in=*/200,
+      /*embedding_stats_in=*/
+      {.num_unquantized_embeddings_scored = 1,
+       .num_quantized_embeddings_scored = 1,
+       .unquantized_shards_read = {0},
+       .quantized_shards_read{0}});  // arbitrary value
   auto third_iter = std::make_unique<DocHitInfoIteratorDummy>();
   third_iter->SetCallStats(third_iter_call_stats);
 
@@ -474,7 +505,12 @@
       /*num_leaf_advance_calls_main_index_in=*/400,
       /*num_leaf_advance_calls_integer_index_in=*/100,
       /*num_leaf_advance_calls_no_index_in=*/20,
-      /*num_blocks_inspected_in=*/50);  // arbitrary value
+      /*num_blocks_inspected_in=*/50,
+      /*embedding_stats_in=*/
+      {.num_unquantized_embeddings_scored = 10,
+       .num_quantized_embeddings_scored = 10,
+       .unquantized_shards_read = {5, 6, 7},
+       .quantized_shards_read{9, 10, 11}});  // arbitrary value
   auto fourth_iter = std::make_unique<DocHitInfoIteratorDummy>();
   fourth_iter->SetCallStats(fourth_iter_call_stats);
 
@@ -630,6 +666,32 @@
   EXPECT_FALSE(and_iter.Advance().ok());
 }
 
+TEST(DocHitInfoIteratorAndNaryTest, GetChildren) {
+  std::vector<DocHitInfo> first_vector = {DocHitInfo(2), DocHitInfo(1),
+                                          DocHitInfo(0)};
+  std::vector<DocHitInfo> second_vector = {DocHitInfo(2), DocHitInfo(1)};
+  std::vector<DocHitInfo> third_vector = {DocHitInfo(2)};
+
+  std::vector<std::unique_ptr<DocHitInfoIterator>> iterators;
+  iterators.push_back(std::make_unique<DocHitInfoIteratorDummy>(first_vector));
+  iterators.push_back(std::make_unique<DocHitInfoIteratorDummy>(second_vector));
+  iterators.push_back(
+      std::make_unique<DocHitInfoIteratorDummy>(third_vector, "term", 10));
+
+  std::vector<DocHitInfoIterator*> iterator_ptrs;
+  for (const auto& iter : iterators) {
+    iterator_ptrs.push_back(iter.get());
+  }
+
+  std::unique_ptr<DocHitInfoIterator> iter =
+      std::make_unique<DocHitInfoIteratorAndNary>(std::move(iterators));
+
+  EXPECT_THAT(iter->GetChildren(),
+              ElementsAre(Pointee(Pointer(iterator_ptrs[0])),
+                          Pointee(Pointer(iterator_ptrs[1])),
+                          Pointee(Pointer(iterator_ptrs[2]))));
+}
+
 }  // namespace
 
 }  // namespace lib

diff --git a/icing/index/iterator/doc-hit-info-iterator-by-uri.h b/icing/index/iterator/doc-hit-info-iterator-by-uri.h
index d136535..9c8bb9f 100644
--- a/icing/index/iterator/doc-hit-info-iterator-by-uri.h
+++ b/icing/index/iterator/doc-hit-info-iterator-by-uri.h

@@ -31,7 +31,8 @@
 namespace icing {
 namespace lib {
 
-class DocHitInfoIteratorByUri : public DocHitInfoIterator {
+class DocHitInfoIteratorByUri
+    : public DocHitInfoIteratorSectionRestrictionNotApplicable {
  public:
   // Creates a DocHitInfoIteratorByUri based on the given search_spec.
   //
@@ -51,7 +52,9 @@
         "DocHitInfoIteratorByUri should not be used in suggestion.");
   }
 
-  void MapChildren(const ChildrenMapper& mapper) override {}
+  std::vector<std::unique_ptr<DocHitInfoIterator>*> GetChildren() override {
+    return {};
+  }
 
   CallStats GetCallStats() const override {
     return CallStats(
@@ -59,7 +62,8 @@
         /*num_leaf_advance_calls_main_index_in=*/0,
         /*num_leaf_advance_calls_integer_index_in=*/0,
         /*num_leaf_advance_calls_no_index_in=*/num_advance_calls_,
-        /*num_blocks_inspected_in=*/0);
+        /*num_blocks_inspected_in=*/0,
+        /*embedding_stats_in=*/{});
   }
 
   std::string ToString() const override { return "uri_iterator"; }

diff --git a/icing/index/iterator/doc-hit-info-iterator-by-uri_test.cc b/icing/index/iterator/doc-hit-info-iterator-by-uri_test.cc
index b783c48..058d13e 100644
--- a/icing/index/iterator/doc-hit-info-iterator-by-uri_test.cc
+++ b/icing/index/iterator/doc-hit-info-iterator-by-uri_test.cc

@@ -26,6 +26,7 @@
 #include "icing/file/filesystem.h"
 #include "icing/file/portable-file-backed-proto-log.h"
 #include "icing/index/iterator/doc-hit-info-iterator-test-util.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/schema.pb.h"
 #include "icing/schema-builder.h"
@@ -36,6 +37,7 @@
 #include "icing/testing/fake-clock.h"
 #include "icing/testing/test-feature-flags.h"
 #include "icing/testing/tmp-directory.h"
+#include "icing/util/document-util.h"
 
 namespace icing {
 namespace lib {
@@ -66,14 +68,18 @@
 
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
-        DocumentStore::Create(&filesystem_, test_dir_, &fake_clock_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, test_dir_, &fake_clock_, schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     document_store_ = std::move(create_result.document_store);
   }
 
@@ -121,11 +127,16 @@
                                 .SetKey("namespace", "email/3")
                                 .SetSchema("email")
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1)));
   DocumentId document_id1 = put_result.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(put_result, document_store_->Put(document2));
-  ICING_ASSERT_OK_AND_ASSIGN(put_result, document_store_->Put(document3));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document2)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document3)));
   DocumentId document_id3 = put_result.new_document_id;
 
   // Create a search spec with uri filters that only match document1 and
@@ -158,12 +169,17 @@
                                 .SetKey("namespace", "email/3")
                                 .SetSchema("email")
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1)));
   DocumentId document_id1 = put_result.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(put_result, document_store_->Put(document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document2)));
   DocumentId document_id2 = put_result.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(put_result, document_store_->Put(document3));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document3)));
   DocumentId document_id3 = put_result.new_document_id;
 
   // Create a search spec with uri filters that match all documents.
@@ -196,12 +212,17 @@
                                 .SetKey("namespace", "email/3")
                                 .SetSchema("email")
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1)));
   DocumentId document_id1 = put_result.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(put_result, document_store_->Put(document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document2)));
   DocumentId document_id2 = put_result.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(put_result, document_store_->Put(document3));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document3)));
   DocumentId document_id3 = put_result.new_document_id;
 
   // Create a search spec with a nonexistent uri in uri filters.
@@ -235,10 +256,15 @@
                                 .SetKey("namespace", "email/3")
                                 .SetSchema("email")
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document1));
-  ICING_ASSERT_OK_AND_ASSIGN(put_result, document_store_->Put(document2));
-  ICING_ASSERT_OK_AND_ASSIGN(put_result, document_store_->Put(document3));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document2)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document3)));
 
   // Create a search spec with all nonexistent uris.
   SearchSpecProto search_spec;
@@ -269,11 +295,16 @@
                                 .SetKey("namespace2", "email/3")
                                 .SetSchema("email")
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1)));
   DocumentId document_id1 = put_result.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(put_result, document_store_->Put(document2));
-  ICING_ASSERT_OK_AND_ASSIGN(put_result, document_store_->Put(document3));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document2)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document3)));
   DocumentId document_id3 = put_result.new_document_id;
 
   // Create a search spec with uri filters that match document1 and document3 in
@@ -314,14 +345,21 @@
                                 .SetKey("namespace", "email/4")
                                 .SetSchema("email")
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1)));
   DocumentId document_id1 = put_result.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(put_result, document_store_->Put(document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document2)));
   DocumentId document_id2 = put_result.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(put_result, document_store_->Put(document3));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document3)));
   DocumentId document_id3 = put_result.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(put_result, document_store_->Put(document4));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document4)));
   DocumentId document_id4 = put_result.new_document_id;
 
   // Create a search spec with duplicated uri filters. The result document ids

diff --git a/icing/index/iterator/doc-hit-info-iterator-data-holder.h b/icing/index/iterator/doc-hit-info-iterator-data-holder.h
new file mode 100644
index 0000000..f8c5ede
--- /dev/null
+++ b/icing/index/iterator/doc-hit-info-iterator-data-holder.h

@@ -0,0 +1,81 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef ICING_INDEX_ITERATOR_DOC_HIT_INFO_ITERATOR_DATA_HOLDER_H_
+#define ICING_INDEX_ITERATOR_DOC_HIT_INFO_ITERATOR_DATA_HOLDER_H_
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "icing/text_classifier/lib3/utils/base/status.h"
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
+#include "icing/index/iterator/doc-hit-info-iterator.h"
+#include "icing/schema/section.h"
+#include "icing/util/status-macros.h"
+
+namespace icing {
+namespace lib {
+
+// An iterator that simply takes ownership of an object.
+template <typename T>
+class DocHitInfoIteratorDataHolder
+    : public DocHitInfoIteratorSectionRestrictionApplyToChildren {
+ public:
+  explicit DocHitInfoIteratorDataHolder(
+      std::unique_ptr<DocHitInfoIterator> delegate, std::unique_ptr<T> data)
+      : delegate_(std::move(delegate)), data_(std::move(data)) {}
+
+  libtextclassifier3::Status Advance() override {
+    auto result = delegate_->Advance();
+    doc_hit_info_ = delegate_->doc_hit_info();
+    return result;
+  }
+
+  libtextclassifier3::StatusOr<TrimmedNode> TrimRightMostNode() && override {
+    ICING_ASSIGN_OR_RETURN(TrimmedNode trimmed_delegate,
+                           std::move(*delegate_).TrimRightMostNode());
+    if (trimmed_delegate.iterator_ != nullptr) {
+      trimmed_delegate.iterator_ =
+          std::make_unique<DocHitInfoIteratorDataHolder>(
+              std::move(trimmed_delegate.iterator_), std::move(data_));
+    }
+    return trimmed_delegate;
+  }
+
+  std::vector<std::unique_ptr<DocHitInfoIterator>*> GetChildren() override {
+    return {&delegate_};
+  }
+
+  CallStats GetCallStats() const override { return delegate_->GetCallStats(); }
+
+  std::string ToString() const override { return delegate_->ToString(); }
+
+  void PopulateMatchedTermsStats(
+      std::vector<TermMatchInfo>* matched_terms_stats,
+      SectionIdMask filtering_section_mask) const override {
+    return delegate_->PopulateMatchedTermsStats(matched_terms_stats,
+                                                filtering_section_mask);
+  }
+
+ private:
+  std::unique_ptr<DocHitInfoIterator> delegate_;
+  std::unique_ptr<T> data_;
+};
+
+}  // namespace lib
+}  // namespace icing
+
+#endif  // ICING_INDEX_ITERATOR_DOC_HIT_INFO_ITERATOR_DATA_HOLDER_H_

diff --git a/icing/index/iterator/doc-hit-info-iterator-filter.cc b/icing/index/iterator/doc-hit-info-iterator-filter.cc
index 3bc3b48..477077d 100644
--- a/icing/index/iterator/doc-hit-info-iterator-filter.cc
+++ b/icing/index/iterator/doc-hit-info-iterator-filter.cc

@@ -14,11 +14,8 @@
 
 #include "icing/index/iterator/doc-hit-info-iterator-filter.h"
 
-#include <cstdint>
 #include <memory>
 #include <string>
-#include <string_view>
-#include <unordered_set>
 #include <utility>
 
 #include "icing/text_classifier/lib3/utils/base/status.h"
@@ -26,53 +23,63 @@
 #include "icing/absl_ports/canonical_errors.h"
 #include "icing/index/hit/doc-hit-info.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
-#include "icing/schema/schema-store.h"
-#include "icing/store/document-filter-data.h"
+#include "icing/index/iterator/document-filter-predicate.h"
 #include "icing/store/document-id.h"
-#include "icing/store/document-store.h"
 #include "icing/util/status-macros.h"
 
 namespace icing {
 namespace lib {
 
-DocHitInfoIteratorFilter::DocHitInfoIteratorFilter(
-    std::unique_ptr<DocHitInfoIterator> delegate,
-    const DocumentStore* document_store, const SchemaStore* schema_store,
-    const Options& options, int64_t current_time_ms)
-    : delegate_(std::move(delegate)),
-      document_store_(*document_store),
-      schema_store_(*schema_store),
-      options_(options),
-      current_time_ms_(current_time_ms) {}
+// Some children iterators (currently only DocHitInfoIteratorEmbedding) can
+// internally handle the filter predicate to accelerate the search. Let's dfs to
+// all children iterators to try to provide the filter predicate.
+//
+// Returns true if all children iterators have accepted the filter predicate, in
+// which case we no longer need to apply the filter at the top level.
+bool PassFilterPredicateToChildrenAndHandle(
+    DocHitInfoIterator* iterator, const DocumentFilterPredicate* predicate) {
+  // If the predicate is accepted, we can stop here.
+  if (iterator->HandleFilter(predicate)) {
+    return true;
+  }
+
+  // Now, we know that the iterator cannot internally handle the filter
+  // predicate.
+  // If it has no children or the iterator cannot pass the filter predicate
+  // through (e.g., NOT iterator), return false to indicate that we should apply
+  // the filter at the top level.
+  if (iterator->GetChildren().empty() ||
+      !iterator->CanPassFilterPredicateThrough()) {
+    return false;
+  }
+
+  // Continue to pass the filter predicate to children iterators.
+  bool all_children_accepted = true;
+  for (std::unique_ptr<DocHitInfoIterator>* child : iterator->GetChildren()) {
+    all_children_accepted &=
+        PassFilterPredicateToChildrenAndHandle(child->get(), predicate);
+  }
+  return all_children_accepted;
+}
+
+/* static */ std::unique_ptr<DocHitInfoIterator>
+DocHitInfoIteratorFilter::ApplyFilter(
+    std::unique_ptr<DocHitInfoIterator> iterator,
+    const DocumentFilterPredicate* predicate,
+    bool enable_passing_filter_to_children) {
+  if (enable_passing_filter_to_children &&
+      PassFilterPredicateToChildrenAndHandle(iterator.get(), predicate)) {
+    return iterator;
+  }
+  return std::unique_ptr<DocHitInfoIteratorFilter>(
+      new DocHitInfoIteratorFilter(std::move(iterator), predicate));
+}
 
 libtextclassifier3::Status DocHitInfoIteratorFilter::Advance() {
   while (delegate_->Advance().ok()) {
-    // Try to get the DocumentFilterData
-    auto document_filter_data_optional =
-        document_store_.GetAliveDocumentFilterData(
-            delegate_->doc_hit_info().document_id(), current_time_ms_);
-    if (!document_filter_data_optional) {
-      // Didn't find the DocumentFilterData in the filter cache. This could be
-      // because the Document doesn't exist or the DocumentId isn't valid or the
-      // filter cache is in some invalid state. This is bad, but not the query's
-      // responsibility to fix, so just skip this result for now.
+    if (!(*predicate_)(delegate_->doc_hit_info().document_id())) {
       continue;
     }
-    // We should be guaranteed that this exists now.
-    DocumentFilterData data = document_filter_data_optional.value();
-
-    if (options_.filter_by_namespace_id_enabled &&
-        options_.target_namespace_ids.count(data.namespace_id()) == 0) {
-      // Doesn't match one of the specified namespaces. Keep searching
-      continue;
-    }
-
-    if (options_.filter_by_schema_type_id_enabled &&
-        options_.target_schema_type_ids.count(data.schema_type_id()) == 0) {
-      // Doesn't match one of the specified schema types. Keep searching
-      continue;
-    }
-
     // Satisfied all our specified filters
     doc_hit_info_ = delegate_->doc_hit_info();
     return libtextclassifier3::Status::OK;
@@ -88,9 +95,9 @@
   ICING_ASSIGN_OR_RETURN(TrimmedNode trimmed_delegate,
                          std::move(*delegate_).TrimRightMostNode());
   if (trimmed_delegate.iterator_ != nullptr) {
-    trimmed_delegate.iterator_ = std::make_unique<DocHitInfoIteratorFilter>(
-        std::move(trimmed_delegate.iterator_), &document_store_, &schema_store_,
-        options_, current_time_ms_);
+    trimmed_delegate.iterator_ =
+        std::unique_ptr<DocHitInfoIteratorFilter>(new DocHitInfoIteratorFilter(
+            std::move(trimmed_delegate.iterator_), predicate_));
   }
   return trimmed_delegate;
 }

diff --git a/icing/index/iterator/doc-hit-info-iterator-filter.h b/icing/index/iterator/doc-hit-info-iterator-filter.h
index b8e70e8..b937c22 100644
--- a/icing/index/iterator/doc-hit-info-iterator-filter.h
+++ b/icing/index/iterator/doc-hit-info-iterator-filter.h

@@ -15,62 +15,38 @@
 #ifndef ICING_INDEX_ITERATOR_DOC_HIT_INFO_ITERATOR_FILTER_H_
 #define ICING_INDEX_ITERATOR_DOC_HIT_INFO_ITERATOR_FILTER_H_
 
-#include <cstdint>
 #include <memory>
 #include <string>
-#include <unordered_set>
-#include <utility>
 #include <vector>
 
 #include "icing/text_classifier/lib3/utils/base/status.h"
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
-#include "icing/schema/schema-store.h"
+#include "icing/index/iterator/document-filter-predicate.h"
 #include "icing/schema/section.h"
-#include "icing/store/document-filter-data.h"
-#include "icing/store/document-store.h"
-#include "icing/store/namespace-id.h"
 
 namespace icing {
 namespace lib {
 
-// A iterator that helps filter out DocHitInfos associated with non-existing
-// document ids.
-class DocHitInfoIteratorFilter : public DocHitInfoIterator {
+// A iterator that helps filter out DocHitInfos by a given predicate.
+//
+// To maintain the correct semantics of section restrictions, it implements
+// DocHitInfoIteratorSectionRestrictionApplyToChildren to pass down section
+// restrictions to child iterators.
+class DocHitInfoIteratorFilter
+    : public DocHitInfoIteratorSectionRestrictionApplyToChildren {
  public:
-  struct Options {
-    // List of namespace ids that documents must have.
-    // filter_by_namespace_id_enabled=false means that all namespaces are valid,
-    // and no documents will be filtered out.
-    //
-    // Note that if we want to reference the strings in namespaces later, ensure
-    // that the caller who passed the Options class outlives the
-    // DocHitInfoIteratorFilter.
-    bool filter_by_namespace_id_enabled = false;
-    std::unordered_set<NamespaceId> target_namespace_ids;
-
-    // List of schema type ids that documents must have.
-    // filter_by_schema_type_id_enabled=false means that all schema types are
-    // valid, and no documents will be filtered out.
-    //
-    // Note that if we want to reference the strings in schema types later,
-    // ensure that the caller who passed the Options class outlives the
-    // DocHitInfoIteratorFilter.
-    bool filter_by_schema_type_id_enabled = false;
-    std::unordered_set<SchemaTypeId> target_schema_type_ids;
-  };
-
-  explicit DocHitInfoIteratorFilter(
-      std::unique_ptr<DocHitInfoIterator> delegate,
-      const DocumentStore* document_store, const SchemaStore* schema_store,
-      const Options& options, int64_t current_time_ms);
+  static std::unique_ptr<DocHitInfoIterator> ApplyFilter(
+      std::unique_ptr<DocHitInfoIterator> iterator,
+      const DocumentFilterPredicate* predicate,
+      bool enable_passing_filter_to_children);
 
   libtextclassifier3::Status Advance() override;
 
   libtextclassifier3::StatusOr<TrimmedNode> TrimRightMostNode() && override;
 
-  void MapChildren(const ChildrenMapper& mapper) override {
-    delegate_ = mapper(std::move(delegate_));
+  std::vector<std::unique_ptr<DocHitInfoIterator>*> GetChildren() override {
+    return {&delegate_};
   }
 
   CallStats GetCallStats() const override { return delegate_->GetCallStats(); }
@@ -85,11 +61,13 @@
   }
 
  private:
+  explicit DocHitInfoIteratorFilter(
+      std::unique_ptr<DocHitInfoIterator> delegate,
+      const DocumentFilterPredicate* predicate)
+      : delegate_(std::move(delegate)), predicate_(predicate) {}
+
   std::unique_ptr<DocHitInfoIterator> delegate_;
-  const DocumentStore& document_store_;
-  const SchemaStore& schema_store_;
-  const Options options_;
-  int64_t current_time_ms_;
+  const DocumentFilterPredicate* predicate_;
 };
 
 }  // namespace lib

diff --git a/icing/index/iterator/doc-hit-info-iterator-filter_test.cc b/icing/index/iterator/doc-hit-info-iterator-filter_test.cc
index 00cdb71..1f7d3b2 100644
--- a/icing/index/iterator/doc-hit-info-iterator-filter_test.cc
+++ b/icing/index/iterator/doc-hit-info-iterator-filter_test.cc

@@ -30,8 +30,12 @@
 #include "icing/file/portable-file-backed-proto-log.h"
 #include "icing/index/hit/doc-hit-info.h"
 #include "icing/index/iterator/doc-hit-info-iterator-and.h"
+#include "icing/index/iterator/doc-hit-info-iterator-not.h"
+#include "icing/index/iterator/doc-hit-info-iterator-or.h"
 #include "icing/index/iterator/doc-hit-info-iterator-test-util.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
+#include "icing/index/iterator/document-filter-predicate.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/schema.pb.h"
 #include "icing/query/query-utils.h"
@@ -45,6 +49,7 @@
 #include "icing/testing/test-feature-flags.h"
 #include "icing/testing/tmp-directory.h"
 #include "icing/util/clock.h"
+#include "icing/util/document-util.h"
 
 namespace icing {
 namespace lib {
@@ -54,6 +59,10 @@
 using ::testing::ElementsAre;
 using ::testing::Eq;
 using ::testing::IsEmpty;
+using ::testing::Ne;
+using ::testing::Pointee;
+using ::testing::Pointer;
+using ::testing::WhenDynamicCastTo;
 
 libtextclassifier3::StatusOr<DocumentStore::CreateResult> CreateDocumentStore(
     const Filesystem* filesystem, const std::string& base_dir,
@@ -64,9 +73,41 @@
       /*force_recovery_and_revalidate_documents=*/false,
       /*pre_mapping_fbv=*/false, /*use_persistent_hash_map=*/true,
       PortableFileBackedProtoLog<DocumentWrapper>::kDefaultCompressionLevel,
+      PortableFileBackedProtoLog<
+          DocumentWrapper>::kDefaultCompressionThresholdBytes,
+      protobuf_ports::kDefaultMemLevel,
       /*initialize_stats=*/nullptr);
 }
 
+class DocHitInfoIteratorDummyHandlingFilter
+    : public DocHitInfoIteratorHandlingFilter {
+ public:
+  libtextclassifier3::Status Advance() override {
+    return libtextclassifier3::Status::OK;
+  }
+
+  libtextclassifier3::StatusOr<TrimmedNode> TrimRightMostNode() && override {
+    TrimmedNode node = {nullptr, /*term=*/"", /*term_start_index_=*/0,
+                        /*unnormalized_term_length_=*/0};
+    return node;
+  }
+
+  std::vector<std::unique_ptr<DocHitInfoIterator>*> GetChildren() override {
+    return {};
+  }
+
+  CallStats GetCallStats() const override { return CallStats(); }
+
+  std::string ToString() const override {
+    return "DocHitInfoIteratorDummyHandlingFilter";
+  }
+
+  std::vector<const DocumentFilterPredicate*> document_filter_predicates()
+      const {
+    return document_filter_predicates_;
+  }
+};
+
 class DocHitInfoIteratorDeletedFilterTest : public ::testing::Test {
  protected:
   DocHitInfoIteratorDeletedFilterTest()
@@ -97,6 +138,10 @@
         CreateDocumentStore(&filesystem_, test_dir_, &fake_clock_,
                             schema_store_.get(), *feature_flags_));
     document_store_ = std::move(create_result.document_store);
+
+    predicate_ = GetFilterPredicateBySchemaAndNamespace(
+        SearchSpecProto::default_instance(), *document_store_, *schema_store_,
+        fake_clock_.GetSystemTimeMilliseconds());
   }
 
   void TearDown() override {
@@ -116,31 +161,39 @@
   DocumentProto test_document1_;
   DocumentProto test_document2_;
   DocumentProto test_document3_;
-  DocHitInfoIteratorFilter::Options options_;
+  std::unique_ptr<DocumentFilterPredicate> predicate_;
 };
 
 TEST_F(DocHitInfoIteratorDeletedFilterTest, EmptyOriginalIterator) {
-  ICING_ASSERT_OK(document_store_->Put(test_document1_));
+  ICING_ASSERT_OK(document_store_->Put(
+      document_util::CreateDocumentWrapper(test_document1_)));
 
   std::unique_ptr<DocHitInfoIterator> original_iterator_empty =
       std::make_unique<DocHitInfoIteratorDummy>();
 
-  DocHitInfoIteratorFilter filtered_iterator(
-      std::move(original_iterator_empty), document_store_.get(),
-      schema_store_.get(), options_, fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator_empty), predicate_.get(),
+          feature_flags_->enable_passing_filter_to_children());
 
-  EXPECT_THAT(GetDocumentIds(&filtered_iterator), IsEmpty());
+  EXPECT_THAT(GetDocumentIds(filtered_iterator.get()), IsEmpty());
 }
 
 TEST_F(DocHitInfoIteratorDeletedFilterTest, DeletedDocumentsAreFiltered) {
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(test_document1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(test_document1_)));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(test_document2_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(test_document2_)));
   DocumentId document_id2 = put_result2.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             document_store_->Put(test_document3_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(test_document3_)));
   DocumentId document_id3 = put_result3.new_document_id;
 
   // Deletes test document 2
@@ -154,23 +207,30 @@
   std::unique_ptr<DocHitInfoIterator> original_iterator =
       std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos);
 
-  DocHitInfoIteratorFilter filtered_iterator(
-      std::move(original_iterator), document_store_.get(), schema_store_.get(),
-      options_, fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator), predicate_.get(),
+          feature_flags_->enable_passing_filter_to_children());
 
-  EXPECT_THAT(GetDocumentIds(&filtered_iterator),
+  EXPECT_THAT(GetDocumentIds(filtered_iterator.get()),
               ElementsAre(document_id1, document_id3));
 }
 
 TEST_F(DocHitInfoIteratorDeletedFilterTest, NonExistingDocumentsAreFiltered) {
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(test_document1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(test_document1_)));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(test_document2_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(test_document2_)));
   DocumentId document_id2 = put_result2.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             document_store_->Put(test_document3_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(test_document3_)));
   DocumentId document_id3 = put_result3.new_document_id;
 
   // Document ids 7, 8, 9 are not existing
@@ -183,11 +243,12 @@
   std::unique_ptr<DocHitInfoIterator> original_iterator =
       std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos);
 
-  DocHitInfoIteratorFilter filtered_iterator(
-      std::move(original_iterator), document_store_.get(), schema_store_.get(),
-      options_, fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator), predicate_.get(),
+          feature_flags_->enable_passing_filter_to_children());
 
-  EXPECT_THAT(GetDocumentIds(&filtered_iterator),
+  EXPECT_THAT(GetDocumentIds(filtered_iterator.get()),
               ElementsAre(document_id1, document_id2, document_id3));
 }
 
@@ -196,11 +257,12 @@
   std::unique_ptr<DocHitInfoIterator> original_iterator =
       std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos);
 
-  DocHitInfoIteratorFilter filtered_iterator(
-      std::move(original_iterator), document_store_.get(), schema_store_.get(),
-      options_, fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator), predicate_.get(),
+          feature_flags_->enable_passing_filter_to_children());
 
-  EXPECT_THAT(filtered_iterator.Advance(),
+  EXPECT_THAT(filtered_iterator->Advance(),
               StatusIs(libtextclassifier3::StatusCode::RESOURCE_EXHAUSTED));
 }
 
@@ -210,11 +272,12 @@
   std::unique_ptr<DocHitInfoIterator> original_iterator =
       std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos);
 
-  DocHitInfoIteratorFilter filtered_iterator(
-      std::move(original_iterator), document_store_.get(), schema_store_.get(),
-      options_, fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator), predicate_.get(),
+          feature_flags_->enable_passing_filter_to_children());
 
-  EXPECT_THAT(filtered_iterator.Advance(),
+  EXPECT_THAT(filtered_iterator->Advance(),
               StatusIs(libtextclassifier3::StatusCode::RESOURCE_EXHAUSTED));
 }
 
@@ -227,11 +290,12 @@
   std::unique_ptr<DocHitInfoIterator> original_iterator =
       std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos);
 
-  DocHitInfoIteratorFilter filtered_iterator(
-      std::move(original_iterator), document_store_.get(), schema_store_.get(),
-      options_, fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator), predicate_.get(),
+          feature_flags_->enable_passing_filter_to_children());
 
-  EXPECT_THAT(filtered_iterator.Advance(),
+  EXPECT_THAT(filtered_iterator->Advance(),
               StatusIs(libtextclassifier3::StatusCode::RESOURCE_EXHAUSTED));
 }
 
@@ -305,19 +369,24 @@
       std::make_unique<DocHitInfoIteratorDummy>();
 
   SearchSpecProto search_spec;
-  DocHitInfoIteratorFilter::Options options =
-      GetFilterOptions(search_spec, *document_store_, *schema_store_);
-  DocHitInfoIteratorFilter filtered_iterator(
-      std::move(original_iterator_empty), document_store_.get(),
-      schema_store_.get(), options, fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocumentFilterPredicate> predicate =
+      GetFilterPredicateBySchemaAndNamespace(
+          search_spec, *document_store_, *schema_store_,
+          fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator_empty), predicate.get(),
+          feature_flags_->enable_passing_filter_to_children());
 
-  EXPECT_THAT(GetDocumentIds(&filtered_iterator), IsEmpty());
+  EXPECT_THAT(GetDocumentIds(filtered_iterator.get()), IsEmpty());
 }
 
 TEST_F(DocHitInfoIteratorNamespaceFilterTest,
        NonexistentNamespacesReturnsEmpty) {
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document1_namespace1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(document1_namespace1_)));
   DocumentId document_id1 = put_result1.new_document_id;
   std::vector<DocHitInfo> doc_hit_infos = {DocHitInfo(document_id1)};
 
@@ -326,18 +395,23 @@
 
   SearchSpecProto search_spec;
   search_spec.add_namespace_filters("nonexistent_namespace");
-  DocHitInfoIteratorFilter::Options options =
-      GetFilterOptions(search_spec, *document_store_, *schema_store_);
-  DocHitInfoIteratorFilter filtered_iterator(
-      std::move(original_iterator), document_store_.get(), schema_store_.get(),
-      options, fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocumentFilterPredicate> predicate =
+      GetFilterPredicateBySchemaAndNamespace(
+          search_spec, *document_store_, *schema_store_,
+          fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator), predicate.get(),
+          feature_flags_->enable_passing_filter_to_children());
 
-  EXPECT_THAT(GetDocumentIds(&filtered_iterator), IsEmpty());
+  EXPECT_THAT(GetDocumentIds(filtered_iterator.get()), IsEmpty());
 }
 
 TEST_F(DocHitInfoIteratorNamespaceFilterTest, NoNamespacesReturnsAll) {
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document1_namespace1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(document1_namespace1_)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   std::vector<DocHitInfo> doc_hit_infos = {DocHitInfo(document_id1)};
@@ -346,25 +420,35 @@
       std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos);
 
   SearchSpecProto search_spec;
-  DocHitInfoIteratorFilter::Options options =
-      GetFilterOptions(search_spec, *document_store_, *schema_store_);
-  DocHitInfoIteratorFilter filtered_iterator(
-      std::move(original_iterator), document_store_.get(), schema_store_.get(),
-      options, fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocumentFilterPredicate> predicate =
+      GetFilterPredicateBySchemaAndNamespace(
+          search_spec, *document_store_, *schema_store_,
+          fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator), predicate.get(),
+          feature_flags_->enable_passing_filter_to_children());
 
-  EXPECT_THAT(GetDocumentIds(&filtered_iterator), ElementsAre(document_id1));
+  EXPECT_THAT(GetDocumentIds(filtered_iterator.get()),
+              ElementsAre(document_id1));
 }
 
 TEST_F(DocHitInfoIteratorNamespaceFilterTest,
        FilterOutExistingDocumentFromDifferentNamespace) {
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document1_namespace1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(document1_namespace1_)));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document2_namespace1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(document2_namespace1_)));
   DocumentId document_id2 = put_result2.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             document_store_->Put(document1_namespace2_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(document1_namespace2_)));
   DocumentId document_id3 = put_result3.new_document_id;
 
   std::vector<DocHitInfo> doc_hit_infos = {DocHitInfo(document_id1),
@@ -376,28 +460,39 @@
 
   SearchSpecProto search_spec;
   search_spec.add_namespace_filters(namespace1_);
-  DocHitInfoIteratorFilter::Options options =
-      GetFilterOptions(search_spec, *document_store_, *schema_store_);
-  DocHitInfoIteratorFilter filtered_iterator(
-      std::move(original_iterator), document_store_.get(), schema_store_.get(),
-      options, fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocumentFilterPredicate> predicate =
+      GetFilterPredicateBySchemaAndNamespace(
+          search_spec, *document_store_, *schema_store_,
+          fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator), predicate.get(),
+          feature_flags_->enable_passing_filter_to_children());
 
-  EXPECT_THAT(GetDocumentIds(&filtered_iterator),
+  EXPECT_THAT(GetDocumentIds(filtered_iterator.get()),
               ElementsAre(document_id1, document_id2));
 }
 
 TEST_F(DocHitInfoIteratorNamespaceFilterTest, FilterForMultipleNamespacesOk) {
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document1_namespace1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(document1_namespace1_)));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document2_namespace1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(document2_namespace1_)));
   DocumentId document_id2 = put_result2.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             document_store_->Put(document1_namespace2_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(document1_namespace2_)));
   DocumentId document_id3 = put_result3.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result4,
-                             document_store_->Put(document1_namespace3_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result4,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(document1_namespace3_)));
   DocumentId document_id4 = put_result4.new_document_id;
 
   std::vector<DocHitInfo> doc_hit_infos = {
@@ -410,13 +505,16 @@
   SearchSpecProto search_spec;
   search_spec.add_namespace_filters(namespace1_);
   search_spec.add_namespace_filters(namespace3_);
-  DocHitInfoIteratorFilter::Options options =
-      GetFilterOptions(search_spec, *document_store_, *schema_store_);
-  DocHitInfoIteratorFilter filtered_iterator(
-      std::move(original_iterator), document_store_.get(), schema_store_.get(),
-      options, fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocumentFilterPredicate> predicate =
+      GetFilterPredicateBySchemaAndNamespace(
+          search_spec, *document_store_, *schema_store_,
+          fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator), predicate.get(),
+          feature_flags_->enable_passing_filter_to_children());
 
-  EXPECT_THAT(GetDocumentIds(&filtered_iterator),
+  EXPECT_THAT(GetDocumentIds(filtered_iterator.get()),
               ElementsAre(document_id1, document_id2, document_id4));
 }
 
@@ -501,19 +599,24 @@
       std::make_unique<DocHitInfoIteratorDummy>();
 
   SearchSpecProto search_spec;
-  DocHitInfoIteratorFilter::Options options =
-      GetFilterOptions(search_spec, *document_store_, *schema_store_);
-  DocHitInfoIteratorFilter filtered_iterator(
-      std::move(original_iterator_empty), document_store_.get(),
-      schema_store_.get(), options, fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocumentFilterPredicate> predicate =
+      GetFilterPredicateBySchemaAndNamespace(
+          search_spec, *document_store_, *schema_store_,
+          fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator_empty), predicate.get(),
+          feature_flags_->enable_passing_filter_to_children());
 
-  EXPECT_THAT(GetDocumentIds(&filtered_iterator), IsEmpty());
+  EXPECT_THAT(GetDocumentIds(filtered_iterator.get()), IsEmpty());
 }
 
 TEST_F(DocHitInfoIteratorSchemaTypeFilterTest,
        NonexistentSchemaTypeReturnsEmpty) {
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document1_schema1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(document1_schema1_)));
   DocumentId document_id1 = put_result1.new_document_id;
   std::vector<DocHitInfo> doc_hit_infos = {DocHitInfo(document_id1)};
 
@@ -522,18 +625,23 @@
 
   SearchSpecProto search_spec;
   search_spec.add_schema_type_filters("nonexistent_schema_type");
-  DocHitInfoIteratorFilter::Options options =
-      GetFilterOptions(search_spec, *document_store_, *schema_store_);
-  DocHitInfoIteratorFilter filtered_iterator(
-      std::move(original_iterator), document_store_.get(), schema_store_.get(),
-      options, fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocumentFilterPredicate> predicate =
+      GetFilterPredicateBySchemaAndNamespace(
+          search_spec, *document_store_, *schema_store_,
+          fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator), predicate.get(),
+          feature_flags_->enable_passing_filter_to_children());
 
-  EXPECT_THAT(GetDocumentIds(&filtered_iterator), IsEmpty());
+  EXPECT_THAT(GetDocumentIds(filtered_iterator.get()), IsEmpty());
 }
 
 TEST_F(DocHitInfoIteratorSchemaTypeFilterTest, NoSchemaTypesReturnsAll) {
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document1_schema1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(document1_schema1_)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   std::vector<DocHitInfo> doc_hit_infos = {DocHitInfo(document_id1)};
@@ -542,22 +650,30 @@
       std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos);
 
   SearchSpecProto search_spec;
-  DocHitInfoIteratorFilter::Options options =
-      GetFilterOptions(search_spec, *document_store_, *schema_store_);
-  DocHitInfoIteratorFilter filtered_iterator(
-      std::move(original_iterator), document_store_.get(), schema_store_.get(),
-      options, fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocumentFilterPredicate> predicate =
+      GetFilterPredicateBySchemaAndNamespace(
+          search_spec, *document_store_, *schema_store_,
+          fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator), predicate.get(),
+          feature_flags_->enable_passing_filter_to_children());
 
-  EXPECT_THAT(GetDocumentIds(&filtered_iterator), ElementsAre(document_id1));
+  EXPECT_THAT(GetDocumentIds(filtered_iterator.get()),
+              ElementsAre(document_id1));
 }
 
 TEST_F(DocHitInfoIteratorSchemaTypeFilterTest,
        FilterOutExistingDocumentFromDifferentSchemaTypes) {
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document1_schema1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(document1_schema1_)));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document2_schema2_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(document2_schema2_)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   std::vector<DocHitInfo> doc_hit_infos = {DocHitInfo(document_id1),
@@ -568,24 +684,34 @@
 
   SearchSpecProto search_spec;
   search_spec.add_schema_type_filters(std::string(kSchema1));
-  DocHitInfoIteratorFilter::Options options =
-      GetFilterOptions(search_spec, *document_store_, *schema_store_);
-  DocHitInfoIteratorFilter filtered_iterator(
-      std::move(original_iterator), document_store_.get(), schema_store_.get(),
-      options, fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocumentFilterPredicate> predicate =
+      GetFilterPredicateBySchemaAndNamespace(
+          search_spec, *document_store_, *schema_store_,
+          fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator), predicate.get(),
+          feature_flags_->enable_passing_filter_to_children());
 
-  EXPECT_THAT(GetDocumentIds(&filtered_iterator), ElementsAre(document_id1));
+  EXPECT_THAT(GetDocumentIds(filtered_iterator.get()),
+              ElementsAre(document_id1));
 }
 
 TEST_F(DocHitInfoIteratorSchemaTypeFilterTest, FilterForMultipleSchemaTypesOk) {
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document1_schema1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(document1_schema1_)));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document2_schema2_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(document2_schema2_)));
   DocumentId document_id2 = put_result2.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             document_store_->Put(document3_schema3_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(document3_schema3_)));
   DocumentId document_id3 = put_result3.new_document_id;
   std::vector<DocHitInfo> doc_hit_infos = {DocHitInfo(document_id1),
                                            DocHitInfo(document_id2),
@@ -597,41 +723,50 @@
   SearchSpecProto search_spec;
   search_spec.add_schema_type_filters(std::string(kSchema2));
   search_spec.add_schema_type_filters(std::string(kSchema3));
-  DocHitInfoIteratorFilter::Options options =
-      GetFilterOptions(search_spec, *document_store_, *schema_store_);
-  DocHitInfoIteratorFilter filtered_iterator(
-      std::move(original_iterator), document_store_.get(), schema_store_.get(),
-      options, fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocumentFilterPredicate> predicate =
+      GetFilterPredicateBySchemaAndNamespace(
+          search_spec, *document_store_, *schema_store_,
+          fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator), predicate.get(),
+          feature_flags_->enable_passing_filter_to_children());
 
-  EXPECT_THAT(GetDocumentIds(&filtered_iterator),
+  EXPECT_THAT(GetDocumentIds(filtered_iterator.get()),
               ElementsAre(document_id2, document_id3));
 }
 
 TEST_F(DocHitInfoIteratorSchemaTypeFilterTest,
        FilterIsExactForSchemaTypePolymorphism) {
   // Add some irrelevant documents.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document1_schema1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(document1_schema1_)));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document2_schema2_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(document2_schema2_)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // Create a person document and an artist document, where the artist should be
   // able to be interpreted as a person by polymorphism.
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult person_put_result,
-      document_store_->Put(DocumentBuilder()
-                               .SetKey("namespace", "person")
-                               .SetSchema("person")
-                               .Build()));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          DocumentBuilder()
+              .SetKey("namespace", "person")
+              .SetSchema("person")
+              .Build())));
   DocumentId person_document_id = person_put_result.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult artist_put_result,
-      document_store_->Put(DocumentBuilder()
-                               .SetKey("namespace", "artist")
-                               .SetSchema("artist")
-                               .Build()));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          DocumentBuilder()
+              .SetKey("namespace", "artist")
+              .SetSchema("artist")
+              .Build())));
   DocumentId artist_document_id = artist_put_result.new_document_id;
 
   std::vector<DocHitInfo> doc_hit_infos = {
@@ -644,23 +779,29 @@
       std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos);
   SearchSpecProto search_spec_1;
   search_spec_1.add_schema_type_filters("person");
-  DocHitInfoIteratorFilter::Options options =
-      GetFilterOptions(search_spec_1, *document_store_, *schema_store_);
-  DocHitInfoIteratorFilter filtered_iterator_1(
-      std::move(original_iterator), document_store_.get(), schema_store_.get(),
-      options, fake_clock_.GetSystemTimeMilliseconds());
-  EXPECT_THAT(GetDocumentIds(&filtered_iterator_1),
+  std::unique_ptr<DocumentFilterPredicate> predicate =
+      GetFilterPredicateBySchemaAndNamespace(
+          search_spec_1, *document_store_, *schema_store_,
+          fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator_1 =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator), predicate.get(),
+          feature_flags_->enable_passing_filter_to_children());
+  EXPECT_THAT(GetDocumentIds(filtered_iterator_1.get()),
               ElementsAre(person_document_id));
 
   // Filters for the "artist" type should not include the "person" type.
   original_iterator = std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos);
   SearchSpecProto search_spec_2;
   search_spec_2.add_schema_type_filters("artist");
-  options = GetFilterOptions(search_spec_2, *document_store_, *schema_store_);
-  DocHitInfoIteratorFilter filtered_iterator_2(
-      std::move(original_iterator), document_store_.get(), schema_store_.get(),
-      options, fake_clock_.GetSystemTimeMilliseconds());
-  EXPECT_THAT(GetDocumentIds(&filtered_iterator_2),
+  predicate = GetFilterPredicateBySchemaAndNamespace(
+      search_spec_2, *document_store_, *schema_store_,
+      fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator_2 =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator), predicate.get(),
+          feature_flags_->enable_passing_filter_to_children());
+  EXPECT_THAT(GetDocumentIds(filtered_iterator_2.get()),
               ElementsAre(artist_document_id));
 }
 
@@ -669,27 +810,30 @@
   // Create an email and a message document.
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult email_put_result,
-      document_store_->Put(DocumentBuilder()
-                               .SetKey("namespace", "email")
-                               .SetSchema("email")
-                               .Build()));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "email")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId email_document_id = email_put_result.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult message_put_result,
-      document_store_->Put(DocumentBuilder()
-                               .SetKey("namespace", "message")
-                               .SetSchema("message")
-                               .Build()));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          DocumentBuilder()
+              .SetKey("namespace", "message")
+              .SetSchema("message")
+              .Build())));
   DocumentId message_document_id = message_put_result.new_document_id;
 
   // Create a emailMessage document, which the should be able to be interpreted
   // as both an email and a message by polymorphism.
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult email_message_put_result,
-      document_store_->Put(DocumentBuilder()
-                               .SetKey("namespace", "emailMessage")
-                               .SetSchema("emailMessage")
-                               .Build()));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          DocumentBuilder()
+              .SetKey("namespace", "emailMessage")
+              .SetSchema("emailMessage")
+              .Build())));
   DocumentId email_message_document_id =
       email_message_put_result.new_document_id;
 
@@ -703,12 +847,15 @@
       std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos);
   SearchSpecProto search_spec_1;
   search_spec_1.add_schema_type_filters("email");
-  DocHitInfoIteratorFilter::Options options =
-      GetFilterOptions(search_spec_1, *document_store_, *schema_store_);
-  DocHitInfoIteratorFilter filtered_iterator_1(
-      std::move(original_iterator), document_store_.get(), schema_store_.get(),
-      options, fake_clock_.GetSystemTimeMilliseconds());
-  EXPECT_THAT(GetDocumentIds(&filtered_iterator_1),
+  std::unique_ptr<DocumentFilterPredicate> predicate =
+      GetFilterPredicateBySchemaAndNamespace(
+          search_spec_1, *document_store_, *schema_store_,
+          fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator_1 =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator), predicate.get(),
+          feature_flags_->enable_passing_filter_to_children());
+  EXPECT_THAT(GetDocumentIds(filtered_iterator_1.get()),
               ElementsAre(email_document_id));
 
   // Filters for the "message" type should NOT include the "emailMessage" type,
@@ -716,22 +863,28 @@
   original_iterator = std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos);
   SearchSpecProto search_spec_2;
   search_spec_2.add_schema_type_filters("message");
-  options = GetFilterOptions(search_spec_2, *document_store_, *schema_store_);
-  DocHitInfoIteratorFilter filtered_iterator_2(
-      std::move(original_iterator), document_store_.get(), schema_store_.get(),
-      options, fake_clock_.GetSystemTimeMilliseconds());
-  EXPECT_THAT(GetDocumentIds(&filtered_iterator_2),
+  predicate = GetFilterPredicateBySchemaAndNamespace(
+      search_spec_2, *document_store_, *schema_store_,
+      fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator_2 =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator), predicate.get(),
+          feature_flags_->enable_passing_filter_to_children());
+  EXPECT_THAT(GetDocumentIds(filtered_iterator_2.get()),
               ElementsAre(message_document_id));
 
   // Filters for a irrelevant type should return nothing.
   original_iterator = std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos);
   SearchSpecProto search_spec_3;
   search_spec_3.add_schema_type_filters("person");
-  options = GetFilterOptions(search_spec_3, *document_store_, *schema_store_);
-  DocHitInfoIteratorFilter filtered_iterator_3(
-      std::move(original_iterator), document_store_.get(), schema_store_.get(),
-      options, fake_clock_.GetSystemTimeMilliseconds());
-  EXPECT_THAT(GetDocumentIds(&filtered_iterator_3), IsEmpty());
+  predicate = GetFilterPredicateBySchemaAndNamespace(
+      search_spec_3, *document_store_, *schema_store_,
+      fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator_3 =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator), predicate.get(),
+          feature_flags_->enable_passing_filter_to_children());
+  EXPECT_THAT(GetDocumentIds(filtered_iterator_3.get()), IsEmpty());
 }
 
 class DocHitInfoIteratorExpirationFilterTest : public ::testing::Test {
@@ -758,6 +911,10 @@
         CreateDocumentStore(&filesystem_, test_dir_, &fake_clock_,
                             schema_store_.get(), *feature_flags_));
     document_store_ = std::move(create_result.document_store);
+
+    predicate_ = GetFilterPredicateBySchemaAndNamespace(
+        SearchSpecProto::default_instance(), *document_store_, *schema_store_,
+        fake_clock_.GetSystemTimeMilliseconds());
   }
 
   void TearDown() override {
@@ -775,12 +932,15 @@
   const Filesystem filesystem_;
   const std::string test_dir_;
   const std::string email_schema_ = "email";
-  DocHitInfoIteratorFilter::Options options_;
+  std::unique_ptr<DocumentFilterPredicate> predicate_;
 };
 
 TEST_F(DocHitInfoIteratorExpirationFilterTest, TtlZeroIsntFilteredOut) {
   // Arbitrary value
   fake_clock_.SetSystemTimeMilliseconds(100);
+  predicate_ = GetFilterPredicateBySchemaAndNamespace(
+      SearchSpecProto::default_instance(), *document_store_, *schema_store_,
+      fake_clock_.GetSystemTimeMilliseconds());
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::CreateResult create_result,
@@ -796,24 +956,30 @@
                                .SetCreationTimestampMs(0)
                                .SetTtlMs(0)
                                .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   std::vector<DocHitInfo> doc_hit_infos = {DocHitInfo(document_id1)};
   std::unique_ptr<DocHitInfoIterator> original_iterator =
       std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos);
 
-  DocHitInfoIteratorFilter filtered_iterator(
-      std::move(original_iterator), document_store.get(), schema_store_.get(),
-      options_, fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator), predicate_.get(),
+          feature_flags_->enable_passing_filter_to_children());
 
-  EXPECT_THAT(GetDocumentIds(&filtered_iterator), ElementsAre(document_id1));
+  EXPECT_THAT(GetDocumentIds(filtered_iterator.get()),
+              ElementsAre(document_id1));
 }
 
 TEST_F(DocHitInfoIteratorExpirationFilterTest, BeforeTtlNotFilteredOut) {
   // Arbitrary value, but must be less than document's creation_timestamp + ttl
   fake_clock_.SetSystemTimeMilliseconds(50);
+  predicate_ = GetFilterPredicateBySchemaAndNamespace(
+      SearchSpecProto::default_instance(), *document_store_, *schema_store_,
+      fake_clock_.GetSystemTimeMilliseconds());
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::CreateResult create_result,
@@ -829,24 +995,30 @@
                                .SetCreationTimestampMs(1)
                                .SetTtlMs(100)
                                .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   std::vector<DocHitInfo> doc_hit_infos = {DocHitInfo(document_id1)};
   std::unique_ptr<DocHitInfoIterator> original_iterator =
       std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos);
 
-  DocHitInfoIteratorFilter filtered_iterator(
-      std::move(original_iterator), document_store.get(), schema_store_.get(),
-      options_, fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator), predicate_.get(),
+          feature_flags_->enable_passing_filter_to_children());
 
-  EXPECT_THAT(GetDocumentIds(&filtered_iterator), ElementsAre(document_id1));
+  EXPECT_THAT(GetDocumentIds(filtered_iterator.get()),
+              ElementsAre(document_id1));
 }
 
 TEST_F(DocHitInfoIteratorExpirationFilterTest, EqualTtlFilteredOut) {
   // Current time is exactly the document's creation_timestamp + ttl
   fake_clock_.SetSystemTimeMilliseconds(150);
+  predicate_ = GetFilterPredicateBySchemaAndNamespace(
+      SearchSpecProto::default_instance(), *document_store_, *schema_store_,
+      fake_clock_.GetSystemTimeMilliseconds());
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::CreateResult create_result,
@@ -862,25 +1034,30 @@
                                .SetCreationTimestampMs(50)
                                .SetTtlMs(100)
                                .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   std::vector<DocHitInfo> doc_hit_infos = {DocHitInfo(document_id1)};
   std::unique_ptr<DocHitInfoIterator> original_iterator =
       std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos);
 
-  DocHitInfoIteratorFilter filtered_iterator(
-      std::move(original_iterator), document_store.get(), schema_store_.get(),
-      options_, fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator), predicate_.get(),
+          feature_flags_->enable_passing_filter_to_children());
 
-  EXPECT_THAT(GetDocumentIds(&filtered_iterator), IsEmpty());
+  EXPECT_THAT(GetDocumentIds(filtered_iterator.get()), IsEmpty());
 }
 
 TEST_F(DocHitInfoIteratorExpirationFilterTest, PastTtlFilteredOut) {
   // Arbitrary value, but must be greater than the document's
   // creation_timestamp + ttl
   fake_clock_.SetSystemTimeMilliseconds(151);
+  predicate_ = GetFilterPredicateBySchemaAndNamespace(
+      SearchSpecProto::default_instance(), *document_store_, *schema_store_,
+      fake_clock_.GetSystemTimeMilliseconds());
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::CreateResult create_result,
@@ -896,19 +1073,21 @@
                                .SetCreationTimestampMs(50)
                                .SetTtlMs(100)
                                .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   std::vector<DocHitInfo> doc_hit_infos = {DocHitInfo(document_id1)};
   std::unique_ptr<DocHitInfoIterator> original_iterator =
       std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos);
 
-  DocHitInfoIteratorFilter filtered_iterator(
-      std::move(original_iterator), document_store.get(), schema_store_.get(),
-      options_, fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator), predicate_.get(),
+          feature_flags_->enable_passing_filter_to_children());
 
-  EXPECT_THAT(GetDocumentIds(&filtered_iterator), IsEmpty());
+  EXPECT_THAT(GetDocumentIds(filtered_iterator.get()), IsEmpty());
 }
 
 class DocHitInfoIteratorFilterTest : public ::testing::Test {
@@ -1005,23 +1184,28 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result1,
-      document_store->Put(document1_namespace1_schema1_));
+      document_store->Put(
+          document_util::CreateDocumentWrapper(document1_namespace1_schema1_)));
   DocumentId document_id1 = put_result1.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result2,
-      document_store->Put(document2_namespace1_schema1_));
+      document_store->Put(
+          document_util::CreateDocumentWrapper(document2_namespace1_schema1_)));
   DocumentId document_id2 = put_result2.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result3,
-      document_store->Put(document3_namespace2_schema1_));
+      document_store->Put(
+          document_util::CreateDocumentWrapper(document3_namespace2_schema1_)));
   DocumentId document_id3 = put_result3.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result4,
-      document_store->Put(document4_namespace1_schema2_));
+      document_store->Put(
+          document_util::CreateDocumentWrapper(document4_namespace1_schema2_)));
   DocumentId document_id4 = put_result4.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result5,
-      document_store->Put(document5_namespace1_schema1_));
+      document_store->Put(
+          document_util::CreateDocumentWrapper(document5_namespace1_schema1_)));
   DocumentId document_id5 = put_result5.new_document_id;
 
   // Deletes document2, causing it to be filtered out
@@ -1043,28 +1227,35 @@
   search_spec.add_namespace_filters(namespace1_);
   // Filters out document4 by schema type
   search_spec.add_schema_type_filters(schema1_);
-  DocHitInfoIteratorFilter::Options options =
-      GetFilterOptions(search_spec, *document_store, *schema_store_);
+  std::unique_ptr<DocumentFilterPredicate> predicate =
+      GetFilterPredicateBySchemaAndNamespace(
+          search_spec, *document_store, *schema_store_,
+          fake_clock_.GetSystemTimeMilliseconds());
 
-  DocHitInfoIteratorFilter filtered_iterator(
-      std::move(original_iterator), document_store.get(), schema_store_.get(),
-      options, fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator), predicate.get(),
+          feature_flags_->enable_passing_filter_to_children());
 
-  EXPECT_THAT(GetDocumentIds(&filtered_iterator), ElementsAre(document_id1));
+  EXPECT_THAT(GetDocumentIds(filtered_iterator.get()),
+              ElementsAre(document_id1));
 }
 
 TEST_F(DocHitInfoIteratorFilterTest, SectionIdMasksArePopulatedCorrectly) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result1,
-      document_store_->Put(document1_namespace1_schema1_));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(document1_namespace1_schema1_)));
   DocumentId document_id1 = put_result1.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result2,
-      document_store_->Put(document2_namespace1_schema1_));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(document2_namespace1_schema1_)));
   DocumentId document_id2 = put_result2.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result3,
-      document_store_->Put(document3_namespace2_schema1_));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(document3_namespace2_schema1_)));
   DocumentId document_id3 = put_result3.new_document_id;
 
   SectionIdMask section_id_mask1 = 0b01001001;  // hits in sections 0, 3, 6
@@ -1081,12 +1272,16 @@
   std::unique_ptr<DocHitInfoIterator> original_iterator =
       std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos);
 
-  DocHitInfoIteratorFilter::Options options;
-  DocHitInfoIteratorFilter filtered_iterator(
-      std::move(original_iterator), document_store_.get(), schema_store_.get(),
-      options, fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocumentFilterPredicate> predicate =
+      GetFilterPredicateBySchemaAndNamespace(
+          SearchSpecProto::default_instance(), *document_store_, *schema_store_,
+          fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator), predicate.get(),
+          feature_flags_->enable_passing_filter_to_children());
 
-  EXPECT_THAT(GetDocHitInfos(&filtered_iterator),
+  EXPECT_THAT(GetDocHitInfos(filtered_iterator.get()),
               ElementsAre(EqualsDocHitInfo(document_id1, section_ids1),
                           EqualsDocHitInfo(document_id2, section_ids2),
                           EqualsDocHitInfo(document_id3, section_ids3)));
@@ -1098,30 +1293,42 @@
       /*num_leaf_advance_calls_main_index_in=*/5,
       /*num_leaf_advance_calls_integer_index_in=*/3,
       /*num_leaf_advance_calls_no_index_in=*/1,
-      /*num_blocks_inspected_in=*/4);  // arbitrary value
+      /*num_blocks_inspected_in=*/4,
+      /*embedding_stats_in=*/
+      {.num_unquantized_embeddings_scored = 2,
+       .num_quantized_embeddings_scored = 3,
+       .unquantized_shards_read = {1, 2},
+       .quantized_shards_read{3, 4}});  // arbitrary value
   auto original_iterator = std::make_unique<DocHitInfoIteratorDummy>();
   original_iterator->SetCallStats(original_call_stats);
 
-  DocHitInfoIteratorFilter::Options options;
-  DocHitInfoIteratorFilter filtered_iterator(
-      std::move(original_iterator), document_store_.get(), schema_store_.get(),
-      options, fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocumentFilterPredicate> predicate =
+      GetFilterPredicateBySchemaAndNamespace(
+          SearchSpecProto::default_instance(), *document_store_, *schema_store_,
+          fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator), predicate.get(),
+          feature_flags_->enable_passing_filter_to_children());
 
-  EXPECT_THAT(filtered_iterator.GetCallStats(), Eq(original_call_stats));
+  EXPECT_THAT(filtered_iterator->GetCallStats(), Eq(original_call_stats));
 }
 
 TEST_F(DocHitInfoIteratorFilterTest, TrimFilterIterator) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result1,
-      document_store_->Put(document1_namespace1_schema1_));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(document1_namespace1_schema1_)));
   DocumentId document_id1 = put_result1.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result2,
-      document_store_->Put(document2_namespace1_schema1_));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(document2_namespace1_schema1_)));
   DocumentId document_id2 = put_result2.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result3,
-      document_store_->Put(document3_namespace2_schema1_));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(document3_namespace2_schema1_)));
   DocumentId document_id3 = put_result3.new_document_id;
 
   // Build an interator tree like:
@@ -1146,24 +1353,549 @@
   SearchSpecProto search_spec;
   // Filters out document3 by namespace
   search_spec.add_namespace_filters(namespace1_);
-  DocHitInfoIteratorFilter::Options options =
-      GetFilterOptions(search_spec, *document_store_, *schema_store_);
-  DocHitInfoIteratorFilter filtered_iterator(
-      std::move(original_iterator), document_store_.get(), schema_store_.get(),
-      options, fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocumentFilterPredicate> predicate =
+      GetFilterPredicateBySchemaAndNamespace(
+          search_spec, *document_store_, *schema_store_,
+          fake_clock_.GetSystemTimeMilliseconds());
+  std::unique_ptr<DocHitInfoIterator> filtered_iterator =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_iterator), predicate.get(),
+          feature_flags_->enable_passing_filter_to_children());
 
   // The trimmed tree.
   //          Filter
   //             |
   //          {1, 3}
   ICING_ASSERT_OK_AND_ASSIGN(DocHitInfoIterator::TrimmedNode trimmed_node,
-                             std::move(filtered_iterator).TrimRightMostNode());
+                             std::move(*filtered_iterator).TrimRightMostNode());
   EXPECT_THAT(trimmed_node.term_, Eq("term"));
   EXPECT_THAT(trimmed_node.term_start_index_, Eq(10));
   EXPECT_THAT(GetDocumentIds(trimmed_node.iterator_.get()),
               ElementsAre(document_id1));
 }
 
+TEST_F(DocHitInfoIteratorFilterTest,
+       ApplyFilterReturnsOriginalWhenRootHandles) {
+  SearchSpecProto search_spec_1;
+  search_spec_1.add_schema_type_filters(schema1_);
+  search_spec_1.add_namespace_filters(namespace1_);
+  std::unique_ptr<DocumentFilterPredicate> predicate_1 =
+      GetFilterPredicateBySchemaAndNamespace(
+          search_spec_1, *document_store_, *schema_store_,
+          fake_clock_.GetSystemTimeMilliseconds());
+  SearchSpecProto search_spec_2;
+  search_spec_2.add_schema_type_filters(schema2_);
+  search_spec_2.add_namespace_filters(namespace2_);
+  std::unique_ptr<DocumentFilterPredicate> predicate_2 =
+      GetFilterPredicateBySchemaAndNamespace(
+          search_spec_2, *document_store_, *schema_store_,
+          fake_clock_.GetSystemTimeMilliseconds());
+
+  // Apply predicate_1 to an iterator that can handle filter internally.
+  std::unique_ptr<DocHitInfoIteratorDummyHandlingFilter> original_root =
+      std::make_unique<DocHitInfoIteratorDummyHandlingFilter>();
+  DocHitInfoIteratorDummyHandlingFilter* original_root_ptr =
+      original_root.get();
+  std::unique_ptr<DocHitInfoIterator> new_root =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_root), predicate_1.get(),
+          feature_flags_->enable_passing_filter_to_children());
+  // The filter is applied internally, so the pointer should not change.
+  EXPECT_THAT(new_root.get(), Eq(original_root_ptr));
+  // new_root is not a DocHitInfoIteratorFilter.
+  EXPECT_THAT(new_root.get(),
+              WhenDynamicCastTo<DocHitInfoIteratorFilter*>(nullptr));
+  // Check that the predicate is applied correctly.
+  EXPECT_THAT(original_root_ptr->document_filter_predicates(),
+              ElementsAre(predicate_1.get()));
+
+  // Apply predicate_2 to the same iterator.
+  new_root = DocHitInfoIteratorFilter::ApplyFilter(
+      std::move(new_root), predicate_2.get(),
+      feature_flags_->enable_passing_filter_to_children());
+  // The pointer should not change for the same reason.
+  EXPECT_THAT(new_root.get(), Eq(original_root_ptr));
+  // new_root is not a DocHitInfoIteratorFilter.
+  EXPECT_THAT(new_root.get(),
+              WhenDynamicCastTo<DocHitInfoIteratorFilter*>(nullptr));
+  // Check that the predicates are applied correctly.
+  EXPECT_THAT(original_root_ptr->document_filter_predicates(),
+              ElementsAre(predicate_1.get(), predicate_2.get()));
+}
+
+TEST_F(DocHitInfoIteratorFilterTest,
+       ApplyFilterReturnsOriginalWhenAllChildrenHandle) {
+  // Create an iterator tree:
+  //                        AND (root)
+  //                         /    \
+  //       DummyHandlingFilter    DummyHandlingFilter
+  std::unique_ptr<DocHitInfoIteratorDummyHandlingFilter> child1 =
+      std::make_unique<DocHitInfoIteratorDummyHandlingFilter>();
+  std::unique_ptr<DocHitInfoIteratorDummyHandlingFilter> child2 =
+      std::make_unique<DocHitInfoIteratorDummyHandlingFilter>();
+  DocHitInfoIteratorDummyHandlingFilter* child1_ptr = child1.get();
+  DocHitInfoIteratorDummyHandlingFilter* child2_ptr = child2.get();
+  std::unique_ptr<DocHitInfoIterator> original_root =
+      std::make_unique<DocHitInfoIteratorAnd>(std::move(child1),
+                                              std::move(child2));
+  DocHitInfoIterator* original_root_ptr = original_root.get();
+
+  // Test that AND iterator can pass predicate through, and since all children
+  // can handle the filter, the original iterator is returned.
+  SearchSpecProto search_spec;
+  search_spec.add_schema_type_filters(schema1_);
+  search_spec.add_namespace_filters(namespace1_);
+  std::unique_ptr<DocumentFilterPredicate> predicate =
+      GetFilterPredicateBySchemaAndNamespace(
+          search_spec, *document_store_, *schema_store_,
+          fake_clock_.GetSystemTimeMilliseconds());
+  // We should get the same tree:
+  //                        AND (root)
+  //                         /    \
+  //       DummyHandlingFilter    DummyHandlingFilter
+  std::unique_ptr<DocHitInfoIterator> new_root =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_root), predicate.get(),
+          feature_flags_->enable_passing_filter_to_children());
+
+  // The filter should be handled by children, returning the original iterator.
+  EXPECT_THAT(new_root.get(), Eq(original_root_ptr));
+  EXPECT_THAT(new_root.get(),
+              WhenDynamicCastTo<DocHitInfoIteratorFilter*>(nullptr));
+  // Verify the predicate was passed down to children.
+  EXPECT_THAT(child1_ptr->document_filter_predicates(),
+              ElementsAre(predicate.get()));
+  EXPECT_THAT(child2_ptr->document_filter_predicates(),
+              ElementsAre(predicate.get()));
+}
+
+TEST_F(DocHitInfoIteratorFilterTest,
+       ApplyFilterReturnsNewWhenSomeChildrenHandle) {
+  // Create an iterator tree:
+  //                        AND (root)
+  //                         /    \
+  //       DummyHandlingFilter    Dummy
+  std::unique_ptr<DocHitInfoIteratorDummyHandlingFilter> child1 =
+      std::make_unique<DocHitInfoIteratorDummyHandlingFilter>();
+  std::unique_ptr<DocHitInfoIteratorDummy> child2 =
+      std::make_unique<DocHitInfoIteratorDummy>();
+  DocHitInfoIteratorDummyHandlingFilter* child1_ptr = child1.get();
+  std::unique_ptr<DocHitInfoIterator> original_root =
+      std::make_unique<DocHitInfoIteratorAnd>(std::move(child1),
+                                              std::move(child2));
+  DocHitInfoIterator* original_root_ptr = original_root.get();
+
+  // Create a predicate.
+  SearchSpecProto search_spec;
+  search_spec.add_schema_type_filters(schema1_);
+  search_spec.add_namespace_filters(namespace1_);
+  std::unique_ptr<DocumentFilterPredicate> predicate =
+      GetFilterPredicateBySchemaAndNamespace(
+          search_spec, *document_store_, *schema_store_,
+          fake_clock_.GetSystemTimeMilliseconds());
+
+  // Apply the filter. Since not all children can handle the filter, a new
+  // filter iterator should be created at the top. However, DummyHandlingFilter
+  // should still get the predicate and try to handle it internally.
+  //
+  // We should get a new tree:
+  //                      Filter (root)
+  //                           |
+  //                          AND
+  //                         /   \
+  //       DummyHandlingFilter   Dummy
+  std::unique_ptr<DocHitInfoIterator> new_root =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_root), predicate.get(),
+          feature_flags_->enable_passing_filter_to_children());
+
+  // The filter should be applied at the top level, so the pointer should
+  // change and be of type DocHitInfoIteratorFilter.
+  EXPECT_THAT(new_root.get(), Ne(original_root_ptr));
+  ASSERT_THAT(new_root.get(),
+              WhenDynamicCastTo<DocHitInfoIteratorFilter*>(Ne(nullptr)));
+  // Check that the original iterator is now the delegate.
+  EXPECT_THAT(new_root->GetChildren(),
+              ElementsAre(Pointee(Pointer(original_root_ptr))));
+  // Verify the predicate was passed down to the child that can handle it.
+  EXPECT_THAT(child1_ptr->document_filter_predicates(),
+              ElementsAre(predicate.get()));
+  // The other child (Dummy) doesn't store predicates, so no check needed there.
+}
+
+TEST_F(DocHitInfoIteratorFilterTest,
+       ApplyFilterReturnsNewWhenNoChildrenHandle) {
+  // Create an iterator tree:
+  //                        AND (root)
+  //                         /    \
+  //                     Dummy    Dummy
+  std::unique_ptr<DocHitInfoIterator> original_root =
+      std::make_unique<DocHitInfoIteratorAnd>(
+          std::make_unique<DocHitInfoIteratorDummy>(),
+          std::make_unique<DocHitInfoIteratorDummy>());
+  DocHitInfoIterator* original_root_ptr = original_root.get();
+
+  // Create a predicate.
+  SearchSpecProto search_spec;
+  search_spec.add_schema_type_filters(schema1_);
+  search_spec.add_namespace_filters(namespace1_);
+  std::unique_ptr<DocumentFilterPredicate> predicate =
+      GetFilterPredicateBySchemaAndNamespace(
+          search_spec, *document_store_, *schema_store_,
+          fake_clock_.GetSystemTimeMilliseconds());
+
+  // Apply the filter. Since no children can handle the filter, a new filter
+  // iterator should be created at the top.
+  //
+  // We should get a new tree:
+  //                      Filter (root)
+  //                           |
+  //                          AND
+  //                         /   \
+  //                     Dummy   Dummy
+  std::unique_ptr<DocHitInfoIterator> new_root =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_root), predicate.get(),
+          feature_flags_->enable_passing_filter_to_children());
+
+  // The filter should be applied at the top level, so the pointer should
+  // change and be of type DocHitInfoIteratorFilter.
+  EXPECT_THAT(new_root.get(), Ne(original_root_ptr));
+  ASSERT_THAT(new_root.get(),
+              WhenDynamicCastTo<DocHitInfoIteratorFilter*>(Ne(nullptr)));
+  // Check that the original iterator is now the delegate.
+  EXPECT_THAT(new_root->GetChildren(),
+              ElementsAre(Pointee(Pointer(original_root_ptr))));
+}
+
+TEST_F(DocHitInfoIteratorFilterTest,
+       ApplyFilterReturnsNewWhenRootHasNoChildren) {
+  // Create a root iterator that cannot handle filters and has no children.
+  std::unique_ptr<DocHitInfoIterator> original_root =
+      std::make_unique<DocHitInfoIteratorDummy>();
+  DocHitInfoIterator* original_root_ptr = original_root.get();
+
+  // Create a predicate.
+  SearchSpecProto search_spec;
+  search_spec.add_schema_type_filters(schema1_);
+  std::unique_ptr<DocumentFilterPredicate> predicate =
+      GetFilterPredicateBySchemaAndNamespace(
+          search_spec, *document_store_, *schema_store_,
+          fake_clock_.GetSystemTimeMilliseconds());
+
+  // Apply the filter. Since the root cannot handle the filter and has no
+  // children, a new filter iterator should be created at the top.
+  //
+  // We should get a new tree:
+  //                      Filter (root)
+  //                           |
+  //                         Dummy
+  std::unique_ptr<DocHitInfoIterator> new_root =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_root), predicate.get(),
+          feature_flags_->enable_passing_filter_to_children());
+
+  // The filter should be applied at the top level, so the pointer should
+  // change and be of type DocHitInfoIteratorFilter.
+  EXPECT_THAT(new_root.get(), Ne(original_root_ptr));
+  ASSERT_THAT(new_root.get(),
+              WhenDynamicCastTo<DocHitInfoIteratorFilter*>(Ne(nullptr)));
+  // Check that the original iterator is now the delegate.
+  EXPECT_THAT(new_root->GetChildren(),
+              ElementsAre(Pointee(Pointer(original_root_ptr))));
+}
+
+TEST_F(DocHitInfoIteratorFilterTest,
+       ApplyFilterToIteratorThatCannotPassPredicateThrough) {
+  // Create an iterator tree:
+  //                NOT (root)
+  //                 |
+  //         DummyHandlingFilter
+  std::unique_ptr<DocHitInfoIteratorDummyHandlingFilter>
+      iterator_handling_filter =
+          std::make_unique<DocHitInfoIteratorDummyHandlingFilter>();
+  DocHitInfoIteratorDummyHandlingFilter* iterator_handling_filter_ptr =
+      iterator_handling_filter.get();
+  std::unique_ptr<DocHitInfoIterator> original_root =
+      std::make_unique<DocHitInfoIteratorNot>(
+          std::move(iterator_handling_filter),
+          /*document_id_limit=*/100);
+  DocHitInfoIterator* original_root_ptr = original_root.get();
+
+  // Test that NOT iterator cannot pass predicate through, and the filter is
+  // applied at the top level.
+  SearchSpecProto search_spec;
+  search_spec.add_schema_type_filters(schema1_);
+  search_spec.add_namespace_filters(namespace1_);
+  std::unique_ptr<DocumentFilterPredicate> predicate =
+      GetFilterPredicateBySchemaAndNamespace(
+          search_spec, *document_store_, *schema_store_,
+          fake_clock_.GetSystemTimeMilliseconds());
+  // We should get a new tree:
+  //            Filter (root)
+  //                 |
+  //                NOT
+  //                 |
+  //         DummyHandlingFilter
+  std::unique_ptr<DocHitInfoIterator> new_root =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_root), predicate.get(),
+          feature_flags_->enable_passing_filter_to_children());
+
+  // The filter should be applied at the top level, so the pointer should
+  // change and is type of DocHitInfoIteratorFilter.
+  EXPECT_THAT(new_root.get(), Ne(original_root_ptr));
+  ASSERT_THAT(new_root.get(),
+              WhenDynamicCastTo<DocHitInfoIteratorFilter*>(Ne(nullptr)));
+  // Check that the original iterator is now the delegate.
+  EXPECT_THAT(new_root->GetChildren(),
+              ElementsAre(Pointee(Pointer(original_root_ptr))));
+  // Check that the predicate is not applied to the original iterator.
+  EXPECT_THAT(iterator_handling_filter_ptr->document_filter_predicates(),
+              IsEmpty());
+}
+
+TEST_F(DocHitInfoIteratorFilterTest, ApplyFilterDeepTreeMixedHandling) {
+  // Create an iterator tree:
+  //                        AND (root)
+  //                         /    \
+  //       DummyHandlingFilter    AND
+  //                             /   \
+  //                           Dummy  DummyHandlingFilter
+  std::unique_ptr<DocHitInfoIteratorDummyHandlingFilter> child1 =
+      std::make_unique<DocHitInfoIteratorDummyHandlingFilter>();
+  std::unique_ptr<DocHitInfoIteratorDummy> grandchild1 =
+      std::make_unique<DocHitInfoIteratorDummy>();
+  std::unique_ptr<DocHitInfoIteratorDummyHandlingFilter> grandchild2 =
+      std::make_unique<DocHitInfoIteratorDummyHandlingFilter>();
+
+  DocHitInfoIteratorDummyHandlingFilter* child1_ptr = child1.get();
+  DocHitInfoIteratorDummyHandlingFilter* grandchild2_ptr = grandchild2.get();
+
+  std::unique_ptr<DocHitInfoIterator> child2 =
+      std::make_unique<DocHitInfoIteratorAnd>(std::move(grandchild1),
+                                              std::move(grandchild2));
+  std::unique_ptr<DocHitInfoIterator> original_root =
+      std::make_unique<DocHitInfoIteratorAnd>(std::move(child1),
+                                              std::move(child2));
+  DocHitInfoIterator* original_root_ptr = original_root.get();
+
+  // Create a predicate.
+  SearchSpecProto search_spec;
+  search_spec.add_schema_type_filters(schema1_);
+  std::unique_ptr<DocumentFilterPredicate> predicate =
+      GetFilterPredicateBySchemaAndNamespace(
+          search_spec, *document_store_, *schema_store_,
+          fake_clock_.GetSystemTimeMilliseconds());
+
+  // Apply the filter. Since grandchild1 (Dummy) cannot handle the filter, a
+  // new filter iterator should be created at the top. However, child1 and
+  // grandchild2 should still get the predicate.
+  //
+  // We should get a new tree:
+  //                      Filter (root)
+  //                           |
+  //                          AND
+  //                         /    \
+  //       DummyHandlingFilter    AND
+  //                             /   \
+  //                           Dummy  DummyHandlingFilter
+  std::unique_ptr<DocHitInfoIterator> new_root =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_root), predicate.get(),
+          feature_flags_->enable_passing_filter_to_children());
+
+  // The filter should be applied at the top level.
+  EXPECT_THAT(new_root.get(), Ne(original_root_ptr));
+  ASSERT_THAT(new_root.get(),
+              WhenDynamicCastTo<DocHitInfoIteratorFilter*>(Ne(nullptr)));
+  // Check that the original iterator is now the delegate.
+  EXPECT_THAT(new_root->GetChildren(),
+              ElementsAre(Pointee(Pointer(original_root_ptr))));
+  // Verify the predicate was passed down to the children that can handle it.
+  EXPECT_THAT(child1_ptr->document_filter_predicates(),
+              ElementsAre(predicate.get()));
+  EXPECT_THAT(grandchild2_ptr->document_filter_predicates(),
+              ElementsAre(predicate.get()));
+}
+
+TEST_F(DocHitInfoIteratorFilterTest, ApplyFilterNestedNonPassingIterator) {
+  // Create an iterator tree:
+  //                        AND (root)
+  //                         /    \
+  //       DummyHandlingFilter    NOT
+  //                               |
+  //                       DummyHandlingFilter
+  std::unique_ptr<DocHitInfoIteratorDummyHandlingFilter> child1 =
+      std::make_unique<DocHitInfoIteratorDummyHandlingFilter>();
+  std::unique_ptr<DocHitInfoIteratorDummyHandlingFilter> grandchild1 =
+      std::make_unique<DocHitInfoIteratorDummyHandlingFilter>();
+
+  DocHitInfoIteratorDummyHandlingFilter* child1_ptr = child1.get();
+  DocHitInfoIteratorDummyHandlingFilter* grandchild1_ptr = grandchild1.get();
+
+  std::unique_ptr<DocHitInfoIterator> child2 =
+      std::make_unique<DocHitInfoIteratorNot>(std::move(grandchild1),
+                                              /*document_id_limit=*/100);
+  std::unique_ptr<DocHitInfoIterator> original_root =
+      std::make_unique<DocHitInfoIteratorAnd>(std::move(child1),
+                                              std::move(child2));
+  DocHitInfoIterator* original_root_ptr = original_root.get();
+
+  // Create a predicate.
+  SearchSpecProto search_spec;
+  search_spec.add_namespace_filters(namespace1_);
+  std::unique_ptr<DocumentFilterPredicate> predicate =
+      GetFilterPredicateBySchemaAndNamespace(
+          search_spec, *document_store_, *schema_store_,
+          fake_clock_.GetSystemTimeMilliseconds());
+
+  // Apply the filter. Since child2 (NOT) cannot pass the filter through, a
+  // new filter iterator should be created at the top. child1 should still get
+  // the predicate, but grandchild1 (under NOT) should not.
+  //
+  // We should get a new tree:
+  //                      Filter (root)
+  //                           |
+  //                          AND
+  //                         /    \
+  //       DummyHandlingFilter    NOT
+  //                               |
+  //                       DummyHandlingFilter
+  std::unique_ptr<DocHitInfoIterator> new_root =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_root), predicate.get(),
+          feature_flags_->enable_passing_filter_to_children());
+
+  // The filter should be applied at the top level.
+  EXPECT_THAT(new_root.get(), Ne(original_root_ptr));
+  ASSERT_THAT(new_root.get(),
+              WhenDynamicCastTo<DocHitInfoIteratorFilter*>(Ne(nullptr)));
+  // Check that the original iterator is now the delegate.
+  EXPECT_THAT(new_root->GetChildren(),
+              ElementsAre(Pointee(Pointer(original_root_ptr))));
+  // Verify the predicate was passed down only to child1.
+  EXPECT_THAT(child1_ptr->document_filter_predicates(),
+              ElementsAre(predicate.get()));
+  // Verify the predicate was *not* passed down to grandchild1.
+  EXPECT_THAT(grandchild1_ptr->document_filter_predicates(), IsEmpty());
+}
+
+TEST_F(DocHitInfoIteratorFilterTest, ApplyFilterComplexTreeAllHandling) {
+  // Create leaf nodes that can handle filters.
+  auto leaf1 = std::make_unique<DocHitInfoIteratorDummyHandlingFilter>();
+  auto leaf2 = std::make_unique<DocHitInfoIteratorDummyHandlingFilter>();
+  auto leaf3 = std::make_unique<DocHitInfoIteratorDummyHandlingFilter>();
+  auto leaf4 = std::make_unique<DocHitInfoIteratorDummyHandlingFilter>();
+  DocHitInfoIteratorDummyHandlingFilter* leaf1_ptr = leaf1.get();
+  DocHitInfoIteratorDummyHandlingFilter* leaf2_ptr = leaf2.get();
+  DocHitInfoIteratorDummyHandlingFilter* leaf3_ptr = leaf3.get();
+  DocHitInfoIteratorDummyHandlingFilter* leaf4_ptr = leaf4.get();
+
+  // Build a complex tree: AND(OR(leaf1, leaf2), AND(leaf3, leaf4))
+  auto or_node = std::make_unique<DocHitInfoIteratorOr>(std::move(leaf1),
+                                                        std::move(leaf2));
+
+  auto and_node2 = std::make_unique<DocHitInfoIteratorAnd>(std::move(leaf3),
+                                                           std::move(leaf4));
+
+  std::unique_ptr<DocHitInfoIterator> original_root =
+      std::make_unique<DocHitInfoIteratorAnd>(std::move(or_node),
+                                              std::move(and_node2));
+  DocHitInfoIterator* original_root_ptr = original_root.get();
+
+  // Create a predicate.
+  SearchSpecProto search_spec;
+  search_spec.add_namespace_filters(namespace1_);
+  search_spec.add_schema_type_filters(schema1_);
+  std::unique_ptr<DocumentFilterPredicate> predicate =
+      GetFilterPredicateBySchemaAndNamespace(
+          search_spec, *document_store_, *schema_store_,
+          fake_clock_.GetSystemTimeMilliseconds());
+
+  // Apply the filter.
+  std::unique_ptr<DocHitInfoIterator> new_root =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_root), predicate.get(),
+          feature_flags_->enable_passing_filter_to_children());
+
+  // Since all leaf nodes can handle the filter, the original root should be
+  // returned without a Filter wrapper.
+  EXPECT_THAT(new_root.get(), Eq(original_root_ptr));
+  EXPECT_THAT(new_root.get(),
+              WhenDynamicCastTo<DocHitInfoIteratorFilter*>(nullptr));
+
+  // Verify the predicate was passed down to all leaf nodes.
+  EXPECT_THAT(leaf1_ptr->document_filter_predicates(),
+              ElementsAre(predicate.get()));
+  EXPECT_THAT(leaf2_ptr->document_filter_predicates(),
+              ElementsAre(predicate.get()));
+  EXPECT_THAT(leaf3_ptr->document_filter_predicates(),
+              ElementsAre(predicate.get()));
+  EXPECT_THAT(leaf4_ptr->document_filter_predicates(),
+              ElementsAre(predicate.get()));
+}
+
+TEST_F(DocHitInfoIteratorFilterTest, ApplyMultipleFiltersSequentially) {
+  // Create an iterator tree:
+  //                        AND (root)
+  //                         /    \
+  //       DummyHandlingFilter    Dummy
+  std::unique_ptr<DocHitInfoIteratorDummyHandlingFilter> child1 =
+      std::make_unique<DocHitInfoIteratorDummyHandlingFilter>();
+  std::unique_ptr<DocHitInfoIteratorDummy> child2 =
+      std::make_unique<DocHitInfoIteratorDummy>();
+  DocHitInfoIteratorDummyHandlingFilter* child1_ptr = child1.get();
+  std::unique_ptr<DocHitInfoIterator> original_root =
+      std::make_unique<DocHitInfoIteratorAnd>(std::move(child1),
+                                              std::move(child2));
+
+  // Create predicates.
+  SearchSpecProto search_spec_1;
+  search_spec_1.add_schema_type_filters(schema1_);
+  std::unique_ptr<DocumentFilterPredicate> predicate_1 =
+      GetFilterPredicateBySchemaAndNamespace(
+          search_spec_1, *document_store_, *schema_store_,
+          fake_clock_.GetSystemTimeMilliseconds());
+  SearchSpecProto search_spec_2;
+  search_spec_2.add_namespace_filters(namespace1_);
+  std::unique_ptr<DocumentFilterPredicate> predicate_2 =
+      GetFilterPredicateBySchemaAndNamespace(
+          search_spec_2, *document_store_, *schema_store_,
+          fake_clock_.GetSystemTimeMilliseconds());
+
+  // Apply the first filter. Since child2 cannot handle it, a Filter iterator
+  // is added. child1 gets predicate_1.
+  // Tree: Filter1(AND(DummyHandlingFilter, Dummy))
+  std::unique_ptr<DocHitInfoIterator> current_root =
+      DocHitInfoIteratorFilter::ApplyFilter(
+          std::move(original_root), predicate_1.get(),
+          feature_flags_->enable_passing_filter_to_children());
+  DocHitInfoIterator* filter1_ptr = current_root.get();
+  ASSERT_THAT(filter1_ptr,
+              WhenDynamicCastTo<DocHitInfoIteratorFilter*>(Ne(nullptr)));
+  EXPECT_THAT(child1_ptr->document_filter_predicates(),
+              ElementsAre(predicate_1.get()));
+
+  // Apply the second filter. A new Filter2 iterator is added on top. child1
+  // gets predicate_2. Tree: Filter2(Filter1(AND(DummyHandlingFilter, Dummy)))
+  current_root = DocHitInfoIteratorFilter::ApplyFilter(
+      std::move(current_root), predicate_2.get(),
+      feature_flags_->enable_passing_filter_to_children());
+  DocHitInfoIterator* filter2_ptr = current_root.get();
+  ASSERT_THAT(filter2_ptr,
+              WhenDynamicCastTo<DocHitInfoIteratorFilter*>(Ne(nullptr)));
+  EXPECT_THAT(filter2_ptr, Ne(filter1_ptr));  // New filter added
+
+  // Check delegate of Filter2 is Filter1
+  ASSERT_THAT(current_root->GetChildren(),
+              ElementsAre(Pointee(Pointer(filter1_ptr))));
+
+  // Verify child1 now has both predicates.
+  EXPECT_THAT(child1_ptr->document_filter_predicates(),
+              ElementsAre(predicate_1.get(), predicate_2.get()));
+}
+
 }  // namespace
 
 }  // namespace lib

diff --git a/icing/index/iterator/doc-hit-info-iterator-match-score-expression.h b/icing/index/iterator/doc-hit-info-iterator-match-score-expression.h
index f46b8da..9c1ece6 100644
--- a/icing/index/iterator/doc-hit-info-iterator-match-score-expression.h
+++ b/icing/index/iterator/doc-hit-info-iterator-match-score-expression.h

@@ -33,7 +33,8 @@
 namespace icing {
 namespace lib {
 
-class DocHitInfoIteratorMatchScoreExpression : public DocHitInfoIterator {
+class DocHitInfoIteratorMatchScoreExpression
+    : public DocHitInfoIteratorSectionRestrictionNotApplicable {
  public:
   explicit DocHitInfoIteratorMatchScoreExpression(
       DocumentId last_added_document_id,
@@ -56,8 +57,8 @@
         "supported");
   }
 
-  void MapChildren(const ChildrenMapper& mapper) override {
-    delegate_ = mapper(std::move(delegate_));
+  std::vector<std::unique_ptr<DocHitInfoIterator>*> GetChildren() override {
+    return {&delegate_};
   }
 
   CallStats GetCallStats() const override { return delegate_->GetCallStats(); }

diff --git a/icing/index/iterator/doc-hit-info-iterator-none.h b/icing/index/iterator/doc-hit-info-iterator-none.h
index c2853f1..4d8bb05 100644
--- a/icing/index/iterator/doc-hit-info-iterator-none.h
+++ b/icing/index/iterator/doc-hit-info-iterator-none.h

@@ -15,10 +15,12 @@
 #ifndef ICING_INDEX_ITERATOR_DOC_HIT_INFO_ITERATOR_NONE_H_
 #define ICING_INDEX_ITERATOR_DOC_HIT_INFO_ITERATOR_NONE_H_
 
-#include <cstdint>
+#include <memory>
 #include <string>
+#include <vector>
 
 #include "icing/text_classifier/lib3/utils/base/status.h"
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "icing/absl_ports/canonical_errors.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
 
@@ -26,7 +28,8 @@
 namespace lib {
 
 // Iterator that will return no results.
-class DocHitInfoIteratorNone : public DocHitInfoIterator {
+class DocHitInfoIteratorNone
+    : public DocHitInfoIteratorSectionRestrictionNotApplicable {
  public:
   libtextclassifier3::Status Advance() override {
     return absl_ports::ResourceExhaustedError(
@@ -39,7 +42,9 @@
     return node;
   }
 
-  void MapChildren(const ChildrenMapper& mapper) override {}
+  std::vector<std::unique_ptr<DocHitInfoIterator>*> GetChildren() override {
+    return {};
+  }
 
   CallStats GetCallStats() const override { return CallStats(); }
 

diff --git a/icing/index/iterator/doc-hit-info-iterator-not.cc b/icing/index/iterator/doc-hit-info-iterator-not.cc
index 10a8292..706a5bd 100644
--- a/icing/index/iterator/doc-hit-info-iterator-not.cc
+++ b/icing/index/iterator/doc-hit-info-iterator-not.cc

@@ -14,11 +14,13 @@
 
 #include "icing/index/iterator/doc-hit-info-iterator-not.h"
 
-#include <cstdint>
 #include <memory>
+#include <string>
 #include <utility>
+#include <vector>
 
 #include "icing/text_classifier/lib3/utils/base/status.h"
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "icing/absl_ports/canonical_errors.h"
 #include "icing/absl_ports/str_cat.h"
 #include "icing/index/hit/doc-hit-info.h"
@@ -69,8 +71,9 @@
       "Cannot generate suggestion if the last term is NOT operator.");
 }
 
-void DocHitInfoIteratorNot::MapChildren(const ChildrenMapper& mapper) {
-  to_be_excluded_ = mapper(std::move(to_be_excluded_));
+std::vector<std::unique_ptr<DocHitInfoIterator>*>
+DocHitInfoIteratorNot::GetChildren() {
+  return {&to_be_excluded_};
 }
 
 std::string DocHitInfoIteratorNot::ToString() const {

diff --git a/icing/index/iterator/doc-hit-info-iterator-not.h b/icing/index/iterator/doc-hit-info-iterator-not.h
index 11575fb..2b02b37 100644
--- a/icing/index/iterator/doc-hit-info-iterator-not.h
+++ b/icing/index/iterator/doc-hit-info-iterator-not.h

@@ -15,11 +15,12 @@
 #ifndef ICING_INDEX_ITERATOR_DOC_HIT_INFO_ITERATOR_NOT_H_
 #define ICING_INDEX_ITERATOR_DOC_HIT_INFO_ITERATOR_NOT_H_
 
-#include <cstdint>
 #include <memory>
 #include <string>
+#include <vector>
 
 #include "icing/text_classifier/lib3/utils/base/status.h"
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "icing/index/iterator/doc-hit-info-iterator-all-document-id.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
 #include "icing/store/document-id.h"
@@ -36,7 +37,8 @@
 // having been chosen because it's term was in a specific section. Since we
 // don't know anything about the sections for the Document, the
 // doc_hit_info.hit_section_ids_mask() is always kSectionIdMaskNone.
-class DocHitInfoIteratorNot : public DocHitInfoIterator {
+class DocHitInfoIteratorNot
+    : public DocHitInfoIteratorSectionRestrictionApplyToChildren {
  public:
   // to_be_excluded_iterator: The results of this iterator will be excluded
   //     from this iterator's results.
@@ -44,7 +46,7 @@
   //     Document to the DocumentStore
   explicit DocHitInfoIteratorNot(
       std::unique_ptr<DocHitInfoIterator> to_be_excluded_iterator,
-      const DocumentId document_id_limit);
+      DocumentId document_id_limit);
 
   libtextclassifier3::Status Advance() override;
 
@@ -53,7 +55,9 @@
   // to NOT operator.
   libtextclassifier3::StatusOr<TrimmedNode> TrimRightMostNode() && override;
 
-  void MapChildren(const ChildrenMapper& mapper) override;
+  std::vector<std::unique_ptr<DocHitInfoIterator>*> GetChildren() override;
+
+  bool CanPassFilterPredicateThrough() const override { return false; }
 
   CallStats GetCallStats() const override {
     return to_be_excluded_->GetCallStats() +

diff --git a/icing/index/iterator/doc-hit-info-iterator-not_test.cc b/icing/index/iterator/doc-hit-info-iterator-not_test.cc
index a8c835f..6b02cd3 100644
--- a/icing/index/iterator/doc-hit-info-iterator-not_test.cc
+++ b/icing/index/iterator/doc-hit-info-iterator-not_test.cc

@@ -108,7 +108,12 @@
       /*num_leaf_advance_calls_main_index_in=*/5,
       /*num_leaf_advance_calls_integer_index_in=*/3,
       /*num_leaf_advance_calls_no_index_in=*/1,
-      /*num_blocks_inspected_in=*/4);  // arbitrary value
+      /*num_blocks_inspected_in=*/4,
+      /*embedding_stats_in=*/
+      {.num_unquantized_embeddings_scored = 2,
+       .num_quantized_embeddings_scored = 3,
+       .unquantized_shards_read = {1, 2},
+       .quantized_shards_read{3, 4}});  // arbitrary value
   auto to_be_excluded_iterator = std::make_unique<DocHitInfoIteratorDummy>();
   to_be_excluded_iterator->SetCallStats(to_be_excluded_iterator_call_stats);
 
@@ -134,7 +139,8 @@
               .num_leaf_advance_calls_integer_index,
           to_be_excluded_iterator_call_stats.num_leaf_advance_calls_no_index +
               all_leaf_advance_calls,
-          to_be_excluded_iterator_call_stats.num_blocks_inspected));
+          to_be_excluded_iterator_call_stats.num_blocks_inspected,
+          to_be_excluded_iterator_call_stats.embedding_stats));
 }
 
 TEST(DocHitInfoIteratorNotTest, SectionIdsAlwaysNone) {

diff --git a/icing/index/iterator/doc-hit-info-iterator-or.h b/icing/index/iterator/doc-hit-info-iterator-or.h
index 8c0427b..c2da3eb 100644
--- a/icing/index/iterator/doc-hit-info-iterator-or.h
+++ b/icing/index/iterator/doc-hit-info-iterator-or.h

@@ -15,12 +15,16 @@
 #ifndef ICING_INDEX_ITERATOR_DOC_HIT_INFO_ITERATOR_OR_H_
 #define ICING_INDEX_ITERATOR_DOC_HIT_INFO_ITERATOR_OR_H_
 
-#include <cstdint>
+#include <cstddef>
 #include <memory>
 #include <string>
-#include <utility>
+#include <vector>
 
+#include "icing/text_classifier/lib3/utils/base/status.h"
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
+#include "icing/schema/section.h"
+#include "icing/store/document-id.h"
 
 namespace icing {
 namespace lib {
@@ -31,7 +35,8 @@
     std::vector<std::unique_ptr<DocHitInfoIterator>> iterators);
 
 // Iterate over a logical OR of two child iterators.
-class DocHitInfoIteratorOr : public DocHitInfoIterator {
+class DocHitInfoIteratorOr
+    : public DocHitInfoIteratorSectionRestrictionApplyToChildren {
  public:
   explicit DocHitInfoIteratorOr(std::unique_ptr<DocHitInfoIterator> left_it,
                                 std::unique_ptr<DocHitInfoIterator> right_it);
@@ -46,9 +51,8 @@
 
   std::string ToString() const override;
 
-  void MapChildren(const ChildrenMapper &mapper) override {
-    left_ = mapper(std::move(left_));
-    right_ = mapper(std::move(right_));
+  std::vector<std::unique_ptr<DocHitInfoIterator> *> GetChildren() override {
+    return {&left_, &right_};
   }
 
   void PopulateMatchedTermsStats(
@@ -81,7 +85,8 @@
 // Iterate over a logical OR of multiple child iterators.
 //
 // NOTE: DocHitInfoIteratorOr is a faster alternative to OR exactly 2 iterators.
-class DocHitInfoIteratorOrNary : public DocHitInfoIterator {
+class DocHitInfoIteratorOrNary
+    : public DocHitInfoIteratorSectionRestrictionApplyToChildren {
  public:
   explicit DocHitInfoIteratorOrNary(
       std::vector<std::unique_ptr<DocHitInfoIterator>> iterators);
@@ -94,10 +99,13 @@
 
   std::string ToString() const override;
 
-  void MapChildren(const ChildrenMapper &mapper) override {
+  std::vector<std::unique_ptr<DocHitInfoIterator> *> GetChildren() override {
+    std::vector<std::unique_ptr<DocHitInfoIterator> *> children;
+    children.reserve(iterators_.size());
     for (int i = 0; i < iterators_.size(); ++i) {
-      iterators_[i] = mapper(std::move(iterators_[i]));
+      children.push_back(&iterators_[i]);
     }
+    return children;
   }
 
   void PopulateMatchedTermsStats(

diff --git a/icing/index/iterator/doc-hit-info-iterator-or_test.cc b/icing/index/iterator/doc-hit-info-iterator-or_test.cc
index d198b53..2632b73 100644
--- a/icing/index/iterator/doc-hit-info-iterator-or_test.cc
+++ b/icing/index/iterator/doc-hit-info-iterator-or_test.cc

@@ -14,11 +14,16 @@
 
 #include "icing/index/iterator/doc-hit-info-iterator-or.h"
 
-#include <string>
+#include <memory>
+#include <unordered_map>
+#include <utility>
 #include <vector>
 
+#include "icing/text_classifier/lib3/utils/base/status.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
+#include "icing/index/hit/doc-hit-info.h"
+#include "icing/index/hit/hit.h"
 #include "icing/index/iterator/doc-hit-info-iterator-test-util.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
 #include "icing/schema/section.h"
@@ -33,6 +38,8 @@
 using ::testing::ElementsAre;
 using ::testing::Eq;
 using ::testing::IsEmpty;
+using ::testing::Pointee;
+using ::testing::Pointer;
 
 TEST(CreateAndIteratorTest, Or) {
   // Basic test that we can create a working Or iterator. Further testing of
@@ -83,7 +90,12 @@
       /*num_leaf_advance_calls_main_index_in=*/5,
       /*num_leaf_advance_calls_integer_index_in=*/3,
       /*num_leaf_advance_calls_no_index_in=*/1,
-      /*num_blocks_inspected_in=*/4);  // arbitrary value
+      /*num_blocks_inspected_in=*/4,
+      /*embedding_stats_in=*/
+      {.num_unquantized_embeddings_scored = 2,
+       .num_quantized_embeddings_scored = 3,
+       .unquantized_shards_read = {1, 2},
+       .quantized_shards_read{3, 4}});  // arbitrary value
   auto first_iter = std::make_unique<DocHitInfoIteratorDummy>();
   first_iter->SetCallStats(first_iter_call_stats);
 
@@ -92,7 +104,12 @@
       /*num_leaf_advance_calls_main_index_in=*/2,
       /*num_leaf_advance_calls_integer_index_in=*/10,
       /*num_leaf_advance_calls_no_index_in=*/3,
-      /*num_blocks_inspected_in=*/7);  // arbitrary value
+      /*num_blocks_inspected_in=*/7,
+      /*embedding_stats_in=*/
+      {.num_unquantized_embeddings_scored = 4,
+       .num_quantized_embeddings_scored = 5,
+       .unquantized_shards_read = {5, 6},
+       .quantized_shards_read{7}});  // arbitrary value
   auto second_iter = std::make_unique<DocHitInfoIteratorDummy>();
   second_iter->SetCallStats(second_iter_call_stats);
 
@@ -375,7 +392,12 @@
       /*num_leaf_advance_calls_main_index_in=*/5,
       /*num_leaf_advance_calls_integer_index_in=*/3,
       /*num_leaf_advance_calls_no_index_in=*/1,
-      /*num_blocks_inspected_in=*/4);  // arbitrary value
+      /*num_blocks_inspected_in=*/4,
+      /*embedding_stats_in=*/
+      {.num_unquantized_embeddings_scored = 2,
+       .num_quantized_embeddings_scored = 3,
+       .unquantized_shards_read = {1, 2},
+       .quantized_shards_read{3, 4}});  // arbitrary value
   auto first_iter = std::make_unique<DocHitInfoIteratorDummy>();
   first_iter->SetCallStats(first_iter_call_stats);
 
@@ -384,7 +406,12 @@
       /*num_leaf_advance_calls_main_index_in=*/2,
       /*num_leaf_advance_calls_integer_index_in=*/10,
       /*num_leaf_advance_calls_no_index_in=*/3,
-      /*num_blocks_inspected_in=*/7);  // arbitrary value
+      /*num_blocks_inspected_in=*/7,
+      /*embedding_stats_in=*/
+      {.num_unquantized_embeddings_scored = 4,
+       .num_quantized_embeddings_scored = 5,
+       .unquantized_shards_read = {5, 6},
+       .quantized_shards_read{7}});  // arbitrary value
   auto second_iter = std::make_unique<DocHitInfoIteratorDummy>();
   second_iter->SetCallStats(second_iter_call_stats);
 
@@ -393,7 +420,12 @@
       /*num_leaf_advance_calls_main_index_in=*/2000,
       /*num_leaf_advance_calls_integer_index_in=*/3000,
       /*num_leaf_advance_calls_no_index_in=*/0,
-      /*num_blocks_inspected_in=*/200);  // arbitrary value
+      /*num_blocks_inspected_in=*/200,
+      /*embedding_stats_in=*/
+      {.num_unquantized_embeddings_scored = 1,
+       .num_quantized_embeddings_scored = 1,
+       .unquantized_shards_read = {0},
+       .quantized_shards_read{0}});  // arbitrary value
   auto third_iter = std::make_unique<DocHitInfoIteratorDummy>();
   third_iter->SetCallStats(third_iter_call_stats);
 
@@ -402,7 +434,12 @@
       /*num_leaf_advance_calls_main_index_in=*/400,
       /*num_leaf_advance_calls_integer_index_in=*/100,
       /*num_leaf_advance_calls_no_index_in=*/20,
-      /*num_blocks_inspected_in=*/50);  // arbitrary value
+      /*num_blocks_inspected_in=*/50,
+      /*embedding_stats_in=*/
+      {.num_unquantized_embeddings_scored = 10,
+       .num_quantized_embeddings_scored = 10,
+       .unquantized_shards_read = {5, 6, 7},
+       .quantized_shards_read{9, 10, 11}});  // arbitrary value
   auto fourth_iter = std::make_unique<DocHitInfoIteratorDummy>();
   fourth_iter->SetCallStats(fourth_iter_call_stats);
 
@@ -578,6 +615,32 @@
   EXPECT_FALSE(or_iter.Advance().ok());
 }
 
+TEST(DocHitInfoIteratorOrNaryTest, GetChildren) {
+  std::vector<DocHitInfo> first_vector = {DocHitInfo(2), DocHitInfo(1),
+                                          DocHitInfo(0)};
+  std::vector<DocHitInfo> second_vector = {DocHitInfo(2), DocHitInfo(1)};
+  std::vector<DocHitInfo> third_vector = {DocHitInfo(2)};
+
+  std::vector<std::unique_ptr<DocHitInfoIterator>> iterators;
+  iterators.push_back(std::make_unique<DocHitInfoIteratorDummy>(first_vector));
+  iterators.push_back(std::make_unique<DocHitInfoIteratorDummy>(second_vector));
+  iterators.push_back(
+      std::make_unique<DocHitInfoIteratorDummy>(third_vector, "term", 10));
+
+  std::vector<DocHitInfoIterator*> iterator_ptrs;
+  for (const auto& iter : iterators) {
+    iterator_ptrs.push_back(iter.get());
+  }
+
+  std::unique_ptr<DocHitInfoIterator> iter =
+      std::make_unique<DocHitInfoIteratorOrNary>(std::move(iterators));
+
+  EXPECT_THAT(iter->GetChildren(),
+              ElementsAre(Pointee(Pointer(iterator_ptrs[0])),
+                          Pointee(Pointer(iterator_ptrs[1])),
+                          Pointee(Pointer(iterator_ptrs[2]))));
+}
+
 }  // namespace
 
 }  // namespace lib

diff --git a/icing/index/iterator/doc-hit-info-iterator-property-in-document.h b/icing/index/iterator/doc-hit-info-iterator-property-in-document.h
index bb2c97a..5a5724e 100644
--- a/icing/index/iterator/doc-hit-info-iterator-property-in-document.h
+++ b/icing/index/iterator/doc-hit-info-iterator-property-in-document.h

@@ -15,10 +15,8 @@
 #ifndef ICING_INDEX_ITERATOR_DOC_HIT_INFO_ITERATOR_PROPERTY_IN_DOCUMENT_H_
 #define ICING_INDEX_ITERATOR_DOC_HIT_INFO_ITERATOR_PROPERTY_IN_DOCUMENT_H_
 
-#include <cstdint>
 #include <memory>
 #include <string>
-#include <utility>
 #include <vector>
 
 #include "icing/text_classifier/lib3/utils/base/status.h"
@@ -34,10 +32,8 @@
 // post-processes metadata hits added by PropertyExistenceIndexingHandler.
 // Specifically, it filters out hits that are not recognized as metadata, and
 // always set hit_section_ids_mask to 0.
-//
-// It is marked as a subclass of DocHitInfoLeafIterator because section
-// restriction should not be passed down to meta_hit_iterator.
-class DocHitInfoIteratorPropertyInDocument : public DocHitInfoLeafIterator {
+class DocHitInfoIteratorPropertyInDocument
+    : public DocHitInfoIteratorSectionRestrictionNotApplicable {
  public:
   explicit DocHitInfoIteratorPropertyInDocument(
       std::unique_ptr<DocHitInfoIterator> meta_hit_iterator);
@@ -46,6 +42,10 @@
 
   libtextclassifier3::StatusOr<TrimmedNode> TrimRightMostNode() && override;
 
+  std::vector<std::unique_ptr<DocHitInfoIterator>*> GetChildren() override {
+    return {&meta_hit_iterator_};
+  }
+
   CallStats GetCallStats() const override {
     return meta_hit_iterator_->GetCallStats();
   }

diff --git a/icing/index/iterator/doc-hit-info-iterator-property-in-schema.h b/icing/index/iterator/doc-hit-info-iterator-property-in-schema.h
index d766712..5450594 100644
--- a/icing/index/iterator/doc-hit-info-iterator-property-in-schema.h
+++ b/icing/index/iterator/doc-hit-info-iterator-property-in-schema.h

@@ -19,7 +19,6 @@
 #include <memory>
 #include <set>
 #include <string>
-#include <utility>
 #include <vector>
 
 #include "icing/text_classifier/lib3/utils/base/status.h"
@@ -35,7 +34,8 @@
 
 // An iterator that helps filter for DocHitInfos whose schemas define the
 // properties named in target_properties_.
-class DocHitInfoIteratorPropertyInSchema : public DocHitInfoIterator {
+class DocHitInfoIteratorPropertyInSchema
+    : public DocHitInfoIteratorSectionRestrictionNotApplicable {
  public:
   // Does not take any ownership, and all pointers must refer to valid objects
   // that outlive the one constructed. The delegate should be at minimum be
@@ -50,8 +50,8 @@
 
   libtextclassifier3::StatusOr<TrimmedNode> TrimRightMostNode() && override;
 
-  void MapChildren(const ChildrenMapper& mapper) override {
-    delegate_ = mapper(std::move(delegate_));
+  std::vector<std::unique_ptr<DocHitInfoIterator>*> GetChildren() override {
+    return {&delegate_};
   }
 
   CallStats GetCallStats() const override { return delegate_->GetCallStats(); }

diff --git a/icing/index/iterator/doc-hit-info-iterator-property-in-schema_test.cc b/icing/index/iterator/doc-hit-info-iterator-property-in-schema_test.cc
index 3c29c1e..36c68af 100644
--- a/icing/index/iterator/doc-hit-info-iterator-property-in-schema_test.cc
+++ b/icing/index/iterator/doc-hit-info-iterator-property-in-schema_test.cc

@@ -19,6 +19,7 @@
 #include <utility>
 #include <vector>
 
+#include "icing/text_classifier/lib3/utils/base/status.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include "icing/document-builder.h"
@@ -29,6 +30,7 @@
 #include "icing/index/iterator/doc-hit-info-iterator-all-document-id.h"
 #include "icing/index/iterator/doc-hit-info-iterator-test-util.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/schema.pb.h"
 #include "icing/schema-builder.h"
@@ -40,6 +42,7 @@
 #include "icing/testing/fake-clock.h"
 #include "icing/testing/test-feature-flags.h"
 #include "icing/testing/tmp-directory.h"
+#include "icing/util/document-util.h"
 
 namespace icing {
 namespace lib {
@@ -100,14 +103,18 @@
 
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
-        DocumentStore::Create(&filesystem_, test_dir_, &fake_clock_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, test_dir_, &fake_clock_, schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     document_store_ = std::move(create_result.document_store);
   }
 
@@ -134,8 +141,9 @@
 TEST_F(DocHitInfoIteratorPropertyInSchemaTest,
        AdvanceToDocumentWithIndexedProperty) {
   // Populate the DocumentStore's FilterCache with this document's data
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1_)));
   DocumentId document_id = put_result.new_document_id;
 
   auto original_iterator = std::make_unique<DocHitInfoIteratorAllDocumentId>(
@@ -155,8 +163,9 @@
 TEST_F(DocHitInfoIteratorPropertyInSchemaTest,
        AdvanceToDocumentWithUnindexedProperty) {
   // Populate the DocumentStore's FilterCache with this document's data
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1_)));
   DocumentId document_id = put_result.new_document_id;
 
   auto original_iterator = std::make_unique<DocHitInfoIteratorAllDocumentId>(
@@ -174,7 +183,8 @@
 }
 
 TEST_F(DocHitInfoIteratorPropertyInSchemaTest, NoMatchWithUndefinedProperty) {
-  ICING_EXPECT_OK(document_store_->Put(document1_));
+  ICING_EXPECT_OK(
+      document_store_->Put(document_util::CreateDocumentWrapper(document1_)));
 
   auto original_iterator = std::make_unique<DocHitInfoIteratorAllDocumentId>(
       document_store_->num_documents());
@@ -189,8 +199,9 @@
 TEST_F(DocHitInfoIteratorPropertyInSchemaTest,
        CorrectlySetsSectionIdMasksAndPopulatesTermMatchInfo) {
   // Populate the DocumentStore's FilterCache with this document's data
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1_)));
   DocumentId document_id = put_result.new_document_id;
 
   // Arbitrary section ids for the documents in the DocHitInfoIterators.
@@ -252,11 +263,13 @@
 
 TEST_F(DocHitInfoIteratorPropertyInSchemaTest,
        FindPropertyDefinedByMultipleTypes) {
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1_)));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document2_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document2_)));
   DocumentId document_id2 = put_result2.new_document_id;
   auto original_iterator = std::make_unique<DocHitInfoIteratorAllDocumentId>(
       document_store_->num_documents());

diff --git a/icing/index/iterator/doc-hit-info-iterator-section-restrict.cc b/icing/index/iterator/doc-hit-info-iterator-section-restrict.cc
index f6e440d..7726736 100644
--- a/icing/index/iterator/doc-hit-info-iterator-section-restrict.cc
+++ b/icing/index/iterator/doc-hit-info-iterator-section-restrict.cc

@@ -21,7 +21,6 @@
 #include <string_view>
 #include <unordered_map>
 #include <utility>
-#include <vector>
 
 #include "icing/text_classifier/lib3/utils/base/status.h"
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
@@ -29,12 +28,12 @@
 #include "icing/absl_ports/str_cat.h"
 #include "icing/absl_ports/str_join.h"
 #include "icing/index/hit/doc-hit-info.h"
+#include "icing/index/iterator/doc-hit-info-iterator-data-holder.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
 #include "icing/index/iterator/section-restrict-data.h"
 #include "icing/proto/search.pb.h"
 #include "icing/schema/schema-store.h"
 #include "icing/schema/section.h"
-#include "icing/store/document-filter-data.h"
 #include "icing/store/document-id.h"
 #include "icing/store/document-store.h"
 #include "icing/util/status-macros.h"
@@ -42,51 +41,6 @@
 namespace icing {
 namespace lib {
 
-// An iterator that simply takes ownership of SectionRestrictData.
-class SectionRestrictDataHolderIterator : public DocHitInfoIterator {
- public:
-  explicit SectionRestrictDataHolderIterator(
-      std::unique_ptr<DocHitInfoIterator> delegate,
-      std::unique_ptr<SectionRestrictData> data)
-      : delegate_(std::move(delegate)), data_(std::move(data)) {}
-
-  libtextclassifier3::Status Advance() override {
-    auto result = delegate_->Advance();
-    doc_hit_info_ = delegate_->doc_hit_info();
-    return result;
-  }
-
-  libtextclassifier3::StatusOr<TrimmedNode> TrimRightMostNode() && override {
-    ICING_ASSIGN_OR_RETURN(TrimmedNode trimmed_delegate,
-                           std::move(*delegate_).TrimRightMostNode());
-    if (trimmed_delegate.iterator_ != nullptr) {
-      trimmed_delegate.iterator_ =
-          std::make_unique<SectionRestrictDataHolderIterator>(
-              std::move(trimmed_delegate.iterator_), std::move(data_));
-    }
-    return trimmed_delegate;
-  }
-
-  void MapChildren(const ChildrenMapper& mapper) override {
-    delegate_ = mapper(std::move(delegate_));
-  }
-
-  CallStats GetCallStats() const override { return delegate_->GetCallStats(); }
-
-  std::string ToString() const override { return delegate_->ToString(); }
-
-  void PopulateMatchedTermsStats(
-      std::vector<TermMatchInfo>* matched_terms_stats,
-      SectionIdMask filtering_section_mask) const override {
-    return delegate_->PopulateMatchedTermsStats(matched_terms_stats,
-                                                filtering_section_mask);
-  }
-
- private:
-  std::unique_ptr<DocHitInfoIterator> delegate_;
-  std::unique_ptr<SectionRestrictData> data_;
-};
-
 DocHitInfoIteratorSectionRestrict::DocHitInfoIteratorSectionRestrict(
     std::unique_ptr<DocHitInfoIterator> delegate, SectionRestrictData* data)
     : delegate_(std::move(delegate)), data_(data) {}
@@ -103,8 +57,8 @@
       document_store, schema_store, current_time_ms, type_property_filters);
   std::unique_ptr<DocHitInfoIterator> result =
       ApplyRestrictions(std::move(iterator), data.get());
-  return std::make_unique<SectionRestrictDataHolderIterator>(std::move(result),
-                                                             std::move(data));
+  return std::make_unique<DocHitInfoIteratorDataHolder<SectionRestrictData>>(
+      std::move(result), std::move(data));
 }
 
 std::unique_ptr<DocHitInfoIterator>
@@ -124,27 +78,37 @@
       document_store, schema_store, current_time_ms, type_property_filters);
   std::unique_ptr<DocHitInfoIterator> result =
       ApplyRestrictions(std::move(iterator), data.get());
-  return std::make_unique<SectionRestrictDataHolderIterator>(std::move(result),
-                                                             std::move(data));
+  return std::make_unique<DocHitInfoIteratorDataHolder<SectionRestrictData>>(
+      std::move(result), std::move(data));
 }
 
 std::unique_ptr<DocHitInfoIterator>
 DocHitInfoIteratorSectionRestrict::ApplyRestrictions(
     std::unique_ptr<DocHitInfoIterator> iterator, SectionRestrictData* data) {
-  ChildrenMapper mapper;
-  mapper = [&data, &mapper](std::unique_ptr<DocHitInfoIterator> iterator)
-      -> std::unique_ptr<DocHitInfoIterator> {
-    if (iterator->HandleSectionRestriction(data)) {
-      return iterator;
-    } else if (iterator->is_leaf()) {
-      return std::make_unique<DocHitInfoIteratorSectionRestrict>(
-          std::move(iterator), data);
-    } else {
-      iterator->MapChildren(mapper);
-      return iterator;
-    }
-  };
-  return mapper(std::move(iterator));
+  // If the iterator does not respect section restrictions, just return it.
+  if (iterator->SectionRestrictionNotApplicable()) {
+    return iterator;
+  }
+
+  // If the iterator can internally handle the section restriction, apply it and
+  // return the iterator.
+  if (iterator->HandleSectionRestriction(data)) {
+    return iterator;
+  }
+
+  // If the iterator accepts section restriction, but does not want to pass it
+  // down to its children, return a new iterator with the section restriction
+  // applied at the top.
+  if (!iterator->SectionRestrictionShouldApplyToChildren()) {
+    return std::make_unique<DocHitInfoIteratorSectionRestrict>(
+        std::move(iterator), data);
+  }
+
+  // Otherwise, apply the section restriction to its children.
+  for (std::unique_ptr<DocHitInfoIterator>* child : iterator->GetChildren()) {
+    *child = ApplyRestrictions(std::move(*child), data);
+  }
+  return iterator;
 }
 
 libtextclassifier3::Status DocHitInfoIteratorSectionRestrict::Advance() {

diff --git a/icing/index/iterator/doc-hit-info-iterator-section-restrict.h b/icing/index/iterator/doc-hit-info-iterator-section-restrict.h
index 387ff52..d67c95a 100644
--- a/icing/index/iterator/doc-hit-info-iterator-section-restrict.h
+++ b/icing/index/iterator/doc-hit-info-iterator-section-restrict.h

@@ -41,18 +41,18 @@
 // That class is meant to be applied to the root of a query tree and filter over
 // all results at the end. This class is more used in the limited scope of a
 // term or a small group of terms.
-class DocHitInfoIteratorSectionRestrict : public DocHitInfoLeafIterator {
+class DocHitInfoIteratorSectionRestrict : public DocHitInfoIterator {
  public:
   // Does not take any ownership, and all pointers must refer to valid objects
   // that outlive the one constructed.
   explicit DocHitInfoIteratorSectionRestrict(
       std::unique_ptr<DocHitInfoIterator> delegate, SectionRestrictData* data);
 
-  // Methods that apply section restrictions to all DocHitInfoLeafIterator nodes
-  // inside the provided iterator tree, and return the root of the tree
-  // afterwards. These methods do not take any ownership for the raw pointer
-  // parameters, which must refer to valid objects that outlive the iterator
-  // returned.
+  // Methods that apply section restrictions to all applicable nodes at the
+  // lowest level inside the provided iterator tree, and return the root of the
+  // tree afterwards. These methods do not take any ownership for the raw
+  // pointer parameters, which must refer to valid objects that outlive the
+  // iterator returned.
   static std::unique_ptr<DocHitInfoIterator> ApplyRestrictions(
       std::unique_ptr<DocHitInfoIterator> iterator,
       const DocumentStore* document_store, const SchemaStore* schema_store,
@@ -72,9 +72,9 @@
 
   std::string ToString() const override;
 
-  // Note that the DocHitInfoIteratorSectionRestrict can only be applied at
-  // DocHitInfoLeafIterator, which can be a term iterator or another
-  // DocHitInfoIteratorSectionRestrict.
+  // Note that the DocHitInfoIteratorSectionRestrict will eventually be applied
+  // at the lowest level of the iterator tree, which can be a term iterator or
+  // another DocHitInfoIteratorSectionRestrict.
   //
   // To filter the matching sections, filtering_section_mask should be set to
   // doc_hit_info_.hit_section_ids_mask() held in the outermost
@@ -96,6 +96,10 @@
             doc_hit_info_.hit_section_ids_mask());
   }
 
+  std::vector<std::unique_ptr<DocHitInfoIterator>*> GetChildren() override {
+    return {&delegate_};
+  }
+
  private:
   std::unique_ptr<DocHitInfoIterator> delegate_;
   // Does not own.

diff --git a/icing/index/iterator/doc-hit-info-iterator-section-restrict_test.cc b/icing/index/iterator/doc-hit-info-iterator-section-restrict_test.cc
index a97f2a8..1c414c1 100644
--- a/icing/index/iterator/doc-hit-info-iterator-section-restrict_test.cc
+++ b/icing/index/iterator/doc-hit-info-iterator-section-restrict_test.cc

@@ -17,6 +17,7 @@
 #include <memory>
 #include <set>
 #include <string>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
@@ -27,9 +28,11 @@
 #include "icing/file/filesystem.h"
 #include "icing/file/portable-file-backed-proto-log.h"
 #include "icing/index/hit/doc-hit-info.h"
+#include "icing/index/hit/hit.h"
 #include "icing/index/iterator/doc-hit-info-iterator-and.h"
 #include "icing/index/iterator/doc-hit-info-iterator-test-util.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/schema.pb.h"
 #include "icing/proto/term.pb.h"
@@ -42,6 +45,7 @@
 #include "icing/testing/fake-clock.h"
 #include "icing/testing/test-feature-flags.h"
 #include "icing/testing/tmp-directory.h"
+#include "icing/util/document-util.h"
 
 namespace icing {
 namespace lib {
@@ -105,14 +109,18 @@
 
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
-        DocumentStore::Create(&filesystem_, test_dir_, &fake_clock_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, test_dir_, &fake_clock_, schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     document_store_ = std::move(create_result.document_store);
   }
 
@@ -139,8 +147,9 @@
 TEST_F(DocHitInfoIteratorSectionRestrictTest,
        PopulateMatchedTermsStats_IncludesHitWithMatchingSection) {
   // Populate the DocumentStore's FilterCache with this document's data
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1_)));
   DocumentId document_id = put_result.new_document_id;
 
   // Arbitrary section ids for the documents in the DocHitInfoIterators.
@@ -205,8 +214,9 @@
 
 TEST_F(DocHitInfoIteratorSectionRestrictTest, IncludesHitWithMatchingSection) {
   // Populate the DocumentStore's FilterCache with this document's data
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1_)));
   DocumentId document_id = put_result.new_document_id;
 
   SectionIdMask section_id_mask = 1U << kIndexedSectionId0;
@@ -233,8 +243,9 @@
 TEST_F(DocHitInfoIteratorSectionRestrictTest,
        IncludesHitWithMultipleMatchingSectionsWithMultipleSectionRestricts) {
   // Populate the DocumentStore's FilterCache with this document's data
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1_)));
   DocumentId document_id = put_result.new_document_id;
 
   SectionIdMask section_id_mask = 1U << kIndexedSectionId0;
@@ -265,8 +276,9 @@
 TEST_F(DocHitInfoIteratorSectionRestrictTest,
        IncludesHitWithMultipleMatchingSectionsWithSingleSectionRestrict) {
   // Populate the DocumentStore's FilterCache with this document's data
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1_)));
   DocumentId document_id = put_result.new_document_id;
 
   SectionIdMask section_id_mask = 1U << kIndexedSectionId0;
@@ -296,8 +308,9 @@
 TEST_F(DocHitInfoIteratorSectionRestrictTest,
        IncludesHitWithSingleMatchingSectionsWithMultiSectionRestrict) {
   // Populate the DocumentStore's FilterCache with this document's data
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1_)));
   DocumentId document_id = put_result.new_document_id;
 
   SectionIdMask section_id_mask = 1U << kIndexedSectionId1;
@@ -346,8 +359,9 @@
 TEST_F(DocHitInfoIteratorSectionRestrictTest,
        DoesntIncludeHitWithWrongSectionName) {
   // Populate the DocumentStore's FilterCache with this document's data
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1_)));
   DocumentId document_id = put_result.new_document_id;
 
   SectionIdMask section_id_mask = 1U << kIndexedSectionId0;
@@ -376,8 +390,9 @@
 TEST_F(DocHitInfoIteratorSectionRestrictTest,
        DoesntIncludeHitWithNoSectionIds) {
   // Populate the DocumentStore's FilterCache with this document's data
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1_)));
   DocumentId document_id = put_result.new_document_id;
 
   // Create a hit that doesn't exist in any sections, so it shouldn't match any
@@ -404,8 +419,9 @@
 TEST_F(DocHitInfoIteratorSectionRestrictTest,
        DoesntIncludeHitWithDifferentSectionId) {
   // Populate the DocumentStore's FilterCache with this document's data
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1_)));
   DocumentId document_id = put_result.new_document_id;
 
   // Anything that's not 0, which is the indexed property
@@ -438,7 +454,12 @@
       /*num_leaf_advance_calls_main_index_in=*/5,
       /*num_leaf_advance_calls_integer_index_in=*/3,
       /*num_leaf_advance_calls_no_index_in=*/1,
-      /*num_blocks_inspected_in=*/4);  // arbitrary value
+      /*num_blocks_inspected_in=*/4,
+      /*embedding_stats_in=*/
+      {.num_unquantized_embeddings_scored = 2,
+       .num_quantized_embeddings_scored = 3,
+       .unquantized_shards_read = {1, 2},
+       .quantized_shards_read{3, 4}});  // arbitrary value
   auto original_iterator = std::make_unique<DocHitInfoIteratorDummy>();
   original_iterator->SetCallStats(original_call_stats);
 
@@ -455,14 +476,17 @@
 TEST_F(DocHitInfoIteratorSectionRestrictTest,
        TrimSectionRestrictIterator_TwoLayer) {
   // Populate the DocumentStore's FilterCache with this document's data
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1_)));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document2_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document2_)));
   DocumentId document_id2 = put_result2.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             document_store_->Put(document3_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      document_store_->Put(document_util::CreateDocumentWrapper(document3_)));
   DocumentId document_id3 = put_result3.new_document_id;
 
   // 0 is the indexed property
@@ -516,11 +540,13 @@
 
 TEST_F(DocHitInfoIteratorSectionRestrictTest, TrimSectionRestrictIterator) {
   // Populate the DocumentStore's FilterCache with this document's data
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1_)));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document2_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document2_)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // 0 is the indexed property

diff --git a/icing/index/iterator/doc-hit-info-iterator-test-util.h b/icing/index/iterator/doc-hit-info-iterator-test-util.h
index cb6258d..78822c5 100644
--- a/icing/index/iterator/doc-hit-info-iterator-test-util.h
+++ b/icing/index/iterator/doc-hit-info-iterator-test-util.h

@@ -19,6 +19,7 @@
 #include <cinttypes>
 #include <cstdint>
 #include <cstring>
+#include <memory>
 #include <string>
 #include <utility>
 #include <vector>
@@ -84,7 +85,7 @@
 // will then proceed to return the doc_hit_infos in order as Advance's are
 // called. After all doc_hit_infos are returned, Advance will return a NotFound
 // error (also like normal DocHitInfoIterators).
-class DocHitInfoIteratorDummy : public DocHitInfoLeafIterator {
+class DocHitInfoIteratorDummy : public DocHitInfoIterator {
  public:
   DocHitInfoIteratorDummy() = default;
   explicit DocHitInfoIteratorDummy(
@@ -121,6 +122,10 @@
     return node;
   }
 
+  std::vector<std::unique_ptr<DocHitInfoIterator>*> GetChildren() override {
+    return {};
+  }
+
   // Imitates behavior of DocHitInfoIteratorTermMain/DocHitInfoIteratorTermLite
   void PopulateMatchedTermsStats(
       std::vector<TermMatchInfo>* matched_terms_stats,

diff --git a/icing/index/iterator/doc-hit-info-iterator.h b/icing/index/iterator/doc-hit-info-iterator.h
index f4f15bb..9a85403 100644
--- a/icing/index/iterator/doc-hit-info-iterator.h
+++ b/icing/index/iterator/doc-hit-info-iterator.h

@@ -21,6 +21,7 @@
 #include <memory>
 #include <string>
 #include <string_view>
+#include <unordered_set>
 #include <utility>
 #include <vector>
 
@@ -36,6 +37,7 @@
 namespace lib {
 
 class SectionRestrictData;
+class DocumentFilterPredicate;
 
 // Data structure that maps a single matched query term to its section mask
 // and the list of term frequencies.
@@ -116,18 +118,60 @@
     // - Internal nodes: should aggregate values from all children.
     int32_t num_blocks_inspected;
 
+    // Stats related to embedding index scoring.
+    struct EmbeddingStats {
+      // The number of unquantized embeddings that have been scored.
+      int32_t num_unquantized_embeddings_scored = 0;
+      // The number of quantized embeddings that have been scored.
+      int32_t num_quantized_embeddings_scored = 0;
+      // The set of shards that have been read for unquantized embeddings.
+      std::unordered_set<uint32_t> unquantized_shards_read;
+      // The set of shards that have been read for quantized embeddings.
+      std::unordered_set<uint32_t> quantized_shards_read;
+      // The number of raw embedding bytes read.
+      int64_t num_embedding_bytes_read = 0;
+
+      bool operator==(const EmbeddingStats& other) const {
+        return num_unquantized_embeddings_scored ==
+                   other.num_unquantized_embeddings_scored &&
+               num_quantized_embeddings_scored ==
+                   other.num_quantized_embeddings_scored &&
+               unquantized_shards_read == other.unquantized_shards_read &&
+               quantized_shards_read == other.quantized_shards_read &&
+               num_embedding_bytes_read == other.num_embedding_bytes_read;
+      }
+
+      EmbeddingStats operator+(const EmbeddingStats& other) const {
+        EmbeddingStats result = *this;
+        result.num_unquantized_embeddings_scored +=
+            other.num_unquantized_embeddings_scored;
+        result.num_quantized_embeddings_scored +=
+            other.num_quantized_embeddings_scored;
+        result.unquantized_shards_read.insert(
+            other.unquantized_shards_read.begin(),
+            other.unquantized_shards_read.end());
+        result.quantized_shards_read.insert(other.quantized_shards_read.begin(),
+                                            other.quantized_shards_read.end());
+        result.num_embedding_bytes_read += other.num_embedding_bytes_read;
+        return result;
+      }
+    };
+    EmbeddingStats embedding_stats;
+
     explicit CallStats()
         : CallStats(/*num_leaf_advance_calls_lite_index_in=*/0,
                     /*num_leaf_advance_calls_main_index_in=*/0,
                     /*num_leaf_advance_calls_integer_index_in=*/0,
                     /*num_leaf_advance_calls_no_index_in=*/0,
-                    /*num_blocks_inspected_in=*/0) {}
+                    /*num_blocks_inspected_in=*/0,
+                    /*embedding_stats_in=*/{}) {}
 
     explicit CallStats(int32_t num_leaf_advance_calls_lite_index_in,
                        int32_t num_leaf_advance_calls_main_index_in,
                        int32_t num_leaf_advance_calls_integer_index_in,
                        int32_t num_leaf_advance_calls_no_index_in,
-                       int32_t num_blocks_inspected_in)
+                       int32_t num_blocks_inspected_in,
+                       EmbeddingStats embedding_stats_in)
         : num_leaf_advance_calls_lite_index(
               num_leaf_advance_calls_lite_index_in),
           num_leaf_advance_calls_main_index(
@@ -135,7 +179,8 @@
           num_leaf_advance_calls_integer_index(
               num_leaf_advance_calls_integer_index_in),
           num_leaf_advance_calls_no_index(num_leaf_advance_calls_no_index_in),
-          num_blocks_inspected(num_blocks_inspected_in) {}
+          num_blocks_inspected(num_blocks_inspected_in),
+          embedding_stats(std::move(embedding_stats_in)) {}
 
     int32_t num_leaf_advance_calls() const {
       return num_leaf_advance_calls_lite_index +
@@ -153,7 +198,8 @@
                  other.num_leaf_advance_calls_integer_index &&
              num_leaf_advance_calls_no_index ==
                  other.num_leaf_advance_calls_no_index &&
-             num_blocks_inspected == other.num_blocks_inspected;
+             num_blocks_inspected == other.num_blocks_inspected &&
+             embedding_stats == other.embedding_stats;
     }
 
     CallStats operator+(const CallStats& other) const {
@@ -165,7 +211,8 @@
                            other.num_leaf_advance_calls_integer_index,
                        num_leaf_advance_calls_no_index +
                            other.num_leaf_advance_calls_no_index,
-                       num_blocks_inspected + other.num_blocks_inspected);
+                       num_blocks_inspected + other.num_blocks_inspected,
+                       embedding_stats + other.embedding_stats);
     }
 
     CallStats& operator+=(const CallStats& other) {
@@ -212,21 +259,94 @@
   //   INVALID_ARGUMENT if the right-most node is not suppose to be trimmed.
   virtual libtextclassifier3::StatusOr<TrimmedNode> TrimRightMostNode() && = 0;
 
-  // Map all direct children of this iterator according to the passed mapper.
-  virtual void MapChildren(const ChildrenMapper& mapper) = 0;
+  // Returns raw pointers to the direct children of this iterator. Empty if this
+  // iterator has no children.
+  //
+  // This allows modifying the iterator tree structure, for example, by
+  // modifying the child iterators directly or even replacing them with new
+  // ones. The lifetime of the returned raw pointers is tied to this iterator
+  // object.
+  virtual std::vector<std::unique_ptr<DocHitInfoIterator>*> GetChildren() = 0;
 
-  virtual bool is_leaf() { return false; }
+  // Returns true if section restrictions are not applicable to this iterator.
+  //
+  // Several iterators do **not** need to respect section restrictions, since it
+  // does not have any section information. For example:
+  // - DocHitInfoIteratorAllDocumentId
+  // - DocHitInfoIteratorByUri
+  // - DocHitInfoIteratorMatchScoreExpression
+  // - DocHitInfoIteratorPropertyInSchema
+  // - DocHitInfoIteratorPropertyInDocument
+  //
+  // Unless DocHitInfoIteratorSectionRestrictionNotApplicable is extended, let's
+  // assume the iterator should respect section restrictions.
+  virtual bool SectionRestrictionNotApplicable() const { return false; }
+
+  // If not SectionRestrictionNotApplicable(), whether section restrictions
+  // should be passed down to the children iterators.
+  //
+  // Several iterators need to pass down section restrictions to their
+  // children to maintain the correct semantics of section restrictions. Check
+  // go/icing-section-restrict-fix for more details. For example:
+  // - DocHitInfoIteratorAnd
+  // - DocHitInfoIteratorOr
+  // - DocHitInfoIteratorFilter
+  //
+  // However, several iterators do respect section restrictions, but do not need
+  // to or cannot pass down section restrictions to their children. For example:
+  // - DocHitInfoIteratorSectionRestrict, since it's a section restriction
+  //   iterator itself. A new section restriction should be chained, instead of
+  //   passing down.
+  // - DocHitInfoIteratorTermLite, since it does not have any children. Section
+  //   restriction should be applied at the top of this iterator directly.
+  // - DocHitInfoIteratorEmbedding, since it does not have any children, and
+  //   in addition, it can internally handle the section restriction logic.
+  //
+  // Unless DocHitInfoIteratorSectionRestrictionApplyToChildren is extended,
+  // let's assume this is false, which means section restrictions should be
+  // applied at the top of this iterator directly or handled internally.
+  virtual bool SectionRestrictionShouldApplyToChildren() const { return false; }
 
   // Try to internally handle the provided section restriction in the iterator.
   //
   // Returns:
   //   - false if the iterator does not support handling section restriction.
   //   - true if the iterator supports handling section restriction, and the
-  //     section restriction has been applied.
+  //     section restriction has been applied. For example,
+  //     DocHitInfoIteratorEmbedding can internally handle the section
+  //     restriction logic.
   virtual bool HandleSectionRestriction(SectionRestrictData* other_data) {
     return false;
   }
 
+  // Whether a filter predicate can be passed through this iterator.
+  //
+  // Currently all iterators except for DocHitInfoIteratorNot are able to pass
+  // filter predicates through, while maintaining the same semantics.
+  virtual bool CanPassFilterPredicateThrough() const { return true; }
+
+  // Try to internally handle the provided filter in the iterator.
+  //
+  // Returns:
+  //   - false if the iterator does not support handling filter.
+  //   - true if the iterator supports handling filter, and the filter has been
+  //     applied.
+  virtual bool HandleFilter(const DocumentFilterPredicate* predicate) {
+    return false;
+  }
+
+  // Whether this iterator can adopt a delegate iterator.
+  //
+  // If true, then AdoptDelegate can be called to adopt a delegate iterator.
+  virtual bool CanAdoptDelegate() const { return false; }
+
+  // If CanAdoptDelegate returns false, then this method will have no effect.
+  //
+  // This iterator instance will then filter all of its hits to only include
+  // documents that are returned by the delegate iterator.
+  virtual void AdoptDelegate(std::unique_ptr<DocHitInfoIterator> delegate,
+                     bool delegate_node_is_right_most) {}
+
   virtual ~DocHitInfoIterator() = default;
 
   // Returns:
@@ -279,14 +399,16 @@
   }
 };
 
-// A leaf node is a term node or a chain of section restriction node applied on
-// a term node.
-class DocHitInfoLeafIterator : public DocHitInfoIterator {
+class DocHitInfoIteratorSectionRestrictionNotApplicable
+    : public DocHitInfoIterator {
  public:
-  bool is_leaf() override { return true; }
+  bool SectionRestrictionNotApplicable() const override { return true; }
+};
 
-  // Calling MapChildren on leaf node does not make sense, and will do nothing.
-  void MapChildren(const ChildrenMapper& mapper) override {}
+class DocHitInfoIteratorSectionRestrictionApplyToChildren
+    : public DocHitInfoIterator {
+ public:
+  bool SectionRestrictionShouldApplyToChildren() const override { return true; }
 };
 
 }  // namespace lib

diff --git a/icing/index/iterator/document-filter-predicate.h b/icing/index/iterator/document-filter-predicate.h
new file mode 100644
index 0000000..59e0d05
--- /dev/null
+++ b/icing/index/iterator/document-filter-predicate.h

@@ -0,0 +1,71 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef ICING_INDEX_ITERATOR_DOCUMENT_FILTER_PREDICATE_H_
+#define ICING_INDEX_ITERATOR_DOCUMENT_FILTER_PREDICATE_H_
+
+#include <vector>
+
+#include "icing/index/iterator/doc-hit-info-iterator.h"
+#include "icing/store/document-id.h"
+
+namespace icing {
+namespace lib {
+
+// An interface for a predicate that determines whether a document, identified
+// by its DocumentId, should be included in a result set.
+class DocumentFilterPredicate {
+ public:
+  virtual ~DocumentFilterPredicate() = default;
+
+  // Evaluates whether the given document_id satisfies the predicate. Returns
+  // true if the document satisfies the predicate and should be included, false
+  // otherwise.
+  virtual bool operator()(DocumentId document_id) const = 0;
+};
+
+// Indicate that the iterator can internally handle filtering logic by itself.
+//
+// This is helpful when some iterators want to have better control for
+// optimization. For example, embedding iterator will be able to filter out
+// embedding hits from unwanted documents to avoid retrieving unnecessary
+// vectors and calculate scores for them.
+class DocHitInfoIteratorHandlingFilter : virtual public DocHitInfoIterator {
+ protected:
+  // After accepting a filter predicate, the iterator will behave equivalently
+  // as if we had applied a filter with this predicate at the top of the
+  // iterator.
+  bool HandleFilter(const DocumentFilterPredicate* predicate) override {
+    document_filter_predicates_.push_back(predicate);
+    return true;
+  }
+
+  bool DoesDocumentPassAllFilters(DocumentId document_id) const {
+    for (const DocumentFilterPredicate* predicate :
+         document_filter_predicates_) {
+      if (!(*predicate)(document_id)) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  // Does not own the pointers.
+  std::vector<const DocumentFilterPredicate*> document_filter_predicates_;
+};
+
+}  // namespace lib
+}  // namespace icing
+
+#endif  // ICING_INDEX_ITERATOR_DOCUMENT_FILTER_PREDICATE_H_

diff --git a/icing/index/iterator/section-restrict-data.h b/icing/index/iterator/section-restrict-data.h
index 38d7158..208d629 100644
--- a/icing/index/iterator/section-restrict-data.h
+++ b/icing/index/iterator/section-restrict-data.h

@@ -114,7 +114,7 @@
 // embedding hits from unwanted sections to avoid retrieving unnecessary vectors
 // and calculate scores for them.
 class DocHitInfoIteratorHandlingSectionRestrict
-    : public DocHitInfoLeafIterator {
+    : virtual public DocHitInfoIterator {
  protected:
   bool HandleSectionRestriction(SectionRestrictData* other_data) override {
     section_restrict_data_.push_back(other_data);

diff --git a/icing/index/lite/doc-hit-info-iterator-term-lite.cc b/icing/index/lite/doc-hit-info-iterator-term-lite.cc
index c356203..c45cb7a 100644
--- a/icing/index/lite/doc-hit-info-iterator-term-lite.cc
+++ b/icing/index/lite/doc-hit-info-iterator-term-lite.cc

@@ -16,10 +16,8 @@
 
 #include <algorithm>
 #include <cstdint>
-#include <cstring>
 #include <numeric>
 #include <string>
-#include <utility>
 #include <vector>
 
 #include "icing/text_classifier/lib3/utils/base/status.h"
@@ -33,6 +31,7 @@
 #include "icing/index/term-id-codec.h"
 #include "icing/schema/section.h"
 #include "icing/util/logging.h"
+#include "icing/util/math-util.h"
 #include "icing/util/status-macros.h"
 
 namespace icing {
@@ -153,31 +152,9 @@
     // Now indices is a map from sorted index to current index. In other words,
     // the sorted cached_hits_[i] should be the current cached_hits_[indices[i]]
     // for every valid i.
-    std::vector<bool> done(indices.size());
-    // Apply permutation
-    for (int i = 0; i < indices.size(); ++i) {
-      if (done[i]) {
-        continue;
-      }
-      done[i] = true;
-      int curr = i;
-      int next = indices[i];
-      // Since every finite permutation is formed by disjoint cycles, we can
-      // start with the current element, at index i, and swap the element at
-      // this position with whatever element that *should* be here. Then,
-      // continue to swap the original element, at its updated positions, with
-      // the element that should be occupying that position until the original
-      // element has reached *its* correct position. This completes applying the
-      // single cycle in the permutation.
-      while (next != i) {
-        std::swap(cached_hits_[curr], cached_hits_[next]);
-        std::swap(cached_hit_term_frequency_[curr],
-                  cached_hit_term_frequency_[next]);
-        done[next] = true;
-        curr = next;
-        next = indices[next];
-      }
-    }
+
+    math_util::ApplyPermutation(indices, cached_hits_,
+                                cached_hit_term_frequency_);
   }
 }
 

diff --git a/icing/index/lite/doc-hit-info-iterator-term-lite.h b/icing/index/lite/doc-hit-info-iterator-term-lite.h
index a63d050..ff286d3 100644
--- a/icing/index/lite/doc-hit-info-iterator-term-lite.h
+++ b/icing/index/lite/doc-hit-info-iterator-term-lite.h

@@ -16,6 +16,7 @@
 #define ICING_INDEX_ITERATOR_DOC_HIT_INFO_ITERATOR_TERM_LITE_H_
 
 #include <array>
+#include <memory>
 #include <string>
 #include <utility>
 #include <vector>
@@ -32,7 +33,7 @@
 namespace icing {
 namespace lib {
 
-class DocHitInfoIteratorTermLite : public DocHitInfoLeafIterator {
+class DocHitInfoIteratorTermLite : public DocHitInfoIterator {
  public:
   explicit DocHitInfoIteratorTermLite(const TermIdCodec* term_id_codec,
                                       LiteIndex* lite_index,
@@ -55,13 +56,18 @@
 
   libtextclassifier3::StatusOr<TrimmedNode> TrimRightMostNode() && override;
 
+  std::vector<std::unique_ptr<DocHitInfoIterator>*> GetChildren() override {
+    return {};
+  }
+
   CallStats GetCallStats() const override {
     return CallStats(
         /*num_leaf_advance_calls_lite_index_in=*/num_advance_calls_,
         /*num_leaf_advance_calls_main_index_in=*/0,
         /*num_leaf_advance_calls_integer_index_in=*/0,
         /*num_leaf_advance_calls_no_index_in=*/0,
-        /*num_blocks_inspected_in=*/0);
+        /*num_blocks_inspected_in=*/0,
+        /*embedding_stats_in=*/{});
   }
 
   void PopulateMatchedTermsStats(

diff --git a/icing/index/lite/lite-index-header.h b/icing/index/lite/lite-index-header.h
index 1aba130..1677d69 100644
--- a/icing/index/lite/lite-index-header.h
+++ b/icing/index/lite/lite-index-header.h

@@ -22,6 +22,7 @@
 #include "icing/legacy/core/icing-string-util.h"
 #include "icing/store/document-id.h"
 #include "icing/util/crc32.h"
+#include "icing/util/logging.h"
 
 namespace icing {
 namespace lib {
@@ -55,7 +56,7 @@
 class LiteIndex_HeaderImpl : public LiteIndex_Header {
  public:
   struct HeaderData {
-    static const uint32_t kMagic = 0xC2EAD682;
+    static constexpr uint32_t kMagic = 0xC2EAD682;
 
     uint32_t lite_index_crc;
     uint32_t magic;
@@ -74,6 +75,10 @@
   explicit LiteIndex_HeaderImpl(HeaderData *hdr) : hdr_(hdr) {}
 
   bool check_magic() const override {
+    if (hdr_->magic != HeaderData::kMagic) {
+      ICING_LOG(ERROR) << "Invalid header magic for LiteIndex. Expected: "
+                       << HeaderData::kMagic << ", actual: " << hdr_->magic;
+    }
     return hdr_->magic == HeaderData::kMagic;
   }
 

diff --git a/icing/index/lite/lite-index.cc b/icing/index/lite/lite-index.cc
index 3862206..ee7b51f 100644
--- a/icing/index/lite/lite-index.cc
+++ b/icing/index/lite/lite-index.cc

@@ -195,7 +195,10 @@
 
     // Check integrity.
     if (!header_->check_magic()) {
-      status = absl_ports::InternalError("Lite index header magic mismatch");
+      ICING_LOG(ERROR) << "Invalid header magic for LiteIndex "
+                       << options_.filename_base;
+      status = absl_ports::InternalError(absl_ports::StrCat(
+          "Invalid header magic for LiteIndex: ", options_.filename_base));
       goto error;
     }
     Crc32 expected_crc(header_->lite_index_crc());
@@ -353,12 +356,9 @@
 
   TermIdHitPair term_id_hit_pair(term_id, hit);
   uint32_t cur_size = header_->cur_size();
-  TermIdHitPair::Value* valp =
-      hit_buffer_.GetMutableMem<TermIdHitPair::Value>(cur_size, 1);
-  if (valp == nullptr) {
-    return absl_ports::ResourceExhaustedError(
-        "Allocating more space in hit buffer failed!");
-  }
+  ICING_ASSIGN_OR_RETURN(
+      TermIdHitPair::Value * valp,
+      hit_buffer_.GetMutableMem<TermIdHitPair::Value>(cur_size, 1));
   *valp = term_id_hit_pair.value();
   header_->set_cur_size(cur_size + 1);
 
@@ -473,7 +473,12 @@
     // after need_sort_at_querying is evaluated.
     // We check need_sort_at_querying to improve query concurrency as threads
     // can avoid acquiring the unique lock if no sorting is needed.
-    SortHitsImpl();
+    libtextclassifier3::Status status = SortHitsImpl();
+    if (!status.ok()) {
+      // Log this error and continue.
+      ICING_LOG(ERROR) << "Failed to sort HitBuffer: "
+                       << status.error_message();
+    }
 
     if (options_.hit_buffer_sort_at_indexing) {
       // This is the second case for sort. Log as this should be a very rare
@@ -624,16 +629,26 @@
   return storage_info;
 }
 
-void LiteIndex::SortHitsImpl() {
+libtextclassifier3::Status LiteIndex::SortHitsImpl() {
   // Make searchable by sorting by hit buffer.
   uint32_t need_sort_len = GetHitBufferUnsortedSizeImpl();
   if (need_sort_len <= 0) {
-    return;
+    return libtextclassifier3::Status::OK;
   }
   IcingTimer timer;
 
-  TermIdHitPair::Value* array_start =
+  auto array_start_or =
       hit_buffer_.GetMutableMem<TermIdHitPair::Value>(0, header_->cur_size());
+  if (!array_start_or.ok()) {
+    // The error here means an allocation of the hit buffer array failed, but
+    // this should NEVER happen since the hit buffer size should be at least
+    // header_->cur_size() at this moment.
+    ICING_LOG(ERROR)
+        << "GetMutableMem failed in SortHitsImpl, which should never happen: "
+        << array_start_or.status().error_message();
+    return array_start_or.status();
+  }
+  TermIdHitPair::Value* array_start = array_start_or.ValueOrDie();
   TermIdHitPair::Value* sort_start = array_start + header_->searchable_end();
   std::sort(sort_start, array_start + header_->cur_size());
 
@@ -653,6 +668,7 @@
 
   // Update crc in-line.
   UpdateChecksumInternal();
+  return libtextclassifier3::Status::OK;
 }
 
 libtextclassifier3::Status LiteIndex::Optimize(
@@ -665,7 +681,7 @@
   }
   // Sort the hits so that hits with the same term id will be grouped together,
   // which helps later to determine which terms will be unused after compaction.
-  SortHitsImpl();
+  ICING_RETURN_IF_ERROR(SortHitsImpl());
   uint32_t new_size = 0;
   uint32_t curr_term_id = 0;
   uint32_t curr_tvi = 0;
@@ -711,8 +727,9 @@
     // new_size is weakly less than idx so we are okay to overwrite the entry at
     // new_size, and valp should never be nullptr since it is within the already
     // allocated region of hit_buffer_.
-    TermIdHitPair::Value* valp =
-        hit_buffer_.GetMutableMem<TermIdHitPair::Value>(new_size++, 1);
+    ICING_ASSIGN_OR_RETURN(
+        TermIdHitPair::Value * valp,
+        hit_buffer_.GetMutableMem<TermIdHitPair::Value>(new_size++, 1));
     if (valp == nullptr) {
       // This really shouldn't happen since we are only writing to the already
       // allocated region of hit_buffer_. But just in case, we log and return an
@@ -741,9 +758,11 @@
     // saves an unnecessary search through the hit buffer). This is acceptable
     // because the free space will eventually be reclaimed the next time that
     // the lite index is merged with the main index.
-    if (!lexicon_.Delete(term)) {
-      return absl_ports::InternalError(
-          "Could not delete invalid terms in lite lexicon during compaction.");
+    libtextclassifier3::Status status = lexicon_.Delete(term);
+    if (!status.ok()) {
+      return absl_ports::InternalError(absl_ports::StrCat(
+          "Could not delete invalid terms in lite lexicon during compaction: ",
+          status.error_message()));
     }
   }
   return libtextclassifier3::Status::OK;

diff --git a/icing/index/lite/lite-index.h b/icing/index/lite/lite-index.h
index 280df88..e174556 100644
--- a/icing/index/lite/lite-index.h
+++ b/icing/index/lite/lite-index.h

@@ -418,7 +418,8 @@
   }
 
   // Non-locking implementation for SortHits().
-  void SortHitsImpl() ICING_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
+  libtextclassifier3::Status SortHitsImpl()
+      ICING_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
 
   // Calculates and adds the score for a fetched hit to total_score_out, while
   // updating last_document_id (which keeps track of the last added docId so

diff --git a/icing/index/main/doc-hit-info-iterator-term-main.h b/icing/index/main/doc-hit-info-iterator-term-main.h
index e32db2a..559412e 100644
--- a/icing/index/main/doc-hit-info-iterator-term-main.h
+++ b/icing/index/main/doc-hit-info-iterator-term-main.h

@@ -15,7 +15,7 @@
 #ifndef ICING_INDEX_ITERATOR_DOC_HIT_INFO_ITERATOR_TERM_MAIN_H_
 #define ICING_INDEX_ITERATOR_DOC_HIT_INFO_ITERATOR_TERM_MAIN_H_
 
-#include <cstdint>
+#include <array>
 #include <memory>
 #include <optional>
 #include <string>
@@ -23,6 +23,7 @@
 #include <vector>
 
 #include "icing/text_classifier/lib3/utils/base/status.h"
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "icing/index/hit/doc-hit-info.h"
 #include "icing/index/hit/hit.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
@@ -33,7 +34,7 @@
 namespace icing {
 namespace lib {
 
-class DocHitInfoIteratorTermMain : public DocHitInfoLeafIterator {
+class DocHitInfoIteratorTermMain : public DocHitInfoIterator {
  public:
   struct DocHitInfoAndTermFrequencyArray {
     DocHitInfo doc_hit_info;
@@ -70,13 +71,18 @@
 
   libtextclassifier3::StatusOr<TrimmedNode> TrimRightMostNode() && override;
 
+  std::vector<std::unique_ptr<DocHitInfoIterator>*> GetChildren() override {
+    return {};
+  }
+
   CallStats GetCallStats() const override {
     return CallStats(
         /*num_leaf_advance_calls_lite_index_in=*/0,
         /*num_leaf_advance_calls_main_index_in=*/num_advance_calls_,
         /*num_leaf_advance_calls_integer_index_in=*/0,
         /*num_leaf_advance_calls_no_index_in=*/0,
-        /*num_blocks_inspected_in=*/num_blocks_inspected_);
+        /*num_blocks_inspected_in=*/num_blocks_inspected_,
+        /*embedding_stats_in=*/{});
   }
 
   void PopulateMatchedTermsStats(

diff --git a/icing/index/main/main-index-merger.cc b/icing/index/main/main-index-merger.cc
index 8397e2c..0e5016e 100644
--- a/icing/index/main/main-index-merger.cc
+++ b/icing/index/main/main-index-merger.cc

@@ -31,6 +31,7 @@
 #include "icing/index/term-id-codec.h"
 #include "icing/legacy/core/icing-string-util.h"
 #include "icing/util/logging.h"
+#include "icing/util/math-util.h"
 #include "icing/util/status-macros.h"
 
 namespace icing {
@@ -142,15 +143,37 @@
   TermIdHitPair prev_;
 };
 
-class HitComparator {
- public:
-  explicit HitComparator(
-      const TermIdCodec& term_id_codec,
-      const std::unordered_map<uint32_t, int>& main_tvi_to_block_index)
-      : term_id_codec_(&term_id_codec),
-        main_tvi_to_block_index_(&main_tvi_to_block_index) {}
+int GetIndexBlock(
+    uint32_t term_id,
+    const std::unordered_map<uint32_t, int>& main_tvi_to_block_index,
+    const TermIdCodec& term_id_codec) {
+  auto term_info_or = term_id_codec.DecodeTermInfo(term_id);
+  if (!term_info_or.ok()) {
+    ICING_LOG(WARNING)
+        << "Unable to decode term-info during merge. This shouldn't happen.";
+    return kInvalidBlockIndex;
+  }
+  TermIdCodec::DecodedTermInfo term_info = std::move(term_info_or).ValueOrDie();
+  auto itr = main_tvi_to_block_index.find(term_info.tvi);
+  if (itr == main_tvi_to_block_index.end()) {
+    return kInvalidBlockIndex;
+  }
+  return itr->second;
+}
 
-  bool operator()(const TermIdHitPair& lhs, const TermIdHitPair& rhs) const {
+void SortHits(std::vector<TermIdHitPair>& hits,
+              const std::unordered_map<uint32_t, int>& main_tvi_to_block_index,
+              const TermIdCodec& term_id_codec) {
+  std::vector<int> indices;
+  indices.reserve(hits.size());
+  std::vector<int> index_blocks;
+  index_blocks.reserve(hits.size());
+  for (int i = 0; i < hits.size(); ++i) {
+    indices.push_back(i);
+    index_blocks.push_back(GetIndexBlock(
+        hits[i].term_id(), main_tvi_to_block_index, term_id_codec));
+  }
+  std::sort(indices.begin(), indices.end(), [&](int i, int j) {
     // Primary sort by index block. This achieves two things:
     // 1. It reduces the number of flash writes by grouping together new hits
     // for terms whose posting lists might share the same index block.
@@ -158,35 +181,14 @@
     // will be populated first (because all newly added terms have an invalid
     // block index of 0) before any new hits are added to the postings lists
     // that they backfill from.
-    int lhs_index_block = GetIndexBlock(lhs.term_id());
-    int rhs_index_block = GetIndexBlock(rhs.term_id());
-    if (lhs_index_block == rhs_index_block) {
-      // Secondary sort by term_id and hit.
-      return lhs.value() < rhs.value();
+    if (index_blocks[i] == index_blocks[j]) {
+      return hits[i].value() < hits[j].value();
     }
-    return lhs_index_block < rhs_index_block;
-  }
+    return index_blocks[i] < index_blocks[j];
+  });
 
- private:
-  int GetIndexBlock(uint32_t term_id) const {
-    auto term_info_or = term_id_codec_->DecodeTermInfo(term_id);
-    if (!term_info_or.ok()) {
-      ICING_LOG(WARNING)
-          << "Unable to decode term-info during merge. This shouldn't happen.";
-      return kInvalidBlockIndex;
-    }
-    TermIdCodec::DecodedTermInfo term_info =
-        std::move(term_info_or).ValueOrDie();
-    auto itr = main_tvi_to_block_index_->find(term_info.tvi);
-    if (itr == main_tvi_to_block_index_->end()) {
-      return kInvalidBlockIndex;
-    }
-    return itr->second;
-  }
-
-  const TermIdCodec* term_id_codec_;
-  const std::unordered_map<uint32_t, int>* main_tvi_to_block_index_;
-};
+  math_util::ApplyPermutation(indices, hits);
+}
 
 // A helper function to dedupe hits stored in hits. Suppose that the lite index
 // contained a single document with two hits in a single prefix section: "foot"
@@ -211,8 +213,7 @@
     const std::unordered_map<uint32_t, int>& main_tvi_to_block_index) {
   // Now all terms are grouped together and all hits for a term are sorted.
   // Merge equivalent hits into one.
-  std::sort(hits->begin(), hits->end(),
-            HitComparator(term_id_codec, main_tvi_to_block_index));
+  SortHits(*hits, main_tvi_to_block_index, term_id_codec);
   size_t current_offset = 0;
   HitSelector hit_selector;
   for (const TermIdHitPair& term_id_hit_pair : *hits) {

diff --git a/icing/index/main/main-index.cc b/icing/index/main/main-index.cc
index 12e8ea9..1956be4 100644
--- a/icing/index/main/main-index.cc
+++ b/icing/index/main/main-index.cc

@@ -343,8 +343,9 @@
        other_term_itr.IsValid(); other_term_itr.Advance()) {
     // If term were inserted in the main lexicon, what new branching would it
     // create? (It always creates at most one.)
-    int prefix_len = main_lexicon_->FindNewBranchingPrefixLength(
-        other_term_itr.GetKey(), /*utf8=*/true);
+    ICING_ASSIGN_OR_RETURN(int prefix_len,
+                           main_lexicon_->FindNewBranchingPrefixLength(
+                               other_term_itr.GetKey(), /*utf8=*/true));
     if (prefix_len <= 0) {
       continue;
     }
@@ -449,8 +450,9 @@
 
     // Get prefixes where there is already a branching point in the main
     // lexicon. We skip prefixes which don't already have a branching point.
-    std::vector<int> prefix_lengths = main_lexicon_->FindBranchingPrefixLengths(
-        other_term_itr.GetKey(), /*utf8=*/true);
+    ICING_ASSIGN_OR_RETURN(std::vector<int> prefix_lengths,
+                           main_lexicon_->FindBranchingPrefixLengths(
+                               other_term_itr.GetKey(), /*utf8=*/true));
 
     int buf_start = outputs.prefix_tvis_buf.size();
     // Add prefixes.
@@ -592,7 +594,8 @@
     PostingListAccessor::FinalizeResult result =
         std::move(*hit_accum).Finalize();
     if (result.id.is_valid()) {
-      main_lexicon_->SetValueAtIndex(other_tvi_main_tvi_pair.first, &result.id);
+      ICING_RETURN_IF_ERROR(main_lexicon_->SetValueAtIndex(
+          other_tvi_main_tvi_pair.first, &result.id));
     }
   }
   flash_index_storage_->set_last_indexed_docid(last_added_document_id);
@@ -648,7 +651,7 @@
   PostingListAccessor::FinalizeResult result =
       std::move(*pl_accessor).Finalize();
   if (result.id.is_valid()) {
-    main_lexicon_->SetValueAtIndex(tvi, &result.id);
+    ICING_RETURN_IF_ERROR(main_lexicon_->SetValueAtIndex(tvi, &result.id));
   }
   return libtextclassifier3::Status::OK;
 }
@@ -803,8 +806,9 @@
   // A term without exact hits indicates that it is a purely backfill term. If
   // the term is not branching in the new trie, it means backfilling is no
   // longer necessary, so that we can skip.
-  if (new_hits.empty() ||
-      (has_no_exact_hits && !new_index->main_lexicon_->IsBranchingTerm(term))) {
+  ICING_ASSIGN_OR_RETURN(bool is_branching_term,
+                         new_index->main_lexicon_->IsBranchingTerm(term));
+  if (new_hits.empty() || (has_no_exact_hits && !is_branching_term)) {
     return largest_document_id;
   }
 

diff --git a/icing/index/numeric/doc-hit-info-iterator-numeric.h b/icing/index/numeric/doc-hit-info-iterator-numeric.h
index 7cdb230..d8f83b4 100644
--- a/icing/index/numeric/doc-hit-info-iterator-numeric.h
+++ b/icing/index/numeric/doc-hit-info-iterator-numeric.h

@@ -20,16 +20,18 @@
 #include <vector>
 
 #include "icing/text_classifier/lib3/utils/base/status.h"
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "icing/absl_ports/canonical_errors.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
 #include "icing/index/numeric/numeric-index.h"
+#include "icing/schema/section.h"
 #include "icing/util/status-macros.h"
 
 namespace icing {
 namespace lib {
 
 template <typename T>
-class DocHitInfoIteratorNumeric : public DocHitInfoLeafIterator {
+class DocHitInfoIteratorNumeric : public DocHitInfoIterator {
  public:
   explicit DocHitInfoIteratorNumeric(
       std::unique_ptr<typename NumericIndex<T>::Iterator> numeric_index_iter)
@@ -53,6 +55,10 @@
         "Cannot generate suggestion if the last term is numeric operator.");
   }
 
+  std::vector<std::unique_ptr<DocHitInfoIterator>*> GetChildren() override {
+    return {};
+  }
+
   CallStats GetCallStats() const override {
     if (numeric_index_iter_ == nullptr) {
       return CallStats();
@@ -64,7 +70,8 @@
                      numeric_index_iter_->GetNumAdvanceCalls(),
                      /*num_leaf_advance_calls_no_index_in=*/0,
                      /*num_blocks_inspected_in=*/
-                     numeric_index_iter_->GetNumBlocksInspected());
+                     numeric_index_iter_->GetNumBlocksInspected(),
+                     /*embedding_stats_in=*/{});
   }
 
   std::string ToString() const override { return "test"; }

diff --git a/icing/index/numeric/integer-index-storage.cc b/icing/index/numeric/integer-index-storage.cc
index 9566737..033c64a 100644
--- a/icing/index/numeric/integer-index-storage.cc
+++ b/icing/index/numeric/integer-index-storage.cc

@@ -46,6 +46,7 @@
 #include "icing/schema/section.h"
 #include "icing/store/document-id.h"
 #include "icing/util/crc32.h"
+#include "icing/util/logging.h"
 #include "icing/util/status-macros.h"
 
 namespace icing {
@@ -928,7 +929,13 @@
   // Validate other values of info and options.
   // Magic should be consistent with the codebase.
   if (integer_index_storage->info().magic != Info::kMagic) {
-    return absl_ports::FailedPreconditionError("Incorrect magic value");
+    ICING_LOG(ERROR) << "Invalid header magic for IntegerIndexStorage "
+                     << integer_index_storage->working_path_
+                     << ". Expected: " << Info::kMagic
+                     << ", actual: " << integer_index_storage->info().magic;
+    return absl_ports::FailedPreconditionError(
+        absl_ports::StrCat("Invalid header magic for IntegerIndexStorage: ",
+                           integer_index_storage->working_path_));
   }
 
   return integer_index_storage;

diff --git a/icing/index/numeric/integer-index-storage_test.cc b/icing/index/numeric/integer-index-storage_test.cc
index 62c8e87..65643da 100644
--- a/icing/index/numeric/integer-index-storage_test.cc
+++ b/icing/index/numeric/integer-index-storage_test.cc

@@ -309,15 +309,17 @@
 
   // Check info section
   Info info;
-  ASSERT_TRUE(filesystem_.PRead(metadata_sfd.get(), &info, sizeof(Info),
-                                IntegerIndexStorage::kInfoMetadataFileOffset));
+  ASSERT_THAT(filesystem_.PRead(metadata_sfd.get(), &info, sizeof(Info),
+                                IntegerIndexStorage::kInfoMetadataFileOffset),
+              Eq(sizeof(Info)));
   EXPECT_THAT(info.magic, Eq(Info::kMagic));
   EXPECT_THAT(info.num_data, Eq(0));
 
   // Check crcs section
   Crcs crcs;
-  ASSERT_TRUE(filesystem_.PRead(metadata_sfd.get(), &crcs, sizeof(Crcs),
-                                IntegerIndexStorage::kCrcsMetadataFileOffset));
+  ASSERT_THAT(filesystem_.PRead(metadata_sfd.get(), &crcs, sizeof(Crcs),
+                                IntegerIndexStorage::kCrcsMetadataFileOffset),
+              Eq(sizeof(Crcs)));
   // # of elements in sorted_buckets should be 1, so it should have non-zero
   // all storages crc value.
   EXPECT_THAT(crcs.component_crcs.storages_crc, Ne(0));
@@ -518,8 +520,9 @@
   ASSERT_TRUE(metadata_sfd.is_valid());
 
   Crcs crcs;
-  ASSERT_TRUE(filesystem_.PRead(metadata_sfd.get(), &crcs, sizeof(Crcs),
-                                IntegerIndexStorage::kCrcsMetadataFileOffset));
+  ASSERT_THAT(filesystem_.PRead(metadata_sfd.get(), &crcs, sizeof(Crcs),
+                                IntegerIndexStorage::kCrcsMetadataFileOffset),
+              Eq(sizeof(Crcs)));
 
   // Manually corrupt all_crc
   crcs.all_crc += kCorruptedValueOffset;
@@ -567,8 +570,9 @@
   ASSERT_TRUE(metadata_sfd.is_valid());
 
   Info info;
-  ASSERT_TRUE(filesystem_.PRead(metadata_sfd.get(), &info, sizeof(Info),
-                                IntegerIndexStorage::kInfoMetadataFileOffset));
+  ASSERT_THAT(filesystem_.PRead(metadata_sfd.get(), &info, sizeof(Info),
+                                IntegerIndexStorage::kInfoMetadataFileOffset),
+              Eq(sizeof(Info)));
 
   // Modify info, but don't update the checksum. This would be similar to
   // corruption of info.
@@ -1513,13 +1517,14 @@
   while (iter1->Advance().ok()) {
     // Advance all hits.
   }
-  EXPECT_THAT(
-      iter1->GetCallStats(),
-      EqualsDocHitInfoIteratorCallStats(
-          /*num_leaf_advance_calls_lite_index=*/0,
-          /*num_leaf_advance_calls_main_index=*/0,
-          /*num_leaf_advance_calls_integer_index=*/5,
-          /*num_leaf_advance_calls_no_index=*/0, /*num_blocks_inspected=*/2));
+  EXPECT_THAT(iter1->GetCallStats(),
+              EqualsDocHitInfoIteratorCallStats(
+                  /*num_leaf_advance_calls_lite_index=*/0,
+                  /*num_leaf_advance_calls_main_index=*/0,
+                  /*num_leaf_advance_calls_integer_index=*/5,
+                  /*num_leaf_advance_calls_no_index=*/0,
+                  /*num_blocks_inspected=*/2,
+                  DocHitInfoIterator::CallStats::EmbeddingStats()));
 
   // GetIterator for range [-1000, -100] and Advance all. Since we only have to
   // read bucket (-1000,-100), there will be 3 advance calls and 1 block
@@ -1530,13 +1535,14 @@
   while (iter2->Advance().ok()) {
     // Advance all hits.
   }
-  EXPECT_THAT(
-      iter2->GetCallStats(),
-      EqualsDocHitInfoIteratorCallStats(
-          /*num_leaf_advance_calls_lite_index=*/0,
-          /*num_leaf_advance_calls_main_index=*/0,
-          /*num_leaf_advance_calls_integer_index=*/3,
-          /*num_leaf_advance_calls_no_index=*/0, /*num_blocks_inspected=*/1));
+  EXPECT_THAT(iter2->GetCallStats(),
+              EqualsDocHitInfoIteratorCallStats(
+                  /*num_leaf_advance_calls_lite_index=*/0,
+                  /*num_leaf_advance_calls_main_index=*/0,
+                  /*num_leaf_advance_calls_integer_index=*/3,
+                  /*num_leaf_advance_calls_no_index=*/0,
+                  /*num_blocks_inspected=*/1,
+                  DocHitInfoIterator::CallStats::EmbeddingStats()));
 }
 
 TEST_P(IntegerIndexStorageTest, IteratorCallStatsSingleBucketChainedBlocks) {
@@ -1570,13 +1576,14 @@
   while (iter1->Advance().ok()) {
     // Advance all hits.
   }
-  EXPECT_THAT(iter1->GetCallStats(),
-              EqualsDocHitInfoIteratorCallStats(
-                  /*num_leaf_advance_calls_lite_index=*/0,
-                  /*num_leaf_advance_calls_main_index=*/0,
-                  /*num_leaf_advance_calls_integer_index=*/num_keys_to_add,
-                  /*num_leaf_advance_calls_no_index=*/0,
-                  expected_num_blocks_inspected));
+  EXPECT_THAT(
+      iter1->GetCallStats(),
+      EqualsDocHitInfoIteratorCallStats(
+          /*num_leaf_advance_calls_lite_index=*/0,
+          /*num_leaf_advance_calls_main_index=*/0,
+          /*num_leaf_advance_calls_integer_index=*/num_keys_to_add,
+          /*num_leaf_advance_calls_no_index=*/0, expected_num_blocks_inspected,
+          DocHitInfoIterator::CallStats::EmbeddingStats()));
 
   // GetIterator for range [1, 1] and Advance all. Although there is only 1
   // relevant data, we still have to inspect the entire bucket and its posting
@@ -1587,13 +1594,14 @@
   while (iter2->Advance().ok()) {
     // Advance all hits.
   }
-  EXPECT_THAT(iter2->GetCallStats(),
-              EqualsDocHitInfoIteratorCallStats(
-                  /*num_leaf_advance_calls_lite_index=*/0,
-                  /*num_leaf_advance_calls_main_index=*/0,
-                  /*num_leaf_advance_calls_integer_index=*/num_keys_to_add,
-                  /*num_leaf_advance_calls_no_index=*/0,
-                  expected_num_blocks_inspected));
+  EXPECT_THAT(
+      iter2->GetCallStats(),
+      EqualsDocHitInfoIteratorCallStats(
+          /*num_leaf_advance_calls_lite_index=*/0,
+          /*num_leaf_advance_calls_main_index=*/0,
+          /*num_leaf_advance_calls_integer_index=*/num_keys_to_add,
+          /*num_leaf_advance_calls_no_index=*/0, expected_num_blocks_inspected,
+          DocHitInfoIterator::CallStats::EmbeddingStats()));
 }
 
 TEST_P(IntegerIndexStorageTest, SplitBuckets) {

diff --git a/icing/index/numeric/integer-index.cc b/icing/index/numeric/integer-index.cc
index 62fb215..36148d4 100644
--- a/icing/index/numeric/integer-index.cc
+++ b/icing/index/numeric/integer-index.cc

@@ -35,6 +35,7 @@
 #include "icing/index/numeric/posting-list-integer-index-serializer.h"
 #include "icing/store/document-id.h"
 #include "icing/util/crc32.h"
+#include "icing/util/logging.h"
 #include "icing/util/status-macros.h"
 
 namespace icing {
@@ -494,7 +495,13 @@
 
   // Validate magic.
   if (integer_index->info().magic != Info::kMagic) {
-    return absl_ports::FailedPreconditionError("Incorrect magic value");
+    ICING_LOG(ERROR) << "Invalid header magic for IntegerIndex "
+                     << integer_index->working_path_
+                     << ". Expected: " << Info::kMagic
+                     << ", actual: " << integer_index->info().magic;
+    return absl_ports::FailedPreconditionError(
+        absl_ports::StrCat("Invalid header magic for IntegerIndex: ",
+                           integer_index->working_path_));
   }
 
   // If num_data_threshold_for_bucket_split mismatches, then return error to let

diff --git a/icing/index/numeric/integer-index_test.cc b/icing/index/numeric/integer-index_test.cc
index 619abc4..ef48314 100644
--- a/icing/index/numeric/integer-index_test.cc
+++ b/icing/index/numeric/integer-index_test.cc

@@ -20,15 +20,20 @@
 #include <string>
 #include <string_view>
 #include <type_traits>
+#include <unordered_set>
+#include <utility>
 #include <vector>
 
 #include "icing/text_classifier/lib3/utils/base/status.h"
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
+#include "icing/absl_ports/canonical_errors.h"
+#include "icing/absl_ports/str_cat.h"
 #include "icing/document-builder.h"
 #include "icing/feature-flags.h"
 #include "icing/file/filesystem.h"
+#include "icing/file/persistent-storage.h"
 #include "icing/file/portable-file-backed-proto-log.h"
 #include "icing/index/hit/doc-hit-info.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
@@ -36,6 +41,7 @@
 #include "icing/index/numeric/integer-index-storage.h"
 #include "icing/index/numeric/numeric-index.h"
 #include "icing/index/numeric/posting-list-integer-index-serializer.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/schema.pb.h"
 #include "icing/schema-builder.h"
@@ -46,6 +52,10 @@
 #include "icing/testing/common-matchers.h"
 #include "icing/testing/test-feature-flags.h"
 #include "icing/testing/tmp-directory.h"
+#include "icing/util/clock.h"
+#include "icing/util/crc32.h"
+#include "icing/util/document-util.h"
+#include "icing/util/status-macros.h"
 
 namespace icing {
 namespace lib {
@@ -92,14 +102,18 @@
         filesystem_.CreateDirectoryRecursively(document_store_dir.c_str()));
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult doc_store_create_result,
-        DocumentStore::Create(&filesystem_, document_store_dir, &clock_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, document_store_dir, &clock_, schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     doc_store_ = std::move(doc_store_create_result.document_store);
   }
 
@@ -167,14 +181,18 @@
 
     ICING_ASSIGN_OR_RETURN(
         DocumentStore::CreateResult doc_store_create_result,
-        DocumentStore::Create(&filesystem_, document_store_dir, &clock_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, document_store_dir, &clock_, schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     doc_store_ = std::move(doc_store_create_result.document_store);
     return std::move(doc_store_optimize_result.document_id_old_to_new);
   }
@@ -413,49 +431,51 @@
   // Put 11 docs of "TypeA" into the document store.
   DocumentProto doc =
       DocumentBuilder().SetKey("ns1", "uri0").SetSchema("TypeA").Build();
-  ICING_ASSERT_OK(this->doc_store_->Put(doc));
   ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri1").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri2").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri3").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri4").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri5").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri6").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri7").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri8").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri9").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri10").Build()));
+      this->doc_store_->Put(document_util::CreateDocumentWrapper(doc)));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri1").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri2").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri3").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri4").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri5").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri6").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri7").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri8").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri9").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri10").Build())));
 
   // Put 5 docs of "TypeB" into the document store.
   doc = DocumentBuilder(doc).SetUri("uri11").SetSchema("TypeB").Build();
-  ICING_ASSERT_OK(this->doc_store_->Put(doc));
   ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri12").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri13").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri14").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri15").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri16").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri17").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri18").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri19").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri20").Build()));
+      this->doc_store_->Put(document_util::CreateDocumentWrapper(doc)));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri12").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri13").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri14").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri15").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri16").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri17").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri18").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri19").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri20").Build())));
 
   // Ids are assigned alphabetically, so the property ids are:
   // TypeA.desiredProperty = 0
@@ -1180,8 +1200,9 @@
 
   // Check info section
   Info info;
-  ASSERT_TRUE(filesystem_.PRead(metadata_sfd.get(), &info, sizeof(Info),
-                                IntegerIndex::kInfoMetadataFileOffset));
+  ASSERT_THAT(filesystem_.PRead(metadata_sfd.get(), &info, sizeof(Info),
+                                IntegerIndex::kInfoMetadataFileOffset),
+              Eq(sizeof(Info)));
   EXPECT_THAT(info.magic, Eq(Info::kMagic));
   EXPECT_THAT(info.last_added_document_id, Eq(kInvalidDocumentId));
   EXPECT_THAT(info.num_data_threshold_for_bucket_split,
@@ -1189,8 +1210,9 @@
 
   // Check crcs section
   Crcs crcs;
-  ASSERT_TRUE(filesystem_.PRead(metadata_sfd.get(), &crcs, sizeof(Crcs),
-                                IntegerIndex::kCrcsMetadataFileOffset));
+  ASSERT_THAT(filesystem_.PRead(metadata_sfd.get(), &crcs, sizeof(Crcs),
+                                IntegerIndex::kCrcsMetadataFileOffset),
+              Eq(sizeof(Crcs)));
   // There are no storages initially, so storages_crc should be 0.
   EXPECT_THAT(crcs.component_crcs.storages_crc, Eq(0));
   EXPECT_THAT(crcs.component_crcs.info_crc,
@@ -1376,8 +1398,9 @@
   ASSERT_TRUE(metadata_sfd.is_valid());
 
   Crcs crcs;
-  ASSERT_TRUE(filesystem_.PRead(metadata_sfd.get(), &crcs, sizeof(Crcs),
-                                IntegerIndex::kCrcsMetadataFileOffset));
+  ASSERT_THAT(filesystem_.PRead(metadata_sfd.get(), &crcs, sizeof(Crcs),
+                                IntegerIndex::kCrcsMetadataFileOffset),
+              Eq(sizeof(Crcs)));
 
   // Manually corrupt all_crc
   crcs.all_crc += kCorruptedValueOffset;
@@ -1425,8 +1448,9 @@
   ASSERT_TRUE(metadata_sfd.is_valid());
 
   Info info;
-  ASSERT_TRUE(filesystem_.PRead(metadata_sfd.get(), &info, sizeof(Info),
-                                IntegerIndex::kInfoMetadataFileOffset));
+  ASSERT_THAT(filesystem_.PRead(metadata_sfd.get(), &info, sizeof(Info),
+                                IntegerIndex::kInfoMetadataFileOffset),
+              Eq(sizeof(Info)));
 
   // Modify info, but don't update the checksum. This would be similar to
   // corruption of info.
@@ -1656,49 +1680,51 @@
   // Put 11 docs of "TypeA" into the document store.
   DocumentProto doc =
       DocumentBuilder().SetKey("ns1", "uri0").SetSchema("TypeA").Build();
-  ICING_ASSERT_OK(this->doc_store_->Put(doc));
   ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri1").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri2").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri3").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri4").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri5").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri6").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri7").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri8").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri9").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri10").Build()));
+      this->doc_store_->Put(document_util::CreateDocumentWrapper(doc)));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri1").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri2").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri3").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri4").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri5").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri6").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri7").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri8").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri9").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri10").Build())));
 
   // Put 10 docs of "TypeB" into the document store.
   doc = DocumentBuilder(doc).SetUri("uri11").SetSchema("TypeB").Build();
-  ICING_ASSERT_OK(this->doc_store_->Put(doc));
   ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri12").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri13").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri14").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri15").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri16").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri17").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri18").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri19").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri20").Build()));
+      this->doc_store_->Put(document_util::CreateDocumentWrapper(doc)));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri12").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri13").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri14").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri15").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri16").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri17").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri18").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri19").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri20").Build())));
 
   {
     ICING_ASSERT_OK_AND_ASSIGN(
@@ -2040,49 +2066,51 @@
   // Put 11 docs of "TypeA" into the document store.
   DocumentProto doc =
       DocumentBuilder().SetKey("ns1", "uri0").SetSchema("TypeA").Build();
-  ICING_ASSERT_OK(this->doc_store_->Put(doc));
   ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri1").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri2").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri3").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri4").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri5").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri6").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri7").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri8").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri9").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri10").Build()));
+      this->doc_store_->Put(document_util::CreateDocumentWrapper(doc)));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri1").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri2").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri3").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri4").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri5").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri6").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri7").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri8").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri9").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri10").Build())));
 
   // Put 10 docs of "TypeB" into the document store.
   doc = DocumentBuilder(doc).SetUri("uri11").SetSchema("TypeB").Build();
-  ICING_ASSERT_OK(this->doc_store_->Put(doc));
   ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri12").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri13").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri14").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri15").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri16").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri17").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri18").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri19").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri20").Build()));
+      this->doc_store_->Put(document_util::CreateDocumentWrapper(doc)));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri12").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri13").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri14").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri15").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri16").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri17").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri18").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri19").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri20").Build())));
 
   {
     ICING_ASSERT_OK_AND_ASSIGN(
@@ -2330,27 +2358,28 @@
   // Put 11 docs of "TypeA" into the document store.
   DocumentProto doc =
       DocumentBuilder().SetKey("ns1", "uri0").SetSchema("TypeA").Build();
-  ICING_ASSERT_OK(this->doc_store_->Put(doc));
   ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri1").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri2").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri3").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri4").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri5").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri6").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri7").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri8").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri9").Build()));
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri10").Build()));
+      this->doc_store_->Put(document_util::CreateDocumentWrapper(doc)));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri1").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri2").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri3").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri4").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri5").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri6").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri7").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri8").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri9").Build())));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri10").Build())));
 
   {
     ICING_ASSERT_OK_AND_ASSIGN(
@@ -2446,7 +2475,8 @@
   // Add a new doc (docid==5) and a hit for desiredProperty. This should still
   // be placed into the wildcard integer storage.
   doc = DocumentBuilder().SetKey("ns1", "uri11").SetSchema("TypeA").Build();
-  ICING_ASSERT_OK(this->doc_store_->Put(doc));
+  ICING_ASSERT_OK(
+      this->doc_store_->Put(document_util::CreateDocumentWrapper(doc)));
   Index(integer_index.get(), desired_property, /*document_id=*/5,
         typea_desired_prop_id, /*keys=*/{12});
   EXPECT_THAT(integer_index->num_property_indices(), Eq(1));
@@ -2458,8 +2488,8 @@
 
   // Add a new doc (docid==6) and a hit for undesiredProperty. This should still
   // be placed into the wildcard integer storage.
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri12").Build()));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri12").Build())));
   Index(integer_index.get(), undesired_property, /*document_id=*/6,
         typea_undesired_prop_id, /*keys=*/{3});
   EXPECT_THAT(integer_index->num_property_indices(), Eq(1));
@@ -2472,8 +2502,8 @@
 
   // Add a new doc (docid==7) and a hit for otherProperty1. This should be given
   // its own individual storage.
-  ICING_ASSERT_OK(
-      this->doc_store_->Put(DocumentBuilder(doc).SetUri("uri13").Build()));
+  ICING_ASSERT_OK(this->doc_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder(doc).SetUri("uri13").Build())));
   Index(integer_index.get(), other_property_1, /*document_id=*/7,
         typea_other1_prop_id, /*keys=*/{3});
   EXPECT_THAT(integer_index->num_property_indices(), Eq(2));
@@ -2519,7 +2549,8 @@
                   /*num_leaf_advance_calls_main_index=*/0,
                   /*num_leaf_advance_calls_integer_index=*/1,
                   /*num_leaf_advance_calls_no_index=*/0,
-                  /*num_blocks_inspected=*/1));
+                  /*num_blocks_inspected=*/1,
+                  DocHitInfoIterator::CallStats::EmbeddingStats()));
 
   // 1st Advance().
   ICING_ASSERT_OK(iter->Advance());
@@ -2529,7 +2560,8 @@
                   /*num_leaf_advance_calls_main_index=*/0,
                   /*num_leaf_advance_calls_integer_index=*/2,
                   /*num_leaf_advance_calls_no_index=*/0,
-                  /*num_blocks_inspected=*/1));
+                  /*num_blocks_inspected=*/1,
+                  DocHitInfoIterator::CallStats::EmbeddingStats()));
 
   // 2nd Advance().
   ICING_ASSERT_OK(iter->Advance());
@@ -2539,7 +2571,8 @@
                   /*num_leaf_advance_calls_main_index=*/0,
                   /*num_leaf_advance_calls_integer_index=*/3,
                   /*num_leaf_advance_calls_no_index=*/0,
-                  /*num_blocks_inspected=*/1));
+                  /*num_blocks_inspected=*/1,
+                  DocHitInfoIterator::CallStats::EmbeddingStats()));
 
   // 3rd Advance().
   ICING_ASSERT_OK(iter->Advance());
@@ -2549,7 +2582,8 @@
                   /*num_leaf_advance_calls_main_index=*/0,
                   /*num_leaf_advance_calls_integer_index=*/4,
                   /*num_leaf_advance_calls_no_index=*/0,
-                  /*num_blocks_inspected=*/1));
+                  /*num_blocks_inspected=*/1,
+                  DocHitInfoIterator::CallStats::EmbeddingStats()));
 
   // 4th Advance().
   ICING_ASSERT_OK(iter->Advance());
@@ -2559,7 +2593,8 @@
                   /*num_leaf_advance_calls_main_index=*/0,
                   /*num_leaf_advance_calls_integer_index=*/4,
                   /*num_leaf_advance_calls_no_index=*/0,
-                  /*num_blocks_inspected=*/1));
+                  /*num_blocks_inspected=*/1,
+                  DocHitInfoIterator::CallStats::EmbeddingStats()));
 
   // 5th Advance().
   ASSERT_THAT(iter->Advance(),
@@ -2570,7 +2605,8 @@
                   /*num_leaf_advance_calls_main_index=*/0,
                   /*num_leaf_advance_calls_integer_index=*/4,
                   /*num_leaf_advance_calls_no_index=*/0,
-                  /*num_blocks_inspected=*/1));
+                  /*num_blocks_inspected=*/1,
+                  DocHitInfoIterator::CallStats::EmbeddingStats()));
 }
 
 TEST_P(IntegerIndexTest, IteratorCallStatsNonExistingProperty) {
@@ -2603,7 +2639,8 @@
                   /*num_leaf_advance_calls_main_index=*/0,
                   /*num_leaf_advance_calls_integer_index=*/0,
                   /*num_leaf_advance_calls_no_index=*/0,
-                  /*num_blocks_inspected=*/0));
+                  /*num_blocks_inspected=*/0,
+                  DocHitInfoIterator::CallStats::EmbeddingStats()));
 
   // 1st Advance().
   ASSERT_THAT(iter->Advance(),
@@ -2614,7 +2651,8 @@
                   /*num_leaf_advance_calls_main_index=*/0,
                   /*num_leaf_advance_calls_integer_index=*/0,
                   /*num_leaf_advance_calls_no_index=*/0,
-                  /*num_blocks_inspected=*/0));
+                  /*num_blocks_inspected=*/0,
+                  DocHitInfoIterator::CallStats::EmbeddingStats()));
 }
 
 INSTANTIATE_TEST_SUITE_P(

diff --git a/icing/index/property-existence-indexing-handler.cc b/icing/index/property-existence-indexing-handler.cc
index 0b4d87f..408897a 100644
--- a/icing/index/property-existence-indexing-handler.cc
+++ b/icing/index/property-existence-indexing-handler.cc

@@ -92,7 +92,8 @@
       index_.Edit(document_id, /*section_id=*/0, /*namespace_id=*/0);
   std::unordered_set<std::string> meta_tokens;
   ConstructPropertyExistenceMetaToken(
-      /*current_path=*/"", tokenized_document.document(), meta_tokens);
+      /*current_path=*/"", tokenized_document.document_wrapper().document(),
+      meta_tokens);
   for (const std::string& meta_token : meta_tokens) {
     status = editor.BufferTerm(meta_token, TermMatchType::EXACT_ONLY);
     if (!status.ok()) {

diff --git a/icing/index/property-existence-indexing-handler_test.cc b/icing/index/property-existence-indexing-handler_test.cc
index 8905bc5..208aebf 100644
--- a/icing/index/property-existence-indexing-handler_test.cc
+++ b/icing/index/property-existence-indexing-handler_test.cc

@@ -35,6 +35,7 @@
 #include "icing/index/index.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
 #include "icing/legacy/index/icing-filesystem.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/portable/platform.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/document_wrapper.pb.h"
@@ -160,14 +161,18 @@
         filesystem_.CreateDirectoryRecursively(document_store_dir_.c_str()));
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult doc_store_create_result,
-        DocumentStore::Create(&filesystem_, document_store_dir_, &fake_clock_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, document_store_dir_, &fake_clock_,
+            schema_store_.get(), feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     document_store_ = std::move(doc_store_create_result.document_store);
   }
 
@@ -219,7 +224,8 @@
                          /*lite_index_sort_size=*/1024 * 8);
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<Index> index,
-      Index::Create(options, &filesystem_, &icing_filesystem_));
+      Index::Create(options, &filesystem_, &icing_filesystem_,
+                    feature_flags_.get()));
 
   // Create a document with every property.
   DocumentProto document0 =
@@ -249,27 +255,33 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document0,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document0)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document0)));
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document1,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document1)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document1)));
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document2,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document2)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document2)));
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result0,
-      document_store_->Put(tokenized_document0.document()));
+      document_store_->Put(tokenized_document0.document_wrapper()));
   DocumentId document_id0 = put_result0.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result1,
-      document_store_->Put(tokenized_document1.document()));
+      document_store_->Put(tokenized_document1.document_wrapper()));
   DocumentId document_id1 = put_result1.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result2,
-      document_store_->Put(tokenized_document2.document()));
+      document_store_->Put(tokenized_document2.document_wrapper()));
   DocumentId document_id2 = put_result2.new_document_id;
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -321,7 +333,8 @@
                          /*lite_index_sort_size=*/1024 * 8);
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<Index> index,
-      Index::Create(options, &filesystem_, &icing_filesystem_));
+      Index::Create(options, &filesystem_, &icing_filesystem_,
+                    feature_flags_.get()));
 
   // Create a complex nested root_document with the following property paths.
   // - name
@@ -380,11 +393,13 @@
   // Handle root_document
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_root_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(root_document)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(root_document)));
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result,
-      document_store_->Put(tokenized_root_document.document()));
+      document_store_->Put(tokenized_root_document.document_wrapper()));
   DocumentId document_id = put_result.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<PropertyExistenceIndexingHandler> handler,
@@ -460,7 +475,8 @@
                          /*lite_index_sort_size=*/1024 * 8);
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<Index> index,
-      Index::Create(options, &filesystem_, &icing_filesystem_));
+      Index::Create(options, &filesystem_, &icing_filesystem_,
+                    feature_flags_.get()));
 
   // Create a document with one empty body.
   DocumentProto document0 =
@@ -486,27 +502,33 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document0,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document0)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document0)));
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document1,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document1)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document1)));
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document2,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document2)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document2)));
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result0,
-      document_store_->Put(tokenized_document0.document()));
+      document_store_->Put(tokenized_document0.document_wrapper()));
   DocumentId document_id0 = put_result0.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result1,
-      document_store_->Put(tokenized_document1.document()));
+      document_store_->Put(tokenized_document1.document_wrapper()));
   DocumentId document_id1 = put_result1.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result2,
-      document_store_->Put(tokenized_document2.document()));
+      document_store_->Put(tokenized_document2.document_wrapper()));
   DocumentId document_id2 = put_result2.new_document_id;
 
   ICING_ASSERT_OK_AND_ASSIGN(

diff --git a/icing/index/term-indexing-handler_test.cc b/icing/index/term-indexing-handler_test.cc
index 577e80b..02f8164 100644
--- a/icing/index/term-indexing-handler_test.cc
+++ b/icing/index/term-indexing-handler_test.cc

@@ -39,6 +39,7 @@
 #include "icing/index/iterator/doc-hit-info-iterator.h"
 #include "icing/index/property-existence-indexing-handler.h"
 #include "icing/legacy/index/icing-filesystem.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/portable/platform.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/document_wrapper.pb.h"
@@ -176,14 +177,18 @@
         filesystem_.CreateDirectoryRecursively(document_store_dir_.c_str()));
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult doc_store_create_result,
-        DocumentStore::Create(&filesystem_, document_store_dir_, &fake_clock_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, document_store_dir_, &fake_clock_,
+            schema_store_.get(), feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     document_store_ = std::move(doc_store_create_result.document_store);
   }
 
@@ -249,7 +254,8 @@
                          /*lite_index_sort_size=*/1024 * 8);
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<Index> index,
-      Index::Create(options, &filesystem_, &icing_filesystem_));
+      Index::Create(options, &filesystem_, &icing_filesystem_,
+                    feature_flags_.get()));
 
   DocumentProto document =
       DocumentBuilder()
@@ -261,12 +267,14 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document)));
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result,
-      document_store_->Put(tokenized_document.document()));
+      document_store_->Put(tokenized_document.document_wrapper()));
   DocumentId document_id = put_result.new_document_id;
 
   EXPECT_THAT(index->last_added_document_id(), Eq(kInvalidDocumentId));
@@ -314,7 +322,8 @@
                          /*lite_index_sort_size=*/1024 * 8);
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<Index> index,
-      Index::Create(options, &filesystem_, &icing_filesystem_));
+      Index::Create(options, &filesystem_, &icing_filesystem_,
+                    feature_flags_.get()));
 
   DocumentProto document =
       DocumentBuilder()
@@ -326,12 +335,14 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document)));
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result,
-      document_store_->Put(tokenized_document.document()));
+      document_store_->Put(tokenized_document.document_wrapper()));
   DocumentId document_id = put_result.new_document_id;
 
   EXPECT_THAT(index->last_added_document_id(), Eq(kInvalidDocumentId));
@@ -378,7 +389,8 @@
                          /*lite_index_sort_size=*/64);
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<Index> index,
-      Index::Create(options, &filesystem_, &icing_filesystem_));
+      Index::Create(options, &filesystem_, &icing_filesystem_,
+                    feature_flags_.get()));
 
   DocumentProto document0 =
       DocumentBuilder()
@@ -405,29 +417,35 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document0,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document0)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document0)));
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result0,
-      document_store_->Put(tokenized_document0.document()));
+      document_store_->Put(tokenized_document0.document_wrapper()));
   DocumentId document_id0 = put_result0.new_document_id;
 
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document1,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document1)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document1)));
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result1,
-      document_store_->Put(tokenized_document1.document()));
+      document_store_->Put(tokenized_document1.document_wrapper()));
   DocumentId document_id1 = put_result1.new_document_id;
 
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document2,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document2)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document2)));
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result2,
-      document_store_->Put(tokenized_document2.document()));
+      document_store_->Put(tokenized_document2.document_wrapper()));
   DocumentId document_id2 = put_result2.new_document_id;
   EXPECT_THAT(index->last_added_document_id(), Eq(kInvalidDocumentId));
 
@@ -526,7 +544,8 @@
                          /*lite_index_sort_size=*/64);
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<Index> index,
-      Index::Create(options, &filesystem_, &icing_filesystem_));
+      Index::Create(options, &filesystem_, &icing_filesystem_,
+                    feature_flags_.get()));
 
   DocumentProto document0 =
       DocumentBuilder()
@@ -553,29 +572,35 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document0,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document0)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document0)));
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result0,
-      document_store_->Put(tokenized_document0.document()));
+      document_store_->Put(tokenized_document0.document_wrapper()));
   DocumentId document_id0 = put_result0.new_document_id;
 
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document1,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document1)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document1)));
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result1,
-      document_store_->Put(tokenized_document1.document()));
+      document_store_->Put(tokenized_document1.document_wrapper()));
   DocumentId document_id1 = put_result1.new_document_id;
 
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document2,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document2)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document2)));
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result2,
-      document_store_->Put(tokenized_document2.document()));
+      document_store_->Put(tokenized_document2.document_wrapper()));
   DocumentId document_id2 = put_result2.new_document_id;
   EXPECT_THAT(index->last_added_document_id(), Eq(kInvalidDocumentId));
 
@@ -612,7 +637,8 @@
                            /*lite_index_sort_at_indexing=*/true,
                            /*lite_index_sort_size=*/64);
   ICING_ASSERT_OK_AND_ASSIGN(
-      index, Index::Create(options, &filesystem_, &icing_filesystem_));
+      index, Index::Create(options, &filesystem_, &icing_filesystem_,
+                           feature_flags_.get()));
 
   // Verify that the HitBuffer has been sorted after initializing with
   // sort_at_indexing enabled.

diff --git a/icing/jni/icing-search-engine-jni.cc b/icing/jni/icing-search-engine-jni.cc
index c17ea0b..dfa28b6 100644
--- a/icing/jni/icing-search-engine-jni.cc
+++ b/icing/jni/icing-search-engine-jni.cc

@@ -14,6 +14,8 @@
 
 #include <jni.h>
 
+#include <cstdint>
+#include <memory>
 #include <string_view>
 #include <utility>
 
@@ -29,8 +31,10 @@
 #include "icing/proto/schema.pb.h"
 #include "icing/proto/scoring.pb.h"
 #include "icing/proto/search.pb.h"
+#include "icing/proto/status.pb.h"
 #include "icing/proto/storage.pb.h"
 #include "icing/proto/usage.pb.h"
+#include "icing/util/clock.h"
 #include "icing/util/logging.h"
 #include <google/protobuf/message_lite.h>
 
@@ -126,10 +130,8 @@
 
   return SerializeProtoToJniByteArray(env, set_schema_result_proto);
 }
-// TODO : b/337913932 - pre-register this API once Jetpack build is dropped back
-// into g3
-JNIEXPORT jbyteArray JNICALL
-Java_com_google_android_icing_IcingSearchEngineImpl_nativeSetSchemaWithRequestProto(
+
+jbyteArray nativeSetSchemaWithRequestProto(
     JNIEnv* env, jclass clazz, jobject object,
     jbyteArray set_schema_request_bytes) {
   icing::lib::IcingSearchEngine* icing =
@@ -202,10 +204,8 @@
   return SerializeProtoToJniByteArray(env, put_result_proto);
 }
 
-JNIEXPORT jbyteArray JNICALL
-Java_com_google_android_icing_IcingSearchEngineImpl_nativeBatchPut(
-    JNIEnv* env, jclass clazz, jobject object,
-    jbyteArray put_document_request_bytes) {
+jbyteArray nativeBatchPut(JNIEnv* env, jclass clazz, jobject object,
+                          jbyteArray put_document_request_bytes) {
   icing::lib::IcingSearchEngine* icing =
       GetIcingSearchEnginePointer(env, object);
 
@@ -219,6 +219,8 @@
 
   icing::lib::BatchPutResultProto batch_put_result_proto =
       icing->BatchPut(std::move(put_document_request));
+  batch_put_result_proto.mutable_status()->set_code(
+      icing::lib::StatusProto::OK);
 
   return SerializeProtoToJniByteArray(env, batch_put_result_proto);
 }
@@ -244,6 +246,24 @@
   return SerializeProtoToJniByteArray(env, get_result_proto);
 }
 
+jbyteArray nativeBatchGet(JNIEnv* env, jclass clazz, jobject object,
+                          jbyteArray result_spec_bytes) {
+  icing::lib::IcingSearchEngine* icing =
+      GetIcingSearchEnginePointer(env, object);
+
+  icing::lib::GetResultSpecProto get_result_spec;
+  if (!ParseProtoFromJniByteArray(env, result_spec_bytes, &get_result_spec)) {
+    ICING_LOG(icing::lib::ERROR)
+        << "Failed to parse GetResultSpecProto in nativeGet";
+    return nullptr;
+  }
+
+  icing::lib::BatchGetResultProto batch_get_result_proto =
+      icing->BatchGet(std::move(get_result_spec));
+
+  return SerializeProtoToJniByteArray(env, batch_get_result_proto);
+}
+
 jbyteArray nativeReportUsage(JNIEnv* env, jclass clazz, jobject object,
                              jbyteArray usage_report_bytes) {
   icing::lib::IcingSearchEngine* icing =
@@ -295,14 +315,57 @@
   return SerializeProtoToJniByteArray(env, next_page_result_proto);
 }
 
+// TODO: b/417644758 - pre-register this method.
+JNIEXPORT jbyteArray JNICALL
+Java_com_google_android_icing_IcingSearchEngineImpl_nativeGetNextPageWithRequestProto(
+    JNIEnv* env, jclass clazz, jobject object,
+    jbyteArray get_next_page_request_bytes,
+    jlong java_to_native_start_timestamp_ms) {
+  icing::lib::IcingSearchEngine* icing =
+      GetIcingSearchEnginePointer(env, object);
+
+  const std::unique_ptr<const icing::lib::Clock> clock =
+      std::make_unique<icing::lib::Clock>();
+  int32_t java_to_native_jni_latency_ms =
+      clock->GetSystemTimeMilliseconds() - java_to_native_start_timestamp_ms;
+
+  icing::lib::GetNextPageRequestProto get_next_page_request_proto;
+  if (!ParseProtoFromJniByteArray(env, get_next_page_request_bytes,
+                                  &get_next_page_request_proto)) {
+    ICING_LOG(icing::lib::ERROR) << "Failed to parse GetNextPageRequestProto "
+                                    "in nativeGetNextPageWithRequestProto";
+    return nullptr;
+  }
+  icing::lib::SearchResultProto next_page_result_proto =
+      icing->GetNextPage(std::move(get_next_page_request_proto));
+
+  icing::lib::QueryStatsProto* query_stats =
+      next_page_result_proto.mutable_query_stats();
+  query_stats->set_java_to_native_jni_latency_ms(java_to_native_jni_latency_ms);
+  query_stats->set_native_to_java_start_timestamp_ms(
+      clock->GetSystemTimeMilliseconds());
+
+  return SerializeProtoToJniByteArray(env, next_page_result_proto);
+}
+
 void nativeInvalidateNextPageToken(JNIEnv* env, jclass clazz, jobject object,
                                    jlong next_page_token) {
   icing::lib::IcingSearchEngine* icing =
       GetIcingSearchEnginePointer(env, object);
 
   icing->InvalidateNextPageToken(next_page_token);
+}
 
-  return;
+// TODO(b/384947619) - pre-register this method.
+JNIEXPORT jbyteArray JNICALL
+Java_com_google_android_icing_IcingSearchEngineImpl_nativeHandleExpiredDocuments(
+    JNIEnv* env, jclass clazz, jobject object) {
+  icing::lib::IcingSearchEngine* icing =
+      GetIcingSearchEnginePointer(env, object);
+
+  icing::lib::HandleExpiredDocumentsResultProto result_proto =
+      icing->HandleExpiredDocuments();
+  return SerializeProtoToJniByteArray(env, result_proto);
 }
 
 // TODO(b/273591938): Change this API back to the pre-registered API.
@@ -374,6 +437,41 @@
   return SerializeProtoToJniByteArray(env, blob_result_proto);
 }
 
+// TODO : b/434206770 - pre-register this API once Jetpack build is dropped back
+// into g3
+JNIEXPORT jbyteArray JNICALL
+Java_com_google_android_icing_IcingSearchEngineImpl_nativeGetAllBlobInfos(JNIEnv* env, jclass clazz,
+                                                  jobject object) {
+  icing::lib::IcingSearchEngine* icing =
+      GetIcingSearchEnginePointer(env, object);
+
+  icing::lib::BlobProto blob_result_proto = icing->GetAllBlobInfos();
+
+  return SerializeProtoToJniByteArray(env, blob_result_proto);
+}
+
+// TODO : b/434206770 - pre-register this API once Jetpack build is dropped back
+// into g3
+JNIEXPORT jbyteArray JNICALL
+Java_com_google_android_icing_IcingSearchEngineImpl_nativePutBlobInfos(JNIEnv* env, jclass clazz,
+                                                jobject object,
+                                                jbyteArray blob_bytes) {
+  icing::lib::IcingSearchEngine* icing =
+      GetIcingSearchEnginePointer(env, object);
+
+  icing::lib::BlobProto blob_proto;
+  if (!ParseProtoFromJniByteArray(env, blob_bytes, &blob_proto)) {
+    ICING_LOG(icing::lib::ERROR)
+        << "Failed to parse BlobInfo in nativePutBlobInfos";
+    return nullptr;
+  }
+
+  icing::lib::BlobProto blob_result_proto =
+      icing->PutBlobInfos(std::move(blob_proto));
+
+  return SerializeProtoToJniByteArray(env, blob_result_proto);
+}
+
 jbyteArray nativeSearch(JNIEnv* env, jclass clazz, jobject object,
                         jbyteArray search_spec_bytes,
                         jbyteArray scoring_spec_bytes,
@@ -532,6 +630,18 @@
   return SerializeProtoToJniByteArray(env, reset_result_proto);
 }
 
+JNIEXPORT jbyteArray JNICALL
+Java_com_google_android_icing_IcingSearchEngineImpl_nativeClearAndDestroy(
+    JNIEnv* env, jclass clazz, jobject object) {
+  icing::lib::IcingSearchEngine* icing =
+      GetIcingSearchEnginePointer(env, object);
+
+  icing::lib::ResetResultProto clear_and_destroy_proto =
+      icing->ClearAndDestroy();
+
+  return SerializeProtoToJniByteArray(env, clear_and_destroy_proto);
+}
+
 jbyteArray nativeSearchSuggestions(JNIEnv* env, jclass clazz, jobject object,
                                    jbyteArray suggestion_spec_bytes) {
   icing::lib::IcingSearchEngine* icing =
@@ -623,6 +733,9 @@
       {"nativeSetSchema",
        "(Lcom/google/android/icing/IcingSearchEngineImpl;[BZ)[B",
        reinterpret_cast<void*>(nativeSetSchema)},
+      {"nativeSetSchemaWithRequestProto",
+       "(Lcom/google/android/icing/IcingSearchEngineImpl;[B)[B",
+       reinterpret_cast<void*>(nativeSetSchemaWithRequestProto)},
       {"nativeGetSchema",
        "(Lcom/google/android/icing/IcingSearchEngineImpl;)[B",
        reinterpret_cast<void*>(nativeGetSchema)},
@@ -634,15 +747,16 @@
        reinterpret_cast<void*>(nativeGetSchemaType)},
       {"nativePut", "(Lcom/google/android/icing/IcingSearchEngineImpl;[B)[B",
        reinterpret_cast<void*>(nativePut)},
-      // TODO(b/394875109): uncomment when Jetpack library is updated with this
-      // change and syned to google3.
-      // {"nativeBatchPut",
-      //  "(Lcom/google/android/icing/IcingSearchEngineImpl;[B)[B",
-      //  reinterpret_cast<void*>(nativeBatchPut)},
+      {"nativeBatchPut",
+       "(Lcom/google/android/icing/IcingSearchEngineImpl;[B)[B",
+       reinterpret_cast<void*>(nativeBatchPut)},
       {"nativeGet",
        "(Lcom/google/android/icing/IcingSearchEngineImpl;Ljava/lang/"
        "String;Ljava/lang/String;[B)[B",
        reinterpret_cast<void*>(nativeGet)},
+      {"nativeBatchGet",
+       "(Lcom/google/android/icing/IcingSearchEngineImpl;[B)[B",
+       reinterpret_cast<void*>(nativeBatchGet)},
       {"nativeReportUsage",
        "(Lcom/google/android/icing/IcingSearchEngineImpl;[B)[B",
        reinterpret_cast<void*>(nativeReportUsage)},

diff --git a/icing/join/delete-propagation-handler.cc b/icing/join/delete-propagation-handler.cc
new file mode 100644
index 0000000..b2943d4
--- /dev/null
+++ b/icing/join/delete-propagation-handler.cc

@@ -0,0 +1,161 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "icing/join/delete-propagation-handler.h"
+
+#include <cstdint>
+#include <deque>
+#include <optional>
+#include <queue>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
+#include "icing/absl_ports/canonical_errors.h"
+#include "icing/join/document-join-id-pair.h"
+#include "icing/join/qualified-id-join-index.h"
+#include "icing/schema/joinable-property.h"
+#include "icing/schema/schema-store.h"
+#include "icing/store/document-filter-data.h"
+#include "icing/store/document-id.h"
+#include "icing/store/document-store.h"
+#include "icing/util/logging.h"
+#include "icing/util/status-macros.h"
+
+namespace icing {
+namespace lib {
+
+/* static */ libtextclassifier3::StatusOr<DeletePropagationHandler>
+DeletePropagationHandler::Create(
+    const SchemaStore* schema_store,
+    const QualifiedIdJoinIndex* qualified_id_join_index,
+    DocumentStore* document_store, int64_t current_time_ms) {
+  ICING_RETURN_ERROR_IF_NULL(schema_store);
+  ICING_RETURN_ERROR_IF_NULL(qualified_id_join_index);
+  ICING_RETURN_ERROR_IF_NULL(document_store);
+
+  if (qualified_id_join_index->version() !=
+      QualifiedIdJoinIndex::Version::kV3) {
+    return absl_ports::FailedPreconditionError(
+        "Delete propagation is enabled but qualified id join index v3 is not "
+        "used.");
+  }
+
+  return DeletePropagationHandler(schema_store, qualified_id_join_index,
+                                  document_store, current_time_ms);
+}
+
+libtextclassifier3::StatusOr<std::vector<DocumentStore::DocumentMetadata>>
+DeletePropagationHandler::Handle(
+    const std::unordered_set<DocumentId>& parent_doc_ids) {
+  ICING_ASSIGN_OR_RETURN(
+      std::unordered_set<DocumentId> propagated_child_doc_ids,
+      GetPropagatedChildDocumentIds(parent_doc_ids));
+
+  // Delete all propagated child documents.
+  std::vector<DocumentStore::DocumentMetadata> deleted_doc_metadata_list;
+  for (DocumentId child_doc_id : propagated_child_doc_ids) {
+    auto deleted_doc_metadata_or = document_store_.ForceDelete(child_doc_id);
+    if (!deleted_doc_metadata_or.ok()) {
+      if (absl_ports::IsNotFound(deleted_doc_metadata_or.status())) {
+        // The child document has already been deleted or expired, so skip the
+        // error. This should not happen, but let's check and skip it just in
+        // case.
+        continue;
+      }
+
+      // Real error.
+      return std::move(deleted_doc_metadata_or).status();
+    }
+    deleted_doc_metadata_list.push_back(
+        std::move(deleted_doc_metadata_or).ValueOrDie());
+  }
+
+  return deleted_doc_metadata_list;
+}
+
+libtextclassifier3::StatusOr<std::unordered_set<DocumentId>>
+DeletePropagationHandler::GetPropagatedChildDocumentIds(
+    const std::unordered_set<DocumentId>& parent_doc_ids) {
+  std::unordered_set<DocumentId> propagated_child_doc_ids;
+
+  // BFS traverse to find all propagated child documents.
+  std::queue<DocumentId> que(
+      std::deque(parent_doc_ids.begin(), parent_doc_ids.end()));
+  while (!que.empty()) {
+    DocumentId doc_id_to_expand = que.front();
+    que.pop();
+
+    ICING_ASSIGN_OR_RETURN(
+        QualifiedIdJoinIndex::DocumentJoinIdPairArrayView
+            child_join_id_pairs_array_view,
+        qualified_id_join_index_.GetDocumentJoinIdPairArrayView(
+            doc_id_to_expand));
+    for (const DocumentJoinIdPair& child_join_id_pair :
+         child_join_id_pairs_array_view) {
+      if (propagated_child_doc_ids.find(child_join_id_pair.document_id()) !=
+              propagated_child_doc_ids.end() ||
+          parent_doc_ids.find(child_join_id_pair.document_id()) !=
+              parent_doc_ids.end()) {
+        // Already added into the propagated set or in the parent set (happens
+        // only when there is a cycle back to the parent or traversed document
+        // in the join relation). Skip it.
+        continue;
+      }
+
+      // Get DocumentFilterData of the child document to look up its schema type
+      // id.
+      // - Skip if the child document has been deleted, since delete propagation
+      //   should've been done to all its children when deleting it previously.
+      // - Otherwise, we have to handle this child document and propagate to the
+      //   grandchildren when it is alive OR expired but not deleted.
+      std::optional<DocumentFilterData> child_filter_data =
+          document_store_.GetNonDeletedDocumentFilterData(
+              child_join_id_pair.document_id());
+      if (!child_filter_data) {
+        // The child document has been deleted. Skip.
+        continue;
+      }
+
+      auto metadata_or = schema_store_.GetJoinablePropertyMetadata(
+          child_filter_data->schema_type_id(),
+          child_join_id_pair.joinable_property_id());
+      if (!metadata_or.ok() || metadata_or.ValueOrDie() == nullptr) {
+        // This shouldn't happen because we've validated it during indexing and
+        // only put valid DocumentJoinIdPair into qualified id join index.
+        // Log and skip it.
+        ICING_LOG(ERROR) << "Failed to get metadata for schema type id "
+                         << child_filter_data->schema_type_id()
+                         << ", joinable property id "
+                         << static_cast<int>(
+                                child_join_id_pair.joinable_property_id());
+        continue;
+      }
+      const JoinablePropertyMetadata* metadata = metadata_or.ValueOrDie();
+
+      if (metadata->value_type == JoinableConfig::ValueType::QUALIFIED_ID &&
+          metadata->delete_propagation_type ==
+              JoinableConfig::DeletePropagationType::PROPAGATE_FROM) {
+        propagated_child_doc_ids.insert(child_join_id_pair.document_id());
+        que.push(child_join_id_pair.document_id());
+      }
+    }
+  }
+
+  return propagated_child_doc_ids;
+}
+
+}  // namespace lib
+}  // namespace icing

diff --git a/icing/join/delete-propagation-handler.h b/icing/join/delete-propagation-handler.h
new file mode 100644
index 0000000..c8d3394
--- /dev/null
+++ b/icing/join/delete-propagation-handler.h

@@ -0,0 +1,85 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef ICING_JOIN_DELETE_PROPAGATION_HANDLER_H_
+#define ICING_JOIN_DELETE_PROPAGATION_HANDLER_H_
+
+#include <cstdint>
+#include <unordered_set>
+#include <vector>
+
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
+#include "icing/join/qualified-id-join-index.h"
+#include "icing/schema/schema-store.h"
+#include "icing/store/document-id.h"
+#include "icing/store/document-store.h"
+
+namespace icing {
+namespace lib {
+
+// A class to handle delete propagation related logics.
+class DeletePropagationHandler {
+ public:
+  // Creates a DeletePropagationHandler.
+  //
+  // Returns:
+  //   - DeletePropagationHandler on success.
+  //   - FAILED_PRECONDITION_ERROR if any pointer is nullptr or qualified id
+  //     join index version is not v3.
+  static libtextclassifier3::StatusOr<DeletePropagationHandler> Create(
+      const SchemaStore* schema_store,
+      const QualifiedIdJoinIndex* qualified_id_join_index,
+      DocumentStore* document_store, int64_t current_time_ms);
+
+  // Handles delete propagation for the given parent document ids.
+  //
+  // Note: this function DOES NOT handle parent_doc_ids' deletion. Instead, it
+  //   only deletes the propagated child documents.
+  //
+  // Returns:
+  //   - A vector of deleted child document metadata on success.
+  //   - INTERNAL_ERROR on any I/O errors.
+  libtextclassifier3::StatusOr<std::vector<DocumentStore::DocumentMetadata>>
+  Handle(const std::unordered_set<DocumentId>& parent_doc_ids);
+
+ private:
+  explicit DeletePropagationHandler(
+      const SchemaStore* schema_store,
+      const QualifiedIdJoinIndex* qualified_id_join_index,
+      DocumentStore* document_store, int64_t current_time_ms)
+      : schema_store_(*schema_store),
+        qualified_id_join_index_(*qualified_id_join_index),
+        document_store_(*document_store),
+        current_time_ms_(current_time_ms) {}
+
+  // Helper function to get all child document ids propagated from the given
+  // parent document ids via join relations with delete propagation enabled.
+  //
+  // Returns:
+  //   - A set of propagated child document ids on success.
+  //   - INTERNAL_ERROR on any I/O errors.
+  libtextclassifier3::StatusOr<std::unordered_set<DocumentId>>
+  GetPropagatedChildDocumentIds(
+      const std::unordered_set<DocumentId>& parent_doc_ids);
+
+  const SchemaStore& schema_store_;                      // Does not own.
+  const QualifiedIdJoinIndex& qualified_id_join_index_;  // Does not own.
+  DocumentStore& document_store_;                        // Does not own.
+  int64_t current_time_ms_;
+};
+
+}  // namespace lib
+}  // namespace icing
+
+#endif  // ICING_JOIN_DELETE_PROPAGATION_HANDLER_H_

diff --git a/icing/join/delete-propagation-handler_test.cc b/icing/join/delete-propagation-handler_test.cc
new file mode 100644
index 0000000..1746937
--- /dev/null
+++ b/icing/join/delete-propagation-handler_test.cc

@@ -0,0 +1,935 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "icing/join/delete-propagation-handler.h"
+
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <string>
+#include <unordered_set>
+#include <utility>
+
+#include "icing/text_classifier/lib3/utils/base/status.h"
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+#include "icing/document-builder.h"
+#include "icing/feature-flags.h"
+#include "icing/file/filesystem.h"
+#include "icing/file/portable-file-backed-proto-log.h"
+#include "icing/join/document-join-id-pair.h"
+#include "icing/join/qualified-id-join-index-impl-v2.h"
+#include "icing/join/qualified-id-join-index-impl-v3.h"
+#include "icing/join/qualified-id-join-index.h"
+#include "icing/join/qualified-id-join-indexing-handler.h"
+#include "icing/portable/gzip_stream.h"
+#include "icing/portable/platform.h"
+#include "icing/proto/document.pb.h"
+#include "icing/proto/document_wrapper.pb.h"
+#include "icing/proto/schema.pb.h"
+#include "icing/proto/scoring.pb.h"
+#include "icing/proto/search.pb.h"
+#include "icing/schema-builder.h"
+#include "icing/schema/schema-store.h"
+#include "icing/store/document-id.h"
+#include "icing/store/document-store.h"
+#include "icing/testing/common-matchers.h"
+#include "icing/testing/fake-clock.h"
+#include "icing/testing/test-data.h"
+#include "icing/testing/test-feature-flags.h"
+#include "icing/testing/tmp-directory.h"
+#include "icing/tokenization/language-segmenter-factory.h"
+#include "icing/tokenization/language-segmenter.h"
+#include "icing/util/icu-data-file-helper.h"
+#include "icing/util/status-macros.h"
+#include "icing/util/tokenized-document.h"
+#include "unicode/uloc.h"
+
+namespace icing {
+namespace lib {
+
+namespace {
+
+using ::testing::ElementsAre;
+using ::testing::Eq;
+using ::testing::HasSubstr;
+using ::testing::IsEmpty;
+using ::testing::IsTrue;
+using ::testing::Ne;
+using ::testing::UnorderedElementsAre;
+
+class DeletePropagationHandlerTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    feature_flags_ = std::make_unique<FeatureFlags>(GetTestFeatureFlags());
+    fake_clock_.SetSystemTimeMilliseconds(123);
+
+    test_dir_ = GetTestTempDir() + "/icing_delete_propagation_handler_test";
+    ASSERT_THAT(filesystem_.CreateDirectoryRecursively(test_dir_.c_str()),
+                IsTrue());
+
+    schema_store_dir_ = test_dir_ + "/schema_store";
+    doc_store_dir_ = test_dir_ + "/doc_store";
+    qualified_id_join_index_dir_ = test_dir_ + "/qualified_id_join_index";
+
+    if (!IsCfStringTokenization() && !IsReverseJniTokenization()) {
+      ICING_ASSERT_OK(
+          // File generated via icu_data_file rule in //icing/BUILD.
+          icu_data_file_helper::SetUpIcuDataFile(
+              GetTestFilePath("icing/icu.dat")));
+    }
+
+    language_segmenter_factory::SegmenterOptions options(ULOC_US);
+    ICING_ASSERT_OK_AND_ASSIGN(
+        lang_segmenter_,
+        language_segmenter_factory::Create(std::move(options)));
+
+    ASSERT_THAT(
+        filesystem_.CreateDirectoryRecursively(schema_store_dir_.c_str()),
+        IsTrue());
+    ICING_ASSERT_OK_AND_ASSIGN(
+        schema_store_, SchemaStore::Create(&filesystem_, schema_store_dir_,
+                                           &fake_clock_, feature_flags_.get()));
+
+    SchemaProto schema =
+        SchemaBuilder()
+            .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+                PropertyConfigBuilder()
+                    .SetName("Name")
+                    .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                    .SetCardinality(CARDINALITY_OPTIONAL)))
+            .AddType(
+                SchemaTypeConfigBuilder()
+                    .SetType("Email")
+                    .AddProperty(PropertyConfigBuilder()
+                                     .SetName("subject")
+                                     .SetDataTypeString(TERM_MATCH_EXACT,
+                                                        TOKENIZER_PLAIN)
+                                     .SetCardinality(CARDINALITY_OPTIONAL))
+                    .AddProperty(PropertyConfigBuilder()
+                                     .SetName("sender")
+                                     .SetDataTypeJoinableString(
+                                         JOINABLE_VALUE_TYPE_QUALIFIED_ID)
+                                     .SetCardinality(CARDINALITY_OPTIONAL)))
+            .AddType(
+                SchemaTypeConfigBuilder()
+                    .SetType("Message")
+                    .AddProperty(PropertyConfigBuilder()
+                                     .SetName("content")
+                                     .SetDataTypeString(TERM_MATCH_EXACT,
+                                                        TOKENIZER_PLAIN)
+                                     .SetCardinality(CARDINALITY_OPTIONAL))
+                    .AddProperty(PropertyConfigBuilder()
+                                     .SetName("receiver")
+                                     .SetDataTypeJoinableString(
+                                         JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                         DELETE_PROPAGATION_TYPE_NONE)
+                                     .SetCardinality(CARDINALITY_OPTIONAL))
+                    .AddProperty(PropertyConfigBuilder()
+                                     .SetName("reporter")
+                                     .SetDataTypeJoinableString(
+                                         JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                         DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                                     .SetCardinality(CARDINALITY_OPTIONAL))
+                    .AddProperty(PropertyConfigBuilder()
+                                     .SetName("sender")
+                                     .SetDataTypeJoinableString(
+                                         JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                         DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                                     .SetCardinality(CARDINALITY_OPTIONAL)))
+            .AddType(
+                SchemaTypeConfigBuilder()
+                    .SetType("Label")
+                    .AddProperty(PropertyConfigBuilder()
+                                     .SetName("text")
+                                     .SetDataTypeString(TERM_MATCH_EXACT,
+                                                        TOKENIZER_PLAIN)
+                                     .SetCardinality(CARDINALITY_OPTIONAL))
+                    .AddProperty(PropertyConfigBuilder()
+                                     .SetName("object")
+                                     .SetDataTypeJoinableString(
+                                         JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                         DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                                     .SetCardinality(CARDINALITY_OPTIONAL))
+                    .AddProperty(PropertyConfigBuilder()
+                                     .SetName("softLink")
+                                     .SetDataTypeJoinableString(
+                                         JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                         DELETE_PROPAGATION_TYPE_NONE)
+                                     .SetCardinality(CARDINALITY_OPTIONAL)))
+
+            .Build();
+    ASSERT_THAT(schema_store_->SetSchema(
+                    schema, /*ignore_errors_and_delete_documents=*/false),
+                IsOk());
+
+    ASSERT_THAT(filesystem_.CreateDirectoryRecursively(doc_store_dir_.c_str()),
+                IsTrue());
+    ICING_ASSERT_OK_AND_ASSIGN(
+        DocumentStore::CreateResult create_result,
+        DocumentStore::Create(
+            &filesystem_, doc_store_dir_, &fake_clock_, schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
+    doc_store_ = std::move(create_result.document_store);
+
+    ICING_ASSERT_OK_AND_ASSIGN(
+        qualified_id_join_index_,
+        QualifiedIdJoinIndexImplV3::Create(
+            filesystem_, qualified_id_join_index_dir_, *feature_flags_));
+  }
+
+  void TearDown() override {
+    qualified_id_join_index_.reset();
+    doc_store_.reset();
+    schema_store_.reset();
+    lang_segmenter_.reset();
+
+    filesystem_.DeleteDirectoryRecursively(test_dir_.c_str());
+  }
+
+  libtextclassifier3::StatusOr<DocumentId> PutAndIndexDocument(
+      DocumentProto document) {
+    ICING_ASSIGN_OR_RETURN(
+        TokenizedDocument tokenized_document,
+        TokenizedDocument::Create(
+            schema_store_.get(), lang_segmenter_.get(),
+            /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+            std::move(document)));
+    ICING_ASSIGN_OR_RETURN(
+        DocumentStore::PutResult put_result,
+        doc_store_->Put(tokenized_document.document_wrapper()));
+
+    ICING_ASSIGN_OR_RETURN(
+        std::unique_ptr<QualifiedIdJoinIndexingHandler> handler,
+        QualifiedIdJoinIndexingHandler::Create(&fake_clock_, doc_store_.get(),
+                                               qualified_id_join_index_.get(),
+                                               feature_flags_.get()));
+    ICING_RETURN_IF_ERROR(
+        handler->Handle(tokenized_document, put_result.new_document_id,
+                        put_result.old_document_id, /*recovery_mode=*/false,
+                        /*put_document_stats=*/nullptr));
+    return put_result.new_document_id;
+  }
+
+  std::unique_ptr<FeatureFlags> feature_flags_;
+  Filesystem filesystem_;
+  std::string test_dir_;
+  std::string schema_store_dir_;
+  std::string doc_store_dir_;
+  std::string qualified_id_join_index_dir_;
+
+  std::unique_ptr<LanguageSegmenter> lang_segmenter_;
+  std::unique_ptr<SchemaStore> schema_store_;
+  std::unique_ptr<DocumentStore> doc_store_;
+  std::unique_ptr<QualifiedIdJoinIndex> qualified_id_join_index_;
+
+  FakeClock fake_clock_;
+};
+
+TEST_F(DeletePropagationHandlerTest, Create_shouldFailWithNullptr) {
+  EXPECT_THAT(DeletePropagationHandler::Create(
+                  /*schema_store=*/nullptr, qualified_id_join_index_.get(),
+                  doc_store_.get(), fake_clock_.GetSystemTimeMilliseconds()),
+              StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
+
+  EXPECT_THAT(DeletePropagationHandler::Create(
+                  schema_store_.get(), /*qualified_id_join_index=*/nullptr,
+                  doc_store_.get(), fake_clock_.GetSystemTimeMilliseconds()),
+              StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
+
+  EXPECT_THAT(
+      DeletePropagationHandler::Create(
+          schema_store_.get(), qualified_id_join_index_.get(),
+          /*document_store=*/nullptr, fake_clock_.GetSystemTimeMilliseconds()),
+      StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
+}
+
+TEST_F(DeletePropagationHandlerTest,
+       Create_shouldFailWithJoinIndexVersionNotV3) {
+  std::string join_index_v2_dir = test_dir_ + "/join_index_v2";
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<QualifiedIdJoinIndex> join_index_v2,
+      QualifiedIdJoinIndexImplV2::Create(filesystem_,
+                                         std::move(join_index_v2_dir),
+                                         /*pre_mapping_fbv=*/false));
+  EXPECT_THAT(DeletePropagationHandler::Create(
+                  schema_store_.get(), join_index_v2.get(), doc_store_.get(),
+                  fake_clock_.GetSystemTimeMilliseconds()),
+              StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION,
+                       HasSubstr("Delete propagation is enabled but qualified "
+                                 "id join index v3 is not used")));
+}
+
+TEST_F(DeletePropagationHandlerTest, Handle) {
+  DocumentProto person = DocumentBuilder()
+                             .SetKey("pkg$db/namespace", "person")
+                             .SetSchema("Person")
+                             .AddStringProperty("Name", "Alice")
+                             .Build();
+
+  // "sender" (joinable property id 2) has
+  // DELETE_PROPAGATION_TYPE_PROPAGATE_FROM.
+  DocumentProto message =
+      DocumentBuilder()
+          .SetKey("pkg$db/namespace", "message")
+          .SetSchema("Message")
+          .AddStringProperty("content", "test content")
+          .AddStringProperty("sender", "pkg$db/namespace#person")
+          .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId person_doc_id,
+                             PutAndIndexDocument(person));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId message_doc_id,
+                             PutAndIndexDocument(message));
+  // Sanity check.
+  ASSERT_THAT(
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(person_doc_id),
+      IsOkAndHolds(ElementsAre(
+          DocumentJoinIdPair(message_doc_id, /*joinable_property_id=*/2))));
+
+  // Deleting the parent document should propagate the delete to its child
+  // document via the joinable property with
+  // DELETE_PROPAGATION_TYPE_PROPAGATE_FROM.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DeletePropagationHandler delete_propagation_handler,
+      DeletePropagationHandler::Create(
+          schema_store_.get(), qualified_id_join_index_.get(), doc_store_.get(),
+          fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_THAT(
+      delete_propagation_handler.Handle(/*parent_doc_ids=*/{person_doc_id}),
+      IsOkAndHolds(UnorderedElementsAre(EqualsDocumentMetadata(
+          "Message", "pkg$db/namespace", "message", message_doc_id))));
+  EXPECT_THAT(doc_store_->GetNonDeletedDocumentFilterData(message_doc_id),
+              Eq(std::nullopt));
+}
+
+TEST_F(DeletePropagationHandlerTest,
+       Handle_shouldNotPropagateToChildDocumentsWithPropagateDeleteDisabled) {
+  DocumentProto person = DocumentBuilder()
+                             .SetKey("pkg$db/namespace", "person")
+                             .SetSchema("Person")
+                             .AddStringProperty("Name", "Alice")
+                             .Build();
+
+  // "receiver" (joinable property id 0) has DELETE_PROPAGATION_TYPE_NONE.
+  DocumentProto message =
+      DocumentBuilder()
+          .SetKey("pkg$db/namespace", "message")
+          .SetSchema("Message")
+          .AddStringProperty("content", "test content")
+          .AddStringProperty("receiver", "pkg$db/namespace#person")
+          .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId person_doc_id,
+                             PutAndIndexDocument(person));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId message_doc_id,
+                             PutAndIndexDocument(message));
+  // Sanity check.
+  ASSERT_THAT(
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(person_doc_id),
+      IsOkAndHolds(ElementsAre(
+          DocumentJoinIdPair(message_doc_id, /*joinable_property_id=*/0))));
+
+  // Deleting the parent document should not propagate the delete to its child
+  // document via the joinable property with DELETE_PROPAGATION_TYPE_NONE.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DeletePropagationHandler delete_propagation_handler,
+      DeletePropagationHandler::Create(
+          schema_store_.get(), qualified_id_join_index_.get(), doc_store_.get(),
+          fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_THAT(
+      delete_propagation_handler.Handle(/*parent_doc_ids=*/{person_doc_id}),
+      IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(doc_store_->GetNonDeletedDocumentFilterData(message_doc_id),
+              Ne(std::nullopt));
+}
+
+TEST_F(DeletePropagationHandlerTest,
+       Handle_shouldNotPropagateToNonJoinableChildDocuments) {
+  DocumentProto person = DocumentBuilder()
+                             .SetKey("pkg$db/namespace", "person")
+                             .SetSchema("Person")
+                             .AddStringProperty("Name", "Alice")
+                             .Build();
+
+  // Put person's qualified id string in a non-joinable property.
+  DocumentProto message =
+      DocumentBuilder()
+          .SetKey("pkg$db/namespace", "message")
+          .SetSchema("Message")
+          .AddStringProperty("content", "pkg$db/namespace#person")
+          .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId person_doc_id,
+                             PutAndIndexDocument(person));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId message_doc_id,
+                             PutAndIndexDocument(message));
+  // Sanity check.
+  ASSERT_THAT(
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(person_doc_id),
+      IsOkAndHolds(IsEmpty()));
+
+  // Deleting the parent document should not propagate the delete to
+  // non-joinable child documents.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DeletePropagationHandler delete_propagation_handler,
+      DeletePropagationHandler::Create(
+          schema_store_.get(), qualified_id_join_index_.get(), doc_store_.get(),
+          fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_THAT(
+      delete_propagation_handler.Handle(/*parent_doc_ids=*/{person_doc_id}),
+      IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(doc_store_->GetNonDeletedDocumentFilterData(message_doc_id),
+              Ne(std::nullopt));
+}
+
+TEST_F(DeletePropagationHandlerTest, Handle_propagateViaMultipleProperties) {
+  DocumentProto person = DocumentBuilder()
+                             .SetKey("pkg$db/namespace", "person")
+                             .SetSchema("Person")
+                             .AddStringProperty("Name", "Alice")
+                             .Build();
+
+  // - "sender" (joinable property id 2) has
+  //   DELETE_PROPAGATION_TYPE_PROPAGATE_FROM.
+  // - "receiver" (joinable property id 0) has DELETE_PROPAGATION_TYPE_NONE.
+  DocumentProto message =
+      DocumentBuilder()
+          .SetKey("pkg$db/namespace", "message")
+          .SetSchema("Message")
+          .AddStringProperty("content", "test content")
+          .AddStringProperty("sender", "pkg$db/namespace#person")
+          .AddStringProperty("receiver", "pkg$db/namespace#person")
+          .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId person_doc_id,
+                             PutAndIndexDocument(person));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId message_doc_id,
+                             PutAndIndexDocument(message));
+  // Sanity check.
+  ASSERT_THAT(
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(person_doc_id),
+      IsOkAndHolds(ElementsAre(
+          DocumentJoinIdPair(message_doc_id, /*joinable_property_id=*/0),
+          DocumentJoinIdPair(message_doc_id, /*joinable_property_id=*/2))));
+
+  // Deleting the parent document should propagate the delete to its child
+  // document when there is at least one property with
+  // DELETE_PROPAGATION_TYPE_PROPAGATE_FROM.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DeletePropagationHandler delete_propagation_handler,
+      DeletePropagationHandler::Create(
+          schema_store_.get(), qualified_id_join_index_.get(), doc_store_.get(),
+          fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_THAT(
+      delete_propagation_handler.Handle(/*parent_doc_ids=*/{person_doc_id}),
+      IsOkAndHolds(UnorderedElementsAre(EqualsDocumentMetadata(
+          "Message", "pkg$db/namespace", "message", message_doc_id))));
+  EXPECT_THAT(doc_store_->GetNonDeletedDocumentFilterData(message_doc_id),
+              Eq(std::nullopt));
+}
+
+TEST_F(DeletePropagationHandlerTest, Handle_propagateToMultipleChildren) {
+  DocumentProto person = DocumentBuilder()
+                             .SetKey("pkg$db/namespace", "person")
+                             .SetSchema("Person")
+                             .AddStringProperty("Name", "Alice")
+                             .Build();
+
+  // "sender" (joinable property id 2) has
+  // DELETE_PROPAGATION_TYPE_PROPAGATE_FROM.
+  DocumentProto message1 =
+      DocumentBuilder()
+          .SetKey("pkg$db/namespace", "message1")
+          .SetSchema("Message")
+          .AddStringProperty("content", "test content")
+          .AddStringProperty("sender", "pkg$db/namespace#person")
+          .Build();
+  DocumentProto message2 =
+      DocumentBuilder()
+          .SetKey("pkg$db/namespace", "message2")
+          .SetSchema("Message")
+          .AddStringProperty("content", "test content")
+          .AddStringProperty("sender", "pkg$db/namespace#person")
+          .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId person_doc_id,
+                             PutAndIndexDocument(person));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId message1_doc_id,
+                             PutAndIndexDocument(message1));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId message2_doc_id,
+                             PutAndIndexDocument(message2));
+  // Sanity check.
+  ASSERT_THAT(
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(person_doc_id),
+      IsOkAndHolds(ElementsAre(
+          DocumentJoinIdPair(message1_doc_id, /*joinable_property_id=*/2),
+          DocumentJoinIdPair(message2_doc_id, /*joinable_property_id=*/2))));
+
+  // Deleting the parent document should propagate the delete to all of its
+  // child documents via the joinable property with
+  // DELETE_PROPAGATION_TYPE_PROPAGATE_FROM.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DeletePropagationHandler delete_propagation_handler,
+      DeletePropagationHandler::Create(
+          schema_store_.get(), qualified_id_join_index_.get(), doc_store_.get(),
+          fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_THAT(
+      delete_propagation_handler.Handle(/*parent_doc_ids=*/{person_doc_id}),
+      IsOkAndHolds(UnorderedElementsAre(
+          EqualsDocumentMetadata("Message", "pkg$db/namespace", "message1",
+                                 message1_doc_id),
+          EqualsDocumentMetadata("Message", "pkg$db/namespace", "message2",
+                                 message2_doc_id))));
+  EXPECT_THAT(doc_store_->GetNonDeletedDocumentFilterData(message1_doc_id),
+              Eq(std::nullopt));
+  EXPECT_THAT(doc_store_->GetNonDeletedDocumentFilterData(message2_doc_id),
+              Eq(std::nullopt));
+}
+
+TEST_F(DeletePropagationHandlerTest, Handle_propagateFromMultipleProperties) {
+  DocumentProto person1 = DocumentBuilder()
+                              .SetKey("pkg$db/namespace", "person1")
+                              .SetSchema("Person")
+                              .AddStringProperty("Name", "Alice")
+                              .Build();
+  DocumentProto person2 = DocumentBuilder()
+                              .SetKey("pkg$db/namespace", "person2")
+                              .SetSchema("Person")
+                              .AddStringProperty("Name", "Bob")
+                              .Build();
+
+  // "sender" (joinable property id 2) and "reporter" (joinable property id 1)
+  // have DELETE_PROPAGATION_TYPE_PROPAGATE_FROM.
+  DocumentProto message =
+      DocumentBuilder()
+          .SetKey("pkg$db/namespace", "message")
+          .SetSchema("Message")
+          .AddStringProperty("content", "test content")
+          .AddStringProperty("sender", "pkg$db/namespace#person1")
+          .AddStringProperty("reporter", "pkg$db/namespace#person2")
+          .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId person1_doc_id,
+                             PutAndIndexDocument(person1));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId person2_doc_id,
+                             PutAndIndexDocument(person2));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId message_doc_id,
+                             PutAndIndexDocument(message));
+  // Sanity check.
+  ASSERT_THAT(
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(person1_doc_id),
+      IsOkAndHolds(ElementsAre(
+          DocumentJoinIdPair(message_doc_id, /*joinable_property_id=*/2))));
+  ASSERT_THAT(
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(person2_doc_id),
+      IsOkAndHolds(ElementsAre(
+          DocumentJoinIdPair(message_doc_id, /*joinable_property_id=*/1))));
+
+  // message document should be propagated to be deleted from both person1 and
+  // person2 via "sender" and "reporter" properties respectively.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DeletePropagationHandler delete_propagation_handler,
+      DeletePropagationHandler::Create(
+          schema_store_.get(), qualified_id_join_index_.get(), doc_store_.get(),
+          fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_THAT(delete_propagation_handler.Handle(
+                  /*parent_doc_ids=*/{person1_doc_id, person2_doc_id}),
+              IsOkAndHolds(UnorderedElementsAre(EqualsDocumentMetadata(
+                  "Message", "pkg$db/namespace", "message", message_doc_id))));
+  EXPECT_THAT(doc_store_->GetNonDeletedDocumentFilterData(message_doc_id),
+              Eq(std::nullopt));
+}
+
+TEST_F(DeletePropagationHandlerTest, Handle_propagateToGrandChildren) {
+  // Create the following relations:
+  //
+  //                         ("object") - label1
+  //                        /
+  //             message1 <-
+  //            /           \
+  //      ("sender")         ("softLink") - label2
+  //          /
+  // person <-
+  //          \
+  //      ("receiver")       ("object") - label3
+  //            \           /
+  //             message2 <-
+  //                        \
+  //                         ("softLink") - label4
+  //
+  // Note: "sender" and "object" have DELETE_PROPAGATION_TYPE_PROPAGATE_FROM,
+  //       while "receiver" and "softLink" have DELETE_PROPAGATION_TYPE_NONE.
+  DocumentProto person = DocumentBuilder()
+                             .SetKey("pkg$db/namespace", "person")
+                             .SetSchema("Person")
+                             .AddStringProperty("Name", "Alice")
+                             .Build();
+
+  DocumentProto message1 =
+      DocumentBuilder()
+          .SetKey("pkg$db/namespace", "message1")
+          .SetSchema("Message")
+          .AddStringProperty("content", "test content")
+          .AddStringProperty("sender", "pkg$db/namespace#person")
+          .Build();
+  DocumentProto message2 =
+      DocumentBuilder()
+          .SetKey("pkg$db/namespace", "message2")
+          .SetSchema("Message")
+          .AddStringProperty("content", "test content")
+          .AddStringProperty("receiver", "pkg$db/namespace#person")
+          .Build();
+
+  DocumentProto label1 =
+      DocumentBuilder()
+          .SetKey("pkg$db/namespace", "label1")
+          .SetSchema("Label")
+          .AddStringProperty("text", " label1")
+          .AddStringProperty("object", "pkg$db/namespace#message1")
+          .Build();
+  DocumentProto label2 =
+      DocumentBuilder()
+          .SetKey("pkg$db/namespace", "label2")
+          .SetSchema("Label")
+          .AddStringProperty("text", " label2")
+          .AddStringProperty("softLink", "pkg$db/namespace#message1")
+          .Build();
+  DocumentProto label3 =
+      DocumentBuilder()
+          .SetKey("pkg$db/namespace", "label3")
+          .SetSchema("Label")
+          .AddStringProperty("text", " label3")
+          .AddStringProperty("object", "pkg$db/namespace#message2")
+          .Build();
+  DocumentProto label4 =
+      DocumentBuilder()
+          .SetKey("pkg$db/namespace", "label4")
+          .SetSchema("Label")
+          .AddStringProperty("text", " label4")
+          .AddStringProperty("softLink", "pkg$db/namespace#message2")
+          .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId person_doc_id,
+                             PutAndIndexDocument(person));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId message1_doc_id,
+                             PutAndIndexDocument(message1));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId message2_doc_id,
+                             PutAndIndexDocument(message2));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId label1_doc_id,
+                             PutAndIndexDocument(label1));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId label2_doc_id,
+                             PutAndIndexDocument(label2));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId label3_doc_id,
+                             PutAndIndexDocument(label3));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId label4_doc_id,
+                             PutAndIndexDocument(label4));
+
+  // - For children with type "Message", only message1 should be propagated to
+  //   be deleted.
+  // - For grand children with type "Label", only label1 should be propagated to
+  //   be deleted.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DeletePropagationHandler delete_propagation_handler,
+      DeletePropagationHandler::Create(
+          schema_store_.get(), qualified_id_join_index_.get(), doc_store_.get(),
+          fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_THAT(
+      delete_propagation_handler.Handle(/*parent_doc_ids=*/{person_doc_id}),
+      IsOkAndHolds(UnorderedElementsAre(
+          EqualsDocumentMetadata("Message", "pkg$db/namespace", "message1",
+                                 message1_doc_id),
+          EqualsDocumentMetadata("Label", "pkg$db/namespace", "label1",
+                                 label1_doc_id))));
+
+  // message1 and label1 should be deleted, while message2 and label2/3/4
+  // should not be deleted.
+  EXPECT_THAT(doc_store_->GetNonDeletedDocumentFilterData(message1_doc_id),
+              Eq(std::nullopt));
+  EXPECT_THAT(doc_store_->GetNonDeletedDocumentFilterData(label1_doc_id),
+              Eq(std::nullopt));
+
+  EXPECT_THAT(doc_store_->GetNonDeletedDocumentFilterData(message2_doc_id),
+              Ne(std::nullopt));
+  EXPECT_THAT(doc_store_->GetNonDeletedDocumentFilterData(label2_doc_id),
+              Ne(std::nullopt));
+  EXPECT_THAT(doc_store_->GetNonDeletedDocumentFilterData(label3_doc_id),
+              Ne(std::nullopt));
+  EXPECT_THAT(doc_store_->GetNonDeletedDocumentFilterData(label4_doc_id),
+              Ne(std::nullopt));
+}
+
+TEST_F(DeletePropagationHandlerTest, Handle_cycleReference) {
+  // Create the following relations:
+  //
+  // label1 <- label2 <- label3
+  //   |                   ^
+  //   |                   |
+  //   +-------------------+
+  DocumentProto label1 =
+      DocumentBuilder()
+          .SetKey("pkg$db/namespace", "label1")
+          .SetSchema("Label")
+          .AddStringProperty("text", " label1")
+          .AddStringProperty("object", "pkg$db/namespace#label3")
+          .Build();
+  DocumentProto label2 =
+      DocumentBuilder()
+          .SetKey("pkg$db/namespace", "label2")
+          .SetSchema("Label")
+          .AddStringProperty("text", " label2")
+          .AddStringProperty("object", "pkg$db/namespace#label1")
+          .Build();
+  DocumentProto label3 =
+      DocumentBuilder()
+          .SetKey("pkg$db/namespace", "label3")
+          .SetSchema("Label")
+          .AddStringProperty("text", " label3")
+          .AddStringProperty("object", "pkg$db/namespace#label2")
+          .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId label1_doc_id,
+                             PutAndIndexDocument(label1));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId label2_doc_id,
+                             PutAndIndexDocument(label2));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId label3_doc_id,
+                             PutAndIndexDocument(label3));
+
+  // label1 should get children = [label2_doc_id].
+  ASSERT_THAT(
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(label1_doc_id),
+      IsOkAndHolds(ElementsAre(
+          DocumentJoinIdPair(label2_doc_id, /*joinable_property_id=*/0))));
+  // label2 should get children = [label3_doc_id].
+  ASSERT_THAT(
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(label2_doc_id),
+      IsOkAndHolds(ElementsAre(
+          DocumentJoinIdPair(label3_doc_id, /*joinable_property_id=*/0))));
+  // label3 should get children = [label1_doc_id]
+  ASSERT_THAT(
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(label3_doc_id),
+      IsOkAndHolds(
+          ElementsAre(DocumentJoinIdPair(label1_doc_id,
+                                         /*joinable_property_id=*/0))));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DeletePropagationHandler delete_propagation_handler,
+      DeletePropagationHandler::Create(
+          schema_store_.get(), qualified_id_join_index_.get(), doc_store_.get(),
+          fake_clock_.GetSystemTimeMilliseconds()));
+
+  // Handle {label1_doc_id}:
+  // - Propagate to label2_doc_id from label1_doc_id.
+  // - Propagate to label3_doc_id from label2_doc_id.
+  // - When trying to propage label3_doc_id to its children =
+  //   [label1_doc_id]:
+  //   - label1_doc_id is already deleted, so it should not be propagated
+  //     again.
+  //   - There should be no infinite propagation loop.
+  EXPECT_THAT(
+      delete_propagation_handler.Handle(/*parent_doc_ids=*/{label1_doc_id}),
+      IsOkAndHolds(UnorderedElementsAre(
+          EqualsDocumentMetadata("Label", "pkg$db/namespace", "label2",
+                                 label2_doc_id),
+          EqualsDocumentMetadata("Label", "pkg$db/namespace", "label3",
+                                 label3_doc_id))));
+  EXPECT_THAT(doc_store_->GetNonDeletedDocumentFilterData(label2_doc_id),
+              Eq(std::nullopt));
+  EXPECT_THAT(doc_store_->GetNonDeletedDocumentFilterData(label3_doc_id),
+              Eq(std::nullopt));
+}
+
+TEST_F(DeletePropagationHandlerTest, Handle_selfCycleReference) {
+  DocumentProto label =
+      DocumentBuilder()
+          .SetKey("pkg$db/namespace", "label")
+          .SetSchema("Label")
+          .AddStringProperty("text", " label")
+          .AddStringProperty("object", "pkg$db/namespace#label")
+          .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId label_doc_id,
+                             PutAndIndexDocument(label));
+
+  // Sanity check.
+  ASSERT_THAT(
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(label_doc_id),
+      IsOkAndHolds(ElementsAre(
+          DocumentJoinIdPair(label_doc_id, /*joinable_property_id=*/0))));
+
+  // Handle {label_doc_id}: should delete nothing since label_doc_id is already
+  // in the deleted set. Also there should be no infinite propagation loop.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DeletePropagationHandler delete_propagation_handler,
+      DeletePropagationHandler::Create(
+          schema_store_.get(), qualified_id_join_index_.get(), doc_store_.get(),
+          fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_THAT(
+      delete_propagation_handler.Handle(/*parent_doc_ids=*/{label_doc_id}),
+      IsOkAndHolds(IsEmpty()));
+}
+
+TEST_F(DeletePropagationHandlerTest, Handle_shouldPropagateToExpiredDocuments) {
+  constexpr int64_t kCreationTimestampMs = 1000;
+  constexpr int64_t kShortTtlMs = 1000;
+  constexpr int64_t kLongTtlMs = 10000;
+
+  // Create the following relations:
+  //
+  // label1 <- label2 <- label3
+  DocumentProto label1 = DocumentBuilder()
+                             .SetKey("pkg$db/namespace", "label1")
+                             .SetSchema("Label")
+                             .AddStringProperty("text", " label1")
+                             .SetCreationTimestampMs(kCreationTimestampMs)
+                             .SetTtlMs(kLongTtlMs)
+                             .Build();
+  DocumentProto label2 =
+      DocumentBuilder()
+          .SetKey("pkg$db/namespace", "label2")
+          .SetSchema("Label")
+          .AddStringProperty("text", " label2")
+          .AddStringProperty("object", "pkg$db/namespace#label1")
+          .SetCreationTimestampMs(kCreationTimestampMs)
+          .SetTtlMs(kShortTtlMs)
+          .Build();
+  DocumentProto label3 =
+      DocumentBuilder()
+          .SetKey("pkg$db/namespace", "label3")
+          .SetSchema("Label")
+          .AddStringProperty("text", " label3")
+          .AddStringProperty("object", "pkg$db/namespace#label2")
+          .SetCreationTimestampMs(kCreationTimestampMs)
+          .SetTtlMs(kLongTtlMs)
+          .Build();
+
+  fake_clock_.SetSystemTimeMilliseconds(kCreationTimestampMs);
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId label1_doc_id,
+                             PutAndIndexDocument(label1));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId label2_doc_id,
+                             PutAndIndexDocument(label2));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId label3_doc_id,
+                             PutAndIndexDocument(label3));
+
+  // Adjust the clock to expire label2.
+  int64_t current_time_ms = kCreationTimestampMs + kShortTtlMs + 100;
+  fake_clock_.SetSystemTimeMilliseconds(current_time_ms);
+
+  // Sanity check: label2 is expired (but not deleted), and label1 and label3
+  // are not expired.
+  ASSERT_THAT(
+      doc_store_->GetAliveDocumentFilterData(label1_doc_id, current_time_ms),
+      Ne(std::nullopt));
+  ASSERT_THAT(
+      doc_store_->GetAliveDocumentFilterData(label2_doc_id, current_time_ms),
+      Eq(std::nullopt));
+  ASSERT_THAT(doc_store_->GetNonDeletedDocumentFilterData(label2_doc_id),
+              Ne(std::nullopt));
+  ASSERT_THAT(
+      doc_store_->GetAliveDocumentFilterData(label3_doc_id, current_time_ms),
+      Ne(std::nullopt));
+  ASSERT_THAT(
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(label1_doc_id),
+      IsOkAndHolds(ElementsAre(
+          DocumentJoinIdPair(label2_doc_id, /*joinable_property_id=*/0))));
+  ASSERT_THAT(
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(label2_doc_id),
+      IsOkAndHolds(ElementsAre(
+          DocumentJoinIdPair(label3_doc_id, /*joinable_property_id=*/0))));
+  ASSERT_THAT(
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(label3_doc_id),
+      IsOkAndHolds(IsEmpty()));
+
+  // Handle {label1_doc_id}: should still propagate to label2_doc_id and
+  // label3_doc_id even though label2_doc_id is expired.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DeletePropagationHandler delete_propagation_handler,
+      DeletePropagationHandler::Create(
+          schema_store_.get(), qualified_id_join_index_.get(), doc_store_.get(),
+          fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_THAT(
+      delete_propagation_handler.Handle(/*parent_doc_ids=*/{label1_doc_id}),
+      IsOkAndHolds(UnorderedElementsAre(
+          EqualsDocumentMetadata("Label", "pkg$db/namespace", "label2",
+                                 label2_doc_id),
+          EqualsDocumentMetadata("Label", "pkg$db/namespace", "label3",
+                                 label3_doc_id))));
+  EXPECT_THAT(doc_store_->GetNonDeletedDocumentFilterData(label2_doc_id),
+              Eq(std::nullopt));
+  EXPECT_THAT(doc_store_->GetNonDeletedDocumentFilterData(label3_doc_id),
+              Eq(std::nullopt));
+}
+
+TEST_F(DeletePropagationHandlerTest,
+       Handle_shouldNotPropagateToDeletedDocuments) {
+  // Create the following relations:
+  //
+  // label1 <- label2
+  DocumentProto label1 = DocumentBuilder()
+                             .SetKey("pkg$db/namespace", "label1")
+                             .SetSchema("Label")
+                             .AddStringProperty("text", " label1")
+                             .Build();
+  DocumentProto label2 =
+      DocumentBuilder()
+          .SetKey("pkg$db/namespace", "label2")
+          .SetSchema("Label")
+          .AddStringProperty("text", " label2")
+          .AddStringProperty("object", "pkg$db/namespace#label1")
+          .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId label1_doc_id,
+                             PutAndIndexDocument(label1));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentId label2_doc_id,
+                             PutAndIndexDocument(label2));
+
+  // Delete label2.
+  ICING_ASSERT_OK(doc_store_->Delete(label2_doc_id,
+                                     fake_clock_.GetSystemTimeMilliseconds()));
+
+  // Sanity check: label2 is deleted, but join data is still present.
+  ASSERT_THAT(doc_store_->GetNonDeletedDocumentFilterData(label2_doc_id),
+              Eq(std::nullopt));
+  ASSERT_THAT(
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(label1_doc_id),
+      IsOkAndHolds(ElementsAre(
+          DocumentJoinIdPair(label2_doc_id, /*joinable_property_id=*/0))));
+  ASSERT_THAT(
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(label2_doc_id),
+      IsOkAndHolds(IsEmpty()));
+
+  // Handle {label1_doc_id}: should not propagate to label2_doc_id.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DeletePropagationHandler delete_propagation_handler,
+      DeletePropagationHandler::Create(
+          schema_store_.get(), qualified_id_join_index_.get(), doc_store_.get(),
+          fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_THAT(
+      delete_propagation_handler.Handle(/*parent_doc_ids=*/{label1_doc_id}),
+      IsOkAndHolds(IsEmpty()));
+}
+
+}  // namespace
+}  // namespace lib
+}  // namespace icing

diff --git a/icing/join/document-dependency-processor.cc b/icing/join/document-dependency-processor.cc
new file mode 100644
index 0000000..8889734
--- /dev/null
+++ b/icing/join/document-dependency-processor.cc

@@ -0,0 +1,141 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "icing/join/document-dependency-processor.h"
+
+#include <cstdint>
+#include <string_view>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "icing/text_classifier/lib3/utils/base/status.h"
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
+#include "icing/absl_ports/canonical_errors.h"
+#include "icing/absl_ports/str_cat.h"
+#include "icing/join/qualified-id.h"
+#include "icing/proto/document.pb.h"
+#include "icing/schema/joinable-property.h"
+#include "icing/store/document-store.h"
+#include "icing/util/status-macros.h"
+#include "icing/util/timestamp-util.h"
+#include "icing/util/tokenized-document.h"
+
+namespace icing {
+namespace lib {
+
+/* static */ libtextclassifier3::StatusOr<DocumentDependencyProcessor>
+DocumentDependencyProcessor::Create(
+    const DocumentStore* document_store,
+    const std::vector<TokenizedDocument>& batch_documents_to_add,
+    int64_t current_time_ms) {
+  ICING_RETURN_ERROR_IF_NULL(document_store);
+
+  std::unordered_map<QualifiedId, int, QualifiedId::Hasher>
+      qualified_id_to_batch_idx;
+  for (int i = 0; i < batch_documents_to_add.size(); ++i) {
+    const TokenizedDocument& tokenized_document = batch_documents_to_add[i];
+
+    // Ensure that the new document is not expired.
+    int64_t expiration_timestamp_ms =
+        timestamp_util::CalculateRawExpirationTimestampMs(
+            tokenized_document.document_wrapper()
+                .document()
+                .creation_timestamp_ms(),
+            tokenized_document.document_wrapper().document().ttl_ms());
+    if (expiration_timestamp_ms <= current_time_ms) {
+      return absl_ports::InvalidArgumentError("The new document is expired.");
+    }
+
+    QualifiedId qualified_id(
+        tokenized_document.document_wrapper().document().namespace_(),
+        tokenized_document.document_wrapper().document().uri());
+    qualified_id_to_batch_idx.insert({std::move(qualified_id), i});
+  }
+
+  return DocumentDependencyProcessor(document_store, batch_documents_to_add,
+                                     std::move(qualified_id_to_batch_idx),
+                                     current_time_ms);
+}
+
+libtextclassifier3::Status DocumentDependencyProcessor::Evaluate() {
+  // Validate the dependencies and construct the dependent graph.
+  for (int i = 0; i < batch_documents_to_add_.size(); ++i) {
+    const TokenizedDocument& tokenized_document = batch_documents_to_add_[i];
+
+    // Iterate through all qualified id joinable properties of the tokenized
+    // document.
+    for (const JoinableProperty<std::string_view>& dep_qualified_id_prop :
+         tokenized_document.qualified_id_join_properties()) {
+      if (dep_qualified_id_prop.metadata.delete_propagation_type ==
+          JoinableConfig::DeletePropagationType::NONE) {
+        // If delete propagation is NONE, then the referenced documents are
+        // not required to be present, so skip the check for this joinable
+        // property.
+        continue;
+      }
+
+      // Otherwise, the referenced documents are required to be present.
+      // For each of the qualified id string:
+      // - Validate and check it should match a document in either the same
+      //   batch of new documents to add or the document store.
+      // - Add the dependency document id (out of the batch) into result.
+      for (std::string_view dep_qualified_id_str :
+           dep_qualified_id_prop.values) {
+        ICING_RETURN_IF_ERROR(ValidateDependency(dep_qualified_id_str));
+      }
+    }
+  }
+  return libtextclassifier3::Status::OK;
+}
+
+libtextclassifier3::Status DocumentDependencyProcessor::ValidateDependency(
+    std::string_view dep_qualified_id_str) const {
+  if (dep_qualified_id_str.empty()) {
+    // Allow empty qualified id.
+    return libtextclassifier3::Status::OK;
+  }
+
+  // Attempt to parse the qualified id string.
+  auto dep_qualified_id_or = QualifiedId::Parse(dep_qualified_id_str);
+  if (!dep_qualified_id_or.ok()) {
+    // Incorrect format of qualified id string. Return INVALID_ARGUMENT_ERROR
+    // for unsatisfied dependency.
+    return absl_ports::InvalidArgumentError(absl_ports::StrCat(
+        "Invalid qualified id string: ", dep_qualified_id_str));
+  }
+  QualifiedId dep_qualified_id = std::move(dep_qualified_id_or).ValueOrDie();
+
+  // Case 1: check if the dependency document is in the same batch of new
+  //         documents.
+  auto itr = qualified_id_to_batch_idx_.find(dep_qualified_id);
+  if (itr != qualified_id_to_batch_idx_.end()) {
+    // We've already validated that the document in the batch is not expired, so
+    // we don't need to check it again here.
+    return libtextclassifier3::Status::OK;
+  }
+
+  // Case 2: check if the dependency document is alive in the document store.
+  if (!document_store_.IsDocumentAlive(dep_qualified_id.name_space(),
+                                       dep_qualified_id.uri(),
+                                       current_time_ms_)) {
+    return absl_ports::InvalidArgumentError(absl_ports::StrCat(
+        "A dependency document is not alive: ", dep_qualified_id_str));
+  }
+
+  return libtextclassifier3::Status::OK;
+}
+
+}  // namespace lib
+}  // namespace icing

diff --git a/icing/join/document-dependency-processor.h b/icing/join/document-dependency-processor.h
new file mode 100644
index 0000000..01e63f8
--- /dev/null
+++ b/icing/join/document-dependency-processor.h

@@ -0,0 +1,100 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef ICING_JOIN_DOCUMENT_DEPENDENCY_PROCESSOR_H_
+#define ICING_JOIN_DOCUMENT_DEPENDENCY_PROCESSOR_H_
+
+#include <cstdint>
+#include <string_view>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "icing/text_classifier/lib3/utils/base/status.h"
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
+#include "icing/join/qualified-id.h"
+#include "icing/proto/document.pb.h"
+#include "icing/store/document-store.h"
+#include "icing/util/tokenized-document.h"
+
+namespace icing {
+namespace lib {
+
+// This class evaluates the dependency of the documents to be added. Currently,
+// dependencies are defined solely by delete propagation.
+class DocumentDependencyProcessor {
+ public:
+  // Creates a DocumentDependencyProcessor for the given batch of documents to
+  // add.
+  //
+  // Returns:
+  //   - A DocumentDependencyProcessor on success.
+  //   - INVALID_ARGUMENT_ERROR if any of the document in the batch is expired.
+  static libtextclassifier3::StatusOr<DocumentDependencyProcessor> Create(
+      const DocumentStore* document_store,
+      const std::vector<TokenizedDocument>& batch_documents_to_add,
+      int64_t current_time_ms);
+
+  // Evaluates the document dependencies. For each document in the batch, its
+  // dependencies (parent documents with delete propagation enabled in the
+  // schema) must be alive in either the same batch of new documents or the
+  // document store.
+  //
+  // Returns:
+  //   - OK on success.
+  //   - INVALID_ARGUMENT_ERROR if the validation fails, e.g. any of the
+  //     dependencies (referenced parent documents) are not present.
+  //   - Any error from document store or schema store.
+  libtextclassifier3::Status Evaluate();
+
+ private:
+  explicit DocumentDependencyProcessor(
+      const DocumentStore* document_store,
+      const std::vector<TokenizedDocument>& batch_documents_to_add,
+      std::unordered_map<QualifiedId, int, QualifiedId::Hasher>
+          qualified_id_to_batch_idx,
+      int64_t current_time_ms)
+      : document_store_(*document_store),
+        batch_documents_to_add_(batch_documents_to_add),
+        qualified_id_to_batch_idx_(std::move(qualified_id_to_batch_idx)),
+        current_time_ms_(current_time_ms) {}
+
+  // Helper function to validate a dependency's qualified id string:
+  // - Is valid or not. Note that empty qualified id string is allowed and will
+  //   be skipped.
+  // - Satisfies the dependency: matches a document in either the same batch of
+  //   new documents to add or the document store.
+  //
+  // Returns:
+  //   - OK on success.
+  //   - INVALID_ARGUMENT_ERROR if dep_qualified_id_str is invalid or the
+  //     document referenced by the qualified id is not present.
+  //   - Any error from document store.
+  libtextclassifier3::Status ValidateDependency(
+      std::string_view dep_qualified_id_str) const;
+
+  const DocumentStore& document_store_;
+  const std::vector<TokenizedDocument>& batch_documents_to_add_;
+
+  // A map for mapping qualified id to the index of batch_documents_to_add_.
+  std::unordered_map<QualifiedId, int, QualifiedId::Hasher>
+      qualified_id_to_batch_idx_;
+
+  int64_t current_time_ms_;
+};
+
+}  // namespace lib
+}  // namespace icing
+
+#endif  // ICING_JOIN_DOCUMENT_DEPENDENCY_PROCESSOR_H_

diff --git a/icing/join/document-dependency-processor_test.cc b/icing/join/document-dependency-processor_test.cc
new file mode 100644
index 0000000..16243fd
--- /dev/null
+++ b/icing/join/document-dependency-processor_test.cc

@@ -0,0 +1,948 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "icing/join/document-dependency-processor.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "icing/text_classifier/lib3/utils/base/status.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+#include "icing/document-builder.h"
+#include "icing/feature-flags.h"
+#include "icing/file/filesystem.h"
+#include "icing/file/portable-file-backed-proto-log.h"
+#include "icing/portable/gzip_stream.h"
+#include "icing/portable/platform.h"
+#include "icing/proto/document.pb.h"
+#include "icing/proto/document_wrapper.pb.h"
+#include "icing/proto/schema.pb.h"
+#include "icing/proto/scoring.pb.h"
+#include "icing/proto/search.pb.h"
+#include "icing/schema-builder.h"
+#include "icing/schema/schema-store.h"
+#include "icing/store/document-store.h"
+#include "icing/testing/common-matchers.h"
+#include "icing/testing/fake-clock.h"
+#include "icing/testing/test-data.h"
+#include "icing/testing/test-feature-flags.h"
+#include "icing/testing/tmp-directory.h"
+#include "icing/tokenization/language-segmenter-factory.h"
+#include "icing/tokenization/language-segmenter.h"
+#include "icing/util/document-util.h"
+#include "icing/util/icu-data-file-helper.h"
+#include "icing/util/tokenized-document.h"
+#include "unicode/uloc.h"
+
+namespace icing {
+namespace lib {
+
+namespace {
+
+using ::testing::HasSubstr;
+using ::testing::IsTrue;
+
+class DocumentDependencyProcessorTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    feature_flags_ = std::make_unique<FeatureFlags>(GetTestFeatureFlags());
+    ASSERT_THAT(feature_flags_->enable_repeated_field_joins(), IsTrue());
+
+    test_dir_ = GetTestTempDir() + "/document_dependency_processor_test";
+    ASSERT_THAT(filesystem_.CreateDirectoryRecursively(test_dir_.c_str()),
+                IsTrue());
+
+    schema_store_dir_ = test_dir_ + "/schema_store";
+    doc_store_dir_ = test_dir_ + "/doc_store";
+
+    if (!IsCfStringTokenization() && !IsReverseJniTokenization()) {
+      ICING_ASSERT_OK(
+          // File generated via icu_data_file rule in //icing/BUILD.
+          icu_data_file_helper::SetUpIcuDataFile(
+              GetTestFilePath("icing/icu.dat")));
+    }
+
+    language_segmenter_factory::SegmenterOptions options(ULOC_US);
+    ICING_ASSERT_OK_AND_ASSIGN(
+        lang_segmenter_,
+        language_segmenter_factory::Create(std::move(options)));
+
+    ASSERT_THAT(
+        filesystem_.CreateDirectoryRecursively(schema_store_dir_.c_str()),
+        IsTrue());
+    ICING_ASSERT_OK_AND_ASSIGN(
+        schema_store_, SchemaStore::Create(&filesystem_, schema_store_dir_,
+                                           &fake_clock_, feature_flags_.get()));
+
+    SchemaProto schema =
+        SchemaBuilder()
+            .AddType(SchemaTypeConfigBuilder().SetType("Person").AddProperty(
+                PropertyConfigBuilder()
+                    .SetName("Name")
+                    .SetDataTypeString(TERM_MATCH_PREFIX, TOKENIZER_PLAIN)
+                    .SetCardinality(CARDINALITY_OPTIONAL)))
+            .AddType(
+                SchemaTypeConfigBuilder()
+                    .SetType("Email")
+                    .AddProperty(PropertyConfigBuilder()
+                                     .SetName("receiver")
+                                     .SetDataTypeJoinableString(
+                                         JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                         DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                                     .SetCardinality(CARDINALITY_REPEATED))
+                    .AddProperty(PropertyConfigBuilder()
+                                     .SetName("sender")
+                                     .SetDataTypeJoinableString(
+                                         JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                         DELETE_PROPAGATION_TYPE_NONE)
+                                     .SetCardinality(CARDINALITY_REPEATED)))
+            .AddType(
+                SchemaTypeConfigBuilder()
+                    .SetType("Label")
+                    .AddProperty(PropertyConfigBuilder()
+                                     .SetName("target")
+                                     .SetDataTypeJoinableString(
+                                         JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                         DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                                     .SetCardinality(CARDINALITY_REPEATED))
+                    .AddProperty(PropertyConfigBuilder()
+                                     .SetName("softTarget")
+                                     .SetDataTypeJoinableString(
+                                         JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                         DELETE_PROPAGATION_TYPE_NONE)
+                                     .SetCardinality(CARDINALITY_REPEATED)))
+
+            .Build();
+    ASSERT_THAT(schema_store_->SetSchema(
+                    schema, /*ignore_errors_and_delete_documents=*/false),
+                IsOk());
+
+    ASSERT_THAT(filesystem_.CreateDirectoryRecursively(doc_store_dir_.c_str()),
+                IsTrue());
+    ICING_ASSERT_OK_AND_ASSIGN(
+        DocumentStore::CreateResult create_result,
+        DocumentStore::Create(
+            &filesystem_, doc_store_dir_, &fake_clock_, schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
+    doc_store_ = std::move(create_result.document_store);
+  }
+
+  void TearDown() override {
+    doc_store_.reset();
+    schema_store_.reset();
+    lang_segmenter_.reset();
+
+    filesystem_.DeleteDirectoryRecursively(test_dir_.c_str());
+  }
+
+  std::unique_ptr<FeatureFlags> feature_flags_;
+  Filesystem filesystem_;
+  FakeClock fake_clock_;
+  std::string test_dir_;
+  std::string schema_store_dir_;
+  std::string doc_store_dir_;
+
+  std::unique_ptr<LanguageSegmenter> lang_segmenter_;
+  std::unique_ptr<SchemaStore> schema_store_;
+  std::unique_ptr<DocumentStore> doc_store_;
+};
+
+TEST_F(DocumentDependencyProcessorTest,
+       Create_alreadyExpiredDocumentShouldFail) {
+  fake_clock_.SetSystemTimeMilliseconds(500);
+
+  DocumentProto email1 = DocumentBuilder()
+                             .SetCreationTimestampMs(100)
+                             .SetTtlMs(400)
+                             .SetKey("namespace", "email")
+                             .SetSchema("Email")
+                             .Build();
+  DocumentProto email2 = DocumentBuilder()
+                             .SetCreationTimestampMs(100)
+                             .SetTtlMs(300)
+                             .SetKey("namespace", "email")
+                             .SetSchema("Email")
+                             .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_email1,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(), email1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_email2,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(), email2));
+
+  std::vector<TokenizedDocument> batch_documents_to_add1;
+  batch_documents_to_add1.push_back(std::move(tokenized_doc_email1));
+
+  EXPECT_THAT(DocumentDependencyProcessor::Create(
+                  doc_store_.get(), batch_documents_to_add1,
+                  /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()),
+              StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT,
+                       HasSubstr("The new document is expired.")));
+
+  std::vector<TokenizedDocument> batch_documents_to_add2;
+  batch_documents_to_add2.push_back(std::move(tokenized_doc_email2));
+
+  EXPECT_THAT(DocumentDependencyProcessor::Create(
+                  doc_store_.get(), batch_documents_to_add2,
+                  /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()),
+              StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT,
+                       HasSubstr("The new document is expired.")));
+}
+
+TEST_F(DocumentDependencyProcessorTest, Evaluate) {
+  // This is a general test case to evaluate a batch of documents with all the
+  // possible scenarios:
+  // - Replace existing documents (i.e. replace documents that were previously
+  //   added into the document store).
+  // - Add new documents.
+  //
+  // And both of them have dependencies on existing, replaced and new documents.
+
+  // Create person1, person2, email1, label1, label2, label3 with the following
+  // relation:
+  //
+  // person1        person2 --> email1        label1 --> label2 --> label3
+  DocumentProto person1 = DocumentBuilder()
+                              .SetKey("namespace", "person1")
+                              .SetSchema("Person")
+                              .AddStringProperty("Name", "Alice")
+                              .Build();
+  DocumentProto person2 = DocumentBuilder()
+                              .SetKey("namespace", "person2")
+                              .SetSchema("Person")
+                              .AddStringProperty("Name", "Bob")
+                              .Build();
+  DocumentProto email1 = DocumentBuilder()
+                             .SetKey("namespace", "email1")
+                             .SetSchema("Email")
+                             .AddStringProperty("receiver", "namespace#person2")
+                             .Build();
+  DocumentProto label1 = DocumentBuilder()
+                             .SetKey("namespace", "label1")
+                             .SetSchema("Label")
+                             .Build();
+  DocumentProto label2 = DocumentBuilder()
+                             .SetKey("namespace", "label2")
+                             .SetSchema("Label")
+                             .AddStringProperty("target", "namespace#label1")
+                             .Build();
+  DocumentProto label3 = DocumentBuilder()
+                             .SetKey("namespace", "label3")
+                             .SetSchema("Label")
+                             .AddStringProperty("target", "namespace#label2")
+                             .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_person1,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          person1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_person2,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          person2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_email1,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(), email1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_label1,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(), label1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_label2,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(), label2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_label3,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(), label3));
+
+  std::vector<TokenizedDocument> batch_documents_to_add1;
+  batch_documents_to_add1.push_back(std::move(tokenized_doc_person1));
+  batch_documents_to_add1.push_back(std::move(tokenized_doc_person2));
+  batch_documents_to_add1.push_back(std::move(tokenized_doc_email1));
+  batch_documents_to_add1.push_back(std::move(tokenized_doc_label1));
+  batch_documents_to_add1.push_back(std::move(tokenized_doc_label2));
+  batch_documents_to_add1.push_back(std::move(tokenized_doc_label3));
+
+  // Evaluate all of them together should succeed.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentDependencyProcessor processor1,
+      DocumentDependencyProcessor::Create(
+          doc_store_.get(), batch_documents_to_add1,
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_THAT(processor1.Evaluate(), IsOk());
+
+  // Put them into the document store.
+  ICING_ASSERT_OK(doc_store_->Put(
+      batch_documents_to_add1[0].document_wrapper()));  // person1
+  ICING_ASSERT_OK(doc_store_->Put(
+      batch_documents_to_add1[1].document_wrapper()));  // person2
+  ICING_ASSERT_OK(doc_store_->Put(
+      batch_documents_to_add1[2].document_wrapper()));  // email1
+  ICING_ASSERT_OK(doc_store_->Put(
+      batch_documents_to_add1[3].document_wrapper()));  // label1
+  ICING_ASSERT_OK(doc_store_->Put(
+      batch_documents_to_add1[4].document_wrapper()));  // label2
+  ICING_ASSERT_OK(doc_store_->Put(
+      batch_documents_to_add1[5].document_wrapper()));  // label3
+
+  // Replace existing person2, email1, label1 and add new label4, label5 to make
+  // the following relation:
+  // (person1, label2, label3 are not changed).
+  //
+  // person1 -------+              +--------> label4 <--------+
+  // (unchanged)    |              |          (NEW)           |
+  //                v              |             |            |
+  //              email1 -----> label5           |         label2 -----> label3
+  //              (REPLACED)     (NEW)           |         (unchanged)
+  //                ^              ^             |            ^
+  //                |              |             v            |
+  // person2 -------+              +--------- label1 ---------+
+  // (REPLACED)                               (REPLACED)
+  DocumentProto person2_to_replace = DocumentBuilder()
+                                         .SetKey("namespace", "person2")
+                                         .SetSchema("Person")
+                                         .AddStringProperty("Name", "Robert")
+                                         .Build();
+  DocumentProto email1_to_replace =
+      DocumentBuilder()
+          .SetKey("namespace", "email1")
+          .SetSchema("Email")
+          .AddStringProperty("receiver", "namespace#person1",
+                             "namespace#person2")
+          .Build();
+  DocumentProto label1_to_replace =
+      DocumentBuilder()
+          .SetKey("namespace", "label1")
+          .SetSchema("Label")
+          .AddStringProperty("target", "namespace#label4")
+          .Build();
+  DocumentProto label4 =
+      DocumentBuilder()
+          .SetKey("namespace", "label4")
+          .SetSchema("Label")
+          .AddStringProperty("target", "namespace#label2", "namespace#label5")
+          .Build();
+  DocumentProto label5 =
+      DocumentBuilder()
+          .SetKey("namespace", "label5")
+          .SetSchema("Label")
+          .AddStringProperty("target", "namespace#label1", "namespace#email1")
+          .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_person2_to_replace,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          person2_to_replace));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_email1_to_replace,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          email1_to_replace));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_label1_to_replace,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          label1_to_replace));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_label4,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(), label4));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_label5,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(), label5));
+
+  std::vector<TokenizedDocument> batch_documents_to_add2;
+  batch_documents_to_add2.push_back(
+      std::move(tokenized_doc_person2_to_replace));
+  batch_documents_to_add2.push_back(std::move(tokenized_doc_email1_to_replace));
+  batch_documents_to_add2.push_back(std::move(tokenized_doc_label1_to_replace));
+  batch_documents_to_add2.push_back(std::move(tokenized_doc_label4));
+  batch_documents_to_add2.push_back(std::move(tokenized_doc_label5));
+
+  // Evaluate all of them together should succeed.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentDependencyProcessor processor2,
+      DocumentDependencyProcessor::Create(
+          doc_store_.get(), batch_documents_to_add2,
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_THAT(processor2.Evaluate(), IsOk());
+}
+
+TEST_F(DocumentDependencyProcessorTest,
+       Evaluate_singleDocumentWithoutDependency) {
+  // Set the current time to 0.
+  fake_clock_.SetSystemTimeMilliseconds(0);
+
+  // Create a person document with expiration timestamp 2000.
+  DocumentProto person = DocumentBuilder()
+                             .SetCreationTimestampMs(0)
+                             .SetTtlMs(2000)
+                             .SetKey("namespace", "person")
+                             .SetSchema("Person")
+                             .AddStringProperty("Name", "Alice")
+                             .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_person,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(), person));
+
+  std::vector<TokenizedDocument> batch_documents_to_add1;
+  batch_documents_to_add1.push_back(std::move(tokenized_doc_person));
+
+  // Evaluate person should succeed.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentDependencyProcessor processor,
+      DocumentDependencyProcessor::Create(
+          doc_store_.get(), batch_documents_to_add1,
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_THAT(processor.Evaluate(), IsOk());
+}
+
+TEST_F(DocumentDependencyProcessorTest,
+       Evaluate_nonExistentReferencedDocumentShouldFail) {
+  DocumentProto email = DocumentBuilder()
+                            .SetKey("namespace", "email")
+                            .SetSchema("Email")
+                            .AddStringProperty("receiver", "namespace#person")
+                            .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_email,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(), email));
+
+  std::vector<TokenizedDocument> batch_documents_to_add;
+  batch_documents_to_add.push_back(std::move(tokenized_doc_email));
+
+  // Evaluate should fail since email's referenced document ("namespace#person")
+  // with delete propagation enabled doesn't exist.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentDependencyProcessor processor,
+      DocumentDependencyProcessor::Create(
+          doc_store_.get(), batch_documents_to_add,
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_THAT(processor.Evaluate(),
+              StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT,
+                       HasSubstr("A dependency document is not alive")));
+}
+
+TEST_F(DocumentDependencyProcessorTest,
+       Evaluate_invalidReferencedQualifiedIdShouldFail) {
+  DocumentProto email =
+      DocumentBuilder()
+          .SetKey("namespace", "email")
+          .SetSchema("Email")
+          .AddStringProperty("receiver", "invalid_qualified_id")
+          .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_email,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(), email));
+
+  std::vector<TokenizedDocument> batch_documents_to_add;
+  batch_documents_to_add.push_back(std::move(tokenized_doc_email));
+
+  // Evaluate should fail since email contains an invalid qualified id in a
+  // joinable property with delete propagation enabled.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentDependencyProcessor processor,
+      DocumentDependencyProcessor::Create(
+          doc_store_.get(), batch_documents_to_add,
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_THAT(processor.Evaluate(),
+              StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT,
+                       HasSubstr("Invalid qualified id string")));
+}
+
+TEST_F(DocumentDependencyProcessorTest,
+       Evaluate_emptyQualifiedIdStringShouldSucceed) {
+  DocumentProto email = DocumentBuilder()
+                            .SetKey("namespace", "email")
+                            .SetSchema("Email")
+                            .AddStringProperty("receiver", "")
+                            .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_email,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(), email));
+
+  std::vector<TokenizedDocument> batch_documents_to_add;
+  batch_documents_to_add.push_back(std::move(tokenized_doc_email));
+
+  // Evaluate should succeed since empty qualified id string is allowed.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentDependencyProcessor processor,
+      DocumentDependencyProcessor::Create(
+          doc_store_.get(), batch_documents_to_add,
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_THAT(processor.Evaluate(), IsOk());
+}
+
+TEST_F(DocumentDependencyProcessorTest,
+       Evaluate_allReferencedDocumentsInBatchShouldSucceed) {
+  // Create person1, person2, email with the following relation:
+  //
+  // person1 -------+
+  //                |
+  //                v
+  //              email
+  //                ^
+  //                |
+  // person2 -------+
+  //
+  // (email has 2 parent documents person1 and person2)
+  DocumentProto person1 = DocumentBuilder()
+                              .SetKey("namespace", "person1")
+                              .SetSchema("Person")
+                              .AddStringProperty("Name", "Alice")
+                              .Build();
+  DocumentProto person2 = DocumentBuilder()
+                              .SetKey("namespace", "person2")
+                              .SetSchema("Person")
+                              .AddStringProperty("Name", "Bob")
+                              .Build();
+  DocumentProto email = DocumentBuilder()
+                            .SetKey("namespace", "email")
+                            .SetSchema("Email")
+                            .AddStringProperty("receiver", "namespace#person1",
+                                               "namespace#person2")
+                            .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_person1,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          person1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_person2,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          person2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_email,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(), email));
+
+  std::vector<TokenizedDocument> batch_documents_to_add;
+  batch_documents_to_add.push_back(std::move(tokenized_doc_person1));
+  batch_documents_to_add.push_back(std::move(tokenized_doc_person2));
+  batch_documents_to_add.push_back(std::move(tokenized_doc_email));
+
+  // Evaluate all of them together should succeed.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentDependencyProcessor processor,
+      DocumentDependencyProcessor::Create(
+          doc_store_.get(), batch_documents_to_add,
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_THAT(processor.Evaluate(), IsOk());
+}
+
+TEST_F(DocumentDependencyProcessorTest,
+       Evaluate_selfReferenceInBatchShouldSucceed) {
+  // Create label document having a self reference on "target" property with
+  // delete propagation enabled.
+  DocumentProto label = DocumentBuilder()
+                            .SetKey("namespace", "label")
+                            .SetSchema("Label")
+                            .AddStringProperty("target", "namespace#label")
+                            .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_label,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(), label));
+
+  std::vector<TokenizedDocument> batch_documents_to_add;
+  batch_documents_to_add.push_back(std::move(tokenized_doc_label));
+
+  // Evaluate should succeed since the referenced document is present in the
+  // same batch of documents to add.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentDependencyProcessor processor,
+      DocumentDependencyProcessor::Create(
+          doc_store_.get(), batch_documents_to_add,
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_THAT(processor.Evaluate(), IsOk());
+}
+
+TEST_F(DocumentDependencyProcessorTest,
+       Evaluate_cycleReferenceInBatchShouldSucceed) {
+  // Create label1, label2, label3 with the following relation:
+  //
+  // label1 -> label2 -> label3
+  //   ^                   |
+  //   |                   |
+  //   +-------------------+
+  DocumentProto label1 = DocumentBuilder()
+                             .SetKey("namespace", "label1")
+                             .SetSchema("Label")
+                             .AddStringProperty("target", "namespace#label3")
+                             .Build();
+  DocumentProto label2 = DocumentBuilder()
+                             .SetKey("namespace", "label2")
+                             .SetSchema("Label")
+                             .AddStringProperty("target", "namespace#label1")
+                             .Build();
+  DocumentProto label3 = DocumentBuilder()
+                             .SetKey("namespace", "label3")
+                             .SetSchema("Label")
+                             .AddStringProperty("target", "namespace#label2")
+                             .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_label1,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(), label1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_label2,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(), label2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_label3,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(), label3));
+
+  std::vector<TokenizedDocument> batch_documents_to_add;
+  batch_documents_to_add.push_back(std::move(tokenized_doc_label1));
+  batch_documents_to_add.push_back(std::move(tokenized_doc_label2));
+  batch_documents_to_add.push_back(std::move(tokenized_doc_label3));
+
+  // Evaluate should succeed since the all referenced documents are present in
+  // the same batch of documents to add.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentDependencyProcessor processor,
+      DocumentDependencyProcessor::Create(
+          doc_store_.get(), batch_documents_to_add,
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_THAT(processor.Evaluate(), IsOk());
+}
+
+TEST_F(DocumentDependencyProcessorTest,
+       Evaluate_allReferencedDocumentsInDocumentStoreShouldSucceed) {
+  // Create person1, person2, email with the following relation:
+  //
+  // person1 -------+
+  //                |
+  //                v
+  //              email
+  //                ^
+  //                |
+  // person2 -------+
+  //
+  // (email has 2 parent documents person1 and person2)
+  DocumentProto person1 = DocumentBuilder()
+                              .SetKey("namespace", "person1")
+                              .SetSchema("Person")
+                              .AddStringProperty("Name", "Alice")
+                              .Build();
+  DocumentProto person2 = DocumentBuilder()
+                              .SetKey("namespace", "person2")
+                              .SetSchema("Person")
+                              .AddStringProperty("Name", "Bob")
+                              .Build();
+  DocumentProto email = DocumentBuilder()
+                            .SetKey("namespace", "email")
+                            .SetSchema("Email")
+                            .AddStringProperty("receiver", "namespace#person1",
+                                               "namespace#person2")
+                            .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_email,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(), email));
+
+  // Put person1, person2 into the document store.
+  ICING_ASSERT_OK(
+      doc_store_->Put(document_util::CreateDocumentWrapper(person1)));
+  ICING_ASSERT_OK(
+      doc_store_->Put(document_util::CreateDocumentWrapper(person2)));
+
+  std::vector<TokenizedDocument> batch_documents_to_add;
+  batch_documents_to_add.push_back(std::move(tokenized_doc_email));
+
+  // Evaluate email should succeed.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentDependencyProcessor processor,
+      DocumentDependencyProcessor::Create(
+          doc_store_.get(), batch_documents_to_add,
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_THAT(processor.Evaluate(), IsOk());
+}
+
+TEST_F(DocumentDependencyProcessorTest,
+       Evaluate_expiredReferencedDocumentInDocumentStoreShouldFail) {
+  fake_clock_.SetSystemTimeMilliseconds(500);
+
+  // Create person1, person2, email with the following relation:
+  //
+  // person1 -------+
+  //                |
+  //                v
+  //              email
+  //                ^
+  //                |
+  // person2 -------+
+  //
+  // (email has 2 parent documents person1 and person2)
+  DocumentProto person1 = DocumentBuilder()
+                              .SetKey("namespace", "person1")
+                              .SetTtlMs(1000)
+                              .SetCreationTimestampMs(500)
+                              .SetSchema("Person")
+                              .AddStringProperty("Name", "Alice")
+                              .Build();
+  DocumentProto person2 = DocumentBuilder()
+                              .SetKey("namespace", "person2")
+                              .SetTtlMs(100)
+                              .SetCreationTimestampMs(500)
+                              .SetSchema("Person")
+                              .AddStringProperty("Name", "Bob")
+                              .Build();
+  DocumentProto email = DocumentBuilder()
+                            .SetKey("namespace", "email")
+                            .SetSchema("Email")
+                            .AddStringProperty("receiver", "namespace#person1",
+                                               "namespace#person2")
+                            .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_email,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(), email));
+
+  // Put person1, person2 into the document store.
+  ICING_ASSERT_OK(
+      doc_store_->Put(document_util::CreateDocumentWrapper(person1)));
+  ICING_ASSERT_OK(
+      doc_store_->Put(document_util::CreateDocumentWrapper(person2)));
+
+  // Adjust the current time to make person2 expired, but person1 is still
+  // alive.
+  fake_clock_.SetSystemTimeMilliseconds(1000);
+
+  std::vector<TokenizedDocument> batch_documents_to_add;
+  batch_documents_to_add.push_back(std::move(tokenized_doc_email));
+
+  // Evaluate email should fail since one of email's referenced documents
+  // (person2) with delete propagation enabled is expired.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentDependencyProcessor processor,
+      DocumentDependencyProcessor::Create(
+          doc_store_.get(), batch_documents_to_add,
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_THAT(processor.Evaluate(),
+              StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT,
+                       HasSubstr("A dependency document is not alive")));
+}
+
+TEST_F(DocumentDependencyProcessorTest,
+       Evaluate_deletedReferencedDocumentInDocumentStoreShouldFail) {
+  // Create person1, person2, email with the following relation:
+  //
+  // person1 -------+
+  //                |
+  //                v
+  //              email
+  //                ^
+  //                |
+  // person2 -------+
+  //
+  // (email has 2 parent documents person1 and person2)
+  DocumentProto person1 = DocumentBuilder()
+                              .SetKey("namespace", "person1")
+                              .SetSchema("Person")
+                              .AddStringProperty("Name", "Alice")
+                              .Build();
+  DocumentProto person2 = DocumentBuilder()
+                              .SetKey("namespace", "person2")
+                              .SetSchema("Person")
+                              .AddStringProperty("Name", "Bob")
+                              .Build();
+  DocumentProto email = DocumentBuilder()
+                            .SetKey("namespace", "email")
+                            .SetSchema("Email")
+                            .AddStringProperty("receiver", "namespace#person1",
+                                               "namespace#person2")
+                            .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_email,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(), email));
+
+  // Put person1, person2 into the document store.
+  ICING_ASSERT_OK(
+      doc_store_->Put(document_util::CreateDocumentWrapper(person1)));
+  ICING_ASSERT_OK(
+      doc_store_->Put(document_util::CreateDocumentWrapper(person2)));
+
+  // Delete person2.
+  ICING_ASSERT_OK(doc_store_->Delete(
+      "namespace", "person2",
+      /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+
+  // Evaluate email should fail since one of email's referenced documents
+  // (person2) with delete propagation enabled is deleted.
+  std::vector<TokenizedDocument> batch_documents_to_add;
+  batch_documents_to_add.push_back(std::move(tokenized_doc_email));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentDependencyProcessor processor,
+      DocumentDependencyProcessor::Create(
+          doc_store_.get(), batch_documents_to_add,
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_THAT(processor.Evaluate(),
+              StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT,
+                       HasSubstr("A dependency document is not alive")));
+}
+
+TEST_F(DocumentDependencyProcessorTest,
+       Evaluate_withoutDeletePropagationShouldAlwaysSucceed) {
+  fake_clock_.SetSystemTimeMilliseconds(500);
+
+  // Create email document having an invalid qualified id string on "sender"
+  // property with delete propagation disabled. Evaluate should succeed since
+  // Icing will ignore invalid qualified id with delete propagation disabled.
+  DocumentProto email1 =
+      DocumentBuilder()
+          .SetKey("namespace", "email")
+          .SetSchema("Email")
+          .AddStringProperty("sender", "invalid_qualified_id")
+          .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_email1,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(), email1));
+  std::vector<TokenizedDocument> batch_documents_to_add1;
+  batch_documents_to_add1.push_back(std::move(tokenized_doc_email1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentDependencyProcessor processor1,
+      DocumentDependencyProcessor::Create(
+          doc_store_.get(), batch_documents_to_add1,
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_THAT(processor1.Evaluate(), IsOk());
+
+  // Create email document having a valid qualified id string on "sender"
+  // property with delete propagation disabled, but the referenced document
+  // doesn't exist. Evaluate should succeed since Icing will ignore non-existent
+  // referenced document with delete propagation disabled.
+  DocumentProto email2 = DocumentBuilder()
+                             .SetKey("namespace", "email")
+                             .SetSchema("Email")
+                             .AddStringProperty("sender", "namespace#person")
+                             .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_email2_1,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(), email2));
+  std::vector<TokenizedDocument> batch_documents_to_add2_1;
+  batch_documents_to_add2_1.push_back(std::move(tokenized_doc_email2_1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentDependencyProcessor processor2_1,
+      DocumentDependencyProcessor::Create(
+          doc_store_.get(), batch_documents_to_add2_1,
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_THAT(processor2_1.Evaluate(), IsOk());
+
+  // Add person document into the document store to make email2's referenced
+  // document exist.
+  DocumentProto person = DocumentBuilder()
+                             .SetTtlMs(100)
+                             .SetCreationTimestampMs(500)
+                             .SetKey("namespace", "person")
+                             .SetSchema("Person")
+                             .AddStringProperty("Name", "Test Name")
+                             .Build();
+  ICING_ASSERT_OK(
+      doc_store_->Put(document_util::CreateDocumentWrapper(person)));
+
+  // Evaluate email2 again. Should succeed.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_email2_2,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(), email2));
+  std::vector<TokenizedDocument> batch_documents_to_add2_2;
+  batch_documents_to_add2_2.push_back(std::move(tokenized_doc_email2_2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentDependencyProcessor processor2_2,
+      DocumentDependencyProcessor::Create(
+          doc_store_.get(), batch_documents_to_add2_2,
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_THAT(processor2_2.Evaluate(), IsOk());
+
+  // Evaluate email2 again with a different current time which makes person
+  // document expired. Since delete propagation is disabled, Evaluate should
+  // still succeed.
+  fake_clock_.SetSystemTimeMilliseconds(1000);
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_doc_email2_3,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(), email2));
+  std::vector<TokenizedDocument> batch_documents_to_add2_3;
+  batch_documents_to_add2_3.push_back(std::move(tokenized_doc_email2_3));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentDependencyProcessor processor2_3,
+      DocumentDependencyProcessor::Create(
+          doc_store_.get(), batch_documents_to_add2_3,
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_THAT(processor2_3.Evaluate(), IsOk());
+}
+
+}  // namespace
+
+}  // namespace lib
+}  // namespace icing

diff --git a/icing/join/document-dependent-graph.cc b/icing/join/document-dependent-graph.cc
new file mode 100644
index 0000000..6665b83
--- /dev/null
+++ b/icing/join/document-dependent-graph.cc

@@ -0,0 +1,127 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "icing/join/document-dependent-graph.h"
+
+#include <memory>
+#include <optional>
+#include <utility>
+
+#include "icing/text_classifier/lib3/utils/base/status.h"
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
+#include "icing/absl_ports/canonical_errors.h"
+#include "icing/graph/graph-interface.h"
+#include "icing/join/qualified-id-join-index.h"
+#include "icing/schema/joinable-property.h"
+#include "icing/schema/schema-store.h"
+#include "icing/store/document-filter-data.h"
+#include "icing/store/document-id.h"
+#include "icing/store/document-store.h"
+#include "icing/util/status-macros.h"
+
+namespace icing {
+namespace lib {
+
+/* static */ libtextclassifier3::StatusOr<
+    std::unique_ptr<DocumentDependentGraph>>
+DocumentDependentGraph::Create(const SchemaStore* schema_store,
+                               const DocumentStore* doc_store,
+                               const QualifiedIdJoinIndex* join_index) {
+  ICING_RETURN_ERROR_IF_NULL(schema_store);
+  ICING_RETURN_ERROR_IF_NULL(doc_store);
+  ICING_RETURN_ERROR_IF_NULL(join_index);
+
+  if (join_index->version() != QualifiedIdJoinIndex::Version::kV3) {
+    return absl_ports::InvalidArgumentError(
+        "DocumentDependentGraph only supports QualifiedIdJoinIndex version "
+        "V3.");
+  }
+
+  return std::unique_ptr<DocumentDependentGraph>(
+      new DocumentDependentGraph(schema_store, doc_store, join_index));
+}
+
+int DocumentDependentGraph::GetNumNodes() const {
+  DocumentId last_stored_doc_id = doc_store_.last_added_document_id();
+  if (last_stored_doc_id == kInvalidDocumentId) {
+    return 0;
+  }
+  // There are documents from 0 to last_stored_doc_id, so num nodes
+  // should be last_stored_doc_id + 1.
+  return last_stored_doc_id + 1;
+}
+
+libtextclassifier3::StatusOr<
+    std::unique_ptr<typename graph::GraphInterface<DocumentId>::EdgeIteratorIf>>
+DocumentDependentGraph::GetEdgesIterator(int node_id) const {
+  if (node_id < 0 || node_id >= GetNumNodes()) {
+    return absl_ports::InvalidArgumentError("Invalid node id.");
+  }
+
+  DocumentId doc_id = static_cast<DocumentId>(node_id);
+  if (!doc_store_.GetNonDeletedDocumentFilterData(doc_id).has_value()) {
+    // Return an iterator with no edge to advance to for a deleted document.
+    return std::make_unique<EdgeIterator>(
+        schema_store_, doc_store_,
+        QualifiedIdJoinIndex::DocumentJoinIdPairArrayView(/*data=*/nullptr,
+                                                          /*size=*/0));
+  }
+
+  ICING_ASSIGN_OR_RETURN(
+      QualifiedIdJoinIndex::DocumentJoinIdPairArrayView data_array_view,
+      join_index_.GetDocumentJoinIdPairArrayView(doc_id));
+  return std::make_unique<EdgeIterator>(schema_store_, doc_store_,
+                                        std::move(data_array_view));
+}
+
+libtextclassifier3::Status DocumentDependentGraph::EdgeIterator::Advance() {
+  while (++curr_idx_ < join_data_array_view_.size()) {
+    DocumentId next_doc_id = join_data_array_view_[curr_idx_].document_id();
+    JoinablePropertyId next_joinable_property_id =
+        join_data_array_view_[curr_idx_].joinable_property_id();
+
+    // Dedupe document id.
+    if (next_doc_id == curr_) {
+      continue;
+    }
+
+    // Lookup the schema type id of the next document.
+    std::optional<DocumentFilterData> next_doc_filter_data =
+        doc_store_.GetNonDeletedDocumentFilterData(next_doc_id);
+    if (!next_doc_filter_data.has_value() ||
+        next_doc_filter_data->schema_type_id() == kInvalidSchemaTypeId) {
+      continue;
+    }
+
+    // Get the joinable property metadata.
+    ICING_ASSIGN_OR_RETURN(
+        const JoinablePropertyMetadata* metadata,
+        schema_store_.GetJoinablePropertyMetadata(
+            next_doc_filter_data->schema_type_id(), next_joinable_property_id));
+
+    if (metadata->delete_propagation_type ==
+        JoinableConfig::DeletePropagationType::PROPAGATE_FROM) {
+      // Found a joinable property hit with delete propagation PROPAGATE_FROM.
+      // It means the next document is a dependent of the current document, so
+      // stop here, record the doc id and return success.
+      curr_ = next_doc_id;
+      return libtextclassifier3::Status::OK;
+    }
+  }
+
+  return absl_ports::ResourceExhaustedError("No more edges to advance to.");
+}
+
+}  // namespace lib
+}  // namespace icing

diff --git a/icing/join/document-dependent-graph.h b/icing/join/document-dependent-graph.h
new file mode 100644
index 0000000..f516400
--- /dev/null
+++ b/icing/join/document-dependent-graph.h

@@ -0,0 +1,102 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef ICING_JOIN_DOCUMENT_DEPENDENT_GRAPH_H_
+#define ICING_JOIN_DOCUMENT_DEPENDENT_GRAPH_H_
+
+#include <memory>
+#include <utility>
+
+#include "icing/text_classifier/lib3/utils/base/status.h"
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
+#include "icing/graph/graph-interface.h"
+#include "icing/join/qualified-id-join-index.h"
+#include "icing/schema/schema-store.h"
+#include "icing/store/document-id.h"
+#include "icing/store/document-store.h"
+
+namespace icing {
+namespace lib {
+
+// Document dependent graph using qualified id join index v3 as the data source.
+// This class is just an interface for the dependent graph and does not own any
+// data. Instead, it reads data from the given data sources and returns
+// dependent relations.
+//
+// Dependent graph:
+// - The nodes are the documents, and the node id type is DocumentId.
+// - The edges are directed and represent the dependencies between documents.
+//   For example, an edge A -> B means:
+//   - A is a dependency of B.
+//   - B is a dependent of A.
+//   - If A is deleted or expired, then B should be deleted or expired as well.
+class DocumentDependentGraph : public graph::GraphInterface<DocumentId> {
+ public:
+  // Creates a document dependent graph object.
+  //
+  // Returns:
+  //   - Non-null unique pointer of DocumentDependentGraph on success.
+  //   - FAILED_PRECONDITION_ERROR if any data source is null.
+  //   - INVALID_ARGUMENT_ERROR if join_index is not version V3.
+  //   - Any errors from the underlying data source.
+  static libtextclassifier3::StatusOr<std::unique_ptr<DocumentDependentGraph>>
+  Create(const SchemaStore* schema_store, const DocumentStore* doc_store,
+         const QualifiedIdJoinIndex* join_index);
+
+  int GetNumNodes() const override;
+
+  libtextclassifier3::StatusOr<std::unique_ptr<EdgeIteratorIf>>
+  GetEdgesIterator(int node_id) const override;
+
+ private:
+  class EdgeIterator : public EdgeIteratorIf {
+   public:
+    explicit EdgeIterator(
+        const SchemaStore& schema_store, const DocumentStore& doc_store,
+        QualifiedIdJoinIndex::DocumentJoinIdPairArrayView join_data_array_view)
+        : schema_store_(schema_store),
+          doc_store_(doc_store),
+          join_data_array_view_(std::move(join_data_array_view)),
+          curr_idx_(-1),
+          curr_(kInvalidDocumentId) {}
+
+    libtextclassifier3::Status Advance() override;
+
+    const DocumentId& Get() const override { return curr_; }
+
+   private:
+    const SchemaStore& schema_store_;  // Does not own.
+    const DocumentStore& doc_store_;   // Does not own.
+
+    QualifiedIdJoinIndex::DocumentJoinIdPairArrayView join_data_array_view_;
+    int curr_idx_;
+    DocumentId curr_;
+  };
+
+  explicit DocumentDependentGraph(const SchemaStore* schema_store,
+                                  const DocumentStore* doc_store,
+                                  const QualifiedIdJoinIndex* join_index)
+      : schema_store_(*schema_store),
+        doc_store_(*doc_store),
+        join_index_(*join_index) {}
+
+  const SchemaStore& schema_store_;         // Does not own.
+  const DocumentStore& doc_store_;          // Does not own.
+  const QualifiedIdJoinIndex& join_index_;  // Does not own.
+};
+
+}  // namespace lib
+}  // namespace icing
+
+#endif  // ICING_JOIN_DOCUMENT_DEPENDENT_GRAPH_H_

diff --git a/icing/join/document-dependent-graph_test.cc b/icing/join/document-dependent-graph_test.cc
new file mode 100644
index 0000000..c08e67c
--- /dev/null
+++ b/icing/join/document-dependent-graph_test.cc

@@ -0,0 +1,858 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "icing/join/document-dependent-graph.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "icing/text_classifier/lib3/utils/base/status.h"
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+#include "icing/document-builder.h"
+#include "icing/feature-flags.h"
+#include "icing/file/filesystem.h"
+#include "icing/file/portable-file-backed-proto-log.h"
+#include "icing/graph/graph-interface.h"
+#include "icing/join/document-join-id-pair.h"
+#include "icing/join/qualified-id-join-index-impl-v2.h"
+#include "icing/join/qualified-id-join-index-impl-v3.h"
+#include "icing/join/qualified-id-join-indexing-handler.h"
+#include "icing/portable/gzip_stream.h"
+#include "icing/portable/platform.h"
+#include "icing/proto/document.pb.h"
+#include "icing/proto/document_wrapper.pb.h"
+#include "icing/proto/schema.pb.h"
+#include "icing/proto/scoring.pb.h"
+#include "icing/proto/search.pb.h"
+#include "icing/schema-builder.h"
+#include "icing/schema/schema-store.h"
+#include "icing/store/document-id.h"
+#include "icing/store/document-store.h"
+#include "icing/testing/common-matchers.h"
+#include "icing/testing/fake-clock.h"
+#include "icing/testing/test-data.h"
+#include "icing/testing/test-feature-flags.h"
+#include "icing/testing/tmp-directory.h"
+#include "icing/tokenization/language-segmenter-factory.h"
+#include "icing/tokenization/language-segmenter.h"
+#include "icing/util/icu-data-file-helper.h"
+#include "icing/util/status-macros.h"
+#include "icing/util/tokenized-document.h"
+#include "unicode/uloc.h"
+
+namespace icing {
+namespace lib {
+
+namespace {
+
+using ::testing::ElementsAre;
+using ::testing::Eq;
+using ::testing::HasSubstr;
+using ::testing::IsEmpty;
+using ::testing::IsFalse;
+using ::testing::IsTrue;
+using ::testing::NotNull;
+
+class DocumentDependentGraphTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    feature_flags_ = std::make_unique<FeatureFlags>(GetTestFeatureFlags());
+    ASSERT_THAT(feature_flags_->enable_repeated_field_joins(), IsTrue());
+
+    test_dir_ = GetTestTempDir() + "/document_dependent_graph_test";
+    ASSERT_THAT(filesystem_.CreateDirectoryRecursively(test_dir_.c_str()),
+                IsTrue());
+
+    schema_store_dir_ = test_dir_ + "/schema_store";
+    doc_store_dir_ = test_dir_ + "/doc_store";
+    join_index_dir_ = test_dir_ + "/join_index";
+
+    if (!IsCfStringTokenization() && !IsReverseJniTokenization()) {
+      ICING_ASSERT_OK(
+          // File generated via icu_data_file rule in //icing/BUILD.
+          icu_data_file_helper::SetUpIcuDataFile(
+              GetTestFilePath("icing/icu.dat")));
+    }
+
+    language_segmenter_factory::SegmenterOptions options(ULOC_US);
+    ICING_ASSERT_OK_AND_ASSIGN(
+        lang_segmenter_,
+        language_segmenter_factory::Create(std::move(options)));
+
+    ASSERT_THAT(
+        filesystem_.CreateDirectoryRecursively(schema_store_dir_.c_str()),
+        IsTrue());
+    ICING_ASSERT_OK_AND_ASSIGN(
+        schema_store_, SchemaStore::Create(&filesystem_, schema_store_dir_,
+                                           &fake_clock_, feature_flags_.get()));
+
+    SchemaProto schema =
+        SchemaBuilder()
+            .AddType(
+                SchemaTypeConfigBuilder()
+                    .SetType("Label")
+                    .AddProperty(PropertyConfigBuilder()
+                                     .SetName("target")
+                                     .SetDataTypeJoinableString(
+                                         JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                         DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                                     .SetCardinality(CARDINALITY_REPEATED))
+                    .AddProperty(PropertyConfigBuilder()
+                                     .SetName("target2")
+                                     .SetDataTypeJoinableString(
+                                         JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                         DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                                     .SetCardinality(CARDINALITY_REPEATED))
+                    .AddProperty(PropertyConfigBuilder()
+                                     .SetName("softTarget")
+                                     .SetDataTypeJoinableString(
+                                         JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                                         DELETE_PROPAGATION_TYPE_NONE)
+                                     .SetCardinality(CARDINALITY_REPEATED)))
+
+            .Build();
+    ICING_ASSERT_OK(schema_store_->SetSchema(
+        schema, /*ignore_errors_and_delete_documents=*/false));
+
+    ASSERT_THAT(filesystem_.CreateDirectoryRecursively(doc_store_dir_.c_str()),
+                IsTrue());
+    ICING_ASSERT_OK_AND_ASSIGN(
+        DocumentStore::CreateResult create_result,
+        DocumentStore::Create(
+            &filesystem_, doc_store_dir_, &fake_clock_, schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
+    doc_store_ = std::move(create_result.document_store);
+
+    ICING_ASSERT_OK_AND_ASSIGN(
+        join_index_, QualifiedIdJoinIndexImplV3::Create(
+                         filesystem_, join_index_dir_, *feature_flags_));
+    ICING_ASSERT_OK_AND_ASSIGN(join_indexing_handler_,
+                               QualifiedIdJoinIndexingHandler::Create(
+                                   &fake_clock_, doc_store_.get(),
+                                   join_index_.get(), feature_flags_.get()));
+  }
+
+  void TearDown() override {
+    join_indexing_handler_.reset();
+
+    join_index_.reset();
+    doc_store_.reset();
+    schema_store_.reset();
+    lang_segmenter_.reset();
+
+    filesystem_.DeleteDirectoryRecursively(test_dir_.c_str());
+  }
+
+  // Helper function to batch add documents.
+  libtextclassifier3::Status AddDocuments(
+      std::vector<DocumentProto> documents) {
+    // Tokenize all documents.
+    std::vector<TokenizedDocument> tokenized_documents;
+    tokenized_documents.reserve(documents.size());
+    for (DocumentProto& document : documents) {
+      ICING_ASSIGN_OR_RETURN(
+          TokenizedDocument tokenized_document,
+          TokenizedDocument::Create(
+              schema_store_.get(), lang_segmenter_.get(),
+              /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+              std::move(document)));
+      tokenized_documents.push_back(std::move(tokenized_document));
+    }
+
+    // Put all documents into the document store and get document ids.
+    std::vector<DocumentStore::PutResult> put_results;
+    put_results.reserve(documents.size());
+    for (const TokenizedDocument& tokenized_document : tokenized_documents) {
+      ICING_ASSIGN_OR_RETURN(
+          DocumentStore::PutResult put_result,
+          doc_store_->Put(tokenized_document.document_wrapper()));
+      put_results.push_back(std::move(put_result));
+    }
+
+    // Index all documents.
+    for (int i = 0; i < tokenized_documents.size(); ++i) {
+      ICING_RETURN_IF_ERROR(join_indexing_handler_->Handle(
+          tokenized_documents[i], put_results[i].new_document_id,
+          put_results[i].old_document_id, /*recovery_mode=*/false,
+          /*put_document_stats=*/nullptr));
+    }
+    return libtextclassifier3::Status::OK;
+  }
+
+  libtextclassifier3::StatusOr<std::vector<DocumentId>> GetEdgesOfNode(
+      const DocumentDependentGraph& graph, int node_id) {
+    ICING_ASSIGN_OR_RETURN(
+        std::unique_ptr<
+            typename graph::GraphInterface<DocumentId>::EdgeIteratorIf>
+            edge_itr,
+        graph.GetEdgesIterator(node_id));
+    std::vector<DocumentId> edges;
+    while (edge_itr->Advance().ok()) {
+      edges.push_back(edge_itr->Get());
+    }
+    return edges;
+  }
+
+  std::unique_ptr<FeatureFlags> feature_flags_;
+  Filesystem filesystem_;
+  FakeClock fake_clock_;
+  std::string test_dir_;
+  std::string schema_store_dir_;
+  std::string doc_store_dir_;
+  std::string join_index_dir_;
+
+  std::unique_ptr<LanguageSegmenter> lang_segmenter_;
+  std::unique_ptr<SchemaStore> schema_store_;
+  std::unique_ptr<DocumentStore> doc_store_;
+  std::unique_ptr<QualifiedIdJoinIndexImplV3> join_index_;
+
+  std::unique_ptr<QualifiedIdJoinIndexingHandler> join_indexing_handler_;
+};
+
+TEST_F(DocumentDependentGraphTest, CreationWithNullPointerShouldFail) {
+  EXPECT_THAT(
+      DocumentDependentGraph::Create(/*schema_store=*/nullptr, doc_store_.get(),
+                                     join_index_.get()),
+      StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
+  EXPECT_THAT(
+      DocumentDependentGraph::Create(schema_store_.get(), /*doc_store=*/nullptr,
+                                     join_index_.get()),
+      StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
+  EXPECT_THAT(
+      DocumentDependentGraph::Create(schema_store_.get(), doc_store_.get(),
+                                     /*join_index=*/nullptr),
+      StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
+}
+
+TEST_F(DocumentDependentGraphTest,
+       CreationWithWrongJoinIndexVersionShouldFail) {
+  std::string join_index_v2_dir = test_dir_ + "/join_index_v2";
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<QualifiedIdJoinIndexImplV2> join_index_v2,
+      QualifiedIdJoinIndexImplV2::Create(filesystem_, join_index_v2_dir,
+                                         /*pre_mapping_fbv=*/false));
+
+  EXPECT_THAT(DocumentDependentGraph::Create(
+                  schema_store_.get(), doc_store_.get(), join_index_v2.get()),
+              StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT,
+                       HasSubstr("DocumentDependentGraph only supports "
+                                 "QualifiedIdJoinIndex version V3.")));
+}
+
+TEST_F(DocumentDependentGraphTest, GetNumNodes) {
+  // Put and index 4 documents.
+  DocumentProto doc0 =
+      DocumentBuilder().SetKey("namespace", "uri/0").SetSchema("Label").Build();
+  DocumentProto doc1 =
+      DocumentBuilder().SetKey("namespace", "uri/1").SetSchema("Label").Build();
+  DocumentProto doc2 =
+      DocumentBuilder().SetKey("namespace", "uri/2").SetSchema("Label").Build();
+  DocumentProto doc3 =
+      DocumentBuilder().SetKey("namespace", "uri/3").SetSchema("Label").Build();
+  ICING_ASSERT_OK(AddDocuments({doc0, doc1, doc2, doc3}));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<DocumentDependentGraph> graph,
+      DocumentDependentGraph::Create(schema_store_.get(), doc_store_.get(),
+                                     join_index_.get()));
+  EXPECT_THAT(graph->GetNumNodes(), Eq(4));
+}
+
+TEST_F(DocumentDependentGraphTest, GetNumNodes_emptyStorage) {
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<DocumentDependentGraph> graph,
+      DocumentDependentGraph::Create(schema_store_.get(), doc_store_.get(),
+                                     join_index_.get()));
+  EXPECT_THAT(graph->GetNumNodes(), Eq(0));
+}
+
+TEST_F(DocumentDependentGraphTest, GetNumNodes_withReplacedDocuments) {
+  // Put and index 1 document.
+  DocumentProto doc =
+      DocumentBuilder().SetKey("namespace", "uri").SetSchema("Label").Build();
+  ICING_ASSERT_OK(AddDocuments({doc}));
+
+  // Replace the document with new content.
+  DocumentProto doc_replaced = DocumentBuilder()
+                                   .SetKey("namespace", "uri")
+                                   .SetSchema("Label")
+                                   .AddStringProperty("target", "namespace#uri")
+                                   .Build();
+  ICING_ASSERT_OK(AddDocuments({doc_replaced}));
+
+  // Sanity check.
+  ASSERT_THAT(doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/0),
+              IsFalse());
+  ASSERT_THAT(doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/1),
+              IsTrue());
+
+  // Even though document 0 is replaced, num nodes should still be 2 since it is
+  // computed based on last stored doc id.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<DocumentDependentGraph> graph,
+      DocumentDependentGraph::Create(schema_store_.get(), doc_store_.get(),
+                                     join_index_.get()));
+  EXPECT_THAT(graph->GetNumNodes(), Eq(2));
+}
+
+TEST_F(DocumentDependentGraphTest, GetNumNodes_withDeletedDocuments) {
+  // Put and index 1 document.
+  DocumentProto doc =
+      DocumentBuilder().SetKey("namespace", "uri").SetSchema("Label").Build();
+  ICING_ASSERT_OK(AddDocuments({doc}));
+
+  // Delete the document.
+  ICING_ASSERT_OK(doc_store_->Delete(
+      /*document_id=*/0,
+      /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+
+  // Sanity check.
+  ASSERT_THAT(doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/0),
+              IsFalse());
+
+  // Even though document 0 is deleted, num nodes should still be 1 since it is
+  // computed based on last stored doc id.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<DocumentDependentGraph> graph,
+      DocumentDependentGraph::Create(schema_store_.get(), doc_store_.get(),
+                                     join_index_.get()));
+  EXPECT_THAT(graph->GetNumNodes(), Eq(1));
+}
+
+TEST_F(DocumentDependentGraphTest, GetNumNodes_withExpiredDocuments) {
+  fake_clock_.SetSystemTimeMilliseconds(0);
+
+  // Put and index 1 document which expires at 1000 ms.
+  DocumentProto doc = DocumentBuilder()
+                          .SetCreationTimestampMs(0)
+                          .SetTtlMs(1000)
+                          .SetKey("namespace", "uri")
+                          .SetSchema("Label")
+                          .Build();
+  ICING_ASSERT_OK(AddDocuments({doc}));
+
+  // Adjust the clock to expire the document.
+  fake_clock_.SetSystemTimeMilliseconds(2000);
+
+  // Sanity check.
+  ASSERT_THAT(doc_store_->GetAliveDocumentFilterData(
+                  /*document_id=*/0,
+                  /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()),
+              IsFalse());
+  ASSERT_THAT(doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/0),
+              IsTrue());
+
+  // Even though document 0 is expired, num nodes should still be 1 since it is
+  // computed based on last stored doc id.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<DocumentDependentGraph> graph,
+      DocumentDependentGraph::Create(schema_store_.get(), doc_store_.get(),
+                                     join_index_.get()));
+  EXPECT_THAT(graph->GetNumNodes(), Eq(1));
+}
+
+TEST_F(DocumentDependentGraphTest, GetEdgesIterator) {
+  // Put and index 5 documents with the following relations:
+  //
+  // doc0 ---+
+  //         |
+  //         +---> Doc2 --> Doc3
+  //         |
+  // doc1 ---+
+  //   |
+  //   +---------> Doc4
+  DocumentProto doc0 =
+      DocumentBuilder().SetKey("namespace", "uri/0").SetSchema("Label").Build();
+  DocumentProto doc1 =
+      DocumentBuilder().SetKey("namespace", "uri/1").SetSchema("Label").Build();
+  DocumentProto doc2 =
+      DocumentBuilder()
+          .SetKey("namespace", "uri/2")
+          .SetSchema("Label")
+          .AddStringProperty("target", "namespace#uri/0", "namespace#uri/1")
+          .Build();
+  DocumentProto doc3 = DocumentBuilder()
+                           .SetKey("namespace", "uri/3")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/2")
+                           .Build();
+  DocumentProto doc4 = DocumentBuilder()
+                           .SetKey("namespace", "uri/4")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/1")
+                           .Build();
+  ICING_ASSERT_OK(AddDocuments({doc0, doc1, doc2, doc3, doc4}));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<DocumentDependentGraph> graph,
+      DocumentDependentGraph::Create(schema_store_.get(), doc_store_.get(),
+                                     join_index_.get()));
+  ASSERT_THAT(graph->GetNumNodes(), Eq(5));
+
+  EXPECT_THAT(GetEdgesOfNode(*graph, /*node_id=*/0),
+              IsOkAndHolds(ElementsAre(2)));
+  EXPECT_THAT(GetEdgesOfNode(*graph, /*node_id=*/1),
+              IsOkAndHolds(ElementsAre(2, 4)));
+  EXPECT_THAT(GetEdgesOfNode(*graph, /*node_id=*/2),
+              IsOkAndHolds(ElementsAre(3)));
+  EXPECT_THAT(GetEdgesOfNode(*graph, /*node_id=*/3), IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(GetEdgesOfNode(*graph, /*node_id=*/4), IsOkAndHolds(IsEmpty()));
+}
+
+TEST_F(DocumentDependentGraphTest,
+       GetEdgesIterator_withoutDeletePropagationShouldNotBeIncludedIntoEdges) {
+  // Put and index 2 documents with the following relations:
+  //
+  // doc0 --(no delete propagation)--> doc1
+  DocumentProto doc0 =
+      DocumentBuilder().SetKey("namespace", "uri/0").SetSchema("Label").Build();
+  DocumentProto doc1 = DocumentBuilder()
+                           .SetKey("namespace", "uri/1")
+                           .SetSchema("Label")
+                           .AddStringProperty("softTarget", "namespace#uri/0")
+                           .Build();
+  ICING_ASSERT_OK(AddDocuments({doc0, doc1}));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<DocumentDependentGraph> graph,
+      DocumentDependentGraph::Create(schema_store_.get(), doc_store_.get(),
+                                     join_index_.get()));
+  ASSERT_THAT(graph->GetNumNodes(), Eq(2));
+
+  // Even though doc0 and doc1 have join relation, since the delete propagation
+  // type of "softTarget" is NONE, it should not be included into the dependent
+  // edges.
+  ASSERT_THAT(
+      join_index_->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/0),
+      IsOkAndHolds(ElementsAre(
+          DocumentJoinIdPair(/*document_id=*/1, /*joinable_property_id=*/0))));
+  EXPECT_THAT(GetEdgesOfNode(*graph, /*node_id=*/0), IsOkAndHolds(IsEmpty()));
+}
+
+TEST_F(DocumentDependentGraphTest, GetEdgesIterator_shouldDedupeDocumentIds) {
+  // Put and index 2 documents with the following relations:
+  //
+  // doc0 -> doc1
+  //
+  // And doc0, doc1 can be joined by multiple joinable properties with delete
+  // propagation enabled.
+  DocumentProto doc0 =
+      DocumentBuilder().SetKey("namespace", "uri/0").SetSchema("Label").Build();
+  DocumentProto doc1 = DocumentBuilder()
+                           .SetKey("namespace", "uri/1")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/0")
+                           .AddStringProperty("target2", "namespace#uri/0")
+                           .Build();
+  ICING_ASSERT_OK(AddDocuments({doc0, doc1}));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<DocumentDependentGraph> graph,
+      DocumentDependentGraph::Create(schema_store_.get(), doc_store_.get(),
+                                     join_index_.get()));
+  ASSERT_THAT(graph->GetNumNodes(), Eq(2));
+
+  // doc0 should have 2 join relations with doc1, but doc1 should be returned
+  // only once by the iterator.
+  ASSERT_THAT(
+      join_index_->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/0),
+      IsOkAndHolds(ElementsAre(
+          DocumentJoinIdPair(/*document_id=*/1, /*joinable_property_id=*/1),
+          DocumentJoinIdPair(/*document_id=*/1, /*joinable_property_id=*/2))));
+  EXPECT_THAT(GetEdgesOfNode(*graph, /*node_id=*/0),
+              IsOkAndHolds(ElementsAre(1)));
+}
+
+TEST_F(DocumentDependentGraphTest,
+       GetEdgesIterator_replacedDocument_removeEdges) {
+  // Put and index 3 documents with the following relations:
+  //
+  // doc0 -> doc1 -> doc2
+  DocumentProto doc0 =
+      DocumentBuilder().SetKey("namespace", "uri/0").SetSchema("Label").Build();
+  DocumentProto doc1 = DocumentBuilder()
+                           .SetKey("namespace", "uri/1")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/0")
+                           .Build();
+  DocumentProto doc2 = DocumentBuilder()
+                           .SetKey("namespace", "uri/2")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/1")
+                           .Build();
+  ICING_ASSERT_OK(AddDocuments({doc0, doc1, doc2}));
+
+  // Replace doc1 with new relations.
+  // doc0  doc1_replaced -> doc2
+  DocumentProto doc1_replaced =
+      DocumentBuilder().SetKey("namespace", "uri/1").SetSchema("Label").Build();
+  ICING_ASSERT_OK(AddDocuments({doc1_replaced}));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<DocumentDependentGraph> graph,
+      DocumentDependentGraph::Create(schema_store_.get(), doc_store_.get(),
+                                     join_index_.get()));
+  ASSERT_THAT(graph->GetNumNodes(), Eq(4));
+
+  // Sanity check: join indexing handler handled migrate parent, so doc id 1
+  // should be migrated to 3. Doc 0 should still contain the original relation
+  // with doc 1.
+  ASSERT_THAT(
+      join_index_->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/0),
+      IsOkAndHolds(ElementsAre(
+          DocumentJoinIdPair(/*document_id=*/1, /*joinable_property_id=*/1))));
+  ASSERT_THAT(
+      join_index_->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/1),
+      IsOkAndHolds(IsEmpty()));
+  ASSERT_THAT(
+      join_index_->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/2),
+      IsOkAndHolds(IsEmpty()));
+  ASSERT_THAT(
+      join_index_->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/3),
+      IsOkAndHolds(ElementsAre(
+          DocumentJoinIdPair(/*document_id=*/2, /*joinable_property_id=*/1))));
+
+  // Doc 0 should skip the edge to doc 1, since doc 1 was replaced.
+  EXPECT_THAT(GetEdgesOfNode(*graph, /*node_id=*/0), IsOkAndHolds(IsEmpty()));
+  // Doc 1 should return empty iterator since it was replaced.
+  EXPECT_THAT(GetEdgesOfNode(*graph, /*node_id=*/1), IsOkAndHolds(IsEmpty()));
+  // Doc 2 should remain the same.
+  EXPECT_THAT(GetEdgesOfNode(*graph, /*node_id=*/2), IsOkAndHolds(IsEmpty()));
+  // Doc 3 should contain 2.
+  EXPECT_THAT(GetEdgesOfNode(*graph, /*node_id=*/3),
+              IsOkAndHolds(ElementsAre(2)));
+}
+
+TEST_F(DocumentDependentGraphTest, GetEdgesIterator_replacedDocument_addEdges) {
+  // Put and index 4 documents with the following relations:
+  //
+  // doc0 -> doc1 -> doc2 -> doc3
+  DocumentProto doc0 =
+      DocumentBuilder().SetKey("namespace", "uri/0").SetSchema("Label").Build();
+  DocumentProto doc1 = DocumentBuilder()
+                           .SetKey("namespace", "uri/1")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/0")
+                           .Build();
+  DocumentProto doc2 = DocumentBuilder()
+                           .SetKey("namespace", "uri/2")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/1")
+                           .Build();
+  DocumentProto doc3 = DocumentBuilder()
+                           .SetKey("namespace", "uri/3")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/2")
+                           .Build();
+  ICING_ASSERT_OK(AddDocuments({doc0, doc1, doc2, doc3}));
+
+  // Replace doc2 with new relations.
+  // doc0 -> doc1 -> doc2_replaced -> doc3
+  //   |                  ^
+  //   |                  |
+  //   +------------------+
+  DocumentProto doc2_replaced =
+      DocumentBuilder()
+          .SetKey("namespace", "uri/2")
+          .SetSchema("Label")
+          .AddStringProperty("target", "namespace#uri/1", "namespace#uri/0")
+          .Build();
+  ICING_ASSERT_OK(AddDocuments({doc2_replaced}));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<DocumentDependentGraph> graph,
+      DocumentDependentGraph::Create(schema_store_.get(), doc_store_.get(),
+                                     join_index_.get()));
+  ASSERT_THAT(graph->GetNumNodes(), Eq(5));
+
+  // Sanity check: join indexing handler handled migrate parent, so doc id 2
+  // should be migrated to 4. Doc 1 should still contain the original relation
+  // with doc 2.
+  ASSERT_THAT(
+      join_index_->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/0),
+      IsOkAndHolds(ElementsAre(
+          DocumentJoinIdPair(/*document_id=*/1, /*joinable_property_id=*/1),
+          DocumentJoinIdPair(/*document_id=*/4, /*joinable_property_id=*/1))));
+  ASSERT_THAT(
+      join_index_->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/1),
+      IsOkAndHolds(ElementsAre(
+          DocumentJoinIdPair(/*document_id=*/2, /*joinable_property_id=*/1),
+          DocumentJoinIdPair(/*document_id=*/4, /*joinable_property_id=*/1))));
+  ASSERT_THAT(
+      join_index_->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/2),
+      IsOkAndHolds(IsEmpty()));
+  ASSERT_THAT(
+      join_index_->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/3),
+      IsOkAndHolds(IsEmpty()));
+  ASSERT_THAT(
+      join_index_->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/4),
+      IsOkAndHolds(ElementsAre(
+          DocumentJoinIdPair(/*document_id=*/3, /*joinable_property_id=*/1))));
+
+  // Doc 0 should contain both edges to doc 1 and doc 4.
+  EXPECT_THAT(GetEdgesOfNode(*graph, /*node_id=*/0),
+              IsOkAndHolds(ElementsAre(1, 4)));
+  // Doc 1 should skip the edge to doc 2, since doc 2 was replaced.
+  EXPECT_THAT(GetEdgesOfNode(*graph, /*node_id=*/1),
+              IsOkAndHolds(ElementsAre(4)));
+  // Doc 2 should return empty iterator since it was replaced.
+  EXPECT_THAT(GetEdgesOfNode(*graph, /*node_id=*/2), IsOkAndHolds(IsEmpty()));
+  // Doc 3 should remain the same.
+  EXPECT_THAT(GetEdgesOfNode(*graph, /*node_id=*/3), IsOkAndHolds(IsEmpty()));
+  // Doc 4 should contain 3.
+  EXPECT_THAT(GetEdgesOfNode(*graph, /*node_id=*/4),
+              IsOkAndHolds(ElementsAre(3)));
+}
+
+TEST_F(DocumentDependentGraphTest,
+       GetEdgesIterator_deletedParentShouldReturnEmptyIterator) {
+  // Put and index 3 documents with the following relations:
+  //
+  // doc0 -> doc1
+  //  |
+  //  +----> doc2
+  DocumentProto doc0 =
+      DocumentBuilder().SetKey("namespace", "uri/0").SetSchema("Label").Build();
+  DocumentProto doc1 = DocumentBuilder()
+                           .SetKey("namespace", "uri/1")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/0")
+                           .Build();
+  DocumentProto doc2 = DocumentBuilder()
+                           .SetKey("namespace", "uri/2")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/0")
+                           .Build();
+  ICING_ASSERT_OK(AddDocuments({doc0, doc1, doc2}));
+
+  // Delete doc0.
+  ICING_ASSERT_OK(doc_store_->Delete(
+      /*document_id=*/0,
+      /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<DocumentDependentGraph> graph,
+      DocumentDependentGraph::Create(schema_store_.get(), doc_store_.get(),
+                                     join_index_.get()));
+  ASSERT_THAT(graph->GetNumNodes(), Eq(3));
+
+  ASSERT_THAT(
+      join_index_->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/0),
+      IsOkAndHolds(ElementsAre(
+          DocumentJoinIdPair(/*document_id=*/1, /*joinable_property_id=*/1),
+          DocumentJoinIdPair(/*document_id=*/2, /*joinable_property_id=*/1))));
+  // Since doc0 is deleted, the iterator should return empty.
+  EXPECT_THAT(GetEdgesOfNode(*graph, /*node_id=*/0), IsOkAndHolds(IsEmpty()));
+}
+
+TEST_F(DocumentDependentGraphTest,
+       GetEdgesIterator_deletedChildShouldBeSkipped) {
+  // Put and index 3 documents with the following relations:
+  //
+  // doc0 -> doc1
+  //  |
+  //  +----> doc2
+  DocumentProto doc0 =
+      DocumentBuilder().SetKey("namespace", "uri/0").SetSchema("Label").Build();
+  DocumentProto doc1 = DocumentBuilder()
+                           .SetKey("namespace", "uri/1")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/0")
+                           .Build();
+  DocumentProto doc2 = DocumentBuilder()
+                           .SetKey("namespace", "uri/2")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/0")
+                           .Build();
+  ICING_ASSERT_OK(AddDocuments({doc0, doc1, doc2}));
+
+  // Delete doc1.
+  ICING_ASSERT_OK(doc_store_->Delete(
+      /*document_id=*/1,
+      /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<DocumentDependentGraph> graph,
+      DocumentDependentGraph::Create(schema_store_.get(), doc_store_.get(),
+                                     join_index_.get()));
+  ASSERT_THAT(graph->GetNumNodes(), Eq(3));
+
+  ASSERT_THAT(
+      join_index_->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/0),
+      IsOkAndHolds(ElementsAre(
+          DocumentJoinIdPair(/*document_id=*/1, /*joinable_property_id=*/1),
+          DocumentJoinIdPair(/*document_id=*/2, /*joinable_property_id=*/1))));
+  // Since doc1 is deleted, the iterator of doc0 should skip doc1.
+  EXPECT_THAT(GetEdgesOfNode(*graph, /*node_id=*/0),
+              IsOkAndHolds(ElementsAre(2)));
+}
+
+TEST_F(DocumentDependentGraphTest,
+       GetEdgesIterator_expiredParentShouldNotBeSkipped) {
+  fake_clock_.SetSystemTimeMilliseconds(0);
+
+  // Put and index 3 documents with the following relations:
+  //
+  // doc0 -> doc1
+  //  |
+  //  +----> doc2
+  DocumentProto doc0 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(1000)
+                           .SetKey("namespace", "uri/0")
+                           .SetSchema("Label")
+                           .Build();
+  DocumentProto doc1 = DocumentBuilder()
+                           .SetKey("namespace", "uri/1")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/0")
+                           .Build();
+  DocumentProto doc2 = DocumentBuilder()
+                           .SetKey("namespace", "uri/2")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/0")
+                           .Build();
+  ICING_ASSERT_OK(AddDocuments({doc0, doc1, doc2}));
+
+  // Adjust the clock to expire doc0.
+  fake_clock_.SetSystemTimeMilliseconds(2000);
+  ASSERT_THAT(doc_store_->GetAliveDocumentFilterData(
+                  /*document_id=*/0,
+                  /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()),
+              IsFalse());
+  ASSERT_THAT(doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/0),
+              IsTrue());
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<DocumentDependentGraph> graph,
+      DocumentDependentGraph::Create(schema_store_.get(), doc_store_.get(),
+                                     join_index_.get()));
+  ASSERT_THAT(graph->GetNumNodes(), Eq(3));
+
+  ASSERT_THAT(
+      join_index_->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/0),
+      IsOkAndHolds(ElementsAre(
+          DocumentJoinIdPair(/*document_id=*/1, /*joinable_property_id=*/1),
+          DocumentJoinIdPair(/*document_id=*/2, /*joinable_property_id=*/1))));
+  // We should still be able to get all of doc0's edges.
+  EXPECT_THAT(GetEdgesOfNode(*graph, /*node_id=*/0),
+              IsOkAndHolds(ElementsAre(1, 2)));
+}
+
+TEST_F(DocumentDependentGraphTest,
+       GetEdgesIterator_expiredChildShouldNotBeSkipped) {
+  fake_clock_.SetSystemTimeMilliseconds(0);
+
+  // Put and index 3 documents with the following relations:
+  //
+  // doc0 -> doc1
+  //  |
+  //  +----> doc2
+  DocumentProto doc0 =
+      DocumentBuilder().SetKey("namespace", "uri/0").SetSchema("Label").Build();
+  DocumentProto doc1 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(1000)
+                           .SetKey("namespace", "uri/1")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/0")
+                           .Build();
+  DocumentProto doc2 = DocumentBuilder()
+                           .SetKey("namespace", "uri/2")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/0")
+                           .Build();
+  ICING_ASSERT_OK(AddDocuments({doc0, doc1, doc2}));
+
+  // Adjust the clock to expire doc1.
+  fake_clock_.SetSystemTimeMilliseconds(2000);
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<DocumentDependentGraph> graph,
+      DocumentDependentGraph::Create(schema_store_.get(), doc_store_.get(),
+                                     join_index_.get()));
+  ASSERT_THAT(graph->GetNumNodes(), Eq(3));
+
+  ASSERT_THAT(
+      join_index_->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/0),
+      IsOkAndHolds(ElementsAre(
+          DocumentJoinIdPair(/*document_id=*/1, /*joinable_property_id=*/1),
+          DocumentJoinIdPair(/*document_id=*/2, /*joinable_property_id=*/1))));
+  // We should still be able to get doc1 from doc0's edge iterator.
+  EXPECT_THAT(GetEdgesOfNode(*graph, /*node_id=*/0),
+              IsOkAndHolds(ElementsAre(1, 2)));
+}
+
+TEST_F(DocumentDependentGraphTest, GetEdgesIterator_invalidNodeIdShouldFail) {
+  ASSERT_THAT(doc_store_->last_added_document_id(), Eq(kInvalidDocumentId));
+  ASSERT_THAT(join_index_->last_added_document_id(), Eq(kInvalidDocumentId));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<DocumentDependentGraph> graph,
+      DocumentDependentGraph::Create(schema_store_.get(), doc_store_.get(),
+                                     join_index_.get()));
+
+  EXPECT_THAT(graph->GetNumNodes(), Eq(0));
+  EXPECT_THAT(graph->GetEdgesIterator(/*node_id=*/-2),
+              StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
+  EXPECT_THAT(graph->GetEdgesIterator(/*node_id=*/-1),
+              StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
+  EXPECT_THAT(graph->GetEdgesIterator(/*node_id=*/0),
+              StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
+  EXPECT_THAT(graph->GetEdgesIterator(/*node_id=*/1),
+              StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
+  EXPECT_THAT(graph->GetEdgesIterator(/*node_id=*/2),
+              StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
+  EXPECT_THAT(graph->GetEdgesIterator(/*node_id=*/kInvalidDocumentId),
+              StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
+
+  // Put and index 2 documents.
+  DocumentProto doc0 =
+      DocumentBuilder().SetKey("namespace", "uri/0").SetSchema("Label").Build();
+  DocumentProto doc1 =
+      DocumentBuilder().SetKey("namespace", "uri/1").SetSchema("Label").Build();
+  ICING_ASSERT_OK(AddDocuments({doc0, doc1}));
+  ASSERT_THAT(doc_store_->last_added_document_id(), Eq(1));
+  ASSERT_THAT(join_index_->last_added_document_id(), Eq(1));
+
+  EXPECT_THAT(graph->GetNumNodes(), Eq(2));
+  EXPECT_THAT(graph->GetEdgesIterator(/*node_id=*/-2),
+              StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
+  EXPECT_THAT(graph->GetEdgesIterator(/*node_id=*/-1),
+              StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
+  EXPECT_THAT(graph->GetEdgesIterator(/*node_id=*/0), IsOkAndHolds(NotNull()));
+  EXPECT_THAT(graph->GetEdgesIterator(/*node_id=*/1), IsOkAndHolds(NotNull()));
+  EXPECT_THAT(graph->GetEdgesIterator(/*node_id=*/2),
+              StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
+  EXPECT_THAT(graph->GetEdgesIterator(/*node_id=*/kInvalidDocumentId),
+              StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
+}
+
+}  // namespace
+
+}  // namespace lib
+}  // namespace icing

diff --git a/icing/join/expiration-timestamp-util.cc b/icing/join/expiration-timestamp-util.cc
new file mode 100644
index 0000000..ed518d4
--- /dev/null
+++ b/icing/join/expiration-timestamp-util.cc

@@ -0,0 +1,159 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "icing/join/expiration-timestamp-util.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <limits>
+#include <memory>
+#include <optional>
+#include <queue>
+#include <unordered_set>
+#include <utility>
+
+#include "icing/text_classifier/lib3/utils/base/status.h"
+#include "icing/absl_ports/canonical_errors.h"
+#include "icing/graph/graph-interface.h"
+#include "icing/join/document-dependent-graph.h"
+#include "icing/join/qualified-id-join-index.h"
+#include "icing/schema/schema-store.h"
+#include "icing/store/document-filter-data.h"
+#include "icing/store/document-id.h"
+#include "icing/store/document-store.h"
+#include "icing/util/logging.h"
+#include "icing/util/status-macros.h"
+
+namespace icing {
+namespace lib {
+
+/* static */ libtextclassifier3::Status
+ExpirationTimestampUtil::SingleDocumentPropagation(
+    DocumentId document_id,
+    const std::unordered_set<DocumentId>& dependency_doc_ids,
+    const SchemaStore& schema_store,
+    const QualifiedIdJoinIndex& qualified_id_join_index,
+    DocumentStore& document_store, int64_t current_time_ms) {
+  // This function implements the algorithm to:
+  // - Update the given document's expiration timestamp by its given
+  //   dependencies. Mostly its dependencies are determined by the dependency
+  //   evaluation. See DocumentDependencyProcessor for more details.
+  // - Propagate the expiration timestamp to all of the given document's
+  //   non-deleted dependents (and their dependents, and so on), based on the
+  //   dependent graph.
+  //
+  // For example:
+  // - dependency_doc_ids = {depcy1, depcy2}
+  // - dependent_graph contains the edges of the following graph.
+  //
+  //                 +--------> dept1 ------+
+  //                 |                      |
+  //                 |                      v
+  // depcy1 --> document_id --> dept2 --> dept3
+  //                 ^            |
+  //                 |            v
+  // depcy2 ---------+          dept4
+  //
+  // The algorithm will:
+  // - Update document_id's expiration timestamp by min(depcy1_exp_ts,
+  //   depcy2_exp_ts) if this min value is smaller than the document's current
+  //   expiration timestamp.
+  // - BFS traverse: propagate the expiration timestamp to dept1, dept2, dept3,
+  //   and dept4. Early stop if any of the dependents' expiration timestamp is
+  //   not smaller than the propagated value.
+
+  // Step 1.0: get the min expiration timestamp of the dependencies.
+  int64_t min_depcy_exp_ts_ms = std::numeric_limits<int64_t>::max();
+  for (DocumentId depcy_doc_id : dependency_doc_ids) {
+    std::optional<DocumentFilterData> doc_filter_data =
+        document_store.GetAliveDocumentFilterData(depcy_doc_id,
+                                                  current_time_ms);
+    if (!doc_filter_data.has_value()) {
+      // This really shouldn't happen since they were validated in the caller
+      // side by DocumentDependencyProcessor.
+      ICING_LOG(ERROR) << "A dependency document is not alive after dependency "
+                          "evaluation. This should never happen.";
+      return absl_ports::InternalError(
+          "A dependency document is not alive after dependency evaluation.");
+    }
+    min_depcy_exp_ts_ms = std::min(min_depcy_exp_ts_ms,
+                                   doc_filter_data->expiration_timestamp_ms());
+  }
+
+  // Step 1.1: update min expiration timestamp from dependencies to the
+  //           document.
+  //
+  // Note:
+  // - UpdateDocumentExpirationTimestamp only overwrites the new expiration
+  //   timestamp if the new value is smaller than the existing one, so it is
+  //   safe to call this method directly.
+  // - final_exp_ts is the final expiration timestamp of the document.
+  //   - If min_depcy_exp_ts_ms is smaller than the document's current (raw)
+  //     expiration timestamp, final_exp_ts will be min_depcy_exp_ts_ms.
+  //   - Otherwise, final_exp_ts will be the document's (raw) expiration
+  //     timestamp.
+  ICING_ASSIGN_OR_RETURN(
+      DocumentStore::UpdateDocumentExpirationTimestampResult update_result,
+      document_store.UpdateDocumentExpirationTimestamp(document_id,
+                                                       min_depcy_exp_ts_ms));
+  int64_t final_exp_ts = update_result.final_expiration_timestamp_ms;
+
+  // Step 2: run BFS to propagate final_exp_ts to all dependents. In most cases,
+  //         a new document won't have any dependents at this point. Only alive
+  //         document replacement will have dependents to update.
+  ICING_ASSIGN_OR_RETURN(
+      std::unique_ptr<DocumentDependentGraph> dependent_graph,
+      DocumentDependentGraph::Create(&schema_store, &document_store,
+                                     &qualified_id_join_index));
+
+  std::unordered_set<DocumentId> visited_doc_ids;
+  std::queue<DocumentId> que;
+  visited_doc_ids.insert(document_id);
+  que.push(document_id);
+  while (!que.empty()) {
+    DocumentId curr_doc_id = que.front();
+    que.pop();
+
+    ICING_ASSIGN_OR_RETURN(
+        std::unique_ptr<graph::GraphInterface<DocumentId>::EdgeIteratorIf> itr,
+        dependent_graph->GetEdgesIterator(curr_doc_id));
+    while (itr->Advance().ok()) {
+      DocumentId next_doc_id = itr->Get();
+      if (!visited_doc_ids.insert(next_doc_id).second) {
+        // Already visited.
+        continue;
+      }
+
+      // Update the next document's expiration timestamp. Push into the queue
+      // only if the expiration timestamp is updated.
+      //
+      // Note: it is safe to call UpdateDocumentExpirationTimestamp directly and
+      //   return if getting an error here, because DocumentDependentGraph
+      //   returns edges to non deleted documents and at this point next_doc_id
+      //   is guaranteed to be valid and non-deleted.
+      ICING_ASSIGN_OR_RETURN(
+          DocumentStore::UpdateDocumentExpirationTimestampResult
+              dependent_update_result,
+          document_store.UpdateDocumentExpirationTimestamp(next_doc_id,
+                                                           final_exp_ts));
+      if (dependent_update_result.was_updated) {
+        que.push(next_doc_id);
+      }
+    }
+  }
+  return libtextclassifier3::Status::OK;
+}
+
+}  // namespace lib
+}  // namespace icing

diff --git a/icing/join/expiration-timestamp-util.h b/icing/join/expiration-timestamp-util.h
new file mode 100644
index 0000000..759772a
--- /dev/null
+++ b/icing/join/expiration-timestamp-util.h

@@ -0,0 +1,82 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef ICING_JOIN_EXPIRATION_TIMESTAMP_UTIL_H_
+#define ICING_JOIN_EXPIRATION_TIMESTAMP_UTIL_H_
+
+#include <cstdint>
+#include <unordered_set>
+
+#include "icing/text_classifier/lib3/utils/base/status.h"
+#include "icing/join/qualified-id-join-index.h"
+#include "icing/schema/schema-store.h"
+#include "icing/store/document-id.h"
+#include "icing/store/document-store.h"
+
+namespace icing {
+namespace lib {
+
+class ExpirationTimestampUtil {
+ public:
+  // Updates the given single document's expiration timestamp by its
+  // dependencies. Also propagates the expiration timestamp to all of its
+  // "non-deleted" dependents (and their dependents, and so on), based on the
+  // dependent graph stored in the join index.
+  //
+  // Note:
+  // - The implementation is designed for updating a single document in (single)
+  //   Put API.
+  // - To make it efficient, currently it only allows a smaller expiration
+  //   timestamp to be propagated down. IOW, if the new expiration timestamp is
+  //   larger than the existing one, the existing one will be kept and stop
+  //   propagation.
+  // - Therefore, we can run simple BFS instead of Bellman-Ford algorithm, and
+  //   the time complexity is better.
+  // - In the future, if we decide to support updating (1) multiple documents on
+  //   the subgraph for batch API (2) larger expiration timestamp, then:
+  //   - Running K times BFS (K = number of documents to update) is not
+  //     efficient.
+  //   - Bellman-Ford algorithm is not efficient enough either, especially when
+  //     there are cycles in the dependent graph.
+  //   - We should consider running the SCC + topological sort algorithm to
+  //     propagate the expiration timestamp in linear time complexity.
+  //   - For optimization, subgraph SCC and topological sort can be performed,
+  //     but this requires to store the reverse dependent edges for determining
+  //     the subgraph that requires updating.
+  //
+  // Parameters:
+  //   - document_id: the single document id to update.
+  //   - dependency_doc_ids: dependencies of the document_id.
+  //   - schema_store: the schema store.
+  //   - qualified_id_join_index: the qualified id join index.
+  //   - document_store.
+  //   - current_time_ms.
+  //
+  // Returns:
+  //   - OK on success, and the expiration timestamps (which are stored in
+  //     DocumentStore -> DocumentFilterData cache) of the document and its
+  //     dependents are updated.
+  //   - Any error from document store, schema store, or join index.
+  static libtextclassifier3::Status SingleDocumentPropagation(
+      DocumentId document_id,
+      const std::unordered_set<DocumentId>& dependency_doc_ids,
+      const SchemaStore& schema_store,
+      const QualifiedIdJoinIndex& qualified_id_join_index,
+      DocumentStore& document_store, int64_t current_time_ms);
+};
+
+}  // namespace lib
+}  // namespace icing
+
+#endif  // ICING_JOIN_EXPIRATION_TIMESTAMP_UTIL_H_

diff --git a/icing/join/expiration-timestamp-util_test.cc b/icing/join/expiration-timestamp-util_test.cc
new file mode 100644
index 0000000..739a566
--- /dev/null
+++ b/icing/join/expiration-timestamp-util_test.cc

@@ -0,0 +1,948 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "icing/join/expiration-timestamp-util.h"
+
+#include <memory>
+#include <optional>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "icing/text_classifier/lib3/utils/base/status.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+#include "icing/document-builder.h"
+#include "icing/feature-flags.h"
+#include "icing/file/filesystem.h"
+#include "icing/file/portable-file-backed-proto-log.h"
+#include "icing/join/qualified-id-join-index-impl-v3.h"
+#include "icing/join/qualified-id-join-indexing-handler.h"
+#include "icing/portable/gzip_stream.h"
+#include "icing/portable/platform.h"
+#include "icing/proto/document.pb.h"
+#include "icing/proto/document_wrapper.pb.h"
+#include "icing/proto/schema.pb.h"
+#include "icing/proto/scoring.pb.h"
+#include "icing/proto/search.pb.h"
+#include "icing/schema-builder.h"
+#include "icing/schema/schema-store.h"
+#include "icing/store/document-filter-data.h"
+#include "icing/store/document-store.h"
+#include "icing/testing/common-matchers.h"
+#include "icing/testing/fake-clock.h"
+#include "icing/testing/test-data.h"
+#include "icing/testing/test-feature-flags.h"
+#include "icing/testing/tmp-directory.h"
+#include "icing/tokenization/language-segmenter-factory.h"
+#include "icing/tokenization/language-segmenter.h"
+#include "icing/util/icu-data-file-helper.h"
+#include "icing/util/status-macros.h"
+#include "icing/util/tokenized-document.h"
+#include "unicode/uloc.h"
+
+namespace icing {
+namespace lib {
+
+namespace {
+
+using ::testing::Eq;
+using ::testing::IsFalse;
+using ::testing::IsTrue;
+using ::testing::Optional;
+using ::testing::Property;
+
+class ExpirationTimestampUtilTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    feature_flags_ = std::make_unique<FeatureFlags>(GetTestFeatureFlags());
+    ASSERT_THAT(feature_flags_->enable_repeated_field_joins(), IsTrue());
+
+    test_dir_ = GetTestTempDir() + "/expiration_timestamp_util_test";
+    ASSERT_THAT(filesystem_.CreateDirectoryRecursively(test_dir_.c_str()),
+                IsTrue());
+
+    schema_store_dir_ = test_dir_ + "/schema_store";
+    doc_store_dir_ = test_dir_ + "/doc_store";
+    qualified_id_join_index_dir_ = test_dir_ + "/qualified_id_join_index";
+
+    if (!IsCfStringTokenization() && !IsReverseJniTokenization()) {
+      ICING_ASSERT_OK(
+          // File generated via icu_data_file rule in //icing/BUILD.
+          icu_data_file_helper::SetUpIcuDataFile(
+              GetTestFilePath("icing/icu.dat")));
+    }
+
+    language_segmenter_factory::SegmenterOptions options(ULOC_US);
+    ICING_ASSERT_OK_AND_ASSIGN(
+        lang_segmenter_,
+        language_segmenter_factory::Create(std::move(options)));
+
+    ASSERT_THAT(
+        filesystem_.CreateDirectoryRecursively(schema_store_dir_.c_str()),
+        IsTrue());
+    ICING_ASSERT_OK_AND_ASSIGN(
+        schema_store_, SchemaStore::Create(&filesystem_, schema_store_dir_,
+                                           &fake_clock_, feature_flags_.get()));
+
+    SchemaProto schema =
+        SchemaBuilder()
+            .AddType(SchemaTypeConfigBuilder().SetType("Label").AddProperty(
+                PropertyConfigBuilder()
+                    .SetName("target")
+                    .SetDataTypeJoinableString(
+                        JOINABLE_VALUE_TYPE_QUALIFIED_ID,
+                        DELETE_PROPAGATION_TYPE_PROPAGATE_FROM)
+                    .SetCardinality(CARDINALITY_REPEATED)))
+
+            .Build();
+    ASSERT_THAT(schema_store_->SetSchema(
+                    schema, /*ignore_errors_and_delete_documents=*/false),
+                IsOk());
+
+    ASSERT_THAT(filesystem_.CreateDirectoryRecursively(doc_store_dir_.c_str()),
+                IsTrue());
+    ICING_ASSERT_OK_AND_ASSIGN(
+        DocumentStore::CreateResult create_result,
+        DocumentStore::Create(
+            &filesystem_, doc_store_dir_, &fake_clock_, schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
+    doc_store_ = std::move(create_result.document_store);
+
+    ICING_ASSERT_OK_AND_ASSIGN(
+        qualified_id_join_index_,
+        QualifiedIdJoinIndexImplV3::Create(
+            filesystem_, qualified_id_join_index_dir_, *feature_flags_));
+
+    ICING_ASSERT_OK_AND_ASSIGN(
+        qualified_id_join_indexing_handler_,
+        QualifiedIdJoinIndexingHandler::Create(&fake_clock_, doc_store_.get(),
+                                               qualified_id_join_index_.get(),
+                                               feature_flags_.get()));
+  }
+
+  void TearDown() override {
+    qualified_id_join_indexing_handler_.reset();
+
+    qualified_id_join_index_.reset();
+    doc_store_.reset();
+    schema_store_.reset();
+    lang_segmenter_.reset();
+
+    filesystem_.DeleteDirectoryRecursively(test_dir_.c_str());
+  }
+
+  // Helper function to batch add documents.
+  libtextclassifier3::Status AddDocuments(
+      std::vector<DocumentProto> documents) {
+    // Tokenize all documents.
+    std::vector<TokenizedDocument> tokenized_documents;
+    tokenized_documents.reserve(documents.size());
+    for (DocumentProto& document : documents) {
+      ICING_ASSIGN_OR_RETURN(
+          TokenizedDocument tokenized_document,
+          TokenizedDocument::Create(
+              schema_store_.get(), lang_segmenter_.get(),
+              /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+              std::move(document)));
+      tokenized_documents.push_back(std::move(tokenized_document));
+    }
+
+    // Put all documents into the document store and get document ids.
+    std::vector<DocumentStore::PutResult> put_results;
+    put_results.reserve(documents.size());
+    for (const TokenizedDocument& tokenized_document : tokenized_documents) {
+      ICING_ASSIGN_OR_RETURN(
+          DocumentStore::PutResult put_result,
+          doc_store_->Put(tokenized_document.document_wrapper()));
+      put_results.push_back(std::move(put_result));
+    }
+
+    // Index all documents.
+    for (int i = 0; i < tokenized_documents.size(); ++i) {
+      ICING_RETURN_IF_ERROR(qualified_id_join_indexing_handler_->Handle(
+          tokenized_documents[i], put_results[i].new_document_id,
+          put_results[i].old_document_id, /*recovery_mode=*/false,
+          /*put_document_stats=*/nullptr));
+    }
+    return libtextclassifier3::Status::OK;
+  }
+
+  std::unique_ptr<FeatureFlags> feature_flags_;
+  Filesystem filesystem_;
+  FakeClock fake_clock_;
+  std::string test_dir_;
+  std::string schema_store_dir_;
+  std::string doc_store_dir_;
+  std::string qualified_id_join_index_dir_;
+
+  std::unique_ptr<LanguageSegmenter> lang_segmenter_;
+  std::unique_ptr<SchemaStore> schema_store_;
+  std::unique_ptr<DocumentStore> doc_store_;
+  std::unique_ptr<QualifiedIdJoinIndexImplV3> qualified_id_join_index_;
+
+  std::unique_ptr<QualifiedIdJoinIndexingHandler>
+      qualified_id_join_indexing_handler_;
+};
+
+TEST_F(ExpirationTimestampUtilTest, SingleDocumentPropagation) {
+  // General test for SingleDocumentPropagation.
+  fake_clock_.SetSystemTimeMilliseconds(0);
+
+  // Create docs with the following (raw) expiration timestamps and relations:
+  //
+  //                     +------> dept1 (30) -------+
+  //                     |                          |
+  //                     |                          v
+  // depcy1 (10) ---> doc (5) --> dept2 (30) --> dept3 (30)
+  //                     ^           |
+  //                     |           v
+  // depcy2 (2) ---------+        dept4 (30)
+  DocumentProto depcy1 = DocumentBuilder()
+                             .SetCreationTimestampMs(0)
+                             .SetTtlMs(10)
+                             .SetKey("namespace", "depcy1")
+                             .SetSchema("Label")
+                             .Build();
+  DocumentProto depcy2 = DocumentBuilder()
+                             .SetCreationTimestampMs(0)
+                             .SetTtlMs(2)
+                             .SetKey("namespace", "depcy2")
+                             .SetSchema("Label")
+                             .Build();
+  DocumentProto doc =
+      DocumentBuilder()
+          .SetCreationTimestampMs(0)
+          .SetTtlMs(5)
+          .SetKey("namespace", "doc")
+          .SetSchema("Label")
+          .AddStringProperty("target", "namespace#depcy1", "namespace#depcy2")
+          .Build();
+  DocumentProto dept1 = DocumentBuilder()
+                            .SetCreationTimestampMs(0)
+                            .SetTtlMs(30)
+                            .SetKey("namespace", "dept1")
+                            .SetSchema("Label")
+                            .AddStringProperty("target", "namespace#doc")
+                            .Build();
+  DocumentProto dept2 = DocumentBuilder()
+                            .SetCreationTimestampMs(0)
+                            .SetTtlMs(30)
+                            .SetKey("namespace", "dept2")
+                            .SetSchema("Label")
+                            .AddStringProperty("target", "namespace#doc")
+                            .Build();
+  DocumentProto dept3 =
+      DocumentBuilder()
+          .SetCreationTimestampMs(0)
+          .SetTtlMs(30)
+          .SetKey("namespace", "dept3")
+          .SetSchema("Label")
+          .AddStringProperty("target", "namespace#dept1", "namespace#dept2")
+          .Build();
+  DocumentProto dept4 = DocumentBuilder()
+                            .SetCreationTimestampMs(0)
+                            .SetTtlMs(30)
+                            .SetKey("namespace", "dept4")
+                            .SetSchema("Label")
+                            .AddStringProperty("target", "namespace#dept2")
+                            .Build();
+  // Add all documents. Note that they will have document ids 0, 1, 2, 3, 4,
+  // 5, 6.
+  ICING_ASSERT_OK(
+      AddDocuments({depcy1, depcy2, doc, dept1, dept2, dept3, dept4}));
+
+  // Run the expiration timestamp propagation on doc (id 2) with its
+  // dependencies depcy1 and depcy2 (id 0 and 1).
+  //
+  // The expiration timestamps should become:
+  //                     +------> dept1 (2) -------+
+  //                     |                         |
+  //                     |                         v
+  // depcy1 (10) ---> doc (2) --> dept2 (2) --> dept3 (2)
+  //                     ^           |
+  //                     |           v
+  // depcy2 (2) ---------+        dept4 (2)
+  EXPECT_THAT(ExpirationTimestampUtil::SingleDocumentPropagation(
+                  /*document_id=*/2, /*dependency_doc_ids=*/{0, 1},
+                  *schema_store_, *qualified_id_join_index_, *doc_store_,
+                  /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()),
+              IsOk());
+
+  // depcy1 and depcy2 should not be updated.
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/0),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(10))));
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/1),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(2))));
+
+  // Verify doc, dept1, dept2, dept3, and dept4's expiration timestamps are
+  // updated.
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/2),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(2))));
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/3),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(2))));
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/4),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(2))));
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/5),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(2))));
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/6),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(2))));
+}
+
+TEST_F(
+    ExpirationTimestampUtilTest,
+    SingleDocumentPropagation_updateFromDependency_shouldUpdateToSmallerExpTs) {
+  fake_clock_.SetSystemTimeMilliseconds(0);
+
+  // Create doc0 and doc1 with raw expiration timestamps 10 and 2.
+  DocumentProto doc0 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(10)
+                           .SetKey("namespace", "uri/0")
+                           .SetSchema("Label")
+                           .Build();
+  DocumentProto doc1 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(2)
+                           .SetKey("namespace", "uri/1")
+                           .SetSchema("Label")
+                           .Build();
+  ICING_ASSERT_OK(AddDocuments({doc0, doc1}));
+
+  // Create doc2 with raw expiration timestamp 5 and with the following
+  // relations:
+  //
+  // doc0 (10) --+
+  //             |
+  //             +---> doc2 (5)
+  //             |
+  // doc1 (2) ---+
+  DocumentProto doc2 =
+      DocumentBuilder()
+          .SetCreationTimestampMs(0)
+          .SetTtlMs(5)
+          .SetKey("namespace", "uri/2")
+          .SetSchema("Label")
+          .AddStringProperty("target", "namespace#uri/0", "namespace#uri/1")
+          .Build();
+  ICING_ASSERT_OK(AddDocuments({doc2}));
+
+  // Sanity check on the document filter data before propagation.
+  ASSERT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/0),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(10))));
+  ASSERT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/1),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(2))));
+  ASSERT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/2),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(5))));
+
+  // Run the expiration timestamp propagation on doc2 with its dependencies doc0
+  // and doc1. Doc2's expiration timestamp should be updated to 2.
+  EXPECT_THAT(ExpirationTimestampUtil::SingleDocumentPropagation(
+                  /*document_id=*/2, /*dependency_doc_ids=*/{0, 1},
+                  *schema_store_, *qualified_id_join_index_, *doc_store_,
+                  /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()),
+              IsOk());
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/2),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(2))));
+}
+
+TEST_F(
+    ExpirationTimestampUtilTest,
+    SingleDocumentPropagation_updateFromDependency_shouldIgnoreGreaterExpTs) {
+  fake_clock_.SetSystemTimeMilliseconds(0);
+
+  // Create doc0 and doc1 with raw expiration timestamps 10 and 15.
+  DocumentProto doc0 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(10)
+                           .SetKey("namespace", "uri/0")
+                           .SetSchema("Label")
+                           .Build();
+  DocumentProto doc1 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(15)
+                           .SetKey("namespace", "uri/1")
+                           .SetSchema("Label")
+                           .Build();
+  ICING_ASSERT_OK(AddDocuments({doc0, doc1}));
+
+  // Create doc2 with raw expiration timestamp 5 and with the following
+  // relations:
+  //
+  // doc0 (10) --+
+  //             |
+  //             +---> doc2 (5)
+  //             |
+  // doc1 (15) --+
+  DocumentProto doc2 =
+      DocumentBuilder()
+          .SetCreationTimestampMs(0)
+          .SetTtlMs(5)
+          .SetKey("namespace", "uri/2")
+          .SetSchema("Label")
+          .AddStringProperty("target", "namespace#uri/0", "namespace#uri/1")
+          .Build();
+  ICING_ASSERT_OK(AddDocuments({doc2}));
+
+  // Sanity check on the document filter data before propagation.
+  ASSERT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/0),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(10))));
+  ASSERT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/1),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(15))));
+  ASSERT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/2),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(5))));
+
+  // Run the expiration timestamp propagation on doc2 with its dependencies doc0
+  // and doc1. Since all of its dependencies have greater expiration timestamps,
+  // doc2's expiration timestamp should remain 5 after propagation.
+  EXPECT_THAT(ExpirationTimestampUtil::SingleDocumentPropagation(
+                  /*document_id=*/2, /*dependency_doc_ids=*/{0, 1},
+                  *schema_store_, *qualified_id_join_index_, *doc_store_,
+                  /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()),
+              IsOk());
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/2),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(5))));
+}
+
+TEST_F(ExpirationTimestampUtilTest,
+       SingleDocumentPropagation_updateFromDependency_selfCycle) {
+  fake_clock_.SetSystemTimeMilliseconds(0);
+
+  // Create a document with a self cycle relation.
+  DocumentProto doc0 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(10)
+                           .SetKey("namespace", "uri/0")
+                           .AddStringProperty("target", "namespace#uri/0")
+                           .SetSchema("Label")
+                           .Build();
+  ICING_ASSERT_OK(AddDocuments({doc0}));
+
+  // Run the expiration timestamp propagation on doc0.
+  EXPECT_THAT(ExpirationTimestampUtil::SingleDocumentPropagation(
+                  /*document_id=*/0, /*dependency_doc_ids=*/{0}, *schema_store_,
+                  *qualified_id_join_index_, *doc_store_,
+                  /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()),
+              IsOk());
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/0),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(10))));
+}
+
+TEST_F(ExpirationTimestampUtilTest,
+       SingleDocumentPropagation_propgateToDependents_smallerExpTs) {
+  // Note: in the actual integration with IcingSearchEngine, the only case to
+  //   propagate to dependent is when replacing an existing parent document, but
+  //   for testing purpose, we can just simply create all documents at once and
+  //   run the algorithm on the target document directly without creating
+  //   replacement scenario.
+  fake_clock_.SetSystemTimeMilliseconds(0);
+
+  // Create documents with the following (raw) expiration timestamps and
+  // relations:
+  //
+  // doc0 (2) ---> doc1 (5) ----> doc3 (1)
+  //    |
+  //    +--------> doc2 (10) ---> doc4 (3)
+  DocumentProto doc0 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(2)
+                           .SetKey("namespace", "uri/0")
+                           .SetSchema("Label")
+                           .Build();
+  DocumentProto doc1 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(5)
+                           .SetKey("namespace", "uri/1")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/0")
+                           .Build();
+  DocumentProto doc2 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(10)
+                           .SetKey("namespace", "uri/2")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/0")
+                           .Build();
+  DocumentProto doc3 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(1)
+                           .SetKey("namespace", "uri/3")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/1")
+                           .Build();
+  DocumentProto doc4 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(3)
+                           .SetKey("namespace", "uri/4")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/2")
+                           .Build();
+  ICING_ASSERT_OK(AddDocuments({doc0, doc1, doc2, doc3, doc4}));
+
+  // Run the expiration timestamp propagation on doc0. It should propagate to
+  // doc1, doc2, and doc4.
+  //
+  // Final graph:
+  // doc0 (2) ----> doc1 (2) ----> doc3 (1)
+  //    |
+  //    +---------> doc2 (2) ----> doc4 (2)
+  EXPECT_THAT(ExpirationTimestampUtil::SingleDocumentPropagation(
+                  /*document_id=*/0, /*dependency_doc_ids=*/{}, *schema_store_,
+                  *qualified_id_join_index_, *doc_store_,
+                  /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()),
+              IsOk());
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/0),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(2))));
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/1),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(2))));
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/2),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(2))));
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/3),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(1))));
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/4),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(2))));
+}
+
+TEST_F(ExpirationTimestampUtilTest,
+       SingleDocumentPropagation_propgateToDependents_greaterExpTs) {
+  // Note: in the actual integration with IcingSearchEngine, the only case to
+  //   propagate to dependent is when replacing an existing parent document, but
+  //   for testing purpose, we can just simply create all documents at once and
+  //   run the algorithm on the target document directly without creating
+  //   replacement scenario.
+  fake_clock_.SetSystemTimeMilliseconds(0);
+
+  // Create documents with the following (raw) expiration timestamps and
+  // relations:
+  //
+  // doc0 (15) ---> doc1 (5) ----> doc3 (1)
+  //    |
+  //    +---------> doc2 (10) ---> doc4 (3)
+  DocumentProto doc0 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(15)
+                           .SetKey("namespace", "uri/0")
+                           .SetSchema("Label")
+                           .Build();
+  DocumentProto doc1 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(5)
+                           .SetKey("namespace", "uri/1")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/0")
+                           .Build();
+  DocumentProto doc2 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(10)
+                           .SetKey("namespace", "uri/2")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/0")
+                           .Build();
+  DocumentProto doc3 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(1)
+                           .SetKey("namespace", "uri/3")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/1")
+                           .Build();
+  DocumentProto doc4 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(3)
+                           .SetKey("namespace", "uri/4")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/2")
+                           .Build();
+  ICING_ASSERT_OK(AddDocuments({doc0, doc1, doc2, doc3, doc4}));
+
+  // Run the expiration timestamp propagation on doc0. Since it has greater
+  // expiration timestamp than all of its dependents, no dependents' expiration
+  // timestamps should be updated.
+  //
+  // Final graph:
+  // doc0 (15) ---> doc1 (5) ----> doc3 (1)
+  //    |
+  //    +---------> doc2 (10) ---> doc4 (3)
+  EXPECT_THAT(ExpirationTimestampUtil::SingleDocumentPropagation(
+                  /*document_id=*/0, /*dependency_doc_ids=*/{}, *schema_store_,
+                  *qualified_id_join_index_, *doc_store_,
+                  /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()),
+              IsOk());
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/0),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(15))));
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/1),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(5))));
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/2),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(10))));
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/3),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(1))));
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/4),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(3))));
+}
+
+TEST_F(ExpirationTimestampUtilTest,
+       SingleDocumentPropagation_propgateToDependents_expTsShouldOnlyDecrease) {
+  fake_clock_.SetSystemTimeMilliseconds(0);
+
+  // Create documents with the following raw expiration timestamps and
+  // relations:
+  //
+  // doc0 (raw 10) --+
+  //                 |
+  //                 +---> doc2 (raw 8)
+  //                 |
+  // doc1 (raw 5) ---+
+  DocumentProto doc0 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(10)
+                           .SetKey("namespace", "uri/0")
+                           .SetSchema("Label")
+                           .Build();
+  DocumentProto doc1 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(5)
+                           .SetKey("namespace", "uri/1")
+                           .SetSchema("Label")
+                           .Build();
+  DocumentProto doc2 =
+      DocumentBuilder()
+          .SetCreationTimestampMs(0)
+          .SetTtlMs(8)
+          .SetKey("namespace", "uri/2")
+          .SetSchema("Label")
+          .AddStringProperty("target", "namespace#uri/0", "namespace#uri/1")
+          .Build();
+  ICING_ASSERT_OK(AddDocuments({doc0, doc1, doc2}));
+
+  // Run the expiration timestamp propagation on doc2. It should have
+  // (expiration ts, raw expiration ts) = (5, 8).
+  //
+  // doc0 (10, 10) --+
+  //                 |
+  //                 +---> doc2 (5, 8)
+  //                 |
+  // doc1 (5, 5) ----+
+  ASSERT_THAT(ExpirationTimestampUtil::SingleDocumentPropagation(
+                  /*document_id=*/2, /*dependency_doc_ids=*/{0, 1},
+                  *schema_store_, *qualified_id_join_index_, *doc_store_,
+                  /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()),
+              IsOk());
+  std::optional<DocumentFilterData> filter_data =
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/2);
+  ASSERT_THAT(filter_data->expiration_timestamp_ms(), Eq(5));
+  ASSERT_THAT(filter_data->raw_expiration_timestamp_ms(), Eq(8));
+
+  // Create doc3 with raw expiration timestamp 7 and the same key as doc1.
+  // Use it to replace doc1.
+  DocumentProto doc3 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(7)
+                           .SetKey("namespace", "uri/1")
+                           .SetSchema("Label")
+                           .Build();
+  ICING_ASSERT_OK(AddDocuments({doc3}));
+
+  // Run the expiration timestamp propagation on doc3.
+  //
+  // doc0 (10, 10) --+
+  //                 |
+  //                 +---> doc2 (_, 8)
+  //                 |
+  // doc3 (7, 7) ----+
+  //
+  // According to the graph, if doc2 had been added after doc0 and doc3, then it
+  // should've had expiration timestamp 7 after propagation. But since doc3 is a
+  // replacement update now and the BFS algorithm only decreases the expiration
+  // timestamp of doc2 without considering potential increase caused by
+  // dependency replacement, doc2 should still have expiration timestamp 5.
+  EXPECT_THAT(ExpirationTimestampUtil::SingleDocumentPropagation(
+                  /*document_id=*/3, /*dependency_doc_ids=*/{}, *schema_store_,
+                  *qualified_id_join_index_, *doc_store_,
+                  /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()),
+              IsOk());
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/2),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(5))));
+}
+
+TEST_F(ExpirationTimestampUtilTest,
+       SingleDocumentPropagation_propgateToDependents_cycle) {
+  // Note: in the actual integration with IcingSearchEngine, the only case to
+  //   propagate to dependent is when replacing an existing parent document, but
+  //   for testing purpose, we can just simply create all documents at once and
+  //   run the algorithm on the target document directly without creating
+  //   replacement scenario.
+  fake_clock_.SetSystemTimeMilliseconds(0);
+
+  // Create documents with the following raw expiration timestamps and
+  // relations:
+  //
+  // doc0 (raw 2) ----> doc1 (raw 5)
+  //    ^                   |
+  //    |                   |
+  //    |                   v
+  // doc3 (raw 7) <---- doc2 (raw 8)
+  DocumentProto doc0 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(2)
+                           .SetKey("namespace", "uri/0")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/3")
+                           .Build();
+  DocumentProto doc1 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(5)
+                           .SetKey("namespace", "uri/1")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/0")
+                           .Build();
+  DocumentProto doc2 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(8)
+                           .SetKey("namespace", "uri/2")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/1")
+                           .Build();
+  DocumentProto doc3 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(7)
+                           .SetKey("namespace", "uri/3")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/2")
+                           .Build();
+  ICING_ASSERT_OK(AddDocuments({doc0, doc1, doc2, doc3}));
+
+  // Run the expiration timestamp propagation on doc0.
+  EXPECT_THAT(ExpirationTimestampUtil::SingleDocumentPropagation(
+                  /*document_id=*/0, /*dependency_doc_ids=*/{3}, *schema_store_,
+                  *qualified_id_join_index_, *doc_store_,
+                  /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()),
+              IsOk());
+
+  // All of them should have expiration timestamp 2.
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/0),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(2))));
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/1),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(2))));
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/2),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(2))));
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/3),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(2))));
+}
+
+TEST_F(ExpirationTimestampUtilTest,
+       SingleDocumentPropagation_propgateToDependents_shouldSkipDeleted) {
+  fake_clock_.SetSystemTimeMilliseconds(0);
+
+  // Create documents with the following (raw) expiration timestamps and
+  // relations:
+  //
+  // doc0 (5) ---> doc1 (10)
+  DocumentProto doc0 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(5)
+                           .SetKey("namespace", "uri/0")
+                           .SetSchema("Label")
+                           .Build();
+  DocumentProto doc1 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(10)
+                           .SetKey("namespace", "uri/1")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/0")
+                           .Build();
+  ICING_ASSERT_OK(AddDocuments({doc0, doc1}));
+
+  // Delete doc1.
+  ICING_ASSERT_OK(doc_store_->Delete(
+      /*document_id=*/1,
+      /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+
+  // Run the expiration timestamp propagation on doc0. doc1 should be skipped
+  // since it's deleted.
+  EXPECT_THAT(ExpirationTimestampUtil::SingleDocumentPropagation(
+                  /*document_id=*/0, /*dependency_doc_ids=*/{}, *schema_store_,
+                  *qualified_id_join_index_, *doc_store_,
+                  /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()),
+              IsOk());
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/0),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(5))));
+  EXPECT_THAT(doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/1),
+              IsFalse());
+}
+
+TEST_F(ExpirationTimestampUtilTest,
+       SingleDocumentPropagation_propgateToDependents_shouldSkipExpired) {
+  fake_clock_.SetSystemTimeMilliseconds(0);
+
+  // Create documents with the following (raw) expiration timestamps and
+  // relations:
+  //
+  // doc0 (30) ---> doc1 (10)
+  DocumentProto doc0 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(30)
+                           .SetKey("namespace", "uri/0")
+                           .SetSchema("Label")
+                           .Build();
+  DocumentProto doc1 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(10)
+                           .SetKey("namespace", "uri/1")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/0")
+                           .Build();
+  ICING_ASSERT_OK(AddDocuments({doc0, doc1}));
+
+  // Adjust the current time to 20. It makes doc1 expired.
+  fake_clock_.SetSystemTimeMilliseconds(20);
+
+  // Run the expiration timestamp propagation on doc0. doc1 is still traversed
+  // and we attempt to set the propagated expiration timestamp, but the new
+  // value is always larger than the existing one, so it's skipped.
+  EXPECT_THAT(ExpirationTimestampUtil::SingleDocumentPropagation(
+                  /*document_id=*/0, /*dependency_doc_ids=*/{}, *schema_store_,
+                  *qualified_id_join_index_, *doc_store_,
+                  /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()),
+              IsOk());
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/0),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(30))));
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/1),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(10))));
+}
+
+TEST_F(ExpirationTimestampUtilTest,
+       SingleDocumentPropagation_propgateToDependents_noBellmanFord) {
+  // Note: in the actual integration with IcingSearchEngine, this update
+  //   scenario is not possible. But for testing purpose, we want to make sure
+  //   the algorithm doesn't traverse nodes on a cycle for multiple times, which
+  //   has bad time complexity of O(V*E).
+  fake_clock_.SetSystemTimeMilliseconds(0);
+
+  // Create documents with the following raw expiration timestamps and
+  // relations:
+  //
+  // doc0 (raw 4) ----> doc1 (raw 8) <---------+
+  //                         |                 |
+  //                         |                 |
+  //                         v                 |
+  //                    doc2 (raw 5) ---> doc3 (raw 2)
+  DocumentProto doc0 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(4)
+                           .SetKey("namespace", "uri/0")
+                           .SetSchema("Label")
+                           .Build();
+  DocumentProto doc1 =
+      DocumentBuilder()
+          .SetCreationTimestampMs(0)
+          .SetTtlMs(8)
+          .SetKey("namespace", "uri/1")
+          .SetSchema("Label")
+          .AddStringProperty("target", "namespace#uri/0", "namespace#uri/3")
+          .Build();
+  DocumentProto doc2 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(5)
+                           .SetKey("namespace", "uri/2")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/1")
+                           .Build();
+  DocumentProto doc3 = DocumentBuilder()
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(2)
+                           .SetKey("namespace", "uri/3")
+                           .SetSchema("Label")
+                           .AddStringProperty("target", "namespace#uri/2")
+                           .Build();
+  ICING_ASSERT_OK(AddDocuments({doc0, doc1, doc2, doc3}));
+
+  // Run the expiration timestamp propagation on doc0.
+  EXPECT_THAT(ExpirationTimestampUtil::SingleDocumentPropagation(
+                  /*document_id=*/0, /*dependency_doc_ids=*/{}, *schema_store_,
+                  *qualified_id_join_index_, *doc_store_,
+                  /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()),
+              IsOk());
+
+  // - According to the definition of the dependency, doc1, doc2, and doc3
+  //   should've had expiration timestamp 2.
+  // - But here, we're not running Bellman-Ford algorithm, so when starting from
+  //   doc0, the expiration timestamp of doc3 is not propagated. So doc3 should
+  //   remain 2, and doc1 and doc2 should be updated to 4.
+  //
+  // Note: in reality, when we added doc3, another round of propagation
+  //   should've propagated doc3's expiration timestamp to others, so we will
+  //   never get this incorrect scenario in production. If we really need batch
+  //   update for doc0 and doc3, then we should consider SCC + topological sort
+  //   (linear) algorithm, as mentioned in the docstring of
+  //   SingleDocumentPropagation.
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/0),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(4))));
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/1),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(4))));
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/2),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(4))));
+  EXPECT_THAT(
+      doc_store_->GetNonDeletedDocumentFilterData(/*document_id=*/3),
+      Optional(Property(&DocumentFilterData::expiration_timestamp_ms, Eq(2))));
+}
+
+}  // namespace
+
+}  // namespace lib
+}  // namespace icing

diff --git a/icing/join/join-children-fetcher-impl-v3.cc b/icing/join/join-children-fetcher-impl-v3.cc
index 456c3ba..52a6373 100644
--- a/icing/join/join-children-fetcher-impl-v3.cc
+++ b/icing/join/join-children-fetcher-impl-v3.cc

@@ -114,8 +114,9 @@
   // Otherwise, fetch the children from the qualified id join index and cache
   // the result.
   ICING_ASSIGN_OR_RETURN(
-      std::vector<DocumentJoinIdPair> child_document_join_id_pairs,
-      qualified_id_join_index_.Get(parent_doc_id));
+      QualifiedIdJoinIndex::DocumentJoinIdPairArrayView
+          child_join_id_pairs_array_view,
+      qualified_id_join_index_.GetDocumentJoinIdPairArrayView(parent_doc_id));
 
   // Filter and construct child_scored_document_hits for the given parent doc
   // id.
@@ -140,10 +141,10 @@
   //     set, or the joinable property does not match the one specified in the
   //     join spec.
   std::vector<ScoredDocumentHit> child_scored_document_hits;
-  for (const DocumentJoinIdPair& child_document_join_id_pair :
-       child_document_join_id_pairs) {
+  for (const DocumentJoinIdPair& child_join_id_pair :
+       child_join_id_pairs_array_view) {
     if (auto filter_itr = child_join_id_pair_to_scored_document_hit_map_.find(
-            child_document_join_id_pair);
+            child_join_id_pair);
         filter_itr != child_join_id_pair_to_scored_document_hit_map_.end()) {
       child_scored_document_hits.push_back(filter_itr->second);
     }

diff --git a/icing/join/join-children-fetcher-impl-v3_test.cc b/icing/join/join-children-fetcher-impl-v3_test.cc
index 9a39157..8cc13e5 100644
--- a/icing/join/join-children-fetcher-impl-v3_test.cc
+++ b/icing/join/join-children-fetcher-impl-v3_test.cc

@@ -32,6 +32,7 @@
 #include "icing/join/qualified-id-join-index-impl-v2.h"
 #include "icing/join/qualified-id-join-index-impl-v3.h"
 #include "icing/join/qualified-id-join-indexing-handler.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/portable/platform.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/document_wrapper.pb.h"
@@ -69,6 +70,8 @@
  protected:
   void SetUp() override {
     feature_flags_ = std::make_unique<FeatureFlags>(GetTestFeatureFlags());
+    fake_clock_.SetSystemTimeMilliseconds(123);
+
     base_dir_ = GetTestTempDir() + "/icing_test";
     ASSERT_THAT(filesystem_.CreateDirectoryRecursively(base_dir_.c_str()),
                 IsTrue());
@@ -131,14 +134,18 @@
                 IsTrue());
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
-        DocumentStore::Create(&filesystem_, doc_store_dir_, &fake_clock_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, doc_store_dir_, &fake_clock_, schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     doc_store_ = std::move(create_result.document_store);
 
     ICING_ASSERT_OK_AND_ASSIGN(
@@ -158,17 +165,21 @@
 
   libtextclassifier3::StatusOr<DocumentId> PutAndIndexDocument(
       const DocumentProto& document) {
-    ICING_ASSIGN_OR_RETURN(DocumentStore::PutResult put_result,
-                           doc_store_->Put(document));
     ICING_ASSIGN_OR_RETURN(
         TokenizedDocument tokenized_document,
-        TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                  document));
+        TokenizedDocument::Create(
+            schema_store_.get(), lang_segmenter_.get(),
+            /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+            std::move(document)));
+    ICING_ASSIGN_OR_RETURN(
+        DocumentStore::PutResult put_result,
+        doc_store_->Put(tokenized_document.document_wrapper()));
 
     ICING_ASSIGN_OR_RETURN(
         std::unique_ptr<QualifiedIdJoinIndexingHandler> handler,
         QualifiedIdJoinIndexingHandler::Create(&fake_clock_, doc_store_.get(),
-                                               qualified_id_join_index_.get()));
+                                               qualified_id_join_index_.get(),
+                                               feature_flags_.get()));
     ICING_RETURN_IF_ERROR(
         handler->Handle(tokenized_document, put_result.new_document_id,
                         put_result.old_document_id, /*recovery_mode=*/false,
@@ -489,12 +500,14 @@
   ICING_ASSERT_OK_AND_ASSIGN(DocumentId message3_id,
                              PutAndIndexDocument(message3));
   // Sanity check for the join index.
-  ASSERT_THAT(qualified_id_join_index_->Get(person1_id),
-              IsOkAndHolds(ElementsAre(DocumentJoinIdPair(message1_id, 0),
-                                       DocumentJoinIdPair(message2_id, 1),
-                                       DocumentJoinIdPair(message3_id, 0))));
-  ASSERT_THAT(qualified_id_join_index_->Get(person2_id),
-              IsOkAndHolds(ElementsAre(DocumentJoinIdPair(message2_id, 0))));
+  ASSERT_THAT(
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(person1_id),
+      IsOkAndHolds(ElementsAre(DocumentJoinIdPair(message1_id, 0),
+                               DocumentJoinIdPair(message2_id, 1),
+                               DocumentJoinIdPair(message3_id, 0))));
+  ASSERT_THAT(
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(person2_id),
+      IsOkAndHolds(ElementsAre(DocumentJoinIdPair(message2_id, 0))));
 
   ScoredDocumentHit scored_doc_hit_message1(message1_id, kSectionIdMaskNone,
                                             /*score=*/1.0);
@@ -596,13 +609,16 @@
   ICING_ASSERT_OK_AND_ASSIGN(DocumentId email3_id, PutAndIndexDocument(email3));
   ICING_ASSERT_OK_AND_ASSIGN(DocumentId email4_id, PutAndIndexDocument(email4));
   // Sanity check for the join index.
-  ASSERT_THAT(qualified_id_join_index_->Get(person1_id),
-              IsOkAndHolds(ElementsAre(DocumentJoinIdPair(email1_id, 0),
-                                       DocumentJoinIdPair(email2_id, 0))));
-  ASSERT_THAT(qualified_id_join_index_->Get(person2_id),
-              IsOkAndHolds(ElementsAre(DocumentJoinIdPair(email3_id, 0))));
-  ASSERT_THAT(qualified_id_join_index_->Get(person3_id),
-              IsOkAndHolds(ElementsAre(DocumentJoinIdPair(email4_id, 0))));
+  ASSERT_THAT(
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(person1_id),
+      IsOkAndHolds(ElementsAre(DocumentJoinIdPair(email1_id, 0),
+                               DocumentJoinIdPair(email2_id, 0))));
+  ASSERT_THAT(
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(person2_id),
+      IsOkAndHolds(ElementsAre(DocumentJoinIdPair(email3_id, 0))));
+  ASSERT_THAT(
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(person3_id),
+      IsOkAndHolds(ElementsAre(DocumentJoinIdPair(email4_id, 0))));
 
   ScoredDocumentHit scored_doc_hit_email1(email1_id, kSectionIdMaskNone,
                                           /*score=*/1.0);

diff --git a/icing/join/join-processor.cc b/icing/join/join-processor.cc
index 9b6aa2a..00efc91 100644
--- a/icing/join/join-processor.cc
+++ b/icing/join/join-processor.cc

@@ -15,14 +15,10 @@
 #include "icing/join/join-processor.h"
 
 #include <algorithm>
-#include <deque>
 #include <memory>
 #include <optional>
-#include <queue>
-#include <string>
 #include <string_view>
 #include <unordered_map>
-#include <unordered_set>
 #include <utility>
 #include <vector>
 
@@ -31,12 +27,10 @@
 #include "icing/absl_ports/canonical_errors.h"
 #include "icing/absl_ports/str_cat.h"
 #include "icing/join/aggregation-scorer.h"
-#include "icing/join/document-join-id-pair.h"
 #include "icing/join/join-children-fetcher-impl-deprecated.h"
 #include "icing/join/join-children-fetcher-impl-v3.h"
 #include "icing/join/join-children-fetcher.h"
 #include "icing/join/qualified-id-join-index.h"
-#include "icing/join/qualified-id.h"
 #include "icing/proto/schema.pb.h"
 #include "icing/proto/scoring.pb.h"
 #include "icing/proto/search.pb.h"
@@ -45,7 +39,6 @@
 #include "icing/store/document-filter-data.h"
 #include "icing/store/document-id.h"
 #include "icing/store/namespace-id-fingerprint.h"
-#include "icing/util/logging.h"
 #include "icing/util/status-macros.h"
 
 namespace icing {
@@ -209,80 +202,5 @@
   return joined_scored_document_hits;
 }
 
-libtextclassifier3::StatusOr<std::unordered_set<DocumentId>>
-JoinProcessor::GetPropagatedChildDocumentsToDelete(
-    const std::unordered_set<DocumentId>& deleted_document_ids) {
-  // Sanity check: join index should be V3.
-  if (qualified_id_join_index_->version() !=
-      QualifiedIdJoinIndex::Version::kV3) {
-    return absl_ports::UnimplementedError(
-        "QualifiedIdJoinIndex version must be V3 to support delete "
-        "propagation.");
-  }
-
-  // BFS traverse to find all child documents to propagate delete.
-  std::queue<DocumentId> que(
-      std::deque(deleted_document_ids.begin(), deleted_document_ids.end()));
-  std::unordered_set<DocumentId> child_documents_to_delete;
-  while (!que.empty()) {
-    DocumentId doc_id_to_expand = que.front();
-    que.pop();
-
-    ICING_ASSIGN_OR_RETURN(std::vector<DocumentJoinIdPair> child_join_id_pairs,
-                           qualified_id_join_index_->Get(doc_id_to_expand));
-    for (const DocumentJoinIdPair& child_join_id_pair : child_join_id_pairs) {
-      if (child_documents_to_delete.find(child_join_id_pair.document_id()) !=
-              child_documents_to_delete.end() ||
-          deleted_document_ids.find(child_join_id_pair.document_id()) !=
-              deleted_document_ids.end()) {
-        // Already added into the set to delete or already deleted (happens only
-        // when there is a cycle back to the deleted or traversed document in
-        // the join relation). Skip it.
-        continue;
-      }
-
-      // Get DocumentFilterData of the child document to look up its schema type
-      // id.
-      // - Skip if the child document has been deleted, since delete propagation
-      //   should've been done to all its children when deleting it previously.
-      // - Otherwise, we have to handle this child document and propagate delete
-      //   to the grandchildren, even if it is expired.
-      std::optional<DocumentFilterData> child_filter_data =
-          doc_store_->GetNonDeletedDocumentFilterData(
-              child_join_id_pair.document_id());
-      if (!child_filter_data) {
-        // The child document has been deleted. Skip.
-        continue;
-      }
-
-      libtextclassifier3::StatusOr<const JoinablePropertyMetadata*>
-          metadata_or = schema_store_->GetJoinablePropertyMetadata(
-              child_filter_data->schema_type_id(),
-              child_join_id_pair.joinable_property_id());
-      if (!metadata_or.ok() || metadata_or.ValueOrDie() == nullptr) {
-        // This shouldn't happen because we've validated it during indexing and
-        // only put valid DocumentJoinIdPair into qualified id join index.
-        // Log and skip it.
-        ICING_LOG(ERROR) << "Failed to get metadata for schema type id "
-                         << child_filter_data->schema_type_id()
-                         << ", joinable property id "
-                         << static_cast<int>(
-                                child_join_id_pair.joinable_property_id());
-        continue;
-      }
-      const JoinablePropertyMetadata* metadata = metadata_or.ValueOrDie();
-
-      if (metadata->value_type == JoinableConfig::ValueType::QUALIFIED_ID &&
-          metadata->delete_propagation_type ==
-              JoinableConfig::DeletePropagationType::PROPAGATE_FROM) {
-        child_documents_to_delete.insert(child_join_id_pair.document_id());
-        que.push(child_join_id_pair.document_id());
-      }
-    }
-  }
-
-  return child_documents_to_delete;
-}
-
 }  // namespace lib
 }  // namespace icing

diff --git a/icing/join/join-processor.h b/icing/join/join-processor.h
index aa74395..11af2f3 100644
--- a/icing/join/join-processor.h
+++ b/icing/join/join-processor.h

@@ -17,19 +17,15 @@
 
 #include <cstdint>
 #include <memory>
-#include <string>
 #include <string_view>
-#include <unordered_set>
 #include <vector>
 
-#include "icing/text_classifier/lib3/utils/base/status.h"
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "icing/join/join-children-fetcher.h"
 #include "icing/join/qualified-id-join-index.h"
 #include "icing/proto/search.pb.h"
 #include "icing/schema/schema-store.h"
 #include "icing/scoring/scored-document-hit.h"
-#include "icing/store/document-id.h"
 #include "icing/store/document-store.h"
 
 namespace icing {
@@ -65,16 +61,6 @@
       std::vector<ScoredDocumentHit>&& parent_scored_document_hits,
       const JoinChildrenFetcher& join_children_fetcher);
 
-  // Gets all child documents to delete, propagated from the given deleted
-  // documents.
-  //
-  // Returns:
-  //   - On success, a set of child document ids to delete.
-  //   - Any other errors.
-  libtextclassifier3::StatusOr<std::unordered_set<DocumentId>>
-  GetPropagatedChildDocumentsToDelete(
-      const std::unordered_set<DocumentId>& deleted_document_ids);
-
  private:
   // TODO(b/275121148): deprecate v2 after rollout v3.
 

diff --git a/icing/join/join-processor_test.cc b/icing/join/join-processor_test.cc
index dd88bd1..8fb4b80 100644
--- a/icing/join/join-processor_test.cc
+++ b/icing/join/join-processor_test.cc

@@ -14,15 +14,11 @@
 
 #include "icing/join/join-processor.h"
 
-#include <cstdint>
 #include <memory>
-#include <optional>
 #include <string>
-#include <unordered_set>
 #include <utility>
 #include <vector>
 
-#include "icing/text_classifier/lib3/utils/base/status.h"
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
@@ -31,12 +27,12 @@
 #include "icing/feature-flags.h"
 #include "icing/file/filesystem.h"
 #include "icing/file/portable-file-backed-proto-log.h"
-#include "icing/join/document-join-id-pair.h"
 #include "icing/join/join-children-fetcher.h"
 #include "icing/join/qualified-id-join-index-impl-v2.h"
 #include "icing/join/qualified-id-join-index-impl-v3.h"
 #include "icing/join/qualified-id-join-index.h"
 #include "icing/join/qualified-id-join-indexing-handler.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/portable/platform.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/document_wrapper.pb.h"
@@ -83,6 +79,8 @@
 
   void SetUp() override {
     feature_flags_ = std::make_unique<FeatureFlags>(GetTestFeatureFlags());
+    fake_clock_.SetSystemTimeMilliseconds(123);
+
     test_dir_ = GetTestTempDir() + "/icing_join_processor_test";
     ASSERT_THAT(filesystem_.CreateDirectoryRecursively(test_dir_.c_str()),
                 IsTrue());
@@ -186,14 +184,18 @@
                 IsTrue());
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
-        DocumentStore::Create(&filesystem_, doc_store_dir_, &fake_clock_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, doc_store_dir_, &fake_clock_, schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     doc_store_ = std::move(create_result.document_store);
 
     ICING_ASSERT_OK_AND_ASSIGN(qualified_id_join_index_,
@@ -230,18 +232,22 @@
   }
 
   libtextclassifier3::StatusOr<DocumentId> PutAndIndexDocument(
-      const DocumentProto& document) {
-    ICING_ASSIGN_OR_RETURN(DocumentStore::PutResult put_result,
-                           doc_store_->Put(document));
+      DocumentProto document) {
     ICING_ASSIGN_OR_RETURN(
         TokenizedDocument tokenized_document,
-        TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                  document));
+        TokenizedDocument::Create(
+            schema_store_.get(), lang_segmenter_.get(),
+            /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+            std::move(document)));
+    ICING_ASSIGN_OR_RETURN(
+        DocumentStore::PutResult put_result,
+        doc_store_->Put(tokenized_document.document_wrapper()));
 
     ICING_ASSIGN_OR_RETURN(
         std::unique_ptr<QualifiedIdJoinIndexingHandler> handler,
         QualifiedIdJoinIndexingHandler::Create(&fake_clock_, doc_store_.get(),
-                                               qualified_id_join_index_.get()));
+                                               qualified_id_join_index_.get(),
+                                               feature_flags_.get()));
     ICING_RETURN_IF_ERROR(
         handler->Handle(tokenized_document, put_result.new_document_id,
                         put_result.old_document_id, /*recovery_mode=*/false,
@@ -963,679 +969,6 @@
                       /*child_scored_document_hits=*/{scored_doc_hit6}))));
 }
 
-TYPED_TEST(JoinProcessorTest, GetPropagatedChildDocumentsToDelete) {
-  DocumentProto person = DocumentBuilder()
-                             .SetKey("pkg$db/namespace", "person")
-                             .SetSchema("Person")
-                             .AddStringProperty("Name", "Alice")
-                             .Build();
-
-  // "sender" (joinable property id 2) has
-  // DELETE_PROPAGATION_TYPE_PROPAGATE_FROM.
-  DocumentProto message =
-      DocumentBuilder()
-          .SetKey("pkg$db/namespace", "message")
-          .SetSchema("Message")
-          .AddStringProperty("content", "test content")
-          .AddStringProperty("sender", "pkg$db/namespace#person")
-          .Build();
-
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentId person_doc_id,
-                             this->PutAndIndexDocument(person));
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentId message_doc_id,
-                             this->PutAndIndexDocument(message));
-
-  JoinProcessor join_processor(
-      this->doc_store_.get(), this->schema_store_.get(),
-      this->qualified_id_join_index_.get(),
-      /*current_time_ms=*/this->fake_clock_.GetSystemTimeMilliseconds());
-
-  if (!std::is_same_v<typename TestFixture::JOIN_INDEX_TYPE,
-                      QualifiedIdJoinIndexImplV3>) {
-    EXPECT_THAT(
-        join_processor.GetPropagatedChildDocumentsToDelete({person_doc_id}),
-        StatusIs(libtextclassifier3::StatusCode::UNIMPLEMENTED,
-                 HasSubstr("QualifiedIdJoinIndex version must be V3 to "
-                           "support delete propagation")));
-  } else {
-    // Sanity check.
-    ASSERT_THAT(this->qualified_id_join_index_->Get(person_doc_id),
-                IsOkAndHolds(ElementsAre(DocumentJoinIdPair(
-                    message_doc_id, /*joinable_property_id=*/2))));
-
-    // Deleting the parent document should propagate the delete to its child
-    // document via the joinable property with
-    // DELETE_PROPAGATION_TYPE_PROPAGATE_FROM.
-    EXPECT_THAT(
-        join_processor.GetPropagatedChildDocumentsToDelete({person_doc_id}),
-        IsOkAndHolds(UnorderedElementsAre(message_doc_id)));
-  }
-}
-
-TYPED_TEST(
-    JoinProcessorTest,
-    GetPropagatedChildDocumentsToDelete_shouldNotPropagateToChildDocumentsWithPropagateDeleteDisabled) {
-  if (!std::is_same_v<typename TestFixture::JOIN_INDEX_TYPE,
-                      QualifiedIdJoinIndexImplV3>) {
-    GTEST_SKIP() << "Skipping test for non v3 QualifiedIdJoinIndex";
-  }
-
-  DocumentProto person = DocumentBuilder()
-                             .SetKey("pkg$db/namespace", "person")
-                             .SetSchema("Person")
-                             .AddStringProperty("Name", "Alice")
-                             .Build();
-
-  // "receiver" (joinable property id 0) has DELETE_PROPAGATION_TYPE_NONE.
-  DocumentProto message =
-      DocumentBuilder()
-          .SetKey("pkg$db/namespace", "message")
-          .SetSchema("Message")
-          .AddStringProperty("content", "test content")
-          .AddStringProperty("receiver", "pkg$db/namespace#person")
-          .Build();
-
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentId person_doc_id,
-                             this->PutAndIndexDocument(person));
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentId message_doc_id,
-                             this->PutAndIndexDocument(message));
-
-  // Sanity check.
-  ASSERT_THAT(this->qualified_id_join_index_->Get(person_doc_id),
-              IsOkAndHolds(ElementsAre(DocumentJoinIdPair(
-                  message_doc_id, /*joinable_property_id=*/0))));
-
-  JoinProcessor join_processor(
-      this->doc_store_.get(), this->schema_store_.get(),
-      this->qualified_id_join_index_.get(),
-      /*current_time_ms=*/this->fake_clock_.GetSystemTimeMilliseconds());
-
-  // Deleting the parent document should not propagate the delete to its child
-  // document via the joinable property with DELETE_PROPAGATION_TYPE_NONE.
-  EXPECT_THAT(
-      join_processor.GetPropagatedChildDocumentsToDelete({person_doc_id}),
-      IsOkAndHolds(IsEmpty()));
-}
-
-TYPED_TEST(
-    JoinProcessorTest,
-    GetPropagatedChildDocumentsToDelete_shouldNotPropagateToNonJoinableChildDocuments) {
-  if (!std::is_same_v<typename TestFixture::JOIN_INDEX_TYPE,
-                      QualifiedIdJoinIndexImplV3>) {
-    GTEST_SKIP() << "Skipping test for non v3 QualifiedIdJoinIndex";
-  }
-
-  DocumentProto person = DocumentBuilder()
-                             .SetKey("pkg$db/namespace", "person")
-                             .SetSchema("Person")
-                             .AddStringProperty("Name", "Alice")
-                             .Build();
-
-  // Put person's qualified id string in a non-joinable property.
-  DocumentProto message =
-      DocumentBuilder()
-          .SetKey("pkg$db/namespace", "message")
-          .SetSchema("Message")
-          .AddStringProperty("content", "pkg$db/namespace#person")
-          .Build();
-
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentId person_doc_id,
-                             this->PutAndIndexDocument(person));
-  ICING_ASSERT_OK(this->PutAndIndexDocument(message));
-
-  // Sanity check.
-  ASSERT_THAT(this->qualified_id_join_index_->Get(person_doc_id),
-              IsOkAndHolds(IsEmpty()));
-
-  JoinProcessor join_processor(
-      this->doc_store_.get(), this->schema_store_.get(),
-      this->qualified_id_join_index_.get(),
-      /*current_time_ms=*/this->fake_clock_.GetSystemTimeMilliseconds());
-
-  // Deleting the parent document should not propagate the delete to
-  // non-joinable child documents.
-  EXPECT_THAT(
-      join_processor.GetPropagatedChildDocumentsToDelete({person_doc_id}),
-      IsOkAndHolds(IsEmpty()));
-}
-
-TYPED_TEST(JoinProcessorTest,
-           GetPropagatedChildDocumentsToDelete_propagateViaMultipleProperties) {
-  if (!std::is_same_v<typename TestFixture::JOIN_INDEX_TYPE,
-                      QualifiedIdJoinIndexImplV3>) {
-    GTEST_SKIP() << "Skipping test for non v3 QualifiedIdJoinIndex";
-  }
-
-  DocumentProto person = DocumentBuilder()
-                             .SetKey("pkg$db/namespace", "person")
-                             .SetSchema("Person")
-                             .AddStringProperty("Name", "Alice")
-                             .Build();
-
-  // - "sender" (joinable property id 2) has
-  //   DELETE_PROPAGATION_TYPE_PROPAGATE_FROM.
-  // - "receiver" (joinable property id 0) has DELETE_PROPAGATION_TYPE_NONE.
-  DocumentProto message =
-      DocumentBuilder()
-          .SetKey("pkg$db/namespace", "message")
-          .SetSchema("Message")
-          .AddStringProperty("content", "test content")
-          .AddStringProperty("sender", "pkg$db/namespace#person")
-          .AddStringProperty("receiver", "pkg$db/namespace#person")
-          .Build();
-
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentId person_doc_id,
-                             this->PutAndIndexDocument(person));
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentId message_doc_id,
-                             this->PutAndIndexDocument(message));
-
-  // Sanity check.
-  ASSERT_THAT(
-      this->qualified_id_join_index_->Get(person_doc_id),
-      IsOkAndHolds(ElementsAre(
-          DocumentJoinIdPair(message_doc_id, /*joinable_property_id=*/0),
-          DocumentJoinIdPair(message_doc_id, /*joinable_property_id=*/2))));
-
-  JoinProcessor join_processor(
-      this->doc_store_.get(), this->schema_store_.get(),
-      this->qualified_id_join_index_.get(),
-      /*current_time_ms=*/this->fake_clock_.GetSystemTimeMilliseconds());
-
-  // Deleting the parent document should propagate the delete to its child
-  // document when there is at least one property with
-  // DELETE_PROPAGATION_TYPE_PROPAGATE_FROM.
-  EXPECT_THAT(
-      join_processor.GetPropagatedChildDocumentsToDelete({person_doc_id}),
-      IsOkAndHolds(UnorderedElementsAre(message_doc_id)));
-}
-
-TYPED_TEST(JoinProcessorTest,
-           GetPropagatedChildDocumentsToDelete_propagateToMultipleChildren) {
-  if (!std::is_same_v<typename TestFixture::JOIN_INDEX_TYPE,
-                      QualifiedIdJoinIndexImplV3>) {
-    GTEST_SKIP() << "Skipping test for non v3 QualifiedIdJoinIndex";
-  }
-
-  DocumentProto person = DocumentBuilder()
-                             .SetKey("pkg$db/namespace", "person")
-                             .SetSchema("Person")
-                             .AddStringProperty("Name", "Alice")
-                             .Build();
-
-  // "sender" (joinable property id 2) has
-  // DELETE_PROPAGATION_TYPE_PROPAGATE_FROM.
-  DocumentProto message1 =
-      DocumentBuilder()
-          .SetKey("pkg$db/namespace", "message1")
-          .SetSchema("Message")
-          .AddStringProperty("content", "test content")
-          .AddStringProperty("sender", "pkg$db/namespace#person")
-          .Build();
-  DocumentProto message2 =
-      DocumentBuilder()
-          .SetKey("pkg$db/namespace", "message2")
-          .SetSchema("Message")
-          .AddStringProperty("content", "test content")
-          .AddStringProperty("sender", "pkg$db/namespace#person")
-          .Build();
-
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentId person_doc_id,
-                             this->PutAndIndexDocument(person));
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentId message1_doc_id,
-                             this->PutAndIndexDocument(message1));
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentId message2_doc_id,
-                             this->PutAndIndexDocument(message2));
-
-  // Sanity check.
-  ASSERT_THAT(
-      this->qualified_id_join_index_->Get(person_doc_id),
-      IsOkAndHolds(ElementsAre(
-          DocumentJoinIdPair(message1_doc_id, /*joinable_property_id=*/2),
-          DocumentJoinIdPair(message2_doc_id, /*joinable_property_id=*/2))));
-
-  JoinProcessor join_processor(
-      this->doc_store_.get(), this->schema_store_.get(),
-      this->qualified_id_join_index_.get(),
-      /*current_time_ms=*/this->fake_clock_.GetSystemTimeMilliseconds());
-
-  // Deleting the parent document should propagate the delete to all of its
-  // child documents via the joinable property with
-  // DELETE_PROPAGATION_TYPE_PROPAGATE_FROM.
-  EXPECT_THAT(
-      join_processor.GetPropagatedChildDocumentsToDelete({person_doc_id}),
-      IsOkAndHolds(UnorderedElementsAre(message1_doc_id, message2_doc_id)));
-}
-
-TYPED_TEST(
-    JoinProcessorTest,
-    GetPropagatedChildDocumentsToDelete_propagateFromMultipleProperties) {
-  if (!std::is_same_v<typename TestFixture::JOIN_INDEX_TYPE,
-                      QualifiedIdJoinIndexImplV3>) {
-    GTEST_SKIP() << "Skipping test for non v3 QualifiedIdJoinIndex";
-  }
-
-  DocumentProto person1 = DocumentBuilder()
-                              .SetKey("pkg$db/namespace", "person1")
-                              .SetSchema("Person")
-                              .AddStringProperty("Name", "Alice")
-                              .Build();
-  DocumentProto person2 = DocumentBuilder()
-                              .SetKey("pkg$db/namespace", "person2")
-                              .SetSchema("Person")
-                              .AddStringProperty("Name", "Bob")
-                              .Build();
-
-  // "sender" (joinable property id 2) and "reporter" (joinable property id 1)
-  // have DELETE_PROPAGATION_TYPE_PROPAGATE_FROM.
-  DocumentProto message =
-      DocumentBuilder()
-          .SetKey("pkg$db/namespace", "message")
-          .SetSchema("Message")
-          .AddStringProperty("content", "test content")
-          .AddStringProperty("sender", "pkg$db/namespace#person1")
-          .AddStringProperty("reporter", "pkg$db/namespace#person2")
-          .Build();
-
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentId person1_doc_id,
-                             this->PutAndIndexDocument(person1));
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentId person2_doc_id,
-                             this->PutAndIndexDocument(person2));
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentId message_doc_id,
-                             this->PutAndIndexDocument(message));
-
-  // Sanity check.
-  ASSERT_THAT(this->qualified_id_join_index_->Get(person1_doc_id),
-              IsOkAndHolds(ElementsAre(DocumentJoinIdPair(
-                  message_doc_id, /*joinable_property_id=*/2))));
-  ASSERT_THAT(this->qualified_id_join_index_->Get(person2_doc_id),
-              IsOkAndHolds(ElementsAre(DocumentJoinIdPair(
-                  message_doc_id, /*joinable_property_id=*/1))));
-
-  JoinProcessor join_processor(
-      this->doc_store_.get(), this->schema_store_.get(),
-      this->qualified_id_join_index_.get(),
-      /*current_time_ms=*/this->fake_clock_.GetSystemTimeMilliseconds());
-
-  // message document should be propagated to be deleted from both person1 and
-  // person2 via "sender" and "reporter" properties respectively.
-  EXPECT_THAT(join_processor.GetPropagatedChildDocumentsToDelete(
-                  {person1_doc_id, person2_doc_id}),
-              IsOkAndHolds(UnorderedElementsAre(message_doc_id)));
-}
-
-TYPED_TEST(JoinProcessorTest,
-           GetPropagatedChildDocumentsToDelete_propagateToGrandChildren) {
-  if (!std::is_same_v<typename TestFixture::JOIN_INDEX_TYPE,
-                      QualifiedIdJoinIndexImplV3>) {
-    GTEST_SKIP() << "Skipping test for non v3 QualifiedIdJoinIndex";
-  }
-
-  // Create the following relations:
-  //
-  //                         ("object") - label1
-  //                        /
-  //             message1 <-
-  //            /           \
-  //      ("sender")         ("softLink") - label2
-  //          /
-  // person <-
-  //          \
-  //      ("receiver")       ("object") - label3
-  //            \           /
-  //             message2 <-
-  //                        \
-  //                         ("softLink") - label4
-  //
-  // Note: "sender" and "object" have DELETE_PROPAGATION_TYPE_PROPAGATE_FROM,
-  //       while "receiver" and "softLink" have DELETE_PROPAGATION_TYPE_NONE.
-  DocumentProto person = DocumentBuilder()
-                             .SetKey("pkg$db/namespace", "person")
-                             .SetSchema("Person")
-                             .AddStringProperty("Name", "Alice")
-                             .Build();
-
-  DocumentProto message1 =
-      DocumentBuilder()
-          .SetKey("pkg$db/namespace", "message1")
-          .SetSchema("Message")
-          .AddStringProperty("content", "test content")
-          .AddStringProperty("sender", "pkg$db/namespace#person")
-          .Build();
-  DocumentProto message2 =
-      DocumentBuilder()
-          .SetKey("pkg$db/namespace", "message2")
-          .SetSchema("Message")
-          .AddStringProperty("content", "test content")
-          .AddStringProperty("receiver", "pkg$db/namespace#person")
-          .Build();
-
-  DocumentProto label1 =
-      DocumentBuilder()
-          .SetKey("pkg$db/namespace", "label1")
-          .SetSchema("Label")
-          .AddStringProperty("text", " label1")
-          .AddStringProperty("object", "pkg$db/namespace#message1")
-          .Build();
-  DocumentProto label2 =
-      DocumentBuilder()
-          .SetKey("pkg$db/namespace", "label2")
-          .SetSchema("Label")
-          .AddStringProperty("text", " label2")
-          .AddStringProperty("softLink", "pkg$db/namespace#message1")
-          .Build();
-  DocumentProto label3 =
-      DocumentBuilder()
-          .SetKey("pkg$db/namespace", "label3")
-          .SetSchema("Label")
-          .AddStringProperty("text", " label3")
-          .AddStringProperty("object", "pkg$db/namespace#message2")
-          .Build();
-  DocumentProto label4 =
-      DocumentBuilder()
-          .SetKey("pkg$db/namespace", "label4")
-          .SetSchema("Label")
-          .AddStringProperty("text", " label4")
-          .AddStringProperty("softLink", "pkg$db/namespace#message2")
-          .Build();
-
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentId person_doc_id,
-                             this->PutAndIndexDocument(person));
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentId message1_doc_id,
-                             this->PutAndIndexDocument(message1));
-  ICING_ASSERT_OK(this->PutAndIndexDocument(message2));
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentId label1_doc_id,
-                             this->PutAndIndexDocument(label1));
-  ICING_ASSERT_OK(this->PutAndIndexDocument(label2));
-  ICING_ASSERT_OK(this->PutAndIndexDocument(label3));
-  ICING_ASSERT_OK(this->PutAndIndexDocument(label4));
-
-  JoinProcessor join_processor(
-      this->doc_store_.get(), this->schema_store_.get(),
-      this->qualified_id_join_index_.get(),
-      /*current_time_ms=*/this->fake_clock_.GetSystemTimeMilliseconds());
-
-  // - For children with type "Message", only message1 should be propagated to
-  //   be deleted.
-  // - For grand children with type "Label", only label1 should be propagated to
-  //   be deleted.
-  EXPECT_THAT(
-      join_processor.GetPropagatedChildDocumentsToDelete({person_doc_id}),
-      IsOkAndHolds(UnorderedElementsAre(message1_doc_id, label1_doc_id)));
-}
-
-TYPED_TEST(JoinProcessorTest,
-           GetPropagatedChildDocumentsToDelete_cycleReference) {
-  if (!std::is_same_v<typename TestFixture::JOIN_INDEX_TYPE,
-                      QualifiedIdJoinIndexImplV3>) {
-    GTEST_SKIP() << "Skipping test for non v3 QualifiedIdJoinIndex";
-  }
-
-  // Create the following relations:
-  //
-  // label1 <- label2 <- label3
-  //   |                   ^
-  //   |                   |
-  //   +-------------------+
-  DocumentProto label1 =
-      DocumentBuilder()
-          .SetKey("pkg$db/namespace", "label1")
-          .SetSchema("Label")
-          .AddStringProperty("text", " label1")
-          .AddStringProperty("object", "pkg$db/namespace#label3")
-          .Build();
-  DocumentProto label2 =
-      DocumentBuilder()
-          .SetKey("pkg$db/namespace", "label2")
-          .SetSchema("Label")
-          .AddStringProperty("text", " label2")
-          .AddStringProperty("object", "pkg$db/namespace#label1")
-          .Build();
-  DocumentProto label3 =
-      DocumentBuilder()
-          .SetKey("pkg$db/namespace", "label3")
-          .SetSchema("Label")
-          .AddStringProperty("text", " label3")
-          .AddStringProperty("object", "pkg$db/namespace#label2")
-          .Build();
-
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentId label1_doc_id_old,
-                             this->PutAndIndexDocument(label1));
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentId label2_doc_id,
-                             this->PutAndIndexDocument(label2));
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentId label3_doc_id,
-                             this->PutAndIndexDocument(label3));
-  // Put label1 again, due to the requirement of the join index: parent document
-  // must be present before the child document. This will make the relation data
-  // "label1 -> label3" present in the join index.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentId label1_doc_id_new,
-                             this->PutAndIndexDocument(label1));
-
-  // Sanity check for migration: put label1 2nd time.
-  ASSERT_THAT(label1_doc_id_new, Ne(label1_doc_id_old));
-  // Old label1's doc id should get children = [].
-  ASSERT_THAT(this->qualified_id_join_index_->Get(label1_doc_id_old),
-              IsOkAndHolds(IsEmpty()));
-  // New label1's doc id should get children = [label2_doc_id].
-  ASSERT_THAT(this->qualified_id_join_index_->Get(label1_doc_id_new),
-              IsOkAndHolds(ElementsAre(DocumentJoinIdPair(
-                  label2_doc_id, /*joinable_property_id=*/0))));
-  // label2 should get children = [label3_doc_id].
-  ASSERT_THAT(this->qualified_id_join_index_->Get(label2_doc_id),
-              IsOkAndHolds(ElementsAre(DocumentJoinIdPair(
-                  label3_doc_id, /*joinable_property_id=*/0))));
-  // label3 should get children = [label1_doc_id_new]. label1_doc_id_old will
-  // not be returned because when putting label1 for the 1st time, label3 was
-  // not present yet, and label1_doc_id_old will not be added to label3's
-  // children list.
-  ASSERT_THAT(this->qualified_id_join_index_->Get(label3_doc_id),
-              IsOkAndHolds(
-                  ElementsAre(DocumentJoinIdPair(label1_doc_id_new,
-                                                 /*joinable_property_id=*/0))));
-
-  JoinProcessor join_processor(
-      this->doc_store_.get(), this->schema_store_.get(),
-      this->qualified_id_join_index_.get(),
-      /*current_time_ms=*/this->fake_clock_.GetSystemTimeMilliseconds());
-
-  // Call GetPropagatedChildDocumentsToDelete with {label1_doc_id_new}:
-  // - Propagate to label2_doc_id from label1_doc_id_new.
-  // - Propagate to label3_doc_id from label2_doc_id.
-  // - When trying to propage label3_doc_id to its children =
-  //   [label1_doc_id_new]:
-  //   - label1_doc_id_new is already deleted, so it should not be propagated
-  //     again.
-  //   - There should be no infinite propagation loop.
-  EXPECT_THAT(
-      join_processor.GetPropagatedChildDocumentsToDelete({label1_doc_id_new}),
-      IsOkAndHolds(UnorderedElementsAre(label2_doc_id, label3_doc_id)));
-
-  // Call GetPropagatedChildDocumentsToDelete with {label2_doc_id}.
-  EXPECT_THAT(
-      join_processor.GetPropagatedChildDocumentsToDelete({label2_doc_id}),
-      IsOkAndHolds(UnorderedElementsAre(label1_doc_id_new, label3_doc_id)));
-
-  // Call GetPropagatedChildDocumentsToDelete with {label3_doc_id}.
-  EXPECT_THAT(
-      join_processor.GetPropagatedChildDocumentsToDelete({label3_doc_id}),
-      IsOkAndHolds(UnorderedElementsAre(label1_doc_id_new, label2_doc_id)));
-}
-
-TYPED_TEST(JoinProcessorTest,
-           GetPropagatedChildDocumentsToDelete_selfCycleReference) {
-  if (!std::is_same_v<typename TestFixture::JOIN_INDEX_TYPE,
-                      QualifiedIdJoinIndexImplV3>) {
-    GTEST_SKIP() << "Skipping test for non v3 QualifiedIdJoinIndex";
-  }
-
-  DocumentProto label =
-      DocumentBuilder()
-          .SetKey("pkg$db/namespace", "label")
-          .SetSchema("Label")
-          .AddStringProperty("text", " label")
-          .AddStringProperty("object", "pkg$db/namespace#label")
-          .Build();
-
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentId label_doc_id,
-                             this->PutAndIndexDocument(label));
-
-  // Sanity check.
-  ASSERT_THAT(this->qualified_id_join_index_->Get(label_doc_id),
-              IsOkAndHolds(ElementsAre(DocumentJoinIdPair(
-                  label_doc_id, /*joinable_property_id=*/0))));
-
-  JoinProcessor join_processor(
-      this->doc_store_.get(), this->schema_store_.get(),
-      this->qualified_id_join_index_.get(),
-      /*current_time_ms=*/this->fake_clock_.GetSystemTimeMilliseconds());
-
-  // Call GetPropagatedChildDocumentsToDelete with {label_doc_id}: should get
-  // nothing since label_doc_id is already in the deleted set. Also there
-  // should be no infinite propagation loop.
-  EXPECT_THAT(
-      join_processor.GetPropagatedChildDocumentsToDelete({label_doc_id}),
-      IsOkAndHolds(IsEmpty()));
-}
-
-TYPED_TEST(
-    JoinProcessorTest,
-    GetPropagatedChildDocumentsToDelete_shouldPropagateToExpiredDocuments) {
-  if (!std::is_same_v<typename TestFixture::JOIN_INDEX_TYPE,
-                      QualifiedIdJoinIndexImplV3>) {
-    GTEST_SKIP() << "Skipping test for non v3 QualifiedIdJoinIndex";
-  }
-
-  constexpr int64_t kCreationTimestampMs = 1000;
-  constexpr int64_t kShortTtlMs = 1000;
-  constexpr int64_t kLongTtlMs = 10000;
-
-  // Create the following relations:
-  //
-  // label1 <- label2 <- label3
-  DocumentProto label1 = DocumentBuilder()
-                             .SetKey("pkg$db/namespace", "label1")
-                             .SetSchema("Label")
-                             .AddStringProperty("text", " label1")
-                             .SetCreationTimestampMs(kCreationTimestampMs)
-                             .SetTtlMs(kLongTtlMs)
-                             .Build();
-  DocumentProto label2 =
-      DocumentBuilder()
-          .SetKey("pkg$db/namespace", "label2")
-          .SetSchema("Label")
-          .AddStringProperty("text", " label2")
-          .AddStringProperty("object", "pkg$db/namespace#label1")
-          .SetCreationTimestampMs(kCreationTimestampMs)
-          .SetTtlMs(kShortTtlMs)
-          .Build();
-  DocumentProto label3 =
-      DocumentBuilder()
-          .SetKey("pkg$db/namespace", "label3")
-          .SetSchema("Label")
-          .AddStringProperty("text", " label3")
-          .AddStringProperty("object", "pkg$db/namespace#label2")
-          .SetCreationTimestampMs(kCreationTimestampMs)
-          .SetTtlMs(kLongTtlMs)
-          .Build();
-
-  this->fake_clock_.SetSystemTimeMilliseconds(kCreationTimestampMs);
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentId label1_doc_id,
-                             this->PutAndIndexDocument(label1));
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentId label2_doc_id,
-                             this->PutAndIndexDocument(label2));
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentId label3_doc_id,
-                             this->PutAndIndexDocument(label3));
-
-  // Adjust the clock to expire label2.
-  int64_t current_time_ms = kCreationTimestampMs + kShortTtlMs + 100;
-  this->fake_clock_.SetSystemTimeMilliseconds(current_time_ms);
-
-  // Sanity check: label2 is expired, and label1 and label3 are not expired.
-  ASSERT_THAT(this->doc_store_->GetAliveDocumentFilterData(label1_doc_id,
-                                                           current_time_ms),
-              Ne(std::nullopt));
-  ASSERT_THAT(this->doc_store_->GetAliveDocumentFilterData(label2_doc_id,
-                                                           current_time_ms),
-              Eq(std::nullopt));
-  ASSERT_THAT(this->doc_store_->GetAliveDocumentFilterData(label3_doc_id,
-                                                           current_time_ms),
-              Ne(std::nullopt));
-  ASSERT_THAT(this->qualified_id_join_index_->Get(label1_doc_id),
-              IsOkAndHolds(ElementsAre(DocumentJoinIdPair(
-                  label2_doc_id, /*joinable_property_id=*/0))));
-  ASSERT_THAT(this->qualified_id_join_index_->Get(label2_doc_id),
-              IsOkAndHolds(ElementsAre(DocumentJoinIdPair(
-                  label3_doc_id, /*joinable_property_id=*/0))));
-  ASSERT_THAT(this->qualified_id_join_index_->Get(label3_doc_id),
-              IsOkAndHolds(IsEmpty()));
-
-  JoinProcessor join_processor(
-      this->doc_store_.get(), this->schema_store_.get(),
-      this->qualified_id_join_index_.get(),
-      /*current_time_ms=*/this->fake_clock_.GetSystemTimeMilliseconds());
-
-  // Call GetPropagatedChildDocumentsToDelete with {label1_doc_id}: should still
-  // propagate to label2_doc_id and label3_doc_id even though label2_doc_id is
-  // expired.
-  EXPECT_THAT(
-      join_processor.GetPropagatedChildDocumentsToDelete({label1_doc_id}),
-      IsOkAndHolds(UnorderedElementsAre(label2_doc_id, label3_doc_id)));
-}
-
-TYPED_TEST(
-    JoinProcessorTest,
-    GetPropagatedChildDocumentsToDelete_shouldNotPropagateToDeletedDocuments) {
-  if (!std::is_same_v<typename TestFixture::JOIN_INDEX_TYPE,
-                      QualifiedIdJoinIndexImplV3>) {
-    GTEST_SKIP() << "Skipping test for non v3 QualifiedIdJoinIndex";
-  }
-
-  // Create the following relations:
-  //
-  // label1 <- label2
-  DocumentProto label1 = DocumentBuilder()
-                             .SetKey("pkg$db/namespace", "label1")
-                             .SetSchema("Label")
-                             .AddStringProperty("text", " label1")
-                             .Build();
-  DocumentProto label2 =
-      DocumentBuilder()
-          .SetKey("pkg$db/namespace", "label2")
-          .SetSchema("Label")
-          .AddStringProperty("text", " label2")
-          .AddStringProperty("object", "pkg$db/namespace#label1")
-          .Build();
-
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentId label1_doc_id,
-                             this->PutAndIndexDocument(label1));
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentId label2_doc_id,
-                             this->PutAndIndexDocument(label2));
-
-  // Delete label2.
-  ICING_ASSERT_OK(this->doc_store_->Delete(
-      label2_doc_id, this->fake_clock_.GetSystemTimeMilliseconds()));
-
-  // Sanity check: label2 is deleted, but join data is still present.
-  ASSERT_THAT(this->doc_store_->GetAliveDocumentFilterData(
-                  label2_doc_id, this->fake_clock_.GetSystemTimeMilliseconds()),
-              Eq(std::nullopt));
-  ASSERT_THAT(this->qualified_id_join_index_->Get(label1_doc_id),
-              IsOkAndHolds(ElementsAre(DocumentJoinIdPair(
-                  label2_doc_id, /*joinable_property_id=*/0))));
-  ASSERT_THAT(this->qualified_id_join_index_->Get(label2_doc_id),
-              IsOkAndHolds(IsEmpty()));
-
-  JoinProcessor join_processor(
-      this->doc_store_.get(), this->schema_store_.get(),
-      this->qualified_id_join_index_.get(),
-      /*current_time_ms=*/this->fake_clock_.GetSystemTimeMilliseconds());
-
-  // Call GetPropagatedChildDocumentsToDelete with {label1_doc_id}: should not
-  // propagate to label2_doc_id.
-  EXPECT_THAT(
-      join_processor.GetPropagatedChildDocumentsToDelete({label1_doc_id}),
-      IsOkAndHolds(IsEmpty()));
-}
-
 // TODO(b/256022027): add unit tests for non-joinable property. If joinable
 //                    value type is unset, then qualifed id join should not
 //                    include the child document even if it contains a valid

diff --git a/icing/join/qualified-id-join-index-impl-v2.cc b/icing/join/qualified-id-join-index-impl-v2.cc
index fb5bd33..1d1b2ed 100644
--- a/icing/join/qualified-id-join-index-impl-v2.cc
+++ b/icing/join/qualified-id-join-index-impl-v2.cc

@@ -38,6 +38,7 @@
 #include "icing/schema/joinable-property.h"
 #include "icing/store/document-filter-data.h"
 #include "icing/store/document-id.h"
+#include "icing/store/document-store.h"
 #include "icing/store/key-mapper.h"
 #include "icing/store/namespace-id-fingerprint.h"
 #include "icing/store/namespace-id.h"
@@ -296,6 +297,7 @@
 }
 
 libtextclassifier3::Status QualifiedIdJoinIndexImplV2::Optimize(
+    const DocumentStore* /*document_store*/,
     const std::vector<DocumentId>& document_id_old_to_new,
     const std::vector<NamespaceId>& namespace_id_old_to_new,
     DocumentId new_last_added_document_id) {
@@ -336,9 +338,9 @@
   }
 
   // Reinitialize qualified id join index.
-  if (!filesystem_.PRead(GetMetadataFilePath(working_path_).c_str(),
-                         metadata_buffer_.get(), kMetadataFileSize,
-                         /*offset=*/0)) {
+  if (filesystem_.PRead(GetMetadataFilePath(working_path_).c_str(),
+                        metadata_buffer_.get(), kMetadataFileSize,
+                        /*offset=*/0) != kMetadataFileSize) {
     return absl_ports::InternalError("Fail to read metadata file");
   }
   ICING_ASSIGN_OR_RETURN(
@@ -461,9 +463,9 @@
     bool pre_mapping_fbv) {
   // PRead metadata file.
   auto metadata_buffer = std::make_unique<uint8_t[]>(kMetadataFileSize);
-  if (!filesystem.PRead(GetMetadataFilePath(working_path).c_str(),
-                        metadata_buffer.get(), kMetadataFileSize,
-                        /*offset=*/0)) {
+  if (filesystem.PRead(GetMetadataFilePath(working_path).c_str(),
+                       metadata_buffer.get(), kMetadataFileSize,
+                       /*offset=*/0) != kMetadataFileSize) {
     return absl_ports::InternalError("Fail to read metadata file");
   }
 

diff --git a/icing/join/qualified-id-join-index-impl-v2.h b/icing/join/qualified-id-join-index-impl-v2.h
index 379f0b8..08b3e94 100644
--- a/icing/join/qualified-id-join-index-impl-v2.h
+++ b/icing/join/qualified-id-join-index-impl-v2.h

@@ -34,9 +34,11 @@
 #include "icing/join/posting-list-join-data-accessor.h"
 #include "icing/join/posting-list-join-data-serializer.h"
 #include "icing/join/qualified-id-join-index.h"
+#include "icing/join/qualified-id.h"
 #include "icing/schema/joinable-property.h"
 #include "icing/store/document-filter-data.h"
 #include "icing/store/document-id.h"
+#include "icing/store/document-store.h"
 #include "icing/store/key-mapper.h"
 #include "icing/store/namespace-id-fingerprint.h"
 #include "icing/store/namespace-id.h"
@@ -158,8 +160,16 @@
   }
 
   // v3 only API. Returns UNIMPLEMENTED_ERROR.
-  libtextclassifier3::StatusOr<std::vector<DocumentJoinIdPair>> Get(
-      DocumentId parent_document_id) const override {
+  libtextclassifier3::Status Put(
+      const DocumentStore* document_store,
+      const DocumentJoinIdPair& child_document_join_id_pair,
+      std::vector<QualifiedId>&& parent_qualified_ids) override {
+    return absl_ports::UnimplementedError("This API is not supported in V2");
+  }
+
+  // v3 only API. Returns UNIMPLEMENTED_ERROR.
+  libtextclassifier3::StatusOr<DocumentJoinIdPairArrayView>
+  GetDocumentJoinIdPairArrayView(DocumentId parent_document_id) const override {
     return absl_ports::UnimplementedError("This API is not supported in V2");
   }
 
@@ -180,7 +190,16 @@
     return libtextclassifier3::Status::OK;
   }
 
+  // No-op since v2 stores parent information in (namespace_id,
+  // fingerprint(uri)) format and does not require parent migration.
+  libtextclassifier3::Status MigrateParent(
+      const QualifiedId& parent_qualified_id,
+      DocumentId new_document_id) override {
+    return libtextclassifier3::Status::OK;
+  }
+
   libtextclassifier3::Status Optimize(
+      const DocumentStore* document_store,
       const std::vector<DocumentId>& document_id_old_to_new,
       const std::vector<NamespaceId>& namespace_id_old_to_new,
       DocumentId new_last_added_document_id) override;

diff --git a/icing/join/qualified-id-join-index-impl-v2_test.cc b/icing/join/qualified-id-join-index-impl-v2_test.cc
index 01c972e..6e89a26 100644
--- a/icing/join/qualified-id-join-index-impl-v2_test.cc
+++ b/icing/join/qualified-id-join-index-impl-v2_test.cc

@@ -131,7 +131,7 @@
       filesystem_.PRead(metadata_file_path.c_str(), metadata_buffer.get(),
                         QualifiedIdJoinIndexImplV2::kMetadataFileSize,
                         /*offset=*/0),
-      IsTrue());
+      Eq(QualifiedIdJoinIndexImplV2::kMetadataFileSize));
 
   // Check info section
   const Info* info = reinterpret_cast<const Info*>(
@@ -372,7 +372,7 @@
     ASSERT_THAT(filesystem_.PRead(metadata_sfd.get(), metadata_buffer.get(),
                                   QualifiedIdJoinIndexImplV2::kMetadataFileSize,
                                   /*offset=*/0),
-                IsTrue());
+                Eq(QualifiedIdJoinIndexImplV2::kMetadataFileSize));
 
     // Manually change magic and update checksum
     Crcs* crcs = reinterpret_cast<Crcs*>(
@@ -425,7 +425,7 @@
     ASSERT_THAT(filesystem_.PRead(metadata_sfd.get(), metadata_buffer.get(),
                                   QualifiedIdJoinIndexImplV2::kMetadataFileSize,
                                   /*offset=*/0),
-                IsTrue());
+                Eq(QualifiedIdJoinIndexImplV2::kMetadataFileSize));
 
     // Manually corrupt all_crc
     Crcs* crcs = reinterpret_cast<Crcs*>(
@@ -474,7 +474,7 @@
     ASSERT_THAT(filesystem_.PRead(metadata_sfd.get(), metadata_buffer.get(),
                                   QualifiedIdJoinIndexImplV2::kMetadataFileSize,
                                   /*offset=*/0),
-                IsTrue());
+                Eq(QualifiedIdJoinIndexImplV2::kMetadataFileSize));
 
     // Modify info, but don't update the checksum. This would be similar to
     // corruption of info.
@@ -850,9 +850,10 @@
   namespace_id_old_to_new[4] = 0;
 
   DocumentId new_last_added_document_id = 2;
-  EXPECT_THAT(index->Optimize(document_id_old_to_new, namespace_id_old_to_new,
-                              new_last_added_document_id),
-              IsOk());
+  EXPECT_THAT(
+      index->Optimize(/*document_store=*/nullptr, document_id_old_to_new,
+                      namespace_id_old_to_new, new_last_added_document_id),
+      IsOk());
   EXPECT_THAT(index, Pointee(SizeIs(3)));
   EXPECT_THAT(index->last_added_document_id(), Eq(new_last_added_document_id));
 
@@ -961,9 +962,10 @@
   std::vector<NamespaceId> namespace_id_old_to_new = {0, 1};
 
   DocumentId new_last_added_document_id = 2;
-  EXPECT_THAT(index->Optimize(document_id_old_to_new, namespace_id_old_to_new,
-                              new_last_added_document_id),
-              IsOk());
+  EXPECT_THAT(
+      index->Optimize(/*document_store=*/nullptr, document_id_old_to_new,
+                      namespace_id_old_to_new, new_last_added_document_id),
+      IsOk());
   EXPECT_THAT(index, Pointee(SizeIs(4)));
   EXPECT_THAT(index->last_added_document_id(), Eq(new_last_added_document_id));
 
@@ -1033,7 +1035,8 @@
   std::vector<NamespaceId> namespace_id_old_to_new = {0, 1};
 
   EXPECT_THAT(
-      index->Optimize(document_id_old_to_new, namespace_id_old_to_new,
+      index->Optimize(/*document_store=*/nullptr, document_id_old_to_new,
+                      namespace_id_old_to_new,
                       /*new_last_added_document_id=*/kInvalidDocumentId),
       StatusIs(libtextclassifier3::StatusCode::INTERNAL,
                HasSubstr("Qualified id join index data document id is out of "
@@ -1089,7 +1092,8 @@
   std::vector<NamespaceId> namespace_id_old_to_new = {0, 1};
 
   EXPECT_THAT(
-      index->Optimize(document_id_old_to_new, namespace_id_old_to_new,
+      index->Optimize(/*document_store=*/nullptr, document_id_old_to_new,
+                      namespace_id_old_to_new,
                       /*new_last_added_document_id=*/kInvalidDocumentId),
       IsOk());
   EXPECT_THAT(index->last_added_document_id(), Eq(kInvalidDocumentId));
@@ -1159,9 +1163,10 @@
   namespace_id_old_to_new[5] = 0;
 
   DocumentId new_last_added_document_id = 21;
-  EXPECT_THAT(index->Optimize(document_id_old_to_new, namespace_id_old_to_new,
-                              new_last_added_document_id),
-              IsOk());
+  EXPECT_THAT(
+      index->Optimize(/*document_store=*/nullptr, document_id_old_to_new,
+                      namespace_id_old_to_new, new_last_added_document_id),
+      IsOk());
   EXPECT_THAT(index, Pointee(SizeIs(4)));
   EXPECT_THAT(index->last_added_document_id(), Eq(new_last_added_document_id));
 
@@ -1250,9 +1255,10 @@
   std::vector<NamespaceId> namespace_id_old_to_new = {2, 0, 1};
 
   DocumentId new_last_added_document_id = 1;
-  EXPECT_THAT(index->Optimize(document_id_old_to_new, namespace_id_old_to_new,
-                              new_last_added_document_id),
-              IsOk());
+  EXPECT_THAT(
+      index->Optimize(/*document_store=*/nullptr, document_id_old_to_new,
+                      namespace_id_old_to_new, new_last_added_document_id),
+      IsOk());
   EXPECT_THAT(index, Pointee(SizeIs(4)));
   EXPECT_THAT(index->last_added_document_id(), Eq(new_last_added_document_id));
 
@@ -1303,7 +1309,8 @@
   std::vector<NamespaceId> namespace_id_old_to_new = {kInvalidNamespaceId};
 
   EXPECT_THAT(
-      index->Optimize(document_id_old_to_new, namespace_id_old_to_new,
+      index->Optimize(/*document_store=*/nullptr, document_id_old_to_new,
+                      namespace_id_old_to_new,
                       /*new_last_added_document_id=*/kInvalidDocumentId),
       StatusIs(libtextclassifier3::StatusCode::INTERNAL,
                HasSubstr("Qualified id join index data ref namespace id is out "
@@ -1348,7 +1355,8 @@
   std::vector<NamespaceId> namespace_id_old_to_new(3, kInvalidNamespaceId);
 
   EXPECT_THAT(
-      index->Optimize(document_id_old_to_new, namespace_id_old_to_new,
+      index->Optimize(/*document_store=*/nullptr, document_id_old_to_new,
+                      namespace_id_old_to_new,
                       /*new_last_added_document_id=*/kInvalidDocumentId),
       IsOk());
   EXPECT_THAT(index->last_added_document_id(), Eq(kInvalidDocumentId));

diff --git a/icing/join/qualified-id-join-index-impl-v3.cc b/icing/join/qualified-id-join-index-impl-v3.cc
index bff564e..7771fc7 100644
--- a/icing/join/qualified-id-join-index-impl-v3.cc
+++ b/icing/join/qualified-id-join-index-impl-v3.cc

@@ -34,8 +34,12 @@
 #include "icing/file/memory-mapped-file.h"
 #include "icing/join/document-join-id-pair.h"
 #include "icing/join/qualified-id-join-index.h"
+#include "icing/join/qualified-id.h"
 #include "icing/store/document-id.h"
+#include "icing/store/document-store.h"
+#include "icing/store/key-mapper.h"
 #include "icing/store/namespace-id.h"
+#include "icing/store/persistent-hash-map-key-mapper.h"
 #include "icing/util/crc32.h"
 #include "icing/util/logging.h"
 #include "icing/util/math-util.h"
@@ -49,6 +53,18 @@
 
 namespace {
 
+// This is the same as the max number of entries of document_key_mapper_ in
+// DocumentStore.
+constexpr uint32_t kQualifiedIdMapperMaxNumEntries = kMaxDocumentId + 1;
+
+// - Key (QualifiedId): 22 bytes
+//   - namespace (10 bytes)
+//   - '#' (1 byte)
+//   - uri (10 bytes)
+//   - '\0' (1 byte)
+// - Value (ArrayInfo): 12 bytes
+constexpr uint32_t kQualifiedIdMapperKVByteSize = 34;
+
 std::string MakeMetadataFilePath(std::string_view working_path) {
   return absl_ports::StrCat(working_path, "/metadata");
 }
@@ -59,11 +75,28 @@
                             "/parent_document_id_to_child_array_info");
 }
 
+std::string MakeParentQualifiedIdToChildArrayInfoFilePath(
+    std::string_view working_path) {
+  return absl_ports::StrCat(working_path,
+                            "/parent_qualified_id_to_child_array_info");
+}
+
 std::string MakeChildDocumentJoinIdPairArrayFilePath(
     std::string_view working_path) {
   return absl_ports::StrCat(working_path, "/child_document_join_id_pair_array");
 }
 
+DocumentId GetDocumentId(const DocumentStore& document_store,
+                         const QualifiedId& parent_qualified_id) {
+  libtextclassifier3::StatusOr<DocumentId> parent_doc_id_or =
+      document_store.GetDocumentId(parent_qualified_id.name_space(),
+                                   parent_qualified_id.uri());
+  if (!parent_doc_id_or.ok()) {
+    return kInvalidDocumentId;
+  }
+  return parent_doc_id_or.ValueOrDie();
+}
+
 }  // namespace
 
 /* static */ libtextclassifier3::StatusOr<
@@ -76,10 +109,14 @@
   bool parent_document_id_to_child_array_info_file_exists =
       filesystem.FileExists(
           MakeParentDocumentIdToChildArrayInfoFilePath(working_path).c_str());
+  bool parent_qualified_id_to_child_array_info_file_exists =
+      filesystem.DirectoryExists(
+          MakeParentQualifiedIdToChildArrayInfoFilePath(working_path).c_str());
   bool child_document_join_id_pair_array_file_exists = filesystem.FileExists(
       MakeChildDocumentJoinIdPairArrayFilePath(working_path).c_str());
 
-  // If all files exist, initialize from existing files.
+  // If all files exist (except parent_qualified_id_to_child_array_info),
+  // initialize from existing files.
   if (metadata_file_exists &&
       parent_document_id_to_child_array_info_file_exists &&
       child_document_join_id_pair_array_file_exists) {
@@ -90,6 +127,7 @@
   // If all files don't exist, initialize new files.
   if (!metadata_file_exists &&
       !parent_document_id_to_child_array_info_file_exists &&
+      !parent_qualified_id_to_child_array_info_file_exists &&
       !child_document_join_id_pair_array_file_exists) {
     return InitializeNewFiles(filesystem, std::move(working_path),
                               feature_flags);
@@ -109,6 +147,7 @@
   }
 }
 
+// Deprecated
 libtextclassifier3::Status QualifiedIdJoinIndexImplV3::Put(
     const DocumentJoinIdPair& child_document_join_id_pair,
     std::vector<DocumentId>&& parent_document_ids) {
@@ -145,15 +184,51 @@
   return libtextclassifier3::Status::OK;
 }
 
-libtextclassifier3::StatusOr<std::vector<DocumentJoinIdPair>>
-QualifiedIdJoinIndexImplV3::Get(DocumentId parent_document_id) const {
+libtextclassifier3::Status QualifiedIdJoinIndexImplV3::Put(
+    const DocumentStore* document_store,
+    const DocumentJoinIdPair& child_document_join_id_pair,
+    std::vector<QualifiedId>&& parent_qualified_ids) {
+  if (!feature_flags_.enable_non_existent_qualified_id_join()) {
+    return absl_ports::InvalidArgumentError(
+        "Put with QualifiedId is not enabled yet!");
+  }
+  ICING_RETURN_ERROR_IF_NULL(document_store);
+
+  if (!child_document_join_id_pair.is_valid()) {
+    return absl_ports::InvalidArgumentError("Invalid child DocumentJoinIdPair");
+  }
+  if (parent_qualified_ids.empty()) {
+    return libtextclassifier3::Status::OK;
+  }
+  SetDirty();
+
+  // Sort and dedupe parent qualified ids.
+  std::sort(parent_qualified_ids.begin(), parent_qualified_ids.end());
+  auto last =
+      std::unique(parent_qualified_ids.begin(), parent_qualified_ids.end());
+  parent_qualified_ids.erase(last, parent_qualified_ids.end());
+
+  // Append child_document_join_id_pair to each parent's DocumentJoinIdPair
+  // array.
+  for (const QualifiedId& parent_qualified_id : parent_qualified_ids) {
+    ICING_RETURN_IF_ERROR(AppendChildDocumentJoinIdPairsForParent(
+        parent_qualified_id,
+        GetDocumentId(*document_store, parent_qualified_id),
+        {child_document_join_id_pair}));
+  }
+  return libtextclassifier3::Status::OK;
+}
+
+libtextclassifier3::StatusOr<QualifiedIdJoinIndex::DocumentJoinIdPairArrayView>
+QualifiedIdJoinIndexImplV3::GetDocumentJoinIdPairArrayView(
+    DocumentId parent_document_id) const {
   if (parent_document_id < 0 || parent_document_id == kInvalidDocumentId) {
     return absl_ports::InvalidArgumentError("Invalid parent document id");
   }
 
   if (parent_document_id >=
       parent_document_id_to_child_array_info_->num_elements()) {
-    return std::vector<DocumentJoinIdPair>();
+    return DocumentJoinIdPairArrayView(/*data=*/nullptr, /*len=*/0);
   }
 
   // Get the child array info for the parent.
@@ -161,7 +236,7 @@
       const ArrayInfo* array_info,
       parent_document_id_to_child_array_info_->Get(parent_document_id));
   if (!array_info->IsValid()) {
-    return std::vector<DocumentJoinIdPair>();
+    return DocumentJoinIdPairArrayView(/*data=*/nullptr, /*len=*/0);
   }
 
   // Safe check to avoid out-of-bound access. This should never happen unless
@@ -174,11 +249,11 @@
         std::to_string(child_document_join_id_pair_array_->num_elements())));
   }
 
-  // Get the DocumentJoinIdPair array and return the child DocumentJoinIdPairs.
+  // Get the DocumentJoinIdPair array ptr and return the array view.
   ICING_ASSIGN_OR_RETURN(
       const DocumentJoinIdPair* ptr,
       child_document_join_id_pair_array_->Get(array_info->index));
-  return std::vector<DocumentJoinIdPair>(ptr, ptr + array_info->used_length);
+  return DocumentJoinIdPairArrayView(ptr, array_info->used_length);
 }
 
 libtextclassifier3::Status QualifiedIdJoinIndexImplV3::MigrateParent(
@@ -205,6 +280,10 @@
     return libtextclassifier3::Status::OK;
   }
 
+  // Set dirty for the entire storage here once we make sure there are children
+  // to migrate for this parent doc.
+  SetDirty();
+
   ICING_ASSIGN_OR_RETURN(
       bool is_extended,
       ExtendParentDocumentIdToChildArrayInfoIfNecessary(new_document_id));
@@ -224,10 +303,41 @@
   return libtextclassifier3::Status::OK;
 }
 
+libtextclassifier3::Status QualifiedIdJoinIndexImplV3::MigrateParent(
+    const QualifiedId& parent_qualified_id, DocumentId new_document_id) {
+  if (!feature_flags_.enable_non_existent_qualified_id_join()) {
+    return absl_ports::InvalidArgumentError(
+        "QualifiedId entry is not enabled yet in the index!");
+  }
+
+  if (!IsDocumentIdValid(new_document_id)) {
+    return absl_ports::InvalidArgumentError(
+        "Invalid new parent document id to migrate to");
+  }
+
+  std::string key = parent_qualified_id.ToString();
+  auto array_info_or = parent_qualified_id_to_child_array_info_->Get(key);
+  if (absl_ports::IsNotFound(array_info_or.status())) {
+    // The qualified id does not exist in the index, nothing to migrate.
+    return libtextclassifier3::Status::OK;
+  }
+  ICING_RETURN_IF_ERROR(array_info_or.status());
+
+  SetDirty();
+  ICING_RETURN_IF_ERROR(
+      ExtendParentDocumentIdToChildArrayInfoIfNecessary(new_document_id));
+  ICING_RETURN_IF_ERROR(parent_document_id_to_child_array_info_->Set(
+      new_document_id, array_info_or.ValueOrDie()));
+  return libtextclassifier3::Status::OK;
+}
+
 libtextclassifier3::Status QualifiedIdJoinIndexImplV3::Optimize(
+    const DocumentStore* document_store,
     const std::vector<DocumentId>& document_id_old_to_new,
-    const std::vector<NamespaceId>& namespace_id_old_to_new,
+    const std::vector<NamespaceId>& /*namespace_id_old_to_new*/,
     DocumentId new_last_added_document_id) {
+  ICING_RETURN_ERROR_IF_NULL(document_store);
+
   std::string temp_working_path = working_path_ + "_temp";
   ICING_RETURN_IF_ERROR(
       QualifiedIdJoinIndex::Discard(filesystem_, temp_working_path));
@@ -247,8 +357,8 @@
     ICING_ASSIGN_OR_RETURN(
         std::unique_ptr<QualifiedIdJoinIndexImplV3> new_index,
         Create(filesystem_, temp_working_path_ddir.dir(), feature_flags_));
-    ICING_RETURN_IF_ERROR(
-        TransferIndex(document_id_old_to_new, new_index.get()));
+    ICING_RETURN_IF_ERROR(TransferIndex(*document_store, document_id_old_to_new,
+                                        new_index.get()));
     new_index->set_last_added_document_id(new_last_added_document_id);
     new_index->SetDirty();
 
@@ -257,6 +367,7 @@
 
   // Destruct current index's storage instances to safely swap directories.
   child_document_join_id_pair_array_.reset();
+  parent_qualified_id_to_child_array_info_.reset();
   parent_document_id_to_child_array_info_.reset();
   metadata_mmapped_file_.reset();
   if (!filesystem_.SwapFiles(temp_working_path_ddir.dir().c_str(),
@@ -286,6 +397,15 @@
           /*pre_mapping_mmap_size=*/0));
 
   ICING_ASSIGN_OR_RETURN(
+      parent_qualified_id_to_child_array_info_,
+      PersistentHashMapKeyMapper<ArrayInfo>::Create(
+          filesystem_,
+          MakeParentQualifiedIdToChildArrayInfoFilePath(working_path_),
+          /*pre_mapping_fbv=*/false,
+          /*max_num_entries=*/kQualifiedIdMapperMaxNumEntries,
+          /*average_kv_byte_size=*/kQualifiedIdMapperKVByteSize));
+
+  ICING_ASSIGN_OR_RETURN(
       child_document_join_id_pair_array_,
       FileBackedVector<DocumentJoinIdPair>::Create(
           filesystem_, MakeChildDocumentJoinIdPairArrayFilePath(working_path_),
@@ -317,6 +437,25 @@
           FileBackedVector<ArrayInfo>::kMaxFileSize,
           /*pre_mapping_mmap_size=*/0));
 
+  parent_qualified_id_to_child_array_info_.reset();
+  // Discard and reinitialize parent_qualified_id_to_child_array_info_.
+  std::string parent_qualified_id_to_child_array_info_file_path =
+      MakeParentQualifiedIdToChildArrayInfoFilePath(working_path_);
+  if (!filesystem_.DeleteDirectoryRecursively(
+          parent_qualified_id_to_child_array_info_file_path.c_str())) {
+    return absl_ports::InternalError(absl_ports::StrCat(
+        "Failed to clear parent qualified id to child array info: ",
+        parent_qualified_id_to_child_array_info_file_path));
+  }
+  ICING_ASSIGN_OR_RETURN(
+      parent_qualified_id_to_child_array_info_,
+      PersistentHashMapKeyMapper<ArrayInfo>::Create(
+          filesystem_,
+          MakeParentQualifiedIdToChildArrayInfoFilePath(working_path_),
+          /*pre_mapping_fbv=*/false,
+          /*max_num_entries=*/kQualifiedIdMapperMaxNumEntries,
+          /*average_kv_byte_size=*/kQualifiedIdMapperKVByteSize));
+
   child_document_join_id_pair_array_.reset();
   // Discard and reinitialize child_document_join_id_pair_array.
   std::string child_document_join_id_pair_array_file_path =
@@ -378,6 +517,17 @@
           FileBackedVector<ArrayInfo>::kMaxFileSize,
           /*pre_mapping_mmap_size=*/0));
 
+  // Initialize parent_qualified_id_to_child_array_info.
+  ICING_ASSIGN_OR_RETURN(
+      std::unique_ptr<KeyMapper<ArrayInfo>>
+          parent_qualified_id_to_child_array_info,
+      PersistentHashMapKeyMapper<ArrayInfo>::Create(
+          filesystem,
+          MakeParentQualifiedIdToChildArrayInfoFilePath(working_path),
+          /*pre_mapping_fbv=*/false,
+          /*max_num_entries=*/kQualifiedIdMapperMaxNumEntries,
+          /*average_kv_byte_size=*/kQualifiedIdMapperKVByteSize));
+
   // Initialize child_document_join_id_pair_array.
   ICING_ASSIGN_OR_RETURN(
       std::unique_ptr<FileBackedVector<DocumentJoinIdPair>>
@@ -394,6 +544,7 @@
           filesystem, std::move(working_path),
           std::make_unique<MemoryMappedFile>(std::move(metadata_mmapped_file)),
           std::move(parent_document_id_to_child_array_info),
+          std::move(parent_qualified_id_to_child_array_info),
           std::move(child_document_join_id_pair_array), feature_flags));
   // Initialize info content.
   new_join_index->info().magic = Info::kMagic;
@@ -436,6 +587,17 @@
           FileBackedVector<ArrayInfo>::kMaxFileSize,
           /*pre_mapping_mmap_size=*/0));
 
+  // Initialize parent_qualified_id_to_child_array_info.
+  ICING_ASSIGN_OR_RETURN(
+      std::unique_ptr<KeyMapper<ArrayInfo>>
+          parent_qualified_id_to_child_array_info,
+      PersistentHashMapKeyMapper<ArrayInfo>::Create(
+          filesystem,
+          MakeParentQualifiedIdToChildArrayInfoFilePath(working_path),
+          /*pre_mapping_fbv=*/false,
+          /*max_num_entries=*/kQualifiedIdMapperMaxNumEntries,
+          /*average_kv_byte_size=*/kQualifiedIdMapperKVByteSize));
+
   // Initialize child_document_join_id_pair_array. Set mmap pre-mapping size to
   // 0, but MemoryMappedFile will still mmap to the file size.
   ICING_ASSIGN_OR_RETURN(
@@ -453,6 +615,7 @@
           filesystem, std::move(working_path),
           std::make_unique<MemoryMappedFile>(std::move(metadata_mmapped_file)),
           std::move(parent_document_id_to_child_array_info),
+          std::move(parent_qualified_id_to_child_array_info),
           std::move(child_document_join_id_pair_array), feature_flags));
 
   // Initialize existing PersistentStorage. Checksums will be validated.
@@ -460,12 +623,19 @@
 
   // Validate magic.
   if (join_index->info().magic != Info::kMagic) {
-    return absl_ports::FailedPreconditionError("Incorrect magic value");
+    ICING_LOG(ERROR) << "Invalid header magic for QualifiedIdJoinIndexImplV3 "
+                     << join_index->working_path_
+                     << ". Expected: " << Info::kMagic
+                     << ", actual: " << join_index->info().magic;
+    return absl_ports::FailedPreconditionError(absl_ports::StrCat(
+        "Invalid header magic for QualifiedIdJoinIndexImplV3: ",
+        join_index->working_path_));
   }
 
   return join_index;
 }
 
+// Deprecated
 libtextclassifier3::Status
 QualifiedIdJoinIndexImplV3::AppendChildDocumentJoinIdPairsForParent(
     DocumentId parent_document_id,
@@ -486,9 +656,25 @@
       const ArrayInfo* array_info,
       parent_document_id_to_child_array_info_->Get(parent_document_id));
   ICING_ASSIGN_OR_RETURN(
+      ArrayInfo new_array_info,
+      AppendToChildArray(*array_info, std::move(child_document_join_id_pairs)));
+
+  // Set ArrayInfo back to parent_document_id_to_child_array_info_.
+  ICING_RETURN_IF_ERROR(parent_document_id_to_child_array_info_->Set(
+      parent_document_id, new_array_info));
+
+  return libtextclassifier3::Status::OK;
+}
+
+libtextclassifier3::StatusOr<QualifiedIdJoinIndexImplV3::ArrayInfo>
+QualifiedIdJoinIndexImplV3::AppendToChildArray(
+    const ArrayInfo& current_array_info,
+    std::vector<DocumentJoinIdPair>&& child_document_join_id_pairs) {
+  ICING_ASSIGN_OR_RETURN(
       GetMutableAndExtendResult result,
       GetMutableAndExtendChildDocumentJoinIdPairArrayIfNecessary(
-          *array_info, /*num_to_add=*/child_document_join_id_pairs.size()));
+          current_array_info,
+          /*num_to_add=*/child_document_join_id_pairs.size()));
   // - [0, result.array_info.used_length) contain valid elements.
   // - We will write new elements starting from index
   //   result.array_info.used_length, and update the used_length.
@@ -497,13 +683,45 @@
                               /*len=*/child_document_join_id_pairs.size());
   result.array_info.used_length += child_document_join_id_pairs.size();
 
-  // Set ArrayInfo back to parent_document_id_to_child_array_info_.
-  ICING_RETURN_IF_ERROR(parent_document_id_to_child_array_info_->Set(
-      parent_document_id, result.array_info));
-
   // Update header.
   info().num_data += child_document_join_id_pairs.size();
 
+  return result.array_info;
+}
+
+libtextclassifier3::Status
+QualifiedIdJoinIndexImplV3::AppendChildDocumentJoinIdPairsForParent(
+    std::string_view parent_qualified_id_str, DocumentId parent_document_id,
+    std::vector<DocumentJoinIdPair>&& child_document_join_id_pairs) {
+  if (child_document_join_id_pairs.empty()) {
+    return libtextclassifier3::Status::OK;
+  }
+
+  auto array_info_or =
+      parent_qualified_id_to_child_array_info_->Get(parent_qualified_id_str);
+  ArrayInfo array_info = kInvalidArrayInfo;
+  if (array_info_or.ok()) {
+    array_info = std::move(array_info_or).ValueOrDie();
+  } else if (!absl_ports::IsNotFound(array_info_or.status())) {
+    return array_info_or.status();
+  }
+
+  ICING_ASSIGN_OR_RETURN(
+      ArrayInfo new_array_info,
+      AppendToChildArray(array_info, std::move(child_document_join_id_pairs)));
+
+  ICING_RETURN_IF_ERROR(parent_qualified_id_to_child_array_info_->Put(
+      parent_qualified_id_str, new_array_info));
+
+  // Update parent_document_id_to_child_array_info_ if the parent document has a
+  // valid document id.
+  if (IsDocumentIdValid(parent_document_id)) {
+    ICING_RETURN_IF_ERROR(
+        ExtendParentDocumentIdToChildArrayInfoIfNecessary(parent_document_id));
+    ICING_RETURN_IF_ERROR(parent_document_id_to_child_array_info_->Set(
+        parent_document_id, new_array_info));
+  }
+
   return libtextclassifier3::Status::OK;
 }
 
@@ -685,66 +903,101 @@
 }
 
 libtextclassifier3::Status QualifiedIdJoinIndexImplV3::TransferIndex(
+    const DocumentStore& document_store,
     const std::vector<DocumentId>& document_id_old_to_new,
     QualifiedIdJoinIndexImplV3* new_index) const {
-  for (DocumentId old_parent_doc_id = 0;
-       old_parent_doc_id <
-       parent_document_id_to_child_array_info_->num_elements();
-       ++old_parent_doc_id) {
-    if (old_parent_doc_id < 0 ||
-        old_parent_doc_id >= document_id_old_to_new.size()) {
-      // If it happens, then the data is corrupted. Return error and let the
-      // caller rebuild everything.
-      return absl_ports::InternalError(
-          "Qualified id join index data parent document id is out of range. "
-          "The index may have been corrupted.");
-    }
-
-    if (document_id_old_to_new[old_parent_doc_id] == kInvalidDocumentId) {
-      // Skip if the old parent document id is invalid after optimization.
-      continue;
-    }
-
-    ICING_ASSIGN_OR_RETURN(
-        const ArrayInfo* array_info,
-        parent_document_id_to_child_array_info_->Get(old_parent_doc_id));
-    if (!array_info->IsValid()) {
-      continue;
-    }
-    ICING_ASSIGN_OR_RETURN(
-        const DocumentJoinIdPair* ptr,
-        child_document_join_id_pair_array_->Get(array_info->index));
-
-    // Get all child DocumentJoinIdPairs and assign new child document ids.
-    std::vector<DocumentJoinIdPair> new_child_doc_join_id_pairs;
-    new_child_doc_join_id_pairs.reserve(array_info->length);
-    for (int i = 0; i < array_info->used_length; ++i) {
-      DocumentId old_child_doc_id = ptr[i].document_id();
-      if (old_child_doc_id < 0 ||
-          old_child_doc_id >= document_id_old_to_new.size()) {
+  // If the flag is disabled, use the old logic to transfer the index.
+  if (!feature_flags_.enable_non_existent_qualified_id_join()) {
+    for (DocumentId old_parent_doc_id = 0;
+         old_parent_doc_id <
+         parent_document_id_to_child_array_info_->num_elements();
+         ++old_parent_doc_id) {
+      if (old_parent_doc_id < 0 ||
+          old_parent_doc_id >= document_id_old_to_new.size()) {
         // If it happens, then the data is corrupted. Return error and let the
         // caller rebuild everything.
         return absl_ports::InternalError(
-            "Qualified id join index data child document id is out of range. "
+            "Qualified id join index data parent document id is out of range. "
             "The index may have been corrupted.");
       }
 
-      DocumentId new_child_doc_id = document_id_old_to_new[old_child_doc_id];
-      if (new_child_doc_id == kInvalidDocumentId) {
+      if (document_id_old_to_new[old_parent_doc_id] == kInvalidDocumentId) {
+        // Skip if the old parent document id is invalid after optimization.
         continue;
       }
 
-      new_child_doc_join_id_pairs.push_back(
-          DocumentJoinIdPair(new_child_doc_id, ptr[i].joinable_property_id()));
+      ICING_ASSIGN_OR_RETURN(
+          const ArrayInfo* array_info,
+          parent_document_id_to_child_array_info_->Get(old_parent_doc_id));
+      ICING_ASSIGN_OR_RETURN(
+          std::vector<DocumentJoinIdPair> new_child_doc_join_id_pairs,
+          GetTransferredChildDocumentJoinIdPairs(*array_info,
+                                                 document_id_old_to_new));
+      ICING_RETURN_IF_ERROR(new_index->AppendChildDocumentJoinIdPairsForParent(
+          document_id_old_to_new[old_parent_doc_id],
+          std::move(new_child_doc_join_id_pairs)));
     }
+    return libtextclassifier3::Status::OK;
+  }
 
+  // The parent qualified id map is the source of truth. Iterate through it to
+  // transfer the index.
+  auto iter = parent_qualified_id_to_child_array_info_->GetIterator();
+  while (iter->Advance()) {
+    ArrayInfo array_info = iter->GetValue();
+    ICING_ASSIGN_OR_RETURN(
+        std::vector<DocumentJoinIdPair> new_child_doc_join_id_pairs,
+        GetTransferredChildDocumentJoinIdPairs(array_info,
+                                               document_id_old_to_new));
+    ICING_ASSIGN_OR_RETURN(QualifiedId parent_qualified_id,
+                           QualifiedId::Parse(iter->GetKey()));
+    // This will automatically transfer the parent document id map (for existing
+    // documents) to the new index.
     ICING_RETURN_IF_ERROR(new_index->AppendChildDocumentJoinIdPairsForParent(
-        document_id_old_to_new[old_parent_doc_id],
+        iter->GetKey(), GetDocumentId(document_store, parent_qualified_id),
         std::move(new_child_doc_join_id_pairs)));
   }
+
   return libtextclassifier3::Status::OK;
 }
 
+libtextclassifier3::StatusOr<std::vector<DocumentJoinIdPair>>
+QualifiedIdJoinIndexImplV3::GetTransferredChildDocumentJoinIdPairs(
+    const QualifiedIdJoinIndexImplV3::ArrayInfo& array_info,
+    const std::vector<DocumentId>& document_id_old_to_new) const {
+  if (!array_info.IsValid()) {
+    return std::vector<DocumentJoinIdPair>();
+  }
+
+  ICING_ASSIGN_OR_RETURN(
+      const DocumentJoinIdPair* ptr,
+      child_document_join_id_pair_array_->Get(array_info.index));
+
+  // Get all child DocumentJoinIdPairs and assign new child document ids.
+  std::vector<DocumentJoinIdPair> new_child_doc_join_id_pairs;
+  new_child_doc_join_id_pairs.reserve(array_info.used_length);
+  for (int i = 0; i < array_info.used_length; ++i) {
+    DocumentId old_child_doc_id = ptr[i].document_id();
+    if (old_child_doc_id < 0 ||
+        old_child_doc_id >= document_id_old_to_new.size()) {
+      // If it happens, then the data is corrupted. Return error and let the
+      // caller rebuild everything.
+      return absl_ports::InternalError(
+          "Qualified id join index data child document id is out of range. "
+          "The index may have been corrupted.");
+    }
+
+    DocumentId new_child_doc_id = document_id_old_to_new[old_child_doc_id];
+    if (new_child_doc_id == kInvalidDocumentId) {
+      continue;
+    }
+
+    new_child_doc_join_id_pairs.push_back(
+        DocumentJoinIdPair(new_child_doc_id, ptr[i].joinable_property_id()));
+  }
+  return new_child_doc_join_id_pairs;
+}
+
 libtextclassifier3::Status QualifiedIdJoinIndexImplV3::PersistMetadataToDisk() {
   // We can skip persisting metadata to disk only if both info and storage are
   // clean.
@@ -767,6 +1020,8 @@
 
   ICING_RETURN_IF_ERROR(
       parent_document_id_to_child_array_info_->PersistToDisk());
+  ICING_RETURN_IF_ERROR(
+      parent_qualified_id_to_child_array_info_->PersistToDisk());
   ICING_RETURN_IF_ERROR(child_document_join_id_pair_array_->PersistToDisk());
 
   is_storage_dirty_ = false;
@@ -783,10 +1038,20 @@
   ICING_ASSIGN_OR_RETURN(
       Crc32 parent_document_id_to_child_array_info_crc,
       parent_document_id_to_child_array_info_->UpdateChecksum());
+  ICING_ASSIGN_OR_RETURN(
+      Crc32 parent_qualified_id_to_child_array_info_crc,
+      parent_qualified_id_to_child_array_info_->UpdateChecksum());
+  if (!feature_flags_.enable_non_existent_qualified_id_join()) {
+    // We have still called UpdateChecksum() for this storage even if the flag
+    // is off. However, for flag guarding purposes, we set the local variable to
+    // 0 to exclude it from the overall checksum.
+    parent_qualified_id_to_child_array_info_crc = Crc32(0);
+  }
   ICING_ASSIGN_OR_RETURN(Crc32 child_document_join_id_pair_array_crc,
                          child_document_join_id_pair_array_->UpdateChecksum());
 
   return Crc32(parent_document_id_to_child_array_info_crc.Get() ^
+               parent_qualified_id_to_child_array_info_crc.Get() ^
                child_document_join_id_pair_array_crc.Get());
 }
 
@@ -808,10 +1073,17 @@
   // Get checksums for all components.
   Crc32 parent_document_id_to_child_array_info_crc =
       parent_document_id_to_child_array_info_->GetChecksum();
+  Crc32 parent_qualified_id_to_child_array_info_crc(0);
+  if (feature_flags_.enable_non_existent_qualified_id_join()) {
+    ICING_ASSIGN_OR_RETURN(
+        parent_qualified_id_to_child_array_info_crc,
+        parent_qualified_id_to_child_array_info_->GetChecksum());
+  }
   Crc32 child_document_join_id_pair_array_crc =
       child_document_join_id_pair_array_->GetChecksum();
 
   return Crc32(parent_document_id_to_child_array_info_crc.Get() ^
+               parent_qualified_id_to_child_array_info_crc.Get() ^
                child_document_join_id_pair_array_crc.Get());
 }
 

diff --git a/icing/join/qualified-id-join-index-impl-v3.h b/icing/join/qualified-id-join-index-impl-v3.h
index 30483df..63f484c 100644
--- a/icing/join/qualified-id-join-index-impl-v3.h
+++ b/icing/join/qualified-id-join-index-impl-v3.h

@@ -32,9 +32,12 @@
 #include "icing/file/persistent-storage.h"
 #include "icing/join/document-join-id-pair.h"
 #include "icing/join/qualified-id-join-index.h"
+#include "icing/join/qualified-id.h"
 #include "icing/schema/joinable-property.h"
 #include "icing/store/document-filter-data.h"
 #include "icing/store/document-id.h"
+#include "icing/store/document-store.h"
+#include "icing/store/key-mapper.h"
 #include "icing/store/namespace-id-fingerprint.h"
 #include "icing/store/namespace-id.h"
 #include "icing/util/crc32.h"
@@ -143,8 +146,8 @@
 
   ~QualifiedIdJoinIndexImplV3() override;
 
-  // Puts new join data into the index: adds a new child document and its
-  // referenced parent documents into the join index.
+  // (Deprecated) Puts new join data into the index: adds a new child document
+  // and its referenced parent documents into the join index.
   //
   // Returns:
   //   - OK on success
@@ -154,13 +157,31 @@
       const DocumentJoinIdPair& child_document_join_id_pair,
       std::vector<DocumentId>&& parent_document_ids) override;
 
-  // Gets the list of joinable children for the given parent document id.
+  // Puts new join data into the index using qualified ids: adds a new child
+  // document and its referenced parent documents into the join index.
+  //
+  // Internally, this function will update both the qualified id map and the
+  // document id map (if the parent document has a valid document id).
   //
   // Returns:
-  //   - A list of children's DocumentJoinIdPair on success
+  //   - OK on success
+  //   - INVALID_ARGUMENT_ERROR if child_document_join_id_pair is invalid
+  //   - Any FileBackedVector or KeyMapper errors
+  libtextclassifier3::Status Put(
+      const DocumentStore* document_store,
+      const DocumentJoinIdPair& child_document_join_id_pair,
+      std::vector<QualifiedId>&& parent_qualified_ids) override;
+
+  // Gets an array view of all joinable children for the given parent document
+  // id.
+  //
+  // Returns:
+  //   - A DocumentJoinIdPairArrayView object on success. If there is no edge
+  //     for a valid node_id, then an array view with data() == nullptr and
+  //     size() == 0 will be returned
   //   - Any FileBackedVector errors
-  libtextclassifier3::StatusOr<std::vector<DocumentJoinIdPair>> Get(
-      DocumentId parent_document_id) const override;
+  libtextclassifier3::StatusOr<DocumentJoinIdPairArrayView>
+  GetDocumentJoinIdPairArrayView(DocumentId parent_document_id) const override;
 
   // Migrates existing join data for a parent document from old_document_id to
   // new_document_id.
@@ -176,6 +197,22 @@
   libtextclassifier3::Status MigrateParent(DocumentId old_document_id,
                                            DocumentId new_document_id) override;
 
+  // Migrates existing join data for a parent document from its qualified id to
+  // a new document id. This is used when a document that was previously
+  // referenced as a "missing parent" (by its qualified id) is indexed and
+  // assigned a document id.
+  //
+  // Note: The entry in the parent qualified id map will be kept as is, since
+  // the qualified id is the source of truth.
+  //
+  // Returns:
+  //   - OK on success
+  //   - INVALID_ARGUMENT_ERROR if new_document_id is invalid
+  //   - Any FileBackedVector or KeyMapper errors
+  libtextclassifier3::Status MigrateParent(
+      const QualifiedId& parent_qualified_id,
+      DocumentId new_document_id) override;
+
   // v2 only API. Returns UNIMPLEMENTED_ERROR.
   libtextclassifier3::Status Put(
       SchemaTypeId schema_type_id, JoinablePropertyId joinable_property_id,
@@ -193,6 +230,7 @@
   }
 
   libtextclassifier3::Status Optimize(
+      const DocumentStore* document_store,
       const std::vector<DocumentId>& document_id_old_to_new,
       const std::vector<NamespaceId>& namespace_id_old_to_new,
       DocumentId new_last_added_document_id) override;
@@ -235,6 +273,8 @@
       std::unique_ptr<MemoryMappedFile> metadata_mmapped_file,
       std::unique_ptr<FileBackedVector<ArrayInfo>>
           parent_document_id_to_child_array_info,
+      std::unique_ptr<KeyMapper<ArrayInfo>>
+          parent_qualified_id_to_child_array_info,
       std::unique_ptr<FileBackedVector<DocumentJoinIdPair>>
           child_document_join_id_pair_array,
       const FeatureFlags& feature_flags)
@@ -242,6 +282,8 @@
         metadata_mmapped_file_(std::move(metadata_mmapped_file)),
         parent_document_id_to_child_array_info_(
             std::move(parent_document_id_to_child_array_info)),
+        parent_qualified_id_to_child_array_info_(
+            std::move(parent_qualified_id_to_child_array_info)),
         child_document_join_id_pair_array_(
             std::move(child_document_join_id_pair_array)),
         feature_flags_(feature_flags),
@@ -259,10 +301,11 @@
                           std::string working_path,
                           const FeatureFlags& feature_flags);
 
-  // Appends a list of new child DocumentJoinIdPair to the parent's
+  // (Deprecated) Appends a list of new child DocumentJoinIdPair to the parent's
   // DocumentJoinIdPair (extensible) array. If the array is invalid or doesn't
   // have enough space for new elements, then extend it and set the new array
-  // info.
+  // info. This function then updates (and only updates) the
+  // parent_document_id_to_child_array_info_ map.
   //
   // Returns:
   //   - OK on success
@@ -273,6 +316,41 @@
       DocumentId parent_document_id,
       std::vector<DocumentJoinIdPair>&& child_document_join_id_pairs);
 
+  // Appends a list of new child DocumentJoinIdPair to the parent's (by its
+  // qualified id) DocumentJoinIdPair (extensible) array. If the array is
+  // invalid or doesn't have enough space for new elements, then extend it and
+  // set the new array info. This function then updates the
+  // parent_qualified_id_to_child_array_info_ map. If the parent document has a
+  // valid document id, it also updates the
+  // parent_document_id_to_child_array_info_ map.
+  //
+  // Returns:
+  //   - OK on success
+  //   - RESOURCE_EXHAUSTED_ERROR if the new # of elements exceed
+  //     kMaxNumChildrenPerParent
+  //   - Any FileBackedVector or KeyMapper errors
+  libtextclassifier3::Status AppendChildDocumentJoinIdPairsForParent(
+      std::string_view parent_qualified_id_str, DocumentId parent_document_id,
+      std::vector<DocumentJoinIdPair>&& child_document_join_id_pairs);
+
+  libtextclassifier3::Status AppendChildDocumentJoinIdPairsForParent(
+      const QualifiedId& parent_qualified_id, DocumentId parent_document_id,
+      std::vector<DocumentJoinIdPair>&& child_document_join_id_pairs) {
+    return AppendChildDocumentJoinIdPairsForParent(
+        parent_qualified_id.ToString(), parent_document_id,
+        std::move(child_document_join_id_pairs));
+  }
+
+  // Appends a list of new child DocumentJoinIdPair to an extensible array
+  // and returns the new ArrayInfo.
+  //
+  // Returns:
+  //   - New ArrayInfo on success
+  //   - Any FileBackedVector errors
+  libtextclassifier3::StatusOr<ArrayInfo> AppendToChildArray(
+      const ArrayInfo& current_array_info,
+      std::vector<DocumentJoinIdPair>&& child_document_join_id_pairs);
+
   // Extends the parent document id to child array info if necessary, according
   // to the new parent document id.
   //
@@ -326,9 +404,24 @@
   //   - OK on success
   //   - INTERNAL_ERROR on I/O error
   libtextclassifier3::Status TransferIndex(
+      const DocumentStore& document_store,
       const std::vector<DocumentId>& document_id_old_to_new,
       QualifiedIdJoinIndexImplV3* new_index) const;
 
+  // Gets all child DocumentJoinIdPairs according to array_info, transfers child
+  // document ids according to document_id_old_to_new, and returns a new vector
+  // of DocumentJoinIdPairs. Invalid child document ids will be skipped.
+  //
+  // Returns:
+  //   - A vector of remapped DocumentJoinIdPairs on success
+  //   - INTERNAL_ERROR if any child document id is out of range of
+  //     document_id_old_to_new
+  //   - Any FileBackedVector errors
+  libtextclassifier3::StatusOr<std::vector<DocumentJoinIdPair>>
+  GetTransferredChildDocumentJoinIdPairs(
+      const QualifiedIdJoinIndexImplV3::ArrayInfo& array_info,
+      const std::vector<DocumentId>& document_id_old_to_new) const;
+
   libtextclassifier3::Status PersistMetadataToDisk() override;
 
   libtextclassifier3::Status PersistStoragesToDisk() override;
@@ -384,6 +477,17 @@
   std::unique_ptr<FileBackedVector<ArrayInfo>>
       parent_document_id_to_child_array_info_;
 
+  // Storage for mapping a parent's qualified id to the ArrayInfo. With the
+  // feature flag `enable_non_existent_qualified_id_join` on, this is the
+  // source of truth for the join index. The ArrayInfo points to an extensible
+  // array stored in child_document_join_id_pair_array_, and the extensible
+  // array contains the parent's joinable children information.
+  //
+  // The parent_document_id_to_child_array_info_ map shares the same ArrayInfo
+  // with this mapper, for documents with valid document ids.
+  std::unique_ptr<KeyMapper<ArrayInfo>>
+      parent_qualified_id_to_child_array_info_;
+
   // Storage for DocumentJoinIdPair.
   // - It is a collection of multiple extensible arrays for parents.
   // - Each extensible array contains a list of child DocumentJoinIdPair.

diff --git a/icing/join/qualified-id-join-index-impl-v3_test.cc b/icing/join/qualified-id-join-index-impl-v3_test.cc
index 114cab7..adba00c 100644
--- a/icing/join/qualified-id-join-index-impl-v3_test.cc
+++ b/icing/join/qualified-id-join-index-impl-v3_test.cc

@@ -25,17 +25,26 @@
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include "icing/absl_ports/str_cat.h"
+#include "icing/document-builder.h"
 #include "icing/feature-flags.h"
 #include "icing/file/file-backed-vector.h"
 #include "icing/file/filesystem.h"
 #include "icing/file/memory-mapped-file.h"
 #include "icing/file/persistent-storage.h"
+#include "icing/file/portable-file-backed-proto-log.h"
 #include "icing/join/document-join-id-pair.h"
+#include "icing/join/qualified-id-join-index.h"
+#include "icing/join/qualified-id.h"
+#include "icing/portable/gzip_stream.h"
+#include "icing/schema-builder.h"
+#include "icing/schema/schema-store.h"
 #include "icing/store/document-id.h"
+#include "icing/store/document-store.h"
 #include "icing/testing/common-matchers.h"
-#include "icing/testing/test-feature-flags.h"
 #include "icing/testing/tmp-directory.h"
+#include "icing/util/clock.h"
 #include "icing/util/crc32.h"
+#include "icing/util/document-util.h"
 
 namespace icing {
 namespace lib {
@@ -43,15 +52,18 @@
 namespace {
 
 using ::testing::ElementsAre;
+using ::testing::ElementsAreArray;
 using ::testing::Eq;
 using ::testing::Gt;
 using ::testing::HasSubstr;
 using ::testing::IsEmpty;
 using ::testing::IsFalse;
+using ::testing::IsNull;
 using ::testing::IsTrue;
 using ::testing::Lt;
 using ::testing::Ne;
 using ::testing::Not;
+using ::testing::NotNull;
 using ::testing::Pointee;
 using ::testing::SizeIs;
 
@@ -59,36 +71,122 @@
 using Info = QualifiedIdJoinIndexImplV3::Info;
 using ArrayInfo = QualifiedIdJoinIndexImplV3::ArrayInfo;
 
-class QualifiedIdJoinIndexImplV3Test : public ::testing::Test {
+class QualifiedIdJoinIndexImplV3Test : public ::testing::TestWithParam<bool> {
  protected:
   void SetUp() override {
-    feature_flags_ = std::make_unique<FeatureFlags>(GetTestFeatureFlags());
+    feature_flags_ = std::make_unique<FeatureFlags>(
+        /*enable_circular_schema_definitions=*/true,
+        /*enable_scorable_properties=*/true,
+        /*enable_embedding_quantization=*/true,
+        /*enable_repeated_field_joins=*/true,
+        /*enable_embedding_backup_generation=*/true,
+        /*enable_schema_database=*/true,
+        /*release_backup_schema_file_if_overlay_present=*/true,
+        /*enable_strict_page_byte_size_limit=*/true,
+        /*enable_smaller_decompression_buffer_size=*/true,
+        /*enable_eigen_embedding_scoring=*/true,
+        /*enable_passing_filter_to_children=*/true,
+        /*enable_proto_log_new_header_format=*/true,
+        /*enable_embedding_iterator_v2=*/true,
+        /*enable_reusable_decompression_buffer=*/true,
+        /*enable_schema_type_id_optimization=*/true,
+        /*enable_optimize_improvements=*/true,
+        /*expired_document_purge_threshold_ms=*/0,
+        /*enable_non_existent_qualified_id_join=*/GetParam(),
+        /*enable_skip_set_schema_type_equality_check=*/true,
+        /*enable_embed_query_optimization=*/true,
+        /*enable_schema_definition_deduping=*/true);
 
     base_dir_ = GetTestTempDir() + "/icing";
+    working_path_ = base_dir_ + "/qualified_id_join_index_impl_v3";
+    document_store_dir_ = base_dir_ + "/document_store";
+    schema_store_dir_ = base_dir_ + "/schema_store";
     ASSERT_THAT(filesystem_.CreateDirectoryRecursively(base_dir_.c_str()),
                 IsTrue());
+    filesystem_.CreateDirectoryRecursively(document_store_dir_.c_str());
+    filesystem_.CreateDirectoryRecursively(schema_store_dir_.c_str());
+    ICING_ASSERT_OK_AND_ASSIGN(
+        schema_store_, SchemaStore::Create(&filesystem_, schema_store_dir_,
+                                           &clock_, feature_flags_.get()));
+    ICING_ASSERT_OK(schema_store_->SetSchema(
+        SchemaBuilder()
+            .AddType(SchemaTypeConfigBuilder().SetType("type"))
+            .Build(),
+        /*ignore_errors_and_delete_documents=*/false));
+    ASSERT_NO_FATAL_FAILURE(CreateDocumentStore());
+  }
 
-    working_path_ = base_dir_ + "/qualified_id_join_index_impl_v3";
+  void CreateDocumentStore() {
+    ICING_ASSERT_OK_AND_ASSIGN(
+        DocumentStore::CreateResult create_result,
+        DocumentStore::Create(
+            &filesystem_, document_store_dir_, &clock_, schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
+    document_store_ = std::move(create_result.document_store);
+  }
+
+  void OptimizeDocumentStore() {
+    std::string optimized_document_store_dir =
+        base_dir_ + "/document_store_optimized";
+    ASSERT_THAT(filesystem_.CreateDirectoryRecursively(
+                    optimized_document_store_dir.c_str()),
+                IsTrue());
+    ICING_ASSERT_OK_AND_ASSIGN(
+        DocumentStore::OptimizeResult optimize_result,
+        document_store_->OptimizeInto(
+            optimized_document_store_dir, /*lang_segmenter=*/nullptr,
+            /*potentially_optimizable_blob_handles=*/{}));
+    document_store_.reset();
+    ASSERT_THAT(filesystem_.SwapFiles(document_store_dir_.c_str(),
+                                      optimized_document_store_dir.c_str()),
+                IsTrue());
+    ASSERT_NO_FATAL_FAILURE(CreateDocumentStore());
+  }
+
+  void FillDocumentStore(int num_documents) {
+    for (int i = 0; i < num_documents; ++i) {
+      ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+          DocumentBuilder()
+              .SetKey("ns", absl_ports::StrCat("uri", std::to_string(i)))
+              .SetSchema("type")
+              .Build())));
+    }
   }
 
   void TearDown() override {
+    document_store_.reset();
+    schema_store_.reset();
     filesystem_.DeleteDirectoryRecursively(base_dir_.c_str());
   }
 
   std::unique_ptr<FeatureFlags> feature_flags_;
+  Clock clock_;
   Filesystem filesystem_;
   std::string base_dir_;
   std::string working_path_;
+  std::string document_store_dir_;
+  std::string schema_store_dir_;
+  std::unique_ptr<SchemaStore> schema_store_;
+  std::unique_ptr<DocumentStore> document_store_;
 };
 
-TEST_F(QualifiedIdJoinIndexImplV3Test, InvalidWorkingPath) {
+TEST_P(QualifiedIdJoinIndexImplV3Test, InvalidWorkingPath) {
   EXPECT_THAT(QualifiedIdJoinIndexImplV3::Create(
                   filesystem_, "/dev/null/qualified_id_join_index_impl_v3",
                   *feature_flags_),
               StatusIs(libtextclassifier3::StatusCode::INTERNAL));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test, InitializeNewFiles) {
+TEST_P(QualifiedIdJoinIndexImplV3Test, InitializeNewFiles) {
   {
     // Create new qualified id join index
     ASSERT_FALSE(filesystem_.DirectoryExists(working_path_.c_str()));
@@ -111,7 +209,7 @@
       filesystem_.PRead(metadata_file_path.c_str(), metadata_buffer.get(),
                         QualifiedIdJoinIndexImplV3::kMetadataFileSize,
                         /*offset=*/0),
-      IsTrue());
+      Eq(QualifiedIdJoinIndexImplV3::kMetadataFileSize));
 
   // Check info section
   const Info* info = reinterpret_cast<const Info*>(
@@ -125,8 +223,10 @@
   const Crcs* crcs = reinterpret_cast<const Crcs*>(
       metadata_buffer.get() +
       QualifiedIdJoinIndexImplV3::kCrcsMetadataFileOffset);
-  // There are no data in FileBackedVectors, so storages_crc should be zero.
-  EXPECT_THAT(crcs->component_crcs.storages_crc, Eq(0));
+  if (!feature_flags_->enable_non_existent_qualified_id_join()) {
+    // There are no data in FileBackedVectors, so storages_crc should be zero.
+    EXPECT_THAT(crcs->component_crcs.storages_crc, Eq(0));
+  }
   EXPECT_THAT(crcs->component_crcs.info_crc,
               Eq(Crc32(std::string_view(reinterpret_cast<const char*>(info),
                                         sizeof(Info)))
@@ -138,7 +238,7 @@
                      .Get()));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test,
+TEST_P(QualifiedIdJoinIndexImplV3Test,
        InitializationShouldFailIfMissingMetadataFile) {
   {
     DocumentJoinIdPair child_join_id_pair1(/*document_id=*/100,
@@ -179,7 +279,7 @@
           HasSubstr("Inconsistent state of qualified id join index (v3)")));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test,
+TEST_P(QualifiedIdJoinIndexImplV3Test,
        InitializationShouldFailIfMissingParentDocumentIdToChildArrayInfoFile) {
   {
     DocumentJoinIdPair child_join_id_pair1(/*document_id=*/100,
@@ -220,7 +320,7 @@
           HasSubstr("Inconsistent state of qualified id join index (v3)")));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test,
+TEST_P(QualifiedIdJoinIndexImplV3Test,
        InitializationShouldFailIfMissingChildDocumentJoinIdPairArrayFile) {
   {
     DocumentJoinIdPair child_join_id_pair1(/*document_id=*/100,
@@ -261,7 +361,45 @@
           HasSubstr("Inconsistent state of qualified id join index (v3)")));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test,
+TEST_P(QualifiedIdJoinIndexImplV3Test,
+       InitializationShouldFailIfMissingQualifiedIdMapperFile) {
+  if (!feature_flags_->enable_non_existent_qualified_id_join()) {
+    return;  // This test is only relevant when the feature is enabled.
+  }
+
+  {
+    // Create new qualified id join index
+    ICING_ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
+        QualifiedIdJoinIndexImplV3::Create(filesystem_, working_path_,
+                                           *feature_flags_));
+
+    // Insert some data with a qualified id.
+    ICING_ASSERT_OK(index->Put(
+        document_store_.get(),
+        DocumentJoinIdPair(/*document_id=*/100, /*joinable_property_id=*/20),
+        /*parent_qualified_ids=*/
+        std::vector<QualifiedId>{QualifiedId("namespace", "uri")}));
+    ASSERT_THAT(index, Pointee(SizeIs(1)));
+
+    ICING_ASSERT_OK(index->PersistToDisk());
+  }
+
+  // Manually delete the parent_qualified_id_to_child_array_info directory.
+  const std::string dir_path = absl_ports::StrCat(
+      working_path_, "/parent_qualified_id_to_child_array_info");
+  ASSERT_THAT(filesystem_.DeleteDirectoryRecursively(dir_path.c_str()),
+              IsTrue());
+
+  // Attempt to create the qualified id join index. This should fail because of
+  // checksum mismatch.
+  EXPECT_THAT(QualifiedIdJoinIndexImplV3::Create(filesystem_, working_path_,
+                                                 *feature_flags_),
+              StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION,
+                       HasSubstr("Invalid storages crc")));
+}
+
+TEST_P(QualifiedIdJoinIndexImplV3Test,
        InitializationShouldFailWithoutPersistToDiskOrDestruction) {
   DocumentJoinIdPair child_join_id_pair1(/*document_id=*/100,
                                          /*joinable_property_id=*/20);
@@ -292,7 +430,7 @@
               StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test,
+TEST_P(QualifiedIdJoinIndexImplV3Test,
        InitializationShouldSucceedWithUpdateChecksums) {
   DocumentJoinIdPair child_join_id_pair1(/*document_id=*/100,
                                          /*joinable_property_id=*/20);
@@ -324,13 +462,13 @@
                              QualifiedIdJoinIndexImplV3::Create(
                                  filesystem_, working_path_, *feature_flags_));
   EXPECT_THAT(index2, Pointee(SizeIs(2)));
-  EXPECT_THAT(index2->Get(/*parent_document_id=*/0),
+  EXPECT_THAT(index2->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/0),
               IsOkAndHolds(ElementsAre(child_join_id_pair1)));
-  EXPECT_THAT(index2->Get(/*parent_document_id=*/1),
+  EXPECT_THAT(index2->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/1),
               IsOkAndHolds(ElementsAre(child_join_id_pair2)));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test,
+TEST_P(QualifiedIdJoinIndexImplV3Test,
        InitializationShouldSucceedWithPersistToDisk) {
   DocumentJoinIdPair child_join_id_pair1(/*document_id=*/100,
                                          /*joinable_property_id=*/20);
@@ -360,13 +498,13 @@
                              QualifiedIdJoinIndexImplV3::Create(
                                  filesystem_, working_path_, *feature_flags_));
   EXPECT_THAT(index2, Pointee(SizeIs(2)));
-  EXPECT_THAT(index2->Get(/*parent_document_id=*/0),
+  EXPECT_THAT(index2->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/0),
               IsOkAndHolds(ElementsAre(child_join_id_pair1)));
-  EXPECT_THAT(index2->Get(/*parent_document_id=*/1),
+  EXPECT_THAT(index2->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/1),
               IsOkAndHolds(ElementsAre(child_join_id_pair2)));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test,
+TEST_P(QualifiedIdJoinIndexImplV3Test,
        InitializationShouldSucceedAfterDestruction) {
   DocumentJoinIdPair child_join_id_pair1(/*document_id=*/100,
                                          /*joinable_property_id=*/20);
@@ -400,14 +538,14 @@
         QualifiedIdJoinIndexImplV3::Create(filesystem_, working_path_,
                                            *feature_flags_));
     EXPECT_THAT(index, Pointee(SizeIs(2)));
-    EXPECT_THAT(index->Get(/*parent_document_id=*/0),
+    EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/0),
                 IsOkAndHolds(ElementsAre(child_join_id_pair1)));
-    EXPECT_THAT(index->Get(/*parent_document_id=*/1),
+    EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/1),
                 IsOkAndHolds(ElementsAre(child_join_id_pair2)));
   }
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test,
+TEST_P(QualifiedIdJoinIndexImplV3Test,
        InitializeExistingFilesWithDifferentMagicShouldFail) {
   {
     // Create new qualified id join index
@@ -433,7 +571,7 @@
     ASSERT_THAT(filesystem_.PRead(metadata_sfd.get(), metadata_buffer.get(),
                                   QualifiedIdJoinIndexImplV3::kMetadataFileSize,
                                   /*offset=*/0),
-                IsTrue());
+                Eq(QualifiedIdJoinIndexImplV3::kMetadataFileSize));
 
     // Manually change magic and update checksum
     Crcs* crcs = reinterpret_cast<Crcs*>(
@@ -453,13 +591,15 @@
 
   // Attempt to create the qualified id join index with different magic. This
   // should fail.
-  EXPECT_THAT(QualifiedIdJoinIndexImplV3::Create(filesystem_, working_path_,
-                                                 *feature_flags_),
-              StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION,
-                       HasSubstr("Incorrect magic value")));
+  EXPECT_THAT(
+      QualifiedIdJoinIndexImplV3::Create(filesystem_, working_path_,
+                                         *feature_flags_),
+      StatusIs(
+          libtextclassifier3::StatusCode::FAILED_PRECONDITION,
+          HasSubstr("Invalid header magic for QualifiedIdJoinIndexImplV3")));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test,
+TEST_P(QualifiedIdJoinIndexImplV3Test,
        InitializeExistingFilesWithWrongAllCrcShouldFail) {
   {
     // Create new qualified id join index
@@ -485,7 +625,7 @@
     ASSERT_THAT(filesystem_.PRead(metadata_sfd.get(), metadata_buffer.get(),
                                   QualifiedIdJoinIndexImplV3::kMetadataFileSize,
                                   /*offset=*/0),
-                IsTrue());
+                Eq(QualifiedIdJoinIndexImplV3::kMetadataFileSize));
 
     // Manually corrupt all_crc
     Crcs* crcs = reinterpret_cast<Crcs*>(
@@ -507,7 +647,7 @@
                        HasSubstr("Invalid all crc")));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test,
+TEST_P(QualifiedIdJoinIndexImplV3Test,
        InitializeExistingFilesWithCorruptedInfoShouldFail) {
   {
     // Create new qualified id join index
@@ -533,7 +673,7 @@
     ASSERT_THAT(filesystem_.PRead(metadata_sfd.get(), metadata_buffer.get(),
                                   QualifiedIdJoinIndexImplV3::kMetadataFileSize,
                                   /*offset=*/0),
-                IsTrue());
+                Eq(QualifiedIdJoinIndexImplV3::kMetadataFileSize));
 
     // Modify info, but don't update the checksum. This would be similar to
     // corruption of info.
@@ -556,7 +696,7 @@
                        HasSubstr("Invalid info crc")));
 }
 
-TEST_F(
+TEST_P(
     QualifiedIdJoinIndexImplV3Test,
     InitializeExistingFilesWithCorruptedParentDocumentIdToChildArrayInfoShouldFail) {
   {
@@ -599,7 +739,7 @@
                        HasSubstr("Invalid storages crc")));
 }
 
-TEST_F(
+TEST_P(
     QualifiedIdJoinIndexImplV3Test,
     InitializeExistingFilesWithCorruptedChildDocumentJoinIdPairArrayShouldFail) {
   {
@@ -641,7 +781,7 @@
                        HasSubstr("Invalid storages crc")));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test, Put) {
+TEST_P(QualifiedIdJoinIndexImplV3Test, Put) {
   // Create new qualified id join index
   ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
                              QualifiedIdJoinIndexImplV3::Create(
@@ -682,20 +822,22 @@
   EXPECT_THAT(index, Pointee(SizeIs(6)));
 
   // Verify Get API.
-  EXPECT_THAT(index->Get(/*parent_document_id=*/0),
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/0),
               IsOkAndHolds(ElementsAre(child_join_id_pair4)));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/1),
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/1),
               IsOkAndHolds(ElementsAre(child_join_id_pair1, child_join_id_pair2,
                                        child_join_id_pair5)));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/2),
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/2),
               IsOkAndHolds(ElementsAre(child_join_id_pair3)));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/3), IsOkAndHolds(IsEmpty()));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/4), IsOkAndHolds(IsEmpty()));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/5),
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/3),
+              IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/4),
+              IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/5),
               IsOkAndHolds(ElementsAre(child_join_id_pair6)));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test,
+TEST_P(QualifiedIdJoinIndexImplV3Test,
        Put_multipleParentsInASingleJoinableProperty) {
   // Create new qualified id join index
   ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
@@ -726,33 +868,35 @@
   EXPECT_THAT(index, Pointee(SizeIs(13)));
 
   // Verify Get API.
-  EXPECT_THAT(index->Get(/*parent_document_id=*/0),
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/0),
               IsOkAndHolds(ElementsAre(child_join_id_pair2)));
   EXPECT_THAT(
-      index->Get(/*parent_document_id=*/1),
+      index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/1),
       IsOkAndHolds(ElementsAre(child_join_id_pair1, child_join_id_pair2)));
   EXPECT_THAT(
-      index->Get(/*parent_document_id=*/2),
+      index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/2),
       IsOkAndHolds(ElementsAre(child_join_id_pair2, child_join_id_pair3)));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/3),
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/3),
               IsOkAndHolds(ElementsAre(child_join_id_pair2)));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/4),
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/4),
               IsOkAndHolds(ElementsAre(child_join_id_pair1)));
   EXPECT_THAT(
-      index->Get(/*parent_document_id=*/5),
+      index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/5),
       IsOkAndHolds(ElementsAre(child_join_id_pair2, child_join_id_pair3)));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/6), IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/6),
+              IsOkAndHolds(IsEmpty()));
   EXPECT_THAT(
-      index->Get(/*parent_document_id=*/7),
+      index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/7),
       IsOkAndHolds(ElementsAre(child_join_id_pair1, child_join_id_pair3)));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/8),
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/8),
               IsOkAndHolds(ElementsAre(child_join_id_pair2)));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/9), IsOkAndHolds(IsEmpty()));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/10),
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/9),
+              IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/10),
               IsOkAndHolds(ElementsAre(child_join_id_pair1)));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test,
+TEST_P(QualifiedIdJoinIndexImplV3Test,
        Put_multipleParentsInMultipleJoinableProperties) {
   // Create new qualified id join index
   ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
@@ -784,33 +928,35 @@
   EXPECT_THAT(index, Pointee(SizeIs(13)));
 
   // Verify Get API.
-  EXPECT_THAT(index->Get(/*parent_document_id=*/0),
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/0),
               IsOkAndHolds(ElementsAre(child_join_id_pair2)));
   EXPECT_THAT(
-      index->Get(/*parent_document_id=*/1),
+      index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/1),
       IsOkAndHolds(ElementsAre(child_join_id_pair1, child_join_id_pair2)));
   EXPECT_THAT(
-      index->Get(/*parent_document_id=*/2),
+      index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/2),
       IsOkAndHolds(ElementsAre(child_join_id_pair2, child_join_id_pair3)));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/3),
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/3),
               IsOkAndHolds(ElementsAre(child_join_id_pair2)));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/4),
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/4),
               IsOkAndHolds(ElementsAre(child_join_id_pair1)));
   EXPECT_THAT(
-      index->Get(/*parent_document_id=*/5),
+      index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/5),
       IsOkAndHolds(ElementsAre(child_join_id_pair2, child_join_id_pair3)));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/6), IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/6),
+              IsOkAndHolds(IsEmpty()));
   EXPECT_THAT(
-      index->Get(/*parent_document_id=*/7),
+      index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/7),
       IsOkAndHolds(ElementsAre(child_join_id_pair1, child_join_id_pair3)));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/8),
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/8),
               IsOkAndHolds(ElementsAre(child_join_id_pair2)));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/9), IsOkAndHolds(IsEmpty()));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/10),
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/9),
+              IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/10),
               IsOkAndHolds(ElementsAre(child_join_id_pair1)));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test,
+TEST_P(QualifiedIdJoinIndexImplV3Test,
        PutShouldResizeParentDocumentIdToChildArrayInfo) {
   // Create new qualified id join index
   ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
@@ -834,14 +980,15 @@
   // Get API should return empty result for document 0 to 9.
   for (DocumentId parent_doc_id = 0; parent_doc_id < kParentDocumentId;
        ++parent_doc_id) {
-    EXPECT_THAT(index->Get(parent_doc_id), IsOkAndHolds(IsEmpty()));
+    EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(parent_doc_id),
+                IsOkAndHolds(IsEmpty()));
   }
   // Get API should return the child document for document 10.
-  EXPECT_THAT(index->Get(kParentDocumentId),
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(kParentDocumentId),
               IsOkAndHolds(ElementsAre(child_join_id_pair)));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test,
+TEST_P(QualifiedIdJoinIndexImplV3Test,
        PutShouldExtendChildDocumentJoinIdPairArray) {
   // Create new qualified id join index
   ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
@@ -885,12 +1032,13 @@
   }
   EXPECT_THAT(index, Pointee(SizeIs(102)));
 
-  EXPECT_THAT(index->Get(parent1), IsOkAndHolds(child_join_id_pairs));
-  EXPECT_THAT(index->Get(parent2),
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(parent1),
+              IsOkAndHolds(ElementsAreArray(child_join_id_pairs)));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(parent2),
               IsOkAndHolds(ElementsAre(child_join_id_pair2)));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test,
+TEST_P(QualifiedIdJoinIndexImplV3Test,
        PutLargeParentShouldHandleAddressCorrectlyForRemap) {
   // Create new qualified id join index
   ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
@@ -932,7 +1080,7 @@
   EXPECT_THAT(file_size_after, Gt(file_size_before));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test,
+TEST_P(QualifiedIdJoinIndexImplV3Test,
        PutLargeNumberOfDataShouldHandleRemapAddressCorrectly) {
   // Create new qualified id join index
   ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
@@ -974,7 +1122,7 @@
   EXPECT_THAT(index, Pointee(SizeIs(kNumChildrenToFillFbv + 1)));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test, PutShouldSkipInvalidParentDocumentId) {
+TEST_P(QualifiedIdJoinIndexImplV3Test, PutShouldSkipInvalidParentDocumentId) {
   // Create new qualified id join index
   ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
                              QualifiedIdJoinIndexImplV3::Create(
@@ -995,15 +1143,15 @@
   EXPECT_THAT(index, Pointee(SizeIs(3)));
 
   // Verify Get API.
-  EXPECT_THAT(index->Get(/*parent_document_id=*/1),
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/1),
               IsOkAndHolds(ElementsAre(child_join_id_pair)));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/2),
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/2),
               IsOkAndHolds(ElementsAre(child_join_id_pair)));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/3),
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/3),
               IsOkAndHolds(ElementsAre(child_join_id_pair)));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test,
+TEST_P(QualifiedIdJoinIndexImplV3Test,
        PutShouldReturnInvalidArgumentErrorForInvalidChild) {
   // Create new qualified id join index
   ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
@@ -1020,24 +1168,105 @@
   EXPECT_THAT(index, Pointee(IsEmpty()));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test, GetEmptyIndex) {
+TEST_P(QualifiedIdJoinIndexImplV3Test, DocumentJoinIdPairArrayView) {
+  // Create new qualified id join index
+  ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
+                             QualifiedIdJoinIndexImplV3::Create(
+                                 filesystem_, working_path_, *feature_flags_));
+
+  // Add 2 children for parent document 0.
+  DocumentJoinIdPair child_join_id_pair1(/*document_id=*/100,
+                                         /*joinable_property_id=*/20);
+  DocumentJoinIdPair child_join_id_pair2(/*document_id=*/101,
+                                         /*joinable_property_id=*/2);
+  EXPECT_THAT(index->Put(child_join_id_pair1,
+                         /*parent_document_ids=*/std::vector<DocumentId>{0}),
+              IsOk());
+  EXPECT_THAT(index->Put(child_join_id_pair2,
+                         /*parent_document_ids=*/std::vector<DocumentId>{0}),
+              IsOk());
+
+  EXPECT_THAT(index, Pointee(SizeIs(2)));
+
+  // Get array view. Test each STL style method.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      QualifiedIdJoinIndex::DocumentJoinIdPairArrayView array_view1,
+      index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/0));
+  EXPECT_THAT(array_view1, Not(IsEmpty()));
+  EXPECT_THAT(array_view1.data(), NotNull());
+  EXPECT_THAT(array_view1, SizeIs(2));
+  EXPECT_THAT(array_view1.begin(), NotNull());
+  EXPECT_THAT(array_view1.end(), NotNull());
+  EXPECT_THAT(array_view1,
+              ElementsAre(child_join_id_pair1, child_join_id_pair2));
+
+  // Add 1 more child for parent document 0.
+  DocumentJoinIdPair child_join_id_pair3(/*document_id=*/102,
+                                         /*joinable_property_id=*/2);
+  EXPECT_THAT(index->Put(child_join_id_pair3,
+                         /*parent_document_ids=*/std::vector<DocumentId>{0}),
+              IsOk());
+
+  // Get array view again. Test each STL style method.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      QualifiedIdJoinIndex::DocumentJoinIdPairArrayView array_view2,
+      index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/0));
+  EXPECT_THAT(array_view2, Not(IsEmpty()));
+  EXPECT_THAT(array_view2.data(), NotNull());
+  EXPECT_THAT(array_view2, SizeIs(3));
+  EXPECT_THAT(array_view2.begin(), NotNull());
+  EXPECT_THAT(array_view2.end(), NotNull());
+  EXPECT_THAT(array_view2, ElementsAre(child_join_id_pair1, child_join_id_pair2,
+                                       child_join_id_pair3));
+}
+
+TEST_P(QualifiedIdJoinIndexImplV3Test, EmptyDocumentJoinIdPairArrayView) {
   // Create new qualified id join index
   ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
                              QualifiedIdJoinIndexImplV3::Create(
                                  filesystem_, working_path_, *feature_flags_));
   EXPECT_THAT(index, Pointee(IsEmpty()));
 
-  EXPECT_THAT(index->Get(/*parent_document_id=*/0), IsOkAndHolds(IsEmpty()));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/1), IsOkAndHolds(IsEmpty()));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/2), IsOkAndHolds(IsEmpty()));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/3), IsOkAndHolds(IsEmpty()));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/4), IsOkAndHolds(IsEmpty()));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/5), IsOkAndHolds(IsEmpty()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      QualifiedIdJoinIndex::DocumentJoinIdPairArrayView array_view,
+      index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/0));
+  EXPECT_THAT(array_view, IsEmpty());
+  EXPECT_THAT(array_view.data(), IsNull());
+  EXPECT_THAT(array_view.size(), Eq(0));
+  EXPECT_THAT(array_view.begin(), IsNull());
+  EXPECT_THAT(array_view.end(), IsNull());
+
+  // Use colon to iterate the array_view. There should be no crash and no-op.
+  for (const DocumentJoinIdPair& _ : array_view) {
+    ADD_FAILURE() << "Unexpectedly iterated the empty array_view.";
+  }
 }
 
-TEST_F(
+TEST_P(QualifiedIdJoinIndexImplV3Test,
+       GetDocumentJoinIdPairArrayView_emptyIndex) {
+  // Create new qualified id join index
+  ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
+                             QualifiedIdJoinIndexImplV3::Create(
+                                 filesystem_, working_path_, *feature_flags_));
+  EXPECT_THAT(index, Pointee(IsEmpty()));
+
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/0),
+              IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/1),
+              IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/2),
+              IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/3),
+              IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/4),
+              IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/5),
+              IsOkAndHolds(IsEmpty()));
+}
+
+TEST_P(
     QualifiedIdJoinIndexImplV3Test,
-    GetShouldReturnEmptyResultWithoutAccessingArrayForNonExistingLargeParent) {
+    GetDocumentJoinIdPairArrayView_shouldReturnEmptyArrayViewForNonExistingLargeParent) {
   // Create new qualified id join index
   ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
                              QualifiedIdJoinIndexImplV3::Create(
@@ -1049,33 +1278,65 @@
       child_join_id_pair, /*parent_document_ids=*/std::vector<DocumentId>{1}));
   EXPECT_THAT(index, Pointee(SizeIs(1)));
 
-  EXPECT_THAT(index->Get(/*parent_document_id=*/0), IsOkAndHolds(IsEmpty()));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/1),
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/0),
+              IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/1),
               IsOkAndHolds(ElementsAre(child_join_id_pair)));
 
   // Now, only parent document id 1 is in the index, so the FileBackedVector has
   // been resized to fit parent document id 1.
   // Get API for parent document id greater than 1 should return empty result
   // without accessing the FileBackedVector.
-  EXPECT_THAT(index->Get(/*parent_document_id=*/2), IsOkAndHolds(IsEmpty()));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/3), IsOkAndHolds(IsEmpty()));
-  EXPECT_THAT(index->Get(kMaxDocumentId), IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/2),
+              IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/3),
+              IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(kMaxDocumentId),
+              IsOkAndHolds(IsEmpty()));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test,
-       GetShouldReturnInvalidArgumentErrorForInvalidParentDocumentId) {
+TEST_P(
+    QualifiedIdJoinIndexImplV3Test,
+    GetDocumentJoinIdPairArrayView_shouldReturnEmptyArrayViewForParentWithNoChildren) {
   // Create new qualified id join index
   ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
                              QualifiedIdJoinIndexImplV3::Create(
                                  filesystem_, working_path_, *feature_flags_));
 
-  EXPECT_THAT(index->Get(/*parent_document_id=*/-1),
+  // Add a child for parent document id 2.
+  DocumentJoinIdPair child_join_id_pair(/*document_id=*/100,
+                                        /*joinable_property_id=*/20);
+  ICING_ASSERT_OK(index->Put(
+      child_join_id_pair, /*parent_document_ids=*/std::vector<DocumentId>{2}));
+  EXPECT_THAT(index, Pointee(SizeIs(1)));
+
+  // Since parent array info FBV is resized to fit parent document id 2, parent
+  // document 0 and 1 should also have array info element with invalid data
+  // index for the 2nd FBV.
+  //
+  // Getting array view for parent document 0 and 1 should return empty result
+  // when seeing invalid data index.
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/0),
+              IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/1),
+              IsOkAndHolds(IsEmpty()));
+}
+
+TEST_P(
+    QualifiedIdJoinIndexImplV3Test,
+    GetDocumentJoinIdPairArrayView_shouldReturnInvalidArgumentErrorForInvalidParentDocumentId) {
+  // Create new qualified id join index
+  ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
+                             QualifiedIdJoinIndexImplV3::Create(
+                                 filesystem_, working_path_, *feature_flags_));
+
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/-1),
               StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
-  EXPECT_THAT(index->Get(kInvalidDocumentId),
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(kInvalidDocumentId),
               StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test, MigrateParent) {
+TEST_P(QualifiedIdJoinIndexImplV3Test, MigrateParent) {
   // Create new qualified id join index
   ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
                              QualifiedIdJoinIndexImplV3::Create(
@@ -1099,19 +1360,21 @@
   // Sanity check.
   ASSERT_THAT(index, Pointee(SizeIs(2)));
   ASSERT_THAT(
-      index->Get(parent_doc_id1),
+      index->GetDocumentJoinIdPairArrayView(parent_doc_id1),
       IsOkAndHolds(ElementsAre(child_join_id_pair1, child_join_id_pair2)));
-  ASSERT_THAT(index->Get(parent_doc_id2), IsOkAndHolds(IsEmpty()));
+  ASSERT_THAT(index->GetDocumentJoinIdPairArrayView(parent_doc_id2),
+              IsOkAndHolds(IsEmpty()));
 
   // Migrate parent document id 1 to 1024.
   EXPECT_THAT(index->MigrateParent(parent_doc_id1, parent_doc_id2), IsOk());
-  EXPECT_THAT(index->Get(parent_doc_id1), IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(parent_doc_id1),
+              IsOkAndHolds(IsEmpty()));
   EXPECT_THAT(
-      index->Get(parent_doc_id2),
+      index->GetDocumentJoinIdPairArrayView(parent_doc_id2),
       IsOkAndHolds(ElementsAre(child_join_id_pair1, child_join_id_pair2)));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test,
+TEST_P(QualifiedIdJoinIndexImplV3Test,
        MigrateParentToLargeIdShouldHandleAddressCorrectlyForRemap) {
   // Create new qualified id join index
   ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
@@ -1142,18 +1405,20 @@
   // Sanity check.
   ASSERT_THAT(index, Pointee(SizeIs(2)));
   ASSERT_THAT(
-      index->Get(parent_doc_id1),
+      index->GetDocumentJoinIdPairArrayView(parent_doc_id1),
       IsOkAndHolds(ElementsAre(child_join_id_pair1, child_join_id_pair2)));
-  ASSERT_THAT(index->Get(parent_doc_id2), IsOkAndHolds(IsEmpty()));
+  ASSERT_THAT(index->GetDocumentJoinIdPairArrayView(parent_doc_id2),
+              IsOkAndHolds(IsEmpty()));
 
   // Migrate parent document id 1 to 30000. This will
   // cause parent_document_id_to_child_array_info being extended and remap. The
   // test verifies that addresses after remap are handled correctly without
   // crashing.
   EXPECT_THAT(index->MigrateParent(parent_doc_id1, parent_doc_id2), IsOk());
-  EXPECT_THAT(index->Get(parent_doc_id1), IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(parent_doc_id1),
+              IsOkAndHolds(IsEmpty()));
   EXPECT_THAT(
-      index->Get(parent_doc_id2),
+      index->GetDocumentJoinIdPairArrayView(parent_doc_id2),
       IsOkAndHolds(ElementsAre(child_join_id_pair1, child_join_id_pair2)));
   int64_t file_size_after = filesystem_.GetFileSize(array_working_path.c_str());
   ASSERT_THAT(file_size_after, Ne(Filesystem::kBadFileSize));
@@ -1162,7 +1427,126 @@
   EXPECT_THAT(file_size_after, Gt(file_size_before));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test, SetLastAddedDocumentId) {
+TEST_P(QualifiedIdJoinIndexImplV3Test, MigrateParentShouldSetDirty) {
+  // Create new qualified id join index
+  ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
+                             QualifiedIdJoinIndexImplV3::Create(
+                                 filesystem_, working_path_, *feature_flags_));
+
+  DocumentId parent_doc_id1 = 1;
+  DocumentId parent_doc_id2 = 1024;
+
+  // Add 2 children with their parents to the index.
+  DocumentJoinIdPair child_join_id_pair1(/*document_id=*/100,
+                                         /*joinable_property_id=*/0);
+  DocumentJoinIdPair child_join_id_pair2(/*document_id=*/101,
+                                         /*joinable_property_id=*/0);
+  ICING_ASSERT_OK(index->Put(
+      child_join_id_pair1,
+      /*parent_document_ids=*/std::vector<DocumentId>{parent_doc_id1}));
+  ICING_ASSERT_OK(index->Put(
+      child_join_id_pair2,
+      /*parent_document_ids=*/std::vector<DocumentId>{parent_doc_id1}));
+
+  // Sanity check.
+  ASSERT_THAT(index, Pointee(SizeIs(2)));
+  ASSERT_THAT(
+      index->GetDocumentJoinIdPairArrayView(parent_doc_id1),
+      IsOkAndHolds(ElementsAre(child_join_id_pair1, child_join_id_pair2)));
+  ASSERT_THAT(index->GetDocumentJoinIdPairArrayView(parent_doc_id2),
+              IsOkAndHolds(IsEmpty()));
+  // PersistToDisk after putting data and get the checksum. This will reset the
+  // dirty flag.
+  ICING_ASSERT_OK(index->PersistToDisk());
+  ICING_ASSERT_OK_AND_ASSIGN(Crc32 crc1, index->GetChecksum());
+
+  // Migrate parent document id 1 to 1024.
+  ICING_ASSERT_OK(index->MigrateParent(parent_doc_id1, parent_doc_id2));
+
+  // Call UpdateChecksums(). The checksum should be recomputed and be different
+  // from the previous one. This validates that MigrateParent() should set the
+  // dirty flag.
+  ICING_ASSERT_OK_AND_ASSIGN(Crc32 crc2, index->UpdateChecksums());
+  EXPECT_THAT(crc2, Ne(crc1));
+
+  // Create another qualified id join index instance with the same file. It
+  // should succeed and GetChecksum() should return the same checksum as the
+  // previous one.
+  ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index2,
+                             QualifiedIdJoinIndexImplV3::Create(
+                                 filesystem_, working_path_, *feature_flags_));
+  ICING_ASSERT_OK_AND_ASSIGN(Crc32 crc3, index2->GetChecksum());
+  EXPECT_THAT(crc3, Eq(crc2));
+}
+
+TEST_P(QualifiedIdJoinIndexImplV3Test, PutAndMigrateQualifiedId) {
+  // Create new qualified id join index
+  ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
+                             QualifiedIdJoinIndexImplV3::Create(
+                                 filesystem_, working_path_, *feature_flags_));
+
+  DocumentJoinIdPair child_join_id_pair1(/*document_id=*/100,
+                                         /*joinable_property_id=*/0);
+  DocumentJoinIdPair child_join_id_pair2(/*document_id=*/101,
+                                         /*joinable_property_id=*/0);
+  QualifiedId parent_qualified_id1("namespace", "uri1");
+  QualifiedId parent_qualified_id2("namespace", "uri2");
+  DocumentId parent_doc_id1 = 1;
+  DocumentId parent_doc_id2 = 2;
+  DocumentId parent_doc_id3 = 3;
+
+  auto put_status1 = index->Put(
+      document_store_.get(), child_join_id_pair1,
+      /*parent_qualified_ids=*/std::vector<QualifiedId>{parent_qualified_id1});
+  auto put_status2 = index->Put(
+      document_store_.get(), child_join_id_pair2,
+      /*parent_qualified_ids=*/
+      std::vector<QualifiedId>{parent_qualified_id1, parent_qualified_id2});
+
+  if (!feature_flags_->enable_non_existent_qualified_id_join()) {
+    // If the flag is not enabled, Put should fail.
+    EXPECT_THAT(put_status1,
+                StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
+    EXPECT_THAT(put_status2,
+                StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
+    EXPECT_THAT(index, Pointee(IsEmpty()));
+    return;
+  }
+
+  ICING_ASSERT_OK(put_status1);
+  ICING_ASSERT_OK(put_status2);
+
+  // Sanity check after Put.
+  EXPECT_THAT(index, Pointee(SizeIs(3)));
+
+  // Migrate parent_qualified_id1 to parent_doc_id1.
+  ICING_ASSERT_OK(index->MigrateParent(parent_qualified_id1, parent_doc_id1));
+
+  // Verify that children are migrated to parent_doc_id1.
+  EXPECT_THAT(
+      index->GetDocumentJoinIdPairArrayView(parent_doc_id1),
+      IsOkAndHolds(ElementsAre(child_join_id_pair1, child_join_id_pair2)));
+
+  // Migrate parent_qualified_id2 to parent_doc_id2.
+  ICING_ASSERT_OK(index->MigrateParent(parent_qualified_id2, parent_doc_id2));
+
+  // Verify that children are migrated to parent_doc_id2.
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(parent_doc_id2),
+              IsOkAndHolds(ElementsAre(child_join_id_pair2)));
+
+  // parent_qualified_id1 and parent_qualified_id2 should still exist in the
+  // qualified id mapper. Migrating them again to another parent_doc_id should
+  // succeed.
+  ICING_ASSERT_OK(index->MigrateParent(parent_qualified_id1, parent_doc_id3));
+  EXPECT_THAT(
+      index->GetDocumentJoinIdPairArrayView(parent_doc_id3),
+      IsOkAndHolds(ElementsAre(child_join_id_pair1, child_join_id_pair2)));
+  ICING_ASSERT_OK(index->MigrateParent(parent_qualified_id2, parent_doc_id3));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(parent_doc_id3),
+              IsOkAndHolds(ElementsAre(child_join_id_pair2)));
+}
+
+TEST_P(QualifiedIdJoinIndexImplV3Test, SetLastAddedDocumentId) {
   ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
                              QualifiedIdJoinIndexImplV3::Create(
                                  filesystem_, working_path_, *feature_flags_));
@@ -1178,7 +1562,7 @@
   EXPECT_THAT(index->last_added_document_id(), Eq(kNextDocumentId));
 }
 
-TEST_F(
+TEST_P(
     QualifiedIdJoinIndexImplV3Test,
     SetLastAddedDocumentIdShouldIgnoreNewDocumentIdNotGreaterThanTheCurrent) {
   ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
@@ -1196,12 +1580,15 @@
   EXPECT_THAT(index->last_added_document_id(), Eq(kDocumentId));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test, Optimize) {
+TEST_P(QualifiedIdJoinIndexImplV3Test, Optimize) {
   // Create new qualified id join index
   ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
                              QualifiedIdJoinIndexImplV3::Create(
                                  filesystem_, working_path_, *feature_flags_));
 
+  // Add documents for parent doc ids to the document store.
+  ASSERT_NO_FATAL_FAILURE(FillDocumentStore(/*num_documents=*/5));
+
   // Create 4 parent and 7 child documents (with N to N joins):
   // - Document 1: 101, 103, 104, 105, 107
   // - Document 2: 102, 103, 105
@@ -1222,27 +1609,62 @@
                                          /*joinable_property_id=*/0);
   DocumentJoinIdPair child_join_id_pair7(/*document_id=*/107,
                                          /*joinable_property_id=*/0);
-  ICING_ASSERT_OK(
-      index->Put(child_join_id_pair1,
-                 /*parent_document_ids=*/std::vector<DocumentId>{1, 3}));
-  ICING_ASSERT_OK(
-      index->Put(child_join_id_pair2,
-                 /*parent_document_ids=*/std::vector<DocumentId>{2}));
-  ICING_ASSERT_OK(
-      index->Put(child_join_id_pair3,
-                 /*parent_document_ids=*/std::vector<DocumentId>{1, 2, 4}));
-  ICING_ASSERT_OK(
-      index->Put(child_join_id_pair4,
-                 /*parent_document_ids=*/std::vector<DocumentId>{1}));
-  ICING_ASSERT_OK(
-      index->Put(child_join_id_pair5,
-                 /*parent_document_ids=*/std::vector<DocumentId>{1, 2}));
-  ICING_ASSERT_OK(
-      index->Put(child_join_id_pair6,
-                 /*parent_document_ids=*/std::vector<DocumentId>{3}));
-  ICING_ASSERT_OK(
-      index->Put(child_join_id_pair7,
-                 /*parent_document_ids=*/std::vector<DocumentId>{1}));
+
+  if (!feature_flags_->enable_non_existent_qualified_id_join()) {
+    ICING_ASSERT_OK(
+        index->Put(child_join_id_pair1,
+                   /*parent_document_ids=*/std::vector<DocumentId>{1, 3}));
+    ICING_ASSERT_OK(
+        index->Put(child_join_id_pair2,
+                   /*parent_document_ids=*/std::vector<DocumentId>{2}));
+    ICING_ASSERT_OK(
+        index->Put(child_join_id_pair3,
+                   /*parent_document_ids=*/std::vector<DocumentId>{1, 2, 4}));
+    ICING_ASSERT_OK(
+        index->Put(child_join_id_pair4,
+                   /*parent_document_ids=*/std::vector<DocumentId>{1}));
+    ICING_ASSERT_OK(
+        index->Put(child_join_id_pair5,
+                   /*parent_document_ids=*/std::vector<DocumentId>{1, 2}));
+    ICING_ASSERT_OK(
+        index->Put(child_join_id_pair6,
+                   /*parent_document_ids=*/std::vector<DocumentId>{3}));
+    ICING_ASSERT_OK(
+        index->Put(child_join_id_pair7,
+                   /*parent_document_ids=*/std::vector<DocumentId>{1}));
+  } else {
+    // With the dual-tracking implementation, we need to use the Put overload
+    // that takes qualified ids to populate the qualified id mapper, which is
+    // the source of truth during Optimize().
+    QualifiedId parent_qid1("ns", "uri1");
+    QualifiedId parent_qid2("ns", "uri2");
+    QualifiedId parent_qid3("ns", "uri3");
+    QualifiedId parent_qid4("ns", "uri4");
+    ICING_ASSERT_OK(
+        index->Put(document_store_.get(), child_join_id_pair1,
+                   /*parent_qualified_ids=*/
+                   std::vector<QualifiedId>{parent_qid1, parent_qid3}));
+    ICING_ASSERT_OK(index->Put(
+        document_store_.get(), child_join_id_pair2,
+        /*parent_qualified_ids=*/std::vector<QualifiedId>{parent_qid2}));
+    ICING_ASSERT_OK(index->Put(
+        document_store_.get(), child_join_id_pair3,
+        /*parent_qualified_ids=*/
+        std::vector<QualifiedId>{parent_qid1, parent_qid2, parent_qid4}));
+    ICING_ASSERT_OK(index->Put(
+        document_store_.get(), child_join_id_pair4,
+        /*parent_qualified_ids=*/std::vector<QualifiedId>{parent_qid1}));
+    ICING_ASSERT_OK(
+        index->Put(document_store_.get(), child_join_id_pair5,
+                   /*parent_qualified_ids=*/
+                   std::vector<QualifiedId>{parent_qid1, parent_qid2}));
+    ICING_ASSERT_OK(index->Put(
+        document_store_.get(), child_join_id_pair6,
+        /*parent_qualified_ids=*/std::vector<QualifiedId>{parent_qid3}));
+    ICING_ASSERT_OK(index->Put(
+        document_store_.get(), child_join_id_pair7,
+        /*parent_qualified_ids=*/std::vector<QualifiedId>{parent_qid1}));
+  }
 
   ASSERT_THAT(index, Pointee(SizeIs(11)));
   index->set_last_added_document_id(107);
@@ -1260,21 +1682,32 @@
   document_id_old_to_new[105] = 14;
   document_id_old_to_new[106] = 15;
 
+  // Update the document store so that Join index can get the correct document
+  // ids for parents.
+  ICING_ASSERT_OK(document_store_->Delete("ns", "uri0", /*current_time_ms=*/0));
+  ICING_ASSERT_OK(document_store_->Delete("ns", "uri3", /*current_time_ms=*/0));
+  ASSERT_NO_FATAL_FAILURE(OptimizeDocumentStore());
+
   // Note: namespace_id_old_to_new is not used in
   // QualifiedIdJoinIndexImplV3::Optimize.
   DocumentId new_last_added_document_id = 15;
-  EXPECT_THAT(
-      index->Optimize(document_id_old_to_new, /*namespace_id_old_to_new=*/{},
-                      new_last_added_document_id),
-      IsOk());
-  EXPECT_THAT(index, Pointee(SizeIs(5)));
+  EXPECT_THAT(index->Optimize(document_store_.get(), document_id_old_to_new,
+                              /*namespace_id_old_to_new=*/{},
+                              new_last_added_document_id),
+              IsOk());
+  if (!feature_flags_->enable_non_existent_qualified_id_join()) {
+    EXPECT_THAT(index, Pointee(SizeIs(5)));
+  } else {
+    // Join index will no longer drop join relations for non-existent parents.
+    EXPECT_THAT(index, Pointee(SizeIs(7)));
+  }
   EXPECT_THAT(index->last_added_document_id(), Eq(new_last_added_document_id));
 
   // Verify document 0 (originally document 1)
   // - Child docs 101, 104, 105 become 11, 13, 14.
   // - Child docs 103, 107 are deleted.
   EXPECT_THAT(
-      index->Get(/*parent_document_id=*/0),
+      index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/0),
       IsOkAndHolds(ElementsAre(
           DocumentJoinIdPair(/*document_id=*/11, /*joinable_property_id=*/0),
           DocumentJoinIdPair(/*document_id=*/13, /*joinable_property_id=*/0),
@@ -1284,20 +1717,23 @@
   // - Child docs 102, 105 become 12, 14.
   // - Child doc 103 is deleted.
   EXPECT_THAT(
-      index->Get(/*parent_document_id=*/1),
+      index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/1),
       IsOkAndHolds(ElementsAre(
           DocumentJoinIdPair(/*document_id=*/12, /*joinable_property_id=*/0),
           DocumentJoinIdPair(/*document_id=*/14, /*joinable_property_id=*/0))));
 
   // Verify document 2 (originally document 4)
   // - Child doc 103 is deleted.
-  EXPECT_THAT(index->Get(/*parent_document_id=*/2), IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/2),
+              IsOkAndHolds(IsEmpty()));
 
   // Verify document 3 and 4:
   // - These 2 doc ids don't exist after optimize.
   // - The relations for the original document 3 and 4 should be deleted.
-  EXPECT_THAT(index->Get(/*parent_document_id=*/3), IsOkAndHolds(IsEmpty()));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/4), IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/3),
+              IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/4),
+              IsOkAndHolds(IsEmpty()));
 
   // Verify Put API should work normally after Optimize().
   DocumentJoinIdPair another_child_join_id_pair(/*document_id=*/16,
@@ -1308,10 +1744,15 @@
       IsOk());
   index->set_last_added_document_id(16);
 
-  EXPECT_THAT(index, Pointee(SizeIs(8)));
+  if (!feature_flags_->enable_non_existent_qualified_id_join()) {
+    EXPECT_THAT(index, Pointee(SizeIs(8)));
+  } else {
+    // Join index will no longer drop join relations for non-existent parents.
+    EXPECT_THAT(index, Pointee(SizeIs(10)));
+  }
   EXPECT_THAT(index->last_added_document_id(), Eq(16));
   EXPECT_THAT(
-      index->Get(/*parent_document_id=*/0),
+      index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/0),
       IsOkAndHolds(ElementsAre(DocumentJoinIdPair(/*document_id=*/11,
                                                   /*joinable_property_id=*/0),
                                DocumentJoinIdPair(/*document_id=*/13,
@@ -1319,13 +1760,21 @@
                                DocumentJoinIdPair(/*document_id=*/14,
                                                   /*joinable_property_id=*/0),
                                another_child_join_id_pair)));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/2),
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/2),
               IsOkAndHolds(ElementsAre(another_child_join_id_pair)));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/3),
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/3),
               IsOkAndHolds(ElementsAre(another_child_join_id_pair)));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test, OptimizeOutOfRangeParentDocumentId) {
+TEST_P(QualifiedIdJoinIndexImplV3Test, OptimizeOutOfRangeParentDocumentId) {
+  if (feature_flags_->enable_non_existent_qualified_id_join()) {
+    // Not applicable for the dual-tracking implementation. Optimize() will use
+    // the qualified id mapper as the source of truth, and parent document ids
+    // will be read from the document store directly using qualified ids,
+    // instead of document_id_old_to_new.
+    return;
+  }
+
   // Create new qualified id join index
   ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
                              QualifiedIdJoinIndexImplV3::Create(
@@ -1369,19 +1818,23 @@
   // QualifiedIdJoinIndexImplV3::Optimize.
   DocumentId new_last_added_document_id = 12;
   EXPECT_THAT(
-      index->Optimize(document_id_old_to_new, /*namespace_id_old_to_new=*/{},
+      index->Optimize(document_store_.get(), document_id_old_to_new,
+                      /*namespace_id_old_to_new=*/{},
                       new_last_added_document_id),
       StatusIs(libtextclassifier3::StatusCode::INTERNAL,
                HasSubstr("Qualified id join index data parent document id is "
                          "out of range. The index may have been corrupted.")));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test, OptimizeOutOfRangeChildDocumentId) {
+TEST_P(QualifiedIdJoinIndexImplV3Test, OptimizeOutOfRangeChildDocumentId) {
   // Create new qualified id join index
   ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
                              QualifiedIdJoinIndexImplV3::Create(
                                  filesystem_, working_path_, *feature_flags_));
 
+  // Add documents for parent doc ids to the document store.
+  ASSERT_NO_FATAL_FAILURE(FillDocumentStore(/*num_documents=*/3));
+
   // Create 2 parent and 3 child documents (with N to N joins):
   // - Document 1: 101, 106, 108
   // - Document 120: 101
@@ -1392,15 +1845,30 @@
                                          /*joinable_property_id=*/0);
   DocumentJoinIdPair child_join_id_pair3(/*document_id=*/108,
                                          /*joinable_property_id=*/0);
-  ICING_ASSERT_OK(
-      index->Put(child_join_id_pair1,
-                 /*parent_document_ids=*/std::vector<DocumentId>{1, 2}));
-  ICING_ASSERT_OK(
-      index->Put(child_join_id_pair2,
-                 /*parent_document_ids=*/std::vector<DocumentId>{1}));
-  ICING_ASSERT_OK(
-      index->Put(child_join_id_pair3,
-                 /*parent_document_ids=*/std::vector<DocumentId>{1}));
+  if (!feature_flags_->enable_non_existent_qualified_id_join()) {
+    ICING_ASSERT_OK(
+        index->Put(child_join_id_pair1,
+                   /*parent_document_ids=*/std::vector<DocumentId>{1, 2}));
+    ICING_ASSERT_OK(
+        index->Put(child_join_id_pair2,
+                   /*parent_document_ids=*/std::vector<DocumentId>{1}));
+    ICING_ASSERT_OK(
+        index->Put(child_join_id_pair3,
+                   /*parent_document_ids=*/std::vector<DocumentId>{1}));
+  } else {
+    QualifiedId parent_qid1("ns", "uri1");
+    QualifiedId parent_qid2("ns", "uri2");
+    ICING_ASSERT_OK(
+        index->Put(document_store_.get(), child_join_id_pair1,
+                   /*parent_qualified_ids=*/
+                   std::vector<QualifiedId>{parent_qid1, parent_qid2}));
+    ICING_ASSERT_OK(index->Put(
+        document_store_.get(), child_join_id_pair2,
+        /*parent_qualified_ids=*/std::vector<QualifiedId>{parent_qid1}));
+    ICING_ASSERT_OK(index->Put(
+        document_store_.get(), child_join_id_pair3,
+        /*parent_qualified_ids=*/std::vector<QualifiedId>{parent_qid1}));
+  }
 
   ASSERT_THAT(index, Pointee(SizeIs(4)));
   index->set_last_added_document_id(120);
@@ -1416,23 +1884,32 @@
   document_id_old_to_new[101] = 11;
   document_id_old_to_new[106] = 12;
 
+  // Update the document store so that Join index can get the correct document
+  // ids for parents.
+  ICING_ASSERT_OK(document_store_->Delete("ns", "uri0", /*current_time_ms=*/0));
+  ASSERT_NO_FATAL_FAILURE(OptimizeDocumentStore());
+
   // Note: namespace_id_old_to_new is not used in
   // QualifiedIdJoinIndexImplV3::Optimize.
   DocumentId new_last_added_document_id = 12;
   EXPECT_THAT(
-      index->Optimize(document_id_old_to_new, /*namespace_id_old_to_new=*/{},
+      index->Optimize(document_store_.get(), document_id_old_to_new,
+                      /*namespace_id_old_to_new=*/{},
                       new_last_added_document_id),
       StatusIs(libtextclassifier3::StatusCode::INTERNAL,
                HasSubstr("Qualified id join index data child document id is "
                          "out of range. The index may have been corrupted.")));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test, OptimizeDeleteAllDocuments) {
+TEST_P(QualifiedIdJoinIndexImplV3Test, OptimizeDeleteAllDocuments) {
   // Create new qualified id join index
   ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
                              QualifiedIdJoinIndexImplV3::Create(
                                  filesystem_, working_path_, *feature_flags_));
 
+  // Add documents for parent doc ids to the document store.
+  ASSERT_NO_FATAL_FAILURE(FillDocumentStore(/*num_documents=*/5));
+
   // Create 4 parent and 7 child documents (with N to N joins):
   // - Document 1: 101, 103, 104, 105, 107
   // - Document 2: 102, 103, 105
@@ -1453,27 +1930,59 @@
                                          /*joinable_property_id=*/0);
   DocumentJoinIdPair child_join_id_pair7(/*document_id=*/107,
                                          /*joinable_property_id=*/0);
-  ICING_ASSERT_OK(
-      index->Put(child_join_id_pair1,
-                 /*parent_document_ids=*/std::vector<DocumentId>{1, 3}));
-  ICING_ASSERT_OK(
-      index->Put(child_join_id_pair2,
-                 /*parent_document_ids=*/std::vector<DocumentId>{2}));
-  ICING_ASSERT_OK(
-      index->Put(child_join_id_pair3,
-                 /*parent_document_ids=*/std::vector<DocumentId>{1, 2, 4}));
-  ICING_ASSERT_OK(
-      index->Put(child_join_id_pair4,
-                 /*parent_document_ids=*/std::vector<DocumentId>{1}));
-  ICING_ASSERT_OK(
-      index->Put(child_join_id_pair5,
-                 /*parent_document_ids=*/std::vector<DocumentId>{1, 2}));
-  ICING_ASSERT_OK(
-      index->Put(child_join_id_pair6,
-                 /*parent_document_ids=*/std::vector<DocumentId>{3}));
-  ICING_ASSERT_OK(
-      index->Put(child_join_id_pair7,
-                 /*parent_document_ids=*/std::vector<DocumentId>{1}));
+
+  if (!feature_flags_->enable_non_existent_qualified_id_join()) {
+    ICING_ASSERT_OK(
+        index->Put(child_join_id_pair1,
+                   /*parent_document_ids=*/std::vector<DocumentId>{1, 3}));
+    ICING_ASSERT_OK(
+        index->Put(child_join_id_pair2,
+                   /*parent_document_ids=*/std::vector<DocumentId>{2}));
+    ICING_ASSERT_OK(
+        index->Put(child_join_id_pair3,
+                   /*parent_document_ids=*/std::vector<DocumentId>{1, 2, 4}));
+    ICING_ASSERT_OK(
+        index->Put(child_join_id_pair4,
+                   /*parent_document_ids=*/std::vector<DocumentId>{1}));
+    ICING_ASSERT_OK(
+        index->Put(child_join_id_pair5,
+                   /*parent_document_ids=*/std::vector<DocumentId>{1, 2}));
+    ICING_ASSERT_OK(
+        index->Put(child_join_id_pair6,
+                   /*parent_document_ids=*/std::vector<DocumentId>{3}));
+    ICING_ASSERT_OK(
+        index->Put(child_join_id_pair7,
+                   /*parent_document_ids=*/std::vector<DocumentId>{1}));
+  } else {
+    QualifiedId parent_qid1("ns", "uri1");
+    QualifiedId parent_qid2("ns", "uri2");
+    QualifiedId parent_qid3("ns", "uri3");
+    QualifiedId parent_qid4("ns", "uri4");
+    ICING_ASSERT_OK(
+        index->Put(document_store_.get(), child_join_id_pair1,
+                   /*parent_qualified_ids=*/
+                   std::vector<QualifiedId>{parent_qid1, parent_qid3}));
+    ICING_ASSERT_OK(index->Put(
+        document_store_.get(), child_join_id_pair2,
+        /*parent_qualified_ids=*/std::vector<QualifiedId>{parent_qid2}));
+    ICING_ASSERT_OK(index->Put(
+        document_store_.get(), child_join_id_pair3,
+        /*parent_qualified_ids=*/
+        std::vector<QualifiedId>{parent_qid1, parent_qid2, parent_qid4}));
+    ICING_ASSERT_OK(index->Put(
+        document_store_.get(), child_join_id_pair4,
+        /*parent_qualified_ids=*/std::vector<QualifiedId>{parent_qid1}));
+    ICING_ASSERT_OK(
+        index->Put(document_store_.get(), child_join_id_pair5,
+                   /*parent_qualified_ids=*/
+                   std::vector<QualifiedId>{parent_qid1, parent_qid2}));
+    ICING_ASSERT_OK(index->Put(
+        document_store_.get(), child_join_id_pair6,
+        /*parent_qualified_ids=*/std::vector<QualifiedId>{parent_qid3}));
+    ICING_ASSERT_OK(index->Put(
+        document_store_.get(), child_join_id_pair7,
+        /*parent_qualified_ids=*/std::vector<QualifiedId>{parent_qid1}));
+  }
 
   ASSERT_THAT(index, Pointee(SizeIs(11)));
   index->set_last_added_document_id(107);
@@ -1482,22 +1991,262 @@
   // Delete all documents.
   std::vector<DocumentId> document_id_old_to_new(108, kInvalidDocumentId);
 
+  // Update the document store so that Join index can get the correct document
+  // ids for parents.
+  ICING_ASSERT_OK(document_store_->Delete("ns", "uri0", /*current_time_ms=*/0));
+  ICING_ASSERT_OK(document_store_->Delete("ns", "uri1", /*current_time_ms=*/0));
+  ICING_ASSERT_OK(document_store_->Delete("ns", "uri2", /*current_time_ms=*/0));
+  ICING_ASSERT_OK(document_store_->Delete("ns", "uri3", /*current_time_ms=*/0));
+  ICING_ASSERT_OK(document_store_->Delete("ns", "uri4", /*current_time_ms=*/0));
+  ASSERT_NO_FATAL_FAILURE(OptimizeDocumentStore());
+
   // Note: namespace_id_old_to_new is not used in
   // QualifiedIdJoinIndexImplV3::Optimize.
   DocumentId new_last_added_document_id = kInvalidDocumentId;
-  EXPECT_THAT(
-      index->Optimize(document_id_old_to_new, /*namespace_id_old_to_new=*/{},
-                      new_last_added_document_id),
-      IsOk());
+  EXPECT_THAT(index->Optimize(document_store_.get(), document_id_old_to_new,
+                              /*namespace_id_old_to_new=*/{},
+                              new_last_added_document_id),
+              IsOk());
   EXPECT_THAT(index, Pointee(IsEmpty()));
   EXPECT_THAT(index->last_added_document_id(), Eq(new_last_added_document_id));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/0), IsOkAndHolds(IsEmpty()));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/1), IsOkAndHolds(IsEmpty()));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/2), IsOkAndHolds(IsEmpty()));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/3), IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/0),
+              IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/1),
+              IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/2),
+              IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/3),
+              IsOkAndHolds(IsEmpty()));
 }
 
-TEST_F(QualifiedIdJoinIndexImplV3Test, Clear) {
+TEST_P(QualifiedIdJoinIndexImplV3Test, OptimizeWithMissingParents) {
+  // Create new qualified id join index
+  ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
+                             QualifiedIdJoinIndexImplV3::Create(
+                                 filesystem_, working_path_, *feature_flags_));
+
+  // Add a document for the existing parent.
+  // Doc id 0: dummy
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
+  // Doc id 1
+  ICING_ASSERT_OK(document_store_->Put(
+      document_util::CreateDocumentWrapper(DocumentBuilder()
+                                               .SetKey("ns", "uri1_existing")
+                                               .SetSchema("type")
+                                               .Build())));
+
+  // Add join data with both regular and missing parents.
+  DocumentId parent_doc_id1 = 1;
+  QualifiedId parent_qid1("ns", "uri1_existing");
+  QualifiedId missing_parent_qualified_id1("namespace", "uri1_missing");
+  QualifiedId missing_parent_qualified_id2("namespace", "uri2_missing");
+
+  DocumentJoinIdPair child_join_id_pair1(/*document_id=*/101,
+                                         /*joinable_property_id=*/0);
+  DocumentJoinIdPair child_join_id_pair2(/*document_id=*/102,
+                                         /*joinable_property_id=*/0);
+  DocumentJoinIdPair child_join_id_pair3(/*document_id=*/103,
+                                         /*joinable_property_id=*/0);
+
+  if (!feature_flags_->enable_non_existent_qualified_id_join()) {
+    ICING_ASSERT_OK(index->Put(
+        child_join_id_pair1,
+        /*parent_document_ids=*/std::vector<DocumentId>{parent_doc_id1}));
+  } else {
+    ICING_ASSERT_OK(index->Put(
+        document_store_.get(), child_join_id_pair1,
+        /*parent_qualified_ids=*/std::vector<QualifiedId>{parent_qid1}));
+  }
+  auto missing_parent1_put_status =
+      index->Put(document_store_.get(), child_join_id_pair2,
+                 /*parent_qualified_ids=*/
+                 std::vector<QualifiedId>{missing_parent_qualified_id1});
+  auto missing_parent2_put_status =
+      index->Put(document_store_.get(), child_join_id_pair3,
+                 /*parent_qualified_ids=*/
+                 std::vector<QualifiedId>{missing_parent_qualified_id1,
+                                          missing_parent_qualified_id2});
+
+  if (!feature_flags_->enable_non_existent_qualified_id_join()) {
+    // If the flag is not enabled, put should fail.
+    EXPECT_THAT(missing_parent1_put_status,
+                StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
+    EXPECT_THAT(missing_parent2_put_status,
+                StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
+    ASSERT_THAT(index, Pointee(SizeIs(1)));
+  } else {
+    ICING_ASSERT_OK(missing_parent1_put_status);
+    ICING_ASSERT_OK(missing_parent2_put_status);
+    ASSERT_THAT(index, Pointee(SizeIs(4)));
+  }
+  index->set_last_added_document_id(103);
+  ASSERT_THAT(index->last_added_document_id(), Eq(103));
+
+  // Remap doc ids. Delete child 102.
+  std::vector<DocumentId> document_id_old_to_new(104, kInvalidDocumentId);
+  document_id_old_to_new[1] = 0;     // parent_doc_id1 -> 0
+  document_id_old_to_new[101] = 11;  // child_join_id_pair1 -> 11
+  document_id_old_to_new[103] = 13;  // child_join_id_pair3 -> 13
+
+  // Update the document store so that Join index can get the correct document
+  // ids for parents.
+  ICING_ASSERT_OK(document_store_->Delete("ns", "uri0", /*current_time_ms=*/0));
+  ASSERT_NO_FATAL_FAILURE(OptimizeDocumentStore());
+
+  DocumentId new_last_added_document_id = 13;
+  ICING_ASSERT_OK(index->Optimize(document_store_.get(), document_id_old_to_new,
+                                  /*namespace_id_old_to_new=*/{},
+                                  new_last_added_document_id));
+
+  // Verify parent_doc_id1 (now 0).
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/0),
+              IsOkAndHolds(ElementsAre(DocumentJoinIdPair(
+                  /*document_id=*/11, /*joinable_property_id=*/0))));
+
+  if (!feature_flags_->enable_non_existent_qualified_id_join()) {
+    // Nothing more to check if the flag is disabled.
+    EXPECT_THAT(index, Pointee(SizeIs(1)));
+    return;
+  }
+
+  EXPECT_THAT(index, Pointee(SizeIs(3)));  // 1 from regular, 2 from missing
+
+  // Migrate missing parents and verify.
+  DocumentId new_parent_doc_id1 = 1;
+  DocumentId new_parent_doc_id2 = 2;
+  ICING_ASSERT_OK(
+      index->MigrateParent(missing_parent_qualified_id1, new_parent_doc_id1));
+  ICING_ASSERT_OK(
+      index->MigrateParent(missing_parent_qualified_id2, new_parent_doc_id2));
+
+  // child 102 was deleted, so missing_parent1 should only have child 103 (now
+  // 13).
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(new_parent_doc_id1),
+              IsOkAndHolds(ElementsAre(DocumentJoinIdPair(
+                  /*document_id=*/13, /*joinable_property_id=*/0))));
+  // missing_parent2 had child 103 (now 13).
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(new_parent_doc_id2),
+              IsOkAndHolds(ElementsAre(DocumentJoinIdPair(
+                  /*document_id=*/13, /*joinable_property_id=*/0))));
+}
+
+TEST_P(QualifiedIdJoinIndexImplV3Test,
+       OptimizeShouldKeepJoinDataOfDeletedParents) {
+  if (!feature_flags_->enable_non_existent_qualified_id_join()) {
+    return;
+  }
+
+  // Index a parent document.
+  ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
+                             QualifiedIdJoinIndexImplV3::Create(
+                                 filesystem_, working_path_, *feature_flags_));
+
+  QualifiedId parent_qualified_id("ns", "uri0");
+  // Doc id 0
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
+  // Doc id 1
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build())));
+
+  auto doc_id_or = document_store_->GetDocumentId("ns", "uri0");
+  ASSERT_THAT(doc_id_or, IsOk());
+  DocumentId parent_doc_id = doc_id_or.ValueOrDie();
+  ASSERT_THAT(parent_doc_id, Eq(0));
+
+  // Index a join relation.
+  DocumentJoinIdPair child_join_id_pair(/*document_id=*/100,
+                                        /*joinable_property_id=*/0);
+  ICING_ASSERT_OK(index->Put(
+      document_store_.get(), child_join_id_pair,
+      /*parent_qualified_ids=*/std::vector<QualifiedId>{parent_qualified_id}));
+
+  // Verify the join relation.
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(parent_doc_id),
+              IsOkAndHolds(ElementsAre(child_join_id_pair)));
+
+  // Delete the parent document.
+  ICING_ASSERT_OK(document_store_->Delete("ns", "uri0", /*current_time_ms=*/0));
+
+  // Run Optimize.
+  // After optimize, doc "ns", "uri1" will be doc 0.
+  std::vector<DocumentId> document_id_old_to_new(101, kInvalidDocumentId);
+  document_id_old_to_new[1] = 0;     // "ns", "uri1" moved to docid 0.
+  document_id_old_to_new[100] = 99;  // child doc 100 becomes 99.
+  ASSERT_NO_FATAL_FAILURE(OptimizeDocumentStore());
+
+  DocumentId new_last_added_document_id = 99;
+  ICING_ASSERT_OK(index->Optimize(document_store_.get(), document_id_old_to_new,
+                                  /*namespace_id_old_to_new=*/{},
+                                  new_last_added_document_id));
+
+  // The new document with docid 0 is the old document 1, which was not a
+  // parent.
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(0),
+              IsOkAndHolds(IsEmpty()));
+
+  // Migrating the parent to a new doc id and check the child is still there.
+  DocumentId new_parent_doc_id = 1;
+  ICING_ASSERT_OK(index->MigrateParent(parent_qualified_id, new_parent_doc_id));
+  DocumentJoinIdPair new_child_join_id_pair(/*document_id=*/99,
+                                            /*joinable_property_id=*/0);
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(new_parent_doc_id),
+              IsOkAndHolds(ElementsAre(new_child_join_id_pair)));
+}
+
+TEST_P(QualifiedIdJoinIndexImplV3Test,
+       OptimizeShouldAllowAddingChildrenToDeletedParents) {
+  if (!feature_flags_->enable_non_existent_qualified_id_join()) {
+    return;
+  }
+
+  ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
+                             QualifiedIdJoinIndexImplV3::Create(
+                                 filesystem_, working_path_, *feature_flags_));
+
+  QualifiedId parent_qualified_id("ns", "uri0");
+  // Doc id 0
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
+  ASSERT_THAT(document_store_->GetDocumentId("ns", "uri0"), IsOkAndHolds(0));
+
+  // Index a join relation.
+  DocumentJoinIdPair child1(/*document_id=*/100, /*joinable_property_id=*/0);
+  ICING_ASSERT_OK(index->Put(
+      document_store_.get(), child1,
+      /*parent_qualified_ids=*/std::vector<QualifiedId>{parent_qualified_id}));
+
+  // Delete the parent document.
+  ICING_ASSERT_OK(document_store_->Delete("ns", "uri0", /*current_time_ms=*/0));
+
+  // Run Optimize.
+  std::vector<DocumentId> document_id_old_to_new(101, kInvalidDocumentId);
+  document_id_old_to_new[100] = 99;  // child doc 100 becomes 99.
+  ASSERT_NO_FATAL_FAILURE(OptimizeDocumentStore());
+
+  DocumentId new_last_added_document_id = 99;
+  ICING_ASSERT_OK(index->Optimize(document_store_.get(), document_id_old_to_new,
+                                  /*namespace_id_old_to_new=*/{},
+                                  new_last_added_document_id));
+
+  // Index a new child for the deleted parent.
+  DocumentJoinIdPair child2(/*document_id=*/101, /*joinable_property_id=*/0);
+  ICING_ASSERT_OK(index->Put(
+      document_store_.get(), child2,
+      /*parent_qualified_ids=*/std::vector<QualifiedId>{parent_qualified_id}));
+
+  // Migrating the parent to a new doc id and check both children are there.
+  DocumentId new_parent_doc_id = 1;
+  ICING_ASSERT_OK(index->MigrateParent(parent_qualified_id, new_parent_doc_id));
+  DocumentJoinIdPair remapped_child1(/*document_id=*/99,
+                                     /*joinable_property_id=*/0);
+  ICING_ASSERT_OK_AND_ASSIGN(
+      auto view, index->GetDocumentJoinIdPairArrayView(new_parent_doc_id));
+  EXPECT_THAT(view, ElementsAre(remapped_child1, child2));
+}
+
+TEST_P(QualifiedIdJoinIndexImplV3Test, Clear) {
   // Create new qualified id join index
   ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
                              QualifiedIdJoinIndexImplV3::Create(
@@ -1534,9 +2283,12 @@
   EXPECT_THAT(index->Clear(), IsOk());
   EXPECT_THAT(index, Pointee(IsEmpty()));
   EXPECT_THAT(index->last_added_document_id(), Eq(kInvalidDocumentId));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/0), IsOkAndHolds(IsEmpty()));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/1), IsOkAndHolds(IsEmpty()));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/2), IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/0),
+              IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/1),
+              IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/2),
+              IsOkAndHolds(IsEmpty()));
 
   // Join index should be able to work normally after Clear().
   EXPECT_THAT(index->Put(child_join_id_pair4,
@@ -1546,7 +2298,7 @@
 
   EXPECT_THAT(index, Pointee(SizeIs(1)));
   EXPECT_THAT(index->last_added_document_id(), Eq(105));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/5),
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/5),
               IsOkAndHolds(ElementsAre(child_join_id_pair4)));
 
   ICING_ASSERT_OK(index->PersistToDisk());
@@ -1558,10 +2310,30 @@
                                                 *feature_flags_));
   EXPECT_THAT(index, Pointee(SizeIs(1)));
   EXPECT_THAT(index->last_added_document_id(), Eq(105));
-  EXPECT_THAT(index->Get(/*parent_document_id=*/5),
+  EXPECT_THAT(index->GetDocumentJoinIdPairArrayView(/*parent_document_id=*/5),
               IsOkAndHolds(ElementsAre(child_join_id_pair4)));
 }
 
+TEST_P(QualifiedIdJoinIndexImplV3Test, V2ApiShouldBeUnimplemented) {
+  ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<QualifiedIdJoinIndexImplV3> index,
+                             QualifiedIdJoinIndexImplV3::Create(
+                                 filesystem_, working_path_, *feature_flags_));
+
+  EXPECT_THAT(
+      index->Put(/*schema_type_id=*/0, /*joinable_property_id=*/0,
+                 /*document_id=*/0, /*ref_namespace_id_uri_fingerprints=*/{}),
+      StatusIs(libtextclassifier3::StatusCode::UNIMPLEMENTED));
+
+  EXPECT_THAT(index->GetIterator(/*schema_type_id=*/0,
+                                 /*joinable_property_id=*/0),
+              StatusIs(libtextclassifier3::StatusCode::UNIMPLEMENTED));
+}
+
+INSTANTIATE_TEST_SUITE_P(QualifiedIdJoinIndexImplV3Test,
+                         QualifiedIdJoinIndexImplV3Test,
+                         // Parameter: enable_non_existent_qualified_id_join
+                         testing::Values(true, false));
+
 }  // namespace
 
 }  // namespace lib

diff --git a/icing/join/qualified-id-join-index.h b/icing/join/qualified-id-join-index.h
index 55427c1..0349a15 100644
--- a/icing/join/qualified-id-join-index.h
+++ b/icing/join/qualified-id-join-index.h

@@ -18,7 +18,6 @@
 #include <cstdint>
 #include <memory>
 #include <string>
-#include <string_view>
 #include <utility>
 #include <vector>
 
@@ -28,9 +27,11 @@
 #include "icing/file/persistent-storage.h"
 #include "icing/join/document-id-to-join-info.h"
 #include "icing/join/document-join-id-pair.h"
+#include "icing/join/qualified-id.h"
 #include "icing/schema/joinable-property.h"
 #include "icing/store/document-filter-data.h"
 #include "icing/store/document-id.h"
+#include "icing/store/document-store.h"
 #include "icing/store/namespace-id-fingerprint.h"
 #include "icing/store/namespace-id.h"
 #include "icing/util/crc32.h"
@@ -52,6 +53,37 @@
         const = 0;
   };
 
+  // An STL style array view object of DocumentJoinIdPairs for a given parent
+  // document id. This is only used in V3.
+  class DocumentJoinIdPairArrayView {
+   public:
+    using value_type = DocumentJoinIdPair;
+    using iterator = const DocumentJoinIdPair*;
+    using const_iterator = const DocumentJoinIdPair*;
+
+    explicit DocumentJoinIdPairArrayView(const DocumentJoinIdPair* data,
+                                         int size)
+        : data_(data), size_(size) {}
+
+    const DocumentJoinIdPair* data() const { return data_; }
+
+    int size() const { return size_; }
+
+    bool empty() const { return size_ == 0 || data_ == nullptr; }
+
+    iterator begin() { return data_; }
+    iterator end() { return data_ + size_; }
+
+    const_iterator begin() const { return data_; }
+    const_iterator end() const { return data_ + size_; }
+
+    const DocumentJoinIdPair& operator[](int idx) const { return data_[idx]; }
+
+   private:
+    const DocumentJoinIdPair* data_;
+    int size_;
+  };
+
   enum class Version { kV2, kV3 };
 
   static constexpr WorkingPathType kWorkingPathType =
@@ -84,8 +116,8 @@
       std::vector<NamespaceIdFingerprint>&&
           ref_namespace_id_uri_fingerprints) = 0;
 
-  // (v3 only) Puts a new child document and its referenced parent documents
-  // into the join index.
+  // (v3 only, deprecated) Puts a new child document and its referenced parent
+  // documents into the join index.
   //
   // Returns:
   //   - OK on success
@@ -95,6 +127,18 @@
       const DocumentJoinIdPair& child_document_join_id_pair,
       std::vector<DocumentId>&& parent_document_ids) = 0;
 
+  // (v3 only) Puts new join data into the index using qualified ids: adds a new
+  // child document and its referenced parent documents into the join index.
+  //
+  // Returns:
+  //   - OK on success
+  //   - INVALID_ARGUMENT_ERROR if child_document_join_id_pair is invalid
+  //   - Any FileBackedVector or KeyMapper errors
+  virtual libtextclassifier3::Status Put(
+      const DocumentStore* document_store,
+      const DocumentJoinIdPair& child_document_join_id_pair,
+      std::vector<QualifiedId>&& parent_qualified_ids) = 0;
+
   // (v2 only) Returns a JoinDataIterator for iterating through all join data of
   // the specified (schema_type_id, joinable_property_id).
   //
@@ -107,14 +151,16 @@
   GetIterator(SchemaTypeId schema_type_id,
               JoinablePropertyId joinable_property_id) const = 0;
 
-  // (v3 only) Gets the list of joinable children for the given parent document
-  // id.
+  // (v3 only) Gets an array view of all joinable children for the given parent
+  // document id.
   //
   // Returns:
-  //   - A list of children's DocumentJoinIdPair on success
+  //   - A DocumentJoinIdPairArrayView object on success. If there is no edge
+  //     for a valid node_id, then an array view with data() == nullptr and
+  //     size() == 0 will be returned
   //   - Any FileBackedVector errors
-  virtual libtextclassifier3::StatusOr<std::vector<DocumentJoinIdPair>> Get(
-      DocumentId parent_document_id) const = 0;
+  virtual libtextclassifier3::StatusOr<DocumentJoinIdPairArrayView>
+  GetDocumentJoinIdPairArrayView(DocumentId parent_document_id) const = 0;
 
   // Migrates existing join data for a parent document from old_document_id to
   // new_document_id if necessary.
@@ -126,6 +172,18 @@
   virtual libtextclassifier3::Status MigrateParent(
       DocumentId old_document_id, DocumentId new_document_id) = 0;
 
+  // Migrates existing join data for a parent document from its qualified id to
+  // a new document id. This is used when a document that was previously
+  // referenced as a "missing parent" (by its qualified id) is indexed and
+  // assigned a document id.
+  //
+  // Returns:
+  //   - OK on success
+  //   - INVALID_ARGUMENT_ERROR if new_document_id is invalid
+  //   - Any errors, depending on the implementation
+  virtual libtextclassifier3::Status MigrateParent(
+      const QualifiedId& parent_qualified_id, DocumentId new_document_id) = 0;
+
   // Reduces internal file sizes by reclaiming space and ids of deleted
   // documents. Qualified id type joinable index will convert all entries to the
   // new document ids.
@@ -144,6 +202,7 @@
   //     an invalid state and the caller should handle it properly (e.g. discard
   //     and rebuild)
   virtual libtextclassifier3::Status Optimize(
+      const DocumentStore* document_store,
       const std::vector<DocumentId>& document_id_old_to_new,
       const std::vector<NamespaceId>& namespace_id_old_to_new,
       DocumentId new_last_added_document_id) = 0;

diff --git a/icing/join/qualified-id-join-indexing-handler-v2_test.cc b/icing/join/qualified-id-join-indexing-handler-v2_test.cc
index b131efc..d1de824 100644
--- a/icing/join/qualified-id-join-indexing-handler-v2_test.cc
+++ b/icing/join/qualified-id-join-indexing-handler-v2_test.cc

@@ -32,6 +32,7 @@
 #include "icing/join/qualified-id-join-index.h"
 #include "icing/join/qualified-id-join-indexing-handler.h"
 #include "icing/join/qualified-id.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/portable/platform.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/schema.pb.h"
@@ -153,14 +154,18 @@
                 IsTrue());
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
-        DocumentStore::Create(&filesystem_, doc_store_dir_, &fake_clock_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, doc_store_dir_, &fake_clock_, schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     doc_store_ = std::move(create_result.document_store);
 
     // Get FakeType related ids.
@@ -242,20 +247,20 @@
 
 TEST_F(QualifiedIdJoinIndexingHandlerV2Test,
        CreationWithNullPointerShouldFail) {
-  EXPECT_THAT(
-      QualifiedIdJoinIndexingHandler::Create(
-          /*clock=*/nullptr, doc_store_.get(), qualified_id_join_index_.get()),
-      StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
+  EXPECT_THAT(QualifiedIdJoinIndexingHandler::Create(
+                  /*clock=*/nullptr, doc_store_.get(),
+                  qualified_id_join_index_.get(), feature_flags_.get()),
+              StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
 
-  EXPECT_THAT(
-      QualifiedIdJoinIndexingHandler::Create(
-          &fake_clock_, /*doc_store=*/nullptr, qualified_id_join_index_.get()),
-      StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
+  EXPECT_THAT(QualifiedIdJoinIndexingHandler::Create(
+                  &fake_clock_, /*doc_store=*/nullptr,
+                  qualified_id_join_index_.get(), feature_flags_.get()),
+              StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
 
-  EXPECT_THAT(
-      QualifiedIdJoinIndexingHandler::Create(
-          &fake_clock_, doc_store_.get(), /*qualified_id_join_index=*/nullptr),
-      StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
+  EXPECT_THAT(QualifiedIdJoinIndexingHandler::Create(
+                  &fake_clock_, doc_store_.get(),
+                  /*qualified_id_join_index=*/nullptr, feature_flags_.get()),
+              StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
 }
 
 TEST_F(QualifiedIdJoinIndexingHandlerV2Test, HandleJoinableProperty) {
@@ -267,14 +272,25 @@
           .SetSchema(std::string(kReferencedType))
           .AddStringProperty(std::string(kPropertyName), "one")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             doc_store_->Put(referenced_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument ref_tokenized_document,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(referenced_document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store_->Put(ref_tokenized_document.document_wrapper()));
   DocumentId ref_doc_id = put_result.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       NamespaceId ref_doc_ns_id,
-      doc_store_->GetNamespaceId(referenced_document.namespace_()));
+      doc_store_->GetNamespaceId(
+          ref_tokenized_document.document_wrapper().document().namespace_()));
   NamespaceIdFingerprint ref_doc_nsid_uri_fingerprint(
-      /*namespace_id=*/ref_doc_ns_id, /*target_str=*/referenced_document.uri());
+      /*namespace_id=*/ref_doc_ns_id,
+      /*target_str=*/ref_tokenized_document.document_wrapper()
+          .document()
+          .uri());
   ASSERT_THAT(doc_store_->GetDocumentId(ref_doc_nsid_uri_fingerprint),
               IsOkAndHolds(ref_doc_id));
 
@@ -286,12 +302,15 @@
           .AddStringProperty(std::string(kPropertyQualifiedId),
                              "pkg$db/ns#ref_type/1")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(put_result, doc_store_->Put(document));
-  DocumentId doc_id = put_result.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      put_result, doc_store_->Put(tokenized_document.document_wrapper()));
+  DocumentId doc_id = put_result.new_document_id;
 
   // Handle document.
   ASSERT_THAT(qualified_id_join_index_->last_added_document_id(),
@@ -299,7 +318,8 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<QualifiedIdJoinIndexingHandler> handler,
       QualifiedIdJoinIndexingHandler::Create(&fake_clock_, doc_store_.get(),
-                                             qualified_id_join_index_.get()));
+                                             qualified_id_join_index_.get(),
+                                             feature_flags_.get()));
   EXPECT_THAT(
       handler->Handle(tokenized_document, doc_id, put_result.old_document_id,
                       /*recovery_mode=*/false, /*put_document_stats=*/nullptr),
@@ -326,15 +346,25 @@
           .SetSchema(std::string(kReferencedType))
           .AddStringProperty(std::string(kPropertyName), "one")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             doc_store_->Put(referenced_document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument ref_tokenized_document1,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(referenced_document1)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      doc_store_->Put(ref_tokenized_document1.document_wrapper()));
   DocumentId ref_doc_id1 = put_result1.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       NamespaceId ref_doc_ns_id1,
-      doc_store_->GetNamespaceId(referenced_document1.namespace_()));
+      doc_store_->GetNamespaceId(
+          ref_tokenized_document1.document_wrapper().document().namespace_()));
   NamespaceIdFingerprint ref_doc_nsid_uri_fingerprint1(
       /*namespace_id=*/ref_doc_ns_id1,
-      /*target_str=*/referenced_document1.uri());
+      /*target_str=*/ref_tokenized_document1.document_wrapper()
+          .document()
+          .uri());
   ASSERT_THAT(doc_store_->GetDocumentId(ref_doc_nsid_uri_fingerprint1),
               IsOkAndHolds(ref_doc_id1));
 
@@ -346,15 +376,25 @@
           .SetSchema(std::string(kReferencedType))
           .AddStringProperty(std::string(kPropertyName), "two")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             doc_store_->Put(referenced_document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument ref_tokenized_document2,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(referenced_document2)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      doc_store_->Put(ref_tokenized_document2.document_wrapper()));
   DocumentId ref_doc_id2 = put_result2.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       NamespaceId ref_doc_ns_id2,
-      doc_store_->GetNamespaceId(referenced_document2.namespace_()));
+      doc_store_->GetNamespaceId(
+          ref_tokenized_document2.document_wrapper().document().namespace_()));
   NamespaceIdFingerprint ref_doc_nsid_uri_fingerprint2(
       /*namespace_id=*/ref_doc_ns_id2,
-      /*target_str=*/referenced_document2.uri());
+      /*target_str=*/ref_tokenized_document2.document_wrapper()
+          .document()
+          .uri());
   ASSERT_THAT(doc_store_->GetDocumentId(ref_doc_nsid_uri_fingerprint2),
               IsOkAndHolds(ref_doc_id2));
 
@@ -378,25 +418,29 @@
           .AddStringProperty(std::string(kPropertyQualifiedId2),
                              "pkg$db/ns#ref_type/1")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             doc_store_->Put(nested_document));
-  DocumentId doc_id = put_result.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
-      TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                nested_document));
+      TokenizedDocument nested_tokenized_document,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(nested_document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store_->Put(nested_tokenized_document.document_wrapper()));
+  DocumentId doc_id = put_result.new_document_id;
 
-  // Handle nested_document.
+  // Handle nested_tokenized_document.
   ASSERT_THAT(qualified_id_join_index_->last_added_document_id(),
               Eq(kInvalidDocumentId));
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<QualifiedIdJoinIndexingHandler> handler,
       QualifiedIdJoinIndexingHandler::Create(&fake_clock_, doc_store_.get(),
-                                             qualified_id_join_index_.get()));
-  EXPECT_THAT(
-      handler->Handle(tokenized_document, doc_id, put_result.old_document_id,
-                      /*recovery_mode=*/false, /*put_document_stats=*/nullptr),
-      IsOk());
+                                             qualified_id_join_index_.get(),
+                                             feature_flags_.get()));
+  EXPECT_THAT(handler->Handle(
+                  nested_tokenized_document, doc_id, put_result.old_document_id,
+                  /*recovery_mode=*/false, /*put_document_stats=*/nullptr),
+              IsOk());
 
   // Verify the state of qualified_id_join_index_ after Handle().
   EXPECT_THAT(qualified_id_join_index_->last_added_document_id(), Eq(doc_id));
@@ -440,13 +484,16 @@
           .AddStringProperty(std::string(kPropertyQualifiedId),
                              std::string(kInvalidFormatQualifiedId))
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             doc_store_->Put(document));
-  DocumentId doc_id = put_result.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store_->Put(tokenized_document.document_wrapper()));
+  DocumentId doc_id = put_result.new_document_id;
 
   // Handle document. Should ignore invalid format qualified id.
   ASSERT_THAT(qualified_id_join_index_->last_added_document_id(),
@@ -454,7 +501,8 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<QualifiedIdJoinIndexingHandler> handler,
       QualifiedIdJoinIndexingHandler::Create(&fake_clock_, doc_store_.get(),
-                                             qualified_id_join_index_.get()));
+                                             qualified_id_join_index_.get(),
+                                             feature_flags_.get()));
   EXPECT_THAT(
       handler->Handle(tokenized_document, doc_id, put_result.old_document_id,
                       /*recovery_mode=*/false, /*put_document_stats=*/nullptr),
@@ -484,13 +532,16 @@
               std::string(kPropertyQualifiedId),
               absl_ports::StrCat(kUnknownNamespace, "#", "ref_type/1"))
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             doc_store_->Put(document));
-  DocumentId doc_id = put_result.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store_->Put(tokenized_document.document_wrapper()));
+  DocumentId doc_id = put_result.new_document_id;
 
   // Handle document.
   ASSERT_THAT(qualified_id_join_index_->last_added_document_id(),
@@ -498,7 +549,8 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<QualifiedIdJoinIndexingHandler> handler,
       QualifiedIdJoinIndexingHandler::Create(&fake_clock_, doc_store_.get(),
-                                             qualified_id_join_index_.get()));
+                                             qualified_id_join_index_.get(),
+                                             feature_flags_.get()));
   EXPECT_THAT(
       handler->Handle(tokenized_document, doc_id, put_result.old_document_id,
                       /*recovery_mode=*/false, /*put_document_stats=*/nullptr),
@@ -520,13 +572,17 @@
                                .SetKey("icing", "fake_type/1")
                                .SetSchema(std::string(kFakeType))
                                .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             doc_store_->Put(document));
-  DocumentId doc_id = put_result.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store_->Put(tokenized_document.document_wrapper()));
+  DocumentId doc_id = put_result.new_document_id;
+
   ASSERT_THAT(tokenized_document.qualified_id_join_properties(), IsEmpty());
 
   // Handle document.
@@ -535,7 +591,8 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<QualifiedIdJoinIndexingHandler> handler,
       QualifiedIdJoinIndexingHandler::Create(&fake_clock_, doc_store_.get(),
-                                             qualified_id_join_index_.get()));
+                                             qualified_id_join_index_.get(),
+                                             feature_flags_.get()));
   EXPECT_THAT(
       handler->Handle(tokenized_document, doc_id, put_result.old_document_id,
                       /*recovery_mode=*/false, /*put_document_stats=*/nullptr),
@@ -562,14 +619,25 @@
           .SetSchema(std::string(kReferencedType))
           .AddStringProperty(std::string(kPropertyName), "one")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             doc_store_->Put(referenced_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument ref_tokenized_document,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(referenced_document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store_->Put(ref_tokenized_document.document_wrapper()));
   DocumentId ref_doc_id = put_result.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       NamespaceId ref_doc_ns_id,
-      doc_store_->GetNamespaceId(referenced_document.namespace_()));
+      doc_store_->GetNamespaceId(
+          ref_tokenized_document.document_wrapper().document().namespace_()));
   NamespaceIdFingerprint ref_doc_nsid_uri_fingerprint(
-      /*namespace_id=*/ref_doc_ns_id, /*target_str=*/referenced_document.uri());
+      /*namespace_id=*/ref_doc_ns_id,
+      /*target_str=*/ref_tokenized_document.document_wrapper()
+          .document()
+          .uri());
   ASSERT_THAT(doc_store_->GetDocumentId(ref_doc_nsid_uri_fingerprint),
               IsOkAndHolds(ref_doc_id));
 
@@ -581,11 +649,13 @@
           .AddStringProperty(std::string(kPropertyQualifiedId),
                              "pkg$db/ns#ref_type/1")
           .Build();
-  ICING_ASSERT_OK(doc_store_->Put(document));
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document)));
+  ICING_ASSERT_OK(doc_store_->Put(tokenized_document.document_wrapper()));
 
   qualified_id_join_index_->set_last_added_document_id(ref_doc_id);
   ASSERT_THAT(qualified_id_join_index_->last_added_document_id(),
@@ -594,7 +664,8 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<QualifiedIdJoinIndexingHandler> handler,
       QualifiedIdJoinIndexingHandler::Create(&fake_clock_, doc_store_.get(),
-                                             qualified_id_join_index_.get()));
+                                             qualified_id_join_index_.get(),
+                                             feature_flags_.get()));
 
   // Handling document with kInvalidDocumentId should cause a failure.
   EXPECT_THAT(
@@ -637,14 +708,25 @@
           .SetSchema(std::string(kReferencedType))
           .AddStringProperty(std::string(kPropertyName), "one")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             doc_store_->Put(referenced_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument ref_tokenized_document,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(referenced_document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store_->Put(ref_tokenized_document.document_wrapper()));
   DocumentId ref_doc_id = put_result.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       NamespaceId ref_doc_ns_id,
-      doc_store_->GetNamespaceId(referenced_document.namespace_()));
+      doc_store_->GetNamespaceId(
+          ref_tokenized_document.document_wrapper().document().namespace_()));
   NamespaceIdFingerprint ref_doc_nsid_uri_fingerprint(
-      /*namespace_id=*/ref_doc_ns_id, /*target_str=*/referenced_document.uri());
+      /*namespace_id=*/ref_doc_ns_id,
+      /*target_str=*/ref_tokenized_document.document_wrapper()
+          .document()
+          .uri());
   ASSERT_THAT(doc_store_->GetDocumentId(ref_doc_nsid_uri_fingerprint),
               IsOkAndHolds(ref_doc_id));
 
@@ -656,17 +738,21 @@
           .AddStringProperty(std::string(kPropertyQualifiedId),
                              "pkg$db/ns#ref_type/1")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(put_result, doc_store_->Put(document));
-  DocumentId doc_id = put_result.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      put_result, doc_store_->Put(tokenized_document.document_wrapper()));
+  DocumentId doc_id = put_result.new_document_id;
 
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<QualifiedIdJoinIndexingHandler> handler,
       QualifiedIdJoinIndexingHandler::Create(&fake_clock_, doc_store_.get(),
-                                             qualified_id_join_index_.get()));
+                                             qualified_id_join_index_.get(),
+                                             feature_flags_.get()));
 
   // Handling document with document_id == last_added_document_id should cause a
   // failure.
@@ -715,14 +801,25 @@
           .SetSchema(std::string(kReferencedType))
           .AddStringProperty(std::string(kPropertyName), "one")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             doc_store_->Put(referenced_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument ref_tokenized_document,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(referenced_document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store_->Put(ref_tokenized_document.document_wrapper()));
   DocumentId ref_doc_id = put_result.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       NamespaceId ref_doc_ns_id,
-      doc_store_->GetNamespaceId(referenced_document.namespace_()));
+      doc_store_->GetNamespaceId(
+          ref_tokenized_document.document_wrapper().document().namespace_()));
   NamespaceIdFingerprint ref_doc_nsid_uri_fingerprint(
-      /*namespace_id=*/ref_doc_ns_id, /*target_str=*/referenced_document.uri());
+      /*namespace_id=*/ref_doc_ns_id,
+      /*target_str=*/ref_tokenized_document.document_wrapper()
+          .document()
+          .uri());
   ASSERT_THAT(doc_store_->GetDocumentId(ref_doc_nsid_uri_fingerprint),
               IsOkAndHolds(ref_doc_id));
 
@@ -734,17 +831,21 @@
           .AddStringProperty(std::string(kPropertyQualifiedId),
                              "pkg$db/ns#ref_type/1")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(put_result, doc_store_->Put(document));
-  DocumentId doc_id = put_result.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      put_result, doc_store_->Put(tokenized_document.document_wrapper()));
+  DocumentId doc_id = put_result.new_document_id;
 
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<QualifiedIdJoinIndexingHandler> handler,
       QualifiedIdJoinIndexingHandler::Create(&fake_clock_, doc_store_.get(),
-                                             qualified_id_join_index_.get()));
+                                             qualified_id_join_index_.get(),
+                                             feature_flags_.get()));
 
   // Handle document with document_id > last_added_document_id in recovery mode.
   // The handler should index this document and update last_added_document_id.
@@ -774,14 +875,25 @@
           .SetSchema(std::string(kReferencedType))
           .AddStringProperty(std::string(kPropertyName), "one")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             doc_store_->Put(referenced_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument ref_tokenized_document,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(referenced_document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store_->Put(ref_tokenized_document.document_wrapper()));
   DocumentId ref_doc_id = put_result.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       NamespaceId ref_doc_ns_id,
-      doc_store_->GetNamespaceId(referenced_document.namespace_()));
+      doc_store_->GetNamespaceId(
+          ref_tokenized_document.document_wrapper().document().namespace_()));
   NamespaceIdFingerprint ref_doc_nsid_uri_fingerprint(
-      /*namespace_id=*/ref_doc_ns_id, /*target_str=*/referenced_document.uri());
+      /*namespace_id=*/ref_doc_ns_id,
+      /*target_str=*/ref_tokenized_document.document_wrapper()
+          .document()
+          .uri());
   ASSERT_THAT(doc_store_->GetDocumentId(ref_doc_nsid_uri_fingerprint),
               IsOkAndHolds(ref_doc_id));
 
@@ -793,17 +905,21 @@
           .AddStringProperty(std::string(kPropertyQualifiedId),
                              "pkg$db/ns#ref_type/1")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(put_result, doc_store_->Put(document));
-  DocumentId doc_id = put_result.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      put_result, doc_store_->Put(tokenized_document.document_wrapper()));
+  DocumentId doc_id = put_result.new_document_id;
 
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<QualifiedIdJoinIndexingHandler> handler,
       QualifiedIdJoinIndexingHandler::Create(&fake_clock_, doc_store_.get(),
-                                             qualified_id_join_index_.get()));
+                                             qualified_id_join_index_.get(),
+                                             feature_flags_.get()));
 
   // Handle document with document_id == last_added_document_id in recovery
   // mode. We should not get any error, but the handler should ignore the

diff --git a/icing/join/qualified-id-join-indexing-handler-v3_test.cc b/icing/join/qualified-id-join-indexing-handler-v3_test.cc
index 7649424..cf70689 100644
--- a/icing/join/qualified-id-join-indexing-handler-v3_test.cc
+++ b/icing/join/qualified-id-join-indexing-handler-v3_test.cc

@@ -29,6 +29,7 @@
 #include "icing/join/qualified-id-join-index-impl-v3.h"
 #include "icing/join/qualified-id-join-indexing-handler.h"
 #include "icing/join/qualified-id.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/portable/platform.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/schema.pb.h"
@@ -77,10 +78,32 @@
 static constexpr std::string_view kPropertyNestedDoc = "nested";
 static constexpr std::string_view kPropertyQualifiedId2 = "qualifiedId2";
 
-class QualifiedIdJoinIndexingHandlerV3Test : public ::testing::Test {
+class QualifiedIdJoinIndexingHandlerV3Test
+    : public ::testing::TestWithParam<bool> {
  protected:
   void SetUp() override {
-    feature_flags_ = std::make_unique<FeatureFlags>(GetTestFeatureFlags());
+    feature_flags_ = std::make_unique<FeatureFlags>(
+        /*enable_circular_schema_definitions=*/true,
+        /*enable_scorable_properties=*/true,
+        /*enable_embedding_quantization=*/true,
+        /*enable_repeated_field_joins=*/true,
+        /*enable_embedding_backup_generation=*/true,
+        /*enable_schema_database=*/true,
+        /*release_backup_schema_file_if_overlay_present=*/true,
+        /*enable_strict_page_byte_size_limit=*/true,
+        /*enable_smaller_decompression_buffer_size=*/true,
+        /*enable_eigen_embedding_scoring=*/true,
+        /*enable_passing_filter_to_children=*/true,
+        /*enable_proto_log_new_header_format=*/true,
+        /*enable_embedding_iterator_v2=*/true,
+        /*enable_reusable_decompression_buffer=*/true,
+        /*enable_schema_type_id_optimization=*/true,
+        /*enable_optimize_improvements=*/true,
+        /*expired_document_purge_threshold_ms=*/0,
+        /*enable_non_existent_qualified_id_join=*/GetParam(),
+        /*enable_skip_set_schema_type_equality_check=*/true,
+        /*enable_embed_query_optimization=*/true,
+        /*enable_schema_definition_deduping=*/true);
 
     if (!IsCfStringTokenization() && !IsReverseJniTokenization()) {
       ICING_ASSERT_OK(
@@ -127,7 +150,7 @@
                 PropertyConfigBuilder()
                     .SetName(kPropertyQualifiedId)
                     .SetDataTypeJoinableString(JOINABLE_VALUE_TYPE_QUALIFIED_ID)
-                    .SetCardinality(CARDINALITY_OPTIONAL)))
+                    .SetCardinality(CARDINALITY_REPEATED)))
             .AddType(
                 SchemaTypeConfigBuilder()
                     .SetType(kNestedType)
@@ -150,14 +173,18 @@
                 IsTrue());
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
-        DocumentStore::Create(&filesystem_, doc_store_dir_, &fake_clock_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, doc_store_dir_, &fake_clock_, schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     doc_store_ = std::move(create_result.document_store);
 
     // Get FakeType related ids.
@@ -220,25 +247,30 @@
   JoinablePropertyId nested_type_joinable_property_id_;
 };
 
-TEST_F(QualifiedIdJoinIndexingHandlerV3Test,
+TEST_P(QualifiedIdJoinIndexingHandlerV3Test,
        CreationWithNullPointerShouldFail) {
-  EXPECT_THAT(
-      QualifiedIdJoinIndexingHandler::Create(
-          /*clock=*/nullptr, doc_store_.get(), qualified_id_join_index_.get()),
-      StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
+  EXPECT_THAT(QualifiedIdJoinIndexingHandler::Create(
+                  /*clock=*/nullptr, doc_store_.get(),
+                  qualified_id_join_index_.get(), feature_flags_.get()),
+              StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
 
-  EXPECT_THAT(
-      QualifiedIdJoinIndexingHandler::Create(
-          &fake_clock_, /*doc_store=*/nullptr, qualified_id_join_index_.get()),
-      StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
+  EXPECT_THAT(QualifiedIdJoinIndexingHandler::Create(
+                  &fake_clock_, /*doc_store=*/nullptr,
+                  qualified_id_join_index_.get(), feature_flags_.get()),
+              StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
 
-  EXPECT_THAT(
-      QualifiedIdJoinIndexingHandler::Create(
-          &fake_clock_, doc_store_.get(), /*qualified_id_join_index=*/nullptr),
-      StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
+  EXPECT_THAT(QualifiedIdJoinIndexingHandler::Create(
+                  &fake_clock_, doc_store_.get(),
+                  /*qualified_id_join_index=*/nullptr, feature_flags_.get()),
+              StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
+
+  EXPECT_THAT(QualifiedIdJoinIndexingHandler::Create(
+                  &fake_clock_, doc_store_.get(),
+                  qualified_id_join_index_.get(), /*feature_flags=*/nullptr),
+              StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
 }
 
-TEST_F(QualifiedIdJoinIndexingHandlerV3Test, HandleJoinableProperty) {
+TEST_P(QualifiedIdJoinIndexingHandlerV3Test, HandleJoinableProperty) {
   // Create and put parent document.
   DocumentProto parent_document =
       DocumentBuilder()
@@ -246,8 +278,15 @@
           .SetSchema(std::string(kReferencedType))
           .AddStringProperty(std::string(kPropertyName), "one")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult parent_put_result,
-                             doc_store_->Put(parent_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument parent_tokenized_document,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(parent_document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult parent_put_result,
+      doc_store_->Put(parent_tokenized_document.document_wrapper()));
 
   // Create and put child document. Also tokenize it.
   DocumentProto child_document =
@@ -257,12 +296,15 @@
           .AddStringProperty(std::string(kPropertyQualifiedId),
                              "pkg$db/ns#ref_type/1")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult child_put_result,
-                             doc_store_->Put(child_document));
   ICING_ASSERT_OK_AND_ASSIGN(
-      TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(child_document)));
+      TokenizedDocument child_tokenized_document,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(child_document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult child_put_result,
+      doc_store_->Put(child_tokenized_document.document_wrapper()));
 
   // Handle document.
   ASSERT_THAT(qualified_id_join_index_->last_added_document_id(),
@@ -270,24 +312,26 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<QualifiedIdJoinIndexingHandler> handler,
       QualifiedIdJoinIndexingHandler::Create(&fake_clock_, doc_store_.get(),
-                                             qualified_id_join_index_.get()));
-  EXPECT_THAT(
-      handler->Handle(tokenized_document, child_put_result.new_document_id,
-                      child_put_result.old_document_id, /*recovery_mode=*/false,
-                      /*put_document_stats=*/nullptr),
-      IsOk());
+                                             qualified_id_join_index_.get(),
+                                             feature_flags_.get()));
+  EXPECT_THAT(handler->Handle(
+                  child_tokenized_document, child_put_result.new_document_id,
+                  child_put_result.old_document_id, /*recovery_mode=*/false,
+                  /*put_document_stats=*/nullptr),
+              IsOk());
 
   // Verify the state of qualified_id_join_index_ after Handle().
   EXPECT_THAT(qualified_id_join_index_->last_added_document_id(),
               Eq(child_put_result.new_document_id));
   EXPECT_THAT(qualified_id_join_index_, Pointee(SizeIs(1)));
   EXPECT_THAT(
-      qualified_id_join_index_->Get(parent_put_result.new_document_id),
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(
+          parent_put_result.new_document_id),
       IsOkAndHolds(ElementsAre(DocumentJoinIdPair(
           child_put_result.new_document_id, fake_type_joinable_property_id_))));
 }
 
-TEST_F(QualifiedIdJoinIndexingHandlerV3Test, HandleNestedJoinableProperty) {
+TEST_P(QualifiedIdJoinIndexingHandlerV3Test, HandleNestedJoinableProperty) {
   // Create and put parent document1. Get its document id and namespace id.
   DocumentProto parent_document1 =
       DocumentBuilder()
@@ -295,8 +339,15 @@
           .SetSchema(std::string(kReferencedType))
           .AddStringProperty(std::string(kPropertyName), "one")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult parent_put_result1,
-                             doc_store_->Put(parent_document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument parent_tokenized_document1,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(parent_document1)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult parent_put_result1,
+      doc_store_->Put(parent_tokenized_document1.document_wrapper()));
 
   // Create and put parent document2.
   DocumentProto parent_document2 =
@@ -305,8 +356,15 @@
           .SetSchema(std::string(kReferencedType))
           .AddStringProperty(std::string(kPropertyName), "two")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult parent_put_result2,
-                             doc_store_->Put(parent_document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument parent_tokenized_document2,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(parent_document2)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult parent_put_result2,
+      doc_store_->Put(parent_tokenized_document2.document_wrapper()));
 
   // Create and put child document:
   // - kPropertyNestedDoc.kPropertyQualifiedId refers to parent_document2.
@@ -328,12 +386,15 @@
           .AddStringProperty(std::string(kPropertyQualifiedId2),
                              "pkg$db/ns#ref_type/1")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult child_put_result,
-                             doc_store_->Put(child_document));
   ICING_ASSERT_OK_AND_ASSIGN(
-      TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                child_document));
+      TokenizedDocument child_tokenized_document,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(child_document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult child_put_result,
+      doc_store_->Put(child_tokenized_document.document_wrapper()));
 
   // Handle nested_document.
   ASSERT_THAT(qualified_id_join_index_->last_added_document_id(),
@@ -341,28 +402,31 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<QualifiedIdJoinIndexingHandler> handler,
       QualifiedIdJoinIndexingHandler::Create(&fake_clock_, doc_store_.get(),
-                                             qualified_id_join_index_.get()));
-  EXPECT_THAT(
-      handler->Handle(tokenized_document, child_put_result.new_document_id,
-                      child_put_result.old_document_id, /*recovery_mode=*/false,
-                      /*put_document_stats=*/nullptr),
-      IsOk());
+                                             qualified_id_join_index_.get(),
+                                             feature_flags_.get()));
+  EXPECT_THAT(handler->Handle(
+                  child_tokenized_document, child_put_result.new_document_id,
+                  child_put_result.old_document_id, /*recovery_mode=*/false,
+                  /*put_document_stats=*/nullptr),
+              IsOk());
 
   // Verify the state of qualified_id_join_index_ after Handle().
   EXPECT_THAT(qualified_id_join_index_->last_added_document_id(),
               Eq(child_put_result.new_document_id));
   EXPECT_THAT(qualified_id_join_index_, Pointee(SizeIs(2)));
-  EXPECT_THAT(qualified_id_join_index_->Get(parent_put_result1.new_document_id),
+  EXPECT_THAT(qualified_id_join_index_->GetDocumentJoinIdPairArrayView(
+                  parent_put_result1.new_document_id),
               IsOkAndHolds(ElementsAre(
                   DocumentJoinIdPair(child_put_result.new_document_id,
                                      nested_type_joinable_property_id_))));
-  EXPECT_THAT(qualified_id_join_index_->Get(parent_put_result2.new_document_id),
+  EXPECT_THAT(qualified_id_join_index_->GetDocumentJoinIdPairArrayView(
+                  parent_put_result2.new_document_id),
               IsOkAndHolds(ElementsAre(DocumentJoinIdPair(
                   child_put_result.new_document_id,
                   nested_type_nested_joinable_property_id_))));
 }
 
-TEST_F(QualifiedIdJoinIndexingHandlerV3Test,
+TEST_P(QualifiedIdJoinIndexingHandlerV3Test,
        HandleShouldSkipInvalidFormatQualifiedId) {
   static constexpr std::string_view kInvalidFormatQualifiedId =
       "invalid_format_qualified_id";
@@ -378,12 +442,15 @@
           .AddStringProperty(std::string(kPropertyQualifiedId),
                              std::string(kInvalidFormatQualifiedId))
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult child_put_result,
-                             doc_store_->Put(document));
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult child_put_result,
+      doc_store_->Put(tokenized_document.document_wrapper()));
 
   // Handle document. Should ignore invalid format qualified id.
   ASSERT_THAT(qualified_id_join_index_->last_added_document_id(),
@@ -391,7 +458,8 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<QualifiedIdJoinIndexingHandler> handler,
       QualifiedIdJoinIndexingHandler::Create(&fake_clock_, doc_store_.get(),
-                                             qualified_id_join_index_.get()));
+                                             qualified_id_join_index_.get(),
+                                             feature_flags_.get()));
   EXPECT_THAT(
       handler->Handle(tokenized_document, child_put_result.new_document_id,
                       child_put_result.old_document_id, /*recovery_mode=*/false,
@@ -406,19 +474,22 @@
   EXPECT_THAT(qualified_id_join_index_, Pointee(IsEmpty()));
 }
 
-TEST_F(QualifiedIdJoinIndexingHandlerV3Test, HandleShouldSkipEmptyQualifiedId) {
+TEST_P(QualifiedIdJoinIndexingHandlerV3Test, HandleShouldSkipEmptyQualifiedId) {
   // Create and put child document without any qualified id. Also tokenize it.
   DocumentProto document = DocumentBuilder()
                                .SetKey("icing", "fake_type/1")
                                .SetSchema(std::string(kFakeType))
                                .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult child_put_result,
-                             doc_store_->Put(document));
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(document)));
   ASSERT_THAT(tokenized_document.qualified_id_join_properties(), IsEmpty());
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult child_put_result,
+      doc_store_->Put(tokenized_document.document_wrapper()));
 
   // Handle document.
   ASSERT_THAT(qualified_id_join_index_->last_added_document_id(),
@@ -426,7 +497,8 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<QualifiedIdJoinIndexingHandler> handler,
       QualifiedIdJoinIndexingHandler::Create(&fake_clock_, doc_store_.get(),
-                                             qualified_id_join_index_.get()));
+                                             qualified_id_join_index_.get(),
+                                             feature_flags_.get()));
   EXPECT_THAT(
       handler->Handle(tokenized_document, child_put_result.new_document_id,
                       child_put_result.old_document_id,
@@ -441,7 +513,7 @@
   EXPECT_THAT(qualified_id_join_index_, Pointee(IsEmpty()));
 }
 
-TEST_F(QualifiedIdJoinIndexingHandlerV3Test, HandleShouldMigrateParent) {
+TEST_P(QualifiedIdJoinIndexingHandlerV3Test, HandleShouldMigrateParent) {
   // Create and put parent document.
   DocumentProto parent_document =
       DocumentBuilder()
@@ -449,8 +521,15 @@
           .SetSchema(std::string(kReferencedType))
           .AddStringProperty(std::string(kPropertyName), "one")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult parent_put_result,
-                             doc_store_->Put(parent_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument parent_tokenized_document,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(parent_document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult parent_put_result,
+      doc_store_->Put(parent_tokenized_document.document_wrapper()));
 
   // Create and put child and grandchild document with relations:
   // parent_document <- child_document <- grandchild_document
@@ -469,33 +548,43 @@
           .AddStringProperty(std::string(kPropertyQualifiedId),
                              "icing#fake_type/1")
           .Build();
+  DocumentProto child_document_replaced = child_document;
 
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<QualifiedIdJoinIndexingHandler> handler,
       QualifiedIdJoinIndexingHandler::Create(&fake_clock_, doc_store_.get(),
-                                             qualified_id_join_index_.get()));
+                                             qualified_id_join_index_.get(),
+                                             feature_flags_.get()));
 
   // Put and index child document.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult child_put_result,
-                             doc_store_->Put(child_document));
-  ASSERT_THAT(child_put_result.old_document_id, Eq(kInvalidDocumentId));
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument child_tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                child_document));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(child_document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult child_put_result,
+      doc_store_->Put(child_tokenized_document.document_wrapper()));
+  ASSERT_THAT(child_put_result.old_document_id, Eq(kInvalidDocumentId));
+
   ICING_ASSERT_OK(handler->Handle(
       child_tokenized_document, child_put_result.new_document_id,
       child_put_result.old_document_id,
       /*recovery_mode=*/false, /*put_document_stats=*/nullptr));
 
   // Put and index grandchild document.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult grandchild_put_result,
-                             doc_store_->Put(grandchild_document));
-  ASSERT_THAT(grandchild_put_result.old_document_id, Eq(kInvalidDocumentId));
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument grandchild_tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(grandchild_document)));
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(grandchild_document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult grandchild_put_result,
+      doc_store_->Put(grandchild_tokenized_document.document_wrapper()));
+  ASSERT_THAT(grandchild_put_result.old_document_id, Eq(kInvalidDocumentId));
+
   ICING_ASSERT_OK(handler->Handle(
       grandchild_tokenized_document, grandchild_put_result.new_document_id,
       grandchild_put_result.old_document_id,
@@ -507,46 +596,59 @@
               Eq(grandchild_put_result.new_document_id));
   ASSERT_THAT(qualified_id_join_index_, Pointee(SizeIs(2)));
   ASSERT_THAT(
-      qualified_id_join_index_->Get(parent_put_result.new_document_id),
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(
+          parent_put_result.new_document_id),
       IsOkAndHolds(ElementsAre(DocumentJoinIdPair(
           child_put_result.new_document_id, fake_type_joinable_property_id_))));
-  ASSERT_THAT(qualified_id_join_index_->Get(child_put_result.new_document_id),
+  ASSERT_THAT(qualified_id_join_index_->GetDocumentJoinIdPairArrayView(
+                  child_put_result.new_document_id),
               IsOkAndHolds(ElementsAre(
                   DocumentJoinIdPair(grandchild_put_result.new_document_id,
                                      fake_type_joinable_property_id_))));
 
   // Update the child document and index it again.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult child_put_result2,
-                             doc_store_->Put(child_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument child_tokenized_document_replaced,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(child_document_replaced)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult child_put_result2,
+      doc_store_->Put(child_tokenized_document_replaced.document_wrapper()));
   ASSERT_THAT(child_put_result2.old_document_id,
               Eq(child_put_result.new_document_id));
 
   // Handle should migrate.
-  EXPECT_THAT(handler->Handle(
-                  child_tokenized_document, child_put_result2.new_document_id,
-                  child_put_result2.old_document_id, /*recovery_mode=*/false,
-                  /*put_document_stats=*/nullptr),
-              IsOk());
+  EXPECT_THAT(
+      handler->Handle(child_tokenized_document_replaced,
+                      child_put_result2.new_document_id,
+                      child_put_result2.old_document_id,
+                      /*recovery_mode=*/false, /*put_document_stats=*/nullptr),
+      IsOk());
   EXPECT_THAT(qualified_id_join_index_, Pointee(SizeIs(3)));
   // Get() with parent document id should return DocumentJoinIdPairs for both
   // old and new child document id.
-  EXPECT_THAT(qualified_id_join_index_->Get(parent_put_result.new_document_id),
+  EXPECT_THAT(qualified_id_join_index_->GetDocumentJoinIdPairArrayView(
+                  parent_put_result.new_document_id),
               IsOkAndHolds(ElementsAre(
                   DocumentJoinIdPair(child_put_result.new_document_id,
                                      fake_type_joinable_property_id_),
                   DocumentJoinIdPair(child_put_result2.new_document_id,
                                      fake_type_joinable_property_id_))));
   // Get() with old child document id should return empty list.
-  EXPECT_THAT(qualified_id_join_index_->Get(child_put_result.new_document_id),
+  EXPECT_THAT(qualified_id_join_index_->GetDocumentJoinIdPairArrayView(
+                  child_put_result.new_document_id),
               IsOkAndHolds(IsEmpty()));
   // Get() with new child document id should return grandchild join id pair.
-  EXPECT_THAT(qualified_id_join_index_->Get(child_put_result2.new_document_id),
+  EXPECT_THAT(qualified_id_join_index_->GetDocumentJoinIdPairArrayView(
+                  child_put_result2.new_document_id),
               IsOkAndHolds(ElementsAre(
                   DocumentJoinIdPair(grandchild_put_result.new_document_id,
                                      fake_type_joinable_property_id_))));
 }
 
-TEST_F(QualifiedIdJoinIndexingHandlerV3Test,
+TEST_P(QualifiedIdJoinIndexingHandlerV3Test,
        HandleInvalidNewDocumentIdShouldReturnInvalidArgumentError) {
   // Create and put parent document.
   DocumentProto parent_document =
@@ -555,23 +657,33 @@
           .SetSchema(std::string(kReferencedType))
           .AddStringProperty(std::string(kPropertyName), "one")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult parent_put_result,
-                             doc_store_->Put(parent_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument parent_tokenized_document,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(parent_document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult parent_put_result,
+      doc_store_->Put(parent_tokenized_document.document_wrapper()));
 
   // Create and put child document. Also tokenize it.
-  DocumentProto document =
+  DocumentProto child_document =
       DocumentBuilder()
           .SetKey("icing", "fake_type/1")
           .SetSchema(std::string(kFakeType))
           .AddStringProperty(std::string(kPropertyQualifiedId),
                              "pkg$db/ns#ref_type/1")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult child_put_result,
-                             doc_store_->Put(document));
   ICING_ASSERT_OK_AND_ASSIGN(
-      TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document)));
+      TokenizedDocument child_tokenized_document,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(child_document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult child_put_result,
+      doc_store_->Put(child_tokenized_document.document_wrapper()));
 
   qualified_id_join_index_->set_last_added_document_id(
       parent_put_result.new_document_id);
@@ -581,11 +693,12 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<QualifiedIdJoinIndexingHandler> handler,
       QualifiedIdJoinIndexingHandler::Create(&fake_clock_, doc_store_.get(),
-                                             qualified_id_join_index_.get()));
+                                             qualified_id_join_index_.get(),
+                                             feature_flags_.get()));
 
   // Handling document with kInvalidDocumentId should cause a failure.
   EXPECT_THAT(
-      handler->Handle(tokenized_document, kInvalidDocumentId,
+      handler->Handle(child_tokenized_document, kInvalidDocumentId,
                       child_put_result.old_document_id, /*recovery_mode=*/false,
                       /*put_document_stats=*/nullptr),
       StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
@@ -597,18 +710,19 @@
 
   // Recovery mode should get the same result.
   EXPECT_THAT(
-      handler->Handle(tokenized_document, kInvalidDocumentId,
+      handler->Handle(child_tokenized_document, kInvalidDocumentId,
                       child_put_result.old_document_id, /*recovery_mode=*/true,
                       /*put_document_stats=*/nullptr),
       StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
   EXPECT_THAT(qualified_id_join_index_->last_added_document_id(),
               Eq(parent_put_result.new_document_id));
   EXPECT_THAT(qualified_id_join_index_, Pointee(IsEmpty()));
-  EXPECT_THAT(qualified_id_join_index_->Get(parent_put_result.new_document_id),
+  EXPECT_THAT(qualified_id_join_index_->GetDocumentJoinIdPairArrayView(
+                  parent_put_result.new_document_id),
               IsOkAndHolds(IsEmpty()));
 }
 
-TEST_F(QualifiedIdJoinIndexingHandlerV3Test,
+TEST_P(QualifiedIdJoinIndexingHandlerV3Test,
        HandleOutOfOrderDocumentIdShouldReturnInvalidArgumentError) {
   // Create and put parent document.
   DocumentProto parent_document =
@@ -617,28 +731,39 @@
           .SetSchema(std::string(kReferencedType))
           .AddStringProperty(std::string(kPropertyName), "one")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult parent_put_result,
-                             doc_store_->Put(parent_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument parent_tokenized_document,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(parent_document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult parent_put_result,
+      doc_store_->Put(parent_tokenized_document.document_wrapper()));
 
   // Create and put child document. Also tokenize it.
-  DocumentProto document =
+  DocumentProto child_document =
       DocumentBuilder()
           .SetKey("icing", "fake_type/1")
           .SetSchema(std::string(kFakeType))
           .AddStringProperty(std::string(kPropertyQualifiedId),
                              "pkg$db/ns#ref_type/1")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult child_put_result,
-                             doc_store_->Put(document));
   ICING_ASSERT_OK_AND_ASSIGN(
-      TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document)));
+      TokenizedDocument child_tokenized_document,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(child_document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult child_put_result,
+      doc_store_->Put(child_tokenized_document.document_wrapper()));
 
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<QualifiedIdJoinIndexingHandler> handler,
       QualifiedIdJoinIndexingHandler::Create(&fake_clock_, doc_store_.get(),
-                                             qualified_id_join_index_.get()));
+                                             qualified_id_join_index_.get(),
+                                             feature_flags_.get()));
 
   // Handling document with document_id == last_added_document_id should cause a
   // failure.
@@ -646,17 +771,18 @@
       child_put_result.new_document_id);
   ASSERT_THAT(qualified_id_join_index_->last_added_document_id(),
               Eq(child_put_result.new_document_id));
-  EXPECT_THAT(
-      handler->Handle(tokenized_document, child_put_result.new_document_id,
-                      child_put_result.old_document_id, /*recovery_mode=*/false,
-                      /*put_document_stats=*/nullptr),
-      StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
+  EXPECT_THAT(handler->Handle(
+                  child_tokenized_document, child_put_result.new_document_id,
+                  child_put_result.old_document_id, /*recovery_mode=*/false,
+                  /*put_document_stats=*/nullptr),
+              StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
   // Verify the state of qualified_id_join_index_ after Handle(). Both index
   // data and last_added_document_id should remain unchanged.
   EXPECT_THAT(qualified_id_join_index_->last_added_document_id(),
               Eq(child_put_result.new_document_id));
   EXPECT_THAT(qualified_id_join_index_, Pointee(IsEmpty()));
-  EXPECT_THAT(qualified_id_join_index_->Get(parent_put_result.new_document_id),
+  EXPECT_THAT(qualified_id_join_index_->GetDocumentJoinIdPairArrayView(
+                  parent_put_result.new_document_id),
               IsOkAndHolds(IsEmpty()));
 
   // Handling document with document_id < last_added_document_id should cause a
@@ -665,21 +791,22 @@
       child_put_result.new_document_id + 1);
   ASSERT_THAT(qualified_id_join_index_->last_added_document_id(),
               Eq(child_put_result.new_document_id + 1));
-  EXPECT_THAT(
-      handler->Handle(tokenized_document, child_put_result.new_document_id,
-                      child_put_result.old_document_id, /*recovery_mode=*/false,
-                      /*put_document_stats=*/nullptr),
-      StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
+  EXPECT_THAT(handler->Handle(
+                  child_tokenized_document, child_put_result.new_document_id,
+                  child_put_result.old_document_id, /*recovery_mode=*/false,
+                  /*put_document_stats=*/nullptr),
+              StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
   // Verify the state of qualified_id_join_index_ after Handle(). Both index
   // data and last_added_document_id should remain unchanged.
   EXPECT_THAT(qualified_id_join_index_->last_added_document_id(),
               Eq(child_put_result.new_document_id + 1));
   EXPECT_THAT(qualified_id_join_index_, Pointee(IsEmpty()));
-  EXPECT_THAT(qualified_id_join_index_->Get(parent_put_result.new_document_id),
+  EXPECT_THAT(qualified_id_join_index_->GetDocumentJoinIdPairArrayView(
+                  parent_put_result.new_document_id),
               IsOkAndHolds(IsEmpty()));
 }
 
-TEST_F(QualifiedIdJoinIndexingHandlerV3Test,
+TEST_P(QualifiedIdJoinIndexingHandlerV3Test,
        HandleRecoveryModeShouldIndexDocsGtLastAddedDocId) {
   // Create and put parent document.
   DocumentProto parent_document =
@@ -688,28 +815,39 @@
           .SetSchema(std::string(kReferencedType))
           .AddStringProperty(std::string(kPropertyName), "one")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult parent_put_result,
-                             doc_store_->Put(parent_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument parent_tokenized_document,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(parent_document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult parent_put_result,
+      doc_store_->Put(parent_tokenized_document.document_wrapper()));
 
   // Create and put child document. Also tokenize it.
-  DocumentProto document =
+  DocumentProto child_document =
       DocumentBuilder()
           .SetKey("icing", "fake_type/1")
           .SetSchema(std::string(kFakeType))
           .AddStringProperty(std::string(kPropertyQualifiedId),
                              "pkg$db/ns#ref_type/1")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult child_put_result,
-                             doc_store_->Put(document));
   ICING_ASSERT_OK_AND_ASSIGN(
-      TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document)));
+      TokenizedDocument child_tokenized_document,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(child_document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult child_put_result,
+      doc_store_->Put(child_tokenized_document.document_wrapper()));
 
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<QualifiedIdJoinIndexingHandler> handler,
       QualifiedIdJoinIndexingHandler::Create(&fake_clock_, doc_store_.get(),
-                                             qualified_id_join_index_.get()));
+                                             qualified_id_join_index_.get(),
+                                             feature_flags_.get()));
 
   // Handle document with document_id > last_added_document_id in recovery mode.
   // The handler should index this document and update last_added_document_id.
@@ -717,21 +855,22 @@
       child_put_result.new_document_id - 1);
   ASSERT_THAT(qualified_id_join_index_->last_added_document_id(),
               Eq(child_put_result.new_document_id - 1));
-  EXPECT_THAT(
-      handler->Handle(tokenized_document, child_put_result.new_document_id,
-                      child_put_result.old_document_id, /*recovery_mode=*/true,
-                      /*put_document_stats=*/nullptr),
-      IsOk());
+  EXPECT_THAT(handler->Handle(
+                  child_tokenized_document, child_put_result.new_document_id,
+                  child_put_result.old_document_id, /*recovery_mode=*/true,
+                  /*put_document_stats=*/nullptr),
+              IsOk());
   EXPECT_THAT(qualified_id_join_index_->last_added_document_id(),
               Eq(child_put_result.new_document_id));
   EXPECT_THAT(qualified_id_join_index_, Pointee(SizeIs(1)));
   EXPECT_THAT(
-      qualified_id_join_index_->Get(parent_put_result.new_document_id),
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(
+          parent_put_result.new_document_id),
       IsOkAndHolds(ElementsAre(DocumentJoinIdPair(
           child_put_result.new_document_id, fake_type_joinable_property_id_))));
 }
 
-TEST_F(QualifiedIdJoinIndexingHandlerV3Test,
+TEST_P(QualifiedIdJoinIndexingHandlerV3Test,
        HandleRecoveryModeShouldIgnoreDocsLeLastAddedDocId) {
   // Create and put parent document.
   DocumentProto parent_document =
@@ -740,28 +879,39 @@
           .SetSchema(std::string(kReferencedType))
           .AddStringProperty(std::string(kPropertyName), "one")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult parent_put_result,
-                             doc_store_->Put(parent_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument parent_tokenized_document,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(parent_document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult parent_put_result,
+      doc_store_->Put(parent_tokenized_document.document_wrapper()));
 
   // Create and put child document. Also tokenize it.
-  DocumentProto document =
+  DocumentProto child_document =
       DocumentBuilder()
           .SetKey("icing", "fake_type/1")
           .SetSchema(std::string(kFakeType))
           .AddStringProperty(std::string(kPropertyQualifiedId),
                              "pkg$db/ns#ref_type/1")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult child_put_result,
-                             doc_store_->Put(document));
   ICING_ASSERT_OK_AND_ASSIGN(
-      TokenizedDocument tokenized_document,
-      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                std::move(document)));
+      TokenizedDocument child_tokenized_document,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(child_document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult child_put_result,
+      doc_store_->Put(child_tokenized_document.document_wrapper()));
 
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<QualifiedIdJoinIndexingHandler> handler,
       QualifiedIdJoinIndexingHandler::Create(&fake_clock_, doc_store_.get(),
-                                             qualified_id_join_index_.get()));
+                                             qualified_id_join_index_.get(),
+                                             feature_flags_.get()));
 
   // Handle document with document_id == last_added_document_id in recovery
   // mode. We should not get any error, but the handler should ignore the
@@ -771,15 +921,16 @@
       child_put_result.new_document_id);
   ASSERT_THAT(qualified_id_join_index_->last_added_document_id(),
               Eq(child_put_result.new_document_id));
-  EXPECT_THAT(
-      handler->Handle(tokenized_document, child_put_result.new_document_id,
-                      child_put_result.old_document_id, /*recovery_mode=*/true,
-                      /*put_document_stats=*/nullptr),
-      IsOk());
+  EXPECT_THAT(handler->Handle(
+                  child_tokenized_document, child_put_result.new_document_id,
+                  child_put_result.old_document_id, /*recovery_mode=*/true,
+                  /*put_document_stats=*/nullptr),
+              IsOk());
   EXPECT_THAT(qualified_id_join_index_->last_added_document_id(),
               Eq(child_put_result.new_document_id));
   EXPECT_THAT(qualified_id_join_index_, Pointee(IsEmpty()));
-  EXPECT_THAT(qualified_id_join_index_->Get(parent_put_result.new_document_id),
+  EXPECT_THAT(qualified_id_join_index_->GetDocumentJoinIdPairArrayView(
+                  parent_put_result.new_document_id),
               IsOkAndHolds(IsEmpty()));
 
   // Handle document with document_id < last_added_document_id in recovery mode.
@@ -789,18 +940,427 @@
       child_put_result.new_document_id + 1);
   ASSERT_THAT(qualified_id_join_index_->last_added_document_id(),
               Eq(child_put_result.new_document_id + 1));
-  EXPECT_THAT(
-      handler->Handle(tokenized_document, child_put_result.new_document_id,
-                      child_put_result.old_document_id, /*recovery_mode=*/true,
-                      /*put_document_stats=*/nullptr),
-      IsOk());
+  EXPECT_THAT(handler->Handle(
+                  child_tokenized_document, child_put_result.new_document_id,
+                  child_put_result.old_document_id, /*recovery_mode=*/true,
+                  /*put_document_stats=*/nullptr),
+              IsOk());
   EXPECT_THAT(qualified_id_join_index_->last_added_document_id(),
               Eq(child_put_result.new_document_id + 1));
   EXPECT_THAT(qualified_id_join_index_, Pointee(IsEmpty()));
-  EXPECT_THAT(qualified_id_join_index_->Get(parent_put_result.new_document_id),
+  EXPECT_THAT(qualified_id_join_index_->GetDocumentJoinIdPairArrayView(
+                  parent_put_result.new_document_id),
               IsOkAndHolds(IsEmpty()));
 }
 
+TEST_P(QualifiedIdJoinIndexingHandlerV3Test, HandleJoinWithMissingParent) {
+  if (!feature_flags_->enable_non_existent_qualified_id_join()) {
+    return;
+  }
+
+  // Create and index a child document that refers to a non-existent parent.
+  DocumentProto child_document =
+      DocumentBuilder()
+          .SetKey("icing", "fake_type/1")
+          .SetSchema(std::string(kFakeType))
+          .AddStringProperty(std::string(kPropertyQualifiedId),
+                             "pkg$db/ns#ref_type/1")  // This parent doesn't
+                                                      // exist yet
+          .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument child_tokenized_document,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(child_document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult child_put_result,
+      doc_store_->Put(child_tokenized_document.document_wrapper()));
+
+  // Handle the child document.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<QualifiedIdJoinIndexingHandler> handler,
+      QualifiedIdJoinIndexingHandler::Create(&fake_clock_, doc_store_.get(),
+                                             qualified_id_join_index_.get(),
+                                             feature_flags_.get()));
+  EXPECT_THAT(handler->Handle(
+                  child_tokenized_document, child_put_result.new_document_id,
+                  child_put_result.old_document_id, /*recovery_mode=*/false,
+                  /*put_document_stats=*/nullptr),
+              IsOk());
+
+  // At this point, there's no parent document, so no entry in the doc id map.
+  // The qualified id map should have an entry, and the index size should be 1.
+  EXPECT_THAT(qualified_id_join_index_, Pointee(SizeIs(1)));
+
+  // Now, create and index the parent document.
+  DocumentProto parent_document =
+      DocumentBuilder()
+          .SetKey("pkg$db/ns", "ref_type/1")
+          .SetSchema(std::string(kReferencedType))
+          .AddStringProperty(std::string(kPropertyName), "one")
+          .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument parent_tokenized_document,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(parent_document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult parent_put_result,
+      doc_store_->Put(parent_tokenized_document.document_wrapper()));
+
+  // Handle the parent document. This should trigger migration.
+  EXPECT_THAT(handler->Handle(
+                  parent_tokenized_document, parent_put_result.new_document_id,
+                  parent_put_result.old_document_id, /*recovery_mode=*/false,
+                  /*put_document_stats=*/nullptr),
+              IsOk());
+
+  // Verify that the join is now established.
+  EXPECT_THAT(
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(
+          parent_put_result.new_document_id),
+      IsOkAndHolds(ElementsAre(DocumentJoinIdPair(
+          child_put_result.new_document_id, fake_type_joinable_property_id_))));
+
+  // The index size should still be 1.
+  EXPECT_THAT(qualified_id_join_index_, Pointee(SizeIs(1)));
+
+  // Update the parent document and index it again.
+  DocumentProto parent_document_replaced =
+      DocumentBuilder()
+          .SetKey("pkg$db/ns", "ref_type/1")
+          .SetSchema(std::string(kReferencedType))
+          .AddStringProperty(std::string(kPropertyName), "one_replaced")
+          .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument parent_tokenized_document_replaced,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(parent_document_replaced)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult parent_put_result2,
+      doc_store_->Put(parent_tokenized_document_replaced.document_wrapper()));
+  ASSERT_THAT(parent_put_result2.old_document_id,
+              Eq(parent_put_result.new_document_id));
+
+  // Handle the updated parent.
+  EXPECT_THAT(
+      handler->Handle(parent_tokenized_document_replaced,
+                      parent_put_result2.new_document_id,
+                      parent_put_result2.old_document_id,
+                      /*recovery_mode=*/false, /*put_document_stats=*/nullptr),
+      IsOk());
+
+  // Verify migration.
+  // Get() with old parent document id should return empty list.
+  EXPECT_THAT(qualified_id_join_index_->GetDocumentJoinIdPairArrayView(
+                  parent_put_result.new_document_id),
+              IsOkAndHolds(IsEmpty()));
+  // Get() with new parent document id should return child join id pair.
+  EXPECT_THAT(
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(
+          parent_put_result2.new_document_id),
+      IsOkAndHolds(ElementsAre(DocumentJoinIdPair(
+          child_put_result.new_document_id, fake_type_joinable_property_id_))));
+  // Index size is still 1.
+  EXPECT_THAT(qualified_id_join_index_, Pointee(SizeIs(1)));
+}
+
+TEST_P(QualifiedIdJoinIndexingHandlerV3Test,
+       HandleJoinWithMultipleMissingParents) {
+  if (!feature_flags_->enable_non_existent_qualified_id_join()) {
+    return;
+  }
+
+  // Create and index a child document that refers to two non-existent parents.
+  DocumentProto child_document =
+      DocumentBuilder()
+          .SetKey("icing", "fake_type/1")
+          .SetSchema(std::string(kFakeType))
+          .AddStringProperty(std::string(kPropertyQualifiedId),
+                             "pkg$db/ns#ref_type/1", "pkg$db/ns#ref_type/2")
+          .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument child_tokenized_document,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(child_document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult child_put_result,
+      doc_store_->Put(child_tokenized_document.document_wrapper()));
+
+  // Handle the child document.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<QualifiedIdJoinIndexingHandler> handler,
+      QualifiedIdJoinIndexingHandler::Create(&fake_clock_, doc_store_.get(),
+                                             qualified_id_join_index_.get(),
+                                             feature_flags_.get()));
+  EXPECT_THAT(handler->Handle(
+                  child_tokenized_document, child_put_result.new_document_id,
+                  child_put_result.old_document_id, /*recovery_mode=*/false,
+                  /*put_document_stats=*/nullptr),
+              IsOk());
+
+  // The qualified id map should have two entries, so the index size should be
+  // 2.
+  EXPECT_THAT(qualified_id_join_index_, Pointee(SizeIs(2)));
+
+  // Now, create and index the parent documents.
+  DocumentProto parent_document1 =
+      DocumentBuilder()
+          .SetKey("pkg$db/ns", "ref_type/1")
+          .SetSchema(std::string(kReferencedType))
+          .AddStringProperty(std::string(kPropertyName), "one")
+          .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument parent_tokenized_document1,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(parent_document1)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult parent_put_result1,
+      doc_store_->Put(parent_tokenized_document1.document_wrapper()));
+  DocumentProto parent_document2 =
+      DocumentBuilder()
+          .SetKey("pkg$db/ns", "ref_type/2")
+          .SetSchema(std::string(kReferencedType))
+          .AddStringProperty(std::string(kPropertyName), "two")
+          .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument parent_tokenized_document2,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(parent_document2)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult parent_put_result2,
+      doc_store_->Put(parent_tokenized_document2.document_wrapper()));
+
+  // Handle the parent documents.
+  EXPECT_THAT(handler->Handle(parent_tokenized_document1,
+                              parent_put_result1.new_document_id,
+                              parent_put_result1.old_document_id,
+                              /*recovery_mode=*/false,
+                              /*put_document_stats=*/nullptr),
+              IsOk());
+  EXPECT_THAT(handler->Handle(parent_tokenized_document2,
+                              parent_put_result2.new_document_id,
+                              parent_put_result2.old_document_id,
+                              /*recovery_mode=*/false,
+                              /*put_document_stats=*/nullptr),
+              IsOk());
+
+  // Verify that the joins are now established for both parents.
+  EXPECT_THAT(
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(
+          parent_put_result1.new_document_id),
+      IsOkAndHolds(ElementsAre(DocumentJoinIdPair(
+          child_put_result.new_document_id, fake_type_joinable_property_id_))));
+  EXPECT_THAT(
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(
+          parent_put_result2.new_document_id),
+      IsOkAndHolds(ElementsAre(DocumentJoinIdPair(
+          child_put_result.new_document_id, fake_type_joinable_property_id_))));
+
+  // Index size should still be 2.
+  EXPECT_THAT(qualified_id_join_index_, Pointee(SizeIs(2)));
+}
+
+TEST_P(QualifiedIdJoinIndexingHandlerV3Test, HandleJoinWithMixedParents) {
+  if (!feature_flags_->enable_non_existent_qualified_id_join()) {
+    return;
+  }
+
+  // Create and index one parent document.
+  DocumentProto existing_parent =
+      DocumentBuilder()
+          .SetKey("pkg$db/ns", "ref_type/1")
+          .SetSchema(std::string(kReferencedType))
+          .AddStringProperty(std::string(kPropertyName), "one")
+          .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument existing_parent_tokenized,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(existing_parent)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult existing_parent_put_result,
+      doc_store_->Put(existing_parent_tokenized.document_wrapper()));
+
+  // Create the handler.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<QualifiedIdJoinIndexingHandler> handler,
+      QualifiedIdJoinIndexingHandler::Create(&fake_clock_, doc_store_.get(),
+                                             qualified_id_join_index_.get(),
+                                             feature_flags_.get()));
+
+  // Create and index a child document that refers to the existing parent and
+  // a missing parent.
+  DocumentProto child_document =
+      DocumentBuilder()
+          .SetKey("icing", "fake_type/1")
+          .SetSchema(std::string(kFakeType))
+          .AddStringProperty(std::string(kPropertyQualifiedId),
+                             // existing
+                             "pkg$db/ns#ref_type/1",
+                             // missing
+                             "pkg$db/ns#ref_type/2")
+          .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument child_tokenized_document,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(child_document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult child_put_result,
+      doc_store_->Put(child_tokenized_document.document_wrapper()));
+
+  // Handle the child document.
+  EXPECT_THAT(handler->Handle(
+                  child_tokenized_document, child_put_result.new_document_id,
+                  child_put_result.old_document_id, /*recovery_mode=*/false,
+                  /*put_document_stats=*/nullptr),
+              IsOk());
+
+  // Index size should be 2.
+  EXPECT_THAT(qualified_id_join_index_, Pointee(SizeIs(2)));
+
+  // Verify join for existing parent.
+  EXPECT_THAT(
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(
+          existing_parent_put_result.new_document_id),
+      IsOkAndHolds(ElementsAre(DocumentJoinIdPair(
+          child_put_result.new_document_id, fake_type_joinable_property_id_))));
+
+  // Index the missing parent.
+  DocumentProto missing_parent =
+      DocumentBuilder()
+          .SetKey("pkg$db/ns", "ref_type/2")
+          .SetSchema(std::string(kReferencedType))
+          .AddStringProperty(std::string(kPropertyName), "two")
+          .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument missing_parent_tokenized,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(missing_parent)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult missing_parent_put_result,
+      doc_store_->Put(missing_parent_tokenized.document_wrapper()));
+
+  // Handle the missing parent.
+  EXPECT_THAT(
+      handler->Handle(missing_parent_tokenized,
+                      missing_parent_put_result.new_document_id,
+                      missing_parent_put_result.old_document_id,
+                      /*recovery_mode=*/false, /*put_document_stats=*/nullptr),
+      IsOk());
+
+  // Verify join for the now-existing parent.
+  EXPECT_THAT(
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(
+          missing_parent_put_result.new_document_id),
+      IsOkAndHolds(ElementsAre(DocumentJoinIdPair(
+          child_put_result.new_document_id, fake_type_joinable_property_id_))));
+
+  // Verify join for the existing parent is still there.
+  EXPECT_THAT(
+      qualified_id_join_index_->GetDocumentJoinIdPairArrayView(
+          existing_parent_put_result.new_document_id),
+      IsOkAndHolds(ElementsAre(DocumentJoinIdPair(
+          child_put_result.new_document_id, fake_type_joinable_property_id_))));
+
+  // Index size should still be 2.
+  EXPECT_THAT(qualified_id_join_index_, Pointee(SizeIs(2)));
+}
+
+TEST_P(QualifiedIdJoinIndexingHandlerV3Test,
+       HandleJoinWithMissingParent_FeatureDisabled) {
+  if (feature_flags_->enable_non_existent_qualified_id_join()) {
+    return;
+  }
+
+  // Create and index a child document that refers to a non-existent parent.
+  DocumentProto child_document =
+      DocumentBuilder()
+          .SetKey("icing", "fake_type/1")
+          .SetSchema(std::string(kFakeType))
+          .AddStringProperty(
+              std::string(kPropertyQualifiedId),
+              "pkg$db/ns#ref_type/1")  // This parent doesn't exist
+          .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument child_tokenized_document,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(child_document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult child_put_result,
+      doc_store_->Put(child_tokenized_document.document_wrapper()));
+
+  // Create the handler with the feature flag disabled.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<QualifiedIdJoinIndexingHandler> handler,
+      QualifiedIdJoinIndexingHandler::Create(&fake_clock_, doc_store_.get(),
+                                             qualified_id_join_index_.get(),
+                                             feature_flags_.get()));
+
+  // Handle the child document.
+  EXPECT_THAT(handler->Handle(
+                  child_tokenized_document, child_put_result.new_document_id,
+                  child_put_result.old_document_id, /*recovery_mode=*/false,
+                  /*put_document_stats=*/nullptr),
+              IsOk());
+
+  // With the feature disabled, no join data should be indexed for a missing
+  // parent.
+  EXPECT_THAT(qualified_id_join_index_, Pointee(IsEmpty()));
+
+  // Now, create and index the parent document.
+  DocumentProto parent_document =
+      DocumentBuilder()
+          .SetKey("pkg$db/ns", "ref_type/1")
+          .SetSchema(std::string(kReferencedType))
+          .AddStringProperty(std::string(kPropertyName), "one")
+          .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument parent_tokenized_document,
+      TokenizedDocument::Create(
+          schema_store_.get(), lang_segmenter_.get(),
+          /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds(),
+          std::move(parent_document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult parent_put_result,
+      doc_store_->Put(parent_tokenized_document.document_wrapper()));
+
+  // Handle the parent document.
+  EXPECT_THAT(handler->Handle(
+                  parent_tokenized_document, parent_put_result.new_document_id,
+                  parent_put_result.old_document_id, /*recovery_mode=*/false,
+                  /*put_document_stats=*/nullptr),
+              IsOk());
+
+  // Verify that no join has been established.
+  EXPECT_THAT(qualified_id_join_index_->GetDocumentJoinIdPairArrayView(
+                  parent_put_result.new_document_id),
+              IsOkAndHolds(IsEmpty()));
+
+  EXPECT_THAT(qualified_id_join_index_, Pointee(IsEmpty()));
+}
+
+INSTANTIATE_TEST_SUITE_P(QualifiedIdJoinIndexingHandlerV3Test,
+                         QualifiedIdJoinIndexingHandlerV3Test,
+                         // Parameter: enable_non_existent_qualified_id_join
+                         testing::Values(true, false));
 }  // namespace
 
 }  // namespace lib

diff --git a/icing/join/qualified-id-join-indexing-handler.cc b/icing/join/qualified-id-join-indexing-handler.cc
index b90227f..4061cae 100644
--- a/icing/join/qualified-id-join-indexing-handler.cc
+++ b/icing/join/qualified-id-join-indexing-handler.cc

@@ -25,6 +25,7 @@
 #include "icing/text_classifier/lib3/utils/base/status.h"
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "icing/absl_ports/canonical_errors.h"
+#include "icing/feature-flags.h"
 #include "icing/join/document-join-id-pair.h"
 #include "icing/join/qualified-id-join-index.h"
 #include "icing/join/qualified-id.h"
@@ -48,14 +49,16 @@
     std::unique_ptr<QualifiedIdJoinIndexingHandler>>
 QualifiedIdJoinIndexingHandler::Create(
     const Clock* clock, const DocumentStore* doc_store,
-    QualifiedIdJoinIndex* qualified_id_join_index) {
+    QualifiedIdJoinIndex* qualified_id_join_index,
+    const FeatureFlags* feature_flags) {
   ICING_RETURN_ERROR_IF_NULL(clock);
   ICING_RETURN_ERROR_IF_NULL(doc_store);
   ICING_RETURN_ERROR_IF_NULL(qualified_id_join_index);
+  ICING_RETURN_ERROR_IF_NULL(feature_flags);
 
   return std::unique_ptr<QualifiedIdJoinIndexingHandler>(
-      new QualifiedIdJoinIndexingHandler(clock, doc_store,
-                                         qualified_id_join_index));
+      new QualifiedIdJoinIndexingHandler(
+          clock, doc_store, qualified_id_join_index, feature_flags));
 }
 
 libtextclassifier3::Status QualifiedIdJoinIndexingHandler::Handle(
@@ -165,6 +168,16 @@
   if (IsDocumentIdValid(old_document_id)) {
     ICING_RETURN_IF_ERROR(
         qualified_id_join_index_.MigrateParent(old_document_id, document_id));
+  } else if (feature_flags_.enable_non_existent_qualified_id_join()) {
+    // This means that the document is new, which could also have child
+    // documents since child documents are allowed to reference non-existing
+    // parent documents. Now, we try to migrate its child documents to the new
+    // document id.
+    const DocumentProto& document =
+        tokenized_document.document_wrapper().document();
+    QualifiedId qualified_id_str(document.namespace_(), document.uri());
+    ICING_RETURN_IF_ERROR(
+        qualified_id_join_index_.MigrateParent(qualified_id_str, document_id));
   }
 
   // (Child perspective)
@@ -178,9 +191,14 @@
     DocumentJoinIdPair child_doc_join_id_pair(
         document_id, qualified_id_property.metadata.id);
 
-    // Extract parent qualified ids and lookup their corresponding document ids.
-    std::vector<DocumentId> parent_doc_ids;
-    parent_doc_ids.reserve(qualified_id_property.values.size());
+    // Extract parent qualified ids.
+    std::vector<DocumentId> parent_doc_ids;  // Deprecated.
+    std::vector<QualifiedId> parent_qualified_ids;
+    if (feature_flags_.enable_non_existent_qualified_id_join()) {
+      parent_qualified_ids.reserve(qualified_id_property.values.size());
+    } else {
+      parent_doc_ids.reserve(qualified_id_property.values.size());
+    }
     for (std::string_view parent_qualified_id_str :
          qualified_id_property.values) {
       libtextclassifier3::StatusOr<QualifiedId> parent_qualified_id_or =
@@ -191,6 +209,10 @@
       }
       QualifiedId parent_qualified_id =
           std::move(parent_qualified_id_or).ValueOrDie();
+      if (feature_flags_.enable_non_existent_qualified_id_join()) {
+        parent_qualified_ids.push_back(std::move(parent_qualified_id));
+        continue;
+      }
 
       // Lookup document store to get the parent document id.
       libtextclassifier3::StatusOr<DocumentId> parent_doc_id_or =
@@ -203,10 +225,17 @@
       }
       parent_doc_ids.push_back(parent_doc_id_or.ValueOrDie());
     }
+    libtextclassifier3::Status status;
+    if (feature_flags_.enable_non_existent_qualified_id_join()) {
+      // Index join relationship by parent qualified ids.
+      status = qualified_id_join_index_.Put(&doc_store_, child_doc_join_id_pair,
+                                            std::move(parent_qualified_ids));
 
-    // Add all parent document ids to the index.
-    libtextclassifier3::Status status = qualified_id_join_index_.Put(
-        child_doc_join_id_pair, std::move(parent_doc_ids));
+    } else {
+      // Index join relationship by parent document ids.
+      status = qualified_id_join_index_.Put(child_doc_join_id_pair,
+                                            std::move(parent_doc_ids));
+    }
     if (!status.ok()) {
       ICING_LOG(WARNING)
           << "Failed to add data into qualified id join index due to: "

diff --git a/icing/join/qualified-id-join-indexing-handler.h b/icing/join/qualified-id-join-indexing-handler.h
index f321f4d..9abf249 100644
--- a/icing/join/qualified-id-join-indexing-handler.h
+++ b/icing/join/qualified-id-join-indexing-handler.h

@@ -19,6 +19,7 @@
 
 #include "icing/text_classifier/lib3/utils/base/status.h"
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
+#include "icing/feature-flags.h"
 #include "icing/index/data-indexing-handler.h"
 #include "icing/join/qualified-id-join-index.h"
 #include "icing/proto/logging.pb.h"
@@ -42,7 +43,8 @@
   static libtextclassifier3::StatusOr<
       std::unique_ptr<QualifiedIdJoinIndexingHandler>>
   Create(const Clock* clock, const DocumentStore* doc_store,
-         QualifiedIdJoinIndex* qualified_id_join_index);
+         QualifiedIdJoinIndex* qualified_id_join_index,
+         const FeatureFlags* feature_flags);
 
   ~QualifiedIdJoinIndexingHandler() override = default;
 
@@ -72,10 +74,12 @@
  private:
   explicit QualifiedIdJoinIndexingHandler(
       const Clock* clock, const DocumentStore* doc_store,
-      QualifiedIdJoinIndex* qualified_id_join_index)
+      QualifiedIdJoinIndex* qualified_id_join_index,
+      const FeatureFlags* feature_flags)
       : DataIndexingHandler(clock),
         doc_store_(*doc_store),
-        qualified_id_join_index_(*qualified_id_join_index) {}
+        qualified_id_join_index_(*qualified_id_join_index),
+        feature_flags_(*feature_flags) {}
 
   // TODO(b/275121148): deprecate v2 after rollout v3.
 
@@ -90,6 +94,7 @@
 
   const DocumentStore& doc_store_;                 // Does not own.
   QualifiedIdJoinIndex& qualified_id_join_index_;  // Does not own.
+  const FeatureFlags& feature_flags_;              // Does not own.
 };
 
 }  // namespace lib

diff --git a/icing/join/qualified-id.cc b/icing/join/qualified-id.cc
index 42e080c..370bbcb 100644
--- a/icing/join/qualified-id.cc
+++ b/icing/join/qualified-id.cc

@@ -14,8 +14,10 @@
 
 #include "icing/join/qualified-id.h"
 
+#include <cstddef>
 #include <string>
 #include <string_view>
+#include <utility>
 
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "icing/absl_ports/canonical_errors.h"
@@ -33,6 +35,19 @@
          c == QualifiedId::kNamespaceUriSeparator;
 }
 
+// Helper function to escape the content.
+std::string Escape(std::string_view content) {
+  std::string escaped_content;
+  escaped_content.reserve(content.length());
+  for (char c : content) {
+    if (IsSpecialCharacter(c)) {
+      escaped_content += QualifiedId::kEscapeChar;
+    }
+    escaped_content += c;
+  }
+  return escaped_content;
+}
+
 // Helper function to verify the format (check the escape format and make sure
 // number of separator '#' is 1) and find the position of the unique separator.
 //
@@ -106,5 +121,9 @@
   return QualifiedId(std::move(name_space), std::move(uri));
 }
 
+std::string QualifiedId::ToString() const {
+  return Escape(name_space_) + kNamespaceUriSeparator + Escape(uri_);
+}
+
 }  // namespace lib
 }  // namespace icing

diff --git a/icing/join/qualified-id.h b/icing/join/qualified-id.h
index eb6606a..d3e8d39 100644
--- a/icing/join/qualified-id.h
+++ b/icing/join/qualified-id.h

@@ -15,8 +15,11 @@
 #ifndef ICING_JOIN_QUALIFIED_ID_H_
 #define ICING_JOIN_QUALIFIED_ID_H_
 
+#include <cstddef>
+#include <functional>
 #include <string>
 #include <string_view>
+#include <utility>
 
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
 
@@ -34,6 +37,13 @@
 // - Raw namespace and uri cannot be empty.
 class QualifiedId {
  public:
+  struct Hasher {
+    std::size_t operator()(const QualifiedId& qualified_id) const {
+      return std::hash<std::string>()(qualified_id.name_space_) ^
+             std::hash<std::string>()(qualified_id.uri_);
+    }
+  };
+
   static constexpr char kEscapeChar = '\\';
   static constexpr char kNamespaceUriSeparator = '#';
 
@@ -51,9 +61,24 @@
   explicit QualifiedId(std::string name_space, std::string uri)
       : name_space_(std::move(name_space)), uri_(std::move(uri)) {}
 
+  bool operator==(const QualifiedId& other) const {
+    return name_space_ == other.name_space_ && uri_ == other.uri_;
+  }
+
+  bool operator<(const QualifiedId& other) const {
+    if (name_space_ != other.name_space_) {
+      return name_space_ < other.name_space_;
+    }
+    return uri_ < other.uri_;
+  }
+
   const std::string& name_space() const { return name_space_; }
   const std::string& uri() const { return uri_; }
 
+  // Unparses the QualifiedId to its string representation:
+  // "<escaped(namespace)>#<escaped(uri)>".
+  std::string ToString() const;
+
  private:
   std::string name_space_;
   std::string uri_;

diff --git a/icing/join/qualified-id_test.cc b/icing/join/qualified-id_test.cc
index 92bf63e..59e9e51 100644
--- a/icing/join/qualified-id_test.cc
+++ b/icing/join/qualified-id_test.cc

@@ -153,6 +153,43 @@
               StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
 }
 
+TEST(QualifiedIdTest, ToStringAndParse) {
+  // "namespace" + "uri" -> "namespace#uri"
+  QualifiedId id1("namespace", "uri");
+  std::string str1 = id1.ToString();
+  EXPECT_THAT(str1, Eq("namespace#uri"));
+  ICING_ASSERT_OK_AND_ASSIGN(QualifiedId parsed_id1, QualifiedId::Parse(str1));
+  EXPECT_THAT(parsed_id1, Eq(id1));
+
+  // "namespace\" + "uri" -> "namespace\\#uri"
+  QualifiedId id2("namespace\\", "uri");
+  std::string str2 = id2.ToString();
+  EXPECT_THAT(str2, Eq("namespace\\\\#uri"));
+  ICING_ASSERT_OK_AND_ASSIGN(QualifiedId parsed_id2, QualifiedId::Parse(str2));
+  EXPECT_THAT(parsed_id2, Eq(id2));
+
+  // "namespace#" + "uri" -> "namespace\##uri"
+  QualifiedId id3("namespace#", "uri");
+  std::string str3 = id3.ToString();
+  EXPECT_THAT(str3, Eq("namespace\\##uri"));
+  ICING_ASSERT_OK_AND_ASSIGN(QualifiedId parsed_id3, QualifiedId::Parse(str3));
+  EXPECT_THAT(parsed_id3, Eq(id3));
+
+  // "namespace" + "#uri" -> "namespace#\#uri"
+  QualifiedId id4("namespace", "#uri");
+  std::string str4 = id4.ToString();
+  EXPECT_THAT(str4, Eq("namespace#\\#uri"));
+  ICING_ASSERT_OK_AND_ASSIGN(QualifiedId parsed_id4, QualifiedId::Parse(str4));
+  EXPECT_THAT(parsed_id4, Eq(id4));
+
+  // "namespace\#" + "#\uri" -> "namespace\\\##\#\\uri"
+  QualifiedId id5("namespace\\#", "#\\uri");
+  std::string str5 = id5.ToString();
+  EXPECT_THAT(str5, Eq("namespace\\\\\\##\\#\\\\uri"));
+  ICING_ASSERT_OK_AND_ASSIGN(QualifiedId parsed_id5, QualifiedId::Parse(str5));
+  EXPECT_THAT(parsed_id5, Eq(id5));
+}
+
 }  // namespace
 
 }  // namespace lib

diff --git a/icing/legacy/index/icing-array-storage.h b/icing/legacy/index/icing-array-storage.h
index 6932f8b..88dc59f 100644
--- a/icing/legacy/index/icing-array-storage.h
+++ b/icing/legacy/index/icing-array-storage.h

@@ -20,10 +20,13 @@
 #ifndef ICING_LEGACY_INDEX_ICING_ARRAY_STORAGE_H_
 #define ICING_LEGACY_INDEX_ICING_ARRAY_STORAGE_H_
 
+#include <cstddef>
 #include <cstdint>
 #include <string>
 #include <vector>
 
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
+#include "icing/absl_ports/canonical_errors.h"
 #include "icing/legacy/index/icing-filesystem.h"
 #include "icing/legacy/index/icing-mmapper.h"
 #include "icing/util/crc32.h"
@@ -71,11 +74,19 @@
   // Make array empty again.
   void Clear();
 
-  // Intent to write memory at (elt_idx, elt_idx + elt_len). Returns
-  // NULL if file cannot be grown to accommodate that offset.
+  // Intent to write memory at (elt_idx, elt_idx + elt_len).
+  //
+  // Returns:
+  //   The pointer to the memory on success.
+  //   RESOURCE_EXHAUSTED if file cannot be grown to accommodate that offset.
   template <class T>
-  T *GetMutableMem(uint32_t elt_idx, uint32_t elt_len) {
-    return static_cast<T *>(GetMutableMemInternal(elt_idx, elt_len));
+  libtextclassifier3::StatusOr<T *> GetMutableMem(uint32_t elt_idx,
+                                                  uint32_t elt_len) {
+    void *mem = GetMutableMemInternal(elt_idx, elt_len);
+    if (mem == nullptr) {
+      return absl_ports::ResourceExhaustedError("Failed to allocate memory");
+    }
+    return static_cast<T *>(mem);
   }
 
   // Resizes to first elt_len elements.
@@ -134,6 +145,7 @@
   static_assert(8 == sizeof(Change), "sizeof(Change) != 8");
   static_assert(4 == alignof(Change), "alignof(Change) != 4");
 
+  // Returns a pointer to the memory. Returns nullptr on failure.
   void *GetMutableMemInternal(uint32_t elt_idx, uint32_t elt_len);
 
   bool GrowIfNecessary(uint32_t num_elts);

diff --git a/icing/legacy/index/icing-array-storage_test.cc b/icing/legacy/index/icing-array-storage_test.cc
index 690c356..37c3f09 100644
--- a/icing/legacy/index/icing-array-storage_test.cc
+++ b/icing/legacy/index/icing-array-storage_test.cc

@@ -20,6 +20,7 @@
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include "icing/legacy/index/icing-filesystem.h"
+#include "icing/testing/common-matchers.h"
 #include "icing/testing/tmp-directory.h"
 #include "icing/util/crc32.h"
 
@@ -61,7 +62,8 @@
   EXPECT_THAT(storage.UpdateCrc(), Eq(Crc32()));
   EXPECT_THAT(storage.GetCrc(), Eq(Crc32()));
 
-  uint32_t* val = storage.GetMutableMem<uint32_t>(0, 1);
+  ICING_ASSERT_OK_AND_ASSIGN(uint32_t* val,
+                             storage.GetMutableMem<uint32_t>(0, 1));
   *val = 5;
 
   // Because there is no crc_ptr, the crc should remain 0 even though we have
@@ -85,7 +87,8 @@
   EXPECT_THAT(storage.UpdateCrc(), Eq(Crc32()));
   EXPECT_THAT(storage.GetCrc(), Eq(Crc32()));
 
-  uint32_t* val = storage.GetMutableMem<uint32_t>(0, 1);
+  ICING_ASSERT_OK_AND_ASSIGN(uint32_t* val,
+                             storage.GetMutableMem<uint32_t>(0, 1));
   *val = 5;
 
   EXPECT_THAT(storage.GetCrc(), Eq(Crc32(937357362)));
@@ -110,7 +113,8 @@
     EXPECT_THAT(storage_one.UpdateCrc(), Eq(Crc32()));
     EXPECT_THAT(storage_one.GetCrc(), Eq(Crc32()));
 
-    uint32_t* val = storage_one.GetMutableMem<uint32_t>(0, 1);
+    ICING_ASSERT_OK_AND_ASSIGN(uint32_t* val,
+                               storage_one.GetMutableMem<uint32_t>(0, 1));
     *val = 5;
 
     EXPECT_THAT(storage_one.GetCrc(), Eq(Crc32(937357362)));
@@ -144,7 +148,8 @@
     EXPECT_THAT(storage_one.UpdateCrc(), Eq(Crc32()));
     EXPECT_THAT(storage_one.GetCrc(), Eq(Crc32()));
 
-    uint32_t* val = storage_one.GetMutableMem<uint32_t>(0, 1);
+    ICING_ASSERT_OK_AND_ASSIGN(uint32_t* val,
+                               storage_one.GetMutableMem<uint32_t>(0, 1));
     *val = 5;
 
     EXPECT_THAT(storage_one.GetCrc(), Eq(Crc32(937357362)));
@@ -167,4 +172,4 @@
 }  // namespace
 
 }  // namespace lib
-}  // namespace icing
\ No newline at end of file
+}  // namespace icing

diff --git a/icing/legacy/index/icing-dynamic-trie.cc b/icing/legacy/index/icing-dynamic-trie.cc
index e665272..f385c94 100644
--- a/icing/legacy/index/icing-dynamic-trie.cc
+++ b/icing/legacy/index/icing-dynamic-trie.cc

@@ -71,6 +71,7 @@
 #include <cerrno>
 #include <cinttypes>
 #include <cstdint>
+#include <cstdlib>
 #include <cstring>
 #include <memory>
 #include <ostream>
@@ -80,6 +81,7 @@
 #include <vector>
 
 #include "icing/text_classifier/lib3/utils/base/status.h"
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "icing/absl_ports/canonical_errors.h"
 #include "icing/legacy/core/icing-string-util.h"
 #include "icing/legacy/core/icing-timer.h"
@@ -87,12 +89,15 @@
 #include "icing/legacy/index/icing-filesystem.h"
 #include "icing/legacy/index/icing-flash-bitmap.h"
 #include "icing/legacy/index/icing-mmapper.h"
+#include "icing/legacy/index/icing-storage.h"
 #include "icing/legacy/index/proto/icing-dynamic-trie-header.pb.h"
 #include "icing/util/crc32.h"
 #include "icing/util/i18n-utils.h"
 #include "icing/util/logging.h"
 #include "icing/util/math-util.h"
 #include "icing/util/status-macros.h"
+#include "unicode/utf8.h"
+#include "unicode/utypes.h"
 
 namespace icing {
 namespace lib {
@@ -285,6 +290,9 @@
     return &array_storage_[NODE].array_cast<Node>()[idx];
   }
 
+  // Returns the number of elements in the node array.
+  uint32_t GetNodeArraySize() const { return array_storage_[NODE].num_elts(); }
+
   // REQUIRES: !empty(). Otherwise node 0 could contain invalid data
   //   (next_index, is_leaf, log2_num_children).
   const Node *GetRootNode() const { return GetNode(0); }
@@ -308,9 +316,11 @@
   // returns a writable element or array within and sets
   // dirty_pages_[array_type] as a side effect, assuming the mutable
   // area will get written to.
-  Node *GetMutableNode(uint32_t idx);
-  Next *GetMutableNextArray(uint32_t idx, uint32_t len);
-  char *GetMutableSuffix(uint32_t idx, uint32_t len);
+  libtextclassifier3::StatusOr<Node *> GetMutableNode(uint32_t idx);
+  libtextclassifier3::StatusOr<Next *> GetMutableNextArray(uint32_t idx,
+                                                           uint32_t len);
+  libtextclassifier3::StatusOr<char *> GetMutableSuffix(uint32_t idx,
+                                                        uint32_t len);
 
   // Update crcs based on current contents. Returns all_crc or kNoCrc.
   Crc32 UpdateCrc();
@@ -324,13 +334,14 @@
   uint32_t suffixes_left() const;
 
   // REQUIRES: nodes_left() > 0.
-  Node *AllocNode();
+  libtextclassifier3::StatusOr<Node *> AllocNode();
   // REQUIRES: nexts_left() >= kMaxNextArraySize.
   libtextclassifier3::StatusOr<Next *> AllocNextArray(int size);
   void FreeNextArray(Next *next, int log2_size);
   // REQUIRES: suffixes_left() >= strlen(suffix) + 1 + value_size()
-  uint32_t MakeSuffix(std::string_view suffix, const void *value,
-                      uint32_t *value_index);
+  libtextclassifier3::StatusOr<uint32_t> MakeSuffix(std::string_view suffix,
+                                                    const void *value,
+                                                    uint32_t *value_index);
 
   const IcingDynamicTrieHeader &hdr() const { return hdr_.hdr; }
 
@@ -535,7 +546,7 @@
   if (!hdr_.Init(hdr_mmapper_.address(),
                  IcingMMapper::system_page_size() - sizeof(Crcs)) ||
       !hdr_.Verify()) {
-    ICING_LOG(ERROR) << "Trie reading header failed";
+    ICING_LOG(ERROR) << "Trie reading header failed. Path: " << file_basename_;
     goto failed;
   }
 
@@ -757,9 +768,10 @@
   return total;
 }
 
-IcingDynamicTrie::Node *IcingDynamicTrie::IcingDynamicTrieStorage::AllocNode() {
+libtextclassifier3::StatusOr<IcingDynamicTrie::Node *>
+IcingDynamicTrie::IcingDynamicTrieStorage::AllocNode() {
   if (nodes_left() == 0) {
-    ICING_LOG(FATAL) << "No allocated nodes left";
+    return absl_ports::ResourceExhaustedError("No allocated nodes left");
   }
 
   hdr_.hdr.set_num_nodes(hdr_.hdr.num_nodes() + 1);
@@ -774,7 +786,7 @@
   }
 
   if (nexts_left() < static_cast<uint32_t>(kMaxNextArraySize)) {
-    ICING_LOG(FATAL) << "'next' buffer not enough";
+    return absl_ports::ResourceExhaustedError("'next' buffer not enough");
   }
 
   // Compute ceil(log2(size)).
@@ -786,15 +798,17 @@
   // Look in free list.
   Next *ret;
   if (hdr_.hdr.free_lists(log2_size) != kInvalidNextIndex) {
-    ret = GetMutableNextArray(hdr_.hdr.free_lists(log2_size), aligned_size);
+    ICING_ASSIGN_OR_RETURN(
+        ret, GetMutableNextArray(hdr_.hdr.free_lists(log2_size), aligned_size));
     uint32_t next_link = ret->next_index();
     if (next_link != kInvalidNextIndex && next_link >= hdr_.hdr.max_nexts()) {
-      ICING_LOG(FATAL) << "'next' index is out of range";
+      return absl_ports::InternalError("'next' index is out of range");
     }
     hdr_.hdr.set_free_lists(log2_size, next_link);
   } else {
     // Allocate a new one.
-    ret = GetMutableNextArray(hdr_.hdr.num_nexts(), aligned_size);
+    ICING_ASSIGN_OR_RETURN(
+        ret, GetMutableNextArray(hdr_.hdr.num_nexts(), aligned_size));
     hdr_.hdr.set_num_nexts(hdr_.hdr.num_nexts() + aligned_size);
   }
 
@@ -816,14 +830,17 @@
   hdr_.hdr.set_free_lists(log2_size, GetNextArrayIndex(next));
 }
 
-uint32_t IcingDynamicTrie::IcingDynamicTrieStorage::MakeSuffix(
-    std::string_view suffix, const void *value, uint32_t *value_index) {
+libtextclassifier3::StatusOr<uint32_t>
+IcingDynamicTrie::IcingDynamicTrieStorage::MakeSuffix(std::string_view suffix,
+                                                      const void *value,
+                                                      uint32_t *value_index) {
   if (suffixes_left() < suffix.size() + 1 + value_size()) {
-    ICING_LOG(FATAL) << "'suffix' buffer not enough";
+    return absl_ports::ResourceExhaustedError("'suffix' buffer not enough");
   }
 
-  char *start = GetMutableSuffix(hdr_.hdr.suffixes_size(),
-                                 suffix.size() + 1 + value_size());
+  ICING_ASSIGN_OR_RETURN(char *start,
+                         GetMutableSuffix(hdr_.hdr.suffixes_size(),
+                                          suffix.size() + 1 + value_size()));
   // Copy suffix.
   memcpy(start, suffix.data(), suffix.size());
   // Set a '\0' after suffix.
@@ -908,19 +925,20 @@
   return hdr_.SerializeToArray(hdr_mmapper_.address(), serialized_header_max());
 }
 
-IcingDynamicTrie::Node *
+libtextclassifier3::StatusOr<IcingDynamicTrie::Node *>
 IcingDynamicTrie::IcingDynamicTrieStorage::GetMutableNode(uint32_t idx) {
   return array_storage_[NODE].GetMutableMem<Node>(idx, 1);
 }
 
-IcingDynamicTrie::Next *
+libtextclassifier3::StatusOr<IcingDynamicTrie::Next *>
 IcingDynamicTrie::IcingDynamicTrieStorage::GetMutableNextArray(uint32_t idx,
                                                                uint32_t len) {
   return array_storage_[NEXT].GetMutableMem<Next>(idx, len);
 }
 
-char *IcingDynamicTrie::IcingDynamicTrieStorage::GetMutableSuffix(
-    uint32_t idx, uint32_t len) {
+libtextclassifier3::StatusOr<char *>
+IcingDynamicTrie::IcingDynamicTrieStorage::GetMutableSuffix(uint32_t idx,
+                                                            uint32_t len) {
   return array_storage_[SUFFIX].GetMutableMem<char>(idx, len);
 }
 
@@ -958,7 +976,8 @@
   uint32_t magic;
   memcpy(&magic, buf, sizeof(magic));
   if (magic != kMagic) {
-    ICING_LOG(ERROR) << "Trie header magic mismatch";
+    ICING_LOG(ERROR) << "Invalid header magic for IcingDynamicTrie. Expected: "
+                     << kMagic << ", actual: " << magic;
     return false;
   }
   uint32_t len;
@@ -1558,49 +1577,42 @@
   deleted_bitmap_->Truncate(0);
 }
 
-bool IcingDynamicTrie::ClearSuffixAndValue(uint32_t suffix_value_index) {
+libtextclassifier3::Status IcingDynamicTrie::ClearSuffixAndValue(
+    uint32_t suffix_value_index) {
   // The size 1 below is for a '\0' between the suffix and the value.
   size_t suffix_and_value_length =
       strlen(this->storage_->GetSuffix(suffix_value_index)) + 1 +
       this->value_size();
-  char *mutable_suffix_and_value = this->storage_->GetMutableSuffix(
-      suffix_value_index, suffix_and_value_length);
-
-  if (mutable_suffix_and_value == nullptr) {
-    return false;
-  }
+  ICING_ASSIGN_OR_RETURN(char *mutable_suffix_and_value,
+                         this->storage_->GetMutableSuffix(
+                             suffix_value_index, suffix_and_value_length));
 
   memset(mutable_suffix_and_value, 0, suffix_and_value_length);
-  return true;
+  return libtextclassifier3::Status::OK;
 }
 
-bool IcingDynamicTrie::ResetNext(uint32_t next_index) {
-  Next *mutable_next =
-      this->storage_->GetMutableNextArray(next_index, /*len=*/1);
+libtextclassifier3::Status IcingDynamicTrie::ResetNext(uint32_t next_index) {
+  ICING_ASSIGN_OR_RETURN(
+      Next * mutable_next,
+      this->storage_->GetMutableNextArray(next_index, /*len=*/1));
 
-  if (mutable_next == nullptr) {
-    return false;
-  }
   ResetMutableNext(*mutable_next);
-  return true;
+  return libtextclassifier3::Status::OK;
 }
 
-bool IcingDynamicTrie::SortNextArray(const Node *node) {
+libtextclassifier3::Status IcingDynamicTrie::SortNextArray(const Node *node) {
   if (node == nullptr) {
     // Nothing to sort, return success directly.
-    return true;
+    return libtextclassifier3::Status::OK;
   }
 
   uint32_t next_array_buffer_size = 1u << node->log2_num_children();
-  Next *next_array_start = this->storage_->GetMutableNextArray(
-      node->next_index(), next_array_buffer_size);
-
-  if (next_array_start == nullptr) {
-    return false;
-  }
+  ICING_ASSIGN_OR_RETURN(Next * next_array_start,
+                         this->storage_->GetMutableNextArray(
+                             node->next_index(), next_array_buffer_size));
 
   std::sort(next_array_start, next_array_start + next_array_buffer_size);
-  return true;
+  return libtextclassifier3::Status::OK;
 }
 
 libtextclassifier3::Status IcingDynamicTrie::Insert(std::string_view key,
@@ -1609,7 +1621,7 @@
                                                     bool replace,
                                                     bool *pnew_key) {
   if (!is_initialized()) {
-    ICING_LOG(FATAL) << "DynamicTrie not initialized";
+    return absl_ports::FailedPreconditionError("DynamicTrie not initialized");
   }
 
   if (pnew_key) *pnew_key = false;
@@ -1632,22 +1644,27 @@
 
   uint32_t best_node_index;
   int key_offset;
-  FindBestNode(key, &best_node_index, &key_offset, false);
+  libtextclassifier3::Status status =
+      FindBestNode(key, &best_node_index, &key_offset, false);
 
   // A negative key_offset indicates that storage_ is empty
-  if (key_offset < 0) {
+  if (key_offset < 0 || !status.ok()) {
     // First key.
     if (!storage_->empty()) {
-      ICING_LOG(FATAL) << "Key offset is negative but storage is not empty, "
-                          "there're inconsistencies in dynamic trie.";
+      return absl_ports::InternalError(
+          "Key offset is negative but storage is not empty, "
+          "there're inconsistencies in dynamic trie.");
     }
-    Node *node = storage_->AllocNode();
-    node->set_next_index(storage_->MakeSuffix(key, value, value_index));
+    ICING_ASSIGN_OR_RETURN(Node * node, storage_->AllocNode());
+    ICING_ASSIGN_OR_RETURN(uint32_t next_index,
+                           storage_->MakeSuffix(key, value, value_index));
+    node->set_next_index(next_index);
     node->set_is_leaf(true);
     node->set_log2_num_children(0);
   } else if (storage_->GetNode(best_node_index)->is_leaf()) {
     // Prefix in the trie. Split at leaf.
-    Node *split_node = storage_->GetMutableNode(best_node_index);
+    ICING_ASSIGN_OR_RETURN(Node * split_node,
+                           storage_->GetMutableNode(best_node_index));
     const char *prev_suffix = storage_->GetSuffix(split_node->next_index());
 
     // Find the common prefix length starting from prev_suffix[0] and
@@ -1672,9 +1689,11 @@
       }
       // Update value if replace == true and return.
       if (replace) {
-        char *mutable_prev_suffix_cur = storage_->GetMutableSuffix(
-            storage_->GetSuffixIndex(prev_suffix + common_prefix_len + 1),
-            value_size());
+        ICING_ASSIGN_OR_RETURN(
+            char *mutable_prev_suffix_cur,
+            storage_->GetMutableSuffix(
+                storage_->GetSuffixIndex(prev_suffix + common_prefix_len + 1),
+                value_size()));
         memcpy(mutable_prev_suffix_cur, value, value_size());
       }
       return libtextclassifier3::Status::OK;
@@ -1689,7 +1708,7 @@
       split_node->set_next_index(storage_->GetNextArrayIndex(split_next));
       split_node->set_is_leaf(false);
       split_node->set_log2_num_children(0);
-      Node *child_node = storage_->AllocNode();
+      ICING_ASSIGN_OR_RETURN(Node * child_node, storage_->AllocNode());
       split_next[0].set_val(*(prev_suffix + i));
       split_next[0].set_node_index(storage_->GetNodeIndex(child_node));
 
@@ -1701,8 +1720,8 @@
     split_node->set_next_index(storage_->GetNextArrayIndex(split_next));
     split_node->set_is_leaf(false);
     split_node->set_log2_num_children(1);
-    Node *prev_suffix_node = storage_->AllocNode();
-    Node *key_node = storage_->AllocNode();
+    ICING_ASSIGN_OR_RETURN(Node * prev_suffix_node, storage_->AllocNode());
+    ICING_ASSIGN_OR_RETURN(Node * key_node, storage_->AllocNode());
     split_next[0].set_val(*(prev_suffix + common_prefix_len));
     split_next[0].set_node_index(storage_->GetNodeIndex(prev_suffix_node));
     if (*(prev_suffix + common_prefix_len)) {
@@ -1721,12 +1740,16 @@
     split_next[1].set_val(next_val);
     split_next[1].set_node_index(storage_->GetNodeIndex(key_node));
     if (next_val != '\0') {
-      uint32_t next_index = storage_->MakeSuffix(
-          key.substr(key_offset + common_prefix_len + 1), value, value_index);
+      ICING_ASSIGN_OR_RETURN(
+          uint32_t next_index,
+          storage_->MakeSuffix(key.substr(key_offset + common_prefix_len + 1),
+                               value, value_index));
       key_node->set_next_index(next_index);
     } else {
-      uint32_t next_index = storage_->MakeSuffix(
-          key.substr(key_offset + common_prefix_len), value, value_index);
+      ICING_ASSIGN_OR_RETURN(
+          uint32_t next_index,
+          storage_->MakeSuffix(key.substr(key_offset + common_prefix_len),
+                               value, value_index));
       key_node->set_next_index(next_index);
     }
     key_node->set_is_leaf(true);
@@ -1738,14 +1761,16 @@
     const Node *best_node = storage_->GetNode(best_node_index);
 
     // Add our value as a node + suffix.
-    Node *new_leaf_node = storage_->AllocNode();
+    ICING_ASSIGN_OR_RETURN(Node * new_leaf_node, storage_->AllocNode());
     if (key_offset < key.size()) {
-      uint32_t next_index =
-          storage_->MakeSuffix(key.substr(key_offset + 1), value, value_index);
+      ICING_ASSIGN_OR_RETURN(
+          uint32_t next_index,
+          storage_->MakeSuffix(key.substr(key_offset + 1), value, value_index));
       new_leaf_node->set_next_index(next_index);
     } else {
-      uint32_t next_index =
-          storage_->MakeSuffix(key.substr(key_offset), value, value_index);
+      ICING_ASSIGN_OR_RETURN(
+          uint32_t next_index,
+          storage_->MakeSuffix(key.substr(key_offset), value, value_index));
       new_leaf_node->set_next_index(next_index);
     }
     new_leaf_node->set_is_leaf(true);
@@ -1753,8 +1778,9 @@
 
     // Figure out the real length of the existing next array.
     uint32_t next_array_buffer_size = 1u << best_node->log2_num_children();
-    Next *cur_next = storage_->GetMutableNextArray(best_node->next_index(),
-                                                   next_array_buffer_size);
+    ICING_ASSIGN_OR_RETURN(
+        Next * cur_next, storage_->GetMutableNextArray(best_node->next_index(),
+                                                       next_array_buffer_size));
     int next_len = GetValidNextsSize(cur_next, next_array_buffer_size);
     Next *new_next = cur_next;
     if (next_len == (next_array_buffer_size)) {
@@ -1772,8 +1798,9 @@
     // If this was new, update the parent node and free the old next
     // array.
     if (new_next != cur_next) {
-      Node *mutable_best_node =
-          storage_->GetMutableNode(storage_->GetNodeIndex(best_node));
+      ICING_ASSIGN_OR_RETURN(
+          Node * mutable_best_node,
+          storage_->GetMutableNode(storage_->GetNodeIndex(best_node)));
       mutable_best_node->set_next_index(storage_->GetNextArrayIndex(new_next));
       mutable_best_node->set_is_leaf(false);
       uint8_t log2_num_children = mutable_best_node->log2_num_children();
@@ -1806,18 +1833,21 @@
   return static_cast<const void *>(storage_->GetSuffix(value_index));
 }
 
-void IcingDynamicTrie::SetValueAtIndex(uint32_t value_index,
-                                       const void *value) {
+libtextclassifier3::Status IcingDynamicTrie::SetValueAtIndex(
+    uint32_t value_index, const void *value) {
   if (!is_initialized()) {
-    ICING_LOG(FATAL) << "DynamicTrie not initialized";
+    return absl_ports::FailedPreconditionError("DynamicTrie not initialized");
   }
 
   if (value_index > storage_->hdr().max_suffixes_size() - value_size()) {
-    ICING_LOG(FATAL) << "Value index is out of range";
+    return absl_ports::OutOfRangeError("Value index is out of range");
   }
 
-  memcpy(storage_->GetMutableSuffix(value_index, value_size()), value,
-         value_size());
+  ICING_ASSIGN_OR_RETURN(char *mutable_suffix,
+                         storage_->GetMutableSuffix(value_index, value_size()));
+  memcpy(mutable_suffix, value, value_size());
+
+  return libtextclassifier3::Status::OK;
 }
 
 bool IcingDynamicTrie::Find(std::string_view key, void *value,
@@ -1832,9 +1862,10 @@
 
   uint32_t best_node_index;
   int key_offset;
-  FindBestNode(key, &best_node_index, &key_offset, false);
+  libtextclassifier3::Status status =
+      FindBestNode(key, &best_node_index, &key_offset, false);
 
-  if (key_offset < 0) {
+  if (key_offset < 0 || !status.ok()) {
     return false;
   }
 
@@ -1929,7 +1960,8 @@
   // Find node matching prefix.
   uint32_t node_index;
   int key_offset;
-  trie_.FindBestNode(cur_key_, &node_index, &key_offset, true);
+  libtextclassifier3::Status status =
+      trie_.FindBestNode(cur_key_, &node_index, &key_offset, true);
 
   // Two cases/states:
   //
@@ -1940,9 +1972,8 @@
   //   prefix. Check that suffix matches the prefix. Then we set
   //   single_leaf_match_ = true and apply different logic for
   //   Advance.
-  if (key_offset < 0) {
-    // A negative key_offset indicates that trie_.storage_ is empty
-    ICING_LOG(FATAL) << "Trie storage is empty";
+  if (key_offset < 0 || !status.ok()) {
+    ICING_LOG(FATAL) << "FindBestNode failed: " << status.error_message();
   }
 
   const Node *best_node = trie_.storage_->GetNode(node_index);
@@ -2179,8 +2210,8 @@
 
 bool IcingDynamicTrie::Utf8Iterator::IsValid() const { return cur_len_ > 0; }
 
-const IcingDynamicTrie::Next *IcingDynamicTrie::GetNextByChar(
-    const Node *node, uint8_t key_char) const {
+libtextclassifier3::StatusOr<const IcingDynamicTrie::Next*>
+IcingDynamicTrie::GetNextByChar(const Node* node, uint8_t key_char) const {
   const Next *next_start = storage_->GetNext(node->next_index(), 0);
   const Next *next_end = next_start + (1 << node->log2_num_children());
 
@@ -2190,6 +2221,11 @@
     return nullptr;
   }
 
+  if (found->node_index() >= storage_->GetNodeArraySize()) {
+    return absl_ports::InternalError(
+        "Node index is out of bounds. The index may be corrupted.");
+  }
+
   return found;
 }
 
@@ -2229,9 +2265,9 @@
   }
 }
 
-void IcingDynamicTrie::FindBestNode(std::string_view key,
-                                    uint32_t *best_node_index, int *key_offset,
-                                    bool prefix, bool utf8) const {
+libtextclassifier3::Status IcingDynamicTrie::FindBestNode(
+    std::string_view key, uint32_t* best_node_index, int* key_offset,
+    bool prefix, bool utf8) const {
   // Find the best node such that:
   //
   // - If key is NOT in the trie, key[0..key_offset) is a prefix to
@@ -2247,7 +2283,7 @@
   if (storage_->empty()) {
     *best_node_index = 0;
     *key_offset = -1;
-    return;
+    return absl_ports::InternalError("Trie is empty.");
   }
 
   const Node *cur_node = storage_->GetRootNode();
@@ -2256,7 +2292,8 @@
   const Node *utf8_node = cur_node;
   while (!cur_node->is_leaf()) {
     char cur_char = GetCharOrNull(key, cur_key_idx);
-    const Next *found = GetNextByChar(cur_node, cur_char);
+    ICING_ASSIGN_OR_RETURN(const Next* found,
+                           GetNextByChar(cur_node, cur_char));
     if (!found) break;
 
     if (prefix && found->val() == 0) {
@@ -2287,10 +2324,13 @@
 
   *best_node_index = storage_->GetNodeIndex(cur_node);
   *key_offset = cur_key_idx;
+
+  return libtextclassifier3::Status::OK;
 }
 
-int IcingDynamicTrie::FindNewBranchingPrefixLength(std::string_view key,
-                                                   bool utf8) const {
+libtextclassifier3::StatusOr<int>
+IcingDynamicTrie::FindNewBranchingPrefixLength(std::string_view key,
+                                               bool utf8) const {
   if (!IsKeyValid(key)) {
     return kNoBranchFound;
   }
@@ -2301,7 +2341,8 @@
 
   uint32_t best_node_index;
   int key_offset;
-  FindBestNode(key, &best_node_index, &key_offset, /*prefix=*/true, utf8);
+  ICING_RETURN_IF_ERROR(
+      FindBestNode(key, &best_node_index, &key_offset, /*prefix=*/true, utf8));
   if (key_offset < 0) {
     return kNoBranchFound;
   }
@@ -2353,8 +2394,9 @@
   return kNoBranchFound;
 }
 
-std::vector<int> IcingDynamicTrie::FindBranchingPrefixLengths(
-    std::string_view key, bool utf8) const {
+libtextclassifier3::StatusOr<std::vector<int>>
+IcingDynamicTrie::FindBranchingPrefixLengths(std::string_view key,
+                                             bool utf8) const {
   std::vector<int> prefix_lengths;
 
   if (!IsKeyValid(key)) {
@@ -2383,7 +2425,8 @@
     }
 
     // Move to next.
-    const Next *found = GetNextByChar(cur_node, key[idx]);
+    ICING_ASSIGN_OR_RETURN(const Next* found,
+                           GetNextByChar(cur_node, key[idx]));
     if (found == nullptr) {
       break;
     }
@@ -2394,7 +2437,8 @@
   return prefix_lengths;
 }
 
-bool IcingDynamicTrie::IsBranchingTerm(std::string_view key) const {
+libtextclassifier3::StatusOr<bool> IcingDynamicTrie::IsBranchingTerm(
+    std::string_view key) const {
   if (!is_initialized()) {
     ICING_LOG(FATAL) << "DynamicTrie not initialized";
   }
@@ -2409,7 +2453,8 @@
 
   uint32_t best_node_index;
   int key_offset;
-  FindBestNode(key, &best_node_index, &key_offset, /*prefix=*/true);
+  ICING_RETURN_IF_ERROR(
+      FindBestNode(key, &best_node_index, &key_offset, /*prefix=*/true));
   if (key_offset < 0) {
     return false;
   }
@@ -2428,7 +2473,8 @@
   // Found key as an intermediate node, but key is not a valid term stored in
   // the trie. In this case, we need at least two children for key to be a
   // branching term.
-  if (GetNextByChar(cur_node, '\0') == nullptr) {
+  ICING_ASSIGN_OR_RETURN(const Next* next, GetNextByChar(cur_node, '\0'));
+  if (next == nullptr) {
     return cur_node->log2_num_children() >= 1;
   }
 
@@ -2620,19 +2666,18 @@
 //        element will cause a crash or fetch incorrect data.
 //    - So we must reset the trie state to make sure storage_->empty() works
 //      correctly and prevents trie APIs from accessing the root node.
-bool IcingDynamicTrie::Delete(std::string_view key) {
+libtextclassifier3::Status IcingDynamicTrie::Delete(std::string_view key) {
   if (!is_initialized()) {
-    ICING_LOG(ERROR) << "DynamicTrie not initialized";
-    return false;
+    return absl_ports::FailedPreconditionError("DynamicTrie not initialized");
   }
 
   if (!IsKeyValid(key)) {
-    return false;
+    return absl_ports::InvalidArgumentError("Invalid key");
   }
 
   if (storage_->empty()) {
     // Nothing to delete.
-    return true;
+    return libtextclassifier3::Status::OK;
   }
 
   // Tries to find the key in the trie, starting from the root.
@@ -2656,7 +2701,7 @@
       // Leaf node, now check the suffix.
       if (key.substr(i) != storage_->GetSuffix(current_node->next_index())) {
         // Key does not exist in the trie, nothing to delete.
-        return true;
+        return libtextclassifier3::Status::OK;
       }
       // Otherwise, key is found.
       break;
@@ -2667,27 +2712,29 @@
     if (i == key.length()) {
       // When we're at the end of the key, the next char is the termination char
       // '\0'.
-      next = GetNextByChar(current_node, '\0');
+      ICING_ASSIGN_OR_RETURN(next, GetNextByChar(current_node, '\0'));
     } else {
-      next = GetNextByChar(current_node, key[i]);
+      ICING_ASSIGN_OR_RETURN(next, GetNextByChar(current_node, key[i]));
     }
 
     if (next == nullptr) {
       // Key does not exist in the trie, nothing to delete.
-      return true;
+      return libtextclassifier3::Status::OK;
     }
 
     // Checks the real size of next array.
     uint32_t next_array_buffer_size = 1u << current_node->log2_num_children();
-    Next *next_array_start = storage_->GetMutableNextArray(
-        current_node->next_index(), next_array_buffer_size);
+    ICING_ASSIGN_OR_RETURN(
+        Next * next_array_start,
+        storage_->GetMutableNextArray(current_node->next_index(),
+                                      next_array_buffer_size));
     int valid_next_array_size =
         GetValidNextsSize(next_array_start, next_array_buffer_size);
     if (valid_next_array_size == 0) {
       // Key does not exist in the trie, nothing to delete.
       // This shouldn't happen, but we put a sanity check here in case something
       // is wrong.
-      return true;
+      return libtextclassifier3::Status::OK;
     } else if (valid_next_array_size == 1) {
       // Single-child branch will be deleted.
       nexts_to_reset.push_back(storage_->GetNextArrayIndex(next));
@@ -2704,25 +2751,28 @@
   }
   // Now we've found the key in the trie.
 
-  ClearSuffixAndValue(current_node->next_index());
+  ICING_RETURN_IF_ERROR(ClearSuffixAndValue(current_node->next_index()));
 
   // Resets nexts to remove key information.
   for (uint32_t next_index : nexts_to_reset) {
-    ResetNext(next_index);
+    ICING_RETURN_IF_ERROR(ResetNext(next_index));
   }
 
   if (last_multichild_node != nullptr) {
-    SortNextArray(last_multichild_node);
+    ICING_RETURN_IF_ERROR(SortNextArray(last_multichild_node));
     uint32_t next_array_buffer_size =
         1u << last_multichild_node->log2_num_children();
-    Next *next_array_start = this->storage_->GetMutableNextArray(
-        last_multichild_node->next_index(), next_array_buffer_size);
+    ICING_ASSIGN_OR_RETURN(
+        Next * next_array_start,
+        this->storage_->GetMutableNextArray(last_multichild_node->next_index(),
+                                            next_array_buffer_size));
     uint32_t num_children =
         GetValidNextsSize(next_array_start, next_array_buffer_size);
     // Shrink the next array if we can.
     if (num_children == next_array_buffer_size / 2) {
-      Node *mutable_node = storage_->GetMutableNode(
-          storage_->GetNodeIndex(last_multichild_node));
+      ICING_ASSIGN_OR_RETURN(Node * mutable_node,
+                             storage_->GetMutableNode(
+                                 storage_->GetNodeIndex(last_multichild_node)));
       mutable_node->set_log2_num_children(mutable_node->log2_num_children() -
                                           1);
       // Add the unused second half of the next array to the free list.
@@ -2744,7 +2794,7 @@
     Clear();
   }
 
-  return true;
+  return libtextclassifier3::Status::OK;
 }
 
 bool IcingDynamicTrie::ClearPropertyForAllValues(uint32_t property_id) {

diff --git a/icing/legacy/index/icing-dynamic-trie.h b/icing/legacy/index/icing-dynamic-trie.h
index 53d9bad..a741264 100644
--- a/icing/legacy/index/icing-dynamic-trie.h
+++ b/icing/legacy/index/icing-dynamic-trie.h

@@ -35,23 +35,21 @@
 #ifndef ICING_LEGACY_INDEX_ICING_DYNAMIC_TRIE_H_
 #define ICING_LEGACY_INDEX_ICING_DYNAMIC_TRIE_H_
 
+#include <cstddef>
 #include <cstdint>
 #include <memory>
+#include <ostream>
 #include <string>
 #include <string_view>
-#include <unordered_map>
 #include <vector>
 
 #include "icing/text_classifier/lib3/utils/base/status.h"
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
-#include "icing/legacy/core/icing-compat.h"
 #include "icing/legacy/core/icing-packed-pod.h"
 #include "icing/legacy/index/icing-filesystem.h"
-#include "icing/legacy/index/icing-mmapper.h"
 #include "icing/legacy/index/icing-storage.h"
 #include "icing/legacy/index/proto/icing-dynamic-trie-header.pb.h"
 #include "icing/util/crc32.h"
-#include "icing/util/i18n-utils.h"
 #include "unicode/utf8.h"
 
 namespace icing {
@@ -299,15 +297,15 @@
   // Empty out the trie without closing or removing.
   void Clear();
 
-  // Clears the suffix and value at the given index. Returns true on success.
-  bool ClearSuffixAndValue(uint32_t suffix_value_index);
+  // Clears the suffix and value at the given index. Returns OK on success.
+  libtextclassifier3::Status ClearSuffixAndValue(uint32_t suffix_value_index);
 
   // Resets the next at the given index so that it points to no node.
-  // Returns true on success.
-  bool ResetNext(uint32_t next_index);
+  // Returns OK on success.
+  libtextclassifier3::Status ResetNext(uint32_t next_index);
 
-  // Sorts the next array of the node. Returns true on success.
-  bool SortNextArray(const Node *node);
+  // Sorts the next array of the node. Returns OK on success.
+  libtextclassifier3::Status SortNextArray(const Node *node);
 
   // Sync to disk.
   bool Sync() override;
@@ -348,7 +346,8 @@
   // value.
   //
   // REQUIRES: value a buffer of size value_size()
-  void SetValueAtIndex(uint32_t value_index, const void *value);
+  libtextclassifier3::Status SetValueAtIndex(uint32_t value_index,
+                                             const void *value);
 
   // Returns true if key is found and sets value. If value_index is
   // not NULL, returns value_index (see Insert discussion above).
@@ -384,19 +383,21 @@
   // Return prefix of any new branches created if key were inserted. If utf8 is
   // true, does not cut key mid-utf8. Returns kNoBranchFound if no branches
   // would be created.
-  int FindNewBranchingPrefixLength(std::string_view key, bool utf8) const;
+  libtextclassifier3::StatusOr<int> FindNewBranchingPrefixLength(
+      std::string_view key, bool utf8) const;
 
   // Find all prefixes of key where the trie branches. Excludes the key
   // itself. If utf8 is true, does not cut key mid-utf8.
-  std::vector<int> FindBranchingPrefixLengths(std::string_view key,
-                                              bool utf8) const;
+  libtextclassifier3::StatusOr<std::vector<int>> FindBranchingPrefixLengths(
+      std::string_view key, bool utf8) const;
 
   // Check if key is a branching term.
   //
   // key is a branching term, if and only if there exists terms s1 and s2 in the
   // trie such that key is the maximum common prefix of s1 and s2, but s1 and s2
   // are not prefixes of each other.
-  bool IsBranchingTerm(std::string_view key) const;
+  libtextclassifier3::StatusOr<bool> IsBranchingTerm(
+      std::string_view key) const;
 
   void GetDebugInfo(int verbosity, std::string *out) const override;
 
@@ -431,8 +432,8 @@
   bool ClearDeleted(uint32_t value_index);
 
   // Deletes the entry associated with the key. Data can not be recovered after
-  // the deletion. Returns true on success.
-  bool Delete(std::string_view key);
+  // the deletion. Returns OK on success.
+  libtextclassifier3::Status Delete(std::string_view key);
 
   // Clear a specific property id from all values.  For each value that has this
   // property cleared, also check to see if it was the only property set;  if
@@ -630,14 +631,17 @@
                              uint32_t depth = 0) const;
 
   // Helpers for Find and Insert.
-  const Next *GetNextByChar(const Node *node, uint8_t key_char) const;
+  libtextclassifier3::StatusOr<const Next*> GetNextByChar(
+      const Node* node, uint8_t key_char) const;
   const Next *LowerBound(const Next *start, const Next *end, uint8_t key_char,
                          uint32_t node_index = 0) const;
   // Returns the number of valid nexts in the array.
   int GetValidNextsSize(const IcingDynamicTrie::Next *next_array_start,
                         int next_array_length) const;
-  void FindBestNode(std::string_view key, uint32_t *best_node_index,
-                    int *key_offset, bool prefix, bool utf8 = false) const;
+  libtextclassifier3::Status FindBestNode(std::string_view key,
+                                          uint32_t* best_node_index,
+                                          int* key_offset, bool prefix,
+                                          bool utf8 = false) const;
 
   // For value properties.  This truncates the data by clearing it, but leaving
   // the storage intact.

diff --git a/icing/legacy/index/icing-dynamic-trie_test.cc b/icing/legacy/index/icing-dynamic-trie_test.cc
index 778a1a9..050ff32 100644
--- a/icing/legacy/index/icing-dynamic-trie_test.cc
+++ b/icing/legacy/index/icing-dynamic-trie_test.cc

@@ -14,14 +14,20 @@
 
 #include "icing/legacy/index/icing-dynamic-trie.h"
 
+#include <algorithm>
 #include <cstddef>
 #include <cstdint>
 #include <cstdio>
+#include <cstring>
+#include <random>
+#include <sstream>
 #include <string>
+#include <string_view>
 #include <unordered_set>
 #include <utility>
 #include <vector>
 
+#include "icing/text_classifier/lib3/utils/base/status.h"
 #include "icing/text_classifier/lib3/utils/hash/farmhash.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
@@ -30,6 +36,7 @@
 #include "icing/testing/common-matchers.h"
 #include "icing/testing/random-string.h"
 #include "icing/testing/tmp-directory.h"
+#include "icing/util/crc32.h"
 #include "icing/util/logging.h"
 
 namespace icing {
@@ -952,7 +959,7 @@
   ASSERT_THAT(trie, Not(IsEmpty()));
 
   // Deletes the key.
-  EXPECT_TRUE(trie.Delete("foo"));
+  ICING_EXPECT_OK(trie.Delete("foo"));
   EXPECT_FALSE(trie.Find("foo", &value));
   EXPECT_THAT(trie, SizeIs(0));  // Explicitly test size() method.
   EXPECT_THAT(trie, IsEmpty());
@@ -982,7 +989,7 @@
   ASSERT_THAT(trie, Not(IsEmpty()));
 
   // Deletes "bar". "r" is a leaf node in the trie.
-  EXPECT_TRUE(trie.Delete("bar"));
+  ICING_EXPECT_OK(trie.Delete("bar"));
   EXPECT_FALSE(trie.Find("bar", &value));
   EXPECT_TRUE(trie.Find("ba", &value));
   EXPECT_THAT(trie, SizeIs(1));
@@ -1013,7 +1020,7 @@
   ASSERT_THAT(trie, Not(IsEmpty()));
 
   // Deletes "ba" which is a key with termination node in the trie.
-  EXPECT_TRUE(trie.Delete("ba"));
+  ICING_EXPECT_OK(trie.Delete("ba"));
   EXPECT_FALSE(trie.Find("ba", &value));
   EXPECT_TRUE(trie.Find("bar", &value));
   EXPECT_THAT(trie, SizeIs(1));
@@ -1046,7 +1053,7 @@
   ASSERT_THAT(trie, Not(IsEmpty()));
 
   // Deletes "bc".
-  EXPECT_TRUE(trie.Delete("bc"));
+  ICING_EXPECT_OK(trie.Delete("bc"));
   EXPECT_FALSE(trie.Find("bc", &value));
   EXPECT_TRUE(trie.Find("ba", &value));
   EXPECT_TRUE(trie.Find("bb", &value));
@@ -1087,7 +1094,7 @@
   ASSERT_THAT(trie, Not(IsEmpty()));
 
   // Deletes "batter".
-  EXPECT_TRUE(trie.Delete("batter"));
+  ICING_EXPECT_OK(trie.Delete("batter"));
   EXPECT_FALSE(trie.Find("batter", &value));
   EXPECT_TRUE(trie.Find("battle", &value));
   EXPECT_TRUE(trie.Find("bar", &value));
@@ -1115,9 +1122,9 @@
   ASSERT_THAT(trie, Not(IsEmpty()));
 
   // Delete "foo", "bar", "baz".
-  EXPECT_TRUE(trie.Delete("foo"));
-  EXPECT_TRUE(trie.Delete("bar"));
-  EXPECT_TRUE(trie.Delete("baz"));
+  ICING_EXPECT_OK(trie.Delete("foo"));
+  ICING_EXPECT_OK(trie.Delete("bar"));
+  ICING_EXPECT_OK(trie.Delete("baz"));
 
   EXPECT_THAT(trie, SizeIs(0));  // Explicitly test size() method.
   EXPECT_THAT(trie, IsEmpty());
@@ -1142,7 +1149,7 @@
   ASSERT_THAT(trie.Insert("foo", &value), IsOk());
 
   // Deletes a key
-  ASSERT_TRUE(trie.Delete("bed"));
+  ICING_ASSERT_OK(trie.Delete("bed"));
   ASSERT_FALSE(trie.Find("bed", &value));
 
   // Inserts after deletion
@@ -1166,7 +1173,7 @@
   ASSERT_THAT(trie.Insert("foo", &value), IsOk());
 
   // Deletes a key
-  ASSERT_TRUE(trie.Delete("bed"));
+  ICING_ASSERT_OK(trie.Delete("bed"));
 
   // Iterates through all keys
   IcingDynamicTrie::Iterator iterator_all(trie, "");
@@ -1199,9 +1206,9 @@
   ASSERT_THAT(trie.Insert("foo", &value), IsOk());
 
   // Deletes all keys
-  ASSERT_TRUE(trie.Delete("bar"));
-  ASSERT_TRUE(trie.Delete("bed"));
-  ASSERT_TRUE(trie.Delete("foo"));
+  ICING_ASSERT_OK(trie.Delete("bar"));
+  ICING_ASSERT_OK(trie.Delete("bed"));
+  ICING_ASSERT_OK(trie.Delete("foo"));
 
   EXPECT_THAT(trie, IsEmpty());
 
@@ -1235,8 +1242,8 @@
   ASSERT_THAT(trie.Insert("bed", &value), IsOk());
 
   // "ba" and bedroom are not keys in the trie.
-  EXPECT_TRUE(trie.Delete("ba"));
-  EXPECT_TRUE(trie.Delete("bedroom"));
+  ICING_EXPECT_OK(trie.Delete("ba"));
+  ICING_EXPECT_OK(trie.Delete("bedroom"));
 
   // The original keys are not affected.
   EXPECT_TRUE(trie.Find("bar", &value));
@@ -1258,7 +1265,7 @@
   ASSERT_THAT(trie.Insert("fjord", &value), IsOk());
 
   // Delete the third child
-  EXPECT_TRUE(trie.Delete("foul"));
+  ICING_EXPECT_OK(trie.Delete("foul"));
 
   std::vector<std::string> remaining;
   for (IcingDynamicTrie::Iterator term_iter(trie, /*prefix=*/"");
@@ -1282,7 +1289,7 @@
   ASSERT_THAT(trie.Insert("fudge", &value), IsOk());
 
   // Delete the second child
-  EXPECT_TRUE(trie.Delete("foul"));
+  ICING_EXPECT_OK(trie.Delete("foul"));
 
   std::vector<std::string> remaining;
   for (IcingDynamicTrie::Iterator term_iter(trie, /*prefix=*/"");
@@ -1314,7 +1321,7 @@
   std::shuffle(terms.begin(), terms.end(), random);
   for (int i = 0; i < 1024; ++i) {
     exp_remaining.erase(terms[i]);
-    ASSERT_TRUE(trie.Delete(terms[i]));
+    ICING_ASSERT_OK(trie.Delete(terms[i]));
   }
 
   // Check that the iterator still works, and the remaining terms are correct.
@@ -1523,68 +1530,68 @@
   uint32_t value = 1;
 
   ASSERT_THAT(trie.Insert("", &value), IsOk());
-  EXPECT_FALSE(trie.IsBranchingTerm(""));
+  EXPECT_FALSE(trie.IsBranchingTerm("").ValueOrDie());
 
   ASSERT_THAT(trie.Insert("ab", &value), IsOk());
-  EXPECT_FALSE(trie.IsBranchingTerm(""));
-  EXPECT_FALSE(trie.IsBranchingTerm("ab"));
+  EXPECT_FALSE(trie.IsBranchingTerm("").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("ab").ValueOrDie());
 
   ASSERT_THAT(trie.Insert("ac", &value), IsOk());
   // "" is a prefix of "ab" and "ac", but it is not a branching term.
-  EXPECT_FALSE(trie.IsBranchingTerm(""));
-  EXPECT_FALSE(trie.IsBranchingTerm("ab"));
-  EXPECT_FALSE(trie.IsBranchingTerm("ac"));
+  EXPECT_FALSE(trie.IsBranchingTerm("").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("ab").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("ac").ValueOrDie());
 
   ASSERT_THAT(trie.Insert("ba", &value), IsOk());
   // "" now branches to "ba"
-  EXPECT_TRUE(trie.IsBranchingTerm(""));
-  EXPECT_FALSE(trie.IsBranchingTerm("ab"));
-  EXPECT_FALSE(trie.IsBranchingTerm("ac"));
-  EXPECT_FALSE(trie.IsBranchingTerm("ba"));
+  EXPECT_TRUE(trie.IsBranchingTerm("").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("ab").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("ac").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("ba").ValueOrDie());
 
   ASSERT_THAT(trie.Insert("a", &value), IsOk());
-  EXPECT_TRUE(trie.IsBranchingTerm(""));
+  EXPECT_TRUE(trie.IsBranchingTerm("").ValueOrDie());
   // "a" branches to "ab" and "ac"
-  EXPECT_TRUE(trie.IsBranchingTerm("a"));
-  EXPECT_FALSE(trie.IsBranchingTerm("ab"));
-  EXPECT_FALSE(trie.IsBranchingTerm("ac"));
-  EXPECT_FALSE(trie.IsBranchingTerm("ba"));
+  EXPECT_TRUE(trie.IsBranchingTerm("a").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("ab").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("ac").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("ba").ValueOrDie());
 
   ASSERT_THAT(trie.Insert("abc", &value), IsOk());
   ASSERT_THAT(trie.Insert("acd", &value), IsOk());
-  EXPECT_TRUE(trie.IsBranchingTerm(""));
-  EXPECT_TRUE(trie.IsBranchingTerm("a"));
+  EXPECT_TRUE(trie.IsBranchingTerm("").ValueOrDie());
+  EXPECT_TRUE(trie.IsBranchingTerm("a").ValueOrDie());
   // "ab" is a prefix of "abc", but it is not a branching term.
-  EXPECT_FALSE(trie.IsBranchingTerm("ab"));
+  EXPECT_FALSE(trie.IsBranchingTerm("ab").ValueOrDie());
   // "ac" is a prefix of "acd", but it is not a branching term.
-  EXPECT_FALSE(trie.IsBranchingTerm("ac"));
-  EXPECT_FALSE(trie.IsBranchingTerm("ba"));
-  EXPECT_FALSE(trie.IsBranchingTerm("abc"));
-  EXPECT_FALSE(trie.IsBranchingTerm("acd"));
+  EXPECT_FALSE(trie.IsBranchingTerm("ac").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("ba").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("abc").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("acd").ValueOrDie());
 
   ASSERT_THAT(trie.Insert("abcd", &value), IsOk());
-  EXPECT_TRUE(trie.IsBranchingTerm(""));
-  EXPECT_TRUE(trie.IsBranchingTerm("a"));
+  EXPECT_TRUE(trie.IsBranchingTerm("").ValueOrDie());
+  EXPECT_TRUE(trie.IsBranchingTerm("a").ValueOrDie());
   // "ab" is a prefix of "abc" and "abcd", but it is not a branching term.
-  EXPECT_FALSE(trie.IsBranchingTerm("ab"));
-  EXPECT_FALSE(trie.IsBranchingTerm("ac"));
-  EXPECT_FALSE(trie.IsBranchingTerm("ba"));
+  EXPECT_FALSE(trie.IsBranchingTerm("ab").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("ac").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("ba").ValueOrDie());
   // "abc" is a prefix of "abcd", but it is not a branching term.
-  EXPECT_FALSE(trie.IsBranchingTerm("abc"));
-  EXPECT_FALSE(trie.IsBranchingTerm("acd"));
-  EXPECT_FALSE(trie.IsBranchingTerm("abcd"));
+  EXPECT_FALSE(trie.IsBranchingTerm("abc").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("acd").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("abcd").ValueOrDie());
 
   ASSERT_THAT(trie.Insert("abd", &value), IsOk());
-  EXPECT_TRUE(trie.IsBranchingTerm(""));
-  EXPECT_TRUE(trie.IsBranchingTerm("a"));
+  EXPECT_TRUE(trie.IsBranchingTerm("").ValueOrDie());
+  EXPECT_TRUE(trie.IsBranchingTerm("a").ValueOrDie());
   // "ab" branches to "abc" and "abd"
-  EXPECT_TRUE(trie.IsBranchingTerm("ab"));
-  EXPECT_FALSE(trie.IsBranchingTerm("ac"));
-  EXPECT_FALSE(trie.IsBranchingTerm("ba"));
-  EXPECT_FALSE(trie.IsBranchingTerm("abc"));
-  EXPECT_FALSE(trie.IsBranchingTerm("acd"));
-  EXPECT_FALSE(trie.IsBranchingTerm("abcd"));
-  EXPECT_FALSE(trie.IsBranchingTerm("abd"));
+  EXPECT_TRUE(trie.IsBranchingTerm("ab").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("ac").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("ba").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("abc").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("acd").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("abcd").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("abd").ValueOrDie());
 }
 
 TEST_F(IcingDynamicTrieTest, IsBranchingTermShouldWorkForNonExistingTerms) {
@@ -1596,58 +1603,58 @@
 
   uint32_t value = 1;
 
-  EXPECT_FALSE(trie.IsBranchingTerm(""));
-  EXPECT_FALSE(trie.IsBranchingTerm("a"));
-  EXPECT_FALSE(trie.IsBranchingTerm("ab"));
-  EXPECT_FALSE(trie.IsBranchingTerm("abc"));
+  EXPECT_FALSE(trie.IsBranchingTerm("").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("a").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("ab").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("abc").ValueOrDie());
 
   ASSERT_THAT(trie.Insert("aa", &value), IsOk());
-  EXPECT_FALSE(trie.IsBranchingTerm(""));
-  EXPECT_FALSE(trie.IsBranchingTerm("a"));
-  EXPECT_FALSE(trie.IsBranchingTerm("ab"));
-  EXPECT_FALSE(trie.IsBranchingTerm("abc"));
+  EXPECT_FALSE(trie.IsBranchingTerm("").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("a").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("ab").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("abc").ValueOrDie());
 
   ASSERT_THAT(trie.Insert("ac", &value), IsOk());
-  EXPECT_FALSE(trie.IsBranchingTerm(""));
+  EXPECT_FALSE(trie.IsBranchingTerm("").ValueOrDie());
   // "a" does not exist in the trie, but now it branches to "aa" and "ac".
-  EXPECT_TRUE(trie.IsBranchingTerm("a"));
-  EXPECT_FALSE(trie.IsBranchingTerm("ab"));
-  EXPECT_FALSE(trie.IsBranchingTerm("abc"));
+  EXPECT_TRUE(trie.IsBranchingTerm("a").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("ab").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("abc").ValueOrDie());
 
   ASSERT_THAT(trie.Insert("ad", &value), IsOk());
-  EXPECT_FALSE(trie.IsBranchingTerm(""));
-  EXPECT_TRUE(trie.IsBranchingTerm("a"));
-  EXPECT_FALSE(trie.IsBranchingTerm("ab"));
-  EXPECT_FALSE(trie.IsBranchingTerm("abc"));
+  EXPECT_FALSE(trie.IsBranchingTerm("").ValueOrDie());
+  EXPECT_TRUE(trie.IsBranchingTerm("a").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("ab").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("abc").ValueOrDie());
 
   ASSERT_THAT(trie.Insert("abcd", &value), IsOk());
-  EXPECT_FALSE(trie.IsBranchingTerm(""));
-  EXPECT_TRUE(trie.IsBranchingTerm("a"));
-  EXPECT_FALSE(trie.IsBranchingTerm("ab"));
-  EXPECT_FALSE(trie.IsBranchingTerm("abc"));
+  EXPECT_FALSE(trie.IsBranchingTerm("").ValueOrDie());
+  EXPECT_TRUE(trie.IsBranchingTerm("a").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("ab").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("abc").ValueOrDie());
 
   ASSERT_THAT(trie.Insert("abd", &value), IsOk());
-  EXPECT_FALSE(trie.IsBranchingTerm(""));
-  EXPECT_TRUE(trie.IsBranchingTerm("a"));
+  EXPECT_FALSE(trie.IsBranchingTerm("").ValueOrDie());
+  EXPECT_TRUE(trie.IsBranchingTerm("a").ValueOrDie());
   // "ab" does not exist in the trie, but now it branches to "abcd" and "abd".
-  EXPECT_TRUE(trie.IsBranchingTerm("ab"));
-  EXPECT_FALSE(trie.IsBranchingTerm("abc"));
+  EXPECT_TRUE(trie.IsBranchingTerm("ab").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("abc").ValueOrDie());
 
   ASSERT_THAT(trie.Insert("abce", &value), IsOk());
-  EXPECT_FALSE(trie.IsBranchingTerm(""));
-  EXPECT_TRUE(trie.IsBranchingTerm("a"));
-  EXPECT_TRUE(trie.IsBranchingTerm("ab"));
+  EXPECT_FALSE(trie.IsBranchingTerm("").ValueOrDie());
+  EXPECT_TRUE(trie.IsBranchingTerm("a").ValueOrDie());
+  EXPECT_TRUE(trie.IsBranchingTerm("ab").ValueOrDie());
   // "abc" does not exist in the trie, but now it branches to "abcd" and "abce".
-  EXPECT_TRUE(trie.IsBranchingTerm("abc"));
+  EXPECT_TRUE(trie.IsBranchingTerm("abc").ValueOrDie());
 
   ASSERT_THAT(trie.Insert("abc_suffix", &value), IsOk());
-  EXPECT_FALSE(trie.IsBranchingTerm(""));
-  EXPECT_TRUE(trie.IsBranchingTerm("a"));
-  EXPECT_TRUE(trie.IsBranchingTerm("ab"));
-  EXPECT_TRUE(trie.IsBranchingTerm("abc"));
-  EXPECT_FALSE(trie.IsBranchingTerm("abc_s"));
-  EXPECT_FALSE(trie.IsBranchingTerm("abc_su"));
-  EXPECT_FALSE(trie.IsBranchingTerm("abc_suffi"));
+  EXPECT_FALSE(trie.IsBranchingTerm("").ValueOrDie());
+  EXPECT_TRUE(trie.IsBranchingTerm("a").ValueOrDie());
+  EXPECT_TRUE(trie.IsBranchingTerm("ab").ValueOrDie());
+  EXPECT_TRUE(trie.IsBranchingTerm("abc").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("abc_s").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("abc_su").ValueOrDie());
+  EXPECT_FALSE(trie.IsBranchingTerm("abc_suffi").ValueOrDie());
 }
 
 }  // namespace lib

diff --git a/icing/legacy/index/icing-filesystem.cc b/icing/legacy/index/icing-filesystem.cc
index 175e075..168361d 100644
--- a/icing/legacy/index/icing-filesystem.cc
+++ b/icing/legacy/index/icing-filesystem.cc

@@ -129,15 +129,21 @@
                            const std::unordered_set<std::string> &exclude,
                            bool recursive, const char *prefix,
                            std::vector<std::string> *entries) {
-  DIR *dir = opendir(dir_name);
-  if (!dir) {
+  auto closer = [](DIR *dir) {
+    if (closedir(dir) != 0) {
+      ICING_LOG(ERROR) << "Error closing dir (" << errno << ") "
+                       << strerror(errno);
+    }
+  };
+  std::unique_ptr<DIR, decltype(closer)> dir(opendir(dir_name), closer);
+  if (dir == nullptr) {
     LogOpenError("Unable to open directory ", dir_name, ": ", errno);
     return false;
   }
 
   dirent *p;
   // readdir's implementation seems to be thread safe.
-  while ((p = readdir(dir)) != nullptr) {
+  while ((p = readdir(dir.get())) != nullptr) {
     std::string file_name(p->d_name);
     if (file_name == "." || file_name == ".." ||
         exclude.find(file_name) != exclude.end()) {
@@ -156,9 +162,6 @@
       }
     }
   }
-  if (closedir(dir) != 0) {
-    ICING_LOG(ERROR) << "Error closing " << dir_name << ": " << strerror(errno);
-  }
   return true;
 }
 
@@ -445,9 +448,12 @@
 bool IcingFilesystem::Write(int fd, const void *data, size_t data_size) const {
   size_t write_len = data_size;
   do {
-    // Don't try to write too much at once.
-    size_t chunk_size = std::min<size_t>(write_len, 64u * 1024);
-    ssize_t wrote = write(fd, data, chunk_size);
+#ifdef __APPLE__
+    // TEMP_FAILURE_RETRY is not defined in unistd.h on iOS.
+    ssize_t wrote = write(fd, data, write_len);
+#else  // __APPLE__
+    ssize_t wrote = TEMP_FAILURE_RETRY(write(fd, data, write_len));
+#endif // __APPLE__
     if (wrote < 0) {
       ICING_LOG(ERROR) << "Bad write: " << strerror(errno);
       return false;
@@ -462,9 +468,12 @@
                              size_t data_size) const {
   size_t write_len = data_size;
   do {
-    // Don't try to write too much at once.
-    size_t chunk_size = std::min<size_t>(write_len, 64u * 1024);
-    ssize_t wrote = pwrite(fd, data, chunk_size, offset);
+#ifdef __APPLE__
+    // TEMP_FAILURE_RETRY is not defined in unistd.h on iOS.
+    ssize_t wrote = pwrite(fd, data, write_len, offset);
+#else  // __APPLE__
+    ssize_t wrote = TEMP_FAILURE_RETRY(pwrite(fd, data, write_len, offset));
+#endif // __APPLE__
     if (wrote < 0) {
       ICING_LOG(ERROR) << "Bad write: " << strerror(errno);
       return false;

diff --git a/icing/monkey_test/icing-monkey-test-runner.cc b/icing/monkey_test/icing-monkey-test-runner.cc
index 14b4e02..c16a5d9 100644
--- a/icing/monkey_test/icing-monkey-test-runner.cc
+++ b/icing/monkey_test/icing-monkey-test-runner.cc

@@ -20,6 +20,7 @@
 #include <functional>
 #include <iomanip>
 #include <ios>
+#include <limits>
 #include <memory>
 #include <random>
 #include <sstream>
@@ -30,7 +31,6 @@
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include "icing/absl_ports/str_cat.h"
-#include "icing/file/destructible-directory.h"
 #include "icing/icing-search-engine.h"
 #include "icing/monkey_test/in-memory-icing-search-engine.h"
 #include "icing/monkey_test/monkey-test-generators.h"
@@ -67,6 +67,11 @@
   return dist(*random) == 1;
 }
 
+int GetRandomInt(MonkeyTestRandomEngine* random, int min, int max) {
+  std::uniform_int_distribution<> dist(min, max);
+  return dist(*random);
+}
+
 SearchSpecProto GenerateRandomSearchSpecProto(
     MonkeyTestRandomEngine* random,
     MonkeyDocumentGenerator* document_generator) {
@@ -118,6 +123,15 @@
           query);
     }
   }
+  // %50 chance of getting one type filter
+  // %25 chance of getting two type filters
+  // %25 chance of getting no type filters
+  for (int i = 0; i < 2; ++i) {
+    if (GetRandomBoolean(random)) {
+      search_spec.add_schema_type_filters(
+          document_generator->GetType().schema_type());
+    }
+  }
   search_spec.set_query(query);
   return search_spec;
 }
@@ -215,9 +229,7 @@
       schema_generator_(
           std::make_unique<MonkeySchemaGenerator>(&random_, &config_)) {
   ICING_LOG(INFO) << "Monkey test runner started with seed: " << config_.seed;
-  std::string dir = GetTestTempDir() + "/icing/monkey";
-  filesystem_.DeleteDirectoryRecursively(dir.c_str());
-  icing_dir_ = std::make_unique<DestructibleDirectory>(&filesystem_, dir);
+  icing_dir_ = GetTestTempDir() + "/icing/monkey";
 }
 
 void IcingMonkeyTestRunner::Run(uint32_t num) {
@@ -230,7 +242,7 @@
     frequency_sum += schedule.second;
   }
   std::uniform_int_distribution<> dist(0, frequency_sum - 1);
-  for (; num; --num) {
+  for (uint32_t i = 0; i < num; ++i) {
     int p = dist(random_);
     for (const auto& schedule : config_.monkey_api_schedules) {
       if (p < schedule.second) {
@@ -239,7 +251,7 @@
       }
       p -= schedule.second;
     }
-    ICING_LOG(INFO) << "Completed Run #" << num
+    ICING_LOG(INFO) << "Completed Run #" << i
                     << ". Documents in the in-memory icing: "
                     << in_memory_icing_->GetNumAliveDocuments();
   }
@@ -254,12 +266,20 @@
 }
 
 void IcingMonkeyTestRunner::Initialize() {
-  ASSERT_NO_FATAL_FAILURE(CreateIcingSearchEngine());
+  if (config_.initialize_by_existing_data) {
+    ICING_LOG(INFO) << "Initializing icing by existing data";
 
-  SchemaProto schema = schema_generator_->GenerateSchema();
-  ICING_LOG(DBG) << "Schema Generated: " << schema.DebugString();
+    ASSERT_NO_FATAL_FAILURE(CreateIcingSearchEngine());
+    ASSERT_NO_FATAL_FAILURE(ReloadInMemoryIcing());
+  } else {
+    ICING_LOG(INFO) << "Initializing icing by empty data";
 
-  ASSERT_THAT(SetSchema(std::move(schema)).status(), ProtoIsOk());
+    filesystem_.DeleteDirectoryRecursively(icing_dir_.c_str());
+    ASSERT_NO_FATAL_FAILURE(CreateIcingSearchEngine());
+    SchemaProto schema = schema_generator_->GenerateSchema();
+    ICING_LOG(DBG) << "Schema Generated: " << schema.DebugString();
+    ASSERT_THAT(SetSchema(std::move(schema)).status(), ProtoIsOk());
+  }
 }
 
 void IcingMonkeyTestRunner::DoUpdateSchema() {
@@ -544,7 +564,7 @@
 
   IcingSearchEngineOptions icing_options;
   icing_options.set_index_merge_size(config_.index_merge_size);
-  icing_options.set_base_dir(icing_dir_->dir());
+  icing_options.set_base_dir(icing_dir_);
   icing_options.set_optimize_rebuild_index_threshold(
       optimize_rebuild_index_threshold);
   // The method will be called every time when we ReloadFromDisk(), so randomly
@@ -552,10 +572,68 @@
   icing_options.set_document_store_namespace_id_fingerprint(
       GetRandomBoolean(&random_));
   icing_options.set_enable_embedding_index(true);
-  icing_options.set_enable_embedding_quantization(GetRandomBoolean(&random_));
+  icing_options.set_enable_embedding_quantization(true);
+  icing_options.set_compression_threshold_bytes(
+      GetRandomInt(&random_, /*min=*/0, /*max=*/10000));
+  icing_options.set_enable_eigen_embedding_scoring(GetRandomBoolean(&random_));
+  icing_options.set_enable_passing_filter_to_children(
+      GetRandomBoolean(&random_));
+  icing_options.set_enable_embedding_iterator_v2(GetRandomBoolean(&random_));
+  // Randomly choose the number of shards from 1, 2, 4, 8, 16, 32.
+  uint32_t num_shards = 1 << GetRandomInt(&random_, /*min=*/0, /*max=*/5);
+  icing_options.set_embedding_index_num_shards(num_shards);
+  icing_options.set_enable_schema_database(GetRandomBoolean(&random_));
+  icing_options.set_enable_schema_type_id_optimization(
+      GetRandomBoolean(&random_));
+  icing_options.set_enable_skip_set_schema_type_equality_check(
+      GetRandomBoolean(&random_));
+  icing_options.set_enable_embed_query_optimization(GetRandomBoolean(&random_));
   icing_ = std::make_unique<IcingSearchEngine>(icing_options);
   ASSERT_THAT(icing_->Initialize().status(), ProtoIsOk());
 }
 
+void IcingMonkeyTestRunner::ReloadInMemoryIcing() {
+  // Reload schema
+  GetSchemaResultProto get_schema_result = icing_->GetSchema();
+  ASSERT_THAT(get_schema_result.status(), ProtoIsOk());
+  in_memory_icing_->SetSchema(get_schema_result.schema());
+  ASSERT_THAT(*in_memory_icing_->GetSchema(),
+              EqualsProto(get_schema_result.schema()));
+  document_generator_ = std::make_unique<MonkeyDocumentGenerator>(
+      &random_, in_memory_icing_->GetSchema(), &config_);
+
+  // Reload documents
+  SearchSpecProto search_spec;
+  search_spec.set_term_match_type(TermMatchType::PREFIX);
+  search_spec.set_query("");
+  ResultSpecProto result_spec;
+  result_spec.set_num_to_score(std::numeric_limits<int32_t>::max());
+  result_spec.set_num_per_page(100);
+  int num_results = 0;
+  int max_uri = 0;
+  SearchResultProto search_result = icing_->Search(
+      search_spec, ScoringSpecProto::default_instance(), result_spec);
+  ASSERT_THAT(search_result.status(), ProtoIsOk());
+  while (true) {
+    num_results += search_result.results_size();
+    for (const SearchResultProto::ResultProto& doc : search_result.results()) {
+      in_memory_icing_->Put(MonkeyTokenizedDocument::Reload(doc.document()));
+      max_uri = std::max(
+          max_uri, std::stoi(doc.document().uri().substr(
+                       MonkeyDocumentGenerator::kDocumentUriPrefix.size())));
+    }
+    if (search_result.next_page_token() == kInvalidNextPageToken) {
+      break;
+    }
+    search_result = icing_->GetNextPage(search_result.next_page_token());
+    ASSERT_THAT(search_result.status(), ProtoIsOk());
+  }
+  ICING_LOG(INFO) << "Reloaded " << num_results << " documents";
+
+  // Reload generators
+  schema_generator_->ReloadPreviousStatus(*in_memory_icing_->GetSchema());
+  document_generator_->ReloadPreviousStatus(max_uri);
+}
+
 }  // namespace lib
 }  // namespace icing

diff --git a/icing/monkey_test/icing-monkey-test-runner.h b/icing/monkey_test/icing-monkey-test-runner.h
index 10be60c..dcb81f4 100644
--- a/icing/monkey_test/icing-monkey-test-runner.h
+++ b/icing/monkey_test/icing-monkey-test-runner.h

@@ -17,8 +17,8 @@
 
 #include <cstdint>
 #include <memory>
+#include <string>
 
-#include "icing/file/destructible-directory.h"
 #include "icing/file/filesystem.h"
 #include "icing/icing-search-engine.h"
 #include "icing/monkey_test/in-memory-icing-search-engine.h"
@@ -63,7 +63,7 @@
   IcingMonkeyTestRunnerConfiguration config_;
   MonkeyTestRandomEngine random_;
   Filesystem filesystem_;
-  std::unique_ptr<DestructibleDirectory> icing_dir_;
+  std::string icing_dir_;
   std::unique_ptr<InMemoryIcingSearchEngine> in_memory_icing_;
   std::unique_ptr<IcingSearchEngine> icing_;
 
@@ -71,6 +71,8 @@
   std::unique_ptr<MonkeyDocumentGenerator> document_generator_;
 
   void CreateIcingSearchEngine();
+  // Reloads the in-memory icing from the disk.
+  void ReloadInMemoryIcing();
 };
 
 }  // namespace lib

diff --git a/icing/monkey_test/icing-search-engine_monkey_crash_simulation_test.cc b/icing/monkey_test/icing-search-engine_monkey_crash_simulation_test.cc
new file mode 100644
index 0000000..f5547cf
--- /dev/null
+++ b/icing/monkey_test/icing-search-engine_monkey_crash_simulation_test.cc

@@ -0,0 +1,145 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <signal.h>
+
+#include <cstdint>
+#include <limits>
+#include <memory>
+#include <random>
+
+#include "gtest/gtest.h"
+#include "icing/monkey_test/icing-monkey-test-runner.h"
+#include "icing/monkey_test/monkey-test-util.h"
+#include "icing/portable/platform.h"
+#include "icing/proto/debug.pb.h"
+#include "icing/schema/section.h"
+#include "icing/util/logging.h"
+
+namespace icing {
+namespace lib {
+
+constexpr uint32_t kNumCrashesToSimulate = 5;
+constexpr uint32_t kMinRunningTimePerCrashSeconds = 5;
+constexpr uint32_t kMaxRunningTimePerCrashSeconds = 15;
+
+void RunMonkeyTest(uint32_t seed, bool is_first_run, uint32_t num_iterations) {
+  IcingMonkeyTestRunnerConfiguration config(
+      seed,
+      /*num_types=*/30,
+      /*num_namespaces=*/100,
+      /*num_uris=*/1000,
+      /*index_merge_size=*/1024 * 1024,
+      /*initialize_by_existing_data=*/!is_first_run);
+  config.possible_num_properties = {0,
+                                    1,
+                                    2,
+                                    4,
+                                    8,
+                                    16,
+                                    kTotalNumSections / 2,
+                                    kTotalNumSections,
+                                    kTotalNumSections + 1,
+                                    kTotalNumSections * 2};
+  config.possible_num_tokens = {0, 1, 4, 16, 64, 256};
+  config.possible_num_vectors = {0, 1, 4};
+  config.possible_vector_dimensions = {128, 512, 768};
+  config.monkey_api_schedules = {
+      {&IcingMonkeyTestRunner::DoPut, 500},
+      {&IcingMonkeyTestRunner::DoSearch, 200},
+      {&IcingMonkeyTestRunner::DoGet, 70},
+      {&IcingMonkeyTestRunner::DoGetAllNamespaces, 50},
+      {&IcingMonkeyTestRunner::DoDelete, 50},
+      {&IcingMonkeyTestRunner::DoDeleteByNamespace, 50},
+      {&IcingMonkeyTestRunner::DoDeleteBySchemaType, 45},
+      {&IcingMonkeyTestRunner::DoDeleteByQuery, 20},
+      {&IcingMonkeyTestRunner::DoOptimize, 5},
+      {&IcingMonkeyTestRunner::DoUpdateSchema, 5},
+      {&IcingMonkeyTestRunner::ReloadFromDisk, 5}};
+
+  std::unique_ptr<IcingMonkeyTestRunner> runner =
+      std::make_unique<IcingMonkeyTestRunner>(config);
+  ASSERT_NO_FATAL_FAILURE(runner->Initialize());
+  ASSERT_NO_FATAL_FAILURE(runner->Run(num_iterations));
+  runner.reset();
+}
+
+TEST(IcingSearchEngineMonkeyCrashSimulationTest, MonkeyTest) {
+  uint32_t seed = std::random_device()();
+  MonkeyTestRandomEngine random(seed);
+  ICING_LOG(INFO) << "Monkey test crash simulation started with seed: " << seed;
+
+  bool is_first_run = true;
+  for (int i = 0; i < kNumCrashesToSimulate; ++i) {
+    uint32_t child_seed = random();
+    pid_t pid = fork();
+    ASSERT_TRUE(pid >= 0) << "fork() failed!";
+
+    // Child process
+    if (pid == 0) {
+      // Run indefinitely until the parent process kills the child process.
+      ICING_LOG(INFO) << "Child process " << i << " started running";
+      ASSERT_NO_FATAL_FAILURE(
+          RunMonkeyTest(child_seed, is_first_run, /*num_iterations=*/
+                        std::numeric_limits<uint32_t>::max()));
+      _exit(0);
+    }
+
+    // Randomly generate the running time for each crash.
+    std::uniform_int_distribution<uint32_t>
+        running_time_per_crash_seconds_distribution(
+            kMinRunningTimePerCrashSeconds, kMaxRunningTimePerCrashSeconds);
+    sleep(running_time_per_crash_seconds_distribution(random));
+
+    // Check if the child process has exited.
+    int status = 0;
+    pid_t ret = waitpid(pid, &status, WNOHANG);
+    if (ret == pid) {
+      // The child process has already exited for some reason, which is not
+      // expected. We should check the reason and fail the overall test.
+      if (WIFEXITED(status)) {
+        int code = WEXITSTATUS(status);
+        FAIL() << "Monkey test in the child process exited with code " << code;
+      } else if (WIFSIGNALED(status)) {
+        FAIL()
+            << "Monkey test in the child process was signaled early with sig="
+            << WTERMSIG(status);
+      }
+
+      FAIL() << "Monkey test in the child process ended in unknown reason";
+    } else if (ret != 0) {
+      FAIL() << "waitpid() error!";
+    }
+
+    // The child process is still running, which is expected. We can now kill
+    // the child process to simulate a crash.
+    if (kill(pid, SIGKILL) != 0) {
+      FAIL() << "Failed to kill child process " << i;
+    }
+    waitpid(pid, &status, 0);
+    ICING_LOG(INFO) << "Child process killed.";
+
+    is_first_run = false;
+  }
+
+  // Run the monkey test after the crash. Icing search engine should be able to
+  // recover from the crash and continue running normally.
+  ICING_LOG(INFO) << "Running final monkey test after crash simulation.";
+  uint32_t num_iterations = IsAndroidArm() ? 200 : 1000;
+  ASSERT_NO_FATAL_FAILURE(
+      RunMonkeyTest(/*seed=*/random(), is_first_run, num_iterations));
+}
+
+}  // namespace lib
+}  // namespace icing

diff --git a/icing/monkey_test/icing-search-engine_monkey_test.cc b/icing/monkey_test/icing-search-engine_monkey_test.cc
index 6d87ba9..7764dd2 100644
--- a/icing/monkey_test/icing-search-engine_monkey_test.cc
+++ b/icing/monkey_test/icing-search-engine_monkey_test.cc

@@ -33,7 +33,8 @@
       /*num_types=*/30,
       /*num_namespaces=*/100,
       /*num_uris=*/1000,
-      /*index_merge_size=*/1024 * 1024);
+      /*index_merge_size=*/1024 * 1024,
+      /*initialize_by_existing_data=*/false);
   config.possible_num_properties = {0,
                                     1,
                                     2,
@@ -45,8 +46,8 @@
                                     kTotalNumSections + 1,
                                     kTotalNumSections * 2};
   config.possible_num_tokens = {0, 1, 4, 16, 64, 256};
-  config.possible_num_vectors = {0, 1, 4};
-  config.possible_vector_dimensions = {128, 512, 768};
+  config.possible_num_vectors = {0, 1, 4, 8};
+  config.possible_vector_dimensions = {8, 16, 128, 512, 768};
   config.monkey_api_schedules = {
       {&IcingMonkeyTestRunner::DoPut, 500},
       {&IcingMonkeyTestRunner::DoSearch, 200},
@@ -71,7 +72,8 @@
       /*num_types=*/30,
       /*num_namespaces=*/200,
       /*num_uris=*/100000,
-      /*index_merge_size=*/1024 * 1024);
+      /*index_merge_size=*/1024 * 1024,
+      /*initialize_by_existing_data=*/false);
 
   // Due to the large amount of documents, we need to make each document smaller
   // to finish the test.

diff --git a/icing/monkey_test/in-memory-icing-search-engine.cc b/icing/monkey_test/in-memory-icing-search-engine.cc
index 417153c..262ba32 100644
--- a/icing/monkey_test/in-memory-icing-search-engine.cc
+++ b/icing/monkey_test/in-memory-icing-search-engine.cc

@@ -30,6 +30,7 @@
 #include "icing/absl_ports/str_cat.h"
 #include "icing/absl_ports/str_join.h"
 #include "icing/index/embed/embedding-scorer.h"
+#include "icing/index/embed/quantizer.h"
 #include "icing/monkey_test/monkey-tokenized-document.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/schema.pb.h"
@@ -77,16 +78,37 @@
   return std::make_pair(values[0], values[1]);
 }
 
-bool DoesVectorsMatch(const EmbeddingScorer *scorer,
-                      std::pair<double, double> embedding_search_range,
-                      const PropertyProto::VectorProto &vector1,
-                      const PropertyProto::VectorProto &vector2) {
-  if (vector1.model_signature() != vector2.model_signature() ||
-      vector1.values_size() != vector2.values_size()) {
+bool DoesVectorsMatch(
+    const EmbeddingScorer* scorer,
+    std::pair<double, double> embedding_search_range,
+    EmbeddingIndexingConfig::QuantizationType::Code quantization_type,
+    const PropertyProto::VectorProto& query,
+    const PropertyProto::VectorProto& candidate) {
+  if (query.model_signature() != candidate.model_signature() ||
+      query.values_size() != candidate.values_size()) {
     return false;
   }
-  float score = scorer->Score(vector1.values_size(), vector1.values().data(),
-                              vector2.values().data());
+
+  const int dimension = query.values_size();
+  float score;
+  if (quantization_type == EmbeddingIndexingConfig::QuantizationType::NONE) {
+    score = scorer->Score(dimension, query.values().data(),
+                          candidate.values().data());
+  } else {
+    // Quantize the candidate vector.
+    auto minmax_pair = std::minmax_element(candidate.values().begin(),
+                                           candidate.values().end());
+    Quantizer quantizer =
+        Quantizer::Create(*minmax_pair.first, *minmax_pair.second).ValueOrDie();
+    std::vector<uint8_t> quantized_candidate;
+    quantized_candidate.reserve(dimension);
+    for (float value : candidate.values()) {
+      quantized_candidate.push_back(quantizer.Quantize(value));
+    }
+    // Score the quantized candidate against the original query.
+    score = scorer->Score(dimension, query.values().data(),
+                          quantized_candidate.data(), quantizer);
+  }
   return embedding_search_range.first <= score &&
          score <= embedding_search_range.second;
 }
@@ -137,7 +159,10 @@
       bool indexable =
           prop->embedding_indexing_config().embedding_indexing_type() !=
           EmbeddingIndexingConfig::EmbeddingIndexingType::UNKNOWN;
-      return PropertyIndexInfo{indexable};
+      EmbeddingIndexingConfig::QuantizationType::Code quantization_type =
+          prop->embedding_indexing_config().quantization_type();
+      return PropertyIndexInfo{indexable, TermMatchType::UNKNOWN,
+                               quantization_type};
     }
 
     if (prop->data_type() != PropertyConfigProto::DataType::DOCUMENT) {
@@ -174,6 +199,14 @@
 InMemoryIcingSearchEngine::DoesDocumentMatchQuery(
     const MonkeyTokenizedDocument &document,
     const SearchSpecProto &search_spec) const {
+  // Check schema type filter.
+  const auto &schema_type_filters = search_spec.schema_type_filters();
+  if (!schema_type_filters.empty() &&
+      std::find(schema_type_filters.begin(), schema_type_filters.end(),
+                document.document.schema()) == schema_type_filters.end()) {
+    return false;
+  }
+
   std::string_view query = search_spec.query();
   std::vector<std::string_view> strs = absl_ports::StrSplit(query, ":");
   std::string_view section_restrict;
@@ -210,6 +243,7 @@
            section.embedding_vectors) {
         if (DoesVectorsMatch(embedding_scorer.get(),
                              embedding_search_range_or.ValueOrDie(),
+                             property_index_info.quantization_type,
                              search_spec.embedding_query_vectors(0), vector)) {
           return true;
         }
@@ -231,7 +265,7 @@
   return false;
 }
 
-void InMemoryIcingSearchEngine::SetSchema(SchemaProto &&schema) {
+void InMemoryIcingSearchEngine::SetSchema(SchemaProto schema) {
   schema_ = std::make_unique<SchemaProto>(std::move(schema));
   property_config_map_.clear();
   for (const SchemaTypeConfigProto &type_config : schema_->types()) {

diff --git a/icing/monkey_test/in-memory-icing-search-engine.h b/icing/monkey_test/in-memory-icing-search-engine.h
index 8e92f10..c8a4b1a 100644
--- a/icing/monkey_test/in-memory-icing-search-engine.h
+++ b/icing/monkey_test/in-memory-icing-search-engine.h

@@ -52,7 +52,7 @@
 
   const SchemaProto *GetSchema() const { return schema_.get(); }
 
-  void SetSchema(SchemaProto &&schema);
+  void SetSchema(SchemaProto schema);
 
   // Randomly pick a document from the in-memory Icing for monkey testing.
   //
@@ -164,6 +164,9 @@
     // The term match type if the property is of type string.
     TermMatchType::Code term_match_type =
         TermMatchType::Code::TermMatchType_Code_UNKNOWN;
+    // The quantization type if the property is of type vector.
+    EmbeddingIndexingConfig::QuantizationType::Code quantization_type =
+        EmbeddingIndexingConfig::QuantizationType::NONE;
   };
   libtextclassifier3::StatusOr<PropertyIndexInfo> GetPropertyIndexInfo(
       const std::string &schema_type,

diff --git a/icing/monkey_test/monkey-test-generators.cc b/icing/monkey_test/monkey-test-generators.cc
index 9f98d15..4dd6a41 100644
--- a/icing/monkey_test/monkey-test-generators.cc
+++ b/icing/monkey_test/monkey-test-generators.cc

@@ -14,6 +14,7 @@
 
 #include "icing/monkey_test/monkey-test-generators.h"
 
+#include <algorithm>
 #include <array>
 #include <cstdint>
 #include <random>
@@ -92,6 +93,10 @@
     property.mutable_embedding_indexing_config()->set_embedding_indexing_type(
         EmbeddingIndexingConfig::EmbeddingIndexingType::LINEAR_SEARCH);
   }
+  if (GetRandomBoolean(random)) {
+    property.mutable_embedding_indexing_config()->set_quantization_type(
+        EmbeddingIndexingConfig::QuantizationType::QUANTIZE_8_BIT);
+  }
 }
 
 }  // namespace
@@ -142,12 +147,34 @@
   return result;
 }
 
+void MonkeySchemaGenerator::ReloadPreviousStatus(const SchemaProto& schema) {
+  int max_schema_id = 0;
+  for (const SchemaTypeConfigProto& type_config : schema.types()) {
+    max_schema_id =
+        std::max(max_schema_id, std::stoi(type_config.schema_type().substr(
+                                    kSchemaTypeNamePrefix.size())));
+
+    // To reset num_properties_generated_ according to the previous run, we use
+    // the maximum property_id + 1 as an estimate.
+    int max_property_id = 0;
+    for (const PropertyConfigProto& property : type_config.properties()) {
+      max_property_id =
+          std::max(max_property_id, std::stoi(property.property_name().substr(
+                                        kSchemaPropertyNamePrefix.size())));
+    }
+    num_properties_generated_[type_config.schema_type()] = max_property_id + 1;
+  }
+  // To reset num_types_generated_ according to the previous run, we use the
+  // maximum schema_id + 1 as an estimate.
+  num_types_generated_ = max_schema_id + 1;
+}
+
 PropertyConfigProto MonkeySchemaGenerator::GenerateProperty(
     const SchemaTypeConfigProto& type_config,
     PropertyConfigProto::Cardinality::Code cardinality, bool indexable) {
   PropertyConfigProto prop;
   prop.set_property_name(
-      "MonkeyTestProp" +
+      std::string(kSchemaPropertyNamePrefix) +
       std::to_string(num_properties_generated_[type_config.schema_type()]++));
   // TODO: Perhaps in future iterations we will want to generate more types of
   // properties.
@@ -201,7 +228,14 @@
       index_incompatible = true;
     }
   } else if (property.data_type() == PropertyConfigProto::DataType::VECTOR) {
+    EmbeddingIndexingConfig::QuantizationType::Code old_quantization_type =
+        property.embedding_indexing_config().quantization_type();
     SetEmbeddingIndexingConfig(random_, property, new_indexable);
+    EmbeddingIndexingConfig::QuantizationType::Code new_quantization_type =
+        property.embedding_indexing_config().quantization_type();
+    if (old_quantization_type != new_quantization_type) {
+      index_incompatible = true;
+    }
   }
   if (index_incompatible) {
     result.schema_types_index_incompatible.insert(type_config.schema_type());
@@ -210,7 +244,7 @@
 
 SchemaTypeConfigProto MonkeySchemaGenerator::GenerateType() {
   SchemaTypeConfigProto type_config;
-  type_config.set_schema_type("MonkeyTestType" +
+  type_config.set_schema_type(std::string(kSchemaTypeNamePrefix) +
                               std::to_string(num_types_generated_++));
   std::uniform_int_distribution<> possible_num_properties_dist(
       0, config_->possible_num_properties.size() - 1);
@@ -325,7 +359,7 @@
     std::uniform_int_distribution<> dist(0, config_->num_uris - 1);
     uri = dist(*random_);
   }
-  return absl_ports::StrCat("uri", std::to_string(uri));
+  return absl_ports::StrCat(kDocumentUriPrefix, std::to_string(uri));
 }
 
 int MonkeyDocumentGenerator::GetNumTokens() const {

diff --git a/icing/monkey_test/monkey-test-generators.h b/icing/monkey_test/monkey-test-generators.h
index 49ad530..5834022 100644
--- a/icing/monkey_test/monkey-test-generators.h
+++ b/icing/monkey_test/monkey-test-generators.h

@@ -35,6 +35,10 @@
 // A random schema generator used for monkey testing.
 class MonkeySchemaGenerator {
  public:
+  static constexpr std::string_view kSchemaTypeNamePrefix = "MonkeyTestType";
+  static constexpr std::string_view kSchemaPropertyNamePrefix =
+      "MonkeyTestProp";
+
   struct UpdateSchemaResult {
     SchemaProto schema;
     bool is_invalid_schema;
@@ -52,6 +56,9 @@
 
   UpdateSchemaResult UpdateSchema(const SchemaProto& schema);
 
+  // Reload the previous status of the schema generator.
+  void ReloadPreviousStatus(const SchemaProto& schema);
+
  private:
   PropertyConfigProto GenerateProperty(
       const SchemaTypeConfigProto& type_config,
@@ -81,6 +88,8 @@
 // Same for num_namespaces.
 class MonkeyDocumentGenerator {
  public:
+  static constexpr std::string_view kDocumentUriPrefix = "uri";
+
   explicit MonkeyDocumentGenerator(
       MonkeyTestRandomEngine* random, const SchemaProto* schema,
       const IcingMonkeyTestRunnerConfiguration* config)
@@ -117,6 +126,13 @@
 
   MonkeyTokenizedDocument GenerateDocument();
 
+  // Reload the previous status of the document generator.
+  void ReloadPreviousStatus(uint32_t max_uri) {
+    // To reset num_docs_generated_ according to the previous run, we use the
+    // maximum uri + 1 as an estimate.
+    num_docs_generated_ = max_uri + 1;
+  }
+
  private:
   MonkeyTestRandomEngine* random_;                    // Does not own.
   const SchemaProto* schema_;                         // Does not own.

diff --git a/icing/monkey_test/monkey-test-util.h b/icing/monkey_test/monkey-test-util.h
index 919640c..2a6f002 100644
--- a/icing/monkey_test/monkey-test-util.h
+++ b/icing/monkey_test/monkey-test-util.h

@@ -31,12 +31,14 @@
 struct IcingMonkeyTestRunnerConfiguration {
   explicit IcingMonkeyTestRunnerConfiguration(uint32_t seed, int num_types,
                                               int num_namespaces, int num_uris,
-                                              int index_merge_size)
+                                              int index_merge_size,
+                                              bool initialize_by_existing_data)
       : seed(seed),
         num_types(num_types),
         num_namespaces(num_namespaces),
         num_uris(num_uris),
-        index_merge_size(index_merge_size) {}
+        index_merge_size(index_merge_size),
+        initialize_by_existing_data(initialize_by_existing_data) {}
 
   uint32_t seed;
   int num_types;
@@ -44,6 +46,12 @@
   int num_uris;
   int index_merge_size;
 
+  // Whether to initialize Icing with existing data. If true, the test will
+  // start from the state of the existing Icing testing data stored in
+  // GetTestTempDir() + "/icing/monkey". Otherwise, the test will start with an
+  // empty Icing.
+  bool initialize_by_existing_data;
+
   // To ensure that the random schema is generated with the best quality, the
   // number of properties for each type will only be randomly picked from this
   // list, instead of picking it from a range. For example, a vector of

diff --git a/icing/monkey_test/monkey-tokenized-document.h b/icing/monkey_test/monkey-tokenized-document.h
index cad838e..ce356c1 100644
--- a/icing/monkey_test/monkey-tokenized-document.h
+++ b/icing/monkey_test/monkey-tokenized-document.h

@@ -16,13 +16,21 @@
 #define ICING_MONKEY_TEST_MONKEY_TOKENIZED_DOCUMENT_H_
 
 #include <string>
+#include <string_view>
+#include <utility>
 #include <vector>
 
+#include "icing/absl_ports/str_cat.h"
+#include "icing/absl_ports/str_join.h"
 #include "icing/proto/document.pb.h"
 
 namespace icing {
 namespace lib {
 
+// No matter whether the property is indexable, we have to create a section for
+// it since a non-indexable property can become indexable after a schema type
+// change. The in-memory icing will automatically skip sections that are
+// non-indexable at the time of search requests.
 struct MonkeyTokenizedSection {
   std::string path;
   std::vector<std::string> token_sequence;
@@ -32,6 +40,50 @@
 struct MonkeyTokenizedDocument {
   DocumentProto document;
   std::vector<MonkeyTokenizedSection> tokenized_sections;
+
+  static MonkeyTokenizedDocument Reload(DocumentProto document) {
+    MonkeyTokenizedDocument tokenized_document;
+    tokenized_document.document = std::move(document);
+    ExtractTokenizedSections(tokenized_document.document, /*curr_path=*/"",
+                             tokenized_document.tokenized_sections);
+    return tokenized_document;
+  }
+
+ private:
+  static void ExtractTokenizedSections(
+      const DocumentProto& document, std::string curr_path,
+      std::vector<MonkeyTokenizedSection>& tokenized_sections) {
+    for (const PropertyProto& property : document.properties()) {
+      std::string new_path =
+          curr_path.empty()
+              ? property.name()
+              : absl_ports::StrCat(curr_path, ".", property.name());
+      if (property.string_values_size() > 0) {
+        std::vector<std::string> token_sequence;
+        for (const std::string& value : property.string_values()) {
+          for (std::string_view token : absl_ports::StrSplit(value, " ")) {
+            token_sequence.push_back(std::string(token));
+          }
+        }
+        tokenized_sections.push_back(MonkeyTokenizedSection{
+            .path = new_path, .token_sequence = std::move(token_sequence)});
+      } else if (property.vector_values_size() > 0) {
+        std::vector<PropertyProto::VectorProto> embedding_vectors;
+        for (const PropertyProto::VectorProto& vector :
+             property.vector_values()) {
+          embedding_vectors.push_back(vector);
+        }
+        tokenized_sections.push_back(MonkeyTokenizedSection{
+            .path = new_path,
+            .embedding_vectors = std::move(embedding_vectors)});
+      } else if (property.document_values_size() > 0) {
+        for (const DocumentProto& document_value : property.document_values()) {
+          ExtractTokenizedSections(document_value, new_path,
+                                   tokenized_sections);
+        }
+      }
+    }
+  }
 };
 
 }  // namespace lib

diff --git a/icing/portable/gzip_stream.cc b/icing/portable/gzip_stream.cc
index f00a993..ea5094a 100644
--- a/icing/portable/gzip_stream.cc
+++ b/icing/portable/gzip_stream.cc

@@ -18,17 +18,23 @@
 // smaller libprotobuf-lite instead.
 
 #include "icing/portable/gzip_stream.h"
+
+#include <cstddef>
+
 #include "icing/util/logging.h"
 
 namespace icing {
 namespace lib {
 namespace protobuf_ports {
 
-static const int kDefaultBufferSize = 65536;
-
 GzipInputStream::GzipInputStream(ZeroCopyInputStream* sub_stream, Format format,
-                                 int buffer_size)
-    : format_(format), sub_stream_(sub_stream), zerror_(Z_OK), byte_count_(0) {
+                                 void* output_buffer, size_t buffer_size)
+    : format_(format),
+      sub_stream_(sub_stream),
+      zerror_(Z_OK),
+      output_buffer_(output_buffer),
+      output_buffer_length_(buffer_size),
+      byte_count_(0) {
   zcontext_.state = Z_NULL;
   zcontext_.zalloc = Z_NULL;
   zcontext_.zfree = Z_NULL;
@@ -38,18 +44,12 @@
   zcontext_.avail_in = 0;
   zcontext_.total_in = 0;
   zcontext_.msg = NULL;
-  if (buffer_size == -1) {
-    output_buffer_length_ = kDefaultBufferSize;
-  } else {
-    output_buffer_length_ = buffer_size;
-  }
-  output_buffer_ = operator new(output_buffer_length_);
   zcontext_.next_out = static_cast<Bytef*>(output_buffer_);
   zcontext_.avail_out = output_buffer_length_;
   output_position_ = output_buffer_;
 }
+
 GzipInputStream::~GzipInputStream() {
-  operator delete(output_buffer_);
   zerror_ = inflateEnd(&zcontext_);
 }
 
@@ -179,7 +179,8 @@
     : format(GZIP),
       buffer_size(kDefaultBufferSize),
       compression_level(Z_DEFAULT_COMPRESSION),
-      compression_strategy(Z_DEFAULT_STRATEGY) {}
+      compression_strategy(Z_DEFAULT_STRATEGY),
+      mem_level(kDefaultMemLevel) {}
 
 GzipOutputStream::GzipOutputStream(ZeroCopyOutputStream* sub_stream) {
   Init(sub_stream, Options());
@@ -214,10 +215,10 @@
   if (options.format == ZLIB) {
     windowBitsFormat = 0;
   }
-  zerror_ =
-      deflateInit2(&zcontext_, options.compression_level, Z_DEFLATED,
-                   /* windowBits */ 15 | windowBitsFormat,
-                   /* memLevel (default) */ 8, options.compression_strategy);
+  zerror_ = deflateInit2(&zcontext_, options.compression_level, Z_DEFLATED,
+                         /* windowBits */ 15 | windowBitsFormat,
+                         /* memLevel */ options.mem_level,
+                         options.compression_strategy);
 }
 
 GzipOutputStream::~GzipOutputStream() {

diff --git a/icing/portable/gzip_stream.h b/icing/portable/gzip_stream.h
index 8008a55..a43f9e2 100644
--- a/icing/portable/gzip_stream.h
+++ b/icing/portable/gzip_stream.h

@@ -27,6 +27,9 @@
 #ifndef GOOGLE3_ICING_PORTABLE_GZIP_STREAM_H_
 #define GOOGLE3_ICING_PORTABLE_GZIP_STREAM_H_
 
+#include <cstddef>
+#include <cstdint>
+
 #include "icing/portable/zlib.h"
 #include <google/protobuf/io/zero_copy_stream_impl_lite.h>
 
@@ -34,6 +37,9 @@
 namespace lib {
 namespace protobuf_ports {
 
+static constexpr size_t kDefaultBufferSize = 64 * 1024;  // 64kb
+static constexpr int kDefaultMemLevel = 8;
+
 // A ZeroCopyInputStream that reads compressed data through zlib
 class GzipInputStream : public google::protobuf::io::ZeroCopyInputStream {
  public:
@@ -51,7 +57,8 @@
 
   // buffer_size and format may be -1 for default of 64kB and GZIP format
   explicit GzipInputStream(google::protobuf::io::ZeroCopyInputStream* sub_stream,
-                           Format format = AUTO, int buffer_size = -1);
+                           Format format, void* output_buffer,
+                           size_t buffer_size);
   virtual ~GzipInputStream();
 
   // Return last error message or NULL if no error.
@@ -67,12 +74,12 @@
  private:
   Format format_;
 
-  google::protobuf::io::ZeroCopyInputStream* sub_stream_;
+  google::protobuf::io::ZeroCopyInputStream* sub_stream_;  // Does not own.
 
   z_stream zcontext_;
   int zerror_;
 
-  void* output_buffer_;
+  void* output_buffer_;  // Does not own.
   void* output_position_;
   size_t output_buffer_length_;
   int64_t byte_count_;
@@ -108,6 +115,16 @@
     // zlib.h for definitions of these constants.
     int compression_strategy;
 
+    // Used for deflateInit2.
+    // A number between 1 and 9, specifying how much memory should be allocated
+    // for the internal compression state.
+    // memLevel=1 uses minimum memory but is slow and reduces compression ratio;
+    // memLevel=9 uses maximum memory for optimal speed.
+    // The default value is 8.
+    //
+    // See the documentation for deflateInit2 in zlib.h for more details.
+    int mem_level;
+
     Options();  // Initializes with default values.
   };
 

diff --git a/icing/query/advanced_query_parser/abstract-syntax-tree.h b/icing/query/advanced_query_parser/abstract-syntax-tree.h
index 3260084..071dd30 100644
--- a/icing/query/advanced_query_parser/abstract-syntax-tree.h
+++ b/icing/query/advanced_query_parser/abstract-syntax-tree.h

@@ -99,6 +99,26 @@
   bool is_prefix_;
 };
 
+class FunctionNode : public Node {
+ public:
+  explicit FunctionNode(std::string function_name)
+      : FunctionNode(std::move(function_name), {}) {}
+  explicit FunctionNode(std::string function_name,
+                        std::vector<std::unique_ptr<Node>> args)
+      : function_name_(std::move(function_name)),
+        args_(std::move(args)) {}
+
+  void Accept(AbstractSyntaxTreeVisitor* visitor) const override {
+    visitor->VisitFunction(this);
+  }
+  const std::string& function_name() const { return function_name_; }
+  const std::vector<std::unique_ptr<Node>>& args() const { return args_; }
+
+ private:
+  std::string function_name_;
+  std::vector<std::unique_ptr<Node>> args_;
+};
+
 class MemberNode : public Node {
  public:
   explicit MemberNode(std::vector<std::unique_ptr<TextNode>> children,
@@ -120,26 +140,6 @@
   std::unique_ptr<FunctionNode> function_;
 };
 
-class FunctionNode : public Node {
- public:
-  explicit FunctionNode(std::string function_name)
-      : FunctionNode(std::move(function_name), {}) {}
-  explicit FunctionNode(std::string function_name,
-                        std::vector<std::unique_ptr<Node>> args)
-      : function_name_(std::move(function_name)),
-        args_(std::move(args)) {}
-
-  void Accept(AbstractSyntaxTreeVisitor* visitor) const override {
-    visitor->VisitFunction(this);
-  }
-  const std::string& function_name() const { return function_name_; }
-  const std::vector<std::unique_ptr<Node>>& args() const { return args_; }
-
- private:
-  std::string function_name_;
-  std::vector<std::unique_ptr<Node>> args_;
-};
-
 class UnaryOperatorNode : public Node {
  public:
   explicit UnaryOperatorNode(std::string operator_text,

diff --git a/icing/query/advanced_query_parser/optimizer/query-optimization-util.cc b/icing/query/advanced_query_parser/optimizer/query-optimization-util.cc
new file mode 100644
index 0000000..6abc4b3
--- /dev/null
+++ b/icing/query/advanced_query_parser/optimizer/query-optimization-util.cc

@@ -0,0 +1,68 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "icing/query/advanced_query_parser/optimizer/query-optimization-util.h"
+
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "icing/feature-flags.h"
+#include "icing/index/embed/doc-hit-info-iterator-embedding-v2.h"
+#include "icing/index/iterator/doc-hit-info-iterator-and.h"
+#include "icing/index/iterator/doc-hit-info-iterator.h"
+
+namespace icing {
+namespace lib {
+namespace query_optimization_util {
+
+std::unique_ptr<DocHitInfoIterator> OptimizeAndIteratorsIfPossible(
+    std::vector<std::unique_ptr<DocHitInfoIterator>>&& iterators,
+    const FeatureFlags& feature_flags) {
+  std::unique_ptr<DocHitInfoIterator> embed_iterator;
+  bool delegate_node_is_right_most = true;
+  if (feature_flags.enable_embed_query_optimization()) {
+    // Find the first embedding iterator and remove it from the vector.
+    int embed_iterator_index = 0;
+    for (int i = embed_iterator_index; embed_iterator_index < iterators.size();
+         ++embed_iterator_index) {
+      if (iterators.at(i)->CanAdoptDelegate()) {
+        embed_iterator = std::move(iterators.at(i));
+        if (i == iterators.size() - 1) {
+          // If this embedding iterator is the last iterator, then the node is
+          // the right most node and the delegate that we're going to create
+          // would not be right most.
+          delegate_node_is_right_most = false;
+        }
+        iterators.erase(iterators.begin() + i);
+        break;
+      }
+    }
+  }
+  // If we found an embedding iterator, then put all other iterators into an
+  // AND iterator together and pass it to this embedding iterator as a delegate.
+  // Otherwise, just return the original AND iterator.
+  std::unique_ptr<DocHitInfoIterator> and_iterator =
+      CreateAndIterator(std::move(iterators));
+  if (embed_iterator == nullptr) {
+    return and_iterator;
+  }
+  embed_iterator->AdoptDelegate(std::move(and_iterator),
+                                    delegate_node_is_right_most);
+  return embed_iterator;
+}
+
+}  // namespace query_optimization_util
+}  // namespace lib
+}  // namespace icing

diff --git a/icing/query/advanced_query_parser/optimizer/query-optimization-util.h b/icing/query/advanced_query_parser/optimizer/query-optimization-util.h
new file mode 100644
index 0000000..c24a97c
--- /dev/null
+++ b/icing/query/advanced_query_parser/optimizer/query-optimization-util.h

@@ -0,0 +1,40 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef ICING_QUERY_ADVANCED_QUERY_PARSER_OPTIMIZER_QUERY_OPTIMIZATION_UTIL_H_
+#define ICING_QUERY_ADVANCED_QUERY_PARSER_OPTIMIZER_QUERY_OPTIMIZATION_UTIL_H_
+
+#include <memory>
+#include <vector>
+
+#include "icing/feature-flags.h"
+#include "icing/index/iterator/doc-hit-info-iterator.h"
+
+namespace icing {
+namespace lib {
+namespace query_optimization_util {
+
+// Optimizes the given iterators that are all intended to be AND'd together.
+// If an optimization is possible, then this node will be rewritten to an
+// equivalent but more efficient iterator sub-tree. Otherwise, the iterators
+// will be returned as a normal AND iterator.
+std::unique_ptr<DocHitInfoIterator> OptimizeAndIteratorsIfPossible(
+    std::vector<std::unique_ptr<DocHitInfoIterator>>&& iterators,
+    const FeatureFlags& feature_flags);
+
+}  // namespace query_optimization_util
+}  // namespace lib
+}  // namespace icing
+
+#endif  // ICING_QUERY_ADVANCED_QUERY_PARSER_OPTIMIZER_QUERY_OPTIMIZATION_UTIL_H_

diff --git a/icing/query/advanced_query_parser/query-visitor.cc b/icing/query/advanced_query_parser/query-visitor.cc
index 04624b7..492a416 100644
--- a/icing/query/advanced_query_parser/query-visitor.cc
+++ b/icing/query/advanced_query_parser/query-visitor.cc

@@ -33,7 +33,8 @@
 #include "icing/absl_ports/canonical_errors.h"
 #include "icing/absl_ports/str_cat.h"
 #include "icing/absl_ports/str_join.h"
-#include "icing/index/embed/doc-hit-info-iterator-embedding.h"
+#include "icing/index/embed/doc-hit-info-iterator-embedding-v1.h"
+#include "icing/index/embed/doc-hit-info-iterator-embedding-v2.h"
 #include "icing/index/embed/embedding-query-results.h"
 #include "icing/index/iterator/doc-hit-info-iterator-all-document-id.h"
 #include "icing/index/iterator/doc-hit-info-iterator-and.h"
@@ -50,6 +51,7 @@
 #include "icing/query/advanced_query_parser/abstract-syntax-tree.h"
 #include "icing/query/advanced_query_parser/function.h"
 #include "icing/query/advanced_query_parser/lexer.h"
+#include "icing/query/advanced_query_parser/optimizer/query-optimization-util.h"
 #include "icing/query/advanced_query_parser/param.h"
 #include "icing/query/advanced_query_parser/parser.h"
 #include "icing/query/advanced_query_parser/pending-value.h"
@@ -200,9 +202,9 @@
                              search_spec_.term_match_type(),
                              needs_term_frequency_info_));
       query_term_iterators_[query_term.term] =
-          std::make_unique<DocHitInfoIteratorFilter>(
-              std::move(term_iterator), &document_store_, &schema_store_,
-              filter_options_, current_time_ms_);
+          DocHitInfoIteratorFilter::ApplyFilter(
+              std::move(term_iterator), filter_predicate_,
+              feature_flags_.enable_passing_filter_to_children());
     }
   }
 
@@ -351,7 +353,7 @@
     QueryVisitor query_visitor(
         &index_, &numeric_index_, &embedding_index_, &document_store_,
         &schema_store_, &normalizer_, &tokenizer_, join_children_fetcher_,
-        search_spec_, filter_options_, needs_term_frequency_info_,
+        search_spec_, filter_predicate_, needs_term_frequency_info_,
         get_embedding_match_info_, &feature_flags_, pending_property_restricts_,
         processing_not_, current_time_ms_);
     tree_root->Accept(&query_visitor);
@@ -462,14 +464,34 @@
   }
 
   // Create and return iterator.
-  EmbeddingQueryResults::EmbeddingQueryMatchInfoMap* info_map =
-      &embedding_query_results_.result_infos[vector_index][metric_type];
   ICING_ASSIGN_OR_RETURN(
-      std::unique_ptr<DocHitInfoIterator> iterator,
-      DocHitInfoIteratorEmbedding::Create(
-          &search_spec_.embedding_query_vectors(vector_index), metric_type, low,
-          high, get_embedding_match_info_, info_map, &embedding_index_,
-          &document_store_, &schema_store_, current_time_ms_));
+      EmbeddingQueryResults::EmbeddingQueryMatchInfoMap * info_map,
+      embedding_query_results_.GetOrCreateMatchInfoMap(vector_index,
+                                                       metric_type));
+  std::unique_ptr<DocHitInfoIterator> iterator;
+  if (feature_flags_.enable_embedding_iterator_v2()) {
+    ICING_ASSIGN_OR_RETURN(
+        iterator,
+        DocHitInfoIteratorEmbeddingV2::Create(
+            &search_spec_.embedding_query_vectors(vector_index), metric_type,
+            low, high, info_map, embedding_query_results_.global_scores.get(),
+            get_embedding_match_info_
+                ? embedding_query_results_.global_section_infos.get()
+                : nullptr,
+            &embedding_index_, &document_store_, &schema_store_,
+            current_time_ms_));
+  } else {
+    ICING_ASSIGN_OR_RETURN(
+        iterator,
+        DocHitInfoIteratorEmbeddingV1::Create(
+            &search_spec_.embedding_query_vectors(vector_index), metric_type,
+            low, high, info_map, embedding_query_results_.global_scores.get(),
+            get_embedding_match_info_
+                ? embedding_query_results_.global_section_infos.get()
+                : nullptr,
+            &embedding_index_, &document_store_, &schema_store_,
+            current_time_ms_));
+  }
   return PendingValue(std::move(iterator));
 }
 
@@ -565,15 +587,16 @@
   std::vector<std::unique_ptr<DocHitInfoIterator>> iterators;
   // raw_text is the portion of text_value.raw_term that hasn't yet been
   // matched to any of the tokens that we've processed. escaped_token will
-  // hold the portion of raw_text that corresponds to the current token that
+  // holds the portion of raw_text that corresponds to the current token that
   // is being processed.
   std::string_view raw_text = text_value.raw_term;
+  std::vector<Token> tokens;
   std::string_view raw_token;
   bool reached_final_token = !token_itr->Advance();
   // If the term is different then the raw_term, then there must have been some
   // escaped characters that we will need to handle.
   while (!reached_final_token) {
-    std::vector<Token> tokens = token_itr->GetTokens();
+    token_itr->GetTokens(&tokens);
     if (tokens.size() > 1) {
       // The tokenizer iterator iterates between token groups. In practice,
       // the tokenizer used with QueryVisitor (PlainTokenizer) will always
@@ -717,7 +740,8 @@
   ICING_ASSIGN_OR_RETURN(
       std::vector<std::unique_ptr<DocHitInfoIterator>> iterators,
       PopAllPendingIterators());
-  return PendingValue(CreateAndIterator(std::move(iterators)));
+  return PendingValue(query_optimization_util::OptimizeAndIteratorsIfPossible(
+      std::move(iterators), feature_flags_));
 }
 
 libtextclassifier3::StatusOr<PendingValue> QueryVisitor::ProcessOrOperator(

diff --git a/icing/query/advanced_query_parser/query-visitor.h b/icing/query/advanced_query_parser/query-visitor.h
index 746d921..b1fd3c6 100644
--- a/icing/query/advanced_query_parser/query-visitor.h
+++ b/icing/query/advanced_query_parser/query-visitor.h

@@ -31,8 +31,8 @@
 #include "icing/index/embed/embedding-index.h"
 #include "icing/index/embed/embedding-query-results.h"
 #include "icing/index/index.h"
-#include "icing/index/iterator/doc-hit-info-iterator-filter.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
+#include "icing/index/iterator/document-filter-predicate.h"
 #include "icing/index/numeric/numeric-index.h"
 #include "icing/join/join-children-fetcher.h"
 #include "icing/query/advanced_query_parser/abstract-syntax-tree.h"
@@ -60,12 +60,12 @@
       const Normalizer* normalizer, const Tokenizer* tokenizer,
       const JoinChildrenFetcher* join_children_fetcher,
       const SearchSpecProto& search_spec,
-      DocHitInfoIteratorFilter::Options filter_options,
+      const DocumentFilterPredicate* filter_predicate,
       bool needs_term_frequency_info, bool get_embedding_match_info,
       const FeatureFlags* feature_flags, int64_t current_time_ms)
       : QueryVisitor(index, numeric_index, embedding_index, document_store,
                      schema_store, normalizer, tokenizer, join_children_fetcher,
-                     search_spec, filter_options, needs_term_frequency_info,
+                     search_spec, filter_predicate, needs_term_frequency_info,
                      get_embedding_match_info, feature_flags,
                      PendingPropertyRestricts(),
                      /*processing_not=*/false, current_time_ms) {}
@@ -119,12 +119,13 @@
       const Normalizer* normalizer, const Tokenizer* tokenizer,
       const JoinChildrenFetcher* join_children_fetcher,
       const SearchSpecProto& search_spec,
-      DocHitInfoIteratorFilter::Options filter_options,
+      const DocumentFilterPredicate* filter_predicate,
       bool needs_term_frequency_info, bool get_embedding_match_info,
       const FeatureFlags* feature_flags,
       PendingPropertyRestricts pending_property_restricts, bool processing_not,
       int64_t current_time_ms)
-      : index_(*index),
+      : embedding_query_results_(search_spec.embedding_query_vectors_size()),
+        index_(*index),
         numeric_index_(*numeric_index),
         embedding_index_(*embedding_index),
         document_store_(*document_store),
@@ -133,7 +134,7 @@
         tokenizer_(*tokenizer),
         join_children_fetcher_(join_children_fetcher),
         search_spec_(search_spec),
-        filter_options_(std::move(filter_options)),
+        filter_predicate_(filter_predicate),
         needs_term_frequency_info_(needs_term_frequency_info),
         get_embedding_match_info_(get_embedding_match_info),
         feature_flags_(*feature_flags),
@@ -381,9 +382,9 @@
   // parent query for a join query, in which case child scores are available.
   const JoinChildrenFetcher* join_children_fetcher_;  // Does not own.
 
-  const SearchSpecProto& search_spec_;
+  const SearchSpecProto& search_spec_;  // Does not own.
 
-  DocHitInfoIteratorFilter::Options filter_options_;
+  const DocumentFilterPredicate* filter_predicate_;  // Does not own.
 
   // Whether or not term_frequency information is needed. This affects:
   //  - how DocHitInfoIteratorTerms are constructed

diff --git a/icing/query/advanced_query_parser/query-visitor_test.cc b/icing/query/advanced_query_parser/query-visitor_test.cc
index 3c50b4a..3d1a903 100644
--- a/icing/query/advanced_query_parser/query-visitor_test.cc
+++ b/icing/query/advanced_query_parser/query-visitor_test.cc

@@ -39,14 +39,15 @@
 #include "icing/index/embed/embedding-query-results.h"
 #include "icing/index/hit/hit.h"
 #include "icing/index/index.h"
-#include "icing/index/iterator/doc-hit-info-iterator-filter.h"
 #include "icing/index/iterator/doc-hit-info-iterator-test-util.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
+#include "icing/index/iterator/document-filter-predicate.h"
 #include "icing/index/numeric/dummy-numeric-index.h"
 #include "icing/index/numeric/numeric-index.h"
 #include "icing/index/property-existence-indexing-handler.h"
 #include "icing/jni/jni-cache.h"
 #include "icing/legacy/index/icing-filesystem.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/portable/platform.h"
 #include "icing/proto/search.pb.h"
 #include "icing/query/advanced_query_parser/abstract-syntax-tree.h"
@@ -54,9 +55,11 @@
 #include "icing/query/advanced_query_parser/parser.h"
 #include "icing/query/query-features.h"
 #include "icing/query/query-results.h"
+#include "icing/query/query-utils.h"
 #include "icing/schema-builder.h"
 #include "icing/schema/schema-store.h"
 #include "icing/schema/section.h"
+#include "icing/scoring/advanced_scoring/double-list.h"
 #include "icing/store/document-id.h"
 #include "icing/store/document-store.h"
 #include "icing/store/namespace-id.h"
@@ -74,6 +77,7 @@
 #include "icing/transform/normalizer-options.h"
 #include "icing/transform/normalizer.h"
 #include "icing/util/clock.h"
+#include "icing/util/document-util.h"
 #include "icing/util/icu-data-file-helper.h"
 #include "icing/util/status-macros.h"
 #include "unicode/uloc.h"
@@ -88,7 +92,6 @@
 using ::testing::ElementsAre;
 using ::testing::IsEmpty;
 using ::testing::IsNull;
-using ::testing::Pointee;
 using ::testing::UnorderedElementsAre;
 
 constexpr float kEps = 0.000001;
@@ -165,20 +168,25 @@
       /*embedding_query_vectors=*/{}, EMBEDDING_METRIC_UNKNOWN);
 }
 
-bool ContainsMatchInfoEntry(const EmbeddingMatchInfos* match_info, double score,
+bool ContainsMatchInfoEntry(EmbeddingQueryResults& embedding_query_results,
+                            const EmbeddingMatchInfos* match_info, double score,
                             int position_in_section, SectionId section_id) {
-  if (match_info == nullptr || match_info->section_infos == nullptr ||
-      match_info->scores.empty()) {
-    return false;
-  }
-  if (match_info->scores.size() != match_info->section_infos->size()) {
+  if (match_info == nullptr ||
+      embedding_query_results.global_section_infos->empty()) {
     return false;
   }
 
-  for (int i = 0; i < match_info->scores.size(); ++i) {
-    if (std::fabs(match_info->scores[i] - score) < kEps &&
-        match_info->section_infos->at(i).position == position_in_section &&
-        match_info->section_infos->at(i).section_id == section_id) {
+  DoubleList matched_scores =
+      embedding_query_results.GetMatchedScoresFromEmbeddingMatchInfos(
+          *match_info);
+
+  for (int i = 0; i < matched_scores.size(); ++i) {
+    const EmbeddingMatchInfos::EmbeddingMatchSectionInfo& section_info =
+        embedding_query_results.global_section_infos->at(
+            i + match_info->score_start_index);
+    if (std::fabs(matched_scores.data()[i] - score) < kEps &&
+        section_info.position == position_in_section &&
+        section_info.section_id == section_id) {
       return true;
     }
   }
@@ -193,18 +201,45 @@
 struct QueryVisitorTestParams {
   QueryType query_type;
   bool get_embedding_match_info;
+  bool enable_embedding_iterator_v2;
+  bool enable_embed_query_optimization;
 
   explicit QueryVisitorTestParams(QueryType query_type,
-                                  bool get_embedding_match_info)
+                                  bool get_embedding_match_info,
+                                  bool enable_embedding_iterator_v2,
+                                  bool enable_embed_query_optimization)
       : query_type(query_type),
-        get_embedding_match_info(get_embedding_match_info) {}
+        get_embedding_match_info(get_embedding_match_info),
+        enable_embedding_iterator_v2(enable_embedding_iterator_v2),
+        enable_embed_query_optimization(enable_embed_query_optimization) {}
 };
 
 class QueryVisitorTest
     : public ::testing::TestWithParam<QueryVisitorTestParams> {
  protected:
   void SetUp() override {
-    feature_flags_ = std::make_unique<FeatureFlags>(GetTestFeatureFlags());
+    feature_flags_ = std::make_unique<FeatureFlags>(
+        FeatureFlags(/*enable_circular_schema_definitions=*/true,
+                     /*enable_scorable_properties=*/true,
+                     /*enable_embedding_quantization=*/true,
+                     /*enable_repeated_field_joins=*/true,
+                     /*enable_embedding_backup_generation=*/true,
+                     /*enable_schema_database=*/true,
+                     /*release_backup_schema_file_if_overlay_present=*/true,
+                     /*enable_strict_page_byte_size_limit=*/true,
+                     /*enable_smaller_decompression_buffer_size=*/true,
+                     /*enable_eigen_embedding_scoring=*/true,
+                     /*enable_passing_filter_to_children=*/true,
+                     /*enable_proto_log_new_header_format=*/true,
+                     GetParam().enable_embedding_iterator_v2,
+                     /*enable_reusable_decompression_buffer=*/true,
+                     /*enable_schema_type_id_optimization=*/true,
+                     /*enable_optimize_improvements=*/true,
+                     /*expired_document_purge_threshold_ms=*/0,
+                     /*enable_non_existent_qualified_id_join=*/true,
+                     /*enable_skip_set_schema_type_equality_check=*/true,
+                     /*enable_embed_query_optimization=*/true,
+                     /*enable_schema_definition_deduping=*/true));
     test_dir_ = GetTestTempDir() + "/icing";
     index_dir_ = test_dir_ + "/index";
     numeric_index_dir_ = test_dir_ + "/numeric_index";
@@ -236,14 +271,18 @@
 
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
-        DocumentStore::Create(&filesystem_, store_dir_, &clock_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, store_dir_, &clock_, schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     document_store_ = std::move(create_result.document_store);
 
     Index::Options options(index_dir_.c_str(),
@@ -251,7 +290,8 @@
                            /*lite_index_sort_at_indexing=*/true,
                            /*lite_index_sort_size=*/1024 * 8);
     ICING_ASSERT_OK_AND_ASSIGN(
-        index_, Index::Create(options, &filesystem_, &icing_filesystem_));
+        index_, Index::Create(options, &filesystem_, &icing_filesystem_,
+                              feature_flags_.get()));
 
     ICING_ASSERT_OK_AND_ASSIGN(
         numeric_index_,
@@ -260,7 +300,8 @@
     ICING_ASSERT_OK_AND_ASSIGN(
         embedding_index_,
         EmbeddingIndex::Create(&filesystem_, embedding_index_dir_, &clock_,
-                               feature_flags_.get()));
+                               feature_flags_.get(),
+                               /*num_shards=*/32));
 
     NormalizerOptions normalizer_options(
         /*max_term_byte_size=*/std::numeric_limits<int32_t>::max());
@@ -277,6 +318,10 @@
                                tokenizer_factory::CreateIndexingTokenizer(
                                    StringIndexingConfig::TokenizerType::PLAIN,
                                    language_segmenter_.get()));
+
+    document_filter_predicate_ = GetFilterPredicateBySchemaAndNamespace(
+        SearchSpecProto::default_instance(), *document_store_, *schema_store_,
+        clock_.GetSystemTimeMilliseconds());
   }
 
   libtextclassifier3::StatusOr<std::unique_ptr<Node>> ParseQueryHelper(
@@ -295,7 +340,7 @@
         index_.get(), numeric_index_.get(), embedding_index_.get(),
         document_store_.get(), schema_store_.get(), normalizer_.get(),
         tokenizer_.get(), /*join_children_fetcher=*/nullptr, search_spec,
-        DocHitInfoIteratorFilter::Options(),
+        document_filter_predicate_.get(),
         /*needs_term_frequency_info=*/true, get_embedding_match_info,
         feature_flags_.get(), clock_.GetSystemTimeMilliseconds());
     root_node->Accept(&query_visitor);
@@ -364,6 +409,7 @@
   std::unique_ptr<Normalizer> normalizer_;
   std::unique_ptr<LanguageSegmenter> language_segmenter_;
   std::unique_ptr<Tokenizer> tokenizer_;
+  std::unique_ptr<DocumentFilterPredicate> document_filter_predicate_;
   std::unique_ptr<const JniCache> jni_cache_;
 };
 
@@ -769,7 +815,7 @@
       index_.get(), numeric_index_.get(), embedding_index_.get(),
       document_store_.get(), schema_store_.get(), normalizer_.get(),
       tokenizer_.get(), /*join_children_fetcher=*/nullptr, search_spec,
-      DocHitInfoIteratorFilter::Options(),
+      document_filter_predicate_.get(),
       /*needs_term_frequency_info=*/true, /*get_embedding_match_info=*/false,
       feature_flags_.get(), clock_.GetSystemTimeMilliseconds());
   EXPECT_THAT(std::move(query_visitor).ConsumeResults(),
@@ -838,8 +884,8 @@
   // - Doc0: ["-2", "-1", "1", "2"] and [-2, -1, 1, 2]
   // - Doc1: [-1]
   // - Doc2: ["2"] and [-1]
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
   std::unique_ptr<NumericIndex<int64_t>::Editor> editor =
       numeric_index_->Edit("price", kDocumentId0, kSectionId0);
   ICING_ASSERT_OK(editor->BufferKey(-2));
@@ -855,14 +901,14 @@
   ICING_ASSERT_OK(term_editor.BufferTerm("2", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(term_editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build())));
   editor = numeric_index_->Edit("price", kDocumentId1, kSectionId0);
   ICING_ASSERT_OK(editor->BufferKey(-1));
   ICING_ASSERT_OK(std::move(*editor).IndexAllBufferedKeys());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build())));
   editor = numeric_index_->Edit("price", kDocumentId2, kSectionId0);
   ICING_ASSERT_OK(editor->BufferKey(-1));
   ICING_ASSERT_OK(std::move(*editor).IndexAllBufferedKeys());
@@ -966,7 +1012,7 @@
       index_.get(), numeric_index_.get(), embedding_index_.get(),
       document_store_.get(), schema_store_.get(), normalizer_.get(),
       tokenizer_.get(), /*join_children_fetcher=*/nullptr, search_spec,
-      DocHitInfoIteratorFilter::Options(),
+      document_filter_predicate_.get(),
       /*needs_term_frequency_info=*/false, /*get_embedding_match_info=*/false,
       feature_flags_.get(), clock_.GetSystemTimeMilliseconds());
   root_node->Accept(&query_visitor);
@@ -1455,22 +1501,22 @@
           .Build(),
       /*ignore_errors_and_delete_documents=*/false));
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
   Index::Editor editor = index_->Edit(kDocumentId0, kSectionId1,
                                       /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId1, kSectionId1,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId2, kSectionId1,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("bar", TERM_MATCH_PREFIX));
@@ -1499,22 +1545,22 @@
           .Build(),
       /*ignore_errors_and_delete_documents=*/false));
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
   Index::Editor editor = index_->Edit(kDocumentId0, kSectionId1,
                                       /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId1, kSectionId1,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId2, kSectionId1,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("bar", TERM_MATCH_PREFIX));
@@ -1539,8 +1585,8 @@
           .Build(),
       /*ignore_errors_and_delete_documents=*/false));
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
   Index::Editor editor = index_->Edit(kDocumentId0, kSectionId1,
                                       /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -1548,16 +1594,16 @@
   ICING_ASSERT_OK(editor.BufferTerm("baz", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId1, kSectionId1,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.BufferTerm("baz", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId2, kSectionId1,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("bar", TERM_MATCH_PREFIX));
@@ -1587,8 +1633,8 @@
           .Build(),
       /*ignore_errors_and_delete_documents=*/false));
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
   Index::Editor editor = index_->Edit(kDocumentId0, kSectionId1,
                                       /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -1596,16 +1642,16 @@
   ICING_ASSERT_OK(editor.BufferTerm("baz", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId1, kSectionId1,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.BufferTerm("baz", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId2, kSectionId1,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("bar", TERM_MATCH_PREFIX));
@@ -1819,23 +1865,23 @@
           .Build(),
       /*ignore_errors_and_delete_documents=*/false));
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
   Index::Editor editor = index_->Edit(kDocumentId0, kSectionId1,
                                       /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId1, kSectionId1,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.BufferTerm("bar", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId2, kSectionId1,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -1889,22 +1935,22 @@
   SectionId prop1_section_id = 0;
   SectionId prop2_section_id = 1;
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
   Index::Editor editor = index_->Edit(kDocumentId0, prop1_section_id,
                                       /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId1, prop1_section_id,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId2, prop2_section_id,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -1955,22 +2001,22 @@
   SectionId prop2_section_id = 1;
   SectionId prop3_section_id = 2;
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
   Index::Editor editor = index_->Edit(kDocumentId0, prop1_section_id,
                                       /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId1, prop2_section_id,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId2, prop3_section_id,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -2035,22 +2081,22 @@
   SectionId prop1_section_id = 0;
   SectionId prop2_section_id = 1;
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
   Index::Editor editor = index_->Edit(kDocumentId0, prop1_section_id,
                                       /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId1, prop1_section_id,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId2, prop2_section_id,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -2095,22 +2141,22 @@
   SectionId prop1_section_id = 0;
   SectionId prop2_section_id = 1;
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
   Index::Editor editor = index_->Edit(kDocumentId0, prop1_section_id,
                                       /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("bar", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId1, prop1_section_id,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId2, prop2_section_id,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -2153,22 +2199,22 @@
   SectionId prop1_section_id = 0;
   SectionId prop2_section_id = 1;
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
   Index::Editor editor = index_->Edit(kDocumentId0, prop1_section_id,
                                       /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("bar", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId1, prop1_section_id,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId2, prop2_section_id,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -2222,22 +2268,22 @@
   SectionId prop1_section_id = 0;
   SectionId prop2_section_id = 1;
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
   Index::Editor editor = index_->Edit(kDocumentId0, prop1_section_id,
                                       /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("bar", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId1, prop1_section_id,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId2, prop2_section_id,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -2287,22 +2333,22 @@
   SectionId prop1_section_id = 0;
   SectionId prop2_section_id = 1;
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
   Index::Editor editor = index_->Edit(kDocumentId0, prop1_section_id,
                                       /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("bar", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId1, prop1_section_id,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId2, prop2_section_id,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -2367,22 +2413,22 @@
   //   Doc2:
   //     prop1: ""
   //     prop2: "foo"
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
   Index::Editor editor = index_->Edit(kDocumentId0, prop1_section_id,
                                       /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("bar", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId1, prop1_section_id,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId2, prop2_section_id,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -2452,8 +2498,8 @@
   // ICU segmentation will break this into "每天" and "上班".
   // CFStringTokenizer (ios) will break this into "每", "天" and "上班"
   std::string query = CreateQuery("每天上班");
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
   Index::Editor editor = index_->Edit(kDocumentId0, prop1_section_id,
                                       /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("上班", TERM_MATCH_PREFIX));
@@ -2468,15 +2514,15 @@
   }
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId1, prop1_section_id,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("上班", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId2, prop2_section_id,
                         /*namespace_id=*/0);
   if (IsCfStringTokenization()) {
@@ -2537,8 +2583,9 @@
   // - Doc 0: Contains 'val0', 'val1', 'val2' in 'prop0'. Shouldn't match.
   DocumentProto doc =
       DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result0,
-                             document_store_->Put(doc));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result0,
+      document_store_->Put(document_util::CreateDocumentWrapper(doc)));
   DocumentId docid0 = put_result0.new_document_id;
   Index::Editor editor = index_->Edit(docid0, prop0_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("val0", TERM_MATCH_PREFIX));
@@ -2548,8 +2595,9 @@
 
   // - Doc 1: Contains 'val0', 'val1', 'val2' in 'prop1'. Should match.
   doc = DocumentBuilder(doc).SetUri("uri1").Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(doc));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(doc)));
   DocumentId docid1 = put_result1.new_document_id;
   editor = index_->Edit(docid1, prop1_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("val0", TERM_MATCH_PREFIX));
@@ -2559,8 +2607,9 @@
 
   // - Doc 2: Contains 'val0', 'val1', 'val2' in 'prop2'. Shouldn't match.
   doc = DocumentBuilder(doc).SetUri("uri2").Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(doc));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(doc)));
   DocumentId docid2 = put_result2.new_document_id;
   editor = index_->Edit(docid2, prop2_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("val0", TERM_MATCH_PREFIX));
@@ -2570,8 +2619,9 @@
 
   // - Doc 3: Contains 'val0' in 'prop0', 'val1' in 'prop1' etc. Should match.
   doc = DocumentBuilder(doc).SetUri("uri3").Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             document_store_->Put(doc));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      document_store_->Put(document_util::CreateDocumentWrapper(doc)));
   DocumentId docid3 = put_result3.new_document_id;
   editor = index_->Edit(docid3, prop0_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("val0", TERM_MATCH_PREFIX));
@@ -2586,8 +2636,9 @@
   // - Doc 4: Contains 'val1' in 'prop0', 'val2' in 'prop1', 'val0' in 'prop2'.
   //          Shouldn't match.
   doc = DocumentBuilder(doc).SetUri("uri4").Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result4,
-                             document_store_->Put(doc));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result4,
+      document_store_->Put(document_util::CreateDocumentWrapper(doc)));
   DocumentId docid4 = put_result4.new_document_id;
   editor = index_->Edit(docid4, prop0_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("val1", TERM_MATCH_PREFIX));
@@ -2651,8 +2702,9 @@
   // - Doc 0: Contains 'val0', 'val1', 'val2' in 'prop0'. Should match.
   DocumentProto doc =
       DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result0,
-                             document_store_->Put(doc));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result0,
+      document_store_->Put(document_util::CreateDocumentWrapper(doc)));
   DocumentId docid0 = put_result0.new_document_id;
   Index::Editor editor = index_->Edit(docid0, prop0_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("val0", TERM_MATCH_PREFIX));
@@ -2662,8 +2714,9 @@
 
   // - Doc 1: Contains 'val0', 'val1', 'val2' in 'prop1'. Shouldn't match.
   doc = DocumentBuilder(doc).SetUri("uri1").Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(doc));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(doc)));
   DocumentId docid1 = put_result1.new_document_id;
   editor = index_->Edit(docid1, prop1_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("val0", TERM_MATCH_PREFIX));
@@ -2673,8 +2726,9 @@
 
   // - Doc 2: Contains 'val0', 'val1', 'val2' in 'prop2'. Should match.
   doc = DocumentBuilder(doc).SetUri("uri2").Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(doc));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(doc)));
   DocumentId docid2 = put_result2.new_document_id;
   editor = index_->Edit(docid2, prop2_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("val0", TERM_MATCH_PREFIX));
@@ -2684,8 +2738,9 @@
 
   // - Doc 3: Contains 'val0' in 'prop0', 'val1' in 'prop1' etc. Should match.
   doc = DocumentBuilder(doc).SetUri("uri3").Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             document_store_->Put(doc));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      document_store_->Put(document_util::CreateDocumentWrapper(doc)));
   DocumentId docid3 = put_result3.new_document_id;
   editor = index_->Edit(docid3, prop0_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("val0", TERM_MATCH_PREFIX));
@@ -2700,8 +2755,9 @@
   // - Doc 4: Contains 'val1' in 'prop0', 'val2' in 'prop1', 'val0' in 'prop2'.
   //          Shouldn't match.
   doc = DocumentBuilder(doc).SetUri("uri4").Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result4,
-                             document_store_->Put(doc));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result4,
+      document_store_->Put(document_util::CreateDocumentWrapper(doc)));
   DocumentId docid4 = put_result4.new_document_id;
   editor = index_->Edit(docid4, prop0_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("val1", TERM_MATCH_PREFIX));
@@ -2807,22 +2863,22 @@
   // Section ids are assigned alphabetically.
   SectionId prop1_section_id = 0;
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
   Index::Editor editor = index_->Edit(kDocumentId0, prop1_section_id,
                                       /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("bar", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId1, prop1_section_id,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri2").SetSchema("type").Build())));
   editor = index_->Edit(kDocumentId2, prop1_section_id,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -2916,8 +2972,9 @@
   NamespaceId ns_id = 0;
   DocumentProto doc =
       DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result0,
-                             document_store_->Put(doc));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result0,
+      document_store_->Put(document_util::CreateDocumentWrapper(doc)));
   DocumentId docid0 = put_result0.new_document_id;
   Index::Editor editor = index_->Edit(kDocumentId0, prop0_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -2925,7 +2982,8 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result1,
-      document_store_->Put(DocumentBuilder(doc).SetUri("uri1").Build()));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          DocumentBuilder(doc).SetUri("uri1").Build())));
   DocumentId docid1 = put_result1.new_document_id;
   editor = index_->Edit(docid1, prop1_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -2933,7 +2991,8 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result2,
-      document_store_->Put(DocumentBuilder(doc).SetUri("uri2").Build()));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          DocumentBuilder(doc).SetUri("uri2").Build())));
   DocumentId docid2 = put_result2.new_document_id;
   editor = index_->Edit(docid2, prop2_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -2941,7 +3000,8 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result3,
-      document_store_->Put(DocumentBuilder(doc).SetUri("uri3").Build()));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          DocumentBuilder(doc).SetUri("uri3").Build())));
   DocumentId docid3 = put_result3.new_document_id;
   editor = index_->Edit(docid3, prop3_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -2949,7 +3009,8 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result4,
-      document_store_->Put(DocumentBuilder(doc).SetUri("uri4").Build()));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          DocumentBuilder(doc).SetUri("uri4").Build())));
   DocumentId docid4 = put_result4.new_document_id;
   editor = index_->Edit(docid4, prop4_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -2957,7 +3018,8 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result5,
-      document_store_->Put(DocumentBuilder(doc).SetUri("uri5").Build()));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          DocumentBuilder(doc).SetUri("uri5").Build())));
   DocumentId docid5 = put_result5.new_document_id;
   editor = index_->Edit(docid5, prop5_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -2965,7 +3027,8 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result6,
-      document_store_->Put(DocumentBuilder(doc).SetUri("uri6").Build()));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          DocumentBuilder(doc).SetUri("uri6").Build())));
   DocumentId docid6 = put_result6.new_document_id;
   editor = index_->Edit(docid6, prop6_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -2973,7 +3036,8 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result7,
-      document_store_->Put(DocumentBuilder(doc).SetUri("uri7").Build()));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          DocumentBuilder(doc).SetUri("uri7").Build())));
   DocumentId docid7 = put_result7.new_document_id;
   editor = index_->Edit(docid7, prop7_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -3076,8 +3140,9 @@
   NamespaceId ns_id = 0;
   DocumentProto doc =
       DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result0,
-                             document_store_->Put(doc));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result0,
+      document_store_->Put(document_util::CreateDocumentWrapper(doc)));
   DocumentId docid0 = put_result0.new_document_id;
   Index::Editor editor = index_->Edit(kDocumentId0, prop0_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -3085,7 +3150,8 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result1,
-      document_store_->Put(DocumentBuilder(doc).SetUri("uri1").Build()));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          DocumentBuilder(doc).SetUri("uri1").Build())));
   DocumentId docid1 = put_result1.new_document_id;
   editor = index_->Edit(docid1, prop1_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -3093,7 +3159,8 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result2,
-      document_store_->Put(DocumentBuilder(doc).SetUri("uri2").Build()));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          DocumentBuilder(doc).SetUri("uri2").Build())));
   DocumentId docid2 = put_result2.new_document_id;
   editor = index_->Edit(docid2, prop2_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -3101,7 +3168,8 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result3,
-      document_store_->Put(DocumentBuilder(doc).SetUri("uri3").Build()));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          DocumentBuilder(doc).SetUri("uri3").Build())));
   DocumentId docid3 = put_result3.new_document_id;
   editor = index_->Edit(docid3, prop3_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -3109,7 +3177,8 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result4,
-      document_store_->Put(DocumentBuilder(doc).SetUri("uri4").Build()));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          DocumentBuilder(doc).SetUri("uri4").Build())));
   DocumentId docid4 = put_result4.new_document_id;
   editor = index_->Edit(docid4, prop4_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -3117,7 +3186,8 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result5,
-      document_store_->Put(DocumentBuilder(doc).SetUri("uri5").Build()));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          DocumentBuilder(doc).SetUri("uri5").Build())));
   DocumentId docid5 = put_result5.new_document_id;
   editor = index_->Edit(docid5, prop5_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -3125,7 +3195,8 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result6,
-      document_store_->Put(DocumentBuilder(doc).SetUri("uri6").Build()));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          DocumentBuilder(doc).SetUri("uri6").Build())));
   DocumentId docid6 = put_result6.new_document_id;
   editor = index_->Edit(docid6, prop6_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -3133,7 +3204,8 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result7,
-      document_store_->Put(DocumentBuilder(doc).SetUri("uri7").Build()));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          DocumentBuilder(doc).SetUri("uri7").Build())));
   DocumentId docid7 = put_result7.new_document_id;
   editor = index_->Edit(docid7, prop7_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -3208,8 +3280,9 @@
   NamespaceId ns_id = 0;
   DocumentProto doc =
       DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result0,
-                             document_store_->Put(doc));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result0,
+      document_store_->Put(document_util::CreateDocumentWrapper(doc)));
   DocumentId docid0 = put_result0.new_document_id;
   Index::Editor editor = index_->Edit(kDocumentId0, prop0_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -3218,7 +3291,8 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result1,
-      document_store_->Put(DocumentBuilder(doc).SetUri("uri1").Build()));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          DocumentBuilder(doc).SetUri("uri1").Build())));
   DocumentId docid1 = put_result1.new_document_id;
   editor = index_->Edit(docid1, prop0_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("bar", TERM_MATCH_PREFIX));
@@ -3226,7 +3300,8 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result2,
-      document_store_->Put(DocumentBuilder(doc).SetUri("uri2").Build()));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          DocumentBuilder(doc).SetUri("uri2").Build())));
   DocumentId docid2 = put_result2.new_document_id;
   editor = index_->Edit(docid2, prop0_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -3327,8 +3402,9 @@
   NamespaceId ns_id = 0;
   DocumentProto doc =
       DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result0,
-                             document_store_->Put(doc));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result0,
+      document_store_->Put(document_util::CreateDocumentWrapper(doc)));
   DocumentId docid0 = put_result0.new_document_id;
   Index::Editor editor = index_->Edit(docid0, prop0_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -3337,7 +3413,8 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result1,
-      document_store_->Put(DocumentBuilder(doc).SetUri("uri1").Build()));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          DocumentBuilder(doc).SetUri("uri1").Build())));
   DocumentId docid1 = put_result1.new_document_id;
   editor = index_->Edit(docid1, prop0_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("bar", TERM_MATCH_PREFIX));
@@ -3348,7 +3425,8 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result2,
-      document_store_->Put(DocumentBuilder(doc).SetUri("uri2").Build()));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          DocumentBuilder(doc).SetUri("uri2").Build())));
   DocumentId docid2 = put_result2.new_document_id;
   editor = index_->Edit(docid2, prop0_id, ns_id);
   ICING_ASSERT_OK(editor.BufferTerm("bar", TERM_MATCH_PREFIX));
@@ -3520,17 +3598,21 @@
 
   // Document 0 has the term "foo" and its schema has the url property.
   ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("typeWithUrl").Build()));
+      document_util::CreateDocumentWrapper(DocumentBuilder()
+                                               .SetKey("ns", "uri0")
+                                               .SetSchema("typeWithUrl")
+                                               .Build())));
   Index::Editor editor = index_->Edit(kDocumentId0, kSectionId1,
                                       /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
   // Document 1 has the term "foo" and its schema DOESN'T have the url property.
-  ICING_ASSERT_OK(document_store_->Put(DocumentBuilder()
-                                           .SetKey("ns", "uri1")
-                                           .SetSchema("typeWithoutUrl")
-                                           .Build()));
+  ICING_ASSERT_OK(document_store_->Put(
+      document_util::CreateDocumentWrapper(DocumentBuilder()
+                                               .SetKey("ns", "uri1")
+                                               .SetSchema("typeWithoutUrl")
+                                               .Build())));
   editor = index_->Edit(kDocumentId1, kSectionId1,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -3538,7 +3620,10 @@
 
   // Document 2 has the term "bar" and its schema has the url property.
   ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri2").SetSchema("typeWithUrl").Build()));
+      document_util::CreateDocumentWrapper(DocumentBuilder()
+                                               .SetKey("ns", "uri2")
+                                               .SetSchema("typeWithUrl")
+                                               .Build())));
   editor = index_->Edit(kDocumentId2, kSectionId1,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("bar", TERM_MATCH_PREFIX));
@@ -3569,17 +3654,21 @@
 
   // Document 0 has the term "foo" and its schema has the url property.
   ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("typeWithUrl").Build()));
+      document_util::CreateDocumentWrapper(DocumentBuilder()
+                                               .SetKey("ns", "uri0")
+                                               .SetSchema("typeWithUrl")
+                                               .Build())));
   Index::Editor editor = index_->Edit(kDocumentId0, kSectionId1,
                                       /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
   // Document 1 has the term "foo" and its schema DOESN'T have the url property.
-  ICING_ASSERT_OK(document_store_->Put(DocumentBuilder()
-                                           .SetKey("ns", "uri1")
-                                           .SetSchema("typeWithoutUrl")
-                                           .Build()));
+  ICING_ASSERT_OK(document_store_->Put(
+      document_util::CreateDocumentWrapper(DocumentBuilder()
+                                               .SetKey("ns", "uri1")
+                                               .SetSchema("typeWithoutUrl")
+                                               .Build())));
   editor = index_->Edit(kDocumentId1, kSectionId1,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -3610,17 +3699,21 @@
 
   // Document 0 has the term "foo" and its schema has the url property.
   ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("typeWithUrl").Build()));
+      document_util::CreateDocumentWrapper(DocumentBuilder()
+                                               .SetKey("ns", "uri0")
+                                               .SetSchema("typeWithUrl")
+                                               .Build())));
   Index::Editor editor = index_->Edit(kDocumentId0, kSectionId1,
                                       /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
   // Document 1 has the term "foo" and its schema DOESN'T have the url property.
-  ICING_ASSERT_OK(document_store_->Put(DocumentBuilder()
-                                           .SetKey("ns", "uri1")
-                                           .SetSchema("typeWithoutUrl")
-                                           .Build()));
+  ICING_ASSERT_OK(document_store_->Put(
+      document_util::CreateDocumentWrapper(DocumentBuilder()
+                                               .SetKey("ns", "uri1")
+                                               .SetSchema("typeWithoutUrl")
+                                               .Build())));
   editor = index_->Edit(kDocumentId1, kSectionId1,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -3680,8 +3773,8 @@
       /*ignore_errors_and_delete_documents=*/false));
 
   // Document 0 has the term "foo" and has the "price" property.
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("Simple").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("Simple").Build())));
   Index::Editor editor = index_->Edit(kDocumentId0, kSectionId0,
                                       /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -3691,16 +3784,16 @@
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
   // Document 1 has the term "foo" and doesn't have the "price" property.
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri1").SetSchema("Simple").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri1").SetSchema("Simple").Build())));
   editor = index_->Edit(kDocumentId1, kSectionId0,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
   // Document 2 has the term "bar" and has the "price" property.
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri2").SetSchema("Simple").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri2").SetSchema("Simple").Build())));
   editor = index_->Edit(kDocumentId2, kSectionId0,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("bar", TERM_MATCH_PREFIX));
@@ -3747,8 +3840,8 @@
       /*ignore_errors_and_delete_documents=*/false));
 
   // Document 0 has the term "foo" and has the "price" property.
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("Simple").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("Simple").Build())));
   Index::Editor editor = index_->Edit(kDocumentId0, kSectionId0,
                                       /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -3758,8 +3851,8 @@
   ICING_ASSERT_OK(editor.IndexAllBufferedTerms());
 
   // Document 1 has the term "foo" and doesn't have the "price" property.
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri1").SetSchema("Simple").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri1").SetSchema("Simple").Build())));
   editor = index_->Edit(kDocumentId1, kSectionId0,
                         /*namespace_id=*/0);
   ICING_ASSERT_OK(editor.BufferTerm("foo", TERM_MATCH_PREFIX));
@@ -3968,10 +4061,10 @@
                                         .SetCardinality(CARDINALITY_OPTIONAL)))
           .Build(),
       /*ignore_errors_and_delete_documents=*/false));
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build()));
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build())));
 
   // Index two embedding vectors.
   PropertyProto::VectorProto vector0 =
@@ -3979,9 +4072,11 @@
   PropertyProto::VectorProto vector1 =
       CreateVector("my_model", {-0.1, -0.2, -0.3});
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-      BasicHit(kSectionId0, kDocumentId0), vector0, QUANTIZATION_TYPE_NONE));
+      BasicHit(kSectionId0, kDocumentId0), vector0, QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-      BasicHit(kSectionId0, kDocumentId1), vector1, QUANTIZATION_TYPE_NONE));
+      BasicHit(kSectionId0, kDocumentId1), vector1, QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
 
   // Create an embedding query that has a semantic score of 1 with vector0 and
@@ -4007,22 +4102,19 @@
   EXPECT_THAT(
       query_results.embedding_query_results.GetMatchedScoresForDocument(
           /*query_vector_index=*/0, EMBEDDING_METRIC_COSINE, kDocumentId0),
-      Pointee(UnorderedElementsAre(DoubleNear(1, kEps))));
+      UnorderedElementsAre(DoubleNear(1, kEps)));
   if (GetParam().get_embedding_match_info) {
     // Check section match info for document 0.
     const EmbeddingMatchInfos* match_infos =
         query_results.embedding_query_results.GetMatchedInfosForDocument(
             /*query_vector_index=*/0, EMBEDDING_METRIC_COSINE, kDocumentId0);
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/1,
+    EXPECT_TRUE(ContainsMatchInfoEntry(query_results.embedding_query_results,
+                                       match_infos, /*score=*/1,
                                        /*position_in_section=*/0,
                                        /*section_id=*/kSectionId0));
   } else {
-    EXPECT_THAT(
-        query_results.embedding_query_results
-            .GetMatchedInfosForDocument(
-                /*query_vector_index=*/0, EMBEDDING_METRIC_COSINE, kDocumentId0)
-            ->section_infos,
-        IsNull());
+    EXPECT_THAT(*query_results.embedding_query_results.global_section_infos,
+                IsEmpty());
   }
 
   // The query should match both vector0 and vector1.
@@ -4039,17 +4131,18 @@
   EXPECT_THAT(
       query_results.embedding_query_results.GetMatchedScoresForDocument(
           /*query_vector_index=*/0, EMBEDDING_METRIC_COSINE, kDocumentId0),
-      Pointee(UnorderedElementsAre(DoubleNear(1, kEps))));
+      UnorderedElementsAre(DoubleNear(1, kEps)));
   EXPECT_THAT(
       query_results.embedding_query_results.GetMatchedScoresForDocument(
           /*query_vector_index=*/0, EMBEDDING_METRIC_COSINE, kDocumentId1),
-      Pointee(UnorderedElementsAre(DoubleNear(-1, kEps))));
+      UnorderedElementsAre(DoubleNear(-1, kEps)));
   if (GetParam().get_embedding_match_info) {
     // Check section match info for document 0.
     const EmbeddingMatchInfos* match_infos =
         query_results.embedding_query_results.GetMatchedInfosForDocument(
             /*query_vector_index=*/0, EMBEDDING_METRIC_COSINE, kDocumentId0);
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/1,
+    EXPECT_TRUE(ContainsMatchInfoEntry(query_results.embedding_query_results,
+                                       match_infos, /*score=*/1,
                                        /*position_in_section=*/0,
                                        /*section_id=*/kSectionId0));
 
@@ -4057,23 +4150,13 @@
     match_infos =
         query_results.embedding_query_results.GetMatchedInfosForDocument(
             /*query_vector_index=*/0, EMBEDDING_METRIC_COSINE, kDocumentId1);
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/-1,
+    EXPECT_TRUE(ContainsMatchInfoEntry(query_results.embedding_query_results,
+                                       match_infos, /*score=*/-1,
                                        /*position_in_section=*/0,
                                        /*section_id=*/kSectionId0));
   } else {
-    EXPECT_THAT(
-        query_results.embedding_query_results
-            .GetMatchedInfosForDocument(
-                /*query_vector_index=*/0, EMBEDDING_METRIC_COSINE, kDocumentId0)
-            ->section_infos,
-        IsNull());
-
-    EXPECT_THAT(
-        query_results.embedding_query_results
-            .GetMatchedInfosForDocument(
-                /*query_vector_index=*/0, EMBEDDING_METRIC_COSINE, kDocumentId1)
-            ->section_infos,
-        IsNull());
+    EXPECT_THAT(*query_results.embedding_query_results.global_section_infos,
+                IsEmpty());
   }
 
   // The query should match nothing, since there is no vector with a
@@ -4107,10 +4190,10 @@
                                         .SetCardinality(CARDINALITY_OPTIONAL)))
           .Build(),
       /*ignore_errors_and_delete_documents=*/false));
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build()));
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build())));
 
   // Index two embedding vectors.
   PropertyProto::VectorProto vector0 =
@@ -4118,9 +4201,11 @@
   PropertyProto::VectorProto vector1 =
       CreateVector("my_model", {-0.1, -0.2, -0.3});
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-      BasicHit(kSectionId0, kDocumentId0), vector0, QUANTIZATION_TYPE_NONE));
+      BasicHit(kSectionId0, kDocumentId0), vector0, QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-      BasicHit(kSectionId0, kDocumentId1), vector1, QUANTIZATION_TYPE_NONE));
+      BasicHit(kSectionId0, kDocumentId1), vector1, QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
 
   // Create an embedding query that has a semantic score of 1 with vector0 and
@@ -4146,22 +4231,19 @@
   EXPECT_THAT(
       query_results.embedding_query_results.GetMatchedScoresForDocument(
           /*query_vector_index=*/0, EMBEDDING_METRIC_COSINE, kDocumentId1),
-      Pointee(UnorderedElementsAre(DoubleNear(-1, kEps))));
+      UnorderedElementsAre(DoubleNear(-1, kEps)));
   if (GetParam().get_embedding_match_info) {
     // Check section match info for document 1.
     const EmbeddingMatchInfos* match_infos =
         query_results.embedding_query_results.GetMatchedInfosForDocument(
             /*query_vector_index=*/0, EMBEDDING_METRIC_COSINE, kDocumentId1);
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/-1,
+    EXPECT_TRUE(ContainsMatchInfoEntry(query_results.embedding_query_results,
+                                       match_infos, /*score=*/-1,
                                        /*position_in_section=*/0,
                                        /*section_id=*/kSectionId0));
   } else {
-    EXPECT_THAT(
-        query_results.embedding_query_results
-            .GetMatchedInfosForDocument(
-                /*query_vector_index=*/0, EMBEDDING_METRIC_COSINE, kDocumentId1)
-            ->section_infos,
-        IsNull());
+    EXPECT_THAT(*query_results.embedding_query_results.global_section_infos,
+                IsEmpty());
   }
 
   // The query should match both vector0 and vector1.
@@ -4178,17 +4260,18 @@
   EXPECT_THAT(
       query_results.embedding_query_results.GetMatchedScoresForDocument(
           /*query_vector_index=*/0, EMBEDDING_METRIC_COSINE, kDocumentId0),
-      Pointee(UnorderedElementsAre(DoubleNear(1, kEps))));
+      UnorderedElementsAre(DoubleNear(1, kEps)));
   EXPECT_THAT(
       query_results.embedding_query_results.GetMatchedScoresForDocument(
           /*query_vector_index=*/0, EMBEDDING_METRIC_COSINE, kDocumentId1),
-      Pointee(UnorderedElementsAre(DoubleNear(-1, kEps))));
+      UnorderedElementsAre(DoubleNear(-1, kEps)));
   if (GetParam().get_embedding_match_info) {
     // Check section match info for document 0.
     const EmbeddingMatchInfos* match_infos =
         query_results.embedding_query_results.GetMatchedInfosForDocument(
             /*query_vector_index=*/0, EMBEDDING_METRIC_COSINE, kDocumentId0);
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/1,
+    EXPECT_TRUE(ContainsMatchInfoEntry(query_results.embedding_query_results,
+                                       match_infos, /*score=*/1,
                                        /*position_in_section=*/0,
                                        /*section_id=*/kSectionId0));
 
@@ -4196,23 +4279,14 @@
     match_infos =
         query_results.embedding_query_results.GetMatchedInfosForDocument(
             /*query_vector_index=*/0, EMBEDDING_METRIC_COSINE, kDocumentId1);
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/-1,
+    EXPECT_TRUE(ContainsMatchInfoEntry(query_results.embedding_query_results,
+                                       match_infos, /*score=*/-1,
                                        /*position_in_section=*/0,
                                        /*section_id=*/kSectionId0));
 
   } else {
-    EXPECT_THAT(
-        query_results.embedding_query_results
-            .GetMatchedInfosForDocument(
-                /*query_vector_index=*/0, EMBEDDING_METRIC_COSINE, kDocumentId0)
-            ->section_infos,
-        IsNull());
-    EXPECT_THAT(
-        query_results.embedding_query_results
-            .GetMatchedInfosForDocument(
-                /*query_vector_index=*/0, EMBEDDING_METRIC_COSINE, kDocumentId1)
-            ->section_infos,
-        IsNull());
+    EXPECT_THAT(*query_results.embedding_query_results.global_section_infos,
+                IsEmpty());
   }
 
   // The query should match nothing, since there is no vector with a
@@ -4246,13 +4320,14 @@
                                         .SetCardinality(CARDINALITY_OPTIONAL)))
           .Build(),
       /*ignore_errors_and_delete_documents=*/false));
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
 
   // Index a embedding vector.
   PropertyProto::VectorProto vector = CreateVector("my_model", {0.1, 0.2, 0.3});
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
-      BasicHit(kSectionId0, kDocumentId0), vector, QUANTIZATION_TYPE_NONE));
+      BasicHit(kSectionId0, kDocumentId0), vector, QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
 
   // Create an embedding query that has:
@@ -4282,22 +4357,19 @@
   EXPECT_THAT(
       query_results.embedding_query_results.GetMatchedScoresForDocument(
           /*query_vector_index=*/0, EMBEDDING_METRIC_COSINE, kDocumentId0),
-      Pointee(UnorderedElementsAre(DoubleNear(1, kEps))));
+      UnorderedElementsAre(DoubleNear(1, kEps)));
   if (GetParam().get_embedding_match_info) {
     // Check section match info for document 0.
     const EmbeddingMatchInfos* match_infos =
         query_results.embedding_query_results.GetMatchedInfosForDocument(
             /*query_vector_index=*/0, EMBEDDING_METRIC_COSINE, kDocumentId0);
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/1,
+    EXPECT_TRUE(ContainsMatchInfoEntry(query_results.embedding_query_results,
+                                       match_infos, /*score=*/1,
                                        /*position_in_section=*/0,
                                        /*section_id=*/kSectionId0));
   } else {
-    EXPECT_THAT(
-        query_results.embedding_query_results
-            .GetMatchedInfosForDocument(
-                /*query_vector_index=*/0, EMBEDDING_METRIC_COSINE, kDocumentId0)
-            ->section_infos,
-        IsNull());
+    EXPECT_THAT(*query_results.embedding_query_results.global_section_infos,
+                IsEmpty());
   }
 
   // Create a query that overrides the metric to DOT_PRODUCT.
@@ -4314,24 +4386,21 @@
   EXPECT_THAT(
       query_results.embedding_query_results.GetMatchedScoresForDocument(
           /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT, kDocumentId0),
-      Pointee(UnorderedElementsAre(DoubleNear(0.14, kEps))));
+      UnorderedElementsAre(DoubleNear(0.14, kEps)));
   if (GetParam().get_embedding_match_info) {
     // Check section match info for document 0.
     const EmbeddingMatchInfos* match_infos =
         query_results.embedding_query_results.GetMatchedInfosForDocument(
             /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT,
             kDocumentId0);
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/0.14,
+    EXPECT_TRUE(ContainsMatchInfoEntry(query_results.embedding_query_results,
+                                       match_infos, /*score=*/0.14,
                                        /*position_in_section=*/0,
                                        /*section_id=*/kSectionId0));
   } else {
     // Check section match info for document 0.
-    EXPECT_THAT(query_results.embedding_query_results
-                    .GetMatchedInfosForDocument(
-                        /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT,
-                        kDocumentId0)
-                    ->section_infos,
-                IsNull());
+    EXPECT_THAT(*query_results.embedding_query_results.global_section_infos,
+                IsEmpty());
   }
 
   // Create a query that overrides the metric to EUCLIDEAN.
@@ -4349,22 +4418,19 @@
   EXPECT_THAT(
       query_results.embedding_query_results.GetMatchedScoresForDocument(
           /*query_vector_index=*/0, EMBEDDING_METRIC_EUCLIDEAN, kDocumentId0),
-      Pointee(UnorderedElementsAre(DoubleNear(0, kEps))));
+      UnorderedElementsAre(DoubleNear(0, kEps)));
   if (GetParam().get_embedding_match_info) {
     // Check section match info for document 0.
     const EmbeddingMatchInfos* match_infos =
         query_results.embedding_query_results.GetMatchedInfosForDocument(
             /*query_vector_index=*/0, EMBEDDING_METRIC_EUCLIDEAN, kDocumentId0);
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/0,
+    EXPECT_TRUE(ContainsMatchInfoEntry(query_results.embedding_query_results,
+                                       match_infos, /*score=*/0,
                                        /*position_in_section=*/0,
                                        /*section_id=*/kSectionId0));
   } else {
-    EXPECT_THAT(query_results.embedding_query_results
-                    .GetMatchedInfosForDocument(
-                        /*query_vector_index=*/0, EMBEDDING_METRIC_EUCLIDEAN,
-                        kDocumentId0)
-                    ->section_infos,
-                IsNull());
+    EXPECT_THAT(*query_results.embedding_query_results.global_section_infos,
+                IsEmpty());
   }
 }
 
@@ -4386,46 +4452,56 @@
                                         .SetCardinality(CARDINALITY_REPEATED)))
           .Build(),
       /*ignore_errors_and_delete_documents=*/false));
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build()));
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build())));
 
   // Index vectors for document 0.
   // Section 0
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(kSectionId0, kDocumentId0),
-      CreateVector("my_model1", {1, -2, -3}), QUANTIZATION_TYPE_NONE));
+      CreateVector("my_model1", {1, -2, -3}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(kSectionId0, kDocumentId0),
-      CreateVector("my_model2", {1, -2, 3, -4}), QUANTIZATION_TYPE_NONE));
+      CreateVector("my_model2", {1, -2, 3, -4}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(kSectionId0, kDocumentId0),
-      CreateVector("my_model1", {-1, -2, 3}), QUANTIZATION_TYPE_NONE));
+      CreateVector("my_model1", {-1, -2, 3}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(kSectionId0, kDocumentId0),
-      CreateVector("my_model1", {1, -2, -4}), QUANTIZATION_TYPE_NONE));
+      CreateVector("my_model1", {1, -2, -4}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
   // Section 1
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(kSectionId1, kDocumentId0),
-      CreateVector("my_model1", {-1, -2, -3}), QUANTIZATION_TYPE_NONE));
+      CreateVector("my_model1", {-1, -2, -3}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(kSectionId1, kDocumentId0),
-      CreateVector("my_model1", {-1, 2, 4}), QUANTIZATION_TYPE_NONE));
+      CreateVector("my_model1", {-1, 2, 4}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
 
   // Index embedding vectors for document 1.
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(kSectionId0, kDocumentId1),
-      CreateVector("my_model1", {1, -2, 6}), QUANTIZATION_TYPE_NONE));
+      CreateVector("my_model1", {1, -2, 6}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(kSectionId0, kDocumentId1),
-      CreateVector("my_model2", {1, -2, 3, 4}), QUANTIZATION_TYPE_NONE));
+      CreateVector("my_model2", {1, -2, 3, 4}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(kSectionId0, kDocumentId1),
-      CreateVector("my_model1", {-1, -2, -6}), QUANTIZATION_TYPE_NONE));
+      CreateVector("my_model1", {-1, -2, -6}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(kSectionId0, kDocumentId1),
-      CreateVector("my_model2", {-1, -2, -3, -4}), QUANTIZATION_TYPE_NONE));
+      CreateVector("my_model2", {-1, -2, -3, -4}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
 
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
 
@@ -4486,11 +4562,11 @@
   EXPECT_THAT(
       query_results.embedding_query_results.GetMatchedScoresForDocument(
           /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT, kDocumentId1),
-      Pointee(UnorderedElementsAre(-3, 7)));
+      UnorderedElementsAre(-3, 7));
   EXPECT_THAT(
       query_results.embedding_query_results.GetMatchedScoresForDocument(
           /*query_vector_index=*/1, EMBEDDING_METRIC_DOT_PRODUCT, kDocumentId1),
-      Pointee(UnorderedElementsAre(6)));
+      UnorderedElementsAre(6));
 
   // Check results for document 0.
   ICING_ASSERT_OK(itr->Advance());
@@ -4500,11 +4576,11 @@
   EXPECT_THAT(
       query_results.embedding_query_results.GetMatchedScoresForDocument(
           /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT, kDocumentId0),
-      Pointee(UnorderedElementsAre(-2, 0, -3, 6, 3)));
+      UnorderedElementsAre(-2, 0, -3, 6, 3));
   EXPECT_THAT(
       query_results.embedding_query_results.GetMatchedScoresForDocument(
           /*query_vector_index=*/1, EMBEDDING_METRIC_DOT_PRODUCT, kDocumentId0),
-      Pointee(UnorderedElementsAre(-2)));
+      UnorderedElementsAre(-2));
   EXPECT_THAT(itr->Advance(),
               StatusIs(libtextclassifier3::StatusCode::RESOURCE_EXHAUSTED));
 
@@ -4523,26 +4599,32 @@
         query_results.embedding_query_results.GetMatchedInfosForDocument(
             /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT,
             kDocumentId0);
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/-2,
+    EXPECT_TRUE(ContainsMatchInfoEntry(query_results.embedding_query_results,
+                                       match_infos, /*score=*/-2,
                                        /*position_in_section=*/0,
                                        /*section_id=*/kSectionId0));
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/-3,
+    EXPECT_TRUE(ContainsMatchInfoEntry(query_results.embedding_query_results,
+                                       match_infos, /*score=*/-3,
                                        /*position_in_section=*/2,
                                        /*section_id=*/kSectionId0));
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/6,
+    EXPECT_TRUE(ContainsMatchInfoEntry(query_results.embedding_query_results,
+                                       match_infos, /*score=*/6,
                                        /*position_in_section=*/1,
                                        /*section_id=*/kSectionId0));
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/0,
+    EXPECT_TRUE(ContainsMatchInfoEntry(query_results.embedding_query_results,
+                                       match_infos, /*score=*/0,
                                        /*position_in_section=*/0,
                                        /*section_id=*/kSectionId1));
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/3,
+    EXPECT_TRUE(ContainsMatchInfoEntry(query_results.embedding_query_results,
+                                       match_infos, /*score=*/3,
                                        /*position_in_section=*/1,
                                        /*section_id=*/kSectionId1));
     match_infos =
         query_results.embedding_query_results.GetMatchedInfosForDocument(
             /*query_vector_index=*/1, EMBEDDING_METRIC_DOT_PRODUCT,
             kDocumentId0);
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/-2,
+    EXPECT_TRUE(ContainsMatchInfoEntry(query_results.embedding_query_results,
+                                       match_infos, /*score=*/-2,
                                        /*position_in_section=*/0,
                                        /*section_id=*/kSectionId0));
 
@@ -4557,44 +4639,25 @@
         query_results.embedding_query_results.GetMatchedInfosForDocument(
             /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT,
             kDocumentId1);
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/-3,
+    EXPECT_TRUE(ContainsMatchInfoEntry(query_results.embedding_query_results,
+                                       match_infos, /*score=*/-3,
                                        /*position_in_section=*/1,
                                        /*section_id=*/kSectionId0));
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/7,
+    EXPECT_TRUE(ContainsMatchInfoEntry(query_results.embedding_query_results,
+                                       match_infos, /*score=*/7,
                                        /*position_in_section=*/0,
                                        /*section_id=*/kSectionId0));
     match_infos =
         query_results.embedding_query_results.GetMatchedInfosForDocument(
             /*query_vector_index=*/1, EMBEDDING_METRIC_DOT_PRODUCT,
             kDocumentId1);
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/6,
+    EXPECT_TRUE(ContainsMatchInfoEntry(query_results.embedding_query_results,
+                                       match_infos, /*score=*/6,
                                        /*position_in_section=*/1,
                                        /*section_id=*/kSectionId0));
   } else {
-    EXPECT_THAT(query_results.embedding_query_results
-                    .GetMatchedInfosForDocument(
-                        /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT,
-                        kDocumentId0)
-                    ->section_infos,
-                IsNull());
-    EXPECT_THAT(query_results.embedding_query_results
-                    .GetMatchedInfosForDocument(
-                        /*query_vector_index=*/1, EMBEDDING_METRIC_DOT_PRODUCT,
-                        kDocumentId0)
-                    ->section_infos,
-                IsNull());
-    EXPECT_THAT(query_results.embedding_query_results
-                    .GetMatchedInfosForDocument(
-                        /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT,
-                        kDocumentId1)
-                    ->section_infos,
-                IsNull());
-    EXPECT_THAT(query_results.embedding_query_results
-                    .GetMatchedInfosForDocument(
-                        /*query_vector_index=*/1, EMBEDDING_METRIC_DOT_PRODUCT,
-                        kDocumentId1)
-                    ->section_infos,
-                IsNull());
+    EXPECT_THAT(*query_results.embedding_query_results.global_section_infos,
+                IsEmpty());
   }
 }
 
@@ -4621,28 +4684,33 @@
                                         .SetCardinality(CARDINALITY_OPTIONAL)))
           .Build(),
       /*ignore_errors_and_delete_documents=*/false));
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build()));
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build())));
 
   // Index 3 embedding vectors for document 0.
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(kSectionId0, kDocumentId0),
-      CreateVector("my_model1", {1, -2, -3}), QUANTIZATION_TYPE_NONE));
+      CreateVector("my_model1", {1, -2, -3}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(kSectionId1, kDocumentId0),
-      CreateVector("my_model1", {-1, -2, -3}), QUANTIZATION_TYPE_NONE));
+      CreateVector("my_model1", {-1, -2, -3}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(kSectionId2, kDocumentId0),
-      CreateVector("my_model2", {-1, 2, 3, -4}), QUANTIZATION_TYPE_NONE));
+      CreateVector("my_model2", {-1, 2, 3, -4}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
   // Index 2 embedding vectors for document 1.
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(kSectionId0, kDocumentId1),
-      CreateVector("my_model1", {-1, -2, 3}), QUANTIZATION_TYPE_NONE));
+      CreateVector("my_model1", {-1, -2, 3}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(kSectionId1, kDocumentId1),
-      CreateVector("my_model2", {1, -2, 3, -4}), QUANTIZATION_TYPE_NONE));
+      CreateVector("my_model2", {1, -2, 3, -4}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
 
   // Create two embedding queries.
@@ -4685,11 +4753,11 @@
   EXPECT_THAT(
       query_results.embedding_query_results.GetMatchedScoresForDocument(
           /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT, kDocumentId0),
-      Pointee(UnorderedElementsAre(-2, 0)));
+      UnorderedElementsAre(-2, 0));
   EXPECT_THAT(
       query_results.embedding_query_results.GetMatchedScoresForDocument(
           /*query_vector_index=*/1, EMBEDDING_METRIC_DOT_PRODUCT, kDocumentId0),
-      Pointee(UnorderedElementsAre(4)));
+      UnorderedElementsAre(4));
   EXPECT_THAT(itr->Advance(),
               StatusIs(libtextclassifier3::StatusCode::RESOURCE_EXHAUSTED));
   if (GetParam().get_embedding_match_info) {
@@ -4698,32 +4766,25 @@
         query_results.embedding_query_results.GetMatchedInfosForDocument(
             /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT,
             kDocumentId0);
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/-2,
+    EXPECT_TRUE(ContainsMatchInfoEntry(query_results.embedding_query_results,
+                                       match_infos, /*score=*/-2,
                                        /*position_in_section=*/0,
                                        /*section_id=*/kSectionId0));
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/0,
+    EXPECT_TRUE(ContainsMatchInfoEntry(query_results.embedding_query_results,
+                                       match_infos, /*score=*/0,
                                        /*position_in_section=*/0,
                                        /*section_id=*/kSectionId1));
     match_infos =
         query_results.embedding_query_results.GetMatchedInfosForDocument(
             /*query_vector_index=*/1, EMBEDDING_METRIC_DOT_PRODUCT,
             kDocumentId0);
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/4,
+    EXPECT_TRUE(ContainsMatchInfoEntry(query_results.embedding_query_results,
+                                       match_infos, /*score=*/4,
                                        /*position_in_section=*/0,
                                        /*section_id=*/kSectionId2));
   } else {
-    EXPECT_THAT(query_results.embedding_query_results
-                    .GetMatchedInfosForDocument(
-                        /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT,
-                        kDocumentId0)
-                    ->section_infos,
-                IsNull());
-    EXPECT_THAT(query_results.embedding_query_results
-                    .GetMatchedInfosForDocument(
-                        /*query_vector_index=*/1, EMBEDDING_METRIC_DOT_PRODUCT,
-                        kDocumentId0)
-                    ->section_infos,
-                IsNull());
+    EXPECT_THAT(*query_results.embedding_query_results.global_section_infos,
+                IsEmpty());
   }
 
   // The query can match both document 0 and document 1:
@@ -4756,11 +4817,11 @@
   EXPECT_THAT(
       query_results.embedding_query_results.GetMatchedScoresForDocument(
           /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT, kDocumentId1),
-      Pointee(UnorderedElementsAre(6)));
+      UnorderedElementsAre(6));
   EXPECT_THAT(
       query_results.embedding_query_results.GetMatchedScoresForDocument(
           /*query_vector_index=*/1, EMBEDDING_METRIC_DOT_PRODUCT, kDocumentId1),
-      IsNull());
+      IsEmpty());
   // Check results for document 0.
   ICING_ASSERT_OK(itr->Advance());
   EXPECT_THAT(
@@ -4769,11 +4830,11 @@
   EXPECT_THAT(
       query_results.embedding_query_results.GetMatchedScoresForDocument(
           /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT, kDocumentId0),
-      IsNull());
+      IsEmpty());
   EXPECT_THAT(
       query_results.embedding_query_results.GetMatchedScoresForDocument(
           /*query_vector_index=*/1, EMBEDDING_METRIC_DOT_PRODUCT, kDocumentId0),
-      Pointee(UnorderedElementsAre(4)));
+      UnorderedElementsAre(4));
   EXPECT_THAT(itr->Advance(),
               StatusIs(libtextclassifier3::StatusCode::RESOURCE_EXHAUSTED));
   if (GetParam().get_embedding_match_info) {
@@ -4787,7 +4848,8 @@
         query_results.embedding_query_results.GetMatchedInfosForDocument(
             /*query_vector_index=*/1, EMBEDDING_METRIC_DOT_PRODUCT,
             kDocumentId0);
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/4,
+    EXPECT_TRUE(ContainsMatchInfoEntry(query_results.embedding_query_results,
+                                       match_infos, /*score=*/4,
                                        /*position_in_section=*/0,
                                        /*section_id=*/kSectionId2));
 
@@ -4796,7 +4858,8 @@
         query_results.embedding_query_results.GetMatchedInfosForDocument(
             /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT,
             kDocumentId1);
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/6,
+    EXPECT_TRUE(ContainsMatchInfoEntry(query_results.embedding_query_results,
+                                       match_infos, /*score=*/6,
                                        /*position_in_section=*/0,
                                        /*section_id=*/kSectionId0));
     match_infos =
@@ -4805,18 +4868,8 @@
             kDocumentId1);
     EXPECT_THAT(match_infos, IsNull());
   } else {
-    EXPECT_THAT(query_results.embedding_query_results
-                    .GetMatchedInfosForDocument(
-                        /*query_vector_index=*/1, EMBEDDING_METRIC_DOT_PRODUCT,
-                        kDocumentId0)
-                    ->section_infos,
-                IsNull());
-    EXPECT_THAT(query_results.embedding_query_results
-                    .GetMatchedInfosForDocument(
-                        /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT,
-                        kDocumentId1)
-                    ->section_infos,
-                IsNull());
+    EXPECT_THAT(*query_results.embedding_query_results.global_section_infos,
+                IsEmpty());
   }
 }
 
@@ -4839,22 +4892,25 @@
                                         .SetCardinality(CARDINALITY_OPTIONAL)))
           .Build(),
       /*ignore_errors_and_delete_documents=*/false));
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build()));
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build())));
 
   // Index 3 embedding vectors for document 0.
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(kSectionId0, kDocumentId0),
-      CreateVector("my_model1", {1, -2, -3}), QUANTIZATION_TYPE_NONE));
+      CreateVector("my_model1", {1, -2, -3}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(kSectionId1, kDocumentId0),
-      CreateVector("my_model1", {-1, -2, -3}), QUANTIZATION_TYPE_NONE));
+      CreateVector("my_model1", {-1, -2, -3}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
   // Index 2 embedding vectors for document 1.
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(kSectionId0, kDocumentId1),
-      CreateVector("my_model1", {-1, -2, 3}), QUANTIZATION_TYPE_NONE));
+      CreateVector("my_model1", {-1, -2, 3}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
 
   // Create two embedding queries.
@@ -4889,7 +4945,7 @@
   EXPECT_THAT(
       query_results.embedding_query_results.GetMatchedScoresForDocument(
           /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT, kDocumentId1),
-      Pointee(UnorderedElementsAre(6)));
+      UnorderedElementsAre(6));
   // Check results for document 0.
   ICING_ASSERT_OK(itr->Advance());
   EXPECT_THAT(itr->doc_hit_info(),
@@ -4898,7 +4954,7 @@
   EXPECT_THAT(
       query_results.embedding_query_results.GetMatchedScoresForDocument(
           /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT, kDocumentId0),
-      Pointee(UnorderedElementsAre(-2, 0)));
+      UnorderedElementsAre(-2, 0));
   EXPECT_THAT(itr->Advance(),
               StatusIs(libtextclassifier3::StatusCode::RESOURCE_EXHAUSTED));
   if (GetParam().get_embedding_match_info) {
@@ -4907,10 +4963,12 @@
         query_results.embedding_query_results.GetMatchedInfosForDocument(
             /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT,
             kDocumentId0);
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/-2,
+    EXPECT_TRUE(ContainsMatchInfoEntry(query_results.embedding_query_results,
+                                       match_infos, /*score=*/-2,
                                        /*position_in_section=*/0,
                                        /*section_id=*/kSectionId0));
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/0,
+    EXPECT_TRUE(ContainsMatchInfoEntry(query_results.embedding_query_results,
+                                       match_infos, /*score=*/0,
                                        /*position_in_section=*/0,
                                        /*section_id=*/kSectionId1));
     match_infos =
@@ -4924,7 +4982,8 @@
         query_results.embedding_query_results.GetMatchedInfosForDocument(
             /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT,
             kDocumentId1);
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/6,
+    EXPECT_TRUE(ContainsMatchInfoEntry(query_results.embedding_query_results,
+                                       match_infos, /*score=*/6,
                                        /*position_in_section=*/0,
                                        /*section_id=*/kSectionId0));
     match_infos =
@@ -4933,18 +4992,8 @@
             kDocumentId0);
     EXPECT_THAT(match_infos, IsNull());
   } else {
-    EXPECT_THAT(query_results.embedding_query_results
-                    .GetMatchedInfosForDocument(
-                        /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT,
-                        kDocumentId0)
-                    ->section_infos,
-                IsNull());
-    EXPECT_THAT(query_results.embedding_query_results
-                    .GetMatchedInfosForDocument(
-                        /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT,
-                        kDocumentId1)
-                    ->section_infos,
-                IsNull());
+    EXPECT_THAT(*query_results.embedding_query_results.global_section_infos,
+                IsEmpty());
   }
 
   // The same query appears twice, in which case all the scores in the results
@@ -4968,7 +5017,7 @@
   EXPECT_THAT(
       query_results.embedding_query_results.GetMatchedScoresForDocument(
           /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT, kDocumentId1),
-      Pointee(UnorderedElementsAre(6, 6)));
+      UnorderedElementsAre(6, 6));
   // Check results for document 0.
   ICING_ASSERT_OK(itr->Advance());
   EXPECT_THAT(itr->doc_hit_info(),
@@ -4977,7 +5026,7 @@
   EXPECT_THAT(
       query_results.embedding_query_results.GetMatchedScoresForDocument(
           /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT, kDocumentId0),
-      Pointee(UnorderedElementsAre(-2, 0, -2, 0)));
+      UnorderedElementsAre(-2, 0, -2, 0));
   EXPECT_THAT(itr->Advance(),
               StatusIs(libtextclassifier3::StatusCode::RESOURCE_EXHAUSTED));
 }
@@ -5000,10 +5049,10 @@
                                         .SetCardinality(CARDINALITY_OPTIONAL)))
           .Build(),
       /*ignore_errors_and_delete_documents=*/false));
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build()));
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build())));
 
   // Index terms
   Index::Editor editor = index_->Edit(kDocumentId0, kSectionId1,
@@ -5018,10 +5067,12 @@
   // Index embedding vectors
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(kSectionId0, kDocumentId0),
-      CreateVector("my_model1", {1, -2, -3}), QUANTIZATION_TYPE_NONE));
+      CreateVector("my_model1", {1, -2, -3}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(kSectionId0, kDocumentId1),
-      CreateVector("my_model1", {-1, -2, 3}), QUANTIZATION_TYPE_NONE));
+      CreateVector("my_model1", {-1, -2, 3}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
 
   // Create an embedding query with semantic scores:
@@ -5057,7 +5108,7 @@
   EXPECT_THAT(
       query_results.embedding_query_results.GetMatchedScoresForDocument(
           /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT, kDocumentId1),
-      Pointee(UnorderedElementsAre(6)));
+      UnorderedElementsAre(6));
   // Check results for document 0.
   ICING_ASSERT_OK(itr->Advance());
   EXPECT_THAT(
@@ -5066,7 +5117,7 @@
   EXPECT_THAT(
       query_results.embedding_query_results.GetMatchedScoresForDocument(
           /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT, kDocumentId0),
-      IsNull());
+      IsEmpty());
   EXPECT_THAT(itr->Advance(),
               StatusIs(libtextclassifier3::StatusCode::RESOURCE_EXHAUSTED));
   if (GetParam().get_embedding_match_info) {
@@ -5082,16 +5133,13 @@
         query_results.embedding_query_results.GetMatchedInfosForDocument(
             /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT,
             kDocumentId1);
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/6,
+    EXPECT_TRUE(ContainsMatchInfoEntry(query_results.embedding_query_results,
+                                       match_infos, /*score=*/6,
                                        /*position_in_section=*/0,
                                        /*section_id=*/kSectionId0));
   } else {
-    EXPECT_THAT(query_results.embedding_query_results
-                    .GetMatchedInfosForDocument(
-                        /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT,
-                        kDocumentId1)
-                    ->section_infos,
-                IsNull());
+    EXPECT_THAT(*query_results.embedding_query_results.global_section_infos,
+                IsEmpty());
   }
 
   // Perform another hybrid search:
@@ -5118,7 +5166,7 @@
   EXPECT_THAT(
       query_results.embedding_query_results.GetMatchedScoresForDocument(
           /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT, kDocumentId0),
-      Pointee(UnorderedElementsAre(-2)));
+      UnorderedElementsAre(-2));
   EXPECT_THAT(itr->Advance(),
               StatusIs(libtextclassifier3::StatusCode::RESOURCE_EXHAUSTED));
 
@@ -5128,31 +5176,13 @@
         query_results.embedding_query_results.GetMatchedInfosForDocument(
             /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT,
             kDocumentId0);
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/-2,
-                                       /*position_in_section=*/0,
-                                       /*section_id=*/kSectionId0));
-
-    // Check section match info for document 1.
-    match_infos =
-        query_results.embedding_query_results.GetMatchedInfosForDocument(
-            /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT,
-            kDocumentId1);
-    EXPECT_TRUE(ContainsMatchInfoEntry(match_infos, /*score=*/6,
+    EXPECT_TRUE(ContainsMatchInfoEntry(query_results.embedding_query_results,
+                                       match_infos, /*score=*/-2,
                                        /*position_in_section=*/0,
                                        /*section_id=*/kSectionId0));
   } else {
-    EXPECT_THAT(query_results.embedding_query_results
-                    .GetMatchedInfosForDocument(
-                        /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT,
-                        kDocumentId0)
-                    ->section_infos,
-                IsNull());
-    EXPECT_THAT(query_results.embedding_query_results
-                    .GetMatchedInfosForDocument(
-                        /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT,
-                        kDocumentId1)
-                    ->section_infos,
-                IsNull());
+    EXPECT_THAT(*query_results.embedding_query_results.global_section_infos,
+                IsEmpty());
   }
 }
 
@@ -5175,23 +5205,27 @@
       /*ignore_errors_and_delete_documents=*/false));
 
   // Create two documents.
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build()));
-  ICING_ASSERT_OK(document_store_->Put(
-      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build()));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build())));
   // Add embedding vectors into different sections for the two documents.
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(kSectionId0, kDocumentId0),
-      CreateVector("my_model1", {1, -2, -3}), QUANTIZATION_TYPE_NONE));
+      CreateVector("my_model1", {1, -2, -3}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(kSectionId1, kDocumentId0),
-      CreateVector("my_model1", {-1, -2, 3}), QUANTIZATION_TYPE_NONE));
+      CreateVector("my_model1", {-1, -2, 3}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(kSectionId0, kDocumentId1),
-      CreateVector("my_model1", {-1, 2, 3}), QUANTIZATION_TYPE_NONE));
+      CreateVector("my_model1", {-1, 2, 3}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
   ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
       BasicHit(kSectionId1, kDocumentId1),
-      CreateVector("my_model1", {1, 2, -3}), QUANTIZATION_TYPE_NONE));
+      CreateVector("my_model1", {1, 2, -3}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
   ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
 
   // Create an embedding query with semantic scores:
@@ -5223,7 +5257,7 @@
   EXPECT_THAT(
       query_results.embedding_query_results.GetMatchedScoresForDocument(
           /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT, kDocumentId1),
-      Pointee(UnorderedElementsAre(2)));
+      UnorderedElementsAre(2));
   ICING_ASSERT_OK(itr->Advance());
   EXPECT_THAT(
       itr->doc_hit_info(),
@@ -5231,11 +5265,171 @@
   EXPECT_THAT(
       query_results.embedding_query_results.GetMatchedScoresForDocument(
           /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT, kDocumentId0),
-      Pointee(UnorderedElementsAre(-2)));
+      UnorderedElementsAre(-2));
   EXPECT_THAT(itr->Advance(),
               StatusIs(libtextclassifier3::StatusCode::RESOURCE_EXHAUSTED));
 }
 
+TEST_P(QueryVisitorTest, SemanticSearchFunctionDeletedDocument) {
+  // Set up schema
+  ICING_ASSERT_OK(schema_store_->SetSchema(
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("type").AddProperty(
+              PropertyConfigBuilder()
+                  .SetName("prop1")
+                  .SetDataTypeVector(EMBEDDING_INDEXING_LINEAR_SEARCH)
+                  .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build(),
+      /*ignore_errors_and_delete_documents=*/false));
+
+  // Add two documents
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build())));
+
+  // Index embedding vectors for both documents
+  PropertyProto::VectorProto vector0 =
+      CreateVector("my_model", {0.1, 0.2, 0.3});
+  PropertyProto::VectorProto vector1 =
+      CreateVector("my_model", {0.4, 0.5, 0.6});
+  ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
+      BasicHit(kSectionId0, kDocumentId0), vector0, QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
+  ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
+      BasicHit(kSectionId0, kDocumentId1), vector1, QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
+  ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
+
+  // Delete kDocumentId0
+  ICING_ASSERT_OK(document_store_->Delete("ns", "uri0",
+                                          clock_.GetSystemTimeMilliseconds()));
+
+  // Create an embedding query that would match both documents if kDocumentId0
+  // was not deleted.
+  std::vector<PropertyProto::VectorProto> embedding_query_vectors = {
+      CreateVector("my_model", {1.0, 1.0, 1.0})};
+  std::string query = "semanticSearch(getEmbeddingParameter(0), -1)";
+  ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<Node> root_node,
+                             ParseQueryHelper(query));
+  SearchSpecProto search_spec =
+      CreateSearchSpec(query, TERM_MATCH_PREFIX, embedding_query_vectors,
+                       EMBEDDING_METRIC_DOT_PRODUCT);
+  ICING_ASSERT_OK_AND_ASSIGN(QueryResults query_results,
+                             ProcessQuery(search_spec, root_node.get()));
+
+  EXPECT_THAT(ExtractKeys(query_results.query_term_iterators), IsEmpty());
+  EXPECT_THAT(query_results.query_terms, IsEmpty());
+  EXPECT_THAT(query_results.features_in_use,
+              UnorderedElementsAre(kListFilterQueryLanguageFeature));
+
+  // Only kDocumentId1 should be returned
+  EXPECT_THAT(GetDocumentIds(query_results.root_iterator.get()),
+              ElementsAre(kDocumentId1));
+
+  // Check that scores for kDocumentId1 are present and kDocumentId0 are not.
+  EXPECT_THAT(
+      query_results.embedding_query_results.GetMatchedScoresForDocument(
+          /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT, kDocumentId1),
+      UnorderedElementsAre(DoubleNear(0.4 + 0.5 + 0.6, kEps)));
+  EXPECT_THAT(
+      query_results.embedding_query_results.GetMatchedScoresForDocument(
+          /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT, kDocumentId0),
+      IsEmpty());
+}
+
+TEST_P(QueryVisitorTest, SemanticSearchFunctionCallStats) {
+  if (!GetParam().enable_embedding_iterator_v2) {
+    // Call stats are only populated in embedding iterator v2.
+    return;
+  }
+
+  // Set up
+  ICING_ASSERT_OK(schema_store_->SetSchema(
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("type")
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("prop1")
+                                        .SetDataTypeVector(
+                                            EMBEDDING_INDEXING_LINEAR_SEARCH)
+                                        .SetCardinality(CARDINALITY_OPTIONAL))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("prop2")
+                                        .SetDataTypeVector(
+                                            EMBEDDING_INDEXING_LINEAR_SEARCH,
+                                            QUANTIZATION_TYPE_QUANTIZE_8_BIT)
+                                        .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build(),
+      /*ignore_errors_and_delete_documents=*/false));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri0").SetSchema("type").Build())));
+  ICING_ASSERT_OK(document_store_->Put(document_util::CreateDocumentWrapper(
+      DocumentBuilder().SetKey("ns", "uri1").SetSchema("type").Build())));
+
+  // Index 2 unquantized embedding vectors, and 1 quantized embedding vector.
+  PropertyProto::VectorProto vector0 =
+      CreateVector("my_model", {1., 1., 1.});  // Size: 12 bytes
+  PropertyProto::VectorProto vector1 =
+      CreateVector("my_model", {-1., -1., -1.});  // Size: 12 bytes
+  PropertyProto::VectorProto vector2 = CreateVector(
+      "my_model", {0., 0., 0.});  // Size: 3 + sizeof(Quantizer) = 11 bytes.
+  ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
+      BasicHit(kSectionId0, kDocumentId0), vector0, QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
+  ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
+      BasicHit(kSectionId0, kDocumentId1), vector1, QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"type"));
+  ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
+      BasicHit(kSectionId1, kDocumentId1), vector2,
+      QUANTIZATION_TYPE_QUANTIZE_8_BIT,
+      /*schema_name=*/"type"));
+  ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
+
+  // Create a query that matches all embeddings.
+  std::vector<PropertyProto::VectorProto> embedding_query_vectors = {
+      CreateVector("my_model", {1., 1., 1.})};
+  std::string query = "semanticSearch(getEmbeddingParameter(0))";
+  ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<Node> root_node,
+                             ParseQueryHelper(query));
+  SearchSpecProto search_spec =
+      CreateSearchSpec(query, TERM_MATCH_PREFIX, embedding_query_vectors,
+                       EMBEDDING_METRIC_DOT_PRODUCT);
+  ICING_ASSERT_OK_AND_ASSIGN(QueryResults query_results,
+                             ProcessQuery(search_spec, root_node.get()));
+  EXPECT_THAT(ExtractKeys(query_results.query_term_iterators), IsEmpty());
+  EXPECT_THAT(query_results.query_terms, IsEmpty());
+  EXPECT_THAT(query_results.features_in_use,
+              UnorderedElementsAre(kListFilterQueryLanguageFeature));
+  EXPECT_THAT(GetDocumentIds(query_results.root_iterator.get()),
+              ElementsAre(kDocumentId1, kDocumentId0));
+  EXPECT_THAT(
+      query_results.embedding_query_results.GetMatchedScoresForDocument(
+          /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT, kDocumentId0),
+      UnorderedElementsAre(DoubleNear(3., kEps)));
+  EXPECT_THAT(
+      query_results.embedding_query_results.GetMatchedScoresForDocument(
+          /*query_vector_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT, kDocumentId1),
+      UnorderedElementsAre(DoubleNear(-3., kEps), DoubleNear(0, kEps)));
+
+  // Check call stats.
+  // We should see 2 unquantized embeddings scored and 1 quantized embedding
+  // scored. The size read should be 12 + 12 + 11 = 35 bytes.
+  EXPECT_THAT(query_results.root_iterator->GetCallStats(),
+              EqualsDocHitInfoIteratorCallStats(
+                  /*num_leaf_advance_calls_lite_index=*/2,
+                  /*num_leaf_advance_calls_main_index=*/0,
+                  /*num_leaf_advance_calls_integer_index=*/0,
+                  /*num_leaf_advance_calls_no_index=*/0,
+                  /*num_blocks_inspected=*/0,
+                  DocHitInfoIterator::CallStats::EmbeddingStats{
+                      .num_unquantized_embeddings_scored = 2,
+                      .num_quantized_embeddings_scored = 1,
+                      .unquantized_shards_read = {5},
+                      .quantized_shards_read = {5},
+                      .num_embedding_bytes_read = 35}));
+}
+
 TEST_P(QueryVisitorTest,
        MatchScoreExpressionFunctionWithInvalidRangeReturnsInvalidArgument) {
   // The expression is invalid, since low > high.
@@ -5258,16 +5452,18 @@
           .AddType(SchemaTypeConfigBuilder().SetType("Simple"))
           .Build(),
       /*ignore_errors_and_delete_documents=*/false));
-  ICING_ASSERT_OK(document_store_->Put(DocumentBuilder()
-                                           .SetKey("ns", "uri0")
-                                           .SetSchema("Simple")
-                                           .SetScore(4)
-                                           .Build()));
-  ICING_ASSERT_OK(document_store_->Put(DocumentBuilder()
-                                           .SetKey("ns", "uri1")
-                                           .SetSchema("Simple")
-                                           .SetScore(0)
-                                           .Build()));
+  ICING_ASSERT_OK(document_store_->Put(
+      document_util::CreateDocumentWrapper(DocumentBuilder()
+                                               .SetKey("ns", "uri0")
+                                               .SetSchema("Simple")
+                                               .SetScore(4)
+                                               .Build())));
+  ICING_ASSERT_OK(document_store_->Put(
+      document_util::CreateDocumentWrapper(DocumentBuilder()
+                                               .SetKey("ns", "uri1")
+                                               .SetSchema("Simple")
+                                               .SetScore(0)
+                                               .Build())));
 
   // Test that `matchScoreExpression("this.documentScore()", 0)` matches
   // all documents.
@@ -5309,16 +5505,18 @@
           .AddType(SchemaTypeConfigBuilder().SetType("Simple"))
           .Build(),
       /*ignore_errors_and_delete_documents=*/false));
-  ICING_ASSERT_OK(document_store_->Put(DocumentBuilder()
-                                           .SetKey("ns", "uri0")
-                                           .SetSchema("Simple")
-                                           .SetScore(4)
-                                           .Build()));
-  ICING_ASSERT_OK(document_store_->Put(DocumentBuilder()
-                                           .SetKey("ns", "uri1")
-                                           .SetSchema("Simple")
-                                           .SetScore(0)
-                                           .Build()));
+  ICING_ASSERT_OK(document_store_->Put(
+      document_util::CreateDocumentWrapper(DocumentBuilder()
+                                               .SetKey("ns", "uri0")
+                                               .SetSchema("Simple")
+                                               .SetScore(4)
+                                               .Build())));
+  ICING_ASSERT_OK(document_store_->Put(
+      document_util::CreateDocumentWrapper(DocumentBuilder()
+                                               .SetKey("ns", "uri1")
+                                               .SetSchema("Simple")
+                                               .SetScore(0)
+                                               .Build())));
 
   // Test that `matchScoreExpression("this.documentScore()", -100, 100)` matches
   // all documents.
@@ -5360,24 +5558,27 @@
           .AddType(SchemaTypeConfigBuilder().SetType("Simple"))
           .Build(),
       /*ignore_errors_and_delete_documents=*/false));
-  ICING_ASSERT_OK(document_store_->Put(DocumentBuilder()
-                                           .SetKey("ns", "uri0")
-                                           .SetSchema("Simple")
-                                           .SetScore(4)
-                                           .SetCreationTimestampMs(1)
-                                           .Build()));
-  ICING_ASSERT_OK(document_store_->Put(DocumentBuilder()
-                                           .SetKey("ns", "uri1")
-                                           .SetSchema("Simple")
-                                           .SetScore(2)
-                                           .SetCreationTimestampMs(2)
-                                           .Build()));
-  ICING_ASSERT_OK(document_store_->Put(DocumentBuilder()
-                                           .SetKey("ns", "uri2")
-                                           .SetSchema("Simple")
-                                           .SetScore(0)
-                                           .SetCreationTimestampMs(3)
-                                           .Build()));
+  ICING_ASSERT_OK(document_store_->Put(
+      document_util::CreateDocumentWrapper(DocumentBuilder()
+                                               .SetKey("ns", "uri0")
+                                               .SetSchema("Simple")
+                                               .SetScore(4)
+                                               .SetCreationTimestampMs(1)
+                                               .Build())));
+  ICING_ASSERT_OK(document_store_->Put(
+      document_util::CreateDocumentWrapper(DocumentBuilder()
+                                               .SetKey("ns", "uri1")
+                                               .SetSchema("Simple")
+                                               .SetScore(2)
+                                               .SetCreationTimestampMs(2)
+                                               .Build())));
+  ICING_ASSERT_OK(document_store_->Put(
+      document_util::CreateDocumentWrapper(DocumentBuilder()
+                                               .SetKey("ns", "uri2")
+                                               .SetSchema("Simple")
+                                               .SetScore(0)
+                                               .SetCreationTimestampMs(3)
+                                               .Build())));
 
   // Query with a complex score expression:
   //   `pow(this.creationTimestamp(), 2) + this.documentScore() - 1`.
@@ -5404,16 +5605,18 @@
           .AddType(SchemaTypeConfigBuilder().SetType("Simple"))
           .Build(),
       /*ignore_errors_and_delete_documents=*/false));
-  ICING_ASSERT_OK(document_store_->Put(DocumentBuilder()
-                                           .SetKey("ns", "uri0")
-                                           .SetSchema("Simple")
-                                           .SetScore(4)
-                                           .Build()));
-  ICING_ASSERT_OK(document_store_->Put(DocumentBuilder()
-                                           .SetKey("ns", "uri1")
-                                           .SetSchema("Simple")
-                                           .SetScore(0)
-                                           .Build()));
+  ICING_ASSERT_OK(document_store_->Put(
+      document_util::CreateDocumentWrapper(DocumentBuilder()
+                                               .SetKey("ns", "uri0")
+                                               .SetSchema("Simple")
+                                               .SetScore(4)
+                                               .Build())));
+  ICING_ASSERT_OK(document_store_->Put(
+      document_util::CreateDocumentWrapper(DocumentBuilder()
+                                               .SetKey("ns", "uri1")
+                                               .SetSchema("Simple")
+                                               .SetScore(0)
+                                               .Build())));
 
   // Test that documents with evaluation errors will be filtered out.
   // Specifically, document1 will be filtered out because its document score is
@@ -5431,10 +5634,71 @@
 
 INSTANTIATE_TEST_SUITE_P(
     QueryVisitorTest, QueryVisitorTest,
-    testing::Values(QueryVisitorTestParams(QueryType::kSearch, true),
-                    QueryVisitorTestParams(QueryType::kSearch, false),
-                    QueryVisitorTestParams(QueryType::kPlain, true),
-                    QueryVisitorTestParams(QueryType::kPlain, false)));
+    testing::Values(
+        QueryVisitorTestParams(QueryType::kSearch,
+                               /*get_embedding_match_info=*/true,
+                               /*enable_embedding_iterator_v2=*/true,
+                               /*enable_embed_query_optimization=*/false),
+        QueryVisitorTestParams(QueryType::kSearch,
+                               /*get_embedding_match_info=*/true,
+                               /*enable_embedding_iterator_v2=*/false,
+                               /*enable_embed_query_optimization=*/false),
+        QueryVisitorTestParams(QueryType::kSearch,
+                               /*get_embedding_match_info=*/false,
+                               /*enable_embedding_iterator_v2=*/true,
+                               /*enable_embed_query_optimization=*/false),
+        QueryVisitorTestParams(QueryType::kSearch,
+                               /*get_embedding_match_info=*/false,
+                               /*enable_embedding_iterator_v2=*/false,
+                               /*enable_embed_query_optimization=*/false),
+        QueryVisitorTestParams(QueryType::kPlain,
+                               /*get_embedding_match_info=*/true,
+                               /*enable_embedding_iterator_v2=*/true,
+                               /*enable_embed_query_optimization=*/false),
+        QueryVisitorTestParams(QueryType::kPlain,
+                               /*get_embedding_match_info=*/true,
+                               /*enable_embedding_iterator_v2=*/false,
+                               /*enable_embed_query_optimization=*/false),
+        QueryVisitorTestParams(QueryType::kPlain,
+                               /*get_embedding_match_info=*/false,
+                               /*enable_embedding_iterator_v2=*/true,
+                               /*enable_embed_query_optimization=*/false),
+        QueryVisitorTestParams(QueryType::kPlain,
+                               /*get_embedding_match_info=*/false,
+                               /*enable_embedding_iterator_v2=*/false,
+                               /*enable_embed_query_optimization=*/false),
+        QueryVisitorTestParams(QueryType::kSearch,
+                               /*get_embedding_match_info=*/true,
+                               /*enable_embedding_iterator_v2=*/true,
+                               /*enable_embed_query_optimization=*/true),
+        QueryVisitorTestParams(QueryType::kSearch,
+                               /*get_embedding_match_info=*/true,
+                               /*enable_embedding_iterator_v2=*/false,
+                               /*enable_embed_query_optimization=*/true),
+        QueryVisitorTestParams(QueryType::kSearch,
+                               /*get_embedding_match_info=*/false,
+                               /*enable_embedding_iterator_v2=*/true,
+                               /*enable_embed_query_optimization=*/true),
+        QueryVisitorTestParams(QueryType::kSearch,
+                               /*get_embedding_match_info=*/false,
+                               /*enable_embedding_iterator_v2=*/false,
+                               /*enable_embed_query_optimization=*/true),
+        QueryVisitorTestParams(QueryType::kPlain,
+                               /*get_embedding_match_info=*/true,
+                               /*enable_embedding_iterator_v2=*/true,
+                               /*enable_embed_query_optimization=*/true),
+        QueryVisitorTestParams(QueryType::kPlain,
+                               /*get_embedding_match_info=*/true,
+                               /*enable_embedding_iterator_v2=*/false,
+                               /*enable_embed_query_optimization=*/true),
+        QueryVisitorTestParams(QueryType::kPlain,
+                               /*get_embedding_match_info=*/false,
+                               /*enable_embedding_iterator_v2=*/true,
+                               /*enable_embed_query_optimization=*/true),
+        QueryVisitorTestParams(QueryType::kPlain,
+                               /*get_embedding_match_info=*/false,
+                               /*enable_embedding_iterator_v2=*/false,
+                               /*enable_embed_query_optimization=*/true)));
 
 }  // namespace
 

diff --git a/icing/query/query-processor.cc b/icing/query/query-processor.cc
index 7830dc5..af200b9 100644
--- a/icing/query/query-processor.cc
+++ b/icing/query/query-processor.cc

@@ -29,9 +29,11 @@
 #include "icing/index/iterator/doc-hit-info-iterator-all-document-id.h"
 #include "icing/index/iterator/doc-hit-info-iterator-and.h"
 #include "icing/index/iterator/doc-hit-info-iterator-by-uri.h"
+#include "icing/index/iterator/doc-hit-info-iterator-data-holder.h"
 #include "icing/index/iterator/doc-hit-info-iterator-filter.h"
 #include "icing/index/iterator/doc-hit-info-iterator-section-restrict.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
+#include "icing/index/iterator/document-filter-predicate.h"
 #include "icing/index/numeric/numeric-index.h"
 #include "icing/join/join-children-fetcher.h"
 #include "icing/proto/logging.pb.h"
@@ -103,10 +105,15 @@
     ScoringSpecProto::RankingStrategy::Code ranking_strategy,
     bool get_embedding_match_info, int64_t current_time_ms,
     QueryStatsProto::SearchStats* search_stats) {
-  ICING_ASSIGN_OR_RETURN(QueryResults results,
-                         ParseAdvancedQuery(search_spec, ranking_strategy,
-                                            get_embedding_match_info,
-                                            current_time_ms, search_stats));
+  std::unique_ptr<DocumentFilterPredicate> filter_predicate =
+      GetFilterPredicateBySchemaAndNamespace(search_spec, document_store_,
+                                             schema_store_, current_time_ms);
+
+  ICING_ASSIGN_OR_RETURN(
+      QueryResults results,
+      ParseAdvancedQuery(search_spec, ranking_strategy,
+                         get_embedding_match_info, current_time_ms,
+                         filter_predicate.get(), search_stats));
 
   // Check that all new features used in the search have been enabled in the
   // SearchSpec.
@@ -138,11 +145,12 @@
   }
   results.root_iterator = CreateAndIterator(std::move(iterators));
 
-  DocHitInfoIteratorFilter::Options options =
-      GetFilterOptions(search_spec, document_store_, schema_store_);
-  results.root_iterator = std::make_unique<DocHitInfoIteratorFilter>(
-      std::move(results.root_iterator), &document_store_, &schema_store_,
-      options, current_time_ms);
+  results.root_iterator = DocHitInfoIteratorFilter::ApplyFilter(
+      std::move(results.root_iterator), filter_predicate.get(),
+      feature_flags_.enable_passing_filter_to_children());
+  results.root_iterator =
+      std::make_unique<DocHitInfoIteratorDataHolder<DocumentFilterPredicate>>(
+          std::move(results.root_iterator), std::move(filter_predicate));
   if (!search_spec.type_property_filters().empty()) {
     results.root_iterator =
         DocHitInfoIteratorSectionRestrict::ApplyRestrictions(
@@ -156,6 +164,7 @@
     const SearchSpecProto& search_spec,
     ScoringSpecProto::RankingStrategy::Code ranking_strategy,
     bool get_embedding_match_info, int64_t current_time_ms,
+    const DocumentFilterPredicate* filter_predicate,
     QueryStatsProto::SearchStats* search_stats) const {
   std::unique_ptr<Timer> lexer_timer = clock_.GetNewTimer();
   Lexer lexer(search_spec.query(), Lexer::Language::QUERY);
@@ -182,8 +191,6 @@
       std::unique_ptr<Tokenizer> plain_tokenizer,
       tokenizer_factory::CreateIndexingTokenizer(
           StringIndexingConfig::TokenizerType::PLAIN, &language_segmenter_));
-  DocHitInfoIteratorFilter::Options options =
-      GetFilterOptions(search_spec, document_store_, schema_store_);
   bool needs_term_frequency_info =
       ranking_strategy == ScoringSpecProto::RankingStrategy::RELEVANCE_SCORE;
 
@@ -191,7 +198,7 @@
   QueryVisitor query_visitor(
       &index_, &numeric_index_, &embedding_index_, &document_store_,
       &schema_store_, &normalizer_, plain_tokenizer.get(),
-      join_children_fetcher_, search_spec, std::move(options),
+      join_children_fetcher_, search_spec, filter_predicate,
       needs_term_frequency_info, get_embedding_match_info, &feature_flags_,
       current_time_ms);
   tree_root->Accept(&query_visitor);

diff --git a/icing/query/query-processor.h b/icing/query/query-processor.h
index 03eb10b..69bd75f 100644
--- a/icing/query/query-processor.h
+++ b/icing/query/query-processor.h

@@ -22,6 +22,7 @@
 #include "icing/feature-flags.h"
 #include "icing/index/embed/embedding-index.h"
 #include "icing/index/index.h"
+#include "icing/index/iterator/document-filter-predicate.h"
 #include "icing/index/numeric/numeric-index.h"
 #include "icing/join/join-children-fetcher.h"
 #include "icing/proto/logging.pb.h"
@@ -98,6 +99,7 @@
       const SearchSpecProto& search_spec,
       ScoringSpecProto::RankingStrategy::Code ranking_strategy,
       bool get_embedding_match_info, int64_t current_time_ms,
+      const DocumentFilterPredicate* filter_predicate,
       QueryStatsProto::SearchStats* search_stats) const;
 
   // Parse the query into a one DocHitInfoIterator that represents the root of a

diff --git a/icing/query/query-processor_benchmark.cc b/icing/query/query-processor_benchmark.cc
index ad125c9..1cc3d19 100644
--- a/icing/query/query-processor_benchmark.cc
+++ b/icing/query/query-processor_benchmark.cc

@@ -30,6 +30,7 @@
 #include "icing/index/index.h"
 #include "icing/index/numeric/dummy-numeric-index.h"
 #include "icing/legacy/index/icing-filesystem.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/proto/schema.pb.h"
 #include "icing/proto/search.pb.h"
 #include "icing/proto/term.pb.h"
@@ -49,6 +50,7 @@
 #include "icing/transform/normalizer-options.h"
 #include "icing/transform/normalizer.h"
 #include "icing/util/clock.h"
+#include "icing/util/document-util.h"
 #include "icing/util/icu-data-file-helper.h"
 #include "icing/util/logging.h"
 #include "unicode/uloc.h"
@@ -97,11 +99,13 @@
 
 std::unique_ptr<Index> CreateIndex(const IcingFilesystem& icing_filesystem,
                                    const Filesystem& filesystem,
-                                   const std::string& index_dir) {
+                                   const std::string& index_dir,
+                                   const FeatureFlags& feature_flags) {
   Index::Options options(index_dir, /*index_merge_size=*/1024 * 1024 * 10,
                          /*lite_index_sort_at_indexing=*/true,
                          /*lite_index_sort_size=*/1024 * 8);
-  return Index::Create(options, &filesystem, &icing_filesystem).ValueOrDie();
+  return Index::Create(options, &filesystem, &icing_filesystem, &feature_flags)
+      .ValueOrDie();
 }
 
 std::unique_ptr<Normalizer> CreateNormalizer() {
@@ -119,6 +123,9 @@
       /*force_recovery_and_revalidate_documents=*/false,
       /*pre_mapping_fbv=*/false, /*use_persistent_hash_map=*/true,
       PortableFileBackedProtoLog<DocumentWrapper>::kDefaultCompressionLevel,
+      PortableFileBackedProtoLog<
+          DocumentWrapper>::kDefaultCompressionThresholdBytes,
+      protobuf_ports::kDefaultMemLevel,
       /*initialize_stats=*/nullptr);
 }
 
@@ -147,7 +154,7 @@
   }
 
   std::unique_ptr<Index> index =
-      CreateIndex(icing_filesystem, filesystem, index_dir);
+      CreateIndex(icing_filesystem, filesystem, index_dir, feature_flags);
   // TODO(b/249829533): switch to use persistent numeric index.
   ICING_ASSERT_OK_AND_ASSIGN(
       auto numeric_index,
@@ -178,13 +185,15 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       auto embedding_index,
       EmbeddingIndex::Create(&filesystem, embedding_index_dir, &clock,
-                             &feature_flags));
+                             &feature_flags,
+                             /*num_shards=*/32));
 
   DocumentId document_id = document_store
-                               ->Put(DocumentBuilder()
-                                         .SetKey("icing", "type1")
-                                         .SetSchema("type1")
-                                         .Build())
+                               ->Put(document_util::CreateDocumentWrapper(
+                                   DocumentBuilder()
+                                       .SetKey("icing", "type1")
+                                       .SetSchema("type1")
+                                       .Build()))
                                .ValueOrDie()
                                .new_document_id;
 
@@ -286,7 +295,7 @@
   }
 
   std::unique_ptr<Index> index =
-      CreateIndex(icing_filesystem, filesystem, index_dir);
+      CreateIndex(icing_filesystem, filesystem, index_dir, feature_flags);
   // TODO(b/249829533): switch to use persistent numeric index.
   ICING_ASSERT_OK_AND_ASSIGN(
       auto numeric_index,
@@ -317,13 +326,15 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       auto embedding_index,
       EmbeddingIndex::Create(&filesystem, embedding_index_dir, &clock,
-                             &feature_flags));
+                             &feature_flags,
+                             /*num_shards=*/32));
 
   DocumentId document_id = document_store
-                               ->Put(DocumentBuilder()
-                                         .SetKey("icing", "type1")
-                                         .SetSchema("type1")
-                                         .Build())
+                               ->Put(document_util::CreateDocumentWrapper(
+                                   DocumentBuilder()
+                                       .SetKey("icing", "type1")
+                                       .SetSchema("type1")
+                                       .Build()))
                                .ValueOrDie()
                                .new_document_id;
 
@@ -443,7 +454,7 @@
   }
 
   std::unique_ptr<Index> index =
-      CreateIndex(icing_filesystem, filesystem, index_dir);
+      CreateIndex(icing_filesystem, filesystem, index_dir, feature_flags);
   // TODO(b/249829533): switch to use persistent numeric index.
   ICING_ASSERT_OK_AND_ASSIGN(
       auto numeric_index,
@@ -474,13 +485,15 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       auto embedding_index,
       EmbeddingIndex::Create(&filesystem, embedding_index_dir, &clock,
-                             &feature_flags));
+                             &feature_flags,
+                             /*num_shards=*/32));
 
   DocumentId document_id = document_store
-                               ->Put(DocumentBuilder()
-                                         .SetKey("icing", "type1")
-                                         .SetSchema("type1")
-                                         .Build())
+                               ->Put(document_util::CreateDocumentWrapper(
+                                   DocumentBuilder()
+                                       .SetKey("icing", "type1")
+                                       .SetSchema("type1")
+                                       .Build()))
                                .ValueOrDie()
                                .new_document_id;
 
@@ -585,7 +598,7 @@
   }
 
   std::unique_ptr<Index> index =
-      CreateIndex(icing_filesystem, filesystem, index_dir);
+      CreateIndex(icing_filesystem, filesystem, index_dir, feature_flags);
   // TODO(b/249829533): switch to use persistent numeric index.
   ICING_ASSERT_OK_AND_ASSIGN(
       auto numeric_index,
@@ -616,13 +629,15 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       auto embedding_index,
       EmbeddingIndex::Create(&filesystem, embedding_index_dir, &clock,
-                             &feature_flags));
+                             &feature_flags,
+                             /*num_shards=*/32));
 
   DocumentId document_id = document_store
-                               ->Put(DocumentBuilder()
-                                         .SetKey("icing", "type1")
-                                         .SetSchema("type1")
-                                         .Build())
+                               ->Put(document_util::CreateDocumentWrapper(
+                                   DocumentBuilder()
+                                       .SetKey("icing", "type1")
+                                       .SetSchema("type1")
+                                       .Build()))
                                .ValueOrDie()
                                .new_document_id;
 

diff --git a/icing/query/query-processor_test.cc b/icing/query/query-processor_test.cc
index b5d1f0f..37aa49e 100644
--- a/icing/query/query-processor_test.cc
+++ b/icing/query/query-processor_test.cc

@@ -41,6 +41,7 @@
 #include "icing/index/numeric/numeric-index.h"
 #include "icing/jni/jni-cache.h"
 #include "icing/legacy/index/icing-filesystem.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/portable/platform.h"
 #include "icing/proto/logging.pb.h"
 #include "icing/proto/schema.pb.h"
@@ -54,6 +55,7 @@
 #include "icing/store/document-id.h"
 #include "icing/store/document-store.h"
 #include "icing/testing/common-matchers.h"
+#include "icing/testing/embedding-test-utils.h"
 #include "icing/testing/fake-clock.h"
 #include "icing/testing/jni-test-helpers.h"
 #include "icing/testing/test-data.h"
@@ -65,6 +67,7 @@
 #include "icing/transform/normalizer-options.h"
 #include "icing/transform/normalizer.h"
 #include "icing/util/clock.h"
+#include "icing/util/document-util.h"
 #include "icing/util/icu-data-file-helper.h"
 #include "icing/util/status-macros.h"
 #include "unicode/uloc.h"
@@ -89,6 +92,9 @@
       /*force_recovery_and_revalidate_documents=*/false,
       /*pre_mapping_fbv=*/false, /*use_persistent_hash_map=*/true,
       PortableFileBackedProtoLog<DocumentWrapper>::kDefaultCompressionLevel,
+      PortableFileBackedProtoLog<
+          DocumentWrapper>::kDefaultCompressionThresholdBytes,
+      protobuf_ports::kDefaultMemLevel,
       /*initialize_stats=*/nullptr);
 }
 
@@ -136,7 +142,8 @@
                            /*lite_index_sort_at_indexing=*/true,
                            /*lite_index_sort_size=*/1024 * 8);
     ICING_ASSERT_OK_AND_ASSIGN(
-        index_, Index::Create(options, &filesystem_, &icing_filesystem_));
+        index_, Index::Create(options, &filesystem_, &icing_filesystem_,
+                              feature_flags_.get()));
     // TODO(b/249829533): switch to use persistent numeric index.
     ICING_ASSERT_OK_AND_ASSIGN(
         numeric_index_,
@@ -144,7 +151,8 @@
     ICING_ASSERT_OK_AND_ASSIGN(
         embedding_index_,
         EmbeddingIndex::Create(&filesystem_, embedding_index_dir_, &fake_clock_,
-                               feature_flags_.get()));
+                               feature_flags_.get(),
+                               /*num_shards=*/32));
 
     language_segmenter_factory::SegmenterOptions segmenter_options(
         ULOC_US, jni_cache_.get());
@@ -303,17 +311,21 @@
                   schema, /*ignore_errors_and_delete_documents=*/false),
               IsOk());
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "2")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "2")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // We don't need to insert anything in the index since the empty query will
@@ -347,11 +359,13 @@
   // These documents don't actually match to the tokens in the index. We're
   // inserting the documents to get the appropriate number of documents and
   // namespaces populated.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace1", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace1", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id = put_result.new_document_id;
 
   // Populate the index
@@ -408,11 +422,13 @@
   // These documents don't actually match to the tokens in the index. We're
   // inserting the documents to get the appropriate number of documents and
   // namespaces populated.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace1", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace1", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id = put_result.new_document_id;
 
   // Populate the index
@@ -464,11 +480,13 @@
   // These documents don't actually match to the tokens in the index. We're
   // inserting the documents to get the appropriate number of documents and
   // namespaces populated.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace1", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace1", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id = put_result.new_document_id;
 
   // Populate the index
@@ -522,11 +540,13 @@
   // These documents don't actually match to the tokens in the index. We're
   // inserting the documents to get the appropriate number of documents and
   // namespaces populated.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace1", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace1", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id = put_result.new_document_id;
 
   // Populate the index
@@ -578,11 +598,13 @@
   // These documents don't actually match to the tokens in the index. We're
   // just inserting the documents so that the DocHitInfoIterators will see
   // that the document exists and not filter out the DocumentId as deleted.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace1", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace1", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id = put_result.new_document_id;
 
   // Populate the index
@@ -636,11 +658,13 @@
   // These documents don't actually match to the tokens in the index. We're
   // just inserting the documents so that the DocHitInfoIterators will see
   // that the document exists and not filter out the DocumentId as deleted.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id = put_result.new_document_id;
 
   // Populate the index
@@ -697,11 +721,13 @@
   // These documents don't actually match to the tokens in the index. We're
   // just inserting the documents so that the DocHitInfoIterators will see
   // that the document exists and not filter out the DocumentId as deleted.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id = put_result.new_document_id;
 
   // Populate the index
@@ -755,11 +781,13 @@
   // These documents don't actually match to the tokens in the index. We're
   // just inserting the documents so that the DocHitInfoIterators will see
   // that the document exists and not filter out the DocumentId as deleted.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id = put_result.new_document_id;
 
   // Populate the index
@@ -817,11 +845,13 @@
   // These documents don't actually match to the tokens in the index. We're
   // just inserting the documents so that the DocHitInfoIterators will see
   // that the document exists and not filter out the DocumentId as deleted.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id = put_result.new_document_id;
 
   // Populate the index
@@ -879,17 +909,21 @@
   // These documents don't actually match to the tokens in the index. We're
   // just inserting the documents so that the DocHitInfoIterators will see
   // that the document exists and not filter out the DocumentId as deleted.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "2")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "2")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // Populate the index
@@ -955,17 +989,21 @@
   // These documents don't actually match to the tokens in the index. We're
   // just inserting the documents so that the DocHitInfoIterators will see
   // that the document exists and not filter out the DocumentId as deleted.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "2")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "2")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // Populate the index
@@ -1031,17 +1069,21 @@
   // These documents don't actually match to the tokens in the index. We're
   // just inserting the documents so that the DocHitInfoIterators will see
   // that the document exists and not filter out the DocumentId as deleted.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "2")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "2")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // Populate the index
@@ -1105,17 +1147,21 @@
   // These documents don't actually match to the tokens in the index. We're
   // just inserting the documents so that the DocHitInfoIterators will see
   // that the document exists and not filter out the DocumentId as deleted.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "2")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "2")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id2 = put_result2.new_document_id;
   // Populate the index
   SectionId section_id = 0;
@@ -1280,17 +1326,21 @@
   // These documents don't actually match to the tokens in the index. We're
   // just inserting the documents so that the DocHitInfoIterators will see
   // that the document exists and not filter out the DocumentId as deleted.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "2")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "2")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // Populate the index
@@ -1348,17 +1398,21 @@
   // These documents don't actually match to the tokens in the index. We're
   // just inserting the documents so that the DocHitInfoIterators will see
   // that the document exists and not filter out the DocumentId as deleted.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "2")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "2")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // Populate the index
@@ -1419,17 +1473,21 @@
   // These documents don't actually match to the tokens in the index. We're
   // just inserting the documents so that the DocHitInfoIterators will see
   // that the document exists and not filter out the DocumentId as deleted.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "2")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "2")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // Populate the index
@@ -1487,17 +1545,21 @@
   // These documents don't actually match to the tokens in the index. We're
   // just inserting the documents so that the DocHitInfoIterators will see
   // that the document exists and not filter out the DocumentId as deleted.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "2")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "2")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // Populate the index
@@ -1557,17 +1619,21 @@
   // These documents don't actually match to the tokens in the index. We're
   // just inserting the documents so that they'll bump the
   // last_added_document_id, which will give us the proper exclusion results
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "2")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "2")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // Populate the index
@@ -1613,17 +1679,21 @@
   // These documents don't actually match to the tokens in the index. We're
   // just inserting the documents so that they'll bump the
   // last_added_document_id, which will give us the proper exclusion results
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "2")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "2")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id2 = put_result2.new_document_id;
   // Populate the index
   SectionId section_id = 0;
@@ -1667,17 +1737,21 @@
   // These documents don't actually match to the tokens in the index. We're
   // just inserting the documents so that they'll bump the
   // last_added_document_id, which will give us the proper exclusion results
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "2")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "2")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // Populate the index
@@ -1754,17 +1828,21 @@
   // These documents don't actually match to the tokens in the index. We're
   // just inserting the documents so that they'll bump the
   // last_added_document_id, which will give us the proper exclusion results
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "2")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "2")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // Populate the index
@@ -1844,17 +1922,21 @@
   // These documents don't actually match to the tokens in the index. We're
   // just inserting the documents so that the DocHitInfoIterators will see
   // that the document exists and not filter out the DocumentId as deleted.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "2")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "2")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // Populate the index
@@ -1945,17 +2027,21 @@
   // These documents don't actually match to the tokens in the index. We're
   // inserting the documents to get the appropriate number of documents and
   // namespaces populated.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "2")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "2")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id2 = put_result2.new_document_id;
   EXPECT_THAT(document_store_->Delete("namespace", "1",
                                       fake_clock_.GetSystemTimeMilliseconds()),
@@ -2013,17 +2099,21 @@
   // These documents don't actually match to the tokens in the index. We're
   // inserting the documents to get the appropriate number of documents and
   // namespaces populated.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace1", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace1", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace2", "2")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace2", "2")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // Populate the index
@@ -2081,17 +2171,21 @@
   // These documents don't actually match to the tokens in the index. We're
   // inserting the documents to get the appropriate number of documents and
   // schema types populated.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "2")
-                                                      .SetSchema("message")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "2")
+                                                   .SetSchema("message")
+                                                   .Build())));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // Populate the index
@@ -2150,11 +2244,13 @@
   // These documents don't actually match to the tokens in the index. We're
   // inserting the documents to get the appropriate number of documents and
   // schema types populated.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id = put_result.new_document_id;
 
   // Populate the index
@@ -2223,17 +2319,21 @@
   // These documents don't actually match to the tokens in the index. We're
   // inserting the documents to get the appropriate number of documents and
   // schema types populated.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId email_document_id = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "2")
-                                                      .SetSchema("message")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "2")
+                                                   .SetSchema("message")
+                                                   .Build())));
   DocumentId message_document_id = put_result2.new_document_id;
 
   // Populate the index
@@ -2298,17 +2398,21 @@
   // These documents don't actually match to the tokens in the index. We're
   // inserting the documents to get the appropriate number of documents and
   // schema types populated.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId email_document_id = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "2")
-                                                      .SetSchema("message")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "2")
+                                                   .SetSchema("message")
+                                                   .Build())));
   DocumentId message_document_id = put_result2.new_document_id;
 
   // Populate the index
@@ -2391,11 +2495,13 @@
   // These documents don't actually match to the tokens in the index. We're
   // inserting the documents to get the appropriate number of documents and
   // schema types populated.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId email_document_id = put_result1.new_document_id;
 
   // Populate the index
@@ -2456,17 +2562,21 @@
   // These documents don't actually match to the tokens in the index. We're
   // inserting the documents to get the appropriate number of documents and
   // schema types populated.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId email_document_id = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "2")
-                                                      .SetSchema("message")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "2")
+                                                   .SetSchema("message")
+                                                   .Build())));
   DocumentId message_document_id = put_result2.new_document_id;
 
   // Populate the index
@@ -2521,11 +2631,13 @@
   // These documents don't actually match to the tokens in the index. We're
   // inserting the documents to get the appropriate number of documents and
   // schema types populated.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId email_document_id = put_result1.new_document_id;
 
   // Populate the index
@@ -2579,11 +2691,13 @@
   // These documents don't actually match to the tokens in the index. We're
   // inserting the documents to get the appropriate number of documents and
   // schema types populated.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId email_document_id = put_result1.new_document_id;
 
   // Populate the index
@@ -2640,17 +2754,21 @@
   // These documents don't actually match to the tokens in the index. We're
   // inserting the documents to get the appropriate number of documents and
   // schema types populated.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId email_document_id = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "2")
-                                                      .SetSchema("message")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "2")
+                                                   .SetSchema("message")
+                                                   .Build())));
   DocumentId message_document_id = put_result2.new_document_id;
 
   // Poplate the index
@@ -2750,17 +2868,21 @@
   // These documents don't actually match to the tokens in the index. We're
   // inserting the documents to get the appropriate number of documents and
   // schema types populated.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId email_document_id = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "2")
-                                                      .SetSchema("message")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "2")
+                                                   .SetSchema("message")
+                                                   .Build())));
   DocumentId message_document_id = put_result2.new_document_id;
 
   // Poplate the index
@@ -2881,17 +3003,21 @@
   // These documents don't actually match to the tokens in the index. We're
   // inserting the documents to get the appropriate number of documents and
   // schema types populated.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId email_document_id = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "2")
-                                                      .SetSchema("message")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "2")
+                                                   .SetSchema("message")
+                                                   .Build())));
   DocumentId message_document_id = put_result2.new_document_id;
 
   // Poplate the index
@@ -2957,6 +3083,308 @@
   EXPECT_THAT(results.query_terms["foo"], UnorderedElementsAre("animal"));
 }
 
+TEST_F(QueryProcessorTest, TypePropertyFilterHybridSearchRestrictTerm) {
+  // Create the schema and document store
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("email")
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("foo")
+                                        .SetDataTypeString(TERM_MATCH_EXACT,
+                                                           TOKENIZER_PLAIN)
+                                        .SetCardinality(CARDINALITY_OPTIONAL))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("bar")
+                                        .SetDataTypeString(TERM_MATCH_EXACT,
+                                                           TOKENIZER_PLAIN)
+                                        .SetCardinality(CARDINALITY_OPTIONAL))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("propEmbed0")
+                                        .SetDataTypeVector(
+                                            EMBEDDING_INDEXING_LINEAR_SEARCH)
+                                        .SetCardinality(CARDINALITY_OPTIONAL))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("propEmbed1")
+                                        .SetDataTypeVector(
+                                            EMBEDDING_INDEXING_LINEAR_SEARCH)
+                                        .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+  // SectionIds are assigned in ascending order per schema type,
+  // alphabetically.
+  int email_bar_section_id = 0;
+  int email_foo_section_id = 1;
+  int email_prop_embed0_section_id = 2;
+  int email_prop_embed1_section_id = 3;
+  ASSERT_THAT(schema_store_->SetSchema(
+                  schema, /*ignore_errors_and_delete_documents=*/false),
+              IsOk());
+
+  // These documents don't actually match to the tokens in the index. We're
+  // inserting the documents to get the appropriate number of documents and
+  // schema types populated.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result0,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
+  DocumentId email_document_id0 = put_result0.new_document_id;
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "2")
+                                                   .SetSchema("email")
+                                                   .Build())));
+  DocumentId email_document_id1 = put_result1.new_document_id;
+
+  // Poplate the index
+  TermMatchType::Code term_match_type = TermMatchType::EXACT_ONLY;
+
+  // Email1 has content "animal" in section foo
+  ASSERT_THAT(AddTokenToIndex(email_document_id0, email_foo_section_id,
+                              term_match_type, "animal"),
+              IsOk());
+
+  // Email2 has content "animal" in section bar
+  ASSERT_THAT(AddTokenToIndex(email_document_id1, email_bar_section_id,
+                              term_match_type, "animal"),
+              IsOk());
+
+  // Add embedding vectors into different sections for the two documents.
+  ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
+      BasicHit(email_prop_embed0_section_id, email_document_id0),
+      CreateVector("my_model1", {1, -2, -3}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"email"));
+  ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
+      BasicHit(email_prop_embed1_section_id, email_document_id0),
+      CreateVector("my_model1", {-1, -2, 3}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"email"));
+  ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
+      BasicHit(email_prop_embed0_section_id, email_document_id1),
+      CreateVector("my_model1", {-1, 2, 3}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"email"));
+  ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
+      BasicHit(email_prop_embed1_section_id, email_document_id1),
+      CreateVector("my_model1", {1, 2, -3}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"email"));
+  ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
+
+  // Create an embedding query with semantic scores:
+  // - document 0: -2 (section 0), 6 (section 1)
+  // - document 1: 2 (section 0), -6 (section 1)
+  PropertyProto::VectorProto embedding_query_vector =
+      CreateVector("my_model1", {-1, -1, 1});
+
+  SearchSpecProto search_spec;
+  std::string query = "semanticSearch(getEmbeddingParameter(0), 0) AND animal";
+  search_spec.set_query(query);
+  search_spec.set_term_match_type(term_match_type);
+  search_spec.set_embedding_query_metric_type(
+      SearchSpecProto::EmbeddingQueryMetricType::DOT_PRODUCT);
+  search_spec.add_enabled_features(kListFilterQueryLanguageFeature);
+  *search_spec.add_embedding_query_vectors() =
+      std::move(embedding_query_vector);
+
+  // Add property filters to match both text sections but only the first
+  // embedding section. Therefore, only the second document should be returned.
+  TypePropertyMask* email_mask = search_spec.add_type_property_filters();
+  email_mask->set_schema_type("email");
+  email_mask->add_paths("foo");
+  email_mask->add_paths("bar");
+  email_mask->add_paths("propEmbed0");
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      QueryResults results,
+      query_processor_->ParseSearch(
+          search_spec, ScoringSpecProto::RankingStrategy::RELEVANCE_SCORE,
+          /*get_embedding_match_info=*/false,
+          fake_clock_.GetSystemTimeMilliseconds()));
+
+  // Ordered by descending DocumentId, so message comes first since it was
+  // inserted last
+  DocHitInfo expected_doc_hit_info1(email_document_id1);
+  expected_doc_hit_info1.UpdateSection(email_bar_section_id);
+  expected_doc_hit_info1.UpdateSection(email_prop_embed0_section_id);
+  EXPECT_THAT(GetDocHitInfos(results.root_iterator.get()),
+              ElementsAre(expected_doc_hit_info1));
+  EXPECT_THAT(results.query_term_iterators, SizeIs(1));
+
+  EXPECT_THAT(results.query_terms, SizeIs(1));
+  EXPECT_THAT(results.query_terms[""], UnorderedElementsAre("animal"));
+
+  // The results should be the same even if the order is flipped.
+  query = "animal AND semanticSearch(getEmbeddingParameter(0), 0)";
+  search_spec.set_query(query);
+  ICING_ASSERT_OK_AND_ASSIGN(
+      results,
+      query_processor_->ParseSearch(
+          search_spec, ScoringSpecProto::RankingStrategy::RELEVANCE_SCORE,
+          /*get_embedding_match_info=*/false,
+          fake_clock_.GetSystemTimeMilliseconds()));
+
+  // Ordered by descending DocumentId, so message comes first since it was
+  // inserted last
+  EXPECT_THAT(GetDocHitInfos(results.root_iterator.get()),
+              ElementsAre(expected_doc_hit_info1));
+  EXPECT_THAT(results.query_term_iterators, SizeIs(1));
+
+  EXPECT_THAT(results.query_terms, SizeIs(1));
+  EXPECT_THAT(results.query_terms[""], UnorderedElementsAre("animal"));
+}
+
+TEST_F(QueryProcessorTest, TypePropertyFilterHybridSearchRestrictEmbedding) {
+  // Create the schema and document store
+  // Create the schema and document store
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("email")
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("foo")
+                                        .SetDataTypeString(TERM_MATCH_EXACT,
+                                                           TOKENIZER_PLAIN)
+                                        .SetCardinality(CARDINALITY_OPTIONAL))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("bar")
+                                        .SetDataTypeString(TERM_MATCH_EXACT,
+                                                           TOKENIZER_PLAIN)
+                                        .SetCardinality(CARDINALITY_OPTIONAL))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("propEmbed0")
+                                        .SetDataTypeVector(
+                                            EMBEDDING_INDEXING_LINEAR_SEARCH)
+                                        .SetCardinality(CARDINALITY_OPTIONAL))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("propEmbed1")
+                                        .SetDataTypeVector(
+                                            EMBEDDING_INDEXING_LINEAR_SEARCH)
+                                        .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+  // SectionIds are assigned in ascending order per schema type,
+  // alphabetically.
+  int email_bar_section_id = 0;
+  int email_foo_section_id = 1;
+  int email_prop_embed0_section_id = 2;
+  int email_prop_embed1_section_id = 3;
+  ASSERT_THAT(schema_store_->SetSchema(
+                  schema, /*ignore_errors_and_delete_documents=*/false),
+              IsOk());
+
+  // These documents don't actually match to the tokens in the index. We're
+  // inserting the documents to get the appropriate number of documents and
+  // schema types populated.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result0,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
+  DocumentId email_document_id0 = put_result0.new_document_id;
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "2")
+                                                   .SetSchema("email")
+                                                   .Build())));
+  DocumentId email_document_id1 = put_result1.new_document_id;
+
+  // Poplate the index
+  TermMatchType::Code term_match_type = TermMatchType::EXACT_ONLY;
+
+  // Email1 has content "animal" in section foo
+  ASSERT_THAT(AddTokenToIndex(email_document_id0, email_foo_section_id,
+                              term_match_type, "animal"),
+              IsOk());
+
+  // Email2 has content "animal" in section bar
+  ASSERT_THAT(AddTokenToIndex(email_document_id1, email_bar_section_id,
+                              term_match_type, "animal"),
+              IsOk());
+
+  // Add embedding vectors into different sections for the two documents.
+  ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
+      BasicHit(email_prop_embed0_section_id, email_document_id0),
+      CreateVector("my_model1", {1, -2, -3}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"email"));
+  ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
+      BasicHit(email_prop_embed1_section_id, email_document_id0),
+      CreateVector("my_model1", {-1, -2, 3}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"email"));
+  ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
+      BasicHit(email_prop_embed0_section_id, email_document_id1),
+      CreateVector("my_model1", {-1, 2, 3}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"email"));
+  ICING_ASSERT_OK(embedding_index_->BufferEmbedding(
+      BasicHit(email_prop_embed1_section_id, email_document_id1),
+      CreateVector("my_model1", {1, 2, -3}), QUANTIZATION_TYPE_NONE,
+      /*schema_name=*/"email"));
+  ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex());
+
+  // Create an embedding query with semantic scores:
+  // - document 0: -2 (section 0), 6 (section 1)
+  // - document 1: 2 (section 0), -6 (section 1)
+  PropertyProto::VectorProto embedding_query_vector =
+      CreateVector("my_model1", {-1, -1, 1});
+
+  SearchSpecProto search_spec;
+  std::string query = "semanticSearch(getEmbeddingParameter(0), 0) AND animal";
+  search_spec.set_query(query);
+  search_spec.set_term_match_type(term_match_type);
+  search_spec.set_embedding_query_metric_type(
+      SearchSpecProto::EmbeddingQueryMetricType::DOT_PRODUCT);
+  search_spec.add_enabled_features(kListFilterQueryLanguageFeature);
+  *search_spec.add_embedding_query_vectors() =
+      std::move(embedding_query_vector);
+
+  // Add property filters to match both embeddings sections but only the first
+  // text section. Therefore, only the first document should be returned.
+  TypePropertyMask* email_mask = search_spec.add_type_property_filters();
+  email_mask->set_schema_type("email");
+  email_mask->add_paths("foo");
+  email_mask->add_paths("propEmbed0");
+  email_mask->add_paths("propEmbed1");
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      QueryResults results,
+      query_processor_->ParseSearch(
+          search_spec, ScoringSpecProto::RankingStrategy::RELEVANCE_SCORE,
+          /*get_embedding_match_info=*/false,
+          fake_clock_.GetSystemTimeMilliseconds()));
+
+  // Ordered by descending DocumentId, so message comes first since it was
+  // inserted last
+  DocHitInfo expected_doc_hit_info1(email_document_id0);
+  expected_doc_hit_info1.UpdateSection(email_foo_section_id);
+  expected_doc_hit_info1.UpdateSection(email_prop_embed1_section_id);
+  EXPECT_THAT(GetDocHitInfos(results.root_iterator.get()),
+              ElementsAre(expected_doc_hit_info1));
+  EXPECT_THAT(results.query_term_iterators, SizeIs(1));
+
+  EXPECT_THAT(results.query_terms, SizeIs(1));
+  EXPECT_THAT(results.query_terms[""], UnorderedElementsAre("animal"));
+
+  // The results should be the same even if the order is flipped.
+  query = "animal AND semanticSearch(getEmbeddingParameter(0), 0)";
+  search_spec.set_query(query);
+  ICING_ASSERT_OK_AND_ASSIGN(
+      results,
+      query_processor_->ParseSearch(
+          search_spec, ScoringSpecProto::RankingStrategy::RELEVANCE_SCORE,
+          /*get_embedding_match_info=*/false,
+          fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_THAT(GetDocHitInfos(results.root_iterator.get()),
+              ElementsAre(expected_doc_hit_info1));
+  EXPECT_THAT(results.query_term_iterators, SizeIs(1));
+
+  EXPECT_THAT(results.query_terms, SizeIs(1));
+  EXPECT_THAT(results.query_terms[""], UnorderedElementsAre("animal"));
+}
+
 TEST_F(QueryProcessorTest, DocumentBeforeTtlNotFilteredOut) {
   // Create the schema and document store
   SchemaProto schema = SchemaBuilder()
@@ -2979,12 +3407,13 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result,
-      document_store_->Put(DocumentBuilder()
-                               .SetKey("namespace", "1")
-                               .SetSchema("email")
-                               .SetCreationTimestampMs(10)
-                               .SetTtlMs(100)
-                               .Build()));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .SetCreationTimestampMs(10)
+                                                   .SetTtlMs(100)
+                                                   .Build())));
   DocumentId document_id = put_result.new_document_id;
 
   // Populate the index
@@ -3044,12 +3473,13 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result,
-      document_store_->Put(DocumentBuilder()
-                               .SetKey("namespace", "1")
-                               .SetSchema("email")
-                               .SetCreationTimestampMs(50)
-                               .SetTtlMs(100)
-                               .Build()));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .SetCreationTimestampMs(50)
+                                                   .SetTtlMs(100)
+                                                   .Build())));
   DocumentId document_id = put_result.new_document_id;
 
   // Populate the index
@@ -3107,33 +3537,36 @@
               IsOk());
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result1,
-      document_store_->Put(DocumentBuilder()
-                               .SetKey("namespace", "1")
-                               .SetSchema("transaction")
-                               .AddInt64Property("price", 10)
-                               .Build()));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          DocumentBuilder()
+              .SetKey("namespace", "1")
+              .SetSchema("transaction")
+              .AddInt64Property("price", 10)
+              .Build())));
   DocumentId document_one_id = put_result1.new_document_id;
   ICING_ASSERT_OK(
       AddToNumericIndex(document_one_id, "price", price_section_id, 10));
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result2,
-      document_store_->Put(DocumentBuilder()
-                               .SetKey("namespace", "2")
-                               .SetSchema("transaction")
-                               .AddInt64Property("price", 25)
-                               .Build()));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          DocumentBuilder()
+              .SetKey("namespace", "2")
+              .SetSchema("transaction")
+              .AddInt64Property("price", 25)
+              .Build())));
   DocumentId document_two_id = put_result2.new_document_id;
   ICING_ASSERT_OK(
       AddToNumericIndex(document_two_id, "price", price_section_id, 25));
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result3,
-      document_store_->Put(DocumentBuilder()
-                               .SetKey("namespace", "3")
-                               .SetSchema("transaction")
-                               .AddInt64Property("cost", 2)
-                               .Build()));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "3")
+                                                   .SetSchema("transaction")
+                                                   .AddInt64Property("cost", 2)
+                                                   .Build())));
   DocumentId document_three_id = put_result3.new_document_id;
   ICING_ASSERT_OK(
       AddToNumericIndex(document_three_id, "cost", cost_section_id, 2));
@@ -3212,11 +3645,12 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result1,
-      document_store_->Put(DocumentBuilder()
-                               .SetKey("namespace", "1")
-                               .SetSchema("transaction")
-                               .AddInt64Property("price", 10)
-                               .Build()));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          DocumentBuilder()
+              .SetKey("namespace", "1")
+              .SetSchema("transaction")
+              .AddInt64Property("price", 10)
+              .Build())));
   DocumentId document_one_id = put_result1.new_document_id;
   ICING_ASSERT_OK(
       AddToNumericIndex(document_one_id, "price", price_section_id, 10));
@@ -3268,11 +3702,13 @@
   //   Doc2:
   //     prop1: "foo bar"
   //     prop2: ""
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result0,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "0")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result0,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "0")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id0 = put_result0.new_document_id;
   EXPECT_THAT(
       AddTokenToIndex(document_id0, prop1_section_id, term_match_type, "foo"),
@@ -3281,11 +3717,13 @@
       AddTokenToIndex(document_id0, prop2_section_id, term_match_type, "bar"),
       IsOk());
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id1 = put_result1.new_document_id;
   EXPECT_THAT(
       AddTokenToIndex(document_id1, prop1_section_id, term_match_type, "bar"),
@@ -3294,11 +3732,13 @@
       AddTokenToIndex(document_id1, prop2_section_id, term_match_type, "foo"),
       IsOk());
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "2")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "2")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id2 = put_result2.new_document_id;
   EXPECT_THAT(
       AddTokenToIndex(document_id2, prop1_section_id, term_match_type, "foo"),
@@ -3387,11 +3827,13 @@
   // These documents don't actually match to the tokens in the index. We're
   // inserting the documents to get the appropriate number of documents and
   // namespaces populated.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId document_id = put_result.new_document_id;
 
   // Populate the index
@@ -3436,10 +3878,11 @@
   ASSERT_THAT(schema_store_->SetSchema(
                   schema, /*ignore_errors_and_delete_documents=*/false),
               IsOk());
-  ICING_ASSERT_OK(document_store_->Put(DocumentBuilder()
-                                           .SetKey("namespace", "uri1")
-                                           .SetSchema("email")
-                                           .Build()));
+  ICING_ASSERT_OK(document_store_->Put(
+      document_util::CreateDocumentWrapper(DocumentBuilder()
+                                               .SetKey("namespace", "uri1")
+                                               .SetSchema("email")
+                                               .Build())));
 
   SearchSpecProto search_spec;
   search_spec.set_term_match_type(TermMatchType::PREFIX);

diff --git a/icing/query/query-utils.cc b/icing/query/query-utils.cc
index 0a4cbd3..5cd8d16 100644
--- a/icing/query/query-utils.cc
+++ b/icing/query/query-utils.cc

@@ -14,13 +14,16 @@
 
 #include "icing/query/query-utils.h"
 
+#include <cstdint>
+#include <memory>
 #include <string_view>
 #include <unordered_set>
 
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
-#include "icing/index/iterator/doc-hit-info-iterator-filter.h"
+#include "icing/index/iterator/document-filter-predicate.h"
 #include "icing/schema/schema-store.h"
 #include "icing/store/document-filter-data.h"
+#include "icing/store/document-id.h"
 #include "icing/store/document-store.h"
 #include "icing/store/namespace-id.h"
 
@@ -59,25 +62,85 @@
   return ids;
 }
 
+class DocumentFilterPredicateBySchemaAndNamespace
+    : public DocumentFilterPredicate {
+ public:
+  DocumentFilterPredicateBySchemaAndNamespace(
+      const SearchSpecProto& search_spec, const DocumentStore& document_store,
+      const SchemaStore& schema_store, int64_t current_time_ms)
+      : document_store_(document_store), current_time_ms_(current_time_ms) {
+    // Precompute all the NamespaceIds
+    filter_by_namespace_id_enabled = !search_spec.namespace_filters().empty();
+    target_namespace_ids = ConvertNamespaceToIds(document_store, search_spec);
+
+    // Precompute all the SchemaTypeIds
+    filter_by_schema_type_id_enabled =
+        !search_spec.schema_type_filters().empty();
+    target_schema_type_ids =
+        ConvertExactSchemaTypeToIds(schema_store, search_spec);
+  }
+
+  bool operator()(DocumentId document_id) const override {
+    // Try to get the DocumentFilterData
+    auto document_filter_data = document_store_.GetAliveDocumentFilterData(
+        document_id, current_time_ms_);
+    if (!document_filter_data) {
+      // Didn't find the DocumentFilterData in the filter cache. This could be
+      // because the Document doesn't exist or the DocumentId isn't valid or the
+      // filter cache is in some invalid state. This is bad, but not the query's
+      // responsibility to fix, so just skip this result for now.
+      return false;
+    }
+    // We should be guaranteed that filter data exists now.
+    if (filter_by_namespace_id_enabled &&
+        target_namespace_ids.count(document_filter_data->namespace_id()) == 0) {
+      // Doesn't match one of the specified namespaces.
+      return false;
+    }
+
+    if (filter_by_schema_type_id_enabled &&
+        target_schema_type_ids.count(document_filter_data->schema_type_id()) ==
+            0) {
+      // Doesn't match one of the specified schema types.
+      return false;
+    }
+
+    return true;
+  }
+
+ private:
+  // List of namespace ids that documents must have.
+  // filter_by_namespace_id_enabled=false means that all namespaces are valid,
+  // and no documents will be filtered out.
+  //
+  // Note that if we want to reference the strings in namespaces later, ensure
+  // that the caller who passed the Options class outlives the
+  // DocHitInfoIteratorFilter.
+  std::unordered_set<NamespaceId> target_namespace_ids;
+
+  // List of schema type ids that documents must have.
+  // filter_by_schema_type_id_enabled=false means that all schema types are
+  // valid, and no documents will be filtered out.
+  //
+  // Note that if we want to reference the strings in schema types later,
+  // ensure that the caller who passed the Options class outlives the
+  // DocHitInfoIteratorFilter.
+  std::unordered_set<SchemaTypeId> target_schema_type_ids;
+
+  bool filter_by_schema_type_id_enabled = false;
+  bool filter_by_namespace_id_enabled = false;
+
+  const DocumentStore& document_store_;
+  int64_t current_time_ms_;
+};
+
 }  // namespace
 
-DocHitInfoIteratorFilter::Options GetFilterOptions(
+std::unique_ptr<DocumentFilterPredicate> GetFilterPredicateBySchemaAndNamespace(
     const SearchSpecProto& search_spec, const DocumentStore& document_store,
-    const SchemaStore& schema_store) {
-  DocHitInfoIteratorFilter::Options options;
-
-  // Precompute all the NamespaceIds
-  options.filter_by_namespace_id_enabled =
-      !search_spec.namespace_filters().empty();
-  options.target_namespace_ids =
-      ConvertNamespaceToIds(document_store, search_spec);
-
-  // Precompute all the SchemaTypeIds
-  options.filter_by_schema_type_id_enabled =
-      !search_spec.schema_type_filters().empty();
-  options.target_schema_type_ids =
-      ConvertExactSchemaTypeToIds(schema_store, search_spec);
-  return options;
+    const SchemaStore& schema_store, int64_t current_time_ms) {
+  return std::make_unique<DocumentFilterPredicateBySchemaAndNamespace>(
+      search_spec, document_store, schema_store, current_time_ms);
 }
 
 }  // namespace lib

diff --git a/icing/query/query-utils.h b/icing/query/query-utils.h
index dae75a2..6c18a39 100644
--- a/icing/query/query-utils.h
+++ b/icing/query/query-utils.h

@@ -15,7 +15,10 @@
 #ifndef ICING_QUERY_QUERY_UTILS_H_
 #define ICING_QUERY_QUERY_UTILS_H_
 
-#include "icing/index/iterator/doc-hit-info-iterator-filter.h"
+#include <cstdint>
+#include <memory>
+
+#include "icing/index/iterator/document-filter-predicate.h"
 #include "icing/proto/search.pb.h"
 #include "icing/schema/schema-store.h"
 #include "icing/store/document-store.h"
@@ -23,9 +26,9 @@
 namespace icing {
 namespace lib {
 
-DocHitInfoIteratorFilter::Options GetFilterOptions(
+std::unique_ptr<DocumentFilterPredicate> GetFilterPredicateBySchemaAndNamespace(
     const SearchSpecProto& search_spec, const DocumentStore& document_store,
-    const SchemaStore& schema_store);
+    const SchemaStore& schema_store, int64_t current_time_ms);
 
 }  // namespace lib
 }  // namespace icing

diff --git a/icing/query/suggestion-processor_test.cc b/icing/query/suggestion-processor_test.cc
index 567eeae..2ebaf08 100644
--- a/icing/query/suggestion-processor_test.cc
+++ b/icing/query/suggestion-processor_test.cc

@@ -35,6 +35,7 @@
 #include "icing/index/term-metadata.h"
 #include "icing/jni/jni-cache.h"
 #include "icing/legacy/index/icing-filesystem.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/portable/platform.h"
 #include "icing/query/query-features.h"
 #include "icing/schema-builder.h"
@@ -50,9 +51,10 @@
 #include "icing/testing/tmp-directory.h"
 #include "icing/tokenization/language-segmenter-factory.h"
 #include "icing/tokenization/language-segmenter.h"
-#include "icing/transform/normalizer-options.h"
 #include "icing/transform/normalizer-factory.h"
+#include "icing/transform/normalizer-options.h"
 #include "icing/transform/normalizer.h"
+#include "icing/util/document-util.h"
 #include "icing/util/icu-data-file-helper.h"
 #include "unicode/uloc.h"
 
@@ -110,14 +112,18 @@
 
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
-        DocumentStore::Create(&filesystem_, store_dir_, &fake_clock_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, store_dir_, &fake_clock_, schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     document_store_ = std::move(create_result.document_store);
 
     Index::Options options(index_dir_,
@@ -125,7 +131,8 @@
                            /*lite_index_sort_at_indexing=*/true,
                            /*lite_index_sort_size=*/1024 * 8);
     ICING_ASSERT_OK_AND_ASSIGN(
-        index_, Index::Create(options, &filesystem_, &icing_filesystem_));
+        index_, Index::Create(options, &filesystem_, &icing_filesystem_,
+                              feature_flags_.get()));
     // TODO(b/249829533): switch to use persistent numeric index.
     ICING_ASSERT_OK_AND_ASSIGN(
         numeric_index_,
@@ -133,7 +140,8 @@
     ICING_ASSERT_OK_AND_ASSIGN(
         embedding_index_,
         EmbeddingIndex::Create(&filesystem_, embedding_index_dir_, &fake_clock_,
-                               feature_flags_.get()));
+                               feature_flags_.get(),
+                               /*num_shards=*/32));
 
     language_segmenter_factory::SegmenterOptions segmenter_options(
         ULOC_US, jni_cache_.get());
@@ -208,17 +216,21 @@
   // These documents don't actually match to the tokens in the index. We're
   // inserting the documents to get the appropriate number of documents and
   // namespaces populated.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result0,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace1", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result0,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace1", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId documentId0 = put_result0.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace1", "2")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace1", "2")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId documentId1 = put_result1.new_document_id;
 
   ASSERT_THAT(AddTokenToIndex(documentId0, kSectionId2,
@@ -256,17 +268,21 @@
   // These documents don't actually match to the tokens in the index. We're
   // inserting the documents to get the appropriate number of documents and
   // namespaces populated.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result0,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace1", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result0,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace1", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId documentId0 = put_result0.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace1", "2")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace1", "2")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId documentId1 = put_result1.new_document_id;
 
   ASSERT_THAT(AddTokenToIndex(documentId0, kSectionId2,
@@ -308,17 +324,21 @@
   // These documents don't actually match to the tokens in the index. We're
   // inserting the documents to get the appropriate number of documents and
   // namespaces populated.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result0,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace1", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result0,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace1", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId documentId0 = put_result0.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace1", "2")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace1", "2")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId documentId1 = put_result1.new_document_id;
 
   ASSERT_THAT(AddTokenToIndex(documentId0, kSectionId2,
@@ -362,23 +382,29 @@
   // These documents don't actually match to the tokens in the index. We're
   // inserting the documents to get the appropriate number of documents and
   // namespaces populated.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result0,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace1", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result0,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace1", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId documentId0 = put_result0.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace1", "2")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace1", "2")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId documentId1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace1", "3")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace1", "3")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId documentId2 = put_result2.new_document_id;
 
   ASSERT_THAT(AddTokenToIndex(documentId0, kSectionId2,
@@ -430,17 +456,21 @@
   // These documents don't actually match to the tokens in the index. We're
   // inserting the documents to get the appropriate number of documents and
   // namespaces populated.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result0,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace1", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result0,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace1", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId documentId0 = put_result0.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace1", "2")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace1", "2")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId documentId1 = put_result1.new_document_id;
 
   ASSERT_THAT(AddTokenToIndex(documentId0, kSectionId2,
@@ -493,11 +523,13 @@
   // These documents don't actually match to the tokens in the index. We're
   // inserting the documents to get the appropriate number of documents and
   // namespaces populated.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result0,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace1", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result0,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace1", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId documentId0 = put_result0.new_document_id;
 
   ASSERT_THAT(AddTokenToIndex(documentId0, kSectionId2,
@@ -529,11 +561,13 @@
   // These documents don't actually match to the tokens in the index. We're
   // inserting the documents to get the appropriate number of documents and
   // namespaces populated.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result0,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace1", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result0,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace1", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId documentId0 = put_result0.new_document_id;
 
   ASSERT_THAT(AddTokenToIndex(documentId0, kSectionId2,
@@ -566,11 +600,13 @@
   // inserting the documents to get the appropriate number of documents and
   // namespaces populated.
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result0,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace1", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result0,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace1", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId documentId0 = put_result0.new_document_id;
   ASSERT_THAT(AddTokenToIndex(documentId0, kSectionId2,
                               TermMatchType::EXACT_ONLY, "foo"),
@@ -619,11 +655,13 @@
   // inserting the documents to get the appropriate number of documents and
   // namespaces populated.
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result0,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace1", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result0,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace1", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId documentId0 = put_result0.new_document_id;
   ASSERT_THAT(AddTokenToIndex(documentId0, kSectionId2,
                               TermMatchType::EXACT_ONLY, "foo"),
@@ -666,11 +704,13 @@
   // These documents don't actually match to the tokens in the index. We're
   // inserting the documents to get the appropriate number of documents and
   // namespaces populated.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result0,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace1", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result0,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace1", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId documentId0 = put_result0.new_document_id;
 
   ASSERT_THAT(AddTokenToIndex(documentId0, kSectionId2,
@@ -729,11 +769,13 @@
   // These documents don't actually match to the tokens in the index. We're
   // inserting the documents to get the appropriate number of documents and
   // namespaces populated.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result0,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace1", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result0,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace1", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId documentId0 = put_result0.new_document_id;
 
   ASSERT_THAT(AddTokenToIndex(documentId0, kSectionId2,
@@ -778,11 +820,13 @@
   // These documents don't actually match to the tokens in the index. We're
   // inserting the documents to get the appropriate number of documents and
   // namespaces populated.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result0,
-                             document_store_->Put(DocumentBuilder()
-                                                      .SetKey("namespace1", "1")
-                                                      .SetSchema("email")
-                                                      .Build()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result0,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(DocumentBuilder()
+                                                   .SetKey("namespace1", "1")
+                                                   .SetSchema("email")
+                                                   .Build())));
   DocumentId documentId0 = put_result0.new_document_id;
 
   ASSERT_THAT(AddTokenToIndex(documentId0, kSectionId2,

diff --git a/icing/result/result-adjustment-info.cc b/icing/result/result-adjustment-info.cc
index 6ae911c..82ffb0b 100644
--- a/icing/result/result-adjustment-info.cc
+++ b/icing/result/result-adjustment-info.cc

@@ -29,7 +29,9 @@
 #include "icing/result/snippet-context.h"
 #include "icing/schema/schema-store.h"
 #include "icing/schema/section.h"
+#include "icing/scoring/advanced_scoring/double-list.h"
 #include "icing/store/document-id.h"
+#include "icing/util/embedding-util.h"
 #include "icing/util/logging.h"
 
 namespace icing {
@@ -47,45 +49,57 @@
     const std::unordered_set<DocumentId>& documents_to_include,
     int num_matches_per_property) {
   SnippetContext::DocumentEmbeddingMatchInfoMap result_map;
+  if (embedding_query_results.global_section_infos->empty()) {
+    // No section info indicates that embedding match info is not
+    // enabled.
+    return result_map;
+  }
+  if (embedding_query_results.global_section_infos->size() !=
+      embedding_query_results.global_scores->size()) {
+    // This should never happen.
+    ICING_LOG(ERROR) << "EmbeddingMatchInfos has mismatched section_infos and "
+                        "scores vectors. Global section_infos size: "
+                     << embedding_query_results.global_section_infos->size()
+                     << ", Global scores size: "
+                     << embedding_query_results.global_scores->size();
+    return result_map;
+  }
+
   // Maps from (document_id, section_id) to the match count for that section.
   std::unordered_map<DocumentId, std::unordered_map<SectionId, int>>
       section_match_count;
 
-  for (const auto& [query_vector_index, metric_type_map] :
-       embedding_query_results.result_infos) {
-    for (const auto& [metric_type, info_map] : metric_type_map) {
-      for (const auto& [doc_id, match_infos] : info_map) {
+  for (int query_vector_index = 0;
+       query_vector_index < embedding_query_results.GetNumQueryVectors();
+       ++query_vector_index) {
+    for (SearchSpecProto::EmbeddingQueryMetricType::Code metric_type :
+         embedding_util::kEmbeddingQueryMetricTypes) {
+      const EmbeddingQueryResults::EmbeddingQueryMatchInfoMap* info_map =
+          embedding_query_results.GetMatchInfoMap(query_vector_index,
+                                                  metric_type);
+      if (info_map == nullptr) {
+        continue;
+      }
+      for (const auto& [doc_id, match_infos] : *info_map) {
         if (documents_to_include.find(doc_id) == documents_to_include.end()) {
           continue;
         }
 
-        if (match_infos.section_infos == nullptr) {
-          // No section info indicates that embedding match info is not
-          // enabled.
-          continue;
-        }
-
-        if (match_infos.section_infos->size() != match_infos.scores.size()) {
-          // This should never happen.
-          ICING_LOG(ERROR)
-              << "EmbeddingMatchInfos has mismatched section_infos and "
-                 "scores vectors for document with id "
-              << doc_id
-              << ". Section_infos size: " << match_infos.section_infos->size()
-              << ", scores size: " << match_infos.scores.size();
-          continue;
-        }
-        for (int i = 0; i < match_infos.scores.size(); ++i) {
-          SectionId section_id = match_infos.section_infos->at(i).section_id;
-          if (section_match_count[doc_id][section_id] >=
+        DoubleList scores =
+            embedding_query_results.GetMatchedScoresFromEmbeddingMatchInfos(
+                match_infos);
+        for (int i = 0; i < scores.size(); ++i) {
+          const EmbeddingMatchInfos::EmbeddingMatchSectionInfo& section_info =
+              embedding_query_results.global_section_infos->at(
+                  i + match_infos.score_start_index);
+          if (section_match_count[doc_id][section_info.section_id] >=
               num_matches_per_property) {
             continue;
           }
           result_map[doc_id].push_back(SnippetContext::EmbeddingMatchInfoEntry(
-              match_infos.scores[i], metric_type,
-              match_infos.section_infos->at(i).position, query_vector_index,
-              section_id));
-          ++section_match_count[doc_id][section_id];
+              scores.data()[i], metric_type, section_info.position,
+              query_vector_index, section_info.section_id));
+          ++section_match_count[doc_id][section_info.section_id];
         }
       }
     }

diff --git a/icing/result/result-adjustment-info_test.cc b/icing/result/result-adjustment-info_test.cc
index d539016..0c0c841 100644
--- a/icing/result/result-adjustment-info_test.cc
+++ b/icing/result/result-adjustment-info_test.cc

@@ -147,49 +147,67 @@
       CreateSearchSpec(TermMatchType::EXACT_ONLY, embedding_query_vectors,
                        EMBEDDING_METRIC_DOT_PRODUCT);
 
-  EmbeddingQueryResults embedding_query_results;
+  EmbeddingQueryResults embedding_query_results(/*num_query_vectors=*/2);
   EmbeddingMatchInfos& info_query0_doc0 =
-      embedding_query_results
-          .result_infos[/*query_index=*/0]
-                       [search_spec.embedding_query_metric_type()]
-                       [kDocumentId0];
-  info_query0_doc0.AppendScore(1);
-  info_query0_doc0.AppendScore(1.7);
-  info_query0_doc0.AppendScore(3.3);
-  info_query0_doc0.AppendSectionInfo(/*section_id=*/0, /*position=*/0);
-  info_query0_doc0.AppendSectionInfo(/*section_id=*/0, /*position=*/3);
-  info_query0_doc0.AppendSectionInfo(/*section_id=*/1, /*position=*/1);
+      GetOrCreateEmbeddingMatchInfosForDocument(
+          embedding_query_results, /*query_index=*/0,
+          search_spec.embedding_query_metric_type(), kDocumentId0);
+  info_query0_doc0.AppendScore(*embedding_query_results.global_scores, 1);
+  info_query0_doc0.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/0,
+      /*position=*/0);
+  info_query0_doc0.AppendScore(*embedding_query_results.global_scores, 1.7);
+  info_query0_doc0.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/0,
+      /*position=*/3);
+  info_query0_doc0.AppendScore(*embedding_query_results.global_scores, 3.3);
+  info_query0_doc0.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/1,
+      /*position=*/1);
   EmbeddingMatchInfos& info_query1_doc0 =
-      embedding_query_results
-          .result_infos[/*query_index=*/1]
-                       [search_spec.embedding_query_metric_type()]
-                       [kDocumentId0];
-  info_query1_doc0.AppendScore(2);
-  info_query1_doc0.AppendScore(1.7);
-  info_query1_doc0.AppendSectionInfo(/*section_id=*/0, /*position=*/0);
-  info_query1_doc0.AppendSectionInfo(/*section_id=*/3, /*position=*/2);
+      GetOrCreateEmbeddingMatchInfosForDocument(
+          embedding_query_results, /*query_index=*/1,
+          search_spec.embedding_query_metric_type(), kDocumentId0);
+  info_query1_doc0.AppendScore(*embedding_query_results.global_scores, 2);
+  info_query1_doc0.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/0,
+      /*position=*/0);
+  info_query1_doc0.AppendScore(*embedding_query_results.global_scores, 1.7);
+  info_query1_doc0.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/3,
+      /*position=*/2);
   EmbeddingMatchInfos& info_query1_doc1 =
-      embedding_query_results
-          .result_infos[/*query_index=*/1][EMBEDDING_METRIC_COSINE]
-                       [kDocumentId1];
-  info_query1_doc1.AppendScore(6.66);
-  info_query1_doc1.AppendSectionInfo(/*section_id=*/1, /*position=*/0);
+      GetOrCreateEmbeddingMatchInfosForDocument(
+          embedding_query_results, /*query_index=*/1, EMBEDDING_METRIC_COSINE,
+          kDocumentId1);
+  info_query1_doc1.AppendScore(*embedding_query_results.global_scores, 6.66);
+  info_query1_doc1.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/1,
+      /*position=*/0);
   EmbeddingMatchInfos& info_query0_doc2 =
-      embedding_query_results
-          .result_infos[/*query_index=*/0][EMBEDDING_METRIC_COSINE]
-                       [kDocumentId2];
-  info_query0_doc2.AppendScore(5.25);
-  info_query0_doc2.AppendSectionInfo(/*section_id=*/1, /*position=*/0);
-  info_query0_doc2.AppendScore(1.33);
-  info_query0_doc2.AppendSectionInfo(/*section_id=*/1, /*position=*/4);
+      GetOrCreateEmbeddingMatchInfosForDocument(
+          embedding_query_results, /*query_index=*/0, EMBEDDING_METRIC_COSINE,
+          kDocumentId2);
+  info_query0_doc2.AppendScore(*embedding_query_results.global_scores, 5.25);
+  info_query0_doc2.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/1,
+      /*position=*/0);
+  info_query0_doc2.AppendScore(*embedding_query_results.global_scores, 1.33);
+  info_query0_doc2.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/1,
+      /*position=*/4);
   EmbeddingMatchInfos& info_query1_doc3 =
-      embedding_query_results
-          .result_infos[/*query_index=*/1][EMBEDDING_METRIC_COSINE]
-                       [kDocumentId3];
-  info_query1_doc3.AppendScore(3.25);
-  info_query1_doc3.AppendSectionInfo(/*section_id=*/1, /*position=*/1);
-  info_query1_doc3.AppendScore(2.33);
-  info_query1_doc3.AppendSectionInfo(/*section_id=*/1, /*position=*/2);
+      GetOrCreateEmbeddingMatchInfosForDocument(
+          embedding_query_results, /*query_index=*/1, EMBEDDING_METRIC_COSINE,
+          kDocumentId3);
+  info_query1_doc3.AppendScore(*embedding_query_results.global_scores, 3.25);
+  info_query1_doc3.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/1,
+      /*position=*/1);
+  info_query1_doc3.AppendScore(*embedding_query_results.global_scores, 2.33);
+  info_query1_doc3.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/1,
+      /*position=*/2);
 
   ResultAdjustmentInfo result_adjustment_info(
       search_spec, CreateScoringSpec(/*is_descending_order=*/true), result_spec,
@@ -309,49 +327,67 @@
       CreateSearchSpec(TermMatchType::EXACT_ONLY, embedding_query_vectors,
                        EMBEDDING_METRIC_DOT_PRODUCT);
 
-  EmbeddingQueryResults embedding_query_results;
+  EmbeddingQueryResults embedding_query_results(/*num_query_vectors=*/2);
   EmbeddingMatchInfos& info_query0_doc0 =
-      embedding_query_results
-          .result_infos[/*query_index=*/0]
-                       [search_spec.embedding_query_metric_type()]
-                       [kDocumentId0];
-  info_query0_doc0.AppendScore(1);
-  info_query0_doc0.AppendScore(1.7);
-  info_query0_doc0.AppendScore(3.3);
-  info_query0_doc0.AppendSectionInfo(/*section_id=*/0, /*position=*/0);
-  info_query0_doc0.AppendSectionInfo(/*section_id=*/0, /*position=*/3);
-  info_query0_doc0.AppendSectionInfo(/*section_id=*/1, /*position=*/1);
+      GetOrCreateEmbeddingMatchInfosForDocument(
+          embedding_query_results, /*query_index=*/0,
+          search_spec.embedding_query_metric_type(), kDocumentId0);
+  info_query0_doc0.AppendScore(*embedding_query_results.global_scores, 1);
+  info_query0_doc0.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/0,
+      /*position=*/0);
+  info_query0_doc0.AppendScore(*embedding_query_results.global_scores, 1.7);
+  info_query0_doc0.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/0,
+      /*position=*/3);
+  info_query0_doc0.AppendScore(*embedding_query_results.global_scores, 3.3);
+  info_query0_doc0.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/1,
+      /*position=*/1);
   EmbeddingMatchInfos& info_query1_doc0 =
-      embedding_query_results
-          .result_infos[/*query_index=*/1]
-                       [search_spec.embedding_query_metric_type()]
-                       [kDocumentId0];
-  info_query1_doc0.AppendScore(2);
-  info_query1_doc0.AppendScore(1.7);
-  info_query1_doc0.AppendSectionInfo(/*section_id=*/0, /*position=*/0);
-  info_query1_doc0.AppendSectionInfo(/*section_id=*/3, /*position=*/2);
+      GetOrCreateEmbeddingMatchInfosForDocument(
+          embedding_query_results, /*query_index=*/1,
+          search_spec.embedding_query_metric_type(), kDocumentId0);
+  info_query1_doc0.AppendScore(*embedding_query_results.global_scores, 2);
+  info_query1_doc0.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/0,
+      /*position=*/0);
+  info_query1_doc0.AppendScore(*embedding_query_results.global_scores, 1.7);
+  info_query1_doc0.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/3,
+      /*position=*/2);
   EmbeddingMatchInfos& info_query1_doc1 =
-      embedding_query_results
-          .result_infos[/*query_index=*/1][EMBEDDING_METRIC_COSINE]
-                       [kDocumentId1];
-  info_query1_doc1.AppendScore(6.66);
-  info_query1_doc1.AppendSectionInfo(/*section_id=*/1, /*position=*/0);
+      GetOrCreateEmbeddingMatchInfosForDocument(
+          embedding_query_results, /*query_index=*/1, EMBEDDING_METRIC_COSINE,
+          kDocumentId1);
+  info_query1_doc1.AppendScore(*embedding_query_results.global_scores, 6.66);
+  info_query1_doc1.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/1,
+      /*position=*/0);
   EmbeddingMatchInfos& info_query0_doc2 =
-      embedding_query_results
-          .result_infos[/*query_index=*/0][EMBEDDING_METRIC_COSINE]
-                       [kDocumentId2];
-  info_query0_doc2.AppendScore(5.25);
-  info_query0_doc2.AppendSectionInfo(/*section_id=*/1, /*position=*/0);
-  info_query0_doc2.AppendScore(1.33);
-  info_query0_doc2.AppendSectionInfo(/*section_id=*/1, /*position=*/4);
+      GetOrCreateEmbeddingMatchInfosForDocument(
+          embedding_query_results, /*query_index=*/0, EMBEDDING_METRIC_COSINE,
+          kDocumentId2);
+  info_query0_doc2.AppendScore(*embedding_query_results.global_scores, 5.25);
+  info_query0_doc2.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/1,
+      /*position=*/0);
+  info_query0_doc2.AppendScore(*embedding_query_results.global_scores, 1.33);
+  info_query0_doc2.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/1,
+      /*position=*/4);
   EmbeddingMatchInfos& info_query1_doc3 =
-      embedding_query_results
-          .result_infos[/*query_index=*/1][EMBEDDING_METRIC_COSINE]
-                       [kDocumentId3];
-  info_query1_doc3.AppendScore(3.25);
-  info_query1_doc3.AppendSectionInfo(/*section_id=*/1, /*position=*/1);
-  info_query1_doc3.AppendScore(2.33);
-  info_query1_doc3.AppendSectionInfo(/*section_id=*/1, /*position=*/2);
+      GetOrCreateEmbeddingMatchInfosForDocument(
+          embedding_query_results, /*query_index=*/1, EMBEDDING_METRIC_COSINE,
+          kDocumentId3);
+  info_query1_doc3.AppendScore(*embedding_query_results.global_scores, 3.25);
+  info_query1_doc3.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/1,
+      /*position=*/1);
+  info_query1_doc3.AppendScore(*embedding_query_results.global_scores, 2.33);
+  info_query1_doc3.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/1,
+      /*position=*/2);
 
   ResultAdjustmentInfo result_adjustment_info(
       search_spec, CreateScoringSpec(/*is_descending_order=*/true), result_spec,
@@ -383,8 +419,8 @@
                                     /*section_id=*/0)),
                             EqualsEmbeddingMatchInfoEntry(
                                 SnippetContext::EmbeddingMatchInfoEntry(
-                                    /*score=*/2, EMBEDDING_METRIC_DOT_PRODUCT,
-                                    /*position=*/0, /*query_vector_index=*/1,
+                                    /*score=*/1.7, EMBEDDING_METRIC_DOT_PRODUCT,
+                                    /*position=*/3, /*query_vector_index=*/0,
                                     /*section_id=*/0)),
                             EqualsEmbeddingMatchInfoEntry(
                                 SnippetContext::EmbeddingMatchInfoEntry(
@@ -422,33 +458,43 @@
   SearchSpecProto search_spec =
       CreateSearchSpec(TermMatchType::EXACT_ONLY, embedding_query_vectors,
                        EMBEDDING_METRIC_DOT_PRODUCT);
-  EmbeddingQueryResults embedding_query_results;
+  EmbeddingQueryResults embedding_query_results(/*num_query_vectors=*/2);
   EmbeddingMatchInfos& info_query0_doc0 =
-      embedding_query_results
-          .result_infos[/*query_index=*/0]
-                       [search_spec.embedding_query_metric_type()]
-                       [kDocumentId0];
-  info_query0_doc0.AppendScore(1);
-  info_query0_doc0.AppendScore(1.7);
-  info_query0_doc0.AppendScore(3.3);
-  info_query0_doc0.AppendSectionInfo(/*section_id=*/0, /*position=*/0);
-  info_query0_doc0.AppendSectionInfo(/*section_id=*/0, /*position=*/3);
-  info_query0_doc0.AppendSectionInfo(/*section_id=*/1, /*position=*/1);
+      GetOrCreateEmbeddingMatchInfosForDocument(
+          embedding_query_results, /*query_index=*/0,
+          search_spec.embedding_query_metric_type(), kDocumentId0);
+  info_query0_doc0.AppendScore(*embedding_query_results.global_scores, 1);
+  info_query0_doc0.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/0,
+      /*position=*/0);
+  info_query0_doc0.AppendScore(*embedding_query_results.global_scores, 1.7);
+  info_query0_doc0.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/0,
+      /*position=*/3);
+  info_query0_doc0.AppendScore(*embedding_query_results.global_scores, 3.3);
+  info_query0_doc0.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/1,
+      /*position=*/1);
   EmbeddingMatchInfos& info_query1_doc0 =
-      embedding_query_results
-          .result_infos[/*query_index=*/1]
-                       [search_spec.embedding_query_metric_type()]
-                       [kDocumentId0];
-  info_query1_doc0.AppendScore(2);
-  info_query1_doc0.AppendScore(1.7);
-  info_query1_doc0.AppendSectionInfo(/*section_id=*/0, /*position=*/0);
-  info_query1_doc0.AppendSectionInfo(/*section_id=*/3, /*position=*/2);
+      GetOrCreateEmbeddingMatchInfosForDocument(
+          embedding_query_results, /*query_index=*/1,
+          search_spec.embedding_query_metric_type(), kDocumentId0);
+  info_query1_doc0.AppendScore(*embedding_query_results.global_scores, 2);
+  info_query1_doc0.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/0,
+      /*position=*/0);
+  info_query1_doc0.AppendScore(*embedding_query_results.global_scores, 1.7);
+  info_query1_doc0.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/3,
+      /*position=*/2);
   EmbeddingMatchInfos& info_query1_doc1 =
-      embedding_query_results
-          .result_infos[/*query_index=*/1][EMBEDDING_METRIC_COSINE]
-                       [kDocumentId1];
-  info_query1_doc1.AppendScore(6.66);
-  info_query1_doc1.AppendSectionInfo(/*section_id=*/1, /*position=*/0);
+      GetOrCreateEmbeddingMatchInfosForDocument(
+          embedding_query_results, /*query_index=*/1, EMBEDDING_METRIC_COSINE,
+          kDocumentId1);
+  info_query1_doc1.AppendScore(*embedding_query_results.global_scores, 6.66);
+  info_query1_doc1.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/1,
+      /*position=*/0);
 
   ResultAdjustmentInfo result_adjustment_info(
       search_spec, CreateScoringSpec(/*is_descending_order=*/true), result_spec,
@@ -526,33 +572,43 @@
   SearchSpecProto search_spec =
       CreateSearchSpec(TermMatchType::EXACT_ONLY, embedding_query_vectors,
                        SearchSpecProto::EmbeddingQueryMetricType::DOT_PRODUCT);
-  EmbeddingQueryResults embedding_query_results;
+  EmbeddingQueryResults embedding_query_results(/*num_query_vectors=*/2);
   EmbeddingMatchInfos& info_query0_doc0 =
-      embedding_query_results
-          .result_infos[/*query_index=*/0]
-                       [search_spec.embedding_query_metric_type()]
-                       [kDocumentId0];
-  info_query0_doc0.AppendScore(1);
-  info_query0_doc0.AppendScore(1.7);
-  info_query0_doc0.AppendScore(3.3);
-  info_query0_doc0.AppendSectionInfo(/*section_id=*/0, /*position=*/0);
-  info_query0_doc0.AppendSectionInfo(/*section_id=*/0, /*position=*/3);
-  info_query0_doc0.AppendSectionInfo(/*section_id=*/1, /*position=*/1);
+      GetOrCreateEmbeddingMatchInfosForDocument(
+          embedding_query_results, /*query_index=*/0,
+          search_spec.embedding_query_metric_type(), kDocumentId0);
+  info_query0_doc0.AppendScore(*embedding_query_results.global_scores, 1);
+  info_query0_doc0.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/0,
+      /*position=*/0);
+  info_query0_doc0.AppendScore(*embedding_query_results.global_scores, 1.7);
+  info_query0_doc0.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/0,
+      /*position=*/3);
+  info_query0_doc0.AppendScore(*embedding_query_results.global_scores, 3.3);
+  info_query0_doc0.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/1,
+      /*position=*/1);
   EmbeddingMatchInfos& info_query1_doc0 =
-      embedding_query_results
-          .result_infos[/*query_index=*/1]
-                       [search_spec.embedding_query_metric_type()]
-                       [kDocumentId0];
-  info_query1_doc0.AppendScore(2);
-  info_query1_doc0.AppendScore(1.7);
-  info_query1_doc0.AppendSectionInfo(/*section_id=*/0, /*position=*/0);
-  info_query1_doc0.AppendSectionInfo(/*section_id=*/3, /*position=*/2);
+      GetOrCreateEmbeddingMatchInfosForDocument(
+          embedding_query_results, /*query_index=*/1,
+          search_spec.embedding_query_metric_type(), kDocumentId0);
+  info_query1_doc0.AppendScore(*embedding_query_results.global_scores, 2);
+  info_query1_doc0.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/0,
+      /*position=*/0);
+  info_query1_doc0.AppendScore(*embedding_query_results.global_scores, 1.7);
+  info_query1_doc0.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/3,
+      /*position=*/2);
   EmbeddingMatchInfos& info_query1_doc1 =
-      embedding_query_results
-          .result_infos[/*query_index=*/1][EMBEDDING_METRIC_COSINE]
-                       [kDocumentId1];
-  info_query1_doc1.AppendScore(6.66);
-  info_query1_doc1.AppendSectionInfo(/*section_id=*/1, /*position=*/0);
+      GetOrCreateEmbeddingMatchInfosForDocument(
+          embedding_query_results, /*query_index=*/1, EMBEDDING_METRIC_COSINE,
+          kDocumentId1);
+  info_query1_doc1.AppendScore(*embedding_query_results.global_scores, 6.66);
+  info_query1_doc1.AppendSectionInfo(
+      *embedding_query_results.global_section_infos, /*section_id=*/1,
+      /*position=*/0);
 
   ResultAdjustmentInfo result_adjustment_info(
       search_spec, CreateScoringSpec(/*is_descending_order=*/true), result_spec,

diff --git a/icing/result/result-retriever-v2.cc b/icing/result/result-retriever-v2.cc
index 4d221f3..3d5104b 100644
--- a/icing/result/result-retriever-v2.cc
+++ b/icing/result/result-retriever-v2.cc

@@ -14,9 +14,11 @@
 
 #include "icing/result/result-retriever-v2.h"
 
+#include <algorithm>
 #include <cstddef>
 #include <cstdint>
 #include <memory>
+#include <optional>
 #include <string>
 #include <unordered_map>
 #include <utility>
@@ -24,6 +26,7 @@
 
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "icing/absl_ports/mutex.h"
+#include "icing/feature-flags.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/search.pb.h"
 #include "icing/result/page-result.h"
@@ -97,52 +100,55 @@
 
 }  // namespace
 
-bool GroupResultLimiterV2::ShouldBeRemoved(
+std::optional<int> GroupResultLimiterV2::GetGroupResultLimitsIndex(
     const ScoredDocumentHit& scored_document_hit,
     const std::unordered_map<int32_t, int>& entry_id_group_id_map,
-    const DocumentStore& document_store, std::vector<int>& group_result_limits,
+    const DocumentStore& document_store,
     ResultSpecProto::ResultGroupingType result_group_type,
     int64_t current_time_ms) const {
+  if (result_group_type == ResultSpecProto::ResultGroupingType::
+                               ResultSpecProto_ResultGroupingType_NONE) {
+    // No limit. Return -1 to skip the group result limit.
+    return -1;
+  }
+
   auto document_filter_data_optional =
       document_store.GetAliveDocumentFilterData(
           scored_document_hit.document_id(), current_time_ms);
   if (!document_filter_data_optional) {
-    // The document doesn't exist.
-    return true;
+    // The document doesn't exist. Return std::nullopt to exclude the document.
+    return std::nullopt;
   }
-  NamespaceId namespace_id =
-      document_filter_data_optional.value().namespace_id();
-  SchemaTypeId schema_type_id =
-      document_filter_data_optional.value().schema_type_id();
-  auto entry_id_or = document_store.GetResultGroupingEntryId(
+  NamespaceId namespace_id = document_filter_data_optional->namespace_id();
+  SchemaTypeId schema_type_id = document_filter_data_optional->schema_type_id();
+
+  std::optional<int32_t> entry_id = document_store.GetResultGroupingEntryId(
       result_group_type, namespace_id, schema_type_id);
-  if (!entry_id_or.ok()) {
-    return false;
+  if (!entry_id.has_value()) {
+    // No limit. Return -1 to skip the group result limit.
+    return -1;
   }
-  int32_t entry_id = entry_id_or.ValueOrDie();
-  auto iter = entry_id_group_id_map.find(entry_id);
+
+  auto iter = entry_id_group_id_map.find(*entry_id);
   if (iter == entry_id_group_id_map.end()) {
     // If a ResultGrouping Entry Id isn't found in entry_id_group_id_map, then
     // there are no limits placed on results from this entry id.
-    return false;
+    return -1;
   }
-  int& count = group_result_limits.at(iter->second);
-  if (count <= 0) {
-    return true;
-  }
-  --count;
-  return false;
+  return iter->second;
 }
 
 libtextclassifier3::StatusOr<std::unique_ptr<ResultRetrieverV2>>
 ResultRetrieverV2::Create(
     const DocumentStore* doc_store, const SchemaStore* schema_store,
     const LanguageSegmenter* language_segmenter, const Normalizer* normalizer,
+    const FeatureFlags* feature_flags,
     std::unique_ptr<const GroupResultLimiterV2> group_result_limiter) {
   ICING_RETURN_ERROR_IF_NULL(doc_store);
   ICING_RETURN_ERROR_IF_NULL(schema_store);
   ICING_RETURN_ERROR_IF_NULL(language_segmenter);
   ICING_RETURN_ERROR_IF_NULL(normalizer);
+  ICING_RETURN_ERROR_IF_NULL(feature_flags);
   ICING_RETURN_ERROR_IF_NULL(group_result_limiter);
 
   ICING_ASSIGN_OR_RETURN(
@@ -151,11 +157,12 @@
 
   return std::unique_ptr<ResultRetrieverV2>(
       new ResultRetrieverV2(doc_store, std::move(snippet_retriever),
-                            std::move(group_result_limiter)));
+                            std::move(group_result_limiter), feature_flags));
 }
 
 std::pair<PageResult, bool> ResultRetrieverV2::RetrieveNextPage(
-    ResultStateV2& result_state, int64_t current_time_ms) const {
+    ResultStateV2& result_state, int32_t max_results,
+    int64_t current_time_ms) const {
   absl_ports::unique_lock l(&result_state.mutex);
 
   // For calculating page
@@ -166,103 +173,49 @@
   // Retrieve info
   std::vector<SearchResultProto::ResultProto> results;
   int32_t num_total_bytes = 0;
-  while (results.size() < result_state.num_per_page() &&
+  int desired_page_size = std::min(result_state.num_per_page(), max_results);
+  while (num_total_bytes < result_state.num_total_bytes_per_page_threshold() &&
+         results.size() < desired_page_size &&
          !result_state.scored_document_hits_ranker->empty()) {
-    JoinedScoredDocumentHit next_best_document_hit =
-        result_state.scored_document_hits_ranker->PopNext();
-    if (group_result_limiter_->ShouldBeRemoved(
-            next_best_document_hit.parent_scored_document_hit(),
-            result_state.entry_id_group_id_map(), doc_store_,
-            result_state.group_result_limits, result_state.result_group_type(),
-            current_time_ms)) {
-      continue;
-    }
+    RetrieveResult result = Retrieve(result_state, current_time_ms);
+    if (result.result_proto.has_value()) {
+      if (result.has_parent_snippets) {
+        ++num_results_with_snippets;
+      }
 
-    DocumentId doc_id =
-        next_best_document_hit.parent_scored_document_hit().document_id();
-    libtextclassifier3::StatusOr<DocumentProto> document_or =
-        doc_store_.Get(doc_id);
-    if (!document_or.ok()) {
-      // Skip the document if getting errors.
-      ICING_LOG(WARNING) << "Fail to fetch document from document store: "
-                         << document_or.status().error_message();
-      continue;
-    }
-
-    DocumentProto document = std::move(document_or).ValueOrDie();
-    // Apply parent projection
-    ApplyProjection(result_state.parent_adjustment_info(), &document);
-
-    SearchResultProto::ResultProto result;
-    // Add parent snippet if requested.
-    if (ApplySnippet(result_state.parent_adjustment_info(), *snippet_retriever_,
-                     document, doc_id,
-                     next_best_document_hit.parent_scored_document_hit()
-                         .hit_section_id_mask(),
-                     &result)) {
-      ++num_results_with_snippets;
-    }
-
-    // Add the document, itself.
-    *result.mutable_document() = std::move(document);
-    result.set_score(next_best_document_hit.final_score());
-    const auto* parent_additional_scores =
-        next_best_document_hit.parent_scored_document_hit().additional_scores();
-    if (parent_additional_scores != nullptr) {
-      result.mutable_additional_scores()->Add(parent_additional_scores->begin(),
-                                              parent_additional_scores->end());
-    }
-
-    // Retrieve child documents
-    for (const ScoredDocumentHit& child_scored_document_hit :
-         next_best_document_hit.child_scored_document_hits()) {
-      if (result.joined_results_size() >=
-          result_state.max_joined_children_per_parent_to_return()) {
+      // Apply byte size threshold enforcement only if it is not the first
+      // document. This ensures that at least one document is returned,
+      // otherwise it will get stuck forever since nothing is popped from the
+      // ranker.
+      //
+      // (Use subtraction to avoid integer overflow).
+      size_t result_bytes = result.result_proto->ByteSizeLong();
+      if (feature_flags_.enable_strict_page_byte_size_limit() &&
+          !results.empty() &&
+          result_bytes >= result_state.num_total_bytes_per_page_threshold() -
+                              num_total_bytes) {
+        // Exceeds the byte size threshold, so skip the current document. Also
+        // it remains in the ranker and will be included in the next page.
+        ICING_LOG(INFO) << "Skipping document due to byte size threshold. "
+                           "Current num docs: "
+                        << results.size()
+                        << ", total byte size: " << num_total_bytes
+                        << ", next doc byte size: " << result_bytes
+                        << ", threshold: "
+                        << result_state.num_total_bytes_per_page_threshold();
         break;
       }
 
-      DocumentId child_doc_id = child_scored_document_hit.document_id();
-      libtextclassifier3::StatusOr<DocumentProto> child_document_or =
-          doc_store_.Get(child_doc_id);
-      if (!child_document_or.ok()) {
-        // Skip the document if getting errors.
-        ICING_LOG(WARNING)
-            << "Fail to fetch child document from document store: "
-            << child_document_or.status().error_message();
-        continue;
-      }
+      results.push_back(std::move(*result.result_proto));
+      num_total_bytes += result_bytes;
 
-      DocumentProto child_document = std::move(child_document_or).ValueOrDie();
-      ApplyProjection(result_state.child_adjustment_info(), &child_document);
-
-      SearchResultProto::ResultProto* child_result =
-          result.add_joined_results();
-      // Add child snippet if requested.
-      ApplySnippet(result_state.child_adjustment_info(), *snippet_retriever_,
-                   child_document, child_doc_id,
-                   child_scored_document_hit.hit_section_id_mask(),
-                   child_result);
-
-      *child_result->mutable_document() = std::move(child_document);
-      child_result->set_score(child_scored_document_hit.score());
-      if (child_scored_document_hit.additional_scores() != nullptr) {
-        child_result->mutable_additional_scores()->Add(
-            child_scored_document_hit.additional_scores()->begin(),
-            child_scored_document_hit.additional_scores()->end());
+      // Decrement the counter of the group result limit in ResultState.
+      if (result.group_result_limits_index != -1) {
+        --result_state.group_result_limits[result.group_result_limits_index];
       }
     }
 
-    size_t result_bytes = result.ByteSizeLong();
-    results.push_back(std::move(result));
-
-    // Check if num_total_bytes + result_bytes reaches or exceeds
-    // num_total_bytes_per_page_threshold. Use subtraction to avoid integer
-    // overflow.
-    if (result_bytes >=
-        result_state.num_total_bytes_per_page_threshold() - num_total_bytes) {
-      break;
-    }
-    num_total_bytes += result_bytes;
+    result_state.scored_document_hits_ranker->Pop();
   }
 
   // Update numbers in ResultState
@@ -279,5 +232,113 @@
       has_more_results);
 }
 
+ResultRetrieverV2::RetrieveResult ResultRetrieverV2::Retrieve(
+    ResultStateV2& result_state, int64_t current_time_ms) const {
+  const JoinedScoredDocumentHit& next_best_document_hit =
+      result_state.scored_document_hits_ranker->Top();
+
+  // Get the index of the group result limits for the given scored document hit,
+  // and check if the document should be excluded.
+  std::optional<int> idx = group_result_limiter_->GetGroupResultLimitsIndex(
+      next_best_document_hit.parent_scored_document_hit(),
+      result_state.entry_id_group_id_map(), doc_store_,
+      result_state.result_group_type(), current_time_ms);
+  if (!idx.has_value()) {
+    // Should exclude the document, so return an invalid result.
+    return {.result_proto = std::nullopt,
+            .group_result_limits_index = -1,
+            .has_parent_snippets = false};
+  }
+
+  if (*idx == -1) {
+    // No limit for the result document. Pass.
+  } else if (*idx < 0 || *idx >= result_state.group_result_limits.size()) {
+    // This should not happen. But let's check it anyway and log. Also set the
+    // index to -1 indicating that there is no limit.
+    ICING_LOG(WARNING) << "Get an invalid group result limits index: " << *idx;
+    *idx = -1;
+  } else if (result_state.group_result_limits[*idx] <= 0) {
+    // Should exclude the document since the group has no budget left, so
+    // return an invalid result.
+    return {.result_proto = std::nullopt,
+            .group_result_limits_index = -1,
+            .has_parent_snippets = false};
+  }
+
+  DocumentId doc_id =
+      next_best_document_hit.parent_scored_document_hit().document_id();
+  auto document_or = doc_store_.Get(doc_id);
+  if (!document_or.ok()) {
+    // Exclude the document if getting errors.
+    ICING_LOG(WARNING) << "Fail to fetch document from document store: "
+                       << document_or.status().error_message();
+    return {.result_proto = std::nullopt,
+            .group_result_limits_index = -1,
+            .has_parent_snippets = false};
+  }
+  DocumentProto document = std::move(document_or).ValueOrDie();
+
+  // Apply parent projection
+  ApplyProjection(result_state.parent_adjustment_info(), &document);
+
+  SearchResultProto::ResultProto result;
+  // Add parent snippet if requested.
+  bool has_parent_snippets = ApplySnippet(
+      result_state.parent_adjustment_info(), *snippet_retriever_, document,
+      doc_id,
+      next_best_document_hit.parent_scored_document_hit().hit_section_id_mask(),
+      &result);
+
+  // Add the document, itself.
+  *result.mutable_document() = std::move(document);
+  result.set_score(next_best_document_hit.final_score());
+  const auto* parent_additional_scores =
+      next_best_document_hit.parent_scored_document_hit().additional_scores();
+  if (parent_additional_scores != nullptr) {
+    result.mutable_additional_scores()->Add(parent_additional_scores->begin(),
+                                            parent_additional_scores->end());
+  }
+
+  // Retrieve child documents
+  for (const ScoredDocumentHit& child_scored_document_hit :
+       next_best_document_hit.child_scored_document_hits()) {
+    if (result.joined_results_size() >=
+        result_state.max_joined_children_per_parent_to_return()) {
+      break;
+    }
+
+    DocumentId child_doc_id = child_scored_document_hit.document_id();
+    libtextclassifier3::StatusOr<DocumentProto> child_document_or =
+        doc_store_.Get(child_doc_id);
+    if (!child_document_or.ok()) {
+      // Skip the document if getting errors.
+      ICING_LOG(WARNING) << "Fail to fetch child document from document store: "
+                         << child_document_or.status().error_message();
+      continue;
+    }
+
+    DocumentProto child_document = std::move(child_document_or).ValueOrDie();
+    ApplyProjection(result_state.child_adjustment_info(), &child_document);
+
+    SearchResultProto::ResultProto* child_result = result.add_joined_results();
+    // Add child snippet if requested.
+    ApplySnippet(result_state.child_adjustment_info(), *snippet_retriever_,
+                 child_document, child_doc_id,
+                 child_scored_document_hit.hit_section_id_mask(), child_result);
+
+    *child_result->mutable_document() = std::move(child_document);
+    child_result->set_score(child_scored_document_hit.score());
+    if (child_scored_document_hit.additional_scores() != nullptr) {
+      child_result->mutable_additional_scores()->Add(
+          child_scored_document_hit.additional_scores()->begin(),
+          child_scored_document_hit.additional_scores()->end());
+    }
+  }
+
+  return RetrieveResult{.result_proto = std::make_optional(std::move(result)),
+                        .group_result_limits_index = *idx,
+                        .has_parent_snippets = has_parent_snippets};
+}
+
 }  // namespace lib
 }  // namespace icing

diff --git a/icing/result/result-retriever-v2.h b/icing/result/result-retriever-v2.h
index 7b1a364..0516feb 100644
--- a/icing/result/result-retriever-v2.h
+++ b/icing/result/result-retriever-v2.h

@@ -17,11 +17,13 @@
 
 #include <cstdint>
 #include <memory>
+#include <optional>
 #include <unordered_map>
 #include <utility>
-#include <vector>
 
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
+#include "icing/absl_ports/thread_annotations.h"
+#include "icing/feature-flags.h"
 #include "icing/proto/search.pb.h"
 #include "icing/result/page-result.h"
 #include "icing/result/result-state-v2.h"
@@ -41,12 +43,18 @@
 
   virtual ~GroupResultLimiterV2() = default;
 
-  // Returns true if the scored_document_hit should be removed.
-  virtual bool ShouldBeRemoved(
+  // Gets the index of the group result limits for the given scored document
+  // hit.
+  //
+  // Returns:
+  //   - A valid index of the group result limits.
+  //   - -1 indicating that there is no limit for the result document.
+  //   - std::nullopt if the result document, its namespace, or its schema type
+  //     is not present. The caller should exclude the document from the page.
+  virtual std::optional<int> GetGroupResultLimitsIndex(
       const ScoredDocumentHit& scored_document_hit,
       const std::unordered_map<int32_t, int>& entry_id_group_id_map,
       const DocumentStore& document_store,
-      std::vector<int>& group_result_limits,
       ResultSpecProto::ResultGroupingType result_group_type,
       int64_t current_time_ms) const;
 };
@@ -63,7 +71,7 @@
   static libtextclassifier3::StatusOr<std::unique_ptr<ResultRetrieverV2>>
   Create(const DocumentStore* doc_store, const SchemaStore* schema_store,
          const LanguageSegmenter* language_segmenter,
-         const Normalizer* normalizer,
+         const Normalizer* normalizer, const FeatureFlags* feature_flags,
          std::unique_ptr<const GroupResultLimiterV2> group_result_limiter =
              std::make_unique<const GroupResultLimiterV2>());
 
@@ -72,8 +80,8 @@
   // out the next top rank documents from ResultState, retrieves the documents
   // from storage, updates ResultState, and finally wraps the result + other
   // information into PageResult. The expected number of documents to return is
-  // min(num_per_page, the number of all scored document hits) inside
-  // ResultState.
+  // min(max_results, num_per_page, the number of all scored document hits)
+  // inside ResultState.
   //
   // The number of snippets to return is based on the total number of snippets
   // needed and number of snippets that have already been returned previously
@@ -89,20 +97,54 @@
   // Returns:
   //   std::pair<PageResult, bool>
   std::pair<PageResult, bool> RetrieveNextPage(ResultStateV2& result_state,
-                                               int64_t current_time_ms) const;
+                                               int32_t max_results,
+                                               int64_t current_time_ms) const
+      ICING_LOCKS_EXCLUDED(result_state.mutex);
 
  private:
   explicit ResultRetrieverV2(
       const DocumentStore* doc_store,
       std::unique_ptr<SnippetRetriever> snippet_retriever,
-      std::unique_ptr<const GroupResultLimiterV2> group_result_limiter)
+      std::unique_ptr<const GroupResultLimiterV2> group_result_limiter,
+      const FeatureFlags* feature_flags)
       : doc_store_(*doc_store),
         snippet_retriever_(std::move(snippet_retriever)),
-        group_result_limiter_(std::move(group_result_limiter)) {}
+        group_result_limiter_(std::move(group_result_limiter)),
+        feature_flags_(*feature_flags) {}
+
+  // Helper function to construct a ResultProto by the next best document hit
+  // from the scored document hits ranker.
+  //
+  // REQUIRES: !result_state.scored_document_hits_ranker.empty()
+  struct RetrieveResult {
+    // The constructed result proto. If std::nullopt, then the document should
+    // be skipped.
+    std::optional<SearchResultProto::ResultProto> result_proto;
+
+    // The index of the group result limits for the result. The caller should
+    // decrement the corresponding result limit in
+    // result_state.group_result_limits after deciding to include the result
+    // document in the page.
+    // - It is guaranteed to be -1 or in the range of [0,
+    //   result_state.group_result_limits.size() - 1]. If it is -1, then it
+    //   means there is no limit for the result document.
+    // - Only used when the proto is not std::nullopt.
+    int group_result_limits_index;
+
+    // Whether the (parent) document of the result has snippets. Only used when
+    // the proto is not std::nullopt.
+    bool has_parent_snippets;
+  };
+
+  RetrieveResult Retrieve(ResultStateV2& result_state,
+                          int64_t current_time_ms) const
+      ICING_EXCLUSIVE_LOCKS_REQUIRED(result_state.mutex);
 
   const DocumentStore& doc_store_;
   std::unique_ptr<SnippetRetriever> snippet_retriever_;
   const std::unique_ptr<const GroupResultLimiterV2> group_result_limiter_;
+
+  const FeatureFlags& feature_flags_;
 };
 
 }  // namespace lib

diff --git a/icing/result/result-retriever-v2_group-result-limiter_test.cc b/icing/result/result-retriever-v2_group-result-limiter_test.cc
index 964e92e..50160c1 100644
--- a/icing/result/result-retriever-v2_group-result-limiter_test.cc
+++ b/icing/result/result-retriever-v2_group-result-limiter_test.cc

@@ -15,13 +15,18 @@
 #include <cstdint>
 #include <limits>
 #include <memory>
+#include <string>
+#include <utility>
 #include <vector>
 
 #include "gtest/gtest.h"
+#include "icing/absl_ports/mutex.h"
 #include "icing/document-builder.h"
 #include "icing/feature-flags.h"
+#include "icing/file/filesystem.h"
 #include "icing/file/portable-file-backed-proto-log.h"
 #include "icing/portable/equals-proto.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/portable/platform.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/schema.pb.h"
@@ -41,9 +46,11 @@
 #include "icing/testing/test-feature-flags.h"
 #include "icing/testing/tmp-directory.h"
 #include "icing/tokenization/language-segmenter-factory.h"
+#include "icing/tokenization/language-segmenter.h"
 #include "icing/transform/normalizer-factory.h"
 #include "icing/transform/normalizer-options.h"
 #include "icing/transform/normalizer.h"
+#include "icing/util/document-util.h"
 #include "icing/util/icu-data-file-helper.h"
 #include "unicode/uloc.h"
 
@@ -99,14 +106,18 @@
 
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
-        DocumentStore::Create(&filesystem_, test_dir_, &fake_clock_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, test_dir_, &fake_clock_, schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     document_store_ = std::move(create_result.document_store);
   }
 
@@ -142,8 +153,9 @@
                                 .SetScore(1)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   DocumentProto document2 = DocumentBuilder()
@@ -152,8 +164,9 @@
                                 .SetScore(2)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document2)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   std::vector<ScoredDocumentHit> scored_document_hits = {
@@ -181,12 +194,14 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // Only the top ranked document in "namespace" (document2), should be
   // returned.
   auto [page_result, has_more_results] = result_retriever->RetrieveNextPage(
-      result_state, fake_clock_.GetSystemTimeMilliseconds());
+      result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+      fake_clock_.GetSystemTimeMilliseconds());
   ASSERT_THAT(page_result.results, SizeIs(1));
   EXPECT_THAT(page_result.results.at(0).document(), EqualsProto(document2));
   // Document1 has not been returned due to GroupResultLimiter, but since it was
@@ -204,8 +219,9 @@
                                 .SetScore(1)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   DocumentProto document2 = DocumentBuilder()
@@ -214,8 +230,9 @@
                                 .SetScore(2)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document2)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   std::vector<ScoredDocumentHit> scored_document_hits = {
@@ -243,11 +260,13 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // First page: empty page
   auto [page_result, has_more_results] = result_retriever->RetrieveNextPage(
-      result_state, fake_clock_.GetSystemTimeMilliseconds());
+      result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+      fake_clock_.GetSystemTimeMilliseconds());
   ASSERT_THAT(page_result.results, IsEmpty());
   EXPECT_FALSE(has_more_results);
 }
@@ -262,8 +281,9 @@
                                 .SetScore(1)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   DocumentProto document2 = DocumentBuilder()
@@ -272,8 +292,9 @@
                                 .SetScore(2)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document2)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   DocumentProto document3 = DocumentBuilder()
@@ -282,8 +303,9 @@
                                 .SetScore(3)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             document_store_->Put(document3));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      document_store_->Put(document_util::CreateDocumentWrapper(document3)));
   DocumentId document_id3 = put_result3.new_document_id;
 
   DocumentProto document4 = DocumentBuilder()
@@ -292,8 +314,9 @@
                                 .SetScore(4)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result4,
-                             document_store_->Put(document4));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result4,
+      document_store_->Put(document_util::CreateDocumentWrapper(document4)));
   DocumentId document_id4 = put_result4.new_document_id;
 
   std::vector<ScoredDocumentHit> scored_document_hits = {
@@ -323,11 +346,13 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // First page: document4 and document3 should be returned.
   auto [page_result1, has_more_results1] = result_retriever->RetrieveNextPage(
-      result_state, fake_clock_.GetSystemTimeMilliseconds());
+      result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+      fake_clock_.GetSystemTimeMilliseconds());
   ASSERT_THAT(page_result1.results, SizeIs(2));
   EXPECT_THAT(page_result1.results.at(0).document(), EqualsProto(document4));
   EXPECT_THAT(page_result1.results.at(1).document(), EqualsProto(document3));
@@ -337,7 +362,8 @@
   // them will be filtered out by group result limiter, so we should get an
   // empty page.
   auto [page_result2, has_more_results2] = result_retriever->RetrieveNextPage(
-      result_state, fake_clock_.GetSystemTimeMilliseconds());
+      result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+      fake_clock_.GetSystemTimeMilliseconds());
   EXPECT_THAT(page_result2.results, SizeIs(0));
   EXPECT_FALSE(has_more_results2);
 }
@@ -352,8 +378,9 @@
                                 .SetScore(1)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   DocumentProto document2 = DocumentBuilder()
@@ -362,8 +389,9 @@
                                 .SetScore(2)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document2)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   DocumentProto document3 = DocumentBuilder()
@@ -372,8 +400,9 @@
                                 .SetScore(3)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             document_store_->Put(document3));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      document_store_->Put(document_util::CreateDocumentWrapper(document3)));
   DocumentId document_id3 = put_result3.new_document_id;
 
   DocumentProto document4 = DocumentBuilder()
@@ -382,8 +411,9 @@
                                 .SetScore(4)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result4,
-                             document_store_->Put(document4));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result4,
+      document_store_->Put(document_util::CreateDocumentWrapper(document4)));
   DocumentId document_id4 = put_result4.new_document_id;
 
   std::vector<ScoredDocumentHit> scored_document_hits = {
@@ -414,13 +444,15 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // All documents in "namespace2" should be returned.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(3));
   EXPECT_THAT(page_result.results.at(0).document(), EqualsProto(document4));
@@ -438,8 +470,9 @@
                                 .SetScore(1)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   DocumentProto document2 = DocumentBuilder()
@@ -448,8 +481,9 @@
                                 .SetScore(2)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document2)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   std::vector<ScoredDocumentHit> scored_document_hits = {
@@ -480,15 +514,17 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // Only the top ranked document in "namespace" (document2), should be
   // returned. The presence of "nonexistentNamespace" in the same result
   // grouping should have no effect.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(1));
   EXPECT_THAT(page_result.results.at(0).document(), EqualsProto(document2));
@@ -504,8 +540,9 @@
                                 .SetScore(1)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   DocumentProto document2 = DocumentBuilder()
@@ -514,8 +551,9 @@
                                 .SetScore(2)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document2)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   std::vector<ScoredDocumentHit> scored_document_hits = {
@@ -546,15 +584,17 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // Only the top ranked document in "Document" (document2), should be
   // returned. The presence of "nonexistentNamespace" in the same result
   // grouping should have no effect.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(1));
   EXPECT_THAT(page_result.results.at(0).document(), EqualsProto(document2));
@@ -571,8 +611,9 @@
                                 .SetScore(1)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   DocumentProto document2 = DocumentBuilder()
@@ -581,8 +622,9 @@
                                 .SetScore(2)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document2)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   DocumentProto document3 = DocumentBuilder()
@@ -591,8 +633,9 @@
                                 .SetScore(3)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             document_store_->Put(document3));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      document_store_->Put(document_util::CreateDocumentWrapper(document3)));
   DocumentId document_id3 = put_result3.new_document_id;
 
   DocumentProto document4 = DocumentBuilder()
@@ -601,8 +644,9 @@
                                 .SetScore(4)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result4,
-                             document_store_->Put(document4));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result4,
+      document_store_->Put(document_util::CreateDocumentWrapper(document4)));
   DocumentId document_id4 = put_result4.new_document_id;
 
   DocumentProto document5 = DocumentBuilder()
@@ -611,8 +655,9 @@
                                 .SetScore(5)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result5,
-                             document_store_->Put(document5));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result5,
+      document_store_->Put(document_util::CreateDocumentWrapper(document5)));
   DocumentId document_id5 = put_result5.new_document_id;
 
   DocumentProto document6 = DocumentBuilder()
@@ -621,8 +666,9 @@
                                 .SetScore(6)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result6,
-                             document_store_->Put(document6));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result6,
+      document_store_->Put(document_util::CreateDocumentWrapper(document6)));
   DocumentId document_id6 = put_result6.new_document_id;
 
   std::vector<ScoredDocumentHit> scored_document_hits = {
@@ -661,15 +707,17 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // Only the top-ranked result in "namespace1" (document2) should be returned.
   // Only the top-ranked results across "namespace2" and "namespace3"
   // (document6, document5) should be returned.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(3));
   EXPECT_THAT(page_result.results.at(0).document(), EqualsProto(document6));
@@ -688,8 +736,9 @@
                                 .SetScore(1)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   DocumentProto document2 = DocumentBuilder()
@@ -698,8 +747,9 @@
                                 .SetScore(2)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document2)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   DocumentProto document3 = DocumentBuilder()
@@ -708,8 +758,9 @@
                                 .SetScore(3)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             document_store_->Put(document3));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      document_store_->Put(document_util::CreateDocumentWrapper(document3)));
   DocumentId document_id3 = put_result3.new_document_id;
 
   DocumentProto document4 = DocumentBuilder()
@@ -718,8 +769,9 @@
                                 .SetScore(4)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result4,
-                             document_store_->Put(document4));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result4,
+      document_store_->Put(document_util::CreateDocumentWrapper(document4)));
   DocumentId document_id4 = put_result4.new_document_id;
 
   DocumentProto document5 = DocumentBuilder()
@@ -728,8 +780,9 @@
                                 .SetScore(5)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result5,
-                             document_store_->Put(document5));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result5,
+      document_store_->Put(document_util::CreateDocumentWrapper(document5)));
   DocumentId document_id5 = put_result5.new_document_id;
 
   DocumentProto document6 = DocumentBuilder()
@@ -738,8 +791,9 @@
                                 .SetScore(6)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result6,
-                             document_store_->Put(document6));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result6,
+      document_store_->Put(document_util::CreateDocumentWrapper(document6)));
   DocumentId document_id6 = put_result6.new_document_id;
 
   std::vector<ScoredDocumentHit> scored_document_hits = {
@@ -778,15 +832,17 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // Only the top-ranked result in "Document" (document6) should be returned.
   // Only the top-ranked results across "Message" and "Person"
   // (document5, document3) should be returned.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(3));
   EXPECT_THAT(page_result.results.at(0).document(), EqualsProto(document6));
@@ -805,8 +861,9 @@
                                 .SetScore(1)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   DocumentProto document2 = DocumentBuilder()
@@ -815,8 +872,9 @@
                                 .SetScore(2)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document2)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   DocumentProto document3 = DocumentBuilder()
@@ -825,8 +883,9 @@
                                 .SetScore(3)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             document_store_->Put(document3));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      document_store_->Put(document_util::CreateDocumentWrapper(document3)));
   DocumentId document_id3 = put_result3.new_document_id;
 
   DocumentProto document4 = DocumentBuilder()
@@ -835,8 +894,9 @@
                                 .SetScore(4)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result4,
-                             document_store_->Put(document4));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result4,
+      document_store_->Put(document_util::CreateDocumentWrapper(document4)));
   DocumentId document_id4 = put_result4.new_document_id;
 
   DocumentProto document5 = DocumentBuilder()
@@ -845,8 +905,9 @@
                                 .SetScore(5)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result5,
-                             document_store_->Put(document5));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result5,
+      document_store_->Put(document_util::CreateDocumentWrapper(document5)));
   DocumentId document_id5 = put_result5.new_document_id;
 
   DocumentProto document6 = DocumentBuilder()
@@ -855,8 +916,9 @@
                                 .SetScore(6)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result6,
-                             document_store_->Put(document6));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result6,
+      document_store_->Put(document_util::CreateDocumentWrapper(document6)));
   DocumentId document_id6 = put_result6.new_document_id;
 
   std::vector<ScoredDocumentHit> scored_document_hits = {
@@ -898,7 +960,8 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // Only the top-ranked result in "namespace1xDocument" (document3)
   // should be returned.
@@ -907,8 +970,9 @@
 
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(3));
   EXPECT_THAT(page_result.results.at(0).document(), EqualsProto(document6));
@@ -926,8 +990,9 @@
                                 .SetScore(1)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   DocumentProto document2 = DocumentBuilder()
@@ -936,8 +1001,9 @@
                                 .SetScore(2)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document2)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   std::vector<ScoredDocumentHit> scored_document_hits = {
@@ -966,14 +1032,16 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // All documents in "namespace" should be returned. The presence of
   // "nonexistentNamespace" should have no effect.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(2));
   EXPECT_THAT(page_result.results.at(0).document(), EqualsProto(document2));
@@ -990,8 +1058,9 @@
                                 .SetScore(1)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   DocumentProto document2 = DocumentBuilder()
@@ -1000,8 +1069,9 @@
                                 .SetScore(2)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document2)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   std::vector<ScoredDocumentHit> scored_document_hits = {
@@ -1030,14 +1100,16 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // All documents in "Document" should be returned. The presence of
   // "nonexistentDocument" should have no effect.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(2));
   EXPECT_THAT(page_result.results.at(0).document(), EqualsProto(document2));
@@ -1054,8 +1126,9 @@
                                 .SetScore(1)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document1)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   DocumentProto document2 = DocumentBuilder()
@@ -1064,8 +1137,9 @@
                                 .SetScore(2)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document2)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   DocumentProto document3 = DocumentBuilder()
@@ -1074,8 +1148,9 @@
                                 .SetScore(3)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             document_store_->Put(document3));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      document_store_->Put(document_util::CreateDocumentWrapper(document3)));
   DocumentId document_id3 = put_result3.new_document_id;
 
   DocumentProto document4 = DocumentBuilder()
@@ -1084,8 +1159,9 @@
                                 .SetScore(4)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result4,
-                             document_store_->Put(document4));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result4,
+      document_store_->Put(document_util::CreateDocumentWrapper(document4)));
   DocumentId document_id4 = put_result4.new_document_id;
 
   DocumentProto document5 = DocumentBuilder()
@@ -1094,8 +1170,9 @@
                                 .SetScore(5)
                                 .SetCreationTimestampMs(1000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result5,
-                             document_store_->Put(document5));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result5,
+      document_store_->Put(document_util::CreateDocumentWrapper(document5)));
   DocumentId document_id5 = put_result5.new_document_id;
 
   std::vector<ScoredDocumentHit> scored_document_hits = {
@@ -1122,13 +1199,13 @@
   entry = result_grouping->add_entry_groupings();
   entry->set_namespace_("namespace2");
 
-  // Get corpus ids.
-  ICING_ASSERT_OK_AND_ASSIGN(
-      CorpusId corpus_id1, document_store_->GetResultGroupingEntryId(
-                               result_grouping_type, "namespace1", "Document"));
-  ICING_ASSERT_OK_AND_ASSIGN(
-      CorpusId corpus_id2, document_store_->GetResultGroupingEntryId(
-                               result_grouping_type, "namespace2", "Document"));
+  // Get result grouping entry ids.
+  ICING_ASSERT_HAS_VALUE_AND_ASSIGN(
+      int32_t entry_id1, document_store_->GetResultGroupingEntryId(
+                             result_grouping_type, "namespace1", "Document"));
+  ICING_ASSERT_HAS_VALUE_AND_ASSIGN(
+      int32_t entry_id2, document_store_->GetResultGroupingEntryId(
+                             result_grouping_type, "namespace2", "Document"));
 
   // Creates a ResultState with 5 ScoredDocumentHits.
   ResultStateV2 result_state(
@@ -1141,22 +1218,24 @@
     absl_ports::shared_lock l(&result_state.mutex);
 
     ASSERT_THAT(result_state.entry_id_group_id_map(),
-                UnorderedElementsAre(Pair(corpus_id1, 0), Pair(corpus_id2, 1)));
+                UnorderedElementsAre(Pair(entry_id1, 0), Pair(entry_id2, 1)));
     ASSERT_THAT(result_state.group_result_limits, ElementsAre(3, 1));
   }
 
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // document5, document4, document1 belong to namespace2 (with max_results =
   // 1).
-  // docuemnt3, document2 belong to namespace 1 (with max_results = 3).
+  // document3, document2 belong to namespace 1 (with max_results = 3).
   // Since num_per_page is 2, we expect to get document5 and document3 in the
   // first page.
   auto [page_result1, has_more_results1] = result_retriever->RetrieveNextPage(
-      result_state, fake_clock_.GetSystemTimeMilliseconds());
+      result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+      fake_clock_.GetSystemTimeMilliseconds());
   ASSERT_THAT(page_result1.results, SizeIs(2));
   ASSERT_THAT(page_result1.results.at(0).document(), EqualsProto(document5));
   ASSERT_THAT(page_result1.results.at(1).document(), EqualsProto(document3));
@@ -1179,9 +1258,9 @@
     // round, num_returned should still be 2, since document4 was "filtered out"
     // and should not be counted into num_returned.
     EXPECT_THAT(result_state.num_returned, Eq(2));
-    // corpus_id_group_id_map should be unchanged.
+    // entry_id_group_id_map should be unchanged.
     EXPECT_THAT(result_state.entry_id_group_id_map(),
-                UnorderedElementsAre(Pair(corpus_id1, 0), Pair(corpus_id2, 1)));
+                UnorderedElementsAre(Pair(entry_id1, 0), Pair(entry_id2, 1)));
     // GroupResultLimiter should decrement the # in group_result_limits.
     EXPECT_THAT(result_state.group_result_limits, ElementsAre(2, 0));
   }
@@ -1189,7 +1268,8 @@
   // Although there are document2 and document1 left, since namespace2 has
   // reached its max results, document1 should be excluded from the second page.
   auto [page_result2, has_more_results2] = result_retriever->RetrieveNextPage(
-      result_state, fake_clock_.GetSystemTimeMilliseconds());
+      result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+      fake_clock_.GetSystemTimeMilliseconds());
   ASSERT_THAT(page_result2.results, SizeIs(1));
   ASSERT_THAT(page_result2.results.at(0).document(), EqualsProto(document2));
   ASSERT_FALSE(has_more_results2);
@@ -1203,9 +1283,9 @@
     // since document1 was "filtered out" and should not be counted into
     // num_returned.
     EXPECT_THAT(result_state.num_returned, Eq(3));
-    // corpus_id_group_id_map should be unchanged.
+    // entry_id_group_id_map should be unchanged.
     EXPECT_THAT(result_state.entry_id_group_id_map(),
-                UnorderedElementsAre(Pair(corpus_id1, 0), Pair(corpus_id2, 1)));
+                UnorderedElementsAre(Pair(entry_id1, 0), Pair(entry_id2, 1)));
     // GroupResultLimiter should decrement the # in group_result_limits.
     EXPECT_THAT(result_state.group_result_limits, ElementsAre(1, 0));
   }

diff --git a/icing/result/result-retriever-v2_projection_test.cc b/icing/result/result-retriever-v2_projection_test.cc
index dfe7bb7..40bf19d 100644
--- a/icing/result/result-retriever-v2_projection_test.cc
+++ b/icing/result/result-retriever-v2_projection_test.cc

@@ -27,6 +27,7 @@
 #include "icing/file/portable-file-backed-proto-log.h"
 #include "icing/index/embed/embedding-query-results.h"
 #include "icing/portable/equals-proto.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/portable/platform.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/schema.pb.h"
@@ -55,6 +56,7 @@
 #include "icing/transform/normalizer-factory.h"
 #include "icing/transform/normalizer-options.h"
 #include "icing/transform/normalizer.h"
+#include "icing/util/document-util.h"
 #include "icing/util/icu-data-file-helper.h"
 #include "unicode/uloc.h"
 
@@ -200,14 +202,18 @@
 
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
-        DocumentStore::Create(&filesystem_, test_dir_, &fake_clock_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, test_dir_, &fake_clock_, schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     document_store_ = std::move(create_result.document_store);
   }
 
@@ -282,8 +288,9 @@
           .AddStringProperty(
               "body", "Oh what a beautiful morning! Oh what a beautiful day!")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document_one));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_one)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   DocumentProto document_two =
@@ -295,8 +302,9 @@
           .AddStringProperty("body",
                              "Count all the sheep and tell them 'Hello'.")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document_two));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_two)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // 2. Setup the scored results.
@@ -330,13 +338,15 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // 5. Verify that the returned results only contain the 'name' property.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(2));
 
@@ -380,8 +390,9 @@
           .AddStringProperty(
               "body", "Oh what a beautiful morning! Oh what a beautiful day!")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document_one));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_one)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   DocumentProto document_two =
@@ -400,8 +411,9 @@
           .AddStringProperty("body",
                              "Count all the sheep and tell them 'Hello'.")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document_two));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_two)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // 2. Setup the scored results.
@@ -434,14 +446,16 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // 5. Verify that the returned results only contain the 'sender.name'
   // property.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(2));
 
@@ -495,8 +509,9 @@
           .AddStringProperty(
               "body", "Oh what a beautiful morning! Oh what a beautiful day!")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document_one));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_one)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   DocumentProto document_two =
@@ -515,8 +530,9 @@
           .AddStringProperty("body",
                              "Count all the sheep and tell them 'Hello'.")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document_two));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_two)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // 2. Setup the scored results.
@@ -549,14 +565,16 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // 5. Verify that the returned results only contain the 'sender'
   // property and all of the subproperties of 'sender'.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(2));
 
@@ -613,8 +631,9 @@
           .AddStringProperty(
               "body", "Oh what a beautiful morning! Oh what a beautiful day!")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document_one));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_one)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   DocumentProto document_two =
@@ -633,8 +652,9 @@
           .AddStringProperty("body",
                              "Count all the sheep and tell them 'Hello'.")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document_two));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_two)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // 2. Setup the scored results.
@@ -668,14 +688,16 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // 5. Verify that the returned results only contain the 'sender.name' and
   // 'sender.address' properties.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(2));
 
@@ -724,8 +746,9 @@
           .AddStringProperty(
               "body", "Oh what a beautiful morning! Oh what a beautiful day!")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document_one));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_one)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   DocumentProto document_two =
@@ -737,8 +760,9 @@
           .AddStringProperty("body",
                              "Count all the sheep and tell them 'Hello'.")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document_two));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_two)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // 2. Setup the scored results.
@@ -770,13 +794,15 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // 5. Verify that the returned results contain *no* properties.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(2));
 
@@ -808,8 +834,9 @@
           .AddStringProperty(
               "body", "Oh what a beautiful morning! Oh what a beautiful day!")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document_one));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_one)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   DocumentProto document_two =
@@ -821,8 +848,9 @@
           .AddStringProperty("body",
                              "Count all the sheep and tell them 'Hello'.")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document_two));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_two)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // 2. Setup the scored results.
@@ -855,13 +883,15 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // 5. Verify that the returned results contain *no* properties.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(2));
 
@@ -893,8 +923,9 @@
           .AddStringProperty(
               "body", "Oh what a beautiful morning! Oh what a beautiful day!")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document_one));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_one)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   DocumentProto document_two =
@@ -906,8 +937,9 @@
           .AddStringProperty("body",
                              "Count all the sheep and tell them 'Hello'.")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document_two));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_two)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // 2. Setup the scored results.
@@ -941,13 +973,15 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // 5. Verify that the returned results only contain the 'name' property.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(2));
 
@@ -983,8 +1017,9 @@
           .AddStringProperty(
               "body", "Oh what a beautiful morning! Oh what a beautiful day!")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document_one));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_one)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   DocumentProto document_two =
@@ -995,8 +1030,9 @@
           .AddStringProperty("name", "Joe Fox")
           .AddStringProperty("emailAddress", "ny152@aol.com")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document_two));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_two)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // 2. Setup the scored results.
@@ -1029,14 +1065,16 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // 5. Verify that the returned Email results only contain the 'name'
   // property and the returned Person results have all of their properties.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(2));
 
@@ -1073,8 +1111,9 @@
           .AddStringProperty(
               "body", "Oh what a beautiful morning! Oh what a beautiful day!")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document_one));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_one)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   DocumentProto document_two =
@@ -1085,8 +1124,9 @@
           .AddStringProperty("name", "Joe Fox")
           .AddStringProperty("emailAddress", "ny152@aol.com")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document_two));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_two)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // 2. Setup the scored results.
@@ -1121,14 +1161,16 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // 5. Verify that the returned Email results only contain the 'name'
   // property and the returned Person results only contain the 'name' property.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(2));
 
@@ -1165,8 +1207,9 @@
           .AddStringProperty(
               "body", "Oh what a beautiful morning! Oh what a beautiful day!")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document_one));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_one)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   DocumentProto document_two =
@@ -1177,8 +1220,9 @@
           .AddStringProperty("name", "Joe Fox")
           .AddStringProperty("emailAddress", "ny152@aol.com")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document_two));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_two)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // 2. Setup the scored results.
@@ -1217,14 +1261,16 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // 5. Verify that the returned Email results only contain the 'body'
   // property and the returned Person results  only contain the 'name' property.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(2));
 
@@ -1270,8 +1316,9 @@
                   .AddStringProperty("emailAddress", "mr.body123@gmail.com")
                   .Build())
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document_one));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_one)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   DocumentProto document_two =
@@ -1282,8 +1329,9 @@
           .AddStringProperty("name", "Joe Fox")
           .AddStringProperty("emailAddress", "ny152@aol.com")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document_two));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_two)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // 2. Setup the scored results.
@@ -1322,14 +1370,16 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // 5. Verify that the returned Email results only contain the 'sender.name'
   // property and the returned Person results only contain the 'name' property.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(2));
 
@@ -1379,8 +1429,9 @@
                   .AddStringProperty("emailAddress", "mr.body123@gmail.com")
                   .Build())
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document_one));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_one)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   DocumentProto document_two =
@@ -1391,8 +1442,9 @@
           .AddStringProperty("name", "Joe Fox")
           .AddStringProperty("emailAddress", "ny152@aol.com")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document_two));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_two)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // 2. Setup the scored results.
@@ -1431,14 +1483,16 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // 5. Verify that the returned Email results only contain the 'sender.name'
   // property and the returned Person results contain no properties.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(2));
 
@@ -1476,8 +1530,10 @@
           .AddStringProperty("name", "Joe Fox")
           .AddStringProperty("emailAddress", "ny152@aol.com")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(person_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(person_document)));
   DocumentId person_document_id = put_result1.new_document_id;
 
   // 2. Add two Email documents
@@ -1491,8 +1547,9 @@
               "body", "Oh what a beautiful morning! Oh what a beautiful day!")
           .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(put_result1,
-                             document_store_->Put(email_document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      put_result1, document_store_->Put(
+                       document_util::CreateDocumentWrapper(email_document1)));
   DocumentId email_document_id1 = put_result1.new_document_id;
 
   DocumentProto email_document2 =
@@ -1504,8 +1561,10 @@
           .AddStringProperty("body",
                              "Count all the sheep and tell them 'Hello'.")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(email_document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(email_document2)));
   DocumentId email_document_id2 = put_result2.new_document_id;
 
   // 3. Setup the joined scored results.
@@ -1569,15 +1628,17 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // 7. Verify that the returned results:
   //    - Person docs only contain the "name" property.
   //    - Email docs only contain the "body" property.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(1));
 
@@ -1628,8 +1689,9 @@
           .AddStringProperty("name", "Joe Fox")
           .AddStringProperty("emailAddress", "ny152@aol.com")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document_one));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_one)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   DocumentProto document_two =
@@ -1640,8 +1702,9 @@
           .AddStringProperty("name", "Joe Artist")
           .AddStringProperty("emailAddress", "artist@aol.com")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document_two));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_two)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // 2. Setup the scored results.
@@ -1674,14 +1737,16 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // 5. Verify that the returned Person and Artist results only contain the
   // 'name' property.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(2));
 
@@ -1716,8 +1781,9 @@
           .AddStringProperty("name", "Joe Fox")
           .AddStringProperty("emailAddress", "ny152@aol.com")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document_one));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_one)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   DocumentProto document_two =
@@ -1728,8 +1794,9 @@
           .AddStringProperty("name", "Joe Musician")
           .AddStringProperty("emailAddress", "Musician@aol.com")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document_two));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_two)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // 2. Setup the scored results.
@@ -1762,14 +1829,16 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // 5. Verify that the returned Person and Musician results only contain the
   // 'name' property.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(2));
 
@@ -1804,8 +1873,9 @@
                                .SetSchema("Artist")
                                .AddStringProperty("name", "Joe Artist")
                                .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document)));
   DocumentId document_id = put_result.new_document_id;
 
   // 2. Setup the scored results.
@@ -1838,14 +1908,16 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // 5. Verify that the returned person document does not contain any property,
   // since 'emailAddress' is missing.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(1));
   DocumentProto projected_document = DocumentBuilder()
@@ -1867,8 +1939,9 @@
           .AddStringProperty("name", "Joe Fox")
           .AddStringProperty("emailAddress", "ny152@aol.com")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document_one));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_one)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   DocumentProto document_two =
@@ -1879,8 +1952,9 @@
           .AddStringProperty("name", "Joe Artist")
           .AddStringProperty("emailAddress", "artist@aol.com")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document_two));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(document_util::CreateDocumentWrapper(document_two)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // 2. Setup the scored results.
@@ -1918,15 +1992,17 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // 5. Verify that the returned Person results only contain the 'name'
   // property and the returned Artist results contain both the 'name' and
   // 'emailAddress' properties.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(2));
 
@@ -1963,8 +2039,9 @@
                                .AddStringProperty("phoneNumber", "12345")
                                .AddStringProperty("phoneModel", "pixel")
                                .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document)));
   DocumentId document_id = put_result.new_document_id;
 
   // 2. Setup the scored results.
@@ -2003,14 +2080,16 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // 5. Verify that the returned document only contains the 'name' and the
   // 'phoneNumber' property.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(1));
 

diff --git a/icing/result/result-retriever-v2_snippet_test.cc b/icing/result/result-retriever-v2_snippet_test.cc
index 09b691e..b2bf82c 100644
--- a/icing/result/result-retriever-v2_snippet_test.cc
+++ b/icing/result/result-retriever-v2_snippet_test.cc

@@ -29,6 +29,7 @@
 #include "icing/file/portable-file-backed-proto-log.h"
 #include "icing/index/embed/embedding-query-results.h"
 #include "icing/portable/equals-proto.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/portable/platform.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/schema.pb.h"
@@ -58,6 +59,7 @@
 #include "icing/transform/normalizer-factory.h"
 #include "icing/transform/normalizer-options.h"
 #include "icing/transform/normalizer.h"
+#include "icing/util/document-util.h"
 #include "icing/util/icu-data-file-helper.h"
 #include "icing/util/snippet-helpers.h"
 #include "unicode/uloc.h"
@@ -147,14 +149,18 @@
 
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
-        DocumentStore::Create(&filesystem_, test_dir_, &fake_clock_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, test_dir_, &fake_clock_, schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     document_store_ = std::move(create_result.document_store);
   }
 
@@ -271,31 +277,35 @@
 EmbeddingQueryResults CreateEmailEmbeddingQueryResults(int num_documents) {
   SectionId embedding1_section_id = 1;
   SectionId embedding2_section_id = 2;
-  EmbeddingQueryResults embedding_query_results;
+  EmbeddingQueryResults embedding_query_results(/*num_query_vectors=*/2);
 
   for (int doc_id = 0; doc_id < num_documents; ++doc_id) {
     EmbeddingMatchInfos& info_model1 =
-        embedding_query_results
-            .result_infos[/*query_index=*/0][EMBEDDING_METRIC_DOT_PRODUCT]
-                         [doc_id];
-    info_model1.AppendScore(1.1 + doc_id);
-    info_model1.AppendScore(2.2 + doc_id);
+        GetOrCreateEmbeddingMatchInfosForDocument(
+            embedding_query_results, /*query_index=*/0,
+            EMBEDDING_METRIC_DOT_PRODUCT, doc_id);
+    info_model1.AppendScore(*embedding_query_results.global_scores,
+                            1.1 + doc_id);
     // {2, 3, 4 + doc_id}, position 2
     info_model1.AppendSectionInfo(
-        embedding1_section_id,
+        *embedding_query_results.global_section_infos, embedding1_section_id,
         /*position_in_section_for_dimension_and_signature=*/1);
+    info_model1.AppendScore(*embedding_query_results.global_scores,
+                            2.2 + doc_id);
     // {-1, -2, -6 + doc_id}, position 1
     info_model1.AppendSectionInfo(
-        embedding2_section_id,
+        *embedding_query_results.global_section_infos, embedding2_section_id,
         /*position_in_section_for_dimension_and_signature=*/0);
 
     EmbeddingMatchInfos& info_model2 =
-        embedding_query_results
-            .result_infos[/*query_index=*/1][EMBEDDING_METRIC_COSINE][doc_id];
-    info_model2.AppendScore(3.3 + doc_id);
+        GetOrCreateEmbeddingMatchInfosForDocument(
+            embedding_query_results, /*query_index=*/1, EMBEDDING_METRIC_COSINE,
+            doc_id);
+    info_model2.AppendScore(*embedding_query_results.global_scores,
+                            3.3 + doc_id);
     // {1, 2, 3, 4 + doc_id}, position 1
     info_model2.AppendSectionInfo(
-        embedding1_section_id,
+        *embedding_query_results.global_section_infos, embedding1_section_id,
         /*position_in_section_for_dimension_and_signature=*/0);
   }
   return embedding_query_results;
@@ -305,15 +315,18 @@
        DefaultSnippetSpecShouldDisableSnippeting) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result1,
-      document_store_->Put(CreateEmailDocument(/*id=*/1)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/1))));
   DocumentId document_id1 = put_result1.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result2,
-      document_store_->Put(CreateEmailDocument(/*id=*/2)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/2))));
   DocumentId document_id2 = put_result2.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result3,
-      document_store_->Put(CreateEmailDocument(/*id=*/3)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/3))));
   DocumentId document_id3 = put_result3.new_document_id;
 
   std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "subject"),
@@ -326,7 +339,8 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   ResultSpecProto result_spec = CreateResultSpec(/*num_per_page=*/3);
 
@@ -345,8 +359,10 @@
       /*child_adjustment_info=*/nullptr, result_spec, *document_store_);
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state,
+              /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(3));
   EXPECT_THAT(page_result.results.at(0).snippet(),
@@ -361,15 +377,18 @@
 TEST_F(ResultRetrieverV2SnippetTest, SimpleSnippeted) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result1,
-      document_store_->Put(CreateEmailDocument(/*id=*/1)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/1))));
   DocumentId document_id1 = put_result1.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result2,
-      document_store_->Put(CreateEmailDocument(/*id=*/2)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/2))));
   DocumentId document_id2 = put_result2.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result3,
-      document_store_->Put(CreateEmailDocument(/*id=*/3)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/3))));
   DocumentId document_id3 = put_result3.new_document_id;
 
   std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "subject"),
@@ -382,7 +401,8 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // Create ResultSpec with custom snippet spec.
   ResultSpecProto result_spec = CreateResultSpec(/*num_per_page=*/3);
@@ -404,8 +424,10 @@
 
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state,
+              /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(3));
   EXPECT_THAT(page_result.num_results_with_snippets, Eq(3));
@@ -476,15 +498,18 @@
 TEST_F(ResultRetrieverV2SnippetTest, OnlyOneDocumentSnippeted) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result1,
-      document_store_->Put(CreateEmailDocument(/*id=*/1)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/1))));
   DocumentId document_id1 = put_result1.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result2,
-      document_store_->Put(CreateEmailDocument(/*id=*/2)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/2))));
   DocumentId document_id2 = put_result2.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result3,
-      document_store_->Put(CreateEmailDocument(/*id=*/3)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/3))));
   DocumentId document_id3 = put_result3.new_document_id;
 
   std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "subject"),
@@ -497,7 +522,8 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // Create ResultSpec with custom snippet spec.
   ResultSpecProto::SnippetSpecProto snippet_spec = CreateSnippetSpec();
@@ -521,8 +547,10 @@
 
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state,
+              /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(3));
   EXPECT_THAT(page_result.num_results_with_snippets, Eq(1));
@@ -560,15 +588,18 @@
 TEST_F(ResultRetrieverV2SnippetTest, SnippetWithGetEmbeddingMatchInfo) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result1,
-      document_store_->Put(CreateEmailDocument(/*id=*/1)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/1))));
   DocumentId document_id1 = put_result1.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result2,
-      document_store_->Put(CreateEmailDocument(/*id=*/2)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/2))));
   DocumentId document_id2 = put_result2.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result3,
-      document_store_->Put(CreateEmailDocument(/*id=*/3)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/3))));
   DocumentId document_id3 = put_result3.new_document_id;
 
   std::vector<SectionId> hit_section_ids = {
@@ -582,7 +613,8 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // Create ResultSpec with custom snippet spec.
   ResultSpecProto::SnippetSpecProto snippet_spec = CreateSnippetSpec();
@@ -613,8 +645,10 @@
 
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state,
+              /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(3));
   EXPECT_THAT(page_result.num_results_with_snippets, Eq(3));
@@ -738,15 +772,18 @@
 TEST_F(ResultRetrieverV2SnippetTest, SnippetWithGetEmbeddingMatchInfoDisabled) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result1,
-      document_store_->Put(CreateEmailDocument(/*id=*/1)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/1))));
   DocumentId document_id1 = put_result1.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result2,
-      document_store_->Put(CreateEmailDocument(/*id=*/2)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/2))));
   DocumentId document_id2 = put_result2.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result3,
-      document_store_->Put(CreateEmailDocument(/*id=*/3)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/3))));
   DocumentId document_id3 = put_result3.new_document_id;
 
   std::vector<SectionId> hit_section_ids = {
@@ -760,7 +797,8 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // Create ResultSpec with custom snippet spec.
   ResultSpecProto::SnippetSpecProto snippet_spec = CreateSnippetSpec();
@@ -791,8 +829,10 @@
 
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state,
+              /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(3));
   EXPECT_THAT(page_result.num_results_with_snippets, Eq(3));
@@ -863,15 +903,18 @@
 TEST_F(ResultRetrieverV2SnippetTest, ShouldSnippetSomeResults) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result1,
-      document_store_->Put(CreateEmailDocument(/*id=*/1)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/1))));
   DocumentId document_id1 = put_result1.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result2,
-      document_store_->Put(CreateEmailDocument(/*id=*/2)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/2))));
   DocumentId document_id2 = put_result2.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result3,
-      document_store_->Put(CreateEmailDocument(/*id=*/3)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/3))));
   DocumentId document_id3 = put_result3.new_document_id;
 
   std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "subject"),
@@ -884,7 +927,8 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // Create ResultSpec with custom snippet spec.
   ResultSpecProto::SnippetSpecProto snippet_spec = CreateSnippetSpec();
@@ -914,8 +958,10 @@
 
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state,
+              /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(3));
   EXPECT_THAT(page_result.results.at(0).snippet().entries(), Not(IsEmpty()));
@@ -927,15 +973,18 @@
 TEST_F(ResultRetrieverV2SnippetTest, ShouldNotSnippetAnyResults) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result1,
-      document_store_->Put(CreateEmailDocument(/*id=*/1)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/1))));
   DocumentId document_id1 = put_result1.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result2,
-      document_store_->Put(CreateEmailDocument(/*id=*/2)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/2))));
   DocumentId document_id2 = put_result2.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result3,
-      document_store_->Put(CreateEmailDocument(/*id=*/3)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/3))));
   DocumentId document_id3 = put_result3.new_document_id;
 
   std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "subject"),
@@ -948,7 +997,8 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // Create ResultSpec with custom snippet spec.
   ResultSpecProto::SnippetSpecProto snippet_spec = CreateSnippetSpec();
@@ -979,8 +1029,10 @@
   // We can't return any snippets for this page.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state,
+              /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(3));
   EXPECT_THAT(page_result.results.at(0).snippet().entries(), IsEmpty());
@@ -993,15 +1045,18 @@
        ShouldNotSnippetAnyResultsForNonPositiveNumMatchesPerProperty) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result1,
-      document_store_->Put(CreateEmailDocument(/*id=*/1)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/1))));
   DocumentId document_id1 = put_result1.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result2,
-      document_store_->Put(CreateEmailDocument(/*id=*/2)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/2))));
   DocumentId document_id2 = put_result2.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result3,
-      document_store_->Put(CreateEmailDocument(/*id=*/3)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/3))));
   DocumentId document_id3 = put_result3.new_document_id;
 
   std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "subject"),
@@ -1014,7 +1069,8 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // Create ResultSpec with custom snippet spec.
   ResultSpecProto::SnippetSpecProto snippet_spec = CreateSnippetSpec();
@@ -1047,8 +1103,10 @@
   // We can't return any snippets for this page even though num_to_snippet > 0.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state,
+              /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(3));
   EXPECT_THAT(page_result.results.at(0).snippet().entries(), IsEmpty());
@@ -1060,28 +1118,34 @@
 TEST_F(ResultRetrieverV2SnippetTest, JoinSnippeted) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult person_put_result1,
-      document_store_->Put(CreatePersonDocument(/*id=*/1)));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          CreatePersonDocument(/*id=*/1))));
   DocumentId person_document_id1 = person_put_result1.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult person_put_result2,
-      document_store_->Put(CreatePersonDocument(/*id=*/2)));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          CreatePersonDocument(/*id=*/2))));
   DocumentId person_document_id2 = person_put_result2.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult person_put_result3,
-      document_store_->Put(CreatePersonDocument(/*id=*/3)));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          CreatePersonDocument(/*id=*/3))));
   DocumentId person_document_id3 = person_put_result3.new_document_id;
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult email_put_result1,
-      document_store_->Put(CreateEmailDocument(/*id=*/1)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/1))));
   DocumentId email_document_id1 = email_put_result1.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult email_put_result2,
-      document_store_->Put(CreateEmailDocument(/*id=*/2)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/2))));
   DocumentId email_document_id2 = email_put_result2.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult email_put_result3,
-      document_store_->Put(CreateEmailDocument(/*id=*/3)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/3))));
   DocumentId email_document_id3 = email_put_result3.new_document_id;
 
   std::vector<SectionId> person_hit_section_ids = {
@@ -1124,7 +1188,8 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // Create parent ResultSpec with custom snippet spec.
   ResultSpecProto parent_result_spec = CreateResultSpec(/*num_per_page=*/3);
@@ -1163,8 +1228,10 @@
 
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state,
+              /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(3));
   EXPECT_THAT(page_result.num_results_with_snippets, Eq(3));
@@ -1304,24 +1371,29 @@
 TEST_F(ResultRetrieverV2SnippetTest, ShouldSnippetAllJoinedResults) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult person_put_result1,
-      document_store_->Put(CreatePersonDocument(/*id=*/1)));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          CreatePersonDocument(/*id=*/1))));
   DocumentId person_document_id1 = person_put_result1.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult person_put_result2,
-      document_store_->Put(CreatePersonDocument(/*id=*/2)));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          CreatePersonDocument(/*id=*/2))));
   DocumentId person_document_id2 = person_put_result2.new_document_id;
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult email_put_result1,
-      document_store_->Put(CreateEmailDocument(/*id=*/1)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/1))));
   DocumentId email_document_id1 = email_put_result1.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult email_put_result2,
-      document_store_->Put(CreateEmailDocument(/*id=*/2)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/2))));
   DocumentId email_document_id2 = email_put_result2.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult email_put_result3,
-      document_store_->Put(CreateEmailDocument(/*id=*/3)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/3))));
   DocumentId email_document_id3 = email_put_result3.new_document_id;
 
   std::vector<SectionId> person_hit_section_ids = {
@@ -1359,7 +1431,8 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // Create parent ResultSpec with custom snippet spec.
   ResultSpecProto::SnippetSpecProto parent_snippet_spec = CreateSnippetSpec();
@@ -1403,8 +1476,10 @@
   // should be snippeted.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state,
+              /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(2));
 
@@ -1432,24 +1507,29 @@
 TEST_F(ResultRetrieverV2SnippetTest, ShouldSnippetSomeJoinedResults) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult person_put_result1,
-      document_store_->Put(CreatePersonDocument(/*id=*/1)));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          CreatePersonDocument(/*id=*/1))));
   DocumentId person_document_id1 = person_put_result1.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult person_put_result2,
-      document_store_->Put(CreatePersonDocument(/*id=*/2)));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          CreatePersonDocument(/*id=*/2))));
   DocumentId person_document_id2 = person_put_result2.new_document_id;
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult email_put_result1,
-      document_store_->Put(CreateEmailDocument(/*id=*/1)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/1))));
   DocumentId email_document_id1 = email_put_result1.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult email_put_result2,
-      document_store_->Put(CreateEmailDocument(/*id=*/2)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/2))));
   DocumentId email_document_id2 = email_put_result2.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult email_put_result3,
-      document_store_->Put(CreateEmailDocument(/*id=*/3)));
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(CreateEmailDocument(/*id=*/3))));
   DocumentId email_document_id3 = email_put_result3.new_document_id;
 
   std::vector<SectionId> person_hit_section_ids = {
@@ -1487,7 +1567,8 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // Create parent ResultSpec with custom snippet spec.
   ResultSpecProto::SnippetSpecProto parent_snippet_spec = CreateSnippetSpec();
@@ -1531,8 +1612,10 @@
   // snippeted.
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state,
+              /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result.results, SizeIs(2));
 

diff --git a/icing/result/result-retriever-v2_test.cc b/icing/result/result-retriever-v2_test.cc
index 49899b1..6c3c727 100644
--- a/icing/result/result-retriever-v2_test.cc
+++ b/icing/result/result-retriever-v2_test.cc

@@ -19,6 +19,7 @@
 #include <cstdint>
 #include <limits>
 #include <memory>
+#include <optional>
 #include <string>
 #include <unordered_map>
 #include <utility>
@@ -35,6 +36,7 @@
 #include "icing/file/mock-filesystem.h"
 #include "icing/file/portable-file-backed-proto-log.h"
 #include "icing/portable/equals-proto.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/portable/platform.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/document_wrapper.pb.h"
@@ -53,7 +55,6 @@
 #include "icing/testing/common-matchers.h"
 #include "icing/testing/fake-clock.h"
 #include "icing/testing/test-data.h"
-#include "icing/testing/test-feature-flags.h"
 #include "icing/testing/tmp-directory.h"
 #include "icing/tokenization/language-segmenter-factory.h"
 #include "icing/tokenization/language-segmenter.h"
@@ -61,6 +62,7 @@
 #include "icing/transform/normalizer-options.h"
 #include "icing/transform/normalizer.h"
 #include "icing/util/clock.h"
+#include "icing/util/document-util.h"
 #include "icing/util/icu-data-file-helper.h"
 #include "unicode/uloc.h"
 
@@ -80,28 +82,30 @@
 using ::testing::SizeIs;
 using EntryIdMap = std::unordered_map<int32_t, int>;
 
-// Mock the behavior of GroupResultLimiter::ShouldBeRemoved.
+// Mock the behavior of GroupResultLimiter::GetGroupResultLimitsIndex: get
+// -1 to avoid excluding documents.
 class MockGroupResultLimiter : public GroupResultLimiterV2 {
  public:
   MockGroupResultLimiter() : GroupResultLimiterV2() {
-    ON_CALL(*this, ShouldBeRemoved).WillByDefault(Return(false));
+    ON_CALL(*this, GetGroupResultLimitsIndex).WillByDefault(Return(-1));
   }
 
-  MOCK_METHOD(bool, ShouldBeRemoved,
+  MOCK_METHOD(std::optional<int>, GetGroupResultLimitsIndex,
               (const ScoredDocumentHit&, const EntryIdMap&,
-               const DocumentStore&, std::vector<int>&,
-               ResultSpecProto::ResultGroupingType, int64_t),
+               const DocumentStore&, ResultSpecProto::ResultGroupingType,
+               int64_t),
               (const, override));
 };
 
-class ResultRetrieverV2Test : public ::testing::Test {
+class ResultRetrieverV2Test : public ::testing::TestWithParam<FeatureFlags> {
  protected:
   ResultRetrieverV2Test() : test_dir_(GetTestTempDir() + "/icing") {
     filesystem_.CreateDirectoryRecursively(test_dir_.c_str());
   }
 
   void SetUp() override {
-    feature_flags_ = std::make_unique<FeatureFlags>(GetTestFeatureFlags());
+    feature_flags_ = std::make_unique<FeatureFlags>(GetParam());
+
     if (!IsCfStringTokenization() && !IsReverseJniTokenization()) {
       ICING_ASSERT_OK(
           // File generated via icu_data_file rule in //icing/BUILD.
@@ -231,13 +235,17 @@
       /*force_recovery_and_revalidate_documents=*/false,
       /*pre_mapping_fbv=*/false, /*use_persistent_hash_map=*/true,
       PortableFileBackedProtoLog<DocumentWrapper>::kDefaultCompressionLevel,
+      PortableFileBackedProtoLog<
+          DocumentWrapper>::kDefaultCompressionThresholdBytes,
+      protobuf_ports::kDefaultMemLevel,
       /*initialize_stats=*/nullptr);
 }
 
-TEST_F(ResultRetrieverV2Test, CreationWithNullPointerShouldFail) {
+TEST_P(ResultRetrieverV2Test, CreationWithNullPointerShouldFail) {
   EXPECT_THAT(
-      ResultRetrieverV2::Create(/*doc_store=*/nullptr, schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()),
+      ResultRetrieverV2::Create(
+          /*doc_store=*/nullptr, schema_store_.get(), language_segmenter_.get(),
+          normalizer_.get(), feature_flags_.get()),
       StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -249,19 +257,32 @@
 
   EXPECT_THAT(
       ResultRetrieverV2::Create(doc_store.get(), /*schema_store=*/nullptr,
-                                language_segmenter_.get(), normalizer_.get()),
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()),
+      StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
+  EXPECT_THAT(
+      ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(),
+                                /*language_segmenter=*/nullptr,
+                                normalizer_.get(), feature_flags_.get()),
+      StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
+  EXPECT_THAT(
+      ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(),
+                                language_segmenter_.get(),
+                                /*normalizer=*/nullptr, feature_flags_.get()),
+      StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
+  EXPECT_THAT(
+      ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(),
+                                language_segmenter_.get(), normalizer_.get(),
+                                /*feature_flags=*/nullptr),
       StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
   EXPECT_THAT(ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(),
-                                        /*language_segmenter=*/nullptr,
-                                        normalizer_.get()),
-              StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
-  EXPECT_THAT(ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(),
                                         language_segmenter_.get(),
-                                        /*normalizer=*/nullptr),
+                                        normalizer_.get(), feature_flags_.get(),
+                                        /*group_result_limiter=*/nullptr),
               StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
 }
 
-TEST_F(ResultRetrieverV2Test, ShouldRetrieveSimpleResults) {
+TEST_P(ResultRetrieverV2Test, ShouldRetrieveSimpleResults) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::CreateResult create_result,
       CreateDocumentStore(&filesystem_, test_dir_, &fake_clock_,
@@ -269,20 +290,30 @@
   std::unique_ptr<DocumentStore> doc_store =
       std::move(create_result.document_store);
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             doc_store->Put(CreateDocument(/*id=*/1)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/1))));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             doc_store->Put(CreateDocument(/*id=*/2)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/2))));
   DocumentId document_id2 = put_result2.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             doc_store->Put(CreateDocument(/*id=*/3)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/3))));
   DocumentId document_id3 = put_result3.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result4,
-                             doc_store->Put(CreateDocument(/*id=*/4)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result4,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/4))));
   DocumentId document_id4 = put_result4.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result5,
-                             doc_store->Put(CreateDocument(/*id=*/5)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result5,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/5))));
   DocumentId document_id5 = put_result5.new_document_id;
 
   std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "name"),
@@ -297,7 +328,8 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   SearchResultProto::ResultProto result1;
   *result1.mutable_document() = CreateDocument(/*id=*/1);
@@ -325,7 +357,8 @@
 
   // First page, 2 results
   auto [page_result1, has_more_results1] = result_retriever->RetrieveNextPage(
-      result_state, fake_clock_.GetSystemTimeMilliseconds());
+      result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+      fake_clock_.GetSystemTimeMilliseconds());
   EXPECT_THAT(page_result1.results,
               ElementsAre(EqualsProto(result1), EqualsProto(result2)));
   // num_results_with_snippets is 0 when there is no snippet.
@@ -337,7 +370,8 @@
 
   // Second page, 2 results
   auto [page_result2, has_more_results2] = result_retriever->RetrieveNextPage(
-      result_state, fake_clock_.GetSystemTimeMilliseconds());
+      result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+      fake_clock_.GetSystemTimeMilliseconds());
   EXPECT_THAT(page_result2.results,
               ElementsAre(EqualsProto(result3), EqualsProto(result4)));
   // num_results_with_snippets is 0 when there is no snippet.
@@ -349,7 +383,8 @@
 
   // Third page, 1 result
   auto [page_result3, has_more_results3] = result_retriever->RetrieveNextPage(
-      result_state, fake_clock_.GetSystemTimeMilliseconds());
+      result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+      fake_clock_.GetSystemTimeMilliseconds());
   EXPECT_THAT(page_result3.results, ElementsAre(EqualsProto(result5)));
   // num_results_with_snippets is 0 when there is no snippet.
   EXPECT_THAT(page_result3.num_results_with_snippets, Eq(0));
@@ -359,7 +394,7 @@
   EXPECT_FALSE(has_more_results3);
 }
 
-TEST_F(ResultRetrieverV2Test, ShouldIgnoreNonInternalErrors) {
+TEST_P(ResultRetrieverV2Test, ShouldIgnoreNonInternalErrors) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::CreateResult create_result,
       CreateDocumentStore(&filesystem_, test_dir_, &fake_clock_,
@@ -367,11 +402,15 @@
   std::unique_ptr<DocumentStore> doc_store =
       std::move(create_result.document_store);
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             doc_store->Put(CreateDocument(/*id=*/1)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/1))));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             doc_store->Put(CreateDocument(/*id=*/2)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/2))));
   DocumentId document_id2 = put_result2.new_document_id;
 
   DocumentId invalid_document_id = -1;
@@ -386,6 +425,7 @@
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(),
                                 language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get(),
                                 std::make_unique<MockGroupResultLimiter>()));
 
   SearchResultProto::ResultProto result1;
@@ -405,8 +445,10 @@
       *doc_store);
   PageResult page_result1 =
       result_retriever
-          ->RetrieveNextPage(result_state1,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state1,
+              /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   EXPECT_THAT(page_result1.results,
               ElementsAre(EqualsProto(result1), EqualsProto(result2)));
@@ -426,14 +468,16 @@
       *doc_store);
   PageResult page_result2 =
       result_retriever
-          ->RetrieveNextPage(result_state2,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state2,
+              /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   EXPECT_THAT(page_result2.results,
               ElementsAre(EqualsProto(result1), EqualsProto(result2)));
 }
 
-TEST_F(ResultRetrieverV2Test,
+TEST_P(ResultRetrieverV2Test,
        ShouldLimitNumChildDocumentsByMaxJoinedChildPerParent) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::CreateResult create_result,
@@ -451,8 +495,9 @@
           .AddStringProperty("name", "Joe Fox")
           .AddStringProperty("emailAddress", "ny152@aol.com")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             doc_store->Put(person_document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      doc_store->Put(document_util::CreateDocumentWrapper(person_document1)));
   DocumentId person_document_id1 = put_result1.new_document_id;
 
   DocumentProto person_document2 =
@@ -463,8 +508,9 @@
           .AddStringProperty("name", "Meg Ryan")
           .AddStringProperty("emailAddress", "shopgirl@aol.com")
           .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             doc_store->Put(person_document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      doc_store->Put(document_util::CreateDocumentWrapper(person_document2)));
   DocumentId person_document_id2 = put_result2.new_document_id;
 
   // 2. Add 4 Email documents
@@ -475,7 +521,9 @@
                                       .AddStringProperty("name", "Test 1")
                                       .AddStringProperty("body", "Test 1")
                                       .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(put_result1, doc_store->Put(email_document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      put_result1,
+      doc_store->Put(document_util::CreateDocumentWrapper(email_document1)));
   DocumentId email_document_id1 = put_result1.new_document_id;
 
   DocumentProto email_document2 = DocumentBuilder()
@@ -485,7 +533,9 @@
                                       .AddStringProperty("name", "Test 2")
                                       .AddStringProperty("body", "Test 2")
                                       .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(put_result2, doc_store->Put(email_document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      put_result2,
+      doc_store->Put(document_util::CreateDocumentWrapper(email_document2)));
   DocumentId email_document_id2 = put_result2.new_document_id;
 
   DocumentProto email_document3 = DocumentBuilder()
@@ -495,8 +545,9 @@
                                       .AddStringProperty("name", "Test 3")
                                       .AddStringProperty("body", "Test 3")
                                       .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             doc_store->Put(email_document3));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      doc_store->Put(document_util::CreateDocumentWrapper(email_document3)));
   DocumentId email_document_id3 = put_result3.new_document_id;
 
   DocumentProto email_document4 = DocumentBuilder()
@@ -506,8 +557,9 @@
                                       .AddStringProperty("name", "Test 4")
                                       .AddStringProperty("body", "Test 4")
                                       .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result4,
-                             doc_store->Put(email_document4));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result4,
+      doc_store->Put(document_util::CreateDocumentWrapper(email_document4)));
   DocumentId email_document_id4 = put_result4.new_document_id;
 
   // 3. Setup the joined scored results.
@@ -555,7 +607,8 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
   ResultStateV2 result_state(
       std::make_unique<
           PriorityQueueScoredDocumentHitsRanker<JoinedScoredDocumentHit>>(
@@ -583,14 +636,15 @@
   child3->set_score(3);
 
   auto [page_result, has_more_results] = result_retriever->RetrieveNextPage(
-      result_state, fake_clock_.GetSystemTimeMilliseconds());
+      result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+      fake_clock_.GetSystemTimeMilliseconds());
   EXPECT_THAT(page_result.results,
               ElementsAre(EqualsProto(result1), EqualsProto(result2)));
   // No more results.
   EXPECT_FALSE(has_more_results);
 }
 
-TEST_F(ResultRetrieverV2Test, ShouldIgnoreInternalErrors) {
+TEST_P(ResultRetrieverV2Test, ShouldIgnoreInternalErrors) {
   MockFilesystem mock_filesystem;
   EXPECT_CALL(mock_filesystem,
               PRead(A<int>(), A<void*>(), A<size_t>(), A<off_t>()))
@@ -604,11 +658,15 @@
   std::unique_ptr<DocumentStore> doc_store =
       std::move(create_result.document_store);
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             doc_store->Put(CreateDocument(/*id=*/1)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/1))));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             doc_store->Put(CreateDocument(/*id=*/2)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/2))));
   DocumentId document_id2 = put_result2.new_document_id;
 
   std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "name"),
@@ -622,6 +680,7 @@
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(),
                                 language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get(),
                                 std::make_unique<MockGroupResultLimiter>()));
 
   SearchResultProto::ResultProto result1;
@@ -638,15 +697,17 @@
       *doc_store);
   PageResult page_result =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state,
+              /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   // We mocked mock_filesystem to return an internal error when retrieving doc2,
   // so doc2 should be skipped and doc1 should still be returned.
   EXPECT_THAT(page_result.results, ElementsAre(EqualsProto(result1)));
 }
 
-TEST_F(ResultRetrieverV2Test, ShouldUpdateResultState) {
+TEST_P(ResultRetrieverV2Test, ShouldUpdateResultState) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::CreateResult create_result,
       CreateDocumentStore(&filesystem_, test_dir_, &fake_clock_,
@@ -654,20 +715,30 @@
   std::unique_ptr<DocumentStore> doc_store =
       std::move(create_result.document_store);
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             doc_store->Put(CreateDocument(/*id=*/1)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/1))));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             doc_store->Put(CreateDocument(/*id=*/2)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/2))));
   DocumentId document_id2 = put_result2.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             doc_store->Put(CreateDocument(/*id=*/3)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/3))));
   DocumentId document_id3 = put_result3.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result4,
-                             doc_store->Put(CreateDocument(/*id=*/4)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result4,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/4))));
   DocumentId document_id4 = put_result4.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result5,
-                             doc_store->Put(CreateDocument(/*id=*/5)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result5,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/5))));
   DocumentId document_id5 = put_result5.new_document_id;
 
   std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "name"),
@@ -682,7 +753,8 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   ResultStateV2 result_state(
       std::make_unique<
@@ -696,8 +768,10 @@
   // First page, 2 results
   PageResult page_result1 =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state,
+              /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result1.results, SizeIs(2));
   {
@@ -713,8 +787,10 @@
   // Second page, 2 results
   PageResult page_result2 =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state,
+              /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result2.results, SizeIs(2));
   {
@@ -730,8 +806,10 @@
   // Third page, 1 result
   PageResult page_result3 =
       result_retriever
-          ->RetrieveNextPage(result_state,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              result_state,
+              /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result3.results, SizeIs(1));
   {
@@ -745,7 +823,7 @@
   }
 }
 
-TEST_F(ResultRetrieverV2Test, ShouldUpdateNumTotalHits) {
+TEST_P(ResultRetrieverV2Test, ShouldUpdateNumTotalHits) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::CreateResult create_result,
       CreateDocumentStore(&filesystem_, test_dir_, &fake_clock_,
@@ -757,11 +835,15 @@
                                             GetSectionId("Email", "body")};
   SectionIdMask hit_section_id_mask = CreateSectionIdMask(hit_section_ids);
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             doc_store->Put(CreateDocument(/*id=*/1)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/1))));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             doc_store->Put(CreateDocument(/*id=*/2)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/2))));
   DocumentId document_id2 = put_result2.new_document_id;
   std::vector<ScoredDocumentHit> scored_document_hits1 = {
       {document_id1, hit_section_id_mask, /*score=*/0},
@@ -782,14 +864,20 @@
     ASSERT_THAT(num_total_hits_, Eq(2));
   }
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             doc_store->Put(CreateDocument(/*id=*/3)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/3))));
   DocumentId document_id3 = put_result3.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result4,
-                             doc_store->Put(CreateDocument(/*id=*/4)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result4,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/4))));
   DocumentId document_id4 = put_result4.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result5,
-                             doc_store->Put(CreateDocument(/*id=*/5)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result5,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/5))));
   DocumentId document_id5 = put_result5.new_document_id;
   std::vector<ScoredDocumentHit> scored_document_hits2 = {
       {document_id3, hit_section_id_mask, /*score=*/0},
@@ -814,14 +902,17 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   // Should get 1 doc in the first page of result_state1, and num_total_hits
   // should be decremented by 1.
   PageResult page_result1 =
       result_retriever
-          ->RetrieveNextPage(*result_state1,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              *result_state1,
+              /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result1.results, SizeIs(1));
   EXPECT_THAT(num_total_hits_, Eq(4));
@@ -830,8 +921,10 @@
   // should be decremented by 2.
   PageResult page_result2 =
       result_retriever
-          ->RetrieveNextPage(*result_state2,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              *result_state2,
+              /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result2.results, SizeIs(2));
   EXPECT_THAT(num_total_hits_, Eq(2));
@@ -841,8 +934,10 @@
   // by 1.
   PageResult page_result3 =
       result_retriever
-          ->RetrieveNextPage(*result_state2,
-                             fake_clock_.GetSystemTimeMilliseconds())
+          ->RetrieveNextPage(
+              *result_state2,
+              /*max_results=*/std::numeric_limits<int32_t>::max(),
+              fake_clock_.GetSystemTimeMilliseconds())
           .first;
   ASSERT_THAT(page_result3.results, SizeIs(1));
   EXPECT_THAT(num_total_hits_, Eq(1));
@@ -858,7 +953,7 @@
   EXPECT_THAT(num_total_hits_, Eq(0));
 }
 
-TEST_F(ResultRetrieverV2Test, ShouldLimitNumTotalBytesPerPage) {
+TEST_P(ResultRetrieverV2Test, ShouldLimitNumTotalBytesPerPage) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::CreateResult create_result,
       CreateDocumentStore(&filesystem_, test_dir_, &fake_clock_,
@@ -866,11 +961,15 @@
   std::unique_ptr<DocumentStore> doc_store =
       std::move(create_result.document_store);
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             doc_store->Put(CreateDocument(/*id=*/1)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/1))));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             doc_store->Put(CreateDocument(/*id=*/2)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/2))));
   DocumentId document_id2 = put_result2.new_document_id;
 
   std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "name"),
@@ -882,7 +981,8 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   SearchResultProto::ResultProto result1;
   *result1.mutable_document() = CreateDocument(/*id=*/1);
@@ -906,20 +1006,22 @@
   // num_total_bytes_per_page_threshold and ResultRetriever should terminate
   // early even though # of results is still below num_per_page.
   auto [page_result1, has_more_results1] = result_retriever->RetrieveNextPage(
-      result_state, fake_clock_.GetSystemTimeMilliseconds());
+      result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+      fake_clock_.GetSystemTimeMilliseconds());
   EXPECT_THAT(page_result1.results, ElementsAre(EqualsProto(result1)));
   // Has more results.
   EXPECT_TRUE(has_more_results1);
 
   // Second page, result2.
   auto [page_result2, has_more_results2] = result_retriever->RetrieveNextPage(
-      result_state, fake_clock_.GetSystemTimeMilliseconds());
+      result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+      fake_clock_.GetSystemTimeMilliseconds());
   EXPECT_THAT(page_result2.results, ElementsAre(EqualsProto(result2)));
   // No more results.
   EXPECT_FALSE(has_more_results2);
 }
 
-TEST_F(ResultRetrieverV2Test,
+TEST_P(ResultRetrieverV2Test,
        ShouldReturnSingleLargeResultAboveNumTotalBytesPerPageThreshold) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::CreateResult create_result,
@@ -928,11 +1030,15 @@
   std::unique_ptr<DocumentStore> doc_store =
       std::move(create_result.document_store);
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             doc_store->Put(CreateDocument(/*id=*/1)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/1))));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             doc_store->Put(CreateDocument(/*id=*/2)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/2))));
   DocumentId document_id2 = put_result2.new_document_id;
 
   std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "name"),
@@ -944,7 +1050,8 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   SearchResultProto::ResultProto result1;
   *result1.mutable_document() = CreateDocument(/*id=*/1);
@@ -970,21 +1077,27 @@
   // First page. Should return single result1 even though its byte size exceeds
   // num_total_bytes_per_page_threshold.
   auto [page_result1, has_more_results1] = result_retriever->RetrieveNextPage(
-      result_state, fake_clock_.GetSystemTimeMilliseconds());
+      result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+      fake_clock_.GetSystemTimeMilliseconds());
   EXPECT_THAT(page_result1.results, ElementsAre(EqualsProto(result1)));
   // Has more results.
   EXPECT_TRUE(has_more_results1);
 
   // Second page, result2.
   auto [page_result2, has_more_results2] = result_retriever->RetrieveNextPage(
-      result_state, fake_clock_.GetSystemTimeMilliseconds());
+      result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+      fake_clock_.GetSystemTimeMilliseconds());
   EXPECT_THAT(page_result2.results, ElementsAre(EqualsProto(result2)));
   // No more results.
   EXPECT_FALSE(has_more_results2);
 }
 
-TEST_F(ResultRetrieverV2Test,
+TEST_P(ResultRetrieverV2Test,
        ShouldRetrieveNextResultWhenBelowNumTotalBytesPerPageThreshold) {
+  if (feature_flags_->enable_strict_page_byte_size_limit()) {
+    GTEST_SKIP() << "Test only applies to non-strict page byte size limit.";
+  }
+
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::CreateResult create_result,
       CreateDocumentStore(&filesystem_, test_dir_, &fake_clock_,
@@ -992,11 +1105,15 @@
   std::unique_ptr<DocumentStore> doc_store =
       std::move(create_result.document_store);
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             doc_store->Put(CreateDocument(/*id=*/1)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/1))));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             doc_store->Put(CreateDocument(/*id=*/2)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/2))));
   DocumentId document_id2 = put_result2.new_document_id;
 
   std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "name"),
@@ -1008,7 +1125,8 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ResultRetrieverV2> result_retriever,
       ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(),
-                                language_segmenter_.get(), normalizer_.get()));
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
 
   SearchResultProto::ResultProto result1;
   *result1.mutable_document() = CreateDocument(/*id=*/1);
@@ -1036,13 +1154,225 @@
   // the retrieval process and thus include result2 into this page, even though
   // finally total bytes of result1 + result2 exceed the threshold.
   auto [page_result, has_more_results] = result_retriever->RetrieveNextPage(
-      result_state, fake_clock_.GetSystemTimeMilliseconds());
+      result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+      fake_clock_.GetSystemTimeMilliseconds());
   EXPECT_THAT(page_result.results,
               ElementsAre(EqualsProto(result1), EqualsProto(result2)));
   // No more results.
   EXPECT_FALSE(has_more_results);
 }
 
+TEST_P(ResultRetrieverV2Test,
+       ShouldNotIncludeNextResultIfExceedingNumTotalBytesPerPageThreshold) {
+  if (!feature_flags_->enable_strict_page_byte_size_limit()) {
+    GTEST_SKIP() << "Test only applies to strict page byte size limit.";
+  }
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, test_dir_, &fake_clock_,
+                          schema_store_.get(), *feature_flags_));
+  std::unique_ptr<DocumentStore> doc_store =
+      std::move(create_result.document_store);
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/1))));
+  DocumentId document_id1 = put_result1.new_document_id;
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/2))));
+  DocumentId document_id2 = put_result2.new_document_id;
+
+  std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "name"),
+                                            GetSectionId("Email", "body")};
+  SectionIdMask hit_section_id_mask = CreateSectionIdMask(hit_section_ids);
+  std::vector<ScoredDocumentHit> scored_document_hits = {
+      {document_id1, hit_section_id_mask, /*score=*/5},
+      {document_id2, hit_section_id_mask, /*score=*/0}};
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<ResultRetrieverV2> result_retriever,
+      ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(),
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
+
+  SearchResultProto::ResultProto result1;
+  *result1.mutable_document() = CreateDocument(/*id=*/1);
+  result1.set_score(5);
+  SearchResultProto::ResultProto result2;
+  *result2.mutable_document() = CreateDocument(/*id=*/2);
+  result2.set_score(0);
+
+  // result1.ByteSizeLong() < threshold < result1.ByteSizeLong() +
+  //                                      result2.ByteSizeLong().
+  int threshold = result1.ByteSizeLong() + 1;
+  ASSERT_THAT(result1.ByteSizeLong() + result2.ByteSizeLong(), Gt(threshold));
+
+  ResultSpecProto result_spec =
+      CreateResultSpec(/*num_per_page=*/2, ResultSpecProto::NAMESPACE);
+  result_spec.set_num_total_bytes_per_page_threshold(threshold);
+  ResultStateV2 result_state(
+      std::make_unique<
+          PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
+          std::move(scored_document_hits),
+          /*is_descending=*/true),
+      /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
+      result_spec, *doc_store);
+
+  // After retrieving result1, total bytes are still below the threshold and #
+  // of results is still below num_per_page, so ResultRetriever should continue
+  // the retrieval process. But when serializing result2, the byte size exceeds
+  // the threshold, so result2 should not be included.
+  auto [page_result1, has_more_results1] = result_retriever->RetrieveNextPage(
+      result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+      fake_clock_.GetSystemTimeMilliseconds());
+  EXPECT_THAT(page_result1.results, ElementsAre(EqualsProto(result1)));
+  // More results.
+  EXPECT_TRUE(has_more_results1);
+
+  // Second page. Even though result2 was evaluated in the previous round but
+  // excluded, it should not be popped from the ranker and thus should be
+  // included in the second page.
+  auto [page_result2, has_more_results2] = result_retriever->RetrieveNextPage(
+      result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+      fake_clock_.GetSystemTimeMilliseconds());
+  EXPECT_THAT(page_result2.results, ElementsAre(EqualsProto(result2)));
+  // No more results.
+  EXPECT_FALSE(has_more_results2);
+}
+
+TEST_P(ResultRetrieverV2Test,
+       ResultGroupingShouldDecrementOnlyWhenResultIsIncludedInThePage) {
+  if (!feature_flags_->enable_strict_page_byte_size_limit()) {
+    GTEST_SKIP() << "Test only applies to non-strict page byte size limit.";
+  }
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, test_dir_, &fake_clock_,
+                          schema_store_.get(), *feature_flags_));
+  std::unique_ptr<DocumentStore> doc_store =
+      std::move(create_result.document_store);
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/1))));
+  DocumentId document_id1 = put_result1.new_document_id;
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      doc_store->Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/2))));
+  DocumentId document_id2 = put_result2.new_document_id;
+
+  std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "name"),
+                                            GetSectionId("Email", "body")};
+  SectionIdMask hit_section_id_mask = CreateSectionIdMask(hit_section_ids);
+  std::vector<ScoredDocumentHit> scored_document_hits = {
+      {document_id1, hit_section_id_mask, /*score=*/5},
+      {document_id2, hit_section_id_mask, /*score=*/0}};
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<ResultRetrieverV2> result_retriever,
+      ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(),
+                                language_segmenter_.get(), normalizer_.get(),
+                                feature_flags_.get()));
+
+  SearchResultProto::ResultProto result1;
+  *result1.mutable_document() = CreateDocument(/*id=*/1);
+  result1.set_score(5);
+  SearchResultProto::ResultProto result2;
+  *result2.mutable_document() = CreateDocument(/*id=*/2);
+  result2.set_score(0);
+
+  // Create a ResultSpec that limits namespace "icing" to 10 results and the
+  // total bytes per page to result1.ByteSizeLong() + 1. This will force the
+  // result retriever to move result2 to the next page.
+  ResultSpecProto result_spec =
+      CreateResultSpec(/*num_per_page=*/10, ResultSpecProto::NAMESPACE);
+  result_spec.set_num_total_bytes_per_page_threshold(result1.ByteSizeLong() +
+                                                     1);
+
+  ResultSpecProto::ResultGrouping* result_grouping =
+      result_spec.add_result_groupings();
+  ResultSpecProto::ResultGrouping::Entry* entry =
+      result_grouping->add_entry_groupings();
+  result_grouping->set_max_results(2);
+  entry->set_namespace_("icing");
+
+  // Creates a ResultState with 2 ScoredDocumentHits.
+  ResultStateV2 result_state(
+      std::make_unique<
+          PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
+          std::move(scored_document_hits), /*is_descending=*/true),
+      /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
+      result_spec, *doc_store);
+
+  // First page: result1 should be returned.
+  auto [page_result1, has_more_results1] = result_retriever->RetrieveNextPage(
+      result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+      fake_clock_.GetSystemTimeMilliseconds());
+  EXPECT_THAT(page_result1.results, ElementsAre(EqualsProto(result1)));
+  // Has more results.
+  EXPECT_TRUE(has_more_results1);
+
+  // Second page: result2 should be returned.
+  auto [page_result2, has_more_results2] = result_retriever->RetrieveNextPage(
+      result_state, /*max_results=*/std::numeric_limits<int32_t>::max(),
+      fake_clock_.GetSystemTimeMilliseconds());
+  EXPECT_THAT(page_result2.results, ElementsAre(EqualsProto(result2)));
+  // No more results.
+  EXPECT_FALSE(has_more_results2);
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    ResultRetrieverV2Test, ResultRetrieverV2Test,
+    testing::Values(FeatureFlags(
+                        /*allow_circular_schema_definitions=*/true,
+                        /*enable_scorable_properties=*/true,
+                        /*enable_embedding_quantization=*/true,
+                        /*enable_repeated_field_joins=*/true,
+                        /*enable_embedding_backup_generation=*/true,
+                        /*enable_schema_database=*/true,
+                        /*release_backup_schema_file_if_overlay_present=*/true,
+                        /*enable_strict_page_byte_size_limit=*/false,
+                        /*enable_smaller_decompression_buffer_size=*/true,
+                        /*enable_eigen_embedding_scoring=*/true,
+                        /*enable_passing_filter_to_children=*/true,
+                        /*enable_proto_log_new_header_format=*/true,
+                        /*enable_embedding_iterator_v2=*/true,
+                        /*enable_reusable_decompression_buffer=*/true,
+                        /*enable_schema_type_id_optimization=*/true,
+                        /*enable_optimize_improvements=*/true,
+                        /*expired_document_purge_threshold_ms=*/0,
+                        /*enable_non_existent_qualified_id_join=*/true,
+                        /*enable_skip_set_schema_type_equality_check=*/true,
+                        /*enable_embed_query_optimization=*/true,
+                        /*enable_schema_definition_deduping=*/true),
+                    FeatureFlags(
+                        /*allow_circular_schema_definitions=*/true,
+                        /*enable_scorable_properties=*/true,
+                        /*enable_embedding_quantization=*/true,
+                        /*enable_repeated_field_joins=*/true,
+                        /*enable_embedding_backup_generation=*/true,
+                        /*enable_schema_database=*/true,
+                        /*release_backup_schema_file_if_overlay_present=*/true,
+                        /*enable_strict_page_byte_size_limit=*/true,
+                        /*enable_smaller_decompression_buffer_size=*/true,
+                        /*enable_eigen_embedding_scoring=*/true,
+                        /*enable_passing_filter_to_children=*/true,
+                        /*enable_proto_log_new_header_format=*/true,
+                        /*enable_embedding_iterator_v2=*/true,
+                        /*enable_reusable_decompression_buffer=*/true,
+                        /*enable_schema_type_id_optimization=*/true,
+                        /*enable_optimize_improvements=*/true,
+                        /*expired_document_purge_threshold_ms=*/0,
+                        /*enable_non_existent_qualified_id_join=*/true,
+                        /*enable_skip_set_schema_type_equality_check=*/true,
+                        /*enable_embed_query_optimization=*/true,
+                        /*enable_schema_definition_deduping=*/true)));
+
 }  // namespace
 
 }  // namespace lib

diff --git a/icing/result/result-state-manager.cc b/icing/result/result-state-manager.cc
index 382f7db..ebd0512 100644
--- a/icing/result/result-state-manager.cc
+++ b/icing/result/result-state-manager.cc

@@ -14,18 +14,24 @@
 
 #include "icing/result/result-state-manager.h"
 
+#include <cstdint>
+#include <limits>
 #include <memory>
 #include <queue>
 #include <utility>
 
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
+#include "icing/absl_ports/canonical_errors.h"
+#include "icing/absl_ports/mutex.h"
+#include "icing/proto/logging.pb.h"
 #include "icing/result/page-result.h"
 #include "icing/result/result-adjustment-info.h"
 #include "icing/result/result-retriever-v2.h"
 #include "icing/result/result-state-v2.h"
 #include "icing/scoring/scored-document-hits-ranker.h"
+#include "icing/store/document-store.h"
 #include "icing/util/clock.h"
 #include "icing/util/logging.h"
-#include "icing/util/status-macros.h"
 
 namespace icing {
 namespace lib {
@@ -43,7 +49,8 @@
     std::unique_ptr<ResultAdjustmentInfo> parent_adjustment_info,
     std::unique_ptr<ResultAdjustmentInfo> child_adjustment_info,
     const ResultSpecProto& result_spec, const DocumentStore& document_store,
-    const ResultRetrieverV2& result_retriever, int64_t current_time_ms) {
+    const ResultRetrieverV2& result_retriever, int64_t current_time_ms,
+    QueryStatsProto* query_stats) {
   if (ranker == nullptr) {
     return absl_ports::InvalidArgumentError("Should not provide null ranker");
   }
@@ -56,8 +63,9 @@
 
   // Retrieve docs outside of ResultStateManager critical section.
   // Will enter ResultState critical section inside ResultRetriever.
-  auto [page_result, has_more_results] =
-      result_retriever.RetrieveNextPage(*result_state, current_time_ms);
+  auto [page_result, has_more_results] = result_retriever.RetrieveNextPage(
+      *result_state,
+      /*max_results=*/std::numeric_limits<int32_t>::max(), current_time_ms);
   if (!has_more_results) {
     // No more pages, won't store ResultState, returns directly
     return std::make_pair(kInvalidNextPageToken, std::move(page_result));
@@ -87,7 +95,7 @@
     InternalInvalidateExpiredResultStates(kDefaultResultStateTtlInMs,
                                           current_time_ms);
     // Remove states to make room for this new state.
-    RemoveStatesIfNeeded(num_hits_to_add);
+    RemoveStatesIfNeeded(num_hits_to_add, query_stats);
     // Generate a new unique token and add it into result_state_map_.
     next_page_token = Add(std::move(result_state), current_time_ms);
   }
@@ -107,7 +115,7 @@
 }
 
 libtextclassifier3::StatusOr<std::pair<uint64_t, PageResult>>
-ResultStateManager::GetNextPage(uint64_t next_page_token,
+ResultStateManager::GetNextPage(uint64_t next_page_token, int32_t max_results,
                                 const ResultRetrieverV2& result_retriever,
                                 int64_t current_time_ms) {
   std::shared_ptr<ResultStateV2> result_state = nullptr;
@@ -128,8 +136,8 @@
 
   // Retrieve docs outside of ResultStateManager critical section.
   // Will enter ResultState critical section inside ResultRetriever.
-  auto [page_result, has_more_results] =
-      result_retriever.RetrieveNextPage(*result_state, current_time_ms);
+  auto [page_result, has_more_results] = result_retriever.RetrieveNextPage(
+      *result_state, max_results, current_time_ms);
 
   if (!has_more_results) {
     {
@@ -144,6 +152,14 @@
   return std::make_pair(next_page_token, std::move(page_result));
 }
 
+int ResultStateManager::GetNumActiveResultStates(int64_t current_time_ms) {
+  absl_ports::unique_lock l(&mutex_);
+
+  InternalInvalidateExpiredResultStates(kDefaultResultStateTtlInMs,
+                                        current_time_ms);
+  return result_state_map_.size();
+}
+
 void ResultStateManager::InvalidateResultState(uint64_t next_page_token) {
   if (next_page_token == kInvalidNextPageToken) {
     return;
@@ -181,16 +197,17 @@
   return new_token;
 }
 
-void ResultStateManager::RemoveStatesIfNeeded(int num_hits_to_add) {
+void ResultStateManager::RemoveStatesIfNeeded(int num_hits_to_add,
+                                              QueryStatsProto* query_stats) {
   if (result_state_map_.empty() || token_queue_.empty()) {
     return;
   }
 
   // 1. Check if this new result_state would take up the entire result state
-  // manager budget.
+  //    manager budget.
   if (num_hits_to_add > max_total_hits_) {
     // This single result state will exceed our budget. Drop everything else to
-    // accomodate it.
+    // accommodate it.
     InternalInvalidateAllResultStates();
     return;
   }
@@ -204,18 +221,33 @@
   }
 
   // 3. If we're over budget, remove states from oldest to newest until we fit
-  // into our budget.
+  //    into our budget.
+  //
   // Note: num_total_hits_ may not be decremented immediately after invalidating
   // a result state, since other threads may still hold the shared pointer.
   // Thus, we have to check if token_queue_ is empty or not, since it is
   // possible that num_total_hits_ is non-zero and still greater than
   // max_total_hits_ when token_queue_ is empty. Still "eventually" it will be
   // decremented after the last thread releases the shared pointer.
+  int num_states_evicted = 0;
   while (!token_queue_.empty() && num_total_hits_ > max_total_hits_) {
+    ICING_LOG(WARNING) << "Evicting result state from token_queue_ due to "
+                          "budget limit. Current num_total_hits_: "
+                       << num_total_hits_;
+
     InternalInvalidateResultState(token_queue_.front().first);
+    ++num_states_evicted;
     token_queue_.pop();
   }
   invalidated_token_set_.clear();
+  if (num_states_evicted > 0) {
+    ICING_LOG(WARNING) << "Evicted " << num_states_evicted
+                       << " states. After eviction: " << num_total_hits_
+                       << " hits and " << token_queue_.size() << " states.";
+    if (query_stats != nullptr) {
+      query_stats->set_num_result_states_evicted(num_states_evicted);
+    }
+  }
 }
 
 void ResultStateManager::InternalInvalidateResultState(uint64_t token) {

diff --git a/icing/result/result-state-manager.h b/icing/result/result-state-manager.h
index a64ae2c..b370c59 100644
--- a/icing/result/result-state-manager.h
+++ b/icing/result/result-state-manager.h

@@ -16,21 +16,25 @@
 #define ICING_RESULT_RESULT_STATE_MANAGER_H_
 
 #include <atomic>
+#include <cstdint>
 #include <memory>
 #include <queue>
 #include <random>
 #include <unordered_map>
 #include <unordered_set>
+#include <utility>
 
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "icing/absl_ports/mutex.h"
+#include "icing/absl_ports/thread_annotations.h"
+#include "icing/proto/logging.pb.h"
 #include "icing/proto/search.pb.h"
 #include "icing/result/page-result.h"
 #include "icing/result/result-adjustment-info.h"
 #include "icing/result/result-retriever-v2.h"
 #include "icing/result/result-state-v2.h"
 #include "icing/scoring/scored-document-hits-ranker.h"
-#include "icing/util/clock.h"
+#include "icing/store/document-store.h"
 
 namespace icing {
 namespace lib {
@@ -77,10 +81,11 @@
       std::unique_ptr<ResultAdjustmentInfo> parent_adjustment_info,
       std::unique_ptr<ResultAdjustmentInfo> child_adjustment_info,
       const ResultSpecProto& result_spec, const DocumentStore& document_store,
-      const ResultRetrieverV2& result_retriever, int64_t current_time_ms)
-      ICING_LOCKS_EXCLUDED(mutex_);
+      const ResultRetrieverV2& result_retriever, int64_t current_time_ms,
+      QueryStatsProto* query_stats = nullptr) ICING_LOCKS_EXCLUDED(mutex_);
 
-  // Retrieves and returns PageResult for the next page.
+  // Retrieves and returns PageResult for the next page, retrieving at most
+  // max_results entries from the page.
   // The returned results won't exist in ResultStateManager anymore. If the
   // query has no more pages after this retrieval, the input token will be
   // invalidated.
@@ -93,8 +98,15 @@
   //   A token and PageResult wrapped by std::pair on success
   //   NOT_FOUND if failed to find any more results
   libtextclassifier3::StatusOr<std::pair<uint64_t, PageResult>> GetNextPage(
-      uint64_t next_page_token, const ResultRetrieverV2& result_retriever,
-      int64_t current_time_ms) ICING_LOCKS_EXCLUDED(mutex_);
+      uint64_t next_page_token, int32_t max_results,
+      const ResultRetrieverV2& result_retriever, int64_t current_time_ms)
+      ICING_LOCKS_EXCLUDED(mutex_);
+
+  // Returns the number of active result states currently in ResultStateManager.
+  // Note that this will invalidate expired result states before counting the
+  // number.
+  int GetNumActiveResultStates(int64_t current_time_ms)
+      ICING_LOCKS_EXCLUDED(mutex_);
 
   // Invalidates the result state associated with the given next-page token.
   void InvalidateResultState(uint64_t next_page_token)
@@ -150,7 +162,7 @@
 
   // Helper method to remove old states to make room for incoming states with
   // size num_hits_to_add.
-  void RemoveStatesIfNeeded(int num_hits_to_add)
+  void RemoveStatesIfNeeded(int num_hits_to_add, QueryStatsProto* query_stats)
       ICING_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
 
   // Helper method to remove a result state from result_state_map_, the token

diff --git a/icing/result/result-state-manager_test.cc b/icing/result/result-state-manager_test.cc
index 463005e..45cc0f8 100644
--- a/icing/result/result-state-manager_test.cc
+++ b/icing/result/result-state-manager_test.cc

@@ -32,6 +32,7 @@
 #include "icing/file/portable-file-backed-proto-log.h"
 #include "icing/index/embed/embedding-query-results.h"
 #include "icing/portable/equals-proto.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/portable/platform.h"
 #include "icing/query/query-terms.h"
 #include "icing/result/page-result.h"
@@ -54,6 +55,7 @@
 #include "icing/transform/normalizer-factory.h"
 #include "icing/transform/normalizer-options.h"
 #include "icing/transform/normalizer.h"
+#include "icing/util/document-util.h"
 #include "icing/util/icu-data-file-helper.h"
 #include "unicode/uloc.h"
 
@@ -129,20 +131,25 @@
 
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult result,
-        DocumentStore::Create(&filesystem_, test_dir_, clock_.get(),
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, test_dir_, clock_.get(), schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     document_store_ = std::move(result.document_store);
 
     ICING_ASSERT_OK_AND_ASSIGN(
-        result_retriever_, ResultRetrieverV2::Create(
-                               document_store_.get(), schema_store_.get(),
-                               language_segmenter_.get(), normalizer_.get()));
+        result_retriever_,
+        ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
+                                  language_segmenter_.get(), normalizer_.get(),
+                                  feature_flags_.get()));
   }
 
   void TearDown() override {
@@ -157,10 +164,15 @@
     document.set_uri(std::to_string(document_id));
     document.set_schema("Document");
     document.set_creation_timestamp_ms(1574365086666 + document_id);
-    document_store_->Put(document);
+
+    DocumentWrapper document_wrapper;
+    *document_wrapper.mutable_document() = std::move(document);
+
+    document_store_->Put(document_wrapper);
+
     return std::make_pair(
         ScoredDocumentHit(document_id, kSectionIdMaskNone, /*score=*/1),
-        std::move(document));
+        std::move(*document_wrapper.mutable_document()));
   }
 
   std::pair<std::vector<ScoredDocumentHit>, std::vector<DocumentProto>>
@@ -207,14 +219,20 @@
 };
 
 TEST_F(ResultStateManagerTest, ShouldCacheAndRetrieveFirstPageOnePage) {
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store().Put(CreateDocument(/*id=*/1)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store().Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/1))));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store().Put(CreateDocument(/*id=*/2)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store().Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/2))));
   DocumentId document_id2 = put_result2.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             document_store().Put(CreateDocument(/*id=*/3)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      document_store().Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/3))));
   DocumentId document_id3 = put_result3.new_document_id;
   std::vector<ScoredDocumentHit> scored_document_hits = {
       {document_id1, kSectionIdMaskNone, /*score=*/1},
@@ -249,20 +267,30 @@
 }
 
 TEST_F(ResultStateManagerTest, ShouldCacheAndRetrieveFirstPageMultiplePages) {
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store().Put(CreateDocument(/*id=*/1)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store().Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/1))));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store().Put(CreateDocument(/*id=*/2)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store().Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/2))));
   DocumentId document_id2 = put_result2.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             document_store().Put(CreateDocument(/*id=*/3)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      document_store().Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/3))));
   DocumentId document_id3 = put_result3.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result4,
-                             document_store().Put(CreateDocument(/*id=*/4)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result4,
+      document_store().Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/4))));
   DocumentId document_id4 = put_result4.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result5,
-                             document_store().Put(CreateDocument(/*id=*/5)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result5,
+      document_store().Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/5))));
   DocumentId document_id5 = put_result5.new_document_id;
   std::vector<ScoredDocumentHit> scored_document_hits = {
       {document_id1, kSectionIdMaskNone, /*score=*/1},
@@ -298,8 +326,9 @@
   // Second page, 2 results
   ICING_ASSERT_OK_AND_ASSIGN(
       PageResultInfo page_result_info2,
-      result_state_manager.GetNextPage(next_page_token, result_retriever(),
-                                       clock()->GetSystemTimeMilliseconds()));
+      result_state_manager.GetNextPage(
+          next_page_token, /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   EXPECT_THAT(page_result_info2.first, Eq(next_page_token));
   ASSERT_THAT(page_result_info2.second.results, SizeIs(2));
   EXPECT_THAT(page_result_info2.second.results.at(0).document(),
@@ -310,8 +339,9 @@
   // Third page, 1 result
   ICING_ASSERT_OK_AND_ASSIGN(
       PageResultInfo page_result_info3,
-      result_state_manager.GetNextPage(next_page_token, result_retriever(),
-                                       clock()->GetSystemTimeMilliseconds()));
+      result_state_manager.GetNextPage(
+          next_page_token, /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   EXPECT_THAT(page_result_info3.first, Eq(kInvalidNextPageToken));
   ASSERT_THAT(page_result_info3.second.results, SizeIs(1));
   EXPECT_THAT(page_result_info3.second.results.at(0).document(),
@@ -319,8 +349,9 @@
 
   // No results
   EXPECT_THAT(
-      result_state_manager.GetNextPage(next_page_token, result_retriever(),
-                                       clock()->GetSystemTimeMilliseconds()),
+      result_state_manager.GetNextPage(
+          next_page_token, /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()),
       StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
 }
 
@@ -357,11 +388,15 @@
 }
 
 TEST_F(ResultStateManagerTest, ShouldAllowEmptyFirstPage) {
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store().Put(CreateDocument(/*id=*/1)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store().Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/1))));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store().Put(CreateDocument(/*id=*/2)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store().Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/2))));
   DocumentId document_id2 = put_result2.new_document_id;
   std::vector<ScoredDocumentHit> scored_document_hits = {
       {document_id1, kSectionIdMaskNone, /*score=*/1},
@@ -396,17 +431,25 @@
 }
 
 TEST_F(ResultStateManagerTest, ShouldAllowEmptyLastPage) {
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store().Put(CreateDocument(/*id=*/1)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store().Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/1))));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store().Put(CreateDocument(/*id=*/2)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store().Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/2))));
   DocumentId document_id2 = put_result2.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             document_store().Put(CreateDocument(/*id=*/3)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      document_store().Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/3))));
   DocumentId document_id3 = put_result3.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result4,
-                             document_store().Put(CreateDocument(/*id=*/4)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result4,
+      document_store().Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/4))));
   DocumentId document_id4 = put_result4.new_document_id;
   std::vector<ScoredDocumentHit> scored_document_hits = {
       {document_id1, kSectionIdMaskNone, /*score=*/1},
@@ -450,8 +493,9 @@
   // limiter, so we should get an empty page.
   ICING_ASSERT_OK_AND_ASSIGN(
       PageResultInfo page_result_info2,
-      result_state_manager.GetNextPage(next_page_token, result_retriever(),
-                                       clock()->GetSystemTimeMilliseconds()));
+      result_state_manager.GetNextPage(
+          next_page_token, /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   EXPECT_THAT(page_result_info2.first, Eq(kInvalidNextPageToken));
   EXPECT_THAT(page_result_info2.second.results, IsEmpty());
 }
@@ -518,8 +562,9 @@
   clock()->SetSystemTimeMilliseconds(1000);
   // page_result_info1's token (page_result_info1.first) shouldn't be found.
   EXPECT_THAT(result_state_manager.GetNextPage(
-                  page_result_info1.first, result_retriever(),
-                  clock()->GetSystemTimeMilliseconds()),
+                  page_result_info1.first,
+                  /*max_results=*/std::numeric_limits<int32_t>::max(),
+                  result_retriever(), clock()->GetSystemTimeMilliseconds()),
               StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
 }
 
@@ -567,18 +612,21 @@
   // 3. Then calling GetNextPage() on state 1 shouldn't get anything.
   clock()->SetSystemTimeMilliseconds(kDefaultResultStateTtlInMs + 1000);
   // page_result_info2's token (page_result_info2.first) should be found
-  ICING_ASSERT_OK_AND_ASSIGN(page_result_info2,
-                             result_state_manager.GetNextPage(
-                                 page_result_info2.first, result_retriever(),
-                                 clock()->GetSystemTimeMilliseconds()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      page_result_info2,
+      result_state_manager.GetNextPage(
+          page_result_info2.first,
+          /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   // We test the behavior by setting time back to 2s, to make sure the
   // invalidation of state 1 was done by the previous GetNextPage() instead of
   // the following GetNextPage().
   clock()->SetSystemTimeMilliseconds(2000);
   // page_result_info1's token (page_result_info1.first) shouldn't be found.
   EXPECT_THAT(result_state_manager.GetNextPage(
-                  page_result_info1.first, result_retriever(),
-                  clock()->GetSystemTimeMilliseconds()),
+                  page_result_info1.first,
+                  /*max_results=*/std::numeric_limits<int32_t>::max(),
+                  result_retriever(), clock()->GetSystemTimeMilliseconds()),
               StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
 }
 
@@ -611,29 +659,42 @@
   clock()->SetSystemTimeMilliseconds(kDefaultResultStateTtlInMs + 1000);
   // page_result_info's token (page_result_info.first) shouldn't be found.
   EXPECT_THAT(result_state_manager.GetNextPage(
-                  page_result_info.first, result_retriever(),
-                  clock()->GetSystemTimeMilliseconds()),
+                  page_result_info.first,
+                  /*max_results=*/std::numeric_limits<int32_t>::max(),
+                  result_retriever(), clock()->GetSystemTimeMilliseconds()),
               StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
 }
 
 TEST_F(ResultStateManagerTest, ShouldInvalidateOneToken) {
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store().Put(CreateDocument(/*id=*/1)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store().Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/1))));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store().Put(CreateDocument(/*id=*/2)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store().Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/2))));
   DocumentId document_id2 = put_result2.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             document_store().Put(CreateDocument(/*id=*/3)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      document_store().Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/3))));
   DocumentId document_id3 = put_result3.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result4,
-                             document_store().Put(CreateDocument(/*id=*/4)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result4,
+      document_store().Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/4))));
   DocumentId document_id4 = put_result4.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result5,
-                             document_store().Put(CreateDocument(/*id=*/5)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result5,
+      document_store().Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/5))));
   DocumentId document_id5 = put_result5.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result6,
-                             document_store().Put(CreateDocument(/*id=*/6)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result6,
+      document_store().Put(
+          document_util::CreateDocumentWrapper(CreateDocument(/*id=*/6))));
   DocumentId document_id6 = put_result6.new_document_id;
   std::vector<ScoredDocumentHit> scored_document_hits1 = {
       {document_id1, kSectionIdMaskNone, /*score=*/1},
@@ -674,15 +735,18 @@
 
   // page_result_info1's token (page_result_info1.first) shouldn't be found
   EXPECT_THAT(result_state_manager.GetNextPage(
-                  page_result_info1.first, result_retriever(),
-                  clock()->GetSystemTimeMilliseconds()),
+                  page_result_info1.first,
+                  /*max_results=*/std::numeric_limits<int32_t>::max(),
+                  result_retriever(), clock()->GetSystemTimeMilliseconds()),
               StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
 
   // page_result_info2's token (page_result_info2.first) should still exist
-  ICING_ASSERT_OK_AND_ASSIGN(page_result_info2,
-                             result_state_manager.GetNextPage(
-                                 page_result_info2.first, result_retriever(),
-                                 clock()->GetSystemTimeMilliseconds()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      page_result_info2,
+      result_state_manager.GetNextPage(
+          page_result_info2.first,
+          /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   // Should get docs.
   ASSERT_THAT(page_result_info2.second.results, SizeIs(1));
   EXPECT_THAT(page_result_info2.second.results.at(0).document(),
@@ -724,14 +788,16 @@
 
   // page_result_info1's token (page_result_info1.first) shouldn't be found
   EXPECT_THAT(result_state_manager.GetNextPage(
-                  page_result_info1.first, result_retriever(),
-                  clock()->GetSystemTimeMilliseconds()),
+                  page_result_info1.first,
+                  /*max_results=*/std::numeric_limits<int32_t>::max(),
+                  result_retriever(), clock()->GetSystemTimeMilliseconds()),
               StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
 
   // page_result_info2's token (page_result_info2.first) shouldn't be found
   EXPECT_THAT(result_state_manager.GetNextPage(
-                  page_result_info2.first, result_retriever(),
-                  clock()->GetSystemTimeMilliseconds()),
+                  page_result_info2.first,
+                  /*max_results=*/std::numeric_limits<int32_t>::max(),
+                  result_retriever(), clock()->GetSystemTimeMilliseconds()),
               StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
 }
 
@@ -781,22 +847,27 @@
           clock()->GetSystemTimeMilliseconds()));
 
   EXPECT_THAT(result_state_manager.GetNextPage(
-                  page_result_info1.first, result_retriever(),
-                  clock()->GetSystemTimeMilliseconds()),
+                  page_result_info1.first,
+                  /*max_results=*/std::numeric_limits<int32_t>::max(),
+                  result_retriever(), clock()->GetSystemTimeMilliseconds()),
               StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
 
-  ICING_ASSERT_OK_AND_ASSIGN(page_result_info2,
-                             result_state_manager.GetNextPage(
-                                 page_result_info2.first, result_retriever(),
-                                 clock()->GetSystemTimeMilliseconds()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      page_result_info2,
+      result_state_manager.GetNextPage(
+          page_result_info2.first,
+          /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   ASSERT_THAT(page_result_info2.second.results, SizeIs(1));
   EXPECT_THAT(page_result_info2.second.results.at(0).document(),
               EqualsProto(document_protos2.at(1)));
 
-  ICING_ASSERT_OK_AND_ASSIGN(page_result_info3,
-                             result_state_manager.GetNextPage(
-                                 page_result_info3.first, result_retriever(),
-                                 clock()->GetSystemTimeMilliseconds()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      page_result_info3,
+      result_state_manager.GetNextPage(
+          page_result_info3.first,
+          /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   ASSERT_THAT(page_result_info3.second.results, SizeIs(1));
   EXPECT_THAT(page_result_info3.second.results.at(0).document(),
               EqualsProto(document_protos3.at(1)));
@@ -872,31 +943,38 @@
           document_store(), result_retriever(),
           clock()->GetSystemTimeMilliseconds()));
 
-  ICING_ASSERT_OK_AND_ASSIGN(page_result_info1,
-                             result_state_manager.GetNextPage(
-                                 page_result_info1.first, result_retriever(),
-                                 clock()->GetSystemTimeMilliseconds()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      page_result_info1,
+      result_state_manager.GetNextPage(
+          page_result_info1.first,
+          /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   ASSERT_THAT(page_result_info1.second.results, SizeIs(1));
   EXPECT_THAT(page_result_info1.second.results.at(0).document(),
               EqualsProto(document_protos1.at(1)));
 
   EXPECT_THAT(result_state_manager.GetNextPage(
-                  page_result_info2.first, result_retriever(),
-                  clock()->GetSystemTimeMilliseconds()),
+                  page_result_info2.first,
+                  /*max_results=*/std::numeric_limits<int32_t>::max(),
+                  result_retriever(), clock()->GetSystemTimeMilliseconds()),
               StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
 
-  ICING_ASSERT_OK_AND_ASSIGN(page_result_info3,
-                             result_state_manager.GetNextPage(
-                                 page_result_info3.first, result_retriever(),
-                                 clock()->GetSystemTimeMilliseconds()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      page_result_info3,
+      result_state_manager.GetNextPage(
+          page_result_info3.first,
+          /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   ASSERT_THAT(page_result_info3.second.results, SizeIs(1));
   EXPECT_THAT(page_result_info3.second.results.at(0).document(),
               EqualsProto(document_protos3.at(1)));
 
-  ICING_ASSERT_OK_AND_ASSIGN(page_result_info4,
-                             result_state_manager.GetNextPage(
-                                 page_result_info4.first, result_retriever(),
-                                 clock()->GetSystemTimeMilliseconds()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      page_result_info4,
+      result_state_manager.GetNextPage(
+          page_result_info4.first,
+          /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   ASSERT_THAT(page_result_info4.second.results, SizeIs(1));
   EXPECT_THAT(page_result_info4.second.results.at(0).document(),
               EqualsProto(document_protos4.at(1)));
@@ -999,40 +1077,49 @@
           clock()->GetSystemTimeMilliseconds()));
 
   EXPECT_THAT(result_state_manager.GetNextPage(
-                  page_result_info1.first, result_retriever(),
-                  clock()->GetSystemTimeMilliseconds()),
+                  page_result_info1.first,
+                  /*max_results=*/std::numeric_limits<int32_t>::max(),
+                  result_retriever(), clock()->GetSystemTimeMilliseconds()),
               StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
 
   EXPECT_THAT(result_state_manager.GetNextPage(
-                  page_result_info2.first, result_retriever(),
-                  clock()->GetSystemTimeMilliseconds()),
+                  page_result_info2.first,
+                  /*max_results=*/std::numeric_limits<int32_t>::max(),
+                  result_retriever(), clock()->GetSystemTimeMilliseconds()),
               StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
 
   EXPECT_THAT(result_state_manager.GetNextPage(
-                  page_result_info3.first, result_retriever(),
-                  clock()->GetSystemTimeMilliseconds()),
+                  page_result_info3.first,
+                  /*max_results=*/std::numeric_limits<int32_t>::max(),
+                  result_retriever(), clock()->GetSystemTimeMilliseconds()),
               StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
 
-  ICING_ASSERT_OK_AND_ASSIGN(page_result_info4,
-                             result_state_manager.GetNextPage(
-                                 page_result_info4.first, result_retriever(),
-                                 clock()->GetSystemTimeMilliseconds()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      page_result_info4,
+      result_state_manager.GetNextPage(
+          page_result_info4.first,
+          /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   ASSERT_THAT(page_result_info4.second.results, SizeIs(1));
   EXPECT_THAT(page_result_info4.second.results.at(0).document(),
               EqualsProto(document_protos4.at(1)));
 
-  ICING_ASSERT_OK_AND_ASSIGN(page_result_info5,
-                             result_state_manager.GetNextPage(
-                                 page_result_info5.first, result_retriever(),
-                                 clock()->GetSystemTimeMilliseconds()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      page_result_info5,
+      result_state_manager.GetNextPage(
+          page_result_info5.first,
+          /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   ASSERT_THAT(page_result_info5.second.results, SizeIs(1));
   EXPECT_THAT(page_result_info5.second.results.at(0).document(),
               EqualsProto(document_protos5.at(1)));
 
-  ICING_ASSERT_OK_AND_ASSIGN(page_result_info6,
-                             result_state_manager.GetNextPage(
-                                 page_result_info6.first, result_retriever(),
-                                 clock()->GetSystemTimeMilliseconds()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      page_result_info6,
+      result_state_manager.GetNextPage(
+          page_result_info6.first,
+          /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   ASSERT_THAT(page_result_info6.second.results, SizeIs(1));
   EXPECT_THAT(page_result_info6.second.results.at(0).document(),
               EqualsProto(document_protos6.at(1)));
@@ -1126,35 +1213,43 @@
           clock()->GetSystemTimeMilliseconds()));
 
   EXPECT_THAT(result_state_manager.GetNextPage(
-                  page_result_info1.first, result_retriever(),
-                  clock()->GetSystemTimeMilliseconds()),
+                  page_result_info1.first,
+                  /*max_results=*/std::numeric_limits<int32_t>::max(),
+                  result_retriever(), clock()->GetSystemTimeMilliseconds()),
               StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
 
   EXPECT_THAT(result_state_manager.GetNextPage(
-                  page_result_info2.first, result_retriever(),
-                  clock()->GetSystemTimeMilliseconds()),
+                  page_result_info2.first,
+                  /*max_results=*/std::numeric_limits<int32_t>::max(),
+                  result_retriever(), clock()->GetSystemTimeMilliseconds()),
               StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
 
-  ICING_ASSERT_OK_AND_ASSIGN(page_result_info3,
-                             result_state_manager.GetNextPage(
-                                 page_result_info3.first, result_retriever(),
-                                 clock()->GetSystemTimeMilliseconds()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      page_result_info3,
+      result_state_manager.GetNextPage(
+          page_result_info3.first,
+          /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   ASSERT_THAT(page_result_info3.second.results, SizeIs(1));
   EXPECT_THAT(page_result_info3.second.results.at(0).document(),
               EqualsProto(document_protos3.at(1)));
 
-  ICING_ASSERT_OK_AND_ASSIGN(page_result_info4,
-                             result_state_manager.GetNextPage(
-                                 page_result_info4.first, result_retriever(),
-                                 clock()->GetSystemTimeMilliseconds()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      page_result_info4,
+      result_state_manager.GetNextPage(
+          page_result_info4.first,
+          /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   ASSERT_THAT(page_result_info4.second.results, SizeIs(1));
   EXPECT_THAT(page_result_info4.second.results.at(0).document(),
               EqualsProto(document_protos4.at(1)));
 
-  ICING_ASSERT_OK_AND_ASSIGN(page_result_info5,
-                             result_state_manager.GetNextPage(
-                                 page_result_info5.first, result_retriever(),
-                                 clock()->GetSystemTimeMilliseconds()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      page_result_info5,
+      result_state_manager.GetNextPage(
+          page_result_info5.first,
+          /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   ASSERT_THAT(page_result_info5.second.results, SizeIs(1));
   EXPECT_THAT(page_result_info5.second.results.at(0).document(),
               EqualsProto(document_protos5.at(1)));
@@ -1211,10 +1306,12 @@
 
   // GetNextPage for result state 1 should return its result and decrement the
   // number of cached hits to 2.
-  ICING_ASSERT_OK_AND_ASSIGN(page_result_info1,
-                             result_state_manager.GetNextPage(
-                                 page_result_info1.first, result_retriever(),
-                                 clock()->GetSystemTimeMilliseconds()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      page_result_info1,
+      result_state_manager.GetNextPage(
+          page_result_info1.first,
+          /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   ASSERT_THAT(page_result_info1.second.results, SizeIs(1));
   EXPECT_THAT(page_result_info1.second.results.at(0).document(),
               EqualsProto(document_protos1.at(1)));
@@ -1236,30 +1333,37 @@
           clock()->GetSystemTimeMilliseconds()));
 
   EXPECT_THAT(result_state_manager.GetNextPage(
-                  page_result_info1.first, result_retriever(),
-                  clock()->GetSystemTimeMilliseconds()),
+                  page_result_info1.first,
+                  /*max_results=*/std::numeric_limits<int32_t>::max(),
+                  result_retriever(), clock()->GetSystemTimeMilliseconds()),
               StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
 
-  ICING_ASSERT_OK_AND_ASSIGN(page_result_info2,
-                             result_state_manager.GetNextPage(
-                                 page_result_info2.first, result_retriever(),
-                                 clock()->GetSystemTimeMilliseconds()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      page_result_info2,
+      result_state_manager.GetNextPage(
+          page_result_info2.first,
+          /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   ASSERT_THAT(page_result_info2.second.results, SizeIs(1));
   EXPECT_THAT(page_result_info2.second.results.at(0).document(),
               EqualsProto(document_protos2.at(1)));
 
-  ICING_ASSERT_OK_AND_ASSIGN(page_result_info3,
-                             result_state_manager.GetNextPage(
-                                 page_result_info3.first, result_retriever(),
-                                 clock()->GetSystemTimeMilliseconds()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      page_result_info3,
+      result_state_manager.GetNextPage(
+          page_result_info3.first,
+          /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   ASSERT_THAT(page_result_info3.second.results, SizeIs(1));
   EXPECT_THAT(page_result_info3.second.results.at(0).document(),
               EqualsProto(document_protos3.at(1)));
 
-  ICING_ASSERT_OK_AND_ASSIGN(page_result_info4,
-                             result_state_manager.GetNextPage(
-                                 page_result_info4.first, result_retriever(),
-                                 clock()->GetSystemTimeMilliseconds()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      page_result_info4,
+      result_state_manager.GetNextPage(
+          page_result_info4.first,
+          /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   ASSERT_THAT(page_result_info4.second.results, SizeIs(1));
   EXPECT_THAT(page_result_info4.second.results.at(0).document(),
               EqualsProto(document_protos4.at(1)));
@@ -1317,10 +1421,12 @@
 
   // GetNextPage for result state 1 should return its result and decrement the
   // number of cached hits to 2.
-  ICING_ASSERT_OK_AND_ASSIGN(page_result_info1,
-                             result_state_manager.GetNextPage(
-                                 page_result_info1.first, result_retriever(),
-                                 clock()->GetSystemTimeMilliseconds()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      page_result_info1,
+      result_state_manager.GetNextPage(
+          page_result_info1.first,
+          /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   ASSERT_THAT(page_result_info1.second.results, SizeIs(1));
   EXPECT_THAT(page_result_info1.second.results.at(0).document(),
               EqualsProto(document_protos1.at(1)));
@@ -1358,35 +1464,43 @@
           clock()->GetSystemTimeMilliseconds()));
 
   EXPECT_THAT(result_state_manager.GetNextPage(
-                  page_result_info1.first, result_retriever(),
-                  clock()->GetSystemTimeMilliseconds()),
+                  page_result_info1.first,
+                  /*max_results=*/std::numeric_limits<int32_t>::max(),
+                  result_retriever(), clock()->GetSystemTimeMilliseconds()),
               StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
 
   EXPECT_THAT(result_state_manager.GetNextPage(
-                  page_result_info2.first, result_retriever(),
-                  clock()->GetSystemTimeMilliseconds()),
+                  page_result_info2.first,
+                  /*max_results=*/std::numeric_limits<int32_t>::max(),
+                  result_retriever(), clock()->GetSystemTimeMilliseconds()),
               StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
 
-  ICING_ASSERT_OK_AND_ASSIGN(page_result_info3,
-                             result_state_manager.GetNextPage(
-                                 page_result_info3.first, result_retriever(),
-                                 clock()->GetSystemTimeMilliseconds()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      page_result_info3,
+      result_state_manager.GetNextPage(
+          page_result_info3.first,
+          /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   ASSERT_THAT(page_result_info3.second.results, SizeIs(1));
   EXPECT_THAT(page_result_info3.second.results.at(0).document(),
               EqualsProto(document_protos3.at(1)));
 
-  ICING_ASSERT_OK_AND_ASSIGN(page_result_info4,
-                             result_state_manager.GetNextPage(
-                                 page_result_info4.first, result_retriever(),
-                                 clock()->GetSystemTimeMilliseconds()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      page_result_info4,
+      result_state_manager.GetNextPage(
+          page_result_info4.first,
+          /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   ASSERT_THAT(page_result_info4.second.results, SizeIs(1));
   EXPECT_THAT(page_result_info4.second.results.at(0).document(),
               EqualsProto(document_protos4.at(1)));
 
-  ICING_ASSERT_OK_AND_ASSIGN(page_result_info5,
-                             result_state_manager.GetNextPage(
-                                 page_result_info5.first, result_retriever(),
-                                 clock()->GetSystemTimeMilliseconds()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      page_result_info5,
+      result_state_manager.GetNextPage(
+          page_result_info5.first,
+          /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   ASSERT_THAT(page_result_info5.second.results, SizeIs(1));
   EXPECT_THAT(page_result_info5.second.results.at(0).document(),
               EqualsProto(document_protos5.at(1)));
@@ -1435,6 +1549,7 @@
   auto [scored_document_hits3, document_protos3] = AddScoredDocuments(
       {/*document_id=*/5, /*document_id=*/6, /*document_id=*/7,
        /*document_id=*/8, /*document_id=*/9, /*document_id=*/10});
+  QueryStatsProto query_stats;
   ICING_ASSERT_OK_AND_ASSIGN(
       PageResultInfo page_result_info3,
       result_state_manager.CacheAndRetrieveFirstPage(
@@ -1444,26 +1559,32 @@
           /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
           CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
           document_store(), result_retriever(),
-          clock()->GetSystemTimeMilliseconds()));
+          clock()->GetSystemTimeMilliseconds(), &query_stats));
   EXPECT_THAT(page_result_info3.first, Not(Eq(kInvalidNextPageToken)));
+  // Should set num_result_states_evicted since result state 1 and 2 were
+  // evicted.
+  EXPECT_THAT(query_stats.num_result_states_evicted(), Eq(2));
 
   // GetNextPage for result state 1 and 2 should return NOT_FOUND.
   EXPECT_THAT(result_state_manager.GetNextPage(
-                  page_result_info1.first, result_retriever(),
-                  clock()->GetSystemTimeMilliseconds()),
+                  page_result_info1.first,
+                  /*max_results=*/std::numeric_limits<int32_t>::max(),
+                  result_retriever(), clock()->GetSystemTimeMilliseconds()),
               StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
 
   EXPECT_THAT(result_state_manager.GetNextPage(
-                  page_result_info2.first, result_retriever(),
-                  clock()->GetSystemTimeMilliseconds()),
+                  page_result_info2.first,
+                  /*max_results=*/std::numeric_limits<int32_t>::max(),
+                  result_retriever(), clock()->GetSystemTimeMilliseconds()),
               StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
 
   // Only the next four results in state 3 should be retrievable.
   uint64_t next_page_token3 = page_result_info3.first;
   ICING_ASSERT_OK_AND_ASSIGN(
       page_result_info3,
-      result_state_manager.GetNextPage(next_page_token3, result_retriever(),
-                                       clock()->GetSystemTimeMilliseconds()));
+      result_state_manager.GetNextPage(
+          next_page_token3, /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   EXPECT_THAT(page_result_info3.first, Eq(next_page_token3));
   ASSERT_THAT(page_result_info3.second.results, SizeIs(1));
   EXPECT_THAT(page_result_info3.second.results.at(0).document(),
@@ -1471,8 +1592,9 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       page_result_info3,
-      result_state_manager.GetNextPage(next_page_token3, result_retriever(),
-                                       clock()->GetSystemTimeMilliseconds()));
+      result_state_manager.GetNextPage(
+          next_page_token3, /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   EXPECT_THAT(page_result_info3.first, Eq(next_page_token3));
   ASSERT_THAT(page_result_info3.second.results, SizeIs(1));
   EXPECT_THAT(page_result_info3.second.results.at(0).document(),
@@ -1480,8 +1602,9 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       page_result_info3,
-      result_state_manager.GetNextPage(next_page_token3, result_retriever(),
-                                       clock()->GetSystemTimeMilliseconds()));
+      result_state_manager.GetNextPage(
+          next_page_token3, /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   EXPECT_THAT(page_result_info3.first, Eq(next_page_token3));
   ASSERT_THAT(page_result_info3.second.results, SizeIs(1));
   EXPECT_THAT(page_result_info3.second.results.at(0).document(),
@@ -1489,8 +1612,9 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       page_result_info3,
-      result_state_manager.GetNextPage(next_page_token3, result_retriever(),
-                                       clock()->GetSystemTimeMilliseconds()));
+      result_state_manager.GetNextPage(
+          next_page_token3, /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   // The final document should have been dropped because it exceeded the budget,
   // so the next page token of the second last round should be
   // kInvalidNextPageToken.
@@ -1501,8 +1625,9 @@
 
   // Double check that next_page_token3 is not retrievable anymore.
   EXPECT_THAT(
-      result_state_manager.GetNextPage(next_page_token3, result_retriever(),
-                                       clock()->GetSystemTimeMilliseconds()),
+      result_state_manager.GetNextPage(
+          next_page_token3, /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()),
       StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
 }
 
@@ -1545,14 +1670,17 @@
 
   // state1 should have been evicted and state2 should still be retrievable.
   EXPECT_THAT(result_state_manager.GetNextPage(
-                  page_result_info1.first, result_retriever(),
-                  clock()->GetSystemTimeMilliseconds()),
+                  page_result_info1.first,
+                  /*max_results=*/std::numeric_limits<int32_t>::max(),
+                  result_retriever(), clock()->GetSystemTimeMilliseconds()),
               StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
 
-  ICING_ASSERT_OK_AND_ASSIGN(page_result_info2,
-                             result_state_manager.GetNextPage(
-                                 page_result_info2.first, result_retriever(),
-                                 clock()->GetSystemTimeMilliseconds()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      page_result_info2,
+      result_state_manager.GetNextPage(
+          page_result_info2.first,
+          /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   ASSERT_THAT(page_result_info2.second.results, SizeIs(1));
   EXPECT_THAT(page_result_info2.second.results.at(0).document(),
               EqualsProto(document_protos2.at(1)));
@@ -1596,8 +1724,9 @@
   // Second page, 2 results.
   ICING_ASSERT_OK_AND_ASSIGN(
       PageResultInfo page_result_info2,
-      result_state_manager.GetNextPage(next_page_token, result_retriever(),
-                                       clock()->GetSystemTimeMilliseconds()));
+      result_state_manager.GetNextPage(
+          next_page_token, /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   ASSERT_THAT(page_result_info2.second.results, SizeIs(2));
   EXPECT_THAT(page_result_info2.second.results.at(0).document(),
               EqualsProto(document_protos.at(2)));
@@ -1607,19 +1736,128 @@
   // Third page, 1 result.
   ICING_ASSERT_OK_AND_ASSIGN(
       PageResultInfo page_result_info3,
-      result_state_manager.GetNextPage(next_page_token, result_retriever(),
-                                       clock()->GetSystemTimeMilliseconds()));
+      result_state_manager.GetNextPage(
+          next_page_token, /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()));
   ASSERT_THAT(page_result_info3.second.results, SizeIs(1));
   EXPECT_THAT(page_result_info3.second.results.at(0).document(),
               EqualsProto(document_protos.at(4)));
 
   // Fourth page, 0 results.
   EXPECT_THAT(
-      result_state_manager.GetNextPage(next_page_token, result_retriever(),
-                                       clock()->GetSystemTimeMilliseconds()),
+      result_state_manager.GetNextPage(
+          next_page_token, /*max_results=*/std::numeric_limits<int32_t>::max(),
+          result_retriever(), clock()->GetSystemTimeMilliseconds()),
       StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
 }
 
+TEST_F(ResultStateManagerTest, GetNumActiveResultStates) {
+  auto [scored_document_hits1, document_protos1] = AddScoredDocuments(
+      {/*document_id=*/0, /*document_id=*/1, /*document_id=*/2});
+  auto [scored_document_hits2, document_protos2] = AddScoredDocuments(
+      {/*document_id=*/3, /*document_id=*/4, /*document_id=*/5});
+  auto [scored_document_hits3, document_protos3] = AddScoredDocuments(
+      {/*document_id=*/6, /*document_id=*/7, /*document_id=*/8});
+
+  ResultStateManager result_state_manager(
+      /*max_total_hits=*/std::numeric_limits<int>::max(), document_store());
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      PageResultInfo page_result_info1,
+      result_state_manager.CacheAndRetrieveFirstPage(
+          std::make_unique<
+              PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
+              std::move(scored_document_hits1), /*is_descending=*/true),
+          /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
+          CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
+          document_store(), result_retriever(),
+          clock()->GetSystemTimeMilliseconds()));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      PageResultInfo page_result_info2,
+      result_state_manager.CacheAndRetrieveFirstPage(
+          std::make_unique<
+              PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
+              std::move(scored_document_hits2), /*is_descending=*/true),
+          /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
+          CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
+          document_store(), result_retriever(),
+          clock()->GetSystemTimeMilliseconds()));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      PageResultInfo page_result_info3,
+      result_state_manager.CacheAndRetrieveFirstPage(
+          std::make_unique<
+              PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
+              std::move(scored_document_hits3), /*is_descending=*/true),
+          /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
+          CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
+          document_store(), result_retriever(),
+          clock()->GetSystemTimeMilliseconds()));
+
+  EXPECT_THAT(result_state_manager.GetNumActiveResultStates(
+                  /*current_time_ms=*/clock()->GetSystemTimeMilliseconds()),
+              Eq(3));
+}
+
+TEST_F(ResultStateManagerTest,
+       GetNumActiveResultStatesShouldInvalidateExpiredResultStates) {
+  auto [scored_document_hits1, document_protos1] = AddScoredDocuments(
+      {/*document_id=*/0, /*document_id=*/1, /*document_id=*/2});
+  auto [scored_document_hits2, document_protos2] = AddScoredDocuments(
+      {/*document_id=*/3, /*document_id=*/4, /*document_id=*/5});
+  auto [scored_document_hits3, document_protos3] = AddScoredDocuments(
+      {/*document_id=*/6, /*document_id=*/7, /*document_id=*/8});
+
+  ResultStateManager result_state_manager(
+      /*max_total_hits=*/std::numeric_limits<int>::max(), document_store());
+
+  // Set time as 1s and add state.
+  clock()->SetSystemTimeMilliseconds(1000);
+  ICING_ASSERT_OK_AND_ASSIGN(
+      PageResultInfo page_result_info1,
+      result_state_manager.CacheAndRetrieveFirstPage(
+          std::make_unique<
+              PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
+              std::move(scored_document_hits1), /*is_descending=*/true),
+          /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
+          CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
+          document_store(), result_retriever(),
+          clock()->GetSystemTimeMilliseconds()));
+
+  // Set time as 10s and add state 2, state 3.
+  clock()->SetSystemTimeMilliseconds(10000);
+  ICING_ASSERT_OK_AND_ASSIGN(
+      PageResultInfo page_result_info2,
+      result_state_manager.CacheAndRetrieveFirstPage(
+          std::make_unique<
+              PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
+              std::move(scored_document_hits2), /*is_descending=*/true),
+          /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
+          CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
+          document_store(), result_retriever(),
+          clock()->GetSystemTimeMilliseconds()));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      PageResultInfo page_result_info3,
+      result_state_manager.CacheAndRetrieveFirstPage(
+          std::make_unique<
+              PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
+              std::move(scored_document_hits3), /*is_descending=*/true),
+          /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
+          CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
+          document_store(), result_retriever(),
+          clock()->GetSystemTimeMilliseconds()));
+
+  // 1. Set time as 1hr1s.
+  // 2. Then calling GetNumActiveResultStates() should invalidate state 1 and
+  //    return 2.
+  clock()->SetSystemTimeMilliseconds(kDefaultResultStateTtlInMs + 1000);
+  EXPECT_THAT(result_state_manager.GetNumActiveResultStates(
+                  /*current_time_ms=*/clock()->GetSystemTimeMilliseconds()),
+              Eq(2));
+}
+
 }  // namespace
 }  // namespace lib
 }  // namespace icing

diff --git a/icing/result/result-state-manager_thread-safety_test.cc b/icing/result/result-state-manager_thread-safety_test.cc
index 5578c54..fd82de5 100644
--- a/icing/result/result-state-manager_thread-safety_test.cc
+++ b/icing/result/result-state-manager_thread-safety_test.cc

@@ -17,8 +17,13 @@
 #include <limits>
 #include <memory>
 #include <optional>
+#include <string>
 #include <thread>  // NOLINT
+#include <utility>
+#include <vector>
 
+#include "icing/text_classifier/lib3/utils/base/status.h"
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include "icing/document-builder.h"
@@ -26,11 +31,15 @@
 #include "icing/file/filesystem.h"
 #include "icing/file/portable-file-backed-proto-log.h"
 #include "icing/portable/equals-proto.h"
+#include "icing/portable/gzip_stream.h"
+#include "icing/portable/platform.h"
 #include "icing/result/page-result.h"
 #include "icing/result/result-retriever-v2.h"
 #include "icing/result/result-state-manager.h"
 #include "icing/schema/schema-store.h"
+#include "icing/schema/section.h"
 #include "icing/scoring/priority-queue-scored-document-hits-ranker.h"
+#include "icing/scoring/scored-document-hit.h"
 #include "icing/store/document-store.h"
 #include "icing/testing/common-matchers.h"
 #include "icing/testing/fake-clock.h"
@@ -38,10 +47,10 @@
 #include "icing/testing/test-feature-flags.h"
 #include "icing/testing/tmp-directory.h"
 #include "icing/tokenization/language-segmenter-factory.h"
+#include "icing/tokenization/language-segmenter.h"
 #include "icing/transform/normalizer-factory.h"
 #include "icing/transform/normalizer-options.h"
 #include "icing/transform/normalizer.h"
-#include "icing/util/clock.h"
 #include "icing/util/icu-data-file-helper.h"
 #include "unicode/uloc.h"
 
@@ -61,14 +70,17 @@
   return result_spec;
 }
 
-DocumentProto CreateDocument(int document_id) {
-  return DocumentBuilder()
-      .SetNamespace("namespace")
-      .SetUri(std::to_string(document_id))
-      .SetSchema("Document")
-      .SetCreationTimestampMs(1574365086666 + document_id)
-      .SetScore(document_id)
-      .Build();
+DocumentWrapper CreateDocument(int document_id) {
+  DocumentWrapper document_wrapper;
+  *document_wrapper.mutable_document() =
+      DocumentBuilder()
+          .SetNamespace("namespace")
+          .SetUri(std::to_string(document_id))
+          .SetSchema("Document")
+          .SetCreationTimestampMs(1574365086666 + document_id)
+          .SetScore(document_id)
+          .Build();
+  return document_wrapper;
 }
 
 class ResultStateManagerThreadSafetyTest : public testing::Test {
@@ -109,20 +121,25 @@
 
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult result,
-        DocumentStore::Create(&filesystem_, test_dir_, clock_.get(),
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, test_dir_, clock_.get(), schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     document_store_ = std::move(result.document_store);
 
     ICING_ASSERT_OK_AND_ASSIGN(
-        result_retriever_, ResultRetrieverV2::Create(
-                               document_store_.get(), schema_store_.get(),
-                               language_segmenter_.get(), normalizer_.get()));
+        result_retriever_,
+        ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
+                                  language_segmenter_.get(), normalizer_.get(),
+                                  feature_flags_.get()));
   }
 
   void TearDown() override {
@@ -194,12 +211,14 @@
     ICING_ASSERT_OK_AND_ASSIGN(
         std::unique_ptr<ResultRetrieverV2> result_retriever,
         ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                  language_segmenter_.get(),
-                                  normalizer_.get()));
+                                  language_segmenter_.get(), normalizer_.get(),
+                                  feature_flags_.get()));
     ICING_ASSERT_OK_AND_ASSIGN(
         PageResultInfo page_result_info,
-        result_state_manager.GetNextPage(next_page_token, *result_retriever,
-                                         clock_->GetSystemTimeMilliseconds()));
+        result_state_manager.GetNextPage(
+            next_page_token,
+            /*max_results=*/std::numeric_limits<int32_t>::max(),
+            *result_retriever, clock_->GetSystemTimeMilliseconds()));
     page_results[thread_id] =
         std::make_optional<PageResultInfo>(std::move(page_result_info));
   };
@@ -297,12 +316,14 @@
     ICING_ASSERT_OK_AND_ASSIGN(
         std::unique_ptr<ResultRetrieverV2> result_retriever,
         ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                  language_segmenter_.get(),
-                                  normalizer_.get()));
+                                  language_segmenter_.get(), normalizer_.get(),
+                                  feature_flags_.get()));
 
     libtextclassifier3::StatusOr<PageResultInfo> page_result_info_or =
-        result_state_manager.GetNextPage(next_page_token, *result_retriever,
-                                         clock_->GetSystemTimeMilliseconds());
+        result_state_manager.GetNextPage(
+            next_page_token,
+            /*max_results=*/std::numeric_limits<int32_t>::max(),
+            *result_retriever, clock_->GetSystemTimeMilliseconds());
     if (page_result_info_or.ok()) {
       page_results[thread_id] = std::make_optional<PageResultInfo>(
           std::move(page_result_info_or).ValueOrDie());
@@ -392,8 +413,8 @@
     ICING_ASSERT_OK_AND_ASSIGN(
         std::unique_ptr<ResultRetrieverV2> result_retriever,
         ResultRetrieverV2::Create(document_store_.get(), schema_store_.get(),
-                                  language_segmenter_.get(),
-                                  normalizer_.get()));
+                                  language_segmenter_.get(), normalizer_.get(),
+                                  feature_flags_.get()));
 
     // Retrieve the first page.
     // Documents are ordered by score *ascending*, so the first page should
@@ -428,10 +449,12 @@
     // each thread should retrieve 1, 2, 3, ..., kNumThreads pages.
     int num_subsequent_pages_to_retrieve = thread_id;
     for (int i = 0; i < num_subsequent_pages_to_retrieve; ++i) {
-      ICING_ASSERT_OK_AND_ASSIGN(PageResultInfo page_result_info,
-                                 result_state_manager.GetNextPage(
-                                     next_page_token, *result_retriever,
-                                     clock_->GetSystemTimeMilliseconds()));
+      ICING_ASSERT_OK_AND_ASSIGN(
+          PageResultInfo page_result_info,
+          result_state_manager.GetNextPage(
+              next_page_token,
+              /*max_results=*/std::numeric_limits<int32_t>::max(),
+              *result_retriever, clock_->GetSystemTimeMilliseconds()));
       EXPECT_THAT(page_result_info.second.results, SizeIs(kNumPerPage));
       for (int j = 0; j < kNumPerPage; ++j) {
         EXPECT_THAT(page_result_info.second.results[j].score(),

diff --git a/icing/result/result-state-v2.cc b/icing/result/result-state-v2.cc
index 3aa9359..e7248bf 100644
--- a/icing/result/result-state-v2.cc
+++ b/icing/result/result-state-v2.cc

@@ -17,7 +17,8 @@
 #include <atomic>
 #include <cstdint>
 #include <memory>
-#include <string>
+#include <optional>
+#include <utility>
 #include <vector>
 
 #include "icing/proto/search.pb.h"
@@ -50,15 +51,12 @@
     group_result_limits.push_back(result_grouping.max_results());
     for (const ResultSpecProto::ResultGrouping::Entry& entry :
          result_grouping.entry_groupings()) {
-      const std::string& name_space = entry.namespace_();
-      const std::string& schema = entry.schema();
-      auto entry_id_or = document_store.GetResultGroupingEntryId(
-          result_group_type_, name_space, schema);
-      if (!entry_id_or.ok()) {
+      std::optional<int32_t> entry_id = document_store.GetResultGroupingEntryId(
+          result_group_type_, entry.namespace_(), entry.schema());
+      if (!entry_id.has_value()) {
         continue;
       }
-      int32_t entry_id = entry_id_or.ValueOrDie();
-      entry_id_group_id_map_.insert({entry_id, group_id});
+      entry_id_group_id_map_.insert({*entry_id, group_id});
     }
   }
 }

diff --git a/icing/result/result-state-v2_test.cc b/icing/result/result-state-v2_test.cc
index 6b8b4bd..fb3dea4 100644
--- a/icing/result/result-state-v2_test.cc
+++ b/icing/result/result-state-v2_test.cc

@@ -25,9 +25,11 @@
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include "icing/absl_ports/mutex.h"
+#include "icing/document-builder.h"
 #include "icing/feature-flags.h"
 #include "icing/file/filesystem.h"
 #include "icing/file/portable-file-backed-proto-log.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/document_wrapper.pb.h"
 #include "icing/proto/schema.pb.h"
@@ -79,14 +81,18 @@
     filesystem_.CreateDirectoryRecursively(doc_store_base_dir_.c_str());
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult result,
-        DocumentStore::Create(&filesystem_, doc_store_base_dir_, &clock_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, doc_store_base_dir_, &clock_, schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     document_store_ = std::move(result.document_store);
 
     num_total_hits_ = 0;
@@ -102,7 +108,11 @@
     document.set_namespace_("namespace");
     document.set_uri(std::to_string(document_id));
     document.set_schema("Document");
-    document_store_->Put(std::move(document));
+
+    DocumentWrapper document_wrapper;
+    *document_wrapper.mutable_document() = std::move(document);
+
+    document_store_->Put(document_wrapper);
     return ScoredDocumentHit(document_id, kSectionIdMaskNone, /*score=*/1);
   }
 
@@ -177,23 +187,29 @@
 TEST_F(ResultStateV2Test,
        ShouldConstructNamespaceGroupIdMapAndGroupResultLimitsAccordingToSpecs) {
   // Create 3 docs under namespace1, namespace2, namespace3.
-  DocumentProto document1;
-  document1.set_namespace_("namespace1");
-  document1.set_uri("uri/1");
-  document1.set_schema("Document");
-  ICING_ASSERT_OK(document_store().Put(std::move(document1)));
+  DocumentWrapper document_wrapper1;
+  *document_wrapper1.mutable_document() = DocumentBuilder()
+                                              .SetNamespace("namespace1")
+                                              .SetUri("uri/1")
+                                              .SetSchema("Document")
+                                              .Build();
+  ICING_ASSERT_OK(document_store().Put(document_wrapper1));
 
-  DocumentProto document2;
-  document2.set_namespace_("namespace2");
-  document2.set_uri("uri/2");
-  document2.set_schema("Document");
-  ICING_ASSERT_OK(document_store().Put(std::move(document2)));
+  DocumentWrapper document_wrapper2;
+  *document_wrapper2.mutable_document() = DocumentBuilder()
+                                              .SetNamespace("namespace2")
+                                              .SetUri("uri/2")
+                                              .SetSchema("Document")
+                                              .Build();
+  ICING_ASSERT_OK(document_store().Put(document_wrapper2));
 
-  DocumentProto document3;
-  document3.set_namespace_("namespace3");
-  document3.set_uri("uri/3");
-  document3.set_schema("Document");
-  ICING_ASSERT_OK(document_store().Put(std::move(document3)));
+  DocumentWrapper document_wrapper3;
+  *document_wrapper3.mutable_document() = DocumentBuilder()
+                                              .SetNamespace("namespace3")
+                                              .SetUri("uri/3")
+                                              .SetSchema("Document")
+                                              .Build();
+  ICING_ASSERT_OK(document_store().Put(document_wrapper3));
 
   // Create a ResultSpec that limits "namespace1" to 3 results and limits
   // "namespace2"+"namespace3" to a total of 2 results. Also add
@@ -222,13 +238,13 @@
   entry->set_namespace_("nonexistentNamespace1");
 
   // Get entry ids.
-  ICING_ASSERT_OK_AND_ASSIGN(
+  ICING_ASSERT_HAS_VALUE_AND_ASSIGN(
       int32_t entry_id1, document_store().GetResultGroupingEntryId(
                              result_grouping_type, "namespace1", "Document"));
-  ICING_ASSERT_OK_AND_ASSIGN(
+  ICING_ASSERT_HAS_VALUE_AND_ASSIGN(
       int32_t entry_id2, document_store().GetResultGroupingEntryId(
                              result_grouping_type, "namespace2", "Document"));
-  ICING_ASSERT_OK_AND_ASSIGN(
+  ICING_ASSERT_HAS_VALUE_AND_ASSIGN(
       int32_t entry_id3, document_store().GetResultGroupingEntryId(
                              result_grouping_type, "namespace3", "Document"));
 

diff --git a/icing/result/snippet-retriever.cc b/icing/result/snippet-retriever.cc
index 59552c8..7806ae3 100644
--- a/icing/result/snippet-retriever.cc
+++ b/icing/result/snippet-retriever.cc

@@ -512,9 +512,10 @@
     CharacterIterator start_itr(value);
     CharacterIterator end_itr(value);
     CharacterIterator reset_itr(value);
+    std::vector<Token> batch_tokens;
     bool encountered_error = false;
     while (iterator->Advance()) {
-      std::vector<Token> batch_tokens = iterator->GetTokens();
+      iterator->GetTokens(&batch_tokens);
       if (batch_tokens.empty()) {
         continue;
       }

diff --git a/icing/result/snippet-retriever_test.cc b/icing/result/snippet-retriever_test.cc
index 45ee1ec..927b419 100644
--- a/icing/result/snippet-retriever_test.cc
+++ b/icing/result/snippet-retriever_test.cc

@@ -2224,12 +2224,12 @@
   EXPECT_THAT(snippet.entries(0).property_name(), Eq("embedding1[0]"));
   EXPECT_THAT(
       snippet.entries(0).embedding_matches(),
-      UnorderedElementsAre(EqualsProto(CreateEmbeddingMatchSnippetProto(
+      ElementsAre(EqualsProto(CreateEmbeddingMatchSnippetProto(
           /*score=*/0.5, /*query_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT))));
   EXPECT_THAT(snippet.entries(1).property_name(), Eq("embedding1[2]"));
   EXPECT_THAT(
       snippet.entries(1).embedding_matches(),
-      UnorderedElementsAre(EqualsProto(CreateEmbeddingMatchSnippetProto(
+      ElementsAre(EqualsProto(CreateEmbeddingMatchSnippetProto(
           /*score=*/0.6, /*query_index=*/1, EMBEDDING_METRIC_DOT_PRODUCT))));
   EXPECT_THAT(snippet.entries(2).property_name(), Eq("embedding1[3]"));
   EXPECT_THAT(
@@ -2241,7 +2241,7 @@
               /*score=*/-0.4, /*query_index=*/3, EMBEDDING_METRIC_COSINE))));
   EXPECT_THAT(snippet.entries(3).property_name(), Eq("embedding1[4]"));
   EXPECT_THAT(snippet.entries(3).embedding_matches(),
-              UnorderedElementsAre(EqualsProto(CreateEmbeddingMatchSnippetProto(
+              ElementsAre(EqualsProto(CreateEmbeddingMatchSnippetProto(
                   /*score=*/3, /*query_index=*/2, EMBEDDING_METRIC_COSINE))));
 
   // Section 1 matches
@@ -2256,7 +2256,7 @@
   EXPECT_THAT(snippet.entries(5).property_name(), Eq("embedding2[2]"));
   EXPECT_THAT(
       snippet.entries(5).embedding_matches(),
-      UnorderedElementsAre(EqualsProto(CreateEmbeddingMatchSnippetProto(
+      ElementsAre(EqualsProto(CreateEmbeddingMatchSnippetProto(
           /*score=*/1, /*query_index=*/1, EMBEDDING_METRIC_DOT_PRODUCT))));
 }
 
@@ -2376,12 +2376,12 @@
   EXPECT_THAT(snippet.entries(0).property_name(), Eq("embedding1[0]"));
   EXPECT_THAT(
       snippet.entries(0).embedding_matches(),
-      UnorderedElementsAre(EqualsProto(CreateEmbeddingMatchSnippetProto(
+      ElementsAre(EqualsProto(CreateEmbeddingMatchSnippetProto(
           /*score=*/0.5, /*query_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT))));
   EXPECT_THAT(snippet.entries(1).property_name(), Eq("embedding1[2]"));
   EXPECT_THAT(
       snippet.entries(1).embedding_matches(),
-      UnorderedElementsAre(EqualsProto(CreateEmbeddingMatchSnippetProto(
+      ElementsAre(EqualsProto(CreateEmbeddingMatchSnippetProto(
           /*score=*/0.6, /*query_index=*/1, EMBEDDING_METRIC_DOT_PRODUCT))));
   EXPECT_THAT(snippet.entries(2).property_name(), Eq("embedding1[3]"));
   EXPECT_THAT(
@@ -2393,7 +2393,7 @@
               /*score=*/-0.4, /*query_index=*/3, EMBEDDING_METRIC_COSINE))));
   EXPECT_THAT(snippet.entries(3).property_name(), Eq("embedding1[4]"));
   EXPECT_THAT(snippet.entries(3).embedding_matches(),
-              UnorderedElementsAre(EqualsProto(CreateEmbeddingMatchSnippetProto(
+              ElementsAre(EqualsProto(CreateEmbeddingMatchSnippetProto(
                   /*score=*/3, /*query_index=*/2, EMBEDDING_METRIC_COSINE))));
 
   // Section 1 matches
@@ -2408,7 +2408,7 @@
   EXPECT_THAT(snippet.entries(5).property_name(), Eq("embedding2[2]"));
   EXPECT_THAT(
       snippet.entries(5).embedding_matches(),
-      UnorderedElementsAre(EqualsProto(CreateEmbeddingMatchSnippetProto(
+      ElementsAre(EqualsProto(CreateEmbeddingMatchSnippetProto(
           /*score=*/1, /*query_index=*/1, EMBEDDING_METRIC_DOT_PRODUCT))));
 
   // Document 1 has no matches
@@ -2559,12 +2559,12 @@
   EXPECT_THAT(snippet.entries(1).property_name(), Eq("embedding1[0]"));
   EXPECT_THAT(
       snippet.entries(1).embedding_matches(),
-      UnorderedElementsAre(EqualsProto(CreateEmbeddingMatchSnippetProto(
+      ElementsAre(EqualsProto(CreateEmbeddingMatchSnippetProto(
           /*score=*/0.5, /*query_index=*/0, EMBEDDING_METRIC_DOT_PRODUCT))));
   EXPECT_THAT(snippet.entries(2).property_name(), Eq("embedding1[2]"));
   EXPECT_THAT(
       snippet.entries(2).embedding_matches(),
-      UnorderedElementsAre(EqualsProto(CreateEmbeddingMatchSnippetProto(
+      ElementsAre(EqualsProto(CreateEmbeddingMatchSnippetProto(
           /*score=*/0.6, /*query_index=*/1, EMBEDDING_METRIC_DOT_PRODUCT))));
   EXPECT_THAT(snippet.entries(3).property_name(), Eq("embedding1[3]"));
   EXPECT_THAT(
@@ -2576,7 +2576,7 @@
               /*score=*/-0.4, /*query_index=*/3, EMBEDDING_METRIC_COSINE))));
   EXPECT_THAT(snippet.entries(4).property_name(), Eq("embedding1[4]"));
   EXPECT_THAT(snippet.entries(4).embedding_matches(),
-              UnorderedElementsAre(EqualsProto(CreateEmbeddingMatchSnippetProto(
+              ElementsAre(EqualsProto(CreateEmbeddingMatchSnippetProto(
                   /*score=*/3, /*query_index=*/2, EMBEDDING_METRIC_COSINE))));
 
   // embedding2 matches
@@ -2591,7 +2591,7 @@
   EXPECT_THAT(snippet.entries(6).property_name(), Eq("embedding2[2]"));
   EXPECT_THAT(
       snippet.entries(6).embedding_matches(),
-      UnorderedElementsAre(EqualsProto(CreateEmbeddingMatchSnippetProto(
+      ElementsAre(EqualsProto(CreateEmbeddingMatchSnippetProto(
           /*score=*/1, /*query_index=*/1, EMBEDDING_METRIC_DOT_PRODUCT))));
 
   // 'subject' text matches

diff --git a/icing/schema/backup-schema-producer_test.cc b/icing/schema/backup-schema-producer_test.cc
index d0f609c..8b1d124 100644
--- a/icing/schema/backup-schema-producer_test.cc
+++ b/icing/schema/backup-schema-producer_test.cc

@@ -945,21 +945,50 @@
 
 INSTANTIATE_TEST_SUITE_P(
     BackupSchemaProducerTest, BackupSchemaProducerTest,
-    testing::Values(
-        FeatureFlags(/*allow_circular_schema_definitions=*/true,
-                     /*enable_scorable_properties=*/true,
-                     /*enable_embedding_quantization=*/true,
-                     /*enable_repeated_field_joins=*/true,
-                     /*enable_embedding_backup_generation=*/false,
-                     /*enable_schema_database=*/true,
-                     /*release_backup_schema_file_if_overlay_present=*/true),
-        FeatureFlags(/*allow_circular_schema_definitions=*/true,
-                     /*enable_scorable_properties=*/true,
-                     /*enable_embedding_quantization=*/true,
-                     /*enable_repeated_field_joins=*/true,
-                     /*enable_embedding_backup_generation=*/true,
-                     /*enable_schema_database=*/true,
-                     /*release_backup_schema_file_if_overlay_present=*/true)));
+    testing::Values(FeatureFlags(
+                        /*allow_circular_schema_definitions=*/true,
+                        /*enable_scorable_properties=*/true,
+                        /*enable_embedding_quantization=*/true,
+                        /*enable_repeated_field_joins=*/true,
+                        /*enable_embedding_backup_generation=*/false,
+                        /*enable_schema_database=*/true,
+                        /*release_backup_schema_file_if_overlay_present=*/true,
+                        /*enable_strict_page_byte_size=*/true,
+                        /*enable_smaller_decompression_buffer_size=*/true,
+                        /*enable_eigen_embedding_scoring=*/true,
+                        /*enable_passing_filter_to_children=*/true,
+                        /*enable_proto_log_new_header_format=*/true,
+                        /*enable_embedding_iterator_v2=*/true,
+                        /*enable_reusable_decompression_buffer=*/true,
+                        /*enable_schema_type_id_optimization=*/true,
+                        /*enable_optimize_improvements=*/true,
+                        /*expired_document_purge_threshold_ms=*/0,
+                        /*enable_non_existent_qualified_id_join=*/true,
+                        /*enable_skip_set_schema_type_equality_check=*/true,
+                        /*enable_embed_query_optimization=*/true,
+                        /*enable_schema_definition_deduping=*/true),
+                    FeatureFlags(
+                        /*allow_circular_schema_definitions=*/true,
+                        /*enable_scorable_properties=*/true,
+                        /*enable_embedding_quantization=*/true,
+                        /*enable_repeated_field_joins=*/true,
+                        /*enable_embedding_backup_generation=*/true,
+                        /*enable_schema_database=*/true,
+                        /*release_backup_schema_file_if_overlay_present=*/true,
+                        /*enable_strict_page_byte_size=*/true,
+                        /*enable_smaller_decompression_buffer_size=*/true,
+                        /*enable_eigen_embedding_scoring=*/true,
+                        /*enable_passing_filter_to_children=*/true,
+                        /*enable_proto_log_new_header_format=*/true,
+                        /*enable_embedding_iterator_v2=*/true,
+                        /*enable_reusable_decompression_buffer=*/true,
+                        /*enable_schema_type_id_optimization=*/true,
+                        /*enable_optimize_improvements=*/true,
+                        /*expired_document_purge_threshold_ms=*/0,
+                        /*enable_non_existent_qualified_id_join=*/true,
+                        /*enable_skip_set_schema_type_equality_check=*/true,
+                        /*enable_embed_query_optimization=*/true,
+                        /*enable_schema_definition_deduping=*/true)));
 
 }  // namespace
 

diff --git a/icing/schema/schema-store.cc b/icing/schema/schema-store.cc
index dda7746..537ea6f 100644
--- a/icing/schema/schema-store.cc
+++ b/icing/schema/schema-store.cc

@@ -141,7 +141,7 @@
   size_t db_index = schema_type.find(database_delimeter);
   std::string database;
   if (db_index != std::string::npos) {
-    database = schema_type.substr(0, db_index);
+    database = schema_type.substr(0, db_index + 1);
   }
   return database;
 }
@@ -222,24 +222,32 @@
   int64_t file_size = filesystem->GetFileSize(sfd.get());
   if (file_size == sizeof(LegacyHeader)) {
     LegacyHeader legacy_header;
-    if (!filesystem->Read(sfd.get(), &legacy_header, sizeof(legacy_header))) {
+    if (filesystem->Read(sfd.get(), &legacy_header, sizeof(legacy_header)) !=
+        sizeof(legacy_header)) {
       return absl_ports::InternalError(
           absl_ports::StrCat("Couldn't read: ", path));
     }
     if (legacy_header.magic != Header::kMagic) {
+      ICING_LOG(ERROR)
+          << "Invalid legacy header magic for SchemaStore. Expected: "
+          << Header::kMagic << ", actual: " << legacy_header.magic;
       return absl_ports::InternalError(
-          absl_ports::StrCat("Invalid header kMagic for file: ", path));
+          "Invalid legacy header magic for SchemaStore");
     }
     serialized_header.checksum = legacy_header.checksum;
   } else if (file_size == sizeof(SerializedHeader)) {
-    if (!filesystem->Read(sfd.get(), &serialized_header,
-                          sizeof(serialized_header))) {
+    if (filesystem->Read(sfd.get(), &serialized_header,
+                         sizeof(serialized_header)) !=
+        sizeof(serialized_header)) {
       return absl_ports::InternalError(
           absl_ports::StrCat("Couldn't read: ", path));
     }
     if (serialized_header.magic != Header::kMagic) {
+      ICING_LOG(ERROR)
+          << "Invalid serialized header magic for SchemaStore. Expected: "
+          << Header::kMagic << ", actual: " << serialized_header.magic;
       return absl_ports::InternalError(
-          absl_ports::StrCat("Invalid header kMagic for file: ", path));
+          "Invalid serialized header magic for SchemaStore");
     }
   } else if (file_size != 0) {
     // file is neither the legacy header, the new header nor empty. Something is
@@ -558,9 +566,10 @@
   bool overlay_schema_file_exists =
       filesystem_->FileExists(overlay_schema_filename.c_str());
 
+  // The base schema file will be released at a later point (if necessary),
+  // after InitializeInternal is done.
   libtextclassifier3::Status base_schema_state = schema_file_.Read().status();
   if (!base_schema_state.ok() && !absl_ports::IsNotFound(base_schema_state)) {
-    ResetSchemaFileIfNeeded();
     return base_schema_state;
   }
 
@@ -568,7 +577,6 @@
   // 1. Everything is missing. This is an empty schema store.
   if (!base_schema_state.ok() && !overlay_schema_file_exists &&
       !header_exists) {
-    ResetSchemaFileIfNeeded();
     return libtextclassifier3::Status::OK;
   }
 
@@ -577,7 +585,6 @@
   if (base_schema_state.ok() && !overlay_schema_file_exists && header_exists &&
       !header_->overlay_created()) {
     // Nothing else to do. Just return safely.
-    ResetSchemaFileIfNeeded();
     return libtextclassifier3::Status::OK;
   }
 
@@ -587,7 +594,6 @@
       header_->overlay_created()) {
     overlay_schema_file_ = std::make_unique<FileBackedProto<SchemaProto>>(
         *filesystem_, MakeOverlaySchemaFilename(base_dir_));
-    ResetSchemaFileIfNeeded();
     return libtextclassifier3::Status::OK;
   }
 
@@ -595,7 +601,6 @@
   // Return an error.
   bool overlay_created = header_->overlay_created();
   bool base_schema_exists = base_schema_state.ok();
-  ResetSchemaFileIfNeeded();
   return absl_ports::InternalError(IcingStringUtil::StringPrintf(
       "Unable to properly load schema. Header {exists:%d, overlay_created:%d}, "
       "base schema exists: %d, overlay_schema_exists: %d",
@@ -625,6 +630,7 @@
     initialize_stats->set_num_schema_types(type_config_map_.size());
   }
   has_schema_successfully_set_ = true;
+  ResetSchemaFileIfNeeded();
 
   return libtextclassifier3::Status::OK;
 }
@@ -694,7 +700,6 @@
 
   // Write the header
   ICING_RETURN_IF_ERROR(UpdateChecksum());
-  ResetSchemaFileIfNeeded();
   return libtextclassifier3::Status::OK;
 }
 
@@ -705,6 +710,7 @@
       SchemaUtil::BuildTransitiveInheritanceGraph(*schema_proto));
 
   reverse_schema_type_mapper_.clear();
+  reverse_schema_type_mapper_hash_.clear();
   database_type_map_.clear();
   type_config_map_.clear();
   schema_subtype_id_map_.clear();
@@ -716,6 +722,8 @@
 
     // Build reverse_schema_type_mapper_
     reverse_schema_type_mapper_.insert({type_id, type_name});
+    reverse_schema_type_mapper_hash_.insert(
+        {type_id, GetSchemaNameHash(type_name)});
 
     // Build database_type_map_
     database_type_map_[database].push_back(type_name);
@@ -780,7 +788,8 @@
   // schema (both of which will have a checksum of 0). For existing, but empty
   // schemas, we need to continue with the checksum calculation of the other
   // components.
-  if (schema_checksum == Crc32() && !has_schema_successfully_set_) {
+  if (schema_checksum == Crc32() &&
+      absl_ports::IsNotFound(schema_file_.Read().status())) {
     return schema_checksum;
   }
 
@@ -805,7 +814,8 @@
   // schema (both of which will have a checksum of 0). For existing, but empty
   // schemas, we need to continue with the checksum calculation of the other
   // components.
-  if (schema_checksum == Crc32() && !has_schema_successfully_set_) {
+  if (schema_checksum == Crc32() &&
+      absl_ports::IsNotFound(schema_file_.Read().status())) {
     return schema_checksum;
   }
   Crc32 total_checksum;
@@ -875,7 +885,8 @@
   bool ignore_errors_and_delete_documents =
       set_schema_request.ignore_errors_and_delete_documents();
 
-  if (feature_flags_->enable_schema_database()) {
+  if (feature_flags_->enable_schema_database() &&
+      !set_schema_request.database().empty()) {
     // Step 1: (Only required if schema database is enabled)
     // Do some preliminary checks on the new schema before formal validation and
     // delta computation. This checks that:
@@ -946,11 +957,19 @@
   // Get the full new SchemaProto that is a combination of the existing schema
   // and new_schema. This is needed as we can only write the full proto to the
   // schema file.
-  ICING_ASSIGN_OR_RETURN(
-      SchemaProto full_new_schema,
-      GetFullSchemaProtoWithUpdatedDb(std::move(new_schema), database));
+  SchemaProto full_new_schema;
+  if (feature_flags_->enable_schema_type_id_optimization()) {
+    ICING_ASSIGN_OR_RETURN(
+        full_new_schema,
+        GetFullOptimizedSchemaProto(std::move(new_schema), database));
+  } else {
+    ICING_ASSIGN_OR_RETURN(
+        full_new_schema,
+        GetFullSchemaProtoWithUpdatedDb(std::move(new_schema), database));
+  }
   ICING_RETURN_IF_ERROR(ApplySchemaChange(std::move(full_new_schema)));
   has_schema_successfully_set_ = true;
+  ResetSchemaFileIfNeeded();
 
   return result;
 }
@@ -963,30 +982,32 @@
   SetSchemaResult result;
   result.success = true;
 
-  if (feature_flags_->enable_schema_database()) {
-    // Sanity check to make sure that we're comparing schemas from the same
-    // database.
-    // The new code path ensures that old_schema contains types from exactly one
-    // database since it's obtained using GetSchema(database), which is
-    // guaranteed to only return types from the single provided database.
-    libtextclassifier3::Status validate_old_schema_database =
-        ValidateSchemaDatabase(old_schema, database);
-    if (!validate_old_schema_database.ok()) {
-      return absl_ports::InvalidArgumentError(
-          "Schema database mismatch between new and old schemas. This should "
-          "never happen");
-    }
+  if (!feature_flags_->enable_skip_set_schema_type_equality_check()) {
+    if (feature_flags_->enable_schema_database()) {
+      // Sanity check to make sure that we're comparing schemas from the same
+      // database.
+      // The new code path ensures that old_schema contains types from exactly
+      // one database since it's obtained using GetSchema(database), which is
+      // guaranteed to only return types from the single provided database.
+      libtextclassifier3::Status validate_old_schema_database =
+          ValidateSchemaDatabase(old_schema, database);
+      if (!validate_old_schema_database.ok()) {
+        return absl_ports::InvalidArgumentError(
+            "Schema database mismatch between new and old schemas. This should "
+            "never happen");
+      }
 
-    // Check if the schema types are the same between the new and old schema,
-    // ignoring order.
-    if (AreSchemaTypesEqual(new_schema, old_schema)) {
-      return result;
-    }
-  } else {
-    // Old equality check that is sensitive to type definition order.
-    if (new_schema.SerializeAsString() == old_schema.SerializeAsString()) {
-      // Same schema as before. No need to update anything
-      return result;
+      // Check if the schema types are the same between the new and old schema,
+      // ignoring order.
+      if (AreSchemaTypesEqual(new_schema, old_schema)) {
+        return result;
+      }
+    } else {
+      // Old equality check that is sensitive to type definition order.
+      if (new_schema.SerializeAsString() == old_schema.SerializeAsString()) {
+        // Same schema as before. No need to update anything
+        return result;
+      }
     }
   }
 
@@ -1038,9 +1059,16 @@
   // Get the full new SchemaProto that is a combination of the existing schema
   // and new_schema. This is needed to calculate the updated SchemaTypeIds, and
   // for writing the full proto to the schema file.
-  ICING_ASSIGN_OR_RETURN(
-      SchemaProto full_new_schema,
-      GetFullSchemaProtoWithUpdatedDb(std::move(new_schema), database));
+  SchemaProto full_new_schema;
+  if (feature_flags_->enable_schema_type_id_optimization()) {
+    ICING_ASSIGN_OR_RETURN(
+        full_new_schema,
+        GetFullOptimizedSchemaProto(std::move(new_schema), database));
+  } else {
+    ICING_ASSIGN_OR_RETURN(
+        full_new_schema,
+        GetFullSchemaProtoWithUpdatedDb(std::move(new_schema), database));
+  }
 
   // We still need to update old_schema_type_ids_changed. We need to retrieve
   // the entire old schema for this, as type ids are assigned for the entire
@@ -1059,6 +1087,7 @@
   if (result.success) {
     ICING_RETURN_IF_ERROR(ApplySchemaChange(std::move(full_new_schema)));
     has_schema_successfully_set_ = true;
+    ResetSchemaFileIfNeeded();
   }
 
   // Convert schema types to SchemaTypeIds after the new schema is applied.
@@ -1111,6 +1140,12 @@
       SchemaStore::Create(filesystem_, temp_schema_store_dir.dir(), clock_,
                           feature_flags_, std::move(new_schema)));
 
+  // Call PersistToDisk() to write the new schema file to disk. This is needed
+  // to ensure that all subcomponents of the new schema store (header file,
+  // derived files, etc) are written to disk before we swap the files in the
+  // next step.
+  ICING_RETURN_IF_ERROR(new_schema_store->PersistToDisk());
+
   // Then we swap the new schema file + new derived files with the old files.
   if (!filesystem_->SwapFiles(base_dir_.c_str(),
                               temp_schema_store_dir.dir().c_str())) {
@@ -1121,12 +1156,13 @@
   std::string old_base_dir = std::move(base_dir_);
   *this = std::move(*new_schema_store);
 
-  // After the std::move, the filepaths saved in this instance and in the
-  // schema_file_ instance will still be the one from temp_schema_store_dir
-  // even though they now point to files that are within old_base_dir.
-  // Manually set them to the correct paths.
+  // After the std::move, the filepaths saved in this instance, the header_ and
+  // in the schema_file_ instance will still be the one from
+  // temp_schema_store_dir even though they now point to files that are within
+  // old_base_dir. Manually set them to the correct paths.
   base_dir_ = std::move(old_base_dir);
   schema_file_.SetSwappedFilepath(MakeSchemaFilename(base_dir_));
+  header_->SetSwappedFilepath(MakeHeaderFilename(base_dir_));
   if (overlay_schema_file_ != nullptr) {
     overlay_schema_file_->SetSwappedFilepath(
         MakeOverlaySchemaFilename(base_dir_));
@@ -1417,7 +1453,8 @@
 
 libtextclassifier3::Status SchemaStore::ValidateSchemaDatabase(
     const SchemaProto& new_schema, const std::string& database) const {
-  if (!feature_flags_->enable_schema_database() || new_schema.types().empty()) {
+  if (!feature_flags_->enable_schema_database() || new_schema.types().empty() ||
+      database.empty()) {
     return libtextclassifier3::Status::OK;
   }
 
@@ -1443,9 +1480,11 @@
       auto iter = type_config_map_.find(type_config.schema_type());
       if (iter != type_config_map_.end() &&
           database != iter->second.database()) {
-        return absl_ports::AlreadyExistsError(
-            absl_ports::StrCat("schema_type name: '", type_config.schema_type(),
-                               "' is already in use by a different database."));
+        return absl_ports::AlreadyExistsError(absl_ports::StrCat(
+            "schema_type name: '", type_config.schema_type(),
+            "' is already in use by a different database. Request database: '",
+            database, "' vs. existing database: '", iter->second.database(),
+            "'."));
       }
     }
   }
@@ -1453,12 +1492,143 @@
 }
 
 libtextclassifier3::StatusOr<SchemaProto>
+SchemaStore::GetFullOptimizedSchemaProto(
+    SchemaProto input_database_schema,
+    const std::string& database_to_update) const {
+  libtextclassifier3::StatusOr<const SchemaProto*> schema_proto = GetSchema();
+  if (absl_ports::IsNotFound(schema_proto.status())) {
+    // We don't have a pre-existing schema -- we can return the input database
+    // schema as it's already the full schema.
+    return input_database_schema;
+  }
+
+  if (!schema_proto.ok()) {
+    // Real error.
+    return schema_proto.status();
+  }
+
+  if (!has_schema_successfully_set_) {
+    return absl_ports::InternalError(
+        "Schema store was not initialized properly.");
+  }
+
+  // At this point, we have a pre-existing schema -- we need to merge the
+  // updated database with the existing schema.
+  const SchemaProto* existing_schema = schema_proto.ValueOrDie();
+
+  // Step 1: Do some pre-processing to build maps of existing and input
+  // types, as well as a list of newly added types.
+  std::unordered_set<std::string> existing_type_names;
+  for (const SchemaTypeConfigProto& type_config : existing_schema->types()) {
+    if (database_to_update.empty() ||
+        type_config.database() == database_to_update) {
+      existing_type_names.insert(type_config.schema_type());
+    }
+  }
+  std::unordered_map<std::string, SchemaTypeConfigProto*> input_type_map;
+  std::vector<SchemaTypeConfigProto*> new_input_types;
+  for (SchemaTypeConfigProto& type_config :
+       *input_database_schema.mutable_types()) {
+    input_type_map.insert({type_config.schema_type(), &type_config});
+    if (existing_type_names.find(type_config.schema_type()) ==
+        existing_type_names.end()) {
+      new_input_types.push_back(&type_config);
+    }
+  }
+
+  // Step 2: Iterate through existing schema and add them to the full schema,
+  // replacing existing types with the input types if the database is the one
+  // being updated by the input schema.
+  std::vector<SchemaTypeConfigProto> full_schema_types;
+  full_schema_types.reserve(existing_schema->types().size() +
+                            new_input_types.size());
+  std::vector<int> unfilled_type_ids;
+  int new_input_types_index = 0;
+
+  for (int i = 0; i < existing_schema->types().size(); ++i) {
+    const SchemaTypeConfigProto& existing_type_config =
+        existing_schema->types().at(i);
+
+    if (database_to_update.empty() ||
+        existing_type_config.database() == database_to_update) {
+      auto itr = input_type_map.find(existing_type_config.schema_type());
+      if (itr != input_type_map.end()) {
+        // Existing type is still present in the input types. Add the input
+        // type to the full schema at the same position.
+        full_schema_types.push_back(std::move(*itr->second));
+      } else {
+        // Existing type is no longer present in input types. Add a new type
+        // to replace the existing type if the new types are not exhausted.
+        // This is to try to preserve the type-ids of subsequent existing
+        // types that are still present.
+        if (new_input_types_index < new_input_types.size()) {
+          full_schema_types.push_back(
+              std::move(*new_input_types[new_input_types_index++]));
+        } else {
+          // No more input types to replace the existing type. Add a dummy
+          // type instead. This will be replaced with valid type from the end
+          // of the schema later.
+          full_schema_types.push_back(SchemaTypeConfigProto());
+          unfilled_type_ids.push_back(i);
+        }
+      }
+    } else {
+      // Existing type is from a different database. Add it to the full schema
+      // without any changes.
+      full_schema_types.push_back(existing_type_config);
+    }
+  }
+
+  // Step 3a: If there are remaining new input types, append them to the end
+  // of the SchemaProto. This happens when there are more input types than
+  // existing types for database_to_update.
+  for (; new_input_types_index < new_input_types.size();
+       ++new_input_types_index) {
+    full_schema_types.push_back(
+        std::move(*new_input_types[new_input_types_index]));
+  }
+
+  // Step 3b: Iterate from the end of the full schema types and fill the
+  // unfilled type ids with the last few types in the schema. This happens
+  // when there are fewer input types than existing types for
+  // database_to_update.
+  if (!unfilled_type_ids.empty()) {
+    int num_unfilled_types = static_cast<int>(unfilled_type_ids.size());
+
+    // Find non-dummy types at the end of the schema.
+    int unfilled_index = 0;
+    for (int i = 0; i < num_unfilled_types; ++i) {
+      int idx =
+          static_cast<int>(full_schema_types.size()) - num_unfilled_types + i;
+      if (!full_schema_types[idx].schema_type().empty()) {
+        // Backfill idx to unfilled_type_ids[unfilled_index].
+        full_schema_types[unfilled_type_ids[unfilled_index++]] =
+            std::move(full_schema_types[idx]);
+      }
+    }
+
+    // Resize the vector to remove the types that were moved.
+    full_schema_types.resize(full_schema_types.size() - num_unfilled_types);
+  }
+
+  SchemaProto full_schema;
+  full_schema.mutable_types()->Reserve(
+      static_cast<int>(full_schema_types.size()));
+  for (SchemaTypeConfigProto& type_config : full_schema_types) {
+    *full_schema.add_types() = std::move(type_config);
+  }
+
+  return full_schema;
+}
+
+libtextclassifier3::StatusOr<SchemaProto>
 SchemaStore::GetFullSchemaProtoWithUpdatedDb(
     SchemaProto input_database_schema,
     const std::string& database_to_update) const {
-  if (!feature_flags_->enable_schema_database()) {
-    // If the schema database is not enabled, the input schema is already the
-    // full schema, so we don't need to do any merges.
+  if (!feature_flags_->enable_schema_database() || database_to_update.empty()) {
+    // The schema database is not enabled, or we're updating using the empty
+    // schema database. This means that the input schema is already the full
+    // schema, so we don't need to do any merges.
     return input_database_schema;
   }
 
@@ -1491,20 +1661,9 @@
   const SchemaProto* existing_schema = schema_proto.ValueOrDie();
   SchemaProto full_schema;
 
-  // 1. Add types from the existing schema, replacing existing types with the
-  // input types if the database is the one being updated by the input schema.
-  // - For database_to_update, we replace the existing types with the input
-  //   types. Any existing type not included in input_database_schema is
-  //   deleted.
-  // - If there are more input types than existing types for database_to_update,
-  //   the rest of the input types are appended to the end of the full_schema.
-  // - If there are fewer input types than existing types for
-  //   database_to_update, we shift forward all existing types that appear after
-  //   the last input type.
-  // - For existing types from other databases, we preserve the existing order
-  //   after adding to full_schema. Note that the type-ids of existing types
-  //   might still change if some types are deleted in the database_to_update as
-  //   this will cause all subsequent types ids to shift forward.
+  // This old rewrite does not preserve the type-ids of existing types, but
+  // rather inserts the input types for the database in the order in which
+  // they appear in the input proto.
   int input_schema_index = 0, existing_schema_index = 0;
   while (input_schema_index < input_database_schema.types().size() &&
          existing_schema_index < existing_schema->types().size()) {

diff --git a/icing/schema/schema-store.h b/icing/schema/schema-store.h
index b5d89ce..766f52b 100644
--- a/icing/schema/schema-store.h
+++ b/icing/schema/schema-store.h

@@ -30,6 +30,7 @@
 #include "icing/text_classifier/lib3/utils/base/status.h"
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "icing/absl_ports/canonical_errors.h"
+#include "icing/absl_ports/str_cat.h"
 #include "icing/feature-flags.h"
 #include "icing/file/file-backed-proto.h"
 #include "icing/file/filesystem.h"
@@ -157,6 +158,8 @@
           min_overlay_version_compatibility;
     }
 
+    void SetSwappedFilepath(std::string path) { path_ = std::move(path); }
+
    private:
     explicit Header(SerializedHeader serialized_header, std::string path,
                     ScopedFd header_fd, const Filesystem* filesystem)
@@ -253,8 +256,6 @@
 
   static constexpr std::string_view kSchemaTypeWildcard = "*";
 
-  static constexpr std::string_view kDefaultEmptySchemaDatabase = "";
-
   // Factory function to create a SchemaStore which does not take ownership
   // of any input components, and all pointers must refer to valid objects that
   // outlive the created SchemaStore instance. The base_dir must already exist.
@@ -329,18 +330,22 @@
   // `SetSchema(SetSchemaRequestProto&& set_schema_request)` instead.
   //
   // TODO: b/337913932 - Remove this method once all callers (currently only
-  // used in tests) are migrated to the new SetSchema method.
+  // used in tests) are migrated to the new SetSchema method that takes a
+  // SetSchemaRequestProto.
   libtextclassifier3::StatusOr<SetSchemaResult> SetSchema(
       SchemaProto new_schema, bool ignore_errors_and_delete_documents);
 
   // Update our current schema if it's compatible. Does not accept incompatible
-  // schema or schema with types from multiple databases. Compatibility rules
-  // defined by SchemaUtil::ComputeCompatibilityDelta.
+  // schema or schema subsets with types from multiple databases. Compatibility
+  // rules defined by SchemaUtil::ComputeCompatibilityDelta.
   //
-  // Does not support setting the schema across multiple databases if
-  // `feature_flags_->enable_schema_database()` is true. This means that:
-  // - All types within the new schema must have their `database` field matching
-  //  `set_schema_request.database()`.
+  // This method accepts either a full schema (indicated by an empty database
+  // field) or a schema subset with types from a single database.
+  // - If `set_schema_request.database()` is non-empty, then all types in the
+  //   new schema must have their `database` field matching
+  //   `set_schema_request.database()`.
+  // - If `set_schema_request.database()` is empty, then the new schema will be
+  //   taken as the full schema, and will replace the entire existing schema.
   //
   // If ignore_errors_and_delete_documents is set to true, then incompatible
   // schema are allowed and we'll force set the schema, meaning
@@ -578,6 +583,26 @@
       const google::protobuf::RepeatedPtrField<TypePropertyMask>& type_property_masks)
       const;
 
+  // Returns the hash of a schema name.
+  static uint32_t GetSchemaNameHash(std::string_view schema_name) {
+    return Crc32(schema_name).Get();
+  }
+
+  // Returns the hash of the schema name for the given schema type id.
+  //
+  // Returns:
+  //   - The hash value on success.
+  //   - INVALID_ARGUMENT_ERROR if schema_type_id is invalid.
+  libtextclassifier3::StatusOr<uint32_t> GetSchemaNameHash(
+      SchemaTypeId schema_type_id) const {
+    auto it = reverse_schema_type_mapper_hash_.find(schema_type_id);
+    if (it == reverse_schema_type_mapper_hash_.end()) {
+      return absl_ports::InvalidArgumentError(absl_ports::StrCat(
+          "Invalid SchemaTypeId ", std::to_string(schema_type_id)));
+    }
+    return it->second;
+  }
+
  private:
   // Factory function to create a SchemaStore and set its schema. The created
   // instance does not take ownership of any input components and all pointers
@@ -763,6 +788,7 @@
   // Requires:
   //   - `new_schema` and `database` are valid according to
   //     `ValidateSchemaDatabase(new_schema, database)`
+  //   - `database` is not empty.
   //   - Types in `new_schema` and `old_schema` all belong to the provided
   //     database.
   //     - The old schema is guaranteed to contain types from exactly one
@@ -805,6 +831,59 @@
       const SchemaProto& new_schema, const std::string& database) const;
 
   // Returns a SchemaProto representing the full schema, which is a combination
+  // of the existing schema and the input database schema. The returned
+  // SchemaProto is optimized to preserve as many type ids as possible.
+  //
+  // Note that `database_to_update` could also be the empty string, which means
+  // that the entire schema is being updated. In this case,
+  // `input_database_schema` must contain all types in the full schema. Any
+  // preexisting type not in `input_database_schema` will be deleted.
+  //
+  // For database_to_update, we replace the existing types with the input
+  //   types. Any existing type not included in input_database_schema is
+  //   deleted.
+  // - When possible, existing types are added in the position in which they
+  //   appear in the existing schema so as to preserve the type-ids of
+  //   existing types.
+  // - If there are more input types than existing types for
+  //   database_to_update, added input types are appended to the end of the
+  //   full_schema.
+  // - If there are fewer input types than existing types for
+  //   database_to_update, we use the last few input types to replace the
+  //   deleted existing types, so as to preserve as many old type-ids as
+  //   possible.
+  // - For existing types from other databases, we preserve the existing order
+  //   after adding to full_schema. Note that the type-ids of existing types
+  //   might still change if some types are deleted in the database_to_update
+  //   as this will cause all subsequent types ids to shift forward.
+  // - This means that:
+  //   - When adding types to a database, the type-ids of existing types will
+  //     not change.
+  //   - When types are deleted, we fill their original slots with the last
+  //     valid types in the schema to preserve as many type-ids as possible.
+  // - If input_database_schema is an empty proto, then all types from
+  //   database_to_update are deleted.
+  //
+  // Requires:
+  //   - input_database_schema is valid according to `ValidateSchemaDatabase`.
+  //
+  // Returns:
+  //   - SchemaProto on success
+  //   - INTERNAL_ERROR on any IO errors, or if the schema store was not
+  //     previously initialized properly.
+  //   - INVALID_ARGUMENT_ERROR if the input schema does not match
+  //     database_to_update.
+  libtextclassifier3::StatusOr<SchemaProto> GetFullOptimizedSchemaProto(
+      SchemaProto input_database_schema,
+      const std::string& database_to_update) const;
+
+  // TODO: b/434218554 - Remove this method once schema type id optimization is
+  // fully rolled out.
+  //
+  // This method should only be called when
+  // `feature_flags_->enable_schema_type_id_optimization` is false.
+  //
+  // Returns a SchemaProto representing the full schema, which is a combination
   // of the existing schema and the input database schema. Deletes all types
   // belonging to the specified database if input_database_schema is an empty
   // proto.
@@ -934,6 +1013,12 @@
   // map of schema_type_mapper_.
   std::unordered_map<SchemaTypeId, std::string> reverse_schema_type_mapper_;
 
+  // Maps schema type ids to the hash value of the corresponding schema type
+  // name.
+  // TODO(b/436237337): Consider merging this with reverse_schema_type_mapper_
+  // to save memory.
+  std::unordered_map<SchemaTypeId, uint32_t> reverse_schema_type_mapper_hash_;
+
   // A hash map of (database -> vector of type config names in the database).
   //
   // We use a vector instead of a set because we need to preserve the order of

diff --git a/icing/schema/schema-store_test.cc b/icing/schema/schema-store_test.cc
index 4993408..8afe4a1 100644
--- a/icing/schema/schema-store_test.cc
+++ b/icing/schema/schema-store_test.cc

@@ -14,6 +14,7 @@
 
 #include "icing/schema/schema-store.h"
 
+#include <algorithm>
 #include <cstdint>
 #include <memory>
 #include <optional>
@@ -620,65 +621,66 @@
                           feature_flags_.get(),
                           /*initialize_stats=*/nullptr));
 
-  SchemaProto db1_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db1_email").SetDatabase("db1"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db1_message")
-                       .SetDatabase("db1"))
-          .Build();
+  SchemaProto db1_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db1/email")
+                                            .SetDatabase("db1/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db1/message")
+                                            .SetDatabase("db1/"))
+                               .Build();
   SchemaStore::SetSchemaResult result;
   result.success = true;
-  result.schema_types_new_by_name.insert("db1_email");
-  result.schema_types_new_by_name.insert("db1_message");
+  result.schema_types_new_by_name.insert("db1/email");
+  result.schema_types_new_by_name.insert("db1/message");
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  db1_schema, /*database=*/"db1",
+                  db1_schema, /*database=*/"db1/",
                   /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   EXPECT_THAT(schema_store->GetSchema(),
               IsOkAndHolds(Pointee(EqualsProto(db1_schema))));
-  EXPECT_THAT(schema_store->GetSchema("db1"),
+  EXPECT_THAT(schema_store->GetSchema("db1/"),
               IsOkAndHolds(EqualsProto(db1_schema)));
 
   // Set a schema in a different database
-  SchemaProto db2_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db2_email").SetDatabase("db2"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_message")
-                       .SetDatabase("db2"))
-          .Build();
+  SchemaProto db2_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db2/email")
+                                            .SetDatabase("db2/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db2/message")
+                                            .SetDatabase("db2/"))
+                               .Build();
   result = SchemaStore::SetSchemaResult();
   result.success = true;
-  result.schema_types_new_by_name.insert("db2_email");
-  result.schema_types_new_by_name.insert("db2_message");
+  result.schema_types_new_by_name.insert("db2/email");
+  result.schema_types_new_by_name.insert("db2/message");
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  db2_schema, /*database=*/"db2",
+                  db2_schema, /*database=*/"db2/",
                   /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
 
   // Check the full schema. Databases that are updated last are appended to the
   // schema proto
-  SchemaProto expected_full_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db1_email").SetDatabase("db1"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db1_message")
-                       .SetDatabase("db1"))
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db2_email").SetDatabase("db2"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_message")
-                       .SetDatabase("db2"))
-          .Build();
+  SchemaProto expected_full_schema = SchemaBuilder()
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db1/email")
+                                                      .SetDatabase("db1/"))
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db1/message")
+                                                      .SetDatabase("db1/"))
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db2/email")
+                                                      .SetDatabase("db2/"))
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db2/message")
+                                                      .SetDatabase("db2/"))
+                                         .Build();
   EXPECT_THAT(schema_store->GetSchema(),
               IsOkAndHolds(Pointee(EqualsProto(expected_full_schema))));
-  EXPECT_THAT(schema_store->GetSchema("db1"),
+  EXPECT_THAT(schema_store->GetSchema("db1/"),
               IsOkAndHolds(EqualsProto(db1_schema)));
-  EXPECT_THAT(schema_store->GetSchema("db2"),
+  EXPECT_THAT(schema_store->GetSchema("db2/"),
               IsOkAndHolds(EqualsProto(db2_schema)));
 }
 
@@ -732,6 +734,175 @@
   EXPECT_THAT(schema_store->GetSchema(""), IsOkAndHolds(EqualsProto(schema)));
 }
 
+TEST_F(SchemaStoreTest, SetSchemaWithEmptyDatabaseResetsEntireSchema) {
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<SchemaStore> schema_store,
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_,
+                          feature_flags_.get()));
+
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("email"))
+          .AddType(SchemaTypeConfigBuilder().SetType("message"))
+          .Build();
+  SchemaStore::SetSchemaResult result;
+  result.success = true;
+  result.schema_types_new_by_name.insert("email");
+  result.schema_types_new_by_name.insert("message");
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  schema, /*database=*/"",
+                  /*ignore_errors_and_delete_documents=*/false)),
+              IsOkAndHolds(EqualsSetSchemaResult(result)));
+  EXPECT_THAT(schema_store->GetSchema(),
+              IsOkAndHolds(Pointee(EqualsProto(schema))));
+  EXPECT_THAT(schema_store->GetSchema(""), IsOkAndHolds(EqualsProto(schema)));
+
+  // Reset the schema using an empty database. This should ignore the
+  // database field in the schema type configs and reset the entire schema.
+  schema = SchemaBuilder()
+               .AddType(SchemaTypeConfigBuilder()
+                            .SetType("db1/email_v1")
+                            .SetDatabase("db1/"))
+               .AddType(SchemaTypeConfigBuilder()
+                            .SetType("db1/message_v1")
+                            .SetDatabase("db1/"))
+               .AddType(SchemaTypeConfigBuilder()
+                            .SetType("db2/email_v2")
+                            .SetDatabase("db2/"))
+               .AddType(SchemaTypeConfigBuilder()
+                            .SetType("db2/message_v2")
+                            .SetDatabase("db2/"))
+               .Build();
+  result = SchemaStore::SetSchemaResult();
+  result.success = true;
+  result.schema_types_new_by_name.insert("db1/email_v1");
+  result.schema_types_new_by_name.insert("db1/message_v1");
+  result.schema_types_new_by_name.insert("db2/email_v2");
+  result.schema_types_new_by_name.insert("db2/message_v2");
+  result.schema_types_deleted_by_name.insert("email");
+  result.schema_types_deleted_by_name.insert("message");
+  result.schema_types_deleted_by_id.insert(0);
+  result.schema_types_deleted_by_id.insert(1);
+
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  schema, /*database=*/"",
+                  /*ignore_errors_and_delete_documents=*/true)),
+              IsOkAndHolds(EqualsSetSchemaResult(result)));
+  EXPECT_THAT(schema_store->GetSchema(),
+              IsOkAndHolds(Pointee(EqualsProto(schema))));
+
+  // The empty database should be replaced by db1/ and db2/.
+  EXPECT_THAT(schema_store->GetSchema(""),
+              StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
+  EXPECT_THAT(
+      schema_store->GetSchema("db1/"),
+      IsOkAndHolds(EqualsProto(SchemaBuilder()
+                                   .AddType(SchemaTypeConfigBuilder()
+                                                .SetType("db1/email_v1")
+                                                .SetDatabase("db1/"))
+                                   .AddType(SchemaTypeConfigBuilder()
+                                                .SetType("db1/message_v1")
+                                                .SetDatabase("db1/"))
+                                   .Build())));
+  EXPECT_THAT(
+      schema_store->GetSchema("db2/"),
+      IsOkAndHolds(EqualsProto(SchemaBuilder()
+                                   .AddType(SchemaTypeConfigBuilder()
+                                                .SetType("db2/email_v2")
+                                                .SetDatabase("db2/"))
+                                   .AddType(SchemaTypeConfigBuilder()
+                                                .SetType("db2/message_v2")
+                                                .SetDatabase("db2/"))
+                                   .Build())));
+}
+
+TEST_F(SchemaStoreTest, SetSchemaWithUnpopulatedDatabaseResetsEntireSchema) {
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<SchemaStore> schema_store,
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_,
+                          feature_flags_.get()));
+
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("email"))
+          .AddType(SchemaTypeConfigBuilder().SetType("message"))
+          .Build();
+  SchemaStore::SetSchemaResult result;
+  result.success = true;
+  result.schema_types_new_by_name.insert("email");
+  result.schema_types_new_by_name.insert("message");
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  schema, /*database=*/"",
+                  /*ignore_errors_and_delete_documents=*/false)),
+              IsOkAndHolds(EqualsSetSchemaResult(result)));
+  EXPECT_THAT(schema_store->GetSchema(),
+              IsOkAndHolds(Pointee(EqualsProto(schema))));
+  EXPECT_THAT(schema_store->GetSchema(""), IsOkAndHolds(EqualsProto(schema)));
+
+  schema = SchemaBuilder()
+               .AddType(SchemaTypeConfigBuilder()
+                            .SetType("db1/email_v1")
+                            .SetDatabase("db1/"))
+               .AddType(SchemaTypeConfigBuilder()
+                            .SetType("db1/message_v1")
+                            .SetDatabase("db1/"))
+               .AddType(SchemaTypeConfigBuilder()
+                            .SetType("db2/email_v2")
+                            .SetDatabase("db2/"))
+               .AddType(SchemaTypeConfigBuilder()
+                            .SetType("db2/message_v2")
+                            .SetDatabase("db2/"))
+               .Build();
+  result = SchemaStore::SetSchemaResult();
+  result.success = true;
+  result.schema_types_new_by_name.insert("db1/email_v1");
+  result.schema_types_new_by_name.insert("db1/message_v1");
+  result.schema_types_new_by_name.insert("db2/email_v2");
+  result.schema_types_new_by_name.insert("db2/message_v2");
+  result.schema_types_deleted_by_name.insert("email");
+  result.schema_types_deleted_by_name.insert("message");
+  result.schema_types_deleted_by_id.insert(0);
+  result.schema_types_deleted_by_id.insert(1);
+
+  // Reset the schema without populating the database field in the
+  // SetSchemaRequestProto. This should ignore the database field in the schema
+  // type configs and reset the entire schema.
+  SetSchemaRequestProto set_schema_request;
+  *set_schema_request.mutable_schema() = schema;
+  set_schema_request.clear_database();
+  set_schema_request.set_ignore_errors_and_delete_documents(
+      /*ignore_errors_and_delete_documents=*/true);
+
+  EXPECT_THAT(schema_store->SetSchema(std::move(set_schema_request)),
+              IsOkAndHolds(EqualsSetSchemaResult(result)));
+  EXPECT_THAT(schema_store->GetSchema(),
+              IsOkAndHolds(Pointee(EqualsProto(schema))));
+
+  // The empty database should be replaced by db1/ and db2/.
+  EXPECT_THAT(schema_store->GetSchema(""),
+              StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
+  EXPECT_THAT(
+      schema_store->GetSchema("db1/"),
+      IsOkAndHolds(EqualsProto(SchemaBuilder()
+                                   .AddType(SchemaTypeConfigBuilder()
+                                                .SetType("db1/email_v1")
+                                                .SetDatabase("db1/"))
+                                   .AddType(SchemaTypeConfigBuilder()
+                                                .SetType("db1/message_v1")
+                                                .SetDatabase("db1/"))
+                                   .Build())));
+  EXPECT_THAT(
+      schema_store->GetSchema("db2/"),
+      IsOkAndHolds(EqualsProto(SchemaBuilder()
+                                   .AddType(SchemaTypeConfigBuilder()
+                                                .SetType("db2/email_v2")
+                                                .SetDatabase("db2/"))
+                                   .AddType(SchemaTypeConfigBuilder()
+                                                .SetType("db2/message_v2")
+                                                .SetDatabase("db2/"))
+                                   .Build())));
+}
+
 TEST_F(SchemaStoreTest, SetSameSchemaOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
@@ -767,49 +938,50 @@
                           /*initialize_stats=*/nullptr));
 
   // Set schema for the first time
-  SchemaProto db1_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db1_email").SetDatabase("db1"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db1_message")
-                       .SetDatabase("db1"))
-          .Build();
-  SchemaProto db2_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db2_email").SetDatabase("db2"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_message")
-                       .SetDatabase("db2"))
-          .Build();
-  SchemaProto expected_full_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db1_email").SetDatabase("db1"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db1_message")
-                       .SetDatabase("db1"))
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db2_email").SetDatabase("db2"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_message")
-                       .SetDatabase("db2"))
-          .Build();
+  SchemaProto db1_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db1/email")
+                                            .SetDatabase("db1/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db1/message")
+                                            .SetDatabase("db1/"))
+                               .Build();
+  SchemaProto db2_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db2/email")
+                                            .SetDatabase("db2/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db2/message")
+                                            .SetDatabase("db2/"))
+                               .Build();
+  SchemaProto expected_full_schema = SchemaBuilder()
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db1/email")
+                                                      .SetDatabase("db1/"))
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db1/message")
+                                                      .SetDatabase("db1/"))
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db2/email")
+                                                      .SetDatabase("db2/"))
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db2/message")
+                                                      .SetDatabase("db2/"))
+                                         .Build();
   SchemaStore::SetSchemaResult result;
   result.success = true;
-  result.schema_types_new_by_name.insert("db1_email");
-  result.schema_types_new_by_name.insert("db1_message");
+  result.schema_types_new_by_name.insert("db1/email");
+  result.schema_types_new_by_name.insert("db1/message");
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  db1_schema, /*database=*/"db1",
+                  db1_schema, /*database=*/"db1/",
                   /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   result = SchemaStore::SetSchemaResult();
   result.success = true;
-  result.schema_types_new_by_name.insert("db2_email");
-  result.schema_types_new_by_name.insert("db2_message");
+  result.schema_types_new_by_name.insert("db2/email");
+  result.schema_types_new_by_name.insert("db2/message");
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  db2_schema, /*database=*/"db2",
+                  db2_schema, /*database=*/"db2/",
                   /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   ICING_ASSERT_OK_AND_ASSIGN(const SchemaProto* actual_full_schema,
@@ -820,16 +992,16 @@
   result = SchemaStore::SetSchemaResult();
   result.success = true;
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  db1_schema, /*database=*/"db1",
+                  db1_schema, /*database=*/"db1/",
                   /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
 
   // Check the schema, this should not have changed
   EXPECT_THAT(schema_store->GetSchema(),
               IsOkAndHolds(Pointee(EqualsProto(expected_full_schema))));
-  EXPECT_THAT(schema_store->GetSchema("db1"),
+  EXPECT_THAT(schema_store->GetSchema("db1/"),
               IsOkAndHolds(EqualsProto(db1_schema)));
-  EXPECT_THAT(schema_store->GetSchema("db2"),
+  EXPECT_THAT(schema_store->GetSchema("db2/"),
               IsOkAndHolds(EqualsProto(db2_schema)));
 }
 
@@ -841,74 +1013,76 @@
                           /*initialize_stats=*/nullptr));
 
   // Set schema for the first time
-  SchemaProto db1_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db1_email").SetDatabase("db1"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db1_message")
-                       .SetDatabase("db1"))
-          .Build();
-  SchemaProto db2_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db2_email").SetDatabase("db2"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_message")
-                       .SetDatabase("db2"))
-          .Build();
-  SchemaProto db3_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db3_email").SetDatabase("db3"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db3_message")
-                       .SetDatabase("db3"))
-          .Build();
-  SchemaProto expected_full_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db1_email").SetDatabase("db1"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db1_message")
-                       .SetDatabase("db1"))
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db2_email").SetDatabase("db2"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_message")
-                       .SetDatabase("db2"))
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db3_email").SetDatabase("db3"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db3_message")
-                       .SetDatabase("db3"))
-          .Build();
+  SchemaProto db1_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db1/email")
+                                            .SetDatabase("db1/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db1/message")
+                                            .SetDatabase("db1/"))
+                               .Build();
+  SchemaProto db2_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db2/email")
+                                            .SetDatabase("db2/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db2/message")
+                                            .SetDatabase("db2/"))
+                               .Build();
+  SchemaProto db3_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db3/email")
+                                            .SetDatabase("db3/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db3/message")
+                                            .SetDatabase("db3/"))
+                               .Build();
+  SchemaProto expected_full_schema = SchemaBuilder()
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db1/email")
+                                                      .SetDatabase("db1/"))
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db1/message")
+                                                      .SetDatabase("db1/"))
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db2/email")
+                                                      .SetDatabase("db2/"))
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db2/message")
+                                                      .SetDatabase("db2/"))
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db3/email")
+                                                      .SetDatabase("db3/"))
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db3/message")
+                                                      .SetDatabase("db3/"))
+                                         .Build();
 
   // Set schema for db1
   SchemaStore::SetSchemaResult result;
   result.success = true;
-  result.schema_types_new_by_name.insert("db1_email");
-  result.schema_types_new_by_name.insert("db1_message");
+  result.schema_types_new_by_name.insert("db1/email");
+  result.schema_types_new_by_name.insert("db1/message");
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  db1_schema, /*database=*/"db1",
+                  db1_schema, /*database=*/"db1/",
                   /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   // Set schema for db2
   result = SchemaStore::SetSchemaResult();
   result.success = true;
-  result.schema_types_new_by_name.insert("db2_email");
-  result.schema_types_new_by_name.insert("db2_message");
+  result.schema_types_new_by_name.insert("db2/email");
+  result.schema_types_new_by_name.insert("db2/message");
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  db2_schema, /*database=*/"db2",
+                  db2_schema, /*database=*/"db2/",
                   /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   // Set schema for db3
   result = SchemaStore::SetSchemaResult();
   result.success = true;
-  result.schema_types_new_by_name.insert("db3_email");
-  result.schema_types_new_by_name.insert("db3_message");
+  result.schema_types_new_by_name.insert("db3/email");
+  result.schema_types_new_by_name.insert("db3/message");
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  db3_schema, /*database=*/"db3",
+                  db3_schema, /*database=*/"db3/",
                   /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   // Verify schema.
@@ -918,20 +1092,20 @@
 
   // Reset db2 with the types reordered. This should not change the existing
   // schema in any way.
-  SchemaProto reordered_db2_schema =
-      SchemaBuilder()
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_message")
-                       .SetDatabase("db2"))
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db2_email").SetDatabase("db2"))
-          .Build();
+  SchemaProto reordered_db2_schema = SchemaBuilder()
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db2/message")
+                                                      .SetDatabase("db2/"))
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db2/email")
+                                                      .SetDatabase("db2/"))
+                                         .Build();
   result = SchemaStore::SetSchemaResult();
   result.success = true;
 
   libtextclassifier3::StatusOr<SchemaStore::SetSchemaResult> actual_result =
       schema_store->SetSchema(CreateSetSchemaRequestProto(
-          reordered_db2_schema, /*database=*/"db2",
+          reordered_db2_schema, /*database=*/"db2/",
           /*ignore_errors_and_delete_documents=*/false));
   EXPECT_THAT(actual_result, IsOkAndHolds(EqualsSetSchemaResult(result)));
   EXPECT_THAT(actual_result.ValueOrDie().old_schema_type_ids_changed,
@@ -940,20 +1114,23 @@
   // Check the schema
   EXPECT_THAT(schema_store->GetSchema(),
               IsOkAndHolds(Pointee(EqualsProto(expected_full_schema))));
-  EXPECT_THAT(schema_store->GetSchema("db1"),
+  EXPECT_THAT(schema_store->GetSchema("db1/"),
               IsOkAndHolds(EqualsProto(db1_schema)));
 
   libtextclassifier3::StatusOr<SchemaProto> actual_db2_schema =
-      schema_store->GetSchema("db2");
+      schema_store->GetSchema("db2/");
   EXPECT_THAT(actual_db2_schema, IsOkAndHolds(EqualsProto(db2_schema)));
   EXPECT_THAT(actual_db2_schema.ValueOrDie(),
               Not(EqualsProto(reordered_db2_schema)));
 
-  EXPECT_THAT(schema_store->GetSchema("db3"),
+  EXPECT_THAT(schema_store->GetSchema("db3/"),
               IsOkAndHolds(EqualsProto(db3_schema)));
 }
 
-TEST_F(SchemaStoreTest, SetDatabaseAddedTypesPreservesSchemaTypeIds) {
+TEST_F(SchemaStoreTest, SetDatabaseAddedTypesPreserveSchemaTypeIds) {
+  if (!feature_flags_->enable_schema_type_id_optimization()) {
+    GTEST_SKIP() << "Test for schema type id optimization.";
+  }
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
       SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_,
@@ -961,74 +1138,228 @@
                           /*initialize_stats=*/nullptr));
 
   // Set schema for the first time
-  SchemaProto db1_schema =
+  SchemaProto db1_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db1/email")
+                                            .SetDatabase("db1/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db1/message")
+                                            .SetDatabase("db1/"))
+                               .Build();
+  SchemaProto db2_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db2/email")
+                                            .SetDatabase("db2/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db2/message")
+                                            .SetDatabase("db2/"))
+                               .Build();
+  SchemaProto db3_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db3/email")
+                                            .SetDatabase("db3/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db3/message")
+                                            .SetDatabase("db3/"))
+                               .Build();
+  SchemaProto expected_full_schema = SchemaBuilder()
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db1/email")
+                                                      .SetDatabase("db1/"))
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db1/message")
+                                                      .SetDatabase("db1/"))
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db2/email")
+                                                      .SetDatabase("db2/"))
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db2/message")
+                                                      .SetDatabase("db2/"))
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db3/email")
+                                                      .SetDatabase("db3/"))
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db3/message")
+                                                      .SetDatabase("db3/"))
+                                         .Build();
+
+  // Set schema for db1
+  SchemaStore::SetSchemaResult expected_result;
+  expected_result.success = true;
+  expected_result.schema_types_new_by_name.insert("db1/email");
+  expected_result.schema_types_new_by_name.insert("db1/message");
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db1_schema, /*database=*/"db1/",
+                  /*ignore_errors_and_delete_documents=*/false)),
+              IsOkAndHolds(EqualsSetSchemaResult(expected_result)));
+  // Set schema for db2
+  expected_result = SchemaStore::SetSchemaResult();
+  expected_result.success = true;
+  expected_result.schema_types_new_by_name.insert("db2/email");
+  expected_result.schema_types_new_by_name.insert("db2/message");
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db2_schema, /*database=*/"db2/",
+                  /*ignore_errors_and_delete_documents=*/false)),
+              IsOkAndHolds(EqualsSetSchemaResult(expected_result)));
+  // Set schema for db3
+  expected_result = SchemaStore::SetSchemaResult();
+  expected_result.success = true;
+  expected_result.schema_types_new_by_name.insert("db3/email");
+  expected_result.schema_types_new_by_name.insert("db3/message");
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db3_schema, /*database=*/"db3/",
+                  /*ignore_errors_and_delete_documents=*/false)),
+              IsOkAndHolds(EqualsSetSchemaResult(expected_result)));
+  // Verify schema.
+  ICING_ASSERT_OK_AND_ASSIGN(const SchemaProto* actual_full_schema,
+                             schema_store->GetSchema());
+  EXPECT_THAT(*actual_full_schema, EqualsProto(expected_full_schema));
+
+  // Reset db2 and add a type. The added type should be appended to the end of
+  // the SchemaProto, and SchemaTypeIds for existing types should not change.
+  db2_schema = SchemaBuilder()
+                   .AddType(SchemaTypeConfigBuilder()
+                                .SetType("db2/recipient")
+                                .SetDatabase("db2/"))
+                   .AddType(SchemaTypeConfigBuilder()
+                                .SetType("db2/email")
+                                .SetDatabase("db2/"))
+                   .AddType(SchemaTypeConfigBuilder()
+                                .SetType("db2/message")
+                                .SetDatabase("db2/"))
+                   .Build();
+  expected_full_schema =
       SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db1_email").SetDatabase("db1"))
+          .AddType(SchemaTypeConfigBuilder()  // db1 types
+                       .SetType("db1/email")
+                       .SetDatabase("db1/"))
           .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db1_message")
-                       .SetDatabase("db1"))
+                       .SetType("db1/message")
+                       .SetDatabase("db1/"))
+          .AddType(SchemaTypeConfigBuilder()  // db2 types
+                       .SetType("db2/email")
+                       .SetDatabase("db2/"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db2/message")
+                       .SetDatabase("db2/"))
+          .AddType(SchemaTypeConfigBuilder()  // db3 types
+                       .SetType("db3/email")
+                       .SetDatabase("db3/"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db3/message")
+                       .SetDatabase("db3/"))
+          .AddType(SchemaTypeConfigBuilder()  // Additional db2 type is appended
+                                              // at the end
+                       .SetType("db2/recipient")
+                       .SetDatabase("db2/"))
           .Build();
-  SchemaProto db2_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db2_email").SetDatabase("db2"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_message")
-                       .SetDatabase("db2"))
-          .Build();
-  SchemaProto db3_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db3_email").SetDatabase("db3"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db3_message")
-                       .SetDatabase("db3"))
-          .Build();
-  SchemaProto expected_full_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db1_email").SetDatabase("db1"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db1_message")
-                       .SetDatabase("db1"))
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db2_email").SetDatabase("db2"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_message")
-                       .SetDatabase("db2"))
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db3_email").SetDatabase("db3"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db3_message")
-                       .SetDatabase("db3"))
-          .Build();
+  expected_result = SchemaStore::SetSchemaResult();
+  expected_result.success = true;
+  expected_result.schema_types_new_by_name.insert("db2/recipient");
+  libtextclassifier3::StatusOr<SchemaStore::SetSchemaResult> actual_result =
+      schema_store->SetSchema(CreateSetSchemaRequestProto(
+          db2_schema, /*database=*/"db2/",
+          /*ignore_errors_and_delete_documents=*/false));
+  EXPECT_THAT(actual_result,
+              IsOkAndHolds(EqualsSetSchemaResult(expected_result)));
+  EXPECT_THAT(actual_result.ValueOrDie().old_schema_type_ids_changed,
+              IsEmpty());
+
+  // Check the schema
+  EXPECT_THAT(schema_store->GetSchema(),
+              IsOkAndHolds(Pointee(EqualsProto(expected_full_schema))));
+  EXPECT_THAT(schema_store->GetSchema("db1/"),
+              IsOkAndHolds(EqualsProto(db1_schema)));
+  EXPECT_THAT(
+      schema_store->GetSchema("db2/"),
+      IsOkAndHolds(EqualsProto(SchemaBuilder()
+                                   .AddType(SchemaTypeConfigBuilder()
+                                                .SetType("db2/email")
+                                                .SetDatabase("db2/"))
+                                   .AddType(SchemaTypeConfigBuilder()
+                                                .SetType("db2/message")
+                                                .SetDatabase("db2/"))
+                                   .AddType(SchemaTypeConfigBuilder()
+                                                .SetType("db2/recipient")
+                                                .SetDatabase("db2/"))
+                                   .Build())));
+  EXPECT_THAT(schema_store->GetSchema("db3/"),
+              IsOkAndHolds(EqualsProto(db3_schema)));
+}
+
+TEST_F(SchemaStoreTest, SetDatabaseReorderedSchemaPreservesSchemaTypeIds) {
+  if (!feature_flags_->enable_schema_database()) {
+    GTEST_SKIP() << "Test for schema database id assignment.";
+  }
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<SchemaStore> schema_store,
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_,
+                          feature_flags_.get(),
+                          /*initialize_stats=*/nullptr));
+
+  // Set schema for the first time
+  SchemaProto db1_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db1/email")
+                                            .SetDatabase("db1/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db1/message")
+                                            .SetDatabase("db1/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db1/recipient")
+                                            .SetDatabase("db1/"))
+                               .Build();
+  SchemaProto db2_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db2/email")
+                                            .SetDatabase("db2/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db2/message")
+                                            .SetDatabase("db2/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db2/recipient")
+                                            .SetDatabase("db2/"))
+                               .Build();
+  SchemaProto expected_full_schema = SchemaBuilder()
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db1/email")
+                                                      .SetDatabase("db1/"))
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db1/message")
+                                                      .SetDatabase("db1/"))
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db1/recipient")
+                                                      .SetDatabase("db1/"))
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db2/email")
+                                                      .SetDatabase("db2/"))
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db2/message")
+                                                      .SetDatabase("db2/"))
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db2/recipient")
+                                                      .SetDatabase("db2/"))
+                                         .Build();
 
   // Set schema for db1
   SchemaStore::SetSchemaResult result;
   result.success = true;
-  result.schema_types_new_by_name.insert("db1_email");
-  result.schema_types_new_by_name.insert("db1_message");
+  result.schema_types_new_by_name.insert("db1/email");
+  result.schema_types_new_by_name.insert("db1/message");
+  result.schema_types_new_by_name.insert("db1/recipient");
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  db1_schema, /*database=*/"db1",
+                  db1_schema, /*database=*/"db1/",
                   /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   // Set schema for db2
   result = SchemaStore::SetSchemaResult();
   result.success = true;
-  result.schema_types_new_by_name.insert("db2_email");
-  result.schema_types_new_by_name.insert("db2_message");
+  result.schema_types_new_by_name.insert("db2/email");
+  result.schema_types_new_by_name.insert("db2/message");
+  result.schema_types_new_by_name.insert("db2/recipient");
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  db2_schema, /*database=*/"db2",
-                  /*ignore_errors_and_delete_documents=*/false)),
-              IsOkAndHolds(EqualsSetSchemaResult(result)));
-  // Set schema for db3
-  result = SchemaStore::SetSchemaResult();
-  result.success = true;
-  result.schema_types_new_by_name.insert("db3_email");
-  result.schema_types_new_by_name.insert("db3_message");
-  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  db3_schema, /*database=*/"db3",
+                  db2_schema, /*database=*/"db2/",
                   /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   // Verify schema.
@@ -1036,65 +1367,114 @@
                              schema_store->GetSchema());
   EXPECT_THAT(*actual_full_schema, EqualsProto(expected_full_schema));
 
-  // Reset db2 and add a type. The added type should be appended to the end of
-  // the SchemaProto, and SchemaTypeIds for db1 and db3 should not change.
-  //
-  // Whether or not the SchemaTypeIds for db2 change depends on the order in the
-  // new db2 SchemaProto (in this case, existing type's order and ids do not
-  // change)
-  db2_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db2_email").SetDatabase("db2"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_message")
-                       .SetDatabase("db2"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_recipient")
-                       .SetDatabase("db2"))
-          .Build();
-  expected_full_schema =
-      SchemaBuilder()
-          .AddType(SchemaTypeConfigBuilder()  // db1 types
-                       .SetType("db1_email")
-                       .SetDatabase("db1"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db1_message")
-                       .SetDatabase("db1"))
-          .AddType(SchemaTypeConfigBuilder()  // db2 types
-                       .SetType("db2_email")
-                       .SetDatabase("db2"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_message")
-                       .SetDatabase("db2"))
-          .AddType(SchemaTypeConfigBuilder()  // db3 types
-                       .SetType("db3_email")
-                       .SetDatabase("db3"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db3_message")
-                       .SetDatabase("db3"))
-          .AddType(SchemaTypeConfigBuilder()  // Additional db2 type is appended
-                                              // at the end
-                       .SetType("db2_recipient")
-                       .SetDatabase("db2"))
-          .Build();
+  // Reset db1 with the types reordered. SchemaTypeIds should not change.
+  SchemaProto new_db1_schema = SchemaBuilder()
+                                   .AddType(SchemaTypeConfigBuilder()
+                                                .SetType("db1/email")
+                                                .SetDatabase("db1/"))
+                                   .AddType(SchemaTypeConfigBuilder()
+                                                .SetType("db1/recipient")
+                                                .SetDatabase("db1/"))
+                                   .AddType(SchemaTypeConfigBuilder()
+                                                .SetType("db1/message")
+                                                .SetDatabase("db1/"))
+                                   .Build();
   result = SchemaStore::SetSchemaResult();
   result.success = true;
-  result.schema_types_new_by_name.insert("db2_recipient");
-  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  db2_schema, /*database=*/"db2",
-                  /*ignore_errors_and_delete_documents=*/false)),
-              IsOkAndHolds(EqualsSetSchemaResult(result)));
+  libtextclassifier3::StatusOr<SchemaStore::SetSchemaResult> actual_result =
+      schema_store->SetSchema(CreateSetSchemaRequestProto(
+          db2_schema, /*database=*/"db2/",
+          /*ignore_errors_and_delete_documents=*/false));
+  EXPECT_THAT(actual_result, IsOkAndHolds(EqualsSetSchemaResult(result)));
+  EXPECT_THAT(actual_result.ValueOrDie().old_schema_type_ids_changed,
+              IsEmpty());
 
   // Check the schema
   EXPECT_THAT(schema_store->GetSchema(),
               IsOkAndHolds(Pointee(EqualsProto(expected_full_schema))));
-  EXPECT_THAT(schema_store->GetSchema("db1"),
+  EXPECT_THAT(schema_store->GetSchema("db1/"),
               IsOkAndHolds(EqualsProto(db1_schema)));
-  EXPECT_THAT(schema_store->GetSchema("db2"),
+  EXPECT_THAT(schema_store->GetSchema("db2/"),
               IsOkAndHolds(EqualsProto(db2_schema)));
-  EXPECT_THAT(schema_store->GetSchema("db3"),
-              IsOkAndHolds(EqualsProto(db3_schema)));
+}
+
+TEST_F(SchemaStoreTest,
+       SetDatabaseReorderedSchemaPreservesSchemaTypeIds_largeSchema) {
+  if (!feature_flags_->enable_schema_database()) {
+    GTEST_SKIP() << "Test for schema database id assignment.";
+  }
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<SchemaStore> schema_store,
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_,
+                          feature_flags_.get(),
+                          /*initialize_stats=*/nullptr));
+
+  // 1. Create five databases, each with five types.
+  constexpr int kNumDbs = 5;
+  constexpr int kNumTypesPerDb = 5;
+  SchemaBuilder expected_full_schema_builder;
+
+  for (int i = 0; i < kNumDbs; ++i) {
+    std::string db_name = absl_ports::StrCat("db", std::to_string(i), "/");
+    SchemaBuilder db_schema_builder;
+    SchemaStore::SetSchemaResult expected_result;
+
+    for (int j = 0; j < kNumTypesPerDb; ++j) {
+      std::string type_name =
+          absl_ports::StrCat(db_name, "type", std::to_string(j));
+      db_schema_builder.AddType(
+          SchemaTypeConfigBuilder().SetType(type_name).SetDatabase(db_name));
+      expected_full_schema_builder.AddType(
+          SchemaTypeConfigBuilder().SetType(type_name).SetDatabase(db_name));
+      expected_result.schema_types_new_by_name.insert(type_name);
+    }
+    SchemaProto db_schema = db_schema_builder.Build();
+    expected_result.success = true;
+    EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                    db_schema, db_name,
+                    /*ignore_errors_and_delete_documents=*/false)),
+                IsOkAndHolds(EqualsSetSchemaResult(expected_result)));
+  }
+  SchemaProto expected_full_schema = expected_full_schema_builder.Build();
+
+  // Verify schema.
+  ICING_ASSERT_OK_AND_ASSIGN(const SchemaProto* actual_full_schema,
+                             schema_store->GetSchema());
+  EXPECT_THAT(*actual_full_schema, EqualsProto(expected_full_schema));
+
+  // 2. In a loop, change the ordering of the types in one of the dbs, call
+  // setSchema and verify that type ids remain the same.
+  for (int i = 0; i < kNumDbs; ++i) {
+    std::string db_name = absl_ports::StrCat("db", std::to_string(i), "/");
+    ICING_ASSERT_OK_AND_ASSIGN(SchemaProto db_schema,
+                               schema_store->GetSchema(db_name));
+
+    // Reset the schema with the types reordered.
+    SchemaProto db_schema_reversed = db_schema;
+    std::reverse(db_schema_reversed.mutable_types()->begin(),
+                 db_schema_reversed.mutable_types()->end());
+
+    SchemaStore::SetSchemaResult expected_result =
+        SchemaStore::SetSchemaResult();
+    expected_result.success = true;
+    libtextclassifier3::StatusOr<SchemaStore::SetSchemaResult> actual_result =
+        schema_store->SetSchema(CreateSetSchemaRequestProto(
+            db_schema_reversed, db_name,
+            /*ignore_errors_and_delete_documents=*/false));
+    EXPECT_THAT(actual_result,
+                IsOkAndHolds(EqualsSetSchemaResult(expected_result)));
+
+    // Verify that no type ids changed.
+    EXPECT_THAT(actual_result.ValueOrDie().old_schema_type_ids_changed,
+                IsEmpty());
+
+    // Verify schema
+    EXPECT_THAT(schema_store->GetSchema(db_name),
+                IsOkAndHolds(EqualsProto(db_schema)));
+    EXPECT_THAT(schema_store->GetSchema(),
+                IsOkAndHolds(Pointee(EqualsProto(expected_full_schema))));
+  }
 }
 
 TEST_F(SchemaStoreTest, SetDatabaseDeletedTypesOk) {
@@ -1105,157 +1485,917 @@
                           /*initialize_stats=*/nullptr));
 
   // Set schema for the first time
-  SchemaProto db1_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db1_email").SetDatabase("db1"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db1_message")
-                       .SetDatabase("db1"))
-          .Build();
-  SchemaProto db2_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db2_email").SetDatabase("db2"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_message")
-                       .SetDatabase("db2"))
-          .Build();
-  SchemaProto db3_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db3_email").SetDatabase("db3"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db3_message")
-                       .SetDatabase("db3"))
-          .Build();
-
+  SchemaProto db1_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db1/email")
+                                            .SetDatabase("db1/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db1/message")
+                                            .SetDatabase("db1/"))
+                               .Build();
+  SchemaProto db2_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db2/email")
+                                            .SetDatabase("db2/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db2/message")
+                                            .SetDatabase("db2/"))
+                               .Build();
+  SchemaProto db3_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db3/email")
+                                            .SetDatabase("db3/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db3/message")
+                                            .SetDatabase("db3/"))
+                               .Build();
   // Set schema for db1
   SchemaStore::SetSchemaResult result;
   result.success = true;
-  result.schema_types_new_by_name.insert("db1_email");
-  result.schema_types_new_by_name.insert("db1_message");
+  result.schema_types_new_by_name.insert("db1/email");
+  result.schema_types_new_by_name.insert("db1/message");
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  db1_schema, /*database=*/"db1",
+                  db1_schema, /*database=*/"db1/",
                   /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   // Set schema for db2
   result = SchemaStore::SetSchemaResult();
   result.success = true;
-  result.schema_types_new_by_name.insert("db2_email");
-  result.schema_types_new_by_name.insert("db2_message");
+  result.schema_types_new_by_name.insert("db2/email");
+  result.schema_types_new_by_name.insert("db2/message");
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  db2_schema, /*database=*/"db2",
+                  db2_schema, /*database=*/"db2/",
                   /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   // Set schema for db3
   result = SchemaStore::SetSchemaResult();
   result.success = true;
-  result.schema_types_new_by_name.insert("db3_email");
-  result.schema_types_new_by_name.insert("db3_message");
+  result.schema_types_new_by_name.insert("db3/email");
+  result.schema_types_new_by_name.insert("db3/message");
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  db3_schema, /*database=*/"db3",
+                  db3_schema, /*database=*/"db3/",
                   /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
-  // Set schema again for db2 and add a type. The added type should be appended
-  // to the end of the SchemaProto.
-  db2_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db2_email").SetDatabase("db2"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_message")
-                       .SetDatabase("db2"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_recipient")
-                       .SetDatabase("db2"))
-          .Build();
-  result = SchemaStore::SetSchemaResult();
-  result.success = true;
-  result.schema_types_new_by_name.insert("db2_recipient");
-  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  db2_schema, /*database=*/"db2",
-                  /*ignore_errors_and_delete_documents=*/false)),
-              IsOkAndHolds(EqualsSetSchemaResult(result)));
+
   SchemaProto expected_full_schema =
       SchemaBuilder()
           .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db1_email")  // SchemaTypeId 0
-                       .SetDatabase("db1"))
+                       .SetType("db1/email")  // SchemaTypeId 0
+                       .SetDatabase("db1/"))
           .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db1_message")  // SchemaTypeId 1
-                       .SetDatabase("db1"))
+                       .SetType("db1/message")  // SchemaTypeId 1
+                       .SetDatabase("db1/"))
           .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_email")  // SchemaTypeId 2
-                       .SetDatabase("db2"))
+                       .SetType("db2/email")  // SchemaTypeId 2
+                       .SetDatabase("db2/"))
           .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_message")  // SchemaTypeId 3
-                       .SetDatabase("db2"))
+                       .SetType("db2/message")  // SchemaTypeId 3
+                       .SetDatabase("db2/"))
           .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db3_email")  // SchemaTypeId 4
-                       .SetDatabase("db3"))
+                       .SetType("db3/email")  // SchemaTypeId 4
+                       .SetDatabase("db3/"))
           .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db3_message")  // SchemaTypeId 5
-                       .SetDatabase("db3"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_recipient")  // SchemaTypeId 6
-                       .SetDatabase("db2"))
+                       .SetType("db3/message")  // SchemaTypeId 5
+                       .SetDatabase("db3/"))
           .Build();
   // Verify schema.
   ICING_ASSERT_OK_AND_ASSIGN(const SchemaProto* actual_full_schema,
                              schema_store->GetSchema());
   EXPECT_THAT(*actual_full_schema, EqualsProto(expected_full_schema));
 
-  // Reset db2 and delete some types. All types that were originally added after
-  // db2 should have their type ids changed.
+  // Reset db2 and delete some types. With the type id optimization, only the
+  // last few type ids should be changed.
   db2_schema = SchemaBuilder()
                    .AddType(SchemaTypeConfigBuilder()
-                                .SetType("db2_message")
-                                .SetDatabase("db2"))
+                                .SetType("db2/message")
+                                .SetDatabase("db2/"))
                    .Build();
-  expected_full_schema =
+  if (feature_flags_->enable_schema_type_id_optimization()) {
+    expected_full_schema =
+        SchemaBuilder()
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db1/email")  // SchemaTypeId 0
+                         .SetDatabase("db1/"))
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db1/message")  // SchemaTypeId 1
+                         .SetDatabase("db1/"))
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db3/message")  // SchemaTypeId 5 -> 2
+                         .SetDatabase("db3/"))
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db2/message")  // SchemaTypeId 3 (unchanged)
+                         .SetDatabase("db2/"))
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db3/email")  // SchemaTypeId 4 (unchanged)
+                         .SetDatabase("db3/"))
+            .Build();
+    result = SchemaStore::SetSchemaResult();
+    result.success = true;
+    result.schema_types_deleted_by_name.insert("db2/email");
+    result.schema_types_deleted_by_id.insert(2);   // db2_email
+    result.old_schema_type_ids_changed.insert(5);  // db3_message
+    EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                    db2_schema, /*database=*/"db2/",
+                    /*ignore_errors_and_delete_documents=*/true)),
+                IsOkAndHolds(EqualsSetSchemaResult(result)));
+    // Check the schema
+    EXPECT_THAT(schema_store->GetSchema(),
+                IsOkAndHolds(Pointee(EqualsProto(expected_full_schema))));
+    EXPECT_THAT(schema_store->GetSchema("db1/"),
+                IsOkAndHolds(EqualsProto(db1_schema)));
+    EXPECT_THAT(schema_store->GetSchema("db2/"),
+                IsOkAndHolds(EqualsProto(db2_schema)));
+    EXPECT_THAT(
+        schema_store->GetSchema("db3/"),
+        IsOkAndHolds(EqualsProto(SchemaBuilder()
+                                     .AddType(SchemaTypeConfigBuilder()
+                                                  .SetType("db3/message")
+                                                  .SetDatabase("db3/"))
+                                     .AddType(SchemaTypeConfigBuilder()
+                                                  .SetType("db3/email")
+                                                  .SetDatabase("db3/"))
+                                     .Build())));
+  } else {
+    expected_full_schema =
+        SchemaBuilder()
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db1/email")  // SchemaTypeId 0
+                         .SetDatabase("db1/"))
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db1/message")  // SchemaTypeId 1
+                         .SetDatabase("db1/"))
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db2/message")  // SchemaTypeId 3 -> 2
+                         .SetDatabase("db2/"))
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db3/email")  // SchemaTypeId 4 -> 3
+                         .SetDatabase("db3/"))
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db3/message")  // SchemaTypeId 5 -> 4
+                         .SetDatabase("db3/"))
+            .Build();
+    result = SchemaStore::SetSchemaResult();
+    result.success = true;
+    result.schema_types_deleted_by_name.insert("db2/email");
+    result.schema_types_deleted_by_id.insert(2);   // db2_email
+    result.old_schema_type_ids_changed.insert(3);  // db2_message
+    result.old_schema_type_ids_changed.insert(4);  // db3_email
+    result.old_schema_type_ids_changed.insert(5);  // db3_message
+    EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                    db2_schema, /*database=*/"db2/",
+                    /*ignore_errors_and_delete_documents=*/true)),
+                IsOkAndHolds(EqualsSetSchemaResult(result)));
+    // Check the schema
+    EXPECT_THAT(schema_store->GetSchema(),
+                IsOkAndHolds(Pointee(EqualsProto(expected_full_schema))));
+    EXPECT_THAT(schema_store->GetSchema("db1/"),
+                IsOkAndHolds(EqualsProto(db1_schema)));
+    EXPECT_THAT(schema_store->GetSchema("db2/"),
+                IsOkAndHolds(EqualsProto(db2_schema)));
+    EXPECT_THAT(schema_store->GetSchema("db3/"),
+                IsOkAndHolds(EqualsProto(db3_schema)));
+  }
+}
+
+TEST_F(SchemaStoreTest, SetDatabaseDeletedTypesWithLastTypeRemovedOk) {
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<SchemaStore> schema_store,
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_,
+                          feature_flags_.get(),
+                          /*initialize_stats=*/nullptr));
+
+  // Set schema for the first time
+  SchemaProto full_schema =
       SchemaBuilder()
           .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db1_email")  // SchemaTypeId 0
-                       .SetDatabase("db1"))
+                       .SetType("db1/email")  // SchemaTypeId 0
+                       .SetDatabase("db1/"))
           .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db1_message")  // SchemaTypeId 1
-                       .SetDatabase("db1"))
+                       .SetType("db1/message")  // SchemaTypeId 1
+                       .SetDatabase("db1/"))
           .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_message")  // SchemaTypeId 2
-                       .SetDatabase("db2"))
+                       .SetType("db2/email")  // SchemaTypeId 2
+                       .SetDatabase("db2/"))
           .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db3_email")  // SchemaTypeId 3
-                       .SetDatabase("db3"))
+                       .SetType("db2/message")  // SchemaTypeId 3
+                       .SetDatabase("db2/"))
           .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db3_message")  // SchemaTypeId 4
-                       .SetDatabase("db3"))
+                       .SetType("db3/email")  // SchemaTypeId 4
+                       .SetDatabase("db3/"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db3/message")  // SchemaTypeId 5
+                       .SetDatabase("db3/"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db2/recipient")  // SchemaTypeId 6
+                       .SetDatabase("db2/"))
           .Build();
+
+  // Set full schema using the empty db
+  EXPECT_TRUE(schema_store
+                  ->SetSchema(CreateSetSchemaRequestProto(
+                      full_schema, /*database=*/"",
+                      /*ignore_errors_and_delete_documents=*/false))
+                  .ok());
+  // Verify schema.
+  ICING_ASSERT_OK_AND_ASSIGN(const SchemaProto* actual_full_schema,
+                             schema_store->GetSchema());
+  EXPECT_THAT(*actual_full_schema, EqualsProto(full_schema));
+
+  // Reset db2 and delete some types. With the type id optimization, only the
+  // last few type ids should be changed.
+  SchemaProto db2_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db2/message")
+                                            .SetDatabase("db2/"))
+                               .Build();
+  if (feature_flags_->enable_schema_type_id_optimization()) {
+    SchemaProto expected_full_schema =
+        SchemaBuilder()
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db1/email")  // SchemaTypeId 0
+                         .SetDatabase("db1/"))
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db1/message")  // SchemaTypeId 1
+                         .SetDatabase("db1/"))
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db3/message")  // SchemaTypeId 5 -> 2
+                         .SetDatabase("db3/"))
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db2/message")  // SchemaTypeId 3
+                         .SetDatabase("db2/"))
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db3/email")  // SchemaTypeId 4
+                         .SetDatabase("db3/"))
+            .Build();
+    SchemaStore::SetSchemaResult result = SchemaStore::SetSchemaResult();
+    result.success = true;
+    result.schema_types_deleted_by_name.insert("db2/email");
+    result.schema_types_deleted_by_name.insert("db2/recipient");
+    result.schema_types_deleted_by_id.insert(
+        schema_store->GetSchemaTypeId("db2/email").ValueOrDie());  // 2
+    result.schema_types_deleted_by_id.insert(
+        schema_store->GetSchemaTypeId("db2/recipient").ValueOrDie());  // 6
+    // Only 1 type id is changed.
+    result.old_schema_type_ids_changed.insert(
+        schema_store->GetSchemaTypeId("db3/message").ValueOrDie());  // 5
+    EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                    db2_schema, /*database=*/"db2/",
+                    /*ignore_errors_and_delete_documents=*/true)),
+                IsOkAndHolds(EqualsSetSchemaResult(result)));
+
+    // Check the schema
+    EXPECT_THAT(schema_store->GetSchema(),
+                IsOkAndHolds(Pointee(EqualsProto(expected_full_schema))));
+    EXPECT_THAT(
+        schema_store->GetSchema("db1/"),
+        IsOkAndHolds(EqualsProto(SchemaBuilder()
+                                     .AddType(SchemaTypeConfigBuilder()
+                                                  .SetType("db1/email")
+                                                  .SetDatabase("db1/"))
+                                     .AddType(SchemaTypeConfigBuilder()
+                                                  .SetType("db1/message")
+                                                  .SetDatabase("db1/"))
+                                     .Build())));
+    EXPECT_THAT(
+        schema_store->GetSchema("db2/"),
+        IsOkAndHolds(EqualsProto(SchemaBuilder()
+                                     .AddType(SchemaTypeConfigBuilder()
+                                                  .SetType("db2/message")
+                                                  .SetDatabase("db2/"))
+                                     .Build())));
+    EXPECT_THAT(
+        schema_store->GetSchema("db3/"),
+        IsOkAndHolds(EqualsProto(SchemaBuilder()
+                                     .AddType(SchemaTypeConfigBuilder()
+                                                  .SetType("db3/message")
+                                                  .SetDatabase("db3/"))
+                                     .AddType(SchemaTypeConfigBuilder()
+                                                  .SetType("db3/email")
+                                                  .SetDatabase("db3/"))
+                                     .Build())));
+  } else {
+    SchemaProto expected_full_schema =
+        SchemaBuilder()
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db1/email")  // SchemaTypeId 0
+                         .SetDatabase("db1/"))
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db1/message")  // SchemaTypeId 1
+                         .SetDatabase("db1/"))
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db2/message")  // SchemaTypeId 3 -> 2
+                         .SetDatabase("db2/"))
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db3/email")  // SchemaTypeId 4 -> 3
+                         .SetDatabase("db3/"))
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db3/message")  // SchemaTypeId 5 -> 4
+                         .SetDatabase("db3/"))
+            .Build();
+    SchemaStore::SetSchemaResult result = SchemaStore::SetSchemaResult();
+    result.success = true;
+    result.success = true;
+    result.schema_types_deleted_by_name.insert("db2/email");
+    result.schema_types_deleted_by_name.insert("db2/recipient");
+    result.schema_types_deleted_by_id.insert(
+        schema_store->GetSchemaTypeId("db2/email").ValueOrDie());  // 2
+    result.schema_types_deleted_by_id.insert(
+        schema_store->GetSchemaTypeId("db2/recipient").ValueOrDie());  // 6
+    // 3 type ids are changed
+    result.old_schema_type_ids_changed.insert(
+        schema_store->GetSchemaTypeId("db2/message")
+            .ValueOrDie());  // db2_message
+    result.old_schema_type_ids_changed.insert(
+        schema_store->GetSchemaTypeId("db3/email").ValueOrDie());  // db3_email
+    result.old_schema_type_ids_changed.insert(
+        schema_store->GetSchemaTypeId("db3/message")
+            .ValueOrDie());  // db3_message
+
+    EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                    db2_schema, /*database=*/"db2/",
+                    /*ignore_errors_and_delete_documents=*/true)),
+                IsOkAndHolds(EqualsSetSchemaResult(result)));
+
+    // Check the schema
+    EXPECT_THAT(schema_store->GetSchema(),
+                IsOkAndHolds(Pointee(EqualsProto(expected_full_schema))));
+    EXPECT_THAT(
+        schema_store->GetSchema("db1/"),
+        IsOkAndHolds(EqualsProto(SchemaBuilder()
+                                     .AddType(SchemaTypeConfigBuilder()
+                                                  .SetType("db1/email")
+                                                  .SetDatabase("db1/"))
+                                     .AddType(SchemaTypeConfigBuilder()
+                                                  .SetType("db1/message")
+                                                  .SetDatabase("db1/"))
+                                     .Build())));
+    EXPECT_THAT(
+        schema_store->GetSchema("db2/"),
+        IsOkAndHolds(EqualsProto(SchemaBuilder()
+                                     .AddType(SchemaTypeConfigBuilder()
+                                                  .SetType("db2/message")
+                                                  .SetDatabase("db2/"))
+                                     .Build())));
+    EXPECT_THAT(
+        schema_store->GetSchema("db3/"),
+        IsOkAndHolds(EqualsProto(SchemaBuilder()
+                                     .AddType(SchemaTypeConfigBuilder()
+                                                  .SetType("db3/email")
+                                                  .SetDatabase("db3/"))
+                                     .AddType(SchemaTypeConfigBuilder()
+                                                  .SetType("db3/message")
+                                                  .SetDatabase("db3/"))
+                                     .Build())));
+  }
+}
+
+TEST_F(SchemaStoreTest, SetDatabaseDeletedTypesWithGapsOk) {
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<SchemaStore> schema_store,
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_,
+                          feature_flags_.get(),
+                          /*initialize_stats=*/nullptr));
+
+  // Set schema for the first time
+  SchemaProto expected_full_schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db1/type1")  // SchemaTypeId 0
+                       .SetDatabase("db1/"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db1/type2")  // SchemaTypeId 1
+                       .SetDatabase("db1/"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db2/type1")  // SchemaTypeId 2
+                       .SetDatabase("db2/"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db2/type2")  // SchemaTypeId 3
+                       .SetDatabase("db2/"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db2/type3")  // SchemaTypeId 4
+                       .SetDatabase("db2/"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db3/type1")  // SchemaTypeId 5
+                       .SetDatabase("db3/"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db2/type4")  // SchemaTypeId 6
+                       .SetDatabase("db2/"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db2/type5")  // SchemaTypeId 7
+                       .SetDatabase("db2/"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db3/type2")  // SchemaTypeId 8
+                       .SetDatabase("db3/"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db3/type3")  // SchemaTypeId 9
+                       .SetDatabase("db3/"))
+          .Build();
+
+  // Set full schema using the empty db
+  EXPECT_TRUE(schema_store
+                  ->SetSchema(CreateSetSchemaRequestProto(
+                      expected_full_schema, /*database=*/"",
+                      /*ignore_errors_and_delete_documents=*/false))
+                  .ok());
+  // Verify schema.
+  ICING_ASSERT_OK_AND_ASSIGN(const SchemaProto* actual_full_schema,
+                             schema_store->GetSchema());
+  EXPECT_THAT(*actual_full_schema, EqualsProto(expected_full_schema));
+
+  // Set schema again for db2 and delete some types.
+  SchemaProto db2_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db2/type2")
+                                            .SetDatabase("db2/"))
+                               .Build();
+
+  if (feature_flags_->enable_schema_type_id_optimization()) {
+    expected_full_schema =
+        SchemaBuilder()
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db1/type1")  // SchemaTypeId 0
+                         .SetDatabase("db1/"))
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db1/type2")  // SchemaTypeId 1
+                         .SetDatabase("db1/"))
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db3/type2")  // SchemaTypeId 8 -> 2
+                         .SetDatabase("db3/"))
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db2/type2")  // SchemaTypeId 3
+                         .SetDatabase("db2/"))
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db3/type3")  // SchemaTypeId 9 -> 4
+                         .SetDatabase("db3/"))
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db3/type1")  // SchemaTypeId 5
+                         .SetDatabase("db3/"))
+            .Build();
+
+    SchemaStore::SetSchemaResult result = SchemaStore::SetSchemaResult();
+    result.success = true;
+    result.schema_types_deleted_by_name.insert("db2/type1");
+    result.schema_types_deleted_by_name.insert("db2/type3");
+    result.schema_types_deleted_by_name.insert("db2/type4");
+    result.schema_types_deleted_by_name.insert("db2/type5");
+    result.schema_types_deleted_by_id.insert(
+        schema_store->GetSchemaTypeId("db2/type1").ValueOrDie());  // 2
+    result.schema_types_deleted_by_id.insert(
+        schema_store->GetSchemaTypeId("db2/type3").ValueOrDie());  // 4
+    result.schema_types_deleted_by_id.insert(
+        schema_store->GetSchemaTypeId("db2/type4").ValueOrDie());  // 6
+    result.schema_types_deleted_by_id.insert(
+        schema_store->GetSchemaTypeId("db2/type5").ValueOrDie());  // 7
+
+    // Only 2 schema type ids are changed since we move the fast few types at
+    // the end of the original schema to fill the gaps of the deleted types.
+    result.old_schema_type_ids_changed.insert(
+        schema_store->GetSchemaTypeId("db3/type2").ValueOrDie());  // 8
+    result.old_schema_type_ids_changed.insert(
+        schema_store->GetSchemaTypeId("db3/type3").ValueOrDie());  // 9
+
+    EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                    db2_schema, /*database=*/"db2/",
+                    /*ignore_errors_and_delete_documents=*/true)),
+                IsOkAndHolds(EqualsSetSchemaResult(result)));
+
+    // Check the schema
+    EXPECT_THAT(schema_store->GetSchema(),
+                IsOkAndHolds(Pointee(EqualsProto(expected_full_schema))));
+
+    EXPECT_THAT(schema_store->GetSchema("db1/"),
+                IsOkAndHolds(EqualsProto(SchemaBuilder()
+                                             .AddType(SchemaTypeConfigBuilder()
+                                                          .SetType("db1/type1")
+                                                          .SetDatabase("db1/"))
+                                             .AddType(SchemaTypeConfigBuilder()
+                                                          .SetType("db1/type2")
+                                                          .SetDatabase("db1/"))
+                                             .Build())));
+    EXPECT_THAT(schema_store->GetSchema("db2/"),
+                IsOkAndHolds(EqualsProto(SchemaBuilder()
+                                             .AddType(SchemaTypeConfigBuilder()
+                                                          .SetType("db2/type2")
+                                                          .SetDatabase("db2/"))
+                                             .Build())));
+    EXPECT_THAT(schema_store->GetSchema("db3/"),
+                IsOkAndHolds(EqualsProto(SchemaBuilder()
+                                             .AddType(SchemaTypeConfigBuilder()
+                                                          .SetType("db3/type2")
+                                                          .SetDatabase("db3/"))
+                                             .AddType(SchemaTypeConfigBuilder()
+                                                          .SetType("db3/type3")
+                                                          .SetDatabase("db3/"))
+                                             .AddType(SchemaTypeConfigBuilder()
+                                                          .SetType("db3/type1")
+                                                          .SetDatabase("db3/"))
+                                             .Build())));
+  } else {
+    expected_full_schema =
+        SchemaBuilder()
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db1/type1")  // SchemaTypeId 0
+                         .SetDatabase("db1/"))
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db1/type2")  // SchemaTypeId 1
+                         .SetDatabase("db1/"))
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db2/type2")  // SchemaTypeId 3 -> 2
+                         .SetDatabase("db2/"))
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db3/type1")  // SchemaTypeId 5 -> 3
+                         .SetDatabase("db3/"))
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db3/type2")  // SchemaTypeId 8 -> 4
+                         .SetDatabase("db3/"))
+            .AddType(SchemaTypeConfigBuilder()
+                         .SetType("db3/type3")  // SchemaTypeId 9 -> 5
+                         .SetDatabase("db3/"))
+            .Build();
+    SchemaStore::SetSchemaResult result = SchemaStore::SetSchemaResult();
+    result.success = true;
+    result.schema_types_deleted_by_name.insert("db2/type1");
+    result.schema_types_deleted_by_name.insert("db2/type3");
+    result.schema_types_deleted_by_name.insert("db2/type4");
+    result.schema_types_deleted_by_name.insert("db2/type5");
+    result.schema_types_deleted_by_id.insert(
+        schema_store->GetSchemaTypeId("db2/type1").ValueOrDie());  // 2
+    result.schema_types_deleted_by_id.insert(
+        schema_store->GetSchemaTypeId("db2/type3").ValueOrDie());  // 4
+    result.schema_types_deleted_by_id.insert(
+        schema_store->GetSchemaTypeId("db2/type4").ValueOrDie());  // 6
+    result.schema_types_deleted_by_id.insert(
+        schema_store->GetSchemaTypeId("db2/type5").ValueOrDie());  // 7
+
+    // 4 type ids are changed since the type ids are assigned based on the
+    // order of the types in the existing schema.
+    result.old_schema_type_ids_changed.insert(
+        schema_store->GetSchemaTypeId("db2/type2").ValueOrDie());  // 3
+    result.old_schema_type_ids_changed.insert(
+        schema_store->GetSchemaTypeId("db3/type1").ValueOrDie());  // 5
+    result.old_schema_type_ids_changed.insert(
+        schema_store->GetSchemaTypeId("db3/type2").ValueOrDie());  // 8
+    result.old_schema_type_ids_changed.insert(
+        schema_store->GetSchemaTypeId("db3/type3").ValueOrDie());  // 9
+
+    EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                    db2_schema, /*database=*/"db2/",
+                    /*ignore_errors_and_delete_documents=*/true)),
+                IsOkAndHolds(EqualsSetSchemaResult(result)));
+
+    // Check the schema
+    EXPECT_THAT(schema_store->GetSchema(),
+                IsOkAndHolds(Pointee(EqualsProto(expected_full_schema))));
+
+    EXPECT_THAT(schema_store->GetSchema("db1/"),
+                IsOkAndHolds(EqualsProto(SchemaBuilder()
+                                             .AddType(SchemaTypeConfigBuilder()
+                                                          .SetType("db1/type1")
+                                                          .SetDatabase("db1/"))
+                                             .AddType(SchemaTypeConfigBuilder()
+                                                          .SetType("db1/type2")
+                                                          .SetDatabase("db1/"))
+                                             .Build())));
+    EXPECT_THAT(schema_store->GetSchema("db2/"),
+                IsOkAndHolds(EqualsProto(SchemaBuilder()
+                                             .AddType(SchemaTypeConfigBuilder()
+                                                          .SetType("db2/type2")
+                                                          .SetDatabase("db2/"))
+                                             .Build())));
+    EXPECT_THAT(schema_store->GetSchema("db3/"),
+                IsOkAndHolds(EqualsProto(SchemaBuilder()
+                                             .AddType(SchemaTypeConfigBuilder()
+                                                          .SetType("db3/type1")
+                                                          .SetDatabase("db3/"))
+                                             .AddType(SchemaTypeConfigBuilder()
+                                                          .SetType("db3/type2")
+                                                          .SetDatabase("db3/"))
+                                             .AddType(SchemaTypeConfigBuilder()
+                                                          .SetType("db3/type3")
+                                                          .SetDatabase("db3/"))
+                                             .Build())));
+  }
+}
+
+TEST_F(SchemaStoreTest, SetDatabaseRenamedTypesPreserveTypeIds) {
+  if (!feature_flags_->enable_schema_type_id_optimization()) {
+    GTEST_SKIP() << "Test for schema type id optimization.";
+  }
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<SchemaStore> schema_store,
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_,
+                          feature_flags_.get(),
+                          /*initialize_stats=*/nullptr));
+
+  // Set schema for the first time
+  SchemaProto db1_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db1/email")
+                                            .SetDatabase("db1/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db1/message")
+                                            .SetDatabase("db1/"))
+                               .Build();
+  SchemaProto db2_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db2/email")
+                                            .SetDatabase("db2/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db2/message")
+                                            .SetDatabase("db2/"))
+                               .Build();
+  SchemaProto db3_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db3/email")
+                                            .SetDatabase("db3/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db3/message")
+                                            .SetDatabase("db3/"))
+                               .Build();
+
+  // Set schema for db1
+  SchemaStore::SetSchemaResult result;
+  result.success = true;
+  result.schema_types_new_by_name.insert("db1/email");
+  result.schema_types_new_by_name.insert("db1/message");
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db1_schema, /*database=*/"db1/",
+                  /*ignore_errors_and_delete_documents=*/false)),
+              IsOkAndHolds(EqualsSetSchemaResult(result)));
+  // Set schema for db2
   result = SchemaStore::SetSchemaResult();
   result.success = true;
-  result.schema_types_deleted_by_name.insert("db2_email");
-  result.schema_types_deleted_by_name.insert("db2_recipient");
-  result.schema_types_deleted_by_id.insert(2);   // db2_email
-  result.schema_types_deleted_by_id.insert(6);   // db2_recipient
-  result.old_schema_type_ids_changed.insert(3);  // db2_message
-  result.old_schema_type_ids_changed.insert(4);  // db3_email
-  result.old_schema_type_ids_changed.insert(5);  // db3_message
+  result.schema_types_new_by_name.insert("db2/email");
+  result.schema_types_new_by_name.insert("db2/message");
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  db2_schema, /*database=*/"db2",
+                  db2_schema, /*database=*/"db2/",
+                  /*ignore_errors_and_delete_documents=*/false)),
+              IsOkAndHolds(EqualsSetSchemaResult(result)));
+  // Set schema for db3
+  result = SchemaStore::SetSchemaResult();
+  result.success = true;
+  result.schema_types_new_by_name.insert("db3/email");
+  result.schema_types_new_by_name.insert("db3/message");
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db3_schema, /*database=*/"db3/",
+                  /*ignore_errors_and_delete_documents=*/false)),
+              IsOkAndHolds(EqualsSetSchemaResult(result)));
+
+  // Set schema again for db2 and rename email -> email2. The type ids of other
+  // existing types should not change.
+  db2_schema = SchemaBuilder()
+                   .AddType(SchemaTypeConfigBuilder()
+                                .SetType("db2/message")
+                                .SetDatabase("db2/"))
+                   .AddType(SchemaTypeConfigBuilder()
+                                .SetType("db2/email2")
+                                .SetDatabase("db2/"))
+                   .Build();
+  result = SchemaStore::SetSchemaResult();
+  result.success = true;
+  result.schema_types_deleted_by_name.insert("db2/email");
+  result.schema_types_deleted_by_id.insert(2);  // db2_email
+  result.schema_types_new_by_name.insert("db2/email2");
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db2_schema, /*database=*/"db2/",
                   /*ignore_errors_and_delete_documents=*/true)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
 
-  // Check the schema
+  SchemaProto expected_full_schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db1/email")  // SchemaTypeId 0
+                       .SetDatabase("db1/"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db1/message")  // SchemaTypeId 1
+                       .SetDatabase("db1/"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db2/email2")  // SchemaTypeId 2
+                       .SetDatabase("db2/"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db2/message")  // SchemaTypeId 3
+                       .SetDatabase("db2/"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db3/email")  // SchemaTypeId 4
+                       .SetDatabase("db3/"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db3/message")  // SchemaTypeId 5
+                       .SetDatabase("db3/"))
+          .Build();
+
+  // Verify schema.
+  ICING_ASSERT_OK_AND_ASSIGN(const SchemaProto* actual_full_schema,
+                             schema_store->GetSchema());
+  EXPECT_THAT(*actual_full_schema, EqualsProto(expected_full_schema));
+
+  // Check the schema for each database.
   EXPECT_THAT(schema_store->GetSchema(),
               IsOkAndHolds(Pointee(EqualsProto(expected_full_schema))));
-  EXPECT_THAT(schema_store->GetSchema("db1"),
+  EXPECT_THAT(schema_store->GetSchema("db1/"),
               IsOkAndHolds(EqualsProto(db1_schema)));
-  EXPECT_THAT(schema_store->GetSchema("db2"),
-              IsOkAndHolds(EqualsProto(db2_schema)));
-  EXPECT_THAT(schema_store->GetSchema("db3"),
+  EXPECT_THAT(schema_store->GetSchema("db2/"),
+              IsOkAndHolds(EqualsProto(SchemaBuilder()
+                                           .AddType(SchemaTypeConfigBuilder()
+                                                        .SetType("db2/email2")
+                                                        .SetDatabase("db2/"))
+                                           .AddType(SchemaTypeConfigBuilder()
+                                                        .SetType("db2/message")
+                                                        .SetDatabase("db2/"))
+                                           .Build())));
+  EXPECT_THAT(schema_store->GetSchema("db3/"),
               IsOkAndHolds(EqualsProto(db3_schema)));
 }
 
+TEST_F(SchemaStoreTest, SetDatabaseAddMoreTypesThanDeletedPreserveTypeIds) {
+  if (!feature_flags_->enable_schema_type_id_optimization()) {
+    GTEST_SKIP() << "Test for schema type id optimization.";
+  }
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<SchemaStore> schema_store,
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_,
+                          feature_flags_.get(),
+                          /*initialize_stats=*/nullptr));
+
+  // Set schema for the first time
+  SchemaProto db1_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db1/email")
+                                            .SetDatabase("db1/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db1/message")
+                                            .SetDatabase("db1/"))
+                               .Build();
+  SchemaProto db2_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db2/email")
+                                            .SetDatabase("db2/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db2/message")
+                                            .SetDatabase("db2/"))
+                               .Build();
+  SchemaProto db3_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db3/email")
+                                            .SetDatabase("db3/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db3/message")
+                                            .SetDatabase("db3/"))
+                               .Build();
+
+  // Set schema for db1
+  SchemaStore::SetSchemaResult result;
+  result.success = true;
+  result.schema_types_new_by_name.insert("db1/email");
+  result.schema_types_new_by_name.insert("db1/message");
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db1_schema, /*database=*/"db1/",
+                  /*ignore_errors_and_delete_documents=*/false)),
+              IsOkAndHolds(EqualsSetSchemaResult(result)));
+  // Set schema for db2
+  result = SchemaStore::SetSchemaResult();
+  result.success = true;
+  result.schema_types_new_by_name.insert("db2/email");
+  result.schema_types_new_by_name.insert("db2/message");
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db2_schema, /*database=*/"db2/",
+                  /*ignore_errors_and_delete_documents=*/false)),
+              IsOkAndHolds(EqualsSetSchemaResult(result)));
+  // Set schema for db3
+  result = SchemaStore::SetSchemaResult();
+  result.success = true;
+  result.schema_types_new_by_name.insert("db3/email");
+  result.schema_types_new_by_name.insert("db3/message");
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db3_schema, /*database=*/"db3/",
+                  /*ignore_errors_and_delete_documents=*/false)),
+              IsOkAndHolds(EqualsSetSchemaResult(result)));
+
+  // Set schema again for db2. Delete the email type, and add 2 more types.
+  // The type id of existing types should not change.
+  db2_schema = SchemaBuilder()
+                   .AddType(SchemaTypeConfigBuilder()
+                                .SetType("db2/message")
+                                .SetDatabase("db2/"))
+                   .AddType(SchemaTypeConfigBuilder()
+                                .SetType("db2/new_type1")
+                                .SetDatabase("db2/"))
+                   .AddType(SchemaTypeConfigBuilder()
+                                .SetType("db2/new_type2")
+                                .SetDatabase("db2/"))
+                   .Build();
+  result = SchemaStore::SetSchemaResult();
+  result.success = true;
+  result.schema_types_deleted_by_name.insert("db2/email");
+  result.schema_types_deleted_by_id.insert(2);  // db2_email
+  result.schema_types_new_by_name.insert("db2/new_type1");
+  result.schema_types_new_by_name.insert("db2/new_type2");
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  db2_schema, /*database=*/"db2/",
+                  /*ignore_errors_and_delete_documents=*/true)),
+              IsOkAndHolds(EqualsSetSchemaResult(result)));
+
+  SchemaProto expected_full_schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db1/email")  // SchemaTypeId 0
+                       .SetDatabase("db1/"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db1/message")  // SchemaTypeId 1
+                       .SetDatabase("db1/"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db2/new_type1")  // SchemaTypeId 2
+                       .SetDatabase("db2/"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db2/message")  // SchemaTypeId 3
+                       .SetDatabase("db2/"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db3/email")  // SchemaTypeId 4
+                       .SetDatabase("db3/"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db3/message")  // SchemaTypeId 5
+                       .SetDatabase("db3/"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db2/new_type2")  // SchemaTypeId 6
+                       .SetDatabase("db2/"))
+          .Build();
+
+  // Verify schema.
+  ICING_ASSERT_OK_AND_ASSIGN(const SchemaProto* actual_full_schema,
+                             schema_store->GetSchema());
+  EXPECT_THAT(*actual_full_schema, EqualsProto(expected_full_schema));
+
+  // Check the schema for each database.
+  EXPECT_THAT(schema_store->GetSchema(),
+              IsOkAndHolds(Pointee(EqualsProto(expected_full_schema))));
+  EXPECT_THAT(schema_store->GetSchema("db1/"),
+              IsOkAndHolds(EqualsProto(db1_schema)));
+  EXPECT_THAT(
+      schema_store->GetSchema("db2/"),
+      IsOkAndHolds(EqualsProto(SchemaBuilder()
+                                   .AddType(SchemaTypeConfigBuilder()
+                                                .SetType("db2/new_type1")
+                                                .SetDatabase("db2/"))
+                                   .AddType(SchemaTypeConfigBuilder()
+                                                .SetType("db2/message")
+                                                .SetDatabase("db2/"))
+                                   .AddType(SchemaTypeConfigBuilder()
+                                                .SetType("db2/new_type2")
+                                                .SetDatabase("db2/"))
+                                   .Build())));
+  EXPECT_THAT(schema_store->GetSchema("db3/"),
+              IsOkAndHolds(EqualsProto(db3_schema)));
+}
+
+TEST_F(SchemaStoreTest, SetEmptySchemaOk) {
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<SchemaStore> schema_store,
+      SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_,
+                          feature_flags_.get(),
+                          /*initialize_stats=*/nullptr));
+
+  // Set schema for the first time
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(
+              SchemaTypeConfigBuilder().SetType("db/email").SetDatabase("db/"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db/message")
+                       .SetDatabase("db/"))
+          .Build();
+  SchemaStore::SetSchemaResult result;
+  result.success = true;
+  result.schema_types_new_by_name.insert("db/email");
+  result.schema_types_new_by_name.insert("db/message");
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  schema, /*database=*/"db/",
+                  /*ignore_errors_and_delete_documents=*/false)),
+              IsOkAndHolds(EqualsSetSchemaResult(result)));
+
+  // Verify schema.
+  EXPECT_THAT(schema_store->GetSchema(),
+              IsOkAndHolds(Pointee(EqualsProto(schema))));
+
+  // Reset to an empty schema.
+  result = SchemaStore::SetSchemaResult();
+  result.success = true;
+  result.schema_types_deleted_by_name.insert("db/email");
+  result.schema_types_deleted_by_name.insert("db/message");
+  result.schema_types_deleted_by_id.insert(0);  // email
+  result.schema_types_deleted_by_id.insert(1);  // message
+  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
+                  SchemaProto(), /*database=*/"db/",
+                  /*ignore_errors_and_delete_documents=*/true)),
+              IsOkAndHolds(EqualsSetSchemaResult(result)));
+
+  // Check the schema. It should be empty.
+  EXPECT_THAT(schema_store->GetSchema(),
+              IsOkAndHolds(Pointee(EqualsProto(SchemaProto()))));
+
+  EXPECT_THAT(schema_store->PersistToDisk(), IsOk());
+}
+
 TEST_F(SchemaStoreTest, SetEmptySchemaClearsDatabase) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
@@ -1264,79 +2404,79 @@
                           /*initialize_stats=*/nullptr));
 
   // Set schema for the first time
-  SchemaProto db1_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db1_email").SetDatabase("db1"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db1_message")
-                       .SetDatabase("db1"))
-          .Build();
-  SchemaProto db2_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db2_email").SetDatabase("db2"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_message")
-                       .SetDatabase("db2"))
-          .Build();
-  SchemaProto db3_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db3_email").SetDatabase("db3"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db3_message")
-                       .SetDatabase("db3"))
-          .Build();
+  SchemaProto db1_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db1/email")
+                                            .SetDatabase("db1/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db1/message")
+                                            .SetDatabase("db1/"))
+                               .Build();
+  SchemaProto db2_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db2/email")
+                                            .SetDatabase("db2/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db2/message")
+                                            .SetDatabase("db2/"))
+                               .Build();
+  SchemaProto db3_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db3/email")
+                                            .SetDatabase("db3/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db3/message")
+                                            .SetDatabase("db3/"))
+                               .Build();
 
   // Set schema for db1
   SchemaStore::SetSchemaResult result;
   result.success = true;
-  result.schema_types_new_by_name.insert("db1_email");
-  result.schema_types_new_by_name.insert("db1_message");
+  result.schema_types_new_by_name.insert("db1/email");
+  result.schema_types_new_by_name.insert("db1/message");
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  db1_schema, /*database=*/"db1",
+                  db1_schema, /*database=*/"db1/",
                   /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   // Set schema for db2
   result = SchemaStore::SetSchemaResult();
   result.success = true;
-  result.schema_types_new_by_name.insert("db2_email");
-  result.schema_types_new_by_name.insert("db2_message");
+  result.schema_types_new_by_name.insert("db2/email");
+  result.schema_types_new_by_name.insert("db2/message");
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  db2_schema, /*database=*/"db2",
+                  db2_schema, /*database=*/"db2/",
                   /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   // Set schema for db3
   result = SchemaStore::SetSchemaResult();
   result.success = true;
-  result.schema_types_new_by_name.insert("db3_email");
-  result.schema_types_new_by_name.insert("db3_message");
+  result.schema_types_new_by_name.insert("db3/email");
+  result.schema_types_new_by_name.insert("db3/message");
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  db3_schema, /*database=*/"db3",
+                  db3_schema, /*database=*/"db3/",
                   /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   // Verify schema.
   SchemaProto expected_full_schema =
       SchemaBuilder()
           .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db1_email")  // SchemaTypeId 0
-                       .SetDatabase("db1"))
+                       .SetType("db1/email")  // SchemaTypeId 0
+                       .SetDatabase("db1/"))
           .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db1_message")  // SchemaTypeId 1
-                       .SetDatabase("db1"))
+                       .SetType("db1/message")  // SchemaTypeId 1
+                       .SetDatabase("db1/"))
           .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_email")  // SchemaTypeId 2
-                       .SetDatabase("db2"))
+                       .SetType("db2/email")  // SchemaTypeId 2
+                       .SetDatabase("db2/"))
           .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_message")  // SchemaTypeId 3
-                       .SetDatabase("db2"))
+                       .SetType("db2/message")  // SchemaTypeId 3
+                       .SetDatabase("db2/"))
           .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db3_email")  // SchemaTypeId 4
-                       .SetDatabase("db3"))
+                       .SetType("db3/email")  // SchemaTypeId 4
+                       .SetDatabase("db3/"))
           .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db3_message")  // SchemaTypeId 5
-                       .SetDatabase("db3"))
+                       .SetType("db3/message")  // SchemaTypeId 5
+                       .SetDatabase("db3/"))
           .Build();
   EXPECT_THAT(schema_store->GetSchema(),
               IsOkAndHolds(Pointee(EqualsProto(expected_full_schema))));
@@ -1347,40 +2487,70 @@
   db2_schema = SchemaProto();
   result = SchemaStore::SetSchemaResult();
   result.success = true;
-  result.schema_types_deleted_by_name.insert("db2_email");
-  result.schema_types_deleted_by_name.insert("db2_message");
+  result.schema_types_deleted_by_name.insert("db2/email");
+  result.schema_types_deleted_by_name.insert("db2/message");
   result.schema_types_deleted_by_id.insert(2);   // db2_email
   result.schema_types_deleted_by_id.insert(3);   // db2_message
   result.old_schema_type_ids_changed.insert(4);  // db3_email
   result.old_schema_type_ids_changed.insert(5);  // db3_message
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  db2_schema, /*database=*/"db2",
+                  db2_schema, /*database=*/"db2/",
                   /*ignore_errors_and_delete_documents=*/true)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
 
-  // Check the schema. Schemas for db1 and db3 should be unchanged.
-  EXPECT_THAT(schema_store->GetSchema("db1"),
+  // Check the schema. Schema types for db1 and db3 should be unchanged, but if
+  // schema_type_id_optimization is enabled, then the schema type order for db3
+  // will change since the position of the deleted types from db2 will be
+  // backfilled from the last type in the original schema.
+  EXPECT_THAT(schema_store->GetSchema("db1/"),
               IsOkAndHolds(EqualsProto(db1_schema)));
-  EXPECT_THAT(schema_store->GetSchema("db3"),
+  if (feature_flags_->enable_schema_type_id_optimization()) {
+    db3_schema = SchemaBuilder()
+                     .AddType(SchemaTypeConfigBuilder()
+                                  .SetType("db3/email")
+                                  .SetDatabase("db3/"))
+                     .AddType(SchemaTypeConfigBuilder()
+                                  .SetType("db3/message")
+                                  .SetDatabase("db3/"))
+                     .Build();
+    expected_full_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db1/email")
+                                            .SetDatabase("db1/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db1/message")
+                                            .SetDatabase("db1/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db3/email")
+                                            .SetDatabase("db3/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db3/message")
+                                            .SetDatabase("db3/"))
+                               .Build();
+  } else {
+    expected_full_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db1/email")
+                                            .SetDatabase("db1/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db1/message")
+                                            .SetDatabase("db1/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db3/email")
+                                            .SetDatabase("db3/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db3/message")
+                                            .SetDatabase("db3/"))
+                               .Build();
+  }
+  EXPECT_THAT(schema_store->GetSchema("db3/"),
               IsOkAndHolds(EqualsProto(db3_schema)));
 
   // GetSchema for db2 should return NotFoundError
-  EXPECT_THAT(schema_store->GetSchema("db2"),
+  EXPECT_THAT(schema_store->GetSchema("db2/"),
               StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
 
-  expected_full_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db1_email").SetDatabase("db1"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db1_message")
-                       .SetDatabase("db1"))
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db3_email").SetDatabase("db3"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db3_message")
-                       .SetDatabase("db3"))
-          .Build();
+  // Get full schema.
   EXPECT_THAT(schema_store->GetSchema(),
               IsOkAndHolds(Pointee(EqualsProto(expected_full_schema))));
 }
@@ -1423,49 +2593,50 @@
                           /*initialize_stats=*/nullptr));
 
   // Set schema for the first time
-  SchemaProto db1_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db1_email").SetDatabase("db1"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db1_message")
-                       .SetDatabase("db1"))
-          .Build();
-  SchemaProto db2_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db2_email").SetDatabase("db2"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_message")
-                       .SetDatabase("db2"))
-          .Build();
-  SchemaProto expected_full_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db1_email").SetDatabase("db1"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db1_message")
-                       .SetDatabase("db1"))
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db2_email").SetDatabase("db2"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_message")
-                       .SetDatabase("db2"))
-          .Build();
+  SchemaProto db1_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db1/email")
+                                            .SetDatabase("db1/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db1/message")
+                                            .SetDatabase("db1/"))
+                               .Build();
+  SchemaProto db2_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db2/email")
+                                            .SetDatabase("db2/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db2/message")
+                                            .SetDatabase("db2/"))
+                               .Build();
+  SchemaProto expected_full_schema = SchemaBuilder()
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db1/email")
+                                                      .SetDatabase("db1/"))
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db1/message")
+                                                      .SetDatabase("db1/"))
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db2/email")
+                                                      .SetDatabase("db2/"))
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db2/message")
+                                                      .SetDatabase("db2/"))
+                                         .Build();
   SchemaStore::SetSchemaResult result;
   result.success = true;
-  result.schema_types_new_by_name.insert("db1_email");
-  result.schema_types_new_by_name.insert("db1_message");
+  result.schema_types_new_by_name.insert("db1/email");
+  result.schema_types_new_by_name.insert("db1/message");
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  db1_schema, /*database=*/"db1",
+                  db1_schema, /*database=*/"db1/",
                   /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   result = SchemaStore::SetSchemaResult();
   result.success = true;
-  result.schema_types_new_by_name.insert("db2_email");
-  result.schema_types_new_by_name.insert("db2_message");
+  result.schema_types_new_by_name.insert("db2/email");
+  result.schema_types_new_by_name.insert("db2/message");
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  db2_schema, /*database=*/"db2",
+                  db2_schema, /*database=*/"db2/",
                   /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   ICING_ASSERT_OK_AND_ASSIGN(const SchemaProto* actual_full_schema,
@@ -1475,28 +2646,29 @@
   // Make db2 incompatible by changing a type name
   SchemaProto db2_schema_incompatible =
       SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db2_email").SetDatabase("db2"))
           .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_recipient")
-                       .SetDatabase("db2"))
+                       .SetType("db2/email")
+                       .SetDatabase("db2/"))
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType("db2/recipient")
+                       .SetDatabase("db2/"))
           .Build();
   result = SchemaStore::SetSchemaResult();
   result.success = false;
-  result.schema_types_deleted_by_name.insert("db2_message");
-  result.schema_types_new_by_name.insert("db2_recipient");
+  result.schema_types_deleted_by_name.insert("db2/message");
+  result.schema_types_new_by_name.insert("db2/recipient");
   result.schema_types_deleted_by_id.insert(3);  // db2_message
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  db2_schema_incompatible, /*database=*/"db2",
+                  db2_schema_incompatible, /*database=*/"db2/",
                   /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
 
   // Check the schema, this should not have changed
   EXPECT_THAT(schema_store->GetSchema(),
               IsOkAndHolds(Pointee(EqualsProto(expected_full_schema))));
-  EXPECT_THAT(schema_store->GetSchema("db1"),
+  EXPECT_THAT(schema_store->GetSchema("db1/"),
               IsOkAndHolds(EqualsProto(db1_schema)));
-  EXPECT_THAT(schema_store->GetSchema("db2"),
+  EXPECT_THAT(schema_store->GetSchema("db2/"),
               IsOkAndHolds(EqualsProto(db2_schema)));
 }
 
@@ -1508,49 +2680,50 @@
                           /*initialize_stats=*/nullptr));
 
   // Set schema for the first time
-  SchemaProto db1_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db1_email").SetDatabase("db1"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db1_message")
-                       .SetDatabase("db1"))
-          .Build();
-  SchemaProto db2_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db2_email").SetDatabase("db2"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_message")
-                       .SetDatabase("db2"))
-          .Build();
-  SchemaProto expected_full_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db1_email").SetDatabase("db1"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db1_message")
-                       .SetDatabase("db1"))
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db2_email").SetDatabase("db2"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_message")
-                       .SetDatabase("db2"))
-          .Build();
+  SchemaProto db1_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db1/email")
+                                            .SetDatabase("db1/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db1/message")
+                                            .SetDatabase("db1/"))
+                               .Build();
+  SchemaProto db2_schema = SchemaBuilder()
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db2/email")
+                                            .SetDatabase("db2/"))
+                               .AddType(SchemaTypeConfigBuilder()
+                                            .SetType("db2/message")
+                                            .SetDatabase("db2/"))
+                               .Build();
+  SchemaProto expected_full_schema = SchemaBuilder()
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db1/email")
+                                                      .SetDatabase("db1/"))
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db1/message")
+                                                      .SetDatabase("db1/"))
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db2/email")
+                                                      .SetDatabase("db2/"))
+                                         .AddType(SchemaTypeConfigBuilder()
+                                                      .SetType("db2/message")
+                                                      .SetDatabase("db2/"))
+                                         .Build();
   SchemaStore::SetSchemaResult result;
   result.success = true;
-  result.schema_types_new_by_name.insert("db1_email");
-  result.schema_types_new_by_name.insert("db1_message");
+  result.schema_types_new_by_name.insert("db1/email");
+  result.schema_types_new_by_name.insert("db1/message");
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  db1_schema, /*database=*/"db1",
+                  db1_schema, /*database=*/"db1/",
                   /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   result = SchemaStore::SetSchemaResult();
   result.success = true;
-  result.schema_types_new_by_name.insert("db2_email");
-  result.schema_types_new_by_name.insert("db2_message");
+  result.schema_types_new_by_name.insert("db2/email");
+  result.schema_types_new_by_name.insert("db2/message");
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  db2_schema, /*database=*/"db2",
+                  db2_schema, /*database=*/"db2/",
                   /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   ICING_ASSERT_OK_AND_ASSIGN(const SchemaProto* actual_full_schema,
@@ -1566,23 +2739,23 @@
           .Build();
   SchemaProto db2_schema_incompatible = SchemaBuilder()
                                             .AddType(SchemaTypeConfigBuilder()
-                                                         .SetType("db2_email")
-                                                         .SetDatabase("db2")
+                                                         .SetType("db2/email")
+                                                         .SetDatabase("db2/")
                                                          .AddProperty(prop)
                                                          .AddProperty(prop))
                                             .Build();
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
                   db2_schema_incompatible,
-                  /*database=*/"db2",
+                  /*database=*/"db2/",
                   /*ignore_errors_and_delete_documents=*/false)),
               StatusIs(libtextclassifier3::StatusCode::ALREADY_EXISTS));
 
   // Check the schema, this should not have changed
   EXPECT_THAT(schema_store->GetSchema(),
               IsOkAndHolds(Pointee(EqualsProto(expected_full_schema))));
-  EXPECT_THAT(schema_store->GetSchema("db1"),
+  EXPECT_THAT(schema_store->GetSchema("db1/"),
               IsOkAndHolds(EqualsProto(db1_schema)));
-  EXPECT_THAT(schema_store->GetSchema("db2"),
+  EXPECT_THAT(schema_store->GetSchema("db2/"),
               IsOkAndHolds(EqualsProto(db2_schema)));
 }
 
@@ -1593,21 +2766,22 @@
                           feature_flags_.get(),
                           /*initialize_stats=*/nullptr));
 
-  SchemaProto combined_schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db2_email").SetDatabase("db2"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_message")
-                       .SetDatabase("db2"))
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db1_email").SetDatabase("db1"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db1_message")
-                       .SetDatabase("db1"))
-          .Build();
+  SchemaProto combined_schema = SchemaBuilder()
+                                    .AddType(SchemaTypeConfigBuilder()
+                                                 .SetType("db2/email")
+                                                 .SetDatabase("db2/"))
+                                    .AddType(SchemaTypeConfigBuilder()
+                                                 .SetType("db2/message")
+                                                 .SetDatabase("db2/"))
+                                    .AddType(SchemaTypeConfigBuilder()
+                                                 .SetType("db1/email")
+                                                 .SetDatabase("db1/"))
+                                    .AddType(SchemaTypeConfigBuilder()
+                                                 .SetType("db1/message")
+                                                 .SetDatabase("db1/"))
+                                    .Build();
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  combined_schema, /*database=*/"db1",
+                  combined_schema, /*database=*/"db1/",
                   /*ignore_errors_and_delete_documents=*/false)),
               StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
 }
@@ -1623,47 +2797,29 @@
       SchemaBuilder()
           // This type does not explicitly set its database, so it defaults to
           // the empty database.
-          .AddType(SchemaTypeConfigBuilder().SetType("db1_email"))
+          .AddType(SchemaTypeConfigBuilder().SetType("db1/email"))
           .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db1_message")
-                       .SetDatabase("db1"))
+                       .SetType("db1/message")
+                       .SetDatabase("db1/"))
           .Build();
 
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  schema, /*database=*/"db1",
+                  schema, /*database=*/"db1/",
                   /*ignore_errors_and_delete_documents=*/false)),
               StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
 
-  schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db1_email").SetDatabase("db1"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db1_message")
-                       .SetDatabase("db1"))
-          .Build();
+  schema = SchemaBuilder()
+               .AddType(SchemaTypeConfigBuilder()
+                            .SetType("db1/email")
+                            .SetDatabase("db1/"))
+               .AddType(SchemaTypeConfigBuilder()
+                            .SetType("db1/message")
+                            .SetDatabase("db1/"))
+               .Build();
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
                   schema, /*database=*/"db_mismatch",
                   /*ignore_errors_and_delete_documents=*/false)),
               StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
-
-  schema =
-      SchemaBuilder()
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db1_email").SetDatabase("db1"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db1_message")
-                       .SetDatabase("db1"))
-          .AddType(
-              SchemaTypeConfigBuilder().SetType("db2_email").SetDatabase("db2"))
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_message")
-                       .SetDatabase("db2"))
-          .Build();
-  EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  schema, /*database=*/"",
-                  /*ignore_errors_and_delete_documents=*/false)),
-              StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
 }
 
 TEST_F(SchemaStoreTest, SetSchemaWithDuplicateTypeNameAcrossDifferentDbFails) {
@@ -1677,42 +2833,42 @@
   SchemaProto db1_schema =
       SchemaBuilder()
           .AddType(
-              SchemaTypeConfigBuilder().SetType("email").SetDatabase("db1"))
+              SchemaTypeConfigBuilder().SetType("email").SetDatabase("db1/"))
           .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db1_message")
-                       .SetDatabase("db1"))
+                       .SetType("db1/message")
+                       .SetDatabase("db1/"))
           .Build();
   SchemaStore::SetSchemaResult result;
   result.success = true;
   result.schema_types_new_by_name.insert("email");
-  result.schema_types_new_by_name.insert("db1_message");
+  result.schema_types_new_by_name.insert("db1/message");
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  db1_schema, /*database=*/"db1",
+                  db1_schema, /*database=*/"db1/",
                   /*ignore_errors_and_delete_documents=*/false)),
               IsOkAndHolds(EqualsSetSchemaResult(result)));
   EXPECT_THAT(schema_store->GetSchema(),
               IsOkAndHolds(Pointee(EqualsProto(db1_schema))));
-  EXPECT_THAT(schema_store->GetSchema("db1"),
+  EXPECT_THAT(schema_store->GetSchema("db1/"),
               IsOkAndHolds(EqualsProto(db1_schema)));
 
   // Set schema in db2 with the same type name
   SchemaProto db2_schema =
       SchemaBuilder()
           .AddType(
-              SchemaTypeConfigBuilder().SetType("email").SetDatabase("db2"))
+              SchemaTypeConfigBuilder().SetType("email").SetDatabase("db2/"))
           .AddType(SchemaTypeConfigBuilder()
-                       .SetType("db2_message")
-                       .SetDatabase("db2"))
+                       .SetType("db2/message")
+                       .SetDatabase("db2/"))
           .Build();
   EXPECT_THAT(schema_store->SetSchema(CreateSetSchemaRequestProto(
-                  db2_schema, /*database=*/"db2",
+                  db2_schema, /*database=*/"db2/",
                   /*ignore_errors_and_delete_documents=*/false)),
               StatusIs(libtextclassifier3::StatusCode::ALREADY_EXISTS));
 
   // Check schema, this should not have changed
   EXPECT_THAT(schema_store->GetSchema(),
               IsOkAndHolds(Pointee(EqualsProto(db1_schema))));
-  EXPECT_THAT(schema_store->GetSchema("db1"),
+  EXPECT_THAT(schema_store->GetSchema("db1/"),
               IsOkAndHolds(EqualsProto(db1_schema)));
 }
 
@@ -1860,21 +3016,21 @@
     EXPECT_THAT(*actual_schema, EqualsProto(schema));
   } else {
     // Since we assign SchemaTypeIds based on order in the SchemaProto, this
-    // will
-    // cause SchemaTypeIds to change
+    // will cause SchemaTypeIds to change
     result = SchemaStore::SetSchemaResult();
     result.success = true;
-    result.old_schema_type_ids_changed.emplace(
-        0);  // Old SchemaTypeId of "email"
-    result.old_schema_type_ids_changed.emplace(
-        1);  // Old SchemaTypeId of "message"
+    // Old SchemaTypeId of "email"
+    result.old_schema_type_ids_changed.emplace(0);
+    // Old SchemaTypeId of "message"
+    result.old_schema_type_ids_changed.emplace(1);
 
     // Set the compatible schema
-    EXPECT_THAT(schema_store->SetSchema(
-                    schema, /*ignore_errors_and_delete_documents=*/false),
-                IsOkAndHolds(EqualsSetSchemaResult(result)));
+    EXPECT_THAT(
+        schema_store->SetSchema(reordered_schema,
+                                /*ignore_errors_and_delete_documents=*/false),
+        IsOkAndHolds(EqualsSetSchemaResult(result)));
     ICING_ASSERT_OK_AND_ASSIGN(actual_schema, schema_store->GetSchema());
-    EXPECT_THAT(*actual_schema, EqualsProto(schema));
+    EXPECT_THAT(*actual_schema, EqualsProto(reordered_schema));
   }
 }
 
@@ -2427,8 +3583,8 @@
   new_expected_result.success = true;
   new_expected_result.schema_types_scorable_property_inconsistent_by_id.insert(
       0);
-  new_expected_result.schema_types_changed_fully_compatible_by_name.insert(
-      "email");
+  new_expected_result.schema_types_scorable_property_inconsistent_by_name
+      .insert("email");
   EXPECT_THAT(schema_store->SetSchema(
                   new_schema, /*ignore_errors_and_delete_documents=*/false),
               IsOkAndHolds(EqualsSetSchemaResult(new_expected_result)));
@@ -2437,7 +3593,7 @@
 }
 
 TEST_F(SchemaStoreTest,
-       SetSchemaWithReorderedSchemeTypesAndUpdatedScorablePropertyOk) {
+       SetSchemaWithChangedSchemaTypeIdAndUpdatedScorablePropertyOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
       SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_,
@@ -2459,7 +3615,7 @@
                                         .SetCardinality(CARDINALITY_OPTIONAL)))
           .Build();
   // The new schema updates "score" as scorable_type ENABLED from type "email",
-  // and it also reorders the schema types of "email" and "message".
+  // and it also changes the schema types of "email".
   SchemaProto new_schema =
       SchemaBuilder()
           .AddType(SchemaTypeConfigBuilder()
@@ -2473,7 +3629,6 @@
                                         .SetDataType(TYPE_DOUBLE)
                                         .SetScorableType(SCORABLE_TYPE_ENABLED)
                                         .SetCardinality(CARDINALITY_OPTIONAL)))
-          .AddType(SchemaTypeConfigBuilder().SetType("message"))
           .Build();
 
   // Set old schema
@@ -2495,12 +3650,14 @@
   SchemaTypeId email_schema_type_id = 0;
   new_expected_result.schema_types_scorable_property_inconsistent_by_id.insert(
       email_schema_type_id);
-  new_expected_result.schema_types_changed_fully_compatible_by_name.insert(
-      "email");
-  new_expected_result.old_schema_type_ids_changed.insert(0);
+  new_expected_result.schema_types_scorable_property_inconsistent_by_name
+      .insert("email");
+  // 'email' schema type id is changed. 'message' is deleted.
   new_expected_result.old_schema_type_ids_changed.insert(1);
+  new_expected_result.schema_types_deleted_by_id.insert(0);
+  new_expected_result.schema_types_deleted_by_name.insert("message");
   EXPECT_THAT(schema_store->SetSchema(
-                  new_schema, /*ignore_errors_and_delete_documents=*/false),
+                  new_schema, /*ignore_errors_and_delete_documents=*/true),
               IsOkAndHolds(EqualsSetSchemaResult(new_expected_result)));
   ICING_ASSERT_OK_AND_ASSIGN(actual_schema, schema_store->GetSchema());
   EXPECT_THAT(*actual_schema, EqualsProto(new_schema));
@@ -4928,6 +6085,22 @@
                   EqualsScorablePropertyInfo("scoreInt", TYPE_INT64)))));
 }
 
+TEST_F(SchemaStoreTest, GetSchemaNameHash) {
+  // Hardcode some inputs to the GetSchemaNameHash function, so that we can be
+  // aware of any changes to the hashing function.
+  EXPECT_EQ(SchemaStore::GetSchemaNameHash("schema1"), 2607648336);
+  EXPECT_EQ(SchemaStore::GetSchemaNameHash("schema2"), 40165354);
+  EXPECT_EQ(SchemaStore::GetSchemaNameHash("schema3"), 1969483644);
+  EXPECT_EQ(SchemaStore::GetSchemaNameHash("aa"), 1179847464);
+  EXPECT_EQ(SchemaStore::GetSchemaNameHash("bb"), 4101441873);
+  EXPECT_EQ(SchemaStore::GetSchemaNameHash("abc"), 3395655888);
+  EXPECT_EQ(SchemaStore::GetSchemaNameHash("cba"), 669986194);
+  EXPECT_EQ(SchemaStore::GetSchemaNameHash("aab"), 2521824133);
+  EXPECT_EQ(SchemaStore::GetSchemaNameHash("bba"), 640501669);
+  EXPECT_EQ(SchemaStore::GetSchemaNameHash("1234"), 3131523007);
+  EXPECT_EQ(SchemaStore::GetSchemaNameHash("4321"), 3855245428);
+}
+
 class SchemaStoreTestWithParam
     : public SchemaStoreTest,
       public testing::WithParamInterface<version_util::StateChange> {};
@@ -4967,7 +6140,7 @@
   SchemaTypeConfigProto db1_email_rfc =
       SchemaTypeConfigBuilder()
           .SetType("db1/email")
-          .SetDatabase("db1")
+          .SetDatabase("db1/")
           .AddProperty(
               PropertyConfigBuilder()
                   .SetName("db1Subject")
@@ -4977,7 +6150,7 @@
   SchemaTypeConfigProto db1_email_no_rfc =
       SchemaTypeConfigBuilder()
           .SetType("db1/email")
-          .SetDatabase("db1")
+          .SetDatabase("db1/")
           .AddProperty(PropertyConfigBuilder()
                            .SetName("db1Subject")
                            .SetCardinality(CARDINALITY_OPTIONAL)
@@ -4986,7 +6159,7 @@
   SchemaTypeConfigProto db2_email =
       SchemaTypeConfigBuilder()
           .SetType("db2/email")
-          .SetDatabase("db2")
+          .SetDatabase("db2/")
           .AddProperty(PropertyConfigBuilder()
                            .SetName("db2Subject")
                            .SetDataTypeString(TERM_MATCH_EXACT, TOKENIZER_PLAIN)
@@ -5019,41 +6192,41 @@
       EXPECT_THAT(
           schema_store->GetSchema(),
           IsOkAndHolds(Pointee(EqualsProto(full_schema_with_database_no_rfc))));
-      EXPECT_THAT(schema_store->GetSchema("db1"),
+      EXPECT_THAT(schema_store->GetSchema("db1/"),
                   IsOkAndHolds(EqualsProto(db1_schema_no_rfc)));
-      EXPECT_THAT(schema_store->GetSchema("db2"),
+      EXPECT_THAT(schema_store->GetSchema("db2/"),
                   IsOkAndHolds(EqualsProto(db2_schema)));
     } else {
       EXPECT_THAT(
           schema_store->GetSchema(),
           IsOkAndHolds(Pointee(EqualsProto(full_schema_with_database_rfc))));
-      EXPECT_THAT(schema_store->GetSchema("db1"),
+      EXPECT_THAT(schema_store->GetSchema("db1/"),
                   IsOkAndHolds(EqualsProto(db1_schema_rfc)));
-      EXPECT_THAT(schema_store->GetSchema("db2"),
+      EXPECT_THAT(schema_store->GetSchema("db2/"),
                   IsOkAndHolds(EqualsProto(db2_schema)));
 
       DocumentProto db1_email_doc =
           DocumentBuilder()
               .SetKey("namespace", "uri1")
               .SetSchema("db1/email")
-              .AddStringProperty("db1Subject", "db1_subject")
+              .AddStringProperty("db1Subject", "db1/subject")
               .Build();
       DocumentProto db2_email_doc =
           DocumentBuilder()
               .SetKey("namespace", "uri3")
               .SetSchema("db2/email")
-              .AddStringProperty("db2Subject", "db2_subject")
+              .AddStringProperty("db2Subject", "db2/subject")
               .Build();
 
       // Verify that our in-memory structures are ok
       ICING_ASSERT_OK_AND_ASSIGN(SectionGroup section_group,
                                  schema_store->ExtractSections(db1_email_doc));
       EXPECT_THAT(section_group.string_sections[0].content,
-                  ElementsAre("db1_subject"));
+                  ElementsAre("db1/subject"));
       ICING_ASSERT_OK_AND_ASSIGN(section_group,
                                  schema_store->ExtractSections(db2_email_doc));
       EXPECT_THAT(section_group.string_sections[0].content,
-                  ElementsAre("db2_subject"));
+                  ElementsAre("db2/subject"));
 
       // Verify that our persisted data are ok
       EXPECT_THAT(schema_store->GetSchemaTypeId("db1/email"), IsOkAndHolds(0));

diff --git a/icing/schema/schema-util.cc b/icing/schema/schema-util.cc
index cfc2391..b482e15 100644
--- a/icing/schema/schema-util.cc
+++ b/icing/schema/schema-util.cc

@@ -49,8 +49,30 @@
 
 bool AreDocumentIndexingConfigsEqual(const DocumentIndexingConfig& old_config,
                                      const DocumentIndexingConfig& new_config) {
-  return old_config.index_nested_properties() ==
-         new_config.index_nested_properties();
+  // TODO(b/265304217): This could mark the new schema as incompatible and
+  // generate some unnecessary index rebuilds if the two schemas have an
+  // equivalent set of indexed properties, but changed the way that it is
+  // declared.
+  if (old_config.index_nested_properties() !=
+      new_config.index_nested_properties()) {
+    return false;
+  }
+
+  if (old_config.indexable_nested_properties_list().size() !=
+      new_config.indexable_nested_properties_list().size()) {
+    return false;
+  }
+
+  std::unordered_set<std::string_view> old_indexable_nested_properies_set(
+      old_config.indexable_nested_properties_list().begin(),
+      old_config.indexable_nested_properties_list().end());
+  for (const auto& property : new_config.indexable_nested_properties_list()) {
+    if (old_indexable_nested_properies_set.find(property) ==
+        old_indexable_nested_properies_set.end()) {
+      return false;
+    }
+  }
+  return true;
 }
 
 bool AreIntegerIndexingConfigsEqual(const IntegerIndexingConfig& old_config,
@@ -100,7 +122,7 @@
   if (old_property.cardinality() < new_property.cardinality()) {
     // We allow a new, less restrictive cardinality (i.e. a REQUIRED field
     // can become REPEATED or OPTIONAL, but not the other way around).
-    ICING_VLOG(1) << absl_ports::StrCat(
+    ICING_LOG(INFO) << absl_ports::StrCat(
         "Cardinality is more restrictive than before ",
         PropertyConfigProto::Cardinality::Code_Name(old_property.cardinality()),
         "->",
@@ -117,7 +139,7 @@
     // TODO(cassiewang): Maybe we can be a bit looser with this, e.g. we just
     // string cast an int64_t to a string. But for now, we'll stick with
     // simplistics.
-    ICING_VLOG(1) << absl_ports::StrCat(
+    ICING_LOG(INFO) << absl_ports::StrCat(
         "Data type ",
         PropertyConfigProto::DataType::Code_Name(old_property.data_type()),
         "->",
@@ -130,9 +152,9 @@
 bool IsSchemaTypeCompatible(const PropertyConfigProto& old_property,
                             const PropertyConfigProto& new_property) {
   if (old_property.schema_type() != new_property.schema_type()) {
-    ICING_VLOG(1) << absl_ports::StrCat("Schema type ",
-                                        old_property.schema_type(), "->",
-                                        new_property.schema_type());
+    ICING_LOG(INFO) << absl_ports::StrCat("Schema type ",
+                                          old_property.schema_type(), "->",
+                                          new_property.schema_type());
     return false;
   }
   return true;
@@ -145,53 +167,6 @@
          IsCardinalityCompatible(old_property, new_property);
 }
 
-bool IsTermMatchTypeCompatible(const StringIndexingConfig& old_indexed,
-                               const StringIndexingConfig& new_indexed) {
-  return old_indexed.term_match_type() == new_indexed.term_match_type() &&
-         old_indexed.tokenizer_type() == new_indexed.tokenizer_type();
-}
-
-bool IsIntegerNumericMatchTypeCompatible(
-    const IntegerIndexingConfig& old_indexed,
-    const IntegerIndexingConfig& new_indexed) {
-  return old_indexed.numeric_match_type() == new_indexed.numeric_match_type();
-}
-
-bool IsEmbeddingIndexingCompatible(const EmbeddingIndexingConfig& old_indexed,
-                                   const EmbeddingIndexingConfig& new_indexed) {
-  return old_indexed.embedding_indexing_type() ==
-             new_indexed.embedding_indexing_type() &&
-         old_indexed.quantization_type() == new_indexed.quantization_type();
-}
-
-bool IsDocumentIndexingCompatible(const DocumentIndexingConfig& old_indexed,
-                                  const DocumentIndexingConfig& new_indexed) {
-  // TODO(b/265304217): This could mark the new schema as incompatible and
-  // generate some unnecessary index rebuilds if the two schemas have an
-  // equivalent set of indexed properties, but changed the way that it is
-  // declared.
-  if (old_indexed.index_nested_properties() !=
-      new_indexed.index_nested_properties()) {
-    return false;
-  }
-
-  if (old_indexed.indexable_nested_properties_list().size() !=
-      new_indexed.indexable_nested_properties_list().size()) {
-    return false;
-  }
-
-  std::unordered_set<std::string_view> old_indexable_nested_properies_set(
-      old_indexed.indexable_nested_properties_list().begin(),
-      old_indexed.indexable_nested_properties_list().end());
-  for (const auto& property : new_indexed.indexable_nested_properties_list()) {
-    if (old_indexable_nested_properies_set.find(property) ==
-        old_indexable_nested_properies_set.end()) {
-      return false;
-    }
-  }
-  return true;
-}
-
 void AddIncompatibleChangeToDelta(
     std::unordered_set<std::string>& incompatible_delta,
     const SchemaTypeConfigProto& old_type_config,
@@ -1123,10 +1098,12 @@
 
   // TODO(cassiewang): consider caching property_config_map for some properties,
   // e.g. using LRU cache. Or changing schema.proto to use go/protomap.
-  for (const PropertyConfigProto& property_config : type_config.properties()) {
+  for (int position = 0; position < type_config.properties_size(); ++position) {
+    const PropertyConfigProto& property_config =
+        type_config.properties(position);
     std::string_view property_name = property_config.property_name();
-    parsed_property_configs.property_config_map.emplace(property_name,
-                                                        &property_config);
+    parsed_property_configs.property_config_map.emplace(
+        property_name, PropertyConfigInfo{&property_config, position});
     if (property_config.cardinality() ==
         PropertyConfigProto::Cardinality::REQUIRED) {
       parsed_property_configs.required_properties.insert(property_name);
@@ -1206,7 +1183,10 @@
     bool is_incompatible = false;
     bool is_index_incompatible = false;
     bool is_join_incompatible = false;
-    for (const auto& old_property_config : old_type_config.properties()) {
+    for (int position = 0; position < old_type_config.properties_size();
+         ++position) {
+      const PropertyConfigProto& old_property_config =
+          old_type_config.properties(position);
       std::string_view property_name = old_property_config.property_name();
       if (old_property_config.cardinality() ==
           PropertyConfigProto::Cardinality::REQUIRED) {
@@ -1242,7 +1222,7 @@
       if (new_property_name_and_config ==
           new_parsed_property_configs.property_config_map.end()) {
         // Didn't find the old property
-        ICING_VLOG(1) << absl_ports::StrCat(
+        ICING_LOG(INFO) << absl_ports::StrCat(
             "Previously defined property type '", old_type_config.schema_type(),
             ".", old_property_config.property_name(),
             "' was not defined in new schema");
@@ -1254,38 +1234,45 @@
       }
 
       const PropertyConfigProto* new_property_config =
-          new_property_name_and_config->second;
+          new_property_name_and_config->second.property_config;
+      bool property_order_changed =
+          feature_flags.enable_schema_definition_deduping() &&
+          position != new_property_name_and_config->second.position;
       if (!has_property_changed &&
-          !ArePropertiesEqual(old_property_config, *new_property_config)) {
-        // Finally found a property that changed.
+          (!ArePropertiesEqual(old_property_config, *new_property_config) ||
+           property_order_changed)) {
+        // Found a property that changed. A property change is either a
+        // PropertyConfigProto change or (when schema deduping is enabled) a
+        // change in the property's position in the type config's repeated
+        // properties field.
         has_property_changed = true;
       }
 
       if (!IsPropertyCompatible(old_property_config, *new_property_config)) {
-        ICING_VLOG(1) << absl_ports::StrCat(
+        ICING_LOG(INFO) << absl_ports::StrCat(
             "Property '", old_type_config.schema_type(), ".",
             old_property_config.property_name(), "' is incompatible.");
         is_incompatible = true;
       }
 
       // Any change in the indexed property requires a reindexing
-      if (!IsTermMatchTypeCompatible(
+      if (!AreStringIndexingConfigsEqual(
               old_property_config.string_indexing_config(),
               new_property_config->string_indexing_config()) ||
-          !IsIntegerNumericMatchTypeCompatible(
+          !AreIntegerIndexingConfigsEqual(
               old_property_config.integer_indexing_config(),
               new_property_config->integer_indexing_config()) ||
-          !IsDocumentIndexingCompatible(
+          !AreDocumentIndexingConfigsEqual(
               old_property_config.document_indexing_config(),
               new_property_config->document_indexing_config()) ||
-          !IsEmbeddingIndexingCompatible(
+          !AreEmbeddingIndexingConfigsEqual(
               old_property_config.embedding_indexing_config(),
               new_property_config->embedding_indexing_config())) {
         is_index_incompatible = true;
       }
 
-      if (old_property_config.joinable_config().value_type() !=
-          new_property_config->joinable_config().value_type()) {
+      if (!AreJoinableConfigsEqual(old_property_config.joinable_config(),
+                                   new_property_config->joinable_config())) {
         is_join_incompatible = true;
       }
     }
@@ -1297,7 +1284,7 @@
     // here to detect new required properties.
     if (!IsSubset(new_parsed_property_configs.required_properties,
                   old_required_properties)) {
-      ICING_VLOG(1) << absl_ports::StrCat(
+      ICING_LOG(INFO) << absl_ports::StrCat(
           "New schema '", old_type_config.schema_type(),
           "' has REQUIRED properties that are not "
           "present in the previously defined schema");
@@ -1310,9 +1297,9 @@
     // reindex everything.
     if (!IsSubset(new_parsed_property_configs.indexed_properties,
                   old_indexed_properties)) {
-      ICING_VLOG(1) << "Set of indexed properties in schema type '"
-                    << old_type_config.schema_type()
-                    << "' has changed, required reindexing.";
+      ICING_LOG(INFO) << "Set of indexed properties in schema type '"
+                      << old_type_config.schema_type()
+                      << "' has changed, required reindexing.";
       is_index_incompatible = true;
     }
 
@@ -1327,9 +1314,10 @@
                   old_joinable_properties) ||
         !IsSubset(new_parsed_property_configs.nested_document_properties,
                   old_nested_document_properties)) {
-      ICING_VLOG(1) << "Set of joinable properties in schema type '"
-                    << old_type_config.schema_type()
-                    << "' has changed, required reconstructing joinable cache.";
+      ICING_LOG(INFO)
+          << "Set of joinable properties in schema type '"
+          << old_type_config.schema_type()
+          << "' has changed, required reconstructing joinable cache.";
       is_join_incompatible = true;
     }
 
@@ -1351,8 +1339,16 @@
                                    old_type_config_map, new_type_config_map);
     }
 
+    // Scorable-property inconsistent types are already added to the schema
+    // delta in FindScorablePropertyInconsistentTypes above.
+    bool is_scorable_property_cache_incompatible =
+        !schema_delta.schema_types_scorable_property_inconsistent.empty() &&
+        schema_delta.schema_types_scorable_property_inconsistent.find(
+            old_type_config.schema_type()) !=
+            schema_delta.schema_types_scorable_property_inconsistent.end();
+
     if (!is_incompatible && !is_index_incompatible && !is_join_incompatible &&
-        has_property_changed) {
+        !is_scorable_property_cache_incompatible && has_property_changed) {
       schema_delta.schema_types_changed_fully_compatible.insert(
           old_type_config.schema_type());
     }

diff --git a/icing/schema/schema-util.h b/icing/schema/schema-util.h
index 39e6160..fa2e575 100644
--- a/icing/schema/schema-util.h
+++ b/icing/schema/schema-util.h

@@ -15,6 +15,7 @@
 #ifndef ICING_SCHEMA_SCHEMA_UTIL_H_
 #define ICING_SCHEMA_SCHEMA_UTIL_H_
 
+#include <cstdint>
 #include <string>
 #include <string_view>
 #include <unordered_map>
@@ -116,9 +117,19 @@
     }
   };
 
+  // A struct that stores the information about a property config parsed from
+  // a SchemaTypeConfigProto.
+  struct PropertyConfigInfo {
+    const PropertyConfigProto* property_config;
+
+    // The position of the property in the type config's repeated property
+    // field.
+    int32_t position;
+  };
+
   struct ParsedPropertyConfigs {
     // Mapping of property name to PropertyConfigProto
-    std::unordered_map<std::string_view, const PropertyConfigProto*>
+    std::unordered_map<std::string_view, PropertyConfigInfo>
         property_config_map;
 
     // Properties that have an indexing config

diff --git a/icing/schema/schema-util_test.cc b/icing/schema/schema-util_test.cc
index 6758984..745c7dd 100644
--- a/icing/schema/schema-util_test.cc
+++ b/icing/schema/schema-util_test.cc

@@ -2438,6 +2438,55 @@
               Eq(schema_delta));
 }
 
+TEST_P(SchemaUtilTest, PropertyConfigReorderingIsCompatible) {
+  // Configure old schema
+  SchemaProto old_schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType(kEmailType)
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("prop1")
+                                        .SetDataType(TYPE_STRING)
+                                        .SetCardinality(CARDINALITY_REQUIRED))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("prop2")
+                                        .SetDataType(TYPE_INT64)
+                                        .SetCardinality(CARDINALITY_OPTIONAL))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("prop3")
+                                        .SetDataType(TYPE_STRING)
+                                        .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  // Configure new schema with the properties in a different order
+  SchemaProto new_schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder()
+                       .SetType(kEmailType)
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("prop2")
+                                        .SetDataType(TYPE_INT64)
+                                        .SetCardinality(CARDINALITY_OPTIONAL))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("prop1")
+                                        .SetDataType(TYPE_STRING)
+                                        .SetCardinality(CARDINALITY_REQUIRED))
+                       .AddProperty(PropertyConfigBuilder()
+                                        .SetName("prop3")
+                                        .SetDataType(TYPE_STRING)
+                                        .SetCardinality(CARDINALITY_OPTIONAL)))
+          .Build();
+
+  SchemaUtil::SchemaDelta schema_delta;
+  if (feature_flags_->enable_schema_definition_deduping()) {
+    schema_delta.schema_types_changed_fully_compatible.insert(kEmailType);
+  }
+  SchemaUtil::DependentMap no_dependents_map;
+  EXPECT_THAT(SchemaUtil::ComputeCompatibilityDelta(
+                  old_schema, new_schema, no_dependents_map, *feature_flags_),
+              Eq(schema_delta));
+}
+
 TEST_P(SchemaUtilTest, CompatibilityOfDifferentCardinalityOk) {
   // Configure less restrictive schema based on cardinality
   SchemaProto less_restrictive_schema =
@@ -3617,9 +3666,10 @@
               IsEmpty());
 }
 
-TEST_P(SchemaUtilTest, ChangingJoinablePropertiesPropagateDeleteIsCompatible) {
+TEST_P(SchemaUtilTest,
+       ChangingJoinablePropertiesPropagateDeleteIsJoinIncompatible) {
   // Configure old schema
-  SchemaProto old_schema =
+  SchemaProto old_schema_without_delete_propagation =
       SchemaBuilder()
           .AddType(SchemaTypeConfigBuilder().SetType("MyType").AddProperty(
               PropertyConfigBuilder()
@@ -3631,7 +3681,7 @@
           .Build();
 
   // Configure new schema with delete propagation type PROPAGATE_FROM
-  SchemaProto new_schema_with_optional =
+  SchemaProto new_schema_with_delete_propagation =
       SchemaBuilder()
           .AddType(SchemaTypeConfigBuilder().SetType("MyType").AddProperty(
               PropertyConfigBuilder()
@@ -3642,13 +3692,23 @@
                   .SetCardinality(CARDINALITY_REQUIRED)))
           .Build();
 
-  SchemaUtil::SchemaDelta schema_delta;
-  schema_delta.schema_types_changed_fully_compatible.insert("MyType");
+  SchemaUtil::SchemaDelta expected_schema_delta;
+  expected_schema_delta.schema_types_join_incompatible.insert("MyType");
+
+  // New schema enabled delete propagation.
   SchemaUtil::DependentMap no_dependents_map;
   EXPECT_THAT(SchemaUtil::ComputeCompatibilityDelta(
-                  old_schema, new_schema_with_optional, no_dependents_map,
+                  old_schema_without_delete_propagation,
+                  new_schema_with_delete_propagation, no_dependents_map,
                   *feature_flags_),
-              Eq(schema_delta));
+              Eq(expected_schema_delta));
+
+  // New schema disabled a joinable property.
+  EXPECT_THAT(SchemaUtil::ComputeCompatibilityDelta(
+                  old_schema_without_delete_propagation,
+                  new_schema_with_delete_propagation, no_dependents_map,
+                  *feature_flags_),
+              Eq(expected_schema_delta));
 }
 
 TEST_P(SchemaUtilTest, AddingTypeIsCompatible) {
@@ -4436,7 +4496,21 @@
       /*enable_repeated_field_joins=*/false,
       /*enable_embedding_backup_generation=*/true,
       /*enable_schema_database=*/true,
-      /*release_backup_schema_file_if_overlay_present=*/true);
+      /*release_backup_schema_file_if_overlay_present=*/true,
+      /*enable_strict_page_byte_size_limit=*/true,
+      /*enable_smaller_decompression_buffer_size=*/true,
+      /*enable_eigen_embedding_scoring=*/true,
+      /*enable_passing_filter_to_children=*/true,
+      /*enable_proto_log_new_header_format=*/true,
+      /*enable_embedding_iterator_v2=*/true,
+      /*enable_reusable_decompression_buffer=*/true,
+      /*enable_schema_type_id_optimization=*/true,
+      /*enable_optimize_improvements=*/true,
+      /*expired_document_purge_threshold_ms=*/0,
+      /*enable_non_existent_qualified_id_join=*/true,
+      /*enable_skip_set_schema_type_equality_check=*/true,
+      /*enable_embed_query_optimization=*/true,
+      /*enable_schema_definition_deduping=*/true);
   SchemaProto schema =
       SchemaBuilder()
           .AddType(SchemaTypeConfigBuilder().SetType("MyType").AddProperty(
@@ -4498,7 +4572,21 @@
       /*enable_repeated_field_joins=*/true,
       /*enable_embedding_backup_generation=*/true,
       /*enable_schema_database=*/true,
-      /*release_backup_schema_file_if_overlay_present=*/true);
+      /*release_backup_schema_file_if_overlay_present=*/true,
+      /*enable_strict_page_byte_size_limit=*/true,
+      /*enable_smaller_decompression_buffer_size=*/true,
+      /*enable_eigen_embedding_scoring=*/true,
+      /*enable_passing_filter_to_children=*/true,
+      /*enable_proto_log_new_header_format=*/true,
+      /*enable_embedding_iterator_v2=*/true,
+      /*enable_reusable_decompression_buffer=*/true,
+      /*enable_schema_type_id_optimization=*/true,
+      /*enable_optimize_improvements=*/true,
+      /*expired_document_purge_threshold_ms=*/0,
+      /*enable_non_existent_qualified_id_join=*/true,
+      /*enable_skip_set_schema_type_equality_check=*/true,
+      /*enable_embed_query_optimization=*/true,
+      /*enable_schema_definition_deduping=*/true);
 
   SchemaProto schema =
       SchemaBuilder()
@@ -5818,23 +5906,50 @@
 
 INSTANTIATE_TEST_SUITE_P(
     SchemaUtilTest, SchemaUtilTest,
-    testing::Values(
-        FeatureFlags(
-            /*enable_circular_schema_definitions=*/false,
-            /*enable_scorable_properties=*/true,
-            /*enable_embedding_quantization=*/true,
-            /*enable_repeated_field_joins=*/true,
-            /*enable_embedding_backup_generation=*/true,
-            /*enable_schema_database=*/true,
-            /*release_backup_schema_file_if_overlay_present=*/true),
-        FeatureFlags(
-            /*enable_circular_schema_definitions=*/true,
-            /*enable_scorable_properties=*/true,
-            /*enable_embedding_quantization=*/true,
-            /*enable_repeated_field_joins=*/true,
-            /*enable_embedding_backup_generation=*/true,
-            /*enable_schema_database=*/true,
-            /*release_backup_schema_file_if_overlay_present=*/true)));
+    testing::Values(FeatureFlags(
+                        /*enable_circular_schema_definitions=*/false,
+                        /*enable_scorable_properties=*/true,
+                        /*enable_embedding_quantization=*/true,
+                        /*enable_repeated_field_joins=*/true,
+                        /*enable_embedding_backup_generation=*/true,
+                        /*enable_schema_database=*/true,
+                        /*release_backup_schema_file_if_overlay_present=*/true,
+                        /*enable_strict_page_byte_size_limit=*/true,
+                        /*enable_smaller_decompression_buffer_size=*/true,
+                        /*enable_eigen_embedding_scoring=*/true,
+                        /*enable_passing_filter_to_children=*/true,
+                        /*enable_proto_log_new_header_format=*/true,
+                        /*enable_embedding_iterator_v2=*/true,
+                        /*enable_reusable_decompression_buffer=*/true,
+                        /*enable_schema_type_id_optimization=*/true,
+                        /*enable_optimize_improvements=*/true,
+                        /*expired_document_purge_threshold_ms=*/0,
+                        /*enable_non_existent_qualified_id_join=*/true,
+                        /*enable_skip_set_schema_type_equality_check=*/true,
+                        /*enable_embed_query_optimization=*/true,
+                        /*enable_schema_definition_deduping=*/false),
+                    FeatureFlags(
+                        /*enable_circular_schema_definitions=*/true,
+                        /*enable_scorable_properties=*/true,
+                        /*enable_embedding_quantization=*/true,
+                        /*enable_repeated_field_joins=*/true,
+                        /*enable_embedding_backup_generation=*/true,
+                        /*enable_schema_database=*/true,
+                        /*release_backup_schema_file_if_overlay_present=*/true,
+                        /*enable_strict_page_byte_size_limit=*/true,
+                        /*enable_smaller_decompression_buffer_size=*/true,
+                        /*enable_eigen_embedding_scoring=*/true,
+                        /*enable_passing_filter_to_children=*/true,
+                        /*enable_proto_log_new_header_format=*/true,
+                        /*enable_embedding_iterator_v2=*/true,
+                        /*enable_reusable_decompression_buffer=*/true,
+                        /*enable_schema_type_id_optimization=*/true,
+                        /*enable_optimize_improvements=*/true,
+                        /*expired_document_purge_threshold_ms=*/0,
+                        /*enable_non_existent_qualified_id_join=*/true,
+                        /*enable_skip_set_schema_type_equality_check=*/true,
+                        /*enable_embed_query_optimization=*/true,
+                        /*enable_schema_definition_deduping=*/true)));
 
 struct IsIndexedPropertyTestParam {
   PropertyConfigProto property_config;

diff --git a/icing/scoring/advanced_scoring/advanced-scorer_fuzz_test.cc b/icing/scoring/advanced_scoring/advanced-scorer_fuzz_test.cc
index 765ccdb..c02f06b 100644
--- a/icing/scoring/advanced_scoring/advanced-scorer_fuzz_test.cc
+++ b/icing/scoring/advanced_scoring/advanced-scorer_fuzz_test.cc

@@ -22,6 +22,7 @@
 #include "icing/file/filesystem.h"
 #include "icing/file/portable-file-backed-proto-log.h"
 #include "icing/index/embed/embedding-query-results.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/schema/schema-store.h"
 #include "icing/scoring/advanced_scoring/advanced-scorer.h"
 #include "icing/store/document-store.h"
@@ -55,6 +56,9 @@
           /*force_recovery_and_revalidate_documents=*/false,
           /*pre_mapping_fbv=*/false, /*use_persistent_hash_map=*/true,
           PortableFileBackedProtoLog<DocumentWrapper>::kDefaultCompressionLevel,
+          PortableFileBackedProtoLog<
+              DocumentWrapper>::kDefaultCompressionThresholdBytes,
+          protobuf_ports::kDefaultMemLevel,
           /*initialize_stats=*/nullptr)
           .ValueOrDie()
           .document_store;

diff --git a/icing/scoring/advanced_scoring/advanced-scorer_test.cc b/icing/scoring/advanced_scoring/advanced-scorer_test.cc
index 386d0c2..a1eb58a 100644
--- a/icing/scoring/advanced_scoring/advanced-scorer_test.cc
+++ b/icing/scoring/advanced_scoring/advanced-scorer_test.cc

@@ -35,6 +35,7 @@
 #include "icing/index/embed/embedding-query-results.h"
 #include "icing/index/hit/doc-hit-info.h"
 #include "icing/join/join-children-fetcher-impl-deprecated.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/schema.pb.h"
 #include "icing/proto/scoring.pb.h"
@@ -48,9 +49,11 @@
 #include "icing/store/document-id.h"
 #include "icing/store/document-store.h"
 #include "icing/testing/common-matchers.h"
+#include "icing/testing/embedding-test-utils.h"
 #include "icing/testing/fake-clock.h"
 #include "icing/testing/test-feature-flags.h"
 #include "icing/testing/tmp-directory.h"
+#include "icing/util/document-util.h"
 
 namespace icing {
 namespace lib {
@@ -80,14 +83,18 @@
 
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
-        DocumentStore::Create(&filesystem_, doc_store_dir_, &fake_clock_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, doc_store_dir_, &fake_clock_, schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     document_store_ = std::move(create_result.document_store);
 
     // Creates the schema
@@ -286,7 +293,8 @@
 TEST_F(AdvancedScorerTest, SimpleExpression) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result,
-      document_store_->Put(CreateDocument("namespace", "uri")));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          CreateDocument("namespace", "uri"))));
   DocumentId document_id = put_result.new_document_id;
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -307,7 +315,8 @@
 TEST_F(AdvancedScorerTest, BasicPureArithmeticExpression) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result,
-      document_store_->Put(CreateDocument("namespace", "uri")));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          CreateDocument("namespace", "uri"))));
   DocumentId document_id = put_result.new_document_id;
   DocHitInfo docHitInfo = DocHitInfo(document_id);
 
@@ -376,7 +385,8 @@
 TEST_F(AdvancedScorerTest, BasicMathFunctionExpression) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result,
-      document_store_->Put(CreateDocument("namespace", "uri")));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          CreateDocument("namespace", "uri"))));
   DocumentId document_id = put_result.new_document_id;
   DocHitInfo docHitInfo = DocHitInfo(document_id);
 
@@ -515,9 +525,9 @@
 TEST_F(AdvancedScorerTest, DocumentScoreCreationTimestampFunctionExpression) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result,
-      document_store_->Put(CreateDocument(
+      document_store_->Put(document_util::CreateDocumentWrapper(CreateDocument(
           "namespace", "uri", /*score=*/123,
-          /*creation_timestamp_ms=*/kDefaultCreationTimestampMs)));
+          /*creation_timestamp_ms=*/kDefaultCreationTimestampMs))));
   DocumentId document_id = put_result.new_document_id;
   DocHitInfo docHitInfo = DocHitInfo(document_id);
 
@@ -558,7 +568,8 @@
 TEST_F(AdvancedScorerTest, DocumentUsageFunctionExpression) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result,
-      document_store_->Put(CreateDocument("namespace", "uri")));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          CreateDocument("namespace", "uri"))));
   DocumentId document_id = put_result.new_document_id;
   DocHitInfo docHitInfo = DocHitInfo(document_id);
 
@@ -615,7 +626,8 @@
 TEST_F(AdvancedScorerTest, DocumentUsageFunctionOutOfRange) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result,
-      document_store_->Put(CreateDocument("namespace", "uri")));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          CreateDocument("namespace", "uri"))));
   DocumentId document_id = put_result.new_document_id;
   DocHitInfo docHitInfo = DocHitInfo(document_id);
 
@@ -667,8 +679,10 @@
           .SetCreationTimestampMs(kDefaultCreationTimestampMs)
           .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(test_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(test_document)));
   DocumentId document_id = put_result.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<AdvancedScorer> scorer,
@@ -691,17 +705,20 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result1,
-      document_store_->Put(CreateDocument("namespace", "uri1")));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          CreateDocument("namespace", "uri1"))));
   DocumentId document_id_1 = put_result1.new_document_id;
   DocHitInfo docHitInfo1 = DocHitInfo(document_id_1);
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result2,
-      document_store_->Put(CreateDocument("namespace", "uri2")));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          CreateDocument("namespace", "uri2"))));
   DocumentId document_id_2 = put_result2.new_document_id;
   DocHitInfo docHitInfo2 = DocHitInfo(document_id_2);
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result3,
-      document_store_->Put(CreateDocument("namespace", "uri3")));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          CreateDocument("namespace", "uri3"))));
   DocumentId document_id_3 = put_result3.new_document_id;
   DocHitInfo docHitInfo3 = DocHitInfo(document_id_3);
 
@@ -803,14 +820,20 @@
   DocumentProto test_document_3 =
       DocumentBuilder().SetKey("namespace", "uri3").SetSchema("person").Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(test_document_1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(test_document_1)));
   DocumentId document_id_1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(test_document_2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(test_document_2)));
   DocumentId document_id_2 = put_result2.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             document_store_->Put(test_document_3));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(test_document_3)));
   DocumentId document_id_3 = put_result3.new_document_id;
 
   ScoringSpecProto spec_proto = CreateAdvancedScoringSpec("");
@@ -900,11 +923,15 @@
   DocumentProto test_document_2 =
       DocumentBuilder().SetKey("namespace", "uri2").SetSchema("person").Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(test_document_1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(test_document_1)));
   DocumentId document_id_1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(test_document_2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(test_document_2)));
   DocumentId document_id_2 = put_result2.new_document_id;
 
   ScoringSpecProto spec_proto = CreateAdvancedScoringSpec("");
@@ -1012,8 +1039,8 @@
   const int64_t creation_timestamp_ms = 123;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result,
-      document_store_->Put(CreateDocument("namespace", "uri", /*score=*/123,
-                                          creation_timestamp_ms)));
+      document_store_->Put(document_util::CreateDocumentWrapper(CreateDocument(
+          "namespace", "uri", /*score=*/123, creation_timestamp_ms))));
   DocumentId document_id = put_result.new_document_id;
   DocHitInfo docHitInfo = DocHitInfo(document_id);
 
@@ -1133,7 +1160,8 @@
   // non-constant 0.
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result,
-      document_store_->Put(CreateDocument("namespace", "uri", /*score=*/0)));
+      document_store_->Put(document_util::CreateDocumentWrapper(
+          CreateDocument("namespace", "uri", /*score=*/0))));
   DocumentId document_id = put_result.new_document_id;
   DocHitInfo docHitInfo = DocHitInfo(document_id);
 
@@ -1363,12 +1391,12 @@
 
 TEST_F(AdvancedScorerTest,
        MatchedSemanticScoresFunctionScoreExpressionTypeError) {
-  EmbeddingQueryResults embedding_query_results;
-  embedding_query_results
-      .result_infos[/*query_vector_index=*/0]
-                    [SearchSpecProto::EmbeddingQueryMetricType::COSINE]
-                    [/*document_id=*/0]
-      .AppendScore(/*semantic_score=*/0.1);
+  EmbeddingQueryResults embedding_query_results(/*num_query_vectors=*/1);
+  GetOrCreateEmbeddingMatchInfosForDocument(
+      embedding_query_results, /*query_vector_index=*/0,
+      SearchSpecProto::EmbeddingQueryMetricType::COSINE, /*document_id=*/0)
+      .AppendScore(*embedding_query_results.global_scores,
+                   /*semantic_score=*/0.1);
 
   libtextclassifier3::StatusOr<std::unique_ptr<AdvancedScorer>> scorer_or =
       AdvancedScorer::Create(
@@ -1464,17 +1492,17 @@
 
 TEST_F(AdvancedScorerTest,
        MatchedSemanticScoresFunctionScoreExpressionNotQueried) {
-  EmbeddingQueryResults embedding_query_results;
-  embedding_query_results
-      .result_infos[/*query_vector_index=*/0]
-                    [SearchSpecProto::EmbeddingQueryMetricType::COSINE]
-                    [/*document_id=*/0]
-      .AppendScore(/*semantic_score=*/0.1);
-  embedding_query_results
-      .result_infos[/*query_vector_index=*/1]
-                    [SearchSpecProto::EmbeddingQueryMetricType::DOT_PRODUCT]
-                    [/*document_id=*/1]
-      .AppendScore(/*semantic_score=*/0.2);
+  EmbeddingQueryResults embedding_query_results(/*num_query_vectors=*/2);
+  GetOrCreateEmbeddingMatchInfosForDocument(
+      embedding_query_results, /*query_vector_index=*/0,
+      SearchSpecProto::EmbeddingQueryMetricType::COSINE, /*document_id=*/0)
+      .AppendScore(*embedding_query_results.global_scores,
+                   /*semantic_score=*/0.1);
+  GetOrCreateEmbeddingMatchInfosForDocument(
+      embedding_query_results, /*query_vector_index=*/1,
+      SearchSpecProto::EmbeddingQueryMetricType::DOT_PRODUCT, /*document_id=*/1)
+      .AppendScore(*embedding_query_results.global_scores,
+                   /*semantic_score=*/0.2);
 
   libtextclassifier3::StatusOr<std::unique_ptr<AdvancedScorer>> scorer_or =
       AdvancedScorer::Create(CreateAdvancedScoringSpec(
@@ -1564,9 +1592,10 @@
 }
 
 void AddEntryToEmbeddingQueryScoreMap(
+    std::vector<double>& global_scores,
     EmbeddingQueryResults::EmbeddingQueryMatchInfoMap& score_map,
     double semantic_score, DocumentId document_id) {
-  score_map[document_id].AppendScore(semantic_score);
+  score_map[document_id].AppendScore(global_scores, semantic_score);
 }
 
 TEST_F(AdvancedScorerTest, MatchedSemanticScoresFunctionScoreExpression) {
@@ -1574,7 +1603,7 @@
   DocumentId document_id_1 = 1;
   DocHitInfo doc_hit_info_0(document_id_0);
   DocHitInfo doc_hit_info_1(document_id_1);
-  EmbeddingQueryResults embedding_query_results;
+  EmbeddingQueryResults embedding_query_results(/*num_query_vectors=*/2);
 
   // Let the first query assign the following semantic scores:
   // COSINE:
@@ -1586,41 +1615,57 @@
   // EUCLIDEAN:
   //   Document 0: 0.7
   //   Document 1: 0.8
-  EmbeddingQueryResults::EmbeddingQueryMatchInfoMap* score_map =
-      &embedding_query_results
-           .result_infos[0][SearchSpecProto::EmbeddingQueryMetricType::COSINE];
-  AddEntryToEmbeddingQueryScoreMap(*score_map,
+  ICING_ASSERT_OK_AND_ASSIGN(
+      EmbeddingQueryResults::EmbeddingQueryMatchInfoMap * score_map,
+      embedding_query_results.GetOrCreateMatchInfoMap(
+          /*query_vector_index=*/0,
+          SearchSpecProto::EmbeddingQueryMetricType::COSINE));
+  AddEntryToEmbeddingQueryScoreMap(*embedding_query_results.global_scores,
+                                   *score_map,
                                    /*semantic_score=*/0.1, document_id_0);
-  AddEntryToEmbeddingQueryScoreMap(*score_map,
+  AddEntryToEmbeddingQueryScoreMap(*embedding_query_results.global_scores,
+                                   *score_map,
                                    /*semantic_score=*/0.2, document_id_0);
-  AddEntryToEmbeddingQueryScoreMap(*score_map,
+  AddEntryToEmbeddingQueryScoreMap(*embedding_query_results.global_scores,
+                                   *score_map,
                                    /*semantic_score=*/0.3, document_id_1);
-  AddEntryToEmbeddingQueryScoreMap(*score_map,
+  AddEntryToEmbeddingQueryScoreMap(*embedding_query_results.global_scores,
+                                   *score_map,
                                    /*semantic_score=*/0.4, document_id_1);
-  score_map = &embedding_query_results.result_infos
-                   [0][SearchSpecProto::EmbeddingQueryMetricType::DOT_PRODUCT];
-  AddEntryToEmbeddingQueryScoreMap(*score_map,
+  ICING_ASSERT_OK_AND_ASSIGN(
+      score_map, embedding_query_results.GetOrCreateMatchInfoMap(
+                     /*query_vector_index=*/0,
+                     SearchSpecProto::EmbeddingQueryMetricType::DOT_PRODUCT));
+  AddEntryToEmbeddingQueryScoreMap(*embedding_query_results.global_scores,
+                                   *score_map,
                                    /*semantic_score=*/0.5, document_id_0);
-  AddEntryToEmbeddingQueryScoreMap(*score_map,
+  AddEntryToEmbeddingQueryScoreMap(*embedding_query_results.global_scores,
+                                   *score_map,
                                    /*semantic_score=*/0.6, document_id_1);
-  score_map =
-      &embedding_query_results
-           .result_infos[0]
-                         [SearchSpecProto::EmbeddingQueryMetricType::EUCLIDEAN];
-  AddEntryToEmbeddingQueryScoreMap(*score_map,
+  ICING_ASSERT_OK_AND_ASSIGN(
+      score_map, embedding_query_results.GetOrCreateMatchInfoMap(
+                     /*query_vector_index=*/0,
+                     SearchSpecProto::EmbeddingQueryMetricType::EUCLIDEAN));
+  AddEntryToEmbeddingQueryScoreMap(*embedding_query_results.global_scores,
+                                   *score_map,
                                    /*semantic_score=*/0.7, document_id_0);
-  AddEntryToEmbeddingQueryScoreMap(*score_map,
+  AddEntryToEmbeddingQueryScoreMap(*embedding_query_results.global_scores,
+                                   *score_map,
                                    /*semantic_score=*/0.8, document_id_1);
 
   // Let the second query only assign DOT_PRODUCT scores:
   // DOT_PRODUCT:
   //   Document 0: 0.1
   //   Document 1: 0.2
-  score_map = &embedding_query_results.result_infos
-                   [1][SearchSpecProto::EmbeddingQueryMetricType::DOT_PRODUCT];
-  AddEntryToEmbeddingQueryScoreMap(*score_map,
+  ICING_ASSERT_OK_AND_ASSIGN(
+      score_map, embedding_query_results.GetOrCreateMatchInfoMap(
+                     /*query_vector_index=*/1,
+                     SearchSpecProto::EmbeddingQueryMetricType::DOT_PRODUCT));
+  AddEntryToEmbeddingQueryScoreMap(*embedding_query_results.global_scores,
+                                   *score_map,
                                    /*semantic_score=*/0.1, document_id_0);
-  AddEntryToEmbeddingQueryScoreMap(*score_map,
+  AddEntryToEmbeddingQueryScoreMap(*embedding_query_results.global_scores,
+                                   *score_map,
                                    /*semantic_score=*/0.2, document_id_1);
 
   // Get semantic scores for default metric (DOT_PRODUCT) for the first query.
@@ -1708,6 +1753,75 @@
                         "has not been queried"));
 }
 
+// This test should be very uncommon, but we should still make sure it works.
+TEST_F(
+    AdvancedScorerTest,
+    MatchedSemanticScoresFunctionScoreExpression_NonConstantQueryVectorIndex) {
+  // Create document 0 with document score 0.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(CreateDocument(
+          "namespace", "uri0", /*score=*/0, kDefaultCreationTimestampMs))));
+  DocumentId document_id_0 = put_result.new_document_id;
+  DocHitInfo doc_hit_info_0(document_id_0);
+
+  // Create document 1 with document score 1.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(CreateDocument(
+          "namespace", "uri1", /*score=*/1, kDefaultCreationTimestampMs))));
+  DocumentId document_id_1 = put_result.new_document_id;
+  DocHitInfo doc_hit_info_1(document_id_1);
+
+  EmbeddingQueryResults embedding_query_results(/*num_query_vectors=*/2);
+
+  // Create the first embedding query with scores:
+  // - Document 0: 0.1
+  // - Document 1: 0.2
+  // Create the second embedding query with scores:
+  // - Document 0: -0.1
+  // - Document 1: -0.2
+  ICING_ASSERT_OK_AND_ASSIGN(
+      EmbeddingQueryResults::EmbeddingQueryMatchInfoMap * score_map,
+      embedding_query_results.GetOrCreateMatchInfoMap(
+          /*query_vector_index=*/0,
+          SearchSpecProto::EmbeddingQueryMetricType::COSINE));
+  AddEntryToEmbeddingQueryScoreMap(*embedding_query_results.global_scores,
+                                   *score_map,
+                                   /*semantic_score=*/0.1, document_id_0);
+  AddEntryToEmbeddingQueryScoreMap(*embedding_query_results.global_scores,
+                                   *score_map,
+                                   /*semantic_score=*/0.2, document_id_1);
+  ICING_ASSERT_OK_AND_ASSIGN(
+      score_map, embedding_query_results.GetOrCreateMatchInfoMap(
+                     /*query_vector_index=*/1,
+                     SearchSpecProto::EmbeddingQueryMetricType::COSINE));
+  AddEntryToEmbeddingQueryScoreMap(*embedding_query_results.global_scores,
+                                   *score_map,
+                                   /*semantic_score=*/-0.1, document_id_0);
+  AddEntryToEmbeddingQueryScoreMap(*embedding_query_results.global_scores,
+                                   *score_map,
+                                   /*semantic_score=*/-0.2, document_id_1);
+
+  // Create a weird scoring expression that has a non-constant query vector
+  // index. Specifically, document 0 should get the scores from the first query,
+  // but document 1 should get the scores from the second query.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<Scorer> scorer,
+      AdvancedScorer::Create(
+          CreateAdvancedScoringSpec(
+              "sum(this.matchedSemanticScores("
+              "getEmbeddingParameter(this.documentScore())))"),
+          kDefaultScore, /*default_semantic_metric_type=*/
+          SearchSpecProto::EmbeddingQueryMetricType::COSINE,
+          document_store_.get(), schema_store_.get(),
+          fake_clock_.GetSystemTimeMilliseconds(),
+          /*join_children_fetcher=*/nullptr, &embedding_query_results,
+          feature_flags_.get()));
+  EXPECT_THAT(scorer->GetScore(doc_hit_info_0), DoubleNear(0.1, kEps));
+  EXPECT_THAT(scorer->GetScore(doc_hit_info_1), DoubleNear(-0.2, kEps));
+}
+
 TEST_F(AdvancedScorerTest, ListRelatedFunctions) {
   DocumentId document_id_0 = 0;
   DocHitInfo doc_hit_info_0(document_id_0);
@@ -1717,23 +1831,31 @@
   //   {4, 5, 2, 1, 3}.
   // - this.matchedSemanticScores(getEmbeddingParameter(1)) returns an empty
   //   list.
-  EmbeddingQueryResults embedding_query_results;
-  EmbeddingQueryResults::EmbeddingQueryMatchInfoMap* score_map =
-      &embedding_query_results
-           .result_infos[0][SearchSpecProto::EmbeddingQueryMetricType::COSINE];
-  AddEntryToEmbeddingQueryScoreMap(*score_map,
+  EmbeddingQueryResults embedding_query_results(/*num_query_vectors=*/2);
+  ICING_ASSERT_OK_AND_ASSIGN(
+      EmbeddingQueryResults::EmbeddingQueryMatchInfoMap * score_map,
+      embedding_query_results.GetOrCreateMatchInfoMap(
+          /*query_vector_index=*/0,
+          SearchSpecProto::EmbeddingQueryMetricType::COSINE));
+  AddEntryToEmbeddingQueryScoreMap(*embedding_query_results.global_scores,
+                                   *score_map,
                                    /*semantic_score=*/4, document_id_0);
-  AddEntryToEmbeddingQueryScoreMap(*score_map,
+  AddEntryToEmbeddingQueryScoreMap(*embedding_query_results.global_scores,
+                                   *score_map,
                                    /*semantic_score=*/5, document_id_0);
-  AddEntryToEmbeddingQueryScoreMap(*score_map,
+  AddEntryToEmbeddingQueryScoreMap(*embedding_query_results.global_scores,
+                                   *score_map,
                                    /*semantic_score=*/2, document_id_0);
-  AddEntryToEmbeddingQueryScoreMap(*score_map,
+  AddEntryToEmbeddingQueryScoreMap(*embedding_query_results.global_scores,
+                                   *score_map,
                                    /*semantic_score=*/1, document_id_0);
-  AddEntryToEmbeddingQueryScoreMap(*score_map,
+  AddEntryToEmbeddingQueryScoreMap(*embedding_query_results.global_scores,
+                                   *score_map,
                                    /*semantic_score=*/3, document_id_0);
-  score_map =
-      &embedding_query_results
-           .result_infos[1][SearchSpecProto::EmbeddingQueryMetricType::COSINE];
+  ICING_ASSERT_OK_AND_ASSIGN(
+      score_map, embedding_query_results.GetOrCreateMatchInfoMap(
+                     /*query_vector_index=*/1,
+                     SearchSpecProto::EmbeddingQueryMetricType::COSINE));
 
   // maxOrDefault({4, 5, 2, 1, 3}, 100) = 5
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -1810,8 +1932,8 @@
   const int64_t creation_timestamp_ms = 123;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result,
-      document_store_->Put(CreateDocument("namespace", "uri", /*score=*/123,
-                                          creation_timestamp_ms)));
+      document_store_->Put(document_util::CreateDocumentWrapper(CreateDocument(
+          "namespace", "uri", /*score=*/123, creation_timestamp_ms))));
   DocumentId document_id = put_result.new_document_id;
   DocHitInfo docHitInfo = DocHitInfo(document_id);
 
@@ -1941,8 +2063,9 @@
           .AddDoubleProperty("frequencyScore", 1.0, 2.0, 3.0)
           .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document)));
   DocHitInfo docHitInfo(put_result.new_document_id);
 
   ScoringSpecProto scoring_spec = CreateAdvancedScoringSpec(
@@ -1977,8 +2100,9 @@
           .AddDoubleProperty("frequencyScore", 1.0, 2.0, 3.0)
           .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document)));
   DocHitInfo docHitInfo(put_result.new_document_id);
 
   ScoringSpecProto scoring_spec = CreateAdvancedScoringSpec(
@@ -2013,8 +2137,9 @@
           .AddDoubleProperty("frequencyScore", 1.0, 2.0, 3.0)
           .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document)));
   DocHitInfo docHitInfo(put_result.new_document_id);
   ScoringSpecProto scoring_spec = CreateAdvancedScoringSpec(
       "this.documentScore() + "
@@ -2047,8 +2172,9 @@
                                .SetScore(100)
                                .SetCreationTimestampMs(123)
                                .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document)));
   DocHitInfo docHitInfo(put_result.new_document_id);
 
   ScoringSpecProto scoring_spec = CreateAdvancedScoringSpec(
@@ -2080,8 +2206,9 @@
                                .SetScore(100)
                                .SetCreationTimestampMs(123)
                                .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document)));
   DocHitInfo docHitInfo(put_result.new_document_id);
 
   // getScorableProperty("aliasPerson", "frequencyScore") will return an empty
@@ -2116,8 +2243,9 @@
                                .SetScore(100)
                                .SetCreationTimestampMs(123)
                                .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document)));
   DocHitInfo docHitInfo(put_result.new_document_id);
 
   // getScorableProperty("aliasPerson", "frequencyScore") will return an empty
@@ -2188,12 +2316,16 @@
   double expected_score_doc1 = 100 + (1 + 2 + 3);
   double expected_score_doc2 = 100 + 0;
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store_->Put(document_from_db1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(document_from_db1)));
   DocHitInfo docHitInfo1(put_result1.new_document_id);
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store_->Put(document_from_db2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store_->Put(
+          document_util::CreateDocumentWrapper(document_from_db2)));
   DocHitInfo docHitInfo2(put_result2.new_document_id);
 
   ScoringSpecProto scoring_spec = CreateAdvancedScoringSpec(
@@ -2230,8 +2362,9 @@
           .AddDoubleProperty("frequencyScore", 1.0, 2.0, 3.0)
           .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document)));
   DocHitInfo docHitInfo(put_result.new_document_id);
 
   ScoringSpecProto scoring_spec = CreateAdvancedScoringSpec(
@@ -2265,8 +2398,9 @@
                                .AddInt64Property("contactTimes", 10, 20, 30)
                                .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document)));
   DocHitInfo docHitInfo(put_result.new_document_id);
 
   ScoringSpecProto scoring_spec = CreateAdvancedScoringSpec(
@@ -2301,8 +2435,9 @@
                                .AddInt64Property("contactTimes", 10, 20, 30)
                                .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document)));
   DocHitInfo docHitInfo(put_result.new_document_id);
 
   ScoringSpecProto scoring_spec = CreateAdvancedScoringSpec(
@@ -2340,8 +2475,9 @@
           .AddDoubleProperty("frequencyScore", 1.0, 2.0, 3.0)
           .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document)));
   DocHitInfo docHitInfo(put_result.new_document_id);
 
   ScoringSpecProto scoring_spec = CreateAdvancedScoringSpec(
@@ -2381,8 +2517,9 @@
           .AddDoubleProperty("frequencyScore", 1.0, 2.0, 3.0)
           .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document)));
   DocHitInfo docHitInfo(put_result.new_document_id);
 
   ScoringSpecProto scoring_spec = CreateAdvancedScoringSpec(
@@ -2420,8 +2557,9 @@
                                .SetCreationTimestampMs(123)
                                .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document)));
   DocHitInfo docHitInfo(put_result.new_document_id);
 
   // Expected score will fall back to the default score, because max() throws an
@@ -2456,8 +2594,9 @@
                                .SetScore(100)
                                .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document)));
   DocHitInfo docHitInfo(put_result.new_document_id);
 
   ScoringSpecProto scoring_spec = CreateAdvancedScoringSpec(
@@ -2589,8 +2728,9 @@
           .SetCreationTimestampMs(0)
           .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document)));
   DocHitInfo docHitInfo(put_result.new_document_id);
   ScoringSpecProto scoring_spec1 = CreateAdvancedScoringSpec(
       "this.documentScore() + "
@@ -2666,8 +2806,9 @@
           .AddDoubleProperty("frequencyScore", 1.0, 2.0, 3.0)
           .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store_->Put(document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store_->Put(document_util::CreateDocumentWrapper(document)));
   DocHitInfo docHitInfo(put_result.new_document_id);
 
   ScoringSpecProto scoring_spec = CreateAdvancedScoringSpec(

diff --git a/icing/scoring/advanced_scoring/double-list.h b/icing/scoring/advanced_scoring/double-list.h
new file mode 100644
index 0000000..6abc176
--- /dev/null
+++ b/icing/scoring/advanced_scoring/double-list.h

@@ -0,0 +1,99 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef ICING_SCORING_ADVANCED_SCORING_DOUBLE_LIST_H_
+#define ICING_SCORING_ADVANCED_SCORING_DOUBLE_LIST_H_
+
+#include <cstddef>
+#include <type_traits>
+#include <utility>
+#include <variant>
+#include <vector>
+
+// Represents the kDoubleList type, which can either own its data or provide a
+// non-owning view of existing data.
+class DoubleList {
+ public:
+  using value_type = double;
+
+  // Creates a list by taking ownership of an existing vector's data via move.
+  explicit DoubleList(std::vector<double>&& vec) : storage_(std::move(vec)) {}
+
+  // Creates a non-owning view of an external data buffer.
+  // The caller must ensure the lifetime of "data" exceeds the lifetime of this
+  // DoubleList.
+  explicit DoubleList(const double* data, size_t size)
+      : storage_(DataView(data, size)) {}
+
+  // Creates an empty list. Represents as an empty non-owning view by default.
+  explicit DoubleList() : storage_(DataView()) {}
+
+  // Disallow copy but allow move.
+  DoubleList(const DoubleList&) = delete;
+  DoubleList& operator=(const DoubleList&) = delete;
+  DoubleList(DoubleList&&) = default;
+  DoubleList& operator=(DoubleList&&) = default;
+
+  const double* data() const {
+    return std::visit(
+        [](const auto& arg) -> const double* { return arg.data(); }, storage_);
+  }
+
+  size_t size() const {
+    return std::visit([](const auto& arg) -> size_t { return arg.size(); },
+                      storage_);
+  }
+
+  const double* begin() const { return data(); }
+
+  const double* end() const { return data() + size(); }
+
+  bool empty() const { return size() == 0; }
+
+  // Releases the ownership of the internal vector, if owned.
+  // If not owned, returns a *new* vector containing a copy of the viewed data.
+  std::vector<double> ReleaseVector() && {
+    return std::visit(
+        [](auto&& arg) -> std::vector<double> {
+          using T = std::decay_t<decltype(arg)>;
+          if constexpr (std::is_same_v<T, std::vector<double>>) {
+            return std::forward<decltype(arg)>(arg);
+          } else {
+            return std::vector<double>(arg.data(), arg.data() + arg.size());
+          }
+        },
+        std::move(storage_));
+  }
+
+ private:
+  // Simple class to represent a non-owning view of data
+  class DataView {
+   public:
+    // Default constructor represents an empty non-owning view.
+    DataView() : DataView(nullptr, 0) {}
+    DataView(const double* data, size_t size) : data_(data), size_(size) {}
+
+    const double* data() const { return data_; }
+    size_t size() const { return size_; }
+
+   private:
+    const double* data_ = nullptr;
+    size_t size_ = 0;
+  };
+
+  // Storage can be either an owned vector or a non-owning view
+  std::variant<std::vector<double>, DataView> storage_;
+};
+
+#endif  // ICING_SCORING_ADVANCED_SCORING_DOUBLE_LIST_H_

diff --git a/icing/scoring/advanced_scoring/double-list_test.cc b/icing/scoring/advanced_scoring/double-list_test.cc
new file mode 100644
index 0000000..f29fe69
--- /dev/null
+++ b/icing/scoring/advanced_scoring/double-list_test.cc

@@ -0,0 +1,245 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "icing/scoring/advanced_scoring/double-list.h"
+
+#include <cstddef>
+#include <iterator>
+#include <numeric>
+#include <utility>
+#include <vector>
+
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+namespace icing {
+namespace lib {
+
+namespace {
+
+using ::testing::ElementsAreArray;
+using ::testing::IsEmpty;
+
+std::vector<double> CreateSequenceVec(size_t size, double start = 0.0) {
+  std::vector<double> vec(size);
+  std::iota(vec.begin(), vec.end(), start);
+  return vec;
+}
+
+TEST(DoubleListTest, DefaultConstructorCreatesEmptyView) {
+  DoubleList list;
+  EXPECT_EQ(list.size(), 0);
+  EXPECT_TRUE(list.empty());
+  EXPECT_EQ(list.data(), nullptr);
+  EXPECT_EQ(list.begin(), list.end());
+}
+
+// Test constructor taking ownership of an empty vector
+TEST(DoubleListTest, EmptyVectorMoveConstructor) {
+  std::vector<double> source_vec;
+  const double* expected_data_ptr =
+      source_vec.data();  // Capture pointer before move
+  DoubleList list(std::move(source_vec));
+
+  EXPECT_EQ(list.size(), 0);
+  EXPECT_TRUE(list.empty());
+  EXPECT_EQ(list.data(), expected_data_ptr);
+  EXPECT_EQ(list.begin(), list.end());
+}
+
+// Test constructor creating an empty non-owning view, which is the same
+// behavior of the default constructor.
+TEST(DoubleListTest, EmptyDataViewConstructor) {
+  const double* external_data = nullptr;
+  DoubleList list(external_data, 0);
+  EXPECT_EQ(list.size(), 0);
+  EXPECT_TRUE(list.empty());
+  EXPECT_EQ(list.data(), nullptr);
+  EXPECT_EQ(list.begin(), list.end());
+}
+
+// Test constructor taking ownership via std::vector rvalue reference
+TEST(DoubleListTest, VectorMoveConstructor) {
+  std::vector<double> source_vec = {1.1, 2.2, 3.3};
+  const double* expected_data_ptr =
+      source_vec.data();  // Capture pointer before move
+  const size_t expected_size = source_vec.size();
+
+  DoubleList list(std::move(source_vec));
+
+  EXPECT_EQ(list.size(), expected_size);
+  ASSERT_FALSE(list.empty());
+  EXPECT_EQ(list.data(), expected_data_ptr);
+  EXPECT_THAT(list, ElementsAreArray({1.1, 2.2, 3.3}));
+}
+
+// Test constructor creating a non-owning view
+TEST(DoubleListTest, DataViewConstructor) {
+  const double external_data[] = {4.4, 5.5, 6.6};
+  const size_t external_size = std::size(external_data);
+
+  DoubleList list(external_data, external_size);
+
+  EXPECT_EQ(list.size(), external_size);
+  ASSERT_FALSE(list.empty());
+  EXPECT_EQ(list.data(),
+            external_data);  // Should point directly to external data
+  EXPECT_THAT(list, ElementsAreArray({4.4, 5.5, 6.6}));
+}
+
+// Test accessors on an owned list
+TEST(DoubleListTest, AccessorsOwned) {
+  const DoubleList list(CreateSequenceVec(3, 10.0));  // {10.0, 11.0, 12.0}
+
+  EXPECT_EQ(list.size(), 3);
+  ASSERT_FALSE(list.empty());
+  ASSERT_NE(list.data(), nullptr);
+  EXPECT_EQ(list.begin()[0], 10.0);
+  EXPECT_EQ(list.begin()[1], 11.0);
+  EXPECT_EQ(list.begin()[2], 12.0);
+  EXPECT_EQ(list.end(), list.begin() + 3);
+}
+
+// Test accessors on a non-owning list
+TEST(DoubleListTest, AccessorsView) {
+  const double external_data[] = {20.0, 21.0};
+  const DoubleList list(external_data, 2);
+
+  EXPECT_EQ(list.size(), 2);
+  ASSERT_FALSE(list.empty());
+  ASSERT_EQ(list.data(), external_data);
+  EXPECT_EQ(list.begin()[0], 20.0);
+  EXPECT_EQ(list.begin()[1], 21.0);
+  EXPECT_EQ(list.end(), list.begin() + 2);
+}
+
+// Test move constructor from an owned list
+TEST(DoubleListTest, MoveConstructionFromOwned) {
+  DoubleList list1(CreateSequenceVec(2, 1.0));  // {1.0, 2.0}
+  const double* original_data_ptr = list1.data();
+
+  DoubleList list2(std::move(list1));
+
+  // Check list2 has the data
+  EXPECT_EQ(list2.size(), 2);
+  ASSERT_FALSE(list2.empty());
+  EXPECT_EQ(list2.data(), original_data_ptr);
+  EXPECT_THAT(list2, ElementsAreArray({1.0, 2.0}));
+}
+
+// Test move constructor from a non-owning list
+TEST(DoubleListTest, MoveConstructionFromView) {
+  const double external_data[] = {3.0, 4.0};
+  DoubleList list1(external_data, 2);
+  const size_t original_size = list1.size();
+
+  DoubleList list2(std::move(list1));
+
+  // Check list2 has the view details
+  EXPECT_EQ(list2.size(), original_size);
+  ASSERT_FALSE(list2.empty());
+  EXPECT_EQ(list2.data(), external_data);
+  EXPECT_THAT(list2, ElementsAreArray({3.0, 4.0}));
+}
+
+// Test move assignment from an owned list to another
+TEST(DoubleListTest, MoveAssignmentFromOwned) {
+  DoubleList list1(CreateSequenceVec(3, 5.0));    // {5.0, 6.0, 7.0}
+  DoubleList list2(CreateSequenceVec(1, 100.0));  // {100.0}
+  const double* original_data_ptr_list1 = list1.data();
+
+  list2 = std::move(list1);
+
+  // Check list2 has the data from list1
+  EXPECT_EQ(list2.size(), 3);
+  ASSERT_FALSE(list2.empty());
+  EXPECT_EQ(list2.data(), original_data_ptr_list1);
+  EXPECT_THAT(list2, ElementsAreArray({5.0, 6.0, 7.0}));
+}
+
+// Test move assignment from a non-owning list to another
+TEST(DoubleListTest, MoveAssignmentFromView) {
+  const double external_data[] = {8.0, 9.0};
+  DoubleList list1(external_data, 2);
+  DoubleList list2(CreateSequenceVec(1, 200.0));  // {200.0}
+
+  list2 = std::move(list1);
+
+  // Check list2 has the view details from list1
+  EXPECT_EQ(list2.size(), 2);
+  ASSERT_FALSE(list2.empty());
+  EXPECT_EQ(list2.data(), external_data);
+  EXPECT_THAT(list2, ElementsAreArray({8.0, 9.0}));
+}
+
+// Test ReleaseVector when the list owns the data
+TEST(DoubleListTest, ReleaseVectorOwned) {
+  DoubleList list(CreateSequenceVec(3, 10.0));  // {10.0, 11.0, 12.0}
+  const double* original_data_ptr = list.data();
+
+  // ReleaseVector must be called on an rvalue
+  std::vector<double> released_vec = std::move(list).ReleaseVector();
+
+  // Check the released vector
+  EXPECT_EQ(released_vec.size(), 3);
+  EXPECT_THAT(released_vec, ElementsAreArray({10.0, 11.0, 12.0}));
+  EXPECT_EQ(released_vec.data(), original_data_ptr);
+}
+
+// Test ReleaseVector when the list has a non-owning view
+TEST(DoubleListTest, ReleaseVectorView) {
+  const double external_data[] = {13.0, 14.0};
+  DoubleList list(external_data, 2);
+  const double* original_view_ptr = list.data();
+
+  // ReleaseVector must be called on an rvalue
+  std::vector<double> released_vec = std::move(list).ReleaseVector();
+
+  // Check the released vector - it should be a COPY
+  EXPECT_EQ(released_vec.size(), 2);
+  EXPECT_THAT(released_vec, ElementsAreArray({13.0, 14.0}));
+
+  // IMPORTANT: Verify it's a copy, not pointing to the original external data
+  EXPECT_NE(released_vec.data(), original_view_ptr);
+}
+
+// Test ReleaseVector when the list is empty (default constructed - view)
+TEST(DoubleListTest, ReleaseVectorEmptyDefault) {
+  DoubleList list;  // Empty view by default
+  std::vector<double> released_vec = std::move(list).ReleaseVector();
+  EXPECT_TRUE(released_vec.empty());
+  EXPECT_THAT(released_vec, IsEmpty());
+}
+
+// Test ReleaseVector when the list is empty (constructed from empty view)
+TEST(DoubleListTest, ReleaseVectorEmptyView) {
+  DoubleList list(nullptr, 0);  // Empty view
+  std::vector<double> released_vec = std::move(list).ReleaseVector();
+  EXPECT_TRUE(released_vec.empty());
+  EXPECT_THAT(released_vec, IsEmpty());
+}
+
+// Test ReleaseVector when the list is empty (constructed from empty vector -
+// owned)
+TEST(DoubleListTest, ReleaseVectorEmptyOwned) {
+  DoubleList list(std::vector<double>{});  // Empty owned vector
+  std::vector<double> released_vec = std::move(list).ReleaseVector();
+  EXPECT_TRUE(released_vec.empty());
+  EXPECT_THAT(released_vec, IsEmpty());
+}
+
+}  // namespace
+
+}  // namespace lib
+}  // namespace icing

diff --git a/icing/scoring/advanced_scoring/score-expression-util_test.cc b/icing/scoring/advanced_scoring/score-expression-util_test.cc
index 4e4355b..51b849a 100644
--- a/icing/scoring/advanced_scoring/score-expression-util_test.cc
+++ b/icing/scoring/advanced_scoring/score-expression-util_test.cc

@@ -29,6 +29,7 @@
 #include "icing/file/filesystem.h"
 #include "icing/file/portable-file-backed-proto-log.h"
 #include "icing/index/hit/doc-hit-info.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/schema.pb.h"
 #include "icing/proto/scoring.pb.h"
@@ -42,6 +43,7 @@
 #include "icing/testing/fake-clock.h"
 #include "icing/testing/test-feature-flags.h"
 #include "icing/testing/tmp-directory.h"
+#include "icing/util/document-util.h"
 
 namespace icing {
 namespace lib {
@@ -72,14 +74,18 @@
 
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
-        DocumentStore::Create(&filesystem_, doc_store_dir_, &fake_clock_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, doc_store_dir_, &fake_clock_, schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     document_store_ = std::move(create_result.document_store);
 
     // Creates the schema
@@ -204,8 +210,8 @@
 TEST_F(ScoreExpressionUtilTest, SimpleExpression) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result,
-      document_store_->Put(CreateDocument("namespace", "uri", /*score=*/10,
-                                          /*creation_timestamp_ms=*/100)));
+      document_store_->Put(document_util::CreateDocumentWrapper(CreateDocument(
+          "namespace", "uri", /*score=*/10, /*creation_timestamp_ms=*/100))));
   DocumentId document_id = put_result.new_document_id;
   DocHitInfo doc_hit_info = DocHitInfo(document_id);
 
@@ -228,8 +234,8 @@
 TEST_F(ScoreExpressionUtilTest, DocumentFunctionsWithoutOptionalDependencies) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result,
-      document_store_->Put(CreateDocument("namespace", "uri", /*score=*/10,
-                                          /*creation_timestamp_ms=*/100)));
+      document_store_->Put(document_util::CreateDocumentWrapper(CreateDocument(
+          "namespace", "uri", /*score=*/10, /*creation_timestamp_ms=*/100))));
   DocumentId document_id = put_result.new_document_id;
   DocHitInfo doc_hit_info = DocHitInfo(document_id);
 

diff --git a/icing/scoring/advanced_scoring/score-expression.cc b/icing/scoring/advanced_scoring/score-expression.cc
index 445e1f3..2dcd20d 100644
--- a/icing/scoring/advanced_scoring/score-expression.cc
+++ b/icing/scoring/advanced_scoring/score-expression.cc

@@ -41,6 +41,7 @@
 #include "icing/proto/internal/scorable_property_set.pb.h"
 #include "icing/schema/schema-store.h"
 #include "icing/schema/section.h"
+#include "icing/scoring/advanced_scoring/double-list.h"
 #include "icing/scoring/bm25f-calculator.h"
 #include "icing/scoring/scored-document-hit.h"
 #include "icing/scoring/section-weights.h"
@@ -312,89 +313,106 @@
 libtextclassifier3::StatusOr<double>
 MathFunctionScoreExpression::EvaluateDouble(
     const DocHitInfo& hit_info, const DocHitInfoIterator* query_it) const {
-  std::vector<double> values;
+  DoubleList list_value;
+  std::vector<double> double_values;
   int ind = 0;
   if (args_.at(0)->type() == ScoreExpressionType::kDoubleList) {
-    ICING_ASSIGN_OR_RETURN(values,
+    ICING_ASSIGN_OR_RETURN(list_value,
                            args_.at(0)->EvaluateList(hit_info, query_it));
     ind = 1;
   }
+  double_values.reserve(args_.size() - ind);
   for (; ind < args_.size(); ++ind) {
     ICING_ASSIGN_OR_RETURN(double v,
                            args_.at(ind)->EvaluateDouble(hit_info, query_it));
-    values.push_back(v);
+    double_values.push_back(v);
+  }
+
+  // Double values in variable-argument functions can be treated as a single
+  // list.
+  if (kVariableArgumentsFunctions.count(function_type_) > 0 &&
+      !double_values.empty()) {
+    if (!list_value.empty()) {
+      return absl_ports::InternalError(
+          "Should never reach here, since static type checking should not "
+          "allow variable-argument functions to have both list arguments and "
+          "double arguments.");
+    }
+    list_value = DoubleList(std::move(double_values));
+    double_values.clear();
   }
 
   double res = 0;
   switch (function_type_) {
     case FunctionType::kLog:
-      if (values.size() == 1) {
-        res = log(values[0]);
+      if (double_values.size() == 1) {
+        res = log(double_values[0]);
       } else {
         // argument 0 is log base
         // argument 1 is the value
-        res = log(values[1]) / log(values[0]);
+        res = log(double_values[1]) / log(double_values[0]);
       }
       break;
     case FunctionType::kPow:
-      res = pow(values[0], values[1]);
+      res = pow(double_values[0], double_values[1]);
       break;
+    case FunctionType::kSqrt:
+      res = sqrt(double_values[0]);
+      break;
+    case FunctionType::kAbs:
+      res = abs(double_values[0]);
+      break;
+    case FunctionType::kSin:
+      res = sin(double_values[0]);
+      break;
+    case FunctionType::kCos:
+      res = cos(double_values[0]);
+      break;
+    case FunctionType::kTan:
+      res = tan(double_values[0]);
+      break;
+    // Variable-argument functions
     case FunctionType::kMax:
-      if (values.empty()) {
+      if (list_value.empty()) {
         return absl_ports::InvalidArgumentError(
             "Got an empty parameter set in max function");
       }
-      res = *std::max_element(values.begin(), values.end());
+      res = *std::max_element(list_value.begin(), list_value.end());
       break;
     case FunctionType::kMin:
-      if (values.empty()) {
+      if (list_value.empty()) {
         return absl_ports::InvalidArgumentError(
             "Got an empty parameter set in min function");
       }
-      res = *std::min_element(values.begin(), values.end());
+      res = *std::min_element(list_value.begin(), list_value.end());
       break;
     case FunctionType::kLen:
-      res = values.size();
+      res = list_value.size();
       break;
     case FunctionType::kSum:
-      res = std::reduce(values.begin(), values.end());
+      res = std::reduce(list_value.begin(), list_value.end());
       break;
     case FunctionType::kAvg:
-      if (values.empty()) {
+      if (list_value.empty()) {
         return absl_ports::InvalidArgumentError(
             "Got an empty parameter set in avg function.");
       }
-      res = std::reduce(values.begin(), values.end()) / values.size();
+      res =
+          std::reduce(list_value.begin(), list_value.end()) / list_value.size();
       break;
-    case FunctionType::kSqrt:
-      res = sqrt(values[0]);
-      break;
-    case FunctionType::kAbs:
-      res = abs(values[0]);
-      break;
-    case FunctionType::kSin:
-      res = sin(values[0]);
-      break;
-    case FunctionType::kCos:
-      res = cos(values[0]);
-      break;
-    case FunctionType::kTan:
-      res = tan(values[0]);
-      break;
-    // For the following two functions, the last value is the default value.
-    // If values.size() == 1, then it means the provided list is empty.
+    // For the following two functions, double_values[0] is the default value.
     case FunctionType::kMaxOrDefault:
-      if (values.size() == 1) {
-        res = values[0];
+      if (list_value.empty()) {
+        res = double_values[0];
       } else {
-        res = *std::max_element(values.begin(), values.end() - 1);
+        res = *std::max_element(list_value.begin(), list_value.end());
       }
       break;
     case FunctionType::kMinOrDefault:
-      if (values.size() == 1) {
-        res = values[0];
+      if (list_value.empty()) {
+        res = double_values[0];
       } else {
-        res = *std::min_element(values.begin(), values.end() - 1);
+        res = *std::min_element(list_value.begin(), list_value.end());
       }
       break;
   }
@@ -444,12 +462,12 @@
       new ListOperationFunctionScoreExpression(function_type, std::move(args)));
 }
 
-libtextclassifier3::StatusOr<std::vector<double>>
+libtextclassifier3::StatusOr<DoubleList>
 ListOperationFunctionScoreExpression::EvaluateList(
     const DocHitInfo& hit_info, const DocHitInfoIterator* query_it) const {
   switch (function_type_) {
     case FunctionType::kFilterByRange:
-      ICING_ASSIGN_OR_RETURN(std::vector<double> list_value,
+      ICING_ASSIGN_OR_RETURN(DoubleList list_value,
                              args_.at(0)->EvaluateList(hit_info, query_it));
       ICING_ASSIGN_OR_RETURN(double low,
                              args_.at(1)->EvaluateDouble(hit_info, query_it));
@@ -459,11 +477,13 @@
         return absl_ports::InvalidArgumentError(
             "The lower bound cannot be greater than the upper bound.");
       }
+      // TODO(b/408437387): Consider avoiding a copy if nothing is filtered out.
+      std::vector<double> new_list = std::move(list_value).ReleaseVector();
       auto new_end =
-          std::remove_if(list_value.begin(), list_value.end(),
+          std::remove_if(new_list.begin(), new_list.end(),
                          [low, high](double v) { return v < low || v > high; });
-      list_value.erase(new_end, list_value.end());
-      return list_value;
+      new_list.erase(new_end, new_list.end());
+      return DoubleList(std::move(new_list));
       break;
   }
   return absl_ports::InternalError("Should never reach here.");
@@ -629,7 +649,7 @@
           document_store, *join_children_fetcher, current_time_ms));
 }
 
-libtextclassifier3::StatusOr<std::vector<double>>
+libtextclassifier3::StatusOr<DoubleList>
 ChildrenRankingSignalsFunctionScoreExpression::EvaluateList(
     const DocHitInfo& hit_info, const DocHitInfoIterator* query_it) const {
   ICING_ASSIGN_OR_RETURN(
@@ -640,7 +660,7 @@
   for (const ScoredDocumentHit& child_hit : children_hits) {
     children_scores.push_back(child_hit.score());
   }
-  return std::move(children_scores);
+  return DoubleList(std::move(children_scores));
 }
 
 libtextclassifier3::StatusOr<
@@ -664,7 +684,7 @@
           document_store, section_weights, current_time_ms));
 }
 
-libtextclassifier3::StatusOr<std::vector<double>>
+libtextclassifier3::StatusOr<DoubleList>
 PropertyWeightsFunctionScoreExpression::EvaluateList(
     const DocHitInfo& hit_info, const DocHitInfoIterator*) const {
   std::vector<double> weights;
@@ -678,7 +698,7 @@
     weights.push_back(section_weights_.GetNormalizedSectionWeight(
         schema_type_id, section_id));
   }
-  return weights;
+  return DoubleList(std::move(weights));
 }
 
 libtextclassifier3::StatusOr<std::unique_ptr<ScoreExpression>>
@@ -767,13 +787,16 @@
         metric_type,
         embedding_util::GetEmbeddingQueryMetricTypeFromName(metric));
   }
+  const EmbeddingQueryResults::EmbeddingQueryMatchInfoMap* match_info_map_ =
+      nullptr;
   if (embedding_index_arg->is_constant()) {
     ICING_ASSIGN_OR_RETURN(
         uint32_t embedding_index,
         embedding_index_arg->EvaluateDouble(DocHitInfo(),
                                             /*query_it=*/nullptr));
-    if (embedding_query_results->GetMatchInfoMap(embedding_index,
-                                                 metric_type) == nullptr) {
+    match_info_map_ =
+        embedding_query_results->GetMatchInfoMap(embedding_index, metric_type);
+    if (match_info_map_ == nullptr) {
       return absl_ports::InvalidArgumentError(absl_ports::StrCat(
           "The embedding query index ", std::to_string(embedding_index),
           " with metric type ",
@@ -783,22 +806,26 @@
   }
   return std::unique_ptr<MatchedSemanticScoresFunctionScoreExpression>(
       new MatchedSemanticScoresFunctionScoreExpression(
-          std::move(args), metric_type, *embedding_query_results));
+          std::move(args), metric_type, *embedding_query_results,
+          match_info_map_));
 }
 
-libtextclassifier3::StatusOr<std::vector<double>>
+libtextclassifier3::StatusOr<DoubleList>
 MatchedSemanticScoresFunctionScoreExpression::EvaluateList(
     const DocHitInfo& hit_info, const DocHitInfoIterator* query_it) const {
   ICING_ASSIGN_OR_RETURN(double raw_query_index,
                          args_[1]->EvaluateDouble(hit_info, query_it));
   uint32_t query_index = (uint32_t)raw_query_index;
-  const std::vector<double>* scores =
-      embedding_query_results_.GetMatchedScoresForDocument(
-          query_index, metric_type_, hit_info.document_id());
-  if (scores == nullptr) {
-    return std::vector<double>();
+  if (match_info_map_ != nullptr) {
+    auto info_it = match_info_map_->find(hit_info.document_id());
+    if (info_it == match_info_map_->end()) {
+      return DoubleList();
+    }
+    return embedding_query_results_.GetMatchedScoresFromEmbeddingMatchInfos(
+        info_it->second);
   }
-  return *scores;
+  return embedding_query_results_.GetMatchedScoresForDocument(
+      query_index, metric_type_, hit_info.document_id());
 }
 
 GetScorablePropertyFunctionScoreExpression::
@@ -885,13 +912,13 @@
           std::move(schema_type_ids), property_path));
 }
 
-libtextclassifier3::StatusOr<std::vector<double>>
+libtextclassifier3::StatusOr<DoubleList>
 GetScorablePropertyFunctionScoreExpression::EvaluateList(
     const DocHitInfo& hit_info, const DocHitInfoIterator* query_it) const {
   SchemaTypeId doc_schema_type_id = GetSchemaTypeId(
       hit_info.document_id(), document_store_, current_time_ms_);
   if (schema_type_ids_.find(doc_schema_type_id) == schema_type_ids_.end()) {
-    return std::vector<double>();
+    return DoubleList();
   }
 
   // By this point, the document to be evaluated is guaranteed to have a
@@ -920,17 +947,19 @@
 
   // Converts ScorablePropertyProto to a vector of doubles.
   if (scorable_property_proto->int64_values_size() > 0) {
-    return std::vector<double>(scorable_property_proto->int64_values().begin(),
-                               scorable_property_proto->int64_values().end());
+    return DoubleList(
+        std::vector<double>(scorable_property_proto->int64_values().begin(),
+                            scorable_property_proto->int64_values().end()));
   } else if (scorable_property_proto->double_values_size() > 0) {
-    return std::vector<double>(scorable_property_proto->double_values().begin(),
-                               scorable_property_proto->double_values().end());
+    return DoubleList(
+        std::vector<double>(scorable_property_proto->double_values().begin(),
+                            scorable_property_proto->double_values().end()));
   } else if (scorable_property_proto->boolean_values_size() > 0) {
-    return std::vector<double>(
-        scorable_property_proto->boolean_values().begin(),
-        scorable_property_proto->boolean_values().end());
+    return DoubleList(
+        std::vector<double>(scorable_property_proto->boolean_values().begin(),
+                            scorable_property_proto->boolean_values().end()));
   }
-  return std::vector<double>();
+  return DoubleList();
 }
 
 }  // namespace lib

diff --git a/icing/scoring/advanced_scoring/score-expression.h b/icing/scoring/advanced_scoring/score-expression.h
index 53c5807..5df15ec 100644
--- a/icing/scoring/advanced_scoring/score-expression.h
+++ b/icing/scoring/advanced_scoring/score-expression.h

@@ -31,12 +31,11 @@
 #include "icing/index/iterator/doc-hit-info-iterator.h"
 #include "icing/join/join-children-fetcher.h"
 #include "icing/schema/schema-store.h"
+#include "icing/scoring/advanced_scoring/double-list.h"
 #include "icing/scoring/bm25f-calculator.h"
 #include "icing/scoring/section-weights.h"
 #include "icing/store/document-filter-data.h"
-#include "icing/store/document-id.h"
 #include "icing/store/document-store.h"
-#include "icing/util/status-macros.h"
 
 namespace icing {
 namespace lib {
@@ -78,7 +77,7 @@
         "double. There must be inconsistencies in the static type checking.");
   }
 
-  virtual libtextclassifier3::StatusOr<std::vector<double>> EvaluateList(
+  virtual libtextclassifier3::StatusOr<DoubleList> EvaluateList(
       const DocHitInfo& hit_info, const DocHitInfoIterator* query_it) const {
     if (type() == ScoreExpressionType::kDoubleList) {
       return absl_ports::UnimplementedError(
@@ -269,7 +268,7 @@
       FunctionType function_type,
       std::vector<std::unique_ptr<ScoreExpression>> args);
 
-  libtextclassifier3::StatusOr<std::vector<double>> EvaluateList(
+  libtextclassifier3::StatusOr<DoubleList> EvaluateList(
       const DocHitInfo& hit_info,
       const DocHitInfoIterator* query_it) const override;
 
@@ -381,7 +380,7 @@
          const JoinChildrenFetcher* join_children_fetcher,
          int64_t current_time_ms);
 
-  libtextclassifier3::StatusOr<std::vector<double>> EvaluateList(
+  libtextclassifier3::StatusOr<DoubleList> EvaluateList(
       const DocHitInfo& hit_info,
       const DocHitInfoIterator* query_it) const override;
 
@@ -416,7 +415,7 @@
          const DocumentStore* document_store,
          const SectionWeights* section_weights, int64_t current_time_ms);
 
-  libtextclassifier3::StatusOr<std::vector<double>> EvaluateList(
+  libtextclassifier3::StatusOr<DoubleList> EvaluateList(
       const DocHitInfo& hit_info, const DocHitInfoIterator*) const override;
 
   ScoreExpressionType type() const override {
@@ -477,7 +476,7 @@
          SearchSpecProto::EmbeddingQueryMetricType::Code default_metric_type,
          const EmbeddingQueryResults* embedding_query_results);
 
-  libtextclassifier3::StatusOr<std::vector<double>> EvaluateList(
+  libtextclassifier3::StatusOr<DoubleList> EvaluateList(
       const DocHitInfo& hit_info,
       const DocHitInfoIterator* query_it) const override;
 
@@ -489,14 +488,20 @@
   explicit MatchedSemanticScoresFunctionScoreExpression(
       std::vector<std::unique_ptr<ScoreExpression>> args,
       SearchSpecProto::EmbeddingQueryMetricType::Code metric_type,
-      const EmbeddingQueryResults& embedding_query_results)
+      const EmbeddingQueryResults& embedding_query_results,
+      const EmbeddingQueryResults::EmbeddingQueryMatchInfoMap* match_info_map)
       : args_(std::move(args)),
         metric_type_(metric_type),
-        embedding_query_results_(embedding_query_results) {}
+        embedding_query_results_(embedding_query_results),
+        match_info_map_(match_info_map) {}
 
   std::vector<std::unique_ptr<ScoreExpression>> args_;
   const SearchSpecProto::EmbeddingQueryMetricType::Code metric_type_;
   const EmbeddingQueryResults& embedding_query_results_;
+  // If the embedding vector's index evaluated from args is a constant, this is
+  // the corresponding EmbeddingQueryMatchInfoMap for the embedding query.
+  // Otherwise, this is nullptr.
+  const EmbeddingQueryResults::EmbeddingQueryMatchInfoMap* match_info_map_;
 };
 
 class GetScorablePropertyFunctionScoreExpression : public ScoreExpression {
@@ -522,7 +527,7 @@
     return ScoreExpressionType::kDoubleList;
   }
 
-  libtextclassifier3::StatusOr<std::vector<double>> EvaluateList(
+  libtextclassifier3::StatusOr<DoubleList> EvaluateList(
       const DocHitInfo& hit_info,
       const DocHitInfoIterator* query_it) const override;
 

diff --git a/icing/scoring/advanced_scoring/score-expression_test.cc b/icing/scoring/advanced_scoring/score-expression_test.cc
index cbfed9b..391248d 100644
--- a/icing/scoring/advanced_scoring/score-expression_test.cc
+++ b/icing/scoring/advanced_scoring/score-expression_test.cc

@@ -24,6 +24,7 @@
 #include "gtest/gtest.h"
 #include "icing/index/hit/doc-hit-info.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
+#include "icing/scoring/advanced_scoring/double-list.h"
 #include "icing/testing/common-matchers.h"
 
 namespace icing {
@@ -62,9 +63,9 @@
     return res;
   }
 
-  libtextclassifier3::StatusOr<std::vector<double>> EvaluateList(
+  libtextclassifier3::StatusOr<DoubleList> EvaluateList(
       const DocHitInfo &, const DocHitInfoIterator *) const override {
-    return values;
+    return DoubleList(values.data(), values.size());
   }
 
   ScoreExpressionType type() const override {

diff --git a/icing/scoring/priority-queue-scored-document-hits-ranker.h b/icing/scoring/priority-queue-scored-document-hits-ranker.h
index 68cf921..135e0e3 100644
--- a/icing/scoring/priority-queue-scored-document-hits-ranker.h
+++ b/icing/scoring/priority-queue-scored-document-hits-ranker.h

@@ -15,6 +15,7 @@
 #ifndef ICING_SCORING_PRIORITY_QUEUE_SCORED_DOCUMENT_HITS_RANKER_H_
 #define ICING_SCORING_PRIORITY_QUEUE_SCORED_DOCUMENT_HITS_RANKER_H_
 
+#include <memory>
 #include <queue>
 #include <unordered_set>
 #include <vector>
@@ -37,22 +38,23 @@
 
   ~PriorityQueueScoredDocumentHitsRanker() override = default;
 
+  void Pop() override;
+
   // Note: ranker may store ScoredDocumentHit or JoinedScoredDocumentHit, so we
   // have template for scored_data_pq_.
   // - JoinedScoredDocumentHit is a superset of ScoredDocumentHit, so we unify
-  //   the return type of PopNext to use the superset type
-  //   JoinedScoredDocumentHit in order to make it simple, and rankers storing
-  //   ScoredDocumentHit should convert it to JoinedScoredDocumentHit before
-  //   returning. It makes the implementation simpler, especially for
-  //   ResultRetriever, which now only needs to deal with one single return
-  //   format.
+  //   the return type of Top to use the superset type JoinedScoredDocumentHit
+  //   in order to make it simple, and rankers storing ScoredDocumentHit should
+  //   convert it to JoinedScoredDocumentHit before returning. It makes the
+  //   implementation simpler, especially for ResultRetriever, which now only
+  //   needs to deal with one single return format.
   // - JoinedScoredDocumentHit has ~2x size of ScoredDocumentHit. Since we cache
   //   ranker (which contains a priority queue of data) in ResultState, if we
   //   store the scored hits in JoinedScoredDocumentHit format directly, then it
   //   doubles the memory usage. Therefore, we still keep the flexibility to
-  //   store ScoredDocumentHit or any other types of data, but require PopNext
-  //   to convert it to JoinedScoredDocumentHit.
-  JoinedScoredDocumentHit PopNext() override;
+  //   store ScoredDocumentHit or any other types of data, but require Pop to
+  //   convert it to JoinedScoredDocumentHit and cache it in curr_.
+  const JoinedScoredDocumentHit& Top() const override { return *curr_; }
 
   // Returns DocumentIds of the top K documents according to the ranking policy.
   // - For ScoredDocumentHit, this returns the DocumentIds of the top K
@@ -70,7 +72,7 @@
 
   int size() const override { return scored_data_pq_.size(); }
 
-  bool empty() const override { return scored_data_pq_.empty(); }
+  bool empty() const override { return curr_ == nullptr; }
 
  private:
   // Comparator for std::priority_queue. Since std::priority is a max heap
@@ -95,6 +97,11 @@
     bool is_ascending_;
   };
 
+  // Helper function to refresh the current element (fetch the top element from
+  // the priority queue, convert it to JoinedScoredDocumentHit, and cache it in
+  // curr_).
+  void RefreshCurrent();
+
   Comparator comparator_;
 
   // Use priority queue to get top K hits in O(KlgN) time.
@@ -102,6 +109,8 @@
       scored_data_pq_;
 
   Converter converter_;
+
+  std::unique_ptr<JoinedScoredDocumentHit> curr_;
 };
 
 template <typename ScoredDataType, typename Converter>
@@ -109,14 +118,14 @@
     PriorityQueueScoredDocumentHitsRanker(
         std::vector<ScoredDataType>&& scored_data_vec, bool is_descending)
     : comparator_(/*is_ascending=*/!is_descending),
-      scored_data_pq_(comparator_, std::move(scored_data_vec)) {}
+      scored_data_pq_(comparator_, std::move(scored_data_vec)) {
+  RefreshCurrent();
+}
 
 template <typename ScoredDataType, typename Converter>
-JoinedScoredDocumentHit
-PriorityQueueScoredDocumentHitsRanker<ScoredDataType, Converter>::PopNext() {
-  ScoredDataType next_scored_data = scored_data_pq_.top();
+void PriorityQueueScoredDocumentHitsRanker<ScoredDataType, Converter>::Pop() {
   scored_data_pq_.pop();
-  return converter_(std::move(next_scored_data));
+  RefreshCurrent();
 }
 
 template <typename ScoredDataType, typename Converter>
@@ -192,7 +201,27 @@
     new_pq.push(scored_data_pq_.top());
     scored_data_pq_.pop();
   }
+
+  // Assign back to the class members.
   scored_data_pq_ = std::move(new_pq);
+  RefreshCurrent();
+}
+
+template <typename ScoredDataType, typename Converter>
+void PriorityQueueScoredDocumentHitsRanker<ScoredDataType,
+                                           Converter>::RefreshCurrent() {
+  if (scored_data_pq_.empty()) {
+    curr_ = nullptr;
+  } else {
+    ScoredDataType scored_data = scored_data_pq_.top();
+
+    if (curr_ == nullptr) {
+      curr_ = std::make_unique<JoinedScoredDocumentHit>(
+          converter_(std::move(scored_data)));
+    } else {
+      *curr_ = converter_(std::move(scored_data));
+    }
+  }
 }
 
 }  // namespace lib

diff --git a/icing/scoring/priority-queue-scored-document-hits-ranker_test.cc b/icing/scoring/priority-queue-scored-document-hits-ranker_test.cc
index 19a9ebb..7ec1d68 100644
--- a/icing/scoring/priority-queue-scored-document-hits-ranker_test.cc
+++ b/icing/scoring/priority-queue-scored-document-hits-ranker_test.cc

@@ -48,7 +48,8 @@
     PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>& ranker) {
   std::vector<JoinedScoredDocumentHit> hits;
   while (!ranker.empty()) {
-    hits.push_back(ranker.PopNext());
+    hits.push_back(ranker.Top());
+    ranker.Pop();
   }
   return hits;
 }
@@ -57,7 +58,8 @@
     PriorityQueueScoredDocumentHitsRanker<JoinedScoredDocumentHit>& ranker) {
   std::vector<JoinedScoredDocumentHit> hits;
   while (!ranker.empty()) {
-    hits.push_back(ranker.PopNext());
+    hits.push_back(ranker.Top());
+    ranker.Pop();
   }
   return hits;
 }
@@ -76,15 +78,15 @@
   EXPECT_THAT(ranker.size(), Eq(3));
   EXPECT_FALSE(ranker.empty());
 
-  ranker.PopNext();
+  ranker.Pop();
   EXPECT_THAT(ranker.size(), Eq(2));
   EXPECT_FALSE(ranker.empty());
 
-  ranker.PopNext();
+  ranker.Pop();
   EXPECT_THAT(ranker.size(), Eq(1));
   EXPECT_FALSE(ranker.empty());
 
-  ranker.PopNext();
+  ranker.Pop();
   EXPECT_THAT(ranker.size(), Eq(0));
   EXPECT_TRUE(ranker.empty());
 }

diff --git a/icing/scoring/score-and-rank_benchmark.cc b/icing/scoring/score-and-rank_benchmark.cc
index c8f6c00..006457f 100644
--- a/icing/scoring/score-and-rank_benchmark.cc
+++ b/icing/scoring/score-and-rank_benchmark.cc

@@ -32,6 +32,7 @@
 #include "icing/index/hit/doc-hit-info.h"
 #include "icing/index/iterator/doc-hit-info-iterator-test-util.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/schema.pb.h"
 #include "icing/proto/scoring.pb.h"
@@ -46,6 +47,7 @@
 #include "icing/testing/test-feature-flags.h"
 #include "icing/testing/tmp-directory.h"
 #include "icing/util/clock.h"
+#include "icing/util/document-util.h"
 
 // This is an overall benchmark for ScoringProcessor, Scorer, and Ranker. It
 // shows how performance varies when we score different numbers of document
@@ -85,6 +87,12 @@
   return schema;
 }
 
+DocumentWrapper CreateDocumentWrapper(DocumentProto document,
+                                      int32_t num_string_tokens) {
+  document.mutable_internal_fields()->set_length_in_tokens(num_string_tokens);
+  return document_util::CreateDocumentWrapper(std::move(document));
+}
+
 DocumentProto CreateEmailDocument(int id, int document_score,
                                   uint64_t creation_timestamp_ms) {
   return DocumentBuilder()
@@ -105,6 +113,9 @@
       /*force_recovery_and_revalidate_documents=*/false,
       /*pre_mapping_fbv=*/false, /*use_persistent_hash_map=*/true,
       PortableFileBackedProtoLog<DocumentWrapper>::kDefaultCompressionLevel,
+      PortableFileBackedProtoLog<
+          DocumentWrapper>::kDefaultCompressionThresholdBytes,
+      protobuf_ports::kDefaultMemLevel,
       /*initialize_stats=*/nullptr);
 }
 
@@ -160,9 +171,10 @@
   for (int i = 0; i < num_of_documents; i++) {
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::PutResult put_result,
-        document_store->Put(CreateEmailDocument(
-            /*id=*/i, /*document_score=*/distribution(random_generator),
-            /*creation_timestamp_ms=*/1)));
+        document_store->Put(
+            document_util::CreateDocumentWrapper(CreateEmailDocument(
+                /*id=*/i, /*document_score=*/distribution(random_generator),
+                /*creation_timestamp_ms=*/1))));
     DocumentId document_id = put_result.new_document_id;
     doc_hit_infos.emplace_back(document_id);
   }
@@ -273,9 +285,10 @@
   for (int i = 0; i < num_of_documents; i++) {
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::PutResult put_result,
-        document_store->Put(CreateEmailDocument(
-            /*id=*/i, /*document_score=*/1,
-            /*creation_timestamp_ms=*/distribution(random_generator))));
+        document_store->Put(
+            document_util::CreateDocumentWrapper(CreateEmailDocument(
+                /*id=*/i, /*document_score=*/1,
+                /*creation_timestamp_ms=*/distribution(random_generator)))));
     DocumentId document_id = put_result.new_document_id;
     doc_hit_infos.emplace_back(document_id);
   }
@@ -382,8 +395,9 @@
   for (int i = 0; i < num_of_documents; i++) {
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::PutResult put_result,
-        document_store->Put(CreateEmailDocument(/*id=*/i, /*document_score=*/1,
-                                                /*creation_timestamp_ms=*/1)));
+        document_store->Put(document_util::CreateDocumentWrapper(
+            CreateEmailDocument(/*id=*/i, /*document_score=*/1,
+                                /*creation_timestamp_ms=*/1))));
     DocumentId document_id = put_result.new_document_id;
     doc_hit_infos.emplace_back(document_id);
   }
@@ -497,10 +511,10 @@
   for (int i = 0; i < num_of_documents; i++) {
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::PutResult put_result,
-        document_store->Put(CreateEmailDocument(
-                                /*id=*/i, /*document_score=*/1,
+        document_store->Put(CreateDocumentWrapper(
+            CreateEmailDocument(/*id=*/i, /*document_score=*/1,
                                 /*creation_timestamp_ms=*/1),
-                            /*num_tokens=*/10));
+            /*num_string_tokens=*/10)));
     DocumentId document_id = put_result.new_document_id;
     DocHitInfoTermFrequencyPair doc_hit =
         DocHitInfo(document_id, section_id_mask);

diff --git a/icing/scoring/scored-document-hits-ranker.h b/icing/scoring/scored-document-hits-ranker.h
index 2a7956d..9135ecb 100644
--- a/icing/scoring/scored-document-hits-ranker.h
+++ b/icing/scoring/scored-document-hits-ranker.h

@@ -25,7 +25,7 @@
 
 // TODO(sungyc): re-evaluate other similar implementations (e.g. std::sort +
 //               std::queue/std::vector). Also revisit the capacity shrinking
-//               issue for PopNext().
+//               issue for Pop().
 
 // ScoredDocumentHitsRanker is an interface class for ranking
 // ScoredDocumentHits.
@@ -33,19 +33,24 @@
  public:
   virtual ~ScoredDocumentHitsRanker() = default;
 
-  // Pop the next top JoinedScoredDocumentHit and return. It is undefined to
-  // call PopNext on an empty ranker, so the caller should check if it is not
-  // empty before calling.
+  // Pops the current element and moves to the next one.
+  //
+  // REQUIRES: !empty().
+  virtual void Pop() = 0;
+
+  // Returns the top JoinedScoredDocumentHit.
   //
   // Note: ranker may store ScoredDocumentHit or JoinedScoredDocumentHit. We can
   // add template for this interface, but since JoinedScoredDocumentHit is a
-  // superset of ScoredDocumentHit, we unify the return type of PopNext to use
-  // the superset type JoinedScoredDocumentHit in order to make it simple, and
+  // superset of ScoredDocumentHit, we unify the return type of Top to use the
+  // superset type JoinedScoredDocumentHit in order to make it simple, and
   // rankers storing ScoredDocumentHit should convert it to
   // JoinedScoredDocumentHit before returning. It makes the implementation
   // simpler, especially for ResultRetriever, which now only needs to deal with
   // one single return format.
-  virtual JoinedScoredDocumentHit PopNext() = 0;
+  //
+  // REQUIRES: !empty().
+  virtual const JoinedScoredDocumentHit& Top() const = 0;
 
   // Truncates the remaining ScoredDocumentHits to the given size. The best
   // ScoredDocumentHits (according to the ranking policy) should be kept.

diff --git a/icing/scoring/scorer_test.cc b/icing/scoring/scorer_test.cc
index 98ca18e..e849971 100644
--- a/icing/scoring/scorer_test.cc
+++ b/icing/scoring/scorer_test.cc

@@ -29,6 +29,7 @@
 #include "icing/file/portable-file-backed-proto-log.h"
 #include "icing/index/embed/embedding-query-results.h"
 #include "icing/index/hit/doc-hit-info.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/schema.pb.h"
 #include "icing/proto/scoring.pb.h"
@@ -43,6 +44,7 @@
 #include "icing/testing/fake-clock.h"
 #include "icing/testing/test-feature-flags.h"
 #include "icing/testing/tmp-directory.h"
+#include "icing/util/document-util.h"
 
 namespace icing {
 namespace lib {
@@ -73,14 +75,18 @@
 
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
-        DocumentStore::Create(&filesystem_, doc_store_dir_, &fake_clock1_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, doc_store_dir_, &fake_clock1_, schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     document_store_ = std::move(create_result.document_store);
 
     // Creates a simple email schema
@@ -193,8 +199,10 @@
           .SetCreationTimestampMs(fake_clock1().GetSystemTimeMilliseconds())
           .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store()->Put(test_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store()->Put(
+          document_util::CreateDocumentWrapper(test_document)));
   DocumentId document_id = put_result.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<Scorer> scorer,
@@ -221,8 +229,10 @@
           .SetCreationTimestampMs(fake_clock2().GetSystemTimeMilliseconds())
           .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store()->Put(test_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store()->Put(
+          document_util::CreateDocumentWrapper(test_document)));
   DocumentId document_id = put_result.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<Scorer> scorer,
@@ -251,8 +261,10 @@
           .SetCreationTimestampMs(fake_clock2().GetSystemTimeMilliseconds())
           .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store()->Put(test_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store()->Put(
+          document_util::CreateDocumentWrapper(test_document)));
   DocumentId document_id = put_result.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<Scorer> scorer,
@@ -286,11 +298,15 @@
           .SetCreationTimestampMs(fake_clock2().GetSystemTimeMilliseconds())
           .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store()->Put(test_document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store()->Put(
+          document_util::CreateDocumentWrapper(test_document1)));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store()->Put(test_document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store()->Put(
+          document_util::CreateDocumentWrapper(test_document2)));
   DocumentId document_id2 = put_result2.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<Scorer> scorer,
@@ -320,8 +336,10 @@
           .SetCreationTimestampMs(fake_clock1().GetSystemTimeMilliseconds())
           .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store()->Put(test_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store()->Put(
+          document_util::CreateDocumentWrapper(test_document)));
   DocumentId document_id = put_result.new_document_id;
 
   // Create 3 scorers for 3 different usage types.
@@ -377,8 +395,10 @@
           .SetCreationTimestampMs(fake_clock1().GetSystemTimeMilliseconds())
           .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store()->Put(test_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store()->Put(
+          document_util::CreateDocumentWrapper(test_document)));
   DocumentId document_id = put_result.new_document_id;
 
   // Create 3 scorers for 3 different usage types.
@@ -434,8 +454,10 @@
           .SetCreationTimestampMs(fake_clock1().GetSystemTimeMilliseconds())
           .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store()->Put(test_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store()->Put(
+          document_util::CreateDocumentWrapper(test_document)));
   DocumentId document_id = put_result.new_document_id;
 
   // Create 3 scorers for 3 different usage types.
@@ -491,8 +513,10 @@
           .SetCreationTimestampMs(fake_clock1().GetSystemTimeMilliseconds())
           .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store()->Put(test_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store()->Put(
+          document_util::CreateDocumentWrapper(test_document)));
   DocumentId document_id = put_result.new_document_id;
 
   // Create 3 scorers for 3 different usage types.
@@ -570,8 +594,10 @@
           .SetCreationTimestampMs(fake_clock1().GetSystemTimeMilliseconds())
           .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store()->Put(test_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store()->Put(
+          document_util::CreateDocumentWrapper(test_document)));
   DocumentId document_id = put_result.new_document_id;
 
   // Create 3 scorers for 3 different usage types.
@@ -649,8 +675,10 @@
           .SetCreationTimestampMs(fake_clock1().GetSystemTimeMilliseconds())
           .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store()->Put(test_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store()->Put(
+          document_util::CreateDocumentWrapper(test_document)));
   DocumentId document_id = put_result.new_document_id;
 
   // Create 3 scorers for 3 different usage types.
@@ -764,8 +792,10 @@
           .SetCreationTimestampMs(fake_clock1().GetSystemTimeMilliseconds())
           .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             document_store()->Put(test_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store()->Put(
+          document_util::CreateDocumentWrapper(test_document)));
   DocumentId document_id = put_result.new_document_id;
 
   ICING_ASSERT_OK_AND_ASSIGN(

diff --git a/icing/scoring/scoring-processor.cc b/icing/scoring/scoring-processor.cc
index 4cba8a6..43f5d9b 100644
--- a/icing/scoring/scoring-processor.cc
+++ b/icing/scoring/scoring-processor.cc

@@ -110,6 +110,15 @@
         iterator_call_stats.num_leaf_advance_calls_main_index);
     search_stats->set_num_fetched_hits_integer_index(
         iterator_call_stats.num_leaf_advance_calls_integer_index);
+    search_stats->set_num_unquantized_embeddings_scored(
+        iterator_call_stats.embedding_stats.num_unquantized_embeddings_scored);
+    search_stats->set_num_quantized_embeddings_scored(
+        iterator_call_stats.embedding_stats.num_quantized_embeddings_scored);
+    search_stats->set_num_embedding_shards_read(static_cast<int32_t>(
+        iterator_call_stats.embedding_stats.unquantized_shards_read.size() +
+        iterator_call_stats.embedding_stats.quantized_shards_read.size()));
+    search_stats->set_num_embedding_bytes_read(
+        iterator_call_stats.embedding_stats.num_embedding_bytes_read);
   }
 
   return scored_document_hits;

diff --git a/icing/scoring/scoring-processor_test.cc b/icing/scoring/scoring-processor_test.cc
index f2dc84b..d94b054 100644
--- a/icing/scoring/scoring-processor_test.cc
+++ b/icing/scoring/scoring-processor_test.cc

@@ -33,6 +33,7 @@
 #include "icing/index/hit/doc-hit-info.h"
 #include "icing/index/iterator/doc-hit-info-iterator-test-util.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/proto/document.pb.h"
 #include "icing/proto/schema.pb.h"
 #include "icing/proto/scoring.pb.h"
@@ -49,6 +50,7 @@
 #include "icing/testing/fake-clock.h"
 #include "icing/testing/test-feature-flags.h"
 #include "icing/testing/tmp-directory.h"
+#include "icing/util/document-util.h"
 #include "icing/util/status-macros.h"
 
 namespace icing {
@@ -82,14 +84,18 @@
 
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
-        DocumentStore::Create(&filesystem_, doc_store_dir_, &fake_clock_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/true,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, doc_store_dir_, &fake_clock_, schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/true,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     document_store_ = std::move(create_result.document_store);
 
     // Creates a simple email schema
@@ -158,6 +164,12 @@
       .Build();
 }
 
+DocumentWrapper CreateDocumentWrapper(DocumentProto document,
+                                      int32_t num_string_tokens) {
+  document.mutable_internal_fields()->set_length_in_tokens(num_string_tokens);
+  return document_util::CreateDocumentWrapper(std::move(document));
+}
+
 libtextclassifier3::StatusOr<
     std::pair<std::vector<DocHitInfo>, std::vector<ScoredDocumentHit>>>
 CreateAndInsertsDocumentsWithScores(DocumentStore* document_store,
@@ -165,10 +177,11 @@
   std::vector<DocHitInfo> doc_hit_infos;
   std::vector<ScoredDocumentHit> scored_document_hits;
   for (int i = 0; i < scores.size(); i++) {
-    ICING_ASSIGN_OR_RETURN(DocumentStore::PutResult put_result,
-                           document_store->Put(CreateDocument(
-                               "icing", "email/" + std::to_string(i),
-                               scores.at(i), kDefaultCreationTimestampMs)));
+    ICING_ASSIGN_OR_RETURN(
+        DocumentStore::PutResult put_result,
+        document_store->Put(document_util::CreateDocumentWrapper(
+            CreateDocument("icing", "email/" + std::to_string(i), scores.at(i),
+                           kDefaultCreationTimestampMs))));
     DocumentId document_id = put_result.new_document_id;
     doc_hit_infos.emplace_back(document_id);
     scored_document_hits.emplace_back(document_id, kSectionIdMaskNone,
@@ -268,8 +281,8 @@
   // Sets up documents
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result1,
-      document_store()->Put(CreateDocument("icing", "email/1", /*score=*/1,
-                                           kDefaultCreationTimestampMs)));
+      document_store()->Put(document_util::CreateDocumentWrapper(CreateDocument(
+          "icing", "email/1", /*score=*/1, kDefaultCreationTimestampMs))));
   DocumentId document_id1 = put_result1.new_document_id;
   DocHitInfo doc_hit_info1(document_id1);
 
@@ -379,17 +392,17 @@
       CreateDocument("icing", "email/3", kDefaultScore,
                      /*creation_timestamp_ms=*/kDefaultCreationTimestampMs);
 
-  ICING_ASSERT_OK_AND_ASSIGN(
-      DocumentStore::PutResult put_result1,
-      document_store()->Put(document1, /*num_tokens=*/10));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
+                             document_store()->Put(CreateDocumentWrapper(
+                                 document1, /*num_string_tokens=*/10)));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(
-      DocumentStore::PutResult put_result2,
-      document_store()->Put(document2, /*num_tokens=*/100));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
+                             document_store()->Put(CreateDocumentWrapper(
+                                 document2, /*num_string_tokens=*/100)));
   DocumentId document_id2 = put_result2.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(
-      DocumentStore::PutResult put_result3,
-      document_store()->Put(document3, /*num_tokens=*/50));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
+                             document_store()->Put(CreateDocumentWrapper(
+                                 document3, /*num_string_tokens=*/50)));
   DocumentId document_id3 = put_result3.new_document_id;
 
   DocHitInfoTermFrequencyPair doc_hit_info1 = DocHitInfo(document_id1);
@@ -455,17 +468,17 @@
       CreateDocument("icing", "email/3", kDefaultScore,
                      /*creation_timestamp_ms=*/kDefaultCreationTimestampMs);
 
-  ICING_ASSERT_OK_AND_ASSIGN(
-      DocumentStore::PutResult put_result1,
-      document_store()->Put(document1, /*num_tokens=*/10));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
+                             document_store()->Put(CreateDocumentWrapper(
+                                 document1, /*num_string_tokens=*/10)));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(
-      DocumentStore::PutResult put_result2,
-      document_store()->Put(document2, /*num_tokens=*/10));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
+                             document_store()->Put(CreateDocumentWrapper(
+                                 document2, /*num_string_tokens=*/10)));
   DocumentId document_id2 = put_result2.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(
-      DocumentStore::PutResult put_result3,
-      document_store()->Put(document3, /*num_tokens=*/10));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
+                             document_store()->Put(CreateDocumentWrapper(
+                                 document3, /*num_string_tokens=*/10)));
   DocumentId document_id3 = put_result3.new_document_id;
 
   DocHitInfoTermFrequencyPair doc_hit_info1 = DocHitInfo(document_id1);
@@ -530,17 +543,17 @@
       CreateDocument("icing", "email/3", kDefaultScore,
                      /*creation_timestamp_ms=*/kDefaultCreationTimestampMs);
 
-  ICING_ASSERT_OK_AND_ASSIGN(
-      DocumentStore::PutResult put_result1,
-      document_store()->Put(document1, /*num_tokens=*/10));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
+                             document_store()->Put(CreateDocumentWrapper(
+                                 document1, /*num_string_tokens=*/10)));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(
-      DocumentStore::PutResult put_result2,
-      document_store()->Put(document2, /*num_tokens=*/10));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
+                             document_store()->Put(CreateDocumentWrapper(
+                                 document2, /*num_string_tokens=*/10)));
   DocumentId document_id2 = put_result2.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(
-      DocumentStore::PutResult put_result3,
-      document_store()->Put(document3, /*num_tokens=*/10));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
+                             document_store()->Put(CreateDocumentWrapper(
+                                 document3, /*num_string_tokens=*/10)));
   DocumentId document_id3 = put_result3.new_document_id;
 
   DocHitInfoTermFrequencyPair doc_hit_info1 = DocHitInfo(document_id1);
@@ -609,17 +622,17 @@
       CreateDocument("icing", "email/3", kDefaultScore,
                      /*creation_timestamp_ms=*/kDefaultCreationTimestampMs);
 
-  ICING_ASSERT_OK_AND_ASSIGN(
-      DocumentStore::PutResult put_result1,
-      document_store()->Put(document1, /*num_tokens=*/20));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
+                             document_store()->Put(CreateDocumentWrapper(
+                                 document1, /*num_string_tokens=*/20)));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(
-      DocumentStore::PutResult put_result2,
-      document_store()->Put(document2, /*num_tokens=*/20));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
+                             document_store()->Put(CreateDocumentWrapper(
+                                 document2, /*num_string_tokens=*/20)));
   DocumentId document_id2 = put_result2.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(
-      DocumentStore::PutResult put_result3,
-      document_store()->Put(document3, /*num_tokens=*/20));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
+                             document_store()->Put(CreateDocumentWrapper(
+                                 document3, /*num_string_tokens=*/20)));
   DocumentId document_id3 = put_result3.new_document_id;
 
   // Index 5 terms with total frequencies:
@@ -830,9 +843,9 @@
       CreateDocument("icing", "email/1", kDefaultScore,
                      /*creation_timestamp_ms=*/kDefaultCreationTimestampMs);
 
-  ICING_ASSERT_OK_AND_ASSIGN(
-      DocumentStore::PutResult put_result1,
-      document_store()->Put(document1, /*num_tokens=*/10));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
+                             document_store()->Put(CreateDocumentWrapper(
+                                 document1, /*num_string_tokens=*/10)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   // Document 1 contains the term "foo" 0 times in the "subject" property
@@ -883,13 +896,13 @@
       CreateDocument("icing", "email/2", kDefaultScore,
                      /*creation_timestamp_ms=*/kDefaultCreationTimestampMs);
 
-  ICING_ASSERT_OK_AND_ASSIGN(
-      DocumentStore::PutResult put_result1,
-      document_store()->Put(document1, /*num_tokens=*/1));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
+                             document_store()->Put(CreateDocumentWrapper(
+                                 document1, /*num_string_tokens=*/1)));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(
-      DocumentStore::PutResult put_result2,
-      document_store()->Put(document2, /*num_tokens=*/1));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
+                             document_store()->Put(CreateDocumentWrapper(
+                                 document2, /*num_string_tokens=*/1)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // Document 1 contains the term "foo" 1 time in the "body" property
@@ -961,13 +974,13 @@
       CreateDocument("icing", "email/2", kDefaultScore,
                      /*creation_timestamp_ms=*/kDefaultCreationTimestampMs);
 
-  ICING_ASSERT_OK_AND_ASSIGN(
-      DocumentStore::PutResult put_result1,
-      document_store()->Put(document1, /*num_tokens=*/1));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
+                             document_store()->Put(CreateDocumentWrapper(
+                                 document1, /*num_string_tokens=*/1)));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(
-      DocumentStore::PutResult put_result2,
-      document_store()->Put(document2, /*num_tokens=*/1));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
+                             document_store()->Put(CreateDocumentWrapper(
+                                 document2, /*num_string_tokens=*/1)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // Document 1 contains the term "foo" 1 time in the "body" property
@@ -1039,9 +1052,9 @@
       CreateDocument("icing", "email/2", kDefaultScore,
                      /*creation_timestamp_ms=*/kDefaultCreationTimestampMs);
 
-  ICING_ASSERT_OK_AND_ASSIGN(
-      DocumentStore::PutResult put_result1,
-      document_store()->Put(document1, /*num_tokens=*/1));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
+                             document_store()->Put(CreateDocumentWrapper(
+                                 document1, /*num_string_tokens=*/1)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   // Document 1 contains the term "foo" 1 time in the "body" property
@@ -1135,13 +1148,13 @@
       CreateDocument("icing", "email/2", kDefaultScore,
                      /*creation_timestamp_ms=*/kDefaultCreationTimestampMs);
 
-  ICING_ASSERT_OK_AND_ASSIGN(
-      DocumentStore::PutResult put_result1,
-      document_store()->Put(document1, /*num_tokens=*/1));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
+                             document_store()->Put(CreateDocumentWrapper(
+                                 document1, /*num_string_tokens=*/1)));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(
-      DocumentStore::PutResult put_result2,
-      document_store()->Put(document2, /*num_tokens=*/1));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
+                             document_store()->Put(CreateDocumentWrapper(
+                                 document2, /*num_string_tokens=*/1)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   // Document 1 contains the term "foo" 1 time in the "body" property
@@ -1215,14 +1228,17 @@
       CreateDocument("icing", "email/3", kDefaultScore,
                      /*creation_timestamp_ms=*/1571100003333);
   // Intentionally inserts documents in a different order
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store()->Put(document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store()->Put(document_util::CreateDocumentWrapper(document1)));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store()->Put(document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store()->Put(document_util::CreateDocumentWrapper(document2)));
   DocumentId document_id2 = put_result2.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             document_store()->Put(document3));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      document_store()->Put(document_util::CreateDocumentWrapper(document3)));
   DocumentId document_id3 = put_result3.new_document_id;
   DocHitInfo doc_hit_info1(document_id1);
   DocHitInfo doc_hit_info2(document_id2);
@@ -1270,14 +1286,17 @@
       CreateDocument("icing", "email/3", kDefaultScore,
                      /*creation_timestamp_ms=*/kDefaultCreationTimestampMs);
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store()->Put(document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store()->Put(document_util::CreateDocumentWrapper(document1)));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store()->Put(document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store()->Put(document_util::CreateDocumentWrapper(document2)));
   DocumentId document_id2 = put_result2.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             document_store()->Put(document3));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      document_store()->Put(document_util::CreateDocumentWrapper(document3)));
   DocumentId document_id3 = put_result3.new_document_id;
 
   // Report usage for doc1 once and doc2 twice.
@@ -1337,14 +1356,17 @@
       CreateDocument("icing", "email/3", kDefaultScore,
                      /*creation_timestamp_ms=*/kDefaultCreationTimestampMs);
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store()->Put(document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store()->Put(document_util::CreateDocumentWrapper(document1)));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store()->Put(document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store()->Put(document_util::CreateDocumentWrapper(document2)));
   DocumentId document_id2 = put_result2.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             document_store()->Put(document3));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      document_store()->Put(document_util::CreateDocumentWrapper(document3)));
   DocumentId document_id3 = put_result3.new_document_id;
 
   // Report usage for doc1 and doc2.
@@ -1442,13 +1464,16 @@
                                            kDefaultCreationTimestampMs);
 
   // Intentionally inserts documents in a different order
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store()->Put(document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store()->Put(document_util::CreateDocumentWrapper(document1)));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             document_store()->Put(document3));
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store()->Put(document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      document_store()->Put(document_util::CreateDocumentWrapper(document3)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store()->Put(document_util::CreateDocumentWrapper(document2)));
   DocumentId document_id2 = put_result2.new_document_id;
   DocumentId document_id3 = put_result3.new_document_id;
   DocHitInfo doc_hit_info1(document_id1);

diff --git a/icing/store/blob-store.cc b/icing/store/blob-store.cc
index eb3b686..8baeda0 100644
--- a/icing/store/blob-store.cc
+++ b/icing/store/blob-store.cc

@@ -141,9 +141,11 @@
 libtextclassifier3::StatusOr<BlobStore> BlobStore::Create(
     const Filesystem* filesystem, std::string base_dir, const Clock* clock,
     int64_t orphan_blob_time_to_live_ms, int32_t compression_level,
-    bool manage_blob_files) {
+    int32_t compression_mem_level, bool manage_blob_files,
+    const FeatureFlags* feature_flags) {
   ICING_RETURN_ERROR_IF_NULL(filesystem);
   ICING_RETURN_ERROR_IF_NULL(clock);
+  ICING_RETURN_ERROR_IF_NULL(feature_flags);
 
   // Make sure the blob file directory exists.
   if (!filesystem->CreateDirectoryRecursively(
@@ -173,19 +175,24 @@
       PortableFileBackedProtoLog<BlobInfoProto>::Create(
           filesystem, blob_info_proto_file_name,
           PortableFileBackedProtoLog<BlobInfoProto>::Options(
-              /*compress_in=*/true, constants::kMaxProtoSize,
-              compression_level)));
+              /*compress_in=*/true, constants::kMaxProtoSize, compression_level,
+              /*compression_threshold_bytes=*/0, compression_mem_level,
+              /*enable_smaller_decompression_buffer_size_in=*/false,
+              /*enable_proto_log_new_header_format_in=*/
+              true, feature_flags->enable_reusable_decompression_buffer())));
+  // TODO(b/435513415): pass feature flags object down to BlobStore and use it.
+  //   It is a remaining task to rollout new header format for proto log.
 
   std::unordered_map<std::string, int> blob_handle_to_offset;
   ICING_ASSIGN_OR_RETURN(
       blob_handle_to_offset,
       LoadBlobHandleToOffsetMapper(log_create_result.proto_log.get()));
 
-  return BlobStore(filesystem, std::move(base_dir), clock,
-                   orphan_blob_time_to_live_ms, compression_level,
-                   manage_blob_files, std::move(log_create_result.proto_log),
-                   std::move(blob_handle_to_offset),
-                   std::move(known_file_names));
+  return BlobStore(
+      filesystem, std::move(base_dir), clock, orphan_blob_time_to_live_ms,
+      compression_level, compression_mem_level, manage_blob_files,
+      std::move(log_create_result.proto_log), std::move(blob_handle_to_offset),
+      std::move(known_file_names));
 }
 
 BlobProto BlobStore::OpenWrite(
@@ -383,7 +390,7 @@
     while (prev_total_read_size < file_size) {
       int32_t size_to_read =
           std::min<int32_t>(kReadBufferSize, file_size - prev_total_read_size);
-      if (!filesystem_.Read(sfd.get(), buffer, size_to_read)) {
+      if (filesystem_.Read(sfd.get(), buffer, size_to_read) != size_to_read) {
         return CreateBlobProtoFromError(absl_ports::InternalError(
             absl_ports::StrCat("Failed to read blob file for handle: ",
                                blob_handle.digest())));
@@ -415,6 +422,50 @@
   return blob_proto;
 }
 
+BlobProto BlobStore::GetAllBlobInfos() {
+  BlobProto blob_proto;
+  blob_proto.mutable_status()->set_code(StatusProto::OK);
+  for (auto itr = blob_handle_to_offset_.begin();
+       itr != blob_handle_to_offset_.end(); ++itr) {
+    auto blob_info_proto_or = blob_info_log_->ReadProto(itr->second);
+    if (!blob_info_proto_or.ok()) {
+      continue;
+    }
+    *blob_proto.add_blob_info_protos() =
+        std::move(blob_info_proto_or).ValueOrDie();
+  }
+  return blob_proto;
+}
+
+BlobProto BlobStore::PutBlobInfos(const BlobProto& blob_proto) {
+  if (manage_blob_files_) {
+    ICING_LOG(ERROR)
+        << "Cannot put blob infos to blob store that manages blob files.";
+    return CreateBlobProtoFromError(absl_ports::FailedPreconditionError(
+        "Cannot put blob infos to blob store that manages blob files."));
+  }
+  BlobProto result_blob_proto;
+  result_blob_proto.mutable_status()->set_code(StatusProto::OK);
+  for (const auto& blob_info_proto : blob_proto.blob_info_protos()) {
+    std::string blob_handle_str =
+        BuildBlobHandleStr(blob_info_proto.blob_handle());
+    auto itr = blob_handle_to_offset_.find(blob_handle_str);
+    if (itr != blob_handle_to_offset_.end()) {
+      // The blob info proto already exists in the blob store.
+      continue;
+    }
+
+    auto blob_info_offset_or = blob_info_log_->WriteProto(blob_info_proto);
+    if (!blob_info_offset_or.ok()) {
+      ICING_LOG(ERROR) << blob_info_offset_or.status().error_message()
+                       << "Failed to write blob info";
+      return CreateBlobProtoFromError(blob_info_offset_or.status());
+    }
+    blob_handle_to_offset_[blob_handle_str] = blob_info_offset_or.ValueOrDie();
+  }
+  return result_blob_proto;
+}
+
 libtextclassifier3::Status BlobStore::PersistToDisk() {
   if (has_mutated_) {
     ICING_RETURN_IF_ERROR(blob_info_log_->PersistToDisk());
@@ -484,7 +535,8 @@
 }
 
 libtextclassifier3::StatusOr<std::vector<std::string>> BlobStore::Optimize(
-    const std::unordered_set<std::string>& dead_blob_handles) {
+    const std::unordered_set<std::string>& dead_blob_handles,
+    const FeatureFlags* feature_flags) {
   std::vector<std::string> blob_file_names_to_remove;
   blob_file_names_to_remove.reserve(dead_blob_handles.size());
 
@@ -496,13 +548,20 @@
         "Unable to delete temp file to prepare to build new blob proto file.");
   }
 
-  ICING_ASSIGN_OR_RETURN(PortableFileBackedProtoLog<BlobInfoProto>::CreateResult
-                             temp_log_create_result,
-                         PortableFileBackedProtoLog<BlobInfoProto>::Create(
-                             &filesystem_, temp_blob_info_proto_file_name,
-                             PortableFileBackedProtoLog<BlobInfoProto>::Options(
-                                 /*compress_in=*/true, constants::kMaxProtoSize,
-                                 compression_level_)));
+  ICING_ASSIGN_OR_RETURN(
+      PortableFileBackedProtoLog<BlobInfoProto>::CreateResult
+          temp_log_create_result,
+      PortableFileBackedProtoLog<BlobInfoProto>::Create(
+          &filesystem_, temp_blob_info_proto_file_name,
+          PortableFileBackedProtoLog<BlobInfoProto>::Options(
+              /*compress_in=*/true, constants::kMaxProtoSize,
+              compression_level_, /*compression_threshold_bytes=*/0,
+              compression_mem_level_,
+              /*enable_smaller_decompression_buffer_size_in=*/false,
+              /*enable_proto_log_new_header_format=*/true,
+              feature_flags->enable_reusable_decompression_buffer())));
+  // TODO(b/435513415): pass feature flags object down to BlobStore and use it.
+  //   It is a remaining task to rollout new header format for proto log.
   std::unique_ptr<PortableFileBackedProtoLog<BlobInfoProto>> new_blob_info_log =
       std::move(temp_log_create_result.proto_log);
 
@@ -565,7 +624,13 @@
           &filesystem_, old_blob_info_proto_file_name,
           PortableFileBackedProtoLog<BlobInfoProto>::Options(
               /*compress_in=*/true, constants::kMaxProtoSize,
-              compression_level_)));
+              compression_level_, /*compression_threshold_bytes=*/0,
+              compression_mem_level_,
+              /*enable_smaller_decompression_buffer_size_in=*/false,
+              /*enable_proto_log_new_header_format=*/true,
+              feature_flags->enable_reusable_decompression_buffer())));
+  // TODO(b/435513415): pass feature flags object down to BlobStore and use it.
+  //   It is a remaining task to rollout new header format for proto log.
   blob_info_log_ = std::move(log_create_result.proto_log);
   blob_handle_to_offset_ = std::move(new_blob_handle_to_offset);
   return blob_file_names_to_remove;

diff --git a/icing/store/blob-store.h b/icing/store/blob-store.h
index 77814a3..9864210 100644
--- a/icing/store/blob-store.h
+++ b/icing/store/blob-store.h

@@ -25,6 +25,7 @@
 
 #include "icing/text_classifier/lib3/utils/base/status.h"
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
+#include "icing/feature-flags.h"
 #include "icing/file/filesystem.h"
 #include "icing/file/portable-file-backed-proto-log.h"
 #include "icing/proto/blob.pb.h"
@@ -68,7 +69,8 @@
   static libtextclassifier3::StatusOr<BlobStore> Create(
       const Filesystem* filesystem, std::string base_dir, const Clock* clock,
       int64_t orphan_blob_time_to_live_ms, int32_t compression_level,
-      bool manage_blob_files);
+      int32_t compression_mem_level, bool manage_blob_files,
+      const FeatureFlags* feature_flags);
 
   // Gets or creates a file for write only purpose for the given blob handle.
   // To mark the blob is completed written, CommitBlob must be called. Once
@@ -142,6 +144,21 @@
   //   NotFoundError if the blob is not found
   BlobProto CommitBlob(const PropertyProto::BlobHandleProto& blob_handle);
 
+  // Gets all the blob info from the blob store.
+  //
+  // Returns:
+  //   BlobProto with all the blob info on success
+  //   InternalError on IO error
+  BlobProto GetAllBlobInfos();
+
+  // Puts the blob info protos from the blob proto to the blob info proto log
+  // file.
+  //
+  // Returns:
+  //   BlobProto with all the blob info on success
+  //   InternalError on IO error
+  BlobProto PutBlobInfos(const BlobProto& blob_proto);
+
   // Persists the blobs to disk.
   libtextclassifier3::Status PersistToDisk();
 
@@ -164,7 +181,8 @@
   //   manages blob files, this list will be empty.
   //   INTERNAL_ERROR on IO error
   libtextclassifier3::StatusOr<std::vector<std::string>> Optimize(
-      const std::unordered_set<std::string>& dead_blob_handles);
+      const std::unordered_set<std::string>& dead_blob_handles,
+      const FeatureFlags* feature_flags);
 
   // Calculates the StorageInfo for the Blob Store.
   //
@@ -178,7 +196,7 @@
   explicit BlobStore(
       const Filesystem* filesystem, std::string base_dir, const Clock* clock,
       int64_t orphan_blob_time_to_live_ms, int32_t compression_level,
-      bool manage_blob_files,
+      int32_t compression_mem_level, bool manage_blob_files,
       std::unique_ptr<PortableFileBackedProtoLog<BlobInfoProto>> blob_info_log,
       std::unordered_map<std::string, int32_t> blob_handle_to_offset,
       std::unordered_set<std::string> known_file_names)
@@ -187,6 +205,7 @@
         clock_(*clock),
         orphan_blob_time_to_live_ms_(orphan_blob_time_to_live_ms),
         compression_level_(compression_level),
+        compression_mem_level_(compression_mem_level),
         manage_blob_files_(manage_blob_files),
         blob_info_log_(std::move(blob_info_log)),
         blob_handle_to_offset_(std::move(blob_handle_to_offset)),
@@ -207,6 +226,7 @@
   const Clock& clock_;
   int64_t orphan_blob_time_to_live_ms_;
   int32_t compression_level_;
+  int32_t compression_mem_level_;
   bool manage_blob_files_;
 
   // The ground truth blob info log file, which is used to read/write/erase

diff --git a/icing/store/document-filter-data.h b/icing/store/document-filter-data.h
index 471ac60..789f45d 100644
--- a/icing/store/document-filter-data.h
+++ b/icing/store/document-filter-data.h

@@ -16,7 +16,6 @@
 #define ICING_STORE_DOCUMENT_FILTER_DATA_H_
 
 #include <cstdint>
-#include <type_traits>
 
 #include "icing/legacy/core/icing-packed-pod.h"
 #include "icing/store/namespace-id.h"
@@ -30,25 +29,23 @@
 class DocumentFilterData {
  public:
   explicit DocumentFilterData(NamespaceId namespace_id,
-                              uint64_t uri_fingerprint,
                               SchemaTypeId schema_type_id,
-                              int64_t expiration_timestamp_ms)
+                              int64_t expiration_timestamp_ms,
+                              int64_t raw_expiration_timestamp_ms)
       : expiration_timestamp_ms_(expiration_timestamp_ms),
-        uri_fingerprint_(uri_fingerprint),
+        raw_expiration_timestamp_ms_(raw_expiration_timestamp_ms),
         namespace_id_(namespace_id),
         schema_type_id_(schema_type_id) {}
 
   bool operator==(const DocumentFilterData& other) const {
     return namespace_id_ == other.namespace_id() &&
-           uri_fingerprint_ == other.uri_fingerprint() &&
            schema_type_id_ == other.schema_type_id() &&
-           expiration_timestamp_ms_ == other.expiration_timestamp_ms();
+           expiration_timestamp_ms_ == other.expiration_timestamp_ms() &&
+           raw_expiration_timestamp_ms_ == other.raw_expiration_timestamp_ms();
   }
 
   NamespaceId namespace_id() const { return namespace_id_; }
 
-  uint64_t uri_fingerprint() const { return uri_fingerprint_; }
-
   SchemaTypeId schema_type_id() const { return schema_type_id_; }
   void set_schema_type_id(SchemaTypeId schema_type_id) {
     schema_type_id_ = schema_type_id;
@@ -56,9 +53,31 @@
 
   int64_t expiration_timestamp_ms() const { return expiration_timestamp_ms_; }
 
+  int64_t raw_expiration_timestamp_ms() const {
+    return raw_expiration_timestamp_ms_;
+  }
+
  private:
+  // The expiration timestamp of the document in milliseconds. This is the min
+  // value of:
+  // - The document's raw expiration timestamp (see
+  //   raw_expiration_timestamp_ms_).
+  // - The propagated expiration timestamps according to its dependencies
+  //   specified in join properties with delete propagation enabled.
+  //
+  // Icing will use this value to:
+  // - Determine if a document is alive or expired.
+  // - Propagte to its dependents' expiration timestamps. Note that in most
+  //   cases the propagation should use this value instead of the raw expiration
+  //   timestamp.
   int64_t expiration_timestamp_ms_;
-  uint64_t uri_fingerprint_;
+
+  // The raw expiration timestamp of the document in milliseconds, calculated by
+  // the document's creation timestamp and ttl. This value is cached and mostly
+  // should only be used to reset expiration_timestamp_ms_ when Icing needs to
+  // recompute the propagated expiration timestamps of all documents.
+  int64_t raw_expiration_timestamp_ms_;
+
   NamespaceId namespace_id_;
   SchemaTypeId schema_type_id_;
 } __attribute__((packed));

diff --git a/icing/store/document-log-creator.cc b/icing/store/document-log-creator.cc
index 683e6e1..0354331 100644
--- a/icing/store/document-log-creator.cc
+++ b/icing/store/document-log-creator.cc

@@ -14,6 +14,7 @@
 
 #include "icing/store/document-log-creator.h"
 
+#include <cstdint>
 #include <memory>
 #include <string>
 #include <utility>
@@ -23,6 +24,7 @@
 #include "icing/absl_ports/annotate.h"
 #include "icing/absl_ports/canonical_errors.h"
 #include "icing/absl_ports/str_cat.h"
+#include "icing/feature-flags.h"
 #include "icing/file/constants.h"
 #include "icing/file/file-backed-proto-log.h"
 #include "icing/file/filesystem.h"
@@ -68,7 +70,10 @@
 libtextclassifier3::StatusOr<DocumentLogCreator::CreateResult>
 DocumentLogCreator::Create(const Filesystem* filesystem,
                            const std::string& base_dir,
-                           int32_t compression_level) {
+                           const FeatureFlags* feature_flags,
+                           int32_t compression_level,
+                           uint32_t compression_threshold_bytes,
+                           int32_t compression_mem_level) {
   bool v0_exists =
       filesystem->FileExists(MakeDocumentLogFilenameV0(base_dir).c_str());
   bool v1_exists =
@@ -77,8 +82,9 @@
   bool new_file = false;
   int preexisting_file_version = kCurrentVersion;
   if (v0_exists && !v1_exists) {
-    ICING_RETURN_IF_ERROR(
-        MigrateFromV0ToV1(filesystem, base_dir, compression_level));
+    ICING_RETURN_IF_ERROR(MigrateFromV0ToV1(
+        filesystem, base_dir, feature_flags, compression_level,
+        compression_threshold_bytes, compression_mem_level));
 
     // Need to regenerate derived files since documents may be written to a
     // different file offset in the log.
@@ -97,8 +103,11 @@
       PortableFileBackedProtoLog<DocumentWrapper>::Create(
           filesystem, MakeDocumentLogFilenameV1(base_dir),
           PortableFileBackedProtoLog<DocumentWrapper>::Options(
-              /*compress_in=*/true, constants::kMaxProtoSize,
-              compression_level)));
+              /*compress_in=*/true, constants::kMaxProtoSize, compression_level,
+              compression_threshold_bytes, compression_mem_level,
+              feature_flags->enable_smaller_decompression_buffer_size(),
+              feature_flags->enable_proto_log_new_header_format(),
+              feature_flags->enable_reusable_decompression_buffer())));
 
   CreateResult create_result = {std::move(log_create_result),
                                 preexisting_file_version, new_file};
@@ -107,7 +116,8 @@
 
 libtextclassifier3::Status DocumentLogCreator::MigrateFromV0ToV1(
     const Filesystem* filesystem, const std::string& base_dir,
-    int32_t compression_level) {
+    const FeatureFlags* feature_flags, int32_t compression_level,
+    uint32_t compression_threshold_bytes, int32_t compression_mem_level) {
   ICING_VLOG(1) << "Migrating from v0 to v1 document log.";
 
   // Our v0 proto log was non-portable, create it so we can read protos out from
@@ -132,9 +142,11 @@
           filesystem, MakeDocumentLogFilenameV1(base_dir),
           PortableFileBackedProtoLog<DocumentWrapper>::Options(
               /*compress_in=*/true,
-              /*max_proto_size_in=*/
-              constants::kMaxProtoSize,
-              /*compression_level_in=*/compression_level));
+              /*max_proto_size_in=*/constants::kMaxProtoSize, compression_level,
+              compression_threshold_bytes, compression_mem_level,
+              feature_flags->enable_smaller_decompression_buffer_size(),
+              feature_flags->enable_proto_log_new_header_format(),
+              feature_flags->enable_reusable_decompression_buffer()));
   if (!v1_create_result_or.ok()) {
     return absl_ports::Annotate(
         v1_create_result_or.status(),

diff --git a/icing/store/document-log-creator.h b/icing/store/document-log-creator.h
index 0c2794a..917e6c3 100644
--- a/icing/store/document-log-creator.h
+++ b/icing/store/document-log-creator.h

@@ -15,10 +15,12 @@
 #ifndef ICING_STORE_DOCUMENT_LOG_CREATOR_H_
 #define ICING_STORE_DOCUMENT_LOG_CREATOR_H_
 
+#include <cstdint>
 #include <string>
 
 #include "icing/text_classifier/lib3/utils/base/status.h"
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
+#include "icing/feature-flags.h"
 #include "icing/file/filesystem.h"
 #include "icing/file/portable-file-backed-proto-log.h"
 #include "icing/proto/document_wrapper.pb.h"
@@ -58,7 +60,8 @@
   //   INTERNAL on any I/O error.
   static libtextclassifier3::StatusOr<DocumentLogCreator::CreateResult> Create(
       const Filesystem* filesystem, const std::string& base_dir,
-      int32_t compression_level);
+      const FeatureFlags* feature_flags, int32_t compression_level,
+      uint32_t compression_threshold_bytes, int32_t compression_mem_level);
 
   // Returns the filename of the document log, without any directory prefixes.
   // Used mainly for testing purposes.
@@ -76,7 +79,8 @@
   //   INTERNAL on I/O error.
   static libtextclassifier3::Status MigrateFromV0ToV1(
       const Filesystem* filesystem, const std::string& base_dir,
-      int32_t compression_level);
+      const FeatureFlags* feature_flags, int32_t compression_level,
+      uint32_t compression_threshold_bytes, int32_t compression_mem_level);
 };
 
 }  // namespace lib

diff --git a/icing/store/document-store.cc b/icing/store/document-store.cc
index da0ff4d..fce44d5 100644
--- a/icing/store/document-store.cc
+++ b/icing/store/document-store.cc

@@ -31,7 +31,6 @@
 #include "icing/absl_ports/canonical_errors.h"
 #include "icing/absl_ports/str_cat.h"
 #include "icing/feature-flags.h"
-#include "icing/file/file-backed-proto-log.h"
 #include "icing/file/file-backed-vector.h"
 #include "icing/file/filesystem.h"
 #include "icing/file/memory-mapped-file-backed-proto-log.h"
@@ -68,10 +67,12 @@
 #include "icing/util/clock.h"
 #include "icing/util/crc32.h"
 #include "icing/util/data-loss.h"
+#include "icing/util/document-util.h"
 #include "icing/util/fingerprint-util.h"
 #include "icing/util/logging.h"
 #include "icing/util/scorable_property_set.h"
 #include "icing/util/status-macros.h"
+#include "icing/util/timestamp-util.h"
 #include "icing/util/tokenized-document.h"
 
 namespace icing {
@@ -110,12 +111,6 @@
 constexpr int32_t kNamespaceMapperMaxSize = 3 * 128 * 1024;  // 384 KiB
 constexpr int32_t kCorpusMapperMaxSize = 3 * 128 * 1024;     // 384 KiB
 
-DocumentWrapper CreateDocumentWrapper(DocumentProto&& document) {
-  DocumentWrapper document_wrapper;
-  *document_wrapper.mutable_document() = std::move(document);
-  return document_wrapper;
-}
-
 std::string MakeHeaderFilename(const std::string& base_dir) {
   return absl_ports::StrCat(base_dir, "/", kDocumentStoreHeaderFilename);
 }
@@ -156,27 +151,6 @@
   return absl_ports::StrCat(base_dir, "/", kCorpusIdMapperFilename);
 }
 
-int64_t CalculateExpirationTimestampMs(int64_t creation_timestamp_ms,
-                                       int64_t ttl_ms) {
-  if (ttl_ms == 0) {
-    // Special case where a TTL of 0 indicates the document should never
-    // expire. int64_t max, interpreted as seconds since epoch, represents
-    // some point in the year 292,277,026,596. So we're probably ok to use
-    // this as "never reaching this point".
-    return std::numeric_limits<int64_t>::max();
-  }
-
-  int64_t expiration_timestamp_ms;
-  if (__builtin_add_overflow(creation_timestamp_ms, ttl_ms,
-                             &expiration_timestamp_ms)) {
-    // Overflow detected. Treat overflow as the same behavior of just int64_t
-    // max
-    return std::numeric_limits<int64_t>::max();
-  }
-
-  return expiration_timestamp_ms;
-}
-
 InitializeStatsProto::RecoveryCause GetRecoveryCause(
     const DocumentLogCreator::CreateResult& create_result,
     bool force_recovery_and_revalidate_documents) {
@@ -287,13 +261,12 @@
 
 }  // namespace
 
-DocumentStore::DocumentStore(const Filesystem* filesystem,
-                             const std::string_view base_dir,
-                             const Clock* clock,
-                             const SchemaStore* schema_store,
-                             const FeatureFlags* feature_flags,
-                             bool pre_mapping_fbv, bool use_persistent_hash_map,
-                             int32_t compression_level)
+DocumentStore::DocumentStore(
+    const Filesystem* filesystem, const std::string_view base_dir,
+    const Clock* clock, const SchemaStore* schema_store,
+    const FeatureFlags* feature_flags, bool pre_mapping_fbv,
+    bool use_persistent_hash_map, int32_t compression_level,
+    uint32_t compression_threshold_bytes, int32_t compression_mem_level)
     : filesystem_(filesystem),
       base_dir_(base_dir),
       clock_(*clock),
@@ -302,19 +275,14 @@
       document_validator_(schema_store),
       pre_mapping_fbv_(pre_mapping_fbv),
       use_persistent_hash_map_(use_persistent_hash_map),
-      compression_level_(compression_level) {}
+      compression_level_(compression_level),
+      compression_threshold_bytes_(compression_threshold_bytes),
+      compression_mem_level_(compression_mem_level) {}
 
 libtextclassifier3::StatusOr<DocumentStore::PutResult> DocumentStore::Put(
-    const DocumentProto& document, int32_t num_tokens,
+    const DocumentWrapper& document_wrapper,
     PutDocumentStatsProto* put_document_stats) {
-  return Put(DocumentProto(document), num_tokens, put_document_stats);
-}
-
-libtextclassifier3::StatusOr<DocumentStore::PutResult> DocumentStore::Put(
-    DocumentProto&& document, int32_t num_tokens,
-    PutDocumentStatsProto* put_document_stats) {
-  document.mutable_internal_fields()->set_length_in_tokens(num_tokens);
-  return InternalPut(std::move(document), put_document_stats);
+  return InternalPut(document_wrapper, put_document_stats);
 }
 
 DocumentStore::~DocumentStore() {
@@ -332,6 +300,7 @@
     const FeatureFlags* feature_flags,
     bool force_recovery_and_revalidate_documents, bool pre_mapping_fbv,
     bool use_persistent_hash_map, int32_t compression_level,
+    uint32_t compression_threshold_bytes, int32_t compression_mem_level,
     InitializeStatsProto* initialize_stats) {
   ICING_RETURN_ERROR_IF_NULL(filesystem);
   ICING_RETURN_ERROR_IF_NULL(clock);
@@ -340,7 +309,8 @@
 
   auto document_store = std::unique_ptr<DocumentStore>(new DocumentStore(
       filesystem, base_dir, clock, schema_store, feature_flags, pre_mapping_fbv,
-      use_persistent_hash_map, compression_level));
+      use_persistent_hash_map, compression_level, compression_threshold_bytes,
+      compression_mem_level));
   ICING_ASSIGN_OR_RETURN(
       InitializeResult initialize_result,
       document_store->Initialize(force_recovery_and_revalidate_documents,
@@ -404,8 +374,9 @@
 libtextclassifier3::StatusOr<DocumentStore::InitializeResult>
 DocumentStore::Initialize(bool force_recovery_and_revalidate_documents,
                           InitializeStatsProto* initialize_stats) {
-  auto create_result_or =
-      DocumentLogCreator::Create(filesystem_, base_dir_, compression_level_);
+  auto create_result_or = DocumentLogCreator::Create(
+      filesystem_, base_dir_, &feature_flags_, compression_level_,
+      compression_threshold_bytes_, compression_mem_level_);
 
   // TODO(b/144458732): Implement a more robust version of TC_ASSIGN_OR_RETURN
   // that can support error logging.
@@ -499,15 +470,16 @@
   }
 
   DocumentStore::Header header;
-  if (!filesystem_->Read(MakeHeaderFilename(base_dir_).c_str(), &header,
-                         sizeof(header))) {
+  if (filesystem_->Read(MakeHeaderFilename(base_dir_).c_str(), &header,
+                        sizeof(header)) != sizeof(header)) {
     return absl_ports::InternalError(
         absl_ports::StrCat("Couldn't read: ", MakeHeaderFilename(base_dir_)));
   }
 
-  if (header.magic != DocumentStore::Header::kMagic) {
-    return absl_ports::InternalError(absl_ports::StrCat(
-        "Invalid header kMagic for file: ", MakeHeaderFilename(base_dir_)));
+  if (header.magic != Header::kMagic) {
+    ICING_LOG(ERROR) << "Invalid header magic for DocumentStore. Expected: "
+                     << Header::kMagic << ", actual: " << header.magic;
+    return absl_ports::InternalError("Invalid header magic for DocumentStore");
   }
 
   // TODO(b/144458732): Implement a more robust version of TC_ASSIGN_OR_RETURN
@@ -713,15 +685,20 @@
             scorable_property_cache_index,
             document_wrapper.document().internal_fields().length_in_tokens())));
 
-    int64_t expiration_timestamp_ms = CalculateExpirationTimestampMs(
-        document_wrapper.document().creation_timestamp_ms(),
-        document_wrapper.document().ttl_ms());
+    int64_t raw_expiration_timestamp_ms =
+        timestamp_util::CalculateRawExpirationTimestampMs(
+            document_wrapper.document().creation_timestamp_ms(),
+            document_wrapper.document().ttl_ms());
 
+    // When regenerating derived files, set the expiration timestamp as the raw
+    // expiration timestamp. The dependency should be evaluated by the caller
+    // and updated to the propagated expiration timestamp later.
     ICING_RETURN_IF_ERROR(UpdateFilterCache(
         new_document_id,
-        DocumentFilterData(namespace_id,
-                           new_doc_nsid_uri_fingerprint.fingerprint(),
-                           schema_type_id, expiration_timestamp_ms)));
+        DocumentFilterData(
+            namespace_id, schema_type_id,
+            /*expiration_timestamp_ms=*/raw_expiration_timestamp_ms,
+            raw_expiration_timestamp_ms)));
     iterator_status = iterator.Advance();
   }
 
@@ -1141,36 +1118,31 @@
 }
 
 libtextclassifier3::StatusOr<DocumentStore::PutResult>
-DocumentStore::InternalPut(DocumentProto&& document,
+DocumentStore::InternalPut(const DocumentWrapper& document_wrapper,
                            PutDocumentStatsProto* put_document_stats) {
   std::unique_ptr<Timer> put_timer = clock_.GetNewTimer();
-  ICING_RETURN_IF_ERROR(document_validator_.Validate(document));
 
   if (put_document_stats != nullptr) {
-    put_document_stats->set_document_size(document.ByteSizeLong());
+    put_document_stats->set_document_size(
+        document_wrapper.document().ByteSizeLong());
   }
 
   // Copy fields needed before they are moved
-  std::string name_space = document.namespace_();
-  std::string uri = document.uri();
-  std::string schema = document.schema();
-  int document_score = document.score();
-  int32_t length_in_tokens = document.internal_fields().length_in_tokens();
-  int64_t creation_timestamp_ms = document.creation_timestamp_ms();
-
-  // Sets the creation timestamp if caller hasn't specified.
-  if (document.creation_timestamp_ms() == 0) {
-    creation_timestamp_ms = clock_.GetSystemTimeMilliseconds();
-    document.set_creation_timestamp_ms(creation_timestamp_ms);
-  }
-
-  int64_t expiration_timestamp_ms =
-      CalculateExpirationTimestampMs(creation_timestamp_ms, document.ttl_ms());
+  std::string name_space = document_wrapper.document().namespace_();
+  std::string uri = document_wrapper.document().uri();
+  std::string schema = document_wrapper.document().schema();
+  int document_score = document_wrapper.document().score();
+  int32_t length_in_tokens =
+      document_wrapper.document().internal_fields().length_in_tokens();
+  int64_t creation_timestamp_ms =
+      document_wrapper.document().creation_timestamp_ms();
+  int64_t raw_expiration_timestamp_ms =
+      timestamp_util::CalculateRawExpirationTimestampMs(
+          creation_timestamp_ms, document_wrapper.document().ttl_ms());
 
   // Update ground truth first
   // TODO(b/144458732): Implement a more robust version of TC_ASSIGN_OR_RETURN
   // that can support error logging.
-  DocumentWrapper document_wrapper = CreateDocumentWrapper(std::move(document));
   auto offset_or = document_log_->WriteProto(document_wrapper);
   if (!offset_or.ok()) {
     ICING_LOG(ERROR) << offset_or.status().error_message()
@@ -1195,6 +1167,7 @@
   }
   PutResult put_result;
   put_result.new_document_id = new_document_id;
+  put_result.expiration_timestamp_ms = raw_expiration_timestamp_ms;
 
   // Update namespace maps
   ICING_ASSIGN_OR_RETURN(
@@ -1233,17 +1206,28 @@
                            corpus_id, document_score, creation_timestamp_ms,
                            scorable_property_cache_index, length_in_tokens)));
 
+  // Here we set the expiration timestamp as the raw expiration timestamp,
+  // because we haven't propagated the expiration timestamps yet and it should
+  // be done in the next step by the caller.
   ICING_RETURN_IF_ERROR(UpdateFilterCache(
       new_document_id,
-      DocumentFilterData(namespace_id,
-                         new_doc_nsid_uri_fingerprint.fingerprint(),
-                         schema_type_id, expiration_timestamp_ms)));
+      DocumentFilterData(
+          namespace_id, schema_type_id,
+          /*expiration_timestamp_ms=*/raw_expiration_timestamp_ms,
+          raw_expiration_timestamp_ms)));
 
   if (old_document_id_or.ok()) {
     // The old document exists, copy over the usage scores and delete the old
     // document.
     DocumentId old_document_id = old_document_id_or.ValueOrDie();
 
+    // Currently, old_document_id is only used by join index v3 to
+    // MigrateParent. old_document_id should be set even if the document is
+    // deleted or expired, but we flag guard it for extra safety.
+    if (feature_flags_.enable_non_existent_qualified_id_join()) {
+      put_result.old_document_id = old_document_id;
+    }
+
     ICING_RETURN_IF_ERROR(
         usage_store_->CloneUsageScores(/*from_document_id=*/old_document_id,
                                        /*to_document_id=*/new_document_id));
@@ -1253,9 +1237,9 @@
     auto delete_status =
         Delete(old_document_id, clock_.GetSystemTimeMilliseconds());
     if (delete_status.ok()) {
-      // The old document had existed and was not previously deleted. Return its
-      // document id to mark it as a replacement.
+      // The old document had existed and was not previously deleted.
       put_result.old_document_id = old_document_id;
+      put_result.was_replacement = true;
     } else if (!absl_ports::IsNotFound(delete_status)) {
       // Real error, pass it up.
       return delete_status;
@@ -1472,9 +1456,8 @@
   if (!filter_data_or.ok()) {
     // This would only happen if document_id is out of range of the
     // filter_cache, meaning we got some invalid document_id. Callers should
-    // already have checked that their document_id is valid or used
-    // DoesDocumentExist(WithStatus). Regardless, return std::nullopt since the
-    // document doesn't exist.
+    // already have checked the status or validated their document_id.
+    // Regardless, return std::nullopt since the document doesn't exist.
     return std::nullopt;
   }
 
@@ -1483,14 +1466,140 @@
   return std::move(filter_data_or).ValueOrDie();
 }
 
+bool DocumentStore::IsDocumentAlive(DocumentId document_id,
+                                    int64_t current_time_ms) const {
+  return GetAliveDocumentFilterData(document_id, current_time_ms).has_value();
+}
+
+bool DocumentStore::IsDocumentAlive(std::string_view name_space,
+                                    std::string_view uri,
+                                    int64_t current_time_ms) const {
+  auto document_id_or = GetDocumentId(name_space, uri);
+  if (!document_id_or.ok()) {
+    return false;
+  }
+  return IsDocumentAlive(document_id_or.ValueOrDie(), current_time_ms);
+}
+
+libtextclassifier3::StatusOr<
+    DocumentStore::UpdateDocumentExpirationTimestampResult>
+DocumentStore::UpdateDocumentExpirationTimestamp(
+    DocumentId document_id, int64_t expiration_timestamp_ms) {
+  if (document_id < 0 || document_id >= filter_cache_->num_elements()) {
+    return absl_ports::InvalidArgumentError(IcingStringUtil::StringPrintf(
+        "Document id '%d' invalid.", document_id));
+  }
+
+  std::optional<DocumentFilterData> filter_data =
+      GetNonDeletedDocumentFilterData(document_id);
+  if (!filter_data.has_value()) {
+    return absl_ports::NotFoundError(IcingStringUtil::StringPrintf(
+        "Document id '%d' is not alive", document_id));
+  }
+
+  // Currently we only support decreasing the expiration timestamp.
+  if (expiration_timestamp_ms < filter_data->expiration_timestamp_ms()) {
+    ICING_RETURN_IF_ERROR(filter_cache_->Set(
+        document_id, DocumentFilterData(
+                         filter_data->namespace_id(),
+                         filter_data->schema_type_id(), expiration_timestamp_ms,
+                         filter_data->raw_expiration_timestamp_ms())));
+    return UpdateDocumentExpirationTimestampResult{
+        .final_expiration_timestamp_ms = expiration_timestamp_ms,
+        .was_updated = true};
+  }
+  return UpdateDocumentExpirationTimestampResult{
+      .final_expiration_timestamp_ms = filter_data->expiration_timestamp_ms(),
+      .was_updated = false};
+}
+
+libtextclassifier3::Status
+DocumentStore::ResetAllAliveExpirationTimestampsToRaw(int64_t current_time_ms) {
+  for (DocumentId document_id = 0; document_id < filter_cache_->num_elements();
+       ++document_id) {
+    std::optional<DocumentFilterData> filter_data =
+        GetAliveDocumentFilterData(document_id, current_time_ms);
+    if (!filter_data.has_value()) {
+      continue;
+    }
+
+    if (filter_data->expiration_timestamp_ms() !=
+        filter_data->raw_expiration_timestamp_ms()) {
+      ICING_RETURN_IF_ERROR(filter_cache_->Set(
+          document_id,
+          DocumentFilterData(filter_data->namespace_id(),
+                             filter_data->schema_type_id(),
+                             /*expiration_timestamp_ms=*/
+                             filter_data->raw_expiration_timestamp_ms(),
+                             filter_data->raw_expiration_timestamp_ms())));
+    }
+  }
+  return libtextclassifier3::Status::OK;
+}
+
+libtextclassifier3::StatusOr<std::vector<DocumentStore::DocumentMetadata>>
+DocumentStore::PurgeExpiredDocuments(int64_t current_time_ms) {
+  std::vector<DocumentMetadata> deleted_doc_metadata_list;
+  for (DocumentId document_id = 0; document_id < filter_cache_->num_elements();
+       ++document_id) {
+    std::optional<DocumentFilterData> filter_data =
+        GetNonDeletedDocumentFilterData(document_id);
+    if (!filter_data.has_value()) {
+      continue;
+    }
+
+    // If the document is expired or expires within the threshold, delete it.
+    if (current_time_ms +
+            feature_flags_.expired_document_purge_threshold_ms() >=
+        filter_data->expiration_timestamp_ms()) {
+      // Expired but not deleted yet. Force delete it.
+      auto deleted_doc_metadata_or = ForceDelete(document_id);
+      if (!deleted_doc_metadata_or.ok()) {
+        if (absl_ports::IsNotFound(deleted_doc_metadata_or.status())) {
+          // This should not happen, but let's handle it here anyway.
+          continue;
+        }
+        // Real error. Return it.
+        return std::move(deleted_doc_metadata_or).status();
+      }
+      deleted_doc_metadata_list.push_back(
+          std::move(deleted_doc_metadata_or).ValueOrDie());
+    }
+  }
+  return deleted_doc_metadata_list;
+}
+
+int64_t DocumentStore::GetNextExpiredDocumentTimestampMs(
+    int64_t current_time_ms) {
+  int64_t next_expired_document_ts_ms = -1;
+  for (DocumentId document_id = 0; document_id < filter_cache_->num_elements();
+       ++document_id) {
+    std::optional<DocumentFilterData> filter_data =
+        GetAliveDocumentFilterData(document_id, current_time_ms);
+    if (!filter_data.has_value()) {
+      continue;
+    }
+
+    // Note: std::numeric_limits<int64_t>::max() is used to indicate that a
+    // document never expires, so we should avoid returning it.
+    if (filter_data->expiration_timestamp_ms() !=
+            std::numeric_limits<int64_t>::max() &&
+        (next_expired_document_ts_ms == -1 ||
+         filter_data->expiration_timestamp_ms() <
+             next_expired_document_ts_ms)) {
+      next_expired_document_ts_ms = filter_data->expiration_timestamp_ms();
+    }
+  }
+  return next_expired_document_ts_ms;
+}
+
 bool DocumentStore::IsDeleted(DocumentId document_id) const {
   auto file_offset_or = document_id_mapper_->Get(document_id);
   if (!file_offset_or.ok()) {
     // This would only happen if document_id is out of range of the
     // document_id_mapper, meaning we got some invalid document_id. Callers
-    // should already have checked that their document_id is valid or used
-    // DoesDocumentExist(WithStatus). Regardless, return true since the
-    // document doesn't exist.
+    // should already have checked the status or validated their document_id.
+    // Regardless, return true since the document doesn't exist.
     return true;
   }
   int64_t file_offset = *file_offset_or.ValueOrDie();
@@ -1506,9 +1615,8 @@
   if (!filter_data_or.ok()) {
     // This would only happen if document_id is out of range of the
     // filter_cache, meaning we got some invalid document_id. Callers should
-    // already have checked that their document_id is valid or used
-    // DoesDocumentExist(WithStatus). Regardless, return std::nullopt since the
-    // document doesn't exist.
+    // already have checked the status or validated their document_id.
+    // Regardless, return std::nullopt since the document doesn't exist.
     return std::nullopt;
   }
   DocumentFilterData document_filter_data = filter_data_or.ValueOrDie();
@@ -1560,6 +1668,46 @@
   return ClearDerivedData(document_id);
 }
 
+libtextclassifier3::StatusOr<DocumentStore::DocumentMetadata>
+DocumentStore::ForceDelete(DocumentId document_id) {
+  if (document_id < 0 || document_id >= document_id_mapper_->num_elements()) {
+    return absl_ports::InvalidArgumentError(IcingStringUtil::StringPrintf(
+        "Document id '%d' is invalid to force delete.", document_id));
+  }
+
+  auto document_log_offset_or = document_id_mapper_->Get(document_id);
+  if (!document_log_offset_or.ok()) {
+    return absl_ports::InternalError("Failed to find document offset.");
+  }
+  int64_t document_log_offset = *document_log_offset_or.ValueOrDie();
+
+  if (document_log_offset == kDocDeletedFlag) {
+    // Already deleted.
+    return absl_ports::NotFoundError(IcingStringUtil::StringPrintf(
+        "Document id '%d' was already deleted.", document_id));
+  }
+
+  // Read the document since we need to return DocumentMetadata.
+  auto document_wrapper_or = document_log_->ReadProto(document_log_offset);
+  if (!document_wrapper_or.ok()) {
+    ICING_LOG(ERROR) << document_wrapper_or.status().error_message()
+                     << "Failed to read from document log.";
+    return std::move(document_wrapper_or).status();
+  }
+  DocumentWrapper document_wrapper =
+      std::move(document_wrapper_or).ValueOrDie();
+
+  // Erases document proto.
+  ICING_RETURN_IF_ERROR(document_log_->EraseProto(document_log_offset));
+  ICING_RETURN_IF_ERROR(ClearDerivedData(document_id));
+
+  return DocumentMetadata{
+      .schema_type_name = document_wrapper.document().schema(),
+      .name_space = document_wrapper.document().namespace_(),
+      .uri = document_wrapper.document().uri(),
+      .document_id = document_id};
+}
+
 libtextclassifier3::StatusOr<NamespaceId> DocumentStore::GetNamespaceId(
     std::string_view name_space) const {
   return namespace_mapper_->Get(name_space);
@@ -1573,54 +1721,66 @@
   return corpus_mapper_->Get(corpus_nsid_schema_fp.EncodeToCString());
 }
 
-libtextclassifier3::StatusOr<int32_t> DocumentStore::GetResultGroupingEntryId(
+std::optional<int32_t> DocumentStore::GetResultGroupingEntryId(
     ResultSpecProto::ResultGroupingType result_group_type,
-    const std::string_view name_space, const std::string_view schema) const {
-  auto namespace_id = GetNamespaceId(name_space);
-  auto schema_type_id = schema_store_->GetSchemaTypeId(schema);
-  switch (result_group_type) {
-    case ResultSpecProto::NONE:
-      return absl_ports::InvalidArgumentError(
-          "Cannot group by ResultSpecProto::NONE");
-    case ResultSpecProto::SCHEMA_TYPE:
-      if (schema_type_id.ok()) {
-        return schema_type_id.ValueOrDie();
-      }
-      break;
-    case ResultSpecProto::NAMESPACE:
-      if (namespace_id.ok()) {
-        return namespace_id.ValueOrDie();
-      }
-      break;
-    case ResultSpecProto::NAMESPACE_AND_SCHEMA_TYPE:
-      if (namespace_id.ok() && schema_type_id.ok()) {
-        // TODO(b/258715421): Temporary workaround to get a
-        //                    ResultGroupingEntryId given the Namespace string
-        //                    and Schema string.
-        return namespace_id.ValueOrDie() << 16 | schema_type_id.ValueOrDie();
-      }
-      break;
-  }
-  return absl_ports::NotFoundError("Cannot generate ResultGrouping Entry Id");
+    std::string_view name_space, std::string_view schema) const {
+  auto namespace_id_or = GetNamespaceId(name_space);
+  auto schema_type_id_or = schema_store_->GetSchemaTypeId(schema);
+
+  NamespaceId namespace_id =
+      namespace_id_or.ok() ? namespace_id_or.ValueOrDie() : kInvalidNamespaceId;
+  SchemaTypeId schema_type_id = schema_type_id_or.ok()
+                                    ? schema_type_id_or.ValueOrDie()
+                                    : kInvalidSchemaTypeId;
+  return GetResultGroupingEntryId(result_group_type, namespace_id,
+                                  schema_type_id);
 }
 
-libtextclassifier3::StatusOr<int32_t> DocumentStore::GetResultGroupingEntryId(
+std::optional<int32_t> DocumentStore::GetResultGroupingEntryId(
     ResultSpecProto::ResultGroupingType result_group_type,
-    const NamespaceId namespace_id, const SchemaTypeId schema_type_id) const {
+    NamespaceId namespace_id, SchemaTypeId schema_type_id) const {
+  static_assert(sizeof(NamespaceId) * 8 <= 16,
+                "Current ResultGroupingEntryId encoding only supports "
+                "namespace id up to 16 bits.");
+  static_assert(sizeof(SchemaTypeId) * 8 <= 16,
+                "Current ResultGroupingEntryId encoding only supports schema "
+                "type id up to 16 bits.");
+
+  // Note: this encoding method only works for a single
+  // ResultSpecProto::ResultGroupingType in a single search request. If multiple
+  // types can be used in the same search request, this encoding method needs to
+  // be updated since there will be encoded id collisions for NAMESPACE and
+  // SCHEMA_TYPE.
+
   switch (result_group_type) {
-    case ResultSpecProto::NONE:
-      return absl_ports::InvalidArgumentError(
-          "Cannot group by ResultSpecProto::NONE");
-    case ResultSpecProto::SCHEMA_TYPE:
+    case ResultSpecProto::ResultGroupingType::
+        ResultSpecProto_ResultGroupingType_NONE:
+      return std::nullopt;
+    case ResultSpecProto::ResultGroupingType::
+        ResultSpecProto_ResultGroupingType_SCHEMA_TYPE: {
+      if (schema_type_id == kInvalidSchemaTypeId) {
+        return std::nullopt;
+      }
       return schema_type_id;
-    case ResultSpecProto::NAMESPACE:
+    }
+    case ResultSpecProto::ResultGroupingType::
+        ResultSpecProto_ResultGroupingType_NAMESPACE: {
+      if (namespace_id == kInvalidNamespaceId) {
+        return std::nullopt;
+      }
       return namespace_id;
-    case ResultSpecProto::NAMESPACE_AND_SCHEMA_TYPE:
+    }
+    case ResultSpecProto::ResultGroupingType::
+        ResultSpecProto_ResultGroupingType_NAMESPACE_AND_SCHEMA_TYPE: {
+      if (namespace_id == kInvalidNamespaceId ||
+          schema_type_id == kInvalidSchemaTypeId) {
+        return std::nullopt;
+      }
       // TODO(b/258715421): Temporary workaround to get a ResultGroupingEntryId
       //                    given the Namespace Id and SchemaType Id.
-      return namespace_id << 16 | schema_type_id;
+      return (static_cast<int32_t>(namespace_id) << 16) | schema_type_id;
+    }
   }
-  return absl_ports::NotFoundError("Cannot generate ResultGrouping Entry Id");
 }
 
 libtextclassifier3::StatusOr<DocumentAssociatedScoreData>
@@ -1806,14 +1966,25 @@
 }
 
 libtextclassifier3::Status DocumentStore::PersistToDisk(
-    PersistType::Code persist_type) {
-  ICING_RETURN_IF_ERROR(document_log_->PersistToDisk());
+    PersistType::Code persist_type, PersistToDiskStatsProto* persist_stats) {
+  ICING_RETURN_IF_ERROR(document_log_->PersistToDisk(persist_stats, &clock_));
   if (persist_type == PersistType::LITE) {
     // only persist the document log.
     return libtextclassifier3::Status::OK;
   }
+
+  std::unique_ptr<Timer> overall_timer;
+
+  if (persist_stats) {
+    overall_timer = clock_.GetNewTimer();
+  }
   if (persist_type == PersistType::RECOVERY_PROOF) {
-    return UpdateChecksum().status();
+    libtextclassifier3::Status status = UpdateChecksum().status();
+    if (persist_stats) {
+      persist_stats->set_document_store_checksum_update_latency_ms(
+          overall_timer->GetElapsedMilliseconds());
+    }
+    return status;
   }
   ICING_RETURN_IF_ERROR(document_key_mapper_->PersistToDisk());
   ICING_RETURN_IF_ERROR(document_id_mapper_->PersistToDisk());
@@ -1825,8 +1996,18 @@
   ICING_RETURN_IF_ERROR(corpus_mapper_->PersistToDisk());
   ICING_RETURN_IF_ERROR(corpus_score_cache_->PersistToDisk());
 
+  if (persist_stats) {
+    persist_stats->set_document_store_components_persist_latency_ms(
+        overall_timer->GetElapsedMilliseconds());
+    overall_timer = clock_.GetNewTimer();
+  }
+
   // Update the combined checksum and write to header file.
   ICING_RETURN_IF_ERROR(UpdateChecksum());
+  if (persist_stats) {
+    persist_stats->set_document_store_checksum_update_latency_ms(
+        overall_timer->GetElapsedMilliseconds());
+  }
   return libtextclassifier3::Status::OK;
 }
 
@@ -1968,8 +2149,10 @@
   return CalculateDocumentStatusCounts(std::move(storage_info));
 }
 
-libtextclassifier3::Status DocumentStore::UpdateSchemaStore(
-    const SchemaStore* schema_store) {
+libtextclassifier3::StatusOr<DocumentStore::UpdateSchemaStoreResult>
+DocumentStore::UpdateSchemaStore(const SchemaStore* schema_store) {
+  UpdateSchemaStoreResult update_result;
+
   // Update all references to the SchemaStore
   schema_store_ = schema_store;
   document_validator_.UpdateSchemaStore(schema_store);
@@ -1990,7 +2173,7 @@
     }
 
     // Guaranteed to have a document now.
-    DocumentProto document = document_or.ValueOrDie();
+    DocumentProto document = std::move(document_or).ValueOrDie();
 
     // Revalidate that this document is still compatible
     if (document_validator_.Validate(document).ok()) {
@@ -2001,7 +2184,10 @@
           typename FileBackedVector<DocumentFilterData>::MutableView
               doc_filter_data_view,
           filter_cache_->GetMutable(document_id));
-      doc_filter_data_view.Get().set_schema_type_id(schema_type_id);
+      if (doc_filter_data_view.Get().schema_type_id() != schema_type_id) {
+        doc_filter_data_view.Get().set_schema_type_id(schema_type_id);
+        update_result.derived_files_changed = true;
+      }
     } else {
       // Document is no longer valid with the new SchemaStore. Mark as
       // deleted
@@ -2011,15 +2197,21 @@
         // Real error, pass up
         return delete_status;
       }
+
+      ++update_result.deleted_document_count;
+      update_result.derived_files_changed = true;
     }
   }
 
-  return libtextclassifier3::Status::OK;
+  return update_result;
 }
 
-libtextclassifier3::Status DocumentStore::OptimizedUpdateSchemaStore(
+libtextclassifier3::StatusOr<DocumentStore::UpdateSchemaStoreResult>
+DocumentStore::OptimizedUpdateSchemaStore(
     const SchemaStore* schema_store,
     const SchemaStore::SetSchemaResult& set_schema_result) {
+  UpdateSchemaStoreResult update_result;
+
   if (!set_schema_result.success) {
     // No new schema was set, no work to be done
     return libtextclassifier3::Status::OK;
@@ -2067,7 +2259,12 @@
             typename FileBackedVector<DocumentFilterData>::MutableView
                 doc_filter_data_view,
             filter_cache_->GetMutable(document_id));
+
+        // Since the old schema type id is present in
+        // old_schema_type_ids_changed, we know that (new) schema_type_id is
+        // different, so update it directly without checking.
         doc_filter_data_view.Get().set_schema_type_id(schema_type_id);
+        update_result.derived_files_changed = true;
       }
       if (revalidate_document) {
         delete_document = !document_validator_.Validate(document).ok();
@@ -2077,14 +2274,17 @@
     if (delete_document) {
       // Document is no longer valid with the new SchemaStore. Mark as deleted
       auto delete_status = Delete(document_id, current_time_ms);
-      if (!delete_status.ok() && !absl_ports::IsNotFound(delete_status)) {
+      if (delete_status.ok()) {
+        ++update_result.deleted_document_count;
+        update_result.derived_files_changed = true;
+      } else if (!absl_ports::IsNotFound(delete_status)) {
         // Real error, pass up
         return delete_status;
       }
     }
   }
 
-  return libtextclassifier3::Status::OK;
+  return update_result;
 }
 
 libtextclassifier3::Status DocumentStore::RegenerateScorablePropertyCache(
@@ -2125,11 +2325,6 @@
   return libtextclassifier3::Status::OK;
 }
 
-// TODO(b/121227117): Implement Optimize()
-libtextclassifier3::Status DocumentStore::Optimize() {
-  return libtextclassifier3::Status::OK;
-}
-
 libtextclassifier3::StatusOr<DocumentStore::OptimizeResult>
 DocumentStore::OptimizeInto(
     const std::string& new_directory, const LanguageSegmenter* lang_segmenter,
@@ -2147,6 +2342,7 @@
           filesystem_, new_directory, &clock_, schema_store_, &feature_flags_,
           /*force_recovery_and_revalidate_documents=*/false, pre_mapping_fbv_,
           use_persistent_hash_map_, compression_level_,
+          compression_threshold_bytes_, compression_mem_level_,
           /*initialize_stats=*/nullptr));
   std::unique_ptr<DocumentStore> new_doc_store =
       std::move(doc_store_create_result.document_store);
@@ -2178,23 +2374,42 @@
 
   int64_t current_time_ms = clock_.GetSystemTimeMilliseconds();
   for (DocumentId document_id = 0; document_id < document_cnt; document_id++) {
-    auto document_or = Get(document_id, /*clear_internal_fields=*/false);
-    if (absl_ports::IsNotFound(document_or.status())) {
+    libtextclassifier3::StatusOr<DocumentProto> document_or;
+    if (feature_flags_.enable_optimize_improvements()) {
       if (IsDeleted(document_id)) {
         ++num_deleted_documents;
+        continue;
       } else if (!GetNonExpiredDocumentFilterData(document_id,
                                                   current_time_ms)) {
         ++num_expired_documents;
+        continue;
       }
-      continue;
-    } else if (!document_or.ok()) {
-      // Real error, pass up
+      // The only expected errors for Get are for deleted
+      // or expired documents.
+      document_or = Get(document_id, /*clear_internal_fields=*/false);
+    } else {
+      document_or = Get(document_id, /*clear_internal_fields=*/false);
+      if (absl_ports::IsNotFound(document_or.status())) {
+        if (IsDeleted(document_id)) {
+          ++num_deleted_documents;
+        } else if (!GetNonExpiredDocumentFilterData(document_id,
+                                                    current_time_ms)) {
+          ++num_expired_documents;
+        }
+        continue;
+      }
+    }
+    if (!document_or.ok()) {
+      // Real error, pass up.
       return absl_ports::Annotate(
           document_or.status(),
           IcingStringUtil::StringPrintf(
               "Failed to retrieve Document for DocumentId %d", document_id));
     }
 
+    // TODO(b/405467475): add a new Get method to return DocumentWrapper and
+    // change document_to_keep to DocumentWrapper.
+
     // Guaranteed to have a document now.
     DocumentProto document_to_keep = std::move(document_or).ValueOrDie();
     // Remove blobs that still have reference are removed from the
@@ -2204,8 +2419,10 @@
 
     libtextclassifier3::StatusOr<PutResult> put_result_or;
     if (document_to_keep.internal_fields().length_in_tokens() == 0) {
+      // Create a TokenizedDocument. length_in_tokens will be set there.
       auto tokenized_document_or = TokenizedDocument::Create(
-          schema_store_, lang_segmenter, document_to_keep);
+          schema_store_, lang_segmenter, current_time_ms,
+          std::move(document_to_keep));
       if (!tokenized_document_or.ok()) {
         return absl_ports::Annotate(
             tokenized_document_or.status(),
@@ -2214,12 +2431,12 @@
       }
       TokenizedDocument tokenized_document(
           std::move(tokenized_document_or).ValueOrDie());
-      put_result_or = new_doc_store->Put(
-          std::move(document_to_keep), tokenized_document.num_string_tokens());
+      put_result_or = new_doc_store->Put(tokenized_document.document_wrapper());
     } else {
       // TODO(b/144458732): Implement a more robust version of
       // TC_ASSIGN_OR_RETURN that can support error logging.
-      put_result_or = new_doc_store->InternalPut(std::move(document_to_keep));
+      put_result_or = new_doc_store->InternalPut(
+          document_util::CreateDocumentWrapper(std::move(document_to_keep)));
     }
     if (!put_result_or.ok()) {
       ICING_LOG(ERROR) << put_result_or.status().error_message()
@@ -2397,10 +2614,9 @@
 
   // Resets the filter cache entry
   ICING_RETURN_IF_ERROR(UpdateFilterCache(
-      document_id,
-      DocumentFilterData(kInvalidNamespaceId, /*uri_fingerprint=*/0,
-                         kInvalidSchemaTypeId,
-                         /*expiration_timestamp_ms=*/-1)));
+      document_id, DocumentFilterData(kInvalidNamespaceId, kInvalidSchemaTypeId,
+                                      /*expiration_timestamp_ms=*/-1,
+                                      /*raw_expiration_timestamp_ms=*/-1)));
 
   // Clears the usage scores.
   return usage_store_->DeleteUsageScores(document_id);

diff --git a/icing/store/document-store.h b/icing/store/document-store.h
index 7246235..516d9db 100644
--- a/icing/store/document-store.h
+++ b/icing/store/document-store.h

@@ -16,6 +16,7 @@
 #define ICING_STORE_DOCUMENT_STORE_H_
 
 #include <cstdint>
+#include <limits>
 #include <memory>
 #include <optional>
 #include <string>
@@ -66,8 +67,8 @@
  public:
   struct Header {
     // Previously used magic numbers, please avoid reusing those:
-    // [0x1b99c8b0, 0x3e005b5e]
-    static constexpr int32_t kMagic = 0x8a32cd1f;
+    // [0x1b99c8b0, 0x3e005b5e, 0x8a32cd1f, 0x2eea71de]
+    static constexpr int32_t kMagic = 0x142f84cd;
 
     // Holds the magic as a quick sanity check against file corruption.
     int32_t magic;
@@ -115,6 +116,13 @@
     bool derived_files_regenerated;
   };
 
+  struct DocumentMetadata {
+    std::string schema_type_name;
+    std::string name_space;
+    std::string uri;
+    DocumentId document_id;
+  };
+
   // Not copyable
   DocumentStore(const DocumentStore&) = delete;
   DocumentStore& operator=(const DocumentStore&) = delete;
@@ -151,6 +159,7 @@
       const FeatureFlags* feature_flags,
       bool force_recovery_and_revalidate_documents, bool pre_mapping_fbv,
       bool use_persistent_hash_map, int32_t compression_level,
+      uint32_t compression_threshold_bytes, int32_t compression_mem_level,
       InitializeStatsProto* initialize_stats);
 
   // Discards all derived data in the document store.
@@ -176,7 +185,7 @@
   // of deleted or expired documents.
   int num_documents() const { return document_id_mapper_->num_elements(); }
 
-  // Puts the document into document store.
+  // Puts a single document into document store.
   //
   // If put_document_stats is present, the fields related to DocumentStore will
   // be populated.
@@ -193,16 +202,11 @@
   struct PutResult {
     DocumentId old_document_id = kInvalidDocumentId;
     DocumentId new_document_id = kInvalidDocumentId;
-
-    bool was_replacement() const {
-      return old_document_id != kInvalidDocumentId;
-    }
+    int64_t expiration_timestamp_ms = std::numeric_limits<int64_t>::max();
+    bool was_replacement = false;
   };
   libtextclassifier3::StatusOr<PutResult> Put(
-      const DocumentProto& document, int32_t num_tokens = 0,
-      PutDocumentStatsProto* put_document_stats = nullptr);
-  libtextclassifier3::StatusOr<PutResult> Put(
-      DocumentProto&& document, int32_t num_tokens = 0,
+      const DocumentWrapper& document_wrapper,
       PutDocumentStatsProto* put_document_stats = nullptr);
 
   // Finds and returns the document identified by the given key (namespace +
@@ -247,8 +251,10 @@
   // or expired). Order of namespaces is undefined.
   std::vector<std::string> GetAllNamespaces() const;
 
-  // Deletes the document identified by the given namespace and uri. The
-  // document proto will be erased immediately.
+  // TODO(b/384947619): migrate Delete APIs to return DocumentMetadata.
+
+  // Deletes the document identified by the given namespace and uri, only if it
+  // is still alive. The document proto will be erased immediately.
   //
   // NOTE:
   //    Space is not reclaimed for deleted documents until Optimize() is
@@ -262,8 +268,8 @@
                                     std::string_view uri,
                                     int64_t current_time_ms);
 
-  // Deletes the document identified by the given document_id. The document
-  // proto will be erased immediately.
+  // Deletes the document identified by the given document_id, only if it is
+  // still alive. The document proto will be erased immediately.
   //
   // NOTE:
   //    Space is not reclaimed for deleted documents until Optimize() is
@@ -277,6 +283,24 @@
   libtextclassifier3::Status Delete(DocumentId document_id,
                                     int64_t current_time_ms);
 
+  // Deletes the document identified by the given document_id. The document
+  // proto will be erased immediately.
+  //
+  // Different from Delete(), this method promises that the document proto will
+  // be erased even if the document is expired.
+  //
+  // NOTE:
+  //    Space is not reclaimed for deleted documents until Optimize() is
+  //    called.
+  //
+  // Returns:
+  //   DocumentMetadata of the deleted document on success
+  //   NOT_FOUND_ERROR if the document doesn't exist or has been deleted
+  //   INTERNAL_ERROR on IO error
+  //   INVALID_ARGUMENT_ERROR if document_id is invalid.
+  libtextclassifier3::StatusOr<DocumentMetadata> ForceDelete(
+      DocumentId document_id);
+
   // Returns the NamespaceId of the string namespace
   //
   // Returns:
@@ -289,9 +313,10 @@
   // Helper method to find a DocumentId that is associated with the given
   // namespace and uri.
   //
-  // NOTE: The DocumentId may refer to a invalid document (deleted
-  // or expired). Callers can call DoesDocumentExist(document_id) to ensure it
-  // refers to a valid Document.
+  // NOTE: if succeeded, it always returns a valid DocumentId, but this
+  // DocumentId may refer to a invalid document (deleted or expired). Callers
+  // can call GetAliveDocumentFilterData(document_id, current_time_ms) and check
+  // the return value to ensure it refers to an alive Document.
   //
   // Returns:
   //   A DocumentId on success
@@ -303,9 +328,10 @@
   // Helper method to find a DocumentId that is associated with the given
   // NamespaceIdFingerprint.
   //
-  // NOTE: The DocumentId may refer to a invalid document (deleted
-  // or expired). Callers can call DoesDocumentExist(document_id) to ensure it
-  // refers to a valid Document.
+  // NOTE: if succeeded, it always returns a valid DocumentId, but this
+  // DocumentId may refer to a invalid document (deleted or expired). Callers
+  // can call GetAliveDocumentFilterData(document_id, current_time_ms) and check
+  // the return value to ensure it refers to an alive Document.
   //
   // Returns:
   //   A DocumentId on success
@@ -328,30 +354,31 @@
   //
   // NOTE: ResultGroupingEntryIds that are generated by calls with different
   // ResultGroupingTypes should not be compared. Returned ResultGroupingEntryIds
-  // are only guarenteed to be unique within their own ResultGroupingType.
+  // are only guaranteed to be unique within their own ResultGroupingType.
   //
   // Returns:
-  //   A ResultGroupingEntryId on success
-  //   NOT_FOUND if the key doesn't exist
-  //   INTERNAL_ERROR on IO error
-  libtextclassifier3::StatusOr<int32_t> GetResultGroupingEntryId(
+  //   A valid ResultGroupingEntryId on success
+  //   std::nullopt there is no such result grouping entry id (e.g. result group
+  //     type is not supported, result group type is based on the given
+  //     namespace/schema but the namespace/schema doesn't exist).
+  std::optional<int32_t> GetResultGroupingEntryId(
       ResultSpecProto::ResultGroupingType result_group_type,
-      const std::string_view name_space, const std::string_view schema) const;
+      std::string_view name_space, std::string_view schema) const;
 
   // Returns the ResultGrouping Entry Id associated with the given NamespaceId
   // and SchemaTypeId
   //
   // NOTE: ResultGroupingEntryIds that are generated by calls with different
   // ResultGroupingTypes should not be compared. Returned ResultGroupingEntryIds
-  // are only guarenteed to be unique within their own ResultGroupingType.
+  // are only guaranteed to be unique within their own ResultGroupingType.
   //
   // Returns:
-  //   A ResultGroupingEntryId on success
-  //   NOT_FOUND if the key doesn't exist
-  //   INTERNAL_ERROR on IO error
-  libtextclassifier3::StatusOr<int32_t> GetResultGroupingEntryId(
+  //   A valid ResultGroupingEntryId on success
+  //   std::nullopt if there is no such result grouping entry id (result group
+  //     type is not supported).
+  std::optional<int32_t> GetResultGroupingEntryId(
       ResultSpecProto::ResultGroupingType result_group_type,
-      const NamespaceId namespace_id, const SchemaTypeId schema_type_id) const;
+      NamespaceId namespace_id, SchemaTypeId schema_type_id) const;
 
   // Returns the DocumentAssociatedScoreData of the document specified by the
   // DocumentId.
@@ -397,6 +424,74 @@
   std::optional<DocumentFilterData> GetNonDeletedDocumentFilterData(
       DocumentId document_id) const;
 
+  // Checks if the document (given document_id) is alive or not.
+  bool IsDocumentAlive(DocumentId document_id, int64_t current_time_ms) const;
+
+  // Checks if the document (given namespace, uri) is alive or not.
+  bool IsDocumentAlive(std::string_view name_space, std::string_view uri,
+                       int64_t current_time_ms) const;
+
+  // Updates the expiration timestamp of the given document only if the document
+  // is not deleted (note: expired document's expiration timestamp can be
+  // updated decreasingly).
+  //
+  // Currently we only allow to decrease the expiration timestamp. This prevents
+  // us from making an expired document alive again.
+  //
+  // The input expiration timestamp is usually the propagated expiration
+  // timestamp computed from the document dependencies (expiration and delete
+  // propagation).
+  //
+  // Returns:
+  //   On success, the final value of the expiration timestamp and a boolean
+  //     flag that indicates if the expiration timestamp is updated. If the
+  //     input expiration timestamp is not smaller than the current expiration
+  //     timestamp stored in DocumentFilterData cache, then it is no-op and the
+  //     current expiration timestamp will be returned.
+  //   NOT_FOUND_ERROR if the document is deleted
+  //   INVALID_ARGUMENT_ERROR if document_id is invalid
+  //   Any FileBackedVector error
+  struct UpdateDocumentExpirationTimestampResult {
+    // The final value of the expiration timestamp.
+    int64_t final_expiration_timestamp_ms;
+
+    // Whether the expiration timestamp was updated or not.
+    bool was_updated;
+  };
+  libtextclassifier3::StatusOr<UpdateDocumentExpirationTimestampResult>
+  UpdateDocumentExpirationTimestamp(DocumentId document_id,
+                                    int64_t expiration_timestamp_ms);
+
+  // Resets the expiration timestamp as the raw expiration timestamp of all
+  // alive documents. See DocumentFilterData for more details.
+  //
+  // Note: this is usually called before recomputing propagated expiration
+  // timestamps when Icing detects the document dependency (expiration and
+  // delete propagation) should be re-evaluated.
+  //
+  // Returns:
+  //   OK on success
+  //   Any FileBackedVector error
+  libtextclassifier3::Status ResetAllAliveExpirationTimestampsToRaw(
+      int64_t current_time_ms);
+
+  // Purges all expired documents given the current time.
+  //
+  // Note: documents that expire before
+  //   current_time_ms + feature_flags_.expired_document_purge_threshold_ms()
+  //   will also be purged.
+  //
+  // Returns:
+  //   A vector of purged expired document metadata on success
+  //   INTERNAL_ERROR on IO error
+  libtextclassifier3::StatusOr<std::vector<DocumentMetadata>>
+  PurgeExpiredDocuments(int64_t current_time_ms);
+
+  // Returns the expiration timestamp (in milliseconds) of the document that
+  // will expire next, relative to current_time_ms. If there are no more expired
+  // documents, returns -1.
+  int64_t GetNextExpiredDocumentTimestampMs(int64_t current_time_ms);
+
   // Gets the SchemaTypeId of a document.
   //
   // Returns:
@@ -461,7 +556,9 @@
   // Returns:
   //   OK on success
   //   INTERNAL on I/O error
-  libtextclassifier3::Status PersistToDisk(PersistType::Code persist_type);
+  libtextclassifier3::Status PersistToDisk(
+      PersistType::Code persist_type,
+      PersistToDiskStatsProto* persist_stats = nullptr);
 
   // Calculates the StorageInfo for the Document Store.
   //
@@ -469,6 +566,10 @@
   // that field will be set to -1.
   DocumentStorageInfoProto GetStorageInfo() const;
 
+  struct UpdateSchemaStoreResult {
+    int deleted_document_count = 0;
+    bool derived_files_changed = false;
+  };
   // Update any derived data off of the SchemaStore with the new SchemaStore.
   // This may include pointers, SchemaTypeIds, etc.
   //
@@ -482,18 +583,24 @@
   // OptimizedUpdateSchemaStore.
   //
   // Returns;
-  //   OK on success
+  //   UpdateSchemaStoreResult on success
   //   INTERNAL_ERROR on IO error
-  libtextclassifier3::Status UpdateSchemaStore(const SchemaStore* schema_store);
+  libtextclassifier3::StatusOr<UpdateSchemaStoreResult> UpdateSchemaStore(
+      const SchemaStore* schema_store);
 
-  // Performs the same funtionality as UpdateSchemaStore, but this can be more
+  // Performs the same functionality as UpdateSchemaStore, but this can be more
   // optimized in terms of less disk reads and less work if we know exactly
   // what's changed between the old and new SchemaStore.
   //
+  // NOTE: This function may delete documents. A document may be invalidated by
+  // the new SchemaStore, such as failing validation or having its schema type
+  // deleted from the schema.
+  //
   // Returns;
-  //   OK on success
+  //   UpdateSchemaStoreResult on success
   //   INTERNAL_ERROR on IO error
-  libtextclassifier3::Status OptimizedUpdateSchemaStore(
+  libtextclassifier3::StatusOr<UpdateSchemaStoreResult>
+  OptimizedUpdateSchemaStore(
       const SchemaStore* schema_store,
       const SchemaStore::SetSchemaResult& set_schema_result);
 
@@ -598,7 +705,9 @@
                          const SchemaStore* schema_store,
                          const FeatureFlags* feature_flags,
                          bool pre_mapping_fbv, bool use_persistent_hash_map,
-                         int32_t compression_level);
+                         int32_t compression_level,
+                         uint32_t compression_threshold_bytes,
+                         int32_t compression_mem_level);
 
   const Filesystem* const filesystem_;
   const std::string base_dir_;
@@ -622,6 +731,10 @@
   bool use_persistent_hash_map_;
 
   const int32_t compression_level_;
+  const uint32_t compression_threshold_bytes_;
+
+  // Level of memory usage for compression.
+  const int32_t compression_mem_level_;
 
   // A log used to store all documents, it serves as a ground truth of doc
   // store. key_mapper_ and document_id_mapper_ can be regenerated from it.
@@ -669,7 +782,7 @@
   std::unique_ptr<KeyMapper<NamespaceId>> namespace_mapper_;
 
   // Maps a corpus, i.e. a (namespace, schema type) pair, to a densely-assigned
-  // unique id. A coprus is assigned an
+  // unique id. A corpus is assigned an
   // id when the first document belonging to that corpus is added to the
   // DocumentStore. Corpus ids may be removed from the mapper during compaction.
   std::unique_ptr<
@@ -778,7 +891,7 @@
   bool HeaderExists();
 
   libtextclassifier3::StatusOr<PutResult> InternalPut(
-      DocumentProto&& document,
+      const DocumentWrapper& document_wrapper,
       PutDocumentStatsProto* put_document_stats = nullptr);
 
   // Helper function to do batch deletes. Documents with the given
@@ -811,32 +924,16 @@
   libtextclassifier3::StatusOr<CorpusAssociatedScoreData>
   GetCorpusAssociatedScoreDataToUpdate(CorpusId corpus_id) const;
 
-  // Check if a document exists. Existence means it hasn't been deleted and it
-  // hasn't expired yet.
-  //
-  // Returns:
-  //   OK if the document exists
-  //   INVALID_ARGUMENT if document_id is less than 0 or greater than the
-  //                    maximum value
-  //   NOT_FOUND if the document doesn't exist (i.e. deleted or expired)
-  //   INTERNAL_ERROR on IO error
-  libtextclassifier3::Status DoesDocumentExistWithStatus(
-      DocumentId document_id) const;
-
-  // Checks if a document has been deleted
+  // Checks if a document has been deleted.
   //
   // This is for internal-use only because we assume that the document_id is
-  // already valid. If you're unsure if the document_id is valid, use
-  // DoesDocumentExist(document_id) instead, which will perform those additional
-  // checks.
+  // already valid.
   bool IsDeleted(DocumentId document_id) const;
 
   // Checks if a document has expired.
   //
   // This is for internal-use only because we assume that the document_id is
-  // already valid. If you're unsure if the document_id is valid, use
-  // DoesDocumentExist(document_id) instead, which will perform those additional
-  // checks.
+  // already valid.
 
   // Returns:
   //   True:DocumentFilterData  if the given document isn't expired.

diff --git a/icing/store/document-store_benchmark.cc b/icing/store/document-store_benchmark.cc
index 3782da0..907c01b 100644
--- a/icing/store/document-store_benchmark.cc
+++ b/icing/store/document-store_benchmark.cc

@@ -14,26 +14,21 @@
 
 #include <unistd.h>
 
-#include <fstream>
-#include <iostream>
 #include <memory>
-#include <ostream>
 #include <random>
-#include <sstream>
-#include <stdexcept>
 #include <string>
 #include <string_view>
-#include <unordered_set>
-#include <vector>
+#include <utility>
 
+#include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "testing/base/public/benchmark.h"
-#include "gmock/gmock.h"
-#include "gtest/gtest.h"
 #include "icing/document-builder.h"
 #include "icing/feature-flags.h"
 #include "icing/file/filesystem.h"
 #include "icing/file/portable-file-backed-proto-log.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/proto/document.pb.h"
+#include "icing/proto/document_wrapper.pb.h"
 #include "icing/proto/persist.pb.h"
 #include "icing/proto/schema.pb.h"
 #include "icing/proto/term.pb.h"
@@ -44,6 +39,8 @@
 #include "icing/testing/test-feature-flags.h"
 #include "icing/testing/tmp-directory.h"
 #include "icing/util/clock.h"
+#include "icing/util/document-util.h"
+#include "icing/util/logging.h"
 
 // Run on a Linux workstation:
 //    $ blaze build -c opt --dynamic_mode=off --copt=-gmlt
@@ -84,14 +81,15 @@
   std::string dir_;
 };
 
-DocumentProto CreateDocument(const std::string namespace_,
-                             const std::string uri) {
-  return DocumentBuilder()
-      .SetKey(namespace_, uri)
-      .SetSchema("email")
-      .AddStringProperty("subject", "subject foo")
-      .AddStringProperty("body", "body bar")
-      .Build();
+DocumentWrapper CreateDocument(const std::string namespace_,
+                               const std::string uri) {
+  DocumentProto document = DocumentBuilder()
+                               .SetKey(namespace_, uri)
+                               .SetSchema("email")
+                               .AddStringProperty("subject", "subject foo")
+                               .AddStringProperty("body", "body bar")
+                               .Build();
+  return document_util::CreateDocumentWrapper(std::move(document));
 }
 
 SchemaProto CreateSchema() {
@@ -138,6 +136,9 @@
       /*force_recovery_and_revalidate_documents=*/false,
       /*pre_mapping_fbv=*/false, /*use_persistent_hash_map=*/true,
       PortableFileBackedProtoLog<DocumentWrapper>::kDefaultCompressionLevel,
+      PortableFileBackedProtoLog<
+          DocumentWrapper>::kDefaultCompressionThresholdBytes,
+      protobuf_ports::kDefaultMemLevel,
       /*initialize_stats=*/nullptr);
 }
 
@@ -204,7 +205,7 @@
   std::unique_ptr<DocumentStore> document_store =
       std::move(create_result.document_store);
 
-  DocumentProto document = CreateDocument("namespace", "uri");
+  DocumentWrapper document = CreateDocument("namespace", "uri");
 
   for (auto s : state) {
     // It's ok that this is the same document over and over. We'll create a new
@@ -262,7 +263,7 @@
   std::unique_ptr<DocumentStore> document_store =
       std::move(create_result.document_store);
 
-  DocumentProto document = CreateDocument("namespace", "uri");
+  DocumentWrapper document = CreateDocument("namespace", "uri");
 
   for (auto s : state) {
     state.PauseTiming();
@@ -298,7 +299,7 @@
     std::unique_ptr<DocumentStore> document_store =
         std::move(create_result.document_store);
 
-    DocumentProto document = CreateDocument("namespace", "uri");
+    DocumentWrapper document = CreateDocument("namespace", "uri");
     ICING_ASSERT_OK(document_store->Put(document));
     ICING_ASSERT_OK(document_store->PersistToDisk(PersistType::FULL));
   }
@@ -336,7 +337,7 @@
   std::unique_ptr<DocumentStore> document_store =
       std::move(create_result.document_store);
 
-  DocumentProto document = CreateDocument("namespace", "uri");
+  DocumentWrapper document = CreateDocument("namespace", "uri");
   ICING_ASSERT_OK(document_store->Put(document));
   ICING_ASSERT_OK(document_store->PersistToDisk(PersistType::LITE));
 

diff --git a/icing/store/document-store_test.cc b/icing/store/document-store_test.cc
index a9cc0a5..ff531cf 100644
--- a/icing/store/document-store_test.cc
+++ b/icing/store/document-store_test.cc

@@ -24,7 +24,6 @@
 
 #include "icing/text_classifier/lib3/utils/base/status.h"
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
-#include "icing/text_classifier/lib3/utils/hash/farmhash.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 #include "icing/absl_ports/str_cat.h"
@@ -36,6 +35,7 @@
 #include "icing/file/mock-filesystem.h"
 #include "icing/file/portable-file-backed-proto-log.h"
 #include "icing/portable/equals-proto.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/portable/platform.h"
 #include "icing/proto/debug.pb.h"
 #include "icing/proto/document.pb.h"
@@ -66,6 +66,7 @@
 #include "icing/util/clock.h"
 #include "icing/util/crc32.h"
 #include "icing/util/data-loss.h"
+#include "icing/util/document-util.h"
 #include "icing/util/icu-data-file-helper.h"
 #include "icing/util/logging.h"
 #include "icing/util/scorable_property_set.h"
@@ -88,10 +89,20 @@
 using ::testing::IsTrue;
 using ::testing::Ne;
 using ::testing::Not;
+using ::testing::Optional;
 using ::testing::Pointee;
 using ::testing::Return;
 using ::testing::UnorderedElementsAre;
 
+using ResultSpecProto::ResultGroupingType::
+    ResultSpecProto_ResultGroupingType_NAMESPACE;
+using ResultSpecProto::ResultGroupingType::
+    ResultSpecProto_ResultGroupingType_NAMESPACE_AND_SCHEMA_TYPE;
+using ResultSpecProto::ResultGroupingType::
+    ResultSpecProto_ResultGroupingType_NONE;
+using ResultSpecProto::ResultGroupingType::
+    ResultSpecProto_ResultGroupingType_SCHEMA_TYPE;
+
 const NamespaceStorageInfoProto& GetNamespaceStorageInfo(
     const DocumentStorageInfoProto& storage_info,
     const std::string& name_space) {
@@ -284,6 +295,9 @@
         /*force_recovery_and_revalidate_documents=*/false,
         GetParam().pre_mapping_fbv, GetParam().use_persistent_hash_map,
         PortableFileBackedProtoLog<DocumentWrapper>::kDefaultCompressionLevel,
+        PortableFileBackedProtoLog<
+            DocumentWrapper>::kDefaultCompressionThresholdBytes,
+        protobuf_ports::kDefaultMemLevel,
         /*initialize_stats=*/nullptr);
   }
 
@@ -312,6 +326,12 @@
   const int64_t document2_expiration_timestamp_ = 3;  // creation + ttl
 };
 
+DocumentWrapper CreateDocumentWrapper(DocumentProto document,
+                                      int32_t num_string_tokens) {
+  document.mutable_internal_fields()->set_length_in_tokens(num_string_tokens);
+  return document_util::CreateDocumentWrapper(std::move(document));
+}
+
 TEST_P(DocumentStoreTest, CreationWithNullPointerShouldFail) {
   EXPECT_THAT(CreateDocumentStore(/*filesystem=*/nullptr, document_store_dir_,
                                   &fake_clock_, schema_store_.get()),
@@ -335,6 +355,33 @@
               StatusIs(libtextclassifier3::StatusCode::INTERNAL));
 }
 
+TEST_P(DocumentStoreTest, PutResult) {
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
+                          schema_store_.get()));
+  std::unique_ptr<DocumentStore> doc_store =
+      std::move(create_result.document_store);
+
+  // Put test_document1_ which never expires.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      doc_store->Put(document_util::CreateDocumentWrapper(test_document1_)));
+  EXPECT_THAT(put_result1.old_document_id, Eq(kInvalidDocumentId));
+  EXPECT_FALSE(put_result1.was_replacement);
+  EXPECT_THAT(put_result1.expiration_timestamp_ms,
+              Eq(document1_expiration_timestamp_));
+
+  // Put test_document2_.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      doc_store->Put(document_util::CreateDocumentWrapper(test_document2_)));
+  EXPECT_THAT(put_result2.old_document_id, Eq(kInvalidDocumentId));
+  EXPECT_FALSE(put_result2.was_replacement);
+  EXPECT_THAT(put_result2.expiration_timestamp_ms,
+              Eq(document2_expiration_timestamp_));
+}
+
 TEST_P(DocumentStoreTest, PutAndGetInSameNamespaceOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::CreateResult create_result,
@@ -344,15 +391,17 @@
       std::move(create_result.document_store);
 
   // Both documents have namespace of "icing"
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             doc_store->Put(test_document1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      doc_store->Put(document_util::CreateDocumentWrapper(test_document1_)));
   EXPECT_THAT(put_result1.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result1.was_replacement());
+  EXPECT_FALSE(put_result1.was_replacement);
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             doc_store->Put(test_document2_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      doc_store->Put(document_util::CreateDocumentWrapper(test_document2_)));
   EXPECT_THAT(put_result2.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result2.was_replacement());
+  EXPECT_FALSE(put_result2.was_replacement);
   DocumentId document_id2 = put_result2.new_document_id;
 
   EXPECT_THAT(doc_store->Get(document_id1),
@@ -381,15 +430,17 @@
                                    .SetCreationTimestampMs(0)
                                    .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             doc_store->Put(foo_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      doc_store->Put(document_util::CreateDocumentWrapper(foo_document)));
   EXPECT_THAT(put_result1.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result1.was_replacement());
+  EXPECT_FALSE(put_result1.was_replacement);
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             doc_store->Put(DocumentProto(bar_document)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      doc_store->Put(document_util::CreateDocumentWrapper(bar_document)));
   EXPECT_THAT(put_result2.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result2.was_replacement());
+  EXPECT_FALSE(put_result2.was_replacement);
   DocumentId document_id2 = put_result2.new_document_id;
 
   EXPECT_THAT(doc_store->Get(document_id1),
@@ -412,16 +463,18 @@
   DocumentProto document1 = DocumentProto(test_document1_);
   DocumentProto document2 = DocumentProto(test_document1_);
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             doc_store->Put(document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      doc_store->Put(document_util::CreateDocumentWrapper(document1)));
   EXPECT_THAT(put_result1.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result1.was_replacement());
+  EXPECT_FALSE(put_result1.was_replacement);
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             doc_store->Put(document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      doc_store->Put(document_util::CreateDocumentWrapper(document2)));
   DocumentId document_id2 = put_result2.new_document_id;
   EXPECT_THAT(put_result2.old_document_id, Eq(document_id1));
-  EXPECT_TRUE(put_result2.was_replacement());
+  EXPECT_TRUE(put_result2.was_replacement);
   EXPECT_THAT(document_id1, Not(document_id2));
   // document2 overrides document1, so document_id1 becomes invalid
   EXPECT_THAT(doc_store->Get(document_id1),
@@ -432,11 +485,12 @@
   // Makes sure that old doc ids are not getting reused.
   DocumentProto document3 = DocumentProto(test_document1_);
   document3.set_uri("another/uri/1");
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             doc_store->Put(document3));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      doc_store->Put(document_util::CreateDocumentWrapper(document3)));
   EXPECT_THAT(put_result3.old_document_id, Eq(kInvalidDocumentId));
   EXPECT_THAT(put_result3.new_document_id, Ne(document_id1));
-  EXPECT_FALSE(put_result3.was_replacement());
+  EXPECT_FALSE(put_result3.was_replacement);
 }
 
 TEST_P(DocumentStoreTest, IsDocumentExistingWithoutStatus) {
@@ -447,15 +501,17 @@
   std::unique_ptr<DocumentStore> doc_store =
       std::move(create_result.document_store);
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             doc_store->Put(test_document1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      doc_store->Put(document_util::CreateDocumentWrapper(test_document1_)));
   EXPECT_THAT(put_result1.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result1.was_replacement());
+  EXPECT_FALSE(put_result1.was_replacement);
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             doc_store->Put(test_document2_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      doc_store->Put(document_util::CreateDocumentWrapper(test_document2_)));
   EXPECT_THAT(put_result2.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result2.was_replacement());
+  EXPECT_FALSE(put_result2.was_replacement);
   DocumentId document_id2 = put_result2.new_document_id;
 
   EXPECT_TRUE(doc_store->GetAliveDocumentFilterData(
@@ -489,7 +545,8 @@
   std::unique_ptr<DocumentStore> document_store =
       std::move(create_result.document_store);
 
-  ICING_EXPECT_OK(document_store->Put(DocumentProto(test_document1_)));
+  ICING_EXPECT_OK(document_store->Put(
+      document_util::CreateDocumentWrapper(test_document1_)));
   EXPECT_THAT(
       document_store->Get(test_document1_.namespace_(), test_document1_.uri()),
       IsOkAndHolds(EqualsProto(test_document1_)));
@@ -517,7 +574,8 @@
   std::unique_ptr<DocumentStore> document_store =
       std::move(create_result.document_store);
 
-  ICING_EXPECT_OK(document_store->Put(document));
+  ICING_EXPECT_OK(
+      document_store->Put(document_util::CreateDocumentWrapper(document)));
   EXPECT_THAT(document_store->Get("namespace", "uri"),
               IsOkAndHolds(EqualsProto(document)));
 
@@ -545,10 +603,11 @@
   std::unique_ptr<DocumentStore> doc_store =
       std::move(create_result.document_store);
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             doc_store->Put(test_document1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store->Put(document_util::CreateDocumentWrapper(test_document1_)));
   EXPECT_THAT(put_result.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result.was_replacement());
+  EXPECT_FALSE(put_result.was_replacement);
   DocumentId document_id = put_result.new_document_id;
 
   DocumentId invalid_document_id_negative = -1;
@@ -630,7 +689,8 @@
   std::unique_ptr<DocumentStore> document_store =
       std::move(create_result.document_store);
 
-  ICING_EXPECT_OK(document_store->Put(test_document1_));
+  ICING_EXPECT_OK(document_store->Put(
+      document_util::CreateDocumentWrapper(test_document1_)));
 
   // First time is OK
   ICING_EXPECT_OK(document_store->Delete(
@@ -644,6 +704,109 @@
               StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
 }
 
+TEST_P(DocumentStoreTest, ForceDelete) {
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
+                          schema_store_.get()));
+  std::unique_ptr<DocumentStore> document_store =
+      std::move(create_result.document_store);
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store->Put(
+          document_util::CreateDocumentWrapper(test_document1_)));
+  DocumentId document_id = put_result.new_document_id;
+
+  // Sanity check that the document is alive.
+  ASSERT_TRUE(document_store->GetAliveDocumentFilterData(
+      document_id,
+      /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+  ASSERT_TRUE(document_store->GetNonDeletedDocumentFilterData(document_id));
+
+  EXPECT_THAT(document_store->ForceDelete(document_id),
+              IsOkAndHolds(EqualsDocumentMetadata(
+                  test_document1_.schema(), test_document1_.namespace_(),
+                  test_document1_.uri(), document_id)));
+  EXPECT_FALSE(document_store->GetAliveDocumentFilterData(
+      document_id,
+      /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_FALSE(document_store->GetNonDeletedDocumentFilterData(document_id));
+}
+
+TEST_P(DocumentStoreTest, ForceDeleteAlreadyDeletedDocumentReturnsNotFound) {
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
+                          schema_store_.get()));
+  std::unique_ptr<DocumentStore> document_store =
+      std::move(create_result.document_store);
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store->Put(
+          document_util::CreateDocumentWrapper(test_document1_)));
+  DocumentId document_id = put_result.new_document_id;
+
+  // Sanity check that the document is alive.
+  ASSERT_TRUE(document_store->GetAliveDocumentFilterData(
+      document_id,
+      /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+  ASSERT_TRUE(document_store->GetNonDeletedDocumentFilterData(document_id));
+
+  // First time is OK
+  ICING_ASSERT_OK(document_store->ForceDelete(document_id));
+
+  // Deleting it again should get NOT_FOUND.
+  EXPECT_THAT(document_store->ForceDelete(document_id),
+              StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
+}
+
+TEST_P(DocumentStoreTest, ForceDeleteExpiredDocumentOk) {
+  fake_clock_.SetSystemTimeMilliseconds(0);
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
+                          schema_store_.get()));
+  std::unique_ptr<DocumentStore> document_store =
+      std::move(create_result.document_store);
+
+  DocumentProto document = DocumentBuilder()
+                               .SetKey("namespace", "uri")
+                               .SetSchema("email")
+                               .SetCreationTimestampMs(0)
+                               .SetTtlMs(1000)
+                               .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      document_store->Put(document_util::CreateDocumentWrapper(document)));
+  DocumentId document_id = put_result.new_document_id;
+
+  // Sanity check that the document is alive.
+  ASSERT_TRUE(document_store->GetAliveDocumentFilterData(
+      document_id,
+      /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+  ASSERT_TRUE(document_store->GetNonDeletedDocumentFilterData(document_id));
+
+  // Adjust the clock to make the document expired.
+  fake_clock_.SetSystemTimeMilliseconds(2000);
+  ASSERT_FALSE(document_store->GetAliveDocumentFilterData(
+      document_id,
+      /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+  ASSERT_TRUE(document_store->GetNonDeletedDocumentFilterData(document_id));
+
+  // Deleting an expired document is OK.
+  EXPECT_THAT(document_store->ForceDelete(document_id),
+              IsOkAndHolds(EqualsDocumentMetadata(
+                  document.schema(), document.namespace_(), document.uri(),
+                  document_id)));
+  EXPECT_FALSE(document_store->GetAliveDocumentFilterData(
+      document_id,
+      /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_FALSE(document_store->GetNonDeletedDocumentFilterData(document_id));
+}
+
 TEST_P(DocumentStoreTest, DeleteByNamespaceOk) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::CreateResult create_result,
@@ -655,22 +818,26 @@
   DocumentProto document1 = test_document1_;
   document1.set_namespace_("namespace.1");
   document1.set_uri("uri1");
-  ICING_ASSERT_OK(doc_store->Put(document1));
+  ICING_ASSERT_OK(
+      doc_store->Put(document_util::CreateDocumentWrapper(document1)));
 
   DocumentProto document2 = test_document1_;
   document2.set_namespace_("namespace.2");
   document2.set_uri("uri1");
-  ICING_ASSERT_OK(doc_store->Put(document2));
+  ICING_ASSERT_OK(
+      doc_store->Put(document_util::CreateDocumentWrapper(document2)));
 
   DocumentProto document3 = test_document1_;
   document3.set_namespace_("namespace.3");
   document3.set_uri("uri1");
-  ICING_ASSERT_OK(doc_store->Put(document3));
+  ICING_ASSERT_OK(
+      doc_store->Put(document_util::CreateDocumentWrapper(document3)));
 
   DocumentProto document4 = test_document1_;
   document4.set_namespace_("namespace.1");
   document4.set_uri("uri2");
-  ICING_ASSERT_OK(doc_store->Put(document4));
+  ICING_ASSERT_OK(
+      doc_store->Put(document_util::CreateDocumentWrapper(document4)));
 
   // DELETE namespace.1. document1 and document 4 should be deleted. document2
   // and document3 should still be retrievable.
@@ -721,7 +888,8 @@
   std::unique_ptr<DocumentStore> document_store =
       std::move(create_result.document_store);
 
-  ICING_EXPECT_OK(document_store->Put(test_document1_));
+  ICING_EXPECT_OK(document_store->Put(
+      document_util::CreateDocumentWrapper(test_document1_)));
   ICING_EXPECT_OK(document_store->Delete(
       test_document1_.namespace_(), test_document1_.uri(),
       fake_clock_.GetSystemTimeMilliseconds()));
@@ -760,10 +928,14 @@
     std::unique_ptr<DocumentStore> doc_store =
         std::move(create_result.document_store);
 
-    ICING_ASSERT_OK(doc_store->Put(document1));
-    ICING_ASSERT_OK(doc_store->Put(document2));
-    ICING_ASSERT_OK(doc_store->Put(document3));
-    ICING_ASSERT_OK(doc_store->Put(document4));
+    ICING_ASSERT_OK(
+        doc_store->Put(document_util::CreateDocumentWrapper(document1)));
+    ICING_ASSERT_OK(
+        doc_store->Put(document_util::CreateDocumentWrapper(document2)));
+    ICING_ASSERT_OK(
+        doc_store->Put(document_util::CreateDocumentWrapper(document3)));
+    ICING_ASSERT_OK(
+        doc_store->Put(document_util::CreateDocumentWrapper(document4)));
 
     // DELETE namespace.1. document1 and document 4 should be deleted. document2
     // and document3 should still be retrievable.
@@ -835,10 +1007,12 @@
                                        .SetSchema("email")
                                        .SetCreationTimestampMs(1)
                                        .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult email_put_result1,
-                             document_store->Put(email_document_1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult email_put_result1,
+      document_store->Put(
+          document_util::CreateDocumentWrapper(email_document_1)));
   EXPECT_THAT(email_put_result1.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(email_put_result1.was_replacement());
+  EXPECT_FALSE(email_put_result1.was_replacement);
   DocumentId email_1_document_id = email_put_result1.new_document_id;
 
   DocumentProto email_document_2 = DocumentBuilder()
@@ -846,10 +1020,12 @@
                                        .SetSchema("email")
                                        .SetCreationTimestampMs(1)
                                        .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult email_put_result2,
-                             document_store->Put(email_document_2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult email_put_result2,
+      document_store->Put(
+          document_util::CreateDocumentWrapper(email_document_2)));
   EXPECT_THAT(email_put_result2.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(email_put_result2.was_replacement());
+  EXPECT_FALSE(email_put_result2.was_replacement);
   DocumentId email_2_document_id = email_put_result2.new_document_id;
 
   DocumentProto message_document = DocumentBuilder()
@@ -857,10 +1033,12 @@
                                        .SetSchema("message")
                                        .SetCreationTimestampMs(1)
                                        .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult message_put_result,
-                             document_store->Put(message_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult message_put_result,
+      document_store->Put(
+          document_util::CreateDocumentWrapper(message_document)));
   EXPECT_THAT(message_put_result.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(message_put_result.was_replacement());
+  EXPECT_FALSE(message_put_result.was_replacement);
   DocumentId message_document_id = message_put_result.new_document_id;
 
   DocumentProto person_document = DocumentBuilder()
@@ -868,10 +1046,12 @@
                                       .SetSchema("person")
                                       .SetCreationTimestampMs(1)
                                       .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult person_put_result,
-                             document_store->Put(person_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult person_put_result,
+      document_store->Put(
+          document_util::CreateDocumentWrapper(person_document)));
   EXPECT_THAT(person_put_result.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(person_put_result.was_replacement());
+  EXPECT_FALSE(person_put_result.was_replacement);
   DocumentId person_document_id = person_put_result.new_document_id;
 
   // Delete the "email" type and ensure that it works across both
@@ -937,7 +1117,8 @@
   std::unique_ptr<DocumentStore> document_store =
       std::move(create_result.document_store);
 
-  ICING_EXPECT_OK(document_store->Put(test_document1_));
+  ICING_EXPECT_OK(document_store->Put(
+      document_util::CreateDocumentWrapper(test_document1_)));
   ICING_EXPECT_OK(document_store->Delete(
       test_document1_.namespace_(), test_document1_.uri(),
       fake_clock_.GetSystemTimeMilliseconds()));
@@ -988,15 +1169,19 @@
     std::unique_ptr<DocumentStore> document_store =
         std::move(create_result.document_store);
 
-    ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult email_put_result,
-                               document_store->Put(email_document));
+    ICING_ASSERT_OK_AND_ASSIGN(
+        DocumentStore::PutResult email_put_result,
+        document_store->Put(
+            document_util::CreateDocumentWrapper(email_document)));
     EXPECT_THAT(email_put_result.old_document_id, Eq(kInvalidDocumentId));
-    EXPECT_FALSE(email_put_result.was_replacement());
+    EXPECT_FALSE(email_put_result.was_replacement);
     email_document_id = email_put_result.new_document_id;
-    ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult message_put_result,
-                               document_store->Put(message_document));
+    ICING_ASSERT_OK_AND_ASSIGN(
+        DocumentStore::PutResult message_put_result,
+        document_store->Put(
+            document_util::CreateDocumentWrapper(message_document)));
     EXPECT_THAT(message_put_result.old_document_id, Eq(kInvalidDocumentId));
-    EXPECT_FALSE(message_put_result.was_replacement());
+    EXPECT_FALSE(message_put_result.was_replacement);
     message_document_id = message_put_result.new_document_id;
     // Delete "email". "message" documents should still be retrievable.
     DocumentStore::DeleteByGroupResult group_result =
@@ -1039,11 +1224,13 @@
                           schema_store_.get()));
   std::unique_ptr<DocumentStore> doc_store =
       std::move(create_result.document_store);
-  ICING_EXPECT_OK(doc_store->Put(test_document1_));
+  ICING_EXPECT_OK(
+      doc_store->Put(document_util::CreateDocumentWrapper(test_document1_)));
   ICING_EXPECT_OK(doc_store->Delete(test_document1_.namespace_(),
                                     test_document1_.uri(),
                                     fake_clock_.GetSystemTimeMilliseconds()));
-  ICING_EXPECT_OK(doc_store->Put(test_document1_));
+  ICING_EXPECT_OK(
+      doc_store->Put(document_util::CreateDocumentWrapper(test_document1_)));
 }
 
 TEST_P(DocumentStoreTest, DeletedSchemaTypeFromSchemaStoreRecoversOk) {
@@ -1087,15 +1274,19 @@
     std::unique_ptr<DocumentStore> document_store =
         std::move(create_result.document_store);
 
-    ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult email_put_result,
-                               document_store->Put(email_document));
+    ICING_ASSERT_OK_AND_ASSIGN(
+        DocumentStore::PutResult email_put_result,
+        document_store->Put(
+            document_util::CreateDocumentWrapper(email_document)));
     EXPECT_THAT(email_put_result.old_document_id, Eq(kInvalidDocumentId));
-    EXPECT_FALSE(email_put_result.was_replacement());
+    EXPECT_FALSE(email_put_result.was_replacement);
     email_document_id = email_put_result.new_document_id;
-    ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult message_put_result,
-                               document_store->Put(message_document));
+    ICING_ASSERT_OK_AND_ASSIGN(
+        DocumentStore::PutResult message_put_result,
+        document_store->Put(
+            document_util::CreateDocumentWrapper(message_document)));
     EXPECT_THAT(message_put_result.old_document_id, Eq(kInvalidDocumentId));
-    EXPECT_FALSE(message_put_result.was_replacement());
+    EXPECT_FALSE(message_put_result.was_replacement);
     message_document_id = message_put_result.new_document_id;
 
     // Delete "email". "message" documents should still be retrievable.
@@ -1145,6 +1336,412 @@
               IsOkAndHolds(EqualsProto(message_document)));
 }
 
+TEST_P(DocumentStoreTest, PurgeExpiredDocuments) {
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
+                          schema_store_.get()));
+  std::unique_ptr<DocumentStore> document_store =
+      std::move(create_result.document_store);
+
+  DocumentProto doc1 = DocumentBuilder()
+                           .SetKey("namespace", "uri1")
+                           .SetSchema("email")
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(2000)
+                           .Build();
+  DocumentProto doc2 = DocumentBuilder()
+                           .SetKey("namespace", "uri2")
+                           .SetSchema("email")
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(1200)
+                           .Build();
+  DocumentProto doc3 = DocumentBuilder()
+                           .SetKey("namespace", "uri3")
+                           .SetSchema("email")
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(500)
+                           .Build();
+  DocumentProto doc4 = DocumentBuilder()
+                           .SetKey("namespace", "uri4")
+                           .SetSchema("email")
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(1800)
+                           .Build();
+  DocumentProto doc5 = DocumentBuilder()
+                           .SetKey("namespace", "uri5")
+                           .SetSchema("email")
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(1000)
+                           .Build();
+  // Put doc1 to doc5.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store->Put(document_util::CreateDocumentWrapper(doc1)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store->Put(document_util::CreateDocumentWrapper(doc2)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      document_store->Put(document_util::CreateDocumentWrapper(doc3)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result4,
+      document_store->Put(document_util::CreateDocumentWrapper(doc4)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result5,
+      document_store->Put(document_util::CreateDocumentWrapper(doc5)));
+  DocumentId doc_id1 = put_result1.new_document_id;
+  DocumentId doc_id2 = put_result2.new_document_id;
+  DocumentId doc_id3 = put_result3.new_document_id;
+  DocumentId doc_id4 = put_result4.new_document_id;
+  DocumentId doc_id5 = put_result5.new_document_id;
+
+  // Purge expired documents at t = 1200. doc2, doc3, doc5 should be purged.
+  EXPECT_THAT(document_store->PurgeExpiredDocuments(/*current_time_ms=*/1200),
+              IsOkAndHolds(ElementsAre(
+                  EqualsDocumentMetadata(doc2.schema(), doc2.namespace_(),
+                                         doc2.uri(), doc_id2),
+                  EqualsDocumentMetadata(doc3.schema(), doc3.namespace_(),
+                                         doc3.uri(), doc_id3),
+                  EqualsDocumentMetadata(doc5.schema(), doc5.namespace_(),
+                                         doc5.uri(), doc_id5))));
+  EXPECT_TRUE(document_store->GetNonDeletedDocumentFilterData(doc_id1));
+  EXPECT_FALSE(document_store->GetNonDeletedDocumentFilterData(doc_id2));
+  EXPECT_FALSE(document_store->GetNonDeletedDocumentFilterData(doc_id3));
+  EXPECT_TRUE(document_store->GetNonDeletedDocumentFilterData(doc_id4));
+  EXPECT_FALSE(document_store->GetNonDeletedDocumentFilterData(doc_id5));
+}
+
+TEST_P(DocumentStoreTest, PurgeExpiredDocuments_shouldSkipDeletedDocuments) {
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
+                          schema_store_.get()));
+  std::unique_ptr<DocumentStore> document_store =
+      std::move(create_result.document_store);
+
+  DocumentProto doc1 = DocumentBuilder()
+                           .SetKey("namespace", "uri1")
+                           .SetSchema("email")
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(2000)
+                           .Build();
+  DocumentProto doc2 = DocumentBuilder()
+                           .SetKey("namespace", "uri2")
+                           .SetSchema("email")
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(1200)
+                           .Build();
+  DocumentProto doc3 = DocumentBuilder()
+                           .SetKey("namespace", "uri3")
+                           .SetSchema("email")
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(500)
+                           .Build();
+  DocumentProto doc4 = DocumentBuilder()
+                           .SetKey("namespace", "uri4")
+                           .SetSchema("email")
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(1800)
+                           .Build();
+  DocumentProto doc5 = DocumentBuilder()
+                           .SetKey("namespace", "uri5")
+                           .SetSchema("email")
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(1000)
+                           .Build();
+  // Put doc1 to doc5.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store->Put(document_util::CreateDocumentWrapper(doc1)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store->Put(document_util::CreateDocumentWrapper(doc2)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      document_store->Put(document_util::CreateDocumentWrapper(doc3)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result4,
+      document_store->Put(document_util::CreateDocumentWrapper(doc4)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result5,
+      document_store->Put(document_util::CreateDocumentWrapper(doc5)));
+  DocumentId doc_id1 = put_result1.new_document_id;
+  DocumentId doc_id2 = put_result2.new_document_id;
+  DocumentId doc_id3 = put_result3.new_document_id;
+  DocumentId doc_id4 = put_result4.new_document_id;
+  DocumentId doc_id5 = put_result5.new_document_id;
+
+  // Delete doc3.
+  ICING_ASSERT_OK(document_store->Delete(doc3.namespace_(), doc3.uri(),
+                                         /*current_time_ms=*/0));
+
+  // Purge expired documents at t = 1200. doc2, doc5 should be purged. Since
+  // doc3 was already deleted, it should be skipped.
+  EXPECT_THAT(document_store->PurgeExpiredDocuments(/*current_time_ms=*/1200),
+              IsOkAndHolds(ElementsAre(
+                  EqualsDocumentMetadata(doc2.schema(), doc2.namespace_(),
+                                         doc2.uri(), doc_id2),
+                  EqualsDocumentMetadata(doc5.schema(), doc5.namespace_(),
+                                         doc5.uri(), doc_id5))));
+  EXPECT_TRUE(document_store->GetNonDeletedDocumentFilterData(doc_id1));
+  EXPECT_FALSE(document_store->GetNonDeletedDocumentFilterData(doc_id2));
+  EXPECT_FALSE(document_store->GetNonDeletedDocumentFilterData(doc_id3));
+  EXPECT_TRUE(document_store->GetNonDeletedDocumentFilterData(doc_id4));
+  EXPECT_FALSE(document_store->GetNonDeletedDocumentFilterData(doc_id5));
+}
+
+TEST_P(DocumentStoreTest,
+       PurgeExpiredDocuments_shouldPurgeDocumentsThatExpireWithinThreshold) {
+  auto custom_feature_flags = std::make_unique<FeatureFlags>(
+      /*enable_circular_schema_definitions=*/true,
+      /*enable_scorable_properties=*/true,
+      /*enable_embedding_quantization=*/true,
+      /*enable_repeated_field_joins=*/true,
+      /*enable_embedding_backup_generation=*/true,
+      /*enable_schema_database=*/true,
+      /*release_backup_schema_file_if_overlay_present=*/true,
+      /*enable_strict_page_byte_size_limit=*/true,
+      /*enable_smaller_decompression_buffer_size=*/true,
+      /*enable_eigen_embedding_scoring=*/true,
+      /*enable_passing_filter_to_children=*/true,
+      /*enable_proto_log_new_header_format=*/true,
+      /*enable_embedding_iterator_v2=*/true,
+      /*enable_reusable_decompression_buffer=*/true,
+      /*enable_schema_type_id_optimization=*/true,
+      /*enable_optimize_improvements=*/true,
+      /*expired_document_purge_threshold_ms=*/1000,  // 1 second
+      /*enable_non_existent_qualified_id_join=*/true,
+      /*enable_skip_set_schema_type_equality_check=*/true,
+      /*enable_embed_query_optimization=*/true,
+      /*enable_schema_definition_deduping=*/true);
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      DocumentStore::Create(
+          &filesystem_, document_store_dir_, &fake_clock_, schema_store_.get(),
+          custom_feature_flags.get(),
+          /*force_recovery_and_revalidate_documents=*/false,
+          GetParam().pre_mapping_fbv, GetParam().use_persistent_hash_map,
+          PortableFileBackedProtoLog<DocumentWrapper>::kDefaultCompressionLevel,
+          PortableFileBackedProtoLog<
+              DocumentWrapper>::kDefaultCompressionThresholdBytes,
+          protobuf_ports::kDefaultMemLevel,
+          /*initialize_stats=*/nullptr));
+
+  std::unique_ptr<DocumentStore> document_store =
+      std::move(create_result.document_store);
+
+  DocumentProto doc1 = DocumentBuilder()
+                           .SetKey("namespace", "uri1")
+                           .SetSchema("email")
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(2000)
+                           .Build();
+  DocumentProto doc2 = DocumentBuilder()
+                           .SetKey("namespace", "uri2")
+                           .SetSchema("email")
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(3000)
+                           .Build();
+  DocumentProto doc3 = DocumentBuilder()
+                           .SetKey("namespace", "uri3")
+                           .SetSchema("email")
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(3001)
+                           .Build();
+  DocumentProto doc4 = DocumentBuilder()
+                           .SetKey("namespace", "uri4")
+                           .SetSchema("email")
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(1200)
+                           .Build();
+  // Put doc1 to doc4.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store->Put(document_util::CreateDocumentWrapper(doc1)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store->Put(document_util::CreateDocumentWrapper(doc2)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      document_store->Put(document_util::CreateDocumentWrapper(doc3)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result4,
+      document_store->Put(document_util::CreateDocumentWrapper(doc4)));
+  DocumentId doc_id1 = put_result1.new_document_id;
+  DocumentId doc_id2 = put_result2.new_document_id;
+  DocumentId doc_id3 = put_result3.new_document_id;
+  DocumentId doc_id4 = put_result4.new_document_id;
+
+  // Purge expired documents at t = 2000. doc1, doc2, doc4 should be purged.
+  // Doc3 should still be alive.
+  EXPECT_THAT(document_store->PurgeExpiredDocuments(/*current_time_ms=*/2000),
+              IsOkAndHolds(ElementsAre(
+                  EqualsDocumentMetadata(doc1.schema(), doc1.namespace_(),
+                                         doc1.uri(), doc_id1),
+                  EqualsDocumentMetadata(doc2.schema(), doc2.namespace_(),
+                                         doc2.uri(), doc_id2),
+                  EqualsDocumentMetadata(doc4.schema(), doc4.namespace_(),
+                                         doc4.uri(), doc_id4))));
+  EXPECT_FALSE(document_store->GetNonDeletedDocumentFilterData(doc_id1));
+  EXPECT_FALSE(document_store->GetNonDeletedDocumentFilterData(doc_id2));
+  EXPECT_TRUE(document_store->GetNonDeletedDocumentFilterData(doc_id3));
+  EXPECT_FALSE(document_store->GetNonDeletedDocumentFilterData(doc_id4));
+}
+
+TEST_P(DocumentStoreTest, GetNextExpiredDocumentTimestampMs) {
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
+                          schema_store_.get()));
+  std::unique_ptr<DocumentStore> document_store =
+      std::move(create_result.document_store);
+
+  DocumentProto doc1 = DocumentBuilder()
+                           .SetKey("namespace", "uri1")
+                           .SetSchema("email")
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(2000)  // Expires at 2000 ms.
+                           .Build();
+  DocumentProto doc2 = DocumentBuilder()
+                           .SetKey("namespace", "uri2")
+                           .SetSchema("email")
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(1200)  // Expires at 1200 ms.
+                           .Build();
+  DocumentProto doc3 = DocumentBuilder()
+                           .SetKey("namespace", "uri3")
+                           .SetSchema("email")
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(500)  // Expires at 500 ms.
+                           .Build();
+  DocumentProto doc4 = DocumentBuilder()
+                           .SetKey("namespace", "uri4")
+                           .SetSchema("email")
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(1800)  // Expires at 1800 ms.
+                           .Build();
+  DocumentProto doc5 = DocumentBuilder()
+                           .SetKey("namespace", "uri5")
+                           .SetSchema("email")
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(1000)  // Expires at 1000 ms.
+                           .Build();
+  // Put doc1 to doc5.
+  ICING_ASSERT_OK(
+      document_store->Put(document_util::CreateDocumentWrapper(doc1)));
+  ICING_ASSERT_OK(
+      document_store->Put(document_util::CreateDocumentWrapper(doc2)));
+  ICING_ASSERT_OK(
+      document_store->Put(document_util::CreateDocumentWrapper(doc3)));
+  ICING_ASSERT_OK(
+      document_store->Put(document_util::CreateDocumentWrapper(doc4)));
+  ICING_ASSERT_OK(
+      document_store->Put(document_util::CreateDocumentWrapper(doc5)));
+
+  EXPECT_THAT(
+      document_store->GetNextExpiredDocumentTimestampMs(/*current_time_ms=*/0),
+      Eq(500));
+  EXPECT_THAT(document_store->GetNextExpiredDocumentTimestampMs(
+                  /*current_time_ms=*/500),
+              Eq(1000));
+  EXPECT_THAT(document_store->GetNextExpiredDocumentTimestampMs(
+                  /*current_time_ms=*/1000),
+              Eq(1200));
+  EXPECT_THAT(document_store->GetNextExpiredDocumentTimestampMs(
+                  /*current_time_ms=*/1200),
+              Eq(1800));
+  EXPECT_THAT(document_store->GetNextExpiredDocumentTimestampMs(
+                  /*current_time_ms=*/1800),
+              Eq(2000));
+  EXPECT_THAT(document_store->GetNextExpiredDocumentTimestampMs(
+                  /*current_time_ms=*/2000),
+              Eq(-1));
+}
+
+TEST_P(DocumentStoreTest,
+       GetNextExpiredDocumentTimestampMs_shouldSkipDeletedDocuments) {
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
+                          schema_store_.get()));
+  std::unique_ptr<DocumentStore> document_store =
+      std::move(create_result.document_store);
+
+  DocumentProto doc1 = DocumentBuilder()
+                           .SetKey("namespace", "uri1")
+                           .SetSchema("email")
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(2000)  // Expires at 2000 ms.
+                           .Build();
+  DocumentProto doc2 = DocumentBuilder()
+                           .SetKey("namespace", "uri2")
+                           .SetSchema("email")
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(500)  // Expires at 500 ms.
+                           .Build();
+  DocumentProto doc3 = DocumentBuilder()
+                           .SetKey("namespace", "uri3")
+                           .SetSchema("email")
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(1200)  // Expires at 1200 ms.
+                           .Build();
+
+  // Put doc1 to doc3.
+  ICING_ASSERT_OK(
+      document_store->Put(document_util::CreateDocumentWrapper(doc1)));
+  ICING_ASSERT_OK(
+      document_store->Put(document_util::CreateDocumentWrapper(doc2)));
+  ICING_ASSERT_OK(
+      document_store->Put(document_util::CreateDocumentWrapper(doc3)));
+
+  EXPECT_THAT(
+      document_store->GetNextExpiredDocumentTimestampMs(/*current_time_ms=*/0),
+      Eq(500));
+
+  // Delete doc2 and get the next expired document timestamp again with current
+  // time = 0.
+  ICING_ASSERT_OK(
+      document_store->Delete("namespace", "uri2", /*current_time_ms=*/0));
+  EXPECT_THAT(
+      document_store->GetNextExpiredDocumentTimestampMs(/*current_time_ms=*/0),
+      Eq(1200));
+}
+
+TEST_P(
+    DocumentStoreTest,
+    GetNextExpiredDocumentTimestampMs_shouldReturnNegativeOneIfNoDocumentExpires) {
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
+                          schema_store_.get()));
+  std::unique_ptr<DocumentStore> document_store =
+      std::move(create_result.document_store);
+
+  DocumentProto doc1 = DocumentBuilder()
+                           .SetKey("namespace", "uri1")
+                           .SetSchema("email")
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(0)  // Never expires.
+                           .Build();
+  DocumentProto doc2 = DocumentBuilder()
+                           .SetKey("namespace", "uri2")
+                           .SetSchema("email")
+                           .SetCreationTimestampMs(0)
+                           .SetTtlMs(0)  // Never expires.
+                           .Build();
+
+  // Put doc1 to doc5.
+  ICING_ASSERT_OK(
+      document_store->Put(document_util::CreateDocumentWrapper(doc1)));
+  ICING_ASSERT_OK(
+      document_store->Put(document_util::CreateDocumentWrapper(doc2)));
+
+  EXPECT_THAT(
+      document_store->GetNextExpiredDocumentTimestampMs(/*current_time_ms=*/0),
+      Eq(-1));
+}
+
 TEST_P(DocumentStoreTest, OptimizeIntoSingleNamespace) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::CreateResult create_result,
@@ -1177,9 +1774,12 @@
   // Nothing should have expired yet.
   fake_clock_.SetSystemTimeMilliseconds(100);
 
-  ICING_ASSERT_OK(doc_store->Put(document1));
-  ICING_ASSERT_OK(doc_store->Put(document2));
-  ICING_ASSERT_OK(doc_store->Put(document3));
+  ICING_ASSERT_OK(doc_store->Put(
+      CreateDocumentWrapper(document1, /*num_string_tokens=*/0)));
+  ICING_ASSERT_OK(doc_store->Put(
+      CreateDocumentWrapper(document2, /*num_string_tokens=*/0)));
+  ICING_ASSERT_OK(doc_store->Put(
+      CreateDocumentWrapper(document3, /*num_string_tokens=*/0)));
 
   std::string original_document_log = absl_ports::StrCat(
       document_store_dir_, "/", DocumentLogCreator::GetDocumentLogFilename());
@@ -1326,11 +1926,16 @@
   // Nothing should have expired yet.
   fake_clock_.SetSystemTimeMilliseconds(100);
 
-  ICING_ASSERT_OK(doc_store->Put(document0));
-  ICING_ASSERT_OK(doc_store->Put(document1));
-  ICING_ASSERT_OK(doc_store->Put(document2));
-  ICING_ASSERT_OK(doc_store->Put(document3));
-  ICING_ASSERT_OK(doc_store->Put(document4));
+  ICING_ASSERT_OK(doc_store->Put(
+      CreateDocumentWrapper(document0, /*num_string_tokens=*/0)));
+  ICING_ASSERT_OK(doc_store->Put(
+      CreateDocumentWrapper(document1, /*num_string_tokens=*/0)));
+  ICING_ASSERT_OK(doc_store->Put(
+      CreateDocumentWrapper(document2, /*num_string_tokens=*/0)));
+  ICING_ASSERT_OK(doc_store->Put(
+      CreateDocumentWrapper(document3, /*num_string_tokens=*/0)));
+  ICING_ASSERT_OK(doc_store->Put(
+      CreateDocumentWrapper(document4, /*num_string_tokens=*/0)));
 
   std::string original_document_log = absl_ports::StrCat(
       document_store_dir_, "/", DocumentLogCreator::GetDocumentLogFilename());
@@ -1529,17 +2134,17 @@
     std::unique_ptr<DocumentStore> doc_store =
         std::move(create_result.document_store);
 
-    ICING_ASSERT_OK_AND_ASSIGN(
-        DocumentStore::PutResult put_result1,
-        doc_store->Put(DocumentProto(test_document1_), /*num_tokens=*/4));
+    ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
+                               doc_store->Put(CreateDocumentWrapper(
+                                   test_document1_, /*num_string_tokens=*/4)));
     EXPECT_THAT(put_result1.old_document_id, Eq(kInvalidDocumentId));
-    EXPECT_FALSE(put_result1.was_replacement());
+    EXPECT_FALSE(put_result1.was_replacement);
     document_id1 = put_result1.new_document_id;
-    ICING_ASSERT_OK_AND_ASSIGN(
-        DocumentStore::PutResult put_result2,
-        doc_store->Put(DocumentProto(test_document2_), /*num_tokens=*/4));
+    ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
+                               doc_store->Put(CreateDocumentWrapper(
+                                   test_document2_, /*num_string_tokens=*/4)));
     EXPECT_THAT(put_result2.old_document_id, Eq(kInvalidDocumentId));
-    EXPECT_FALSE(put_result2.was_replacement());
+    EXPECT_FALSE(put_result2.was_replacement);
     document_id2 = put_result2.new_document_id;
     EXPECT_THAT(doc_store->Get(document_id1),
                 IsOkAndHolds(EqualsProto(test_document1_)));
@@ -1615,11 +2220,10 @@
       DocumentFilterData doc_filter_data,
       doc_store->GetAliveDocumentFilterData(
           document_id2, fake_clock_.GetSystemTimeMilliseconds()));
-  EXPECT_THAT(
-      doc_filter_data,
-      Eq(DocumentFilterData(
-          /*namespace_id=*/0, tc3farmhash::Fingerprint64(test_document2_.uri()),
-          /*schema_type_id=*/0, document2_expiration_timestamp_)));
+  EXPECT_THAT(doc_filter_data,
+              Eq(DocumentFilterData(/*namespace_id=*/0, /*schema_type_id=*/0,
+                                    document2_expiration_timestamp_,
+                                    document2_expiration_timestamp_)));
 
   // Checks derived score cache
   EXPECT_THAT(
@@ -1655,17 +2259,17 @@
     std::unique_ptr<DocumentStore> doc_store =
         std::move(create_result.document_store);
 
-    ICING_ASSERT_OK_AND_ASSIGN(
-        DocumentStore::PutResult put_result1,
-        doc_store->Put(DocumentProto(test_document1_), /*num_tokens=*/4));
+    ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
+                               doc_store->Put(CreateDocumentWrapper(
+                                   test_document1_, /*num_string_tokens=*/4)));
     EXPECT_THAT(put_result1.old_document_id, Eq(kInvalidDocumentId));
-    EXPECT_FALSE(put_result1.was_replacement());
+    EXPECT_FALSE(put_result1.was_replacement);
     document_id1 = put_result1.new_document_id;
-    ICING_ASSERT_OK_AND_ASSIGN(
-        DocumentStore::PutResult put_result2,
-        doc_store->Put(DocumentProto(test_document2_), /*num_tokens=*/4));
+    ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
+                               doc_store->Put(CreateDocumentWrapper(
+                                   test_document2_, /*num_string_tokens=*/4)));
     EXPECT_THAT(put_result2.old_document_id, Eq(kInvalidDocumentId));
-    EXPECT_FALSE(put_result2.was_replacement());
+    EXPECT_FALSE(put_result2.was_replacement);
     document_id2 = put_result2.new_document_id;
     EXPECT_THAT(doc_store->Get(document_id1),
                 IsOkAndHolds(EqualsProto(test_document1_)));
@@ -1751,11 +2355,10 @@
       DocumentFilterData doc_filter_data,
       doc_store->GetAliveDocumentFilterData(
           document_id2, fake_clock_.GetSystemTimeMilliseconds()));
-  EXPECT_THAT(
-      doc_filter_data,
-      Eq(DocumentFilterData(
-          /*namespace_id=*/0, tc3farmhash::Fingerprint64(test_document2_.uri()),
-          /*schema_type_id=*/0, document2_expiration_timestamp_)));
+  EXPECT_THAT(doc_filter_data,
+              Eq(DocumentFilterData(/*namespace_id=*/0, /*schema_type_id=*/0,
+                                    document2_expiration_timestamp_,
+                                    document2_expiration_timestamp_)));
 
   // Checks derived score cache
   EXPECT_THAT(
@@ -1801,17 +2404,17 @@
     std::unique_ptr<DocumentStore> doc_store =
         std::move(create_result.document_store);
 
-    ICING_ASSERT_OK_AND_ASSIGN(
-        DocumentStore::PutResult put_result1,
-        doc_store->Put(DocumentProto(test_document1_), /*num_tokens=*/4));
+    ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
+                               doc_store->Put(CreateDocumentWrapper(
+                                   test_document1_, /*num_string_tokens=*/4)));
     EXPECT_THAT(put_result1.old_document_id, Eq(kInvalidDocumentId));
-    EXPECT_FALSE(put_result1.was_replacement());
+    EXPECT_FALSE(put_result1.was_replacement);
     document_id1 = put_result1.new_document_id;
-    ICING_ASSERT_OK_AND_ASSIGN(
-        DocumentStore::PutResult put_result2,
-        doc_store->Put(DocumentProto(test_document2_), /*num_tokens=*/4));
+    ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
+                               doc_store->Put(CreateDocumentWrapper(
+                                   test_document2_, /*num_string_tokens=*/4)));
     EXPECT_THAT(put_result2.old_document_id, Eq(kInvalidDocumentId));
-    EXPECT_FALSE(put_result2.was_replacement());
+    EXPECT_FALSE(put_result2.was_replacement);
     document_id2 = put_result2.new_document_id;
     EXPECT_THAT(doc_store->Get(document_id1),
                 IsOkAndHolds(EqualsProto(test_document1_)));
@@ -1885,11 +2488,10 @@
       DocumentFilterData doc_filter_data,
       doc_store->GetAliveDocumentFilterData(
           document_id2, fake_clock_.GetSystemTimeMilliseconds()));
-  EXPECT_THAT(
-      doc_filter_data,
-      Eq(DocumentFilterData(
-          /*namespace_id=*/0, tc3farmhash::Fingerprint64(test_document2_.uri()),
-          /*schema_type_id=*/0, document2_expiration_timestamp_)));
+  EXPECT_THAT(doc_filter_data,
+              Eq(DocumentFilterData(/*namespace_id=*/0, /*schema_type_id=*/0,
+                                    document2_expiration_timestamp_,
+                                    document2_expiration_timestamp_)));
 
   // Checks derived score cache.
   EXPECT_THAT(
@@ -1935,17 +2537,17 @@
     std::unique_ptr<DocumentStore> doc_store =
         std::move(create_result.document_store);
 
-    ICING_ASSERT_OK_AND_ASSIGN(
-        DocumentStore::PutResult put_result1,
-        doc_store->Put(DocumentProto(test_document1_), /*num_tokens=*/4));
+    ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
+                               doc_store->Put(CreateDocumentWrapper(
+                                   test_document1_, /*num_string_tokens=*/4)));
     EXPECT_THAT(put_result1.old_document_id, Eq(kInvalidDocumentId));
-    EXPECT_FALSE(put_result1.was_replacement());
+    EXPECT_FALSE(put_result1.was_replacement);
     document_id1 = put_result1.new_document_id;
-    ICING_ASSERT_OK_AND_ASSIGN(
-        DocumentStore::PutResult put_result2,
-        doc_store->Put(DocumentProto(test_document2_), /*num_tokens=*/4));
+    ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
+                               doc_store->Put(CreateDocumentWrapper(
+                                   test_document2_, /*num_string_tokens=*/4)));
     EXPECT_THAT(put_result2.old_document_id, Eq(kInvalidDocumentId));
-    EXPECT_FALSE(put_result2.was_replacement());
+    EXPECT_FALSE(put_result2.was_replacement);
     document_id2 = put_result2.new_document_id;
     EXPECT_THAT(doc_store->Get(document_id1),
                 IsOkAndHolds(EqualsProto(test_document1_)));
@@ -2010,11 +2612,10 @@
       DocumentFilterData doc_filter_data,
       doc_store->GetAliveDocumentFilterData(
           document_id2, fake_clock_.GetSystemTimeMilliseconds()));
-  EXPECT_THAT(
-      doc_filter_data,
-      Eq(DocumentFilterData(
-          /*namespace_id=*/0, tc3farmhash::Fingerprint64(test_document2_.uri()),
-          /*schema_type_id=*/0, document2_expiration_timestamp_)));
+  EXPECT_THAT(doc_filter_data,
+              Eq(DocumentFilterData(/*namespace_id=*/0, /*schema_type_id=*/0,
+                                    document2_expiration_timestamp_,
+                                    document2_expiration_timestamp_)));
   // Checks derived score cache
   EXPECT_THAT(
       doc_store->GetDocumentAssociatedScoreData(document_id2),
@@ -2061,7 +2662,8 @@
   // 1 block size. The number 100 is a bit arbitrary, gotten from manually
   // testing.
   for (int i = 0; i < 100; ++i) {
-    ICING_ASSERT_OK(doc_store->Put(document));
+    ICING_ASSERT_OK(
+        doc_store->Put(document_util::CreateDocumentWrapper(document)));
   }
   doc_store_storage_info = doc_store->GetStorageInfo();
   EXPECT_THAT(doc_store_storage_info.document_store_size(),
@@ -2093,10 +2695,11 @@
   // Since the DocumentStore is empty, we get an invalid DocumentId
   EXPECT_THAT(doc_store->last_added_document_id(), Eq(kInvalidDocumentId));
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             doc_store->Put(DocumentProto(test_document1_)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      doc_store->Put(document_util::CreateDocumentWrapper(test_document1_)));
   EXPECT_THAT(put_result1.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result1.was_replacement());
+  EXPECT_FALSE(put_result1.was_replacement);
   DocumentId document_id1 = put_result1.new_document_id;
   EXPECT_THAT(doc_store->last_added_document_id(), Eq(document_id1));
 
@@ -2105,10 +2708,11 @@
                                     fake_clock_.GetSystemTimeMilliseconds()));
   EXPECT_THAT(doc_store->last_added_document_id(), Eq(document_id1));
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             doc_store->Put(DocumentProto(test_document2_)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      doc_store->Put(document_util::CreateDocumentWrapper(test_document2_)));
   EXPECT_THAT(put_result2.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result2.was_replacement());
+  EXPECT_FALSE(put_result2.was_replacement);
   DocumentId document_id2 = put_result2.new_document_id;
   EXPECT_THAT(doc_store->last_added_document_id(), Eq(document_id2));
 }
@@ -2126,8 +2730,10 @@
   DocumentProto document_namespace2 =
       DocumentBuilder().SetKey("namespace2", "2").SetSchema("email").Build();
 
-  ICING_ASSERT_OK(doc_store->Put(DocumentProto(document_namespace1)));
-  ICING_ASSERT_OK(doc_store->Put(DocumentProto(document_namespace2)));
+  ICING_ASSERT_OK(doc_store->Put(
+      document_util::CreateDocumentWrapper(document_namespace1)));
+  ICING_ASSERT_OK(doc_store->Put(
+      document_util::CreateDocumentWrapper(document_namespace2)));
 
   // NamespaceId of 0 since it was the first namespace seen by the DocumentStore
   EXPECT_THAT(doc_store->GetNamespaceId("namespace1"), IsOkAndHolds(Eq(0)));
@@ -2159,8 +2765,10 @@
   DocumentProto document2 =
       DocumentBuilder().SetKey("namespace", "2").SetSchema("email").Build();
 
-  ICING_ASSERT_OK(doc_store->Put(document1));
-  ICING_ASSERT_OK(doc_store->Put(document2));
+  ICING_ASSERT_OK(
+      doc_store->Put(document_util::CreateDocumentWrapper(document1)));
+  ICING_ASSERT_OK(
+      doc_store->Put(document_util::CreateDocumentWrapper(document2)));
 
   // NamespaceId of 0 since it was the first namespace seen by the DocumentStore
   EXPECT_THAT(doc_store->GetNamespaceId("namespace"), IsOkAndHolds(Eq(0)));
@@ -2191,8 +2799,10 @@
   DocumentProto document2 =
       DocumentBuilder().SetKey("namespace", "2").SetSchema("email").Build();
 
-  ICING_ASSERT_OK(doc_store->Put(document1));
-  ICING_ASSERT_OK(doc_store->Put(document2));
+  ICING_ASSERT_OK(
+      doc_store->Put(document_util::CreateDocumentWrapper(document1)));
+  ICING_ASSERT_OK(
+      doc_store->Put(document_util::CreateDocumentWrapper(document2)));
 
   // CorpusId of 0 since it was the first namespace seen by the DocumentStore
   EXPECT_THAT(doc_store->GetCorpusId("namespace", "email"),
@@ -2212,8 +2822,10 @@
   DocumentProto document_corpus2 =
       DocumentBuilder().SetKey("namespace2", "2").SetSchema("email").Build();
 
-  ICING_ASSERT_OK(doc_store->Put(DocumentProto(document_corpus1)));
-  ICING_ASSERT_OK(doc_store->Put(DocumentProto(document_corpus2)));
+  ICING_ASSERT_OK(
+      doc_store->Put(document_util::CreateDocumentWrapper(document_corpus1)));
+  ICING_ASSERT_OK(
+      doc_store->Put(document_util::CreateDocumentWrapper(document_corpus2)));
 
   // CorpusId of 0 since it was the first corpus seen by the DocumentStore
   EXPECT_THAT(doc_store->GetCorpusId("namespace1", "email"),
@@ -2248,7 +2860,8 @@
 
   DocumentProto document_corpus =
       DocumentBuilder().SetKey("namespace1", "1").SetSchema("email").Build();
-  ICING_ASSERT_OK(doc_store->Put(DocumentProto(document_corpus)));
+  ICING_ASSERT_OK(
+      doc_store->Put(document_util::CreateDocumentWrapper(document_corpus)));
 
   EXPECT_THAT(doc_store->GetCorpusId("nonexistent_namespace", "email"),
               StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
@@ -2271,8 +2884,10 @@
   DocumentProto document2 =
       DocumentBuilder().SetKey("namespace", "2").SetSchema("email").Build();
 
-  ICING_ASSERT_OK(doc_store->Put(document1, /*num_tokens=*/5));
-  ICING_ASSERT_OK(doc_store->Put(document2, /*num_tokens=*/7));
+  ICING_ASSERT_OK(doc_store->Put(
+      CreateDocumentWrapper(document1, /*num_string_tokens=*/5)));
+  ICING_ASSERT_OK(doc_store->Put(
+      CreateDocumentWrapper(document2, /*num_string_tokens=*/7)));
 
   // CorpusId of 0 since it was the first namespace seen by the DocumentStore
   EXPECT_THAT(doc_store->GetCorpusAssociatedScoreData(/*corpus_id=*/0),
@@ -2296,10 +2911,10 @@
   DocumentProto document_corpus2 =
       DocumentBuilder().SetKey("namespace2", "2").SetSchema("email").Build();
 
-  ICING_ASSERT_OK(
-      doc_store->Put(DocumentProto(document_corpus1), /*num_tokens=*/5));
-  ICING_ASSERT_OK(
-      doc_store->Put(DocumentProto(document_corpus2), /*num_tokens=*/7));
+  ICING_ASSERT_OK(doc_store->Put(
+      CreateDocumentWrapper(document_corpus1, /*num_string_tokens=*/5)));
+  ICING_ASSERT_OK(doc_store->Put(
+      CreateDocumentWrapper(document_corpus2, /*num_string_tokens=*/7)));
 
   // CorpusId of 0 since it was the first corpus seen by the DocumentStore
   EXPECT_THAT(doc_store->GetCorpusAssociatedScoreData(/*corpus_id=*/0),
@@ -2358,17 +2973,17 @@
               document2_creation_timestamp_)  // A random timestamp
           .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(
-      DocumentStore::PutResult put_result1,
-      doc_store->Put(DocumentProto(document1), /*num_tokens=*/5));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
+                             doc_store->Put(CreateDocumentWrapper(
+                                 document1, /*num_string_tokens=*/5)));
   EXPECT_THAT(put_result1.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result1.was_replacement());
+  EXPECT_FALSE(put_result1.was_replacement);
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(
-      DocumentStore::PutResult put_result2,
-      doc_store->Put(DocumentProto(document2), /*num_tokens=*/7));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
+                             doc_store->Put(CreateDocumentWrapper(
+                                 document2, /*num_string_tokens=*/7)));
   EXPECT_THAT(put_result2.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result2.was_replacement());
+  EXPECT_FALSE(put_result2.was_replacement);
   DocumentId document_id2 = put_result2.new_document_id;
 
   EXPECT_THAT(
@@ -2410,17 +3025,17 @@
               document2_creation_timestamp_)  // A random timestamp
           .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(
-      DocumentStore::PutResult put_result1,
-      doc_store->Put(DocumentProto(document1), /*num_tokens=*/5));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
+                             doc_store->Put(CreateDocumentWrapper(
+                                 document1, /*num_string_tokens=*/5)));
   EXPECT_THAT(put_result1.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result1.was_replacement());
+  EXPECT_FALSE(put_result1.was_replacement);
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(
-      DocumentStore::PutResult put_result2,
-      doc_store->Put(DocumentProto(document2), /*num_tokens=*/7));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
+                             doc_store->Put(CreateDocumentWrapper(
+                                 document2, /*num_string_tokens=*/7)));
   EXPECT_THAT(put_result2.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result2.was_replacement());
+  EXPECT_FALSE(put_result2.was_replacement);
   DocumentId document_id2 = put_result2.new_document_id;
 
   EXPECT_THAT(
@@ -2469,21 +3084,21 @@
   std::unique_ptr<DocumentStore> doc_store =
       std::move(create_result.document_store);
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             doc_store->Put(DocumentProto(test_document1_)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store->Put(document_util::CreateDocumentWrapper(test_document1_)));
   EXPECT_THAT(put_result.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result.was_replacement());
+  EXPECT_FALSE(put_result.was_replacement);
   DocumentId document_id = put_result.new_document_id;
 
   ICING_ASSERT_HAS_VALUE_AND_ASSIGN(
       DocumentFilterData doc_filter_data,
       doc_store->GetAliveDocumentFilterData(
           document_id, fake_clock_.GetSystemTimeMilliseconds()));
-  EXPECT_THAT(
-      doc_filter_data,
-      Eq(DocumentFilterData(
-          /*namespace_id=*/0, tc3farmhash::Fingerprint64(test_document1_.uri()),
-          /*schema_type_id=*/0, document1_expiration_timestamp_)));
+  EXPECT_THAT(doc_filter_data,
+              Eq(DocumentFilterData(/*namespace_id=*/0, /*schema_type_id=*/0,
+                                    document1_expiration_timestamp_,
+                                    document1_expiration_timestamp_)));
 
   ICING_ASSERT_OK(doc_store->Delete("icing", "email/1",
                                     fake_clock_.GetSystemTimeMilliseconds()));
@@ -2500,11 +3115,11 @@
   std::unique_ptr<DocumentStore> doc_store =
       std::move(create_result.document_store);
 
-  ICING_ASSERT_OK_AND_ASSIGN(
-      DocumentStore::PutResult put_result,
-      doc_store->Put(DocumentProto(test_document1_), /*num_tokens=*/4));
+  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
+                             doc_store->Put(CreateDocumentWrapper(
+                                 test_document1_, /*num_string_tokens=*/4)));
   EXPECT_THAT(put_result.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result.was_replacement());
+  EXPECT_FALSE(put_result.was_replacement);
   DocumentId document_id = put_result.new_document_id;
 
   EXPECT_THAT(doc_store->GetDocumentAssociatedScoreData(document_id),
@@ -2536,7 +3151,8 @@
       std::move(create_result.document_store);
 
   ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             doc_store->Put(test_document1_, /*num_tokens=*/4));
+                             doc_store->Put(CreateDocumentWrapper(
+                                 test_document1_, /*num_string_tokens=*/4)));
   DocumentId document_id = put_result1.new_document_id;
   EXPECT_NE(doc_store->GetScorablePropertySet(
                 document_id, fake_clock_.GetSystemTimeMilliseconds()),
@@ -2557,10 +3173,11 @@
   std::unique_ptr<DocumentStore> doc_store =
       std::move(create_result.document_store);
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             doc_store->Put(DocumentProto(test_document1_)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store->Put(document_util::CreateDocumentWrapper(test_document1_)));
   EXPECT_THAT(put_result.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result.was_replacement());
+  EXPECT_FALSE(put_result.was_replacement);
   DocumentId document_id = put_result.new_document_id;
 
   // Report usage with type 1.
@@ -2609,10 +3226,11 @@
                                .SetTtlMs(100)
                                .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             doc_store->Put(document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store->Put(document_util::CreateDocumentWrapper(document)));
   EXPECT_THAT(put_result.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result.was_replacement());
+  EXPECT_FALSE(put_result.was_replacement);
   DocumentId document_id = put_result.new_document_id;
 
   // Some arbitrary time before the document's creation time (10) + ttl (100)
@@ -2645,6 +3263,252 @@
       document_id, fake_clock_.GetSystemTimeMilliseconds()));
 }
 
+TEST_P(DocumentStoreTest, IsDocumentAlive) {
+  DocumentProto document = DocumentBuilder()
+                               .SetKey("namespace", "uri")
+                               .SetSchema("email")
+                               .SetCreationTimestampMs(0)
+                               .SetTtlMs(1000)
+                               .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
+                          schema_store_.get()));
+  std::unique_ptr<DocumentStore> doc_store =
+      std::move(create_result.document_store);
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store->Put(document_util::CreateDocumentWrapper(document)));
+  ASSERT_THAT(put_result.old_document_id, Eq(kInvalidDocumentId));
+  ASSERT_FALSE(put_result.was_replacement);
+  DocumentId document_id = put_result.new_document_id;
+
+  // The expiration timestamp is 1000, so the document is alive at any time
+  // before that.
+  EXPECT_TRUE(doc_store->IsDocumentAlive(document_id, /*current_time_ms=*/0));
+  EXPECT_TRUE(doc_store->IsDocumentAlive(document_id, /*current_time_ms=*/100));
+  EXPECT_TRUE(doc_store->IsDocumentAlive(document_id, /*current_time_ms=*/999));
+  EXPECT_FALSE(
+      doc_store->IsDocumentAlive(document_id, /*current_time_ms=*/1000));
+  EXPECT_FALSE(
+      doc_store->IsDocumentAlive(document_id, /*current_time_ms=*/1001));
+  EXPECT_FALSE(
+      doc_store->IsDocumentAlive(document_id, /*current_time_ms=*/2000));
+
+  EXPECT_TRUE(
+      doc_store->IsDocumentAlive("namespace", "uri", /*current_time_ms=*/0));
+  EXPECT_TRUE(
+      doc_store->IsDocumentAlive("namespace", "uri", /*current_time_ms=*/100));
+  EXPECT_TRUE(
+      doc_store->IsDocumentAlive("namespace", "uri", /*current_time_ms=*/999));
+  EXPECT_FALSE(
+      doc_store->IsDocumentAlive("namespace", "uri", /*current_time_ms=*/1000));
+  EXPECT_FALSE(
+      doc_store->IsDocumentAlive("namespace", "uri", /*current_time_ms=*/1001));
+  EXPECT_FALSE(
+      doc_store->IsDocumentAlive("namespace", "uri", /*current_time_ms=*/2000));
+}
+
+TEST_P(DocumentStoreTest, IsDocumentAlive_deletedDocument) {
+  DocumentProto document = DocumentBuilder()
+                               .SetKey("namespace", "uri")
+                               .SetSchema("email")
+                               .SetCreationTimestampMs(0)
+                               .SetTtlMs(1000)
+                               .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
+                          schema_store_.get()));
+  std::unique_ptr<DocumentStore> doc_store =
+      std::move(create_result.document_store);
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store->Put(document_util::CreateDocumentWrapper(document)));
+  ASSERT_THAT(put_result.old_document_id, Eq(kInvalidDocumentId));
+  ASSERT_FALSE(put_result.was_replacement);
+  DocumentId document_id = put_result.new_document_id;
+
+  // Delete the document.
+  ICING_ASSERT_OK(doc_store->Delete("namespace", "uri",
+                                    fake_clock_.GetSystemTimeMilliseconds()));
+
+  // Should be non-alive regardless of the current timestamp.
+  EXPECT_FALSE(doc_store->IsDocumentAlive(document_id, /*current_time_ms=*/0));
+  EXPECT_FALSE(
+      doc_store->IsDocumentAlive(document_id, /*current_time_ms=*/100));
+  EXPECT_FALSE(
+      doc_store->IsDocumentAlive(document_id, /*current_time_ms=*/1099));
+  EXPECT_FALSE(
+      doc_store->IsDocumentAlive(document_id, /*current_time_ms=*/1100));
+  EXPECT_FALSE(
+      doc_store->IsDocumentAlive(document_id, /*current_time_ms=*/1101));
+  EXPECT_FALSE(
+      doc_store->IsDocumentAlive(document_id, /*current_time_ms=*/2000));
+
+  EXPECT_FALSE(
+      doc_store->IsDocumentAlive("namespace", "uri", /*current_time_ms=*/0));
+  EXPECT_FALSE(
+      doc_store->IsDocumentAlive("namespace", "uri", /*current_time_ms=*/100));
+  EXPECT_FALSE(
+      doc_store->IsDocumentAlive("namespace", "uri", /*current_time_ms=*/1099));
+  EXPECT_FALSE(
+      doc_store->IsDocumentAlive("namespace", "uri", /*current_time_ms=*/1100));
+  EXPECT_FALSE(
+      doc_store->IsDocumentAlive("namespace", "uri", /*current_time_ms=*/1101));
+  EXPECT_FALSE(
+      doc_store->IsDocumentAlive("namespace", "uri", /*current_time_ms=*/2000));
+}
+
+TEST_P(DocumentStoreTest, IsDocumentAlive_infiniteTtlDocument) {
+  DocumentProto document = DocumentBuilder()
+                               .SetKey("namespace", "uri")
+                               .SetSchema("email")
+                               .SetCreationTimestampMs(0)
+                               .SetTtlMs(0)  // infinite TTL
+                               .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
+                          schema_store_.get()));
+  std::unique_ptr<DocumentStore> doc_store =
+      std::move(create_result.document_store);
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store->Put(document_util::CreateDocumentWrapper(document)));
+  ASSERT_THAT(put_result.old_document_id, Eq(kInvalidDocumentId));
+  ASSERT_FALSE(put_result.was_replacement);
+  DocumentId document_id = put_result.new_document_id;
+
+  EXPECT_TRUE(doc_store->IsDocumentAlive(document_id, /*current_time_ms=*/0));
+  EXPECT_TRUE(doc_store->IsDocumentAlive(document_id, /*current_time_ms=*/100));
+  EXPECT_TRUE(
+      doc_store->IsDocumentAlive(document_id, /*current_time_ms=*/1099));
+  EXPECT_TRUE(
+      doc_store->IsDocumentAlive(document_id, /*current_time_ms=*/1100));
+  EXPECT_TRUE(
+      doc_store->IsDocumentAlive(document_id, /*current_time_ms=*/1101));
+  EXPECT_TRUE(
+      doc_store->IsDocumentAlive(document_id, /*current_time_ms=*/2000));
+  EXPECT_TRUE(doc_store->IsDocumentAlive(
+      document_id,
+      /*current_time_ms=*/std::numeric_limits<int64_t>::max() - 1));
+
+  EXPECT_TRUE(
+      doc_store->IsDocumentAlive("namespace", "uri", /*current_time_ms=*/0));
+  EXPECT_TRUE(
+      doc_store->IsDocumentAlive("namespace", "uri", /*current_time_ms=*/100));
+  EXPECT_TRUE(
+      doc_store->IsDocumentAlive("namespace", "uri", /*current_time_ms=*/1099));
+  EXPECT_TRUE(
+      doc_store->IsDocumentAlive("namespace", "uri", /*current_time_ms=*/1100));
+  EXPECT_TRUE(
+      doc_store->IsDocumentAlive("namespace", "uri", /*current_time_ms=*/1101));
+  EXPECT_TRUE(
+      doc_store->IsDocumentAlive("namespace", "uri", /*current_time_ms=*/2000));
+  EXPECT_TRUE(doc_store->IsDocumentAlive(
+      "namespace", "uri",
+      /*current_time_ms=*/std::numeric_limits<int64_t>::max() - 1));
+}
+
+TEST_P(DocumentStoreTest, IsDocumentAlive_invalidDocumentId) {
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
+                          schema_store_.get()));
+  std::unique_ptr<DocumentStore> doc_store =
+      std::move(create_result.document_store);
+
+  ASSERT_THAT(doc_store->last_added_document_id(), Eq(kInvalidDocumentId));
+
+  EXPECT_FALSE(
+      doc_store->IsDocumentAlive(/*document_id=*/-1, /*current_time_ms=*/0));
+  EXPECT_FALSE(
+      doc_store->IsDocumentAlive(/*document_id=*/0, /*current_time_ms=*/0));
+  EXPECT_FALSE(
+      doc_store->IsDocumentAlive(/*document_id=*/10000, /*current_time_ms=*/0));
+  EXPECT_FALSE(doc_store->IsDocumentAlive(/*document_id=*/100000,
+                                          /*current_time_ms=*/0));
+  EXPECT_FALSE(
+      doc_store->IsDocumentAlive(kInvalidDocumentId, /*current_time_ms=*/0));
+}
+
+TEST_P(DocumentStoreTest, IsDocumentAlive_invalidNamespaceUri) {
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
+                          schema_store_.get()));
+  std::unique_ptr<DocumentStore> doc_store =
+      std::move(create_result.document_store);
+
+  ASSERT_THAT(doc_store->last_added_document_id(), Eq(kInvalidDocumentId));
+
+  EXPECT_FALSE(
+      doc_store->IsDocumentAlive("namespace", "uri", /*current_time_ms=*/0));
+}
+
+TEST_P(DocumentStoreTest,
+       GetAliveDocumentFilterDataShouldCheckExpirationTimestamp) {
+  DocumentProto document = DocumentBuilder()
+                               .SetKey("namespace1", "1")
+                               .SetSchema("email")
+                               .SetCreationTimestampMs(100)
+                               .SetTtlMs(1000)
+                               .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
+                          schema_store_.get()));
+  std::unique_ptr<DocumentStore> doc_store =
+      std::move(create_result.document_store);
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store->Put(document_util::CreateDocumentWrapper(document)));
+  EXPECT_THAT(put_result.old_document_id, Eq(kInvalidDocumentId));
+  EXPECT_FALSE(put_result.was_replacement);
+  DocumentId document_id = put_result.new_document_id;
+
+  // The raw expiration timestamp is 1100. Now, set the expiration timestamp to
+  // 500.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::UpdateDocumentExpirationTimestampResult update_result,
+      doc_store->UpdateDocumentExpirationTimestamp(
+          document_id, /*expiration_timestamp_ms=*/500));
+  EXPECT_THAT(update_result.final_expiration_timestamp_ms, Eq(500));
+  EXPECT_THAT(update_result.was_updated, IsTrue());
+
+  // Case 1: call GetAliveDocumentFilterData with a current ts before the
+  //         expiration timestamp.
+  EXPECT_THAT(
+      doc_store->GetAliveDocumentFilterData(document_id,
+                                            /*current_time_ms=*/300),
+      Optional(Eq(DocumentFilterData(/*namespace_id=*/0, /*schema_type_id=*/0,
+                                     /*expiration_timestamp_ms=*/500,
+                                     /*raw_expiration_timestamp_ms=*/1100))));
+
+  // Case 2: call GetAliveDocumentFilterData with a current ts at the
+  //         expiration timestamp.
+  // This should fail even though the current ts is smaller than the raw
+  // expiration timestamp.
+  EXPECT_FALSE(doc_store->GetAliveDocumentFilterData(document_id,
+                                                     /*current_time_ms=*/500));
+
+  // Case 3: call GetAliveDocumentFilterData with a current ts after the
+  //         expiration timestamp.
+  // This should fail even though the current ts is smaller than the raw
+  // expiration timestamp.
+  EXPECT_FALSE(doc_store->GetAliveDocumentFilterData(document_id,
+                                                     /*current_time_ms=*/700));
+}
+
 TEST_P(DocumentStoreTest,
        ExpirationTimestampIsSumOfNonZeroTtlAndCreationTimestamp) {
   DocumentProto document = DocumentBuilder()
@@ -2661,20 +3525,20 @@
   std::unique_ptr<DocumentStore> doc_store =
       std::move(create_result.document_store);
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             doc_store->Put(document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store->Put(document_util::CreateDocumentWrapper(document)));
   EXPECT_THAT(put_result.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result.was_replacement());
+  EXPECT_FALSE(put_result.was_replacement);
   DocumentId document_id = put_result.new_document_id;
   ICING_ASSERT_HAS_VALUE_AND_ASSIGN(
       DocumentFilterData doc_filter_data,
       doc_store->GetAliveDocumentFilterData(
           document_id, fake_clock_.GetSystemTimeMilliseconds()));
-  EXPECT_THAT(
-      doc_filter_data,
-      Eq(DocumentFilterData(
-          /*namespace_id=*/0, tc3farmhash::Fingerprint64(document.uri()),
-          /*schema_type_id=*/0, /*expiration_timestamp_ms=*/1100)));
+  EXPECT_THAT(doc_filter_data, Eq(DocumentFilterData(
+                                   /*namespace_id=*/0, /*schema_type_id=*/0,
+                                   /*expiration_timestamp_ms=*/1100,
+                                   /*raw_expiration_timestamp_ms=*/1100)));
 }
 
 TEST_P(DocumentStoreTest, ExpirationTimestampIsInt64MaxIfTtlIsZero) {
@@ -2692,10 +3556,11 @@
   std::unique_ptr<DocumentStore> doc_store =
       std::move(create_result.document_store);
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             doc_store->Put(document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store->Put(document_util::CreateDocumentWrapper(document)));
   EXPECT_THAT(put_result.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result.was_replacement());
+  EXPECT_FALSE(put_result.was_replacement);
   DocumentId document_id = put_result.new_document_id;
 
   ICING_ASSERT_HAS_VALUE_AND_ASSIGN(
@@ -2706,9 +3571,10 @@
   EXPECT_THAT(
       doc_filter_data,
       Eq(DocumentFilterData(
-          /*namespace_id=*/0, tc3farmhash::Fingerprint64(document.uri()),
-          /*schema_type_id=*/0,
-          /*expiration_timestamp_ms=*/std::numeric_limits<int64_t>::max())));
+          /*namespace_id=*/0, /*schema_type_id=*/0,
+          /*expiration_timestamp_ms=*/std::numeric_limits<int64_t>::max(),
+          /*raw_expiration_timestamp_ms=*/
+          std::numeric_limits<int64_t>::max())));
 }
 
 TEST_P(DocumentStoreTest, ExpirationTimestampIsInt64MaxOnOverflow) {
@@ -2727,10 +3593,11 @@
   std::unique_ptr<DocumentStore> doc_store =
       std::move(create_result.document_store);
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             doc_store->Put(document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store->Put(document_util::CreateDocumentWrapper(document)));
   EXPECT_THAT(put_result.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result.was_replacement());
+  EXPECT_FALSE(put_result.was_replacement);
   DocumentId document_id = put_result.new_document_id;
 
   ICING_ASSERT_HAS_VALUE_AND_ASSIGN(
@@ -2741,23 +3608,20 @@
   EXPECT_THAT(
       doc_filter_data,
       Eq(DocumentFilterData(
-          /*namespace_id=*/0, tc3farmhash::Fingerprint64(document.uri()),
-          /*schema_type_id=*/0,
-          /*expiration_timestamp_ms=*/std::numeric_limits<int64_t>::max())));
+          /*namespace_id=*/0, /*schema_type_id=*/0,
+          /*expiration_timestamp_ms=*/std::numeric_limits<int64_t>::max(),
+          /*raw_expiration_timestamp_ms=*/
+          std::numeric_limits<int64_t>::max())));
 }
 
-TEST_P(DocumentStoreTest, CreationTimestampShouldBePopulated) {
-  // Creates a document without a given creation timestamp
-  DocumentProto document_without_creation_timestamp =
-      DocumentBuilder()
-          .SetKey("icing", "email/1")
-          .SetSchema("email")
-          .AddStringProperty("subject", "subject foo")
-          .AddStringProperty("body", "body bar")
-          .Build();
+TEST_P(DocumentStoreTest, UpdateDocumentExpirationTimestamp) {
+  DocumentProto document = DocumentBuilder()
+                               .SetKey("namespace1", "1")
+                               .SetSchema("email")
+                               .SetCreationTimestampMs(100)
+                               .SetTtlMs(1000)
+                               .Build();
 
-  int64_t fake_real_time = 100;
-  fake_clock_.SetSystemTimeMilliseconds(fake_real_time);
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::CreateResult create_result,
       CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
@@ -2767,17 +3631,475 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult put_result,
-      doc_store->Put(document_without_creation_timestamp));
+      doc_store->Put(document_util::CreateDocumentWrapper(document)));
   EXPECT_THAT(put_result.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result.was_replacement());
+  EXPECT_FALSE(put_result.was_replacement);
   DocumentId document_id = put_result.new_document_id;
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentProto document_with_creation_timestamp,
-                             doc_store->Get(document_id));
+  EXPECT_THAT(
+      doc_store->GetAliveDocumentFilterData(
+          document_id, fake_clock_.GetSystemTimeMilliseconds()),
+      Optional(Eq(DocumentFilterData(/*namespace_id=*/0, /*schema_type_id=*/0,
+                                     /*expiration_timestamp_ms=*/1100,
+                                     /*raw_expiration_timestamp_ms=*/1100))));
 
-  // Now the creation timestamp should be set by document store.
-  EXPECT_THAT(document_with_creation_timestamp.creation_timestamp_ms(),
-              Eq(fake_real_time));
+  // Set the expiration timestamp to 700.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::UpdateDocumentExpirationTimestampResult update_result1,
+      doc_store->UpdateDocumentExpirationTimestamp(
+          document_id, /*expiration_timestamp_ms=*/700));
+  EXPECT_THAT(update_result1.final_expiration_timestamp_ms, Eq(700));
+  EXPECT_THAT(update_result1.was_updated, IsTrue());
+  EXPECT_THAT(
+      doc_store->GetAliveDocumentFilterData(
+          document_id, fake_clock_.GetSystemTimeMilliseconds()),
+      Optional(Eq(DocumentFilterData(/*namespace_id=*/0, /*schema_type_id=*/0,
+                                     /*expiration_timestamp_ms=*/700,
+                                     /*raw_expiration_timestamp_ms=*/1100))));
+
+  // Set the expiration timestamp to 500.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::UpdateDocumentExpirationTimestampResult update_result2,
+      doc_store->UpdateDocumentExpirationTimestamp(
+          document_id, /*expiration_timestamp_ms=*/500));
+  EXPECT_THAT(update_result2.final_expiration_timestamp_ms, Eq(500));
+  EXPECT_THAT(update_result2.was_updated, IsTrue());
+  EXPECT_THAT(
+      doc_store->GetAliveDocumentFilterData(
+          document_id, fake_clock_.GetSystemTimeMilliseconds()),
+      Optional(Eq(DocumentFilterData(/*namespace_id=*/0, /*schema_type_id=*/0,
+                                     /*expiration_timestamp_ms=*/500,
+                                     /*raw_expiration_timestamp_ms=*/1100))));
+}
+
+TEST_P(DocumentStoreTest, UpdateDocumentExpirationTimestamp_invalidDocumentId) {
+  DocumentProto document = DocumentBuilder()
+                               .SetKey("namespace1", "1")
+                               .SetSchema("email")
+                               .SetCreationTimestampMs(100)
+                               .SetTtlMs(1000)
+                               .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
+                          schema_store_.get()));
+  std::unique_ptr<DocumentStore> doc_store =
+      std::move(create_result.document_store);
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store->Put(document_util::CreateDocumentWrapper(document)));
+  ASSERT_THAT(put_result.old_document_id, Eq(kInvalidDocumentId));
+  ASSERT_FALSE(put_result.was_replacement);
+
+  EXPECT_THAT(doc_store->UpdateDocumentExpirationTimestamp(
+                  /*document_id=*/-1, /*expiration_timestamp_ms=*/700),
+              StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
+  EXPECT_THAT(doc_store->UpdateDocumentExpirationTimestamp(
+                  doc_store->last_added_document_id() + 1,
+                  /*expiration_timestamp_ms=*/700),
+              StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
+  EXPECT_THAT(doc_store->UpdateDocumentExpirationTimestamp(
+                  kInvalidDocumentId, /*expiration_timestamp_ms=*/700),
+              StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
+}
+
+TEST_P(DocumentStoreTest,
+       UpdateDocumentExpirationTimestamp_shouldOnlyDecrease) {
+  DocumentProto document = DocumentBuilder()
+                               .SetKey("namespace1", "1")
+                               .SetSchema("email")
+                               .SetCreationTimestampMs(100)
+                               .SetTtlMs(1000)
+                               .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
+                          schema_store_.get()));
+  std::unique_ptr<DocumentStore> doc_store =
+      std::move(create_result.document_store);
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store->Put(document_util::CreateDocumentWrapper(document)));
+  EXPECT_THAT(put_result.old_document_id, Eq(kInvalidDocumentId));
+  EXPECT_FALSE(put_result.was_replacement);
+  DocumentId document_id = put_result.new_document_id;
+
+  // The raw expiration timestamp and the expiration timestamp are 1100.
+  // Setting the expiration timestamp to 1500 should be no-op.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::UpdateDocumentExpirationTimestampResult update_result1,
+      doc_store->UpdateDocumentExpirationTimestamp(
+          document_id, /*expiration_timestamp_ms=*/1500));
+  EXPECT_THAT(update_result1.final_expiration_timestamp_ms, Eq(1100));
+  EXPECT_THAT(update_result1.was_updated, IsFalse());
+  EXPECT_THAT(
+      doc_store->GetAliveDocumentFilterData(
+          document_id, fake_clock_.GetSystemTimeMilliseconds()),
+      Optional(Eq(DocumentFilterData(/*namespace_id=*/0, /*schema_type_id=*/0,
+                                     /*expiration_timestamp_ms=*/1100,
+                                     /*raw_expiration_timestamp_ms=*/1100))));
+
+  // Set the expiration timestamp to 500. This should change the expiration
+  // timestamp to 500.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::UpdateDocumentExpirationTimestampResult update_result2,
+      doc_store->UpdateDocumentExpirationTimestamp(
+          document_id, /*expiration_timestamp_ms=*/500));
+  EXPECT_THAT(update_result2.final_expiration_timestamp_ms, Eq(500));
+  EXPECT_THAT(update_result2.was_updated, IsTrue());
+  EXPECT_THAT(
+      doc_store->GetAliveDocumentFilterData(
+          document_id, fake_clock_.GetSystemTimeMilliseconds()),
+      Optional(Eq(DocumentFilterData(/*namespace_id=*/0, /*schema_type_id=*/0,
+                                     /*expiration_timestamp_ms=*/500,
+                                     /*raw_expiration_timestamp_ms=*/1100))));
+
+  // Setting the expiration timestamp to 700 should be no-op.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::UpdateDocumentExpirationTimestampResult update_result3,
+      doc_store->UpdateDocumentExpirationTimestamp(
+          document_id, /*expiration_timestamp_ms=*/700));
+  EXPECT_THAT(update_result3.final_expiration_timestamp_ms, Eq(500));
+  EXPECT_THAT(update_result3.was_updated, IsFalse());
+  EXPECT_THAT(
+      doc_store->GetAliveDocumentFilterData(
+          document_id, fake_clock_.GetSystemTimeMilliseconds()),
+      Optional(Eq(DocumentFilterData(/*namespace_id=*/0, /*schema_type_id=*/0,
+                                     /*expiration_timestamp_ms=*/500,
+                                     /*raw_expiration_timestamp_ms=*/1100))));
+}
+
+TEST_P(DocumentStoreTest,
+       UpdateDocumentExpirationTimestamp_canUpdateExpiredDocument) {
+  DocumentProto document = DocumentBuilder()
+                               .SetKey("namespace1", "1")
+                               .SetSchema("email")
+                               .SetCreationTimestampMs(100)
+                               .SetTtlMs(1000)
+                               .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
+                          schema_store_.get()));
+  std::unique_ptr<DocumentStore> doc_store =
+      std::move(create_result.document_store);
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store->Put(document_util::CreateDocumentWrapper(document)));
+  EXPECT_THAT(put_result.old_document_id, Eq(kInvalidDocumentId));
+  EXPECT_FALSE(put_result.was_replacement);
+  DocumentId document_id = put_result.new_document_id;
+
+  // Set the expiration timestamp to 500.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::UpdateDocumentExpirationTimestampResult update_result1,
+      doc_store->UpdateDocumentExpirationTimestamp(
+          document_id, /*expiration_timestamp_ms=*/500));
+  ASSERT_THAT(update_result1.final_expiration_timestamp_ms, Eq(500));
+  ASSERT_THAT(update_result1.was_updated, IsTrue());
+  // Sanity check.
+  ASSERT_THAT(
+      doc_store->GetAliveDocumentFilterData(
+          document_id, fake_clock_.GetSystemTimeMilliseconds()),
+      Optional(Eq(DocumentFilterData(/*namespace_id=*/0, /*schema_type_id=*/0,
+                                     /*expiration_timestamp_ms=*/500,
+                                     /*raw_expiration_timestamp_ms=*/1100))));
+
+  // Adjust the current time to 500. The document is expired.
+  fake_clock_.SetSystemTimeMilliseconds(500);
+  ASSERT_THAT(doc_store->IsDocumentAlive(
+                  /*document_id=*/0,
+                  /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()),
+              IsFalse());
+
+  // Updating an expired document is allowed, but since we only allow to
+  // decrease the expiration timestamp, it is a no-op.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::UpdateDocumentExpirationTimestampResult update_result2,
+      doc_store->UpdateDocumentExpirationTimestamp(
+          document_id, /*expiration_timestamp_ms=*/800));
+  EXPECT_THAT(update_result2.final_expiration_timestamp_ms, Eq(500));
+  EXPECT_THAT(update_result2.was_updated, IsFalse());
+  EXPECT_THAT(
+      doc_store->GetNonDeletedDocumentFilterData(document_id),
+      Optional(Eq(DocumentFilterData(/*namespace_id=*/0, /*schema_type_id=*/0,
+                                     /*expiration_timestamp_ms=*/500,
+                                     /*raw_expiration_timestamp_ms=*/1100))));
+
+  // Updating an expired document is allowed, and it only takes effect if the
+  // new expiration timestamp is smaller than the current one, so the document
+  // remains expired.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::UpdateDocumentExpirationTimestampResult update_result3,
+      doc_store->UpdateDocumentExpirationTimestamp(
+          document_id, /*expiration_timestamp_ms=*/400));
+  EXPECT_THAT(update_result3.final_expiration_timestamp_ms, Eq(400));
+  EXPECT_THAT(update_result3.was_updated, IsTrue());
+  EXPECT_THAT(
+      doc_store->GetNonDeletedDocumentFilterData(document_id),
+      Optional(Eq(DocumentFilterData(/*namespace_id=*/0, /*schema_type_id=*/0,
+                                     /*expiration_timestamp_ms=*/400,
+                                     /*raw_expiration_timestamp_ms=*/1100))));
+}
+
+TEST_P(DocumentStoreTest,
+       UpdateDocumentExpirationTimestamp_shouldFailForDeletedDocument) {
+  DocumentProto document = DocumentBuilder()
+                               .SetKey("namespace1", "1")
+                               .SetSchema("email")
+                               .SetCreationTimestampMs(100)
+                               .SetTtlMs(1000)
+                               .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
+                          schema_store_.get()));
+  std::unique_ptr<DocumentStore> doc_store =
+      std::move(create_result.document_store);
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store->Put(document_util::CreateDocumentWrapper(document)));
+  EXPECT_THAT(put_result.old_document_id, Eq(kInvalidDocumentId));
+  EXPECT_FALSE(put_result.was_replacement);
+  DocumentId document_id = put_result.new_document_id;
+
+  // Delete the document.
+  ICING_ASSERT_OK(doc_store->Delete(
+      document_id,
+      /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+
+  EXPECT_THAT(doc_store->UpdateDocumentExpirationTimestamp(
+                  document_id, /*expiration_timestamp_ms=*/500),
+              StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
+}
+
+TEST_P(DocumentStoreTest,
+       UpdateDocumentExpirationTimestamp_canUpdateToAnExpiredTimestamp) {
+  DocumentProto document = DocumentBuilder()
+                               .SetKey("namespace1", "1")
+                               .SetSchema("email")
+                               .SetCreationTimestampMs(100)
+                               .SetTtlMs(1000)
+                               .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
+                          schema_store_.get()));
+  std::unique_ptr<DocumentStore> doc_store =
+      std::move(create_result.document_store);
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store->Put(document_util::CreateDocumentWrapper(document)));
+  EXPECT_THAT(put_result.old_document_id, Eq(kInvalidDocumentId));
+  EXPECT_FALSE(put_result.was_replacement);
+  DocumentId document_id = put_result.new_document_id;
+
+  // Adjust the current time to 500. The document is still alive.
+  fake_clock_.SetSystemTimeMilliseconds(500);
+  ASSERT_THAT(
+      doc_store->GetAliveDocumentFilterData(
+          document_id, fake_clock_.GetSystemTimeMilliseconds()),
+      Optional(Eq(DocumentFilterData(/*namespace_id=*/0, /*schema_type_id=*/0,
+                                     /*expiration_timestamp_ms=*/1100,
+                                     /*raw_expiration_timestamp_ms=*/1100))));
+
+  // Set the expiration timestamp to 400. It is allowed to set the expiration
+  // timestamp to a smaller value than the current timestamp, and the document
+  // will be treated as expired.
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::UpdateDocumentExpirationTimestampResult update_result,
+      doc_store->UpdateDocumentExpirationTimestamp(
+          document_id, /*expiration_timestamp_ms=*/400));
+  EXPECT_THAT(update_result.final_expiration_timestamp_ms, Eq(400));
+  EXPECT_THAT(update_result.was_updated, IsTrue());
+  EXPECT_THAT(doc_store->IsDocumentAlive(
+                  document_id,
+                  /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()),
+              IsFalse());
+}
+
+TEST_P(DocumentStoreTest, ResetAllAliveExpirationTimestampsToRaw) {
+  DocumentProto document1 = DocumentBuilder()
+                                .SetKey("namespace1", "1")
+                                .SetSchema("email")
+                                .SetCreationTimestampMs(100)
+                                .SetTtlMs(1000)
+                                .Build();
+  DocumentProto document2 =
+      DocumentBuilder()
+          .SetKey("namespace1", "2")
+          .SetSchema("email")
+          .SetCreationTimestampMs(300)
+          .SetTtlMs(0)  // No TTL. The expiration timestamp will be INT64_MAX.
+          .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
+                          schema_store_.get()));
+  std::unique_ptr<DocumentStore> doc_store =
+      std::move(create_result.document_store);
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      doc_store->Put(document_util::CreateDocumentWrapper(document1)));
+  EXPECT_THAT(put_result1.old_document_id, Eq(kInvalidDocumentId));
+  EXPECT_FALSE(put_result1.was_replacement);
+  DocumentId document_id1 = put_result1.new_document_id;
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      doc_store->Put(document_util::CreateDocumentWrapper(document2)));
+  EXPECT_THAT(put_result2.old_document_id, Eq(kInvalidDocumentId));
+  EXPECT_FALSE(put_result2.was_replacement);
+  DocumentId document_id2 = put_result2.new_document_id;
+
+  // Update the expiration timestamp of document1 and document2.
+  ICING_ASSERT_OK(doc_store->UpdateDocumentExpirationTimestamp(
+      document_id1, /*expiration_timestamp_ms=*/500));
+  ICING_ASSERT_OK(doc_store->UpdateDocumentExpirationTimestamp(
+      document_id2, /*expiration_timestamp_ms=*/3000));
+  // Sanity check.
+  ASSERT_THAT(
+      doc_store->GetAliveDocumentFilterData(
+          document_id1, fake_clock_.GetSystemTimeMilliseconds()),
+      Optional(Eq(DocumentFilterData(/*namespace_id=*/0, /*schema_type_id=*/0,
+                                     /*expiration_timestamp_ms=*/500,
+                                     /*raw_expiration_timestamp_ms=*/1100))));
+  ASSERT_THAT(
+      doc_store->GetAliveDocumentFilterData(
+          document_id2, fake_clock_.GetSystemTimeMilliseconds()),
+      Optional(Eq(DocumentFilterData(/*namespace_id=*/0, /*schema_type_id=*/0,
+                                     /*expiration_timestamp_ms=*/3000,
+                                     /*raw_expiration_timestamp_ms=*/
+                                     std::numeric_limits<int64_t>::max()))));
+
+  // Reset.
+  EXPECT_THAT(doc_store->ResetAllAliveExpirationTimestampsToRaw(
+                  /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()),
+              IsOk());
+  EXPECT_THAT(
+      doc_store->GetAliveDocumentFilterData(
+          document_id1, fake_clock_.GetSystemTimeMilliseconds()),
+      Optional(Eq(DocumentFilterData(/*namespace_id=*/0, /*schema_type_id=*/0,
+                                     /*expiration_timestamp_ms=*/1100,
+                                     /*raw_expiration_timestamp_ms=*/1100))));
+  EXPECT_THAT(
+      doc_store->GetAliveDocumentFilterData(
+          document_id2, fake_clock_.GetSystemTimeMilliseconds()),
+      Optional(Eq(DocumentFilterData(
+          /*namespace_id=*/0, /*schema_type_id=*/0,
+          /*expiration_timestamp_ms=*/std::numeric_limits<int64_t>::max(),
+          /*raw_expiration_timestamp_ms=*/
+          std::numeric_limits<int64_t>::max()))));
+}
+
+TEST_P(DocumentStoreTest,
+       ResetAllAliveExpirationTimestampsToRaw_shouldSkipExpiredDocuments) {
+  DocumentProto document1 = DocumentBuilder()
+                                .SetKey("namespace1", "1")
+                                .SetSchema("email")
+                                .SetCreationTimestampMs(100)
+                                .SetTtlMs(1000)
+                                .Build();
+  DocumentProto document2 =
+      DocumentBuilder()
+          .SetKey("namespace1", "2")
+          .SetSchema("email")
+          .SetCreationTimestampMs(300)
+          .SetTtlMs(0)  // No TTL. The expiration timestamp will be INT64_MAX.
+          .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
+                          schema_store_.get()));
+  std::unique_ptr<DocumentStore> doc_store =
+      std::move(create_result.document_store);
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      doc_store->Put(document_util::CreateDocumentWrapper(document1)));
+  EXPECT_THAT(put_result1.old_document_id, Eq(kInvalidDocumentId));
+  EXPECT_FALSE(put_result1.was_replacement);
+  DocumentId document_id1 = put_result1.new_document_id;
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      doc_store->Put(document_util::CreateDocumentWrapper(document2)));
+  EXPECT_THAT(put_result2.old_document_id, Eq(kInvalidDocumentId));
+  EXPECT_FALSE(put_result2.was_replacement);
+  DocumentId document_id2 = put_result2.new_document_id;
+
+  // Update the expiration timestamp of document1 and document2.
+  ICING_ASSERT_OK(doc_store->UpdateDocumentExpirationTimestamp(
+      document_id1, /*expiration_timestamp_ms=*/500));
+  ICING_ASSERT_OK(doc_store->UpdateDocumentExpirationTimestamp(
+      document_id2, /*expiration_timestamp_ms=*/3000));
+  // Sanity check.
+  ASSERT_THAT(
+      doc_store->GetAliveDocumentFilterData(
+          document_id1, fake_clock_.GetSystemTimeMilliseconds()),
+      Optional(Eq(DocumentFilterData(/*namespace_id=*/0, /*schema_type_id=*/0,
+                                     /*expiration_timestamp_ms=*/500,
+                                     /*raw_expiration_timestamp_ms=*/1100))));
+  ASSERT_THAT(
+      doc_store->GetAliveDocumentFilterData(
+          document_id2, fake_clock_.GetSystemTimeMilliseconds()),
+      Optional(Eq(DocumentFilterData(/*namespace_id=*/0, /*schema_type_id=*/0,
+                                     /*expiration_timestamp_ms=*/3000,
+                                     /*raw_expiration_timestamp_ms=*/
+                                     std::numeric_limits<int64_t>::max()))));
+
+  // Reset with current timestamp 500. Document1 should be skipped since it is
+  // expired (at t = 500).
+  EXPECT_THAT(doc_store->ResetAllAliveExpirationTimestampsToRaw(
+                  /*current_time_ms=*/500),
+              IsOk());
+  EXPECT_THAT(
+      doc_store->GetAliveDocumentFilterData(
+          document_id1, fake_clock_.GetSystemTimeMilliseconds()),
+      Optional(Eq(DocumentFilterData(/*namespace_id=*/0, /*schema_type_id=*/0,
+                                     /*expiration_timestamp_ms=*/500,
+                                     /*raw_expiration_timestamp_ms=*/1100))));
+  EXPECT_THAT(
+      doc_store->GetAliveDocumentFilterData(
+          document_id2, fake_clock_.GetSystemTimeMilliseconds()),
+      Optional(Eq(DocumentFilterData(
+          /*namespace_id=*/0, /*schema_type_id=*/0,
+          /*expiration_timestamp_ms=*/std::numeric_limits<int64_t>::max(),
+          /*raw_expiration_timestamp_ms=*/
+          std::numeric_limits<int64_t>::max()))));
+
+  // Reset with current timestamp 700. Document1 should be skipped since it is
+  // expired (at t = 500).
+  EXPECT_THAT(doc_store->ResetAllAliveExpirationTimestampsToRaw(
+                  /*current_time_ms=*/700),
+              IsOk());
+  EXPECT_THAT(
+      doc_store->GetAliveDocumentFilterData(
+          document_id1, fake_clock_.GetSystemTimeMilliseconds()),
+      Optional(Eq(DocumentFilterData(/*namespace_id=*/0, /*schema_type_id=*/0,
+                                     /*expiration_timestamp_ms=*/500,
+                                     /*raw_expiration_timestamp_ms=*/1100))));
+  EXPECT_THAT(
+      doc_store->GetAliveDocumentFilterData(
+          document_id2, fake_clock_.GetSystemTimeMilliseconds()),
+      Optional(Eq(DocumentFilterData(
+          /*namespace_id=*/0, /*schema_type_id=*/0,
+          /*expiration_timestamp_ms=*/std::numeric_limits<int64_t>::max(),
+          /*raw_expiration_timestamp_ms=*/
+          std::numeric_limits<int64_t>::max()))));
 }
 
 TEST_P(DocumentStoreTest, ShouldWriteAndReadScoresCorrectly) {
@@ -2801,15 +4123,17 @@
   std::unique_ptr<DocumentStore> doc_store =
       std::move(create_result.document_store);
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             doc_store->Put(document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      doc_store->Put(document_util::CreateDocumentWrapper(document1)));
   EXPECT_THAT(put_result1.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result1.was_replacement());
+  EXPECT_FALSE(put_result1.was_replacement);
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             doc_store->Put(document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      doc_store->Put(document_util::CreateDocumentWrapper(document2)));
   EXPECT_THAT(put_result2.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result2.was_replacement());
+  EXPECT_FALSE(put_result2.was_replacement);
   DocumentId document_id2 = put_result2.new_document_id;
 
   EXPECT_THAT(doc_store->GetDocumentAssociatedScoreData(document_id1),
@@ -2835,7 +4159,8 @@
   std::unique_ptr<DocumentStore> document_store =
       std::move(create_result.document_store);
 
-  ICING_EXPECT_OK(document_store->Put(test_document1_));
+  ICING_EXPECT_OK(document_store->Put(
+      document_util::CreateDocumentWrapper(test_document1_)));
   // GetChecksum should succeed without updating the checksum.
   ICING_EXPECT_OK(document_store->GetChecksum());
 
@@ -2854,7 +4179,8 @@
   std::unique_ptr<DocumentStore> document_store =
       std::move(create_result.document_store);
 
-  ICING_EXPECT_OK(document_store->Put(test_document1_));
+  ICING_EXPECT_OK(document_store->Put(
+      document_util::CreateDocumentWrapper(test_document1_)));
   ICING_ASSERT_OK_AND_ASSIGN(Crc32 checksum, document_store->GetChecksum());
   EXPECT_THAT(document_store->UpdateChecksum(), IsOkAndHolds(checksum));
   EXPECT_THAT(document_store->GetChecksum(), IsOkAndHolds(checksum));
@@ -2880,7 +4206,8 @@
   std::unique_ptr<DocumentStore> document_store =
       std::move(create_result.document_store);
 
-  ICING_EXPECT_OK(document_store->Put(test_document1_));
+  ICING_EXPECT_OK(document_store->Put(
+      document_util::CreateDocumentWrapper(test_document1_)));
   // GetChecksum should return the same value as UpdateChecksum
   ICING_ASSERT_OK_AND_ASSIGN(Crc32 checksum, document_store->GetChecksum());
   EXPECT_THAT(document_store->UpdateChecksum(), IsOkAndHolds(checksum));
@@ -2898,12 +4225,14 @@
   std::unique_ptr<DocumentStore> document_store =
       std::move(create_result.document_store);
 
-  ICING_EXPECT_OK(document_store->Put(test_document1_));
+  ICING_EXPECT_OK(document_store->Put(
+      document_util::CreateDocumentWrapper(test_document1_)));
   ICING_ASSERT_OK_AND_ASSIGN(Crc32 checksum, document_store->GetChecksum());
   EXPECT_THAT(document_store->UpdateChecksum(), IsOkAndHolds(checksum));
   EXPECT_THAT(document_store->GetChecksum(), IsOkAndHolds(checksum));
 
-  ICING_EXPECT_OK(document_store->Put(test_document2_));
+  ICING_EXPECT_OK(document_store->Put(
+      document_util::CreateDocumentWrapper(test_document2_)));
   EXPECT_THAT(document_store->GetChecksum(), IsOkAndHolds(Not(Eq(checksum))));
   EXPECT_THAT(document_store->UpdateChecksum(),
               IsOkAndHolds(Not(Eq(checksum))));
@@ -2917,7 +4246,8 @@
   std::unique_ptr<DocumentStore> document_store =
       std::move(create_result.document_store);
 
-  ICING_EXPECT_OK(document_store->Put(test_document1_));
+  ICING_EXPECT_OK(document_store->Put(
+      document_util::CreateDocumentWrapper(test_document1_)));
   ICING_ASSERT_OK_AND_ASSIGN(Crc32 checksum, document_store->GetChecksum());
   EXPECT_THAT(document_store->UpdateChecksum(), IsOkAndHolds(checksum));
   EXPECT_THAT(document_store->GetChecksum(), IsOkAndHolds(checksum));
@@ -2983,9 +4313,10 @@
     // Insert and verify a "email "document
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::PutResult email_put_result,
-        document_store->Put(DocumentProto(email_document)));
+        document_store->Put(
+            document_util::CreateDocumentWrapper(email_document)));
     EXPECT_THAT(email_put_result.old_document_id, Eq(kInvalidDocumentId));
-    EXPECT_FALSE(email_put_result.was_replacement());
+    EXPECT_FALSE(email_put_result.was_replacement);
     email_document_id = email_put_result.new_document_id;
     EXPECT_THAT(document_store->Get(email_document_id),
                 IsOkAndHolds(EqualsProto(email_document)));
@@ -2993,8 +4324,6 @@
         DocumentFilterData email_data,
         document_store->GetAliveDocumentFilterData(
             email_document_id, fake_clock_.GetSystemTimeMilliseconds()));
-    EXPECT_THAT(email_data.uri_fingerprint(),
-                Eq(tc3farmhash::Fingerprint64(email_document.uri())));
     EXPECT_THAT(email_data.schema_type_id(), Eq(email_schema_type_id));
     email_namespace_id = email_data.namespace_id();
     email_expiration_timestamp = email_data.expiration_timestamp_ms();
@@ -3002,9 +4331,10 @@
     // Insert and verify a "message" document
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::PutResult message_put_result,
-        document_store->Put(DocumentProto(message_document)));
+        document_store->Put(
+            document_util::CreateDocumentWrapper(message_document)));
     EXPECT_THAT(message_put_result.old_document_id, Eq(kInvalidDocumentId));
-    EXPECT_FALSE(message_put_result.was_replacement());
+    EXPECT_FALSE(message_put_result.was_replacement);
     message_document_id = message_put_result.new_document_id;
     EXPECT_THAT(document_store->Get(message_document_id),
                 IsOkAndHolds(EqualsProto(message_document)));
@@ -3012,8 +4342,6 @@
         DocumentFilterData message_data,
         document_store->GetAliveDocumentFilterData(
             message_document_id, fake_clock_.GetSystemTimeMilliseconds()));
-    EXPECT_THAT(message_data.uri_fingerprint(),
-                Eq(tc3farmhash::Fingerprint64(message_document.uri())));
     EXPECT_THAT(message_data.schema_type_id(), Eq(message_schema_type_id));
     message_namespace_id = message_data.namespace_id();
     message_expiration_timestamp = message_data.expiration_timestamp_ms();
@@ -3058,12 +4386,12 @@
       document_store->GetAliveDocumentFilterData(
           email_document_id, fake_clock_.GetSystemTimeMilliseconds()));
   EXPECT_THAT(email_data.schema_type_id(), Eq(email_schema_type_id));
-  // Make sure that all the other fields are stll valid/the same
+  // Make sure that all the other fields are still valid/the same
   EXPECT_THAT(email_data.namespace_id(), Eq(email_namespace_id));
-  EXPECT_THAT(email_data.uri_fingerprint(),
-              Eq(tc3farmhash::Fingerprint64(email_document.uri())));
   EXPECT_THAT(email_data.expiration_timestamp_ms(),
               Eq(email_expiration_timestamp));
+  EXPECT_THAT(email_data.raw_expiration_timestamp_ms(),
+              Eq(email_expiration_timestamp));
 
   // "message" document has an invalid SchemaTypeId
   EXPECT_THAT(document_store->Get(message_document_id),
@@ -3073,12 +4401,12 @@
       document_store->GetAliveDocumentFilterData(
           message_document_id, fake_clock_.GetSystemTimeMilliseconds()));
   EXPECT_THAT(message_data.schema_type_id(), Eq(-1));
-  // Make sure that all the other fields are stll valid/the same
+  // Make sure that all the other fields are still valid/the same
   EXPECT_THAT(message_data.namespace_id(), Eq(message_namespace_id));
-  EXPECT_THAT(message_data.uri_fingerprint(),
-              Eq(tc3farmhash::Fingerprint64(message_document.uri())));
   EXPECT_THAT(message_data.expiration_timestamp_ms(),
               Eq(message_expiration_timestamp));
+  EXPECT_THAT(message_data.raw_expiration_timestamp_ms(),
+              Eq(message_expiration_timestamp));
 }
 
 TEST_P(DocumentStoreTest, UpdateSchemaStoreUpdatesSchemaTypeIds) {
@@ -3089,6 +4417,7 @@
   // Set a schema
   SchemaProto schema =
       SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("dummy"))
           .AddType(SchemaTypeConfigBuilder().SetType("email"))
           .AddType(SchemaTypeConfigBuilder().SetType("message"))
           .Build();
@@ -3127,69 +4456,69 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult email_put_result,
-      document_store->Put(DocumentProto(email_document)));
+      document_store->Put(
+          document_util::CreateDocumentWrapper(email_document)));
   EXPECT_THAT(email_put_result.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(email_put_result.was_replacement());
+  EXPECT_FALSE(email_put_result.was_replacement);
   DocumentId email_document_id = email_put_result.new_document_id;
   ICING_ASSERT_HAS_VALUE_AND_ASSIGN(
       DocumentFilterData email_data,
       document_store->GetAliveDocumentFilterData(
           email_document_id, fake_clock_.GetSystemTimeMilliseconds()));
-  EXPECT_THAT(email_data.uri_fingerprint(),
-              Eq(tc3farmhash::Fingerprint64(email_document.uri())));
   EXPECT_THAT(email_data.schema_type_id(), Eq(old_email_schema_type_id));
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult message_put_result,
-      document_store->Put(DocumentProto(message_document)));
+      document_store->Put(
+          document_util::CreateDocumentWrapper(message_document)));
   EXPECT_THAT(message_put_result.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(message_put_result.was_replacement());
+  EXPECT_FALSE(message_put_result.was_replacement);
   DocumentId message_document_id = message_put_result.new_document_id;
   ICING_ASSERT_HAS_VALUE_AND_ASSIGN(
       DocumentFilterData message_data,
       document_store->GetAliveDocumentFilterData(
           message_document_id, fake_clock_.GetSystemTimeMilliseconds()));
-  EXPECT_THAT(message_data.uri_fingerprint(),
-              Eq(tc3farmhash::Fingerprint64(message_document.uri())));
   EXPECT_THAT(message_data.schema_type_id(), Eq(old_message_schema_type_id));
 
-  // Add a new schema type. Since SchemaTypeId is assigned based on order,
-  // this should change the SchemaTypeIds.
+  // Delete the first type. With schema type id optimization, this should change
+  // the type id of "message" but not "email."
   schema = SchemaBuilder()
-               .AddType(SchemaTypeConfigBuilder().SetType("newType"))
                .AddType(SchemaTypeConfigBuilder().SetType("email"))
                .AddType(SchemaTypeConfigBuilder().SetType("message"))
                .Build();
 
   ICING_EXPECT_OK(schema_store->SetSchema(
-      schema, /*ignore_errors_and_delete_documents=*/false));
+      schema, /*ignore_errors_and_delete_documents=*/true));
 
   ICING_ASSERT_OK_AND_ASSIGN(SchemaTypeId new_email_schema_type_id,
                              schema_store->GetSchemaTypeId("email"));
   ICING_ASSERT_OK_AND_ASSIGN(SchemaTypeId new_message_schema_type_id,
                              schema_store->GetSchemaTypeId("message"));
 
-  // SchemaTypeIds should have changed.
-  EXPECT_NE(old_email_schema_type_id, new_email_schema_type_id);
+  // SchemaTypeId for message should have changed.
   EXPECT_NE(old_message_schema_type_id, new_message_schema_type_id);
+  EXPECT_EQ(old_email_schema_type_id, new_email_schema_type_id);
 
-  ICING_EXPECT_OK(document_store->UpdateSchemaStore(schema_store.get()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::UpdateSchemaStoreResult update_result,
+      document_store->UpdateSchemaStore(schema_store.get()));
+  // - No documents should be deleted since 'dummy' type did not have any
+  //   documents.
+  // - The derived files should be changed.
+  EXPECT_THAT(update_result.deleted_document_count, Eq(0));
+  EXPECT_THAT(update_result.derived_files_changed, IsTrue());
 
   // Check that the FilterCache holds the new SchemaTypeIds
   ICING_ASSERT_HAS_VALUE_AND_ASSIGN(
       email_data,
       document_store->GetAliveDocumentFilterData(
           email_document_id, fake_clock_.GetSystemTimeMilliseconds()));
-  EXPECT_THAT(email_data.uri_fingerprint(),
-              Eq(tc3farmhash::Fingerprint64(email_document.uri())));
   EXPECT_THAT(email_data.schema_type_id(), Eq(new_email_schema_type_id));
 
   ICING_ASSERT_HAS_VALUE_AND_ASSIGN(
       message_data,
       document_store->GetAliveDocumentFilterData(
           message_document_id, fake_clock_.GetSystemTimeMilliseconds()));
-  EXPECT_THAT(message_data.uri_fingerprint(),
-              Eq(tc3farmhash::Fingerprint64(message_document.uri())));
   EXPECT_THAT(message_data.schema_type_id(), Eq(new_message_schema_type_id));
 }
 
@@ -3241,10 +4570,11 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult email_without_subject_put_result,
-      document_store->Put(DocumentProto(email_without_subject)));
+      document_store->Put(
+          document_util::CreateDocumentWrapper(email_without_subject)));
   EXPECT_THAT(email_without_subject_put_result.old_document_id,
               Eq(kInvalidDocumentId));
-  EXPECT_FALSE(email_without_subject_put_result.was_replacement());
+  EXPECT_FALSE(email_without_subject_put_result.was_replacement);
   DocumentId email_without_subject_document_id =
       email_without_subject_put_result.new_document_id;
   EXPECT_THAT(document_store->Get(email_without_subject_document_id),
@@ -3252,10 +4582,11 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult email_with_subject_put_result,
-      document_store->Put(DocumentProto(email_with_subject)));
+      document_store->Put(
+          document_util::CreateDocumentWrapper(email_with_subject)));
   EXPECT_THAT(email_with_subject_put_result.old_document_id,
               Eq(kInvalidDocumentId));
-  EXPECT_FALSE(email_with_subject_put_result.was_replacement());
+  EXPECT_FALSE(email_with_subject_put_result.was_replacement);
   DocumentId email_with_subject_document_id =
       email_with_subject_put_result.new_document_id;
   EXPECT_THAT(document_store->Get(email_with_subject_document_id),
@@ -3269,7 +4600,14 @@
   ICING_EXPECT_OK(schema_store->SetSchema(
       schema, /*ignore_errors_and_delete_documents=*/true));
 
-  ICING_EXPECT_OK(document_store->UpdateSchemaStore(schema_store.get()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::UpdateSchemaStoreResult update_result,
+      document_store->UpdateSchemaStore(schema_store.get()));
+  // - email_without_subject should be deleted since the new required property
+  //   is missing.
+  // - The derived files should be changed.
+  EXPECT_THAT(update_result.deleted_document_count, Eq(1));
+  EXPECT_THAT(update_result.derived_files_changed, IsTrue());
 
   // The email without a subject should be marked as deleted
   EXPECT_THAT(document_store->Get(email_without_subject_document_id),
@@ -3323,18 +4661,22 @@
   std::unique_ptr<DocumentStore> document_store =
       std::move(create_result.document_store);
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult email_put_result,
-                             document_store->Put(email_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult email_put_result,
+      document_store->Put(
+          document_util::CreateDocumentWrapper(email_document)));
   EXPECT_THAT(email_put_result.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(email_put_result.was_replacement());
+  EXPECT_FALSE(email_put_result.was_replacement);
   DocumentId email_document_id = email_put_result.new_document_id;
   EXPECT_THAT(document_store->Get(email_document_id),
               IsOkAndHolds(EqualsProto(email_document)));
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult message_put_result,
-                             document_store->Put(message_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult message_put_result,
+      document_store->Put(
+          document_util::CreateDocumentWrapper(message_document)));
   EXPECT_THAT(message_put_result.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(message_put_result.was_replacement());
+  EXPECT_FALSE(message_put_result.was_replacement);
   DocumentId message_document_id = message_put_result.new_document_id;
   EXPECT_THAT(document_store->Get(message_document_id),
               IsOkAndHolds(EqualsProto(message_document)));
@@ -3348,7 +4690,13 @@
       schema_store->SetSchema(new_schema,
                               /*ignore_errors_and_delete_documents=*/true));
 
-  ICING_EXPECT_OK(document_store->UpdateSchemaStore(schema_store.get()));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::UpdateSchemaStoreResult update_result,
+      document_store->UpdateSchemaStore(schema_store.get()));
+  // - email should be deleted.
+  // - The derived files should be changed.
+  EXPECT_THAT(update_result.deleted_document_count, Eq(1));
+  EXPECT_THAT(update_result.derived_files_changed, IsTrue());
 
   // The "email" type is unknown now, so the "email" document should be deleted
   EXPECT_THAT(document_store->Get(email_document_id),
@@ -3359,6 +4707,80 @@
               IsOkAndHolds(EqualsProto(message_document)));
 }
 
+TEST_P(DocumentStoreTest, UpdateSchemaStoreWithFullyCompatibleDocumentStore) {
+  const std::string schema_store_dir = test_dir_ + "_custom";
+  filesystem_.DeleteDirectoryRecursively(schema_store_dir.c_str());
+  filesystem_.CreateDirectoryRecursively(schema_store_dir.c_str());
+
+  // Set a schema with "email" and "message" types.
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("email"))
+          .AddType(SchemaTypeConfigBuilder().SetType("message"))
+          .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<SchemaStore> schema_store,
+      SchemaStore::Create(&filesystem_, schema_store_dir, &fake_clock_,
+                          feature_flags_.get()));
+  ICING_EXPECT_OK(schema_store->SetSchema(
+      schema, /*ignore_errors_and_delete_documents=*/false));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
+                          schema_store.get()));
+  std::unique_ptr<DocumentStore> document_store =
+      std::move(create_result.document_store);
+
+  // Add "email" document
+  DocumentProto email_document = DocumentBuilder()
+                                     .SetNamespace("namespace")
+                                     .SetUri("email_uri")
+                                     .SetSchema("email")
+                                     .SetCreationTimestampMs(0)
+                                     .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult email_put_result,
+      document_store->Put(
+          document_util::CreateDocumentWrapper(email_document)));
+  EXPECT_THAT(email_put_result.old_document_id, Eq(kInvalidDocumentId));
+  EXPECT_FALSE(email_put_result.was_replacement);
+  DocumentId email_document_id = email_put_result.new_document_id;
+  EXPECT_THAT(document_store->Get(email_document_id),
+              IsOkAndHolds(EqualsProto(email_document)));
+
+  // Delete schema type "message".
+  // - "email" type should have the same schema type id.
+  // - There was no document with "message" schema type, so no document should
+  //   be deleted.
+  //
+  // Therefore, the new schema is fully compatible with the DocumentStore.
+  // Derived files should not be changed.
+  SchemaProto new_schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("email"))
+          .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      SchemaStore::SetSchemaResult set_schema_result,
+      schema_store->SetSchema(new_schema,
+                              /*ignore_errors_and_delete_documents=*/true));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::UpdateSchemaStoreResult update_result,
+      document_store->UpdateSchemaStore(schema_store.get()));
+  // - No document is deleted since there were no documents with "message"
+  //   schema type.
+  // - The derived files should be UNCHANGED.
+  EXPECT_THAT(update_result.deleted_document_count, Eq(0));
+  EXPECT_THAT(update_result.derived_files_changed, IsFalse());
+
+  // The "email" document should be unaffected
+  EXPECT_THAT(document_store->Get(email_document_id),
+              IsOkAndHolds(EqualsProto(email_document)));
+}
+
 TEST_P(DocumentStoreTest, OptimizedUpdateSchemaStoreUpdatesSchemaTypeIds) {
   const std::string schema_store_dir = test_dir_ + "_custom";
   filesystem_.DeleteDirectoryRecursively(schema_store_dir.c_str());
@@ -3367,6 +4789,7 @@
   // Set a schema
   SchemaProto schema =
       SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("dummy"))
           .AddType(SchemaTypeConfigBuilder().SetType("email"))
           .AddType(SchemaTypeConfigBuilder().SetType("message"))
           .Build();
@@ -3403,36 +4826,35 @@
   std::unique_ptr<DocumentStore> document_store =
       std::move(create_result.document_store);
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult email_put_result,
-                             document_store->Put(email_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult email_put_result,
+      document_store->Put(
+          document_util::CreateDocumentWrapper(email_document)));
   EXPECT_THAT(email_put_result.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(email_put_result.was_replacement());
+  EXPECT_FALSE(email_put_result.was_replacement);
   DocumentId email_document_id = email_put_result.new_document_id;
   ICING_ASSERT_HAS_VALUE_AND_ASSIGN(
       DocumentFilterData email_data,
       document_store->GetAliveDocumentFilterData(
           email_document_id, fake_clock_.GetSystemTimeMilliseconds()));
-  EXPECT_THAT(email_data.uri_fingerprint(),
-              Eq(tc3farmhash::Fingerprint64(email_document.uri())));
   EXPECT_THAT(email_data.schema_type_id(), Eq(old_email_schema_type_id));
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult message_put_result,
-                             document_store->Put(message_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult message_put_result,
+      document_store->Put(
+          document_util::CreateDocumentWrapper(message_document)));
   EXPECT_THAT(message_put_result.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(message_put_result.was_replacement());
+  EXPECT_FALSE(message_put_result.was_replacement);
   DocumentId message_document_id = message_put_result.new_document_id;
   ICING_ASSERT_HAS_VALUE_AND_ASSIGN(
       DocumentFilterData message_data,
       document_store->GetAliveDocumentFilterData(
           message_document_id, fake_clock_.GetSystemTimeMilliseconds()));
-  EXPECT_THAT(message_data.uri_fingerprint(),
-              Eq(tc3farmhash::Fingerprint64(message_document.uri())));
   EXPECT_THAT(message_data.schema_type_id(), Eq(old_message_schema_type_id));
 
-  // Add a new schema type. Since SchemaTypeId is assigned based on order,
-  // this should change the SchemaTypeIds.
+  // Delete the first type. With schema type id optimization, this should change
+  // the type id of "message" but not "email."
   schema = SchemaBuilder()
-               .AddType(SchemaTypeConfigBuilder().SetType("newType"))
                .AddType(SchemaTypeConfigBuilder().SetType("email"))
                .AddType(SchemaTypeConfigBuilder().SetType("message"))
                .Build();
@@ -3440,35 +4862,38 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       SchemaStore::SetSchemaResult set_schema_result,
       schema_store->SetSchema(schema,
-                              /*ignore_errors_and_delete_documents=*/false));
+                              /*ignore_errors_and_delete_documents=*/true));
 
   ICING_ASSERT_OK_AND_ASSIGN(SchemaTypeId new_email_schema_type_id,
                              schema_store->GetSchemaTypeId("email"));
   ICING_ASSERT_OK_AND_ASSIGN(SchemaTypeId new_message_schema_type_id,
                              schema_store->GetSchemaTypeId("message"));
 
-  // SchemaTypeIds should have changed.
-  EXPECT_NE(old_email_schema_type_id, new_email_schema_type_id);
+  // SchemaTypeId for 'message' should have changed.
   EXPECT_NE(old_message_schema_type_id, new_message_schema_type_id);
+  EXPECT_EQ(old_email_schema_type_id, new_email_schema_type_id);
 
-  ICING_EXPECT_OK(document_store->OptimizedUpdateSchemaStore(
-      schema_store.get(), set_schema_result));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::UpdateSchemaStoreResult update_result,
+      document_store->OptimizedUpdateSchemaStore(schema_store.get(),
+                                                 set_schema_result));
+  // - No documents should be deleted since 'dummy' type did not have any
+  //   documents.
+  // - The derived files should be changed.
+  EXPECT_THAT(update_result.deleted_document_count, Eq(0));
+  EXPECT_THAT(update_result.derived_files_changed, IsTrue());
 
   // Check that the FilterCache holds the new SchemaTypeIds
   ICING_ASSERT_HAS_VALUE_AND_ASSIGN(
       email_data,
       document_store->GetAliveDocumentFilterData(
           email_document_id, fake_clock_.GetSystemTimeMilliseconds()));
-  EXPECT_THAT(email_data.uri_fingerprint(),
-              Eq(tc3farmhash::Fingerprint64(email_document.uri())));
   EXPECT_THAT(email_data.schema_type_id(), Eq(new_email_schema_type_id));
 
   ICING_ASSERT_HAS_VALUE_AND_ASSIGN(
       message_data,
       document_store->GetAliveDocumentFilterData(
           message_document_id, fake_clock_.GetSystemTimeMilliseconds()));
-  EXPECT_THAT(message_data.uri_fingerprint(),
-              Eq(tc3farmhash::Fingerprint64(message_document.uri())));
   EXPECT_THAT(message_data.schema_type_id(), Eq(new_message_schema_type_id));
 }
 
@@ -3520,10 +4945,11 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult email_without_subject_put_result,
-      document_store->Put(email_without_subject));
+      document_store->Put(
+          document_util::CreateDocumentWrapper(email_without_subject)));
   EXPECT_THAT(email_without_subject_put_result.old_document_id,
               Eq(kInvalidDocumentId));
-  EXPECT_FALSE(email_without_subject_put_result.was_replacement());
+  EXPECT_FALSE(email_without_subject_put_result.was_replacement);
   DocumentId email_without_subject_document_id =
       email_without_subject_put_result.new_document_id;
   EXPECT_THAT(document_store->Get(email_without_subject_document_id),
@@ -3531,10 +4957,11 @@
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::PutResult email_with_subject_put_result,
-      document_store->Put(email_with_subject));
+      document_store->Put(
+          document_util::CreateDocumentWrapper(email_with_subject)));
   EXPECT_THAT(email_with_subject_put_result.old_document_id,
               Eq(kInvalidDocumentId));
-  EXPECT_FALSE(email_with_subject_put_result.was_replacement());
+  EXPECT_FALSE(email_with_subject_put_result.was_replacement);
   DocumentId email_with_subject_document_id =
       email_with_subject_put_result.new_document_id;
   EXPECT_THAT(document_store->Get(email_with_subject_document_id),
@@ -3550,8 +4977,15 @@
       schema_store->SetSchema(schema,
                               /*ignore_errors_and_delete_documents=*/true));
 
-  ICING_EXPECT_OK(document_store->OptimizedUpdateSchemaStore(
-      schema_store.get(), set_schema_result));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::UpdateSchemaStoreResult update_result,
+      document_store->OptimizedUpdateSchemaStore(schema_store.get(),
+                                                 set_schema_result));
+  // - email_without_subject should be deleted since the new required property
+  //   is missing.
+  // - The derived files should be changed.
+  EXPECT_THAT(update_result.deleted_document_count, Eq(1));
+  EXPECT_THAT(update_result.derived_files_changed, IsTrue());
 
   // The email without a subject should be marked as deleted
   EXPECT_THAT(document_store->Get(email_without_subject_document_id),
@@ -3605,18 +5039,22 @@
   std::unique_ptr<DocumentStore> document_store =
       std::move(create_result.document_store);
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult email_put_result,
-                             document_store->Put(email_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult email_put_result,
+      document_store->Put(
+          document_util::CreateDocumentWrapper(email_document)));
   EXPECT_THAT(email_put_result.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(email_put_result.was_replacement());
+  EXPECT_FALSE(email_put_result.was_replacement);
   DocumentId email_document_id = email_put_result.new_document_id;
   EXPECT_THAT(document_store->Get(email_document_id),
               IsOkAndHolds(EqualsProto(email_document)));
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult message_put_result,
-                             document_store->Put(message_document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult message_put_result,
+      document_store->Put(
+          document_util::CreateDocumentWrapper(message_document)));
   EXPECT_THAT(message_put_result.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(message_put_result.was_replacement());
+  EXPECT_FALSE(message_put_result.was_replacement);
   DocumentId message_document_id = message_put_result.new_document_id;
   EXPECT_THAT(document_store->Get(message_document_id),
               IsOkAndHolds(EqualsProto(message_document)));
@@ -3631,8 +5069,14 @@
       schema_store->SetSchema(new_schema,
                               /*ignore_errors_and_delete_documents=*/true));
 
-  ICING_EXPECT_OK(document_store->OptimizedUpdateSchemaStore(
-      schema_store.get(), set_schema_result));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::UpdateSchemaStoreResult update_result,
+      document_store->OptimizedUpdateSchemaStore(schema_store.get(),
+                                                 set_schema_result));
+  // - email should be deleted.
+  // - The derived files should be changed.
+  EXPECT_THAT(update_result.deleted_document_count, Eq(1));
+  EXPECT_THAT(update_result.derived_files_changed, IsTrue());
 
   // The "email" type is unknown now, so the "email" document should be deleted
   EXPECT_THAT(document_store->Get(email_document_id),
@@ -3643,6 +5087,82 @@
               IsOkAndHolds(EqualsProto(message_document)));
 }
 
+TEST_P(DocumentStoreTest,
+       OptimizedUpdateSchemaStoreWithFullyCompatibleDocumentStore) {
+  const std::string schema_store_dir = test_dir_ + "_custom";
+  filesystem_.DeleteDirectoryRecursively(schema_store_dir.c_str());
+  filesystem_.CreateDirectoryRecursively(schema_store_dir.c_str());
+
+  // Set a schema with "email" and "message" types.
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("email"))
+          .AddType(SchemaTypeConfigBuilder().SetType("message"))
+          .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<SchemaStore> schema_store,
+      SchemaStore::Create(&filesystem_, schema_store_dir, &fake_clock_,
+                          feature_flags_.get()));
+  ICING_EXPECT_OK(schema_store->SetSchema(
+      schema, /*ignore_errors_and_delete_documents=*/false));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
+                          schema_store.get()));
+  std::unique_ptr<DocumentStore> document_store =
+      std::move(create_result.document_store);
+
+  // Add "email" document
+  DocumentProto email_document = DocumentBuilder()
+                                     .SetNamespace("namespace")
+                                     .SetUri("email_uri")
+                                     .SetSchema("email")
+                                     .SetCreationTimestampMs(0)
+                                     .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult email_put_result,
+      document_store->Put(
+          document_util::CreateDocumentWrapper(email_document)));
+  EXPECT_THAT(email_put_result.old_document_id, Eq(kInvalidDocumentId));
+  EXPECT_FALSE(email_put_result.was_replacement);
+  DocumentId email_document_id = email_put_result.new_document_id;
+  EXPECT_THAT(document_store->Get(email_document_id),
+              IsOkAndHolds(EqualsProto(email_document)));
+
+  // Delete schema type "message".
+  // - "email" type should have the same schema type id.
+  // - There was no document with "message" schema type, so no document should
+  //   be deleted.
+  //
+  // Therefore, the new schema is fully compatible with the DocumentStore.
+  // Derived files should not be changed.
+  SchemaProto new_schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("email"))
+          .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      SchemaStore::SetSchemaResult set_schema_result,
+      schema_store->SetSchema(new_schema,
+                              /*ignore_errors_and_delete_documents=*/true));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::UpdateSchemaStoreResult update_result,
+      document_store->OptimizedUpdateSchemaStore(schema_store.get(),
+                                                 set_schema_result));
+  // - No document is deleted since there were no documents with "message"
+  //   schema type.
+  // - The derived files should be UNCHANGED.
+  EXPECT_THAT(update_result.deleted_document_count, Eq(0));
+  EXPECT_THAT(update_result.derived_files_changed, IsFalse());
+
+  // The "email" document should be unaffected
+  EXPECT_THAT(document_store->Get(email_document_id),
+              IsOkAndHolds(EqualsProto(email_document)));
+}
+
 TEST_P(DocumentStoreTest, GetOptimizeInfo) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::CreateResult create_result,
@@ -3658,7 +5178,8 @@
   EXPECT_THAT(optimize_info.optimizable_docs, Eq(0));
   EXPECT_THAT(optimize_info.estimated_optimizable_bytes, Eq(0));
 
-  ICING_EXPECT_OK(document_store->Put(DocumentProto(test_document1_)));
+  ICING_EXPECT_OK(document_store->Put(
+      document_util::CreateDocumentWrapper(test_document1_)));
 
   // Adding a document, still nothing is optimizable
   ICING_ASSERT_OK_AND_ASSIGN(optimize_info, document_store->GetOptimizeInfo());
@@ -3733,10 +5254,14 @@
                                  .SetTtlMs(100)
                                  .Build();
 
-  ICING_ASSERT_OK(document_store->Put(namespace1));
-  ICING_ASSERT_OK(document_store->Put(namespace2_uri1));
-  ICING_ASSERT_OK(document_store->Put(namespace2_uri2));
-  ICING_ASSERT_OK(document_store->Put(namespace3));
+  ICING_ASSERT_OK(
+      document_store->Put(document_util::CreateDocumentWrapper(namespace1)));
+  ICING_ASSERT_OK(document_store->Put(
+      document_util::CreateDocumentWrapper(namespace2_uri1)));
+  ICING_ASSERT_OK(document_store->Put(
+      document_util::CreateDocumentWrapper(namespace2_uri2)));
+  ICING_ASSERT_OK(
+      document_store->Put(document_util::CreateDocumentWrapper(namespace3)));
 
   auto get_result = document_store->Get("namespace1", "uri");
   get_result = document_store->Get("namespace2", "uri1");
@@ -3777,10 +5302,12 @@
   std::unique_ptr<DocumentStore> document_store =
       std::move(create_result.document_store);
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store->Put(test_document1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store->Put(
+          document_util::CreateDocumentWrapper(test_document1_)));
   EXPECT_THAT(put_result1.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result1.was_replacement());
+  EXPECT_FALSE(put_result1.was_replacement);
   DocumentId document_id = put_result1.new_document_id;
 
   // Report usage with type 1 and time 1.
@@ -3872,10 +5399,12 @@
   std::unique_ptr<DocumentStore> document_store =
       std::move(create_result.document_store);
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store->Put(test_document1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store->Put(
+          document_util::CreateDocumentWrapper(test_document1_)));
   EXPECT_THAT(put_result1.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result1.was_replacement());
+  EXPECT_FALSE(put_result1.was_replacement);
   DocumentId document_id = put_result1.new_document_id;
 
   // Report usage with type 1.
@@ -3928,10 +5457,12 @@
     std::unique_ptr<DocumentStore> document_store =
         std::move(create_result.document_store);
 
-    ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                               document_store->Put(test_document1_));
+    ICING_ASSERT_OK_AND_ASSIGN(
+        DocumentStore::PutResult put_result1,
+        document_store->Put(
+            document_util::CreateDocumentWrapper(test_document1_)));
     EXPECT_THAT(put_result1.old_document_id, Eq(kInvalidDocumentId));
-    EXPECT_FALSE(put_result1.was_replacement());
+    EXPECT_FALSE(put_result1.was_replacement);
     document_id = put_result1.new_document_id;
 
     // Report usage with type 1.
@@ -3976,10 +5507,12 @@
     std::unique_ptr<DocumentStore> document_store =
         std::move(create_result.document_store);
 
-    ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                               document_store->Put(test_document1_));
+    ICING_ASSERT_OK_AND_ASSIGN(
+        DocumentStore::PutResult put_result1,
+        document_store->Put(
+            document_util::CreateDocumentWrapper(test_document1_)));
     EXPECT_THAT(put_result1.old_document_id, Eq(kInvalidDocumentId));
-    EXPECT_FALSE(put_result1.was_replacement());
+    EXPECT_FALSE(put_result1.was_replacement);
     document_id = put_result1.new_document_id;
 
     // Report usage with type 1.
@@ -4032,10 +5565,12 @@
   std::unique_ptr<DocumentStore> document_store =
       std::move(create_result.document_store);
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store->Put(test_document1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store->Put(
+          document_util::CreateDocumentWrapper(test_document1_)));
   EXPECT_THAT(put_result1.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result1.was_replacement());
+  EXPECT_FALSE(put_result1.was_replacement);
   DocumentId document_id = put_result1.new_document_id;
 
   // Report usage with type 1.
@@ -4053,10 +5588,12 @@
   EXPECT_THAT(actual_scores, Eq(expected_scores));
 
   // Update the document.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store->Put(test_document1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store->Put(
+          document_util::CreateDocumentWrapper(test_document1_)));
   EXPECT_THAT(put_result2.old_document_id, Eq(document_id));
-  EXPECT_TRUE(put_result2.was_replacement());
+  EXPECT_TRUE(put_result2.was_replacement);
   DocumentId updated_document_id = put_result2.new_document_id;
   // We should get a different document id.
   ASSERT_THAT(updated_document_id, Not(Eq(document_id)));
@@ -4077,15 +5614,19 @@
   std::unique_ptr<DocumentStore> document_store =
       std::move(create_result.document_store);
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store->Put(test_document1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store->Put(
+          document_util::CreateDocumentWrapper(test_document1_)));
   EXPECT_THAT(put_result1.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result1.was_replacement());
+  EXPECT_FALSE(put_result1.was_replacement);
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store->Put(test_document2_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store->Put(
+          document_util::CreateDocumentWrapper(test_document2_)));
   EXPECT_THAT(put_result2.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result2.was_replacement());
+  EXPECT_FALSE(put_result2.was_replacement);
   DocumentId document_id2 = put_result2.new_document_id;
   ICING_ASSERT_OK(document_store->Delete(
       document_id1, fake_clock_.GetSystemTimeMilliseconds()));
@@ -4136,10 +5677,12 @@
       std::move(create_result.document_store);
 
   // Add the document.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store->Put(test_document1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store->Put(
+          document_util::CreateDocumentWrapper(test_document1_)));
   EXPECT_THAT(put_result1.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result1.was_replacement());
+  EXPECT_FALSE(put_result1.was_replacement);
   DocumentId document_id = put_result1.new_document_id;
 
   // Delete the document.
@@ -4147,12 +5690,14 @@
       document_id, fake_clock_.GetSystemTimeMilliseconds()));
 
   // Re-add the document.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store->Put(test_document1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store->Put(
+          document_util::CreateDocumentWrapper(test_document1_)));
 
   // Because the document was deleted, it should not be a replacement.
-  EXPECT_THAT(put_result2.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result2.was_replacement());
+  EXPECT_THAT(put_result2.old_document_id, Eq(0));
+  EXPECT_FALSE(put_result2.was_replacement);
   DocumentId updated_document_id = put_result2.new_document_id;
   ASSERT_THAT(updated_document_id, Not(Eq(document_id)));
 }
@@ -4168,10 +5713,11 @@
   // Add the document.
   DocumentProto doc1 = test_document1_;
   doc1.set_ttl_ms(1);
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             document_store->Put(doc1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      document_store->Put(document_util::CreateDocumentWrapper(doc1)));
   EXPECT_THAT(put_result1.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result1.was_replacement());
+  EXPECT_FALSE(put_result1.was_replacement);
   DocumentId document_id = put_result1.new_document_id;
 
   // Expire the document by advancing the clock by two milliseconds.
@@ -4179,12 +5725,14 @@
       fake_clock_.GetSystemTimeMilliseconds() + 2);
 
   // Re-add the document.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             document_store->Put(test_document1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      document_store->Put(
+          document_util::CreateDocumentWrapper(test_document1_)));
 
   // Because the document was expired, it should not be a replacement.
-  EXPECT_THAT(put_result2.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result2.was_replacement());
+  EXPECT_THAT(put_result2.old_document_id, Eq(0));
+  EXPECT_FALSE(put_result2.was_replacement);
   DocumentId updated_document_id = put_result2.new_document_id;
   ASSERT_THAT(updated_document_id, Not(Eq(document_id)));
 }
@@ -4201,10 +5749,11 @@
     EXPECT_THAT(create_result.data_loss, Eq(DataLoss::NONE));
     EXPECT_THAT(create_result.derived_files_regenerated, IsFalse());
 
-    ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                               doc_store->Put(test_document1_));
+    ICING_ASSERT_OK_AND_ASSIGN(
+        DocumentStore::PutResult put_result,
+        doc_store->Put(document_util::CreateDocumentWrapper(test_document1_)));
     EXPECT_THAT(put_result.old_document_id, Eq(kInvalidDocumentId));
-    EXPECT_FALSE(put_result.was_replacement());
+    EXPECT_FALSE(put_result.was_replacement);
     DocumentId document_id = put_result.new_document_id;
     EXPECT_THAT(doc_store->Get(document_id),
                 IsOkAndHolds(EqualsProto(test_document1_)));
@@ -4257,10 +5806,11 @@
     // initialization.
     corruptible_offset = filesystem_.GetFileSize(document_log_file.c_str());
 
-    ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                               doc_store->Put(test_document1_));
+    ICING_ASSERT_OK_AND_ASSIGN(
+        DocumentStore::PutResult put_result,
+        doc_store->Put(document_util::CreateDocumentWrapper(test_document1_)));
     EXPECT_THAT(put_result.old_document_id, Eq(kInvalidDocumentId));
-    EXPECT_FALSE(put_result.was_replacement());
+    EXPECT_FALSE(put_result.was_replacement);
     DocumentId document_id = put_result.new_document_id;
     EXPECT_THAT(doc_store->Get(document_id),
                 IsOkAndHolds(EqualsProto(test_document1_)));
@@ -4289,7 +5839,8 @@
 
     // Set dirty bit to true to reflect that something changed in the log.
     header.SetDirtyFlag(true);
-    header.SetHeaderChecksum(header.CalculateHeaderChecksum());
+    header.UpdateHeaderChecksums(
+        feature_flags_->enable_proto_log_new_header_format());
 
     WriteDocumentLogHeader(filesystem_, document_log_file, header);
   }
@@ -4350,7 +5901,9 @@
           /*force_recovery_and_revalidate_documents=*/false,
           GetParam().pre_mapping_fbv, GetParam().use_persistent_hash_map,
           PortableFileBackedProtoLog<DocumentWrapper>::kDefaultCompressionLevel,
-          &initialize_stats));
+          PortableFileBackedProtoLog<
+              DocumentWrapper>::kDefaultCompressionThresholdBytes,
+          protobuf_ports::kDefaultMemLevel, &initialize_stats));
   std::unique_ptr<DocumentStore> doc_store =
       std::move(create_result.document_store);
   // The document log is using the legacy v0 format so that a migration is
@@ -4375,24 +5928,28 @@
   DocumentProto document1 = test_document1_;
   document1.set_namespace_("namespace.1");
   document1.set_uri("uri1");
-  ICING_ASSERT_OK(doc_store->Put(document1));
+  ICING_ASSERT_OK(
+      doc_store->Put(document_util::CreateDocumentWrapper(document1)));
 
   DocumentProto document2 = test_document1_;
   document2.set_namespace_("namespace.1");
   document2.set_uri("uri2");
   document2.set_creation_timestamp_ms(fake_clock_.GetSystemTimeMilliseconds());
   document2.set_ttl_ms(100);
-  ICING_ASSERT_OK(doc_store->Put(document2));
+  ICING_ASSERT_OK(
+      doc_store->Put(document_util::CreateDocumentWrapper(document2)));
 
   DocumentProto document3 = test_document1_;
   document3.set_namespace_("namespace.1");
   document3.set_uri("uri3");
-  ICING_ASSERT_OK(doc_store->Put(document3));
+  ICING_ASSERT_OK(
+      doc_store->Put(document_util::CreateDocumentWrapper(document3)));
 
   DocumentProto document4 = test_document1_;
   document4.set_namespace_("namespace.2");
   document4.set_uri("uri1");
-  ICING_ASSERT_OK(doc_store->Put(document4));
+  ICING_ASSERT_OK(
+      doc_store->Put(document_util::CreateDocumentWrapper(document4)));
 
   // Report usage with type 1 on document1
   UsageReport usage_report_type1 = CreateUsageReport(
@@ -4492,7 +6049,10 @@
                            .SetDataTypeString(TERM_MATCH_EXACT, TOKENIZER_PLAIN)
                            .SetCardinality(CARDINALITY_OPTIONAL))
           .Build();
-  SchemaProto schema = SchemaBuilder().AddType(email_type_config).Build();
+  SchemaProto schema = SchemaBuilder()
+                           .AddType(SchemaTypeConfigBuilder().SetType("dummy"))
+                           .AddType(email_type_config)
+                           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
       SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_,
@@ -4500,8 +6060,9 @@
   ASSERT_THAT(schema_store->SetSchema(
                   schema, /*ignore_errors_and_delete_documents=*/false),
               IsOk());
-  // The typeid for "email" should be 0.
-  ASSERT_THAT(schema_store->GetSchemaTypeId("email"), IsOkAndHolds(0));
+  // The typeid for "dummy" should be 0 and "email" should be 1.
+  ASSERT_THAT(schema_store->GetSchemaTypeId("dummy"), IsOkAndHolds(0));
+  ASSERT_THAT(schema_store->GetSchemaTypeId("email"), IsOkAndHolds(1));
 
   DocumentId docid = kInvalidDocumentId;
   {
@@ -4524,58 +6085,43 @@
                 document1_creation_timestamp_)  // A random timestamp
             .SetTtlMs(document1_ttl_)
             .Build();
-    ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                               doc_store->Put(doc));
+    ICING_ASSERT_OK_AND_ASSIGN(
+        DocumentStore::PutResult put_result,
+        doc_store->Put(document_util::CreateDocumentWrapper(doc)));
     EXPECT_THAT(put_result.old_document_id, Eq(kInvalidDocumentId));
-    EXPECT_FALSE(put_result.was_replacement());
+    EXPECT_FALSE(put_result.was_replacement);
     docid = put_result.new_document_id;
     ICING_ASSERT_HAS_VALUE_AND_ASSIGN(
         DocumentFilterData filter_data,
         doc_store->GetAliveDocumentFilterData(
             docid, fake_clock_.GetSystemTimeMilliseconds()));
 
-    ASSERT_THAT(filter_data.uri_fingerprint(),
-                Eq(tc3farmhash::Fingerprint64(doc.uri())));
-    ASSERT_THAT(filter_data.schema_type_id(), Eq(0));
+    ASSERT_THAT(filter_data.schema_type_id(), Eq(1));
   }
 
-  // Add another type to the schema before the email type.
-  schema =
-      SchemaBuilder()
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("alarm")
-                       .AddProperty(PropertyConfigBuilder()
-                                        .SetName("name")
-                                        .SetDataTypeString(TERM_MATCH_EXACT,
-                                                           TOKENIZER_PLAIN)
-                                        .SetCardinality(CARDINALITY_OPTIONAL))
-                       .AddProperty(PropertyConfigBuilder()
-                                        .SetName("time")
-                                        .SetDataType(TYPE_INT64)
-                                        .SetCardinality(CARDINALITY_OPTIONAL)))
-          .AddType(email_type_config)
-          .Build();
+  // Delete the dummy type from the schema.
+  schema = SchemaBuilder().AddType(email_type_config).Build();
   ASSERT_THAT(schema_store->SetSchema(
-                  schema, /*ignore_errors_and_delete_documents=*/false),
+                  schema, /*ignore_errors_and_delete_documents=*/true),
               IsOk());
-  // Adding a new type should cause ids to be reassigned. Ids are assigned in
-  // order of appearance so 'alarm' should be 0 and 'email' should be 1.
-  ASSERT_THAT(schema_store->GetSchemaTypeId("alarm"), IsOkAndHolds(0));
-  ASSERT_THAT(schema_store->GetSchemaTypeId("email"), IsOkAndHolds(1));
+  // Deleting the dummy type will cause email's type id to be reassigned to 0.
+  ASSERT_THAT(schema_store->GetSchemaTypeId("email"), IsOkAndHolds(0));
 
   {
     // Create the document store the second time and force recovery
     InitializeStatsProto initialize_stats;
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
-        DocumentStore::Create(&filesystem_, document_store_dir_, &fake_clock_,
-                              schema_store.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/true,
-                              GetParam().pre_mapping_fbv,
-                              GetParam().use_persistent_hash_map,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              &initialize_stats));
+        DocumentStore::Create(
+            &filesystem_, document_store_dir_, &fake_clock_, schema_store.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/true,
+            GetParam().pre_mapping_fbv, GetParam().use_persistent_hash_map,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel, &initialize_stats));
     std::unique_ptr<DocumentStore> doc_store =
         std::move(create_result.document_store);
 
@@ -4584,9 +6130,7 @@
         DocumentFilterData filter_data,
         doc_store->GetAliveDocumentFilterData(
             docid, fake_clock_.GetSystemTimeMilliseconds()));
-    EXPECT_THAT(filter_data.uri_fingerprint(),
-                Eq(tc3farmhash::Fingerprint64(std::string("email/1"))));
-    EXPECT_THAT(filter_data.schema_type_id(), Eq(1));
+    EXPECT_THAT(filter_data.schema_type_id(), Eq(0));
     EXPECT_THAT(initialize_stats.document_store_recovery_cause(),
                 Eq(InitializeStatsProto::SCHEMA_CHANGES_OUT_OF_SYNC));
   }
@@ -4611,7 +6155,10 @@
                            .SetDataTypeString(TERM_MATCH_EXACT, TOKENIZER_PLAIN)
                            .SetCardinality(CARDINALITY_OPTIONAL))
           .Build();
-  SchemaProto schema = SchemaBuilder().AddType(email_type_config).Build();
+  SchemaProto schema = SchemaBuilder()
+                           .AddType(SchemaTypeConfigBuilder().SetType("dummy"))
+                           .AddType(email_type_config)
+                           .Build();
   ICING_ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<SchemaStore> schema_store,
       SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_,
@@ -4619,8 +6166,9 @@
   ASSERT_THAT(schema_store->SetSchema(
                   schema, /*ignore_errors_and_delete_documents=*/false),
               IsOk());
-  // The typeid for "email" should be 0.
-  ASSERT_THAT(schema_store->GetSchemaTypeId("email"), IsOkAndHolds(0));
+  // The typeid for "dummy" should be 0 and "email" should be 1.
+  ASSERT_THAT(schema_store->GetSchemaTypeId("dummy"), IsOkAndHolds(0));
+  ASSERT_THAT(schema_store->GetSchemaTypeId("email"), IsOkAndHolds(1));
 
   DocumentId docid = kInvalidDocumentId;
   {
@@ -4643,44 +6191,28 @@
                 document1_creation_timestamp_)  // A random timestamp
             .SetTtlMs(document1_ttl_)
             .Build();
-    ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                               doc_store->Put(doc));
+    ICING_ASSERT_OK_AND_ASSIGN(
+        DocumentStore::PutResult put_result,
+        doc_store->Put(document_util::CreateDocumentWrapper(doc)));
     EXPECT_THAT(put_result.old_document_id, Eq(kInvalidDocumentId));
-    EXPECT_FALSE(put_result.was_replacement());
+    EXPECT_FALSE(put_result.was_replacement);
     docid = put_result.new_document_id;
     ICING_ASSERT_HAS_VALUE_AND_ASSIGN(
         DocumentFilterData filter_data,
         doc_store->GetAliveDocumentFilterData(
             docid, fake_clock_.GetSystemTimeMilliseconds()));
 
-    EXPECT_THAT(filter_data.uri_fingerprint(),
-                Eq(tc3farmhash::Fingerprint64(doc.uri())));
-    ASSERT_THAT(filter_data.schema_type_id(), Eq(0));
+    ASSERT_THAT(filter_data.schema_type_id(), Eq(1));
   }
 
-  // Add another type to the schema.
-  schema =
-      SchemaBuilder()
-          .AddType(SchemaTypeConfigBuilder()
-                       .SetType("alarm")
-                       .AddProperty(PropertyConfigBuilder()
-                                        .SetName("name")
-                                        .SetDataTypeString(TERM_MATCH_EXACT,
-                                                           TOKENIZER_PLAIN)
-                                        .SetCardinality(CARDINALITY_OPTIONAL))
-                       .AddProperty(PropertyConfigBuilder()
-                                        .SetName("time")
-                                        .SetDataType(TYPE_INT64)
-                                        .SetCardinality(CARDINALITY_OPTIONAL)))
-          .AddType(email_type_config)
-          .Build();
+  // Delete the dummy type from the schema.
+  schema = SchemaBuilder().AddType(email_type_config).Build();
   ASSERT_THAT(schema_store->SetSchema(
-                  schema, /*ignore_errors_and_delete_documents=*/false),
+                  schema, /*ignore_errors_and_delete_documents=*/true),
               IsOk());
-  // Adding a new type should cause ids to be reassigned. Ids are assigned in
-  // order of appearance so 'alarm' should be 0 and 'email' should be 1.
-  ASSERT_THAT(schema_store->GetSchemaTypeId("alarm"), IsOkAndHolds(0));
-  ASSERT_THAT(schema_store->GetSchemaTypeId("email"), IsOkAndHolds(1));
+  // Deleting the dummy type will cause email's type id to be reassigned to 0.
+
+  ASSERT_THAT(schema_store->GetSchemaTypeId("email"), IsOkAndHolds(0));
 
   {
     // Create the document store the second time. Don't force recovery.
@@ -4696,9 +6228,7 @@
         DocumentFilterData filter_data,
         doc_store->GetAliveDocumentFilterData(
             docid, fake_clock_.GetSystemTimeMilliseconds()));
-    EXPECT_THAT(filter_data.uri_fingerprint(),
-                Eq(tc3farmhash::Fingerprint64(std::string("email/1"))));
-    ASSERT_THAT(filter_data.schema_type_id(), Eq(0));
+    ASSERT_THAT(filter_data.schema_type_id(), Eq(1));
   }
 }
 
@@ -4730,7 +6260,7 @@
                   schema, /*ignore_errors_and_delete_documents=*/false),
               IsOk());
 
-  DocumentProto docWithBody =
+  DocumentProto doc_with_body =
       DocumentBuilder()
           .SetKey("icing", "email/1")
           .SetSchema("email")
@@ -4741,7 +6271,7 @@
               document1_creation_timestamp_)  // A random timestamp
           .SetTtlMs(document1_ttl_)
           .Build();
-  DocumentProto docWithoutBody =
+  DocumentProto doc_without_body =
       DocumentBuilder()
           .SetKey("icing", "email/2")
           .SetSchema("email")
@@ -4763,26 +6293,28 @@
         std::move(create_result.document_store);
 
     DocumentId docid = kInvalidDocumentId;
-    ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_with_body_result,
-                               doc_store->Put(docWithBody));
+    ICING_ASSERT_OK_AND_ASSIGN(
+        DocumentStore::PutResult put_with_body_result,
+        doc_store->Put(document_util::CreateDocumentWrapper(doc_with_body)));
     EXPECT_THAT(put_with_body_result.old_document_id, Eq(kInvalidDocumentId));
-    EXPECT_FALSE(put_with_body_result.was_replacement());
+    EXPECT_FALSE(put_with_body_result.was_replacement);
     docid = put_with_body_result.new_document_id;
     ASSERT_NE(docid, kInvalidDocumentId);
     docid = kInvalidDocumentId;
-    ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_without_body_result,
-                               doc_store->Put(docWithoutBody));
+    ICING_ASSERT_OK_AND_ASSIGN(
+        DocumentStore::PutResult put_without_body_result,
+        doc_store->Put(document_util::CreateDocumentWrapper(doc_without_body)));
     EXPECT_THAT(put_without_body_result.old_document_id,
                 Eq(kInvalidDocumentId));
-    EXPECT_FALSE(put_without_body_result.was_replacement());
+    EXPECT_FALSE(put_without_body_result.was_replacement);
     docid = put_without_body_result.new_document_id;
     ASSERT_NE(docid, kInvalidDocumentId);
 
-    ASSERT_THAT(doc_store->Get(docWithBody.namespace_(), docWithBody.uri()),
-                IsOkAndHolds(EqualsProto(docWithBody)));
+    ASSERT_THAT(doc_store->Get(doc_with_body.namespace_(), doc_with_body.uri()),
+                IsOkAndHolds(EqualsProto(doc_with_body)));
     ASSERT_THAT(
-        doc_store->Get(docWithoutBody.namespace_(), docWithoutBody.uri()),
-        IsOkAndHolds(EqualsProto(docWithoutBody)));
+        doc_store->Get(doc_without_body.namespace_(), doc_without_body.uri()),
+        IsOkAndHolds(EqualsProto(doc_without_body)));
   }
 
   // Delete the 'body' property from the 'email' type, making all pre-existing
@@ -4805,22 +6337,25 @@
     CorruptDocStoreHeaderChecksumFile();
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
-        DocumentStore::Create(&filesystem_, document_store_dir_, &fake_clock_,
-                              schema_store.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/true,
-                              GetParam().pre_mapping_fbv,
-                              GetParam().use_persistent_hash_map,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, document_store_dir_, &fake_clock_, schema_store.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/true,
+            GetParam().pre_mapping_fbv, GetParam().use_persistent_hash_map,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     std::unique_ptr<DocumentStore> doc_store =
         std::move(create_result.document_store);
 
-    ASSERT_THAT(doc_store->Get(docWithBody.namespace_(), docWithBody.uri()),
+    ASSERT_THAT(doc_store->Get(doc_with_body.namespace_(), doc_with_body.uri()),
                 StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
     ASSERT_THAT(
-        doc_store->Get(docWithoutBody.namespace_(), docWithoutBody.uri()),
-        IsOkAndHolds(EqualsProto(docWithoutBody)));
+        doc_store->Get(doc_without_body.namespace_(), doc_without_body.uri()),
+        IsOkAndHolds(EqualsProto(doc_without_body)));
   }
 }
 
@@ -4852,7 +6387,7 @@
                   schema, /*ignore_errors_and_delete_documents=*/false),
               IsOk());
 
-  DocumentProto docWithBody =
+  DocumentProto doc_with_body =
       DocumentBuilder()
           .SetKey("icing", "email/1")
           .SetSchema("email")
@@ -4863,7 +6398,7 @@
               document1_creation_timestamp_)  // A random timestamp
           .SetTtlMs(document1_ttl_)
           .Build();
-  DocumentProto docWithoutBody =
+  DocumentProto doc_without_body =
       DocumentBuilder()
           .SetKey("icing", "email/2")
           .SetSchema("email")
@@ -4885,26 +6420,28 @@
         std::move(create_result.document_store);
 
     DocumentId docid = kInvalidDocumentId;
-    ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_with_body_result,
-                               doc_store->Put(docWithBody));
+    ICING_ASSERT_OK_AND_ASSIGN(
+        DocumentStore::PutResult put_with_body_result,
+        doc_store->Put(document_util::CreateDocumentWrapper(doc_with_body)));
     EXPECT_THAT(put_with_body_result.old_document_id, Eq(kInvalidDocumentId));
-    EXPECT_FALSE(put_with_body_result.was_replacement());
+    EXPECT_FALSE(put_with_body_result.was_replacement);
     docid = put_with_body_result.new_document_id;
     ASSERT_NE(docid, kInvalidDocumentId);
     docid = kInvalidDocumentId;
-    ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_without_body_result,
-                               doc_store->Put(docWithoutBody));
+    ICING_ASSERT_OK_AND_ASSIGN(
+        DocumentStore::PutResult put_without_body_result,
+        doc_store->Put(document_util::CreateDocumentWrapper(doc_without_body)));
     EXPECT_THAT(put_without_body_result.old_document_id,
                 Eq(kInvalidDocumentId));
-    EXPECT_FALSE(put_without_body_result.was_replacement());
+    EXPECT_FALSE(put_without_body_result.was_replacement);
     docid = put_without_body_result.new_document_id;
     ASSERT_NE(docid, kInvalidDocumentId);
 
-    ASSERT_THAT(doc_store->Get(docWithBody.namespace_(), docWithBody.uri()),
-                IsOkAndHolds(EqualsProto(docWithBody)));
+    ASSERT_THAT(doc_store->Get(doc_with_body.namespace_(), doc_with_body.uri()),
+                IsOkAndHolds(EqualsProto(doc_with_body)));
     ASSERT_THAT(
-        doc_store->Get(docWithoutBody.namespace_(), docWithoutBody.uri()),
-        IsOkAndHolds(EqualsProto(docWithoutBody)));
+        doc_store->Get(doc_without_body.namespace_(), doc_without_body.uri()),
+        IsOkAndHolds(EqualsProto(doc_without_body)));
   }
 
   // Delete the 'body' property from the 'email' type, making all pre-existing
@@ -4933,11 +6470,11 @@
     std::unique_ptr<DocumentStore> doc_store =
         std::move(create_result.document_store);
 
-    ASSERT_THAT(doc_store->Get(docWithBody.namespace_(), docWithBody.uri()),
-                IsOkAndHolds(EqualsProto(docWithBody)));
+    ASSERT_THAT(doc_store->Get(doc_with_body.namespace_(), doc_with_body.uri()),
+                IsOkAndHolds(EqualsProto(doc_with_body)));
     ASSERT_THAT(
-        doc_store->Get(docWithoutBody.namespace_(), docWithoutBody.uri()),
-        IsOkAndHolds(EqualsProto(docWithoutBody)));
+        doc_store->Get(doc_without_body.namespace_(), doc_without_body.uri()),
+        IsOkAndHolds(EqualsProto(doc_without_body)));
   }
 }
 
@@ -5012,7 +6549,9 @@
           /*force_recovery_and_revalidate_documents=*/false,
           GetParam().pre_mapping_fbv, GetParam().use_persistent_hash_map,
           PortableFileBackedProtoLog<DocumentWrapper>::kDefaultCompressionLevel,
-          &initialize_stats));
+          PortableFileBackedProtoLog<
+              DocumentWrapper>::kDefaultCompressionThresholdBytes,
+          protobuf_ports::kDefaultMemLevel, &initialize_stats));
   std::unique_ptr<DocumentStore> document_store =
       std::move(create_result.document_store);
 
@@ -5116,7 +6655,8 @@
                                 .AddStringProperty("body", "dd ee")
                                 .SetCreationTimestampMs(1)
                                 .Build();
-  ICING_ASSERT_OK(document_store->Put(document1, 5));
+  ICING_ASSERT_OK(document_store->Put(
+      CreateDocumentWrapper(document1, /*num_string_tokens=*/5)));
 
   DocumentProto document2 = DocumentBuilder()
                                 .SetKey("namespace2", "email/2")
@@ -5125,7 +6665,8 @@
                                 .AddStringProperty("body", "cc")
                                 .SetCreationTimestampMs(1)
                                 .Build();
-  ICING_ASSERT_OK(document_store->Put(document2, 3));
+  ICING_ASSERT_OK(document_store->Put(
+      CreateDocumentWrapper(document2, /*num_string_tokens=*/3)));
 
   DocumentProto document3 = DocumentBuilder()
                                 .SetKey("namespace2", "email/3")
@@ -5134,7 +6675,8 @@
                                 .AddStringProperty("body", "")
                                 .SetCreationTimestampMs(1)
                                 .Build();
-  ICING_ASSERT_OK(document_store->Put(document3, 1));
+  ICING_ASSERT_OK(document_store->Put(
+      CreateDocumentWrapper(document3, /*num_string_tokens=*/1)));
 
   DocumentProto document4 = DocumentBuilder()
                                 .SetKey("namespace1", "person/1")
@@ -5142,7 +6684,8 @@
                                 .AddStringProperty("name", "test test")
                                 .SetCreationTimestampMs(1)
                                 .Build();
-  ICING_ASSERT_OK(document_store->Put(document4, 2));
+  ICING_ASSERT_OK(document_store->Put(
+      CreateDocumentWrapper(document4, /*num_string_tokens=*/2)));
 
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentDebugInfoProto out1,
@@ -5246,21 +6789,25 @@
   {
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
-        DocumentStore::Create(&filesystem_, document_store_dir_, &fake_clock_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              GetParam().pre_mapping_fbv,
-                              GetParam().use_persistent_hash_map,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, document_store_dir_, &fake_clock_,
+            schema_store_.get(), feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            GetParam().pre_mapping_fbv, GetParam().use_persistent_hash_map,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
 
     std::unique_ptr<DocumentStore> doc_store =
         std::move(create_result.document_store);
-    ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                               doc_store->Put(test_document1_));
+    ICING_ASSERT_OK_AND_ASSIGN(
+        DocumentStore::PutResult put_result1,
+        doc_store->Put(document_util::CreateDocumentWrapper(test_document1_)));
     EXPECT_THAT(put_result1.old_document_id, Eq(kInvalidDocumentId));
-    EXPECT_FALSE(put_result1.was_replacement());
+    EXPECT_FALSE(put_result1.was_replacement);
     document_id1 = put_result1.new_document_id;
 
     if (GetParam().use_persistent_hash_map) {
@@ -5295,7 +6842,9 @@
             /*use_persistent_hash_map=*/switch_key_mapper_flag,
             PortableFileBackedProtoLog<
                 DocumentWrapper>::kDefaultCompressionLevel,
-            &initialize_stats));
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel, &initialize_stats));
     EXPECT_THAT(initialize_stats.document_store_recovery_cause(),
                 Eq(InitializeStatsProto::IO_ERROR));
 
@@ -5332,21 +6881,25 @@
   {
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
-        DocumentStore::Create(&filesystem_, document_store_dir_, &fake_clock_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              GetParam().pre_mapping_fbv,
-                              GetParam().use_persistent_hash_map,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, document_store_dir_, &fake_clock_,
+            schema_store_.get(), feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            GetParam().pre_mapping_fbv, GetParam().use_persistent_hash_map,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
 
     std::unique_ptr<DocumentStore> doc_store =
         std::move(create_result.document_store);
-    ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                               doc_store->Put(test_document1_));
+    ICING_ASSERT_OK_AND_ASSIGN(
+        DocumentStore::PutResult put_result1,
+        doc_store->Put(document_util::CreateDocumentWrapper(test_document1_)));
     EXPECT_THAT(put_result1.old_document_id, Eq(kInvalidDocumentId));
-    EXPECT_FALSE(put_result1.was_replacement());
+    EXPECT_FALSE(put_result1.was_replacement);
     document_id1 = put_result1.new_document_id;
 
     if (GetParam().use_persistent_hash_map) {
@@ -5371,14 +6924,16 @@
     InitializeStatsProto initialize_stats;
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
-        DocumentStore::Create(&filesystem_, document_store_dir_, &fake_clock_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              GetParam().pre_mapping_fbv,
-                              GetParam().use_persistent_hash_map,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              &initialize_stats));
+        DocumentStore::Create(
+            &filesystem_, document_store_dir_, &fake_clock_,
+            schema_store_.get(), feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            GetParam().pre_mapping_fbv, GetParam().use_persistent_hash_map,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel, &initialize_stats));
     EXPECT_THAT(initialize_stats.document_store_recovery_cause(),
                 Eq(InitializeStatsProto::NONE));
 
@@ -5406,11 +6961,7 @@
   }
 }
 
-TEST_P(DocumentStoreTest, GetDocumentIdByNamespaceIdFingerprint) {
-  std::string dynamic_trie_uri_mapper_dir =
-      document_store_dir_ + "/key_mapper_dir";
-  std::string persistent_hash_map_uri_mapper_dir =
-      document_store_dir_ + "/uri_mapper";
+TEST_P(DocumentStoreTest, GetDocumentId_expiredDocument) {
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentStore::CreateResult create_result,
       DocumentStore::Create(
@@ -5419,14 +6970,93 @@
           /*force_recovery_and_revalidate_documents=*/false,
           GetParam().pre_mapping_fbv, GetParam().use_persistent_hash_map,
           PortableFileBackedProtoLog<DocumentWrapper>::kDefaultCompressionLevel,
+          PortableFileBackedProtoLog<
+              DocumentWrapper>::kDefaultCompressionThresholdBytes,
+          protobuf_ports::kDefaultMemLevel,
+          /*initialize_stats=*/nullptr));
+  std::unique_ptr<DocumentStore> doc_store =
+      std::move(create_result.document_store);
+
+  fake_clock_.SetSystemTimeMilliseconds(0);
+  DocumentProto foo_document = DocumentBuilder()
+                                   .SetCreationTimestampMs(0)
+                                   .SetTtlMs(1000)
+                                   .SetKey("namespace", "uri")
+                                   .SetSchema("email")
+                                   .SetCreationTimestampMs(0)
+                                   .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store->Put(document_util::CreateDocumentWrapper(foo_document)));
+  EXPECT_THAT(put_result.old_document_id, Eq(kInvalidDocumentId));
+  EXPECT_FALSE(put_result.was_replacement);
+  DocumentId document_id = put_result.new_document_id;
+
+  // Adjust the clock to make the document expired. GetDocumentId should still
+  // return the original document id.
+  fake_clock_.SetSystemTimeMilliseconds(2000);
+  EXPECT_THAT(doc_store->GetDocumentId("namespace", "uri"),
+              IsOkAndHolds(document_id));
+}
+
+TEST_P(DocumentStoreTest, GetDocumentId_deletedDocument) {
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      DocumentStore::Create(
+          &filesystem_, document_store_dir_, &fake_clock_, schema_store_.get(),
+          feature_flags_.get(),
+          /*force_recovery_and_revalidate_documents=*/false,
+          GetParam().pre_mapping_fbv, GetParam().use_persistent_hash_map,
+          PortableFileBackedProtoLog<DocumentWrapper>::kDefaultCompressionLevel,
+          PortableFileBackedProtoLog<
+              DocumentWrapper>::kDefaultCompressionThresholdBytes,
+          protobuf_ports::kDefaultMemLevel,
+          /*initialize_stats=*/nullptr));
+  std::unique_ptr<DocumentStore> doc_store =
+      std::move(create_result.document_store);
+
+  DocumentProto foo_document = DocumentBuilder()
+                                   .SetKey("namespace", "uri")
+                                   .SetSchema("email")
+                                   .SetCreationTimestampMs(0)
+                                   .Build();
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store->Put(document_util::CreateDocumentWrapper(foo_document)));
+  EXPECT_THAT(put_result.old_document_id, Eq(kInvalidDocumentId));
+  EXPECT_FALSE(put_result.was_replacement);
+  DocumentId document_id = put_result.new_document_id;
+
+  // Delete the document. GetDocumentId should still return the original
+  // document id.
+  ICING_ASSERT_OK(doc_store->Delete(
+      document_id,
+      /*current_time_ms=*/fake_clock_.GetSystemTimeMilliseconds()));
+  EXPECT_THAT(doc_store->GetDocumentId("namespace", "uri"),
+              IsOkAndHolds(document_id));
+}
+
+TEST_P(DocumentStoreTest, GetDocumentIdByNamespaceIdFingerprint) {
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      DocumentStore::Create(
+          &filesystem_, document_store_dir_, &fake_clock_, schema_store_.get(),
+          feature_flags_.get(),
+          /*force_recovery_and_revalidate_documents=*/false,
+          GetParam().pre_mapping_fbv, GetParam().use_persistent_hash_map,
+          PortableFileBackedProtoLog<DocumentWrapper>::kDefaultCompressionLevel,
+          PortableFileBackedProtoLog<
+              DocumentWrapper>::kDefaultCompressionThresholdBytes,
+          protobuf_ports::kDefaultMemLevel,
           /*initialize_stats=*/nullptr));
 
   std::unique_ptr<DocumentStore> doc_store =
       std::move(create_result.document_store);
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             doc_store->Put(test_document1_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store->Put(document_util::CreateDocumentWrapper(test_document1_)));
   EXPECT_THAT(put_result.old_document_id, Eq(kInvalidDocumentId));
-  EXPECT_FALSE(put_result.was_replacement());
+  EXPECT_FALSE(put_result.was_replacement);
   DocumentId document_id = put_result.new_document_id;
 
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -5477,8 +7107,9 @@
                                .AddStringProperty("subject", "subject foo")
                                .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             doc_store->Put(document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store->Put(document_util::CreateDocumentWrapper(document)));
   DocumentId document_id = put_result.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentAssociatedScoreData score_data,
@@ -5503,8 +7134,9 @@
                                .AddStringProperty("subject", "subject foo")
                                .Build();
 
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             doc_store->Put(document));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store->Put(document_util::CreateDocumentWrapper(document)));
   DocumentId document_id = put_result.new_document_id;
   std::unique_ptr<ScorablePropertySet> scorable_property_set =
       doc_store->GetScorablePropertySet(
@@ -5579,11 +7211,13 @@
                                 .AddInt64Property("scoreInt64", 5)
                                 .SetCreationTimestampMs(0)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             doc_store->Put(document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      doc_store->Put(document_util::CreateDocumentWrapper(document1)));
   DocumentId document_id1 = put_result1.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             doc_store->Put(document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      doc_store->Put(document_util::CreateDocumentWrapper(document2)));
   DocumentId document_id2 = put_result2.new_document_id;
 
   std::unique_ptr<ScorablePropertySet> scorable_property_set_doc1 =
@@ -5625,8 +7259,9 @@
                                 .Build();
   // Add document3 to the document store, it will result in document1 being
   // deleted.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
-                             doc_store->Put(document3));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result3,
+      doc_store->Put(document_util::CreateDocumentWrapper(document3)));
   DocumentId document_id3 = put_result3.new_document_id;
   std::unique_ptr<ScorablePropertySet> scorable_property_set_doc3 =
       doc_store->GetScorablePropertySet(
@@ -5654,8 +7289,9 @@
                                 .Build();
   // Add document4 to the document store, it will result in document3 being
   // deleted.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result4,
-                             doc_store->Put(document4));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result4,
+      doc_store->Put(document_util::CreateDocumentWrapper(document4)));
   DocumentId document_id4 = put_result4.new_document_id;
   std::unique_ptr<ScorablePropertySet> scorable_property_set_doc4 =
       doc_store->GetScorablePropertySet(
@@ -5728,8 +7364,9 @@
                                 .AddDoubleProperty("scoreDouble", 1.5, 2.5)
                                 .SetCreationTimestampMs(0)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
-                             doc_store->Put(document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result1,
+      doc_store->Put(document_util::CreateDocumentWrapper(document1)));
   DocumentId document_id1 = put_result1.new_document_id;
 
   std::unique_ptr<ScorablePropertySet> scorable_property_set_doc1 =
@@ -5760,8 +7397,9 @@
 
   // Add document2 to the document store, it will result in document1 being
   // deleted.
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
-                             doc_store->Put(document2));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result2,
+      doc_store->Put(document_util::CreateDocumentWrapper(document2)));
   DocumentId document_id2 = put_result2.new_document_id;
   std::unique_ptr<ScorablePropertySet> scorable_property_set_doc2 =
       doc_store->GetScorablePropertySet(
@@ -5840,8 +7478,9 @@
                                 .SetCreationTimestampMs(1)
                                 .AddDoubleProperty("income", 10000, 20000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             doc_store->Put(document0));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store->Put(document_util::CreateDocumentWrapper(document0)));
   DocumentId document_id = put_result.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentAssociatedScoreData score_data,
@@ -5914,8 +7553,9 @@
                                 .SetCreationTimestampMs(1)
                                 .AddDoubleProperty("income", 10000, 20000)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             doc_store->Put(document0));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store->Put(document_util::CreateDocumentWrapper(document0)));
   DocumentId document_id = put_result.new_document_id;
   ICING_ASSERT_OK_AND_ASSIGN(
       DocumentAssociatedScoreData score_data,
@@ -5999,10 +7639,13 @@
                                 .SetCreationTimestampMs(1)
                                 .AddDoubleProperty("score", 10, 20)
                                 .Build();
-  ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result,
-                             doc_store->Put(document0));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::PutResult put_result,
+      doc_store->Put(document_util::CreateDocumentWrapper(document0)));
   DocumentId document_id0 = put_result.new_document_id;
-  ICING_ASSERT_OK_AND_ASSIGN(put_result, doc_store->Put(document1));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      put_result,
+      doc_store->Put(document_util::CreateDocumentWrapper(document1)));
   DocumentId document_id1 = put_result.new_document_id;
 
   // Update the schema by rearranging the schema types. Since SchemaTypeId is
@@ -6047,6 +7690,445 @@
       Pointee(EqualsProto(BuildScorablePropertyProtoFromDouble({10, 20}))));
 }
 
+TEST_P(DocumentStoreTest, GetResultGroupingEntryId_getByFilterName) {
+  // Put 2 schema types into the schema store.
+  SchemaProto schema =
+      SchemaBuilder()
+          .AddType(SchemaTypeConfigBuilder().SetType("Email"))
+          .AddType(SchemaTypeConfigBuilder().SetType("Message"))
+          .Build();
+
+  std::string schema_store_dir = schema_store_dir_ + "_custom";
+  filesystem_.DeleteDirectoryRecursively(schema_store_dir.c_str());
+  filesystem_.CreateDirectoryRecursively(schema_store_dir.c_str());
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<SchemaStore> schema_store,
+      SchemaStore::Create(&filesystem_, schema_store_dir, &fake_clock_,
+                          feature_flags_.get()));
+
+  ICING_ASSERT_OK(schema_store->SetSchema(
+      std::move(schema), /*ignore_errors_and_delete_documents=*/false));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
+                          schema_store.get()));
+  std::unique_ptr<DocumentStore> document_store =
+      std::move(create_result.document_store);
+
+  // Put 3 documents into the document store to create 3 different namespaces.
+  DocumentProto document0 = DocumentBuilder()
+                                .SetKey("namespace0", "uri/0")
+                                .SetSchema("Email")
+                                .Build();
+  DocumentProto document1 = DocumentBuilder()
+                                .SetKey("namespace1", "uri/1")
+                                .SetSchema("Message")
+                                .Build();
+  DocumentProto document2 = DocumentBuilder()
+                                .SetKey("namespace2", "uri/2")
+                                .SetSchema("Message")
+                                .Build();
+  ICING_ASSERT_OK(document_store->Put(
+      document_util::CreateDocumentWrapper(std::move(document0))));
+  ICING_ASSERT_OK(document_store->Put(
+      document_util::CreateDocumentWrapper(std::move(document1))));
+  ICING_ASSERT_OK(document_store->Put(
+      document_util::CreateDocumentWrapper(std::move(document2))));
+
+  ASSERT_THAT(document_store->GetNamespaceId("namespace0"), IsOkAndHolds(0));
+  ASSERT_THAT(document_store->GetNamespaceId("namespace1"), IsOkAndHolds(1));
+  ASSERT_THAT(document_store->GetNamespaceId("namespace2"), IsOkAndHolds(2));
+
+  ASSERT_THAT(schema_store->GetSchemaTypeId("Email"), IsOkAndHolds(0));
+  ASSERT_THAT(schema_store->GetSchemaTypeId("Message"), IsOkAndHolds(1));
+
+  // NONE should always return std::nullopt.
+  EXPECT_THAT(
+      document_store->GetResultGroupingEntryId(
+          ResultSpecProto_ResultGroupingType_NONE, "namespace0", "Email"),
+      IsFalse());
+  EXPECT_THAT(
+      document_store->GetResultGroupingEntryId(
+          ResultSpecProto_ResultGroupingType_NONE, "namespace1", "Email"),
+      IsFalse());
+  EXPECT_THAT(
+      document_store->GetResultGroupingEntryId(
+          ResultSpecProto_ResultGroupingType_NONE, "namespace2", "Email"),
+      IsFalse());
+  EXPECT_THAT(
+      document_store->GetResultGroupingEntryId(
+          ResultSpecProto_ResultGroupingType_NONE, "namespace0", "Message"),
+      IsFalse());
+  EXPECT_THAT(
+      document_store->GetResultGroupingEntryId(
+          ResultSpecProto_ResultGroupingType_NONE, "namespace1", "Message"),
+      IsFalse());
+  EXPECT_THAT(
+      document_store->GetResultGroupingEntryId(
+          ResultSpecProto_ResultGroupingType_NONE, "namespace2", "Message"),
+      IsFalse());
+
+  // SCHEMA_TYPE should return id based on the schema type and ignore the
+  // namespace.
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_SCHEMA_TYPE, "namespace0",
+                  "Email"),
+              Optional(Eq(0)));
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_SCHEMA_TYPE, "namespace1",
+                  "Email"),
+              Optional(Eq(0)));
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_SCHEMA_TYPE, "namespace2",
+                  "Email"),
+              Optional(Eq(0)));
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_SCHEMA_TYPE, "namespace0",
+                  "Message"),
+              Optional(Eq(1)));
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_SCHEMA_TYPE, "namespace1",
+                  "Message"),
+              Optional(Eq(1)));
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_SCHEMA_TYPE, "namespace2",
+                  "Message"),
+              Optional(Eq(1)));
+
+  // NAMESPACE should return id based on the namespace and ignore the schema
+  // type.
+  EXPECT_THAT(
+      document_store->GetResultGroupingEntryId(
+          ResultSpecProto_ResultGroupingType_NAMESPACE, "namespace0", "Email"),
+      Optional(Eq(0)));
+  EXPECT_THAT(
+      document_store->GetResultGroupingEntryId(
+          ResultSpecProto_ResultGroupingType_NAMESPACE, "namespace1", "Email"),
+      Optional(Eq(1)));
+  EXPECT_THAT(
+      document_store->GetResultGroupingEntryId(
+          ResultSpecProto_ResultGroupingType_NAMESPACE, "namespace2", "Email"),
+      Optional(Eq(2)));
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE, "namespace0",
+                  "Message"),
+              Optional(Eq(0)));
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE, "namespace1",
+                  "Message"),
+              Optional(Eq(1)));
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE, "namespace2",
+                  "Message"),
+              Optional(Eq(2)));
+
+  // NAMESPACE_AND_SCHEMA_TYPE should return id based on both namespace and
+  // schema type.
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE_AND_SCHEMA_TYPE,
+                  "namespace0", "Email"),
+              Optional(Eq(0x00000000)));
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE_AND_SCHEMA_TYPE,
+                  "namespace1", "Email"),
+              Optional(Eq(0x00010000)));
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE_AND_SCHEMA_TYPE,
+                  "namespace2", "Email"),
+              Optional(Eq(0x00020000)));
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE_AND_SCHEMA_TYPE,
+                  "namespace0", "Message"),
+              Optional(Eq(0x00000001)));
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE_AND_SCHEMA_TYPE,
+                  "namespace1", "Message"),
+              Optional(Eq(0x00010001)));
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE_AND_SCHEMA_TYPE,
+                  "namespace2", "Message"),
+              Optional(Eq(0x00020001)));
+}
+
+TEST_P(DocumentStoreTest, GetResultGroupingEntryId_getByNonExistingFilterName) {
+  // Put 1 schema type into the schema store.
+  SchemaProto schema = SchemaBuilder()
+                           .AddType(SchemaTypeConfigBuilder().SetType("Email"))
+                           .Build();
+
+  std::string schema_store_dir = schema_store_dir_ + "_custom";
+  filesystem_.DeleteDirectoryRecursively(schema_store_dir.c_str());
+  filesystem_.CreateDirectoryRecursively(schema_store_dir.c_str());
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<SchemaStore> schema_store,
+      SchemaStore::Create(&filesystem_, schema_store_dir, &fake_clock_,
+                          feature_flags_.get()));
+
+  ICING_ASSERT_OK(schema_store->SetSchema(
+      std::move(schema), /*ignore_errors_and_delete_documents=*/false));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
+                          schema_store.get()));
+  std::unique_ptr<DocumentStore> document_store =
+      std::move(create_result.document_store);
+
+  // Put 1 document into the document store to create 1 namespace.
+  DocumentProto document0 = DocumentBuilder()
+                                .SetKey("namespace0", "uri/0")
+                                .SetSchema("Email")
+                                .Build();
+  ICING_ASSERT_OK(document_store->Put(
+      document_util::CreateDocumentWrapper(std::move(document0))));
+
+  ASSERT_THAT(document_store->GetNamespaceId("namespace0"), IsOkAndHolds(0));
+  ASSERT_THAT(schema_store->GetSchemaTypeId("Email"), IsOkAndHolds(0));
+
+  // NONE should always return std::nullopt.
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NONE,
+                  "nonExistingNamespace", "nonExistingSchemaType"),
+              IsFalse());
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NONE,
+                  "nonExistingNamespace", "Email"),
+              IsFalse());
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NONE, "namespace0",
+                  "nonExistingSchemaType"),
+              IsFalse());
+
+  // SCHEMA_TYPE should return id based on the schema type and ignore the
+  // namespace. It is ok that the namespace does not exist.
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_SCHEMA_TYPE,
+                  "nonExistingNamespace", "nonExistingSchemaType"),
+              IsFalse());
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_SCHEMA_TYPE,
+                  "nonExistingNamespace", "Email"),
+              Optional(Eq(0)));
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_SCHEMA_TYPE, "namespace0",
+                  "nonExistingSchemaType"),
+              IsFalse());
+
+  // NAMESPACE should return id based on the namespace and ignore the schema
+  // type. It is ok that the schema type does not exist.
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE,
+                  "nonExistingNamespace", "nonExistingSchemaType"),
+              IsFalse());
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE,
+                  "nonExistingNamespace", "Email"),
+              IsFalse());
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE, "namespace0",
+                  "nonExistingSchemaType"),
+              Optional(Eq(0)));
+
+  // NAMESPACE_AND_SCHEMA_TYPE should return id based on both namespace and
+  // schema type. Both namespace and schema type must exist.
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE_AND_SCHEMA_TYPE,
+                  "nonExistingNamespace", "nonExistingSchemaType"),
+              IsFalse());
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE_AND_SCHEMA_TYPE,
+                  "nonExistingNamespace", "Email"),
+              IsFalse());
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE_AND_SCHEMA_TYPE,
+                  "namespace0", "nonExistingSchemaType"),
+              IsFalse());
+}
+
+TEST_P(DocumentStoreTest, GetResultGroupingEntryId_getByIds) {
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
+                          schema_store_.get()));
+  std::unique_ptr<DocumentStore> document_store =
+      std::move(create_result.document_store);
+
+  // GetResultGroupingEntryId() by id only handles the encoding and won't check
+  // if the id exists (except kInvalidNamespaceId and kInvalidSchemaTypeId), so
+  // we don't need to set up schema types and namespaces here.
+
+  // NONE should always return std::nullopt.
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NONE, /*namespace_id=*/0,
+                  /*schema_type_id=*/0),
+              IsFalse());
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NONE, /*namespace_id=*/0,
+                  /*schema_type_id=*/1),
+              IsFalse());
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NONE, /*namespace_id=*/1,
+                  /*schema_type_id=*/0),
+              IsFalse());
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NONE, /*namespace_id=*/1,
+                  /*schema_type_id=*/1),
+              IsFalse());
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NONE,
+                  /*namespace_id=*/std::numeric_limits<NamespaceId>::max(),
+                  /*schema_type_id=*/std::numeric_limits<SchemaTypeId>::max()),
+              IsFalse());
+
+  // SCHEMA_TYPE should return id based on the schema type id and ignore the
+  // namespace id.
+  EXPECT_THAT(
+      document_store->GetResultGroupingEntryId(
+          ResultSpecProto_ResultGroupingType_SCHEMA_TYPE, /*namespace_id=*/0,
+          /*schema_type_id=*/0),
+      Optional(Eq(0)));
+  EXPECT_THAT(
+      document_store->GetResultGroupingEntryId(
+          ResultSpecProto_ResultGroupingType_SCHEMA_TYPE, /*namespace_id=*/0,
+          /*schema_type_id=*/1),
+      Optional(Eq(1)));
+  EXPECT_THAT(
+      document_store->GetResultGroupingEntryId(
+          ResultSpecProto_ResultGroupingType_SCHEMA_TYPE, /*namespace_id=*/1,
+          /*schema_type_id=*/0),
+      Optional(Eq(0)));
+  EXPECT_THAT(
+      document_store->GetResultGroupingEntryId(
+          ResultSpecProto_ResultGroupingType_SCHEMA_TYPE, /*namespace_id=*/1,
+          /*schema_type_id=*/1),
+      Optional(Eq(1)));
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_SCHEMA_TYPE,
+                  /*namespace_id=*/std::numeric_limits<NamespaceId>::max(),
+                  /*schema_type_id=*/std::numeric_limits<SchemaTypeId>::max()),
+              Optional(Eq(std::numeric_limits<SchemaTypeId>::max())));
+
+  // NAMESPACE should return id based on the namespace id and ignore the schema
+  // type id.
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE,
+                  /*namespace_id=*/0, /*schema_type_id=*/0),
+              Optional(Eq(0)));
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE,
+                  /*namespace_id=*/0, /*schema_type_id=*/1),
+              Optional(Eq(0)));
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE,
+                  /*namespace_id=*/1, /*schema_type_id=*/0),
+              Optional(Eq(1)));
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE,
+                  /*namespace_id=*/1, /*schema_type_id=*/1),
+              Optional(Eq(1)));
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE,
+                  /*namespace_id=*/std::numeric_limits<NamespaceId>::max(),
+                  /*schema_type_id=*/std::numeric_limits<SchemaTypeId>::max()),
+              Optional(Eq(std::numeric_limits<NamespaceId>::max())));
+
+  // NAMESPACE_AND_SCHEMA_TYPE should return id based on both namespace and
+  // schema type ids.
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE_AND_SCHEMA_TYPE,
+                  /*namespace_id=*/0, /*schema_type_id=*/0),
+              Optional(Eq(0x00000000)));
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE_AND_SCHEMA_TYPE,
+                  /*namespace_id=*/0, /*schema_type_id=*/1),
+              Optional(Eq(0x00000001)));
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE_AND_SCHEMA_TYPE,
+                  /*namespace_id=*/1, /*schema_type_id=*/0),
+              Optional(Eq(0x00010000)));
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE_AND_SCHEMA_TYPE,
+                  /*namespace_id=*/1, /*schema_type_id=*/1),
+              Optional(Eq(0x00010001)));
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE_AND_SCHEMA_TYPE,
+                  /*namespace_id=*/std::numeric_limits<NamespaceId>::max(),
+                  /*schema_type_id=*/std::numeric_limits<SchemaTypeId>::max()),
+              Optional(Eq(0x7fff7fff)));
+}
+
+TEST_P(DocumentStoreTest, GetResultGroupingEntryId_getByInvalidIds) {
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentStore::CreateResult create_result,
+      CreateDocumentStore(&filesystem_, document_store_dir_, &fake_clock_,
+                          schema_store_.get()));
+  std::unique_ptr<DocumentStore> document_store =
+      std::move(create_result.document_store);
+
+  // GetResultGroupingEntryId() by id only handles the encoding and won't check
+  // if the id exists (except kInvalidNamespaceId and kInvalidSchemaTypeId), so
+  // we don't need to set up schema types and namespaces here.
+
+  // NONE should always return std::nullopt.
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NONE, kInvalidNamespaceId,
+                  kInvalidSchemaTypeId),
+              IsFalse());
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NONE, kInvalidNamespaceId,
+                  /*schema_type_id=*/0),
+              IsFalse());
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NONE, /*namespace_id=*/0,
+                  kInvalidSchemaTypeId),
+              IsFalse());
+
+  // SCHEMA_TYPE should return id based on the schema type id and ignore the
+  // namespace id. It is ok that the namespace id is invalid.
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_SCHEMA_TYPE,
+                  kInvalidNamespaceId, kInvalidSchemaTypeId),
+              IsFalse());
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_SCHEMA_TYPE,
+                  kInvalidNamespaceId, /*schema_type_id=*/0),
+              Optional(Eq(0)));
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_SCHEMA_TYPE,
+                  /*namespace_id=*/0, kInvalidSchemaTypeId),
+              IsFalse());
+
+  // NAMESPACE should return id based on the namespace id and ignore the schema
+  // type id. It is ok that the schema type id is invalid.
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE,
+                  kInvalidNamespaceId, kInvalidSchemaTypeId),
+              IsFalse());
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE,
+                  kInvalidNamespaceId, /*schema_type_id=*/0),
+              IsFalse());
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE,
+                  /*namespace_id=*/0, kInvalidSchemaTypeId),
+              Optional(Eq(0)));
+  // NAMESPACE_AND_SCHEMA_TYPE should return id based on both namespace and
+  // schema type ids. Both namespace and schema type ids must be valid.
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE_AND_SCHEMA_TYPE,
+                  kInvalidNamespaceId, kInvalidSchemaTypeId),
+              IsFalse());
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE_AND_SCHEMA_TYPE,
+                  kInvalidNamespaceId, /*schema_type_id=*/0),
+              IsFalse());
+  EXPECT_THAT(document_store->GetResultGroupingEntryId(
+                  ResultSpecProto_ResultGroupingType_NAMESPACE_AND_SCHEMA_TYPE,
+                  /*namespace_id=*/0, kInvalidSchemaTypeId),
+              IsFalse());
+}
+
 INSTANTIATE_TEST_SUITE_P(
     DocumentStoreTest, DocumentStoreTest,
     testing::Values(

diff --git a/icing/store/dynamic-trie-key-mapper.h b/icing/store/dynamic-trie-key-mapper.h
index f7fd665..e12b44e 100644
--- a/icing/store/dynamic-trie-key-mapper.h
+++ b/icing/store/dynamic-trie-key-mapper.h

@@ -81,7 +81,7 @@
 
   libtextclassifier3::StatusOr<T> Get(std::string_view key) const override;
 
-  bool Delete(std::string_view key) override;
+  libtextclassifier3::Status Delete(std::string_view key) override;
 
   std::unique_ptr<typename KeyMapper<T, Formatter>::Iterator> GetIterator()
       const override;
@@ -277,7 +277,8 @@
 }
 
 template <typename T, typename Formatter>
-bool DynamicTrieKeyMapper<T, Formatter>::Delete(std::string_view key) {
+libtextclassifier3::Status DynamicTrieKeyMapper<T, Formatter>::Delete(
+    std::string_view key) {
   return trie_.Delete(key);
 }
 

diff --git a/icing/store/key-mapper.h b/icing/store/key-mapper.h
index 83f725c..c7d5bb0 100644
--- a/icing/store/key-mapper.h
+++ b/icing/store/key-mapper.h

@@ -80,8 +80,8 @@
   // Returns any encountered IO errors.
   virtual libtextclassifier3::StatusOr<T> Get(std::string_view key) const = 0;
 
-  // Deletes data related to the given key. Returns true on success.
-  virtual bool Delete(std::string_view key) = 0;
+  // Deletes data related to the given key. Returns OK on success.
+  virtual libtextclassifier3::Status Delete(std::string_view key) = 0;
 
   // Returns an iterator of the key mapper.
   //

diff --git a/icing/store/key-mapper_test.cc b/icing/store/key-mapper_test.cc
index 229e519..f75fc37 100644
--- a/icing/store/key-mapper_test.cc
+++ b/icing/store/key-mapper_test.cc

@@ -235,7 +235,7 @@
               UnorderedElementsAre(Pair("foo", 1)));
 
   // Delete "foo".
-  EXPECT_THAT(key_mapper->Delete("foo"), IsTrue());
+  ICING_EXPECT_OK(key_mapper->Delete("foo"));
 
   // Verify num_keys().
   EXPECT_THAT(key_mapper->num_keys(), 0);
@@ -269,7 +269,7 @@
       UnorderedElementsAre(Pair("foo", 1), Pair("bar", 2), Pair("baz", 3)));
 
   // Delete "foo".
-  EXPECT_THAT(key_mapper->Delete("foo"), IsTrue());
+  ICING_EXPECT_OK(key_mapper->Delete("foo"));
 
   // Verify num_keys().
   EXPECT_THAT(key_mapper->num_keys(), 2);
@@ -310,9 +310,9 @@
       UnorderedElementsAre(Pair("foo", 1), Pair("bar", 2), Pair("baz", 3)));
 
   // Delete all keys.
-  EXPECT_THAT(key_mapper->Delete("foo"), IsTrue());
-  EXPECT_THAT(key_mapper->Delete("bar"), IsTrue());
-  EXPECT_THAT(key_mapper->Delete("baz"), IsTrue());
+  ICING_EXPECT_OK(key_mapper->Delete("foo"));
+  ICING_EXPECT_OK(key_mapper->Delete("bar"));
+  ICING_EXPECT_OK(key_mapper->Delete("baz"));
 
   // Verify num_keys().
   EXPECT_THAT(key_mapper->num_keys(), 0);

diff --git a/icing/store/persistent-hash-map-key-mapper.h b/icing/store/persistent-hash-map-key-mapper.h
index e8b5832..21381e4 100644
--- a/icing/store/persistent-hash-map-key-mapper.h
+++ b/icing/store/persistent-hash-map-key-mapper.h

@@ -108,8 +108,8 @@
     return value;
   }
 
-  bool Delete(std::string_view key) override {
-    return persistent_hash_map_->Delete(key).ok();
+  libtextclassifier3::Status Delete(std::string_view key) override {
+    return persistent_hash_map_->Delete(key);
   }
 
   std::unique_ptr<typename KeyMapper<T, Formatter>::Iterator> GetIterator()

diff --git a/icing/testing/common-matchers.h b/icing/testing/common-matchers.h
index d730c83..8e94a7c 100644
--- a/icing/testing/common-matchers.h
+++ b/icing/testing/common-matchers.h

@@ -93,27 +93,47 @@
 }
 
 // Used to match a DocHitInfoIterator::CallStats
-MATCHER_P5(EqualsDocHitInfoIteratorCallStats, num_leaf_advance_calls_lite_index,
+MATCHER_P6(EqualsDocHitInfoIteratorCallStats, num_leaf_advance_calls_lite_index,
            num_leaf_advance_calls_main_index,
            num_leaf_advance_calls_integer_index,
-           num_leaf_advance_calls_no_index, num_blocks_inspected, "") {
+           num_leaf_advance_calls_no_index, num_blocks_inspected,
+           embedding_stats, "") {
   const DocHitInfoIterator::CallStats& actual = arg;
+  auto embedding_stats_to_string =
+      [](const DocHitInfoIterator::CallStats::EmbeddingStats& stats) {
+        return absl_ports::StrCat(
+            "{num_unquantized_embeddings_scored=",
+            std::to_string(stats.num_unquantized_embeddings_scored),
+            ", num_quantized_embeddings_scored=",
+            std::to_string(stats.num_quantized_embeddings_scored),
+            ", unquantized_shards_read=[",
+            absl_ports::StrJoin(stats.unquantized_shards_read, ",",
+                                absl_ports::NumberFormatter()),
+            "], quantized_shards_read=[",
+            absl_ports::StrJoin(stats.quantized_shards_read, ",",
+                                absl_ports::NumberFormatter()),
+            "], num_embedding_bytes_read=",
+            std::to_string(stats.num_embedding_bytes_read), "}");
+      };
   *result_listener << IcingStringUtil::StringPrintf(
       "(actual is {num_leaf_advance_calls_lite_index=%d, "
       "num_leaf_advance_calls_main_index=%d, "
       "num_leaf_advance_calls_integer_index=%d, "
-      "num_leaf_advance_calls_no_index=%d, num_blocks_inspected=%d}, but "
+      "num_leaf_advance_calls_no_index=%d, num_blocks_inspected=%d, "
+      "embedding_stats=%s}, but "
       "expected was {num_leaf_advance_calls_lite_index=%d, "
       "num_leaf_advance_calls_main_index=%d, "
       "num_leaf_advance_calls_integer_index=%d, "
-      "num_leaf_advance_calls_no_index=%d, num_blocks_inspected=%d}.)",
+      "num_leaf_advance_calls_no_index=%d, num_blocks_inspected=%d, "
+      "embedding_stats=%s}.)",
       actual.num_leaf_advance_calls_lite_index,
       actual.num_leaf_advance_calls_main_index,
       actual.num_leaf_advance_calls_integer_index,
       actual.num_leaf_advance_calls_no_index, actual.num_blocks_inspected,
+      embedding_stats_to_string(actual.embedding_stats).c_str(),
       num_leaf_advance_calls_lite_index, num_leaf_advance_calls_main_index,
       num_leaf_advance_calls_integer_index, num_leaf_advance_calls_no_index,
-      num_blocks_inspected);
+      num_blocks_inspected, embedding_stats_to_string(embedding_stats).c_str());
   return actual.num_leaf_advance_calls_lite_index ==
              num_leaf_advance_calls_lite_index &&
          actual.num_leaf_advance_calls_main_index ==
@@ -122,7 +142,8 @@
              num_leaf_advance_calls_integer_index &&
          actual.num_leaf_advance_calls_no_index ==
              num_leaf_advance_calls_no_index &&
-         actual.num_blocks_inspected == num_blocks_inspected;
+         actual.num_blocks_inspected == num_blocks_inspected &&
+         actual.embedding_stats == embedding_stats;
 }
 
 // Used to match a DocumentAssociatedScoreData
@@ -139,6 +160,13 @@
              has_valid_scorable_property_cache_index;
 }
 
+MATCHER_P4(EqualsDocumentMetadata, schema_type_name, name_space, uri,
+           document_id, "") {
+  return arg.schema_type_name == schema_type_name &&
+         arg.name_space == name_space && arg.uri == uri &&
+         arg.document_id == document_id;
+}
+
 // Used to match a ScorablePropertyManager::ScorablePropertyInfo
 MATCHER_P2(EqualsScorablePropertyInfo, property_path, data_type, "") {
   const ScorablePropertyManager::ScorablePropertyInfo& actual = arg;
@@ -612,6 +640,7 @@
   SearchResultProto actual_copy = arg;
   actual_copy.clear_query_stats();
   actual_copy.clear_debug_info();
+  actual_copy.clear_vm_binder_transaction_latency_start_time_ms();
   for (SearchResultProto::ResultProto& result :
        *actual_copy.mutable_results()) {
     // Joined results
@@ -625,6 +654,7 @@
   SearchResultProto expected_copy = expected;
   expected_copy.clear_query_stats();
   expected_copy.clear_debug_info();
+  actual_copy.clear_vm_binder_transaction_latency_start_time_ms();
   for (SearchResultProto::ResultProto& result :
        *expected_copy.mutable_results()) {
     // Joined results

diff --git a/icing/testing/embedding-test-utils.cc b/icing/testing/embedding-test-utils.cc
index 23f4197..db2d622 100644
--- a/icing/testing/embedding-test-utils.cc
+++ b/icing/testing/embedding-test-utils.cc

@@ -24,9 +24,10 @@
 #include "icing/absl_ports/canonical_errors.h"
 #include "icing/index/embed/embedding-hit.h"
 #include "icing/index/embed/embedding-index.h"
-#include "icing/index/embed/posting-list-embedding-hit-accessor.h"
+#include "icing/index/embed/embedding-query-results.h"
 #include "icing/index/embed/quantizer.h"
 #include "icing/proto/document.pb.h"
+#include "icing/store/document-id.h"
 #include "icing/util/status-macros.h"
 
 namespace icing {
@@ -38,18 +39,20 @@
                           std::string_view model_signature) {
   std::vector<EmbeddingHit> hits;
 
-  libtextclassifier3::StatusOr<std::unique_ptr<PostingListEmbeddingHitAccessor>>
-      pl_accessor_or = embedding_index->GetAccessor(dimension, model_signature);
-  if (absl_ports::IsNotFound(pl_accessor_or.status())) {
+  libtextclassifier3::StatusOr<
+      std::unique_ptr<EmbeddingIndex::EmbeddingHitAccessor>>
+      embedding_hit_accessor_or =
+          embedding_index->GetAccessor(dimension, model_signature);
+  if (absl_ports::IsNotFound(embedding_hit_accessor_or.status())) {
     return hits;
   }
-  ICING_ASSIGN_OR_RETURN(
-      std::unique_ptr<PostingListEmbeddingHitAccessor> pl_accessor,
-      std::move(pl_accessor_or));
+  ICING_ASSIGN_OR_RETURN(std::unique_ptr<EmbeddingIndex::EmbeddingHitAccessor>
+                             embedding_hit_accessor,
+                         std::move(embedding_hit_accessor_or));
 
   while (true) {
     ICING_ASSIGN_OR_RETURN(std::vector<EmbeddingHit> batch,
-                           pl_accessor->GetNextHitsBatch());
+                           embedding_hit_accessor->GetNextHitsBatch());
     if (batch.empty()) {
       return hits;
     }
@@ -58,20 +61,24 @@
 }
 
 std::vector<float> GetRawEmbeddingDataFromIndex(
-    const EmbeddingIndex* embedding_index) {
+    const EmbeddingIndex* embedding_index, uint32_t shard_id) {
   ICING_ASSIGN_OR_RETURN(const float* data,
-                         embedding_index->GetRawEmbeddingData(),
+                         embedding_index->GetRawEmbeddingData(shard_id),
                          std::vector<float>());
-  return std::vector<float>(data, data + embedding_index->GetTotalVectorSize());
+  return std::vector<float>(
+      data, data + embedding_index->GetTotalVectorSize(shard_id));
 }
 
 libtextclassifier3::StatusOr<std::vector<float>>
 GetAndRestoreQuantizedEmbeddingVectorFromIndex(
     const EmbeddingIndex* embedding_index, const EmbeddingHit& hit,
-    uint32_t dimension) {
+    uint32_t dimension, std::string_view model_signature,
+    std::string_view schema_name) {
+  uint32_t shard_id =
+      embedding_index->GetShardId(dimension, model_signature, schema_name);
   ICING_ASSIGN_OR_RETURN(
       const char* data,
-      embedding_index->GetQuantizedEmbeddingVector(hit, dimension));
+      embedding_index->GetQuantizedEmbeddingVector(hit, dimension, shard_id));
   Quantizer quantizer(data);
   const uint8_t* quantized_vector =
       reinterpret_cast<const uint8_t*>(data + sizeof(Quantizer));
@@ -83,5 +90,16 @@
   return result;
 }
 
+EmbeddingMatchInfos& GetOrCreateEmbeddingMatchInfosForDocument(
+    EmbeddingQueryResults& embedding_query_results, int query_vector_index,
+    SearchSpecProto::EmbeddingQueryMetricType::Code metric_type,
+    DocumentId doc_id) {
+  EmbeddingQueryResults::EmbeddingQueryMatchInfoMap* info_map =
+      embedding_query_results
+          .GetOrCreateMatchInfoMap(query_vector_index, metric_type)
+          .ValueOrDie();
+  return (*info_map)[doc_id];
+}
+
 }  // namespace lib
 }  // namespace icing

diff --git a/icing/testing/embedding-test-utils.h b/icing/testing/embedding-test-utils.h
index 7f0b0c1..1908561 100644
--- a/icing/testing/embedding-test-utils.h
+++ b/icing/testing/embedding-test-utils.h

@@ -24,7 +24,9 @@
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "icing/index/embed/embedding-hit.h"
 #include "icing/index/embed/embedding-index.h"
+#include "icing/index/embed/embedding-query-results.h"
 #include "icing/proto/document.pb.h"
+#include "icing/store/document-id.h"
 
 namespace icing {
 namespace lib {
@@ -50,14 +52,22 @@
                           uint32_t dimension, std::string_view model_signature);
 
 std::vector<float> GetRawEmbeddingDataFromIndex(
-    const EmbeddingIndex* embedding_index);
+    const EmbeddingIndex* embedding_index, uint32_t shard_id);
 
 // Gets the quantized embedding vector from the index based on the given hit,
 // and returns the dequantized version of the vector.
 libtextclassifier3::StatusOr<std::vector<float>>
 GetAndRestoreQuantizedEmbeddingVectorFromIndex(
     const EmbeddingIndex* embedding_index, const EmbeddingHit& hit,
-    uint32_t dimension);
+    uint32_t dimension, std::string_view model_signature,
+    std::string_view schema_name);
+
+// Gets or creates the EmbeddingMatchInfos in embedding_query_results for the
+// given query_vector_index, metric_type, and document.
+EmbeddingMatchInfos& GetOrCreateEmbeddingMatchInfosForDocument(
+    EmbeddingQueryResults& embedding_query_results, int query_vector_index,
+    SearchSpecProto::EmbeddingQueryMetricType::Code metric_type,
+    DocumentId doc_id);
 
 }  // namespace lib
 }  // namespace icing

diff --git a/icing/testing/fake-clock.h b/icing/testing/fake-clock.h
index f451753..a7783e9 100644
--- a/icing/testing/fake-clock.h
+++ b/icing/testing/fake-clock.h

@@ -15,6 +15,10 @@
 #ifndef ICING_TESTING_FAKE_CLOCK_H_
 #define ICING_TESTING_FAKE_CLOCK_H_
 
+#include <atomic>
+#include <cstdint>
+#include <memory>
+
 #include "icing/util/clock.h"
 
 namespace icing {
@@ -55,7 +59,7 @@
   }
 
  private:
-  int64_t milliseconds_ = 0;
+  std::atomic<int64_t> milliseconds_ = 0;
   FakeTimer fake_timer_;
 };
 

diff --git a/icing/testing/test-feature-flags.cc b/icing/testing/test-feature-flags.cc
index 3d167b4..416cf9c 100644
--- a/icing/testing/test-feature-flags.cc
+++ b/icing/testing/test-feature-flags.cc

@@ -20,13 +20,28 @@
 namespace lib {
 
 FeatureFlags GetTestFeatureFlags() {
-  return FeatureFlags(/*enable_circular_schema_definitions=*/true,
-                      /*enable_scorable_properties=*/true,
-                      /*enable_embedding_quantization=*/true,
-                      /*enable_repeated_field_joins=*/true,
-                      /*enable_embedding_backup_generation=*/true,
-                      /*enable_schema_database=*/true,
-                      /*release_backup_schema_file_if_overlay_present=*/true);
+  return FeatureFlags(
+      /*enable_circular_schema_definitions=*/true,
+      /*enable_scorable_properties=*/true,
+      /*enable_embedding_quantization=*/true,
+      /*enable_repeated_field_joins=*/true,
+      /*enable_embedding_backup_generation=*/true,
+      /*enable_schema_database=*/true,
+      /*release_backup_schema_file_if_overlay_present=*/true,
+      /*enable_strict_page_byte_size_limit=*/true,
+      /*enable_smaller_decompression_buffer_size=*/true,
+      /*enable_eigen_embedding_scoring=*/true,
+      /*enable_passing_filter_to_children=*/true,
+      /*enable_proto_log_new_header_format=*/true,
+      /*enable_embedding_iterator_v2=*/true,
+      /*enable_reusable_decompression_buffer=*/true,
+      /*enable_schema_type_id_optimization=*/true,
+      /*enable_optimize_improvements=*/true,
+      /*expired_document_purge_threshold_ms=*/0,
+      /*enable_non_existent_qualified_id_join=*/true,
+      /*enable_skip_set_schema_type_equality_check=*/true,
+      /*enable_embed_query_optimization=*/true,
+      /*enable_schema_definition_deduping=*/true);
 }
 
 }  // namespace lib

diff --git a/icing/text_classifier/lib3/utils/base/config.h b/icing/text_classifier/lib3/utils/base/config.h
deleted file mode 100644
index 88af90b..0000000
--- a/icing/text_classifier/lib3/utils/base/config.h
+++ /dev/null

@@ -1,41 +0,0 @@
-// Copyright (C) 2019 Google LLC
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//      http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-// Define macros to indicate C++ standard / platform / etc we use.
-
-#ifndef ICING_TEXT_CLASSIFIER_LIB3_UTILS_BASE_CONFIG_H_
-#define ICING_TEXT_CLASSIFIER_LIB3_UTILS_BASE_CONFIG_H_
-
-namespace libtextclassifier3 {
-
-// Define LANG_CXX11 to 1 if current compiler supports C++11.
-//
-// GXX_EXPERIMENTAL_CXX0X is defined by gcc and clang up to at least
-// gcc-4.7 and clang-3.1 (2011-12-13).  __cplusplus was defined to 1
-// in gcc before 4.7 (Crosstool 16) and clang before 3.1, but is
-// defined according to the language version in effect thereafter.
-// Microsoft Visual Studio 14 (2015) sets __cplusplus==199711 despite
-// reasonably good C++11 support, so we set LANG_CXX for it and
-// newer versions (_MSC_VER >= 1900).
-#if (defined(__GXX_EXPERIMENTAL_CXX0X__) || __cplusplus >= 201103L || \
-     (defined(_MSC_VER) && _MSC_VER >= 1900))
-// Define this to 1 if the code is compiled in C++11 mode; leave it
-// undefined otherwise.  Do NOT define it to 0 -- that causes
-// '#ifdef LANG_CXX11' to behave differently from '#if LANG_CXX11'.
-#define LANG_CXX11 1
-#endif
-
-}  // namespace libtextclassifier3
-
-#endif  // ICING_TEXT_CLASSIFIER_LIB3_UTILS_BASE_CONFIG_H_

diff --git a/icing/text_classifier/lib3/utils/base/integral_types.h b/icing/text_classifier/lib3/utils/base/integral_types.h
index 3baa903..242c913 100644
--- a/icing/text_classifier/lib3/utils/base/integral_types.h
+++ b/icing/text_classifier/lib3/utils/base/integral_types.h

@@ -17,7 +17,6 @@
 #ifndef ICING_TEXT_CLASSIFIER_LIB3_UTILS_BASE_INTEGRAL_TYPES_H_
 #define ICING_TEXT_CLASSIFIER_LIB3_UTILS_BASE_INTEGRAL_TYPES_H_
 
-#include "icing/text_classifier/lib3/utils/base/config.h"
 
 namespace libtextclassifier3 {
 
@@ -46,8 +45,6 @@
 #endif  // COMPILER_MSVC
 
 // Some compile-time assertions that our new types have the intended size.
-// static_assert exists only since C++11, so we need an ifdef.
-#ifdef LANG_CXX11
 static_assert(sizeof(int) == 4, "Our typedefs depend on int being 32 bits");
 static_assert(sizeof(uint32) == 4, "wrong size");
 static_assert(sizeof(int32) == 4, "wrong size");
@@ -57,7 +54,6 @@
 static_assert(sizeof(int16) == 2, "wrong size");
 static_assert(sizeof(char32) == 4, "wrong size");
 static_assert(sizeof(int64) == 8, "wrong size");
-#endif  // LANG_CXX11
 
 // There are still some requirements that we build these headers in
 // C-compatibility mode. Unfortunately, -Wall doesn't like c-style

diff --git a/icing/text_classifier/lib3/utils/base/macros.h b/icing/text_classifier/lib3/utils/base/macros.h
index f6642d2..eefaeda 100644
--- a/icing/text_classifier/lib3/utils/base/macros.h
+++ b/icing/text_classifier/lib3/utils/base/macros.h

@@ -15,26 +15,15 @@
 #ifndef ICING_TEXT_CLASSIFIER_LIB3_UTILS_BASE_MACROS_H_
 #define ICING_TEXT_CLASSIFIER_LIB3_UTILS_BASE_MACROS_H_
 
-#include "icing/text_classifier/lib3/utils/base/config.h"
 
 namespace libtextclassifier3 {
 
 #define TC3_ARRAYSIZE(a) \
   ((sizeof(a) / sizeof(*(a))) / (size_t)(!(sizeof(a) % sizeof(*(a)))))
 
-#if LANG_CXX11
 #define TC3_DISALLOW_COPY_AND_ASSIGN(TypeName) \
   TypeName(const TypeName &) = delete;         \
   TypeName &operator=(const TypeName &) = delete
-#else  // C++98 case follows
-
-// Note that these C++98 implementations cannot completely disallow copying,
-// as members and friends can still accidentally make elided copies without
-// triggering a linker error.
-#define TC3_DISALLOW_COPY_AND_ASSIGN(TypeName) \
-  TypeName(const TypeName &);                  \
-  TypeName &operator=(const TypeName &)
-#endif  // LANG_CXX11
 
 // The TC3_FALLTHROUGH_INTENDED macro can be used to annotate implicit
 // fall-through between switch labels:

diff --git a/icing/text_classifier/lib3/utils/hash/farmhash.cc b/icing/text_classifier/lib3/utils/hash/farmhash.cc
index c04df07..10da72a 100644
--- a/icing/text_classifier/lib3/utils/hash/farmhash.cc
+++ b/icing/text_classifier/lib3/utils/hash/farmhash.cc

@@ -29,13 +29,6 @@
 #define FARMHASH_ASSUME_AESNI 1
 #endif
 
-#if !defined(FARMHASH_CAN_USE_CXX11) && defined(LANG_CXX11)
-#define FARMHASH_CAN_USE_CXX11 1
-#else
-#undef FARMHASH_CAN_USE_CXX11
-#define FARMHASH_CAN_USE_CXX11 0
-#endif
-
 // FARMHASH PORTABILITY LAYER: Runtime error if misconfigured
 
 #ifndef FARMHASH_DIE_IF_MISCONFIGURED
@@ -284,11 +277,6 @@
 #endif
 // Building blocks for hash functions
 
-// std::swap() was in <algorithm> but is in <utility> from C++11 on.
-#if !FARMHASH_CAN_USE_CXX11
-#include <algorithm>
-#endif
-
 #undef PERMUTE3
 #define PERMUTE3(a, b, c) do { std::swap(a, b); std::swap(a, c); } while (0)
 

diff --git a/icing/tokenization/combined-tokenizer_test.cc b/icing/tokenization/combined-tokenizer_test.cc
index 520ea47..b769e28 100644
--- a/icing/tokenization/combined-tokenizer_test.cc
+++ b/icing/tokenization/combined-tokenizer_test.cc

@@ -34,6 +34,7 @@
 #include "icing/index/numeric/numeric-index.h"
 #include "icing/jni/jni-cache.h"
 #include "icing/legacy/index/icing-filesystem.h"
+#include "icing/portable/gzip_stream.h"
 #include "icing/portable/platform.h"
 #include "icing/proto/schema.pb.h"
 #include "icing/query/query-processor.h"
@@ -99,14 +100,18 @@
 
     ICING_ASSERT_OK_AND_ASSIGN(
         DocumentStore::CreateResult create_result,
-        DocumentStore::Create(&filesystem_, store_dir_, &fake_clock_,
-                              schema_store_.get(), feature_flags_.get(),
-                              /*force_recovery_and_revalidate_documents=*/false,
-                              /*pre_mapping_fbv=*/false,
-                              /*use_persistent_hash_map=*/false,
-                              PortableFileBackedProtoLog<
-                                  DocumentWrapper>::kDefaultCompressionLevel,
-                              /*initialize_stats=*/nullptr));
+        DocumentStore::Create(
+            &filesystem_, store_dir_, &fake_clock_, schema_store_.get(),
+            feature_flags_.get(),
+            /*force_recovery_and_revalidate_documents=*/false,
+            /*pre_mapping_fbv=*/false,
+            /*use_persistent_hash_map=*/false,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionLevel,
+            PortableFileBackedProtoLog<
+                DocumentWrapper>::kDefaultCompressionThresholdBytes,
+            protobuf_ports::kDefaultMemLevel,
+            /*initialize_stats=*/nullptr));
     document_store_ = std::move(create_result.document_store);
 
     Index::Options options(index_dir_,
@@ -114,7 +119,8 @@
                            /*lite_index_sort_at_indexing=*/true,
                            /*lite_index_sort_size=*/1024 * 8);
     ICING_ASSERT_OK_AND_ASSIGN(
-        index_, Index::Create(options, &filesystem_, &icing_filesystem_));
+        index_, Index::Create(options, &filesystem_, &icing_filesystem_,
+                              feature_flags_.get()));
     // TODO(b/249829533): switch to use persistent numeric index.
     ICING_ASSERT_OK_AND_ASSIGN(
         numeric_index_,
@@ -122,7 +128,8 @@
     ICING_ASSERT_OK_AND_ASSIGN(
         embedding_index_,
         EmbeddingIndex::Create(&filesystem_, embedding_index_dir_, &fake_clock_,
-                               feature_flags_.get()));
+                               feature_flags_.get(),
+                               /*num_shards=*/32));
 
     language_segmenter_factory::SegmenterOptions segmenter_options(
         ULOC_US, jni_cache_.get());

diff --git a/icing/tokenization/icu/icu-language-segmenter-factory.cc b/icing/tokenization/icu/icu-language-segmenter-factory.cc
index 4f8263b..4117733 100644
--- a/icing/tokenization/icu/icu-language-segmenter-factory.cc
+++ b/icing/tokenization/icu/icu-language-segmenter-factory.cc

@@ -12,9 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <utility>
+
 #include "icing/tokenization/icu/icu-language-segmenter.h"
 #include "icing/tokenization/language-segmenter-factory.h"
 #include "icing/util/logging.h"
+#include "icing/util/status-util.h"
 #include "unicode/uloc.h"
 
 namespace icing {
@@ -22,11 +25,17 @@
 
 namespace language_segmenter_factory {
 
+using ::icing::lib::status_util::TransformStatus;
+
 namespace {
 constexpr std::string_view kLocaleAmericanEnglishComputer = "en_US_POSIX";
 }  // namespace
 
-// Creates a language segmenter with the given locale.
+// Creates a language segmenter based on the provided options.
+//
+// @param options: The options for creating the language segmenter.
+// @param icu_segmenter_creation_status: Optional output parameter that will be
+//        populated with the status of IcuLanguageSegmenter.
 //
 // Returns:
 //   A LanguageSegmenter on success
@@ -36,7 +45,7 @@
 // users. Right now illegal locale strings will be ignored by ICU. ICU
 // components will be created with its default locale.
 libtextclassifier3::StatusOr<std::unique_ptr<LanguageSegmenter>> Create(
-    SegmenterOptions options) {
+    SegmenterOptions options, StatusProto* icu_segmenter_creation_status) {
   // Word connector rules for "en_US_POSIX" (American English (Computer)) are
   // different from other locales. E.g. "email.subject" will be split into 3
   // terms in "en_US_POSIX": "email", ".", and "subject", while it's just one
@@ -47,7 +56,12 @@
                        << " not supported. Converting to locale " << ULOC_US;
     options.locale = ULOC_US;
   }
-  return IcuLanguageSegmenter::Create(std::move(options.locale));
+  auto icu_segmenter_or =
+      IcuLanguageSegmenter::Create(std::move(options.locale));
+  if (icu_segmenter_creation_status != nullptr) {
+    TransformStatus(icu_segmenter_or.status(), icu_segmenter_creation_status);
+  }
+  return icu_segmenter_or;
 }
 
 }  // namespace language_segmenter_factory

diff --git a/icing/tokenization/icu/icu-language-segmenter_test.cc b/icing/tokenization/icu/icu-language-segmenter_test.cc
index cde62f2..f1c8f1a 100644
--- a/icing/tokenization/icu/icu-language-segmenter_test.cc
+++ b/icing/tokenization/icu/icu-language-segmenter_test.cc

@@ -61,6 +61,16 @@
 
 }  // namespace
 
+TEST_P(IcuLanguageSegmenterAllLocalesTest, SuccessfulIcuSegmenterSetsOkStatus) {
+  StatusProto icu_segmenter_creation_status;
+  ICING_ASSERT_OK_AND_ASSIGN(auto language_segmenter,
+                             language_segmenter_factory::Create(
+                                 language_segmenter_factory::SegmenterOptions(
+                                     GetLocale(), jni_cache_.get()),
+                                 &icu_segmenter_creation_status));
+  EXPECT_THAT(icu_segmenter_creation_status.code(), Eq(StatusProto::OK));
+}
+
 TEST_P(IcuLanguageSegmenterAllLocalesTest, EmptyText) {
   ICING_ASSERT_OK_AND_ASSIGN(
       auto language_segmenter,

diff --git a/icing/tokenization/language-segmenter-factory.h b/icing/tokenization/language-segmenter-factory.h
index 5e4879c..24d8930 100644
--- a/icing/tokenization/language-segmenter-factory.h
+++ b/icing/tokenization/language-segmenter-factory.h

@@ -21,6 +21,7 @@
 
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "icing/jni/jni-cache.h"
+#include "icing/proto/status.pb.h"
 #include "icing/tokenization/language-segmenter.h"
 
 namespace icing {
@@ -53,13 +54,19 @@
   bool enable_icu_segmenter;
 };
 
-// Creates a language segmenter with the given locale.
+// Creates a language segmenter based on the provided options.
+//
+// @param options: The options for creating the language segmenter.
+// @param icu_segmenter_creation_status: Optional output parameter that will be
+//        populated with the status of IcuLanguageSegmenter if it's creation is
+//        attempted.
 //
 // Returns:
 //   A LanguageSegmenter on success
 //   INVALID_ARGUMENT_ERROR if locale string is invalid
 libtextclassifier3::StatusOr<std::unique_ptr<LanguageSegmenter>> Create(
-    SegmenterOptions options);
+    SegmenterOptions options,
+    StatusProto* icu_segmenter_creation_status = nullptr);
 
 }  // namespace language_segmenter_factory
 

diff --git a/icing/tokenization/plain-tokenizer.cc b/icing/tokenization/plain-tokenizer.cc
index ec7a783..c67b289 100644
--- a/icing/tokenization/plain-tokenizer.cc
+++ b/icing/tokenization/plain-tokenizer.cc

@@ -20,6 +20,7 @@
 
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "icing/tokenization/language-segmenter.h"
+#include "icing/tokenization/token.h"
 #include "icing/util/character-iterator.h"
 #include "icing/util/i18n-utils.h"
 #include "icing/util/status-macros.h"
@@ -67,12 +68,11 @@
     return found_next_valid_term;
   }
 
-  std::vector<Token> GetTokens() const override {
-    std::vector<Token> result;
+  void GetTokens(std::vector<Token>* out_tokens) const override {
+    out_tokens->clear();
     if (!current_term_.empty()) {
-      result.push_back(Token(Token::Type::REGULAR, current_term_));
+      out_tokens->emplace_back(Token::Type::REGULAR, current_term_);
     }
-    return result;
   }
 
   libtextclassifier3::StatusOr<CharacterIterator> CalculateTokenStart()
@@ -143,8 +143,9 @@
   ICING_ASSIGN_OR_RETURN(std::unique_ptr<Tokenizer::Iterator> iterator,
                          Tokenize(text));
   std::vector<Token> tokens;
+  std::vector<Token> batch_tokens;
   while (iterator->Advance()) {
-    std::vector<Token> batch_tokens = iterator->GetTokens();
+    iterator->GetTokens(&batch_tokens);
     tokens.insert(tokens.end(), batch_tokens.begin(), batch_tokens.end());
   }
   return tokens;

diff --git a/icing/tokenization/plain-tokenizer_test.cc b/icing/tokenization/plain-tokenizer_test.cc
index 08d7f3e..20f74c8 100644
--- a/icing/tokenization/plain-tokenizer_test.cc
+++ b/icing/tokenization/plain-tokenizer_test.cc

@@ -72,7 +72,7 @@
                              plain_tokenizer->Tokenize(kText));
 
   // We should get no tokens if we get the token before advancing.
-  EXPECT_THAT(token_iterator->GetTokens(), IsEmpty());
+  EXPECT_THAT(token_iterator->GetTokensForTest(), IsEmpty());
 }
 
 TEST_F(PlainTokenizerTest, LastTokenAfterFullyAdvanced) {
@@ -93,7 +93,7 @@
   while (token_iterator->Advance()) {}
 
   // After advance returns false, GetTokens will stay on the last token.
-  EXPECT_THAT(token_iterator->GetTokens(),
+  EXPECT_THAT(token_iterator->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::REGULAR, "!")));
 }
 
@@ -164,6 +164,27 @@
                                EqualsToken(Token::Type::REGULAR, "World"))));
 }
 
+TEST_F(PlainTokenizerTest, ContinuousWhitespace) {
+  language_segmenter_factory::SegmenterOptions options(ULOC_US,
+                                                       jni_cache_.get());
+  ICING_ASSERT_OK_AND_ASSIGN(
+      auto language_segmenter,
+      language_segmenter_factory::Create(std::move(options)));
+  ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<Tokenizer> plain_tokenizer,
+                             tokenizer_factory::CreateIndexingTokenizer(
+                                 StringIndexingConfig::TokenizerType::PLAIN,
+                                 language_segmenter.get()));
+
+  // Multiple continuous whitespaces will be ignored.
+  const int kNumSeparators = 256;
+  std::string text_with_spaces =
+      absl_ports::StrCat("Hello", std::string(kNumSeparators, ' '), "World");
+  EXPECT_THAT(
+      plain_tokenizer->TokenizeAll(text_with_spaces),
+      IsOkAndHolds(ElementsAre(EqualsToken(Token::Type::REGULAR, "Hello"),
+                               EqualsToken(Token::Type::REGULAR, "World"))));
+}
+
 TEST_F(PlainTokenizerTest, Punctuation) {
   language_segmenter_factory::SegmenterOptions options(ULOC_US,
                                                        jni_cache_.get());
@@ -374,7 +395,7 @@
   auto iterator = plain_tokenizer->Tokenize(kText).ValueOrDie();
 
   EXPECT_TRUE(iterator->ResetToTokenStartingAfter(0));
-  EXPECT_THAT(iterator->GetTokens(),
+  EXPECT_THAT(iterator->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::REGULAR, "b")));
 
   EXPECT_FALSE(iterator->ResetToTokenStartingAfter(2));
@@ -395,7 +416,7 @@
   auto iterator = plain_tokenizer->Tokenize(kText).ValueOrDie();
 
   EXPECT_TRUE(iterator->ResetToTokenEndingBefore(2));
-  EXPECT_THAT(iterator->GetTokens(),
+  EXPECT_THAT(iterator->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::REGULAR, "f")));
 
   EXPECT_FALSE(iterator->ResetToTokenEndingBefore(0));
@@ -441,13 +462,13 @@
 
   auto iterator = plain_tokenizer->Tokenize(kText).ValueOrDie();
   EXPECT_TRUE(iterator->Advance());
-  EXPECT_THAT(iterator->GetTokens(),
+  EXPECT_THAT(iterator->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::REGULAR, "foo")));
   for (int i = 0; i < kText.length(); ++i) {
     if (i < expected_text.size()) {
       EXPECT_TRUE(iterator->ResetToTokenStartingAfter(i));
       EXPECT_THAT(
-          iterator->GetTokens(),
+          iterator->GetTokensForTest(),
           ElementsAre(EqualsToken(Token::Type::REGULAR, expected_text[i])));
     } else {
       EXPECT_FALSE(iterator->ResetToTokenStartingAfter(i));
@@ -495,13 +516,13 @@
 
   auto iterator = plain_tokenizer->Tokenize(kText).ValueOrDie();
   EXPECT_TRUE(iterator->Advance());
-  EXPECT_THAT(iterator->GetTokens(),
+  EXPECT_THAT(iterator->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::REGULAR, "foo")));
   for (int i = kText.length() - 1; i >= 0; --i) {
     int expected_index = kText.length() - 1 - i;
     if (expected_index < expected_text.size()) {
       EXPECT_TRUE(iterator->ResetToTokenEndingBefore(i));
-      EXPECT_THAT(iterator->GetTokens(),
+      EXPECT_THAT(iterator->GetTokensForTest(),
                   ElementsAre(EqualsToken(Token::Type::REGULAR,
                                           expected_text[expected_index])));
     } else {

diff --git a/icing/tokenization/reverse_jni/reverse-jni-language-segmenter-factory.cc b/icing/tokenization/reverse_jni/reverse-jni-language-segmenter-factory.cc
index 1581e13..e6b2e70 100644
--- a/icing/tokenization/reverse_jni/reverse-jni-language-segmenter-factory.cc
+++ b/icing/tokenization/reverse_jni/reverse-jni-language-segmenter-factory.cc

@@ -28,7 +28,11 @@
 constexpr std::string_view kLocaleAmericanEnglishComputer = "en_US_POSIX";
 }  // namespace
 
-// Creates a language segmenter with the given locale.
+// Creates a language segmenter based on the provided options.
+//
+// @param options: The options for creating the language segmenter.
+// @param icu_segmenter_creation_status: Optional output parameter that will be
+//        ignored.
 //
 // Returns:
 //   A LanguageSegmenter on success
@@ -38,7 +42,7 @@
 // users. Right now illegal locale strings will be ignored by ICU. ICU
 // components will be created with its default locale.
 libtextclassifier3::StatusOr<std::unique_ptr<LanguageSegmenter>> Create(
-    SegmenterOptions options) {
+    SegmenterOptions options, StatusProto* /*icu_segmenter_creation_status*/) {
   if (options.jni_cache == nullptr) {
     return absl_ports::InvalidArgumentError(
         "Cannot create Reverse Jni Language Segmenter without a valid JniCache "

diff --git a/icing/tokenization/reverse_jni/reverse-jni-language-segmenter_test.cc b/icing/tokenization/reverse_jni/reverse-jni-language-segmenter_test.cc
index 683c5f9..7e10741 100644
--- a/icing/tokenization/reverse_jni/reverse-jni-language-segmenter_test.cc
+++ b/icing/tokenization/reverse_jni/reverse-jni-language-segmenter_test.cc

@@ -317,34 +317,6 @@
                                        "７", "８", "９")));
 }
 
-TEST_P(ReverseJniLanguageSegmenterTest, ContinuousWhitespaces) {
-  ICING_ASSERT_OK_AND_ASSIGN(
-      auto language_segmenter,
-      language_segmenter_factory::Create(
-          language_segmenter_factory::SegmenterOptions(
-              GetLocale(), jni_cache_.get())));
-  // Multiple continuous whitespaces are treated as one.
-  const int kNumSeparators = 256;
-  std::string text_with_spaces =
-      absl_ports::StrCat("Hello", std::string(kNumSeparators, ' '), "World");
-  EXPECT_THAT(language_segmenter->GetAllTerms(text_with_spaces),
-              IsOkAndHolds(ElementsAre("Hello", " ", "World")));
-
-  // Multiple continuous whitespaces are treated as one. Whitespace at the
-  // beginning of the text doesn't affect the results of GetTerm() after the
-  // iterator is done.
-  text_with_spaces = absl_ports::StrCat(std::string(kNumSeparators, ' '),
-                                        "Hello", " ", "World");
-  ICING_ASSERT_OK_AND_ASSIGN(auto itr,
-                             language_segmenter->Segment(text_with_spaces));
-  std::vector<std::string_view> terms;
-  while (itr->Advance()) {
-    terms.push_back(itr->GetTerm());
-  }
-  EXPECT_THAT(terms, ElementsAre(" ", "Hello", " ", "World"));
-  EXPECT_THAT(itr->GetTerm(), IsEmpty());
-}
-
 TEST_P(ReverseJniLanguageSegmenterTest, CJKT) {
   ICING_ASSERT_OK_AND_ASSIGN(
       auto language_segmenter,
@@ -775,47 +747,6 @@
   EXPECT_THAT(itr->GetTerm(), IsEmpty());
 }
 
-TEST_P(ReverseJniLanguageSegmenterTest,
-       ContinuousWhitespacesResetToTermAfterUtf32) {
-  ICING_ASSERT_OK_AND_ASSIGN(
-      auto language_segmenter,
-      language_segmenter_factory::Create(
-          language_segmenter_factory::SegmenterOptions(
-              GetLocale(), jni_cache_.get())));
-  // Multiple continuous whitespaces are treated as one.
-  constexpr std::string_view kTextWithSpace = "Hello          World";
-  ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<LanguageSegmenter::Iterator> itr,
-                             language_segmenter->Segment(kTextWithSpace));
-
-  // String:      "Hello          World"
-  //               ^    ^         ^
-  // UTF-8 idx:    0    5         15
-  // UTF-32 idx:   0    5         15
-  EXPECT_THAT(itr->ResetToTermStartingAfterUtf32(0), IsOkAndHolds(Eq(5)));
-  EXPECT_THAT(itr->GetTerm(), Eq(" "));
-
-  EXPECT_THAT(itr->ResetToTermStartingAfterUtf32(2), IsOkAndHolds(Eq(5)));
-  EXPECT_THAT(itr->GetTerm(), Eq(" "));
-
-  EXPECT_THAT(itr->ResetToTermStartingAfterUtf32(10), IsOkAndHolds(Eq(15)));
-  EXPECT_THAT(itr->GetTerm(), Eq("World"));
-
-  EXPECT_THAT(itr->ResetToTermStartingAfterUtf32(5), IsOkAndHolds(Eq(15)));
-  EXPECT_THAT(itr->GetTerm(), Eq("World"));
-
-  EXPECT_THAT(itr->ResetToTermStartingAfterUtf32(15),
-              StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
-  EXPECT_THAT(itr->GetTerm(), IsEmpty());
-
-  EXPECT_THAT(itr->ResetToTermStartingAfterUtf32(17),
-              StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
-  EXPECT_THAT(itr->GetTerm(), IsEmpty());
-
-  EXPECT_THAT(itr->ResetToTermStartingAfterUtf32(19),
-              StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
-  EXPECT_THAT(itr->GetTerm(), IsEmpty());
-}
-
 TEST_P(ReverseJniLanguageSegmenterTest, ChineseResetToTermAfterUtf32) {
   ICING_ASSERT_OK_AND_ASSIGN(
       auto language_segmenter,
@@ -1097,46 +1028,6 @@
   EXPECT_THAT(itr->GetTerm(), Eq("です"));
 }
 
-TEST_P(ReverseJniLanguageSegmenterTest,
-       ContinuousWhitespacesResetToTermBeforeUtf32) {
-  ICING_ASSERT_OK_AND_ASSIGN(
-      auto language_segmenter,
-      language_segmenter_factory::Create(
-          language_segmenter_factory::SegmenterOptions(
-              GetLocale(), jni_cache_.get())));
-  // Multiple continuous whitespaces are treated as one.
-  constexpr std::string_view kTextWithSpace = "Hello          World";
-  ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<LanguageSegmenter::Iterator> itr,
-                             language_segmenter->Segment(kTextWithSpace));
-
-  // String:      "Hello          World"
-  //               ^    ^         ^
-  // UTF-8 idx:    0    5         15
-  // UTF-32 idx:   0    5         15
-  EXPECT_THAT(itr->ResetToTermEndingBeforeUtf32(0),
-              StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
-  EXPECT_THAT(itr->GetTerm(), IsEmpty());
-
-  EXPECT_THAT(itr->ResetToTermEndingBeforeUtf32(2),
-              StatusIs(libtextclassifier3::StatusCode::NOT_FOUND));
-  EXPECT_THAT(itr->GetTerm(), IsEmpty());
-
-  EXPECT_THAT(itr->ResetToTermEndingBeforeUtf32(10), IsOkAndHolds(Eq(0)));
-  EXPECT_THAT(itr->GetTerm(), Eq("Hello"));
-
-  EXPECT_THAT(itr->ResetToTermEndingBeforeUtf32(5), IsOkAndHolds(Eq(0)));
-  EXPECT_THAT(itr->GetTerm(), Eq("Hello"));
-
-  EXPECT_THAT(itr->ResetToTermEndingBeforeUtf32(15), IsOkAndHolds(Eq(5)));
-  EXPECT_THAT(itr->GetTerm(), Eq(" "));
-
-  EXPECT_THAT(itr->ResetToTermEndingBeforeUtf32(17), IsOkAndHolds(Eq(5)));
-  EXPECT_THAT(itr->GetTerm(), Eq(" "));
-
-  EXPECT_THAT(itr->ResetToTermEndingBeforeUtf32(19), IsOkAndHolds(Eq(5)));
-  EXPECT_THAT(itr->GetTerm(), Eq(" "));
-}
-
 TEST_P(ReverseJniLanguageSegmenterTest, ChineseResetToTermBeforeUtf32) {
   ICING_ASSERT_OK_AND_ASSIGN(
       auto language_segmenter,

diff --git a/icing/tokenization/rfc822-tokenizer.cc b/icing/tokenization/rfc822-tokenizer.cc
index 554340a..9b740c9 100644
--- a/icing/tokenization/rfc822-tokenizer.cc
+++ b/icing/tokenization/rfc822-tokenizer.cc

@@ -79,16 +79,15 @@
   //   A vector of Tokens on success
   //   An empty vector if the token list is empty
   //   An empty vector if the index is past the end of the token list
-  std::vector<Token> GetTokens() const override {
-    std::vector<Token> result;
+  void GetTokens(std::vector<Token>* out_tokens) const override {
+    out_tokens->clear();
     if (token_index_ < tokens_.size() && token_index_ >= 0) {
       int index = token_index_;
       do {
-        result.push_back(tokens_[index]);
+        out_tokens->push_back(tokens_[index]);
       } while (++index < tokens_.size() &&
                tokens_[index].type != Token::Type::RFC822_TOKEN);
     }
-    return result;
   }
 
   bool ResetToTokenStartingAfter(int32_t utf32_offset) override {
@@ -154,16 +153,16 @@
   // Returns a character iterator to the start of the token.
   libtextclassifier3::StatusOr<CharacterIterator> CalculateTokenStart()
       override {
-    CharacterIterator token_start = iterator_;
-    token_start.MoveToUtf8(GetTokens().at(0).text.begin() - text_.begin());
+    CharacterIterator token_start(text_);
+    token_start.MoveToUtf8(tokens_[token_index_].text.begin() - text_.begin());
     return token_start;
   }
 
   // Returns a character iterator to right after the end of the token.
   libtextclassifier3::StatusOr<CharacterIterator> CalculateTokenEndExclusive()
       override {
-    CharacterIterator token_end = iterator_;
-    token_end.MoveToUtf8(GetTokens().at(0).text.end() - text_.begin());
+    CharacterIterator token_end(text_);
+    token_end.MoveToUtf8(tokens_[token_index_].text.end() - text_.begin());
     return token_end;
   }
 
@@ -787,8 +786,9 @@
   ICING_ASSIGN_OR_RETURN(std::unique_ptr<Tokenizer::Iterator> iterator,
                          Tokenize(text));
   std::vector<Token> tokens;
+  std::vector<Token> batch_tokens;
   while (iterator->Advance()) {
-    std::vector<Token> batch_tokens = iterator->GetTokens();
+    iterator->GetTokens(&batch_tokens);
     tokens.insert(tokens.end(), batch_tokens.begin(), batch_tokens.end());
   }
   return tokens;

diff --git a/icing/tokenization/rfc822-tokenizer_test.cc b/icing/tokenization/rfc822-tokenizer_test.cc
index ee3a95d..d886196 100644
--- a/icing/tokenization/rfc822-tokenizer_test.cc
+++ b/icing/tokenization/rfc822-tokenizer_test.cc

@@ -33,9 +33,9 @@
   std::string text = "a@g.c";
   auto token_iterator = rfc822_tokenizer.Tokenize(text).ValueOrDie();
 
-  ASSERT_THAT(token_iterator->GetTokens(), IsEmpty());
+  ASSERT_THAT(token_iterator->GetTokensForTest(), IsEmpty());
   ASSERT_TRUE(token_iterator->Advance());
-  ASSERT_THAT(token_iterator->GetTokens(), Not(IsEmpty()));
+  ASSERT_THAT(token_iterator->GetTokensForTest(), Not(IsEmpty()));
 }
 
 TEST(Rfc822TokenizerTest, EmptyMiddleToken) {
@@ -967,10 +967,10 @@
   ASSERT_TRUE(token_iterator->Advance());
 
   ASSERT_TRUE(token_iterator->ResetToTokenStartingAfter(-1));
-  EXPECT_THAT(token_iterator->GetTokens().at(0).text, "a@g.c");
+  EXPECT_THAT(token_iterator->GetTokensForTest().at(0).text, "a@g.c");
 
   ASSERT_TRUE(token_iterator->ResetToTokenStartingAfter(5));
-  EXPECT_THAT(token_iterator->GetTokens().at(0).text, "b@g.c");
+  EXPECT_THAT(token_iterator->GetTokensForTest().at(0).text, "b@g.c");
 
   ASSERT_FALSE(token_iterator->ResetToTokenStartingAfter(6));
 }
@@ -982,7 +982,7 @@
   token_iterator->Advance();
 
   ASSERT_TRUE(token_iterator->ResetToTokenEndingBefore(5));
-  EXPECT_THAT(token_iterator->GetTokens().at(0).text, "a@g.c");
+  EXPECT_THAT(token_iterator->GetTokensForTest().at(0).text, "a@g.c");
 
   ASSERT_FALSE(token_iterator->ResetToTokenEndingBefore(4));
 }

diff --git a/icing/tokenization/tokenizer.h b/icing/tokenization/tokenizer.h
index fb7613f..778ba99 100644
--- a/icing/tokenization/tokenizer.h
+++ b/icing/tokenization/tokenizer.h

@@ -44,9 +44,10 @@
   // An iterator helping to get tokens.
   // Example usage:
   //
+  // std::vector<Token> tokens;
   // while (iterator.Advance()) {
-  //   const Token& token = iterator.GetToken();
-  //   // Do something
+  //   iterator.GetTokens(&tokens);
+  //   // Do something for each Token
   // }
   class Iterator {
    public:
@@ -55,10 +56,17 @@
     // Advances to the next token. Returns false if it has reached the end.
     virtual bool Advance() = 0;
 
-    // Returns the current token, maybe with compound tokens as well. It can be
-    // called only when Advance() returns true, otherwise an empty Token vector
-    // may be returned.
-    virtual std::vector<Token> GetTokens() const = 0;
+    // Populates `out_tokens` with the current token, maybe with compound tokens
+    // as well. It can be called only when Advance() returns true. The vector is
+    // cleared before being populated.
+    virtual void GetTokens(std::vector<Token>* out_tokens) const = 0;
+
+    // Make it faster to write/refactor tests.
+    std::vector<Token> GetTokensForTest() const {
+      std::vector<Token> result;
+      this->GetTokens(&result);
+      return result;
+    }
 
     virtual libtextclassifier3::StatusOr<CharacterIterator>
     CalculateTokenStart() {
@@ -80,7 +88,9 @@
     // iterator.ResetToTokenStartingAfter(4);
     // // The first full token starting after position 4 (the 'b' in "bar") is
     // // "baz".
-    // PrintToken(iterator.GetToken());  // prints "baz"
+    // std::vector<Token> tokens;
+    // iterator.GetTokens(&tokens);
+    // PrintToken(tokens[0]);  // prints "baz"
     virtual bool ResetToTokenStartingAfter(int32_t utf32_offset) {
       return false;
     }
@@ -93,7 +103,9 @@
     // iterator.ResetToTokenEndingBefore(4);
     // // The first full token ending before position 4 (the 'b' in "bar") is
     // // "foo".
-    // PrintToken(iterator.GetToken());  // prints "foo"
+    // std::vector<Token> tokens;
+    // iterator.GetTokens(&tokens);
+    // PrintToken(tokens[0]);  // prints "baz"
     virtual bool ResetToTokenEndingBefore(int32_t utf32_offset) {
       return false;
     }

diff --git a/icing/tokenization/trigram-tokenizer.cc b/icing/tokenization/trigram-tokenizer.cc
index 870a75d..745b773 100644
--- a/icing/tokenization/trigram-tokenizer.cc
+++ b/icing/tokenization/trigram-tokenizer.cc

@@ -60,15 +60,16 @@
     return true;
   }
 
-  std::vector<Token> GetTokens() const override {
+  void GetTokens(std::vector<Token>* out_tokens) const override {
     int start_byte_idx = utf8_char_itrs_.front().utf8_index();
     int end_byte_idx =
         utf8_char_itrs_.back().utf8_index() +
         i18n_utils::GetUtf8Length(utf8_char_itrs_.back().GetCurrentChar());
 
-    return {Token(/*type_in=*/Token::Type::TRIGRAM,
-                  /*text_in=*/text_.substr(start_byte_idx,
-                                           end_byte_idx - start_byte_idx))};
+    out_tokens->assign(
+        {Token(/*type_in=*/Token::Type::TRIGRAM,
+               /*text_in=*/text_.substr(start_byte_idx,
+                                        end_byte_idx - start_byte_idx))});
   };
 
   // Returns a character iterator to the start of the current trigram token.
@@ -279,8 +280,9 @@
   ICING_ASSIGN_OR_RETURN(std::unique_ptr<Tokenizer::Iterator> iterator,
                          Tokenize(text));
   std::vector<Token> tokens;
+  std::vector<Token> batch_tokens;
   while (iterator->Advance()) {
-    std::vector<Token> batch_tokens = iterator->GetTokens();
+    iterator->GetTokens(&batch_tokens);
     tokens.insert(tokens.end(), batch_tokens.begin(), batch_tokens.end());
   }
   return tokens;

diff --git a/icing/tokenization/trigram-tokenizer_test.cc b/icing/tokenization/trigram-tokenizer_test.cc
index a1466c5..fbb8831 100644
--- a/icing/tokenization/trigram-tokenizer_test.cc
+++ b/icing/tokenization/trigram-tokenizer_test.cc

@@ -189,7 +189,7 @@
 
   // Advance to "abc".
   ASSERT_THAT(itr->Advance(), IsTrue());
-  ASSERT_THAT(itr->GetTokens(),
+  ASSERT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "abc")));
   ICING_ASSERT_OK_AND_ASSIGN(CharacterIterator start_itr0,
                              itr->CalculateTokenStart());
@@ -200,7 +200,7 @@
 
   // Advance to "bcd".
   ASSERT_THAT(itr->Advance(), IsTrue());
-  ASSERT_THAT(itr->GetTokens(),
+  ASSERT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "bcd")));
   ICING_ASSERT_OK_AND_ASSIGN(CharacterIterator start_itr1,
                              itr->CalculateTokenStart());
@@ -221,7 +221,7 @@
 
   // Advance to "我每天".
   ASSERT_THAT(itr->Advance(), IsTrue());
-  ASSERT_THAT(itr->GetTokens(),
+  ASSERT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "我每天")));
   ICING_ASSERT_OK_AND_ASSIGN(CharacterIterator start_itr0,
                              itr->CalculateTokenStart());
@@ -232,7 +232,7 @@
 
   // Advance to "每天走".
   ASSERT_THAT(itr->Advance(), IsTrue());
-  ASSERT_THAT(itr->GetTokens(),
+  ASSERT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "每天走")));
   ICING_ASSERT_OK_AND_ASSIGN(CharacterIterator start_itr1,
                              itr->CalculateTokenStart());
@@ -257,7 +257,7 @@
   ASSERT_THAT(itr->Advance(), IsTrue());
   ASSERT_THAT(itr->Advance(), IsTrue());
   ASSERT_THAT(itr->Advance(), IsTrue());
-  ASSERT_THAT(itr->GetTokens(),
+  ASSERT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "efg")));
 
   // Reset backward to after utf32_offset 1.
@@ -265,22 +265,22 @@
 
   // Verify the iterator is pointing to "cde" since it is the leftmost trigram
   // token after utf32_offset 1.
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "cde")));
 
   // Verify the iterator is still valid and able to advance to the rest of the
   // tokens.
   // Advance to "def".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "def")));
   // Advance to "efg".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "efg")));
   // Advance to "fgh".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "fgh")));
 
   EXPECT_THAT(itr->Advance(), IsFalse());
@@ -291,26 +291,26 @@
 
   // Verify the iterator is pointing to "bcd" since it is the leftmost trigram
   // token after utf32_offset 0.
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "bcd")));
 
   // Verify the iterator is still valid and able to advance to the rest of the
   // tokens.
   // Advance to "cde".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "cde")));
   // Advance to "def".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "def")));
   // Advance to "efg".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "efg")));
   // Advance to "fgh".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "fgh")));
 
   EXPECT_THAT(itr->Advance(), IsFalse());
@@ -329,7 +329,7 @@
   ASSERT_THAT(itr->Advance(), IsTrue());
   ASSERT_THAT(itr->Advance(), IsTrue());
   ASSERT_THAT(itr->Advance(), IsTrue());
-  ASSERT_THAT(itr->GetTokens(),
+  ASSERT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "路去上")));
 
   // Reset backward to after "每". Note that the utf32_offset of "每" is 1.
@@ -337,22 +337,22 @@
 
   // Verify the iterator is pointing to "天走路" since it is the leftmost
   // trigram token after "每".
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "天走路")));
 
   // Verify the iterator is still valid and able to advance to the rest of the
   // tokens.
   // Advance to "走路去".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "走路去")));
   // Advance to "路去上".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "路去上")));
   // Advance to "去上班".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "去上班")));
 
   EXPECT_THAT(itr->Advance(), IsFalse());
@@ -363,26 +363,26 @@
 
   // Verify the iterator is pointing to "每天走" since it is the leftmost
   // trigram token after utf32_offset 0.
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "每天走")));
 
   // Verify the iterator is still valid and able to advance to the rest of the
   // tokens.
   // Advance to "天走路".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "天走路")));
   // Advance to "走路去".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "走路去")));
   // Advance to "路去上".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "路去上")));
   // Advance to "去上班".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "去上班")));
 
   EXPECT_THAT(itr->Advance(), IsFalse());
@@ -397,7 +397,7 @@
 
   // Advance 1 time to "abc".
   ASSERT_THAT(itr->Advance(), IsTrue());
-  ASSERT_THAT(itr->GetTokens(),
+  ASSERT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "abc")));
 
   // Reset forward to after utf32_offset 2.
@@ -405,18 +405,18 @@
 
   // Verify the iterator is pointing to "def" since it is the leftmost trigram
   // token after utf32_offset 2.
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "def")));
 
   // Verify the iterator is still valid and able to advance to the rest of the
   // tokens.
   // Advance to "efg".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "efg")));
   // Advance to "fgh".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "fgh")));
 
   EXPECT_THAT(itr->Advance(), IsFalse());
@@ -431,7 +431,7 @@
 
   // Advance 1 time to "我每天".
   ASSERT_THAT(itr->Advance(), IsTrue());
-  ASSERT_THAT(itr->GetTokens(),
+  ASSERT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "我每天")));
 
   // Reset forward to after "天". Note that the utf32_offset of "天" is 2.
@@ -439,18 +439,18 @@
 
   // Verify the iterator is pointing to "走路去" since it is the leftmost
   // trigram token after "天".
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "走路去")));
 
   // Verify the iterator is still valid and able to advance to the rest of the
   // tokens.
   // Advance to "路去上".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "路去上")));
   // Advance to "去上班".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "去上班")));
 
   EXPECT_THAT(itr->Advance(), IsFalse());
@@ -467,7 +467,7 @@
   // Advance twice to "bcd".
   ASSERT_THAT(itr->Advance(), IsTrue());
   ASSERT_THAT(itr->Advance(), IsTrue());
-  ASSERT_THAT(itr->GetTokens(),
+  ASSERT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "bcd")));
 
   // Reset to after utf32_offset -1.
@@ -475,18 +475,18 @@
 
   // Verify the iterator is pointing to "abc" since it is the leftmost trigram
   // token after utf32_offset -1.
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "abc")));
 
   // Verify the iterator is still valid and able to advance to the rest of the
   // tokens.
   // Advance to "bcd".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "bcd")));
   // Advance to "cde".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "cde")));
 
   EXPECT_THAT(itr->Advance(), IsFalse());
@@ -496,18 +496,18 @@
 
   // Verify the iterator is pointing to "abc" since it is the leftmost trigram
   // token after utf32_offset -2.
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "abc")));
 
   // Verify the iterator is still valid and able to advance to the rest of the
   // tokens.
   // Advance to "bcd".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "bcd")));
   // Advance to "cde".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "cde")));
 }
 
@@ -522,7 +522,7 @@
   // Advance twice to "bcd".
   ASSERT_THAT(itr->Advance(), IsTrue());
   ASSERT_THAT(itr->Advance(), IsTrue());
-  ASSERT_THAT(itr->GetTokens(),
+  ASSERT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "bcd")));
 
   // Reset to after utf32_offset 2. Since the last token is "cde" and the
@@ -556,7 +556,7 @@
 
   // Advance once to "我每天".
   ASSERT_THAT(itr->Advance(), IsTrue());
-  ASSERT_THAT(itr->GetTokens(),
+  ASSERT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "我每天")));
 
   // Reset to after "天" (with utf32_offset 2). Since the last token is
@@ -595,14 +595,14 @@
 
   // Verify the iterator is pointing to "cde" since it is the leftmost trigram
   // token after utf32_offset 1.
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "cde")));
 
   // Verify the iterator is still valid and able to advance to the rest of the
   // tokens.
   // Advance to "def".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "def")));
 
   EXPECT_THAT(itr->Advance(), IsFalse());
@@ -622,14 +622,14 @@
 
   // Verify the iterator is pointing to "天走路" since it is the leftmost
   // trigram token after utf32_offset 1.
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "天走路")));
 
   // Verify the iterator is still valid and able to advance to the rest of the
   // tokens.
   // Advance to "走路去".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "走路去")));
 
   EXPECT_THAT(itr->Advance(), IsFalse());
@@ -705,14 +705,14 @@
 
   // Verify the iterator is pointing to "cde" since it is the leftmost trigram
   // token after utf32_offset 1.
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "cde")));
 
   // Verify the iterator is still valid and able to advance to the rest of the
   // tokens.
   // Advance to "def".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "def")));
 
   EXPECT_THAT(itr->Advance(), IsFalse());
@@ -739,14 +739,14 @@
 
   // Verify the iterator is pointing to "天走路" since it is the leftmost
   // trigram token after utf32_offset 1.
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "天走路")));
 
   // Verify the iterator is still valid and able to advance to the rest of the
   // tokens.
   // Advance to "走路去".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "走路去")));
 
   EXPECT_THAT(itr->Advance(), IsFalse());
@@ -822,7 +822,7 @@
   ASSERT_THAT(itr->Advance(), IsTrue());
   ASSERT_THAT(itr->Advance(), IsTrue());
   ASSERT_THAT(itr->Advance(), IsTrue());
-  ASSERT_THAT(itr->GetTokens(),
+  ASSERT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "efg")));
 
   // Reset backward to before utf32_offset 5.
@@ -830,22 +830,22 @@
 
   // Verify the iterator is pointing to "cde" since it is the rightmost trigram
   // token before utf32_offset 5.
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "cde")));
 
   // Verify the iterator is still valid and able to advance to the rest of the
   // tokens.
   // Advance to "def".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "def")));
   // Advance to "efg".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "efg")));
   // Advance to "fgh".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "fgh")));
 
   EXPECT_THAT(itr->Advance(), IsFalse());
@@ -864,7 +864,7 @@
   ASSERT_THAT(itr->Advance(), IsTrue());
   ASSERT_THAT(itr->Advance(), IsTrue());
   ASSERT_THAT(itr->Advance(), IsTrue());
-  ASSERT_THAT(itr->GetTokens(),
+  ASSERT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "路去上")));
 
   // Reset backward to before "去". Note that the utf32_offset of "去" is 5.
@@ -872,22 +872,22 @@
 
   // Verify the iterator is pointing to "天走路" since it is the rightmost
   // trigram token before "每".
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "天走路")));
 
   // Verify the iterator is still valid and able to advance to the rest of the
   // tokens.
   // Advance to "走路去".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "走路去")));
   // Advance to "路去上".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "路去上")));
   // Advance to "去上班".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "去上班")));
 
   EXPECT_THAT(itr->Advance(), IsFalse());
@@ -902,7 +902,7 @@
 
   // Advance 1 time to "abc".
   ASSERT_THAT(itr->Advance(), IsTrue());
-  ASSERT_THAT(itr->GetTokens(),
+  ASSERT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "abc")));
 
   // Reset forward to before utf32_offset 6.
@@ -910,18 +910,18 @@
 
   // Verify the iterator is pointing to "def" since it is the rightmost trigram
   // token before utf32_offset 6.
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "def")));
 
   // Verify the iterator is still valid and able to advance to the rest of the
   // tokens.
   // Advance to "efg".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "efg")));
   // Advance to "fgh".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "fgh")));
 
   EXPECT_THAT(itr->Advance(), IsFalse());
@@ -936,7 +936,7 @@
 
   // Advance 1 time to "我每天".
   ASSERT_THAT(itr->Advance(), IsTrue());
-  ASSERT_THAT(itr->GetTokens(),
+  ASSERT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "我每天")));
 
   // Reset forward to before "上". Note that the utf32_offset of "上" is 6.
@@ -944,18 +944,18 @@
 
   // Verify the iterator is pointing to "走路去" since it is the rightmost
   // trigram token before "上".
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "走路去")));
 
   // Verify the iterator is still valid and able to advance to the rest of the
   // tokens.
   // Advance to "路去上".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "路去上")));
   // Advance to "去上班".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "去上班")));
 
   EXPECT_THAT(itr->Advance(), IsFalse());
@@ -972,7 +972,7 @@
 
   // Advance once to "abc".
   ASSERT_THAT(itr->Advance(), IsTrue());
-  ASSERT_THAT(itr->GetTokens(),
+  ASSERT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "abc")));
 
   // Reset to before utf32_offset 5. Since the last token "cde" ends at
@@ -982,7 +982,7 @@
 
   // Verify the iterator is pointing to "cde" since it is the rightmost trigram
   // token before utf32_offset 5.
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "cde")));
 
   // Since it is the last token, the next Advance() should fail.
@@ -998,7 +998,7 @@
 
   // Verify the iterator is pointing to "cde" since it is the rightmost trigram
   // token before utf32_offset 6.
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "cde")));
 
   // Since it is the last token, the next Advance() should fail.
@@ -1016,7 +1016,7 @@
 
   // Advance once to "我每天".
   ASSERT_THAT(itr->Advance(), IsTrue());
-  ASSERT_THAT(itr->GetTokens(),
+  ASSERT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "我每天")));
 
   // Reset to before utf32_offset 5. Since the last token "天走路" ends at
@@ -1026,7 +1026,7 @@
 
   // Verify the iterator is pointing to "天走路" since it is the rightmost
   // trigram token before utf32_offset 5.
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "天走路")));
 
   // Since it is the last token, the next Advance() should fail.
@@ -1042,7 +1042,7 @@
 
   // Verify the iterator is pointing to "天走路" since it is the rightmost
   // trigram token before utf32_offset 6.
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "天走路")));
 
   // Since it is the last token, the next Advance() should fail.
@@ -1060,7 +1060,7 @@
   // Advance twice to "bcd".
   ASSERT_THAT(itr->Advance(), IsTrue());
   ASSERT_THAT(itr->Advance(), IsTrue());
-  ASSERT_THAT(itr->GetTokens(),
+  ASSERT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "bcd")));
 
   // Reset to before utf32_offset 2. Since the first token is "abc" and the
@@ -1091,7 +1091,7 @@
 
   // Advance once to "我每天".
   ASSERT_THAT(itr->Advance(), IsTrue());
-  ASSERT_THAT(itr->GetTokens(),
+  ASSERT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "我每天")));
 
   // Reset to before "天" (with utf32_offset 2). Since the first token is
@@ -1126,14 +1126,14 @@
 
   // Verify the iterator is pointing to "cde" since it is the rightmost trigram
   // token before utf32_offset 5.
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "cde")));
 
   // Verify the iterator is still valid and able to advance to the rest of the
   // tokens.
   // Advance to "def".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "def")));
 
   EXPECT_THAT(itr->Advance(), IsFalse());
@@ -1153,14 +1153,14 @@
 
   // Verify the iterator is pointing to "天走路" since it is the rightmost
   // trigram token before utf32_offset 5.
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "天走路")));
 
   // Verify the iterator is still valid and able to advance to the rest of the
   // tokens.
   // Advance to "走路去".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "走路去")));
 
   EXPECT_THAT(itr->Advance(), IsFalse());
@@ -1240,14 +1240,14 @@
 
   // Verify the iterator is pointing to "cde" since it is the rightmost trigram
   // token before utf32_offset 5.
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "cde")));
 
   // Verify the iterator is still valid and able to advance to the rest of the
   // tokens.
   // Advance to "def".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "def")));
 
   EXPECT_THAT(itr->Advance(), IsFalse());
@@ -1274,14 +1274,14 @@
 
   // Verify the iterator is pointing to "天走路" since it is the rightmost
   // trigram token before utf32_offset 5.
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "天走路")));
 
   // Verify the iterator is still valid and able to advance to the rest of the
   // tokens.
   // Advance to "走路去".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "走路去")));
 
   EXPECT_THAT(itr->Advance(), IsFalse());
@@ -1360,29 +1360,29 @@
   ASSERT_THAT(itr->Advance(), IsTrue());
   ASSERT_THAT(itr->Advance(), IsTrue());
   ASSERT_THAT(itr->Advance(), IsTrue());
-  ASSERT_THAT(itr->GetTokens(),
+  ASSERT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "cde")));
 
   // Call ResetToStart().
   EXPECT_THAT(itr->ResetToStart(), IsTrue());
 
   // Verify the iterator is pointing to the start trigram "abc".
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "abc")));
 
   // Verify the iterator is still valid and able to advance to the rest of the
   // tokens.
   // Advance to "bcd".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "bcd")));
   // Advance to "cde".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "cde")));
   // Advance to "def".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "def")));
 
   EXPECT_THAT(itr->Advance(), IsFalse());
@@ -1399,29 +1399,29 @@
   ASSERT_THAT(itr->Advance(), IsTrue());
   ASSERT_THAT(itr->Advance(), IsTrue());
   ASSERT_THAT(itr->Advance(), IsTrue());
-  ASSERT_THAT(itr->GetTokens(),
+  ASSERT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "天走路")));
 
   // Call ResetToStart().
   EXPECT_THAT(itr->ResetToStart(), IsTrue());
 
   // Verify the iterator is pointing to the start trigram "我每天".
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "我每天")));
 
   // Verify the iterator is still valid and able to advance to the rest of the
   // tokens.
   // Advance to "每天走".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "每天走")));
   // Advance to "天走路".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "天走路")));
   // Advance to "走路去".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "走路去")));
 
   EXPECT_THAT(itr->Advance(), IsFalse());
@@ -1439,18 +1439,18 @@
   EXPECT_THAT(itr->ResetToStart(), IsTrue());
 
   // Verify the iterator is pointing to "abc".
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "abc")));
 
   // Verify the iterator is still valid and able to advance to the rest of the
   // tokens.
   // Advance to "bcd".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "bcd")));
   // Advance to "cde".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "cde")));
 
   EXPECT_THAT(itr->Advance(), IsFalse());
@@ -1468,18 +1468,18 @@
   EXPECT_THAT(itr->ResetToStart(), IsTrue());
 
   // Verify the iterator is pointing to "我每天".
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "我每天")));
 
   // Verify the iterator is still valid and able to advance to the rest of the
   // tokens.
   // Advance to "每天走".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "每天走")));
   // Advance to "天走路".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "天走路")));
 
   EXPECT_THAT(itr->Advance(), IsFalse());
@@ -1543,22 +1543,22 @@
   EXPECT_THAT(itr->ResetToStart(), IsTrue());
 
   // Verify the iterator is pointing to the start trigram "abc".
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "abc")));
 
   // Verify the iterator is still valid and able to advance to the rest of the
   // tokens.
   // Advance to "bcd".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "bcd")));
   // Advance to "cde".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "cde")));
   // Advance to "def".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "def")));
 
   EXPECT_THAT(itr->Advance(), IsFalse());
@@ -1583,22 +1583,22 @@
   EXPECT_THAT(itr->ResetToStart(), IsTrue());
 
   // Verify the iterator is pointing to the start trigram "我每天".
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "我每天")));
 
   // Verify the iterator is still valid and able to advance to the rest of the
   // tokens.
   // Advance to "每天走".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "每天走")));
   // Advance to "天走路".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "天走路")));
   // Advance to "走路去".
   EXPECT_THAT(itr->Advance(), IsTrue());
-  EXPECT_THAT(itr->GetTokens(),
+  EXPECT_THAT(itr->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::TRIGRAM, "走路去")));
 
   EXPECT_THAT(itr->Advance(), IsFalse());

diff --git a/icing/tokenization/verbatim-tokenizer.cc b/icing/tokenization/verbatim-tokenizer.cc
index b7745f8..1b6bd31 100644
--- a/icing/tokenization/verbatim-tokenizer.cc
+++ b/icing/tokenization/verbatim-tokenizer.cc

@@ -44,14 +44,11 @@
     return true;
   }
 
-  std::vector<Token> GetTokens() const override {
-    std::vector<Token> result;
-
+  void GetTokens(std::vector<Token>* out_tokens) const override {
+    out_tokens->clear();
     if (!term_.empty() && has_advanced_to_end_) {
-      result.push_back(Token(Token::Type::VERBATIM, term_));
+      out_tokens->emplace_back(Token::Type::VERBATIM, term_);
     }
-
-    return result;
   }
 
   libtextclassifier3::StatusOr<CharacterIterator> CalculateTokenStart()
@@ -148,9 +145,10 @@
   ICING_ASSIGN_OR_RETURN(std::unique_ptr<Tokenizer::Iterator> iterator,
                          Tokenize(text));
   std::vector<Token> tokens;
+  std::vector<Token> batch_tokens;
   while (iterator->Advance()) {
-    std::vector<Token> batch = iterator->GetTokens();
-    tokens.insert(tokens.end(), batch.begin(), batch.end());
+    iterator->GetTokens(&batch_tokens);
+    tokens.insert(tokens.end(), batch_tokens.begin(), batch_tokens.end());
   }
   return tokens;
 }

diff --git a/icing/tokenization/verbatim-tokenizer_test.cc b/icing/tokenization/verbatim-tokenizer_test.cc
index 8e4418b..3fff6f2 100644
--- a/icing/tokenization/verbatim-tokenizer_test.cc
+++ b/icing/tokenization/verbatim-tokenizer_test.cc

@@ -97,7 +97,7 @@
   auto token_iterator = verbatim_tokenizer->Tokenize(kText).ValueOrDie();
 
   // We should get no tokens if we get the token before advancing.
-  EXPECT_THAT(token_iterator->GetTokens(), IsEmpty());
+  EXPECT_THAT(token_iterator->GetTokensForTest(), IsEmpty());
 }
 
 TEST_F(VerbatimTokenizerTest, ResetToTokenEndingBefore) {
@@ -112,13 +112,13 @@
   // Reset to beginning of verbatim of token. We provide an offset of 13 as it
   // is larger than the final index (12) of the verbatim token.
   EXPECT_TRUE(token_iterator->ResetToTokenEndingBefore(13));
-  EXPECT_THAT(token_iterator->GetTokens(),
+  EXPECT_THAT(token_iterator->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::VERBATIM, "Hello, world!")));
 
   // Ensure our cached character iterator propertly maintains the end of the
   // verbatim token.
   EXPECT_TRUE(token_iterator->ResetToTokenEndingBefore(13));
-  EXPECT_THAT(token_iterator->GetTokens(),
+  EXPECT_THAT(token_iterator->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::VERBATIM, "Hello, world!")));
 
   // We should not be able to reset with an offset before or within
@@ -138,7 +138,7 @@
 
   // Get token without resetting
   EXPECT_TRUE(token_iterator->Advance());
-  EXPECT_THAT(token_iterator->GetTokens(),
+  EXPECT_THAT(token_iterator->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::VERBATIM, "Hello, world!")));
 
   // We expect a sole verbatim token, so it's not possible to reset after the
@@ -148,7 +148,7 @@
   // We expect to be reset to the sole verbatim token when the offset is
   // negative.
   EXPECT_TRUE(token_iterator->ResetToTokenStartingAfter(-1));
-  EXPECT_THAT(token_iterator->GetTokens(),
+  EXPECT_THAT(token_iterator->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::VERBATIM, "Hello, world!")));
 }
 
@@ -163,12 +163,12 @@
 
   // Get token without resetting
   EXPECT_TRUE(token_iterator->Advance());
-  EXPECT_THAT(token_iterator->GetTokens(),
+  EXPECT_THAT(token_iterator->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::VERBATIM, "Hello, world!")));
 
   // Retrieve token again after resetting to start
   EXPECT_TRUE(token_iterator->ResetToStart());
-  EXPECT_THAT(token_iterator->GetTokens(),
+  EXPECT_THAT(token_iterator->GetTokensForTest(),
               ElementsAre(EqualsToken(Token::Type::VERBATIM, "Hello, world!")));
 }
 

diff --git a/icing/transform/icu/icu-normalizer-factory.cc b/icing/transform/icu/icu-normalizer-factory.cc
index 290fbcc..ebc6e97 100644
--- a/icing/transform/icu/icu-normalizer-factory.cc
+++ b/icing/transform/icu/icu-normalizer-factory.cc

@@ -23,25 +23,37 @@
 #include "icing/transform/icu/icu-normalizer.h"
 #include "icing/transform/normalizer-options.h"
 #include "icing/transform/normalizer.h"
+#include "icing/util/status-util.h"
 
 namespace icing {
 namespace lib {
 
 namespace normalizer_factory {
 
-// Creates an ICU-based normalizer.
+using ::icing::lib::status_util::TransformStatus;
+
+// Creates an ICU-based normalizer based on the provided options.
+//
+// @param options: The options for creating the normalizer.
+// @param icu_normalizer_creation_status: Optional output parameter that will be
+//        populated with the status of IcuNormalizer.
 //
 // Returns:
 //   A normalizer on success
 //   INVALID_ARGUMENT_ERROR if options.max_term_byte_size <= 0
 //   INTERNAL_ERROR on errors
 libtextclassifier3::StatusOr<std::unique_ptr<Normalizer>> Create(
-    const NormalizerOptions& options) {
+    const NormalizerOptions& options,
+    StatusProto* icu_normalizer_creation_status) {
   if (options.max_term_byte_size <= 0) {
     return absl_ports::InvalidArgumentError(
         "max_term_byte_size must be greater than zero.");
   }
-  return IcuNormalizer::Create(options.max_term_byte_size);
+  auto icu_normalizer_or = IcuNormalizer::Create(options.max_term_byte_size);
+  if (icu_normalizer_creation_status != nullptr) {
+    TransformStatus(icu_normalizer_or.status(), icu_normalizer_creation_status);
+  }
+  return icu_normalizer_or;
 }
 
 }  // namespace normalizer_factory

diff --git a/icing/transform/icu/icu-normalizer_test.cc b/icing/transform/icu/icu-normalizer_test.cc
index 499c7c1..1d1e63d 100644
--- a/icing/transform/icu/icu-normalizer_test.cc
+++ b/icing/transform/icu/icu-normalizer_test.cc

@@ -46,6 +46,15 @@
   std::unique_ptr<Normalizer> normalizer_;
 };
 
+TEST_F(IcuNormalizerTest, SuccessfulIcuNormalizerSetsOkStatus) {
+  StatusProto icu_normalizer_creation_status;
+  NormalizerOptions options(/*max_term_byte_size=*/1024);
+  ICING_ASSERT_OK_AND_ASSIGN(
+      auto normalizer,
+      normalizer_factory::Create(options, &icu_normalizer_creation_status));
+  EXPECT_THAT(icu_normalizer_creation_status.code(), Eq(StatusProto::OK));
+}
+
 TEST_F(IcuNormalizerTest, Creation) {
   NormalizerOptions options1(/*max_term_byte_size=*/5);
   EXPECT_THAT(normalizer_factory::Create(options1), IsOk());

diff --git a/icing/transform/map/map-normalizer-factory.cc b/icing/transform/map/map-normalizer-factory.cc
index 2c0dc8d..e012303 100644
--- a/icing/transform/map/map-normalizer-factory.cc
+++ b/icing/transform/map/map-normalizer-factory.cc

@@ -16,6 +16,7 @@
 
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "icing/absl_ports/canonical_errors.h"
+#include "icing/proto/status.pb.h"
 #include "icing/transform/map/map-normalizer.h"
 #include "icing/transform/normalizer-options.h"
 #include "icing/transform/normalizer.h"
@@ -25,14 +26,19 @@
 
 namespace normalizer_factory {
 
-// Creates a map-based normalizer.
+// Creates a map-based normalizer based on the provided options.
+//
+// @param options: The options for creating the normalizer.
+// @param icu_normalizer_creation_status: Optional output parameter that will be
+//        ignored.
 //
 // Returns:
 //   A normalizer on success
 //   INVALID_ARGUMENT_ERROR if options.max_term_byte_size <= 0
 //   INTERNAL_ERROR on errors
 libtextclassifier3::StatusOr<std::unique_ptr<Normalizer>> Create(
-    const NormalizerOptions& options) {
+    const NormalizerOptions& options,
+    StatusProto* /*icu_normalizer_creation_status*/) {
   if (options.max_term_byte_size <= 0) {
     return absl_ports::InvalidArgumentError(
         "normalizer_max_term_byte_size must be greater than zero.");

diff --git a/icing/transform/normalizer-factory.h b/icing/transform/normalizer-factory.h
index 6a9709c..0935c8d 100644
--- a/icing/transform/normalizer-factory.h
+++ b/icing/transform/normalizer-factory.h

@@ -18,6 +18,7 @@
 #include <memory>
 
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
+#include "icing/proto/status.pb.h"
 #include "icing/transform/normalizer-options.h"
 #include "icing/transform/normalizer.h"
 
@@ -26,14 +27,20 @@
 
 namespace normalizer_factory {
 
-// Creates a normalizer.
+// Creates a normalizer based on the provided options.
+//
+// @param options: The options for creating the normalizer.
+// @param icu_normalizer_creation_status: Optional output parameter that will be
+//        populated with the status of IcuNormalizer if it's creation is
+//        attempted.
 //
 // Returns:
 //   A normalizer on success
 //   INVALID_ARGUMENT if options.max_term_byte_size <= 0
 //   INTERNAL_ERROR on errors
 libtextclassifier3::StatusOr<std::unique_ptr<Normalizer>> Create(
-    const NormalizerOptions& options);
+    const NormalizerOptions& options,
+    StatusProto* icu_normalizer_creation_status = nullptr);
 
 }  // namespace normalizer_factory
 

diff --git a/icing/util/crc32.cc b/icing/util/crc32.cc
index c182f31..a8a6553 100644
--- a/icing/util/crc32.cc
+++ b/icing/util/crc32.cc

@@ -46,6 +46,11 @@
   return crc_;
 }
 
+uint32_t Crc32::Combine(const Crc32& other_crc, int other_data_size) {
+  crc_ = crc32_combine(crc_, other_crc.Get(), other_data_size);
+  return crc_;
+}
+
 libtextclassifier3::StatusOr<uint32_t> Crc32::UpdateWithXor(
     const std::string_view xored_str, int full_data_size, int position) {
   // For appending, use Append().

diff --git a/icing/util/crc32.h b/icing/util/crc32.h
index dfa7756..ab65eb1 100644
--- a/icing/util/crc32.h
+++ b/icing/util/crc32.h

@@ -62,6 +62,16 @@
   // Crc32(base_crc).Append(str) is not the same as zlib::crc32(base_crc, str);
   uint32_t Append(std::string_view str);
 
+  // Combines other_crc to the current crc for bytes concatenation. The caller
+  // must provide the raw data (byte) size of other_crc.
+  //
+  // NOTE: all the following are equivalent:
+  // 1) crc32.Append(str_a + str_b);
+  // 2) crc32.Append(str_a); crc32.Append(str_b);
+  // 3) crc32.Append(str_a); crc32_b.Append(str_b);
+  //    crc32.Combine(crc32_b, str_b.size());
+  uint32_t Combine(const Crc32& other_crc, int other_data_size);
+
   // Update a string's rolling crc when some content is modified in the middle
   // at an offset. We need the xored_str, which is the new value xored with the
   // original value.

diff --git a/icing/util/crc32_test.cc b/icing/util/crc32_test.cc
index d6f4c46..5351b7a 100644
--- a/icing/util/crc32_test.cc
+++ b/icing/util/crc32_test.cc

@@ -30,7 +30,9 @@
 namespace lib {
 
 namespace {
+
 using ::testing::Eq;
+using ::testing::Ne;
 
 void UpdateAtRandomOffset(std::string* buf, uint32_t* update_xor, int* offset) {
   // The max value of rand() is 2^31 - 1 (2147483647) but the max value of
@@ -79,6 +81,60 @@
   EXPECT_THAT(crc32_foo_and_bar.Get(), Eq(crc32_foobar.Get()));
 }
 
+TEST(Crc32Test, Combine) {
+  Crc32 crc32_foo;
+  crc32_foo.Append("foo");
+
+  Crc32 crc32_barbaz;
+  crc32_barbaz.Append("barbaz");
+
+  Crc32 crc32_foobarbaz;
+  crc32_foobarbaz.Append("foobarbaz");
+
+  EXPECT_THAT(crc32_foo.Combine(crc32_barbaz, 6), Eq(crc32_foobarbaz.Get()));
+  EXPECT_THAT(crc32_foo.Get(), Eq(crc32_foobarbaz.Get()));
+}
+
+TEST(Crc32Test, Combine_incorrectResultWithWrongSize) {
+  Crc32 crc32_foo;
+  crc32_foo.Append("foo");
+
+  Crc32 crc32_barbaz;
+  crc32_barbaz.Append("barbaz");
+
+  Crc32 crc32_foobarbaz;
+  crc32_foobarbaz.Append("foobarbaz");
+
+  int wrong_size = 5;
+  EXPECT_THAT(crc32_foo.Combine(crc32_barbaz, wrong_size),
+              Ne(crc32_foobarbaz.Get()));
+  EXPECT_THAT(crc32_foo.Get(), Ne(crc32_foobarbaz.Get()));
+}
+
+TEST(Crc32Test, Combine_withEmptyString) {
+  Crc32 crc32_foo;
+  crc32_foo.Append("foo");
+  uint32_t crc_foo_val = crc32_foo.Get();
+
+  Crc32 crc32_empty;
+  crc32_empty.Append("");
+
+  EXPECT_THAT(crc32_foo.Combine(crc32_empty, 0), Eq(crc_foo_val));
+  EXPECT_THAT(crc32_foo.Get(), Eq(crc_foo_val));
+}
+
+TEST(Crc32Test, Combine_toEmptyString) {
+  Crc32 crc32_empty;
+  crc32_empty.Append("");
+
+  Crc32 crc32_foo;
+  crc32_foo.Append("foo");
+  uint32_t crc_foo_val = crc32_foo.Get();
+
+  EXPECT_THAT(crc32_empty.Combine(crc32_foo, 3), Eq(crc_foo_val));
+  EXPECT_THAT(crc32_empty.Get(), Eq(crc_foo_val));
+}
+
 TEST(Crc32Test, UpdateAtPosition) {
   std::string buf;
   buf.resize(1000);

diff --git a/icing/util/document-util.cc b/icing/util/document-util.cc
new file mode 100644
index 0000000..d3834e6
--- /dev/null
+++ b/icing/util/document-util.cc

@@ -0,0 +1,36 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "icing/util/document-util.h"
+
+#include <utility>
+
+#include "icing/proto/document.pb.h"
+#include "icing/proto/document_wrapper.pb.h"
+
+namespace icing {
+namespace lib {
+
+namespace document_util {
+
+DocumentWrapper CreateDocumentWrapper(DocumentProto document) {
+  DocumentWrapper document_wrapper;
+  *document_wrapper.mutable_document() = std::move(document);
+  return document_wrapper;
+}
+
+}  // namespace document_util
+
+}  // namespace lib
+}  // namespace icing

diff --git a/icing/util/document-util.h b/icing/util/document-util.h
new file mode 100644
index 0000000..c9e117a
--- /dev/null
+++ b/icing/util/document-util.h

@@ -0,0 +1,34 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef ICING_UTIL_DOCUMENT_UTIL_H_
+#define ICING_UTIL_DOCUMENT_UTIL_H_
+
+#include "icing/proto/document.pb.h"
+#include "icing/proto/document_wrapper.pb.h"
+
+namespace icing {
+namespace lib {
+
+namespace document_util {
+
+// Creates a DocumentWrapper from a DocumentProto.
+DocumentWrapper CreateDocumentWrapper(DocumentProto document);
+
+}  // namespace document_util
+
+}  // namespace lib
+}  // namespace icing
+
+#endif  // ICING_UTIL_DOCUMENT_UTIL_H_

diff --git a/icing/util/document-util_test.cc b/icing/util/document-util_test.cc
new file mode 100644
index 0000000..c93ad6d
--- /dev/null
+++ b/icing/util/document-util_test.cc

@@ -0,0 +1,47 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "icing/util/document-util.h"
+
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+#include "icing/document-builder.h"
+#include "icing/portable/equals-proto.h"
+#include "icing/proto/document_wrapper.pb.h"
+
+namespace icing {
+namespace lib {
+namespace document_util {
+
+namespace {
+
+using ::icing::lib::portable_equals_proto::EqualsProto;
+
+TEST(DocumentUtilTest, CreateDocumentWrapper) {
+  DocumentProto document = DocumentBuilder()
+                               .SetKey("icing", "fake_type/1")
+                               .SetSchema("FakeType")
+                               .AddStringProperty("prop1", "foo", "bar", "baz")
+                               .Build();
+
+  DocumentWrapper document_wrapper = CreateDocumentWrapper(document);
+  EXPECT_THAT(document_wrapper.document(), EqualsProto(document));
+}
+
+}  // namespace
+
+}  // namespace document_util
+
+}  // namespace lib
+}  // namespace icing

diff --git a/icing/util/document-validator.cc b/icing/util/document-validator.cc
index 916b34e..8c68efa 100644
--- a/icing/util/document-validator.cc
+++ b/icing/util/document-validator.cc

@@ -80,7 +80,7 @@
   auto type_config_or = schema_store_->GetSchemaTypeConfig(document.schema());
   if (!type_config_or.ok()) {
     ICING_LOG(ERROR) << type_config_or.status().error_message()
-                     << "Error while validating document ("
+                     << " Error while validating document ("
                      << document.namespace_() << ", " << document.uri() << ")";
     return type_config_or.status();
   }
@@ -114,7 +114,8 @@
           document.namespace_(), ", ", document.uri(),
           ") of type: ", document.schema(), "."));
     }
-    const PropertyConfigProto& property_config = *property_iter->second;
+    const PropertyConfigProto& property_config =
+        *property_iter->second.property_config;
 
     // Get the property value size according to data type.
     int value_size = 0;

diff --git a/icing/util/embedding-util.h b/icing/util/embedding-util.h
index 5026051..1390d4e 100644
--- a/icing/util/embedding-util.h
+++ b/icing/util/embedding-util.h

@@ -15,6 +15,7 @@
 #ifndef ICING_UTIL_EMBEDDING_UTIL_H_
 #define ICING_UTIL_EMBEDDING_UTIL_H_
 
+#include <array>
 #include <string_view>
 
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
@@ -27,6 +28,15 @@
 
 namespace embedding_util {
 
+// The list of embedding query metric types, with the order matching the
+// enum value.
+static const std::array<SearchSpecProto::EmbeddingQueryMetricType::Code, 3>
+    kEmbeddingQueryMetricTypes = {
+        SearchSpecProto::EmbeddingQueryMetricType::COSINE,       // value = 1
+        SearchSpecProto::EmbeddingQueryMetricType::DOT_PRODUCT,  // value = 2
+        SearchSpecProto::EmbeddingQueryMetricType::EUCLIDEAN     // value = 3
+};
+
 inline libtextclassifier3::StatusOr<
     SearchSpecProto::EmbeddingQueryMetricType::Code>
 GetEmbeddingQueryMetricTypeFromName(std::string_view metric_name) {

diff --git a/icing/util/logging.cc b/icing/util/logging.cc
index f60526b..fe14be0 100644
--- a/icing/util/logging.cc
+++ b/icing/util/logging.cc

@@ -64,6 +64,11 @@
 // The last 16 bits represent the current verbosity.
 std::atomic<uint32_t> global_logging_level = DEFAULT_LOGGING_LEVEL;
 
+// TODO(b/401363381): Remove this once we have a better way to log to
+// /dev/hvc2 in isolated storage.
+// Indicate whether we should force logging to /dev/hvc2 for ICING_LOG.
+std::atomic<bool> global_force_debug_logs = false;
+
 }  // namespace
 
 // Whether we should log according to the current logging level.
@@ -102,11 +107,30 @@
   return true;
 }
 
+// TODO(b/401363381): Remove this once we have a better way to log to
+// /dev/hvc2 in isolated storage.
+void SetForceDebugLogging(bool force) {
+  // Using the relaxed order for better performance because we only need to
+  // guarantee the atomicity for this specific statement, without the need to
+  // worry about reordering.
+  global_force_debug_logs.store(force, std::memory_order_relaxed);
+}
+
+// TODO(b/401363381): Remove this once we have a better way to log to
+// /dev/hvc2 in isolated storage.
+bool GetForceDebugLogging() {
+  // Using the relaxed order for better performance because we only need to
+  // guarantee the atomicity for this specific statement, without the need to
+  // worry about reordering.
+  return global_force_debug_logs.load(std::memory_order_relaxed);
+}
+
 LogMessage::LogMessage(LogSeverity::Code severity, uint16_t verbosity,
                        const char *file_name, int line_number)
     : severity_(severity),
       verbosity_(verbosity),
       should_log_(ShouldLog(severity_, verbosity_)),
+      force_debug_logs_(GetForceDebugLogging()),
       stream_(should_log_) {
   if (should_log_) {
     stream_ << JumpToBasename(file_name) << ":" << line_number << ": ";
@@ -115,7 +139,8 @@
 
 LogMessage::~LogMessage() {
   if (should_log_) {
-    LowLevelLogging(severity_, kIcingLoggingTag, stream_.message);
+    LowLevelLogging(severity_, kIcingLoggingTag, stream_.message,
+                    force_debug_logs_);
   }
   if (severity_ == LogSeverity::FATAL) {
     std::terminate();  // Will print a stacktrace (stdout or logcat).

diff --git a/icing/util/logging.h b/icing/util/logging.h
index 23280dc..7eaf776 100644
--- a/icing/util/logging.h
+++ b/icing/util/logging.h

@@ -39,6 +39,16 @@
 // The function will always return false when verbosity is negative.
 bool SetLoggingLevel(LogSeverity::Code severity, int16_t verbosity = 0);
 
+// TODO(b/401363381): Remove this once we have a better way to log to
+// /dev/hvc2 in isolated storage.
+// Indicate whether we should force logging to /dev/hvc2 for ICING_LOG.
+void SetForceDebugLogging(bool force);
+
+// TODO(b/401363381): Remove this once we have a better way to log to
+// /dev/hvc2 in isolated storage.
+// Get whether or not we should force logging to /dev/hvc2 for ICING_LOG.
+bool GetForceDebugLogging();
+
 // A tiny code footprint string stream for assembling log messages.
 struct LoggingStringStream {
   explicit LoggingStringStream(bool should_log) : should_log_(should_log) {}
@@ -68,6 +78,14 @@
 }
 
 inline LoggingStringStream& operator<<(LoggingStringStream& stream,
+                                       char* message) {
+  if (stream.should_log_) {
+    stream.message.append(message);
+  }
+  return stream;
+}
+
+inline LoggingStringStream& operator<<(LoggingStringStream& stream,
                                        const char* message) {
   if (stream.should_log_) {
     stream.message.append(message);
@@ -84,6 +102,14 @@
 }
 
 inline LoggingStringStream& operator<<(LoggingStringStream& stream,
+                                       std::string& message) {
+  if (stream.should_log_) {
+    stream.message.append(message);
+  }
+  return stream;
+}
+
+inline LoggingStringStream& operator<<(LoggingStringStream& stream,
                                        std::string_view message) {
   if (stream.should_log_) {
     stream.message.append(message);
@@ -123,6 +149,7 @@
   const LogSeverity::Code severity_;
   const uint16_t verbosity_;
   const bool should_log_;
+  const bool force_debug_logs_;
 
   // Stream that "prints" all info into a string (not to a file).  We construct
   // here the entire logging message and next print it in one operation.

diff --git a/icing/util/logging_raw.cc b/icing/util/logging_raw.cc
index 44dd000..09d0308 100644
--- a/icing/util/logging_raw.cc
+++ b/icing/util/logging_raw.cc

@@ -50,7 +50,7 @@
 }  // namespace
 
 void LowLevelLogging(LogSeverity::Code severity, const std::string& tag,
-                     const std::string& message) {
+                     const std::string& message, const bool force_debug_logs) {
   const int android_log_level = GetAndroidLogLevel(severity);
 #if __ANDROID_API__ >= 30
   if (!__android_log_is_loggable(android_log_level, tag.c_str(),
@@ -58,7 +58,28 @@
     return;
   }
 #endif  // __ANDROID_API__ >= 30
-  __android_log_write(android_log_level, tag.c_str(), message.c_str());
+  // TODO(b/401363381): Remove this once we have a better way to log to
+  // /dev/hvc2 in isolated storage.
+  if (force_debug_logs) {
+    // When isolated icing storage is enabled, the VM debug level determines
+    // whether icing debug logs are delivered. We want the icing debug logs
+    // to always be present. Thus force logging to /dev/hvc2.
+    const char* file_logger_path = "/dev/hvc2";
+    // "e" opens the file with the O_CLOEXEC flag. Icing should not be starting
+    // any other processes, but it is added as a precaution.
+    static FILE* stream = [&file_logger_path]() {
+      FILE* f = fopen(file_logger_path, "ae");
+      if (f != nullptr) {
+        return f;
+      }
+      fprintf(stderr, "Failed to open /dev/hvc2 for logging. "
+                      "Falling back to stderr.\n");
+      return stderr;
+    }();
+    fprintf(stream, "%s: %s\n", tag.c_str(), message.c_str());
+  } else {
+    __android_log_write(android_log_level, tag.c_str(), message.c_str());
+  }
 }
 
 }  // namespace lib
@@ -91,7 +112,7 @@
 }  // namespace
 
 void LowLevelLogging(LogSeverity::Code severity, const std::string &tag,
-                     const std::string &message) {
+                     const std::string &message, const bool force_debug_logs) {
   // TODO(b/146903474) Do not log to stderr for logs other than FATAL and ERROR.
   fprintf(stderr, "[%s] %s : %s\n", LogSeverityToString(severity), tag.c_str(),
           message.c_str());

diff --git a/icing/util/logging_raw.h b/icing/util/logging_raw.h
index 99dddb6..438e1a8 100644
--- a/icing/util/logging_raw.h
+++ b/icing/util/logging_raw.h

@@ -26,7 +26,7 @@
 // severity.  From android/log.h: "the tag normally corresponds to the component
 // that emits the log message, and should be reasonably small".
 void LowLevelLogging(LogSeverity::Code severity, const std::string &tag,
-                     const std::string &message);
+                     const std::string &message, const bool force_debug_logs);
 
 }  // namespace lib
 }  // namespace icing

diff --git a/icing/util/logging_test.cc b/icing/util/logging_test.cc
index eac018e..29e39b5 100644
--- a/icing/util/logging_test.cc
+++ b/icing/util/logging_test.cc

@@ -13,6 +13,8 @@
 // limitations under the License.
 
 #include "icing/util/logging.h"
+#include <cerrno>
+#include <cstring>
 
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
@@ -153,6 +155,18 @@
   EXPECT_THAT(stream2.message, IsEmpty());
 }
 
+TEST(LoggingTest, LoggingStreamStrErrorTest) {
+  ASSERT_TRUE(SetLoggingLevel(LogSeverity::INFO));
+  char* error_message = strerror(ENOENT);
+  // This one should be logged.
+  LoggingStringStream stream1 = (ICING_LOG(INFO) << error_message);
+  EXPECT_THAT(stream1.message, EndsWith("No such file or directory"));
+
+  // This one should not be logged, thus empty.
+  LoggingStringStream stream2 = (ICING_LOG(DBG) << error_message);
+  EXPECT_THAT(stream2.message, IsEmpty());
+}
+
 }  // namespace
 }  // namespace lib
 }  // namespace icing

diff --git a/icing/util/math-util.h b/icing/util/math-util.h
index 22f429f..8aab34e 100644
--- a/icing/util/math-util.h
+++ b/icing/util/math-util.h

@@ -17,6 +17,8 @@
 
 #include <cstdint>
 #include <limits>
+#include <utility>
+#include <vector>
 
 namespace icing {
 namespace lib {
@@ -91,6 +93,38 @@
   return n;
 }
 
+// Applies a permutation to multiple containers in-place, which permutes
+// elements within the given containers according to the provided permutation
+// vector. permutation[i] specifies the new index for the element originally at
+// index i.
+template <typename... ContainerTypes>
+static void ApplyPermutation(const std::vector<int> &permutation,
+                             ContainerTypes &...values) {
+  std::vector<bool> done(permutation.size());
+  // Apply permutation
+  for (int i = 0; i < permutation.size(); ++i) {
+    if (done[i]) {
+      continue;
+    }
+    done[i] = true;
+    int curr = i;
+    int next = permutation[i];
+    // Since every finite permutation is formed by disjoint cycles, we can
+    // start with the current element, at index i, and swap the element at
+    // this position with whatever element that *should* be here. Then,
+    // continue to swap the original element, at its updated positions, with
+    // the element that should be occupying that position until the original
+    // element has reached *its* correct position. This completes applying the
+    // single cycle in the permutation.
+    while (next != i) {
+      (std::swap(values[curr], values[next]), ...);
+      done[next] = true;
+      curr = next;
+      next = permutation[next];
+    }
+  }
+}
+
 }  // namespace math_util
 
 }  // namespace lib

diff --git a/icing/util/simple-task-scheduler.cc b/icing/util/simple-task-scheduler.cc
new file mode 100644
index 0000000..cc3fae5
--- /dev/null
+++ b/icing/util/simple-task-scheduler.cc

@@ -0,0 +1,230 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "icing/util/simple-task-scheduler.h"
+
+#include <chrono>              // NOLINT
+#include <condition_variable>  // NOLINT
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <mutex>   // NOLINT
+#include <thread>  // NOLINT
+#include <utility>
+
+#include "icing/util/clock.h"
+
+namespace icing {
+namespace lib {
+
+std::unique_ptr<SimpleTaskScheduler> SimpleTaskScheduler::Create(
+    const Clock& clock) {
+  auto scheduler =
+      std::unique_ptr<SimpleTaskScheduler>(new SimpleTaskScheduler(clock));
+  scheduler->Initialize();
+  return scheduler;
+}
+
+SimpleTaskScheduler::~SimpleTaskScheduler() {
+  // Notify the thread to terminate.
+  {
+    std::unique_lock<std::mutex> lk(mutex_);  // NOLINT
+
+    is_terminated_ = true;
+  }
+  cv_.notify_all();
+
+  if (thread_ != nullptr) {
+    thread_->join();
+  }
+}
+
+void SimpleTaskScheduler::ScheduleAt(TaskId task_id, std::function<void()> task,
+                                     int64_t new_execution_timestamp_ms) {
+  if (new_execution_timestamp_ms < 0) {
+    new_execution_timestamp_ms = kInvalidTimestampMs;
+  }
+
+  // - Check if the task id exists.
+  //   - If yes, then overwrite the TaskInfo for the task id.
+  //   - If no, then create a new task info and add it into the map.
+  // - Notify the executor thread blocking at wait() or wait_for() to wake up
+  //   and wait for the new execution time.
+  {
+    std::unique_lock<std::mutex> lk(mutex_);  // NOLINT
+
+    auto itr = tasks_.find(task_id);
+    if (itr != tasks_.end()) {
+      itr->second->task =
+          std::make_shared<std::function<void()>>(std::move(task));
+      itr->second->execution_timestamp_ms = new_execution_timestamp_ms;
+    } else {
+      tasks_.insert(
+          {task_id, std::make_unique<TaskInfo>(std::move(task),
+                                               new_execution_timestamp_ms)});
+    }
+  }
+  cv_.notify_all();
+}
+
+void SimpleTaskScheduler::Cancel(TaskId task_id) {
+  {
+    std::unique_lock<std::mutex> lk(mutex_);  // NOLINT
+
+    auto itr = tasks_.find(task_id);
+    if (itr == tasks_.end()) {
+      return;
+    }
+
+    // Set the timestamp to kInvalidTimestampMs to cancel the task.
+    //
+    // Note: we CANNOT remove the entry from the map here. Suppose the executor
+    //   thread is currently waiting for the execution of THIS task (i.e.
+    //   blocking at wait_for()):
+    // - We need kInvalidTimestampMs as a mark to tell the executor thread not
+    //   to execute the task after waking up.
+    // - Also in this case, the executor thread currently holds a TaskInfo
+    //   pointer to THIS task in a local variable. Removing the task entry from
+    //   the map here will cause pointer instability and crash after waking up.
+    //
+    // Therefore, we keep the entry in the map, and delegate the "cleanup" work
+    // to the executor thread.
+    itr->second->execution_timestamp_ms = kInvalidTimestampMs;
+  }
+  cv_.notify_all();
+}
+
+int64_t SimpleTaskScheduler::GetScheduledTimeMs(TaskId task_id) const {
+  std::unique_lock<std::mutex> lk(mutex_);  // NOLINT
+
+  auto itr = tasks_.find(task_id);
+  if (itr != tasks_.end()) {
+    return itr->second->execution_timestamp_ms;
+  }
+  return kInvalidTimestampMs;
+}
+
+void SimpleTaskScheduler::Initialize() {
+  std::unique_lock<std::mutex> lk(mutex_);  // NOLINT
+
+  thread_ = std::make_unique<std::thread>([this]() {
+    while (true) {
+      std::shared_ptr<std::function<void()>> next_task = nullptr;
+
+      {
+        std::unique_lock<std::mutex> lk(mutex_);  // NOLINT
+
+        if (is_terminated_) {
+          break;
+        }
+
+        CleanUpCompletedTasksLocked();
+
+        int64_t current_time_ms = clock_.GetSystemTimeMilliseconds();
+        TaskInfo* next_task_info = GetNextTaskInfoLocked();
+        if (next_task_info == nullptr) {
+          // No scheduled task. Wait for the next task being scheduled. This
+          // will avoid busy-waiting.
+          cv_.wait(lk);
+        } else if (next_task_info->execution_timestamp_ms > current_time_ms) {
+          // Wait for the execution time. It is possible that wait_for() is
+          // interrupted by ScheduleAt(), Cancel(), or spurious wakeup.
+          //
+          // Note: we use wait_for() instead of wait_until() for the convenience
+          //   of mocking the clock in the unit test.
+          std::cv_status wait_status = cv_.wait_for(
+              lk,
+              std::chrono::milliseconds(next_task_info->execution_timestamp_ms -
+                                        current_time_ms));
+          // If the thread is woken up by timeout, execute the task.
+          // Otherwise, the thread is woken up by ScheduleAt(), Cancel() or
+          // spurious wakeup.
+          if (wait_status != std::cv_status::timeout) {
+            // Set the pointer to nullptr indicating that we need another round
+            // of wait() before executing it.
+            next_task_info = nullptr;
+          }
+        }
+
+        // It is possible that is_terminated_ is modified during wait() or
+        // wait_for(). Check it again here.
+        if (is_terminated_) {
+          break;
+        }
+
+        // It is possible that the task is canceled or rescheduled during
+        // wait_for().
+        // - In most cases, wait_status above will NOT be
+        //   std::cv_status::timeout and we should've already set next_task_info
+        //   to nullptr.
+        // - However, it is possible that the cancel/reschedule operation took
+        //   place right at the wait_for() timeout. Therefore, let's check
+        //   execution_timestamp_ms again here.
+        if (next_task_info != nullptr &&
+            next_task_info->execution_timestamp_ms != kInvalidTimestampMs) {
+          // Reset the timestamp before executing the task.
+          next_task_info->execution_timestamp_ms = kInvalidTimestampMs;
+
+          // Get a shared pointer to the task.
+          //
+          // Note: since we execute the task outside of the critical section,
+          //   it is possible that the task is overwritten between the gap of
+          //   the critical section and the actual execution below (i.e. the
+          //   caller calls ScheduleAt() with the same task id + a different
+          //   task implementation and overwrites the task object). Here, we
+          //   have to execute the original task, so we need to use a shared
+          //   pointer for it.
+          next_task = next_task_info->task;
+        }
+      }
+
+      if (next_task != nullptr) {
+        // It is possible that the task is expensive or attempts to acquire
+        // the mutex again and schedule another task, so execute the current one
+        // outside of the critical section to avoid blocking or deadlock.
+        (*next_task)();
+      }
+    }
+  });
+}
+
+SimpleTaskScheduler::TaskInfo* SimpleTaskScheduler::GetNextTaskInfoLocked()
+    const {
+  TaskInfo* next_task_info = nullptr;
+  for (const auto& [_, task_info] : tasks_) {
+    if (task_info->execution_timestamp_ms == kInvalidTimestampMs) {
+      continue;
+    }
+
+    if (next_task_info == nullptr ||
+        task_info->execution_timestamp_ms <
+            next_task_info->execution_timestamp_ms) {
+      next_task_info = task_info.get();
+    }
+  }
+  return next_task_info;
+}
+
+void SimpleTaskScheduler::CleanUpCompletedTasksLocked() {
+  for (auto itr = tasks_.begin(); itr != tasks_.end();) {
+    if (itr->second->execution_timestamp_ms == kInvalidTimestampMs) {
+      itr = tasks_.erase(itr);
+    } else {
+      ++itr;
+    }
+  }
+}
+
+}  // namespace lib
+}  // namespace icing

diff --git a/icing/util/simple-task-scheduler.h b/icing/util/simple-task-scheduler.h
new file mode 100644
index 0000000..6bbd30e
--- /dev/null
+++ b/icing/util/simple-task-scheduler.h

@@ -0,0 +1,170 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef ICING_UTIL_SIMPLE_TASK_SCHEDULER_H_
+#define ICING_UTIL_SIMPLE_TASK_SCHEDULER_H_
+
+#include <condition_variable>  // NOLINT
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <mutex>   // NOLINT
+#include <thread>  // NOLINT
+#include <unordered_map>
+#include <utility>
+
+#include "icing/util/clock.h"
+
+namespace icing {
+namespace lib {
+
+// A utility class to schedule custom tasks to fire at a specific timestamp.
+//
+// Example:
+//   Clock clock;
+//   std::unique_ptr<SimpleTaskScheduler> scheduler =
+//       SimpleTaskScheduler::Create(clock);
+//
+//   // Schedule task_id1 to fire in 5 seconds.
+//   SimpleTaskScheduler::TaskId task_id1 = 1;
+//   scheduler->ScheduleAt(
+//       task_id1,
+//       [&]() { FooApi(); },
+//       clock.GetSystemTimeMilliseconds() + 5000);
+//
+//   // Schedule task_id2 to fire in 1 second.
+//   SimpleTaskScheduler::TaskId task_id2 = 2;
+//   scheduler->ScheduleAt(
+//       task_id2,
+//       [&]() { BarApi(); },
+//       clock.GetSystemTimeMilliseconds() + 1000);
+//
+// Then, task_id2 will fire first (running BarApi()), and then task_id1 will
+// fire afterwards (running FooApi()).
+//
+// Note: tasks are executed sequentially and there is only one execution at a
+// time.
+// - Tasks may be blocked by other requests to Icing.
+// - A task scheduled for time T will execute at the greater of T or the
+//   earliest time at which no previously scheduled tasks or incoming api calls
+//   are running.
+class SimpleTaskScheduler {
+ public:
+  using TaskId = int32_t;
+
+  // Creates a SimpleTaskScheduler.
+  //
+  // Note: after creation, no task is scheduled yet. It is only active after
+  //   calling ScheduleAt().
+  static std::unique_ptr<SimpleTaskScheduler> Create(const Clock& clock);
+
+  ~SimpleTaskScheduler();
+
+  // Disables copy, move, and assignment.
+  SimpleTaskScheduler(const SimpleTaskScheduler&) = delete;
+  SimpleTaskScheduler& operator=(const SimpleTaskScheduler&) = delete;
+  SimpleTaskScheduler(SimpleTaskScheduler&&) = delete;
+  SimpleTaskScheduler& operator=(SimpleTaskScheduler&&) = delete;
+
+  // (Re)Schedules the task (given the task id) with the given new timestamp in
+  // milliseconds.
+  // - If there is a pending task of the same id, then its execution time will
+  //   be overwritten by new_execution_timestamp_ms and the pending task is
+  //   rescheduled. The task function is also overwritten.
+  // - If new_execution_timestamp_ms is not greater than the current time
+  //   (clock_.GetSystemTimeMilliseconds()), an execution will fire without
+  //   waiting (as long as there is no ongoing execution).
+  // - If the timestamp is invalid (smaller than 0), the previous scheduled task
+  //   of this id will be canceled.
+  // - If it is time to run another new scheduled task but there is still an
+  //   ongoing execution, then the new scheduled task will wait until the
+  //   ongoing execution finishes and fire after that.
+  void ScheduleAt(TaskId task_id, std::function<void()> task,
+                  int64_t new_execution_timestamp_ms);
+
+  // Cancels the task by the task id.
+  // - It only cancels the pending task and does not interrupt an ongoing
+  //   (already started) execution.
+  // - If no task was scheduled or the task id is invalid, then this is a no-op.
+  void Cancel(TaskId task_id);
+
+  // Returns the scheduled time (next execution time) of the task in
+  // milliseconds. If the task id is unregistered or there is no scheduled
+  // execution, then returns -1.
+  int64_t GetScheduledTimeMs(TaskId task_id) const;
+
+ private:
+  static constexpr int64_t kInvalidTimestampMs = -1;
+
+  struct TaskInfo {
+    std::shared_ptr<std::function<void()>> task;
+    int64_t execution_timestamp_ms;
+
+    explicit TaskInfo(std::function<void()> task_in,
+                      int64_t execution_timestamp_ms_in)
+        : task(std::make_shared<std::function<void()>>(std::move(task_in))),
+          execution_timestamp_ms(execution_timestamp_ms_in) {}
+  };
+
+  explicit SimpleTaskScheduler(const Clock& clock)
+      : clock_(clock), is_terminated_(false) {};
+
+  // Helper function to initialize the scheduler.
+  void Initialize();
+
+  // Gets the next task info to execute.
+  //
+  // Returns:
+  //   - nullptr if there is no scheduled task.
+  //   - Otherwise, a valid TaskInfo pointer for the next scheduled task. i.e.
+  //     the task info with the smallest valid execution timestamp of all
+  //     scheduled tasks.
+  SimpleTaskScheduler::TaskInfo* GetNextTaskInfoLocked() const;
+
+  // Helper function to clean up completed or canceled tasks (i.e. tasks with
+  // invalid execution timestamp).
+  //
+  // Note: this function MUST be called only by the executor thread.
+  void CleanUpCompletedTasksLocked();
+
+  // The executor thread for waiting for the execution time and executing the
+  // task.
+  //
+  // Use unique_ptr to delay initialization the thread.
+  std::unique_ptr<std::thread> thread_;
+
+  // Mutex and condition variable for the thread to wait and be notified.
+  mutable std::mutex mutex_;  // NOLINT
+  std::condition_variable cv_;
+
+  // The clock used to get the current time.
+  const Clock& clock_;  // Does not own.
+
+  // Shared memory between the main thread and the executor thread.
+  // - The main thread will use public methods to update the values and notify
+  //   the executor thread via cv_.
+  // - The executor thread will use the values to determine the next action
+  //   after being notified.
+
+  // Mapping scheduled task id to info.
+  std::unordered_map<TaskId, std::unique_ptr<TaskInfo>> tasks_;
+
+  // Whether the scheduler is about to be destroyed or not.
+  bool is_terminated_;
+};
+
+}  // namespace lib
+}  // namespace icing
+
+#endif  // ICING_UTIL_SIMPLE_TASK_SCHEDULER_H_

diff --git a/icing/util/simple-task-scheduler_test.cc b/icing/util/simple-task-scheduler_test.cc
new file mode 100644
index 0000000..197bfab
--- /dev/null
+++ b/icing/util/simple-task-scheduler_test.cc

@@ -0,0 +1,504 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "icing/util/simple-task-scheduler.h"
+
+#include <atomic>
+#include <chrono>  // NOLINT
+#include <cstdint>
+#include <functional>
+#include <memory>
+#include <mutex>
+#include <thread>  // NOLINT
+
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+#include "icing/testing/fake-clock.h"
+#include "icing/util/clock.h"
+
+namespace icing {
+namespace lib {
+
+namespace {
+
+using ::testing::Eq;
+
+constexpr SimpleTaskScheduler::TaskId kTaskId = 0x12345678;
+
+// The epsilon for sleep() in milliseconds. There may be a small delay between
+// the scheduled task firing time and std::this_thread::sleep_for in the unit
+// test, so we need to sleep a little longer to make sure the task completes.
+constexpr int64_t kSleepMsEps = 100;
+
+// Simulate a class holds and uses the scheduler object to fire its own API
+// call.
+class MainClass {
+ public:
+  explicit MainClass(const Clock& clock) : clock_(clock) {}
+
+  void Initialize() {
+    std::unique_lock<std::mutex> lk(mutex_);  // NOLINT
+
+    scheduler_ = SimpleTaskScheduler::Create(clock_);
+    scheduler_->ScheduleAt(kTaskId, CreateTask(),
+                           clock_.GetSystemTimeMilliseconds() + 1000);
+  }
+
+  void FooApi() {
+    std::unique_lock<std::mutex> lk(mutex_);  // NOLINT
+
+    ++counter_;
+    if (scheduler_ != nullptr) {
+      // Re-schedule a new task to fire in 1 second.
+      scheduler_->ScheduleAt(kTaskId, CreateTask(),
+                             clock_.GetSystemTimeMilliseconds() + 1000);
+    }
+  }
+
+  int GetCounter() {
+    std::unique_lock<std::mutex> lk(mutex_);  // NOLINT
+
+    return counter_;
+  }
+
+ private:
+  std::function<void()> CreateTask() {
+    return [this]() { FooApi(); };
+  }
+
+  std::mutex mutex_;  // NOLINT
+
+  const Clock& clock_;  // Does not own.
+  std::unique_ptr<SimpleTaskScheduler> scheduler_;
+  std::atomic<int> counter_ = 0;
+};
+
+TEST(SimpleTaskSchedulerTest, SimpleScheduleAndExecute) {
+  Clock clock;
+  std::atomic<int> counter = 0;
+
+  std::unique_ptr<SimpleTaskScheduler> scheduler =
+      SimpleTaskScheduler::Create(clock);
+
+  // Schedule the task to fire in 1 second.
+  scheduler->ScheduleAt(
+      kTaskId, /*task=*/[&counter]() { ++counter; },
+      clock.GetSystemTimeMilliseconds() + 1000);
+
+  // After 1 second, the task should fire and increment the counter.
+  std::this_thread::sleep_for(std::chrono::milliseconds(1000 + kSleepMsEps));
+  EXPECT_THAT(counter, Eq(1));
+}
+
+TEST(SimpleTaskSchedulerTest, ScheduleAt_rescheduleTask) {
+  Clock clock;
+  std::atomic<int> counter = 0;
+  const auto task = [&counter]() { ++counter; };
+
+  std::unique_ptr<SimpleTaskScheduler> scheduler =
+      SimpleTaskScheduler::Create(clock);
+
+  // Schedule a task to fire in 3 seconds.
+  scheduler->ScheduleAt(kTaskId, task,
+                        clock.GetSystemTimeMilliseconds() + 3000);
+
+  // Sleep for 100 ms and reschedule the task (with the same id) to fire in 1
+  // second.
+  std::this_thread::sleep_for(std::chrono::milliseconds(100));
+  scheduler->ScheduleAt(kTaskId, task,
+                        clock.GetSystemTimeMilliseconds() + 1000);
+
+  // - After 1 second, the task should fire and increment the counter.
+  // - Although ScheduleAt() interrupts wait_for(), the task should not be
+  //   executed, and the counter should be only incremented once.
+  std::this_thread::sleep_for(std::chrono::milliseconds(1000 + kSleepMsEps));
+  EXPECT_THAT(counter, Eq(1));
+
+  // Sleep for another 3 seconds. The counter should remain 1. This verifies
+  // the original scheduled task (3 seconds) is canceled.
+  std::this_thread::sleep_for(std::chrono::milliseconds(3000 + kSleepMsEps));
+  EXPECT_THAT(counter, Eq(1));
+}
+
+TEST(SimpleTaskSchedulerTest, ScheduleAt_overwriteTask) {
+  Clock clock;
+  std::atomic<int> counter = 0;
+  const auto task1 = [&counter]() { ++counter; };
+  const auto task2 = [&counter]() { counter += 10; };
+
+  std::unique_ptr<SimpleTaskScheduler> scheduler =
+      SimpleTaskScheduler::Create(clock);
+
+  // Schedule a task to fire in 3 seconds.
+  scheduler->ScheduleAt(kTaskId, task1,
+                        clock.GetSystemTimeMilliseconds() + 3000);
+
+  // Sleep for 100 ms and reschedule the task to fire in 1 second, with a
+  // different task implementation.
+  std::this_thread::sleep_for(std::chrono::milliseconds(100));
+  scheduler->ScheduleAt(kTaskId, task2,
+                        clock.GetSystemTimeMilliseconds() + 1000);
+
+  // - After 1 second, the task should fire.
+  // - It should execute the 2nd task, so the counter is incremented by 10.
+  std::this_thread::sleep_for(std::chrono::milliseconds(1000 + kSleepMsEps));
+  EXPECT_THAT(counter, Eq(10));
+
+  // Sleep for another 3 seconds. The counter should remain 10. This verifies
+  // the original scheduled task (3 seconds) is not executed.
+  std::this_thread::sleep_for(std::chrono::milliseconds(3000 + kSleepMsEps));
+  EXPECT_THAT(counter, Eq(10));
+}
+
+TEST(SimpleTaskSchedulerTest, ScheduleAt_negativeTimestampShouldCancel) {
+  Clock clock;
+  std::atomic<int> counter = 0;
+  const auto task = [&counter]() { ++counter; };
+
+  std::unique_ptr<SimpleTaskScheduler> scheduler =
+      SimpleTaskScheduler::Create(clock);
+
+  // Schedule a task to fire in 1 second.
+  scheduler->ScheduleAt(kTaskId, task,
+                        clock.GetSystemTimeMilliseconds() + 1000);
+
+  // Schedule again with a negative timestamp to cancel the previous task. The
+  // scheduled time should be updated to -1, which means the task is canceled.
+  scheduler->ScheduleAt(kTaskId, task, -100);
+  EXPECT_THAT(scheduler->GetScheduledTimeMs(kTaskId), Eq(-1));
+
+  // The task should NOT fire.
+  std::this_thread::sleep_for(std::chrono::milliseconds(1000 + kSleepMsEps));
+  EXPECT_THAT(counter, Eq(0));
+}
+
+TEST(SimpleTaskSchedulerTest,
+     ScheduleAt_timestampEarlierThanNowShouldFireImmediately) {
+  Clock clock;
+  std::atomic<int> counter = 0;
+  const auto task = [&counter]() { ++counter; };
+
+  std::unique_ptr<SimpleTaskScheduler> scheduler =
+      SimpleTaskScheduler::Create(clock);
+
+  // Schedule a task with a timestamp earlier than now.
+  scheduler->ScheduleAt(kTaskId, task,
+                        clock.GetSystemTimeMilliseconds() - 1000);
+
+  // The task should fire immediately and increment the counter.
+  std::this_thread::sleep_for(std::chrono::milliseconds(kSleepMsEps));
+  EXPECT_THAT(counter, Eq(1));
+}
+
+TEST(SimpleTaskSchedulerTest, Cancel) {
+  Clock clock;
+  std::atomic<int> counter = 0;
+  const auto task = [&counter]() { ++counter; };
+
+  std::unique_ptr<SimpleTaskScheduler> scheduler =
+      SimpleTaskScheduler::Create(clock);
+
+  // Schedule a task to fire in 1 second.
+  scheduler->ScheduleAt(kTaskId, task,
+                        clock.GetSystemTimeMilliseconds() + 1000);
+
+  // Cancel the task after 100 ms.
+  std::this_thread::sleep_for(std::chrono::milliseconds(100));
+  scheduler->Cancel(kTaskId);
+
+  // - After 1 second, the task should not fire and the counter should remain 0.
+  // - Although Cancel() interrupts wait_for(), the task should not be executed,
+  //   and the counter should not be incremented.
+  std::this_thread::sleep_for(std::chrono::milliseconds(1000 + kSleepMsEps));
+  EXPECT_THAT(counter, Eq(0));
+}
+
+TEST(SimpleTaskSchedulerTest, GetScheduledTimeMs) {
+  FakeClock fake_clock;
+  fake_clock.SetSystemTimeMilliseconds(0);
+  const auto task = []() {
+    // no-op
+  };
+
+  std::unique_ptr<SimpleTaskScheduler> scheduler =
+      SimpleTaskScheduler::Create(fake_clock);
+
+  // Before scheduling, the scheduled time should be -1 (invalid timestamp).
+  EXPECT_THAT(scheduler->GetScheduledTimeMs(kTaskId), Eq(-1));
+
+  scheduler->ScheduleAt(kTaskId, task, 100000);
+  EXPECT_THAT(scheduler->GetScheduledTimeMs(kTaskId), Eq(100000));
+
+  scheduler->ScheduleAt(kTaskId, task, 200000);
+  EXPECT_THAT(scheduler->GetScheduledTimeMs(kTaskId), Eq(200000));
+
+  scheduler->ScheduleAt(kTaskId, task, 300000);
+  EXPECT_THAT(scheduler->GetScheduledTimeMs(kTaskId), Eq(300000));
+
+  scheduler->Cancel(kTaskId);
+  EXPECT_THAT(scheduler->GetScheduledTimeMs(kTaskId), Eq(-1));
+
+  scheduler->ScheduleAt(kTaskId, task, 400000);
+  EXPECT_THAT(scheduler->GetScheduledTimeMs(kTaskId), Eq(400000));
+
+  scheduler->ScheduleAt(kTaskId, task, -1000);
+  EXPECT_THAT(scheduler->GetScheduledTimeMs(kTaskId), Eq(-1));
+
+  // Unregistered task id should return -1.
+  EXPECT_THAT(scheduler->GetScheduledTimeMs(kTaskId + 1), Eq(-1));
+}
+
+TEST(SimpleTaskSchedulerTest, Destructor) {
+  Clock clock;
+  std::atomic<int> counter = 0;
+  const auto task = [&counter]() { ++counter; };
+
+  std::unique_ptr<SimpleTaskScheduler> scheduler =
+      SimpleTaskScheduler::Create(clock);
+
+  // Schedule a task to fire in 1 second.
+  scheduler->ScheduleAt(kTaskId, task,
+                        clock.GetSystemTimeMilliseconds() + 1000);
+
+  // Destruct the scheduler after 100 ms.
+  std::this_thread::sleep_for(std::chrono::milliseconds(100));
+  scheduler.reset();
+
+  // After 1 second, the task should not fire and the counter should remain 0.
+  std::this_thread::sleep_for(std::chrono::milliseconds(1000 + kSleepMsEps));
+  EXPECT_THAT(counter, Eq(0));
+}
+
+TEST(SimpleTaskSchedulerTest, Destructor_shouldCancelAndJoinThread) {
+  Clock clock;
+  std::atomic<int> counter = 0;
+  const auto task = [&counter]() { ++counter; };
+
+  std::unique_ptr<SimpleTaskScheduler> scheduler =
+      SimpleTaskScheduler::Create(clock);
+
+  // Schedule a task to fire in 1 second.
+  scheduler->ScheduleAt(kTaskId, task,
+                        clock.GetSystemTimeMilliseconds() + 1000);
+
+  // Destruct the scheduler object immediately. The background thread should be
+  // joined successfully in the destructor, without any crash or deadlock.
+  scheduler.reset();
+
+  EXPECT_THAT(counter, Eq(0));
+}
+
+TEST(SimpleTaskSchedulerTest, Destructor_noScheduledTask) {
+  Clock clock;
+  std::atomic<int> counter = 0;
+
+  std::unique_ptr<SimpleTaskScheduler> scheduler =
+      SimpleTaskScheduler::Create(clock);
+
+  // Destruct the scheduler object immediately without scheduling any task.
+  // There should be no crash or deadlock.
+  scheduler.reset();
+
+  EXPECT_THAT(counter, Eq(0));
+}
+
+TEST(SimpleTaskSchedulerTest, Destructor_noDeadlock) {
+  Clock clock;
+  std::atomic<int> counter = 0;
+  const auto task = [&counter]() {
+    ++counter;
+    std::this_thread::sleep_for(std::chrono::milliseconds(3000));
+  };
+
+  // Create a scheduler and schedule an expensive task that sleeps for 3
+  // seconds.
+  //
+  // This allows the test to trigger a different execution order (race
+  // condition) and verify the correctness of the conditional variable
+  // notification in the destructor before joining the thread, and
+  // is_terminated_ flag check.
+  std::unique_ptr<SimpleTaskScheduler> scheduler =
+      SimpleTaskScheduler::Create(clock);
+
+  // Schedule a task to fire in 1 second.
+  scheduler->ScheduleAt(kTaskId, task,
+                        clock.GetSystemTimeMilliseconds() + 1000);
+
+  // Sleep for 1 second.
+  // - At this moment, the executor thread is still running the expensive task.
+  // - Destruct the scheduler object at this moment.
+  //   - The destructor will set is_terminated_ to true and notify the executor
+  //     thread to terminate, but the executor thread is still running the
+  //     expensive task, so the notification is no-op.
+  //   - The destructor will join the executor thread.
+  //   - After the executor thread finishes the expensive task, it should not
+  //     call wait() or wait_for() again. Instead, it should terminate and join
+  //     the destructor.
+  std::this_thread::sleep_for(std::chrono::milliseconds(1000 + kSleepMsEps));
+  scheduler.reset();
+
+  EXPECT_THAT(counter, Eq(1));
+}
+
+TEST(SimpleTaskSchedulerTest, MultipleScheduledTasks) {
+  Clock clock;
+  std::atomic<int> counter = 0;
+
+  std::unique_ptr<SimpleTaskScheduler> scheduler =
+      SimpleTaskScheduler::Create(clock);
+
+  SimpleTaskScheduler::TaskId task_id1 = 1;
+  SimpleTaskScheduler::TaskId task_id2 = 2;
+  // Schedule task_id1 to fire in 2 seconds.
+  int64_t current_time_ms = clock.GetSystemTimeMilliseconds();
+  scheduler->ScheduleAt(
+      task_id1, [&counter]() { ++counter; }, current_time_ms + 2000);
+  // Schedule task_id2 to fire in 1 second.
+  scheduler->ScheduleAt(
+      task_id2, [&counter]() { counter += 10; }, current_time_ms + 1000);
+
+  // After 1 second, task_id2 should fire and increment the counter by 10.
+  std::this_thread::sleep_for(std::chrono::milliseconds(1000 + kSleepMsEps));
+  EXPECT_THAT(counter, Eq(10));
+
+  // After another 1 second, task_id1 should fire and increment the counter
+  // by 1.
+  std::this_thread::sleep_for(std::chrono::milliseconds(1000 + kSleepMsEps));
+  EXPECT_THAT(counter, Eq(11));
+}
+
+TEST(SimpleTaskSchedulerTest, ExpensiveTask_nextExecutionShouldWaitAndExecute) {
+  Clock clock;
+  std::atomic<int> counter = 0;
+  const auto task = [&counter]() {
+    ++counter;
+    std::this_thread::sleep_for(std::chrono::milliseconds(2000));
+  };
+
+  // Create a scheduler and schedule an expensive task that sleeps for 2
+  // seconds.
+  //
+  // This allows the test to trigger a different execution order (race
+  // condition) and verify the correctness of the conditional variable
+  // notification in the destructor before joining the thread, and
+  // is_terminated_ flag check.
+  std::unique_ptr<SimpleTaskScheduler> scheduler =
+      SimpleTaskScheduler::Create(clock);
+
+  // Schedule a task to fire immediately. The counter is incremented, but the
+  // 1st execution is still running.
+  scheduler->ScheduleAt(kTaskId, task, clock.GetSystemTimeMilliseconds());
+  std::this_thread::sleep_for(std::chrono::milliseconds(kSleepMsEps));
+  EXPECT_THAT(counter, Eq(1));
+
+  // Schedule another task to fire in 10 ms for the 2nd execution.
+  scheduler->ScheduleAt(kTaskId, task, clock.GetSystemTimeMilliseconds() + 10);
+  // Sleep for 10 ms. The 1st execution is still running, so the 2nd execution
+  // should wait and run after the 1st execution finishes.
+  std::this_thread::sleep_for(std::chrono::milliseconds(10 + kSleepMsEps));
+  EXPECT_THAT(counter, Eq(1));
+
+  // Finally, the 1st execution finishes and the 2nd execution fires.
+  std::this_thread::sleep_for(std::chrono::milliseconds(2000 + kSleepMsEps));
+  EXPECT_THAT(counter, Eq(2));
+}
+
+TEST(SimpleTaskSchedulerTest,
+     ExpensiveTask_multipleSchedulesDuringOngoingExecutionFiresOnceAfterwards) {
+  Clock clock;
+  std::atomic<int> counter = 0;
+  const auto task = [&counter]() {
+    ++counter;
+    std::this_thread::sleep_for(std::chrono::milliseconds(2000));
+  };
+
+  // Create a scheduler and schedule an expensive task that sleeps for 2
+  // seconds.
+  //
+  // This allows the test to trigger a different execution order (race
+  // condition) and verify the correctness of the conditional variable
+  // notification in the destructor before joining the thread, and
+  // is_terminated_ flag check.
+  std::unique_ptr<SimpleTaskScheduler> scheduler =
+      SimpleTaskScheduler::Create(clock);
+
+  // Schedule a task to fire immediately. The counter is incremented, but the
+  // 1st execution is still running.
+  scheduler->ScheduleAt(kTaskId, task, clock.GetSystemTimeMilliseconds());
+  std::this_thread::sleep_for(std::chrono::milliseconds(kSleepMsEps));
+  EXPECT_THAT(counter, Eq(1));
+
+  // Schedule the task for the 2nd time to fire in 10 ms.
+  scheduler->ScheduleAt(kTaskId, task, clock.GetSystemTimeMilliseconds() + 10);
+  // Sleep for 10 ms. The 1st execution is still running, so the next execution
+  // should wait and execute after the 1st execution finishes.
+  std::this_thread::sleep_for(std::chrono::milliseconds(10 + kSleepMsEps));
+  EXPECT_THAT(counter, Eq(1));
+
+  // Schedule the task for the 3rd time to fire in 20 ms.
+  scheduler->ScheduleAt(kTaskId, task, clock.GetSystemTimeMilliseconds() + 20);
+  // Sleep for 10 ms. The 1st execution is still running, so the next execution
+  // should wait and execute after the 1st execution finishes.
+  std::this_thread::sleep_for(std::chrono::milliseconds(10 + kSleepMsEps));
+  EXPECT_THAT(counter, Eq(1));
+
+  // Finally, the 1st execution finishes and another execution fires.
+  std::this_thread::sleep_for(std::chrono::milliseconds(2000 + kSleepMsEps));
+  EXPECT_THAT(counter, Eq(2));
+
+  // Sleep for another 2 seconds. The counter should remain 2 since only one
+  // execution should fire after the 1st execution finishes, even though there
+  // were two executions expected to fire when the 1st execution was still
+  // running.
+  std::this_thread::sleep_for(std::chrono::milliseconds(2000 + kSleepMsEps));
+  EXPECT_THAT(counter, Eq(2));
+}
+
+TEST(SimpleTaskSchedulerTest, MockClock) {
+  FakeClock fake_clock;
+  std::atomic<int> counter = 0;
+  const auto task = [&counter]() { ++counter; };
+
+  std::unique_ptr<SimpleTaskScheduler> scheduler =
+      SimpleTaskScheduler::Create(fake_clock);
+
+  fake_clock.SetSystemTimeMilliseconds(30000);
+  // Schedule a task to fire in 1 second.
+  scheduler->ScheduleAt(kTaskId, task,
+                        fake_clock.GetSystemTimeMilliseconds() + 1000);
+
+  // After 1 second, the task should fire and increment the counter.
+  std::this_thread::sleep_for(std::chrono::milliseconds(1000 + kSleepMsEps));
+  EXPECT_THAT(counter, Eq(1));
+}
+
+TEST(SimpleTaskSchedulerTest, SelfTask) {
+  Clock clock;
+
+  auto main_class = std::make_unique<MainClass>(clock);
+  main_class->Initialize();
+
+  // After 1 second, the task should fire and increment the counter.
+  std::this_thread::sleep_for(std::chrono::milliseconds(1000 + kSleepMsEps));
+  EXPECT_THAT(main_class->GetCounter(), Eq(1));
+
+  // After another 1 second, the task should fire and increment the counter.
+  std::this_thread::sleep_for(std::chrono::milliseconds(1000 + kSleepMsEps));
+  EXPECT_THAT(main_class->GetCounter(), Eq(2));
+}
+
+}  // namespace
+
+}  // namespace lib
+}  // namespace icing

diff --git a/icing/util/timestamp-util.cc b/icing/util/timestamp-util.cc
new file mode 100644
index 0000000..6e50d2f
--- /dev/null
+++ b/icing/util/timestamp-util.cc

@@ -0,0 +1,49 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "icing/util/timestamp-util.h"
+
+#include <cstdint>
+#include <limits>
+
+namespace icing {
+namespace lib {
+
+namespace timestamp_util {
+
+int64_t CalculateRawExpirationTimestampMs(int64_t creation_timestamp_ms,
+                                          int64_t ttl_ms) {
+  if (ttl_ms == 0) {
+    // Special case where a TTL of 0 indicates the document should never
+    // expire. int64_t max, interpreted as seconds since epoch, represents
+    // some point in the year 292,277,026,596. So we're probably ok to use
+    // this as "never reaching this point".
+    return std::numeric_limits<int64_t>::max();
+  }
+
+  int64_t expiration_timestamp_ms;
+  if (__builtin_add_overflow(creation_timestamp_ms, ttl_ms,
+                             &expiration_timestamp_ms)) {
+    // Overflow detected. Treat overflow as the same behavior of just int64_t
+    // max
+    return std::numeric_limits<int64_t>::max();
+  }
+
+  return expiration_timestamp_ms;
+}
+
+}  // namespace timestamp_util
+
+}  // namespace lib
+}  // namespace icing

diff --git a/icing/util/timestamp-util.h b/icing/util/timestamp-util.h
new file mode 100644
index 0000000..6479b4f
--- /dev/null
+++ b/icing/util/timestamp-util.h

@@ -0,0 +1,42 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef ICING_UTIL_TIMESTAMP_UTIL_H_
+#define ICING_UTIL_TIMESTAMP_UTIL_H_
+
+#include <cstdint>
+
+namespace icing {
+namespace lib {
+
+namespace timestamp_util {
+
+// Calculates the (raw) expiration timestamp of a document given its creation
+// timestamp and time-to-live (TTL) in milliseconds.
+//
+// If the TTL is 0, the document should never expire and the function will
+// return INT64_MAX as the expiration timestamp.
+//
+// If an overflow occurs, the function will return INT64_MAX.
+//
+// REQUIRES: creation_timestamp_ms >= 0 && ttl_ms >= 0.
+int64_t CalculateRawExpirationTimestampMs(int64_t creation_timestamp_ms,
+                                          int64_t ttl_ms);
+
+}  // namespace timestamp_util
+
+}  // namespace lib
+}  // namespace icing
+
+#endif  // ICING_UTIL_TIMESTAMP_UTIL_H_

diff --git a/icing/util/timestamp-util_test.cc b/icing/util/timestamp-util_test.cc
new file mode 100644
index 0000000..fd488cd
--- /dev/null
+++ b/icing/util/timestamp-util_test.cc

@@ -0,0 +1,88 @@
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "icing/util/timestamp-util.h"
+
+#include <cstdint>
+#include <limits>
+
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+namespace icing {
+namespace lib {
+namespace timestamp_util {
+
+namespace {
+
+using ::testing::Eq;
+
+TEST(TimestampUtilTest, CalculateRawExpirationTimestampMs) {
+  EXPECT_THAT(CalculateRawExpirationTimestampMs(/*creation_timestamp_ms=*/0,
+                                                /*ttl_ms=*/1000),
+              Eq(1000));
+  EXPECT_THAT(CalculateRawExpirationTimestampMs(/*creation_timestamp_ms=*/1000,
+                                                /*ttl_ms=*/1000),
+              Eq(2000));
+  EXPECT_THAT(CalculateRawExpirationTimestampMs(/*creation_timestamp_ms=*/2000,
+                                                /*ttl_ms=*/1000),
+              Eq(3000));
+  EXPECT_THAT(CalculateRawExpirationTimestampMs(/*creation_timestamp_ms=*/5000,
+                                                /*ttl_ms=*/1000),
+              Eq(6000));
+  EXPECT_THAT(CalculateRawExpirationTimestampMs(/*creation_timestamp_ms=*/5000,
+                                                /*ttl_ms=*/3000),
+              Eq(8000));
+  EXPECT_THAT(CalculateRawExpirationTimestampMs(/*creation_timestamp_ms=*/5000,
+                                                /*ttl_ms=*/10000),
+              Eq(15000));
+}
+
+TEST(TimestampUtilTest,
+     CalculateRawExpirationTimestampMs_zeroTtlShouldReturnInt64Max) {
+  EXPECT_THAT(CalculateRawExpirationTimestampMs(/*creation_timestamp_ms=*/1000,
+                                                /*ttl_ms=*/0),
+              Eq(std::numeric_limits<int64_t>::max()));
+  EXPECT_THAT(CalculateRawExpirationTimestampMs(/*creation_timestamp_ms=*/2000,
+                                                /*ttl_ms=*/0),
+              Eq(std::numeric_limits<int64_t>::max()));
+  EXPECT_THAT(CalculateRawExpirationTimestampMs(/*creation_timestamp_ms=*/5000,
+                                                /*ttl_ms=*/0),
+              Eq(std::numeric_limits<int64_t>::max()));
+}
+
+TEST(TimestampUtilTest,
+     CalculateRawExpirationTimestampMs_shouldPreventOverflow) {
+  EXPECT_THAT(
+      CalculateRawExpirationTimestampMs(
+          /*creation_timestamp_ms=*/std::numeric_limits<int64_t>::max() - 2,
+          /*ttl_ms=*/2),
+      Eq(std::numeric_limits<int64_t>::max()));
+  EXPECT_THAT(
+      CalculateRawExpirationTimestampMs(
+          /*creation_timestamp_ms=*/std::numeric_limits<int64_t>::max() - 2,
+          /*ttl_ms=*/100),
+      Eq(std::numeric_limits<int64_t>::max()));
+  EXPECT_THAT(
+      CalculateRawExpirationTimestampMs(
+          /*creation_timestamp_ms=*/std::numeric_limits<int64_t>::max() - 1,
+          /*ttl_ms=*/100000),
+      Eq(std::numeric_limits<int64_t>::max()));
+}
+
+}  // namespace
+
+}  // namespace timestamp_util
+}  // namespace lib
+}  // namespace icing

diff --git a/icing/util/tokenized-document.cc b/icing/util/tokenized-document.cc
index 576e0a2..61f4396 100644
--- a/icing/util/tokenized-document.cc
+++ b/icing/util/tokenized-document.cc

@@ -14,6 +14,7 @@
 
 #include "icing/util/tokenized-document.h"
 
+#include <cstdint>
 #include <memory>
 #include <string_view>
 #include <utility>
@@ -21,6 +22,7 @@
 
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "icing/proto/document.pb.h"
+#include "icing/proto/document_wrapper.pb.h"
 #include "icing/schema/joinable-property.h"
 #include "icing/schema/schema-store.h"
 #include "icing/schema/section.h"
@@ -28,6 +30,7 @@
 #include "icing/tokenization/token.h"
 #include "icing/tokenization/tokenizer-factory.h"
 #include "icing/tokenization/tokenizer.h"
+#include "icing/util/document-util.h"
 #include "icing/util/document-validator.h"
 #include "icing/util/status-macros.h"
 
@@ -41,6 +44,7 @@
     const LanguageSegmenter* language_segmenter,
     const std::vector<Section<std::string_view>>& string_sections) {
   std::vector<TokenizedSection> tokenized_string_sections;
+  std::vector<Token> batch_tokens;
   for (const Section<std::string_view>& section : string_sections) {
     ICING_ASSIGN_OR_RETURN(std::unique_ptr<Tokenizer> tokenizer,
                            tokenizer_factory::CreateIndexingTokenizer(
@@ -50,7 +54,7 @@
       ICING_ASSIGN_OR_RETURN(std::unique_ptr<Tokenizer::Iterator> itr,
                              tokenizer->Tokenize(subcontent));
       while (itr->Advance()) {
-        std::vector<Token> batch_tokens = itr->GetTokens();
+        itr->GetTokens(&batch_tokens);
         for (const Token& token : batch_tokens) {
           token_sequence.push_back(token.text);
         }
@@ -68,21 +72,28 @@
 /* static */ libtextclassifier3::StatusOr<TokenizedDocument>
 TokenizedDocument::Create(const SchemaStore* schema_store,
                           const LanguageSegmenter* language_segmenter,
-                          DocumentProto document) {
+                          int64_t current_time_ms, DocumentProto document) {
+  // Set the creation timestamp if it is not set.
+  if (document.creation_timestamp_ms() == 0) {
+    document.set_creation_timestamp_ms(current_time_ms);
+  }
+
   // Since there are many std::string_view objects pointing to the document
-  // proto, we should make sure DocumentProto has a fixed address. The simplest
-  // way is to use a unique_ptr.
-  auto document_ptr = std::make_unique<DocumentProto>(std::move(document));
+  // proto, we should make sure DocumentProto in DocumentWrapper has a fixed
+  // address. The simplest way is to use a unique_ptr.
+  auto document_wrapper_ptr = std::make_unique<DocumentWrapper>(
+      document_util::CreateDocumentWrapper(std::move(document)));
 
   DocumentValidator validator(schema_store);
-  ICING_RETURN_IF_ERROR(validator.Validate(*document_ptr));
-
-  ICING_ASSIGN_OR_RETURN(SectionGroup section_group,
-                         schema_store->ExtractSections(*document_ptr));
+  ICING_RETURN_IF_ERROR(validator.Validate(document_wrapper_ptr->document()));
 
   ICING_ASSIGN_OR_RETURN(
-      JoinablePropertyGroup joinable_property_group,
-      schema_store->ExtractJoinableProperties(*document_ptr));
+      SectionGroup section_group,
+      schema_store->ExtractSections(document_wrapper_ptr->document()));
+
+  ICING_ASSIGN_OR_RETURN(JoinablePropertyGroup joinable_property_group,
+                         schema_store->ExtractJoinableProperties(
+                             document_wrapper_ptr->document()));
 
   // Tokenize string sections
   ICING_ASSIGN_OR_RETURN(
@@ -90,11 +101,19 @@
       Tokenize(schema_store, language_segmenter,
                section_group.string_sections));
 
-  return TokenizedDocument(std::move(document_ptr),
-                           std::move(tokenized_string_sections),
-                           std::move(section_group.integer_sections),
-                           std::move(section_group.vector_sections),
-                           std::move(joinable_property_group));
+  TokenizedDocument tokenized_document(
+      std::move(document_wrapper_ptr), std::move(tokenized_string_sections),
+      std::move(section_group.integer_sections),
+      std::move(section_group.vector_sections),
+      std::move(joinable_property_group));
+
+  // Set the num_string_tokens into the document proto.
+  int32_t num_string_tokens = tokenized_document.num_string_tokens();
+  tokenized_document.document_wrapper_->mutable_document()
+      ->mutable_internal_fields()
+      ->set_length_in_tokens(num_string_tokens);
+
+  return tokenized_document;
 }
 
 }  // namespace lib

diff --git a/icing/util/tokenized-document.h b/icing/util/tokenized-document.h
index ae92311..5a7f635 100644
--- a/icing/util/tokenized-document.h
+++ b/icing/util/tokenized-document.h

@@ -23,6 +23,7 @@
 
 #include "icing/text_classifier/lib3/utils/base/statusor.h"
 #include "icing/proto/document.pb.h"
+#include "icing/proto/document_wrapper.pb.h"
 #include "icing/schema/joinable-property.h"
 #include "icing/schema/schema-store.h"
 #include "icing/schema/section.h"
@@ -41,13 +42,29 @@
         token_sequence(std::move(token_sequence_in)) {}
 };
 
+// A wrapper class to hold the input DocumentProto and the tokenized sections.
+// It finalizes the proto (to be added into DocumentStore) and tokens (to be
+// indexed), so after the creation, the proto and tokens are immutable and the
+// caller can simply use them to write.
+//
+// Note: the document dependency is not the responsibility of this class.
+// Additional steps (see DocumentDependencyProcessor) are needed to handle the
+// dependencies.
+//
+// Some essential fields are populated into the DocumentProto. For example:
+// - Creation timestamp: if it is not set by the client, it will be set to the
+//   current time.
+// - Length in tokens: the number of tokens in all string sections.
 class TokenizedDocument {
  public:
   static libtextclassifier3::StatusOr<TokenizedDocument> Create(
       const SchemaStore* schema_store,
-      const LanguageSegmenter* language_segmenter, DocumentProto document);
+      const LanguageSegmenter* language_segmenter, int64_t current_time_ms,
+      DocumentProto document);
 
-  const DocumentProto& document() const { return *document_; }
+  // Due to DocumentStore's internal implementation, we need to wrap
+  // DocumentProto into DocumentWrapper.
+  const DocumentWrapper& document_wrapper() const { return *document_wrapper_; }
 
   int32_t num_string_tokens() const {
     int32_t num_string_tokens = 0;
@@ -78,18 +95,18 @@
  private:
   // Use TokenizedDocument::Create() to instantiate.
   explicit TokenizedDocument(
-      std::unique_ptr<DocumentProto> document,
+      std::unique_ptr<DocumentWrapper> document_wrapper,
       std::vector<TokenizedSection>&& tokenized_string_sections,
       std::vector<Section<int64_t>>&& integer_sections,
       std::vector<Section<PropertyProto::VectorProto>>&& vector_sections,
       JoinablePropertyGroup&& joinable_property_group)
-      : document_(std::move(document)),
+      : document_wrapper_(std::move(document_wrapper)),
         tokenized_string_sections_(std::move(tokenized_string_sections)),
         integer_sections_(std::move(integer_sections)),
         vector_sections_(std::move(vector_sections)),
         joinable_property_group_(std::move(joinable_property_group)) {}
 
-  std::unique_ptr<DocumentProto> document_;
+  std::unique_ptr<DocumentWrapper> document_wrapper_;
   std::vector<TokenizedSection> tokenized_string_sections_;
   std::vector<Section<int64_t>> integer_sections_;
   std::vector<Section<PropertyProto::VectorProto>> vector_sections_;

diff --git a/icing/util/tokenized-document_test.cc b/icing/util/tokenized-document_test.cc
index 513b8d0..84528e7 100644
--- a/icing/util/tokenized-document_test.cc
+++ b/icing/util/tokenized-document_test.cc

@@ -14,6 +14,7 @@
 
 #include "icing/util/tokenized-document.h"
 
+#include <cstdint>
 #include <memory>
 #include <string>
 #include <string_view>
@@ -27,6 +28,7 @@
 #include "icing/file/filesystem.h"
 #include "icing/portable/platform.h"
 #include "icing/proto/document.pb.h"
+#include "icing/proto/document_wrapper.pb.h"
 #include "icing/proto/schema.pb.h"
 #include "icing/schema-builder.h"
 #include "icing/schema/joinable-property.h"
@@ -39,6 +41,7 @@
 #include "icing/testing/tmp-directory.h"
 #include "icing/tokenization/language-segmenter-factory.h"
 #include "icing/tokenization/language-segmenter.h"
+#include "icing/util/document-util.h"
 #include "icing/util/icu-data-file-helper.h"
 #include "unicode/uloc.h"
 
@@ -239,7 +242,15 @@
   std::unique_ptr<SchemaStore> schema_store_;
 };
 
+DocumentWrapper CreateDocumentWrapper(DocumentProto document,
+                                      int32_t num_string_tokens) {
+  document.mutable_internal_fields()->set_length_in_tokens(num_string_tokens);
+  return document_util::CreateDocumentWrapper(std::move(document));
+}
+
 TEST_F(TokenizedDocumentTest, CreateAll) {
+  int64_t current_time_ms = 123;
+
   PropertyProto::VectorProto vector1;
   vector1.set_model_signature("my_model1");
   vector1.add_values(1.0f);
@@ -252,6 +263,7 @@
 
   DocumentProto document =
       DocumentBuilder()
+          .SetCreationTimestampMs(current_time_ms)
           .SetKey("icing", "fake_type/1")
           .SetSchema(std::string(kFakeType))
           .AddStringProperty(std::string(kUnindexedStringProperty),
@@ -273,9 +285,11 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
       TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+                                current_time_ms, document));
 
-  EXPECT_THAT(tokenized_document.document(), EqualsProto(document));
+  EXPECT_THAT(
+      tokenized_document.document_wrapper(),
+      EqualsProto(CreateDocumentWrapper(document, /*num_string_tokens=*/9)));
   EXPECT_THAT(tokenized_document.num_string_tokens(), Eq(9));
 
   // string sections
@@ -326,8 +340,11 @@
 }
 
 TEST_F(TokenizedDocumentTest, CreateNoIndexableIntegerProperties) {
+  int64_t current_time_ms = 123;
+
   DocumentProto document =
       DocumentBuilder()
+          .SetCreationTimestampMs(current_time_ms)
           .SetKey("icing", "fake_type/1")
           .SetSchema(std::string(kFakeType))
           .AddInt64Property(std::string(kUnindexedIntegerProperty), 789)
@@ -336,9 +353,11 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
       TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+                                current_time_ms, document));
 
-  EXPECT_THAT(tokenized_document.document(), EqualsProto(document));
+  EXPECT_THAT(
+      tokenized_document.document_wrapper(),
+      EqualsProto(CreateDocumentWrapper(document, /*num_string_tokens=*/0)));
   EXPECT_THAT(tokenized_document.num_string_tokens(), Eq(0));
 
   // string sections
@@ -355,8 +374,11 @@
 }
 
 TEST_F(TokenizedDocumentTest, CreateMultipleIndexableIntegerProperties) {
+  int64_t current_time_ms = 123;
+
   DocumentProto document =
       DocumentBuilder()
+          .SetCreationTimestampMs(current_time_ms)
           .SetKey("icing", "fake_type/1")
           .SetSchema(std::string(kFakeType))
           .AddInt64Property(std::string(kUnindexedIntegerProperty), 789)
@@ -367,9 +389,11 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
       TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+                                current_time_ms, document));
 
-  EXPECT_THAT(tokenized_document.document(), EqualsProto(document));
+  EXPECT_THAT(
+      tokenized_document.document_wrapper(),
+      EqualsProto(CreateDocumentWrapper(document, /*num_string_tokens=*/0)));
   EXPECT_THAT(tokenized_document.num_string_tokens(), Eq(0));
 
   // string sections
@@ -394,8 +418,11 @@
 }
 
 TEST_F(TokenizedDocumentTest, CreateNoIndexableStringProperties) {
+  int64_t current_time_ms = 123;
+
   DocumentProto document =
       DocumentBuilder()
+          .SetCreationTimestampMs(current_time_ms)
           .SetKey("icing", "fake_type/1")
           .SetSchema(std::string(kFakeType))
           .AddStringProperty(std::string(kUnindexedStringProperty),
@@ -405,9 +432,11 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
       TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+                                current_time_ms, document));
 
-  EXPECT_THAT(tokenized_document.document(), EqualsProto(document));
+  EXPECT_THAT(
+      tokenized_document.document_wrapper(),
+      EqualsProto(CreateDocumentWrapper(document, /*num_string_tokens=*/0)));
   EXPECT_THAT(tokenized_document.num_string_tokens(), Eq(0));
 
   // string sections
@@ -424,8 +453,11 @@
 }
 
 TEST_F(TokenizedDocumentTest, CreateMultipleIndexableStringProperties) {
+  int64_t current_time_ms = 123;
+
   DocumentProto document =
       DocumentBuilder()
+          .SetCreationTimestampMs(current_time_ms)
           .SetKey("icing", "fake_type/1")
           .SetSchema(std::string(kFakeType))
           .AddStringProperty(std::string(kUnindexedStringProperty),
@@ -438,9 +470,11 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
       TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+                                current_time_ms, document));
 
-  EXPECT_THAT(tokenized_document.document(), EqualsProto(document));
+  EXPECT_THAT(
+      tokenized_document.document_wrapper(),
+      EqualsProto(CreateDocumentWrapper(document, /*num_string_tokens=*/9)));
   EXPECT_THAT(tokenized_document.num_string_tokens(), Eq(9));
 
   // string sections
@@ -467,12 +501,15 @@
 }
 
 TEST_F(TokenizedDocumentTest, CreateNoIndexableVectorProperties) {
+  int64_t current_time_ms = 123;
+
   PropertyProto::VectorProto vector;
   vector.set_model_signature("my_model");
   vector.add_values(1.0f);
 
   DocumentProto document =
       DocumentBuilder()
+          .SetCreationTimestampMs(current_time_ms)
           .SetKey("icing", "fake_type/1")
           .SetSchema(std::string(kFakeType))
           .AddVectorProperty(std::string(kUnindexedVectorProperty), vector)
@@ -481,9 +518,11 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
       TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+                                current_time_ms, document));
 
-  EXPECT_THAT(tokenized_document.document(), EqualsProto(document));
+  EXPECT_THAT(
+      tokenized_document.document_wrapper(),
+      EqualsProto(CreateDocumentWrapper(document, /*num_string_tokens=*/0)));
   EXPECT_THAT(tokenized_document.num_string_tokens(), Eq(0));
 
   // string sections
@@ -500,6 +539,8 @@
 }
 
 TEST_F(TokenizedDocumentTest, CreateMultipleIndexableVectorProperties) {
+  int64_t current_time_ms = 123;
+
   PropertyProto::VectorProto vector1;
   vector1.set_model_signature("my_model1");
   vector1.add_values(1.0f);
@@ -512,6 +553,7 @@
 
   DocumentProto document =
       DocumentBuilder()
+          .SetCreationTimestampMs(current_time_ms)
           .SetKey("icing", "fake_type/1")
           .SetSchema(std::string(kFakeType))
           .AddVectorProperty(std::string(kUnindexedVectorProperty), vector1)
@@ -523,9 +565,11 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
       TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+                                current_time_ms, document));
 
-  EXPECT_THAT(tokenized_document.document(), EqualsProto(document));
+  EXPECT_THAT(
+      tokenized_document.document_wrapper(),
+      EqualsProto(CreateDocumentWrapper(document, /*num_string_tokens=*/0)));
   EXPECT_THAT(tokenized_document.num_string_tokens(), Eq(0));
 
   // string sections
@@ -550,8 +594,11 @@
 }
 
 TEST_F(TokenizedDocumentTest, CreateNoJoinQualifiedIdProperties) {
+  int64_t current_time_ms = 123;
+
   DocumentProto document =
       DocumentBuilder()
+          .SetCreationTimestampMs(current_time_ms)
           .SetKey("icing", "fake_type/1")
           .SetSchema(std::string(kFakeType))
           .AddStringProperty(std::string(kUnindexedStringProperty),
@@ -561,9 +608,11 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
       TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+                                current_time_ms, document));
 
-  EXPECT_THAT(tokenized_document.document(), EqualsProto(document));
+  EXPECT_THAT(
+      tokenized_document.document_wrapper(),
+      EqualsProto(CreateDocumentWrapper(document, /*num_string_tokens=*/0)));
   EXPECT_THAT(tokenized_document.num_string_tokens(), Eq(0));
 
   // string sections
@@ -580,8 +629,11 @@
 }
 
 TEST_F(TokenizedDocumentTest, CreateMultipleJoinQualifiedIdProperties) {
+  int64_t current_time_ms = 123;
+
   DocumentProto document =
       DocumentBuilder()
+          .SetCreationTimestampMs(current_time_ms)
           .SetKey("icing", "fake_type/1")
           .SetSchema(std::string(kFakeType))
           .AddStringProperty(std::string(kUnindexedStringProperty),
@@ -593,9 +645,11 @@
   ICING_ASSERT_OK_AND_ASSIGN(
       TokenizedDocument tokenized_document,
       TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
-                                document));
+                                current_time_ms, document));
 
-  EXPECT_THAT(tokenized_document.document(), EqualsProto(document));
+  EXPECT_THAT(
+      tokenized_document.document_wrapper(),
+      EqualsProto(CreateDocumentWrapper(document, /*num_string_tokens=*/0)));
   EXPECT_THAT(tokenized_document.num_string_tokens(), Eq(0));
 
   // string sections
@@ -619,6 +673,40 @@
               ElementsAre("pkg$db/ns#uri2"));
 }
 
+TEST_F(TokenizedDocumentTest,
+       CreateShouldSetCreationTimestampToCurrentTimeIfUnset) {
+  DocumentProto document = DocumentBuilder()
+                               .SetKey("icing", "fake_type/1")
+                               .SetSchema(std::string(kFakeType))
+                               .Build();
+  ASSERT_THAT(document.creation_timestamp_ms(), Eq(0));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_document,
+      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
+                                /*current_time_ms=*/123, document));
+  EXPECT_THAT(
+      tokenized_document.document_wrapper().document().creation_timestamp_ms(),
+      Eq(123));
+}
+
+TEST_F(TokenizedDocumentTest,
+       CreateShouldNotOverwriteExistingCreationTimestamp) {
+  DocumentProto document = DocumentBuilder()
+                               .SetCreationTimestampMs(123)
+                               .SetKey("icing", "fake_type/1")
+                               .SetSchema(std::string(kFakeType))
+                               .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_document,
+      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
+                                /*current_time_ms=*/456, document));
+  EXPECT_THAT(
+      tokenized_document.document_wrapper().document().creation_timestamp_ms(),
+      Eq(123));
+}
+
 }  // namespace
 
 }  // namespace lib

diff --git a/java/src/com/google/android/icing/IcingSearchEngine.java b/java/src/com/google/android/icing/IcingSearchEngine.java
index 9856e09..f44a1f7 100644
--- a/java/src/com/google/android/icing/IcingSearchEngine.java
+++ b/java/src/com/google/android/icing/IcingSearchEngine.java

@@ -14,6 +14,7 @@
 
 package com.google.android.icing;
 
+import com.google.android.icing.proto.BatchGetResultProto;
 import com.google.android.icing.proto.BatchPutResultProto;
 import com.google.android.icing.proto.BlobProto;
 import com.google.android.icing.proto.DebugInfoResultProto;
@@ -24,11 +25,13 @@
 import com.google.android.icing.proto.DeleteResultProto;
 import com.google.android.icing.proto.DocumentProto;
 import com.google.android.icing.proto.GetAllNamespacesResultProto;
+import com.google.android.icing.proto.GetNextPageRequestProto;
 import com.google.android.icing.proto.GetOptimizeInfoResultProto;
 import com.google.android.icing.proto.GetResultProto;
 import com.google.android.icing.proto.GetResultSpecProto;
 import com.google.android.icing.proto.GetSchemaResultProto;
 import com.google.android.icing.proto.GetSchemaTypeResultProto;
+import com.google.android.icing.proto.HandleExpiredDocumentsResultProto;
 import com.google.android.icing.proto.IcingSearchEngineOptions;
 import com.google.android.icing.proto.InitializeResultProto;
 import com.google.android.icing.proto.LogSeverity;
@@ -135,7 +138,7 @@
 
   @Override
   public @NonNull BatchPutResultProto batchPut(@NonNull PutDocumentRequest documents) {
-    return IcingSearchEngineUtils.byteArrayToPutResultProtos(
+    return IcingSearchEngineUtils.byteArrayToBatchPutResultProto(
         icingSearchEngineImpl.batchPut(documents.toByteArray()));
   }
 
@@ -147,6 +150,12 @@
   }
 
   @Override
+  public @NonNull BatchGetResultProto batchGet(@NonNull GetResultSpecProto getResultSpec) {
+    return IcingSearchEngineUtils.byteArrayToBatchGetResultProto(
+        icingSearchEngineImpl.batchGet(getResultSpec.toByteArray()));
+  }
+
+  @Override
   public @NonNull ReportUsageResultProto reportUsage(@NonNull UsageReport usageReport) {
     return IcingSearchEngineUtils.byteArrayToReportUsageResultProto(
         icingSearchEngineImpl.reportUsage(usageReport.toByteArray()));
@@ -175,11 +184,24 @@
   }
 
   @Override
+  public @NonNull SearchResultProto getNextPage(
+      @NonNull GetNextPageRequestProto getNextPageRequest) {
+    return IcingSearchEngineUtils.byteArrayToSearchResultProto(
+        icingSearchEngineImpl.getNextPageWithRequestProto(getNextPageRequest.toByteArray()));
+  }
+
+  @Override
   public void invalidateNextPageToken(long nextPageToken) {
     icingSearchEngineImpl.invalidateNextPageToken(nextPageToken);
   }
 
   @Override
+  public @NonNull HandleExpiredDocumentsResultProto handleExpiredDocuments() {
+    return IcingSearchEngineUtils.byteArrayToHandleExpiredDocumentsResultProto(
+        icingSearchEngineImpl.handleExpiredDocuments());
+  }
+
+  @Override
   public @NonNull BlobProto openWriteBlob(PropertyProto.@NonNull BlobHandleProto blobHandle) {
     return IcingSearchEngineUtils.byteArrayToBlobProto(
         icingSearchEngineImpl.openWriteBlob(blobHandle.toByteArray()));
@@ -204,6 +226,17 @@
   }
 
   @Override
+  public @NonNull BlobProto getAllBlobInfos() {
+    return IcingSearchEngineUtils.byteArrayToBlobProto(icingSearchEngineImpl.getAllBlobInfos());
+  }
+
+  @Override
+  public @NonNull BlobProto putBlobInfos(@NonNull BlobProto blobProto) {
+    return IcingSearchEngineUtils.byteArrayToBlobProto(
+        icingSearchEngineImpl.putBlobInfos(blobProto.toByteArray()));
+  }
+
+  @Override
   public @NonNull DeleteResultProto delete(@NonNull String namespace, @NonNull String uri) {
     return IcingSearchEngineUtils.byteArrayToDeleteResultProto(
         icingSearchEngineImpl.delete(namespace, uri));
@@ -275,6 +308,12 @@
     return IcingSearchEngineUtils.byteArrayToResetResultProto(icingSearchEngineImpl.reset());
   }
 
+  @Override
+  public @NonNull ResetResultProto clearAndDestroy() {
+    return IcingSearchEngineUtils.byteArrayToResetResultProto(
+        icingSearchEngineImpl.clearAndDestroy());
+  }
+
   public static boolean shouldLog(LogSeverity.Code severity) {
     return shouldLog(severity, (short) 0);
   }

diff --git a/java/src/com/google/android/icing/IcingSearchEngineImpl.java b/java/src/com/google/android/icing/IcingSearchEngineImpl.java
index ad5a5aa..0389542 100644
--- a/java/src/com/google/android/icing/IcingSearchEngineImpl.java
+++ b/java/src/com/google/android/icing/IcingSearchEngineImpl.java

@@ -132,6 +132,12 @@
   }
 
   @Nullable
+  public byte[] batchGet(@NonNull byte[] getResultSpecBytes) {
+    throwIfClosed();
+    return nativeBatchGet(this, getResultSpecBytes);
+  }
+
+  @Nullable
   public byte[] reportUsage(@NonNull byte[] usageReportBytes) {
     throwIfClosed();
     return nativeReportUsage(this, usageReportBytes);
@@ -168,11 +174,24 @@
     return nativeGetNextPage(this, nextPageToken, System.currentTimeMillis());
   }
 
+  @Nullable
+  public byte[] getNextPageWithRequestProto(@NonNull byte[] getNextPageRequestBytes) {
+    throwIfClosed();
+    return nativeGetNextPageWithRequestProto(
+        this, getNextPageRequestBytes, System.currentTimeMillis());
+  }
+
   public void invalidateNextPageToken(long nextPageToken) {
     throwIfClosed();
     nativeInvalidateNextPageToken(this, nextPageToken);
   }
 
+  @Nullable
+  public byte[] handleExpiredDocuments() {
+    throwIfClosed();
+    return nativeHandleExpiredDocuments(this);
+  }
+
   @NonNull
   public byte[] openWriteBlob(@NonNull byte[] blobHandleBytes) {
     throwIfClosed();
@@ -197,6 +216,18 @@
     return nativeCommitBlob(this, blobHandleBytes);
   }
 
+  @NonNull
+  public byte[] getAllBlobInfos() {
+    throwIfClosed();
+    return nativeGetAllBlobInfos(this);
+  }
+
+  @NonNull
+  public byte[] putBlobInfos(@NonNull byte[] blobProtoBytes) {
+    throwIfClosed();
+    return nativePutBlobInfos(this, blobProtoBytes);
+  }
+
   @Nullable
   public byte[] delete(@NonNull String namespace, @NonNull String uri) {
     throwIfClosed();
@@ -268,6 +299,12 @@
     return nativeReset(this);
   }
 
+  @Nullable
+  public byte[] clearAndDestroy() {
+    throwIfClosed();
+    return nativeClearAndDestroy(this);
+  }
+
   public static boolean shouldLog(short severity) {
     return shouldLog(severity, (short) 0);
   }
@@ -321,6 +358,9 @@
   private static native byte[] nativeGet(
       IcingSearchEngineImpl instance, String namespace, String uri, byte[] getResultSpecBytes);
 
+  private static native byte[] nativeBatchGet(
+      IcingSearchEngineImpl instance, byte[] getResultSpecBytes);
+
   private static native byte[] nativeReportUsage(
       IcingSearchEngineImpl instance, byte[] usageReportBytes);
 
@@ -336,9 +376,16 @@
   private static native byte[] nativeGetNextPage(
       IcingSearchEngineImpl instance, long nextPageToken, long javaToNativeStartTimestampMs);
 
+  private static native byte[] nativeGetNextPageWithRequestProto(
+      IcingSearchEngineImpl instance,
+      byte[] getNextPageRequestBytes,
+      long javaToNativeStartTimestampMs);
+
   private static native void nativeInvalidateNextPageToken(
       IcingSearchEngineImpl instance, long nextPageToken);
 
+  private static native byte[] nativeHandleExpiredDocuments(IcingSearchEngineImpl instance);
+
   private static native byte[] nativeOpenWriteBlob(
       IcingSearchEngineImpl instance, byte[] blobHandleBytes);
 
@@ -351,6 +398,11 @@
   private static native byte[] nativeCommitBlob(
       IcingSearchEngineImpl instance, byte[] blobHandleBytes);
 
+  private static native byte[] nativeGetAllBlobInfos(IcingSearchEngineImpl instance);
+
+  private static native byte[] nativePutBlobInfos(
+      IcingSearchEngineImpl instance, byte[] blobProtoBytes);
+
   private static native byte[] nativeDelete(
       IcingSearchEngineImpl instance, String namespace, String uri);
 
@@ -373,6 +425,8 @@
 
   private static native byte[] nativeReset(IcingSearchEngineImpl instance);
 
+  private static native byte[] nativeClearAndDestroy(IcingSearchEngineImpl instance);
+
   private static native byte[] nativeSearchSuggestions(
       IcingSearchEngineImpl instance, byte[] suggestionSpecBytes);
 

diff --git a/java/src/com/google/android/icing/IcingSearchEngineInterface.java b/java/src/com/google/android/icing/IcingSearchEngineInterface.java
index c381d7c..bf6ba10 100644
--- a/java/src/com/google/android/icing/IcingSearchEngineInterface.java
+++ b/java/src/com/google/android/icing/IcingSearchEngineInterface.java

@@ -1,5 +1,6 @@
 package com.google.android.icing;
 
+import com.google.android.icing.proto.BatchGetResultProto;
 import com.google.android.icing.proto.BatchPutResultProto;
 import com.google.android.icing.proto.BlobProto;
 import com.google.android.icing.proto.DebugInfoResultProto;
@@ -10,11 +11,13 @@
 import com.google.android.icing.proto.DeleteResultProto;
 import com.google.android.icing.proto.DocumentProto;
 import com.google.android.icing.proto.GetAllNamespacesResultProto;
+import com.google.android.icing.proto.GetNextPageRequestProto;
 import com.google.android.icing.proto.GetOptimizeInfoResultProto;
 import com.google.android.icing.proto.GetResultProto;
 import com.google.android.icing.proto.GetResultSpecProto;
 import com.google.android.icing.proto.GetSchemaResultProto;
 import com.google.android.icing.proto.GetSchemaTypeResultProto;
+import com.google.android.icing.proto.HandleExpiredDocumentsResultProto;
 import com.google.android.icing.proto.InitializeResultProto;
 import com.google.android.icing.proto.OptimizeResultProto;
 import com.google.android.icing.proto.PersistToDiskResultProto;
@@ -105,6 +108,13 @@
    */
   GetResultProto get(String namespace, String uri, GetResultSpecProto getResultSpec);
 
+  /**
+   * Gets a list of documents.
+   *
+   * @param getResultSpec the spec for getting the documents.
+   */
+  BatchGetResultProto batchGet(GetResultSpecProto getResultSpec);
+
   /** Reports usage. */
   ReportUsageResultProto reportUsage(UsageReport usageReport);
 
@@ -120,12 +130,46 @@
   SearchResultProto search(
       SearchSpecProto searchSpec, ScoringSpecProto scoringSpec, ResultSpecProto resultSpec);
 
-  /** Gets the next page. */
+  /**
+   * Gets the next page.
+   *
+   * <p>Note: This method is deprecated. Please use {@link #getNextPage(GetNextPageRequestProto)}
+   * instead.
+   */
   SearchResultProto getNextPage(long nextPageToken);
 
+  /**
+   * Gets the next page.
+   *
+   * @param getNextPageRequest the request proto for getting the next page.
+   */
+  SearchResultProto getNextPage(GetNextPageRequestProto getNextPageRequest);
+
   /** Invalidates the next page token. */
   void invalidateNextPageToken(long nextPageToken);
 
+  /**
+   * Handles expired documents. The operation includes:
+   *
+   * <ul>
+   *   <li>Zero out bytes for all expired documents.
+   *   <li>Propagate delete to child documents with delete propagation enabled. Zero out bytes for
+   *       all propagated documents.
+   * </ul>
+   *
+   * <p>For Icing direct clients, there is no need to call this API since Icing already has a native
+   * background thread for scheduling tasks to call it. This API is exposed only for AppSearch
+   * system service (in Android) since native background thread is not recommended in system
+   * service, and tasks should be scheduled at AppSearch level instead.
+   *
+   * <p>Note: {@link #optimize()} also zeros out bytes for non-alive (delete, expired) documents,
+   * but it also additionally reclaims disk space for storage and rebuilds indices by rearranging
+   * alive documents. On the other hand, {@link #handleExpiredDocuments()} is a cheaper API that
+   * only deals with expired documents and delete propagation for child documents, without
+   * reclaiming disk space or rebuilding indices which requires additional disk I/O.
+   */
+  HandleExpiredDocumentsResultProto handleExpiredDocuments();
+
   /** Gets a file descriptor to write blob data. */
   BlobProto openWriteBlob(PropertyProto.BlobHandleProto blobHandle);
 
@@ -138,6 +182,12 @@
   /** Marks the blob is committed. */
   BlobProto commitBlob(PropertyProto.BlobHandleProto blobHandle);
 
+  /** Gets all the blob info from the blob store. */
+  BlobProto getAllBlobInfos();
+
+  /** Puts the blob info protos from the blob proto to the blob info proto log file. */
+  BlobProto putBlobInfos(BlobProto blobProto);
+
   /**
    * Deletes the document.
    *
@@ -185,6 +235,9 @@
   /** Clears all data from the current icing instance, and reinitializes it. */
   ResetResultProto reset();
 
+  /** Clears all data from the current icing instance. */
+  ResetResultProto clearAndDestroy();
+
   /** Closes the current icing instance. */
   @Override
   void close();

diff --git a/java/src/com/google/android/icing/IcingSearchEngineUtils.java b/java/src/com/google/android/icing/IcingSearchEngineUtils.java
index ca08dd6..1e08262 100644
--- a/java/src/com/google/android/icing/IcingSearchEngineUtils.java
+++ b/java/src/com/google/android/icing/IcingSearchEngineUtils.java

@@ -17,6 +17,7 @@
 import android.util.Log;
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
+import com.google.android.icing.proto.BatchGetResultProto;
 import com.google.android.icing.proto.BatchPutResultProto;
 import com.google.android.icing.proto.BlobProto;
 import com.google.android.icing.proto.DebugInfoResultProto;
@@ -29,6 +30,7 @@
 import com.google.android.icing.proto.GetResultProto;
 import com.google.android.icing.proto.GetSchemaResultProto;
 import com.google.android.icing.proto.GetSchemaTypeResultProto;
+import com.google.android.icing.proto.HandleExpiredDocumentsResultProto;
 import com.google.android.icing.proto.InitializeResultProto;
 import com.google.android.icing.proto.OptimizeResultProto;
 import com.google.android.icing.proto.PersistToDiskResultProto;
@@ -159,18 +161,22 @@
   }
 
   @NonNull
-  public static BatchPutResultProto byteArrayToPutResultProtos(@Nullable byte[] putResultsBytes) {
+  public static BatchPutResultProto byteArrayToBatchPutResultProto(
+      @Nullable byte[] putResultsBytes) {
     if (putResultsBytes == null) {
       Log.e(TAG, "Received null PutResultProtos from native.");
-      // TODO(b/394875109) Or we should create a proto with error status?
-      return BatchPutResultProto.getDefaultInstance();
+      return BatchPutResultProto.newBuilder()
+          .setStatus(StatusProto.newBuilder().setCode(StatusProto.Code.INTERNAL))
+          .build();
     }
 
     try {
       return BatchPutResultProto.parseFrom(putResultsBytes, EXTENSION_REGISTRY_LITE);
     } catch (InvalidProtocolBufferException e) {
       Log.e(TAG, "Error parsing PutResultProtos.", e);
-      return BatchPutResultProto.getDefaultInstance();
+      return BatchPutResultProto.newBuilder()
+          .setStatus(StatusProto.newBuilder().setCode(StatusProto.Code.INTERNAL))
+          .build();
     }
   }
 
@@ -194,6 +200,26 @@
   }
 
   @NonNull
+  public static BatchGetResultProto byteArrayToBatchGetResultProto(
+      @Nullable byte[] batchGetResultBytes) {
+    if (batchGetResultBytes == null) {
+      Log.e(TAG, "Received null BatchGetResultProto from native.");
+      return BatchGetResultProto.newBuilder()
+          .setStatus(StatusProto.newBuilder().setCode(StatusProto.Code.INTERNAL))
+          .build();
+    }
+
+    try {
+      return BatchGetResultProto.parseFrom(batchGetResultBytes, EXTENSION_REGISTRY_LITE);
+    } catch (InvalidProtocolBufferException e) {
+      Log.e(TAG, "Error parsing BatchGetResultProto.", e);
+      return BatchGetResultProto.newBuilder()
+          .setStatus(StatusProto.newBuilder().setCode(StatusProto.Code.INTERNAL))
+          .build();
+    }
+  }
+
+  @NonNull
   public static ReportUsageResultProto byteArrayToReportUsageResultProto(
       @Nullable byte[] reportUsageResultBytes) {
     if (reportUsageResultBytes == null) {
@@ -257,6 +283,33 @@
   }
 
   /**
+   * Converts a byte array to a {@link HandleExpiredDocumentsResultProto}.
+   *
+   * @param handleExpiredDocumentsResultBytes the byte array to convert
+   * @return the {@link HandleExpiredDocumentsResultProto}
+   */
+  @NonNull
+  public static HandleExpiredDocumentsResultProto byteArrayToHandleExpiredDocumentsResultProto(
+      @Nullable byte[] handleExpiredDocumentsResultBytes) {
+    if (handleExpiredDocumentsResultBytes == null) {
+      Log.e(TAG, "Received null HandleExpiredDocumentsResultProto from native.");
+      return HandleExpiredDocumentsResultProto.newBuilder()
+          .setStatus(StatusProto.newBuilder().setCode(StatusProto.Code.INTERNAL))
+          .build();
+    }
+
+    try {
+      return HandleExpiredDocumentsResultProto.parseFrom(
+          handleExpiredDocumentsResultBytes, EXTENSION_REGISTRY_LITE);
+    } catch (InvalidProtocolBufferException e) {
+      Log.e(TAG, "Error parsing HandleExpiredDocumentsResultProto.", e);
+      return HandleExpiredDocumentsResultProto.newBuilder()
+          .setStatus(StatusProto.newBuilder().setCode(StatusProto.Code.INTERNAL))
+          .build();
+    }
+  }
+
+  /**
    * Converts a byte array to a {@link BlobProto}.
    *
    * @param blobBytes the byte array to convert

diff --git a/java/tests/instrumentation/src/com/google/android/icing/IcingSearchEngineTest.java b/java/tests/instrumentation/src/com/google/android/icing/IcingSearchEngineTest.java
index 975699c..cf44b32 100644
--- a/java/tests/instrumentation/src/com/google/android/icing/IcingSearchEngineTest.java
+++ b/java/tests/instrumentation/src/com/google/android/icing/IcingSearchEngineTest.java

@@ -16,8 +16,10 @@
 
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.Truth.assertWithMessage;
+import static org.junit.Assert.assertThrows;
 
 import com.google.android.icing.IcingSearchEngine;
+import com.google.android.icing.proto.BatchGetResultProto;
 import com.google.android.icing.proto.BatchPutResultProto;
 import com.google.android.icing.proto.BlobProto;
 import com.google.android.icing.proto.DebugInfoResultProto;
@@ -28,13 +30,16 @@
 import com.google.android.icing.proto.DeleteResultProto;
 import com.google.android.icing.proto.DocumentProto;
 import com.google.android.icing.proto.GetAllNamespacesResultProto;
+import com.google.android.icing.proto.GetNextPageRequestProto;
 import com.google.android.icing.proto.GetOptimizeInfoResultProto;
 import com.google.android.icing.proto.GetResultProto;
 import com.google.android.icing.proto.GetResultSpecProto;
 import com.google.android.icing.proto.GetSchemaResultProto;
 import com.google.android.icing.proto.GetSchemaTypeResultProto;
+import com.google.android.icing.proto.HandleExpiredDocumentsResultProto;
 import com.google.android.icing.proto.IcingSearchEngineOptions;
 import com.google.android.icing.proto.InitializeResultProto;
+import com.google.android.icing.proto.JoinableConfig;
 import com.google.android.icing.proto.LogSeverity;
 import com.google.android.icing.proto.OptimizeResultProto;
 import com.google.android.icing.proto.PersistToDiskResultProto;
@@ -65,6 +70,7 @@
 import com.google.android.icing.proto.TermMatchType;
 import com.google.android.icing.proto.TermMatchType.Code;
 import com.google.android.icing.proto.UsageReport;
+import com.google.common.collect.ImmutableList;
 import com.google.android.icing.protobuf.ByteString;
 import java.io.File;
 import java.io.FileDescriptor;
@@ -212,9 +218,6 @@
     assertThat(getSchemaTypeResultProto.getSchemaTypeConfig()).isEqualTo(emailTypeConfig);
   }
 
-  // TODO: b/383379132 - Re-enable this test once the JNI API is pre-registered and dropped back
-  // into g3.
-  @Ignore
   @Test
   public void setAndGetSchemaWithDatabase_ok() throws Exception {
     IcingSearchEngineOptions options =
@@ -316,6 +319,7 @@
             .build();
     BatchPutResultProto batchPutResultProto = icingSearchEngine.batchPut(putDocumentRequest);
 
+    assertStatusOk(batchPutResultProto.getStatus());
     assertThat(batchPutResultProto.getPutResultProtos(0).getUri()).isEqualTo("uri1");
     assertStatusOk(batchPutResultProto.getPutResultProtos(0).getStatus());
     assertThat(batchPutResultProto.getPutResultProtos(1).getUri()).isEqualTo("uri2");
@@ -326,16 +330,81 @@
     assertThat(batchPutResultProto.getPersistToDiskResultProto().getStatus().getCode())
         .isEqualTo(StatusProto.Code.UNKNOWN);
 
-    // Check document 1
-    GetResultProto getResultProto =
-        icingSearchEngine.get("namespace", "uri1", GetResultSpecProto.getDefaultInstance());
-    assertStatusOk(getResultProto.getStatus());
-    assertThat(getResultProto.getDocument()).isEqualTo(emailDocument1);
-    // check document 2
-    getResultProto =
-        icingSearchEngine.get("namespace", "uri2", GetResultSpecProto.getDefaultInstance());
-    assertStatusOk(getResultProto.getStatus());
-    assertThat(getResultProto.getDocument()).isEqualTo(emailDocument2);
+    GetResultSpecProto getResultSpecProto =
+        GetResultSpecProto.newBuilder()
+            .setNamespaceRequested("namespace")
+            .addIds("uri1")
+            .addIds("uri2")
+            .build();
+    BatchGetResultProto batchGetResultProto = icingSearchEngine.batchGet(getResultSpecProto);
+
+    assertStatusOk(batchGetResultProto.getStatus());
+    // Check doc1
+    DocumentProto document = batchGetResultProto.getGetResultProtos(0).getDocument();
+    assertStatusOk(batchGetResultProto.getGetResultProtos(0).getStatus());
+    assertThat(document).isEqualTo(emailDocument1);
+    // Check doc2
+    document = batchGetResultProto.getGetResultProtos(1).getDocument();
+    assertStatusOk(batchGetResultProto.getGetResultProtos(1).getStatus());
+    assertThat(document).isEqualTo(emailDocument2);
+  }
+
+  @Test
+  public void testBatchGetWithEmptyResult() throws Exception {
+    assertStatusOk(icingSearchEngine.initialize().getStatus());
+
+    SchemaTypeConfigProto emailTypeConfig = createEmailTypeConfig();
+    SchemaProto schema = SchemaProto.newBuilder().addTypes(emailTypeConfig).build();
+    assertThat(
+            icingSearchEngine
+                .setSchema(schema, /* ignoreErrorsAndDeleteDocuments= */ false)
+                .getStatus()
+                .getCode())
+        .isEqualTo(StatusProto.Code.OK);
+
+    DocumentProto emailDocument1 = createEmailDocument("namespace", "uri1");
+    DocumentProto emailDocument2 = createEmailDocument("namespace", "uri2");
+    PutDocumentRequest putDocumentRequest =
+        PutDocumentRequest.newBuilder()
+            .addDocuments(emailDocument1)
+            .addDocuments(emailDocument2)
+            .build();
+    BatchPutResultProto batchPutResultProto = icingSearchEngine.batchPut(putDocumentRequest);
+    assertStatusOk(batchPutResultProto.getStatus());
+
+    // no ids.
+    GetResultSpecProto getResultSpecProto =
+        GetResultSpecProto.newBuilder().setNamespaceRequested("namespace").build();
+    BatchGetResultProto batchGetResultProto = icingSearchEngine.batchGet(getResultSpecProto);
+
+    // Check no doc returned if no ids are specified.
+    assertStatusOk(batchGetResultProto.getStatus());
+    assertThat(batchGetResultProto.getGetResultProtosList()).isEmpty();
+
+    // empty namespace.
+    getResultSpecProto = GetResultSpecProto.newBuilder().addIds("uri1").build();
+    batchGetResultProto = icingSearchEngine.batchGet(getResultSpecProto);
+    assertStatusOk(batchGetResultProto.getStatus());
+    assertThat(batchGetResultProto.getGetResultProtosList()).hasSize(1);
+    assertThat(batchGetResultProto.getGetResultProtos(0).getStatus().getCode())
+        .isEqualTo(StatusProto.Code.NOT_FOUND);
+
+    // different namespace.
+    getResultSpecProto =
+        GetResultSpecProto.newBuilder()
+            .setNamespaceRequested("otherNameSpace")
+            .addIds("uri1")
+            .addIds("uri2")
+            .build();
+    batchGetResultProto = icingSearchEngine.batchGet(getResultSpecProto);
+
+    // Check not found returned if namespace is different.
+    assertStatusOk(batchGetResultProto.getStatus());
+    assertThat(batchGetResultProto.getGetResultProtosList()).hasSize(2);
+    assertThat(batchGetResultProto.getGetResultProtos(0).getStatus().getCode())
+        .isEqualTo(StatusProto.Code.NOT_FOUND);
+    assertThat(batchGetResultProto.getGetResultProtos(1).getStatus().getCode())
+        .isEqualTo(StatusProto.Code.NOT_FOUND);
   }
 
   @Test
@@ -391,13 +460,18 @@
 
     PutDocumentRequest putDocumentRequest = PutDocumentRequest.getDefaultInstance();
     BatchPutResultProto batchPutResultProto = icingSearchEngine.batchPut(putDocumentRequest);
+    BatchPutResultProto actualProto =
+        batchPutResultProto.toBuilder().clearVmBinderTransactionLatencyStartTimeMs().build();
 
-    BatchPutResultProto expected = BatchPutResultProto.getDefaultInstance();
-    assertThat(batchPutResultProto).isEqualTo(expected);
+    BatchPutResultProto expected =
+        BatchPutResultProto.newBuilder()
+            .setStatus(StatusProto.newBuilder().setCode(StatusProto.Code.OK))
+            .build();
+    assertThat(actualProto).isEqualTo(expected);
 
     // PersistToDiskResultProto should not be set if persist_type is not set in the
     // PutDocumentRequest.
-    assertThat(batchPutResultProto.getPersistToDiskResultProto().getStatus().getCode())
+    assertThat(actualProto.getPersistToDiskResultProto().getStatus().getCode())
         .isEqualTo(StatusProto.Code.UNKNOWN);
   }
 
@@ -593,6 +667,84 @@
     assertThat(searchResultProto.getResultsCount()).isEqualTo(0);
   }
 
+  // TODO: b/417644758 - Re-enable this test once the JNI API is pre-registered and dropped back
+  // into g3.
+  @Ignore
+  @Test
+  public void getNextPageWithRequestProto() throws Exception {
+    assertStatusOk(icingSearchEngine.initialize().getStatus());
+
+    SchemaTypeConfigProto emailTypeConfig = createEmailTypeConfig();
+    SchemaProto schema = SchemaProto.newBuilder().addTypes(emailTypeConfig).build();
+    assertThat(
+            icingSearchEngine
+                .setSchema(schema, /* ignoreErrorsAndDeleteDocuments= */ false)
+                .getStatus()
+                .getCode())
+        .isEqualTo(StatusProto.Code.OK);
+
+    Map<String, DocumentProto> documents = new HashMap<>();
+    for (int i = 0; i < 10; i++) {
+      DocumentProto emailDocument =
+          createEmailDocument("namespace", "uri:" + i).toBuilder()
+              .addProperties(PropertyProto.newBuilder().setName("subject").addStringValues("foo"))
+              .build();
+      documents.put("uri:" + i, emailDocument);
+      assertWithMessage(icingSearchEngine.put(emailDocument).getStatus().getMessage())
+          .that(icingSearchEngine.put(emailDocument).getStatus().getCode())
+          .isEqualTo(StatusProto.Code.OK);
+    }
+
+    SearchSpecProto searchSpec =
+        SearchSpecProto.newBuilder()
+            .setQuery("foo")
+            .setTermMatchType(TermMatchType.Code.PREFIX)
+            .build();
+    ResultSpecProto resultSpecProto = ResultSpecProto.newBuilder().setNumPerPage(2).build();
+
+    SearchResultProto searchResultProto =
+        icingSearchEngine.search(
+            searchSpec, ScoringSpecProto.getDefaultInstance(), resultSpecProto);
+    assertStatusOk(searchResultProto.getStatus());
+    assertThat(searchResultProto.getResultsCount()).isEqualTo(2);
+    DocumentProto resultDocument1 = searchResultProto.getResults(0).getDocument();
+    DocumentProto resultDocument2 = searchResultProto.getResults(1).getDocument();
+    assertThat(resultDocument1).isEqualTo(documents.remove(resultDocument1.getUri()));
+    assertThat(resultDocument2).isEqualTo(documents.remove(resultDocument2.getUri()));
+
+    assertThat(searchResultProto.getQueryStats().hasNativeToJavaStartTimestampMs()).isTrue();
+    assertThat(searchResultProto.getQueryStats().hasNativeToJavaJniLatencyMs()).isTrue();
+    assertThat(searchResultProto.getQueryStats().hasJavaToNativeJniLatencyMs()).isTrue();
+    assertThat(searchResultProto.getQueryStats().getNativeToJavaStartTimestampMs())
+        .isGreaterThan(0);
+    assertThat(searchResultProto.getQueryStats().getNativeToJavaJniLatencyMs()).isAtLeast(0);
+    assertThat(searchResultProto.getQueryStats().getJavaToNativeJniLatencyMs()).isAtLeast(0);
+
+    GetNextPageRequestProto getNextPageRequestProto =
+        GetNextPageRequestProto.newBuilder()
+            .setNextPageToken(searchResultProto.getNextPageToken())
+            .setMaxResultsToRetrieveFromPage(1)
+            .build();
+    // fetch rest pages
+    for (int i = 1; i < 5; i++) {
+      searchResultProto = icingSearchEngine.getNextPage(getNextPageRequestProto);
+      DocumentProto resultDocument = searchResultProto.getResults(0).getDocument();
+      assertWithMessage(searchResultProto.getStatus().getMessage())
+          .that(searchResultProto.getStatus().getCode())
+          .isEqualTo(StatusProto.Code.OK);
+      assertThat(searchResultProto.getResultsCount()).isEqualTo(1);
+      assertThat(resultDocument).isEqualTo(documents.remove(resultDocument.getUri()));
+    }
+
+    // invalidate rest result
+    icingSearchEngine.invalidateNextPageToken(searchResultProto.getNextPageToken());
+
+    searchResultProto = icingSearchEngine.getNextPage(getNextPageRequestProto);
+    assertStatusOk(searchResultProto.getStatus());
+    assertThat(searchResultProto.getResultsCount()).isEqualTo(0);
+  }
+
+
   @Ignore // b/350530146
   @Test
   public void writeAndReadBlob_blobContentMatches() throws Exception {
@@ -606,7 +758,7 @@
     IcingSearchEngine icing = new IcingSearchEngine(options);
     assertStatusOk(icing.initialize().getStatus());
 
-    byte[] data = generateRandomBytes(100); // 10 Bytes
+    byte[] data = generateRandomBytes(100); // 100 Bytes
     byte[] digest = calculateDigest(data);
     PropertyProto.BlobHandleProto blobHandle =
         PropertyProto.BlobHandleProto.newBuilder()
@@ -661,7 +813,7 @@
     IcingSearchEngine icing = new IcingSearchEngine(options);
     assertStatusOk(icing.initialize().getStatus());
 
-    byte[] data = generateRandomBytes(100); // 10 Bytes
+    byte[] data = generateRandomBytes(100); // 100 Bytes
     byte[] digest = calculateDigest(data);
     PropertyProto.BlobHandleProto blobHandle =
         PropertyProto.BlobHandleProto.newBuilder()
@@ -694,6 +846,74 @@
     assertThat(commitBlobProto.getStatus().getCode()).isEqualTo(StatusProto.Code.NOT_FOUND);
   }
 
+  @Ignore // b/350530146
+  @Test
+  public void getAndPutBlobInfo() throws Exception {
+    // 1 Arrange: set up IcingSearchEngine with and blob data
+    File tempDir = temporaryFolder.newFolder();
+    IcingSearchEngineOptions options =
+        IcingSearchEngineOptions.newBuilder()
+            .setBaseDir(tempDir.getCanonicalPath())
+            .setEnableBlobStore(true)
+            .build();
+    IcingSearchEngine icing = new IcingSearchEngine(options);
+    assertStatusOk(icing.initialize().getStatus());
+
+    byte[] data = generateRandomBytes(100); // 100 Bytes
+    byte[] digest = calculateDigest(data);
+    PropertyProto.BlobHandleProto blobHandle =
+        PropertyProto.BlobHandleProto.newBuilder()
+            .setNamespace("ns")
+            .setDigest(ByteString.copyFrom(digest))
+            .build();
+
+    // 2 Act: write the blob
+    BlobProto openWriteBlobProto = icing.openWriteBlob(blobHandle);
+    assertStatusOk(openWriteBlobProto.getStatus());
+    Field field = FileDescriptor.class.getDeclaredField("fd");
+    field.setAccessible(true); // Make the field accessible
+
+    // Create a new FileDescriptor object
+    FileDescriptor writeFd = new FileDescriptor();
+
+    // Set the file descriptor value using reflection
+    field.setInt(writeFd, openWriteBlobProto.getFileDescriptor());
+
+    try (FileOutputStream outputStream = new FileOutputStream(writeFd)) {
+      outputStream.write(data);
+    }
+
+    // Commit and read the blob.
+    BlobProto commitBlobProto = icing.commitBlob(blobHandle);
+    assertStatusOk(commitBlobProto.getStatus());
+
+    // Get the blob info.
+    BlobProto allBlobInfoProto = icing.getAllBlobInfos();
+    assertStatusOk(allBlobInfoProto.getStatus());
+    assertThat(allBlobInfoProto.getBlobInfoProtosCount()).isEqualTo(1);
+
+    // create a new IcingSearchEngine instance.
+    IcingSearchEngineOptions options2 =
+        IcingSearchEngineOptions.newBuilder()
+            .setBaseDir(tempDir.getCanonicalPath())
+            .setEnableBlobStore(true)
+            .setManageBlobFiles(false)
+            .build();
+    IcingSearchEngine icing2 = new IcingSearchEngine(options2);
+    assertStatusOk(icing2.initialize().getStatus());
+
+    // Put the blob info.
+    BlobProto putBlobInfosProto = icing2.putBlobInfos(allBlobInfoProto);
+    assertStatusOk(putBlobInfosProto.getStatus());
+
+    // 3: Verify the blob info.
+    BlobProto allBlobInfoProto2 = icing2.getAllBlobInfos();
+    assertStatusOk(allBlobInfoProto2.getStatus());
+    assertThat(allBlobInfoProto2.getBlobInfoProtosCount()).isEqualTo(1);
+    assertThat(allBlobInfoProto2.getBlobInfoProtos(0))
+        .isEqualTo(allBlobInfoProto2.getBlobInfoProtos(0));
+  }
+
   @Test
   public void testDelete() throws Exception {
     assertStatusOk(icingSearchEngine.initialize().getStatus());
@@ -968,6 +1188,45 @@
   }
 
   @Test
+  public void testClearAndDestroy() throws Exception {
+    assertStatusOk(icingSearchEngine.initialize().getStatus());
+
+    // Simple put and get
+    SchemaTypeConfigProto emailTypeConfig = createEmailTypeConfig();
+    SchemaProto schema = SchemaProto.newBuilder().addTypes(emailTypeConfig).build();
+    assertThat(
+            icingSearchEngine
+                .setSchema(schema, /* ignoreErrorsAndDeleteDocuments= */ false)
+                .getStatus()
+                .getCode())
+        .isEqualTo(StatusProto.Code.OK);
+
+    DocumentProto emailDocument = createEmailDocument("namespace", "uri");
+    PutResultProto putResultProto = icingSearchEngine.put(emailDocument);
+    assertStatusOk(putResultProto.getStatus());
+
+    GetResultProto getResultProto =
+        icingSearchEngine.get("namespace", "uri", GetResultSpecProto.getDefaultInstance());
+    assertStatusOk(getResultProto.getStatus());
+    assertThat(getResultProto.getDocument()).isEqualTo(emailDocument);
+
+    // Clear and destroy
+    ResetResultProto clearAndDestroyResult = icingSearchEngine.clearAndDestroy();
+    assertStatusOk(clearAndDestroyResult.getStatus());
+
+    // Try to put and get again, but it should fail since the instance is
+    // uninitialized after clearAndDestroy().
+    assertThat(icingSearchEngine.put(emailDocument).getStatus().getCode())
+        .isEqualTo(StatusProto.Code.FAILED_PRECONDITION);
+    assertThat(
+            icingSearchEngine
+                .get("namespace", "uri", GetResultSpecProto.getDefaultInstance())
+                .getStatus()
+                .getCode())
+        .isEqualTo(StatusProto.Code.FAILED_PRECONDITION);
+  }
+
+  @Test
   public void testReportUsage() throws Exception {
     assertStatusOk(icingSearchEngine.initialize().getStatus());
 
@@ -1205,4 +1464,44 @@
   private static void assertStatusOk(StatusProto status) {
     assertWithMessage(status.getMessage()).that(status.getCode()).isEqualTo(StatusProto.Code.OK);
   }
+
+  @Test
+  public void throwIfClosed() throws Exception {
+    icingSearchEngine.close();
+    assertThrows(IllegalStateException.class, () -> icingSearchEngine.initialize());
+    assertThrows(
+        IllegalStateException.class,
+        () -> icingSearchEngine.setSchema(SchemaProto.getDefaultInstance()));
+    assertThrows(IllegalStateException.class, () -> icingSearchEngine.getSchema());
+    assertThrows(IllegalStateException.class, () -> icingSearchEngine.getSchemaType("type"));
+    assertThrows(
+        IllegalStateException.class,
+        () -> icingSearchEngine.put(DocumentProto.getDefaultInstance()));
+    assertThrows(
+        IllegalStateException.class,
+        () -> icingSearchEngine.batchPut(PutDocumentRequest.getDefaultInstance()));
+    assertThrows(
+        IllegalStateException.class,
+        () -> icingSearchEngine.get("namespace", "uri", GetResultSpecProto.getDefaultInstance()));
+    assertThrows(
+        IllegalStateException.class,
+        () -> icingSearchEngine.batchGet(GetResultSpecProto.getDefaultInstance()));
+    assertThrows(IllegalStateException.class, () -> icingSearchEngine.optimize());
+    assertThrows(
+        IllegalStateException.class, () -> icingSearchEngine.persistToDisk(PersistType.Code.LITE));
+    assertThrows(
+        IllegalStateException.class,
+        () -> icingSearchEngine.reportUsage(UsageReport.getDefaultInstance()));
+    assertThrows(IllegalStateException.class, () -> icingSearchEngine.getStorageInfo());
+    assertThrows(
+        IllegalStateException.class,
+        () ->
+            icingSearchEngine.search(
+                SearchSpecProto.getDefaultInstance(),
+                ScoringSpecProto.getDefaultInstance(),
+                ResultSpecProto.getDefaultInstance()));
+    assertThrows(
+        IllegalStateException.class,
+        () -> icingSearchEngine.getNextPage(GetNextPageRequestProto.getDefaultInstance()));
+  }
 }

diff --git a/proto/appsearch/account.proto b/proto/appsearch/account.proto
new file mode 100644
index 0000000..0ee8afa
--- /dev/null
+++ b/proto/appsearch/account.proto

@@ -0,0 +1,36 @@
+
+// Copyright (C) 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+syntax = "proto2";
+
+package appsearch;
+
+option java_package = "com.google.android.appsearch.proto";
+option java_multiple_files = true;
+option objc_class_prefix = "APPSEARCH";
+
+// The proto to reprent an AppSearchAccount object
+// Next tag: 6
+message AccountProto {
+  // The namespace the account document belongs to .
+  optional string namespace = 1;
+  // The unique identifier for the document within the namespace.
+  optional string id = 2;
+  // The type of the account. This corresponds to android.accounts.Account.type.
+  optional string account_type = 3;
+  // The name of the account. This corresponds to android.accounts.Account.name.
+  optional string account_name = 4;
+  // An opaque, optional identifier for the account, like gaia id.
+  optional string account_id = 5;
+}

diff --git a/proto/icing/proto/blob.proto b/proto/icing/proto/blob.proto
index b9a8fc2..0df640a 100644
--- a/proto/icing/proto/blob.proto
+++ b/proto/icing/proto/blob.proto

@@ -26,7 +26,7 @@
 // Defines the blob operation result proto that user try to
 // read/write/commit/remove a blob to Icing.
 //
-// Next tag: 4
+// Next tag: 7
 message BlobProto {
   // Status code can be one of:
   //   OK
@@ -43,6 +43,18 @@
   // If Icing does not manage blob files, this is the file name of the blob file
   // from Icing returned by OpenWriteBlob, OpenReadBlob and RemoveBlob.
   optional string file_name = 3;
+
+  // The latency of the GetVm call.
+  optional int32 get_vm_latency_ms = 4;
+
+  // The blob info protos that are stored in the blob info proto log file. This
+  // is used migrate the blob info proto log file to another Icing instance.
+  // This field will only be set via GetAllBlobInfo call,
+  repeated BlobInfoProto blob_info_protos = 5;
+
+  // The start time of the binder transaction latency. This is used to measure
+  // the latency of passing data from the VM via the binder.
+  optional int64 vm_binder_transaction_latency_start_time_ms = 6;
 }
 
 // BlobInfo holds information about a blob. It is used to store in the blob

diff --git a/proto/icing/proto/document.proto b/proto/icing/proto/document.proto
index da82d60..84870d6 100644
--- a/proto/icing/proto/document.proto
+++ b/proto/icing/proto/document.proto

@@ -130,17 +130,30 @@
 }
 
 // Holds a list of PutResultProto.
-// Next tag: 3
+// Next tag: 6
 message BatchPutResultProto {
+  // The overall status of doing the batch put. It is independent of the status
+  // of the individual PutResultProtos. E.g. we may still return OK here even if
+  // all the PutResultProto have errors. Or return INTERNAL_ERROR if there is
+  // some issue parsing this proto, even if all documents are indexed
+  // successfully.
+  optional StatusProto status = 3;
   repeated PutResultProto put_result_protos = 1;
 
   // The result of calling PersistToDisk at the end of the BatchPut request if
   // PutDocumentRequest.persist_type is set.
   optional PersistToDiskResultProto persist_to_disk_result_proto = 2;
+
+  // The latency of the GetVm call.
+  optional int32 get_vm_latency_ms = 4;
+
+  // The start time of the binder transaction latency. This is used to measure
+  // the latency of passing data from the VM via the binder.
+  optional int64 vm_binder_transaction_latency_start_time_ms = 5;
 }
 
 // Result of a call to IcingSearchEngine.Put
-// Next tag: 5
+// Next tag: 8
 message PutResultProto {
   // Status code can be one of:
   //   OK
@@ -166,10 +179,26 @@
 
   // Whether or not the document was a replacement of an existing document.
   optional bool was_replacement = 3;
+
+  // The latency of the GetVm call.
+  optional int32 get_vm_latency_ms = 5;
+
+  // The start time of the binder transaction latency. This is used to measure
+  // the latency of passing data from the VM via the binder.
+  optional int64 vm_binder_transaction_latency_start_time_ms = 6;
+
+  // The expiration timestamp (in milliseconds) of the newly added document.
+  // This field can be used by the caller (mostly AppSearch) for expired
+  // document and expiration propagation handling.
+  // - INT64_MAX means the new document never expires.
+  // - This field is set only when the document is added successfully (i.e.
+  //   status is OK). The caller should ignore this field if the value is 0,
+  //   since it is impossible to have a document expires at t = 0.
+  optional int64 document_expiration_timestamp_ms = 7;
 }
 
 // Result of a call to IcingSearchEngine.Get
-// Next tag: 3
+// Next tag: 6
 message GetResultProto {
   // Status code can be one of:
   //   OK
@@ -183,13 +212,41 @@
   // go/icing-library-apis.
   optional StatusProto status = 1;
 
+  // The uri of the document that was retrieved, or failed to be retrieved.
+  // It won't be set now for a single Get call, as the caller knows the uri if
+  // the Get fails.
+  // TODO(b/404275015) We should always set this field in Get once we refactor
+  // the tests.
+  optional string uri = 3;
+
   // Copy of the Document proto with the specified name_space, uri. Modifying
   // this does not affect the Document in IcingSearchEngine.
   optional DocumentProto document = 2;
+
+  // The latency of the GetVm call.
+  optional int32 get_vm_latency_ms = 4;
+
+  // The start time of the binder transaction latency. This is used to measure
+  // the latency of passing data from the VM via the binder.
+  optional int64 vm_binder_transaction_latency_start_time_ms = 5;
+}
+
+// Result of a call to IcingSearchEngine.BatchGet
+// Next tag: 5
+message BatchGetResultProto {
+  optional StatusProto status = 1;
+  repeated GetResultProto get_result_protos = 2;
+
+  // The latency of the GetVm call.
+  optional int32 get_vm_latency_ms = 3;
+
+  // The start time of the binder transaction latency. This is used to measure
+  // the latency of passing data from the VM via the binder.
+  optional int64 vm_binder_transaction_latency_start_time_ms = 4;
 }
 
 // Result of a call to IcingSearchEngine.GetAllNamespaces
-// Next tag: 3
+// Next tag: 5
 message GetAllNamespacesResultProto {
   // Status code can be one of:
   //   OK
@@ -200,10 +257,17 @@
   // List of namespaces which have at least one existing document in it (not
   // deleted and not expired). Order of namespaces is undefined.
   repeated string namespaces = 2;
+
+  // The latency of the GetVm call.
+  optional int32 get_vm_latency_ms = 3;
+
+  // The start time of the binder transaction latency. This is used to measure
+  // the latency of passing data from the VM via the binder.
+  optional int64 vm_binder_transaction_latency_start_time_ms = 4;
 }
 
 // Result of a call to IcingSearchEngine.Delete
-// Next tag: 3
+// Next tag: 5
 message DeleteResultProto {
   // Status code can be one of:
   //   OK
@@ -219,10 +283,17 @@
 
   // Stats for delete execution performance.
   optional DeleteStatsProto delete_stats = 2;
+
+  // The latency of the GetVm call.
+  optional int32 get_vm_latency_ms = 3;
+
+  // The start time of the binder transaction latency. This is used to measure
+  // the latency of passing data from the VM via the binder.
+  optional int64 vm_binder_transaction_latency_start_time_ms = 4;
 }
 
 // Result of a call to IcingSearchEngine.DeleteByNamespace
-// Next tag: 3
+// Next tag: 5
 message DeleteByNamespaceResultProto {
   // Status code can be one of:
   //   OK
@@ -238,10 +309,17 @@
 
   // Stats for delete execution performance.
   optional DeleteStatsProto delete_stats = 2;
+
+  // The latency of the GetVm call.
+  optional int32 get_vm_latency_ms = 3;
+
+  // The start time of the binder transaction latency. This is used to measure
+  // the latency of passing data from the VM via the binder.
+  optional int64 vm_binder_transaction_latency_start_time_ms = 4;
 }
 
 // Result of a call to IcingSearchEngine.DeleteBySchemaType
-// Next tag: 3
+// Next tag: 5
 message DeleteBySchemaTypeResultProto {
   // Status code can be one of:
   //   OK
@@ -257,10 +335,17 @@
 
   // Stats for delete execution performance.
   optional DeleteStatsProto delete_stats = 2;
+
+  // The latency of the GetVm call.
+  optional int32 get_vm_latency_ms = 3;
+
+  // The start time of the binder transaction latency. This is used to measure
+  // the latency of passing data from the VM via the binder.
+  optional int64 vm_binder_transaction_latency_start_time_ms = 4;
 }
 
 // Result of a call to IcingSearchEngine.DeleteByQuery
-// Next tag: 5
+// Next tag: 7
 message DeleteByQueryResultProto {
   // Status code can be one of:
   //   OK
@@ -291,4 +376,49 @@
   repeated DocumentGroupInfo deleted_documents = 4;
 
   reserved 2;
+
+  // The latency of the GetVm call.
+  optional int32 get_vm_latency_ms = 5;
+
+  // The start time of the binder transaction latency. This is used to measure
+  // the latency of passing data from the VM via the binder.
+  optional int64 vm_binder_transaction_latency_start_time_ms = 6;
+}
+
+// Next tag: 6
+message HandleExpiredDocumentsResultProto {
+  // Status code can be one of:
+  //   OK
+  //   FAILED_PRECONDITION
+  //   NOT_FOUND
+  //   INTERNAL
+  //
+  // See status.proto for more details.
+  //
+  // TODO(b/147699081): Fix error codes: +ABORTED.
+  // go/icing-library-apis.
+  optional StatusProto status = 1;
+
+  // Metadata for the deleted documents, grouped by namespace and schema.
+  message DocumentGroupInfo {
+    optional string name_space = 1;
+    optional string schema = 2;
+    repeated string uris = 3;
+  }
+  // A list of groups of deleted documents. A document was deleted because:
+  // - It was expired.
+  // - OR it was a child document of an expired document with delete propagation
+  //   enabled.
+  repeated DocumentGroupInfo deleted_documents = 2;
+
+  // The number of expired documents that were deleted. This DOES NOT include
+  // the propagated child documents.
+  optional int32 num_expired_documents = 3;
+
+  // The number of deleted documents that were propagated from the expired
+  // documents.
+  optional int32 num_propagated_deleted_documents = 4;
+
+  // The expiration timestamp (in milliseconds) of the next expired document.
+  optional int64 next_expiration_timestamp_ms = 5;
 }

diff --git a/proto/icing/proto/initialize.proto b/proto/icing/proto/initialize.proto
index 5977aa9..40a2af0 100644
--- a/proto/icing/proto/initialize.proto
+++ b/proto/icing/proto/initialize.proto

@@ -17,6 +17,7 @@
 package icing.lib;
 
 import "icing/proto/logging.proto";
+import "icing/proto/persist.proto";
 import "icing/proto/status.proto";
 
 option java_package = "com.google.android.icing.proto";
@@ -71,8 +72,18 @@
     // join index v3 will be rebuilt to replace v2.
     FEATURE_QUALIFIED_ID_JOIN_INDEX_V3 = 6;
 
-    // TODO(b/384947619): decide whether need to add a feature type for delete
-    // propagation.
+    // Feature for flag
+    // IcingSearchEngineOptions::enable_delete_propagation_from.
+    //
+    // This feature covers whether to enable delete propagation PROPAGATE_FROM.
+    FEATURE_DELETE_PROPAGATION_FROM = 7;
+
+    // Feature for flag
+    // IcingSearchEngineOptions::enable_non_existent_qualified_id_join.
+    //
+    // This feature enables handling of cases where a child document is indexed
+    // before its join parent document.
+    FEATURE_NON_EXISTENT_QUALIFIED_ID_JOIN = 8;
   }
 
   // Whether the feature requires the document store to be rebuilt.
@@ -137,7 +148,7 @@
   optional OperationType.Code operation_type = 1;
 }
 
-// Next tag: 35
+// Next tag: 54
 message IcingSearchEngineOptions {
   // Directory to persist files for Icing. Required.
   // If Icing was previously initialized with this directory, it will reload
@@ -193,6 +204,16 @@
   // Optional.
   optional int32 compression_level = 7 [default = 3];
 
+  // Memory level for compression.
+  // 1 uses minimum memory but is slow and reduces compression ratio; 9 uses
+  // maximum memory for optimal speed and compression ratio.
+  // Icing historically used a memLevel of 8. Therefore, that value will be the
+  // default.
+  //
+  // Valid values: [1, 9]
+  // Optional.
+  optional int32 compression_mem_level = 37 [default = 8];
+
   // OPTIONAL: Whether to allow circular references between schema types for
   // the schema definition.
   //
@@ -290,21 +311,19 @@
   // Optional.
   optional int32 blob_store_compression_level = 23 [default = 3];
 
+  // Memory level for compression in blob store.
+  // 1 uses minimum memory but is slow and reduces compression ratio; 9 uses
+  // maximum memory for optimal speed and compression ratio.
+  // Icing historically used a memLevel of 8. Therefore, that value will be the
+  // default.
+  //
+  // Valid values: [1, 9]
+  // Optional.
+  optional int32 blob_store_compression_mem_level = 38 [default = 8];
+
   // Whether to allow repeated fields to have a joinable value type.
   optional bool enable_repeated_field_joins = 24;
 
-  // DEPRECATED (separate them into 2 flags): whether to use qualified id join
-  // index v3 and enable delete propagation PROPAGATE_FROM.
-  //
-  // - If set to true, qualified id join index v3 will be created and delete
-  //   propagation PROPAGATE_FROM will be enabled.
-  // - Otherwise, qualified id join index v2 will be created and delete
-  //   propagation will be disabled.
-  //
-  // The default value is false.
-  optional bool enable_qualified_id_join_index_v3_and_delete_propagate_from =
-      25;
-
   // The absolute path to the ICU data file.
   // If set, ICU will be initialized using this data file.
   optional string icu_data_file_absolute_path = 26;
@@ -361,11 +380,121 @@
   // the overlay exists.
   optional bool release_backup_schema_file_if_overlay_present = 34;
 
-  reserved 2, 20;
+  // Whether to enable strict byte size enforcement on a result page.
+  // - If set to true, then ResultRetrieverV2 will not include the next result
+  //   document if it would exceed the page byte size threshold (except for the
+  //   first document).
+  // - Otherwise, the page byte size enforcement works as a soft limit, i.e. it
+  //   will keep adding documents until exceeding the threshold.
+  optional bool enable_strict_page_byte_size_limit = 35;
+
+  // The threshold in bytes for compressing documents. If a document is larger
+  // than or equal to this threshold, it will be compressed based on
+  // compression_level. 0 means always compress.
+  optional uint32 compression_threshold_bytes = 36;
+
+  // Whether to enable using smaller input buffer size for decompression
+  // If set to false, the default buffer size of 64kb will be used.
+  optional bool enable_smaller_decompression_buffer_size = 39;
+
+  // Whether to enable the Eigen library for embedding scoring.
+  // If set to true **and** Eigen is compiled in (when ICING_DISABLE_EIGEN is
+  // not defined), Eigen will be used for embedding scoring.
+  optional bool enable_eigen_embedding_scoring = 40;
+
+  // Whether to enable passing down filter predicates to child iterators during
+  // DocHitInfo retrieval. If set to true, filter predicates may be passed to
+  // children for optimized retrieval.
+  optional bool enable_passing_filter_to_children = 41;
+
+  // Whether to enable the new header format (refactor legacy format and
+  // introduce unsynced tail checksum) related changes in
+  // PortableFileBackedProtoLog.
+  optional bool enable_proto_log_new_header_format = 42;
+
+  // Whether to enable the new embedding iterator DocHitInfoIteratorEmbeddingV2.
+  // This new iterator optimizes the access order of the embedding data for
+  // sequential reads.
+  optional bool enable_embedding_iterator_v2 = 43;
+
+  // Whether PortableFileBackedProtoLog should retain a decompression buffer
+  // that reads can reuse rather than allocating a new one for each read.
+  optional bool enable_reusable_decompression_buffer = 44;
+
+  // Whether to enable schema-type id optimization for setSchema. When enabled,
+  // the type-ids of existing types will be preserved when possible, and there
+  // will be no schema type-id reassignment for adding new types.
+  optional bool enable_schema_type_id_optimization = 45;
+
+  // The number of shards to use for the embedding index. 1 means no sharding.
+  optional uint32 embedding_index_num_shards = 46 [default = 1];
+
+  // Whether to enable optimize improvements.
+  optional bool enable_optimize_improvements = 47;
+
+  // The time threshold for an expired document to be purged.
+  // - Since we schedule a task to purge expired documents according to the next
+  //   expiration time of the documents, it is possible that some documents
+  //   expire within a small time window and the task executes too frequently.
+  // - Therefore, we use this flag to purge more documents that also expire in a
+  //   short period of time after the current time.
+  //
+  // For example, if the value is 1000 ms and the current time is 10000 ms:
+  // - All documents that are expired before 10000 ms will be purged, since they
+  //   are already expired.
+  // - Additionally, we will also purge documents that expire in the next 1000
+  //   ms, i.e. (10000, 11000] ms.
+  //
+  // The default value is 1 minute (1 * 60 * 1000 ms).
+  optional int64 expired_document_purge_threshold_ms = 48 [default = 60000];
+
+  // Whether to allow a document to reference a join parent (by its qualified
+  // id) that does not yet exist. When enabled, the join index will handle cases
+  // where a child document is indexed before its parent. The join relationship
+  // will be established once the parent document is indexed. If disabled, the
+  // join relationship will be lost if the child is indexed before the parent.
+  optional bool enable_non_existent_qualified_id_join = 49;
+
+  // Whether to skip the schema type equality check during SetSchema. We
+  // serialize the schema proto to strings when doing this check, which is slow.
+  //
+  // This will be set to true in AppSearch.
+  // - AppSearch already checks if schema types are unchanged for a new
+  //   SetSchema request, and skips the Icing interaction entirely if that is
+  //   the case.
+  // - Therefore, the Icing-side equality check in SetSchema is redundant and
+  //   can be skipped if caller is AppSearch.
+  optional bool enable_skip_set_schema_type_equality_check = 50;
+
+  // Whether to enable the background task scheduler. This value should be true
+  // for all Icing lib direct clients (e.g. iOS), and false for AppSearch since
+  // task scheduling must be handled at the AppSearch layer.
+  //
+  // If set to true, the background task scheduler will be enabled. Otherwise,
+  // the background task scheduler will be disabled.
+  optional bool enable_background_task_scheduler =
+      51;  // TODO(b/384947619): set default to true
+
+  // Whether to enable a query optimization that will rewrite embedding query
+  // iterators that are being AND'ed with other iterators such that those other
+  // iterators can be pushed down into the embedding iterator as a delegate.
+  // This allows us to avoid reading and scoring embeddings for documents that
+  // don't match the other requirements of the query.
+  optional bool enable_embed_query_optimization = 52;
+
+  // Whether to enable deduping schema types's property definitions when
+  // storing SchemaTypeConfigs in the schema store.
+  //
+  // Once enabled, type configs with the same property definitions will be
+  // deduped, and only one of the types will be stored with full property
+  // definitions.
+  optional bool enable_schema_definition_deduping = 53;
+
+  reserved 2, 20, 25;
 }
 
 // Result of a call to IcingSearchEngine.Initialize
-// Next tag: 3
+// Next tag: 6
 message InitializeResultProto {
   // Status code can be one of:
   //   OK
@@ -386,6 +515,17 @@
   // logging.proto for details.
   optional InitializeStatsProto initialize_stats = 2;
 
+  // The latency of the GetVm call.
+  optional int32 get_vm_latency_ms = 3;
+
+  // Persist type needed after initialization. If UNKNOWN, then everything is
+  // clean and no need to call persistToDisk.
+  optional PersistType.Code needs_persist_type = 4;
+
+  // The start time of the binder transaction latency. This is used to measure
+  // the latency of passing data from the VM via the binder.
+  optional int64 vm_binder_transaction_latency_start_time_ms = 5;
+
   // TODO(b/147699081): Add a field to indicate lost_schema and lost_documents.
   // go/icing-library-apis.
 }

diff --git a/proto/icing/proto/logging.proto b/proto/icing/proto/logging.proto
index 2e620cd..783111d 100644
--- a/proto/icing/proto/logging.proto
+++ b/proto/icing/proto/logging.proto

@@ -24,7 +24,7 @@
 option objc_class_prefix = "ICNG";
 
 // Stats of the top-level function IcingSearchEngine::Initialize().
-// Next tag: 17
+// Next tag: 21
 message InitializeStatsProto {
   // Overall time used for the function call.
   optional int32 latency_ms = 1;
@@ -150,6 +150,79 @@
 
   // Number of documents that failed to be reindexed during index restoration.
   optional int32 num_failed_reindexed_documents = 16;
+
+  // A number indicating which stage of initialization failed.
+  message FailureStage {
+    enum Code {
+      NONE = 0;
+
+      OPTIONS_VALIDATION = 1;
+
+      BASE_DIRECTORY_CREATION = 2;
+
+      INIT_MARKER_FILE = 3;
+
+      READ_VERSION_FILE = 4;
+
+      MIGRATE_SCHEMA = 5;
+
+      DISCARD_DERIVED_FILES = 6;
+
+      WRITE_VERSION_FILE = 7;
+
+      SCHEMA_STORE_DIRECTORY_CREATION = 8;
+
+      SCHEMA_STORE_INSTANTIATION = 9;
+
+      LANGUAGE_SEGMENTER_CREATION = 10;
+
+      NORMALIZER_CREATION = 11;
+
+      // Discard any components, including delete the directory or call discard
+      // function of an index.
+      DISCARD_COMPONENTS = 12;
+
+      BLOB_STORE_INSTANTIATION = 13;
+
+      DOCUMENT_STORE_DIRECTORY_CREATION = 14;
+
+      DOCUMENT_STORE_INSTANTIATION = 15;
+
+      TERM_INDEX_DIRECTORY = 16;
+
+      TERM_INDEX_INSTANTIATION = 17;
+
+      INTEGER_INDEX_DIRECTORY = 18;
+
+      INTEGER_INDEX_INSTANTIATION = 19;
+
+      QUALIFIED_ID_JOIN_INDEX_DIRECTORY = 20;
+
+      QUALIFIED_ID_JOIN_INDEX_INSTANTIATION = 21;
+
+      EMBEDDING_INDEX_DIRECTORY = 22;
+
+      EMBEDDING_INDEX_INSTANTIATION = 23;
+
+      RESTORE_INDEX = 24;
+    }
+  }
+  optional FailureStage.Code failure_stage = 17;
+
+  // Status of the attempt to create the IcuLanguageSegmenter.
+  // This is only set if the ICU segmenter is enabled. Possible status codes:
+  // - OK: IcuLanguageSegmenter created successfully.
+  // - ABORTED: Failed to create the ICU break iterator.
+  optional StatusProto icu_segmenter_creation_status = 18;
+
+  // Status of the attempt to create the IcuNormalizer.
+  // This is only set if the ICU normalizer is enabled. Possible status codes:
+  // - OK: IcuNormalizer created successfully.
+  // - INTERNAL: Failed to create the ICU Transliterator.
+  optional StatusProto icu_normalizer_creation_status = 19;
+
+  // The expiration timestamp (in milliseconds) of the next expired document.
+  optional int64 next_expiration_timestamp_ms = 20;
 }
 
 // Stats of the top-level function IcingSearchEngine::Put().
@@ -206,7 +279,7 @@
 
 // Stats of the top-level function IcingSearchEngine::Search() and
 // IcingSearchEngine::GetNextPage().
-// Next tag: 28
+// Next tag: 30
 message QueryStatsProto {
   // TODO(b/305098009): deprecate. Use parent_search_stats instead.
   // The UTF-8 length of the query string
@@ -288,7 +361,7 @@
   optional bool is_join_query = 23;
 
   // Stats of the search. Only valid for first page.
-  // Next tag: 16
+  // Next tag: 20
   message SearchStats {
     // The UTF-8 length of the query string
     optional int32 query_length = 1;
@@ -335,6 +408,18 @@
 
     // Time used in QueryVisitor to visit and build (nested) DocHitInfoIterator.
     optional int32 query_processor_query_visitor_latency_ms = 15;
+
+    // Number of unquantized embeddings scored.
+    optional int32 num_unquantized_embeddings_scored = 16;
+
+    // Number of quantized embeddings scored.
+    optional int32 num_quantized_embeddings_scored = 17;
+
+    // Number of shards read from the embedding index.
+    optional int32 num_embedding_shards_read = 18;
+
+    // Number of raw embedding bytes read from the embedding index.
+    optional int64 num_embedding_bytes_read = 19;
   }
 
   // Search stats for parent. Only valid for first page.
@@ -349,6 +434,28 @@
   // Byte size of the unsorted tail of the lite index hit buffer.
   optional int64 lite_index_hit_buffer_unsorted_byte_size = 27;
 
+  message PageTokenType {
+    enum Code {
+      // Default. Usually used when it is the first page.
+      NONE = 0;
+
+      VALID = 1;
+
+      // The current page token is not found in ResultStateManager. This is
+      // usually caused by cache eviction.
+      NOT_FOUND = 2;
+
+      // The current page token is empty (kInvalidNextPageToken).
+      EMPTY = 3;
+    }
+  }
+  // The type of the input page token.
+  optional PageTokenType.Code page_token_type = 28;
+
+  // Number of result states being force-evicted from ResultStateManager due to
+  // budget limit. This doesn't include expired or invalidated states.
+  optional int32 num_result_states_evicted = 29;
+
   reserved 9;
 }
 

diff --git a/proto/icing/proto/optimize.proto b/proto/icing/proto/optimize.proto
index 0ba0a86..12c65dc 100644
--- a/proto/icing/proto/optimize.proto
+++ b/proto/icing/proto/optimize.proto

@@ -16,6 +16,7 @@
 
 package icing.lib;
 
+import "icing/proto/persist.proto";
 import "icing/proto/status.proto";
 
 option java_package = "com.google.android.icing.proto";
@@ -23,7 +24,7 @@
 option objc_class_prefix = "ICNG";
 
 // Result of a call to IcingSearchEngine.Optimize
-// Next tag: 4
+// Next tag: 6
 message OptimizeResultProto {
   // Status code can be one of:
   //   OK
@@ -43,10 +44,17 @@
 
   // TODO(b/147699081): Add a field to indicate lost_schema and lost_documents.
   // go/icing-library-apis.
+
+  // The latency of the GetVm call.
+  optional int32 get_vm_latency_ms = 4;
+
+  // The start time of the binder transaction latency. This is used to measure
+  // the latency of passing data from the VM via the binder.
+  optional int64 vm_binder_transaction_latency_start_time_ms = 5;
 }
 
 // Result of a call to IcingSearchEngine.GetOptimizeInfo
-// Next tag: 5
+// Next tag: 9
 message GetOptimizeInfoResultProto {
   // Status code can be one of:
   //   OK
@@ -66,9 +74,28 @@
 
   // The amount of time since the last optimize ran.
   optional int64 time_since_last_optimize_ms = 4;
+
+  // True if there is no previous optimize info written to disk. This would
+  // indicate that Icing has never run optimize before.
+  optional bool no_previous_optimize_info = 5;
+
+  // The number of active result states in ResultStateManager.
+  //
+  // If this number is not 0, then the caller should avoid running optimize as
+  // much as possible to avoid result state manager reset and interrupted
+  // pagination, unless the last optimize run was a long time ago which exceeds
+  // the threshold.
+  optional int32 num_active_result_states = 6;
+
+  // The latency of the GetVm call.
+  optional int32 get_vm_latency_ms = 7;
+
+  // The start time of the binder transaction latency. This is used to measure
+  // the latency of passing data from the VM via the binder.
+  optional int64 vm_binder_transaction_latency_start_time_ms = 8;
 }
 
-// Next tag: 14
+// Next tag: 16
 message OptimizeStatsProto {
   // Overall time used for the function call.
   optional int32 latency_ms = 1;
@@ -116,4 +143,10 @@
 
   // Number of namespaces deleted.
   optional int32 num_deleted_namespaces = 12;
+
+  // Stats of the PersistToDisk call before optimization.
+  optional PersistToDiskStatsProto before_optimize_persist_stats = 14;
+
+  // Stats of the PersistToDisk call after optimization.
+  optional PersistToDiskStatsProto after_optimize_persist_stats = 15;
 }

diff --git a/proto/icing/proto/persist.proto b/proto/icing/proto/persist.proto
index 802195d..7af2e98 100644
--- a/proto/icing/proto/persist.proto
+++ b/proto/icing/proto/persist.proto

@@ -57,7 +57,7 @@
 }
 
 // Result of a call to IcingSearchEngine.Persist
-// Next tag: 2
+// Next tag: 5
 message PersistToDiskResultProto {
   // Status code can be one of:
   //   OK
@@ -66,4 +66,71 @@
   //
   // See status.proto for more details.
   optional StatusProto status = 1;
+
+  // Stats about the PersistToDisk call.
+  optional PersistToDiskStatsProto persist_stats = 2;
+
+  // The latency of the GetVm call.
+  optional int32 get_vm_latency_ms = 3;
+
+  // The start time of the binder transaction latency. This is used to measure
+  // the latency of passing data from the VM via the binder.
+  optional int64 vm_binder_transaction_latency_start_time_ms = 4;
+}
+
+// Stats of the top-level function IcingSearchEngine::PersistToDisk.
+// Next tag: 14
+message PersistToDiskStatsProto {
+  // The type of persistence guarantee that PersistToDisk should provide.
+  optional PersistType.Code persist_type = 1;
+
+  // Overall time used for the function call.
+  optional int32 latency_ms = 2;
+
+  // Time used to persist the blob store.
+  optional int32 blob_store_persist_latency_ms = 3;
+
+  // Time used to persist the document store.
+  optional int32 document_store_total_persist_latency_ms = 4;
+
+  // Time used to persist the document store components.
+  optional int32 document_store_components_persist_latency_ms = 5;
+
+  // Time used to update the document store checksum.
+  optional int32 document_store_checksum_update_latency_ms = 6;
+
+  // Time used to update the document log checksum.
+  optional int32 document_log_checksum_update_latency_ms = 7;
+
+  // Time used to sync the document log data to disk.
+  optional int32 document_log_data_sync_latency_ms = 8;
+
+  // Time used to persist the schema store. If the persist type is
+  // RECOVERY_PROOF, this will update the checksums of the schema store. If the
+  // persist type is FULL, this will update the checksums internally and
+  // fdatasync/msync.
+  optional int32 schema_store_persist_latency_ms = 9;
+
+  // Time used to persist the index. If the persist type is RECOVERY_PROOF, this
+  // will update the checksums of the index.  If the persist type is FULL, this
+  // will update the checksums internally and fdatasync/msync.
+  optional int32 index_persist_latency_ms = 10;
+
+  // Time used to persist the integer index. If the persist type is
+  // RECOVERY_PROOF, this will update the checksums of the integer index. If the
+  // persist type is FULL, this will update the checksums internally and
+  // fdatasync/msync.
+  optional int32 integer_index_persist_latency_ms = 11;
+
+  // Time used to persist the qualified id join index.  If the persist type is
+  // RECOVERY_PROOF, this will update the checksums of the qualified id join
+  // index. If the persist type is FULL, this will update the checksums
+  // internally and fdatasync/msync.
+  optional int32 qualified_id_join_index_persist_latency_ms = 12;
+
+  // Time used to persist the embedding index. If the persist type is
+  // RECOVERY_PROOF, this will update the checksums of the embedding index. If
+  // the persist type is FULL, this will update the checksums internally and
+  // fdatasync/msync.
+  optional int32 embedding_index_persist_latency_ms = 13;
 }

diff --git a/proto/icing/proto/reset.proto b/proto/icing/proto/reset.proto
index 5e8b9f5..0bc22f3 100644
--- a/proto/icing/proto/reset.proto
+++ b/proto/icing/proto/reset.proto

@@ -20,11 +20,10 @@
 
 option java_package = "com.google.android.icing.proto";
 option java_multiple_files = true;
-
 option objc_class_prefix = "ICNG";
 
 // Result of a call to IcingSearchEngine.Reset
-// Next tag: 2
+// Next tag: 4
 message ResetResultProto {
   // Status code can be one of:
   //   OK
@@ -33,4 +32,11 @@
   //
   // See status.proto for more details.
   optional StatusProto status = 1;
+
+  // The latency of the GetVm call.
+  optional int32 get_vm_latency_ms = 2;
+
+  // The start time of the binder transaction latency. This is used to measure
+  // the latency of passing data from the VM via the binder.
+  optional int64 vm_binder_transaction_latency_start_time_ms = 3;
 }

diff --git a/proto/icing/proto/schema.proto b/proto/icing/proto/schema.proto
index 1f7eb4b..2d3f1b7 100644
--- a/proto/icing/proto/schema.proto
+++ b/proto/icing/proto/schema.proto

@@ -16,6 +16,7 @@
 
 package icing.lib;
 
+import "icing/proto/persist.proto";
 import "icing/proto/status.proto";
 import "icing/proto/term.proto";
 
@@ -34,7 +35,7 @@
 // TODO(cassiewang) Define a sample proto file that can be used by tests and for
 // documentation.
 //
-// Next tag: 9
+// Next tag: 10
 message SchemaTypeConfigProto {
   // REQUIRED: Named type that identifies the structured, logical schema being
   // defined.
@@ -83,6 +84,10 @@
   // TypePropertyMask.
   repeated string parent_types = 6;
 
+  // List of all account properties that need to be managed. All accounts under
+  // this property should be existed in the account manager.
+  repeated string account_properties = 9;
+
   reserved 2, 3;
 }
 
@@ -455,7 +460,7 @@
 }
 
 // Result of a call to IcingSearchEngine.SetSchema
-// Next tag: 9
+// Next tag: 19
 message SetSchemaResultProto {
   // Status code can be one of:
   //   OK
@@ -493,7 +498,9 @@
   repeated string index_incompatible_changed_schema_types = 6;
 
   // Overall time used for the function call.
-  optional int32 latency_ms = 7;
+  //
+  // DEPRECATED: Use SetSchemaStatsProto.overall_latency_ms instead.
+  optional int32 latency_ms = 7 [deprecated = true];
 
   // Schema types that were changed in a way that was backwards compatible, but
   // invalidated the joinable cache.
@@ -502,10 +509,77 @@
   // but changed to joinable in the new definition. In this case, this property
   // will be considered join incompatible when setting new schema.
   repeated string join_incompatible_changed_schema_types = 8;
+
+  // Number of documents that were deleted due to the new schema with
+  // ignore_errors_and_delete_documents enabled.
+  optional int32 deleted_document_count = 9;
+
+  // The latency of the GetVm call.
+  optional int32 get_vm_latency_ms = 10;
+
+  // Schema types that were changed in a way that was backwards compatible, but
+  // invalidated the scorable cache.
+  repeated string scorable_property_incompatible_changed_schema_types = 11;
+
+  // Stats and latencies for the SetSchema call.
+  optional SetSchemaStatsProto set_schema_stats = 12;
+
+  // Whether the term index (lite and main indices) was restored as a result of
+  // the set schema call.
+  optional bool has_term_index_restored = 13;
+
+  // Whether the integer index was restored as a result of the set schema call.
+  optional bool has_integer_index_restored = 14;
+
+  // Whether the embedding index was restored as a result of the set schema
+  // call.
+  optional bool has_embedding_index_restored = 15;
+
+  // Whether the qualified id join index was restored as a result of the set
+  // schema call.
+  optional bool has_qualified_id_join_index_restored = 16;
+
+  // Persist type needed after setting the new schema. If UNKNOWN, then
+  // everything is clean and no need to call persistToDisk.
+  optional PersistType.Code needs_persist_type = 17;
+
+  // The start time of the binder transaction latency. This is used to measure
+  // the latency of passing data from the VM via the binder.
+  optional int64 vm_binder_transaction_latency_start_time_ms = 18;
+}
+
+// Stats and latencies for IcingSearchEngine::SetSchema.
+// Next tag: 7
+message SetSchemaStatsProto {
+  // Overall time used for the function call. This is the sum of the latencies
+  // recorded in the other fields (schema_store_set_schema_latency,
+  // document_store_update_schema_latency,
+  // document_store_optimized_update_schema_latency,
+  // index_restoration_latency,
+  // scorable_property_cache_regeneration_latency)
+  optional int32 overall_latency_ms = 1;
+
+  // Latency for setting the schema in the schema store.
+  optional int32 schema_store_set_schema_latency_ms = 2;
+
+  // Latency for updating the document store's derived data using
+  // DocumentStore::UpdateSchemaStore.
+  optional int32 document_store_update_schema_latency_ms = 3;
+
+  // Latency for updating the document store's derived data using
+  // DocumentStore::OptimizedUpdateSchemaStore.
+  optional int32 document_store_optimized_update_schema_latency_ms = 4;
+
+  // Latency for rebuilding the index following an index-incompatible schema
+  // change.
+  optional int32 index_restoration_latency_ms = 5;
+
+  // Latency for regenerating the scorable property cache.
+  optional int32 scorable_property_cache_regeneration_latency_ms = 6;
 }
 
 // Result of a call to IcingSearchEngine.GetSchema
-// Next tag: 3
+// Next tag: 4
 message GetSchemaResultProto {
   // Status code can be one of:
   //   OK
@@ -522,10 +596,13 @@
   // Copy of the Schema proto. Modifying this does not affect the Schema that
   // IcingSearchEngine holds.
   optional SchemaProto schema = 2;
+
+  // The latency of the GetVm call.
+  optional int32 get_vm_latency_ms = 3;
 }
 
 // Result of a call to IcingSearchEngine.GetSchemaType
-// Next tag: 3
+// Next tag: 5
 message GetSchemaTypeResultProto {
   // Status code can be one of:
   //   OK
@@ -543,4 +620,11 @@
   // Modifying this does not affect the SchemaTypeConfig that IcingSearchEngine
   // holds.
   optional SchemaTypeConfigProto schema_type_config = 2;
+
+  // The latency of the GetVm call.
+  optional int32 get_vm_latency_ms = 3;
+
+  // The start time of the binder transaction latency. This is used to measure
+  // the latency of passing data from the VM via the binder.
+  optional int64 vm_binder_transaction_latency_start_time_ms = 4;
 }

diff --git a/proto/icing/proto/search.proto b/proto/icing/proto/search.proto
index f0f6c5c..d6f5665 100644
--- a/proto/icing/proto/search.proto
+++ b/proto/icing/proto/search.proto

@@ -382,7 +382,7 @@
 }
 
 // Icing lib-supplied results from a search results.
-// Next tag: 6
+// Next tag: 9
 message SearchResultProto {
   // Status code can be one of:
   //   OK
@@ -442,6 +442,21 @@
 
   // Stats for query execution performance.
   optional QueryStatsProto query_stats = 5;
+
+  // This field is used only in the GetNextPage API.
+  //
+  // Whether GetNextPage returning an empty result is due to page token not
+  // found (mostly caused by cache eviction). If the client needs the entire set
+  // of result documents (as their ground truth), then they should restart the
+  // query via Search API when this field is true.
+  optional bool page_token_not_found = 6;
+
+  // The latency of the GetVm call.
+  optional int32 get_vm_latency_ms = 7;
+
+  // The start time of the binder transaction latency. This is used to measure
+  // the latency of passing data from the VM via the binder.
+  optional int64 vm_binder_transaction_latency_start_time_ms = 8;
 }
 
 // Next tag: 3
@@ -459,12 +474,23 @@
   repeated string paths = 2;
 }
 
-// Next tag: 2
+// Next tag: 5
+// TODO(b/394875109): Rename it to GetRequestProto to be consistent with the
+// name in AppSearch.
 message GetResultSpecProto {
+  optional string namespace_requested = 2;
+  repeated string ids = 3;
+
   // How to specify a subset of properties to retrieve. If no type property mask
   // has been specified for a schema type, then *all* properties of that schema
   // type will be retrieved.
   repeated TypePropertyMask type_property_masks = 1;
+
+  // The maximum number of accumulated bytes for the documents to return in the
+  // result. This limit is to prevent the result from being too large, and we
+  // can't send it over VM boundary, which has transaction limit 600KB.
+  optional int32 num_total_document_bytes_to_return = 4
+      [default = 2147483647];  // INT_MAX
 }
 
 // Next tag: 12
@@ -541,7 +567,7 @@
   repeated string document_uris = 2;
 }
 
-// Next tag: 3
+// Next tag: 5
 message SuggestionResponse {
   message Suggestion {
     // The suggested query string for client to search for.
@@ -557,6 +583,13 @@
   optional StatusProto status = 1;
 
   repeated Suggestion suggestions = 2;
+
+  // The latency of the GetVm call.
+  optional int32 get_vm_latency_ms = 3;
+
+  // The start time of the binder transaction latency. This is used to measure
+  // the latency of passing data from the VM via the binder.
+  optional int64 vm_binder_transaction_latency_start_time_ms = 4;
 }
 
 // Specification for a left outer join.
@@ -627,3 +660,23 @@
   }
   optional AggregationScoringStrategy.Code aggregation_scoring_strategy = 5;
 }
+
+// Next tag: 3
+message GetNextPageRequestProto {
+  // REQUIRED: The next page token to use for retrieving the next page of
+  // results.
+  //
+  // This is a token returned from a previous Search or GetNextPage
+  // call. The retrieved SearchResultProto will be empty if this token is
+  // invalid (value of 0, or previously invalidated by Icing)
+  optional int64 next_page_token = 1;
+
+  // The max number of results to return from the page.
+  //
+  // This is only an upper limit on the number of results to fetch from the
+  // page. Fewer results may be returned if there are not enough results
+  // remaining in the page, or if other limits (e.g. page size limit) are
+  // reached.
+  optional int32 max_results_to_retrieve_from_page = 2
+      [default = 2147483647];  // INT_MAX
+}

diff --git a/proto/icing/proto/status.proto b/proto/icing/proto/status.proto
index 06ec6c4..4b97e91 100644
--- a/proto/icing/proto/status.proto
+++ b/proto/icing/proto/status.proto

@@ -18,13 +18,12 @@
 
 option java_package = "com.google.android.icing.proto";
 option java_multiple_files = true;
-
 option objc_class_prefix = "ICNG";
 
 // Canonical status to indicate the results of API calls.
 // Next tag: 3
 message StatusProto {
-  // Next tag: 10
+  // Next tag: 11
   enum Code {
     // A default for all other use-cases. Should never be used in practice. This
     // may happen if there are backwards-compatibility issues.
@@ -71,6 +70,9 @@
     // same name within a type.
     ALREADY_EXISTS = 9;
 
+    // The system is unavailable to process the requested operation.
+    UNAVAILABLE = 10;
+
     // Any future status codes.
   }
   optional Code code = 1;

diff --git a/proto/icing/proto/storage.proto b/proto/icing/proto/storage.proto
index a9eb84a..6b5a5cd 100644
--- a/proto/icing/proto/storage.proto
+++ b/proto/icing/proto/storage.proto

@@ -200,7 +200,7 @@
   repeated NamespaceBlobStorageInfoProto namespace_blob_storage_info = 5;
 }
 
-// Next tag: 3
+// Next tag: 5
 message StorageInfoResultProto {
   // Status code can be one of:
   //   OK
@@ -211,4 +211,11 @@
 
   // Storage information of Icing.
   optional StorageInfoProto storage_info = 2;
+
+  // The latency of the GetVm call.
+  optional int32 get_vm_latency_ms = 3;
+
+  // The start time of the binder transaction latency. This is used to measure
+  // the latency of passing data from the VM via the binder.
+  optional int64 vm_binder_transaction_latency_start_time_ms = 4;
 }

diff --git a/proto/icing/proto/usage.proto b/proto/icing/proto/usage.proto
index eaa2671..3e0da95 100644
--- a/proto/icing/proto/usage.proto
+++ b/proto/icing/proto/usage.proto

@@ -57,7 +57,7 @@
 }
 
 // Result of a call to IcingSearchEngine.ReportUsage
-// Next tag: 2
+// Next tag: 4
 message ReportUsageResultProto {
   // Status code can be one of:
   //   OK
@@ -66,4 +66,11 @@
   //
   // See status.proto for more details.
   optional StatusProto status = 1;
+
+  // The latency of the GetVm call.
+  optional int32 get_vm_latency_ms = 2;
+
+  // The start time of the binder transaction latency. This is used to measure
+  // the latency of passing data from the VM via the binder.
+  optional int64 vm_binder_transaction_latency_start_time_ms = 3;
 }

diff --git a/synced_AOSP_CL_number.txt b/synced_AOSP_CL_number.txt
index ba3548f..abe0202 100644
--- a/synced_AOSP_CL_number.txt
+++ b/synced_AOSP_CL_number.txt

@@ -1 +1 @@
-set(synced_AOSP_CL_number=733916508)
+set(synced_AOSP_CL_number=853041928)