Update Icing from upstream.

Descriptions:
========================================================================
[fixit] Remove incorrect reserve usage for STL data structure
========================================================================
Move LiteIndex::SortHits() from the query path to the indexing path
========================================================================

Bug: 286418010
Bug: 297549761
Bug: 248295995
Change-Id: I46168178e0227451b0f7dbf394bb6cc5db94e1e8

26 files changed, 1661 insertions, 208 deletions

diff --git a/icing/icing-search-engine.cc b/icing/icing-search-engine.cc
index 467c943..6680dae 100644
--- a/icing/icing-search-engine.cc
+++ b/icing/icing-search-engine.cc
@@ -654,7 +654,9 @@ libtextclassifier3::Status IcingSearchEngine::InitializeMembers(
     // We're going to need to build the index from scratch. So just delete its
     // directory now.
     // Discard index directory and instantiate a new one.
-    Index::Options index_options(index_dir, options_.index_merge_size());
+    Index::Options index_options(index_dir, options_.index_merge_size(),
+                                 options_.lite_index_sort_at_indexing(),
+                                 options_.lite_index_sort_size());
     if (!filesystem_->DeleteDirectoryRecursively(index_dir.c_str()) ||
         !filesystem_->CreateDirectoryRecursively(index_dir.c_str())) {
       return absl_ports::InternalError(
@@ -798,7 +800,9 @@ libtextclassifier3::Status IcingSearchEngine::InitializeIndex(
     return absl_ports::InternalError(
         absl_ports::StrCat("Could not create directory: ", index_dir));
   }
-  Index::Options index_options(index_dir, options_.index_merge_size());
+  Index::Options index_options(index_dir, options_.index_merge_size(),
+                               options_.lite_index_sort_at_indexing(),
+                               options_.lite_index_sort_size());
 
   // Term index
   InitializeStatsProto::RecoveryCause index_recovery_cause;
diff --git a/icing/icing-search-engine_initialization_test.cc b/icing/icing-search-engine_initialization_test.cc
index d6c8de8..b4853b4 100644
--- a/icing/icing-search-engine_initialization_test.cc
+++ b/icing/icing-search-engine_initialization_test.cc
@@ -1853,7 +1853,9 @@ TEST_F(IcingSearchEngineInitializationTest, RestoreIndexLoseTermIndex) {
     ICING_ASSERT_OK_AND_ASSIGN(
         std::unique_ptr<Index> index,
         Index::Create(Index::Options(GetIndexDir(),
-                                     /*index_merge_size=*/100),
+                                     /*index_merge_size=*/100,
+                                     /*lite_index_sort_at_indexing=*/true,
+                                     /*lite_index_sort_size=*/50),
                       filesystem(), icing_filesystem()));
     ICING_ASSERT_OK(index->PersistToDisk());
   }
@@ -2463,7 +2465,9 @@ TEST_F(IcingSearchEngineInitializationTest,
         std::unique_ptr<Index> index,
         Index::Create(
             Index::Options(GetIndexDir(),
-                           /*index_merge_size=*/message.ByteSizeLong()),
+                           /*index_merge_size=*/message.ByteSizeLong(),
+                           /*lite_index_sort_at_indexing=*/true,
+                           /*lite_index_sort_size=*/8),
             filesystem(), icing_filesystem()));
     DocumentId original_last_added_doc_id = index->last_added_document_id();
     index->set_last_added_document_id(original_last_added_doc_id + 1);
@@ -2595,7 +2599,9 @@ TEST_F(IcingSearchEngineInitializationTest,
         std::unique_ptr<Index> index,
         Index::Create(
             Index::Options(GetIndexDir(),
-                           /*index_merge_size=*/message.ByteSizeLong()),
+                           /*index_merge_size=*/message.ByteSizeLong(),
+                           /*lite_index_sort_at_indexing=*/true,
+                           /*lite_index_sort_size=*/8),
             filesystem(), icing_filesystem()));
     ICING_ASSERT_OK_AND_ASSIGN(
         std::unique_ptr<DocHitInfoIterator> doc_hit_info_iter,
@@ -2707,7 +2713,9 @@ TEST_F(IcingSearchEngineInitializationTest,
         std::unique_ptr<Index> index,
         Index::Create(
             Index::Options(GetIndexDir(),
-                           /*index_merge_size=*/message.ByteSizeLong()),
+                           /*index_merge_size=*/message.ByteSizeLong(),
+                           /*lite_index_sort_at_indexing=*/true,
+                           /*lite_index_sort_size=*/8),
             filesystem(), icing_filesystem()));
     DocumentId original_last_added_doc_id = index->last_added_document_id();
     index->set_last_added_document_id(original_last_added_doc_id + 1);
@@ -2844,7 +2852,9 @@ TEST_F(IcingSearchEngineInitializationTest,
         std::unique_ptr<Index> index,
         Index::Create(
             Index::Options(GetIndexDir(),
-                           /*index_merge_size=*/message.ByteSizeLong()),
+                           /*index_merge_size=*/message.ByteSizeLong(),
+                           /*lite_index_sort_at_indexing=*/true,
+                           /*lite_index_sort_size=*/8),
             filesystem(), icing_filesystem()));
     ICING_ASSERT_OK_AND_ASSIGN(
         std::unique_ptr<DocHitInfoIterator> doc_hit_info_iter,
@@ -2904,7 +2914,9 @@ TEST_F(IcingSearchEngineInitializationTest,
         Index::Create(
             // index merge size is not important here because we will manually
             // invoke merge below.
-            Index::Options(GetIndexDir(), /*index_merge_size=*/100),
+            Index::Options(GetIndexDir(), /*index_merge_size=*/100,
+                           /*lite_index_sort_at_indexing=*/true,
+                           /*lite_index_sort_size=*/50),
             filesystem(), icing_filesystem()));
     // Add hits for document 0 and merge.
     ASSERT_THAT(index->last_added_document_id(), kInvalidDocumentId);
@@ -2980,7 +2992,9 @@ TEST_F(IcingSearchEngineInitializationTest,
   {
     ICING_ASSERT_OK_AND_ASSIGN(
         std::unique_ptr<Index> index,
-        Index::Create(Index::Options(GetIndexDir(), /*index_merge_size=*/100),
+        Index::Create(Index::Options(GetIndexDir(), /*index_merge_size=*/100,
+                                     /*lite_index_sort_at_indexing=*/true,
+                                     /*lite_index_sort_size=*/50),
                       filesystem(), icing_filesystem()));
     ICING_ASSERT_OK_AND_ASSIGN(
         std::unique_ptr<DocHitInfoIterator> doc_hit_info_iter,
@@ -3096,7 +3110,9 @@ TEST_F(IcingSearchEngineInitializationTest,
         std::unique_ptr<Index> index,
         Index::Create(
             Index::Options(GetIndexDir(),
-                           /*index_merge_size=*/message.ByteSizeLong()),
+                           /*index_merge_size=*/message.ByteSizeLong(),
+                           /*lite_index_sort_at_indexing=*/true,
+                           /*lite_index_sort_size=*/8),
             filesystem(), icing_filesystem()));
     // Add hits for document 4 and merge.
     DocumentId original_last_added_doc_id = index->last_added_document_id();
@@ -3239,7 +3255,9 @@ TEST_F(IcingSearchEngineInitializationTest,
   {
     ICING_ASSERT_OK_AND_ASSIGN(
         std::unique_ptr<Index> index,
-        Index::Create(Index::Options(GetIndexDir(), /*index_merge_size=*/100),
+        Index::Create(Index::Options(GetIndexDir(), /*index_merge_size=*/100,
+                                     /*lite_index_sort_at_indexing=*/true,
+                                     /*lite_index_sort_size=*/50),
                       filesystem(), icing_filesystem()));
     ICING_ASSERT_OK_AND_ASSIGN(
         std::unique_ptr<DocHitInfoIterator> doc_hit_info_iter,
@@ -4412,7 +4430,9 @@ TEST_F(IcingSearchEngineInitializationTest,
     ICING_ASSERT_OK_AND_ASSIGN(
         std::unique_ptr<Index> index,
         Index::Create(Index::Options(GetIndexDir(),
-                                     /*index_merge_size=*/100),
+                                     /*index_merge_size=*/100,
+                                     /*lite_index_sort_at_indexing=*/true,
+                                     /*lite_index_sort_size=*/50),
                       filesystem(), icing_filesystem()));
     ICING_ASSERT_OK(index->PersistToDisk());
   }
@@ -5241,7 +5261,9 @@ TEST_P(IcingSearchEngineInitializationVersionChangeTest,
     ICING_ASSERT_OK_AND_ASSIGN(DocumentId doc_id, document_store->Put(message));
 
     // Index doc_id with incorrect data
-    Index::Options options(GetIndexDir(), /*index_merge_size=*/1024 * 1024);
+    Index::Options options(GetIndexDir(), /*index_merge_size=*/1024 * 1024,
+                           /*lite_index_sort_at_indexing=*/true,
+                           /*lite_index_sort_size=*/1024 * 8);
     ICING_ASSERT_OK_AND_ASSIGN(
         std::unique_ptr<Index> index,
         Index::Create(options, filesystem(), icing_filesystem()));
diff --git a/icing/index/index-processor_benchmark.cc b/icing/index/index-processor_benchmark.cc
index 9fee925..8766f0b 100644
--- a/icing/index/index-processor_benchmark.cc
+++ b/icing/index/index-processor_benchmark.cc
@@ -150,7 +150,9 @@ DocumentProto CreateDocumentWithHiragana(int content_length) {
 std::unique_ptr<Index> CreateIndex(const IcingFilesystem& icing_filesystem,
                                    const Filesystem& filesystem,
                                    const std::string& index_dir) {
-  Index::Options options(index_dir, /*index_merge_size=*/1024 * 1024 * 10);
+  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();
 }
 
diff --git a/icing/index/index-processor_test.cc b/icing/index/index-processor_test.cc
index c2eaf62..ba4ece3 100644
--- a/icing/index/index-processor_test.cc
+++ b/icing/index/index-processor_test.cc
@@ -167,7 +167,9 @@ class IndexProcessorTest : public Test {
     schema_store_dir_ = base_dir_ + "/schema_store";
     doc_store_dir_ = base_dir_ + "/doc_store";
 
-    Index::Options options(index_dir_, /*index_merge_size=*/1024 * 1024);
+    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_));
 
@@ -970,7 +972,9 @@ TEST_F(IndexProcessorTest, IndexingDocAutomaticMerge) {
       TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
                                 document));
   Index::Options options(index_dir_,
-                         /*index_merge_size=*/document.ByteSizeLong() * 100);
+                         /*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_));
 
@@ -1033,7 +1037,9 @@ TEST_F(IndexProcessorTest, IndexingDocMergeFailureResets) {
   // 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.
   Index::Options options(index_dir_,
-                         /*index_merge_size=*/document.ByteSizeLong());
+                         /*index_merge_size=*/document.ByteSizeLong(),
+                         /*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()));
diff --git a/icing/index/index.cc b/icing/index/index.cc
index 19edbb6..31dcc7e 100644
--- a/icing/index/index.cc
+++ b/icing/index/index.cc
@@ -14,31 +14,38 @@
 
 #include "icing/index/index.h"
 
+#include <algorithm>
+#include <cstddef>
 #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/file/filesystem.h"
 #include "icing/index/hit/hit.h"
 #include "icing/index/iterator/doc-hit-info-iterator-or.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
 #include "icing/index/lite/doc-hit-info-iterator-term-lite.h"
 #include "icing/index/lite/lite-index.h"
 #include "icing/index/main/doc-hit-info-iterator-term-main.h"
+#include "icing/index/main/main-index.h"
 #include "icing/index/term-id-codec.h"
-#include "icing/index/term-property-id.h"
+#include "icing/index/term-metadata.h"
 #include "icing/legacy/core/icing-string-util.h"
 #include "icing/legacy/index/icing-dynamic-trie.h"
 #include "icing/legacy/index/icing-filesystem.h"
+#include "icing/proto/scoring.pb.h"
 #include "icing/proto/storage.pb.h"
 #include "icing/proto/term.pb.h"
 #include "icing/schema/section.h"
 #include "icing/scoring/ranker.h"
 #include "icing/store/document-id.h"
+#include "icing/store/suggestion-result-checker.h"
 #include "icing/util/logging.h"
 #include "icing/util/status-macros.h"
 
@@ -59,7 +66,9 @@ libtextclassifier3::StatusOr<LiteIndex::Options> CreateLiteIndexOptions(
         options.index_merge_size));
   }
   return LiteIndex::Options(options.base_dir + "/idx/lite.",
-                            options.index_merge_size);
+                            options.index_merge_size,
+                            options.lite_index_sort_at_indexing,
+                            options.lite_index_sort_size);
 }
 
 std::string MakeMainIndexFilepath(const std::string& base_dir) {
@@ -151,9 +160,17 @@ libtextclassifier3::StatusOr<std::unique_ptr<Index>> Index::Create(
           IcingDynamicTrie::max_value_index(GetMainLexiconOptions()),
           IcingDynamicTrie::max_value_index(
               lite_index_options.lexicon_options)));
+
   ICING_ASSIGN_OR_RETURN(
       std::unique_ptr<LiteIndex> lite_index,
       LiteIndex::Create(lite_index_options, icing_filesystem));
+  // Sort the lite index if we've enabled sorting the HitBuffer at indexing
+  // time, and there's an unsorted tail exceeding the threshold.
+  if (options.lite_index_sort_at_indexing &&
+      lite_index->HasUnsortedHitsExceedingSortThreshold()) {
+    lite_index->SortHits();
+  }
+
   ICING_ASSIGN_OR_RETURN(
       std::unique_ptr<MainIndex> main_index,
       MainIndex::Create(MakeMainIndexFilepath(options.base_dir), filesystem,
diff --git a/icing/index/index.h b/icing/index/index.h
index c170278..32ea97b 100644
--- a/icing/index/index.h
+++ b/icing/index/index.h
@@ -18,8 +18,9 @@
 #include <cstdint>
 #include <memory>
 #include <string>
-#include <unordered_set>
+#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"
@@ -27,6 +28,7 @@
 #include "icing/index/hit/hit.h"
 #include "icing/index/iterator/doc-hit-info-iterator.h"
 #include "icing/index/lite/lite-index.h"
+#include "icing/index/lite/term-id-hit-pair.h"
 #include "icing/index/main/main-index-merger.h"
 #include "icing/index/main/main-index.h"
 #include "icing/index/term-id-codec.h"
@@ -40,7 +42,7 @@
 #include "icing/store/document-id.h"
 #include "icing/store/namespace-id.h"
 #include "icing/store/suggestion-result-checker.h"
-#include "icing/util/crc32.h"
+#include "icing/util/status-macros.h"
 
 namespace icing {
 namespace lib {
@@ -68,11 +70,18 @@ namespace lib {
 class Index {
  public:
   struct Options {
-    explicit Options(const std::string& base_dir, uint32_t index_merge_size)
-        : base_dir(base_dir), index_merge_size(index_merge_size) {}
+    explicit Options(const std::string& base_dir, uint32_t index_merge_size,
+                     bool lite_index_sort_at_indexing,
+                     uint32_t lite_index_sort_size)
+        : base_dir(base_dir),
+          index_merge_size(index_merge_size),
+          lite_index_sort_at_indexing(lite_index_sort_at_indexing),
+          lite_index_sort_size(lite_index_sort_size) {}
 
     std::string base_dir;
     int32_t index_merge_size;
+    bool lite_index_sort_at_indexing;
+    int32_t lite_index_sort_size;
   };
 
   // Creates an instance of Index in the directory pointed by file_dir.
@@ -279,6 +288,19 @@ class Index {
     return lite_index_->Reset();
   }
 
+  // Whether the LiteIndex HitBuffer requires sorting. This is only true if
+  // Icing has enabled sorting during indexing time, and the HitBuffer's
+  // unsorted tail has exceeded the lite_index_sort_size.
+  bool LiteIndexNeedSort() const {
+    return options_.lite_index_sort_at_indexing &&
+           lite_index_->HasUnsortedHitsExceedingSortThreshold();
+  }
+
+  // Sorts the LiteIndex HitBuffer.
+  void SortLiteIndex() {
+    lite_index_->SortHits();
+  }
+
   // Reduces internal file sizes by reclaiming space of deleted documents.
   // new_last_added_document_id will be used to update the last added document
   // id in the lite index.
diff --git a/icing/index/index_test.cc b/icing/index/index_test.cc
index d563bcb..b823535 100644
--- a/icing/index/index_test.cc
+++ b/icing/index/index_test.cc
@@ -58,6 +58,7 @@ using ::testing::Eq;
 using ::testing::Ge;
 using ::testing::Gt;
 using ::testing::IsEmpty;
+using ::testing::IsFalse;
 using ::testing::IsTrue;
 using ::testing::Ne;
 using ::testing::NiceMock;
@@ -75,7 +76,9 @@ class IndexTest : public Test {
  protected:
   void SetUp() override {
     index_dir_ = GetTestTempDir() + "/index_test/";
-    Index::Options options(index_dir_, /*index_merge_size=*/1024 * 1024);
+    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_));
   }
@@ -146,7 +149,9 @@ MATCHER_P2(EqualsTermMetadata, content, hit_count, "") {
 }
 
 TEST_F(IndexTest, CreationWithNullPointerShouldFail) {
-  Index::Options options(index_dir_, /*index_merge_size=*/1024 * 1024);
+  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));
@@ -192,6 +197,36 @@ TEST_F(IndexTest, EmptyIndexAfterMerge) {
               StatusIs(libtextclassifier3::StatusCode::RESOURCE_EXHAUSTED));
 }
 
