path: root/icing/index/embedding-indexing-handler_test.cc

// Copyright (C) 2024 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/embedding-indexing-handler.h"

#include <cstdint>
#include <initializer_list>
#include <memory>
#include <string>
#include <string_view>
#include <utility>
#include <vector>

#include "icing/text_classifier/lib3/utils/base/status.h"
#include "icing/text_classifier/lib3/utils/base/statusor.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "icing/absl_ports/canonical_errors.h"
#include "icing/document-builder.h"
#include "icing/file/filesystem.h"
#include "icing/file/portable-file-backed-proto-log.h"
#include "icing/index/embed/embedding-hit.h"
#include "icing/index/embed/embedding-index.h"
#include "icing/index/embed/posting-list-embedding-hit-accessor.h"
#include "icing/index/hit/hit.h"
#include "icing/portable/platform.h"
#include "icing/proto/document_wrapper.pb.h"
#include "icing/proto/schema.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/embedding-test-utils.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/util/status-macros.h"
#include "icing/util/tokenized-document.h"
#include "unicode/uloc.h"

namespace icing {
namespace lib {

namespace {

using ::testing::ElementsAre;
using ::testing::Eq;
using ::testing::IsEmpty;
using ::testing::IsTrue;

// Indexable properties (section) and section id. Section id is determined by
// the lexicographical order of indexable property paths.
// Schema type with indexable properties: FakeType
// Section id = 0: "body"
// Section id = 1: "bodyEmbedding"
// Section id = 2: "title"
// Section id = 3: "titleEmbedding"
static constexpr std::string_view kFakeType = "FakeType";
static constexpr std::string_view kPropertyBody = "body";
static constexpr std::string_view kPropertyBodyEmbedding = "bodyEmbedding";
static constexpr std::string_view kPropertyTitle = "title";
static constexpr std::string_view kPropertyTitleEmbedding = "titleEmbedding";
static constexpr std::string_view kPropertyNonIndexableEmbedding =
    "nonIndexableEmbedding";

static constexpr SectionId kSectionIdBodyEmbedding = 1;
static constexpr SectionId kSectionIdTitleEmbedding = 3;

// Schema type with nested indexable properties: FakeCollectionType
// Section id = 0: "collection.body"
// Section id = 1: "collection.bodyEmbedding"
// Section id = 2: "collection.title"
// Section id = 3: "collection.titleEmbedding"
// Section id = 4: "fullDocEmbedding"
static constexpr std::string_view kFakeCollectionType = "FakeCollectionType";
static constexpr std::string_view kPropertyCollection = "collection";
static constexpr std::string_view kPropertyFullDocEmbedding =
    "fullDocEmbedding";

static constexpr SectionId kSectionIdNestedBodyEmbedding = 1;
static constexpr SectionId kSectionIdNestedTitleEmbedding = 3;
static constexpr SectionId kSectionIdFullDocEmbedding = 4;

class EmbeddingIndexingHandlerTest : public ::testing::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());

    embedding_index_working_path_ = base_dir_ + "/embedding_index";
    schema_store_dir_ = base_dir_ + "/schema_store";
    document_store_dir_ = base_dir_ + "/document_store";

    ICING_ASSERT_OK_AND_ASSIGN(
        embedding_index_,
        EmbeddingIndex::Create(&filesystem_, embedding_index_working_path_));

    language_segmenter_factory::SegmenterOptions segmenter_options(ULOC_US);
    ICING_ASSERT_OK_AND_ASSIGN(
        lang_segmenter_,
        language_segmenter_factory::Create(std::move(segmenter_options)));

