xref: /external/icing/icing/index/embed/embedding-index.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/embed/embedding-index.h"

#include <algorithm>
#include <cstdint>
#include <cstring>
#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 "icing/absl_ports/canonical_errors.h"
#include "icing/absl_ports/str_cat.h"
#include "icing/feature-flags.h"
#include "icing/file/destructible-directory.h"
#include "icing/file/file-backed-vector.h"
#include "icing/file/filesystem.h"
#include "icing/file/memory-mapped-file.h"
#include "icing/file/posting_list/flash-index-storage.h"
#include "icing/file/posting_list/posting-list-identifier.h"
#include "icing/index/embed/embedding-hit.h"
#include "icing/index/embed/embedding-scorer.h"
#include "icing/index/embed/posting-list-embedding-hit-accessor.h"
#include "icing/index/embed/quantizer.h"
#include "icing/index/hit/hit.h"
#include "icing/schema/schema-store.h"
#include "icing/store/document-filter-data.h"
#include "icing/store/document-id.h"
#include "icing/store/document-store.h"
#include "icing/store/dynamic-trie-key-mapper.h"
#include "icing/store/key-mapper.h"
#include "icing/util/clock.h"
#include "icing/util/crc32.h"
#include "icing/util/encode-util.h"
#include "icing/util/logging.h"
#include "icing/util/status-macros.h"