+TEST_F(IndexTest, CreationWithLiteIndexSortAtIndexingEnabledShouldSort) {
+  // Make the index with lite_index_sort_at_indexing=false and a very small sort
+  // threshold.
+  Index::Options options(index_dir_, /*index_merge_size=*/1024,
+                         /*lite_index_sort_at_indexing=*/false,
+                         /*lite_index_sort_size=*/16);
+  ICING_ASSERT_OK_AND_ASSIGN(
+      index_, Index::Create(options, &filesystem_, &icing_filesystem_));
+
+  Index::Editor edit = index_->Edit(
+      kDocumentId0, kSectionId2, TermMatchType::EXACT_ONLY, /*namespace_id=*/0);
+  ASSERT_THAT(edit.BufferTerm("foo"), IsOk());
+  ASSERT_THAT(edit.BufferTerm("bar"), IsOk());
+  ASSERT_THAT(edit.BufferTerm("baz"), IsOk());
+  ASSERT_THAT(edit.IndexAllBufferedTerms(), IsOk());
+
+  // Persist and recreate the index with lite_index_sort_at_indexing=true
+  ASSERT_THAT(index_->PersistToDisk(), IsOk());
+  options = Index::Options(index_dir_, /*index_merge_size=*/1024,
+                           /*lite_index_sort_at_indexing=*/true,
+                           /*lite_index_sort_size=*/16);
+  ICING_ASSERT_OK_AND_ASSIGN(
+      index_, Index::Create(options, &filesystem_, &icing_filesystem_));
+
+  // Check that the index is sorted after recreating with
+  // lite_index_sort_at_indexing, with the unsorted HitBuffer exceeding the sort
+  // threshold.
+  EXPECT_THAT(index_->LiteIndexNeedSort(), IsFalse());
+}
+
 TEST_F(IndexTest, AdvancePastEnd) {
   Index::Editor edit = index_->Edit(
       kDocumentId0, kSectionId2, TermMatchType::EXACT_ONLY, /*namespace_id=*/0);
@@ -967,7 +1002,9 @@ TEST_F(IndexTest, NonAsciiTermsAfterMerge) {
 
 TEST_F(IndexTest, FullIndex) {
   // Make a smaller index so that it's easier to fill up.
-  Index::Options options(index_dir_, /*index_merge_size=*/1024);
+  Index::Options options(index_dir_, /*index_merge_size=*/1024,
+                         /*lite_index_sort_at_indexing=*/true,
+                         /*lite_index_sort_size=*/64);
   ICING_ASSERT_OK_AND_ASSIGN(
       index_, Index::Create(options, &filesystem_, &icing_filesystem_));
 
@@ -1035,7 +1072,9 @@ TEST_F(IndexTest, FullIndex) {
 
 TEST_F(IndexTest, FullIndexMerge) {
   // Make a smaller index so that it's easier to fill up.
-  Index::Options options(index_dir_, /*index_merge_size=*/1024);
+  Index::Options options(index_dir_, /*index_merge_size=*/1024,
+                         /*lite_index_sort_at_indexing=*/true,
+                         /*lite_index_sort_size=*/64);
   ICING_ASSERT_OK_AND_ASSIGN(
       index_, Index::Create(options, &filesystem_, &icing_filesystem_));
 
@@ -1368,7 +1407,9 @@ TEST_F(IndexTest, IndexCreateIOFailure) {
   NiceMock<IcingMockFilesystem> mock_icing_filesystem;
   ON_CALL(mock_icing_filesystem, CreateDirectoryRecursively)
       .WillByDefault(Return(false));
-  Index::Options options(index_dir_, /*index_merge_size=*/1024 * 1024);
+  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),
               StatusIs(libtextclassifier3::StatusCode::INTERNAL));
 }
@@ -1399,7 +1440,9 @@ TEST_F(IndexTest, IndexCreateCorruptionFailure) {
       IsTrue());
 
   // Recreate the index.
-  Index::Options options(index_dir_, /*index_merge_size=*/1024 * 1024);
+  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_),
               StatusIs(libtextclassifier3::StatusCode::DATA_LOSS));
 }
@@ -1417,7 +1460,9 @@ TEST_F(IndexTest, IndexPersistence) {
   index_.reset();
 
   // Recreate the index.
-  Index::Options options(index_dir_, /*index_merge_size=*/1024 * 1024);
+  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_));
 
@@ -1446,7 +1491,9 @@ TEST_F(IndexTest, IndexPersistenceAfterMerge) {
   index_.reset();
 
   // Recreate the index.
-  Index::Options options(index_dir_, /*index_merge_size=*/1024 * 1024);
+  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_));
 
@@ -1463,7 +1510,8 @@ TEST_F(IndexTest, IndexPersistenceAfterMerge) {
 
 TEST_F(IndexTest, InvalidHitBufferSize) {
   Index::Options options(
-      index_dir_, /*index_merge_size=*/std::numeric_limits<uint32_t>::max());
+      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_),
               StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
 }
diff --git a/icing/index/lite/lite-index-options.cc b/icing/index/lite/lite-index-options.cc
index 29075f8..8780d45 100644
--- a/icing/index/lite/lite-index-options.cc
+++ b/icing/index/lite/lite-index-options.cc
@@ -14,6 +14,8 @@
 
 #include "icing/index/lite/lite-index-options.h"
 
+#include <cstdint>
+
 #include "icing/index/lite/term-id-hit-pair.h"
 
 namespace icing {
@@ -64,9 +66,13 @@ IcingDynamicTrie::Options CalculateTrieOptions(uint32_t hit_buffer_size) {
 }  // namespace
 
 LiteIndexOptions::LiteIndexOptions(const std::string& filename_base,
-                                   uint32_t hit_buffer_want_merge_bytes)
+                                   uint32_t hit_buffer_want_merge_bytes,
+                                   bool hit_buffer_sort_at_indexing,
+                                   uint32_t hit_buffer_sort_threshold_bytes)
     : filename_base(filename_base),
