| // Copyright (C) 2022 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/result/result-retriever-v2.h" |
| |
| #include <atomic> |
| #include <cstddef> |
| #include <cstdint> |
| #include <memory> |
| #include <string> |
| #include <unordered_map> |
| #include <utility> |
| #include <vector> |
| |
| #include "icing/text_classifier/lib3/utils/base/status.h" |
| #include "icing/text_classifier/lib3/utils/base/statusor.h" |
| #include "gmock/gmock.h" |
| #include "gtest/gtest.h" |
| #include "icing/absl_ports/mutex.h" |
| #include "icing/document-builder.h" |
| #include "icing/feature-flags.h" |
| #include "icing/file/filesystem.h" |
| #include "icing/file/mock-filesystem.h" |
| #include "icing/file/portable-file-backed-proto-log.h" |
| #include "icing/portable/equals-proto.h" |
| #include "icing/portable/platform.h" |
| #include "icing/proto/document.pb.h" |
| #include "icing/proto/document_wrapper.pb.h" |
| #include "icing/proto/schema.pb.h" |
| #include "icing/proto/search.pb.h" |
| #include "icing/result/page-result.h" |
| #include "icing/result/result-state-v2.h" |
| #include "icing/schema-builder.h" |
| #include "icing/schema/schema-store.h" |
| #include "icing/schema/section.h" |
| #include "icing/scoring/priority-queue-scored-document-hits-ranker.h" |
| #include "icing/scoring/scored-document-hit.h" |
| #include "icing/store/document-filter-data.h" |
| #include "icing/store/document-id.h" |
| #include "icing/store/document-store.h" |
| #include "icing/testing/common-matchers.h" |
| #include "icing/testing/fake-clock.h" |
| #include "icing/testing/test-data.h" |
| #include "icing/testing/test-feature-flags.h" |
| #include "icing/testing/tmp-directory.h" |
| #include "icing/tokenization/language-segmenter-factory.h" |
| #include "icing/tokenization/language-segmenter.h" |
| #include "icing/transform/normalizer-factory.h" |
| #include "icing/transform/normalizer.h" |
| #include "icing/util/clock.h" |
| #include "icing/util/icu-data-file-helper.h" |
| #include "unicode/uloc.h" |
| |
| namespace icing { |
| namespace lib { |
| |
| namespace { |
| |
| using ::icing::lib::portable_equals_proto::EqualsProto; |
| using ::testing::DoDefault; |
| using ::testing::ElementsAre; |
| using ::testing::Eq; |
| using ::testing::Gt; |
| using ::testing::IsEmpty; |
| using ::testing::Pointee; |
| using ::testing::Return; |
| using ::testing::SizeIs; |
| using EntryIdMap = std::unordered_map<int32_t, int>; |
| |
| // Mock the behavior of GroupResultLimiter::ShouldBeRemoved. |
| class MockGroupResultLimiter : public GroupResultLimiterV2 { |
| public: |
| MockGroupResultLimiter() : GroupResultLimiterV2() { |
| ON_CALL(*this, ShouldBeRemoved).WillByDefault(Return(false)); |
| } |
| |
| MOCK_METHOD(bool, ShouldBeRemoved, |
| (const ScoredDocumentHit&, const EntryIdMap&, |
| const DocumentStore&, std::vector<int>&, |
| ResultSpecProto::ResultGroupingType, int64_t), |
| (const, override)); |
| }; |
| |
| class ResultRetrieverV2Test : public ::testing::Test { |
| protected: |
| ResultRetrieverV2Test() : test_dir_(GetTestTempDir() + "/icing") { |
| filesystem_.CreateDirectoryRecursively(test_dir_.c_str()); |
| } |
| |
| void SetUp() override { |
| feature_flags_ = std::make_unique<FeatureFlags>(GetTestFeatureFlags()); |
| if (!IsCfStringTokenization() && !IsReverseJniTokenization()) { |
| ICING_ASSERT_OK( |
| // File generated via icu_data_file rule in //icing/BUILD. |
| icu_data_file_helper::SetUpIcuDataFile( |
| GetTestFilePath("icing/icu.dat"))); |
| } |
| language_segmenter_factory::SegmenterOptions options(ULOC_US); |
| ICING_ASSERT_OK_AND_ASSIGN( |
| language_segmenter_, |
| language_segmenter_factory::Create(std::move(options))); |
| |
| ICING_ASSERT_OK_AND_ASSIGN( |
| schema_store_, SchemaStore::Create(&filesystem_, test_dir_, |
| &fake_clock_, feature_flags_.get())); |
| ICING_ASSERT_OK_AND_ASSIGN(normalizer_, normalizer_factory::Create( |
| /*max_term_byte_size=*/10000)); |
| |
| SchemaProto schema = |
| SchemaBuilder() |
| .AddType(SchemaTypeConfigBuilder() |
| .SetType("Email") |
| .AddProperty(PropertyConfigBuilder() |
| .SetName("name") |
| .SetDataTypeString(TERM_MATCH_PREFIX, |
| TOKENIZER_PLAIN) |
| .SetCardinality(CARDINALITY_OPTIONAL)) |
| .AddProperty(PropertyConfigBuilder() |
| .SetName("body") |
| .SetDataTypeString(TERM_MATCH_EXACT, |
| TOKENIZER_PLAIN) |
| .SetCardinality(CARDINALITY_OPTIONAL)) |
| .AddProperty( |
| PropertyConfigBuilder() |
| .SetName("sender") |
| .SetDataTypeDocument( |
| "Person", /*index_nested_properties=*/true) |
| .SetCardinality(CARDINALITY_OPTIONAL))) |
| .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))) |
| .Build(); |
| ASSERT_THAT(schema_store_->SetSchema( |
| schema, /*ignore_errors_and_delete_documents=*/false, |
| /*allow_circular_schema_definitions=*/false), |
| IsOk()); |
| |
| num_total_hits_ = 0; |
| } |
| |
| void TearDown() override { |
| filesystem_.DeleteDirectoryRecursively(test_dir_.c_str()); |
| } |
| |
| SectionId GetSectionId(const std::string& type, const std::string& property) { |
| auto type_id_or = schema_store_->GetSchemaTypeId(type); |
| if (!type_id_or.ok()) { |
| return kInvalidSectionId; |
| } |
| SchemaTypeId type_id = type_id_or.ValueOrDie(); |
| for (SectionId section_id = 0; section_id <= kMaxSectionId; ++section_id) { |
| auto metadata_or = schema_store_->GetSectionMetadata(type_id, section_id); |
| if (!metadata_or.ok()) { |
| break; |
| } |
| const SectionMetadata* metadata = metadata_or.ValueOrDie(); |
| if (metadata->path == property) { |
| return metadata->id; |
| } |
| } |
| return kInvalidSectionId; |
| } |
| |
| std::unique_ptr<FeatureFlags> feature_flags_; |
| const Filesystem filesystem_; |
| const std::string test_dir_; |
| std::unique_ptr<LanguageSegmenter> language_segmenter_; |
| std::unique_ptr<SchemaStore> schema_store_; |
| std::unique_ptr<Normalizer> normalizer_; |
| std::atomic<int> num_total_hits_; |
| FakeClock fake_clock_; |
| }; |
| |
| DocumentProto CreateDocument(int id) { |
| return DocumentBuilder() |
| .SetKey("icing", "Email/" + std::to_string(id)) |
| .SetSchema("Email") |
| .AddStringProperty("name", "subject foo " + std::to_string(id)) |
| .AddStringProperty("body", "body bar " + std::to_string(id)) |
| .SetCreationTimestampMs(1574365086666 + id) |
| .Build(); |
| } |
| |
| SectionIdMask CreateSectionIdMask(const std::vector<SectionId>& section_ids) { |
| SectionIdMask mask = 0; |
| for (SectionId section_id : section_ids) { |
| mask |= (UINT64_C(1) << section_id); |
| } |
| return mask; |
| } |
| |
| ResultSpecProto CreateResultSpec( |
| int num_per_page, ResultSpecProto::ResultGroupingType result_group_type) { |
| ResultSpecProto result_spec; |
| result_spec.set_result_group_type(result_group_type); |
| result_spec.set_num_per_page(num_per_page); |
| return result_spec; |
| } |
| |
| libtextclassifier3::StatusOr<DocumentStore::CreateResult> CreateDocumentStore( |
| const Filesystem* filesystem, const std::string& base_dir, |
| const Clock* clock, const SchemaStore* schema_store, |
| const FeatureFlags& feature_flags) { |
| return DocumentStore::Create( |
| filesystem, base_dir, clock, schema_store, &feature_flags, |
| /*force_recovery_and_revalidate_documents=*/false, |
| /*pre_mapping_fbv=*/false, /*use_persistent_hash_map=*/true, |
| PortableFileBackedProtoLog<DocumentWrapper>::kDefaultCompressionLevel, |
| /*initialize_stats=*/nullptr); |
| } |
| |
| TEST_F(ResultRetrieverV2Test, CreationWithNullPointerShouldFail) { |
| EXPECT_THAT( |
| ResultRetrieverV2::Create(/*doc_store=*/nullptr, schema_store_.get(), |
| language_segmenter_.get(), normalizer_.get()), |
| StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION)); |
| |
| ICING_ASSERT_OK_AND_ASSIGN( |
| DocumentStore::CreateResult create_result, |
| CreateDocumentStore(&filesystem_, test_dir_, &fake_clock_, |
| schema_store_.get(), *feature_flags_)); |
| std::unique_ptr<DocumentStore> doc_store = |
| std::move(create_result.document_store); |
| |
| EXPECT_THAT( |
| ResultRetrieverV2::Create(doc_store.get(), /*schema_store=*/nullptr, |
| language_segmenter_.get(), normalizer_.get()), |
| StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION)); |
| EXPECT_THAT(ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(), |
| /*language_segmenter=*/nullptr, |
| normalizer_.get()), |
| StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION)); |
| EXPECT_THAT(ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(), |
| language_segmenter_.get(), |
| /*normalizer=*/nullptr), |
| StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION)); |
| } |
| |
| TEST_F(ResultRetrieverV2Test, ShouldRetrieveSimpleResults) { |
| ICING_ASSERT_OK_AND_ASSIGN( |
| DocumentStore::CreateResult create_result, |
| CreateDocumentStore(&filesystem_, test_dir_, &fake_clock_, |
| schema_store_.get(), *feature_flags_)); |
| std::unique_ptr<DocumentStore> doc_store = |
| std::move(create_result.document_store); |
| |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1, |
| doc_store->Put(CreateDocument(/*id=*/1))); |
| DocumentId document_id1 = put_result1.new_document_id; |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2, |
| doc_store->Put(CreateDocument(/*id=*/2))); |
| DocumentId document_id2 = put_result2.new_document_id; |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3, |
| doc_store->Put(CreateDocument(/*id=*/3))); |
| DocumentId document_id3 = put_result3.new_document_id; |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result4, |
| doc_store->Put(CreateDocument(/*id=*/4))); |
| DocumentId document_id4 = put_result4.new_document_id; |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result5, |
| doc_store->Put(CreateDocument(/*id=*/5))); |
| DocumentId document_id5 = put_result5.new_document_id; |
| |
| std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "name"), |
| GetSectionId("Email", "body")}; |
| SectionIdMask hit_section_id_mask = CreateSectionIdMask(hit_section_ids); |
| std::vector<ScoredDocumentHit> scored_document_hits = { |
| {document_id1, hit_section_id_mask, /*score=*/19}, |
| {document_id2, hit_section_id_mask, /*score=*/12}, |
| {document_id3, hit_section_id_mask, /*score=*/8}, |
| {document_id4, hit_section_id_mask, /*score=*/3}, |
| {document_id5, hit_section_id_mask, /*score=*/1}}; |
| ICING_ASSERT_OK_AND_ASSIGN( |
| std::unique_ptr<ResultRetrieverV2> result_retriever, |
| ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(), |
| language_segmenter_.get(), normalizer_.get())); |
| |
| SearchResultProto::ResultProto result1; |
| *result1.mutable_document() = CreateDocument(/*id=*/1); |
| result1.set_score(19); |
| SearchResultProto::ResultProto result2; |
| *result2.mutable_document() = CreateDocument(/*id=*/2); |
| result2.set_score(12); |
| SearchResultProto::ResultProto result3; |
| *result3.mutable_document() = CreateDocument(/*id=*/3); |
| result3.set_score(8); |
| SearchResultProto::ResultProto result4; |
| *result4.mutable_document() = CreateDocument(/*id=*/4); |
| result4.set_score(3); |
| SearchResultProto::ResultProto result5; |
| *result5.mutable_document() = CreateDocument(/*id=*/5); |
| result5.set_score(1); |
| |
| ResultStateV2 result_state( |
| std::make_unique< |
| PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>( |
| std::move(scored_document_hits), /*is_descending=*/true), |
| /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr, |
| CreateResultSpec(/*num_per_page=*/2, ResultSpecProto::NAMESPACE), |
| *doc_store); |
| |
| // First page, 2 results |
| auto [page_result1, has_more_results1] = result_retriever->RetrieveNextPage( |
| result_state, fake_clock_.GetSystemTimeMilliseconds()); |
| EXPECT_THAT(page_result1.results, |
| ElementsAre(EqualsProto(result1), EqualsProto(result2))); |
| // num_results_with_snippets is 0 when there is no snippet. |
| EXPECT_THAT(page_result1.num_results_with_snippets, Eq(0)); |
| // Requested page size is same as num_per_page. |
| EXPECT_THAT(page_result1.requested_page_size, Eq(2)); |
| // Has more results. |
| EXPECT_TRUE(has_more_results1); |
| |
| // Second page, 2 results |
| auto [page_result2, has_more_results2] = result_retriever->RetrieveNextPage( |
| result_state, fake_clock_.GetSystemTimeMilliseconds()); |
| EXPECT_THAT(page_result2.results, |
| ElementsAre(EqualsProto(result3), EqualsProto(result4))); |
| // num_results_with_snippets is 0 when there is no snippet. |
| EXPECT_THAT(page_result2.num_results_with_snippets, Eq(0)); |
| // Requested page size is same as num_per_page. |
| EXPECT_THAT(page_result2.requested_page_size, Eq(2)); |
| // Has more results. |
| EXPECT_TRUE(has_more_results2); |
| |
| // Third page, 1 result |
| auto [page_result3, has_more_results3] = result_retriever->RetrieveNextPage( |
| result_state, fake_clock_.GetSystemTimeMilliseconds()); |
| EXPECT_THAT(page_result3.results, ElementsAre(EqualsProto(result5))); |
| // num_results_with_snippets is 0 when there is no snippet. |
| EXPECT_THAT(page_result3.num_results_with_snippets, Eq(0)); |
| // Requested page size is same as num_per_page. |
| EXPECT_THAT(page_result3.requested_page_size, Eq(2)); |
| // No more results. |
| EXPECT_FALSE(has_more_results3); |
| } |
| |
| TEST_F(ResultRetrieverV2Test, ShouldIgnoreNonInternalErrors) { |
| ICING_ASSERT_OK_AND_ASSIGN( |
| DocumentStore::CreateResult create_result, |
| CreateDocumentStore(&filesystem_, test_dir_, &fake_clock_, |
| schema_store_.get(), *feature_flags_)); |
| std::unique_ptr<DocumentStore> doc_store = |
| std::move(create_result.document_store); |
| |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1, |
| doc_store->Put(CreateDocument(/*id=*/1))); |
| DocumentId document_id1 = put_result1.new_document_id; |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2, |
| doc_store->Put(CreateDocument(/*id=*/2))); |
| DocumentId document_id2 = put_result2.new_document_id; |
| |
| DocumentId invalid_document_id = -1; |
| std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "name"), |
| GetSectionId("Email", "body")}; |
| SectionIdMask hit_section_id_mask = CreateSectionIdMask(hit_section_ids); |
| std::vector<ScoredDocumentHit> scored_document_hits = { |
| {document_id1, hit_section_id_mask, /*score=*/12}, |
| {document_id2, hit_section_id_mask, /*score=*/4}, |
| {invalid_document_id, hit_section_id_mask, /*score=*/0}}; |
| ICING_ASSERT_OK_AND_ASSIGN( |
| std::unique_ptr<ResultRetrieverV2> result_retriever, |
| ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(), |
| language_segmenter_.get(), normalizer_.get(), |
| std::make_unique<MockGroupResultLimiter>())); |
| |
| SearchResultProto::ResultProto result1; |
| *result1.mutable_document() = CreateDocument(/*id=*/1); |
| result1.set_score(12); |
| SearchResultProto::ResultProto result2; |
| *result2.mutable_document() = CreateDocument(/*id=*/2); |
| result2.set_score(4); |
| |
| ResultStateV2 result_state1( |
| std::make_unique< |
| PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>( |
| std::move(scored_document_hits), |
| /*is_descending=*/true), |
| /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr, |
| CreateResultSpec(/*num_per_page=*/3, ResultSpecProto::NAMESPACE), |
| *doc_store); |
| PageResult page_result1 = |
| result_retriever |
| ->RetrieveNextPage(result_state1, |
| fake_clock_.GetSystemTimeMilliseconds()) |
| .first; |
| EXPECT_THAT(page_result1.results, |
| ElementsAre(EqualsProto(result1), EqualsProto(result2))); |
| |
| DocumentId non_existing_document_id = 4; |
| scored_document_hits = { |
| {non_existing_document_id, hit_section_id_mask, /*score=*/15}, |
| {document_id1, hit_section_id_mask, /*score=*/12}, |
| {document_id2, hit_section_id_mask, /*score=*/4}}; |
| ResultStateV2 result_state2( |
| std::make_unique< |
| PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>( |
| std::move(scored_document_hits), |
| /*is_descending=*/true), |
| /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr, |
| CreateResultSpec(/*num_per_page=*/3, ResultSpecProto::NAMESPACE), |
| *doc_store); |
| PageResult page_result2 = |
| result_retriever |
| ->RetrieveNextPage(result_state2, |
| fake_clock_.GetSystemTimeMilliseconds()) |
| .first; |
| EXPECT_THAT(page_result2.results, |
| ElementsAre(EqualsProto(result1), EqualsProto(result2))); |
| } |
| |
| TEST_F(ResultRetrieverV2Test, |
| ShouldLimitNumChildDocumentsByMaxJoinedChildPerParent) { |
| ICING_ASSERT_OK_AND_ASSIGN( |
| DocumentStore::CreateResult create_result, |
| CreateDocumentStore(&filesystem_, test_dir_, &fake_clock_, |
| schema_store_.get(), *feature_flags_)); |
| std::unique_ptr<DocumentStore> doc_store = |
| std::move(create_result.document_store); |
| |
| // 1. Add 2 Person document |
| DocumentProto person_document1 = |
| DocumentBuilder() |
| .SetKey("namespace", "Person/1") |
| .SetCreationTimestampMs(1000) |
| .SetSchema("Person") |
| .AddStringProperty("name", "Joe Fox") |
| .AddStringProperty("emailAddress", "ny152@aol.com") |
| .Build(); |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1, |
| doc_store->Put(person_document1)); |
| DocumentId person_document_id1 = put_result1.new_document_id; |
| |
| DocumentProto person_document2 = |
| DocumentBuilder() |
| .SetKey("namespace", "Person/2") |
| .SetCreationTimestampMs(1000) |
| .SetSchema("Person") |
| .AddStringProperty("name", "Meg Ryan") |
| .AddStringProperty("emailAddress", "shopgirl@aol.com") |
| .Build(); |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2, |
| doc_store->Put(person_document2)); |
| DocumentId person_document_id2 = put_result2.new_document_id; |
| |
| // 2. Add 4 Email documents |
| DocumentProto email_document1 = DocumentBuilder() |
| .SetKey("namespace", "Email/1") |
| .SetCreationTimestampMs(1000) |
| .SetSchema("Email") |
| .AddStringProperty("name", "Test 1") |
| .AddStringProperty("body", "Test 1") |
| .Build(); |
| ICING_ASSERT_OK_AND_ASSIGN(put_result1, doc_store->Put(email_document1)); |
| DocumentId email_document_id1 = put_result1.new_document_id; |
| |
| DocumentProto email_document2 = DocumentBuilder() |
| .SetKey("namespace", "Email/2") |
| .SetCreationTimestampMs(1000) |
| .SetSchema("Email") |
| .AddStringProperty("name", "Test 2") |
| .AddStringProperty("body", "Test 2") |
| .Build(); |
| ICING_ASSERT_OK_AND_ASSIGN(put_result2, doc_store->Put(email_document2)); |
| DocumentId email_document_id2 = put_result2.new_document_id; |
| |
| DocumentProto email_document3 = DocumentBuilder() |
| .SetKey("namespace", "Email/3") |
| .SetCreationTimestampMs(1000) |
| .SetSchema("Email") |
| .AddStringProperty("name", "Test 3") |
| .AddStringProperty("body", "Test 3") |
| .Build(); |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3, |
| doc_store->Put(email_document3)); |
| DocumentId email_document_id3 = put_result3.new_document_id; |
| |
| DocumentProto email_document4 = DocumentBuilder() |
| .SetKey("namespace", "Email/4") |
| .SetCreationTimestampMs(1000) |
| .SetSchema("Email") |
| .AddStringProperty("name", "Test 4") |
| .AddStringProperty("body", "Test 4") |
| .Build(); |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result4, |
| doc_store->Put(email_document4)); |
| DocumentId email_document_id4 = put_result4.new_document_id; |
| |
| // 3. Setup the joined scored results. |
| std::vector<SectionId> person_hit_section_ids = { |
| GetSectionId("Person", "name")}; |
| std::vector<SectionId> email_hit_section_ids = { |
| GetSectionId("Email", "name"), GetSectionId("Email", "body")}; |
| SectionIdMask person_hit_section_id_mask = |
| CreateSectionIdMask(person_hit_section_ids); |
| SectionIdMask email_hit_section_id_mask = |
| CreateSectionIdMask(email_hit_section_ids); |
| |
| ScoredDocumentHit person1_scored_doc_hit( |
| person_document_id1, person_hit_section_id_mask, /*score=*/1); |
| ScoredDocumentHit person2_scored_doc_hit( |
| person_document_id2, person_hit_section_id_mask, /*score=*/2); |
| ScoredDocumentHit email1_scored_doc_hit( |
| email_document_id1, email_hit_section_id_mask, /*score=*/3); |
| ScoredDocumentHit email2_scored_doc_hit( |
| email_document_id2, email_hit_section_id_mask, /*score=*/4); |
| ScoredDocumentHit email3_scored_doc_hit( |
| email_document_id3, email_hit_section_id_mask, /*score=*/5); |
| ScoredDocumentHit email4_scored_doc_hit( |
| email_document_id4, email_hit_section_id_mask, /*score=*/6); |
| // Create JoinedScoredDocumentHits mapping: |
| // - Person1 to Email1 |
| // - Person2 to Email2, Email3, Email4 |
| std::vector<JoinedScoredDocumentHit> joined_scored_document_hits = { |
| JoinedScoredDocumentHit( |
| /*final_score=*/1, |
| /*parent_scored_document_hit=*/person1_scored_doc_hit, |
| /*child_scored_document_hits=*/{email1_scored_doc_hit}), |
| JoinedScoredDocumentHit( |
| /*final_score=*/3, |
| /*parent_scored_document_hit=*/person2_scored_doc_hit, |
| /*child_scored_document_hits=*/ |
| {email4_scored_doc_hit, email3_scored_doc_hit, |
| email2_scored_doc_hit})}; |
| |
| // 4. Retrieve result with max_joined_children_per_parent_to_return = 2. |
| ResultSpecProto result_spec = |
| CreateResultSpec(/*num_per_page=*/2, ResultSpecProto::NAMESPACE); |
| result_spec.set_max_joined_children_per_parent_to_return(2); |
| |
| ICING_ASSERT_OK_AND_ASSIGN( |
| std::unique_ptr<ResultRetrieverV2> result_retriever, |
| ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(), |
| language_segmenter_.get(), normalizer_.get())); |
| ResultStateV2 result_state( |
| std::make_unique< |
| PriorityQueueScoredDocumentHitsRanker<JoinedScoredDocumentHit>>( |
| std::move(joined_scored_document_hits), /*is_descending=*/true), |
| /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr, |
| result_spec, *doc_store); |
| |
| // Result1: person2 with child docs = [email4, email3] |
| SearchResultProto::ResultProto result1; |
| *result1.mutable_document() = person_document2; |
| result1.set_score(3); |
| SearchResultProto::ResultProto* child1 = result1.add_joined_results(); |
| *child1->mutable_document() = email_document4; |
| child1->set_score(6); |
| SearchResultProto::ResultProto* child2 = result1.add_joined_results(); |
| *child2->mutable_document() = email_document3; |
| child2->set_score(5); |
| |
| // Result2: person1 with child docs = [email1] |
| SearchResultProto::ResultProto result2; |
| *result2.mutable_document() = person_document1; |
| result2.set_score(1); |
| SearchResultProto::ResultProto* child3 = result2.add_joined_results(); |
| *child3->mutable_document() = email_document1; |
| child3->set_score(3); |
| |
| auto [page_result, has_more_results] = result_retriever->RetrieveNextPage( |
| result_state, fake_clock_.GetSystemTimeMilliseconds()); |
| EXPECT_THAT(page_result.results, |
| ElementsAre(EqualsProto(result1), EqualsProto(result2))); |
| // No more results. |
| EXPECT_FALSE(has_more_results); |
| } |
| |
| TEST_F(ResultRetrieverV2Test, ShouldIgnoreInternalErrors) { |
| MockFilesystem mock_filesystem; |
| EXPECT_CALL(mock_filesystem, |
| PRead(A<int>(), A<void*>(), A<size_t>(), A<off_t>())) |
| .WillOnce(Return(false)) |
| .WillRepeatedly(DoDefault()); |
| |
| ICING_ASSERT_OK_AND_ASSIGN( |
| DocumentStore::CreateResult create_result, |
| CreateDocumentStore(&mock_filesystem, test_dir_, &fake_clock_, |
| schema_store_.get(), *feature_flags_)); |
| std::unique_ptr<DocumentStore> doc_store = |
| std::move(create_result.document_store); |
| |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1, |
| doc_store->Put(CreateDocument(/*id=*/1))); |
| DocumentId document_id1 = put_result1.new_document_id; |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2, |
| doc_store->Put(CreateDocument(/*id=*/2))); |
| DocumentId document_id2 = put_result2.new_document_id; |
| |
| std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "name"), |
| GetSectionId("Email", "body")}; |
| SectionIdMask hit_section_id_mask = CreateSectionIdMask(hit_section_ids); |
| std::vector<ScoredDocumentHit> scored_document_hits = { |
| {document_id1, hit_section_id_mask, /*score=*/0}, |
| {document_id2, hit_section_id_mask, /*score=*/0}}; |
| |
| ICING_ASSERT_OK_AND_ASSIGN( |
| std::unique_ptr<ResultRetrieverV2> result_retriever, |
| ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(), |
| language_segmenter_.get(), normalizer_.get(), |
| std::make_unique<MockGroupResultLimiter>())); |
| |
| SearchResultProto::ResultProto result1; |
| *result1.mutable_document() = CreateDocument(/*id=*/1); |
| result1.set_score(0); |
| |
| ResultStateV2 result_state( |
| std::make_unique< |
| PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>( |
| std::move(scored_document_hits), |
| /*is_descending=*/true), |
| /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr, |
| CreateResultSpec(/*num_per_page=*/2, ResultSpecProto::NAMESPACE), |
| *doc_store); |
| PageResult page_result = |
| result_retriever |
| ->RetrieveNextPage(result_state, |
| fake_clock_.GetSystemTimeMilliseconds()) |
| .first; |
| // We mocked mock_filesystem to return an internal error when retrieving doc2, |
| // so doc2 should be skipped and doc1 should still be returned. |
| EXPECT_THAT(page_result.results, ElementsAre(EqualsProto(result1))); |
| } |
| |
| TEST_F(ResultRetrieverV2Test, ShouldUpdateResultState) { |
| ICING_ASSERT_OK_AND_ASSIGN( |
| DocumentStore::CreateResult create_result, |
| CreateDocumentStore(&filesystem_, test_dir_, &fake_clock_, |
| schema_store_.get(), *feature_flags_)); |
| std::unique_ptr<DocumentStore> doc_store = |
| std::move(create_result.document_store); |
| |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1, |
| doc_store->Put(CreateDocument(/*id=*/1))); |
| DocumentId document_id1 = put_result1.new_document_id; |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2, |
| doc_store->Put(CreateDocument(/*id=*/2))); |
| DocumentId document_id2 = put_result2.new_document_id; |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3, |
| doc_store->Put(CreateDocument(/*id=*/3))); |
| DocumentId document_id3 = put_result3.new_document_id; |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result4, |
| doc_store->Put(CreateDocument(/*id=*/4))); |
| DocumentId document_id4 = put_result4.new_document_id; |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result5, |
| doc_store->Put(CreateDocument(/*id=*/5))); |
| DocumentId document_id5 = put_result5.new_document_id; |
| |
| std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "name"), |
| GetSectionId("Email", "body")}; |
| SectionIdMask hit_section_id_mask = CreateSectionIdMask(hit_section_ids); |
| std::vector<ScoredDocumentHit> scored_document_hits = { |
| {document_id1, hit_section_id_mask, /*score=*/0}, |
| {document_id2, hit_section_id_mask, /*score=*/0}, |
| {document_id3, hit_section_id_mask, /*score=*/0}, |
| {document_id4, hit_section_id_mask, /*score=*/0}, |
| {document_id5, hit_section_id_mask, /*score=*/0}}; |
| ICING_ASSERT_OK_AND_ASSIGN( |
| std::unique_ptr<ResultRetrieverV2> result_retriever, |
| ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(), |
| language_segmenter_.get(), normalizer_.get())); |
| |
| ResultStateV2 result_state( |
| std::make_unique< |
| PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>( |
| std::move(scored_document_hits), |
| /*is_descending=*/true), |
| /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr, |
| CreateResultSpec(/*num_per_page=*/2, ResultSpecProto::NAMESPACE), |
| *doc_store); |
| |
| // First page, 2 results |
| PageResult page_result1 = |
| result_retriever |
| ->RetrieveNextPage(result_state, |
| fake_clock_.GetSystemTimeMilliseconds()) |
| .first; |
| ASSERT_THAT(page_result1.results, SizeIs(2)); |
| { |
| absl_ports::shared_lock l(&result_state.mutex); |
| |
| // num_returned = size of first page |
| EXPECT_THAT(result_state.num_returned, Eq(2)); |
| // Should remove the 2 returned docs from scored_document_hits and only |
| // contain the remaining 3. |
| EXPECT_THAT(result_state.scored_document_hits_ranker, Pointee(SizeIs(3))); |
| } |
| |
| // Second page, 2 results |
| PageResult page_result2 = |
| result_retriever |
| ->RetrieveNextPage(result_state, |
| fake_clock_.GetSystemTimeMilliseconds()) |
| .first; |
| ASSERT_THAT(page_result2.results, SizeIs(2)); |
| { |
| absl_ports::shared_lock l(&result_state.mutex); |
| |
| // num_returned = size of first and second pages |
| EXPECT_THAT(result_state.num_returned, Eq(4)); |
| // Should remove the 2 returned docs from scored_document_hits and only |
| // contain the remaining 1. |
| EXPECT_THAT(result_state.scored_document_hits_ranker, Pointee(SizeIs(1))); |
| } |
| |
| // Third page, 1 result |
| PageResult page_result3 = |
| result_retriever |
| ->RetrieveNextPage(result_state, |
| fake_clock_.GetSystemTimeMilliseconds()) |
| .first; |
| ASSERT_THAT(page_result3.results, SizeIs(1)); |
| { |
| absl_ports::shared_lock l(&result_state.mutex); |
| |
| // num_returned = size of first, second and third pages |
| EXPECT_THAT(result_state.num_returned, Eq(5)); |
| // Should remove the 1 returned doc from scored_document_hits and become |
| // empty. |
| EXPECT_THAT(result_state.scored_document_hits_ranker, Pointee(IsEmpty())); |
| } |
| } |
| |
| TEST_F(ResultRetrieverV2Test, ShouldUpdateNumTotalHits) { |
| ICING_ASSERT_OK_AND_ASSIGN( |
| DocumentStore::CreateResult create_result, |
| CreateDocumentStore(&filesystem_, test_dir_, &fake_clock_, |
| schema_store_.get(), *feature_flags_)); |
| std::unique_ptr<DocumentStore> doc_store = |
| std::move(create_result.document_store); |
| |
| std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "name"), |
| GetSectionId("Email", "body")}; |
| SectionIdMask hit_section_id_mask = CreateSectionIdMask(hit_section_ids); |
| |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1, |
| doc_store->Put(CreateDocument(/*id=*/1))); |
| DocumentId document_id1 = put_result1.new_document_id; |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2, |
| doc_store->Put(CreateDocument(/*id=*/2))); |
| DocumentId document_id2 = put_result2.new_document_id; |
| std::vector<ScoredDocumentHit> scored_document_hits1 = { |
| {document_id1, hit_section_id_mask, /*score=*/0}, |
| {document_id2, hit_section_id_mask, /*score=*/0}}; |
| std::shared_ptr<ResultStateV2> result_state1 = |
| std::make_shared<ResultStateV2>( |
| std::make_unique< |
| PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>( |
| std::move(scored_document_hits1), |
| /*is_descending=*/true), |
| /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr, |
| CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE), |
| *doc_store); |
| { |
| absl_ports::unique_lock l(&result_state1->mutex); |
| |
| result_state1->RegisterNumTotalHits(&num_total_hits_); |
| ASSERT_THAT(num_total_hits_, Eq(2)); |
| } |
| |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3, |
| doc_store->Put(CreateDocument(/*id=*/3))); |
| DocumentId document_id3 = put_result3.new_document_id; |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result4, |
| doc_store->Put(CreateDocument(/*id=*/4))); |
| DocumentId document_id4 = put_result4.new_document_id; |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result5, |
| doc_store->Put(CreateDocument(/*id=*/5))); |
| DocumentId document_id5 = put_result5.new_document_id; |
| std::vector<ScoredDocumentHit> scored_document_hits2 = { |
| {document_id3, hit_section_id_mask, /*score=*/0}, |
| {document_id4, hit_section_id_mask, /*score=*/0}, |
| {document_id5, hit_section_id_mask, /*score=*/0}}; |
| std::shared_ptr<ResultStateV2> result_state2 = |
| std::make_shared<ResultStateV2>( |
| std::make_unique< |
| PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>( |
| std::move(scored_document_hits2), |
| /*is_descending=*/true), |
| /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr, |
| CreateResultSpec(/*num_per_page=*/2, ResultSpecProto::NAMESPACE), |
| *doc_store); |
| { |
| absl_ports::unique_lock l(&result_state2->mutex); |
| |
| result_state2->RegisterNumTotalHits(&num_total_hits_); |
| ASSERT_THAT(num_total_hits_, Eq(5)); |
| } |
| |
| ICING_ASSERT_OK_AND_ASSIGN( |
| std::unique_ptr<ResultRetrieverV2> result_retriever, |
| ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(), |
| language_segmenter_.get(), normalizer_.get())); |
| |
| // Should get 1 doc in the first page of result_state1, and num_total_hits |
| // should be decremented by 1. |
| PageResult page_result1 = |
| result_retriever |
| ->RetrieveNextPage(*result_state1, |
| fake_clock_.GetSystemTimeMilliseconds()) |
| .first; |
| ASSERT_THAT(page_result1.results, SizeIs(1)); |
| EXPECT_THAT(num_total_hits_, Eq(4)); |
| |
| // Should get 2 docs in the first page of result_state2, and num_total_hits |
| // should be decremented by 2. |
| PageResult page_result2 = |
| result_retriever |
| ->RetrieveNextPage(*result_state2, |
| fake_clock_.GetSystemTimeMilliseconds()) |
| .first; |
| ASSERT_THAT(page_result2.results, SizeIs(2)); |
| EXPECT_THAT(num_total_hits_, Eq(2)); |
| |
| // Should get 1 doc in the second page of result_state2 (although num_per_page |
| // is 2, there is only 1 doc left), and num_total_hits should be decremented |
| // by 1. |
| PageResult page_result3 = |
| result_retriever |
| ->RetrieveNextPage(*result_state2, |
| fake_clock_.GetSystemTimeMilliseconds()) |
| .first; |
| ASSERT_THAT(page_result3.results, SizeIs(1)); |
| EXPECT_THAT(num_total_hits_, Eq(1)); |
| |
| // Destruct result_state1. There is 1 doc left, so num_total_hits should be |
| // decremented by 1 when destructing it. |
| result_state1.reset(); |
| EXPECT_THAT(num_total_hits_, Eq(0)); |
| |
| // Destruct result_state2. There is 0 doc left, so num_total_hits should be |
| // unchanged when destructing it. |
| result_state1.reset(); |
| EXPECT_THAT(num_total_hits_, Eq(0)); |
| } |
| |
| TEST_F(ResultRetrieverV2Test, ShouldLimitNumTotalBytesPerPage) { |
| ICING_ASSERT_OK_AND_ASSIGN( |
| DocumentStore::CreateResult create_result, |
| CreateDocumentStore(&filesystem_, test_dir_, &fake_clock_, |
| schema_store_.get(), *feature_flags_)); |
| std::unique_ptr<DocumentStore> doc_store = |
| std::move(create_result.document_store); |
| |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1, |
| doc_store->Put(CreateDocument(/*id=*/1))); |
| DocumentId document_id1 = put_result1.new_document_id; |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2, |
| doc_store->Put(CreateDocument(/*id=*/2))); |
| DocumentId document_id2 = put_result2.new_document_id; |
| |
| std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "name"), |
| GetSectionId("Email", "body")}; |
| SectionIdMask hit_section_id_mask = CreateSectionIdMask(hit_section_ids); |
| std::vector<ScoredDocumentHit> scored_document_hits = { |
| {document_id1, hit_section_id_mask, /*score=*/5}, |
| {document_id2, hit_section_id_mask, /*score=*/0}}; |
| ICING_ASSERT_OK_AND_ASSIGN( |
| std::unique_ptr<ResultRetrieverV2> result_retriever, |
| ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(), |
| language_segmenter_.get(), normalizer_.get())); |
| |
| SearchResultProto::ResultProto result1; |
| *result1.mutable_document() = CreateDocument(/*id=*/1); |
| result1.set_score(5); |
| SearchResultProto::ResultProto result2; |
| *result2.mutable_document() = CreateDocument(/*id=*/2); |
| result2.set_score(0); |
| |
| ResultSpecProto result_spec = |
| CreateResultSpec(/*num_per_page=*/2, ResultSpecProto::NAMESPACE); |
| result_spec.set_num_total_bytes_per_page_threshold(result1.ByteSizeLong()); |
| ResultStateV2 result_state( |
| std::make_unique< |
| PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>( |
| std::move(scored_document_hits), |
| /*is_descending=*/true), |
| /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr, |
| result_spec, *doc_store); |
| |
| // First page. Only result1 should be returned, since its byte size meets |
| // num_total_bytes_per_page_threshold and ResultRetriever should terminate |
| // early even though # of results is still below num_per_page. |
| auto [page_result1, has_more_results1] = result_retriever->RetrieveNextPage( |
| result_state, fake_clock_.GetSystemTimeMilliseconds()); |
| EXPECT_THAT(page_result1.results, ElementsAre(EqualsProto(result1))); |
| // Has more results. |
| EXPECT_TRUE(has_more_results1); |
| |
| // Second page, result2. |
| auto [page_result2, has_more_results2] = result_retriever->RetrieveNextPage( |
| result_state, fake_clock_.GetSystemTimeMilliseconds()); |
| EXPECT_THAT(page_result2.results, ElementsAre(EqualsProto(result2))); |
| // No more results. |
| EXPECT_FALSE(has_more_results2); |
| } |
| |
| TEST_F(ResultRetrieverV2Test, |
| ShouldReturnSingleLargeResultAboveNumTotalBytesPerPageThreshold) { |
| ICING_ASSERT_OK_AND_ASSIGN( |
| DocumentStore::CreateResult create_result, |
| CreateDocumentStore(&filesystem_, test_dir_, &fake_clock_, |
| schema_store_.get(), *feature_flags_)); |
| std::unique_ptr<DocumentStore> doc_store = |
| std::move(create_result.document_store); |
| |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1, |
| doc_store->Put(CreateDocument(/*id=*/1))); |
| DocumentId document_id1 = put_result1.new_document_id; |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2, |
| doc_store->Put(CreateDocument(/*id=*/2))); |
| DocumentId document_id2 = put_result2.new_document_id; |
| |
| std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "name"), |
| GetSectionId("Email", "body")}; |
| SectionIdMask hit_section_id_mask = CreateSectionIdMask(hit_section_ids); |
| std::vector<ScoredDocumentHit> scored_document_hits = { |
| {document_id1, hit_section_id_mask, /*score=*/5}, |
| {document_id2, hit_section_id_mask, /*score=*/0}}; |
| ICING_ASSERT_OK_AND_ASSIGN( |
| std::unique_ptr<ResultRetrieverV2> result_retriever, |
| ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(), |
| language_segmenter_.get(), normalizer_.get())); |
| |
| SearchResultProto::ResultProto result1; |
| *result1.mutable_document() = CreateDocument(/*id=*/1); |
| result1.set_score(5); |
| SearchResultProto::ResultProto result2; |
| *result2.mutable_document() = CreateDocument(/*id=*/2); |
| result2.set_score(0); |
| |
| int threshold = 1; |
| ASSERT_THAT(result1.ByteSizeLong(), Gt(threshold)); |
| |
| ResultSpecProto result_spec = |
| CreateResultSpec(/*num_per_page=*/2, ResultSpecProto::NAMESPACE); |
| result_spec.set_num_total_bytes_per_page_threshold(threshold); |
| ResultStateV2 result_state( |
| std::make_unique< |
| PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>( |
| std::move(scored_document_hits), |
| /*is_descending=*/true), |
| /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr, |
| result_spec, *doc_store); |
| |
| // First page. Should return single result1 even though its byte size exceeds |
| // num_total_bytes_per_page_threshold. |
| auto [page_result1, has_more_results1] = result_retriever->RetrieveNextPage( |
| result_state, fake_clock_.GetSystemTimeMilliseconds()); |
| EXPECT_THAT(page_result1.results, ElementsAre(EqualsProto(result1))); |
| // Has more results. |
| EXPECT_TRUE(has_more_results1); |
| |
| // Second page, result2. |
| auto [page_result2, has_more_results2] = result_retriever->RetrieveNextPage( |
| result_state, fake_clock_.GetSystemTimeMilliseconds()); |
| EXPECT_THAT(page_result2.results, ElementsAre(EqualsProto(result2))); |
| // No more results. |
| EXPECT_FALSE(has_more_results2); |
| } |
| |
| TEST_F(ResultRetrieverV2Test, |
| ShouldRetrieveNextResultWhenBelowNumTotalBytesPerPageThreshold) { |
| ICING_ASSERT_OK_AND_ASSIGN( |
| DocumentStore::CreateResult create_result, |
| CreateDocumentStore(&filesystem_, test_dir_, &fake_clock_, |
| schema_store_.get(), *feature_flags_)); |
| std::unique_ptr<DocumentStore> doc_store = |
| std::move(create_result.document_store); |
| |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1, |
| doc_store->Put(CreateDocument(/*id=*/1))); |
| DocumentId document_id1 = put_result1.new_document_id; |
| ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2, |
| doc_store->Put(CreateDocument(/*id=*/2))); |
| DocumentId document_id2 = put_result2.new_document_id; |
| |
| std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "name"), |
| GetSectionId("Email", "body")}; |
| SectionIdMask hit_section_id_mask = CreateSectionIdMask(hit_section_ids); |
| std::vector<ScoredDocumentHit> scored_document_hits = { |
| {document_id1, hit_section_id_mask, /*score=*/5}, |
| {document_id2, hit_section_id_mask, /*score=*/0}}; |
| ICING_ASSERT_OK_AND_ASSIGN( |
| std::unique_ptr<ResultRetrieverV2> result_retriever, |
| ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(), |
| language_segmenter_.get(), normalizer_.get())); |
| |
| SearchResultProto::ResultProto result1; |
| *result1.mutable_document() = CreateDocument(/*id=*/1); |
| result1.set_score(5); |
| SearchResultProto::ResultProto result2; |
| *result2.mutable_document() = CreateDocument(/*id=*/2); |
| result2.set_score(0); |
| |
| int threshold = result1.ByteSizeLong() + 1; |
| ASSERT_THAT(result1.ByteSizeLong() + result2.ByteSizeLong(), Gt(threshold)); |
| |
| ResultSpecProto result_spec = |
| CreateResultSpec(/*num_per_page=*/2, ResultSpecProto::NAMESPACE); |
| result_spec.set_num_total_bytes_per_page_threshold(threshold); |
| ResultStateV2 result_state( |
| std::make_unique< |
| PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>( |
| std::move(scored_document_hits), |
| /*is_descending=*/true), |
| /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr, |
| result_spec, *doc_store); |
| |
| // After retrieving result1, total bytes are still below the threshold and # |
| // of results is still below num_per_page, so ResultRetriever should continue |
| // the retrieval process and thus include result2 into this page, even though |
| // finally total bytes of result1 + result2 exceed the threshold. |
| auto [page_result, has_more_results] = result_retriever->RetrieveNextPage( |
| result_state, fake_clock_.GetSystemTimeMilliseconds()); |
| EXPECT_THAT(page_result.results, |
| ElementsAre(EqualsProto(result1), EqualsProto(result2))); |
| // No more results. |
| EXPECT_FALSE(has_more_results); |
| } |
| |
| } // namespace |
| |
| } // namespace lib |
| } // namespace icing |