    ASSERT_THAT(
        filesystem_.CreateDirectoryRecursively(schema_store_dir_.c_str()),
        IsTrue());
    ICING_ASSERT_OK_AND_ASSIGN(
        schema_store_,
        SchemaStore::Create(&filesystem_, schema_store_dir_, &fake_clock_));
    SchemaProto schema =
        SchemaBuilder()
            .AddType(
                SchemaTypeConfigBuilder()
                    .SetType(kFakeType)
                    .AddProperty(PropertyConfigBuilder()
                                     .SetName(kPropertyTitle)
                                     .SetDataTypeString(TERM_MATCH_EXACT,
                                                        TOKENIZER_PLAIN)
                                     .SetCardinality(CARDINALITY_OPTIONAL))
                    .AddProperty(PropertyConfigBuilder()
                                     .SetName(kPropertyBody)
                                     .SetDataTypeString(TERM_MATCH_EXACT,
                                                        TOKENIZER_PLAIN)
                                     .SetCardinality(CARDINALITY_REPEATED))
                    .AddProperty(
                        PropertyConfigBuilder()
                            .SetName(kPropertyTitleEmbedding)
                            .SetDataTypeVector(
                                EmbeddingIndexingConfig::EmbeddingIndexingType::
                                    LINEAR_SEARCH)
                            .SetCardinality(CARDINALITY_OPTIONAL))
                    .AddProperty(
                        PropertyConfigBuilder()
                            .SetName(kPropertyBodyEmbedding)
                            .SetDataTypeVector(
                                EmbeddingIndexingConfig::EmbeddingIndexingType::
                                    LINEAR_SEARCH)
                            .SetCardinality(CARDINALITY_REPEATED))
                    .AddProperty(PropertyConfigBuilder()
                                     .SetName(kPropertyNonIndexableEmbedding)
                                     .SetDataType(TYPE_VECTOR)
                                     .SetCardinality(CARDINALITY_REPEATED)))
            .AddType(SchemaTypeConfigBuilder()
                         .SetType(kFakeCollectionType)
                         .AddProperty(PropertyConfigBuilder()
                                          .SetName(kPropertyCollection)
                                          .SetDataTypeDocument(
                                              kFakeType,
                                              /*index_nested_properties=*/true)
                                          .SetCardinality(CARDINALITY_REPEATED))
                         .AddProperty(
                             PropertyConfigBuilder()
                                 .SetName(kPropertyFullDocEmbedding)
                                 .SetDataTypeVector(
                                     EmbeddingIndexingConfig::
                                         EmbeddingIndexingType::LINEAR_SEARCH)
                                 .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();
    lang_segmenter_.reset();
    embedding_index_.reset();

    filesystem_.DeleteDirectoryRecursively(base_dir_.c_str());
  }

  libtextclassifier3::StatusOr<std::vector<EmbeddingHit>> GetHits(
      uint32_t dimension, std::string_view model_signature) {
    std::vector<EmbeddingHit> hits;

    libtextclassifier3::StatusOr<
        std::unique_ptr<PostingListEmbeddingHitAccessor>>
        pl_accessor_or =
            embedding_index_->GetAccessor(dimension, model_signature);
    std::unique_ptr<PostingListEmbeddingHitAccessor> pl_accessor;
    if (pl_accessor_or.ok()) {
      pl_accessor = std::move(pl_accessor_or).ValueOrDie();
    } else if (absl_ports::IsNotFound(pl_accessor_or.status())) {
      return hits;
    } else {
      return std::move(pl_accessor_or).status();
    }

    while (true) {
      ICING_ASSIGN_OR_RETURN(std::vector<EmbeddingHit> batch,
                             pl_accessor->GetNextHitsBatch());
      if (batch.empty()) {
        return hits;
      }
      hits.insert(hits.end(), batch.begin(), batch.end());
    }
  }

  std::vector<float> GetRawEmbeddingData() {
    return std::vector<float>(embedding_index_->GetRawEmbeddingData(),
                              embedding_index_->GetRawEmbeddingData() +
                                  embedding_index_->GetTotalVectorSize());
  }

  Filesystem filesystem_;
  FakeClock fake_clock_;
  std::string base_dir_;
  std::string embedding_index_working_path_;
  std::string schema_store_dir_;
  std::string document_store_dir_;

  std::unique_ptr<EmbeddingIndex> embedding_index_;
  std::unique_ptr<LanguageSegmenter> lang_segmenter_;
  std::unique_ptr<SchemaStore> schema_store_;
  std::unique_ptr<DocumentStore> document_store_;
};

}  // namespace

TEST_F(EmbeddingIndexingHandlerTest, CreationWithNullPointerShouldFail) {
  EXPECT_THAT(EmbeddingIndexingHandler::Create(/*clock=*/nullptr,
                                               embedding_index_.get()),
              StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));

