Snap for 11973804 from 9aa7c9b430cff162b7a9d582908c28f06c5010ce to 24Q3-release Change-Id: I65e70b31707b6b727bf8c8132c84916f17eb459d
diff --git a/icing/icing-search-engine_search_test.cc b/icing/icing-search-engine_search_test.cc index 0a66da8..db97ac4 100644 --- a/icing/icing-search-engine_search_test.cc +++ b/icing/icing-search-engine_search_test.cc
@@ -3607,6 +3607,98 @@ EqualsProto(projected_document_one)); } +TEST_F(IcingSearchEngineSearchTest, + SearchWithPolymorphicProjectionAndExactSchemaFilter) { + IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache()); + ASSERT_THAT(icing.Initialize().status(), ProtoIsOk()); + SchemaProto schema = + SchemaBuilder() + .AddType(SchemaTypeConfigBuilder() + .SetType("Person") + .AddProperty(PropertyConfigBuilder() + .SetName("name") + .SetDataTypeString(TERM_MATCH_PREFIX, + TOKENIZER_PLAIN) + .SetCardinality(CARDINALITY_OPTIONAL)) + .AddProperty(PropertyConfigBuilder() + .SetName("emailAddress") + .SetDataTypeString(TERM_MATCH_PREFIX, + TOKENIZER_PLAIN) + .SetCardinality(CARDINALITY_OPTIONAL))) + .AddType(SchemaTypeConfigBuilder() + .SetType("Artist") + .AddParentType("Person") + .AddProperty(PropertyConfigBuilder() + .SetName("name") + .SetDataTypeString(TERM_MATCH_PREFIX, + TOKENIZER_PLAIN) + .SetCardinality(CARDINALITY_OPTIONAL)) + .AddProperty(PropertyConfigBuilder() + .SetName("emailAddress") + .SetDataTypeString(TERM_MATCH_PREFIX, + TOKENIZER_PLAIN) + .SetCardinality(CARDINALITY_OPTIONAL)) + .AddProperty(PropertyConfigBuilder() + .SetName("company") + .SetDataTypeString(TERM_MATCH_PREFIX, + TOKENIZER_PLAIN) + .SetCardinality(CARDINALITY_OPTIONAL))) + .Build(); + ASSERT_THAT(icing.SetSchema(schema).status(), ProtoIsOk()); + + // Add a person document and an artist document + DocumentProto document_person = + DocumentBuilder() + .SetKey("namespace", "uri1") + .SetCreationTimestampMs(1000) + .SetSchema("Person") + .AddStringProperty("name", "Foo Person") + .AddStringProperty("emailAddress", "person@gmail.com") + .Build(); + DocumentProto document_artist = + DocumentBuilder() + .SetKey("namespace", "uri2") + .SetCreationTimestampMs(1000) + .SetSchema("Artist") + .AddStringProperty("name", "Foo Artist") + .AddStringProperty("emailAddress", "artist@gmail.com") + .AddStringProperty("company", "Company") + .Build(); + ASSERT_THAT(icing.Put(document_person).status(), ProtoIsOk()); + ASSERT_THAT(icing.Put(document_artist).status(), ProtoIsOk()); + + // Issue a query with a exact schema filter for "Person", which will **not** + // be expanded to "Artist" via polymorphism, and test that projection works + // for both types even though artist will not be returned at all. + SearchSpecProto search_spec; + search_spec.set_term_match_type(TermMatchType::PREFIX); + search_spec.set_query("Foo"); + search_spec.add_schema_type_filters("Person"); + + ResultSpecProto result_spec; + TypePropertyMask* person_field_mask = result_spec.add_type_property_masks(); + person_field_mask->set_schema_type("Person"); + person_field_mask->add_paths("name"); + TypePropertyMask* artist_field_mask = result_spec.add_type_property_masks(); + artist_field_mask->set_schema_type("Artist"); + artist_field_mask->add_paths("emailAddress"); + + // Verify results + DocumentProto projected_document_person = + DocumentBuilder() + .SetKey("namespace", "uri1") + .SetCreationTimestampMs(1000) + .SetSchema("Person") + .AddStringProperty("name", "Foo Person") + .Build(); + SearchResultProto results = + icing.Search(search_spec, GetDefaultScoringSpec(), result_spec); + EXPECT_THAT(results.status(), ProtoIsOk()); + EXPECT_THAT(results.results(), SizeIs(1)); + EXPECT_THAT(results.results(0).document(), + EqualsProto(projected_document_person)); +} + TEST_F(IcingSearchEngineSearchTest, SearchWithPropertyFilters) { IcingSearchEngine icing(GetDefaultIcingOptions(), GetTestJniCache()); ASSERT_THAT(icing.Initialize().status(), ProtoIsOk());
diff --git a/icing/index/embed/embedding-index.cc b/icing/index/embed/embedding-index.cc index 2e70f0e..63381fd 100644 --- a/icing/index/embed/embedding-index.cc +++ b/icing/index/embed/embedding-index.cc
@@ -16,6 +16,7 @@ #include <algorithm> #include <cstdint> +#include <cstring> #include <memory> #include <string> #include <string_view> @@ -108,6 +109,41 @@ 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)); + + 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())) { @@ -129,30 +165,13 @@ metadata_mmapped_file_ = std::make_unique<MemoryMappedFile>(std::move(metadata_mmapped_file)); - 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)); - 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) { @@ -162,6 +181,7 @@ 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; @@ -186,6 +206,10 @@ 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)); @@ -199,6 +223,7 @@ if (vector.values_size() == 0) { return absl_ports::InvalidArgumentError("Vector dimension is 0"); } + ICING_RETURN_IF_ERROR(MarkIndexNonEmpty()); uint32_t location = embedding_vectors_->num_elements(); uint32_t dimension = vector.values_size(); @@ -217,6 +242,11 @@ } 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()) { @@ -276,6 +306,10 @@ libtextclassifier3::Status EmbeddingIndex::TransferIndex( const std::vector<DocumentId>& document_id_old_to_new, EmbeddingIndex* new_index) const { + if (is_empty()) { + return absl_ports::FailedPreconditionError("EmbeddingIndex is empty"); + } + std::unique_ptr<KeyMapper<PostingListIdentifier>::Iterator> itr = embedding_posting_list_mapper_->GetIterator(); while (itr->Advance()) { @@ -322,6 +356,7 @@ if (new_document_id == kInvalidDocumentId) { continue; } + ICING_RETURN_IF_ERROR(new_index->MarkIndexNonEmpty()); uint32_t new_location = new_index->embedding_vectors_->num_elements(); new_hits.push_back(EmbeddingHit( BasicHit(old_hit.basic_hit().section_id(), new_document_id), @@ -336,7 +371,7 @@ } // No hit needs to be added to the new index. if (new_hits.empty()) { - return libtextclassifier3::Status::OK; + continue; } // Add transferred hits to the new index. ICING_ASSIGN_OR_RETURN( @@ -362,6 +397,11 @@ libtextclassifier3::Status EmbeddingIndex::Optimize( const std::vector<DocumentId>& document_id_old_to_new, DocumentId new_last_added_document_id) { + 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()); @@ -413,6 +453,10 @@ } libtextclassifier3::Status EmbeddingIndex::PersistStoragesToDisk(bool force) { + if (is_empty()) { + return libtextclassifier3::Status::OK; + } + if (!flash_index_storage_->PersistToDisk()) { return absl_ports::InternalError("Fail to persist flash index to disk"); } @@ -428,6 +472,9 @@ libtextclassifier3::StatusOr<Crc32> EmbeddingIndex::ComputeStoragesChecksum( bool force) { + if (is_empty()) { + return Crc32(0); + } ICING_ASSIGN_OR_RETURN(Crc32 embedding_posting_list_mapper_crc, embedding_posting_list_mapper_->ComputeChecksum()); ICING_ASSIGN_OR_RETURN(Crc32 embedding_vectors_crc,
diff --git a/icing/index/embed/embedding-index.h b/icing/index/embed/embedding-index.h index 7318871..bf91a83 100644 --- a/icing/index/embed/embedding-index.h +++ b/icing/index/embed/embedding-index.h
@@ -45,24 +45,29 @@ class EmbeddingIndex : public PersistentStorage { public: struct Info { - static constexpr int32_t kMagic = 0xfbe13cbb; + static constexpr int32_t kMagic = 0x61e7cbf1; int32_t magic; DocumentId last_added_document_id; + bool is_empty; + + static constexpr int kPaddingSize = 1000; + // Padding exists just to reserve space for additional values. + uint8_t padding_[kPaddingSize]; Crc32 ComputeChecksum() const { return Crc32( std::string_view(reinterpret_cast<const char*>(this), sizeof(Info))); } - } __attribute__((packed)); - static_assert(sizeof(Info) == 8, ""); + }; + static_assert(sizeof(Info) == 1012, ""); // Metadata file layout: <Crcs><Info> static constexpr int32_t kCrcsMetadataBufferOffset = 0; static constexpr int32_t kInfoMetadataBufferOffset = static_cast<int32_t>(sizeof(Crcs)); static constexpr int32_t kMetadataFileSize = sizeof(Crcs) + sizeof(Info); - static_assert(kMetadataFileSize == 20, ""); + static_assert(kMetadataFileSize == 1024, ""); static constexpr WorkingPathType kWorkingPathType = WorkingPathType::kDirectory; @@ -145,21 +150,33 @@ const std::vector<DocumentId>& document_id_old_to_new, DocumentId new_last_added_document_id); + // Returns a pointer to the embedding vector for the given hit. + // + // Returns: + // - a pointer to the embedding vector on success. + // - OUT_OF_RANGE error if the referred vector is out of range based on the + // location and dimension. libtextclassifier3::StatusOr<const float*> GetEmbeddingVector( const EmbeddingHit& hit, uint32_t dimension) const { if (static_cast<int64_t>(hit.location()) + dimension > GetTotalVectorSize()) { - return absl_ports::InternalError( + return absl_ports::OutOfRangeError( "Got an embedding hit that refers to a vector out of range."); } return embedding_vectors_->array() + hit.location(); } - const float* GetRawEmbeddingData() const { + libtextclassifier3::StatusOr<const float*> GetRawEmbeddingData() const { + if (is_empty()) { + return absl_ports::NotFoundError("EmbeddingIndex is empty"); + } return embedding_vectors_->array(); } int32_t GetTotalVectorSize() const { + if (is_empty()) { + return 0; + } return embedding_vectors_->num_elements(); } @@ -175,18 +192,40 @@ } } + bool is_empty() const { return info().is_empty; } + private: explicit EmbeddingIndex(const Filesystem& filesystem, std::string working_path) : PersistentStorage(filesystem, std::move(working_path), kWorkingPathType) {} + // Creates the storage data if the index is not empty. This will initialize + // flash_index_storage_, embedding_posting_list_mapper_, embedding_vectors_. + // + // Returns: + // - OK on success + // - Any error from FlashIndexStorage, DynamicTrieKeyMapper, or + // FileBackedVector. + libtextclassifier3::Status CreateStorageDataIfNonEmpty(); + + // Marks the index's header to indicate that the index is non-empty. + // + // If the index is already marked as non-empty, this is a no-op. Otherwise, + // CreateStorageDataIfNonEmpty will be called to create the storage data. + // + // Returns: + // - OK on success + // - Any error when calling CreateStorageDataIfNonEmpty. + libtextclassifier3::Status MarkIndexNonEmpty(); + libtextclassifier3::Status Initialize(); // Transfers embedding data and hits from the current index to new_index. // // Returns: // - OK on success + // - FAILED_PRECONDITION_ERROR if the current index is empty. // - INTERNAL_ERROR on I/O error. This could potentially leave the storages // in an invalid state and the caller should handle it properly (e.g. // discard and rebuild) @@ -254,6 +293,8 @@ std::unique_ptr<PostingListEmbeddingHitSerializer> posting_list_hit_serializer_ = std::make_unique<PostingListEmbeddingHitSerializer>(); + + // null if the index is empty. std::unique_ptr<FlashIndexStorage> flash_index_storage_; // The mapper from embedding keys to the corresponding posting list identifier @@ -261,10 +302,14 @@ // // The key for an embedding hit is a one-to-one encoded string of the ordered // pair (dimension, model_signature) corresponding to the embedding. + // + // null if the index is empty. std::unique_ptr<KeyMapper<PostingListIdentifier>> embedding_posting_list_mapper_; // A single FileBackedVector that holds all embedding vectors. + // + // null if the index is empty. std::unique_ptr<FileBackedVector<float>> embedding_vectors_; };
diff --git a/icing/index/embed/embedding-index_test.cc b/icing/index/embed/embedding-index_test.cc index 5980e82..baa7b94 100644 --- a/icing/index/embed/embedding-index_test.cc +++ b/icing/index/embed/embedding-index_test.cc
@@ -90,9 +90,22 @@ } std::vector<float> GetRawEmbeddingData() { - return std::vector<float>(embedding_index_->GetRawEmbeddingData(), - embedding_index_->GetRawEmbeddingData() + - embedding_index_->GetTotalVectorSize()); + auto data_or = embedding_index_->GetRawEmbeddingData(); + if (!data_or.ok()) { + return std::vector<float>(); + } + return std::vector<float>( + data_or.ValueOrDie(), + data_or.ValueOrDie() + embedding_index_->GetTotalVectorSize()); + } + + libtextclassifier3::StatusOr<bool> IndexContainsMetadataOnly() { + std::vector<std::string> sub_dirs; + if (!filesystem_.ListDirectory(embedding_index_dir_.c_str(), /*exclude=*/{}, + /*recursive=*/true, &sub_dirs)) { + return absl_ports::InternalError("Failed to list directory"); + } + return sub_dirs.size() == 1 && sub_dirs[0] == "metadata"; } Filesystem filesystem_; @@ -100,6 +113,10 @@ std::unique_ptr<EmbeddingIndex> embedding_index_; }; +TEST_F(EmbeddingIndexTest, EmptyIndexContainsMetadataOnly) { + EXPECT_THAT(IndexContainsMetadataOnly(), IsOkAndHolds(true)); +} + TEST_F(EmbeddingIndexTest, AddSingleEmbedding) { PropertyProto::VectorProto vector = CreateVector("model", {0.1, 0.2, 0.3}); ICING_ASSERT_OK(embedding_index_->BufferEmbedding( @@ -254,9 +271,52 @@ EXPECT_THAT(GetRawEmbeddingData(), ElementsAre(0.1, 0.2, 0.3, -0.1, -0.2, -0.3)); EXPECT_EQ(embedding_index_->last_added_document_id(), 1); + EXPECT_FALSE(embedding_index_->is_empty()); + EXPECT_THAT(IndexContainsMetadataOnly(), IsOkAndHolds(false)); // Check that clear works as expected. ICING_ASSERT_OK(embedding_index_->Clear()); + EXPECT_TRUE(embedding_index_->is_empty()); + EXPECT_THAT(IndexContainsMetadataOnly(), IsOkAndHolds(true)); + EXPECT_THAT(GetRawEmbeddingData(), IsEmpty()); + EXPECT_EQ(embedding_index_->last_added_document_id(), kInvalidDocumentId); + } +} + +TEST_F(EmbeddingIndexTest, DiscardIndex) { + // Loop the same logic twice to make sure that Discard works as expected, and + // the index is still valid after discarding. + for (int i = 0; i < 2; i++) { + PropertyProto::VectorProto vector1 = CreateVector("model", {0.1, 0.2, 0.3}); + PropertyProto::VectorProto vector2 = + CreateVector("model", {-0.1, -0.2, -0.3}); + ICING_ASSERT_OK(embedding_index_->BufferEmbedding( + BasicHit(/*section_id=*/1, /*document_id=*/0), vector1)); + ICING_ASSERT_OK(embedding_index_->BufferEmbedding( + BasicHit(/*section_id=*/2, /*document_id=*/1), vector2)); + ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex()); + embedding_index_->set_last_added_document_id(1); + + EXPECT_THAT(GetHits(/*dimension=*/3, /*model_signature=*/"model"), + IsOkAndHolds(ElementsAre( + EmbeddingHit(BasicHit(/*section_id=*/2, /*document_id=*/1), + /*location=*/3), + EmbeddingHit(BasicHit(/*section_id=*/1, /*document_id=*/0), + /*location=*/0)))); + EXPECT_THAT(GetRawEmbeddingData(), + ElementsAre(0.1, 0.2, 0.3, -0.1, -0.2, -0.3)); + EXPECT_EQ(embedding_index_->last_added_document_id(), 1); + EXPECT_FALSE(embedding_index_->is_empty()); + EXPECT_THAT(IndexContainsMetadataOnly(), IsOkAndHolds(false)); + + // Check that Discard works as expected. + embedding_index_.reset(); + EmbeddingIndex::Discard(filesystem_, embedding_index_dir_); + ICING_ASSERT_OK_AND_ASSIGN( + embedding_index_, + EmbeddingIndex::Create(&filesystem_, embedding_index_dir_)); + EXPECT_TRUE(embedding_index_->is_empty()); + EXPECT_THAT(IndexContainsMetadataOnly(), IsOkAndHolds(true)); EXPECT_THAT(GetRawEmbeddingData(), IsEmpty()); EXPECT_EQ(embedding_index_->last_added_document_id(), kInvalidDocumentId); } @@ -264,6 +324,8 @@ TEST_F(EmbeddingIndexTest, EmptyCommitIsOk) { ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex()); + EXPECT_TRUE(embedding_index_->is_empty()); + EXPECT_THAT(IndexContainsMetadataOnly(), IsOkAndHolds(true)); EXPECT_THAT(GetRawEmbeddingData(), IsEmpty()); } @@ -329,6 +391,8 @@ ICING_ASSERT_OK(embedding_index_->Optimize( /*document_id_old_to_new=*/{}, /*new_last_added_document_id=*/kInvalidDocumentId)); + EXPECT_TRUE(embedding_index_->is_empty()); + EXPECT_THAT(IndexContainsMetadataOnly(), IsOkAndHolds(true)); EXPECT_THAT(GetRawEmbeddingData(), IsEmpty()); } @@ -377,6 +441,8 @@ /*new_last_added_document_id=*/0)); EXPECT_THAT(GetHits(/*dimension=*/3, /*model_signature=*/"model"), IsOkAndHolds(IsEmpty())); + EXPECT_TRUE(embedding_index_->is_empty()); + EXPECT_THAT(IndexContainsMetadataOnly(), IsOkAndHolds(true)); EXPECT_THAT(GetRawEmbeddingData(), IsEmpty()); EXPECT_EQ(embedding_index_->last_added_document_id(), 0); } @@ -439,6 +505,8 @@ /*new_last_added_document_id=*/0)); EXPECT_THAT(GetHits(/*dimension=*/3, /*model_signature=*/"model"), IsOkAndHolds(IsEmpty())); + EXPECT_TRUE(embedding_index_->is_empty()); + EXPECT_THAT(IndexContainsMetadataOnly(), IsOkAndHolds(true)); EXPECT_THAT(GetRawEmbeddingData(), IsEmpty()); EXPECT_EQ(embedding_index_->last_added_document_id(), 0); } @@ -577,6 +645,45 @@ EXPECT_EQ(embedding_index_->last_added_document_id(), 0); } +TEST_F(EmbeddingIndexTest, + OptimizeEmbeddingsFromDifferentModelsAndDeleteTheFirst) { + PropertyProto::VectorProto vector1 = CreateVector("model1", {0.1, 0.2}); + PropertyProto::VectorProto vector2 = + CreateVector("model2", {-0.1, -0.2, -0.3}); + ICING_ASSERT_OK(embedding_index_->BufferEmbedding( + BasicHit(/*section_id=*/0, /*document_id=*/0), vector1)); + ICING_ASSERT_OK(embedding_index_->BufferEmbedding( + BasicHit(/*section_id=*/1, /*document_id=*/1), vector2)); + ICING_ASSERT_OK(embedding_index_->CommitBufferToIndex()); + embedding_index_->set_last_added_document_id(1); + + // Before optimize + EXPECT_THAT(GetHits(/*dimension=*/2, /*model_signature=*/"model1"), + IsOkAndHolds(ElementsAre( + EmbeddingHit(BasicHit(/*section_id=*/0, /*document_id=*/0), + /*location=*/0)))); + EXPECT_THAT(GetHits(/*dimension=*/3, /*model_signature=*/"model2"), + IsOkAndHolds(ElementsAre( + EmbeddingHit(BasicHit(/*section_id=*/1, /*document_id=*/1), + /*location=*/2)))); + EXPECT_THAT(GetRawEmbeddingData(), ElementsAre(0.1, 0.2, -0.1, -0.2, -0.3)); + EXPECT_EQ(embedding_index_->last_added_document_id(), 1); + + // Run optimize to delete document 0, and check that the index is + // updated correctly. + ICING_ASSERT_OK(embedding_index_->Optimize( + /*document_id_old_to_new=*/{kInvalidDocumentId, 0}, + /*new_last_added_document_id=*/0)); + EXPECT_THAT(GetHits(/*dimension=*/2, /*model_signature=*/"model1"), + IsOkAndHolds(IsEmpty())); + EXPECT_THAT(GetHits(/*dimension=*/3, /*model_signature=*/"model2"), + IsOkAndHolds(ElementsAre( + EmbeddingHit(BasicHit(/*section_id=*/1, /*document_id=*/0), + /*location=*/0)))); + EXPECT_THAT(GetRawEmbeddingData(), ElementsAre(-0.1, -0.2, -0.3)); + EXPECT_EQ(embedding_index_->last_added_document_id(), 0); +} + } // namespace } // namespace lib } // namespace icing
diff --git a/icing/index/embedding-indexing-handler_test.cc b/icing/index/embedding-indexing-handler_test.cc index c5fe3c7..556ba6e 100644 --- a/icing/index/embedding-indexing-handler_test.cc +++ b/icing/index/embedding-indexing-handler_test.cc
@@ -236,9 +236,13 @@ } std::vector<float> GetRawEmbeddingData() { - return std::vector<float>(embedding_index_->GetRawEmbeddingData(), - embedding_index_->GetRawEmbeddingData() + - embedding_index_->GetTotalVectorSize()); + auto data_or = embedding_index_->GetRawEmbeddingData(); + if (!data_or.ok()) { + return std::vector<float>(); + } + return std::vector<float>( + data_or.ValueOrDie(), + data_or.ValueOrDie() + embedding_index_->GetTotalVectorSize()); } Filesystem filesystem_; @@ -420,6 +424,7 @@ // Check that the embedding index should be empty EXPECT_THAT(GetHits(/*dimension=*/3, /*model_signature=*/"model"), IsOkAndHolds(IsEmpty())); + EXPECT_TRUE(embedding_index_->is_empty()); EXPECT_THAT(GetRawEmbeddingData(), IsEmpty()); // Recovery mode should get the same result. @@ -432,6 +437,7 @@ // Check that the embedding index should be empty EXPECT_THAT(GetHits(/*dimension=*/3, /*model_signature=*/"model"), IsOkAndHolds(IsEmpty())); + EXPECT_TRUE(embedding_index_->is_empty()); EXPECT_THAT(GetRawEmbeddingData(), IsEmpty()); } @@ -477,6 +483,7 @@ // Check that the embedding index should be empty EXPECT_THAT(GetHits(/*dimension=*/3, /*model_signature=*/"model"), IsOkAndHolds(IsEmpty())); + EXPECT_TRUE(embedding_index_->is_empty()); EXPECT_THAT(GetRawEmbeddingData(), IsEmpty()); // Handling document with document_id < last_added_document_id should cause a @@ -493,6 +500,7 @@ // Check that the embedding index should be empty EXPECT_THAT(GetHits(/*dimension=*/3, /*model_signature=*/"model"), IsOkAndHolds(IsEmpty())); + EXPECT_TRUE(embedding_index_->is_empty()); EXPECT_THAT(GetRawEmbeddingData(), IsEmpty()); }
diff --git a/icing/index/iterator/doc-hit-info-iterator-filter.cc b/icing/index/iterator/doc-hit-info-iterator-filter.cc index 82d1ac7..3bc3b48 100644 --- a/icing/index/iterator/doc-hit-info-iterator-filter.cc +++ b/icing/index/iterator/doc-hit-info-iterator-filter.cc
@@ -20,17 +20,17 @@ #include <string_view> #include <unordered_set> #include <utility> -#include <vector> #include "icing/text_classifier/lib3/utils/base/status.h" +#include "icing/text_classifier/lib3/utils/base/statusor.h" #include "icing/absl_ports/canonical_errors.h" #include "icing/index/hit/doc-hit-info.h" #include "icing/index/iterator/doc-hit-info-iterator.h" #include "icing/schema/schema-store.h" -#include "icing/schema/section.h" #include "icing/store/document-filter-data.h" #include "icing/store/document-id.h" #include "icing/store/document-store.h" +#include "icing/util/status-macros.h" namespace icing { namespace lib { @@ -43,32 +43,7 @@ document_store_(*document_store), schema_store_(*schema_store), options_(options), - current_time_ms_(current_time_ms) { - // Precompute all the NamespaceIds - for (std::string_view name_space : options_.namespaces) { - auto namespace_id_or = document_store_.GetNamespaceId(name_space); - - // If we can't find the NamespaceId, just throw it away - if (namespace_id_or.ok()) { - target_namespace_ids_.emplace(namespace_id_or.ValueOrDie()); - } - } - - // Precompute all the SchemaTypeIds - for (std::string_view schema_type : options_.schema_types) { - libtextclassifier3::StatusOr<const std::unordered_set<SchemaTypeId>*> - schema_type_ids_or = - schema_store_.GetSchemaTypeIdsWithChildren(schema_type); - - // If we can't find the SchemaTypeId, just throw it away - if (schema_type_ids_or.ok()) { - const std::unordered_set<SchemaTypeId>* schema_type_ids = - schema_type_ids_or.ValueOrDie(); - target_schema_type_ids_.insert(schema_type_ids->begin(), - schema_type_ids->end()); - } - } -} + current_time_ms_(current_time_ms) {} libtextclassifier3::Status DocHitInfoIteratorFilter::Advance() { while (delegate_->Advance().ok()) { @@ -86,14 +61,14 @@ // We should be guaranteed that this exists now. DocumentFilterData data = document_filter_data_optional.value(); - if (!options_.namespaces.empty() && - target_namespace_ids_.count(data.namespace_id()) == 0) { + if (options_.filter_by_namespace_id_enabled && + options_.target_namespace_ids.count(data.namespace_id()) == 0) { // Doesn't match one of the specified namespaces. Keep searching continue; } - if (!options_.schema_types.empty() && - target_schema_type_ids_.count(data.schema_type_id()) == 0) { + if (options_.filter_by_schema_type_id_enabled && + options_.target_schema_type_ids.count(data.schema_type_id()) == 0) { // Doesn't match one of the specified schema types. Keep searching continue; }
diff --git a/icing/index/iterator/doc-hit-info-iterator-filter.h b/icing/index/iterator/doc-hit-info-iterator-filter.h index 608665e..b8e70e8 100644 --- a/icing/index/iterator/doc-hit-info-iterator-filter.h +++ b/icing/index/iterator/doc-hit-info-iterator-filter.h
@@ -18,14 +18,16 @@ #include <cstdint> #include <memory> #include <string> -#include <string_view> #include <unordered_set> #include <utility> #include <vector> #include "icing/text_classifier/lib3/utils/base/status.h" +#include "icing/text_classifier/lib3/utils/base/statusor.h" #include "icing/index/iterator/doc-hit-info-iterator.h" #include "icing/schema/schema-store.h" +#include "icing/schema/section.h" +#include "icing/store/document-filter-data.h" #include "icing/store/document-store.h" #include "icing/store/namespace-id.h" @@ -37,21 +39,25 @@ class DocHitInfoIteratorFilter : public DocHitInfoIterator { public: struct Options { - // List of namespaces that documents must have. An empty vector means that - // all namespaces are valid, and no documents will be filtered out. + // List of namespace ids that documents must have. + // filter_by_namespace_id_enabled=false means that all namespaces are valid, + // and no documents will be filtered out. // // Note that if we want to reference the strings in namespaces later, ensure // that the caller who passed the Options class outlives the // DocHitInfoIteratorFilter. - std::vector<std::string_view> namespaces; + bool filter_by_namespace_id_enabled = false; + std::unordered_set<NamespaceId> target_namespace_ids; - // List of schema types that documents must have. An empty vector means that - // all schema types are valid, and no documents will be filtered out. + // List of schema type ids that documents must have. + // filter_by_schema_type_id_enabled=false means that all schema types are + // valid, and no documents will be filtered out. // // Note that if we want to reference the strings in schema types later, // ensure that the caller who passed the Options class outlives the // DocHitInfoIteratorFilter. - std::vector<std::string_view> schema_types; + bool filter_by_schema_type_id_enabled = false; + std::unordered_set<SchemaTypeId> target_schema_type_ids; }; explicit DocHitInfoIteratorFilter( @@ -83,8 +89,6 @@ const DocumentStore& document_store_; const SchemaStore& schema_store_; const Options options_; - std::unordered_set<NamespaceId> target_namespace_ids_; - std::unordered_set<SchemaTypeId> target_schema_type_ids_; int64_t current_time_ms_; };
diff --git a/icing/index/iterator/doc-hit-info-iterator-filter_test.cc b/icing/index/iterator/doc-hit-info-iterator-filter_test.cc index 823014f..c7f54b2 100644 --- a/icing/index/iterator/doc-hit-info-iterator-filter_test.cc +++ b/icing/index/iterator/doc-hit-info-iterator-filter_test.cc
@@ -14,23 +14,26 @@ #include "icing/index/iterator/doc-hit-info-iterator-filter.h" -#include <limits> #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/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/iterator/doc-hit-info-iterator-and.h" #include "icing/index/iterator/doc-hit-info-iterator-test-util.h" #include "icing/index/iterator/doc-hit-info-iterator.h" #include "icing/proto/document.pb.h" #include "icing/proto/schema.pb.h" +#include "icing/query/query-utils.h" #include "icing/schema-builder.h" #include "icing/schema/schema-store.h" #include "icing/schema/section.h" @@ -39,6 +42,7 @@ #include "icing/testing/common-matchers.h" #include "icing/testing/fake-clock.h" #include "icing/testing/tmp-directory.h" +#include "icing/util/clock.h" namespace icing { namespace lib { @@ -283,17 +287,18 @@ DocumentProto document2_namespace1_; DocumentProto document1_namespace2_; DocumentProto document1_namespace3_; - DocHitInfoIteratorFilter::Options options_; }; TEST_F(DocHitInfoIteratorNamespaceFilterTest, EmptyOriginalIterator) { std::unique_ptr<DocHitInfoIterator> original_iterator_empty = std::make_unique<DocHitInfoIteratorDummy>(); - options_.namespaces = std::vector<std::string_view>{}; + SearchSpecProto search_spec; + DocHitInfoIteratorFilter::Options options = + GetFilterOptions(search_spec, *document_store_, *schema_store_); DocHitInfoIteratorFilter filtered_iterator( std::move(original_iterator_empty), document_store_.get(), - schema_store_.get(), options_, fake_clock_.GetSystemTimeMilliseconds()); + schema_store_.get(), options, fake_clock_.GetSystemTimeMilliseconds()); EXPECT_THAT(GetDocumentIds(&filtered_iterator), IsEmpty()); } @@ -308,10 +313,13 @@ std::unique_ptr<DocHitInfoIterator> original_iterator = std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos); - options_.namespaces = std::vector<std::string_view>{"nonexistent_namespace"}; + SearchSpecProto search_spec; + search_spec.add_namespace_filters("nonexistent_namespace"); + DocHitInfoIteratorFilter::Options options = + GetFilterOptions(search_spec, *document_store_, *schema_store_); DocHitInfoIteratorFilter filtered_iterator( std::move(original_iterator), document_store_.get(), schema_store_.get(), - options_, fake_clock_.GetSystemTimeMilliseconds()); + options, fake_clock_.GetSystemTimeMilliseconds()); EXPECT_THAT(GetDocumentIds(&filtered_iterator), IsEmpty()); } @@ -325,10 +333,12 @@ std::unique_ptr<DocHitInfoIterator> original_iterator = std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos); - options_.namespaces = std::vector<std::string_view>{}; + SearchSpecProto search_spec; + DocHitInfoIteratorFilter::Options options = + GetFilterOptions(search_spec, *document_store_, *schema_store_); DocHitInfoIteratorFilter filtered_iterator( std::move(original_iterator), document_store_.get(), schema_store_.get(), - options_, fake_clock_.GetSystemTimeMilliseconds()); + options, fake_clock_.GetSystemTimeMilliseconds()); EXPECT_THAT(GetDocumentIds(&filtered_iterator), ElementsAre(document_id1)); } @@ -349,10 +359,13 @@ std::unique_ptr<DocHitInfoIterator> original_iterator = std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos); - options_.namespaces = std::vector<std::string_view>{namespace1_}; + SearchSpecProto search_spec; + search_spec.add_namespace_filters(namespace1_); + DocHitInfoIteratorFilter::Options options = + GetFilterOptions(search_spec, *document_store_, *schema_store_); DocHitInfoIteratorFilter filtered_iterator( std::move(original_iterator), document_store_.get(), schema_store_.get(), - options_, fake_clock_.GetSystemTimeMilliseconds()); + options, fake_clock_.GetSystemTimeMilliseconds()); EXPECT_THAT(GetDocumentIds(&filtered_iterator), ElementsAre(document_id1, document_id2)); @@ -375,10 +388,14 @@ std::unique_ptr<DocHitInfoIterator> original_iterator = std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos); - options_.namespaces = std::vector<std::string_view>{namespace1_, namespace3_}; + SearchSpecProto search_spec; + search_spec.add_namespace_filters(namespace1_); + search_spec.add_namespace_filters(namespace3_); + DocHitInfoIteratorFilter::Options options = + GetFilterOptions(search_spec, *document_store_, *schema_store_); DocHitInfoIteratorFilter filtered_iterator( std::move(original_iterator), document_store_.get(), schema_store_.get(), - options_, fake_clock_.GetSystemTimeMilliseconds()); + options, fake_clock_.GetSystemTimeMilliseconds()); EXPECT_THAT(GetDocumentIds(&filtered_iterator), ElementsAre(document_id1, document_id2, document_id4)); @@ -457,17 +474,18 @@ DocumentProto document2_schema2_; DocumentProto document3_schema3_; DocumentProto document4_schema1_; - DocHitInfoIteratorFilter::Options options_; }; TEST_F(DocHitInfoIteratorSchemaTypeFilterTest, EmptyOriginalIterator) { std::unique_ptr<DocHitInfoIterator> original_iterator_empty = std::make_unique<DocHitInfoIteratorDummy>(); - options_.schema_types = std::vector<std::string_view>{}; + SearchSpecProto search_spec; + DocHitInfoIteratorFilter::Options options = + GetFilterOptions(search_spec, *document_store_, *schema_store_); DocHitInfoIteratorFilter filtered_iterator( std::move(original_iterator_empty), document_store_.get(), - schema_store_.get(), options_, fake_clock_.GetSystemTimeMilliseconds()); + schema_store_.get(), options, fake_clock_.GetSystemTimeMilliseconds()); EXPECT_THAT(GetDocumentIds(&filtered_iterator), IsEmpty()); } @@ -482,11 +500,13 @@ std::unique_ptr<DocHitInfoIterator> original_iterator = std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos); - options_.schema_types = - std::vector<std::string_view>{"nonexistent_schema_type"}; + SearchSpecProto search_spec; + search_spec.add_schema_type_filters("nonexistent_schema_type"); + DocHitInfoIteratorFilter::Options options = + GetFilterOptions(search_spec, *document_store_, *schema_store_); DocHitInfoIteratorFilter filtered_iterator( std::move(original_iterator), document_store_.get(), schema_store_.get(), - options_, fake_clock_.GetSystemTimeMilliseconds()); + options, fake_clock_.GetSystemTimeMilliseconds()); EXPECT_THAT(GetDocumentIds(&filtered_iterator), IsEmpty()); } @@ -500,10 +520,12 @@ std::unique_ptr<DocHitInfoIterator> original_iterator = std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos); - options_.schema_types = std::vector<std::string_view>{}; + SearchSpecProto search_spec; + DocHitInfoIteratorFilter::Options options = + GetFilterOptions(search_spec, *document_store_, *schema_store_); DocHitInfoIteratorFilter filtered_iterator( std::move(original_iterator), document_store_.get(), schema_store_.get(), - options_, fake_clock_.GetSystemTimeMilliseconds()); + options, fake_clock_.GetSystemTimeMilliseconds()); EXPECT_THAT(GetDocumentIds(&filtered_iterator), ElementsAre(document_id1)); } @@ -521,10 +543,13 @@ std::unique_ptr<DocHitInfoIterator> original_iterator = std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos); - options_.schema_types = std::vector<std::string_view>{kSchema1}; + SearchSpecProto search_spec; + search_spec.add_schema_type_filters(std::string(kSchema1)); + DocHitInfoIteratorFilter::Options options = + GetFilterOptions(search_spec, *document_store_, *schema_store_); DocHitInfoIteratorFilter filtered_iterator( std::move(original_iterator), document_store_.get(), schema_store_.get(), - options_, fake_clock_.GetSystemTimeMilliseconds()); + options, fake_clock_.GetSystemTimeMilliseconds()); EXPECT_THAT(GetDocumentIds(&filtered_iterator), ElementsAre(document_id1)); } @@ -544,17 +569,21 @@ std::unique_ptr<DocHitInfoIterator> original_iterator = std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos); - options_.schema_types = std::vector<std::string_view>{kSchema2, kSchema3}; + SearchSpecProto search_spec; + search_spec.add_schema_type_filters(std::string(kSchema2)); + search_spec.add_schema_type_filters(std::string(kSchema3)); + DocHitInfoIteratorFilter::Options options = + GetFilterOptions(search_spec, *document_store_, *schema_store_); DocHitInfoIteratorFilter filtered_iterator( std::move(original_iterator), document_store_.get(), schema_store_.get(), - options_, fake_clock_.GetSystemTimeMilliseconds()); + options, fake_clock_.GetSystemTimeMilliseconds()); EXPECT_THAT(GetDocumentIds(&filtered_iterator), ElementsAre(document_id2, document_id3)); } TEST_F(DocHitInfoIteratorSchemaTypeFilterTest, - FilterForSchemaTypePolymorphismOk) { + FilterIsExactForSchemaTypePolymorphism) { // Add some irrelevant documents. ICING_ASSERT_OK_AND_ASSIGN(DocumentId document_id1, document_store_->Put(document1_schema1_)); @@ -580,28 +609,34 @@ DocHitInfo(document_id1), DocHitInfo(document_id2), DocHitInfo(person_document_id), DocHitInfo(artist_document_id)}; - // Filters for the "person" type should also include the "artist" type. + // Filters for the "person" type should NOT include the "artist" type, since + // schema filters should not expand for polymorphism. std::unique_ptr<DocHitInfoIterator> original_iterator = std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos); - options_.schema_types = {"person"}; + SearchSpecProto search_spec_1; + search_spec_1.add_schema_type_filters("person"); + DocHitInfoIteratorFilter::Options options = + GetFilterOptions(search_spec_1, *document_store_, *schema_store_); DocHitInfoIteratorFilter filtered_iterator_1( std::move(original_iterator), document_store_.get(), schema_store_.get(), - options_, fake_clock_.GetSystemTimeMilliseconds()); + options, fake_clock_.GetSystemTimeMilliseconds()); EXPECT_THAT(GetDocumentIds(&filtered_iterator_1), - ElementsAre(person_document_id, artist_document_id)); + ElementsAre(person_document_id)); // Filters for the "artist" type should not include the "person" type. original_iterator = std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos); - options_.schema_types = {"artist"}; + SearchSpecProto search_spec_2; + search_spec_2.add_schema_type_filters("artist"); + options = GetFilterOptions(search_spec_2, *document_store_, *schema_store_); DocHitInfoIteratorFilter filtered_iterator_2( std::move(original_iterator), document_store_.get(), schema_store_.get(), - options_, fake_clock_.GetSystemTimeMilliseconds()); + options, fake_clock_.GetSystemTimeMilliseconds()); EXPECT_THAT(GetDocumentIds(&filtered_iterator_2), ElementsAre(artist_document_id)); } TEST_F(DocHitInfoIteratorSchemaTypeFilterTest, - FilterForSchemaTypeMultipleParentPolymorphismOk) { + FilterIsExactForSchemaTypeMultipleParentPolymorphism) { // Create an email and a message document. ICING_ASSERT_OK_AND_ASSIGN( DocumentId email_document_id, @@ -629,31 +664,40 @@ DocHitInfo(email_document_id), DocHitInfo(message_document_id), DocHitInfo(email_message_document_id)}; - // Filters for the "email" type should also include the "emailMessage" type. + // Filters for the "email" type should NOT include the "emailMessage" type, + // since schema filters should not expand for polymorphism. std::unique_ptr<DocHitInfoIterator> original_iterator = std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos); - options_.schema_types = std::vector<std::string_view>{"email"}; + SearchSpecProto search_spec_1; + search_spec_1.add_schema_type_filters("email"); + DocHitInfoIteratorFilter::Options options = + GetFilterOptions(search_spec_1, *document_store_, *schema_store_); DocHitInfoIteratorFilter filtered_iterator_1( std::move(original_iterator), document_store_.get(), schema_store_.get(), - options_, fake_clock_.GetSystemTimeMilliseconds()); + options, fake_clock_.GetSystemTimeMilliseconds()); EXPECT_THAT(GetDocumentIds(&filtered_iterator_1), - ElementsAre(email_document_id, email_message_document_id)); + ElementsAre(email_document_id)); - // Filters for the "message" type should also include the "emailMessage" type. + // Filters for the "message" type should NOT include the "emailMessage" type, + // since schema filters should not expand for polymorphism. original_iterator = std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos); - options_.schema_types = std::vector<std::string_view>{"message"}; + SearchSpecProto search_spec_2; + search_spec_2.add_schema_type_filters("message"); + options = GetFilterOptions(search_spec_2, *document_store_, *schema_store_); DocHitInfoIteratorFilter filtered_iterator_2( std::move(original_iterator), document_store_.get(), schema_store_.get(), - options_, fake_clock_.GetSystemTimeMilliseconds()); + options, fake_clock_.GetSystemTimeMilliseconds()); EXPECT_THAT(GetDocumentIds(&filtered_iterator_2), - ElementsAre(message_document_id, email_message_document_id)); + ElementsAre(message_document_id)); // Filters for a irrelevant type should return nothing. original_iterator = std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos); - options_.schema_types = std::vector<std::string_view>{"person"}; + SearchSpecProto search_spec_3; + search_spec_3.add_schema_type_filters("person"); + options = GetFilterOptions(search_spec_3, *document_store_, *schema_store_); DocHitInfoIteratorFilter filtered_iterator_3( std::move(original_iterator), document_store_.get(), schema_store_.get(), - options_, fake_clock_.GetSystemTimeMilliseconds()); + options, fake_clock_.GetSystemTimeMilliseconds()); EXPECT_THAT(GetDocumentIds(&filtered_iterator_3), IsEmpty()); } @@ -950,13 +994,13 @@ std::unique_ptr<DocHitInfoIterator> original_iterator = std::make_unique<DocHitInfoIteratorDummy>(doc_hit_infos); - DocHitInfoIteratorFilter::Options options; - + SearchSpecProto search_spec; // Filters out document3 by namespace - options.namespaces = std::vector<std::string_view>{namespace1_}; - + search_spec.add_namespace_filters(namespace1_); // Filters out document4 by schema type - options.schema_types = std::vector<std::string_view>{schema1_}; + search_spec.add_schema_type_filters(schema1_); + DocHitInfoIteratorFilter::Options options = + GetFilterOptions(search_spec, *document_store, *schema_store_); DocHitInfoIteratorFilter filtered_iterator( std::move(original_iterator), document_store.get(), schema_store_.get(), @@ -1049,9 +1093,11 @@ std::make_unique<DocHitInfoIteratorAnd>(std::move(left_iter), std::move(right_iter)); - DocHitInfoIteratorFilter::Options options; + SearchSpecProto search_spec; // Filters out document3 by namespace - options.namespaces = std::vector<std::string_view>{namespace1_}; + search_spec.add_namespace_filters(namespace1_); + DocHitInfoIteratorFilter::Options options = + GetFilterOptions(search_spec, *document_store_, *schema_store_); DocHitInfoIteratorFilter filtered_iterator( std::move(original_iterator), document_store_.get(), schema_store_.get(), options, fake_clock_.GetSystemTimeMilliseconds());
diff --git a/icing/query/query-processor.cc b/icing/query/query-processor.cc index 35213da..ed3b5ca 100644 --- a/icing/query/query-processor.cc +++ b/icing/query/query-processor.cc
@@ -110,7 +110,8 @@ } } - DocHitInfoIteratorFilter::Options options = GetFilterOptions(search_spec); + DocHitInfoIteratorFilter::Options options = + GetFilterOptions(search_spec, document_store_, schema_store_); results.root_iterator = std::make_unique<DocHitInfoIteratorFilter>( std::move(results.root_iterator), &document_store_, &schema_store_, options, current_time_ms); @@ -155,7 +156,8 @@ std::unique_ptr<Tokenizer> plain_tokenizer, tokenizer_factory::CreateIndexingTokenizer( StringIndexingConfig::TokenizerType::PLAIN, &language_segmenter_)); - DocHitInfoIteratorFilter::Options options = GetFilterOptions(search_spec); + DocHitInfoIteratorFilter::Options options = + GetFilterOptions(search_spec, document_store_, schema_store_); bool needs_term_frequency_info = ranking_strategy == ScoringSpecProto::RankingStrategy::RELEVANCE_SCORE;
diff --git a/icing/query/query-utils.cc b/icing/query/query-utils.cc index 37c3600..0a4cbd3 100644 --- a/icing/query/query-utils.cc +++ b/icing/query/query-utils.cc
@@ -15,26 +15,68 @@ #include "icing/query/query-utils.h" #include <string_view> -#include <vector> +#include <unordered_set> + +#include "icing/text_classifier/lib3/utils/base/statusor.h" +#include "icing/index/iterator/doc-hit-info-iterator-filter.h" +#include "icing/schema/schema-store.h" +#include "icing/store/document-filter-data.h" +#include "icing/store/document-store.h" +#include "icing/store/namespace-id.h" namespace icing { namespace lib { +namespace { + +std::unordered_set<NamespaceId> ConvertNamespaceToIds( + const DocumentStore& document_store, const SearchSpecProto& search_spec) { + std::unordered_set<NamespaceId> ids; + for (std::string_view name_space : search_spec.namespace_filters()) { + auto namespace_id_or = document_store.GetNamespaceId(name_space); + + // If we can't find the NamespaceId, just throw it away + if (namespace_id_or.ok()) { + ids.insert(namespace_id_or.ValueOrDie()); + } + } + return ids; +} + +std::unordered_set<SchemaTypeId> ConvertExactSchemaTypeToIds( + const SchemaStore& schema_store, const SearchSpecProto& search_spec) { + std::unordered_set<SchemaTypeId> ids; + ids.reserve(search_spec.schema_type_filters_size()); + for (std::string_view schema_type : search_spec.schema_type_filters()) { + libtextclassifier3::StatusOr<SchemaTypeId> schema_type_id_or = + schema_store.GetSchemaTypeId(schema_type); + + // If we can't find the SchemaTypeId, just throw it away + if (schema_type_id_or.ok()) { + ids.insert(schema_type_id_or.ValueOrDie()); + } + } + return ids; +} + +} // namespace + DocHitInfoIteratorFilter::Options GetFilterOptions( - const SearchSpecProto& search_spec) { + const SearchSpecProto& search_spec, const DocumentStore& document_store, + const SchemaStore& schema_store) { DocHitInfoIteratorFilter::Options options; - if (search_spec.namespace_filters_size() > 0) { - options.namespaces = - std::vector<std::string_view>(search_spec.namespace_filters().begin(), - search_spec.namespace_filters().end()); - } + // Precompute all the NamespaceIds + options.filter_by_namespace_id_enabled = + !search_spec.namespace_filters().empty(); + options.target_namespace_ids = + ConvertNamespaceToIds(document_store, search_spec); - if (search_spec.schema_type_filters_size() > 0) { - options.schema_types = - std::vector<std::string_view>(search_spec.schema_type_filters().begin(), - search_spec.schema_type_filters().end()); - } + // Precompute all the SchemaTypeIds + options.filter_by_schema_type_id_enabled = + !search_spec.schema_type_filters().empty(); + options.target_schema_type_ids = + ConvertExactSchemaTypeToIds(schema_store, search_spec); return options; }
diff --git a/icing/query/query-utils.h b/icing/query/query-utils.h index d85cf3a..dae75a2 100644 --- a/icing/query/query-utils.h +++ b/icing/query/query-utils.h
@@ -17,12 +17,15 @@ #include "icing/index/iterator/doc-hit-info-iterator-filter.h" #include "icing/proto/search.pb.h" +#include "icing/schema/schema-store.h" +#include "icing/store/document-store.h" namespace icing { namespace lib { DocHitInfoIteratorFilter::Options GetFilterOptions( - const SearchSpecProto& search_spec); + const SearchSpecProto& search_spec, const DocumentStore& document_store, + const SchemaStore& schema_store); } // namespace lib } // namespace icing
diff --git a/icing/store/document-store.cc b/icing/store/document-store.cc index e0a5c68..a039eb1 100644 --- a/icing/store/document-store.cc +++ b/icing/store/document-store.cc
@@ -1083,24 +1083,31 @@ // TODO(b/147231617): Make a better way to replace the error message in an // existing Status. auto document_id_or = GetDocumentId(name_space, uri); - if (absl_ports::IsNotFound(document_id_or.status())) { - ICING_VLOG(1) << document_id_or.status().error_message(); - return libtextclassifier3::Status( - document_id_or.status().CanonicalCode(), - IcingStringUtil::StringPrintf("Document (%s, %s) not found.", - name_space.data(), uri.data())); + if (!document_id_or.ok()) { + if (absl_ports::IsNotFound(document_id_or.status())) { + ICING_VLOG(1) << document_id_or.status().error_message(); + return absl_ports::NotFoundError(absl_ports::StrCat( + "Document (", name_space, ", ", uri, ") not found.")); + } + + // Real error. Log it in error level and pass it up. + ICING_LOG(ERROR) << document_id_or.status().error_message(); + return std::move(document_id_or).status(); } DocumentId document_id = document_id_or.ValueOrDie(); // TODO(b/147231617): Make a better way to replace the error message in an // existing Status. - auto status_or = Get(document_id); - if (absl_ports::IsNotFound(status_or.status())) { - ICING_LOG(ERROR) << document_id_or.status().error_message(); - return libtextclassifier3::Status( - status_or.status().CanonicalCode(), - IcingStringUtil::StringPrintf("Document (%s, %s) not found.", - name_space.data(), uri.data())); + auto status_or = Get(document_id, clear_internal_fields); + if (!status_or.ok()) { + if (absl_ports::IsNotFound(status_or.status())) { + ICING_VLOG(1) << status_or.status().error_message(); + return absl_ports::NotFoundError(absl_ports::StrCat( + "Document (", name_space, ", ", uri, ") not found.")); + } + + // Real error. Log it in error level. + ICING_LOG(ERROR) << status_or.status().error_message(); } return status_or; } @@ -1381,30 +1388,21 @@ libtextclassifier3::StatusOr<CorpusAssociatedScoreData> DocumentStore::GetCorpusAssociatedScoreData(CorpusId corpus_id) const { - auto score_data_or = corpus_score_cache_->GetCopy(corpus_id); - if (!score_data_or.ok()) { - return score_data_or.status(); - } - - CorpusAssociatedScoreData corpus_associated_score_data = - std::move(score_data_or).ValueOrDie(); - return corpus_associated_score_data; + return corpus_score_cache_->GetCopy(corpus_id); } libtextclassifier3::StatusOr<CorpusAssociatedScoreData> DocumentStore::GetCorpusAssociatedScoreDataToUpdate(CorpusId corpus_id) const { auto corpus_scoring_data_or = GetCorpusAssociatedScoreData(corpus_id); - if (corpus_scoring_data_or.ok()) { - return std::move(corpus_scoring_data_or).ValueOrDie(); + if (!corpus_scoring_data_or.ok() && + absl_ports::IsOutOfRange(corpus_scoring_data_or.status())) { + // OUT_OF_RANGE is the StatusCode returned when a corpus id is added to + // corpus_score_cache_ for the first time. Return a default + // CorpusAssociatedScoreData object in this case. + return CorpusAssociatedScoreData(); } - CorpusAssociatedScoreData scoringData; - // OUT_OF_RANGE is the StatusCode returned when a corpus id is added to - // corpus_score_cache_ for the first time. - if (corpus_scoring_data_or.status().CanonicalCode() == - libtextclassifier3::StatusCode::OUT_OF_RANGE) { - return scoringData; - } - return corpus_scoring_data_or.status(); + + return corpus_scoring_data_or; } // TODO(b/273826815): Decide on and adopt a consistent pattern for handling
diff --git a/proto/icing/proto/search.proto b/proto/icing/proto/search.proto index 293c062..3b3a955 100644 --- a/proto/icing/proto/search.proto +++ b/proto/icing/proto/search.proto
@@ -60,7 +60,8 @@ // OPTIONAL: Only search for documents that have the specified schema types. // If unset, the query will search over all schema types. Note that this // applies to the entire 'query'. To issue different queries for different - // schema types, separate Search()'s will need to be made. + // schema types, separate Search()'s will need to be made. Also note that + // schema filters will not be expanded for polymorphism. repeated string schema_type_filters = 4; // Timestamp taken just before sending proto across the JNI boundary from java @@ -451,6 +452,7 @@ // If unset, the suggestion will search over all schema types. Note that this // applies to the entire 'prefix'. To issue different suggestions for // different schema typs, separate RunSuggestion()'s will need to be made. + // Also note that schema filters will not be expanded for polymorphism. repeated string schema_type_filters = 6; // OPTIONAL: Only search for suggestions that under the specified types and