-      hit_buffer_want_merge_bytes(hit_buffer_want_merge_bytes) {
+      hit_buffer_want_merge_bytes(hit_buffer_want_merge_bytes),
+      hit_buffer_sort_at_indexing(hit_buffer_sort_at_indexing),
+      hit_buffer_sort_threshold_bytes(hit_buffer_sort_threshold_bytes) {
   hit_buffer_size = CalculateHitBufferSize(hit_buffer_want_merge_bytes);
   lexicon_options = CalculateTrieOptions(hit_buffer_size);
   display_mappings_options = CalculateTrieOptions(hit_buffer_size);
diff --git a/icing/index/lite/lite-index-options.h b/icing/index/lite/lite-index-options.h
index ae58802..9f8452c 100644
--- a/icing/index/lite/lite-index-options.h
+++ b/icing/index/lite/lite-index-options.h
@@ -27,7 +27,9 @@ struct LiteIndexOptions {
   // hit_buffer_want_merge_bytes and the logic in CalculateHitBufferSize and
   // CalculateTrieOptions.
   LiteIndexOptions(const std::string& filename_base,
-                   uint32_t hit_buffer_want_merge_bytes);
+                   uint32_t hit_buffer_want_merge_bytes,
+                   bool hit_buffer_sort_at_indexing,
+                   uint32_t hit_buffer_sort_threshold_bytes);
 
   IcingDynamicTrie::Options lexicon_options;
   IcingDynamicTrie::Options display_mappings_options;
@@ -35,6 +37,8 @@ struct LiteIndexOptions {
   std::string filename_base;
   uint32_t hit_buffer_want_merge_bytes = 0;
   uint32_t hit_buffer_size = 0;
+  bool hit_buffer_sort_at_indexing = false;
+  uint32_t hit_buffer_sort_threshold_bytes = 0;
 };
 
 }  // namespace lib
diff --git a/icing/index/lite/lite-index.cc b/icing/index/lite/lite-index.cc
index bf54dec..ec7141a 100644
--- a/icing/index/lite/lite-index.cc
+++ b/icing/index/lite/lite-index.cc
@@ -36,6 +36,8 @@
 #include "icing/index/hit/doc-hit-info.h"
 #include "icing/index/hit/hit.h"
 #include "icing/index/lite/lite-index-header.h"
+#include "icing/index/lite/term-id-hit-pair.h"
+#include "icing/index/term-id-codec.h"
 #include "icing/index/term-property-id.h"
 #include "icing/legacy/core/icing-string-util.h"
 #include "icing/legacy/core/icing-timer.h"
@@ -44,10 +46,13 @@
 #include "icing/legacy/index/icing-filesystem.h"
 #include "icing/legacy/index/icing-mmapper.h"
 #include "icing/proto/debug.pb.h"
+#include "icing/proto/scoring.pb.h"
 #include "icing/proto/storage.pb.h"
 #include "icing/proto/term.pb.h"
 #include "icing/schema/section.h"
 #include "icing/store/document-id.h"
+#include "icing/store/namespace-id.h"
+#include "icing/store/suggestion-result-checker.h"
 #include "icing/util/crc32.h"
 #include "icing/util/logging.h"
 #include "icing/util/status-macros.h"
@@ -160,7 +165,7 @@ libtextclassifier3::Status LiteIndex::Initialize() {
     }
 
     // Set up header.
-    header_mmap_.Remap(hit_buffer_fd_.get(), 0, header_size());
+    header_mmap_.Remap(hit_buffer_fd_.get(), kHeaderFileOffset, header_size());
     header_ = std::make_unique<LiteIndex_HeaderImpl>(
         reinterpret_cast<LiteIndex_HeaderImpl::HeaderData*>(
             header_mmap_.address()));
@@ -175,7 +180,7 @@ libtextclassifier3::Status LiteIndex::Initialize() {
 
     UpdateChecksum();
   } else {
-    header_mmap_.Remap(hit_buffer_fd_.get(), 0, header_size());
+    header_mmap_.Remap(hit_buffer_fd_.get(), kHeaderFileOffset, header_size());
     header_ = std::make_unique<LiteIndex_HeaderImpl>(
         reinterpret_cast<LiteIndex_HeaderImpl::HeaderData*>(
             header_mmap_.address()));
@@ -352,6 +357,73 @@ libtextclassifier3::StatusOr<uint32_t> LiteIndex::GetTermId(
   return tvi;
 }
 
+void LiteIndex::ScoreAndAppendFetchedHit(
+    const Hit& hit, SectionIdMask section_id_mask,
+    bool only_from_prefix_sections,
+    SuggestionScoringSpecProto::SuggestionRankingStrategy::Code score_by,
+    const SuggestionResultChecker* suggestion_result_checker,
+    DocumentId& last_document_id, bool& is_last_document_desired,
+    int& total_score_out, std::vector<DocHitInfo>* hits_out,
+    std::vector<Hit::TermFrequencyArray>* term_frequency_out) const {
+  // Check sections.
+  if (((UINT64_C(1) << hit.section_id()) & section_id_mask) == 0) {
+    return;
+  }
+  // Check prefix section only.
+  if (only_from_prefix_sections && !hit.is_in_prefix_section()) {
+    return;
+  }
+  // Check whether this Hit is desired.
+  // TODO(b/230553264) Move common logic into helper function once we support
+  // score term by prefix_hit in lite_index.
+  DocumentId document_id = hit.document_id();
+  bool is_new_document = document_id != last_document_id;
+  if (is_new_document) {
+    last_document_id = document_id;
+    is_last_document_desired =
+        suggestion_result_checker == nullptr ||
+        suggestion_result_checker->BelongsToTargetResults(document_id,
+                                                          hit.section_id());
+  }
+  if (!is_last_document_desired) {
+    // The document is removed or expired or not desired.
+    return;
+  }
+
+  // Score the hit by the strategy
+  switch (score_by) {
+    case SuggestionScoringSpecProto::SuggestionRankingStrategy::NONE:
+      total_score_out = 1;
+      break;
+    case SuggestionScoringSpecProto::SuggestionRankingStrategy::DOCUMENT_COUNT:
+      if (is_new_document) {
+        ++total_score_out;
+      }
+      break;
+    case SuggestionScoringSpecProto::SuggestionRankingStrategy::TERM_FREQUENCY:
+      if (hit.has_term_frequency()) {
+        total_score_out += hit.term_frequency();
+      } else {
+        ++total_score_out;
+      }
+      break;
+  }
+
+  // Append the Hit or update hit section to the output vector.
+  if (is_new_document && hits_out != nullptr) {
+    hits_out->push_back(DocHitInfo(document_id));
+    if (term_frequency_out != nullptr) {
+      term_frequency_out->push_back(Hit::TermFrequencyArray());
+    }
+  }
+  if (hits_out != nullptr) {
+    hits_out->back().UpdateSection(hit.section_id());
+    if (term_frequency_out != nullptr) {
+      term_frequency_out->back()[hit.section_id()] = hit.term_frequency();
+    }
+  }
+}
+
 int LiteIndex::FetchHits(
     uint32_t term_id, SectionIdMask section_id_mask,
     bool only_from_prefix_sections,
@@ -359,19 +431,38 @@ int LiteIndex::FetchHits(
     const SuggestionResultChecker* suggestion_result_checker,
     std::vector<DocHitInfo>* hits_out,
     std::vector<Hit::TermFrequencyArray>* term_frequency_out) {
-  int score = 0;
-  DocumentId last_document_id = kInvalidDocumentId;
-  // Record whether the last document belongs to the given namespaces.
-  bool is_last_document_desired = false;
-
-  if (NeedSort()) {
-    // Transition from shared_lock in NeedSort to unique_lock here is safe
-    // because it doesn't hurt to sort again if sorting was done already by
-    // another thread after NeedSort is evaluated. NeedSort is called before
-    // sorting to improve concurrency as threads can avoid acquiring the unique
-    // lock if no sorting is needed.
+  bool need_sort_at_querying = false;
+  {
+    absl_ports::shared_lock l(&mutex_);
+
+    // We sort here when:
+    // 1. We don't enable sorting at indexing time (i.e. we sort at querying
+    //    time), and there is an unsorted tail portion. OR
+    // 2. The unsorted tail size exceeds the hit_buffer_sort_threshold,
+    //    regardless of whether or not hit_buffer_sort_at_indexing is enabled.
+    //    This is more of a sanity check. We should not really be encountering
+    //    this case.
+    need_sort_at_querying = NeedSortAtQuerying();
+  }
+  if (need_sort_at_querying) {
     absl_ports::unique_lock l(&mutex_);
-    SortHits();
+    IcingTimer timer;
+
+    // Transition from shared_lock to unique_lock is safe here because it
+    // doesn't hurt to sort again if sorting was done already by another thread
+    // 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();
+
+    if (options_.hit_buffer_sort_at_indexing) {
+      // This is the second case for sort. Log as this should be a very rare
+      // occasion.
+      ICING_LOG(WARNING) << "Sorting HitBuffer at querying time when "
+                            "hit_buffer_sort_at_indexing is enabled. Sort and "
+                            "merge HitBuffer in "
+                         << timer.Elapsed() * 1000 << " ms.";
+    }
   }
 
   // This downgrade from an unique_lock to a shared_lock is safe because we're
@@ -379,75 +470,72 @@ int LiteIndex::FetchHits(
   // only in Seek().
   // Any operations that might execute in between the transition of downgrading
   // the lock here are guaranteed not to alter the searchable section (or the
-  // LiteIndex due to a global lock in IcingSearchEngine).
+  // LiteIndex) due to a global lock in IcingSearchEngine.
   absl_ports::shared_lock l(&mutex_);
-  for (uint32_t idx = Seek(term_id); idx < header_->searchable_end(); idx++) {
-    TermIdHitPair term_id_hit_pair =
-        hit_buffer_.array_cast<TermIdHitPair>()[idx];
-    if (term_id_hit_pair.term_id() != term_id) break;
-
-    const Hit& hit = term_id_hit_pair.hit();
-    // Check sections.
-    if (((UINT64_C(1) << hit.section_id()) & section_id_mask) == 0) {
-      continue;
-    }
-    // Check prefix section only.
-    if (only_from_prefix_sections && !hit.is_in_prefix_section()) {
-      continue;
-    }
-    // TODO(b/230553264) Move common logic into helper function once we support
-    //  score term by prefix_hit in lite_index.
-    // Check whether this Hit is desired.
-    DocumentId document_id = hit.document_id();
-    bool is_new_document = document_id != last_document_id;
-    if (is_new_document) {
-      last_document_id = document_id;
-      is_last_document_desired =
-          suggestion_result_checker == nullptr ||
-          suggestion_result_checker->BelongsToTargetResults(document_id,
-                                                            hit.section_id());
-    }
-    if (!is_last_document_desired) {
-      // The document is removed or expired or not desired.
-      continue;
-    }
 
-    // Score the hit by the strategy
-    switch (score_by) {
-      case SuggestionScoringSpecProto::SuggestionRankingStrategy::NONE:
-        score = 1;
-        break;
-      case SuggestionScoringSpecProto::SuggestionRankingStrategy::
-          DOCUMENT_COUNT:
-        if (is_new_document) {
-          ++score;
-        }
-        break;
-      case SuggestionScoringSpecProto::SuggestionRankingStrategy::
-          TERM_FREQUENCY:
-        if (hit.has_term_frequency()) {
-          score += hit.term_frequency();
-        } else {
-          ++score;
-        }
-        break;
-    }
+  // Search in the HitBuffer array for Hits with the corresponding term_id.
+  // Hits are added in increasing order of doc ids, so hits that get appended
+  // later have larger docIds. This means that:
+  // 1. Hits in the unsorted tail will have larger docIds than hits in the
+  //    sorted portion.
+  // 2. Hits at the end of the unsorted tail will have larger docIds than hits
+  //    in the front of the tail.
+  // We want to retrieve hits in descending order of docIds. Therefore we should
+  // search by doing:
+  // 1. Linear search first in reverse iteration order over the unsorted tail
+  //    portion.
+  // 2. Followed by binary search on the sorted portion.
+  const TermIdHitPair* array = hit_buffer_.array_cast<TermIdHitPair>();
 
-    // Append the Hit or update hit section to the output vector.
-    if (is_new_document && hits_out != nullptr) {
-      hits_out->push_back(DocHitInfo(document_id));
-      if (term_frequency_out != nullptr) {
-        term_frequency_out->push_back(Hit::TermFrequencyArray());
+  DocumentId last_document_id = kInvalidDocumentId;
+  // Record whether the last document belongs to the given namespaces.
+  bool is_last_document_desired = false;
+  int total_score = 0;
+
+  // Linear search over unsorted tail in reverse iteration order.
+  // This should only be performed when hit_buffer_sort_at_indexing is enabled.
+  // When disabled, the entire HitBuffer should be sorted already and only
+  // binary search is needed.
+  if (options_.hit_buffer_sort_at_indexing) {
+    uint32_t unsorted_length = header_->cur_size() - header_->searchable_end();
+    for (uint32_t i = 1; i <= unsorted_length; ++i) {
+      TermIdHitPair term_id_hit_pair = array[header_->cur_size() - i];
+      if (term_id_hit_pair.term_id() == term_id) {
+        // We've found a matched hit.
+        const Hit& matched_hit = term_id_hit_pair.hit();
+        // Score the hit and add to total_score. Also add the hits and its term
+        // frequency info to hits_out and term_frequency_out if the two vectors
+        // are non-null.
+        ScoreAndAppendFetchedHit(matched_hit, section_id_mask,
+                                 only_from_prefix_sections, score_by,
+                                 suggestion_result_checker, last_document_id,
+                                 is_last_document_desired, total_score,
+                                 hits_out, term_frequency_out);
       }
     }
-    if (hits_out != nullptr) {
-      hits_out->back().UpdateSection(hit.section_id());
-      if (term_frequency_out != nullptr) {
-        term_frequency_out->back()[hit.section_id()] = hit.term_frequency();
-      }
+  }
+
+  // Do binary search over the sorted section and repeat the above steps.
+  TermIdHitPair target_term_id_hit_pair(
+      term_id, Hit(Hit::kMaxDocumentIdSortValue, Hit::kDefaultTermFrequency));
+  for (const TermIdHitPair* ptr = std::lower_bound(
+           array, array + header_->searchable_end(), target_term_id_hit_pair);
+       ptr < array + header_->searchable_end(); ++ptr) {
+    if (ptr->term_id() != term_id) {
+      // We've processed all matches. Stop iterating further.
+      break;
     }
+
+    const Hit& matched_hit = ptr->hit();
+    // Score the hit and add to total_score. Also add the hits and its term
+    // frequency info to hits_out and term_frequency_out if the two vectors are
+    // non-null.
+    ScoreAndAppendFetchedHit(
+        matched_hit, section_id_mask, only_from_prefix_sections, score_by,
+        suggestion_result_checker, last_document_id, is_last_document_desired,
+        total_score, hits_out, term_frequency_out);
   }
-  return score;
+  return total_score;
 }
 
 libtextclassifier3::StatusOr<int> LiteIndex::ScoreHits(
@@ -455,9 +543,9 @@ libtextclassifier3::StatusOr<int> LiteIndex::ScoreHits(
     SuggestionScoringSpecProto::SuggestionRankingStrategy::Code score_by,
     const SuggestionResultChecker* suggestion_result_checker) {
   return FetchHits(term_id, kSectionIdMaskAll,
-                    /*only_from_prefix_sections=*/false, score_by,
-                    suggestion_result_checker,
-                    /*hits_out=*/nullptr);
+                   /*only_from_prefix_sections=*/false, score_by,
+                   suggestion_result_checker,
+                   /*hits_out=*/nullptr);
 }
 
 bool LiteIndex::is_full() const {
@@ -515,7 +603,7 @@ IndexStorageInfoProto LiteIndex::GetStorageInfo(
   return storage_info;
 }
 
-void LiteIndex::SortHits() {
+void LiteIndex::SortHitsImpl() {
   // Make searchable by sorting by hit buffer.
   uint32_t sort_len = header_->cur_size() - header_->searchable_end();
   if (sort_len <= 0) {
@@ -546,25 +634,6 @@ void LiteIndex::SortHits() {
   UpdateChecksum();
 }
 
-uint32_t LiteIndex::Seek(uint32_t term_id) const {
-  // Binary search for our term_id.  Make sure we get the first
-  // element.  Using kBeginSortValue ensures this for the hit value.
-  TermIdHitPair term_id_hit_pair(
-      term_id, Hit(Hit::kMaxDocumentIdSortValue, Hit::kDefaultTermFrequency));
-
-  const TermIdHitPair::Value* array =
-      hit_buffer_.array_cast<TermIdHitPair::Value>();
-  if (header_->searchable_end() != header_->cur_size()) {
-    ICING_LOG(WARNING) << "Lite index: hit buffer searchable end != current "
-                       << "size during Seek(): "
-                       << header_->searchable_end() << " vs "
-                       << header_->cur_size();
-  }
-  const TermIdHitPair::Value* ptr = std::lower_bound(
-      array, array + header_->searchable_end(), term_id_hit_pair.value());
-  return ptr - array;
-}
-
 libtextclassifier3::Status LiteIndex::Optimize(
     const std::vector<DocumentId>& document_id_old_to_new,
     const TermIdCodec* term_id_codec, DocumentId new_last_added_document_id) {
@@ -575,7 +644,7 @@ libtextclassifier3::Status LiteIndex::Optimize(
   }
   // 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.
-  SortHits();
+  SortHitsImpl();
   uint32_t new_size = 0;
   uint32_t curr_term_id = 0;
   uint32_t curr_tvi = 0;
diff --git a/icing/index/lite/lite-index.h b/icing/index/lite/lite-index.h
index 916a14b..288602a 100644
--- a/icing/index/lite/lite-index.h
+++ b/icing/index/lite/lite-index.h
@@ -20,6 +20,7 @@
 #define ICING_INDEX_LITE_INDEX_H_
 
 #include <cstdint>
+#include <iterator>
 #include <limits>
 #include <memory>
 #include <string>
@@ -48,7 +49,6 @@
 #include "icing/store/document-id.h"
 #include "icing/store/namespace-id.h"
 #include "icing/store/suggestion-result-checker.h"
-#include "icing/util/bit-util.h"
 #include "icing/util/crc32.h"
 
 namespace icing {
@@ -63,6 +63,9 @@ class LiteIndex {
   // An entry in the hit buffer.
   using Options = LiteIndexOptions;
 
+  // Offset for the LiteIndex_Header in the hit buffer mmap.
+  static constexpr uint32_t kHeaderFileOffset = 0;
+
   // Updates checksum of subcomponents.
   ~LiteIndex();
 
@@ -152,8 +155,8 @@ class LiteIndex {
   // Add all hits with term_id from the sections specified in section_id_mask,
   // skipping hits in non-prefix sections if only_from_prefix_sections is true,
   // to hits_out. If hits_out is nullptr, no hits will be added. The
-  // corresponding hit term frequencies will also be added if term_frequency_out
-  // is nullptr.
+  // corresponding hit term frequencies will also not be added if
+  // term_frequency_out is nullptr.
   //
   // Only those hits which belongs to the given namespaces will be counted and
   // fetched. A nullptr namespace checker will disable this check.
@@ -181,15 +184,29 @@ class LiteIndex {
 
   uint32_t size() const ICING_LOCKS_EXCLUDED(mutex_) {
     absl_ports::shared_lock l(&mutex_);
-    return sizeLocked();
+    return size_impl();
   }
 
   bool WantsMerge() const ICING_LOCKS_EXCLUDED(mutex_) {
     absl_ports::shared_lock l(&mutex_);
-    return is_full() || sizeLocked() >= (options_.hit_buffer_want_merge_bytes /
-                        sizeof(TermIdHitPair::Value));
+    return is_full() || size_impl() >= (options_.hit_buffer_want_merge_bytes /
+                                        sizeof(TermIdHitPair::Value));
+  }
+
+  // Whether or not the HitBuffer's unsorted tail size exceeds the sort
+  // threshold.
+  bool HasUnsortedHitsExceedingSortThreshold() const
+      ICING_LOCKS_EXCLUDED(mutex_) {
+    absl_ports::shared_lock l(&mutex_);
+    return HasUnsortedHitsExceedingSortThresholdImpl();
   }
 
+  // Sort hits stored in the index.
+  void SortHits() ICING_LOCKS_EXCLUDED(mutex_) {
+    absl_ports::unique_lock l(&mutex_);
+    SortHitsImpl();
+  };
+
   class const_iterator {
     friend class LiteIndex;
 
@@ -326,17 +343,13 @@ class LiteIndex {
   // Check if the hit buffer has reached its capacity.
   bool is_full() const ICING_SHARED_LOCKS_REQUIRED(mutex_);
 
-  uint32_t sizeLocked() const ICING_SHARED_LOCKS_REQUIRED(mutex_) {
-    return header_->cur_size();
-  }
-
   // Non-locking implementation for empty().
   bool empty_impl() const ICING_SHARED_LOCKS_REQUIRED(mutex_) {
     return size_impl() == 0;
   }
 
   // Non-locking implementation for size().
-  bool size_impl() const ICING_SHARED_LOCKS_REQUIRED(mutex_) {
+  uint32_t size_impl() const ICING_SHARED_LOCKS_REQUIRED(mutex_) {
     return header_->cur_size();
   }
 
@@ -352,18 +365,48 @@ class LiteIndex {
                                                       NamespaceId namespace_id)
       ICING_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
 
-  // Whether or not the HitBuffer requires sorting.
-  bool NeedSort() ICING_LOCKS_EXCLUDED(mutex_) {
-    absl_ports::shared_lock l(&mutex_);
-    return header_->cur_size() - header_->searchable_end() > 0;
+  // We need to sort during querying time when:
+  // 1. Sorting at indexing time is not enabled and there is an unsorted tail
+  //    section in the HitBuffer.
+  // 2. The unsorted tail size exceeds the hit_buffer_sort_threshold, regardless
+  //    of whether or not hit_buffer_sort_at_indexing is enabled. This is to
+  //    prevent performing sequential search on a large unsorted tail section,
+  //    which would result in bad query performance.
+  //    This is more of a sanity check. We should not really be encountering
+  //    this case.
+  bool NeedSortAtQuerying() const ICING_SHARED_LOCKS_REQUIRED(mutex_) {
+    return HasUnsortedHitsExceedingSortThresholdImpl() ||
+           (!options_.hit_buffer_sort_at_indexing &&
+            header_->cur_size() - header_->searchable_end() > 0);
   }
 
-  // Sort hits stored in the index.
-  void SortHits() ICING_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
+  // Non-locking implementation for HasUnsortedHitsExceedingSortThresholdImpl().
+  bool HasUnsortedHitsExceedingSortThresholdImpl() const
+      ICING_SHARED_LOCKS_REQUIRED(mutex_) {
+    return header_->cur_size() - header_->searchable_end() >=
+           (options_.hit_buffer_sort_threshold_bytes /
+            sizeof(TermIdHitPair::Value));
+  }
 
-  // Returns the position of the first element with term_id, or the searchable
-  // end of the hit buffer if term_id is not present.
-  uint32_t Seek(uint32_t term_id) const ICING_SHARED_LOCKS_REQUIRED(mutex_);
+  // Non-locking implementation for SortHits().
+  void 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
+  // far), and is_last_document_desired (which keeps track of whether that last
+  // added docId belongs to the query's desired namespace.)
+  //
+  // Also appends the hit to hits_out and term_frequency_out if the vectors are
+  // not null.
+  void ScoreAndAppendFetchedHit(
+      const Hit& hit, SectionIdMask section_id_mask,
+      bool only_from_prefix_sections,
+      SuggestionScoringSpecProto::SuggestionRankingStrategy::Code score_by,
+      const SuggestionResultChecker* suggestion_result_checker,
+      DocumentId& last_document_id, bool& is_last_document_desired,
+      int& total_score_out, std::vector<DocHitInfo>* hits_out,
+      std::vector<Hit::TermFrequencyArray>* term_frequency_out) const
+      ICING_SHARED_LOCKS_REQUIRED(mutex_);
 
   // File descriptor that points to where the header and hit buffer are written
   // to.
diff --git a/icing/index/lite/lite-index_test.cc b/icing/index/lite/lite-index_test.cc
index 5f141ed..9811fa2 100644
--- a/icing/index/lite/lite-index_test.cc
+++ b/icing/index/lite/lite-index_test.cc
@@ -14,14 +14,27 @@
 
 #include "icing/index/lite/lite-index.h"
 
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <unordered_map>
 #include <vector>
 
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
+#include "icing/file/filesystem.h"
+#include "icing/index/hit/doc-hit-info.h"
+#include "icing/index/hit/hit.h"
+#include "icing/index/iterator/doc-hit-info-iterator.h"
 #include "icing/index/lite/doc-hit-info-iterator-term-lite.h"
+#include "icing/index/lite/lite-index-header.h"
 #include "icing/index/term-id-codec.h"
+#include "icing/legacy/index/icing-dynamic-trie.h"
+#include "icing/legacy/index/icing-filesystem.h"
+#include "icing/proto/scoring.pb.h"
+#include "icing/proto/term.pb.h"
 #include "icing/schema/section.h"
-#include "icing/store/suggestion-result-checker.h"
+#include "icing/store/namespace-id.h"
 #include "icing/testing/always-false-suggestion-result-checker-impl.h"
 #include "icing/testing/common-matchers.h"
 #include "icing/testing/tmp-directory.h"
@@ -34,6 +47,8 @@ namespace {
 using ::testing::ElementsAre;
 using ::testing::Eq;
 using ::testing::IsEmpty;
+using ::testing::IsFalse;
+using ::testing::IsTrue;
 using ::testing::SizeIs;
 
 class LiteIndexTest : public testing::Test {
@@ -41,62 +56,518 @@ class LiteIndexTest : public testing::Test {
   void SetUp() override {
     index_dir_ = GetTestTempDir() + "/test_dir";
     ASSERT_TRUE(filesystem_.CreateDirectoryRecursively(index_dir_.c_str()));
-
-    std::string lite_index_file_name = index_dir_ + "/test_file.lite-idx.index";
-    LiteIndex::Options options(lite_index_file_name,
-                               /*hit_buffer_want_merge_bytes=*/1024 * 1024);
-    ICING_ASSERT_OK_AND_ASSIGN(lite_index_,
-                               LiteIndex::Create(options, &icing_filesystem_));
-
-    ICING_ASSERT_OK_AND_ASSIGN(
-        term_id_codec_,
-        TermIdCodec::Create(
-            IcingDynamicTrie::max_value_index(IcingDynamicTrie::Options()),
-            IcingDynamicTrie::max_value_index(options.lexicon_options)));
   }
 
   void TearDown() override {
     term_id_codec_.reset();
-    lite_index_.reset();
     ASSERT_TRUE(filesystem_.DeleteDirectoryRecursively(index_dir_.c_str()));
   }
 
   std::string index_dir_;
   Filesystem filesystem_;
   IcingFilesystem icing_filesystem_;
-  std::unique_ptr<LiteIndex> lite_index_;
   std::unique_ptr<TermIdCodec> term_id_codec_;
 };
 
 constexpr NamespaceId kNamespace0 = 0;
 
-TEST_F(LiteIndexTest, LiteIndexAppendHits) {
+TEST_F(LiteIndexTest,
+       LiteIndexFetchHits_sortAtQuerying_unsortedHitsBelowSortThreshold) {
+  // Set up LiteIndex and TermIdCodec
+  std::string lite_index_file_name = index_dir_ + "/test_file.lite-idx.index";
+  // At 64 bytes the unsorted tail can contain a max of 8 TermHitPairs.
+  LiteIndex::Options options(lite_index_file_name,
+                             /*hit_buffer_want_merge_bytes=*/1024 * 1024,
+                             /*hit_buffer_sort_at_indexing=*/false,
+                             /*hit_buffer_sort_threshold_bytes=*/64);
+  ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<LiteIndex> lite_index,
+                             LiteIndex::Create(options, &icing_filesystem_));
   ICING_ASSERT_OK_AND_ASSIGN(
-      uint32_t tvi,
-      lite_index_->InsertTerm("foo", TermMatchType::PREFIX, kNamespace0));
+      term_id_codec_,
+      TermIdCodec::Create(
+          IcingDynamicTrie::max_value_index(IcingDynamicTrie::Options()),
+          IcingDynamicTrie::max_value_index(options.lexicon_options)));
+
+  // Add some hits
+  ICING_ASSERT_OK_AND_ASSIGN(
+      uint32_t foo_tvi,
+      lite_index->InsertTerm("foo", TermMatchType::PREFIX, kNamespace0));
   ICING_ASSERT_OK_AND_ASSIGN(uint32_t foo_term_id,
-                             term_id_codec_->EncodeTvi(tvi, TviType::LITE));
-  Hit doc_hit0(/*section_id=*/0, /*document_id=*/0, Hit::kDefaultTermFrequency,
+                             term_id_codec_->EncodeTvi(foo_tvi, TviType::LITE));
+  Hit foo_hit0(/*section_id=*/0, /*document_id=*/1, Hit::kDefaultTermFrequency,
                /*is_in_prefix_section=*/false);
-  Hit doc_hit1(/*section_id=*/1, /*document_id=*/0, Hit::kDefaultTermFrequency,
+  Hit foo_hit1(/*section_id=*/1, /*document_id=*/1, Hit::kDefaultTermFrequency,
                /*is_in_prefix_section=*/false);
-  ICING_ASSERT_OK(lite_index_->AddHit(foo_term_id, doc_hit0));
-  ICING_ASSERT_OK(lite_index_->AddHit(foo_term_id, doc_hit1));
+  ICING_ASSERT_OK(lite_index->AddHit(foo_term_id, foo_hit0));
+  ICING_ASSERT_OK(lite_index->AddHit(foo_term_id, foo_hit1));
 
+  ICING_ASSERT_OK_AND_ASSIGN(
+      uint32_t bar_tvi,
+      lite_index->InsertTerm("bar", TermMatchType::PREFIX, kNamespace0));
+  ICING_ASSERT_OK_AND_ASSIGN(uint32_t bar_term_id,
+                             term_id_codec_->EncodeTvi(bar_tvi, TviType::LITE));
+  Hit bar_hit0(/*section_id=*/0, /*document_id=*/0, Hit::kDefaultTermFrequency,
+               /*is_in_prefix_section=*/false);
+  Hit bar_hit1(/*section_id=*/1, /*document_id=*/0, Hit::kDefaultTermFrequency,
+               /*is_in_prefix_section=*/false);
+  ICING_ASSERT_OK(lite_index->AddHit(bar_term_id, bar_hit0));
+  ICING_ASSERT_OK(lite_index->AddHit(bar_term_id, bar_hit1));
+
+  // Check that unsorted hits does not exceed the sort threshold.
+  EXPECT_THAT(lite_index->HasUnsortedHitsExceedingSortThreshold(), IsFalse());
+
+  // Check that hits are unsorted. Persist the data and pread from
+  // LiteIndexHeader.
+  ASSERT_THAT(lite_index->PersistToDisk(), IsOk());
+  LiteIndex_HeaderImpl::HeaderData header_data;
+  ASSERT_TRUE(filesystem_.PRead((lite_index_file_name + "hb").c_str(),
+                                &header_data, sizeof(header_data),
+                                LiteIndex::kHeaderFileOffset));
+  EXPECT_THAT(header_data.cur_size - header_data.searchable_end, Eq(4));
+
+  // Query the LiteIndex
   std::vector<DocHitInfo> hits1;
-  lite_index_->FetchHits(
+  lite_index->FetchHits(
       foo_term_id, kSectionIdMaskAll,
       /*only_from_prefix_sections=*/false,
       SuggestionScoringSpecProto::SuggestionRankingStrategy::DOCUMENT_COUNT,
       /*namespace_checker=*/nullptr, &hits1);
   EXPECT_THAT(hits1, SizeIs(1));
-  EXPECT_THAT(hits1.back().document_id(), Eq(0));
+  EXPECT_THAT(hits1.back().document_id(), Eq(1));
   // Check that the hits are coming from section 0 and section 1.
   EXPECT_THAT(hits1.back().hit_section_ids_mask(), Eq(0b11));
 
   std::vector<DocHitInfo> hits2;
   AlwaysFalseSuggestionResultCheckerImpl always_false_suggestion_result_checker;
-  lite_index_->FetchHits(
+  lite_index->FetchHits(
+      foo_term_id, kSectionIdMaskAll,
+      /*only_from_prefix_sections=*/false,
+      SuggestionScoringSpecProto::SuggestionRankingStrategy::DOCUMENT_COUNT,
+      &always_false_suggestion_result_checker, &hits2);
+  // Check that no hits are returned because they get skipped by the namespace
+  // checker.
+  EXPECT_THAT(hits2, IsEmpty());
+
+  // Check that hits are sorted after querying LiteIndex. Persist the data and
+  // pread from LiteIndexHeader.
+  ASSERT_THAT(lite_index->PersistToDisk(), IsOk());
+  ASSERT_TRUE(filesystem_.PRead((lite_index_file_name + "hb").c_str(),
+                                &header_data, sizeof(header_data),
+                                LiteIndex::kHeaderFileOffset));
+  EXPECT_THAT(header_data.cur_size - header_data.searchable_end, Eq(0));
+}
+
+TEST_F(LiteIndexTest,
+       LiteIndexFetchHits_sortAtIndexing_unsortedHitsBelowSortThreshold) {
+  // Set up LiteIndex and TermIdCodec
+  std::string lite_index_file_name = index_dir_ + "/test_file.lite-idx.index";
+  // At 64 bytes the unsorted tail can contain a max of 8 TermHitPairs.
+  // However note that in these tests we're unable to sort hits after
+  // indexing, as sorting performed by the string-section-indexing-handler
+  // after indexing all hits in an entire document, rather than after each
+  // AddHits() operation.
+  LiteIndex::Options options(lite_index_file_name,
+                             /*hit_buffer_want_merge_bytes=*/1024 * 1024,
+                             /*hit_buffer_sort_at_indexing=*/true,
+                             /*hit_buffer_sort_threshold_bytes=*/64);
+  ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<LiteIndex> lite_index,
+                             LiteIndex::Create(options, &icing_filesystem_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      term_id_codec_,
+      TermIdCodec::Create(
+          IcingDynamicTrie::max_value_index(IcingDynamicTrie::Options()),
+          IcingDynamicTrie::max_value_index(options.lexicon_options)));
+
+  // Add some hits
+  ICING_ASSERT_OK_AND_ASSIGN(
+      uint32_t foo_tvi,
+      lite_index->InsertTerm("foo", TermMatchType::PREFIX, kNamespace0));
+  ICING_ASSERT_OK_AND_ASSIGN(uint32_t foo_term_id,
+                             term_id_codec_->EncodeTvi(foo_tvi, TviType::LITE));
+  Hit foo_hit0(/*section_id=*/0, /*document_id=*/1, Hit::kDefaultTermFrequency,
+               /*is_in_prefix_section=*/false);
+  Hit foo_hit1(/*section_id=*/1, /*document_id=*/1, Hit::kDefaultTermFrequency,
+               /*is_in_prefix_section=*/false);
+  ICING_ASSERT_OK(lite_index->AddHit(foo_term_id, foo_hit0));
+  ICING_ASSERT_OK(lite_index->AddHit(foo_term_id, foo_hit1));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      uint32_t bar_tvi,
+      lite_index->InsertTerm("bar", TermMatchType::PREFIX, kNamespace0));
+  ICING_ASSERT_OK_AND_ASSIGN(uint32_t bar_term_id,
+                             term_id_codec_->EncodeTvi(bar_tvi, TviType::LITE));
+  Hit bar_hit0(/*section_id=*/0, /*document_id=*/0, Hit::kDefaultTermFrequency,
+               /*is_in_prefix_section=*/false);
+  Hit bar_hit1(/*section_id=*/1, /*document_id=*/0, Hit::kDefaultTermFrequency,
+               /*is_in_prefix_section=*/false);
+  ICING_ASSERT_OK(lite_index->AddHit(bar_term_id, bar_hit0));
+  ICING_ASSERT_OK(lite_index->AddHit(bar_term_id, bar_hit1));
+
+  // Check that unsorted hits does not exceed the sort threshold.
+  EXPECT_THAT(lite_index->HasUnsortedHitsExceedingSortThreshold(), IsFalse());
+
+  // Check that hits are unsorted. Persist the data and pread from
+  // LiteIndexHeader.
+  ASSERT_THAT(lite_index->PersistToDisk(), IsOk());
+  LiteIndex_HeaderImpl::HeaderData header_data;
+  ASSERT_TRUE(filesystem_.PRead((lite_index_file_name + "hb").c_str(),
+                                &header_data, sizeof(header_data),
+                                LiteIndex::kHeaderFileOffset));
+  EXPECT_THAT(header_data.cur_size - header_data.searchable_end, Eq(4));
+
+  // Query the LiteIndex
+  std::vector<DocHitInfo> hits1;
+  lite_index->FetchHits(
+      foo_term_id, kSectionIdMaskAll,
+      /*only_from_prefix_sections=*/false,
+      SuggestionScoringSpecProto::SuggestionRankingStrategy::DOCUMENT_COUNT,
+      /*namespace_checker=*/nullptr, &hits1);
+  EXPECT_THAT(hits1, SizeIs(1));
+  EXPECT_THAT(hits1.back().document_id(), Eq(1));
+  // Check that the hits are coming from section 0 and section 1.
+  EXPECT_THAT(hits1.back().hit_section_ids_mask(), Eq(0b11));
+
+  std::vector<DocHitInfo> hits2;
+  AlwaysFalseSuggestionResultCheckerImpl always_false_suggestion_result_checker;
+  lite_index->FetchHits(
+      foo_term_id, kSectionIdMaskAll,
+      /*only_from_prefix_sections=*/false,
+      SuggestionScoringSpecProto::SuggestionRankingStrategy::DOCUMENT_COUNT,
+      &always_false_suggestion_result_checker, &hits2);
+  // Check that no hits are returned because they get skipped by the namespace
+  // checker.
+  EXPECT_THAT(hits2, IsEmpty());
+
+  // Check that hits are still unsorted after querying LiteIndex because the
+  // HitBuffer unsorted size is still below the sort threshold, and we've
+  // enabled sort_at_indexing.
+  // Persist the data and performing a pread on LiteIndexHeader.
+  ASSERT_THAT(lite_index->PersistToDisk(), IsOk());
+  ASSERT_TRUE(filesystem_.PRead((lite_index_file_name + "hb").c_str(),
+                                &header_data, sizeof(header_data),
+                                LiteIndex::kHeaderFileOffset));
+  EXPECT_THAT(header_data.cur_size - header_data.searchable_end, Eq(4));
+}
+
+TEST_F(
+    LiteIndexTest,
+    LiteIndexFetchHits_sortAtQuerying_unsortedHitsExceedingSortAtIndexThreshold) {
+  // Set up LiteIndex and TermIdCodec
+  std::string lite_index_file_name = index_dir_ + "/test_file.lite-idx.index";
+  // At 64 bytes the unsorted tail can contain a max of 8 TermHitPairs.
+  // However note that in these tests we're unable to sort hits after
+  // indexing, as sorting performed by the string-section-indexing-handler
+  // after indexing all hits in an entire document, rather than after each
+  // AddHits() operation.
+  LiteIndex::Options options(lite_index_file_name,
+                             /*hit_buffer_want_merge_bytes=*/1024 * 1024,
+                             /*hit_buffer_sort_at_indexing=*/false,
+                             /*hit_buffer_sort_threshold_bytes=*/64);
+  ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<LiteIndex> lite_index,
+                             LiteIndex::Create(options, &icing_filesystem_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      term_id_codec_,
+      TermIdCodec::Create(
+          IcingDynamicTrie::max_value_index(IcingDynamicTrie::Options()),
+          IcingDynamicTrie::max_value_index(options.lexicon_options)));
+
+  // Create 4 hits for docs 0-2, and 2 hits for doc 3 -- 14 in total
+  // Doc 0
+  Hit doc0_hit0(/*section_id=*/0, /*document_id=*/0, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  Hit doc0_hit1(/*section_id=*/0, /*document_id=*/0, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  Hit doc0_hit2(/*section_id=*/1, /*document_id=*/0, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  Hit doc0_hit3(/*section_id=*/2, /*document_id=*/0, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  // Doc 1
+  Hit doc1_hit0(/*section_id=*/0, /*document_id=*/1, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  Hit doc1_hit1(/*section_id=*/0, /*document_id=*/1, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  Hit doc1_hit2(/*section_id=*/1, /*document_id=*/1, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  Hit doc1_hit3(/*section_id=*/2, /*document_id=*/1, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  // Doc 2
+  Hit doc2_hit0(/*section_id=*/0, /*document_id=*/2, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  Hit doc2_hit1(/*section_id=*/0, /*document_id=*/2, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  Hit doc2_hit2(/*section_id=*/1, /*document_id=*/2, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  Hit doc2_hit3(/*section_id=*/2, /*document_id=*/2, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  // Doc 3
+  Hit doc3_hit0(/*section_id=*/0, /*document_id=*/3, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  Hit doc3_hit1(/*section_id=*/0, /*document_id=*/3, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+
+  // Create terms
+  // Foo
+  ICING_ASSERT_OK_AND_ASSIGN(
+      uint32_t foo_tvi,
+      lite_index->InsertTerm("foo", TermMatchType::EXACT_ONLY, kNamespace0));
+  ICING_ASSERT_OK_AND_ASSIGN(uint32_t foo_term_id,
+                             term_id_codec_->EncodeTvi(foo_tvi, TviType::LITE));
+  // Bar
+  ICING_ASSERT_OK_AND_ASSIGN(
+      uint32_t bar_tvi,
+      lite_index->InsertTerm("bar", TermMatchType::PREFIX, kNamespace0));
+  ICING_ASSERT_OK_AND_ASSIGN(uint32_t bar_term_id,
+                             term_id_codec_->EncodeTvi(bar_tvi, TviType::LITE));
+  // Baz
+  ICING_ASSERT_OK_AND_ASSIGN(
+      uint32_t baz_tvi,
+      lite_index->InsertTerm("baz", TermMatchType::PREFIX, kNamespace0));
+  ICING_ASSERT_OK_AND_ASSIGN(uint32_t baz_term_id,
+                             term_id_codec_->EncodeTvi(baz_tvi, TviType::LITE));
+  // Qux
+  ICING_ASSERT_OK_AND_ASSIGN(
+      uint32_t qux_tvi,
+      lite_index->InsertTerm("qux", TermMatchType::PREFIX, kNamespace0));
+  ICING_ASSERT_OK_AND_ASSIGN(uint32_t qux_term_id,
+                             term_id_codec_->EncodeTvi(qux_tvi, TviType::LITE));
+
+  // Add 14 hits and make sure that termIds are added in unsorted order.
+  // Documents should be inserted in order as new incoming hits should have
+  // larger document ids.
+  ICING_ASSERT_OK(lite_index->AddHit(foo_term_id, doc0_hit0));
+  ICING_ASSERT_OK(lite_index->AddHit(bar_term_id, doc0_hit1));
+  ICING_ASSERT_OK(lite_index->AddHit(baz_term_id, doc0_hit2));
+  ICING_ASSERT_OK(lite_index->AddHit(qux_term_id, doc0_hit3));
+  ICING_ASSERT_OK(lite_index->AddHit(foo_term_id, doc1_hit0));
+  ICING_ASSERT_OK(lite_index->AddHit(bar_term_id, doc1_hit1));
+  ICING_ASSERT_OK(lite_index->AddHit(foo_term_id, doc1_hit2));
+  ICING_ASSERT_OK(lite_index->AddHit(bar_term_id, doc1_hit3));
+  ICING_ASSERT_OK(lite_index->AddHit(foo_term_id, doc2_hit0));
+  ICING_ASSERT_OK(lite_index->AddHit(baz_term_id, doc2_hit1));
+  ICING_ASSERT_OK(lite_index->AddHit(qux_term_id, doc2_hit2));
+  ICING_ASSERT_OK(lite_index->AddHit(foo_term_id, doc2_hit3));
+  ICING_ASSERT_OK(lite_index->AddHit(foo_term_id, doc3_hit0));
+  ICING_ASSERT_OK(lite_index->AddHit(baz_term_id, doc3_hit1));
+  // Verify that the HitBuffer has not been sorted.
+  EXPECT_THAT(lite_index->HasUnsortedHitsExceedingSortThreshold(), IsTrue());
+
+  // We now have the following in the hit buffer:
+  // <term>: {(docId, sectionId)...}
+  // foo: {(0, 0); (1, 0); (1, 1); (2, 0); (2, 2); (3, 0)}
+  // bar: {(0, 0); (1, 0); (1, 2)}
+  // baz: {(0, 1); (2, 0); (3, 0)}
+  // quz: {(0, 2); (2, 1)}
+
+  // Search over the HitBuffer.
+  std::vector<DocHitInfo> hits1;
+  lite_index->FetchHits(
+      foo_term_id, kSectionIdMaskAll,
+      /*only_from_prefix_sections=*/false,
+      SuggestionScoringSpecProto::SuggestionRankingStrategy::DOCUMENT_COUNT,
+      /*namespace_checker=*/nullptr, &hits1);
+  EXPECT_THAT(hits1, SizeIs(4));
+  // Check that hits are retrieved in descending order of docIds.
+  EXPECT_THAT(hits1[0].document_id(), Eq(3));
+  EXPECT_THAT(hits1[0].hit_section_ids_mask(), Eq(0b1));
+  EXPECT_THAT(hits1[1].document_id(), Eq(2));
+  EXPECT_THAT(hits1[1].hit_section_ids_mask(), Eq(0b101));
+  EXPECT_THAT(hits1[2].document_id(), Eq(1));
+  EXPECT_THAT(hits1[2].hit_section_ids_mask(), Eq(0b11));
+  EXPECT_THAT(hits1[3].document_id(), Eq(0));
+  EXPECT_THAT(hits1[3].hit_section_ids_mask(), Eq(0b1));
+
+  std::vector<DocHitInfo> hits2;
+  AlwaysFalseSuggestionResultCheckerImpl always_false_suggestion_result_checker;
+  lite_index->FetchHits(
+      foo_term_id, kSectionIdMaskAll,
+      /*only_from_prefix_sections=*/false,
+      SuggestionScoringSpecProto::SuggestionRankingStrategy::DOCUMENT_COUNT,
+      &always_false_suggestion_result_checker, &hits2);
+  // Check that no hits are returned because they get skipped by the namespace
+  // checker.
+  EXPECT_THAT(hits2, IsEmpty());
+
+  std::vector<DocHitInfo> hits3;
+  lite_index->FetchHits(
+      bar_term_id, 0b1,
+      /*only_from_prefix_sections=*/false,
+      SuggestionScoringSpecProto::SuggestionRankingStrategy::DOCUMENT_COUNT,
+      /*namespace_checker=*/nullptr, &hits3);
+  EXPECT_THAT(hits3, SizeIs(2));
+  // Check fetching hits with SectionIdMask.
+  EXPECT_THAT(hits3[0].document_id(), Eq(1));
+  EXPECT_THAT(hits3[1].hit_section_ids_mask(), Eq(0b1));
+  EXPECT_THAT(hits3[1].document_id(), Eq(0));
+  EXPECT_THAT(hits3[1].hit_section_ids_mask(), Eq(0b1));
+
+  // Check that the HitBuffer is sorted after the query call.
+  EXPECT_THAT(lite_index->HasUnsortedHitsExceedingSortThreshold(), IsFalse());
+}
+
+TEST_F(
+    LiteIndexTest,
+    LiteIndexFetchHits_sortAtIndexing_unsortedHitsExceedingSortAtIndexThreshold) {
+  // Set up LiteIndex and TermIdCodec
+  std::string lite_index_file_name = index_dir_ + "/test_file.lite-idx.index";
+  // At 64 bytes the unsorted tail can contain a max of 8 TermHitPairs.
+  LiteIndex::Options options(lite_index_file_name,
+                             /*hit_buffer_want_merge_bytes=*/1024 * 1024,
+                             /*hit_buffer_sort_at_indexing=*/true,
+                             /*hit_buffer_sort_threshold_bytes=*/64);
+  ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<LiteIndex> lite_index,
+                             LiteIndex::Create(options, &icing_filesystem_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      term_id_codec_,
+      TermIdCodec::Create(
+          IcingDynamicTrie::max_value_index(IcingDynamicTrie::Options()),
+          IcingDynamicTrie::max_value_index(options.lexicon_options)));
+
+  // Create 4 hits for docs 0-2, and 2 hits for doc 3 -- 14 in total
+  // Doc 0
+  Hit doc0_hit0(/*section_id=*/0, /*document_id=*/0, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  Hit doc0_hit1(/*section_id=*/0, /*document_id=*/0, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  Hit doc0_hit2(/*section_id=*/1, /*document_id=*/0, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  Hit doc0_hit3(/*section_id=*/2, /*document_id=*/0, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  // Doc 1
+  Hit doc1_hit0(/*section_id=*/0, /*document_id=*/1, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  Hit doc1_hit1(/*section_id=*/0, /*document_id=*/1, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  Hit doc1_hit2(/*section_id=*/1, /*document_id=*/1, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  Hit doc1_hit3(/*section_id=*/2, /*document_id=*/1, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  // Doc 2
+  Hit doc2_hit0(/*section_id=*/0, /*document_id=*/2, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  Hit doc2_hit1(/*section_id=*/0, /*document_id=*/2, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  Hit doc2_hit2(/*section_id=*/1, /*document_id=*/2, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  Hit doc2_hit3(/*section_id=*/2, /*document_id=*/2, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  // Doc 3
+  Hit doc3_hit0(/*section_id=*/0, /*document_id=*/3, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  Hit doc3_hit1(/*section_id=*/0, /*document_id=*/3, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  Hit doc3_hit2(/*section_id=*/1, /*document_id=*/3, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  Hit doc3_hit3(/*section_id=*/2, /*document_id=*/3, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  // Doc 4
+  Hit doc4_hit0(/*section_id=*/0, /*document_id=*/4, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  Hit doc4_hit1(/*section_id=*/0, /*document_id=*/4, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  Hit doc4_hit2(/*section_id=*/1, /*document_id=*/4, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+  Hit doc4_hit3(/*section_id=*/2, /*document_id=*/4, Hit::kDefaultTermFrequency,
+                /*is_in_prefix_section=*/false);
+
+  // Create terms
+  // Foo
+  ICING_ASSERT_OK_AND_ASSIGN(
+      uint32_t foo_tvi,
+      lite_index->InsertTerm("foo", TermMatchType::EXACT_ONLY, kNamespace0));
+  ICING_ASSERT_OK_AND_ASSIGN(uint32_t foo_term_id,
+                             term_id_codec_->EncodeTvi(foo_tvi, TviType::LITE));
+  // Bar
+  ICING_ASSERT_OK_AND_ASSIGN(
+      uint32_t bar_tvi,
+      lite_index->InsertTerm("bar", TermMatchType::PREFIX, kNamespace0));
+  ICING_ASSERT_OK_AND_ASSIGN(uint32_t bar_term_id,
+                             term_id_codec_->EncodeTvi(bar_tvi, TviType::LITE));
+  // Baz
+  ICING_ASSERT_OK_AND_ASSIGN(
+      uint32_t baz_tvi,
+      lite_index->InsertTerm("baz", TermMatchType::PREFIX, kNamespace0));
+  ICING_ASSERT_OK_AND_ASSIGN(uint32_t baz_term_id,
+                             term_id_codec_->EncodeTvi(baz_tvi, TviType::LITE));
+  // Qux
+  ICING_ASSERT_OK_AND_ASSIGN(
+      uint32_t qux_tvi,
+      lite_index->InsertTerm("qux", TermMatchType::PREFIX, kNamespace0));
+  ICING_ASSERT_OK_AND_ASSIGN(uint32_t qux_term_id,
+                             term_id_codec_->EncodeTvi(qux_tvi, TviType::LITE));
+
+  // Add hits and make sure that termIds are added in unsorted order.
+  // Documents should be inserted in order as new incoming hits should have
+  // larger document ids.
+  ICING_ASSERT_OK(lite_index->AddHit(foo_term_id, doc0_hit0));
+  ICING_ASSERT_OK(lite_index->AddHit(bar_term_id, doc0_hit1));
+  ICING_ASSERT_OK(lite_index->AddHit(baz_term_id, doc0_hit2));
+  ICING_ASSERT_OK(lite_index->AddHit(qux_term_id, doc0_hit3));
+  ICING_ASSERT_OK(lite_index->AddHit(foo_term_id, doc1_hit0));
+  ICING_ASSERT_OK(lite_index->AddHit(bar_term_id, doc1_hit1));
+  ICING_ASSERT_OK(lite_index->AddHit(foo_term_id, doc1_hit2));
+  ICING_ASSERT_OK(lite_index->AddHit(bar_term_id, doc1_hit3));
+  // Adding 8 hits exceeds the sort threshold. However when sort_at_indexing is
+  // enabled, sorting is done in the string-section-indexing-handler rather than
+  // AddHit() itself, we need to invoke SortHits() manually.
+  EXPECT_THAT(lite_index->HasUnsortedHitsExceedingSortThreshold(), IsTrue());
+  lite_index->SortHits();
+  // Check that the HitBuffer is sorted.
+  ASSERT_THAT(lite_index->PersistToDisk(), IsOk());
+  LiteIndex_HeaderImpl::HeaderData header_data;
+  ASSERT_TRUE(filesystem_.PRead((lite_index_file_name + "hb").c_str(),
+                                &header_data, sizeof(header_data),
+                                LiteIndex::kHeaderFileOffset));
+  EXPECT_THAT(header_data.cur_size - header_data.searchable_end, Eq(0));
+
+  // Add 12 more hits so that sort threshold is exceeded again.
+  ICING_ASSERT_OK(lite_index->AddHit(foo_term_id, doc2_hit0));
+  ICING_ASSERT_OK(lite_index->AddHit(baz_term_id, doc2_hit1));
+  ICING_ASSERT_OK(lite_index->AddHit(qux_term_id, doc2_hit2));
+  ICING_ASSERT_OK(lite_index->AddHit(foo_term_id, doc2_hit3));
+  ICING_ASSERT_OK(lite_index->AddHit(foo_term_id, doc3_hit0));
+  ICING_ASSERT_OK(lite_index->AddHit(baz_term_id, doc3_hit1));
+  ICING_ASSERT_OK(lite_index->AddHit(foo_term_id, doc3_hit2));
+  ICING_ASSERT_OK(lite_index->AddHit(bar_term_id, doc3_hit3));
+  ICING_ASSERT_OK(lite_index->AddHit(baz_term_id, doc4_hit0));
+  ICING_ASSERT_OK(lite_index->AddHit(qux_term_id, doc4_hit1));
+  ICING_ASSERT_OK(lite_index->AddHit(foo_term_id, doc4_hit2));
+  ICING_ASSERT_OK(lite_index->AddHit(bar_term_id, doc4_hit3));
+
+  // Adding these hits exceeds the sort threshold. However when sort_at_indexing
+  // is enabled, sorting is done in the string-section-indexing-handler rather
+  // than AddHit() itself.
+  EXPECT_THAT(lite_index->HasUnsortedHitsExceedingSortThreshold(), IsTrue());
+
+  // We now have the following in the hit buffer:
+  // <term>: {(docId, sectionId)...}
+  // foo: {(0, 0); (1, 0); (1, 1); (2, 0); (2, 2); (3, 0); (3, 1); (4, 1)}
+  // bar: {(0, 0); (1, 0); (1, 2); (3, 2); (4, 2)}
+  // baz: {(0, 1); (2, 0); (3, 0); (4, 0)}
+  // quz: {(0, 2); (2, 1); (4, 0)}
+
+  // Search over the HitBuffer.
+  std::vector<DocHitInfo> hits1;
+  lite_index->FetchHits(
+      foo_term_id, kSectionIdMaskAll,
+      /*only_from_prefix_sections=*/false,
+      SuggestionScoringSpecProto::SuggestionRankingStrategy::DOCUMENT_COUNT,
+      /*namespace_checker=*/nullptr, &hits1);
+  EXPECT_THAT(hits1, SizeIs(5));
+  // Check that hits are retrieved in descending order of docIds.
+  EXPECT_THAT(hits1[0].document_id(), Eq(4));
+  EXPECT_THAT(hits1[0].hit_section_ids_mask(), Eq(0b10));
+  EXPECT_THAT(hits1[1].document_id(), Eq(3));
+  EXPECT_THAT(hits1[1].hit_section_ids_mask(), Eq(0b11));
+  EXPECT_THAT(hits1[2].document_id(), Eq(2));
+  EXPECT_THAT(hits1[2].hit_section_ids_mask(), Eq(0b101));
+  EXPECT_THAT(hits1[3].document_id(), Eq(1));
+  EXPECT_THAT(hits1[3].hit_section_ids_mask(), Eq(0b11));
+  EXPECT_THAT(hits1[4].document_id(), Eq(0));
+  EXPECT_THAT(hits1[4].hit_section_ids_mask(), Eq(0b1));
+
+  std::vector<DocHitInfo> hits2;
+  AlwaysFalseSuggestionResultCheckerImpl always_false_suggestion_result_checker;
+  lite_index->FetchHits(
       foo_term_id, kSectionIdMaskAll,
       /*only_from_prefix_sections=*/false,
       SuggestionScoringSpecProto::SuggestionRankingStrategy::DOCUMENT_COUNT,
@@ -104,13 +575,119 @@ TEST_F(LiteIndexTest, LiteIndexAppendHits) {
   // Check that no hits are returned because they get skipped by the namespace
   // checker.
   EXPECT_THAT(hits2, IsEmpty());
+
+  std::vector<DocHitInfo> hits3;
+  lite_index->FetchHits(
+      bar_term_id, 0b1,
+      /*only_from_prefix_sections=*/false,
+      SuggestionScoringSpecProto::SuggestionRankingStrategy::DOCUMENT_COUNT,
+      /*namespace_checker=*/nullptr, &hits3);
+  EXPECT_THAT(hits3, SizeIs(2));
+  // Check fetching hits with SectionIdMask.
+  EXPECT_THAT(hits3[0].document_id(), Eq(1));
+  EXPECT_THAT(hits3[1].hit_section_ids_mask(), Eq(0b1));
+  EXPECT_THAT(hits3[1].document_id(), Eq(0));
+  EXPECT_THAT(hits3[1].hit_section_ids_mask(), Eq(0b1));
+
+  // Check that the HitBuffer is sorted after the query call. FetchHits should
+  // sort before performing binary search if the HitBuffer unsorted size exceeds
+  // the sort threshold. Regardless of the sort_at_indexing config.
+  EXPECT_THAT(lite_index->HasUnsortedHitsExceedingSortThreshold(), IsFalse());
+  ASSERT_THAT(lite_index->PersistToDisk(), IsOk());
+  ASSERT_TRUE(filesystem_.PRead((lite_index_file_name + "hb").c_str(),
+                                &header_data, sizeof(header_data),
+                                LiteIndex::kHeaderFileOffset));
+  EXPECT_THAT(header_data.cur_size - header_data.searchable_end, Eq(0));
 }
 
 TEST_F(LiteIndexTest, LiteIndexIterator) {
+  // Set up LiteIndex and TermIdCodec
+  std::string lite_index_file_name = index_dir_ + "/test_file.lite-idx.index";
+  // At 64 bytes the unsorted tail can contain a max of 8 TermHitPairs.
+  LiteIndex::Options options(lite_index_file_name,
+                             /*hit_buffer_want_merge_bytes=*/1024 * 1024,
+                             /*hit_buffer_sort_at_indexing=*/true,
+                             /*hit_buffer_sort_threshold_bytes=*/64);
+  ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<LiteIndex> lite_index,
+                             LiteIndex::Create(options, &icing_filesystem_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      term_id_codec_,
+      TermIdCodec::Create(
+          IcingDynamicTrie::max_value_index(IcingDynamicTrie::Options()),
+          IcingDynamicTrie::max_value_index(options.lexicon_options)));
+
+  const std::string term = "foo";
+  ICING_ASSERT_OK_AND_ASSIGN(
+      uint32_t tvi,
+      lite_index->InsertTerm(term, TermMatchType::PREFIX, kNamespace0));
+  ICING_ASSERT_OK_AND_ASSIGN(uint32_t foo_term_id,
+                             term_id_codec_->EncodeTvi(tvi, TviType::LITE));
+  Hit doc0_hit0(/*section_id=*/0, /*document_id=*/0, /*term_frequency=*/3,
+                /*is_in_prefix_section=*/false);
+  Hit doc0_hit1(/*section_id=*/1, /*document_id=*/0, /*term_frequency=*/5,
+                /*is_in_prefix_section=*/false);
+  SectionIdMask doc0_section_id_mask = 0b11;
+  std::unordered_map<SectionId, Hit::TermFrequency>
+      expected_section_ids_tf_map0 = {{0, 3}, {1, 5}};
+  ICING_ASSERT_OK(lite_index->AddHit(foo_term_id, doc0_hit0));
+  ICING_ASSERT_OK(lite_index->AddHit(foo_term_id, doc0_hit1));
+
+  Hit doc1_hit1(/*section_id=*/1, /*document_id=*/1, /*term_frequency=*/7,
+                /*is_in_prefix_section=*/false);
+  Hit doc1_hit2(/*section_id=*/2, /*document_id=*/1, /*term_frequency=*/11,
+                /*is_in_prefix_section=*/false);
+  SectionIdMask doc1_section_id_mask = 0b110;
+  std::unordered_map<SectionId, Hit::TermFrequency>
+      expected_section_ids_tf_map1 = {{1, 7}, {2, 11}};
+  ICING_ASSERT_OK(lite_index->AddHit(foo_term_id, doc1_hit1));
+  ICING_ASSERT_OK(lite_index->AddHit(foo_term_id, doc1_hit2));
+
+  std::unique_ptr<DocHitInfoIteratorTermLiteExact> iter =
+      std::make_unique<DocHitInfoIteratorTermLiteExact>(
+          term_id_codec_.get(), lite_index.get(), term, /*term_start_index=*/0,
+          /*unnormalized_term_length=*/0, kSectionIdMaskAll,
+          /*need_hit_term_frequency=*/true);
+
+  ASSERT_THAT(iter->Advance(), IsOk());
+  EXPECT_THAT(iter->doc_hit_info().document_id(), Eq(1));
+  EXPECT_THAT(iter->doc_hit_info().hit_section_ids_mask(),
+              Eq(doc1_section_id_mask));
+
+  std::vector<TermMatchInfo> matched_terms_stats;
+  iter->PopulateMatchedTermsStats(&matched_terms_stats);
+  EXPECT_THAT(matched_terms_stats, ElementsAre(EqualsTermMatchInfo(
+                                       term, expected_section_ids_tf_map1)));
+
+  ASSERT_THAT(iter->Advance(), IsOk());
+  EXPECT_THAT(iter->doc_hit_info().document_id(), Eq(0));
+  EXPECT_THAT(iter->doc_hit_info().hit_section_ids_mask(),
+              Eq(doc0_section_id_mask));
+  matched_terms_stats.clear();
+  iter->PopulateMatchedTermsStats(&matched_terms_stats);
+  EXPECT_THAT(matched_terms_stats, ElementsAre(EqualsTermMatchInfo(
+                                       term, expected_section_ids_tf_map0)));
+}
+
+TEST_F(LiteIndexTest, LiteIndexIterator_sortAtIndexingDisabled) {
+  // Set up LiteIndex and TermIdCodec
+  std::string lite_index_file_name = index_dir_ + "/test_file.lite-idx.index";
+  // At 64 bytes the unsorted tail can contain a max of 8 TermHitPairs.
+  LiteIndex::Options options(lite_index_file_name,
+                             /*hit_buffer_want_merge_bytes=*/1024 * 1024,
+                             /*hit_buffer_sort_at_indexing=*/false,
+                             /*hit_buffer_sort_threshold_bytes=*/64);
+  ICING_ASSERT_OK_AND_ASSIGN(std::unique_ptr<LiteIndex> lite_index,
+                             LiteIndex::Create(options, &icing_filesystem_));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      term_id_codec_,
+      TermIdCodec::Create(
+          IcingDynamicTrie::max_value_index(IcingDynamicTrie::Options()),
+          IcingDynamicTrie::max_value_index(options.lexicon_options)));
+
   const std::string term = "foo";
   ICING_ASSERT_OK_AND_ASSIGN(
       uint32_t tvi,
-      lite_index_->InsertTerm(term, TermMatchType::PREFIX, kNamespace0));
+      lite_index->InsertTerm(term, TermMatchType::PREFIX, kNamespace0));
   ICING_ASSERT_OK_AND_ASSIGN(uint32_t foo_term_id,
                              term_id_codec_->EncodeTvi(tvi, TviType::LITE));
   Hit doc0_hit0(/*section_id=*/0, /*document_id=*/0, /*term_frequency=*/3,
@@ -120,8 +697,8 @@ TEST_F(LiteIndexTest, LiteIndexIterator) {
   SectionIdMask doc0_section_id_mask = 0b11;
   std::unordered_map<SectionId, Hit::TermFrequency>
       expected_section_ids_tf_map0 = {{0, 3}, {1, 5}};
-  ICING_ASSERT_OK(lite_index_->AddHit(foo_term_id, doc0_hit0));
-  ICING_ASSERT_OK(lite_index_->AddHit(foo_term_id, doc0_hit1));
+  ICING_ASSERT_OK(lite_index->AddHit(foo_term_id, doc0_hit0));
+  ICING_ASSERT_OK(lite_index->AddHit(foo_term_id, doc0_hit1));
 
   Hit doc1_hit1(/*section_id=*/1, /*document_id=*/1, /*term_frequency=*/7,
                 /*is_in_prefix_section=*/false);
@@ -130,12 +707,12 @@ TEST_F(LiteIndexTest, LiteIndexIterator) {
   SectionIdMask doc1_section_id_mask = 0b110;
   std::unordered_map<SectionId, Hit::TermFrequency>
       expected_section_ids_tf_map1 = {{1, 7}, {2, 11}};
-  ICING_ASSERT_OK(lite_index_->AddHit(foo_term_id, doc1_hit1));
-  ICING_ASSERT_OK(lite_index_->AddHit(foo_term_id, doc1_hit2));
+  ICING_ASSERT_OK(lite_index->AddHit(foo_term_id, doc1_hit1));
+  ICING_ASSERT_OK(lite_index->AddHit(foo_term_id, doc1_hit2));
 
   std::unique_ptr<DocHitInfoIteratorTermLiteExact> iter =
       std::make_unique<DocHitInfoIteratorTermLiteExact>(
-          term_id_codec_.get(), lite_index_.get(), term, /*term_start_index=*/0,
+          term_id_codec_.get(), lite_index.get(), term, /*term_start_index=*/0,
           /*unnormalized_term_length=*/0, kSectionIdMaskAll,
           /*need_hit_term_frequency=*/true);
 
diff --git a/icing/index/lite/lite-index_thread-safety_test.cc b/icing/index/lite/lite-index_thread-safety_test.cc
index 7711f92..53aa6cd 100644
--- a/icing/index/lite/lite-index_thread-safety_test.cc
+++ b/icing/index/lite/lite-index_thread-safety_test.cc
@@ -19,12 +19,9 @@
 
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
-#include "icing/index/lite/doc-hit-info-iterator-term-lite.h"
 #include "icing/index/lite/lite-index.h"
 #include "icing/index/term-id-codec.h"
 #include "icing/schema/section.h"
-#include "icing/store/suggestion-result-checker.h"
-#include "icing/testing/always-false-suggestion-result-checker-impl.h"
 #include "icing/testing/common-matchers.h"
 #include "icing/testing/tmp-directory.h"
 
@@ -52,7 +49,9 @@ class LiteIndexThreadSafetyTest : public testing::Test {
     std::string lite_index_file_name =
         index_dir_ + "/test_file.lite-idx-thread-safety.index";
     LiteIndex::Options options(lite_index_file_name,
-                               /*hit_buffer_want_merge_bytes=*/1024 * 1024);
+                               /*hit_buffer_want_merge_bytes=*/1024 * 1024,
+                               /*hit_buffer_sort_at_indexing=*/true,
+                               /*hit_buffer_sort_threshold_bytes=*/64);
     ICING_ASSERT_OK_AND_ASSIGN(lite_index_,
                                LiteIndex::Create(options, &icing_filesystem_));
 
diff --git a/icing/index/lite/term-id-hit-pair.h b/icing/index/lite/term-id-hit-pair.h
index 61ec502..82bd010 100644
--- a/icing/index/lite/term-id-hit-pair.h
+++ b/icing/index/lite/term-id-hit-pair.h
@@ -73,6 +73,8 @@ class TermIdHitPair {
     return value_ == rhs.value_;
   }
 
+  bool operator<(const TermIdHitPair& rhs) const { return value_ < rhs.value_; }
+
  private:
   Value value_;
 };
diff --git a/icing/index/main/main-index-merger_test.cc b/icing/index/main/main-index-merger_test.cc
index 8a2f691..37e14fc 100644
--- a/icing/index/main/main-index-merger_test.cc
+++ b/icing/index/main/main-index-merger_test.cc
@@ -45,7 +45,9 @@ class MainIndexMergerTest : public testing::Test {
 
     std::string lite_index_file_name = index_dir_ + "/test_file.lite-idx.index";
     LiteIndex::Options options(lite_index_file_name,
-                               /*hit_buffer_want_merge_bytes=*/1024 * 1024);
+                               /*hit_buffer_want_merge_bytes=*/1024 * 1024,
+                               /*hit_buffer_sort_at_indexing=*/true,
+                               /*hit_buffer_sort_threshold_bytes=*/1024 * 8);
     ICING_ASSERT_OK_AND_ASSIGN(lite_index_,
                                LiteIndex::Create(options, &icing_filesystem_));
 
diff --git a/icing/index/main/main-index_test.cc b/icing/index/main/main-index_test.cc
index c59a3c8..fa96e6c 100644
--- a/icing/index/main/main-index_test.cc
+++ b/icing/index/main/main-index_test.cc
@@ -91,7 +91,9 @@ class MainIndexTest : public testing::Test {
 
     std::string lite_index_file_name = index_dir_ + "/test_file.lite-idx.index";
     LiteIndex::Options options(lite_index_file_name,
-                               /*hit_buffer_want_merge_bytes=*/1024 * 1024);
+                               /*hit_buffer_want_merge_bytes=*/1024 * 1024,
+                               /*hit_buffer_sort_at_indexing=*/true,
+                               /*hit_buffer_sort_threshold_bytes=*/1024 * 8);
     ICING_ASSERT_OK_AND_ASSIGN(lite_index_,
                                LiteIndex::Create(options, &icing_filesystem_));
 
@@ -362,7 +364,9 @@ TEST_F(MainIndexTest, MergeIndexToPreexisting) {
   // - Doc4 {"four", "foul" is_in_prefix_section=true}
   std::string lite_index_file_name2 = index_dir_ + "/test_file.lite-idx.index2";
   LiteIndex::Options options(lite_index_file_name2,
-                             /*hit_buffer_want_merge_bytes=*/1024 * 1024);
+                             /*hit_buffer_want_merge_bytes=*/1024 * 1024,
+                             /*hit_buffer_sort_at_indexing=*/true,
+                             /*hit_buffer_sort_threshold_bytes=*/1024 * 8);
   ICING_ASSERT_OK_AND_ASSIGN(lite_index_,
                              LiteIndex::Create(options, &icing_filesystem_));
   ICING_ASSERT_OK_AND_ASSIGN(
@@ -632,7 +636,9 @@ TEST_F(MainIndexTest, MergeIndexBackfilling) {
   // - Doc1 {"foot" is_in_prefix_section=false}
   std::string lite_index_file_name2 = index_dir_ + "/test_file.lite-idx.index2";
   LiteIndex::Options options(lite_index_file_name2,
-                             /*hit_buffer_want_merge_bytes=*/1024 * 1024);
+                             /*hit_buffer_want_merge_bytes=*/1024 * 1024,
+                             /*hit_buffer_sort_at_indexing=*/true,
+                             /*hit_buffer_sort_threshold_bytes=*/1024 * 8);
   ICING_ASSERT_OK_AND_ASSIGN(lite_index_,
                              LiteIndex::Create(options, &icing_filesystem_));
   ICING_ASSERT_OK_AND_ASSIGN(
diff --git a/icing/index/string-section-indexing-handler.cc b/icing/index/string-section-indexing-handler.cc
index 69b8889..f5e06ad 100644
--- a/icing/index/string-section-indexing-handler.cc
+++ b/icing/index/string-section-indexing-handler.cc
@@ -122,6 +122,17 @@ libtextclassifier3::Status StringSectionIndexingHandler::Handle(
     }
   }
 
+  // Check and sort the LiteIndex HitBuffer if we're successful.
+  if (status.ok() && index_.LiteIndexNeedSort()) {
+    std::unique_ptr<Timer> sort_timer = clock_.GetNewTimer();
+    index_.SortLiteIndex();
+
+    if (put_document_stats != nullptr) {
+      put_document_stats->set_lite_index_sort_latency_ms(
+          sort_timer->GetElapsedMilliseconds());
+    }
+  }
+
   if (put_document_stats != nullptr) {
     put_document_stats->set_term_index_latency_ms(
         index_timer->GetElapsedMilliseconds());
diff --git a/icing/index/string-section-indexing-handler_test.cc b/icing/index/string-section-indexing-handler_test.cc
new file mode 100644
index 0000000..2c7f5e3
--- /dev/null
+++ b/icing/index/string-section-indexing-handler_test.cc
@@ -0,0 +1,587 @@
+// Copyright (C) 2023 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/string-section-indexing-handler.h"
+
+#include <cstdint>
+#include <limits>
+#include <memory>
+#include <string>
+#include <string_view>
+#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/document-builder.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/hit/hit.h"
+#include "icing/index/index.h"
+#include "icing/index/iterator/doc-hit-info-iterator-test-util.h"
+#include "icing/index/iterator/doc-hit-info-iterator.h"
+#include "icing/legacy/index/icing-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/proto/term.pb.h"
+#include "icing/schema-builder.h"
+#include "icing/schema/schema-store.h"
+#include "icing/schema/section.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/icu-data-file-helper.h"
+#include "icing/testing/test-data.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.h"
+#include "icing/util/tokenized-document.h"
+#include "unicode/uloc.h"
+
+namespace icing {
+namespace lib {
+
+namespace {
+
+using ::testing::ElementsAre;
+using ::testing::Eq;
+using ::testing::IsEmpty;
+using ::testing::IsFalse;
+using ::testing::IsTrue;
+using ::testing::Test;
+
+// Schema type with indexable properties and section Id.
+// Section Id is determined by the lexicographical order of indexable property
+// path.
+// Section id = 0: body
+// Section id = 1: title
+constexpr std::string_view kFakeType = "FakeType";
+constexpr std::string_view kPropertyBody = "body";
+constexpr std::string_view kPropertyTitle = "title";
+
+constexpr SectionId kSectionIdBody = 0;
+constexpr SectionId kSectionIdTitle = 1;
+
+// Schema type with nested indexable properties and section Id.
+// Section id = 0: "name"
+// Section id = 1: "nested.body"
+// Section id = 3: "nested.title"
+// Section id = 4: "subject"
+constexpr std::string_view kNestedType = "NestedType";
+constexpr std::string_view kPropertyName = "name";
+constexpr std::string_view kPropertyNestedDoc = "nested";
+constexpr std::string_view kPropertySubject = "subject";
+
+constexpr SectionId kSectionIdNestedBody = 1;
+
+class StringSectionIndexingHandlerTest : public Test {
+ protected:
+  void SetUp() override {
+    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")));
+    }
+
+    base_dir_ = GetTestTempDir() + "/icing_test";
+    ASSERT_THAT(filesystem_.CreateDirectoryRecursively(base_dir_.c_str()),
+                IsTrue());
+
+    index_dir_ = base_dir_ + "/index";
+    schema_store_dir_ = base_dir_ + "/schema_store";
+    document_store_dir_ = base_dir_ + "/document_store";
+
+    language_segmenter_factory::SegmenterOptions segmenter_options(ULOC_US);
+    ICING_ASSERT_OK_AND_ASSIGN(
+        lang_segmenter_,
+        language_segmenter_factory::Create(std::move(segmenter_options)));
+
+    ICING_ASSERT_OK_AND_ASSIGN(
+        normalizer_,
+        normalizer_factory::Create(
+            /*max_term_byte_size=*/std::numeric_limits<int32_t>::max()));
+
+    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_));
+    SchemaProto schema =
+        SchemaBuilder()
+            .AddType(
+                SchemaTypeConfigBuilder()
+                    .SetType(kFakeType)
+                    .AddProperty(PropertyConfigBuilder()
+                                     .SetName(kPropertyTitle)
+                                     .SetDataTypeString(TERM_MATCH_PREFIX,
+                                                        TOKENIZER_PLAIN)
+                                     .SetCardinality(CARDINALITY_OPTIONAL))
+                    .AddProperty(PropertyConfigBuilder()
+                                     .SetName(kPropertyBody)
+                                     .SetDataTypeString(TERM_MATCH_EXACT,
+                                                        TOKENIZER_PLAIN)
+                                     .SetCardinality(CARDINALITY_OPTIONAL)))
+            .AddType(
+                SchemaTypeConfigBuilder()
+                    .SetType(kNestedType)
+                    .AddProperty(
+                        PropertyConfigBuilder()
+                            .SetName(kPropertyNestedDoc)
+                            .SetDataTypeDocument(
+                                kFakeType, /*index_nested_properties=*/true)
+                            .SetCardinality(CARDINALITY_OPTIONAL))
+                    .AddProperty(PropertyConfigBuilder()
+                                     .SetName(kPropertySubject)
+                                     .SetDataTypeString(TERM_MATCH_EXACT,
+                                                        TOKENIZER_PLAIN)
+                                     .SetCardinality(CARDINALITY_OPTIONAL))
+                    .AddProperty(PropertyConfigBuilder()
+                                     .SetName(kPropertyName)
+                                     .SetDataTypeString(TERM_MATCH_EXACT,
+                                                        TOKENIZER_PLAIN)
+                                     .SetCardinality(CARDINALITY_OPTIONAL)))
+            .Build();
+    ICING_ASSERT_OK(schema_store_->SetSchema(
+        schema, /*ignore_errors_and_delete_documents=*/false,
+        /*allow_circular_schema_definitions=*/false));
+
+    ASSERT_TRUE(
+        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(),
+                              /*force_recovery_and_revalidate_documents=*/false,
+                              /*namespace_id_fingerprint=*/false,
+                              /*pre_mapping_fbv=*/false,
+                              /*use_persistent_hash_map=*/false,
+                              PortableFileBackedProtoLog<
+                                  DocumentWrapper>::kDeflateCompressionLevel,
+                              /*initialize_stats=*/nullptr));
+    document_store_ = std::move(doc_store_create_result.document_store);
+  }
+
+  void TearDown() override {
+    document_store_.reset();
+    schema_store_.reset();
+    normalizer_.reset();
+    lang_segmenter_.reset();
+
+    filesystem_.DeleteDirectoryRecursively(base_dir_.c_str());
+  }
+
+  Filesystem filesystem_;
+  IcingFilesystem icing_filesystem_;
+  FakeClock fake_clock_;
+  std::string base_dir_;
+  std::string index_dir_;
+  std::string schema_store_dir_;
+  std::string document_store_dir_;
+
+  std::unique_ptr<LanguageSegmenter> lang_segmenter_;
+  std::unique_ptr<Normalizer> normalizer_;
+  std::unique_ptr<SchemaStore> schema_store_;
+  std::unique_ptr<DocumentStore> document_store_;
+};
+
+std::vector<DocHitInfo> GetHits(std::unique_ptr<DocHitInfoIterator> iterator) {
+  std::vector<DocHitInfo> infos;
+  while (iterator->Advance().ok()) {
+    infos.push_back(iterator->doc_hit_info());
+  }
+  return infos;
+}
+
+std::vector<DocHitInfoTermFrequencyPair> GetHitsWithTermFrequency(
+    std::unique_ptr<DocHitInfoIterator> iterator) {
+  std::vector<DocHitInfoTermFrequencyPair> infos;
+  while (iterator->Advance().ok()) {
+    std::vector<TermMatchInfo> matched_terms_stats;
+    iterator->PopulateMatchedTermsStats(&matched_terms_stats);
+    for (const TermMatchInfo& term_match_info : matched_terms_stats) {
+      infos.push_back(DocHitInfoTermFrequencyPair(
+          iterator->doc_hit_info(), term_match_info.term_frequencies));
+    }
+  }
+  return infos;
+}
+
+TEST_F(StringSectionIndexingHandlerTest,
+       HandleIntoLiteIndex_sortInIndexingNotTriggered) {
+  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(
+      std::unique_ptr<Index> index,
+      Index::Create(options, &filesystem_, &icing_filesystem_));
+
+  DocumentProto document =
+      DocumentBuilder()
+          .SetKey("icing", "fake_type/1")
+          .SetSchema(std::string(kFakeType))
+          .AddStringProperty(std::string(kPropertyTitle), "foo")
+          .AddStringProperty(std::string(kPropertyBody), "foo bar baz")
+          .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_document,
+      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
+                                std::move(document)));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentId document_id,
+      document_store_->Put(tokenized_document.document()));
+
+  EXPECT_THAT(index->last_added_document_id(), Eq(kInvalidDocumentId));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<StringSectionIndexingHandler> handler,
+      StringSectionIndexingHandler::Create(&fake_clock_, normalizer_.get(),
+                                           index.get()));
+  EXPECT_THAT(
+      handler->Handle(tokenized_document, document_id, /*recovery_mode=*/false,
+                      /*put_document_stats=*/nullptr),
+      IsOk());
+
+  EXPECT_THAT(index->last_added_document_id(), Eq(document_id));
+
+  // Query 'foo'
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<DocHitInfoIterator> itr,
+      index->GetIterator("foo", /*term_start_index=*/0,
+                         /*unnormalized_term_length=*/0, kSectionIdMaskAll,
+                         TermMatchType::EXACT_ONLY));
+  std::vector<DocHitInfoTermFrequencyPair> hits =
+      GetHitsWithTermFrequency(std::move(itr));
+  std::unordered_map<SectionId, Hit::TermFrequency> expected_map{
+      {kSectionIdTitle, 1}, {kSectionIdBody, 1}};
+  EXPECT_THAT(hits, ElementsAre(EqualsDocHitInfoWithTermFrequency(
+                        document_id, expected_map)));
+
+  // Query 'foo' with sectionId mask that masks all results
+  ICING_ASSERT_OK_AND_ASSIGN(
+      itr, index->GetIterator("foo", /*term_start_index=*/0,
+                              /*unnormalized_term_length=*/0, 1U << 2,
+                              TermMatchType::EXACT_ONLY));
+  EXPECT_THAT(GetHits(std::move(itr)), IsEmpty());
+}
+
+TEST_F(StringSectionIndexingHandlerTest,
+       HandleIntoLiteIndex_sortInIndexingTriggered) {
+  // Create the LiteIndex with a smaller sort threshold. At 64 bytes we sort the
+  // HitBuffer after inserting 8 hits
+  Index::Options options(index_dir_,
+                         /*index_merge_size=*/1024 * 1024,
+                         /*lite_index_sort_at_indexing=*/true,
+                         /*lite_index_sort_size=*/64);
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<Index> index,
+      Index::Create(options, &filesystem_, &icing_filesystem_));
+
+  DocumentProto document0 =
+      DocumentBuilder()
+          .SetKey("icing", "fake_type/0")
+          .SetSchema(std::string(kFakeType))
+          .AddStringProperty(std::string(kPropertyTitle), "foo foo foo")
+          .AddStringProperty(std::string(kPropertyBody), "foo bar baz")
+          .Build();
+  DocumentProto document1 =
+      DocumentBuilder()
+          .SetKey("icing", "fake_type/1")
+          .SetSchema(std::string(kFakeType))
+          .AddStringProperty(std::string(kPropertyTitle), "bar baz baz")
+          .AddStringProperty(std::string(kPropertyBody), "foo foo baz")
+          .Build();
+  DocumentProto document2 =
+      DocumentBuilder()
+          .SetKey("icing", "nested_type/0")
+          .SetSchema(std::string(kNestedType))
+          .AddDocumentProperty(std::string(kPropertyNestedDoc), document1)
+          .AddStringProperty(std::string(kPropertyName), "qux")
+          .AddStringProperty(std::string(kPropertySubject), "bar bar")
+          .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_document0,
+      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
+                                std::move(document0)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentId document_id0,
+      document_store_->Put(tokenized_document0.document()));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_document1,
+      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
+                                std::move(document1)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentId document_id1,
+      document_store_->Put(tokenized_document1.document()));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_document2,
+      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
+                                std::move(document2)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentId document_id2,
+      document_store_->Put(tokenized_document2.document()));
+  EXPECT_THAT(index->last_added_document_id(), Eq(kInvalidDocumentId));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<StringSectionIndexingHandler> handler,
+      StringSectionIndexingHandler::Create(&fake_clock_, normalizer_.get(),
+                                           index.get()));
+
+  // Handle doc0 and doc1. The LiteIndex should sort and merge after adding
+  // these
+  EXPECT_THAT(handler->Handle(tokenized_document0, document_id0,
+                              /*recovery_mode=*/false,
+                              /*put_document_stats=*/nullptr),
+              IsOk());
+  EXPECT_THAT(handler->Handle(tokenized_document1, document_id1,
+                              /*recovery_mode=*/false,
+                              /*put_document_stats=*/nullptr),
+              IsOk());
+  EXPECT_THAT(index->last_added_document_id(), Eq(document_id1));
+  EXPECT_THAT(index->LiteIndexNeedSort(), IsFalse());
+
+  // Handle doc2. The LiteIndex should have an unsorted portion after adding
+  EXPECT_THAT(handler->Handle(tokenized_document2, document_id2,
+                              /*recovery_mode=*/false,
+                              /*put_document_stats=*/nullptr),
+              IsOk());
+  EXPECT_THAT(index->last_added_document_id(), Eq(document_id2));
+
+  // Hits in the hit buffer:
+  // <term>: {(docId, sectionId, term_freq)...}
+  // foo: {(0, kSectionIdTitle, 3); (0, kSectionIdBody, 1);
+  //       (1, kSectionIdBody, 2);
+  //       (2, kSectionIdNestedBody, 2)}
+  // bar: {(0, kSectionIdBody, 1);
+  //       (1, kSectionIdTitle, 1);
+  //       (2, kSectionIdNestedTitle, 1); (2, kSectionIdSubject, 2)}
+  // baz: {(0, kSectionIdBody, 1);
+  //       (1, kSectionIdTitle, 2); (1, kSectionIdBody, 1),
+  //       (2, kSectionIdNestedTitle, 2); (2, kSectionIdNestedBody, 1)}
+  // qux: {(2, kSectionIdName, 1)}
+
+  // Query 'foo'
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<DocHitInfoIterator> itr,
+      index->GetIterator("foo", /*term_start_index=*/0,
+                         /*unnormalized_term_length=*/0, kSectionIdMaskAll,
+                         TermMatchType::EXACT_ONLY));
+
+  // Advance the iterator and verify that we're returning hits in the correct
+  // order (i.e. in descending order of DocId)
+  ASSERT_THAT(itr->Advance(), IsOk());
+  EXPECT_THAT(itr->doc_hit_info().document_id(), Eq(2));
+  EXPECT_THAT(itr->doc_hit_info().hit_section_ids_mask(),
+              Eq(1U << kSectionIdNestedBody));
+  std::vector<TermMatchInfo> matched_terms_stats;
+  std::unordered_map<SectionId, Hit::TermFrequency>
+      expected_section_ids_tf_map2 = {{kSectionIdNestedBody, 2}};
+  itr->PopulateMatchedTermsStats(&matched_terms_stats);
+  EXPECT_THAT(matched_terms_stats, ElementsAre(EqualsTermMatchInfo(
+                                       "foo", expected_section_ids_tf_map2)));
+
+  ASSERT_THAT(itr->Advance(), IsOk());
+  EXPECT_THAT(itr->doc_hit_info().document_id(), Eq(1));
+  EXPECT_THAT(itr->doc_hit_info().hit_section_ids_mask(),
+              Eq(1U << kSectionIdBody));
+  std::unordered_map<SectionId, Hit::TermFrequency>
+      expected_section_ids_tf_map1 = {{kSectionIdBody, 2}};
+  matched_terms_stats.clear();
+  itr->PopulateMatchedTermsStats(&matched_terms_stats);
+  EXPECT_THAT(matched_terms_stats, ElementsAre(EqualsTermMatchInfo(
+                                       "foo", expected_section_ids_tf_map1)));
+
+  ASSERT_THAT(itr->Advance(), IsOk());
+  EXPECT_THAT(itr->doc_hit_info().document_id(), Eq(0));
+  EXPECT_THAT(itr->doc_hit_info().hit_section_ids_mask(),
+              Eq(1U << kSectionIdTitle | 1U << kSectionIdBody));
+  std::unordered_map<SectionId, Hit::TermFrequency>
+      expected_section_ids_tf_map0 = {{kSectionIdTitle, 3},
+                                      {kSectionIdBody, 1}};
+  matched_terms_stats.clear();
+  itr->PopulateMatchedTermsStats(&matched_terms_stats);
+  EXPECT_THAT(matched_terms_stats, ElementsAre(EqualsTermMatchInfo(
+                                       "foo", expected_section_ids_tf_map0)));
+}
+
+TEST_F(StringSectionIndexingHandlerTest,
+       HandleIntoLiteIndex_enableSortInIndexing) {
+  // Create the LiteIndex with a smaller sort threshold. At 64 bytes we sort the
+  // HitBuffer after inserting 8 hits
+  Index::Options options(index_dir_,
+                         /*index_merge_size=*/1024 * 1024,
+                         /*lite_index_sort_at_indexing=*/false,
+                         /*lite_index_sort_size=*/64);
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<Index> index,
+      Index::Create(options, &filesystem_, &icing_filesystem_));
+
+  DocumentProto document0 =
+      DocumentBuilder()
+          .SetKey("icing", "fake_type/0")
+          .SetSchema(std::string(kFakeType))
+          .AddStringProperty(std::string(kPropertyTitle), "foo foo foo")
+          .AddStringProperty(std::string(kPropertyBody), "foo bar baz")
+          .Build();
+  DocumentProto document1 =
+      DocumentBuilder()
+          .SetKey("icing", "fake_type/1")
+          .SetSchema(std::string(kFakeType))
+          .AddStringProperty(std::string(kPropertyTitle), "bar baz baz")
+          .AddStringProperty(std::string(kPropertyBody), "foo foo baz")
+          .Build();
+  DocumentProto document2 =
+      DocumentBuilder()
+          .SetKey("icing", "nested_type/0")
+          .SetSchema(std::string(kNestedType))
+          .AddDocumentProperty(std::string(kPropertyNestedDoc), document1)
+          .AddStringProperty(std::string(kPropertyName), "qux")
+          .AddStringProperty(std::string(kPropertySubject), "bar bar")
+          .Build();
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_document0,
+      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
+                                std::move(document0)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentId document_id0,
+      document_store_->Put(tokenized_document0.document()));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_document1,
+      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
+                                std::move(document1)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentId document_id1,
+      document_store_->Put(tokenized_document1.document()));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      TokenizedDocument tokenized_document2,
+      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
+                                std::move(document2)));
+  ICING_ASSERT_OK_AND_ASSIGN(
+      DocumentId document_id2,
+      document_store_->Put(tokenized_document2.document()));
+  EXPECT_THAT(index->last_added_document_id(), Eq(kInvalidDocumentId));
+
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<StringSectionIndexingHandler> handler,
+      StringSectionIndexingHandler::Create(&fake_clock_, normalizer_.get(),
+                                           index.get()));
+
+  // Handle all docs
+  EXPECT_THAT(handler->Handle(tokenized_document0, document_id0,
+                              /*recovery_mode=*/false,
+                              /*put_document_stats=*/nullptr),
+              IsOk());
+  EXPECT_THAT(handler->Handle(tokenized_document1, document_id1,
+                              /*recovery_mode=*/false,
+                              /*put_document_stats=*/nullptr),
+              IsOk());
+  EXPECT_THAT(handler->Handle(tokenized_document2, document_id2,
+                              /*recovery_mode=*/false,
+                              /*put_document_stats=*/nullptr),
+              IsOk());
+  EXPECT_THAT(index->last_added_document_id(), Eq(document_id2));
+
+  // We've disabled sorting during indexing so the HitBuffer's unsorted section
+  // should exceed the sort threshold. PersistToDisk and reinitialize the
+  // LiteIndex with sort_at_indexing=true.
+  ASSERT_THAT(index->PersistToDisk(), IsOk());
+  options = Index::Options(index_dir_,
+                           /*index_merge_size=*/1024 * 1024,
+                           /*lite_index_sort_at_indexing=*/true,
+                           /*lite_index_sort_size=*/64);
+  ICING_ASSERT_OK_AND_ASSIGN(
+      index, Index::Create(options, &filesystem_, &icing_filesystem_));
+
+  // Verify that the HitBuffer has been sorted after initializing with
+  // sort_at_indexing enabled.
+  EXPECT_THAT(index->LiteIndexNeedSort(), IsFalse());
+
+  // Hits in the hit buffer:
+  // <term>: {(docId, sectionId, term_freq)...}
+  // foo: {(0, kSectionIdTitle, 3); (0, kSectionIdBody, 1);
+  //       (1, kSectionIdBody, 2);
+  //       (2, kSectionIdNestedBody, 2)}
+  // bar: {(0, kSectionIdBody, 1);
+  //       (1, kSectionIdTitle, 1);
+  //       (2, kSectionIdNestedTitle, 1); (2, kSectionIdSubject, 2)}
+  // baz: {(0, kSectionIdBody, 1);
+  //       (1, kSectionIdTitle, 2); (1, kSectionIdBody, 1),
+  //       (2, kSectionIdNestedTitle, 2); (2, kSectionIdNestedBody, 1)}
+  // qux: {(2, kSectionIdName, 1)}
+
+  // Query 'foo'
+  ICING_ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<DocHitInfoIterator> itr,
+      index->GetIterator("foo", /*term_start_index=*/0,
+                         /*unnormalized_term_length=*/0, kSectionIdMaskAll,
+                         TermMatchType::EXACT_ONLY));
+
+  // Advance the iterator and verify that we're returning hits in the correct
+  // order (i.e. in descending order of DocId)
+  ASSERT_THAT(itr->Advance(), IsOk());
+  EXPECT_THAT(itr->doc_hit_info().document_id(), Eq(2));
+  EXPECT_THAT(itr->doc_hit_info().hit_section_ids_mask(),
+              Eq(1U << kSectionIdNestedBody));
+  std::vector<TermMatchInfo> matched_terms_stats;
+  std::unordered_map<SectionId, Hit::TermFrequency>
+      expected_section_ids_tf_map2 = {{kSectionIdNestedBody, 2}};
+  itr->PopulateMatchedTermsStats(&matched_terms_stats);
+  EXPECT_THAT(matched_terms_stats, ElementsAre(EqualsTermMatchInfo(
+                                       "foo", expected_section_ids_tf_map2)));
+
+  ASSERT_THAT(itr->Advance(), IsOk());
+  EXPECT_THAT(itr->doc_hit_info().document_id(), Eq(1));
+  EXPECT_THAT(itr->doc_hit_info().hit_section_ids_mask(),
+              Eq(1U << kSectionIdBody));
+  std::unordered_map<SectionId, Hit::TermFrequency>
+      expected_section_ids_tf_map1 = {{kSectionIdBody, 2}};
+  matched_terms_stats.clear();
+  itr->PopulateMatchedTermsStats(&matched_terms_stats);
+  EXPECT_THAT(matched_terms_stats, ElementsAre(EqualsTermMatchInfo(
+                                       "foo", expected_section_ids_tf_map1)));
+
+  ASSERT_THAT(itr->Advance(), IsOk());
+  EXPECT_THAT(itr->doc_hit_info().document_id(), Eq(0));
+  EXPECT_THAT(itr->doc_hit_info().hit_section_ids_mask(),
+              Eq(1U << kSectionIdTitle | 1U << kSectionIdBody));
+  std::unordered_map<SectionId, Hit::TermFrequency>
+      expected_section_ids_tf_map0 = {{kSectionIdTitle, 3},
+                                      {kSectionIdBody, 1}};
+  matched_terms_stats.clear();
+  itr->PopulateMatchedTermsStats(&matched_terms_stats);
+  EXPECT_THAT(matched_terms_stats, ElementsAre(EqualsTermMatchInfo(
+                                       "foo", expected_section_ids_tf_map0)));
+}
+
+}  // namespace
+
+}  // namespace lib
+}  // namespace icing
diff --git a/icing/query/advanced_query_parser/query-visitor_test.cc b/icing/query/advanced_query_parser/query-visitor_test.cc
index d118691..59e924d 100644
--- a/icing/query/advanced_query_parser/query-visitor_test.cc
+++ b/icing/query/advanced_query_parser/query-visitor_test.cc
@@ -125,7 +125,9 @@ class QueryVisitorTest : public ::testing::TestWithParam<QueryType> {
     document_store_ = std::move(create_result.document_store);
 
     Index::Options options(index_dir_.c_str(),
-                           /*index_merge_size=*/1024 * 1024);
+                           /*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_));
 
diff --git a/icing/query/query-processor_benchmark.cc b/icing/query/query-processor_benchmark.cc
index 626f352..025e8e6 100644
--- a/icing/query/query-processor_benchmark.cc
+++ b/icing/query/query-processor_benchmark.cc
@@ -81,7 +81,9 @@ void AddTokenToIndex(Index* index, DocumentId document_id, SectionId section_id,
 std::unique_ptr<Index> CreateIndex(const IcingFilesystem& icing_filesystem,
                                    const Filesystem& filesystem,
                                    const std::string& index_dir) {
-  Index::Options options(index_dir, /*index_merge_size=*/1024 * 1024 * 10);
+  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();
 }
 
diff --git a/icing/query/query-processor_test.cc b/icing/query/query-processor_test.cc
index fa70de1..e64de32 100644
--- a/icing/query/query-processor_test.cc
+++ b/icing/query/query-processor_test.cc
@@ -113,7 +113,9 @@ class QueryProcessorTest
     document_store_ = std::move(create_result.document_store);
 
     Index::Options options(index_dir_,
-                           /*index_merge_size=*/1024 * 1024);
+                           /*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_));
     // TODO(b/249829533): switch to use persistent numeric index.
diff --git a/icing/query/suggestion-processor_test.cc b/icing/query/suggestion-processor_test.cc
index 49f8d43..9f9094d 100644
--- a/icing/query/suggestion-processor_test.cc
+++ b/icing/query/suggestion-processor_test.cc
@@ -96,7 +96,9 @@ class SuggestionProcessorTest : public Test {
     document_store_ = std::move(create_result.document_store);
 
     Index::Options options(index_dir_,
-                           /*index_merge_size=*/1024 * 1024);
+                           /*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_));
     // TODO(b/249829533): switch to use persistent numeric index.
diff --git a/icing/schema/schema-util.cc b/icing/schema/schema-util.cc
index 4390b6a..af6feda 100644
--- a/icing/schema/schema-util.cc
+++ b/icing/schema/schema-util.cc
@@ -280,8 +280,6 @@ libtextclassifier3::Status CalculateTransitiveNestedTypeRelations(
     // 4. "adjacent" has been fully expanded. Add all of its transitive
     // outgoing relations to this type's transitive outgoing relations.
     auto adjacent_expanded_itr = expanded_nested_types_map->find(adjacent_type);
-    expanded_relations.reserve(expanded_relations.size() +
-                               adjacent_expanded_itr->second.size());
     for (const auto& [transitive_reachable, _] :
          adjacent_expanded_itr->second) {
       // Insert a transitive reachable node `transitive_reachable` for `type` if
@@ -345,8 +343,6 @@ libtextclassifier3::Status CalculateAcyclicTransitiveRelations(
     // 3. "adjacent" has been fully expanded. Add all of its transitive outgoing
     // relations to this type's transitive outgoing relations.
     auto adjacent_expanded_itr = expanded_relation_map->find(adjacent);
-    expanded_relations.reserve(expanded_relations.size() +
-                               adjacent_expanded_itr->second.size());
     for (const auto& [transitive_reachable, _] :
          adjacent_expanded_itr->second) {
       // Insert a transitive reachable node `transitive_reachable` for `type`.
@@ -526,7 +522,6 @@ BuildTransitiveDependentGraph(const SchemaProto& schema,
     // Insert the parent_type into the dependent map if it is not present
     // already.
     merged_dependent_map.insert({parent_type, {}});
-    merged_dependent_map[parent_type].reserve(inheritance_relation.size());
     for (const auto& [child_type, _] : inheritance_relation) {
       // Insert the child_type into parent_type's dependent map if it's not
       // present already, in which case the value will be an empty vector.
diff --git a/proto/icing/proto/initialize.proto b/proto/icing/proto/initialize.proto
index 507e745..958767b 100644
--- a/proto/icing/proto/initialize.proto
+++ b/proto/icing/proto/initialize.proto
@@ -23,7 +23,7 @@ option java_package = "com.google.android.icing.proto";
 option java_multiple_files = true;
 option objc_class_prefix = "ICNG";
 
-// Next tag: 12
+// Next tag: 14
 message IcingSearchEngineOptions {
   // Directory to persist files for Icing. Required.
   // If Icing was previously initialized with this directory, it will reload
@@ -107,6 +107,26 @@ message IcingSearchEngineOptions {
   // Integer index bucket split threshold.
   optional int32 integer_index_bucket_split_threshold = 11 [default = 65536];
 
+  // Whether Icing should sort and merge its lite index HitBuffer unsorted tail
+  // at indexing time.
+  //
+  // If set to true, the HitBuffer will be sorted at indexing time after
+  // exceeding the sort threshold. If false, the HifBuffer will be sorted at
+  // querying time, before the first query after inserting new elements into the
+  // HitBuffer.
+  //
+  // The default value is false.
+  optional bool lite_index_sort_at_indexing = 12;
+
+  // Size (in bytes) at which Icing's lite index should sort and merge the
+  // HitBuffer's unsorted tail into the sorted head for sorting at indexing
+  // time. Size specified here is the maximum byte size to allow for the
+  // unsorted tail section.
+  //
+  // Setting a lower sort size reduces querying latency at the expense of
+  // indexing latency.
+  optional int32 lite_index_sort_size = 13 [default = 8192];  // 8 KiB
+
   reserved 2;
 }
 
diff --git a/proto/icing/proto/logging.proto b/proto/icing/proto/logging.proto
index ca795cd..418fc88 100644
--- a/proto/icing/proto/logging.proto
+++ b/proto/icing/proto/logging.proto
@@ -76,7 +76,7 @@ message InitializeStatsProto {
   // Time used to restore the index.
   optional int32 index_restoration_latency_ms = 6;
 
-  // Time used to restore the index.
+  // Time used to restore the schema store.
   optional int32 schema_store_recovery_latency_ms = 7;
 
   // Status regarding how much data is lost during the initialization.
@@ -117,7 +117,7 @@ message InitializeStatsProto {
 }
 
 // Stats of the top-level function IcingSearchEngine::Put().
-// Next tag: 10
+// Next tag: 11
 message PutDocumentStatsProto {
   // Overall time used for the function call.
   optional int32 latency_ms = 1;
@@ -151,6 +151,9 @@ message PutDocumentStatsProto {
 
   // Time used to index all qualified id join strings in the document.
   optional int32 qualified_id_join_index_latency_ms = 9;
+
+  // Time used to sort and merge the LiteIndex's HitBuffer.
+  optional int32 lite_index_sort_latency_ms = 10;
 }
 
 // Stats of the top-level function IcingSearchEngine::Search() and
diff --git a/synced_AOSP_CL_number.txt b/synced_AOSP_CL_number.txt
index e7bee73..bd3f395 100644
--- a/synced_AOSP_CL_number.txt
+++ b/synced_AOSP_CL_number.txt
@@ -1 +1 @@
-set(synced_AOSP_CL_number=559533396)
+set(synced_AOSP_CL_number=561560020)