1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
|
// Copyright (C) 2024 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#ifndef ICING_INDEX_EMBED_EMBEDDING_QUERY_RESULTS_H_
#define ICING_INDEX_EMBED_EMBEDDING_QUERY_RESULTS_H_
#include <unordered_map>
#include <vector>
#include "icing/proto/search.pb.h"
#include "icing/store/document-id.h"
namespace icing {
namespace lib {
// A class to store results generated from embedding queries.
struct EmbeddingQueryResults {
// Maps from DocumentId to the list of matched embedding scores for that
// document, which will be used in the advanced scoring language to
// determine the results for the "this.matchedSemanticScores(...)" function.
using EmbeddingQueryScoreMap =
std::unordered_map<DocumentId, std::vector<double>>;
// Maps from (query_vector_index, metric_type) to EmbeddingQueryScoreMap.
std::unordered_map<
int, std::unordered_map<SearchSpecProto::EmbeddingQueryMetricType::Code,
EmbeddingQueryScoreMap>>
result_scores;
// Returns the matched scores for the given query_vector_index, metric_type,
// and doc_id. Returns nullptr if (query_vector_index, metric_type) does not
// exist in the result_scores map.
const std::vector<double>* GetMatchedScoresForDocument(
int query_vector_index,
SearchSpecProto::EmbeddingQueryMetricType::Code metric_type,
DocumentId doc_id) const {
// Check if a mapping exists for the query_vector_index
auto outer_it = result_scores.find(query_vector_index);
if (outer_it == result_scores.end()) {
return nullptr;
}
// Check if a mapping exists for the metric_type
auto inner_it = outer_it->second.find(metric_type);
if (inner_it == outer_it->second.end()) {
return nullptr;
}
const EmbeddingQueryScoreMap& score_map = inner_it->second;
// Check if the doc_id exists in the score_map
auto scores_it = score_map.find(doc_id);
if (scores_it == score_map.end()) {
return nullptr;
}
return &scores_it->second;
}
};
} // namespace lib
} // namespace icing
#endif // ICING_INDEX_EMBED_EMBEDDING_QUERY_RESULTS_H_
|