namespace icing {
namespace lib {

namespace {

// The maximum size of the embedding hit list mmapper.
// We use 64MiB for 32-bit platforms and 128MiB for 64-bit platforms.
#ifdef ICING_ARCH_BIT_64
  constexpr uint32_t kEmbeddingHitListMapperMaxSize = 128 * 1024 * 1024;
#else
  constexpr uint32_t kEmbeddingHitListMapperMaxSize = 64 * 1024 * 1024;
#endif

// The maximum length returned by encode_util::EncodeIntToCString is 5 for
// uint32_t.
constexpr uint32_t kEncodedDimensionLength = 5;

std::string GetMetadataFilePath(std::string_view working_path) {
  return absl_ports::StrCat(working_path, "/metadata");
}

std::string GetFlashIndexStorageFilePath(std::string_view working_path) {
  return absl_ports::StrCat(working_path, "/flash_index_storage");
}

std::string GetEmbeddingHitListMapperPath(std::string_view working_path) {
  return absl_ports::StrCat(working_path, "/embedding_hit_list_mapper");
}

std::string GetEmbeddingVectorsFilePath(std::string_view working_path) {
  return absl_ports::StrCat(working_path, "/embedding_vectors");
}

std::string GetQuantizedEmbeddingVectorsFilePath(
    std::string_view working_path) {
  return absl_ports::StrCat(working_path, "/quantized_embedding_vectors");
}

// An injective function that maps the ordered pair (dimension, model_signature)
// to a string, which is used to form a key for embedding_posting_list_mapper_.
std::string GetPostingListKey(uint32_t dimension,
                              std::string_view model_signature) {
  std::string encoded_dimension_str =
      encode_util::EncodeIntToCString(dimension);
  // Make encoded_dimension_str to fixed kEncodedDimensionLength bytes.
  while (encoded_dimension_str.size() < kEncodedDimensionLength) {
    // C string cannot contain 0 bytes, so we append it using 1, just like what
    // we do in encode_util::EncodeIntToCString.
    //
    // The reason that this works is because DecodeIntToString decodes a byte
    // value of 0x01 as 0x00. When EncodeIntToCString returns an encoded
    // dimension that is less than 5 bytes, it means that the dimension contains
    // unencoded leading 0x00. So here we're explicitly encoding those bytes as
    // 0x01.
    encoded_dimension_str.push_back(1);
  }
  return absl_ports::StrCat(encoded_dimension_str, model_signature);
}

std::string GetPostingListKey(const PropertyProto::VectorProto& vector) {
  return GetPostingListKey(vector.values().size(), vector.model_signature());
}

libtextclassifier3::StatusOr<Quantizer> CreateQuantizer(
    const PropertyProto::VectorProto& vector) {
  if (vector.values().empty()) {
    return absl_ports::InvalidArgumentError("Vector dimension is 0");
  }
  auto minmax_pair =
      std::minmax_element(vector.values().begin(), vector.values().end());
  return Quantizer::Create(*minmax_pair.first, *minmax_pair.second);
}

}  // namespace

libtextclassifier3::StatusOr<std::unique_ptr<EmbeddingIndex>>
EmbeddingIndex::Create(const Filesystem* filesystem, std::string working_path,
                       const Clock* clock, const FeatureFlags* feature_flags) {
  ICING_RETURN_ERROR_IF_NULL(filesystem);
  ICING_RETURN_ERROR_IF_NULL(clock);

  std::unique_ptr<EmbeddingIndex> index =
      std::unique_ptr<EmbeddingIndex>(new EmbeddingIndex(
          *filesystem, std::move(working_path), clock, feature_flags));
  ICING_RETURN_IF_ERROR(index->Initialize());
  return index;
}

libtextclassifier3::Status EmbeddingIndex::CreateStorageDataIfNonEmpty() {
  if (is_empty()) {
    return libtextclassifier3::Status::OK;
  }

  ICING_ASSIGN_OR_RETURN(FlashIndexStorage flash_index_storage,
                         FlashIndexStorage::Create(
                             GetFlashIndexStorageFilePath(working_path_),
                             &filesystem_, posting_list_hit_serializer_.get()));
  flash_index_storage_ =
      std::make_unique<FlashIndexStorage>(std::move(flash_index_storage));

  ICING_ASSIGN_OR_RETURN(
      embedding_posting_list_mapper_,
      DynamicTrieKeyMapper<PostingListIdentifier>::Create(
          filesystem_, GetEmbeddingHitListMapperPath(working_path_),
          kEmbeddingHitListMapperMaxSize));

  ICING_ASSIGN_OR_RETURN(
      embedding_vectors_,
      FileBackedVector<float>::Create(
          filesystem_, GetEmbeddingVectorsFilePath(working_path_),
          MemoryMappedFile::READ_WRITE_AUTO_SYNC));

  ICING_ASSIGN_OR_RETURN(
      quantized_embedding_vectors_,
      FileBackedVector<char>::Create(
          filesystem_, GetQuantizedEmbeddingVectorsFilePath(working_path_),
          MemoryMappedFile::READ_WRITE_AUTO_SYNC));

  return libtextclassifier3::Status::OK;
}

libtextclassifier3::Status EmbeddingIndex::MarkIndexNonEmpty() {
  if (!is_empty()) {
    return libtextclassifier3::Status::OK;
  }
  info().is_empty = false;
  return CreateStorageDataIfNonEmpty();
}

libtextclassifier3::Status EmbeddingIndex::Initialize() {
  bool is_new = false;
  if (!filesystem_.FileExists(GetMetadataFilePath(working_path_).c_str())) {
    // Create working directory.
    if (!filesystem_.CreateDirectoryRecursively(working_path_.c_str())) {
      return absl_ports::InternalError(
          absl_ports::StrCat("Failed to create directory: ", working_path_));
    }
    is_new = true;
  }

  ICING_ASSIGN_OR_RETURN(
      MemoryMappedFile metadata_mmapped_file,
      MemoryMappedFile::Create(filesystem_, GetMetadataFilePath(working_path_),
                               MemoryMappedFile::Strategy::READ_WRITE_AUTO_SYNC,
                               /*max_file_size=*/kMetadataFileSize,
                               /*pre_mapping_file_offset=*/0,
                               /*pre_mapping_mmap_size=*/kMetadataFileSize));
  metadata_mmapped_file_ =
      std::make_unique<MemoryMappedFile>(std::move(metadata_mmapped_file));

  if (is_new) {
    ICING_RETURN_IF_ERROR(metadata_mmapped_file_->GrowAndRemapIfNecessary(
        /*file_offset=*/0, /*mmap_size=*/kMetadataFileSize));
    info().magic = Info::kMagic;
    info().last_added_document_id = kInvalidDocumentId;
    info().is_empty = true;
    memset(Info().padding_, 0, Info::kPaddingSize);
    ICING_RETURN_IF_ERROR(InitializeNewStorage());
  } else {
    if (metadata_mmapped_file_->available_size() != kMetadataFileSize) {
      return absl_ports::FailedPreconditionError(
          "Incorrect metadata file size");
    }
    if (info().magic != Info::kMagic) {
      return absl_ports::FailedPreconditionError("Incorrect magic value");
    }
    ICING_RETURN_IF_ERROR(CreateStorageDataIfNonEmpty());
    ICING_RETURN_IF_ERROR(InitializeExistingStorage());
  }
  return libtextclassifier3::Status::OK;
}

libtextclassifier3::Status EmbeddingIndex::Clear() {
  pending_embedding_hits_.clear();
  metadata_mmapped_file_.reset();
  flash_index_storage_.reset();
  embedding_posting_list_mapper_.reset();
  embedding_vectors_.reset();
  quantized_embedding_vectors_.reset();
  if (filesystem_.DirectoryExists(working_path_.c_str())) {
    ICING_RETURN_IF_ERROR(Discard(filesystem_, working_path_));
  }
  is_initialized_ = false;
  return Initialize();
}

libtextclassifier3::StatusOr<std::unique_ptr<PostingListEmbeddingHitAccessor>>
EmbeddingIndex::GetAccessor(uint32_t dimension,
                            std::string_view model_signature) const {
  if (dimension == 0) {
    return absl_ports::InvalidArgumentError("Dimension is 0");
  }
  if (is_empty()) {
    return absl_ports::NotFoundError("EmbeddingIndex is empty");
  }

  std::string key = GetPostingListKey(dimension, model_signature);
  ICING_ASSIGN_OR_RETURN(PostingListIdentifier posting_list_id,
                         embedding_posting_list_mapper_->Get(key));
  return PostingListEmbeddingHitAccessor::CreateFromExisting(
      flash_index_storage_.get(), posting_list_hit_serializer_.get(),
      posting_list_id);
}

libtextclassifier3::StatusOr<uint32_t> EmbeddingIndex::AppendEmbeddingVector(
    const PropertyProto::VectorProto& vector,
    EmbeddingIndexingConfig::QuantizationType::Code quantization_type) {
  uint32_t dimension = vector.values().size();
  uint32_t location;
  if (!feature_flags_->enable_embedding_quantization() ||
      quantization_type == EmbeddingIndexingConfig::QuantizationType::NONE) {
    location = embedding_vectors_->num_elements();
    ICING_ASSIGN_OR_RETURN(
        FileBackedVector<float>::MutableArrayView mutable_arr,
        embedding_vectors_->Allocate(dimension));
    mutable_arr.SetArray(/*idx=*/0, vector.values().data(), dimension);
  } else {
    ICING_ASSIGN_OR_RETURN(Quantizer quantizer, CreateQuantizer(vector));
    // Quantize the vector
    std::vector<uint8_t> quantized_values;
    quantized_values.reserve(vector.values().size());
    for (float value : vector.values()) {
      quantized_values.push_back(quantizer.Quantize(value));
    }

    // Store the quantizer and the quantized vector
    location = quantized_embedding_vectors_->num_elements();
    ICING_ASSIGN_OR_RETURN(
        FileBackedVector<char>::MutableArrayView mutable_arr,
        quantized_embedding_vectors_->Allocate(sizeof(Quantizer) + dimension));
    mutable_arr.SetArray(/*idx=*/0, reinterpret_cast<char*>(&quantizer),
                         sizeof(Quantizer));
    mutable_arr.SetArray(/*idx=*/sizeof(Quantizer),
                         reinterpret_cast<char*>(quantized_values.data()),
                         dimension);
  }
  return location;
}

libtextclassifier3::Status EmbeddingIndex::BufferEmbedding(
    const BasicHit& basic_hit, const PropertyProto::VectorProto& vector,
    EmbeddingIndexingConfig::QuantizationType::Code quantization_type) {
  if (vector.values().empty()) {
    return absl_ports::InvalidArgumentError("Vector dimension is 0");
  }
  ICING_RETURN_IF_ERROR(MarkIndexNonEmpty());

  std::string key = GetPostingListKey(vector);
  ICING_ASSIGN_OR_RETURN(uint32_t location,
                         AppendEmbeddingVector(vector, quantization_type));

  // Buffer the embedding hit.
  pending_embedding_hits_.push_back(
      {std::move(key), EmbeddingHit(basic_hit, location)});
  return libtextclassifier3::Status::OK;
}

libtextclassifier3::Status EmbeddingIndex::CommitBufferToIndex() {
  if (pending_embedding_hits_.empty()) {
    return libtextclassifier3::Status::OK;
  }
  ICING_RETURN_IF_ERROR(MarkIndexNonEmpty());

  std::sort(pending_embedding_hits_.begin(), pending_embedding_hits_.end());
  auto iter_curr_key = pending_embedding_hits_.rbegin();
  while (iter_curr_key != pending_embedding_hits_.rend()) {
    // In order to batch putting embedding hits with the same key (dimension,
    // model_signature) to the same posting list, we find the range
    // [iter_curr_key, iter_next_key) of embedding hits with the same key and
    // put them into their corresponding posting list together.
    auto iter_next_key = iter_curr_key;
    while (iter_next_key != pending_embedding_hits_.rend() &&
           iter_next_key->first == iter_curr_key->first) {
      iter_next_key++;
    }

    const std::string& key = iter_curr_key->first;
    libtextclassifier3::StatusOr<PostingListIdentifier> posting_list_id_or =
        embedding_posting_list_mapper_->Get(key);
    std::unique_ptr<PostingListEmbeddingHitAccessor> pl_accessor;
    if (posting_list_id_or.ok()) {
      // Existing posting list.
      ICING_ASSIGN_OR_RETURN(
          pl_accessor,
          PostingListEmbeddingHitAccessor::CreateFromExisting(
              flash_index_storage_.get(), posting_list_hit_serializer_.get(),
              posting_list_id_or.ValueOrDie()));
    } else if (absl_ports::IsNotFound(posting_list_id_or.status())) {
      // New posting list.
      ICING_ASSIGN_OR_RETURN(
          pl_accessor,
          PostingListEmbeddingHitAccessor::Create(
              flash_index_storage_.get(), posting_list_hit_serializer_.get()));
    } else {
      // Errors
      return std::move(posting_list_id_or).status();
    }

    // Adding the embedding hits.
    for (auto iter = iter_curr_key; iter != iter_next_key; ++iter) {
      ICING_RETURN_IF_ERROR(pl_accessor->PrependHit(iter->second));
    }

    // Finalize this posting list and add the posting list id in
    // embedding_posting_list_mapper_.
    PostingListEmbeddingHitAccessor::FinalizeResult result =
        std::move(*pl_accessor).Finalize();
    if (!result.id.is_valid()) {
      return absl_ports::InternalError("Failed to finalize posting list");
    }
    ICING_RETURN_IF_ERROR(embedding_posting_list_mapper_->Put(key, result.id));

    // Advance to the next key.
    iter_curr_key = iter_next_key;
  }
  pending_embedding_hits_.clear();
  return libtextclassifier3::Status::OK;
}

libtextclassifier3::StatusOr<uint32_t> EmbeddingIndex::TransferEmbeddingVector(
    const EmbeddingHit& old_hit, uint32_t dimension,
    EmbeddingIndexingConfig::QuantizationType::Code quantization_type,
    EmbeddingIndex* new_index) const {
  uint32_t new_location;
  if (!feature_flags_->enable_embedding_quantization() ||
      quantization_type == EmbeddingIndexingConfig::QuantizationType::NONE) {
    ICING_ASSIGN_OR_RETURN(const float* old_vector,
                           GetEmbeddingVector(old_hit, dimension));
    new_location = new_index->embedding_vectors_->num_elements();

    // Copy the embedding vector of the hit to the new index.
    ICING_ASSIGN_OR_RETURN(
        FileBackedVector<float>::MutableArrayView mutable_arr,
        new_index->embedding_vectors_->Allocate(dimension));
    mutable_arr.SetArray(/*idx=*/0, old_vector, dimension);
  } else {
    ICING_ASSIGN_OR_RETURN(const char* old_data,
                           GetQuantizedEmbeddingVector(old_hit, dimension));
    new_location = new_index->quantized_embedding_vectors_->num_elements();

    // Copy the embedding vector of the hit to the new index.
    ICING_ASSIGN_OR_RETURN(FileBackedVector<char>::MutableArrayView mutable_arr,
                           new_index->quantized_embedding_vectors_->Allocate(
                               sizeof(Quantizer) + dimension));
    mutable_arr.SetArray(/*idx=*/0, old_data, sizeof(Quantizer) + dimension);
  }
  return new_location;
}

libtextclassifier3::Status EmbeddingIndex::TransferIndex(
    const DocumentStore& document_store, const SchemaStore& schema_store,
    const std::vector<DocumentId>& document_id_old_to_new,
    EmbeddingIndex* new_index) const {
  if (is_empty()) {
    return absl_ports::FailedPreconditionError("EmbeddingIndex is empty");
  }

  const int64_t current_time_ms = clock_.GetSystemTimeMilliseconds();
  std::unique_ptr<KeyMapper<PostingListIdentifier>::Iterator> itr =
      embedding_posting_list_mapper_->GetIterator();
  while (itr->Advance()) {
    std::string_view key = itr->GetKey();
    // This should never happen unless there is an inconsistency, or the index
    // is corrupted.
    if (key.size() < kEncodedDimensionLength) {
      return absl_ports::InternalError(
          "Got invalid key from embedding posting list mapper.");
    }
    uint32_t dimension = encode_util::DecodeIntFromCString(
        std::string_view(key.begin(), kEncodedDimensionLength));

    // Transfer hits
    std::vector<EmbeddingHit> new_hits;
    ICING_ASSIGN_OR_RETURN(
        std::unique_ptr<PostingListEmbeddingHitAccessor> old_pl_accessor,
        PostingListEmbeddingHitAccessor::CreateFromExisting(
            flash_index_storage_.get(), posting_list_hit_serializer_.get(),
            /*existing_posting_list_id=*/itr->GetValue()));
    DocumentId last_new_document_id = kInvalidDocumentId;
    SchemaTypeId schema_type_id = kInvalidSchemaTypeId;
    while (true) {
      ICING_ASSIGN_OR_RETURN(std::vector<EmbeddingHit> batch,
                             old_pl_accessor->GetNextHitsBatch());
      if (batch.empty()) {
        break;
      }
      for (EmbeddingHit& old_hit : batch) {
        // Safety checks to add robustness to the codebase, so to make sure
        // that we never access invalid memory, in case that hit from the
        // posting list is corrupted.
        if (old_hit.basic_hit().document_id() < 0 ||
            old_hit.basic_hit().document_id() >=
                document_id_old_to_new.size()) {
          return absl_ports::InternalError(
              "Embedding hit document id is out of bound. The provided map is "
              "too small, or the index may have been corrupted.");
        }

        // Construct transferred hit and add the embedding vector to the new
        // index.
        DocumentId new_document_id =
            document_id_old_to_new[old_hit.basic_hit().document_id()];
        if (new_document_id == kInvalidDocumentId) {
          continue;
        }
        if (new_document_id != last_new_document_id) {
          schema_type_id =
              document_store.GetSchemaTypeId(new_document_id, current_time_ms);
        }
        last_new_document_id = new_document_id;
        if (schema_type_id == kInvalidSchemaTypeId) {
          // This should not happen, since document store is optimized first,
          // so that new_document_id here should be alive.
          continue;
        }
        ICING_ASSIGN_OR_RETURN(
            EmbeddingIndexingConfig::QuantizationType::Code quantization_type,
            schema_store.GetQuantizationType(schema_type_id,
                                             old_hit.basic_hit().section_id()));
        ICING_RETURN_IF_ERROR(new_index->MarkIndexNonEmpty());

        ICING_ASSIGN_OR_RETURN(
            uint32_t new_location,
            TransferEmbeddingVector(old_hit, dimension, quantization_type,
                                    new_index));
        new_hits.push_back(EmbeddingHit(
            BasicHit(old_hit.basic_hit().section_id(), new_document_id),
            new_location));
      }
    }
    // No hit needs to be added to the new index.
    if (new_hits.empty()) {
      continue;
    }
    // Add transferred hits to the new index.
    ICING_ASSIGN_OR_RETURN(
        std::unique_ptr<PostingListEmbeddingHitAccessor> hit_accum,
        PostingListEmbeddingHitAccessor::Create(
            new_index->flash_index_storage_.get(),
            new_index->posting_list_hit_serializer_.get()));
    for (auto new_hit_itr = new_hits.rbegin(); new_hit_itr != new_hits.rend();
         ++new_hit_itr) {
      ICING_RETURN_IF_ERROR(hit_accum->PrependHit(*new_hit_itr));
    }
    PostingListEmbeddingHitAccessor::FinalizeResult result =
        std::move(*hit_accum).Finalize();
    if (!result.id.is_valid()) {
      return absl_ports::InternalError("Failed to finalize posting list");
    }
    ICING_RETURN_IF_ERROR(
        new_index->embedding_posting_list_mapper_->Put(key, result.id));
  }
  return libtextclassifier3::Status::OK;
}

libtextclassifier3::Status EmbeddingIndex::Optimize(
    const DocumentStore* document_store, const SchemaStore* schema_store,
    const std::vector<DocumentId>& document_id_old_to_new,
    DocumentId new_last_added_document_id) {
  ICING_RETURN_ERROR_IF_NULL(document_store);
  ICING_RETURN_ERROR_IF_NULL(schema_store);
  if (is_empty()) {
    info().last_added_document_id = new_last_added_document_id;
    return libtextclassifier3::Status::OK;
  }

  // This is just for completeness, but this should never be necessary, since we
  // should never have pending hits at the time when Optimize is run.
  ICING_RETURN_IF_ERROR(CommitBufferToIndex());

  std::string temporary_index_working_path = working_path_ + "_temp";
  if (!filesystem_.DeleteDirectoryRecursively(
          temporary_index_working_path.c_str())) {
    ICING_LOG(ERROR) << "Recursively deleting " << temporary_index_working_path;
    return absl_ports::InternalError(
        "Unable to delete temp directory to prepare to build new index.");
  }

  DestructibleDirectory temporary_index_dir(
      &filesystem_, std::move(temporary_index_working_path));
  if (!temporary_index_dir.is_valid()) {
    return absl_ports::InternalError(
        "Unable to create temp directory to build new index.");
  }

  {
    ICING_ASSIGN_OR_RETURN(
        std::unique_ptr<EmbeddingIndex> new_index,
        EmbeddingIndex::Create(&filesystem_, temporary_index_dir.dir(), &clock_,
                               feature_flags_));
    ICING_RETURN_IF_ERROR(TransferIndex(*document_store, *schema_store,
                                        document_id_old_to_new,
                                        new_index.get()));
    new_index->set_last_added_document_id(new_last_added_document_id);
    ICING_RETURN_IF_ERROR(new_index->PersistToDisk());
  }

  // Destruct current storage instances to safely swap directories.
  metadata_mmapped_file_.reset();
  flash_index_storage_.reset();
  embedding_posting_list_mapper_.reset();
  embedding_vectors_.reset();
  quantized_embedding_vectors_.reset();

  if (!filesystem_.SwapFiles(temporary_index_dir.dir().c_str(),
                             working_path_.c_str())) {
    return absl_ports::InternalError(
        "Unable to apply new index due to failed swap!");
  }

  // Reinitialize the index.
  is_initialized_ = false;
  return Initialize();
}

libtextclassifier3::StatusOr<float> EmbeddingIndex::ScoreEmbeddingHit(
    const EmbeddingScorer& scorer, const PropertyProto::VectorProto& query,
    const EmbeddingHit& hit,
    EmbeddingIndexingConfig::QuantizationType::Code quantization_type) const {
  int dimension = query.values().size();
  float semantic_score;
  if (!feature_flags_->enable_embedding_quantization() ||
      quantization_type == EmbeddingIndexingConfig::QuantizationType::NONE) {
    ICING_ASSIGN_OR_RETURN(const float* vector,
                           GetEmbeddingVector(hit, dimension));
    semantic_score = scorer.Score(dimension,
                                  /*v1=*/query.values().data(),
                                  /*v2=*/vector);
  } else {
    ICING_ASSIGN_OR_RETURN(const char* data,
                           GetQuantizedEmbeddingVector(hit, dimension));
    Quantizer quantizer(data);
    const uint8_t* quantized_vector =
        reinterpret_cast<const uint8_t*>(data + sizeof(Quantizer));
    semantic_score = scorer.Score(dimension,
                                  /*v1=*/query.values().data(),
                                  /*v2=*/quantized_vector, quantizer);
  }
  return semantic_score;
}

libtextclassifier3::Status EmbeddingIndex::PersistMetadataToDisk() {
  return metadata_mmapped_file_->PersistToDisk();
}

libtextclassifier3::Status EmbeddingIndex::PersistStoragesToDisk() {
  if (is_empty()) {
    return libtextclassifier3::Status::OK;
  }
  if (!flash_index_storage_->PersistToDisk()) {
    return absl_ports::InternalError("Fail to persist flash index to disk");
  }
  ICING_RETURN_IF_ERROR(embedding_posting_list_mapper_->PersistToDisk());
  ICING_RETURN_IF_ERROR(embedding_vectors_->PersistToDisk());
  ICING_RETURN_IF_ERROR(quantized_embedding_vectors_->PersistToDisk());
  return libtextclassifier3::Status::OK;
}

libtextclassifier3::StatusOr<Crc32> EmbeddingIndex::UpdateStoragesChecksum() {
  if (is_empty()) {
    return Crc32(0);
  }
  ICING_ASSIGN_OR_RETURN(Crc32 embedding_posting_list_mapper_crc,
                         embedding_posting_list_mapper_->UpdateChecksum());
  ICING_ASSIGN_OR_RETURN(Crc32 embedding_vectors_crc,
                         embedding_vectors_->UpdateChecksum());
  ICING_ASSIGN_OR_RETURN(Crc32 quantized_embedding_vectors_crc,
                         quantized_embedding_vectors_->UpdateChecksum());
  return Crc32(embedding_posting_list_mapper_crc.Get() ^
               embedding_vectors_crc.Get() ^
               quantized_embedding_vectors_crc.Get());
}

libtextclassifier3::StatusOr<Crc32> EmbeddingIndex::GetStoragesChecksum()
    const {
  if (is_empty()) {
    return Crc32(0);
  }
  ICING_ASSIGN_OR_RETURN(Crc32 embedding_posting_list_mapper_crc,
                         embedding_posting_list_mapper_->GetChecksum());
  Crc32 embedding_vectors_crc = embedding_vectors_->GetChecksum();
  Crc32 quantized_embedding_vectors_crc =
      quantized_embedding_vectors_->GetChecksum();
  return Crc32(embedding_posting_list_mapper_crc.Get() ^
               embedding_vectors_crc.Get() ^
               quantized_embedding_vectors_crc.Get());
}

}  // namespace lib
}  // namespace icing