  EXPECT_THAT(EmbeddingIndexingHandler::Create(&fake_clock_,
                                               /*embedding_index=*/nullptr),
              StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
}

TEST_F(EmbeddingIndexingHandlerTest, HandleEmbeddingSection) {
  DocumentProto document =
      DocumentBuilder()
          .SetKey("icing", "fake_type/1")
          .SetSchema(std::string(kFakeType))
          .AddStringProperty(std::string(kPropertyTitle), "title")
          .AddVectorProperty(std::string(kPropertyTitleEmbedding),
                             CreateVector("model", {0.1, 0.2, 0.3}))
          .AddStringProperty(std::string(kPropertyBody), "body")
          .AddVectorProperty(std::string(kPropertyBodyEmbedding),
                             CreateVector("model", {0.4, 0.5, 0.6}),
                             CreateVector("model", {0.7, 0.8, 0.9}))
          .AddVectorProperty(std::string(kPropertyNonIndexableEmbedding),
                             CreateVector("model", {1.1, 1.2, 1.3}))
          .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()));

  ASSERT_THAT(embedding_index_->last_added_document_id(),
              Eq(kInvalidDocumentId));
  // Handle document.
  ICING_ASSERT_OK_AND_ASSIGN(
      std::unique_ptr<EmbeddingIndexingHandler> handler,
      EmbeddingIndexingHandler::Create(&fake_clock_, embedding_index_.get()));
  EXPECT_THAT(
      handler->Handle(tokenized_document, document_id, /*recovery_mode=*/false,
                      /*put_document_stats=*/nullptr),
      IsOk());

  // Check index
  EXPECT_THAT(
      GetHits(/*dimension=*/3, /*model_signature=*/"model"),
      IsOkAndHolds(ElementsAre(
          EmbeddingHit(BasicHit(kSectionIdBodyEmbedding, /*document_id=*/0),
                       /*location=*/0),
          EmbeddingHit(BasicHit(kSectionIdBodyEmbedding, /*document_id=*/0),
                       /*location=*/3),
          EmbeddingHit(BasicHit(kSectionIdTitleEmbedding, /*document_id=*/0),
                       /*location=*/6))));
  EXPECT_THAT(GetRawEmbeddingData(),
              ElementsAre(0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.1, 0.2, 0.3));
  EXPECT_THAT(embedding_index_->last_added_document_id(), Eq(document_id));
}

TEST_F(EmbeddingIndexingHandlerTest, HandleNestedEmbeddingSection) {
  DocumentProto document =
      DocumentBuilder()
          .SetKey("icing", "fake_collection_type/1")
          .SetSchema(std::string(kFakeCollectionType))
          .AddDocumentProperty(
              std::string(kPropertyCollection),
              DocumentBuilder()
                  .SetKey("icing", "nested_fake_type/1")
                  .SetSchema(std::string(kFakeType))
                  .AddStringProperty(std::string(kPropertyTitle), "title")
                  .AddVectorProperty(std::string(kPropertyTitleEmbedding),
                                     CreateVector("model", {0.1, 0.2, 0.3}))
                  .AddStringProperty(std::string(kPropertyBody), "body")
                  .AddVectorProperty(std::string(kPropertyBodyEmbedding),
                                     CreateVector("model", {0.4, 0.5, 0.6}),
                                     CreateVector("model", {0.7, 0.8, 0.9}))
                  .AddVectorProperty(
                      std::string(kPropertyNonIndexableEmbedding),
                      CreateVector("model", {1.1, 1.2, 1.3}))
                  .Build())
          .AddVectorProperty(std::string(kPropertyFullDocEmbedding),
                             CreateVector("model", {2.1, 2.2, 2.3}))
          .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()));

  ASSERT_THAT(embedding_index_->last_added_document_id(),
              Eq(kInvalidDocumentId));
  // Handle document.
  ICING_ASSERT_OK_AND_ASSIGN(
      std::unique_ptr<EmbeddingIndexingHandler> handler,
      EmbeddingIndexingHandler::Create(&fake_clock_, embedding_index_.get()));
  EXPECT_THAT(
      handler->Handle(tokenized_document, document_id, /*recovery_mode=*/false,
                      /*put_document_stats=*/nullptr),
      IsOk());

  // Check index
  EXPECT_THAT(
      GetHits(/*dimension=*/3, /*model_signature=*/"model"),
      IsOkAndHolds(ElementsAre(
          EmbeddingHit(
              BasicHit(kSectionIdNestedBodyEmbedding, /*document_id=*/0),
              /*location=*/0),
          EmbeddingHit(
              BasicHit(kSectionIdNestedBodyEmbedding, /*document_id=*/0),
              /*location=*/3),
          EmbeddingHit(
              BasicHit(kSectionIdNestedTitleEmbedding, /*document_id=*/0),
              /*location=*/6),
          EmbeddingHit(BasicHit(kSectionIdFullDocEmbedding, /*document_id=*/0),
                       /*location=*/9))));
  EXPECT_THAT(GetRawEmbeddingData(), ElementsAre(0.4, 0.5, 0.6, 0.7, 0.8, 0.9,
                                                 0.1, 0.2, 0.3, 2.1, 2.2, 2.3));
  EXPECT_THAT(embedding_index_->last_added_document_id(), Eq(document_id));
}

