1*8b6cd535SAndroid Build Coastguard Worker // Copyright (C) 2022 Google LLC
2*8b6cd535SAndroid Build Coastguard Worker //
3*8b6cd535SAndroid Build Coastguard Worker // Licensed under the Apache License, Version 2.0 (the "License");
4*8b6cd535SAndroid Build Coastguard Worker // you may not use this file except in compliance with the License.
5*8b6cd535SAndroid Build Coastguard Worker // You may obtain a copy of the License at
6*8b6cd535SAndroid Build Coastguard Worker //
7*8b6cd535SAndroid Build Coastguard Worker // http://www.apache.org/licenses/LICENSE-2.0
8*8b6cd535SAndroid Build Coastguard Worker //
9*8b6cd535SAndroid Build Coastguard Worker // Unless required by applicable law or agreed to in writing, software
10*8b6cd535SAndroid Build Coastguard Worker // distributed under the License is distributed on an "AS IS" BASIS,
11*8b6cd535SAndroid Build Coastguard Worker // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12*8b6cd535SAndroid Build Coastguard Worker // See the License for the specific language governing permissions and
13*8b6cd535SAndroid Build Coastguard Worker // limitations under the License.
14*8b6cd535SAndroid Build Coastguard Worker
15*8b6cd535SAndroid Build Coastguard Worker #include "icing/result/result-retriever-v2.h"
16*8b6cd535SAndroid Build Coastguard Worker
17*8b6cd535SAndroid Build Coastguard Worker #include <atomic>
18*8b6cd535SAndroid Build Coastguard Worker #include <cstddef>
19*8b6cd535SAndroid Build Coastguard Worker #include <cstdint>
20*8b6cd535SAndroid Build Coastguard Worker #include <memory>
21*8b6cd535SAndroid Build Coastguard Worker #include <string>
22*8b6cd535SAndroid Build Coastguard Worker #include <unordered_map>
23*8b6cd535SAndroid Build Coastguard Worker #include <utility>
24*8b6cd535SAndroid Build Coastguard Worker #include <vector>
25*8b6cd535SAndroid Build Coastguard Worker
26*8b6cd535SAndroid Build Coastguard Worker #include "icing/text_classifier/lib3/utils/base/status.h"
27*8b6cd535SAndroid Build Coastguard Worker #include "icing/text_classifier/lib3/utils/base/statusor.h"
28*8b6cd535SAndroid Build Coastguard Worker #include "gmock/gmock.h"
29*8b6cd535SAndroid Build Coastguard Worker #include "gtest/gtest.h"
30*8b6cd535SAndroid Build Coastguard Worker #include "icing/absl_ports/mutex.h"
31*8b6cd535SAndroid Build Coastguard Worker #include "icing/document-builder.h"
32*8b6cd535SAndroid Build Coastguard Worker #include "icing/feature-flags.h"
33*8b6cd535SAndroid Build Coastguard Worker #include "icing/file/filesystem.h"
34*8b6cd535SAndroid Build Coastguard Worker #include "icing/file/mock-filesystem.h"
35*8b6cd535SAndroid Build Coastguard Worker #include "icing/file/portable-file-backed-proto-log.h"
36*8b6cd535SAndroid Build Coastguard Worker #include "icing/portable/equals-proto.h"
37*8b6cd535SAndroid Build Coastguard Worker #include "icing/portable/platform.h"
38*8b6cd535SAndroid Build Coastguard Worker #include "icing/proto/document.pb.h"
39*8b6cd535SAndroid Build Coastguard Worker #include "icing/proto/document_wrapper.pb.h"
40*8b6cd535SAndroid Build Coastguard Worker #include "icing/proto/schema.pb.h"
41*8b6cd535SAndroid Build Coastguard Worker #include "icing/proto/search.pb.h"
42*8b6cd535SAndroid Build Coastguard Worker #include "icing/result/page-result.h"
43*8b6cd535SAndroid Build Coastguard Worker #include "icing/result/result-state-v2.h"
44*8b6cd535SAndroid Build Coastguard Worker #include "icing/schema-builder.h"
45*8b6cd535SAndroid Build Coastguard Worker #include "icing/schema/schema-store.h"
46*8b6cd535SAndroid Build Coastguard Worker #include "icing/schema/section.h"
47*8b6cd535SAndroid Build Coastguard Worker #include "icing/scoring/priority-queue-scored-document-hits-ranker.h"
48*8b6cd535SAndroid Build Coastguard Worker #include "icing/scoring/scored-document-hit.h"
49*8b6cd535SAndroid Build Coastguard Worker #include "icing/store/document-filter-data.h"
50*8b6cd535SAndroid Build Coastguard Worker #include "icing/store/document-id.h"
51*8b6cd535SAndroid Build Coastguard Worker #include "icing/store/document-store.h"
52*8b6cd535SAndroid Build Coastguard Worker #include "icing/testing/common-matchers.h"
53*8b6cd535SAndroid Build Coastguard Worker #include "icing/testing/fake-clock.h"
54*8b6cd535SAndroid Build Coastguard Worker #include "icing/testing/test-data.h"
55*8b6cd535SAndroid Build Coastguard Worker #include "icing/testing/test-feature-flags.h"
56*8b6cd535SAndroid Build Coastguard Worker #include "icing/testing/tmp-directory.h"
57*8b6cd535SAndroid Build Coastguard Worker #include "icing/tokenization/language-segmenter-factory.h"
58*8b6cd535SAndroid Build Coastguard Worker #include "icing/tokenization/language-segmenter.h"
59*8b6cd535SAndroid Build Coastguard Worker #include "icing/transform/normalizer-factory.h"
60*8b6cd535SAndroid Build Coastguard Worker #include "icing/transform/normalizer.h"
61*8b6cd535SAndroid Build Coastguard Worker #include "icing/util/clock.h"
62*8b6cd535SAndroid Build Coastguard Worker #include "icing/util/icu-data-file-helper.h"
63*8b6cd535SAndroid Build Coastguard Worker #include "unicode/uloc.h"
64*8b6cd535SAndroid Build Coastguard Worker
65*8b6cd535SAndroid Build Coastguard Worker namespace icing {
66*8b6cd535SAndroid Build Coastguard Worker namespace lib {
67*8b6cd535SAndroid Build Coastguard Worker
68*8b6cd535SAndroid Build Coastguard Worker namespace {
69*8b6cd535SAndroid Build Coastguard Worker
70*8b6cd535SAndroid Build Coastguard Worker using ::icing::lib::portable_equals_proto::EqualsProto;
71*8b6cd535SAndroid Build Coastguard Worker using ::testing::DoDefault;
72*8b6cd535SAndroid Build Coastguard Worker using ::testing::ElementsAre;
73*8b6cd535SAndroid Build Coastguard Worker using ::testing::Eq;
74*8b6cd535SAndroid Build Coastguard Worker using ::testing::Gt;
75*8b6cd535SAndroid Build Coastguard Worker using ::testing::IsEmpty;
76*8b6cd535SAndroid Build Coastguard Worker using ::testing::Pointee;
77*8b6cd535SAndroid Build Coastguard Worker using ::testing::Return;
78*8b6cd535SAndroid Build Coastguard Worker using ::testing::SizeIs;
79*8b6cd535SAndroid Build Coastguard Worker using EntryIdMap = std::unordered_map<int32_t, int>;
80*8b6cd535SAndroid Build Coastguard Worker
81*8b6cd535SAndroid Build Coastguard Worker // Mock the behavior of GroupResultLimiter::ShouldBeRemoved.
82*8b6cd535SAndroid Build Coastguard Worker class MockGroupResultLimiter : public GroupResultLimiterV2 {
83*8b6cd535SAndroid Build Coastguard Worker public:
MockGroupResultLimiter()84*8b6cd535SAndroid Build Coastguard Worker MockGroupResultLimiter() : GroupResultLimiterV2() {
85*8b6cd535SAndroid Build Coastguard Worker ON_CALL(*this, ShouldBeRemoved).WillByDefault(Return(false));
86*8b6cd535SAndroid Build Coastguard Worker }
87*8b6cd535SAndroid Build Coastguard Worker
88*8b6cd535SAndroid Build Coastguard Worker MOCK_METHOD(bool, ShouldBeRemoved,
89*8b6cd535SAndroid Build Coastguard Worker (const ScoredDocumentHit&, const EntryIdMap&,
90*8b6cd535SAndroid Build Coastguard Worker const DocumentStore&, std::vector<int>&,
91*8b6cd535SAndroid Build Coastguard Worker ResultSpecProto::ResultGroupingType, int64_t),
92*8b6cd535SAndroid Build Coastguard Worker (const, override));
93*8b6cd535SAndroid Build Coastguard Worker };
94*8b6cd535SAndroid Build Coastguard Worker
95*8b6cd535SAndroid Build Coastguard Worker class ResultRetrieverV2Test : public ::testing::Test {
96*8b6cd535SAndroid Build Coastguard Worker protected:
ResultRetrieverV2Test()97*8b6cd535SAndroid Build Coastguard Worker ResultRetrieverV2Test() : test_dir_(GetTestTempDir() + "/icing") {
98*8b6cd535SAndroid Build Coastguard Worker filesystem_.CreateDirectoryRecursively(test_dir_.c_str());
99*8b6cd535SAndroid Build Coastguard Worker }
100*8b6cd535SAndroid Build Coastguard Worker
SetUp()101*8b6cd535SAndroid Build Coastguard Worker void SetUp() override {
102*8b6cd535SAndroid Build Coastguard Worker feature_flags_ = std::make_unique<FeatureFlags>(GetTestFeatureFlags());
103*8b6cd535SAndroid Build Coastguard Worker if (!IsCfStringTokenization() && !IsReverseJniTokenization()) {
104*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK(
105*8b6cd535SAndroid Build Coastguard Worker // File generated via icu_data_file rule in //icing/BUILD.
106*8b6cd535SAndroid Build Coastguard Worker icu_data_file_helper::SetUpIcuDataFile(
107*8b6cd535SAndroid Build Coastguard Worker GetTestFilePath("icing/icu.dat")));
108*8b6cd535SAndroid Build Coastguard Worker }
109*8b6cd535SAndroid Build Coastguard Worker language_segmenter_factory::SegmenterOptions options(ULOC_US);
110*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(
111*8b6cd535SAndroid Build Coastguard Worker language_segmenter_,
112*8b6cd535SAndroid Build Coastguard Worker language_segmenter_factory::Create(std::move(options)));
113*8b6cd535SAndroid Build Coastguard Worker
114*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(
115*8b6cd535SAndroid Build Coastguard Worker schema_store_, SchemaStore::Create(&filesystem_, test_dir_,
116*8b6cd535SAndroid Build Coastguard Worker &fake_clock_, feature_flags_.get()));
117*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(normalizer_, normalizer_factory::Create(
118*8b6cd535SAndroid Build Coastguard Worker /*max_term_byte_size=*/10000));
119*8b6cd535SAndroid Build Coastguard Worker
120*8b6cd535SAndroid Build Coastguard Worker SchemaProto schema =
121*8b6cd535SAndroid Build Coastguard Worker SchemaBuilder()
122*8b6cd535SAndroid Build Coastguard Worker .AddType(SchemaTypeConfigBuilder()
123*8b6cd535SAndroid Build Coastguard Worker .SetType("Email")
124*8b6cd535SAndroid Build Coastguard Worker .AddProperty(PropertyConfigBuilder()
125*8b6cd535SAndroid Build Coastguard Worker .SetName("name")
126*8b6cd535SAndroid Build Coastguard Worker .SetDataTypeString(TERM_MATCH_PREFIX,
127*8b6cd535SAndroid Build Coastguard Worker TOKENIZER_PLAIN)
128*8b6cd535SAndroid Build Coastguard Worker .SetCardinality(CARDINALITY_OPTIONAL))
129*8b6cd535SAndroid Build Coastguard Worker .AddProperty(PropertyConfigBuilder()
130*8b6cd535SAndroid Build Coastguard Worker .SetName("body")
131*8b6cd535SAndroid Build Coastguard Worker .SetDataTypeString(TERM_MATCH_EXACT,
132*8b6cd535SAndroid Build Coastguard Worker TOKENIZER_PLAIN)
133*8b6cd535SAndroid Build Coastguard Worker .SetCardinality(CARDINALITY_OPTIONAL))
134*8b6cd535SAndroid Build Coastguard Worker .AddProperty(
135*8b6cd535SAndroid Build Coastguard Worker PropertyConfigBuilder()
136*8b6cd535SAndroid Build Coastguard Worker .SetName("sender")
137*8b6cd535SAndroid Build Coastguard Worker .SetDataTypeDocument(
138*8b6cd535SAndroid Build Coastguard Worker "Person", /*index_nested_properties=*/true)
139*8b6cd535SAndroid Build Coastguard Worker .SetCardinality(CARDINALITY_OPTIONAL)))
140*8b6cd535SAndroid Build Coastguard Worker .AddType(
141*8b6cd535SAndroid Build Coastguard Worker SchemaTypeConfigBuilder()
142*8b6cd535SAndroid Build Coastguard Worker .SetType("Person")
143*8b6cd535SAndroid Build Coastguard Worker .AddProperty(PropertyConfigBuilder()
144*8b6cd535SAndroid Build Coastguard Worker .SetName("name")
145*8b6cd535SAndroid Build Coastguard Worker .SetDataTypeString(TERM_MATCH_PREFIX,
146*8b6cd535SAndroid Build Coastguard Worker TOKENIZER_PLAIN)
147*8b6cd535SAndroid Build Coastguard Worker .SetCardinality(CARDINALITY_OPTIONAL))
148*8b6cd535SAndroid Build Coastguard Worker .AddProperty(PropertyConfigBuilder()
149*8b6cd535SAndroid Build Coastguard Worker .SetName("emailAddress")
150*8b6cd535SAndroid Build Coastguard Worker .SetDataTypeString(TERM_MATCH_PREFIX,
151*8b6cd535SAndroid Build Coastguard Worker TOKENIZER_PLAIN)
152*8b6cd535SAndroid Build Coastguard Worker .SetCardinality(CARDINALITY_OPTIONAL)))
153*8b6cd535SAndroid Build Coastguard Worker .Build();
154*8b6cd535SAndroid Build Coastguard Worker ASSERT_THAT(schema_store_->SetSchema(
155*8b6cd535SAndroid Build Coastguard Worker schema, /*ignore_errors_and_delete_documents=*/false,
156*8b6cd535SAndroid Build Coastguard Worker /*allow_circular_schema_definitions=*/false),
157*8b6cd535SAndroid Build Coastguard Worker IsOk());
158*8b6cd535SAndroid Build Coastguard Worker
159*8b6cd535SAndroid Build Coastguard Worker num_total_hits_ = 0;
160*8b6cd535SAndroid Build Coastguard Worker }
161*8b6cd535SAndroid Build Coastguard Worker
TearDown()162*8b6cd535SAndroid Build Coastguard Worker void TearDown() override {
163*8b6cd535SAndroid Build Coastguard Worker filesystem_.DeleteDirectoryRecursively(test_dir_.c_str());
164*8b6cd535SAndroid Build Coastguard Worker }
165*8b6cd535SAndroid Build Coastguard Worker
GetSectionId(const std::string & type,const std::string & property)166*8b6cd535SAndroid Build Coastguard Worker SectionId GetSectionId(const std::string& type, const std::string& property) {
167*8b6cd535SAndroid Build Coastguard Worker auto type_id_or = schema_store_->GetSchemaTypeId(type);
168*8b6cd535SAndroid Build Coastguard Worker if (!type_id_or.ok()) {
169*8b6cd535SAndroid Build Coastguard Worker return kInvalidSectionId;
170*8b6cd535SAndroid Build Coastguard Worker }
171*8b6cd535SAndroid Build Coastguard Worker SchemaTypeId type_id = type_id_or.ValueOrDie();
172*8b6cd535SAndroid Build Coastguard Worker for (SectionId section_id = 0; section_id <= kMaxSectionId; ++section_id) {
173*8b6cd535SAndroid Build Coastguard Worker auto metadata_or = schema_store_->GetSectionMetadata(type_id, section_id);
174*8b6cd535SAndroid Build Coastguard Worker if (!metadata_or.ok()) {
175*8b6cd535SAndroid Build Coastguard Worker break;
176*8b6cd535SAndroid Build Coastguard Worker }
177*8b6cd535SAndroid Build Coastguard Worker const SectionMetadata* metadata = metadata_or.ValueOrDie();
178*8b6cd535SAndroid Build Coastguard Worker if (metadata->path == property) {
179*8b6cd535SAndroid Build Coastguard Worker return metadata->id;
180*8b6cd535SAndroid Build Coastguard Worker }
181*8b6cd535SAndroid Build Coastguard Worker }
182*8b6cd535SAndroid Build Coastguard Worker return kInvalidSectionId;
183*8b6cd535SAndroid Build Coastguard Worker }
184*8b6cd535SAndroid Build Coastguard Worker
185*8b6cd535SAndroid Build Coastguard Worker std::unique_ptr<FeatureFlags> feature_flags_;
186*8b6cd535SAndroid Build Coastguard Worker const Filesystem filesystem_;
187*8b6cd535SAndroid Build Coastguard Worker const std::string test_dir_;
188*8b6cd535SAndroid Build Coastguard Worker std::unique_ptr<LanguageSegmenter> language_segmenter_;
189*8b6cd535SAndroid Build Coastguard Worker std::unique_ptr<SchemaStore> schema_store_;
190*8b6cd535SAndroid Build Coastguard Worker std::unique_ptr<Normalizer> normalizer_;
191*8b6cd535SAndroid Build Coastguard Worker std::atomic<int> num_total_hits_;
192*8b6cd535SAndroid Build Coastguard Worker FakeClock fake_clock_;
193*8b6cd535SAndroid Build Coastguard Worker };
194*8b6cd535SAndroid Build Coastguard Worker
CreateDocument(int id)195*8b6cd535SAndroid Build Coastguard Worker DocumentProto CreateDocument(int id) {
196*8b6cd535SAndroid Build Coastguard Worker return DocumentBuilder()
197*8b6cd535SAndroid Build Coastguard Worker .SetKey("icing", "Email/" + std::to_string(id))
198*8b6cd535SAndroid Build Coastguard Worker .SetSchema("Email")
199*8b6cd535SAndroid Build Coastguard Worker .AddStringProperty("name", "subject foo " + std::to_string(id))
200*8b6cd535SAndroid Build Coastguard Worker .AddStringProperty("body", "body bar " + std::to_string(id))
201*8b6cd535SAndroid Build Coastguard Worker .SetCreationTimestampMs(1574365086666 + id)
202*8b6cd535SAndroid Build Coastguard Worker .Build();
203*8b6cd535SAndroid Build Coastguard Worker }
204*8b6cd535SAndroid Build Coastguard Worker
CreateSectionIdMask(const std::vector<SectionId> & section_ids)205*8b6cd535SAndroid Build Coastguard Worker SectionIdMask CreateSectionIdMask(const std::vector<SectionId>& section_ids) {
206*8b6cd535SAndroid Build Coastguard Worker SectionIdMask mask = 0;
207*8b6cd535SAndroid Build Coastguard Worker for (SectionId section_id : section_ids) {
208*8b6cd535SAndroid Build Coastguard Worker mask |= (UINT64_C(1) << section_id);
209*8b6cd535SAndroid Build Coastguard Worker }
210*8b6cd535SAndroid Build Coastguard Worker return mask;
211*8b6cd535SAndroid Build Coastguard Worker }
212*8b6cd535SAndroid Build Coastguard Worker
CreateResultSpec(int num_per_page,ResultSpecProto::ResultGroupingType result_group_type)213*8b6cd535SAndroid Build Coastguard Worker ResultSpecProto CreateResultSpec(
214*8b6cd535SAndroid Build Coastguard Worker int num_per_page, ResultSpecProto::ResultGroupingType result_group_type) {
215*8b6cd535SAndroid Build Coastguard Worker ResultSpecProto result_spec;
216*8b6cd535SAndroid Build Coastguard Worker result_spec.set_result_group_type(result_group_type);
217*8b6cd535SAndroid Build Coastguard Worker result_spec.set_num_per_page(num_per_page);
218*8b6cd535SAndroid Build Coastguard Worker return result_spec;
219*8b6cd535SAndroid Build Coastguard Worker }
220*8b6cd535SAndroid Build Coastguard Worker
CreateDocumentStore(const Filesystem * filesystem,const std::string & base_dir,const Clock * clock,const SchemaStore * schema_store,const FeatureFlags & feature_flags)221*8b6cd535SAndroid Build Coastguard Worker libtextclassifier3::StatusOr<DocumentStore::CreateResult> CreateDocumentStore(
222*8b6cd535SAndroid Build Coastguard Worker const Filesystem* filesystem, const std::string& base_dir,
223*8b6cd535SAndroid Build Coastguard Worker const Clock* clock, const SchemaStore* schema_store,
224*8b6cd535SAndroid Build Coastguard Worker const FeatureFlags& feature_flags) {
225*8b6cd535SAndroid Build Coastguard Worker return DocumentStore::Create(
226*8b6cd535SAndroid Build Coastguard Worker filesystem, base_dir, clock, schema_store, &feature_flags,
227*8b6cd535SAndroid Build Coastguard Worker /*force_recovery_and_revalidate_documents=*/false,
228*8b6cd535SAndroid Build Coastguard Worker /*pre_mapping_fbv=*/false, /*use_persistent_hash_map=*/true,
229*8b6cd535SAndroid Build Coastguard Worker PortableFileBackedProtoLog<DocumentWrapper>::kDefaultCompressionLevel,
230*8b6cd535SAndroid Build Coastguard Worker /*initialize_stats=*/nullptr);
231*8b6cd535SAndroid Build Coastguard Worker }
232*8b6cd535SAndroid Build Coastguard Worker
TEST_F(ResultRetrieverV2Test,CreationWithNullPointerShouldFail)233*8b6cd535SAndroid Build Coastguard Worker TEST_F(ResultRetrieverV2Test, CreationWithNullPointerShouldFail) {
234*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(
235*8b6cd535SAndroid Build Coastguard Worker ResultRetrieverV2::Create(/*doc_store=*/nullptr, schema_store_.get(),
236*8b6cd535SAndroid Build Coastguard Worker language_segmenter_.get(), normalizer_.get()),
237*8b6cd535SAndroid Build Coastguard Worker StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
238*8b6cd535SAndroid Build Coastguard Worker
239*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(
240*8b6cd535SAndroid Build Coastguard Worker DocumentStore::CreateResult create_result,
241*8b6cd535SAndroid Build Coastguard Worker CreateDocumentStore(&filesystem_, test_dir_, &fake_clock_,
242*8b6cd535SAndroid Build Coastguard Worker schema_store_.get(), *feature_flags_));
243*8b6cd535SAndroid Build Coastguard Worker std::unique_ptr<DocumentStore> doc_store =
244*8b6cd535SAndroid Build Coastguard Worker std::move(create_result.document_store);
245*8b6cd535SAndroid Build Coastguard Worker
246*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(
247*8b6cd535SAndroid Build Coastguard Worker ResultRetrieverV2::Create(doc_store.get(), /*schema_store=*/nullptr,
248*8b6cd535SAndroid Build Coastguard Worker language_segmenter_.get(), normalizer_.get()),
249*8b6cd535SAndroid Build Coastguard Worker StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
250*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(),
251*8b6cd535SAndroid Build Coastguard Worker /*language_segmenter=*/nullptr,
252*8b6cd535SAndroid Build Coastguard Worker normalizer_.get()),
253*8b6cd535SAndroid Build Coastguard Worker StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
254*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(),
255*8b6cd535SAndroid Build Coastguard Worker language_segmenter_.get(),
256*8b6cd535SAndroid Build Coastguard Worker /*normalizer=*/nullptr),
257*8b6cd535SAndroid Build Coastguard Worker StatusIs(libtextclassifier3::StatusCode::FAILED_PRECONDITION));
258*8b6cd535SAndroid Build Coastguard Worker }
259*8b6cd535SAndroid Build Coastguard Worker
TEST_F(ResultRetrieverV2Test,ShouldRetrieveSimpleResults)260*8b6cd535SAndroid Build Coastguard Worker TEST_F(ResultRetrieverV2Test, ShouldRetrieveSimpleResults) {
261*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(
262*8b6cd535SAndroid Build Coastguard Worker DocumentStore::CreateResult create_result,
263*8b6cd535SAndroid Build Coastguard Worker CreateDocumentStore(&filesystem_, test_dir_, &fake_clock_,
264*8b6cd535SAndroid Build Coastguard Worker schema_store_.get(), *feature_flags_));
265*8b6cd535SAndroid Build Coastguard Worker std::unique_ptr<DocumentStore> doc_store =
266*8b6cd535SAndroid Build Coastguard Worker std::move(create_result.document_store);
267*8b6cd535SAndroid Build Coastguard Worker
268*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
269*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(CreateDocument(/*id=*/1)));
270*8b6cd535SAndroid Build Coastguard Worker DocumentId document_id1 = put_result1.new_document_id;
271*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
272*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(CreateDocument(/*id=*/2)));
273*8b6cd535SAndroid Build Coastguard Worker DocumentId document_id2 = put_result2.new_document_id;
274*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
275*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(CreateDocument(/*id=*/3)));
276*8b6cd535SAndroid Build Coastguard Worker DocumentId document_id3 = put_result3.new_document_id;
277*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result4,
278*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(CreateDocument(/*id=*/4)));
279*8b6cd535SAndroid Build Coastguard Worker DocumentId document_id4 = put_result4.new_document_id;
280*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result5,
281*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(CreateDocument(/*id=*/5)));
282*8b6cd535SAndroid Build Coastguard Worker DocumentId document_id5 = put_result5.new_document_id;
283*8b6cd535SAndroid Build Coastguard Worker
284*8b6cd535SAndroid Build Coastguard Worker std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "name"),
285*8b6cd535SAndroid Build Coastguard Worker GetSectionId("Email", "body")};
286*8b6cd535SAndroid Build Coastguard Worker SectionIdMask hit_section_id_mask = CreateSectionIdMask(hit_section_ids);
287*8b6cd535SAndroid Build Coastguard Worker std::vector<ScoredDocumentHit> scored_document_hits = {
288*8b6cd535SAndroid Build Coastguard Worker {document_id1, hit_section_id_mask, /*score=*/19},
289*8b6cd535SAndroid Build Coastguard Worker {document_id2, hit_section_id_mask, /*score=*/12},
290*8b6cd535SAndroid Build Coastguard Worker {document_id3, hit_section_id_mask, /*score=*/8},
291*8b6cd535SAndroid Build Coastguard Worker {document_id4, hit_section_id_mask, /*score=*/3},
292*8b6cd535SAndroid Build Coastguard Worker {document_id5, hit_section_id_mask, /*score=*/1}};
293*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(
294*8b6cd535SAndroid Build Coastguard Worker std::unique_ptr<ResultRetrieverV2> result_retriever,
295*8b6cd535SAndroid Build Coastguard Worker ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(),
296*8b6cd535SAndroid Build Coastguard Worker language_segmenter_.get(), normalizer_.get()));
297*8b6cd535SAndroid Build Coastguard Worker
298*8b6cd535SAndroid Build Coastguard Worker SearchResultProto::ResultProto result1;
299*8b6cd535SAndroid Build Coastguard Worker *result1.mutable_document() = CreateDocument(/*id=*/1);
300*8b6cd535SAndroid Build Coastguard Worker result1.set_score(19);
301*8b6cd535SAndroid Build Coastguard Worker SearchResultProto::ResultProto result2;
302*8b6cd535SAndroid Build Coastguard Worker *result2.mutable_document() = CreateDocument(/*id=*/2);
303*8b6cd535SAndroid Build Coastguard Worker result2.set_score(12);
304*8b6cd535SAndroid Build Coastguard Worker SearchResultProto::ResultProto result3;
305*8b6cd535SAndroid Build Coastguard Worker *result3.mutable_document() = CreateDocument(/*id=*/3);
306*8b6cd535SAndroid Build Coastguard Worker result3.set_score(8);
307*8b6cd535SAndroid Build Coastguard Worker SearchResultProto::ResultProto result4;
308*8b6cd535SAndroid Build Coastguard Worker *result4.mutable_document() = CreateDocument(/*id=*/4);
309*8b6cd535SAndroid Build Coastguard Worker result4.set_score(3);
310*8b6cd535SAndroid Build Coastguard Worker SearchResultProto::ResultProto result5;
311*8b6cd535SAndroid Build Coastguard Worker *result5.mutable_document() = CreateDocument(/*id=*/5);
312*8b6cd535SAndroid Build Coastguard Worker result5.set_score(1);
313*8b6cd535SAndroid Build Coastguard Worker
314*8b6cd535SAndroid Build Coastguard Worker ResultStateV2 result_state(
315*8b6cd535SAndroid Build Coastguard Worker std::make_unique<
316*8b6cd535SAndroid Build Coastguard Worker PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
317*8b6cd535SAndroid Build Coastguard Worker std::move(scored_document_hits), /*is_descending=*/true),
318*8b6cd535SAndroid Build Coastguard Worker /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
319*8b6cd535SAndroid Build Coastguard Worker CreateResultSpec(/*num_per_page=*/2, ResultSpecProto::NAMESPACE),
320*8b6cd535SAndroid Build Coastguard Worker *doc_store);
321*8b6cd535SAndroid Build Coastguard Worker
322*8b6cd535SAndroid Build Coastguard Worker // First page, 2 results
323*8b6cd535SAndroid Build Coastguard Worker auto [page_result1, has_more_results1] = result_retriever->RetrieveNextPage(
324*8b6cd535SAndroid Build Coastguard Worker result_state, fake_clock_.GetSystemTimeMilliseconds());
325*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(page_result1.results,
326*8b6cd535SAndroid Build Coastguard Worker ElementsAre(EqualsProto(result1), EqualsProto(result2)));
327*8b6cd535SAndroid Build Coastguard Worker // num_results_with_snippets is 0 when there is no snippet.
328*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(page_result1.num_results_with_snippets, Eq(0));
329*8b6cd535SAndroid Build Coastguard Worker // Requested page size is same as num_per_page.
330*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(page_result1.requested_page_size, Eq(2));
331*8b6cd535SAndroid Build Coastguard Worker // Has more results.
332*8b6cd535SAndroid Build Coastguard Worker EXPECT_TRUE(has_more_results1);
333*8b6cd535SAndroid Build Coastguard Worker
334*8b6cd535SAndroid Build Coastguard Worker // Second page, 2 results
335*8b6cd535SAndroid Build Coastguard Worker auto [page_result2, has_more_results2] = result_retriever->RetrieveNextPage(
336*8b6cd535SAndroid Build Coastguard Worker result_state, fake_clock_.GetSystemTimeMilliseconds());
337*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(page_result2.results,
338*8b6cd535SAndroid Build Coastguard Worker ElementsAre(EqualsProto(result3), EqualsProto(result4)));
339*8b6cd535SAndroid Build Coastguard Worker // num_results_with_snippets is 0 when there is no snippet.
340*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(page_result2.num_results_with_snippets, Eq(0));
341*8b6cd535SAndroid Build Coastguard Worker // Requested page size is same as num_per_page.
342*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(page_result2.requested_page_size, Eq(2));
343*8b6cd535SAndroid Build Coastguard Worker // Has more results.
344*8b6cd535SAndroid Build Coastguard Worker EXPECT_TRUE(has_more_results2);
345*8b6cd535SAndroid Build Coastguard Worker
346*8b6cd535SAndroid Build Coastguard Worker // Third page, 1 result
347*8b6cd535SAndroid Build Coastguard Worker auto [page_result3, has_more_results3] = result_retriever->RetrieveNextPage(
348*8b6cd535SAndroid Build Coastguard Worker result_state, fake_clock_.GetSystemTimeMilliseconds());
349*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(page_result3.results, ElementsAre(EqualsProto(result5)));
350*8b6cd535SAndroid Build Coastguard Worker // num_results_with_snippets is 0 when there is no snippet.
351*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(page_result3.num_results_with_snippets, Eq(0));
352*8b6cd535SAndroid Build Coastguard Worker // Requested page size is same as num_per_page.
353*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(page_result3.requested_page_size, Eq(2));
354*8b6cd535SAndroid Build Coastguard Worker // No more results.
355*8b6cd535SAndroid Build Coastguard Worker EXPECT_FALSE(has_more_results3);
356*8b6cd535SAndroid Build Coastguard Worker }
357*8b6cd535SAndroid Build Coastguard Worker
TEST_F(ResultRetrieverV2Test,ShouldIgnoreNonInternalErrors)358*8b6cd535SAndroid Build Coastguard Worker TEST_F(ResultRetrieverV2Test, ShouldIgnoreNonInternalErrors) {
359*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(
360*8b6cd535SAndroid Build Coastguard Worker DocumentStore::CreateResult create_result,
361*8b6cd535SAndroid Build Coastguard Worker CreateDocumentStore(&filesystem_, test_dir_, &fake_clock_,
362*8b6cd535SAndroid Build Coastguard Worker schema_store_.get(), *feature_flags_));
363*8b6cd535SAndroid Build Coastguard Worker std::unique_ptr<DocumentStore> doc_store =
364*8b6cd535SAndroid Build Coastguard Worker std::move(create_result.document_store);
365*8b6cd535SAndroid Build Coastguard Worker
366*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
367*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(CreateDocument(/*id=*/1)));
368*8b6cd535SAndroid Build Coastguard Worker DocumentId document_id1 = put_result1.new_document_id;
369*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
370*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(CreateDocument(/*id=*/2)));
371*8b6cd535SAndroid Build Coastguard Worker DocumentId document_id2 = put_result2.new_document_id;
372*8b6cd535SAndroid Build Coastguard Worker
373*8b6cd535SAndroid Build Coastguard Worker DocumentId invalid_document_id = -1;
374*8b6cd535SAndroid Build Coastguard Worker std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "name"),
375*8b6cd535SAndroid Build Coastguard Worker GetSectionId("Email", "body")};
376*8b6cd535SAndroid Build Coastguard Worker SectionIdMask hit_section_id_mask = CreateSectionIdMask(hit_section_ids);
377*8b6cd535SAndroid Build Coastguard Worker std::vector<ScoredDocumentHit> scored_document_hits = {
378*8b6cd535SAndroid Build Coastguard Worker {document_id1, hit_section_id_mask, /*score=*/12},
379*8b6cd535SAndroid Build Coastguard Worker {document_id2, hit_section_id_mask, /*score=*/4},
380*8b6cd535SAndroid Build Coastguard Worker {invalid_document_id, hit_section_id_mask, /*score=*/0}};
381*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(
382*8b6cd535SAndroid Build Coastguard Worker std::unique_ptr<ResultRetrieverV2> result_retriever,
383*8b6cd535SAndroid Build Coastguard Worker ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(),
384*8b6cd535SAndroid Build Coastguard Worker language_segmenter_.get(), normalizer_.get(),
385*8b6cd535SAndroid Build Coastguard Worker std::make_unique<MockGroupResultLimiter>()));
386*8b6cd535SAndroid Build Coastguard Worker
387*8b6cd535SAndroid Build Coastguard Worker SearchResultProto::ResultProto result1;
388*8b6cd535SAndroid Build Coastguard Worker *result1.mutable_document() = CreateDocument(/*id=*/1);
389*8b6cd535SAndroid Build Coastguard Worker result1.set_score(12);
390*8b6cd535SAndroid Build Coastguard Worker SearchResultProto::ResultProto result2;
391*8b6cd535SAndroid Build Coastguard Worker *result2.mutable_document() = CreateDocument(/*id=*/2);
392*8b6cd535SAndroid Build Coastguard Worker result2.set_score(4);
393*8b6cd535SAndroid Build Coastguard Worker
394*8b6cd535SAndroid Build Coastguard Worker ResultStateV2 result_state1(
395*8b6cd535SAndroid Build Coastguard Worker std::make_unique<
396*8b6cd535SAndroid Build Coastguard Worker PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
397*8b6cd535SAndroid Build Coastguard Worker std::move(scored_document_hits),
398*8b6cd535SAndroid Build Coastguard Worker /*is_descending=*/true),
399*8b6cd535SAndroid Build Coastguard Worker /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
400*8b6cd535SAndroid Build Coastguard Worker CreateResultSpec(/*num_per_page=*/3, ResultSpecProto::NAMESPACE),
401*8b6cd535SAndroid Build Coastguard Worker *doc_store);
402*8b6cd535SAndroid Build Coastguard Worker PageResult page_result1 =
403*8b6cd535SAndroid Build Coastguard Worker result_retriever
404*8b6cd535SAndroid Build Coastguard Worker ->RetrieveNextPage(result_state1,
405*8b6cd535SAndroid Build Coastguard Worker fake_clock_.GetSystemTimeMilliseconds())
406*8b6cd535SAndroid Build Coastguard Worker .first;
407*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(page_result1.results,
408*8b6cd535SAndroid Build Coastguard Worker ElementsAre(EqualsProto(result1), EqualsProto(result2)));
409*8b6cd535SAndroid Build Coastguard Worker
410*8b6cd535SAndroid Build Coastguard Worker DocumentId non_existing_document_id = 4;
411*8b6cd535SAndroid Build Coastguard Worker scored_document_hits = {
412*8b6cd535SAndroid Build Coastguard Worker {non_existing_document_id, hit_section_id_mask, /*score=*/15},
413*8b6cd535SAndroid Build Coastguard Worker {document_id1, hit_section_id_mask, /*score=*/12},
414*8b6cd535SAndroid Build Coastguard Worker {document_id2, hit_section_id_mask, /*score=*/4}};
415*8b6cd535SAndroid Build Coastguard Worker ResultStateV2 result_state2(
416*8b6cd535SAndroid Build Coastguard Worker std::make_unique<
417*8b6cd535SAndroid Build Coastguard Worker PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
418*8b6cd535SAndroid Build Coastguard Worker std::move(scored_document_hits),
419*8b6cd535SAndroid Build Coastguard Worker /*is_descending=*/true),
420*8b6cd535SAndroid Build Coastguard Worker /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
421*8b6cd535SAndroid Build Coastguard Worker CreateResultSpec(/*num_per_page=*/3, ResultSpecProto::NAMESPACE),
422*8b6cd535SAndroid Build Coastguard Worker *doc_store);
423*8b6cd535SAndroid Build Coastguard Worker PageResult page_result2 =
424*8b6cd535SAndroid Build Coastguard Worker result_retriever
425*8b6cd535SAndroid Build Coastguard Worker ->RetrieveNextPage(result_state2,
426*8b6cd535SAndroid Build Coastguard Worker fake_clock_.GetSystemTimeMilliseconds())
427*8b6cd535SAndroid Build Coastguard Worker .first;
428*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(page_result2.results,
429*8b6cd535SAndroid Build Coastguard Worker ElementsAre(EqualsProto(result1), EqualsProto(result2)));
430*8b6cd535SAndroid Build Coastguard Worker }
431*8b6cd535SAndroid Build Coastguard Worker
TEST_F(ResultRetrieverV2Test,ShouldLimitNumChildDocumentsByMaxJoinedChildPerParent)432*8b6cd535SAndroid Build Coastguard Worker TEST_F(ResultRetrieverV2Test,
433*8b6cd535SAndroid Build Coastguard Worker ShouldLimitNumChildDocumentsByMaxJoinedChildPerParent) {
434*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(
435*8b6cd535SAndroid Build Coastguard Worker DocumentStore::CreateResult create_result,
436*8b6cd535SAndroid Build Coastguard Worker CreateDocumentStore(&filesystem_, test_dir_, &fake_clock_,
437*8b6cd535SAndroid Build Coastguard Worker schema_store_.get(), *feature_flags_));
438*8b6cd535SAndroid Build Coastguard Worker std::unique_ptr<DocumentStore> doc_store =
439*8b6cd535SAndroid Build Coastguard Worker std::move(create_result.document_store);
440*8b6cd535SAndroid Build Coastguard Worker
441*8b6cd535SAndroid Build Coastguard Worker // 1. Add 2 Person document
442*8b6cd535SAndroid Build Coastguard Worker DocumentProto person_document1 =
443*8b6cd535SAndroid Build Coastguard Worker DocumentBuilder()
444*8b6cd535SAndroid Build Coastguard Worker .SetKey("namespace", "Person/1")
445*8b6cd535SAndroid Build Coastguard Worker .SetCreationTimestampMs(1000)
446*8b6cd535SAndroid Build Coastguard Worker .SetSchema("Person")
447*8b6cd535SAndroid Build Coastguard Worker .AddStringProperty("name", "Joe Fox")
448*8b6cd535SAndroid Build Coastguard Worker .AddStringProperty("emailAddress", "[email protected]")
449*8b6cd535SAndroid Build Coastguard Worker .Build();
450*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
451*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(person_document1));
452*8b6cd535SAndroid Build Coastguard Worker DocumentId person_document_id1 = put_result1.new_document_id;
453*8b6cd535SAndroid Build Coastguard Worker
454*8b6cd535SAndroid Build Coastguard Worker DocumentProto person_document2 =
455*8b6cd535SAndroid Build Coastguard Worker DocumentBuilder()
456*8b6cd535SAndroid Build Coastguard Worker .SetKey("namespace", "Person/2")
457*8b6cd535SAndroid Build Coastguard Worker .SetCreationTimestampMs(1000)
458*8b6cd535SAndroid Build Coastguard Worker .SetSchema("Person")
459*8b6cd535SAndroid Build Coastguard Worker .AddStringProperty("name", "Meg Ryan")
460*8b6cd535SAndroid Build Coastguard Worker .AddStringProperty("emailAddress", "[email protected]")
461*8b6cd535SAndroid Build Coastguard Worker .Build();
462*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
463*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(person_document2));
464*8b6cd535SAndroid Build Coastguard Worker DocumentId person_document_id2 = put_result2.new_document_id;
465*8b6cd535SAndroid Build Coastguard Worker
466*8b6cd535SAndroid Build Coastguard Worker // 2. Add 4 Email documents
467*8b6cd535SAndroid Build Coastguard Worker DocumentProto email_document1 = DocumentBuilder()
468*8b6cd535SAndroid Build Coastguard Worker .SetKey("namespace", "Email/1")
469*8b6cd535SAndroid Build Coastguard Worker .SetCreationTimestampMs(1000)
470*8b6cd535SAndroid Build Coastguard Worker .SetSchema("Email")
471*8b6cd535SAndroid Build Coastguard Worker .AddStringProperty("name", "Test 1")
472*8b6cd535SAndroid Build Coastguard Worker .AddStringProperty("body", "Test 1")
473*8b6cd535SAndroid Build Coastguard Worker .Build();
474*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(put_result1, doc_store->Put(email_document1));
475*8b6cd535SAndroid Build Coastguard Worker DocumentId email_document_id1 = put_result1.new_document_id;
476*8b6cd535SAndroid Build Coastguard Worker
477*8b6cd535SAndroid Build Coastguard Worker DocumentProto email_document2 = DocumentBuilder()
478*8b6cd535SAndroid Build Coastguard Worker .SetKey("namespace", "Email/2")
479*8b6cd535SAndroid Build Coastguard Worker .SetCreationTimestampMs(1000)
480*8b6cd535SAndroid Build Coastguard Worker .SetSchema("Email")
481*8b6cd535SAndroid Build Coastguard Worker .AddStringProperty("name", "Test 2")
482*8b6cd535SAndroid Build Coastguard Worker .AddStringProperty("body", "Test 2")
483*8b6cd535SAndroid Build Coastguard Worker .Build();
484*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(put_result2, doc_store->Put(email_document2));
485*8b6cd535SAndroid Build Coastguard Worker DocumentId email_document_id2 = put_result2.new_document_id;
486*8b6cd535SAndroid Build Coastguard Worker
487*8b6cd535SAndroid Build Coastguard Worker DocumentProto email_document3 = DocumentBuilder()
488*8b6cd535SAndroid Build Coastguard Worker .SetKey("namespace", "Email/3")
489*8b6cd535SAndroid Build Coastguard Worker .SetCreationTimestampMs(1000)
490*8b6cd535SAndroid Build Coastguard Worker .SetSchema("Email")
491*8b6cd535SAndroid Build Coastguard Worker .AddStringProperty("name", "Test 3")
492*8b6cd535SAndroid Build Coastguard Worker .AddStringProperty("body", "Test 3")
493*8b6cd535SAndroid Build Coastguard Worker .Build();
494*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
495*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(email_document3));
496*8b6cd535SAndroid Build Coastguard Worker DocumentId email_document_id3 = put_result3.new_document_id;
497*8b6cd535SAndroid Build Coastguard Worker
498*8b6cd535SAndroid Build Coastguard Worker DocumentProto email_document4 = DocumentBuilder()
499*8b6cd535SAndroid Build Coastguard Worker .SetKey("namespace", "Email/4")
500*8b6cd535SAndroid Build Coastguard Worker .SetCreationTimestampMs(1000)
501*8b6cd535SAndroid Build Coastguard Worker .SetSchema("Email")
502*8b6cd535SAndroid Build Coastguard Worker .AddStringProperty("name", "Test 4")
503*8b6cd535SAndroid Build Coastguard Worker .AddStringProperty("body", "Test 4")
504*8b6cd535SAndroid Build Coastguard Worker .Build();
505*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result4,
506*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(email_document4));
507*8b6cd535SAndroid Build Coastguard Worker DocumentId email_document_id4 = put_result4.new_document_id;
508*8b6cd535SAndroid Build Coastguard Worker
509*8b6cd535SAndroid Build Coastguard Worker // 3. Setup the joined scored results.
510*8b6cd535SAndroid Build Coastguard Worker std::vector<SectionId> person_hit_section_ids = {
511*8b6cd535SAndroid Build Coastguard Worker GetSectionId("Person", "name")};
512*8b6cd535SAndroid Build Coastguard Worker std::vector<SectionId> email_hit_section_ids = {
513*8b6cd535SAndroid Build Coastguard Worker GetSectionId("Email", "name"), GetSectionId("Email", "body")};
514*8b6cd535SAndroid Build Coastguard Worker SectionIdMask person_hit_section_id_mask =
515*8b6cd535SAndroid Build Coastguard Worker CreateSectionIdMask(person_hit_section_ids);
516*8b6cd535SAndroid Build Coastguard Worker SectionIdMask email_hit_section_id_mask =
517*8b6cd535SAndroid Build Coastguard Worker CreateSectionIdMask(email_hit_section_ids);
518*8b6cd535SAndroid Build Coastguard Worker
519*8b6cd535SAndroid Build Coastguard Worker ScoredDocumentHit person1_scored_doc_hit(
520*8b6cd535SAndroid Build Coastguard Worker person_document_id1, person_hit_section_id_mask, /*score=*/1);
521*8b6cd535SAndroid Build Coastguard Worker ScoredDocumentHit person2_scored_doc_hit(
522*8b6cd535SAndroid Build Coastguard Worker person_document_id2, person_hit_section_id_mask, /*score=*/2);
523*8b6cd535SAndroid Build Coastguard Worker ScoredDocumentHit email1_scored_doc_hit(
524*8b6cd535SAndroid Build Coastguard Worker email_document_id1, email_hit_section_id_mask, /*score=*/3);
525*8b6cd535SAndroid Build Coastguard Worker ScoredDocumentHit email2_scored_doc_hit(
526*8b6cd535SAndroid Build Coastguard Worker email_document_id2, email_hit_section_id_mask, /*score=*/4);
527*8b6cd535SAndroid Build Coastguard Worker ScoredDocumentHit email3_scored_doc_hit(
528*8b6cd535SAndroid Build Coastguard Worker email_document_id3, email_hit_section_id_mask, /*score=*/5);
529*8b6cd535SAndroid Build Coastguard Worker ScoredDocumentHit email4_scored_doc_hit(
530*8b6cd535SAndroid Build Coastguard Worker email_document_id4, email_hit_section_id_mask, /*score=*/6);
531*8b6cd535SAndroid Build Coastguard Worker // Create JoinedScoredDocumentHits mapping:
532*8b6cd535SAndroid Build Coastguard Worker // - Person1 to Email1
533*8b6cd535SAndroid Build Coastguard Worker // - Person2 to Email2, Email3, Email4
534*8b6cd535SAndroid Build Coastguard Worker std::vector<JoinedScoredDocumentHit> joined_scored_document_hits = {
535*8b6cd535SAndroid Build Coastguard Worker JoinedScoredDocumentHit(
536*8b6cd535SAndroid Build Coastguard Worker /*final_score=*/1,
537*8b6cd535SAndroid Build Coastguard Worker /*parent_scored_document_hit=*/person1_scored_doc_hit,
538*8b6cd535SAndroid Build Coastguard Worker /*child_scored_document_hits=*/{email1_scored_doc_hit}),
539*8b6cd535SAndroid Build Coastguard Worker JoinedScoredDocumentHit(
540*8b6cd535SAndroid Build Coastguard Worker /*final_score=*/3,
541*8b6cd535SAndroid Build Coastguard Worker /*parent_scored_document_hit=*/person2_scored_doc_hit,
542*8b6cd535SAndroid Build Coastguard Worker /*child_scored_document_hits=*/
543*8b6cd535SAndroid Build Coastguard Worker {email4_scored_doc_hit, email3_scored_doc_hit,
544*8b6cd535SAndroid Build Coastguard Worker email2_scored_doc_hit})};
545*8b6cd535SAndroid Build Coastguard Worker
546*8b6cd535SAndroid Build Coastguard Worker // 4. Retrieve result with max_joined_children_per_parent_to_return = 2.
547*8b6cd535SAndroid Build Coastguard Worker ResultSpecProto result_spec =
548*8b6cd535SAndroid Build Coastguard Worker CreateResultSpec(/*num_per_page=*/2, ResultSpecProto::NAMESPACE);
549*8b6cd535SAndroid Build Coastguard Worker result_spec.set_max_joined_children_per_parent_to_return(2);
550*8b6cd535SAndroid Build Coastguard Worker
551*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(
552*8b6cd535SAndroid Build Coastguard Worker std::unique_ptr<ResultRetrieverV2> result_retriever,
553*8b6cd535SAndroid Build Coastguard Worker ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(),
554*8b6cd535SAndroid Build Coastguard Worker language_segmenter_.get(), normalizer_.get()));
555*8b6cd535SAndroid Build Coastguard Worker ResultStateV2 result_state(
556*8b6cd535SAndroid Build Coastguard Worker std::make_unique<
557*8b6cd535SAndroid Build Coastguard Worker PriorityQueueScoredDocumentHitsRanker<JoinedScoredDocumentHit>>(
558*8b6cd535SAndroid Build Coastguard Worker std::move(joined_scored_document_hits), /*is_descending=*/true),
559*8b6cd535SAndroid Build Coastguard Worker /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
560*8b6cd535SAndroid Build Coastguard Worker result_spec, *doc_store);
561*8b6cd535SAndroid Build Coastguard Worker
562*8b6cd535SAndroid Build Coastguard Worker // Result1: person2 with child docs = [email4, email3]
563*8b6cd535SAndroid Build Coastguard Worker SearchResultProto::ResultProto result1;
564*8b6cd535SAndroid Build Coastguard Worker *result1.mutable_document() = person_document2;
565*8b6cd535SAndroid Build Coastguard Worker result1.set_score(3);
566*8b6cd535SAndroid Build Coastguard Worker SearchResultProto::ResultProto* child1 = result1.add_joined_results();
567*8b6cd535SAndroid Build Coastguard Worker *child1->mutable_document() = email_document4;
568*8b6cd535SAndroid Build Coastguard Worker child1->set_score(6);
569*8b6cd535SAndroid Build Coastguard Worker SearchResultProto::ResultProto* child2 = result1.add_joined_results();
570*8b6cd535SAndroid Build Coastguard Worker *child2->mutable_document() = email_document3;
571*8b6cd535SAndroid Build Coastguard Worker child2->set_score(5);
572*8b6cd535SAndroid Build Coastguard Worker
573*8b6cd535SAndroid Build Coastguard Worker // Result2: person1 with child docs = [email1]
574*8b6cd535SAndroid Build Coastguard Worker SearchResultProto::ResultProto result2;
575*8b6cd535SAndroid Build Coastguard Worker *result2.mutable_document() = person_document1;
576*8b6cd535SAndroid Build Coastguard Worker result2.set_score(1);
577*8b6cd535SAndroid Build Coastguard Worker SearchResultProto::ResultProto* child3 = result2.add_joined_results();
578*8b6cd535SAndroid Build Coastguard Worker *child3->mutable_document() = email_document1;
579*8b6cd535SAndroid Build Coastguard Worker child3->set_score(3);
580*8b6cd535SAndroid Build Coastguard Worker
581*8b6cd535SAndroid Build Coastguard Worker auto [page_result, has_more_results] = result_retriever->RetrieveNextPage(
582*8b6cd535SAndroid Build Coastguard Worker result_state, fake_clock_.GetSystemTimeMilliseconds());
583*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(page_result.results,
584*8b6cd535SAndroid Build Coastguard Worker ElementsAre(EqualsProto(result1), EqualsProto(result2)));
585*8b6cd535SAndroid Build Coastguard Worker // No more results.
586*8b6cd535SAndroid Build Coastguard Worker EXPECT_FALSE(has_more_results);
587*8b6cd535SAndroid Build Coastguard Worker }
588*8b6cd535SAndroid Build Coastguard Worker
TEST_F(ResultRetrieverV2Test,ShouldIgnoreInternalErrors)589*8b6cd535SAndroid Build Coastguard Worker TEST_F(ResultRetrieverV2Test, ShouldIgnoreInternalErrors) {
590*8b6cd535SAndroid Build Coastguard Worker MockFilesystem mock_filesystem;
591*8b6cd535SAndroid Build Coastguard Worker EXPECT_CALL(mock_filesystem,
592*8b6cd535SAndroid Build Coastguard Worker PRead(A<int>(), A<void*>(), A<size_t>(), A<off_t>()))
593*8b6cd535SAndroid Build Coastguard Worker .WillOnce(Return(false))
594*8b6cd535SAndroid Build Coastguard Worker .WillRepeatedly(DoDefault());
595*8b6cd535SAndroid Build Coastguard Worker
596*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(
597*8b6cd535SAndroid Build Coastguard Worker DocumentStore::CreateResult create_result,
598*8b6cd535SAndroid Build Coastguard Worker CreateDocumentStore(&mock_filesystem, test_dir_, &fake_clock_,
599*8b6cd535SAndroid Build Coastguard Worker schema_store_.get(), *feature_flags_));
600*8b6cd535SAndroid Build Coastguard Worker std::unique_ptr<DocumentStore> doc_store =
601*8b6cd535SAndroid Build Coastguard Worker std::move(create_result.document_store);
602*8b6cd535SAndroid Build Coastguard Worker
603*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
604*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(CreateDocument(/*id=*/1)));
605*8b6cd535SAndroid Build Coastguard Worker DocumentId document_id1 = put_result1.new_document_id;
606*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
607*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(CreateDocument(/*id=*/2)));
608*8b6cd535SAndroid Build Coastguard Worker DocumentId document_id2 = put_result2.new_document_id;
609*8b6cd535SAndroid Build Coastguard Worker
610*8b6cd535SAndroid Build Coastguard Worker std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "name"),
611*8b6cd535SAndroid Build Coastguard Worker GetSectionId("Email", "body")};
612*8b6cd535SAndroid Build Coastguard Worker SectionIdMask hit_section_id_mask = CreateSectionIdMask(hit_section_ids);
613*8b6cd535SAndroid Build Coastguard Worker std::vector<ScoredDocumentHit> scored_document_hits = {
614*8b6cd535SAndroid Build Coastguard Worker {document_id1, hit_section_id_mask, /*score=*/0},
615*8b6cd535SAndroid Build Coastguard Worker {document_id2, hit_section_id_mask, /*score=*/0}};
616*8b6cd535SAndroid Build Coastguard Worker
617*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(
618*8b6cd535SAndroid Build Coastguard Worker std::unique_ptr<ResultRetrieverV2> result_retriever,
619*8b6cd535SAndroid Build Coastguard Worker ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(),
620*8b6cd535SAndroid Build Coastguard Worker language_segmenter_.get(), normalizer_.get(),
621*8b6cd535SAndroid Build Coastguard Worker std::make_unique<MockGroupResultLimiter>()));
622*8b6cd535SAndroid Build Coastguard Worker
623*8b6cd535SAndroid Build Coastguard Worker SearchResultProto::ResultProto result1;
624*8b6cd535SAndroid Build Coastguard Worker *result1.mutable_document() = CreateDocument(/*id=*/1);
625*8b6cd535SAndroid Build Coastguard Worker result1.set_score(0);
626*8b6cd535SAndroid Build Coastguard Worker
627*8b6cd535SAndroid Build Coastguard Worker ResultStateV2 result_state(
628*8b6cd535SAndroid Build Coastguard Worker std::make_unique<
629*8b6cd535SAndroid Build Coastguard Worker PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
630*8b6cd535SAndroid Build Coastguard Worker std::move(scored_document_hits),
631*8b6cd535SAndroid Build Coastguard Worker /*is_descending=*/true),
632*8b6cd535SAndroid Build Coastguard Worker /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
633*8b6cd535SAndroid Build Coastguard Worker CreateResultSpec(/*num_per_page=*/2, ResultSpecProto::NAMESPACE),
634*8b6cd535SAndroid Build Coastguard Worker *doc_store);
635*8b6cd535SAndroid Build Coastguard Worker PageResult page_result =
636*8b6cd535SAndroid Build Coastguard Worker result_retriever
637*8b6cd535SAndroid Build Coastguard Worker ->RetrieveNextPage(result_state,
638*8b6cd535SAndroid Build Coastguard Worker fake_clock_.GetSystemTimeMilliseconds())
639*8b6cd535SAndroid Build Coastguard Worker .first;
640*8b6cd535SAndroid Build Coastguard Worker // We mocked mock_filesystem to return an internal error when retrieving doc2,
641*8b6cd535SAndroid Build Coastguard Worker // so doc2 should be skipped and doc1 should still be returned.
642*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(page_result.results, ElementsAre(EqualsProto(result1)));
643*8b6cd535SAndroid Build Coastguard Worker }
644*8b6cd535SAndroid Build Coastguard Worker
TEST_F(ResultRetrieverV2Test,ShouldUpdateResultState)645*8b6cd535SAndroid Build Coastguard Worker TEST_F(ResultRetrieverV2Test, ShouldUpdateResultState) {
646*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(
647*8b6cd535SAndroid Build Coastguard Worker DocumentStore::CreateResult create_result,
648*8b6cd535SAndroid Build Coastguard Worker CreateDocumentStore(&filesystem_, test_dir_, &fake_clock_,
649*8b6cd535SAndroid Build Coastguard Worker schema_store_.get(), *feature_flags_));
650*8b6cd535SAndroid Build Coastguard Worker std::unique_ptr<DocumentStore> doc_store =
651*8b6cd535SAndroid Build Coastguard Worker std::move(create_result.document_store);
652*8b6cd535SAndroid Build Coastguard Worker
653*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
654*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(CreateDocument(/*id=*/1)));
655*8b6cd535SAndroid Build Coastguard Worker DocumentId document_id1 = put_result1.new_document_id;
656*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
657*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(CreateDocument(/*id=*/2)));
658*8b6cd535SAndroid Build Coastguard Worker DocumentId document_id2 = put_result2.new_document_id;
659*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
660*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(CreateDocument(/*id=*/3)));
661*8b6cd535SAndroid Build Coastguard Worker DocumentId document_id3 = put_result3.new_document_id;
662*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result4,
663*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(CreateDocument(/*id=*/4)));
664*8b6cd535SAndroid Build Coastguard Worker DocumentId document_id4 = put_result4.new_document_id;
665*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result5,
666*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(CreateDocument(/*id=*/5)));
667*8b6cd535SAndroid Build Coastguard Worker DocumentId document_id5 = put_result5.new_document_id;
668*8b6cd535SAndroid Build Coastguard Worker
669*8b6cd535SAndroid Build Coastguard Worker std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "name"),
670*8b6cd535SAndroid Build Coastguard Worker GetSectionId("Email", "body")};
671*8b6cd535SAndroid Build Coastguard Worker SectionIdMask hit_section_id_mask = CreateSectionIdMask(hit_section_ids);
672*8b6cd535SAndroid Build Coastguard Worker std::vector<ScoredDocumentHit> scored_document_hits = {
673*8b6cd535SAndroid Build Coastguard Worker {document_id1, hit_section_id_mask, /*score=*/0},
674*8b6cd535SAndroid Build Coastguard Worker {document_id2, hit_section_id_mask, /*score=*/0},
675*8b6cd535SAndroid Build Coastguard Worker {document_id3, hit_section_id_mask, /*score=*/0},
676*8b6cd535SAndroid Build Coastguard Worker {document_id4, hit_section_id_mask, /*score=*/0},
677*8b6cd535SAndroid Build Coastguard Worker {document_id5, hit_section_id_mask, /*score=*/0}};
678*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(
679*8b6cd535SAndroid Build Coastguard Worker std::unique_ptr<ResultRetrieverV2> result_retriever,
680*8b6cd535SAndroid Build Coastguard Worker ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(),
681*8b6cd535SAndroid Build Coastguard Worker language_segmenter_.get(), normalizer_.get()));
682*8b6cd535SAndroid Build Coastguard Worker
683*8b6cd535SAndroid Build Coastguard Worker ResultStateV2 result_state(
684*8b6cd535SAndroid Build Coastguard Worker std::make_unique<
685*8b6cd535SAndroid Build Coastguard Worker PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
686*8b6cd535SAndroid Build Coastguard Worker std::move(scored_document_hits),
687*8b6cd535SAndroid Build Coastguard Worker /*is_descending=*/true),
688*8b6cd535SAndroid Build Coastguard Worker /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
689*8b6cd535SAndroid Build Coastguard Worker CreateResultSpec(/*num_per_page=*/2, ResultSpecProto::NAMESPACE),
690*8b6cd535SAndroid Build Coastguard Worker *doc_store);
691*8b6cd535SAndroid Build Coastguard Worker
692*8b6cd535SAndroid Build Coastguard Worker // First page, 2 results
693*8b6cd535SAndroid Build Coastguard Worker PageResult page_result1 =
694*8b6cd535SAndroid Build Coastguard Worker result_retriever
695*8b6cd535SAndroid Build Coastguard Worker ->RetrieveNextPage(result_state,
696*8b6cd535SAndroid Build Coastguard Worker fake_clock_.GetSystemTimeMilliseconds())
697*8b6cd535SAndroid Build Coastguard Worker .first;
698*8b6cd535SAndroid Build Coastguard Worker ASSERT_THAT(page_result1.results, SizeIs(2));
699*8b6cd535SAndroid Build Coastguard Worker {
700*8b6cd535SAndroid Build Coastguard Worker absl_ports::shared_lock l(&result_state.mutex);
701*8b6cd535SAndroid Build Coastguard Worker
702*8b6cd535SAndroid Build Coastguard Worker // num_returned = size of first page
703*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(result_state.num_returned, Eq(2));
704*8b6cd535SAndroid Build Coastguard Worker // Should remove the 2 returned docs from scored_document_hits and only
705*8b6cd535SAndroid Build Coastguard Worker // contain the remaining 3.
706*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(result_state.scored_document_hits_ranker, Pointee(SizeIs(3)));
707*8b6cd535SAndroid Build Coastguard Worker }
708*8b6cd535SAndroid Build Coastguard Worker
709*8b6cd535SAndroid Build Coastguard Worker // Second page, 2 results
710*8b6cd535SAndroid Build Coastguard Worker PageResult page_result2 =
711*8b6cd535SAndroid Build Coastguard Worker result_retriever
712*8b6cd535SAndroid Build Coastguard Worker ->RetrieveNextPage(result_state,
713*8b6cd535SAndroid Build Coastguard Worker fake_clock_.GetSystemTimeMilliseconds())
714*8b6cd535SAndroid Build Coastguard Worker .first;
715*8b6cd535SAndroid Build Coastguard Worker ASSERT_THAT(page_result2.results, SizeIs(2));
716*8b6cd535SAndroid Build Coastguard Worker {
717*8b6cd535SAndroid Build Coastguard Worker absl_ports::shared_lock l(&result_state.mutex);
718*8b6cd535SAndroid Build Coastguard Worker
719*8b6cd535SAndroid Build Coastguard Worker // num_returned = size of first and second pages
720*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(result_state.num_returned, Eq(4));
721*8b6cd535SAndroid Build Coastguard Worker // Should remove the 2 returned docs from scored_document_hits and only
722*8b6cd535SAndroid Build Coastguard Worker // contain the remaining 1.
723*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(result_state.scored_document_hits_ranker, Pointee(SizeIs(1)));
724*8b6cd535SAndroid Build Coastguard Worker }
725*8b6cd535SAndroid Build Coastguard Worker
726*8b6cd535SAndroid Build Coastguard Worker // Third page, 1 result
727*8b6cd535SAndroid Build Coastguard Worker PageResult page_result3 =
728*8b6cd535SAndroid Build Coastguard Worker result_retriever
729*8b6cd535SAndroid Build Coastguard Worker ->RetrieveNextPage(result_state,
730*8b6cd535SAndroid Build Coastguard Worker fake_clock_.GetSystemTimeMilliseconds())
731*8b6cd535SAndroid Build Coastguard Worker .first;
732*8b6cd535SAndroid Build Coastguard Worker ASSERT_THAT(page_result3.results, SizeIs(1));
733*8b6cd535SAndroid Build Coastguard Worker {
734*8b6cd535SAndroid Build Coastguard Worker absl_ports::shared_lock l(&result_state.mutex);
735*8b6cd535SAndroid Build Coastguard Worker
736*8b6cd535SAndroid Build Coastguard Worker // num_returned = size of first, second and third pages
737*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(result_state.num_returned, Eq(5));
738*8b6cd535SAndroid Build Coastguard Worker // Should remove the 1 returned doc from scored_document_hits and become
739*8b6cd535SAndroid Build Coastguard Worker // empty.
740*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(result_state.scored_document_hits_ranker, Pointee(IsEmpty()));
741*8b6cd535SAndroid Build Coastguard Worker }
742*8b6cd535SAndroid Build Coastguard Worker }
743*8b6cd535SAndroid Build Coastguard Worker
TEST_F(ResultRetrieverV2Test,ShouldUpdateNumTotalHits)744*8b6cd535SAndroid Build Coastguard Worker TEST_F(ResultRetrieverV2Test, ShouldUpdateNumTotalHits) {
745*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(
746*8b6cd535SAndroid Build Coastguard Worker DocumentStore::CreateResult create_result,
747*8b6cd535SAndroid Build Coastguard Worker CreateDocumentStore(&filesystem_, test_dir_, &fake_clock_,
748*8b6cd535SAndroid Build Coastguard Worker schema_store_.get(), *feature_flags_));
749*8b6cd535SAndroid Build Coastguard Worker std::unique_ptr<DocumentStore> doc_store =
750*8b6cd535SAndroid Build Coastguard Worker std::move(create_result.document_store);
751*8b6cd535SAndroid Build Coastguard Worker
752*8b6cd535SAndroid Build Coastguard Worker std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "name"),
753*8b6cd535SAndroid Build Coastguard Worker GetSectionId("Email", "body")};
754*8b6cd535SAndroid Build Coastguard Worker SectionIdMask hit_section_id_mask = CreateSectionIdMask(hit_section_ids);
755*8b6cd535SAndroid Build Coastguard Worker
756*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
757*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(CreateDocument(/*id=*/1)));
758*8b6cd535SAndroid Build Coastguard Worker DocumentId document_id1 = put_result1.new_document_id;
759*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
760*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(CreateDocument(/*id=*/2)));
761*8b6cd535SAndroid Build Coastguard Worker DocumentId document_id2 = put_result2.new_document_id;
762*8b6cd535SAndroid Build Coastguard Worker std::vector<ScoredDocumentHit> scored_document_hits1 = {
763*8b6cd535SAndroid Build Coastguard Worker {document_id1, hit_section_id_mask, /*score=*/0},
764*8b6cd535SAndroid Build Coastguard Worker {document_id2, hit_section_id_mask, /*score=*/0}};
765*8b6cd535SAndroid Build Coastguard Worker std::shared_ptr<ResultStateV2> result_state1 =
766*8b6cd535SAndroid Build Coastguard Worker std::make_shared<ResultStateV2>(
767*8b6cd535SAndroid Build Coastguard Worker std::make_unique<
768*8b6cd535SAndroid Build Coastguard Worker PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
769*8b6cd535SAndroid Build Coastguard Worker std::move(scored_document_hits1),
770*8b6cd535SAndroid Build Coastguard Worker /*is_descending=*/true),
771*8b6cd535SAndroid Build Coastguard Worker /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
772*8b6cd535SAndroid Build Coastguard Worker CreateResultSpec(/*num_per_page=*/1, ResultSpecProto::NAMESPACE),
773*8b6cd535SAndroid Build Coastguard Worker *doc_store);
774*8b6cd535SAndroid Build Coastguard Worker {
775*8b6cd535SAndroid Build Coastguard Worker absl_ports::unique_lock l(&result_state1->mutex);
776*8b6cd535SAndroid Build Coastguard Worker
777*8b6cd535SAndroid Build Coastguard Worker result_state1->RegisterNumTotalHits(&num_total_hits_);
778*8b6cd535SAndroid Build Coastguard Worker ASSERT_THAT(num_total_hits_, Eq(2));
779*8b6cd535SAndroid Build Coastguard Worker }
780*8b6cd535SAndroid Build Coastguard Worker
781*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result3,
782*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(CreateDocument(/*id=*/3)));
783*8b6cd535SAndroid Build Coastguard Worker DocumentId document_id3 = put_result3.new_document_id;
784*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result4,
785*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(CreateDocument(/*id=*/4)));
786*8b6cd535SAndroid Build Coastguard Worker DocumentId document_id4 = put_result4.new_document_id;
787*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result5,
788*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(CreateDocument(/*id=*/5)));
789*8b6cd535SAndroid Build Coastguard Worker DocumentId document_id5 = put_result5.new_document_id;
790*8b6cd535SAndroid Build Coastguard Worker std::vector<ScoredDocumentHit> scored_document_hits2 = {
791*8b6cd535SAndroid Build Coastguard Worker {document_id3, hit_section_id_mask, /*score=*/0},
792*8b6cd535SAndroid Build Coastguard Worker {document_id4, hit_section_id_mask, /*score=*/0},
793*8b6cd535SAndroid Build Coastguard Worker {document_id5, hit_section_id_mask, /*score=*/0}};
794*8b6cd535SAndroid Build Coastguard Worker std::shared_ptr<ResultStateV2> result_state2 =
795*8b6cd535SAndroid Build Coastguard Worker std::make_shared<ResultStateV2>(
796*8b6cd535SAndroid Build Coastguard Worker std::make_unique<
797*8b6cd535SAndroid Build Coastguard Worker PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
798*8b6cd535SAndroid Build Coastguard Worker std::move(scored_document_hits2),
799*8b6cd535SAndroid Build Coastguard Worker /*is_descending=*/true),
800*8b6cd535SAndroid Build Coastguard Worker /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
801*8b6cd535SAndroid Build Coastguard Worker CreateResultSpec(/*num_per_page=*/2, ResultSpecProto::NAMESPACE),
802*8b6cd535SAndroid Build Coastguard Worker *doc_store);
803*8b6cd535SAndroid Build Coastguard Worker {
804*8b6cd535SAndroid Build Coastguard Worker absl_ports::unique_lock l(&result_state2->mutex);
805*8b6cd535SAndroid Build Coastguard Worker
806*8b6cd535SAndroid Build Coastguard Worker result_state2->RegisterNumTotalHits(&num_total_hits_);
807*8b6cd535SAndroid Build Coastguard Worker ASSERT_THAT(num_total_hits_, Eq(5));
808*8b6cd535SAndroid Build Coastguard Worker }
809*8b6cd535SAndroid Build Coastguard Worker
810*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(
811*8b6cd535SAndroid Build Coastguard Worker std::unique_ptr<ResultRetrieverV2> result_retriever,
812*8b6cd535SAndroid Build Coastguard Worker ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(),
813*8b6cd535SAndroid Build Coastguard Worker language_segmenter_.get(), normalizer_.get()));
814*8b6cd535SAndroid Build Coastguard Worker
815*8b6cd535SAndroid Build Coastguard Worker // Should get 1 doc in the first page of result_state1, and num_total_hits
816*8b6cd535SAndroid Build Coastguard Worker // should be decremented by 1.
817*8b6cd535SAndroid Build Coastguard Worker PageResult page_result1 =
818*8b6cd535SAndroid Build Coastguard Worker result_retriever
819*8b6cd535SAndroid Build Coastguard Worker ->RetrieveNextPage(*result_state1,
820*8b6cd535SAndroid Build Coastguard Worker fake_clock_.GetSystemTimeMilliseconds())
821*8b6cd535SAndroid Build Coastguard Worker .first;
822*8b6cd535SAndroid Build Coastguard Worker ASSERT_THAT(page_result1.results, SizeIs(1));
823*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(num_total_hits_, Eq(4));
824*8b6cd535SAndroid Build Coastguard Worker
825*8b6cd535SAndroid Build Coastguard Worker // Should get 2 docs in the first page of result_state2, and num_total_hits
826*8b6cd535SAndroid Build Coastguard Worker // should be decremented by 2.
827*8b6cd535SAndroid Build Coastguard Worker PageResult page_result2 =
828*8b6cd535SAndroid Build Coastguard Worker result_retriever
829*8b6cd535SAndroid Build Coastguard Worker ->RetrieveNextPage(*result_state2,
830*8b6cd535SAndroid Build Coastguard Worker fake_clock_.GetSystemTimeMilliseconds())
831*8b6cd535SAndroid Build Coastguard Worker .first;
832*8b6cd535SAndroid Build Coastguard Worker ASSERT_THAT(page_result2.results, SizeIs(2));
833*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(num_total_hits_, Eq(2));
834*8b6cd535SAndroid Build Coastguard Worker
835*8b6cd535SAndroid Build Coastguard Worker // Should get 1 doc in the second page of result_state2 (although num_per_page
836*8b6cd535SAndroid Build Coastguard Worker // is 2, there is only 1 doc left), and num_total_hits should be decremented
837*8b6cd535SAndroid Build Coastguard Worker // by 1.
838*8b6cd535SAndroid Build Coastguard Worker PageResult page_result3 =
839*8b6cd535SAndroid Build Coastguard Worker result_retriever
840*8b6cd535SAndroid Build Coastguard Worker ->RetrieveNextPage(*result_state2,
841*8b6cd535SAndroid Build Coastguard Worker fake_clock_.GetSystemTimeMilliseconds())
842*8b6cd535SAndroid Build Coastguard Worker .first;
843*8b6cd535SAndroid Build Coastguard Worker ASSERT_THAT(page_result3.results, SizeIs(1));
844*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(num_total_hits_, Eq(1));
845*8b6cd535SAndroid Build Coastguard Worker
846*8b6cd535SAndroid Build Coastguard Worker // Destruct result_state1. There is 1 doc left, so num_total_hits should be
847*8b6cd535SAndroid Build Coastguard Worker // decremented by 1 when destructing it.
848*8b6cd535SAndroid Build Coastguard Worker result_state1.reset();
849*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(num_total_hits_, Eq(0));
850*8b6cd535SAndroid Build Coastguard Worker
851*8b6cd535SAndroid Build Coastguard Worker // Destruct result_state2. There is 0 doc left, so num_total_hits should be
852*8b6cd535SAndroid Build Coastguard Worker // unchanged when destructing it.
853*8b6cd535SAndroid Build Coastguard Worker result_state1.reset();
854*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(num_total_hits_, Eq(0));
855*8b6cd535SAndroid Build Coastguard Worker }
856*8b6cd535SAndroid Build Coastguard Worker
TEST_F(ResultRetrieverV2Test,ShouldLimitNumTotalBytesPerPage)857*8b6cd535SAndroid Build Coastguard Worker TEST_F(ResultRetrieverV2Test, ShouldLimitNumTotalBytesPerPage) {
858*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(
859*8b6cd535SAndroid Build Coastguard Worker DocumentStore::CreateResult create_result,
860*8b6cd535SAndroid Build Coastguard Worker CreateDocumentStore(&filesystem_, test_dir_, &fake_clock_,
861*8b6cd535SAndroid Build Coastguard Worker schema_store_.get(), *feature_flags_));
862*8b6cd535SAndroid Build Coastguard Worker std::unique_ptr<DocumentStore> doc_store =
863*8b6cd535SAndroid Build Coastguard Worker std::move(create_result.document_store);
864*8b6cd535SAndroid Build Coastguard Worker
865*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
866*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(CreateDocument(/*id=*/1)));
867*8b6cd535SAndroid Build Coastguard Worker DocumentId document_id1 = put_result1.new_document_id;
868*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
869*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(CreateDocument(/*id=*/2)));
870*8b6cd535SAndroid Build Coastguard Worker DocumentId document_id2 = put_result2.new_document_id;
871*8b6cd535SAndroid Build Coastguard Worker
872*8b6cd535SAndroid Build Coastguard Worker std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "name"),
873*8b6cd535SAndroid Build Coastguard Worker GetSectionId("Email", "body")};
874*8b6cd535SAndroid Build Coastguard Worker SectionIdMask hit_section_id_mask = CreateSectionIdMask(hit_section_ids);
875*8b6cd535SAndroid Build Coastguard Worker std::vector<ScoredDocumentHit> scored_document_hits = {
876*8b6cd535SAndroid Build Coastguard Worker {document_id1, hit_section_id_mask, /*score=*/5},
877*8b6cd535SAndroid Build Coastguard Worker {document_id2, hit_section_id_mask, /*score=*/0}};
878*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(
879*8b6cd535SAndroid Build Coastguard Worker std::unique_ptr<ResultRetrieverV2> result_retriever,
880*8b6cd535SAndroid Build Coastguard Worker ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(),
881*8b6cd535SAndroid Build Coastguard Worker language_segmenter_.get(), normalizer_.get()));
882*8b6cd535SAndroid Build Coastguard Worker
883*8b6cd535SAndroid Build Coastguard Worker SearchResultProto::ResultProto result1;
884*8b6cd535SAndroid Build Coastguard Worker *result1.mutable_document() = CreateDocument(/*id=*/1);
885*8b6cd535SAndroid Build Coastguard Worker result1.set_score(5);
886*8b6cd535SAndroid Build Coastguard Worker SearchResultProto::ResultProto result2;
887*8b6cd535SAndroid Build Coastguard Worker *result2.mutable_document() = CreateDocument(/*id=*/2);
888*8b6cd535SAndroid Build Coastguard Worker result2.set_score(0);
889*8b6cd535SAndroid Build Coastguard Worker
890*8b6cd535SAndroid Build Coastguard Worker ResultSpecProto result_spec =
891*8b6cd535SAndroid Build Coastguard Worker CreateResultSpec(/*num_per_page=*/2, ResultSpecProto::NAMESPACE);
892*8b6cd535SAndroid Build Coastguard Worker result_spec.set_num_total_bytes_per_page_threshold(result1.ByteSizeLong());
893*8b6cd535SAndroid Build Coastguard Worker ResultStateV2 result_state(
894*8b6cd535SAndroid Build Coastguard Worker std::make_unique<
895*8b6cd535SAndroid Build Coastguard Worker PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
896*8b6cd535SAndroid Build Coastguard Worker std::move(scored_document_hits),
897*8b6cd535SAndroid Build Coastguard Worker /*is_descending=*/true),
898*8b6cd535SAndroid Build Coastguard Worker /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
899*8b6cd535SAndroid Build Coastguard Worker result_spec, *doc_store);
900*8b6cd535SAndroid Build Coastguard Worker
901*8b6cd535SAndroid Build Coastguard Worker // First page. Only result1 should be returned, since its byte size meets
902*8b6cd535SAndroid Build Coastguard Worker // num_total_bytes_per_page_threshold and ResultRetriever should terminate
903*8b6cd535SAndroid Build Coastguard Worker // early even though # of results is still below num_per_page.
904*8b6cd535SAndroid Build Coastguard Worker auto [page_result1, has_more_results1] = result_retriever->RetrieveNextPage(
905*8b6cd535SAndroid Build Coastguard Worker result_state, fake_clock_.GetSystemTimeMilliseconds());
906*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(page_result1.results, ElementsAre(EqualsProto(result1)));
907*8b6cd535SAndroid Build Coastguard Worker // Has more results.
908*8b6cd535SAndroid Build Coastguard Worker EXPECT_TRUE(has_more_results1);
909*8b6cd535SAndroid Build Coastguard Worker
910*8b6cd535SAndroid Build Coastguard Worker // Second page, result2.
911*8b6cd535SAndroid Build Coastguard Worker auto [page_result2, has_more_results2] = result_retriever->RetrieveNextPage(
912*8b6cd535SAndroid Build Coastguard Worker result_state, fake_clock_.GetSystemTimeMilliseconds());
913*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(page_result2.results, ElementsAre(EqualsProto(result2)));
914*8b6cd535SAndroid Build Coastguard Worker // No more results.
915*8b6cd535SAndroid Build Coastguard Worker EXPECT_FALSE(has_more_results2);
916*8b6cd535SAndroid Build Coastguard Worker }
917*8b6cd535SAndroid Build Coastguard Worker
TEST_F(ResultRetrieverV2Test,ShouldReturnSingleLargeResultAboveNumTotalBytesPerPageThreshold)918*8b6cd535SAndroid Build Coastguard Worker TEST_F(ResultRetrieverV2Test,
919*8b6cd535SAndroid Build Coastguard Worker ShouldReturnSingleLargeResultAboveNumTotalBytesPerPageThreshold) {
920*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(
921*8b6cd535SAndroid Build Coastguard Worker DocumentStore::CreateResult create_result,
922*8b6cd535SAndroid Build Coastguard Worker CreateDocumentStore(&filesystem_, test_dir_, &fake_clock_,
923*8b6cd535SAndroid Build Coastguard Worker schema_store_.get(), *feature_flags_));
924*8b6cd535SAndroid Build Coastguard Worker std::unique_ptr<DocumentStore> doc_store =
925*8b6cd535SAndroid Build Coastguard Worker std::move(create_result.document_store);
926*8b6cd535SAndroid Build Coastguard Worker
927*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
928*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(CreateDocument(/*id=*/1)));
929*8b6cd535SAndroid Build Coastguard Worker DocumentId document_id1 = put_result1.new_document_id;
930*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
931*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(CreateDocument(/*id=*/2)));
932*8b6cd535SAndroid Build Coastguard Worker DocumentId document_id2 = put_result2.new_document_id;
933*8b6cd535SAndroid Build Coastguard Worker
934*8b6cd535SAndroid Build Coastguard Worker std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "name"),
935*8b6cd535SAndroid Build Coastguard Worker GetSectionId("Email", "body")};
936*8b6cd535SAndroid Build Coastguard Worker SectionIdMask hit_section_id_mask = CreateSectionIdMask(hit_section_ids);
937*8b6cd535SAndroid Build Coastguard Worker std::vector<ScoredDocumentHit> scored_document_hits = {
938*8b6cd535SAndroid Build Coastguard Worker {document_id1, hit_section_id_mask, /*score=*/5},
939*8b6cd535SAndroid Build Coastguard Worker {document_id2, hit_section_id_mask, /*score=*/0}};
940*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(
941*8b6cd535SAndroid Build Coastguard Worker std::unique_ptr<ResultRetrieverV2> result_retriever,
942*8b6cd535SAndroid Build Coastguard Worker ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(),
943*8b6cd535SAndroid Build Coastguard Worker language_segmenter_.get(), normalizer_.get()));
944*8b6cd535SAndroid Build Coastguard Worker
945*8b6cd535SAndroid Build Coastguard Worker SearchResultProto::ResultProto result1;
946*8b6cd535SAndroid Build Coastguard Worker *result1.mutable_document() = CreateDocument(/*id=*/1);
947*8b6cd535SAndroid Build Coastguard Worker result1.set_score(5);
948*8b6cd535SAndroid Build Coastguard Worker SearchResultProto::ResultProto result2;
949*8b6cd535SAndroid Build Coastguard Worker *result2.mutable_document() = CreateDocument(/*id=*/2);
950*8b6cd535SAndroid Build Coastguard Worker result2.set_score(0);
951*8b6cd535SAndroid Build Coastguard Worker
952*8b6cd535SAndroid Build Coastguard Worker int threshold = 1;
953*8b6cd535SAndroid Build Coastguard Worker ASSERT_THAT(result1.ByteSizeLong(), Gt(threshold));
954*8b6cd535SAndroid Build Coastguard Worker
955*8b6cd535SAndroid Build Coastguard Worker ResultSpecProto result_spec =
956*8b6cd535SAndroid Build Coastguard Worker CreateResultSpec(/*num_per_page=*/2, ResultSpecProto::NAMESPACE);
957*8b6cd535SAndroid Build Coastguard Worker result_spec.set_num_total_bytes_per_page_threshold(threshold);
958*8b6cd535SAndroid Build Coastguard Worker ResultStateV2 result_state(
959*8b6cd535SAndroid Build Coastguard Worker std::make_unique<
960*8b6cd535SAndroid Build Coastguard Worker PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
961*8b6cd535SAndroid Build Coastguard Worker std::move(scored_document_hits),
962*8b6cd535SAndroid Build Coastguard Worker /*is_descending=*/true),
963*8b6cd535SAndroid Build Coastguard Worker /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
964*8b6cd535SAndroid Build Coastguard Worker result_spec, *doc_store);
965*8b6cd535SAndroid Build Coastguard Worker
966*8b6cd535SAndroid Build Coastguard Worker // First page. Should return single result1 even though its byte size exceeds
967*8b6cd535SAndroid Build Coastguard Worker // num_total_bytes_per_page_threshold.
968*8b6cd535SAndroid Build Coastguard Worker auto [page_result1, has_more_results1] = result_retriever->RetrieveNextPage(
969*8b6cd535SAndroid Build Coastguard Worker result_state, fake_clock_.GetSystemTimeMilliseconds());
970*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(page_result1.results, ElementsAre(EqualsProto(result1)));
971*8b6cd535SAndroid Build Coastguard Worker // Has more results.
972*8b6cd535SAndroid Build Coastguard Worker EXPECT_TRUE(has_more_results1);
973*8b6cd535SAndroid Build Coastguard Worker
974*8b6cd535SAndroid Build Coastguard Worker // Second page, result2.
975*8b6cd535SAndroid Build Coastguard Worker auto [page_result2, has_more_results2] = result_retriever->RetrieveNextPage(
976*8b6cd535SAndroid Build Coastguard Worker result_state, fake_clock_.GetSystemTimeMilliseconds());
977*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(page_result2.results, ElementsAre(EqualsProto(result2)));
978*8b6cd535SAndroid Build Coastguard Worker // No more results.
979*8b6cd535SAndroid Build Coastguard Worker EXPECT_FALSE(has_more_results2);
980*8b6cd535SAndroid Build Coastguard Worker }
981*8b6cd535SAndroid Build Coastguard Worker
TEST_F(ResultRetrieverV2Test,ShouldRetrieveNextResultWhenBelowNumTotalBytesPerPageThreshold)982*8b6cd535SAndroid Build Coastguard Worker TEST_F(ResultRetrieverV2Test,
983*8b6cd535SAndroid Build Coastguard Worker ShouldRetrieveNextResultWhenBelowNumTotalBytesPerPageThreshold) {
984*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(
985*8b6cd535SAndroid Build Coastguard Worker DocumentStore::CreateResult create_result,
986*8b6cd535SAndroid Build Coastguard Worker CreateDocumentStore(&filesystem_, test_dir_, &fake_clock_,
987*8b6cd535SAndroid Build Coastguard Worker schema_store_.get(), *feature_flags_));
988*8b6cd535SAndroid Build Coastguard Worker std::unique_ptr<DocumentStore> doc_store =
989*8b6cd535SAndroid Build Coastguard Worker std::move(create_result.document_store);
990*8b6cd535SAndroid Build Coastguard Worker
991*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result1,
992*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(CreateDocument(/*id=*/1)));
993*8b6cd535SAndroid Build Coastguard Worker DocumentId document_id1 = put_result1.new_document_id;
994*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(DocumentStore::PutResult put_result2,
995*8b6cd535SAndroid Build Coastguard Worker doc_store->Put(CreateDocument(/*id=*/2)));
996*8b6cd535SAndroid Build Coastguard Worker DocumentId document_id2 = put_result2.new_document_id;
997*8b6cd535SAndroid Build Coastguard Worker
998*8b6cd535SAndroid Build Coastguard Worker std::vector<SectionId> hit_section_ids = {GetSectionId("Email", "name"),
999*8b6cd535SAndroid Build Coastguard Worker GetSectionId("Email", "body")};
1000*8b6cd535SAndroid Build Coastguard Worker SectionIdMask hit_section_id_mask = CreateSectionIdMask(hit_section_ids);
1001*8b6cd535SAndroid Build Coastguard Worker std::vector<ScoredDocumentHit> scored_document_hits = {
1002*8b6cd535SAndroid Build Coastguard Worker {document_id1, hit_section_id_mask, /*score=*/5},
1003*8b6cd535SAndroid Build Coastguard Worker {document_id2, hit_section_id_mask, /*score=*/0}};
1004*8b6cd535SAndroid Build Coastguard Worker ICING_ASSERT_OK_AND_ASSIGN(
1005*8b6cd535SAndroid Build Coastguard Worker std::unique_ptr<ResultRetrieverV2> result_retriever,
1006*8b6cd535SAndroid Build Coastguard Worker ResultRetrieverV2::Create(doc_store.get(), schema_store_.get(),
1007*8b6cd535SAndroid Build Coastguard Worker language_segmenter_.get(), normalizer_.get()));
1008*8b6cd535SAndroid Build Coastguard Worker
1009*8b6cd535SAndroid Build Coastguard Worker SearchResultProto::ResultProto result1;
1010*8b6cd535SAndroid Build Coastguard Worker *result1.mutable_document() = CreateDocument(/*id=*/1);
1011*8b6cd535SAndroid Build Coastguard Worker result1.set_score(5);
1012*8b6cd535SAndroid Build Coastguard Worker SearchResultProto::ResultProto result2;
1013*8b6cd535SAndroid Build Coastguard Worker *result2.mutable_document() = CreateDocument(/*id=*/2);
1014*8b6cd535SAndroid Build Coastguard Worker result2.set_score(0);
1015*8b6cd535SAndroid Build Coastguard Worker
1016*8b6cd535SAndroid Build Coastguard Worker int threshold = result1.ByteSizeLong() + 1;
1017*8b6cd535SAndroid Build Coastguard Worker ASSERT_THAT(result1.ByteSizeLong() + result2.ByteSizeLong(), Gt(threshold));
1018*8b6cd535SAndroid Build Coastguard Worker
1019*8b6cd535SAndroid Build Coastguard Worker ResultSpecProto result_spec =
1020*8b6cd535SAndroid Build Coastguard Worker CreateResultSpec(/*num_per_page=*/2, ResultSpecProto::NAMESPACE);
1021*8b6cd535SAndroid Build Coastguard Worker result_spec.set_num_total_bytes_per_page_threshold(threshold);
1022*8b6cd535SAndroid Build Coastguard Worker ResultStateV2 result_state(
1023*8b6cd535SAndroid Build Coastguard Worker std::make_unique<
1024*8b6cd535SAndroid Build Coastguard Worker PriorityQueueScoredDocumentHitsRanker<ScoredDocumentHit>>(
1025*8b6cd535SAndroid Build Coastguard Worker std::move(scored_document_hits),
1026*8b6cd535SAndroid Build Coastguard Worker /*is_descending=*/true),
1027*8b6cd535SAndroid Build Coastguard Worker /*parent_adjustment_info=*/nullptr, /*child_adjustment_info=*/nullptr,
1028*8b6cd535SAndroid Build Coastguard Worker result_spec, *doc_store);
1029*8b6cd535SAndroid Build Coastguard Worker
1030*8b6cd535SAndroid Build Coastguard Worker // After retrieving result1, total bytes are still below the threshold and #
1031*8b6cd535SAndroid Build Coastguard Worker // of results is still below num_per_page, so ResultRetriever should continue
1032*8b6cd535SAndroid Build Coastguard Worker // the retrieval process and thus include result2 into this page, even though
1033*8b6cd535SAndroid Build Coastguard Worker // finally total bytes of result1 + result2 exceed the threshold.
1034*8b6cd535SAndroid Build Coastguard Worker auto [page_result, has_more_results] = result_retriever->RetrieveNextPage(
1035*8b6cd535SAndroid Build Coastguard Worker result_state, fake_clock_.GetSystemTimeMilliseconds());
1036*8b6cd535SAndroid Build Coastguard Worker EXPECT_THAT(page_result.results,
1037*8b6cd535SAndroid Build Coastguard Worker ElementsAre(EqualsProto(result1), EqualsProto(result2)));
1038*8b6cd535SAndroid Build Coastguard Worker // No more results.
1039*8b6cd535SAndroid Build Coastguard Worker EXPECT_FALSE(has_more_results);
1040*8b6cd535SAndroid Build Coastguard Worker }
1041*8b6cd535SAndroid Build Coastguard Worker
1042*8b6cd535SAndroid Build Coastguard Worker } // namespace
1043*8b6cd535SAndroid Build Coastguard Worker
1044*8b6cd535SAndroid Build Coastguard Worker } // namespace lib
1045*8b6cd535SAndroid Build Coastguard Worker } // namespace icing
1046