TEST_F(EmbeddingIndexingHandlerTest,
       HandleInvalidDocumentIdShouldReturnInvalidArgumentError) {
  DocumentProto document =
      DocumentBuilder()
          .SetKey("icing", "fake_type/1")
          .SetSchema(std::string(kFakeType))
          .AddStringProperty(std::string(kPropertyTitle), "title")
          .AddVectorProperty(std::string(kPropertyTitleEmbedding),
                             CreateVector("model", {0.1, 0.2, 0.3}))
          .AddStringProperty(std::string(kPropertyBody), "body")
          .AddVectorProperty(std::string(kPropertyBodyEmbedding),
                             CreateVector("model", {0.4, 0.5, 0.6}),
                             CreateVector("model", {0.7, 0.8, 0.9}))
          .AddVectorProperty(std::string(kPropertyNonIndexableEmbedding),
                             CreateVector("model", {1.1, 1.2, 1.3}))
          .Build();
  ICING_ASSERT_OK_AND_ASSIGN(
      TokenizedDocument tokenized_document,
      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
                                std::move(document)));
  ICING_ASSERT_OK(document_store_->Put(tokenized_document.document()));

  static constexpr DocumentId kCurrentDocumentId = 3;
  embedding_index_->set_last_added_document_id(kCurrentDocumentId);
  ASSERT_THAT(embedding_index_->last_added_document_id(),
              Eq(kCurrentDocumentId));

  ICING_ASSERT_OK_AND_ASSIGN(
      std::unique_ptr<EmbeddingIndexingHandler> handler,
      EmbeddingIndexingHandler::Create(&fake_clock_, embedding_index_.get()));

  // Handling document with kInvalidDocumentId should cause a failure, and both
  // index data and last_added_document_id should remain unchanged.
  EXPECT_THAT(
      handler->Handle(tokenized_document, kInvalidDocumentId,
                      /*recovery_mode=*/false, /*put_document_stats=*/nullptr),
      StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
  EXPECT_THAT(embedding_index_->last_added_document_id(),
              Eq(kCurrentDocumentId));
  // Check that the embedding index should be empty
  EXPECT_THAT(GetHits(/*dimension=*/3, /*model_signature=*/"model"),
              IsOkAndHolds(IsEmpty()));
  EXPECT_THAT(GetRawEmbeddingData(), IsEmpty());

  // Recovery mode should get the same result.
  EXPECT_THAT(
      handler->Handle(tokenized_document, kInvalidDocumentId,
                      /*recovery_mode=*/true, /*put_document_stats=*/nullptr),
      StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
  EXPECT_THAT(embedding_index_->last_added_document_id(),
              Eq(kCurrentDocumentId));
  // Check that the embedding index should be empty
  EXPECT_THAT(GetHits(/*dimension=*/3, /*model_signature=*/"model"),
              IsOkAndHolds(IsEmpty()));
  EXPECT_THAT(GetRawEmbeddingData(), IsEmpty());
}

TEST_F(EmbeddingIndexingHandlerTest,
       HandleOutOfOrderDocumentIdShouldReturnInvalidArgumentError) {
  DocumentProto document =
      DocumentBuilder()
          .SetKey("icing", "fake_type/1")
          .SetSchema(std::string(kFakeType))
          .AddStringProperty(std::string(kPropertyTitle), "title")
          .AddVectorProperty(std::string(kPropertyTitleEmbedding),
                             CreateVector("model", {0.1, 0.2, 0.3}))
          .AddStringProperty(std::string(kPropertyBody), "body")
          .AddVectorProperty(std::string(kPropertyBodyEmbedding),
                             CreateVector("model", {0.4, 0.5, 0.6}),
                             CreateVector("model", {0.7, 0.8, 0.9}))
          .AddVectorProperty(std::string(kPropertyNonIndexableEmbedding),
                             CreateVector("model", {1.1, 1.2, 1.3}))
          .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()));

  ICING_ASSERT_OK_AND_ASSIGN(
      std::unique_ptr<EmbeddingIndexingHandler> handler,
      EmbeddingIndexingHandler::Create(&fake_clock_, embedding_index_.get()));

  // Handling document with document_id == last_added_document_id should cause a
  // failure, and both index data and last_added_document_id should remain
  // unchanged.
  embedding_index_->set_last_added_document_id(document_id);
  ASSERT_THAT(embedding_index_->last_added_document_id(), Eq(document_id));
  EXPECT_THAT(
      handler->Handle(tokenized_document, document_id, /*recovery_mode=*/false,
                      /*put_document_stats=*/nullptr),
      StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
  EXPECT_THAT(embedding_index_->last_added_document_id(), Eq(document_id));

  // Check that the embedding index should be empty
  EXPECT_THAT(GetHits(/*dimension=*/3, /*model_signature=*/"model"),
              IsOkAndHolds(IsEmpty()));
  EXPECT_THAT(GetRawEmbeddingData(), IsEmpty());

  // Handling document with document_id < last_added_document_id should cause a
  // failure, and both index data and last_added_document_id should remain
  // unchanged.
  embedding_index_->set_last_added_document_id(document_id + 1);
  ASSERT_THAT(embedding_index_->last_added_document_id(), Eq(document_id + 1));
  EXPECT_THAT(
      handler->Handle(tokenized_document, document_id, /*recovery_mode=*/false,
                      /*put_document_stats=*/nullptr),
      StatusIs(libtextclassifier3::StatusCode::INVALID_ARGUMENT));
  EXPECT_THAT(embedding_index_->last_added_document_id(), Eq(document_id + 1));

  // Check that the embedding index should be empty
  EXPECT_THAT(GetHits(/*dimension=*/3, /*model_signature=*/"model"),
              IsOkAndHolds(IsEmpty()));
  EXPECT_THAT(GetRawEmbeddingData(), IsEmpty());
}

TEST_F(EmbeddingIndexingHandlerTest,
       HandleRecoveryModeShouldIgnoreDocsLELastAddedDocId) {
  DocumentProto document1 =
      DocumentBuilder()
          .SetKey("icing", "fake_type/1")
          .SetSchema(std::string(kFakeType))
          .AddStringProperty(std::string(kPropertyTitle), "title one")
          .AddVectorProperty(std::string(kPropertyTitleEmbedding),
                             CreateVector("model", {0.1, 0.2, 0.3}))
          .AddStringProperty(std::string(kPropertyBody), "body one")
          .AddVectorProperty(std::string(kPropertyBodyEmbedding),
                             CreateVector("model", {0.4, 0.5, 0.6}),
                             CreateVector("model", {0.7, 0.8, 0.9}))
          .AddVectorProperty(std::string(kPropertyNonIndexableEmbedding),
                             CreateVector("model", {1.1, 1.2, 1.3}))
          .Build();
  DocumentProto document2 =
      DocumentBuilder()
          .SetKey("icing", "fake_type/2")
          .SetSchema(std::string(kFakeType))
          .AddStringProperty(std::string(kPropertyTitle), "title two")
          .AddVectorProperty(std::string(kPropertyTitleEmbedding),
                             CreateVector("model", {10.1, 10.2, 10.3}))
          .AddStringProperty(std::string(kPropertyBody), "body two")
          .AddVectorProperty(std::string(kPropertyBodyEmbedding),
                             CreateVector("model", {10.4, 10.5, 10.6}),
                             CreateVector("model", {10.7, 10.8, 10.9}))
          .AddVectorProperty(std::string(kPropertyNonIndexableEmbedding),
                             CreateVector("model", {11.1, 11.2, 11.3}))
          .Build();
  ICING_ASSERT_OK_AND_ASSIGN(
      TokenizedDocument tokenized_document1,
      TokenizedDocument::Create(schema_store_.get(), lang_segmenter_.get(),
                                std::move(document1)));
  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_id1,
      document_store_->Put(tokenized_document1.document()));
  ICING_ASSERT_OK_AND_ASSIGN(
      DocumentId document_id2,
      document_store_->Put(tokenized_document2.document()));

  ICING_ASSERT_OK_AND_ASSIGN(
      std::unique_ptr<EmbeddingIndexingHandler> handler,
      EmbeddingIndexingHandler::Create(&fake_clock_, embedding_index_.get()));

  // Handle document with document_id > last_added_document_id in recovery mode.
  // The handler should index this document and update last_added_document_id.
  EXPECT_THAT(
      handler->Handle(tokenized_document1, document_id1, /*recovery_mode=*/true,
                      /*put_document_stats=*/nullptr),
      IsOk());
  EXPECT_THAT(embedding_index_->last_added_document_id(), Eq(document_id1));

  // Check index
  EXPECT_THAT(
      GetHits(/*dimension=*/3, /*model_signature=*/"model"),
      IsOkAndHolds(ElementsAre(
          EmbeddingHit(BasicHit(kSectionIdBodyEmbedding, /*document_id=*/0),
                       /*location=*/0),
          EmbeddingHit(BasicHit(kSectionIdBodyEmbedding, /*document_id=*/0),
                       /*location=*/3),
          EmbeddingHit(BasicHit(kSectionIdTitleEmbedding, /*document_id=*/0),
                       /*location=*/6))));
  EXPECT_THAT(GetRawEmbeddingData(),
              ElementsAre(0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.1, 0.2, 0.3));

  // Handle document with document_id == last_added_document_id in recovery
  // mode. We should not get any error, but the handler should ignore the
  // document, so both index data and last_added_document_id should remain
  // unchanged.
  embedding_index_->set_last_added_document_id(document_id2);
  ASSERT_THAT(embedding_index_->last_added_document_id(), Eq(document_id2));
  EXPECT_THAT(
      handler->Handle(tokenized_document2, document_id2, /*recovery_mode=*/true,
                      /*put_document_stats=*/nullptr),
      IsOk());
  EXPECT_THAT(embedding_index_->last_added_document_id(), Eq(document_id2));

  // Check index
  EXPECT_THAT(
      GetHits(/*dimension=*/3, /*model_signature=*/"model"),
      IsOkAndHolds(ElementsAre(
          EmbeddingHit(BasicHit(kSectionIdBodyEmbedding, /*document_id=*/0),
                       /*location=*/0),
          EmbeddingHit(BasicHit(kSectionIdBodyEmbedding, /*document_id=*/0),
                       /*location=*/3),
          EmbeddingHit(BasicHit(kSectionIdTitleEmbedding, /*document_id=*/0),
                       /*location=*/6))));
  EXPECT_THAT(GetRawEmbeddingData(),
              ElementsAre(0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.1, 0.2, 0.3));

  // Handle document with document_id < last_added_document_id in recovery mode.
  // We should not get any error, but the handler should ignore the document, so
  // both index data and last_added_document_id should remain unchanged.
  embedding_index_->set_last_added_document_id(document_id2 + 1);
  ASSERT_THAT(embedding_index_->last_added_document_id(), Eq(document_id2 + 1));
  EXPECT_THAT(
      handler->Handle(tokenized_document2, document_id2, /*recovery_mode=*/true,
                      /*put_document_stats=*/nullptr),
      IsOk());
  EXPECT_THAT(embedding_index_->last_added_document_id(), Eq(document_id2 + 1));

  // Check index
  EXPECT_THAT(
      GetHits(/*dimension=*/3, /*model_signature=*/"model"),
      IsOkAndHolds(ElementsAre(
          EmbeddingHit(BasicHit(kSectionIdBodyEmbedding, /*document_id=*/0),
                       /*location=*/0),
          EmbeddingHit(BasicHit(kSectionIdBodyEmbedding, /*document_id=*/0),
                       /*location=*/3),
          EmbeddingHit(BasicHit(kSectionIdTitleEmbedding, /*document_id=*/0),
                       /*location=*/6))));
  EXPECT_THAT(GetRawEmbeddingData(),
              ElementsAre(0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.1, 0.2, 0.3));
}

}  // namespace lib
}  // namespace icing