Skip to content

Commit c1ea67a

Browse files
kderussokderusso
committed
Fix header
1 parent f1a5762 commit c1ea67a

File tree

2 files changed

+240
-5
lines changed

2 files changed

+240
-5
lines changed

‎x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/application/rules/retriever/QueryRuleRetrieverBuilder.java‎

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,8 @@
11
/*
22
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3-
* or more contributor license agreements. Licensed under the "Elastic License
4-
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5-
* Public License v 1"; you may not use this file except in compliance with, at
6-
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7-
* License v3.0 only", or the "Server Side Public License, v 1".
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
86
*/
97

108
package org.elasticsearch.xpack.application.rules.retriever;
@@ -15,6 +13,8 @@
1513
import org.elasticsearch.index.query.BoolQueryBuilder;
1614
import org.elasticsearch.index.query.QueryBuilder;
1715
import org.elasticsearch.license.LicenseUtils;
16+
import org.elasticsearch.logging.LogManager;
17+
import org.elasticsearch.logging.Logger;
1818
import org.elasticsearch.search.builder.PointInTimeBuilder;
1919
import org.elasticsearch.search.builder.SearchSourceBuilder;
2020
import org.elasticsearch.search.fetch.StoredFieldsContext;
@@ -23,6 +23,7 @@
2323
import org.elasticsearch.search.retriever.RetrieverBuilder;
2424
import org.elasticsearch.search.retriever.RetrieverParserContext;
2525
import org.elasticsearch.search.retriever.rankdoc.RankDocsQueryBuilder;
26+
import org.elasticsearch.search.sort.ScoreSortBuilder;
2627
import org.elasticsearch.xcontent.ConstructingObjectParser;
2728
import org.elasticsearch.xcontent.ParseField;
2829
import org.elasticsearch.xcontent.XContentBuilder;
@@ -124,6 +125,7 @@ public String getName() {
124125

125126
@Override
126127
protected SearchSourceBuilder createSearchSourceBuilder(PointInTimeBuilder pit, RetrieverBuilder retrieverBuilder) {
128+
Logger logger = LogManager.getLogger(QueryRuleRetrieverBuilder.class);
127129
var sourceBuilder = new SearchSourceBuilder().pointInTimeBuilder(pit)
128130
.trackTotalHits(false)
129131
.storedFields(new StoredFieldsContext(false))
@@ -149,6 +151,9 @@ protected SearchSourceBuilder createSearchSourceBuilder(PointInTimeBuilder pit,
149151
sourceBuilder.query(newQuery);
150152
}
151153

154+
sourceBuilder.sort(new ScoreSortBuilder());
155+
156+
logger.info("sourceBuilder: " + sourceBuilder);
152157
return sourceBuilder;
153158
}
154159

‎x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/application/rules/retriever/QueryRuleRetrieverBuilder.java.new‎

Lines changed: 230 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,230 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the "Elastic License
4+
* 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.xpack.application.rules.retriever;
11+
12+
/*
13+
* Licensed to the Apache Software Foundation (ASF) under one
14+
* or more contributor license agreements. See the NOTICE file
15+
* distributed with this work for additional information
16+
* regarding copyright ownership. The ASF licenses this file
17+
* to you under the Apache License, Version 2.0 (the
18+
* "License"); you may not use this file except in compliance
19+
* with the License. You may obtain a copy of the License at
20+
*
21+
* http://www.apache.org/licenses/LICENSE-2.0
22+
*
23+
* Unless required by applicable law or agreed to in writing,
24+
* software distributed under the License is distributed on an
25+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
26+
* KIND, either express or implied. See the License for the
27+
* specific language governing permissions and limitations
28+
* under the License.
29+
*/
30+
31+
32+
import org.apache.lucene.search.ScoreDoc;
33+
import org.elasticsearch.common.ParsingException;
34+
import org.elasticsearch.features.NodeFeature;
35+
import org.elasticsearch.index.query.BoolQueryBuilder;
36+
import org.elasticsearch.index.query.QueryBuilder;
37+
import org.elasticsearch.license.LicenseUtils;
38+
import org.elasticsearch.logging.LogManager;
39+
import org.elasticsearch.logging.Logger;
40+
import org.elasticsearch.search.builder.PointInTimeBuilder;
41+
import org.elasticsearch.search.builder.SearchSourceBuilder;
42+
import org.elasticsearch.search.fetch.StoredFieldsContext;
43+
import org.elasticsearch.search.rank.RankDoc;
44+
import org.elasticsearch.search.retriever.CompoundRetrieverBuilder;
45+
import org.elasticsearch.search.retriever.RetrieverBuilder;
46+
import org.elasticsearch.search.retriever.RetrieverParserContext;
47+
import org.elasticsearch.search.retriever.rankdoc.RankDocsQueryBuilder;
48+
import org.elasticsearch.search.sort.ScoreSortBuilder;
49+
import org.elasticsearch.xcontent.ConstructingObjectParser;
50+
import org.elasticsearch.xcontent.ParseField;
51+
import org.elasticsearch.xcontent.XContentBuilder;
52+
import org.elasticsearch.xcontent.XContentParser;
53+
import org.elasticsearch.xpack.application.EnterpriseSearch;
54+
import org.elasticsearch.xpack.application.rules.RuleQueryBuilder;
55+
import org.elasticsearch.xpack.core.XPackPlugin;
56+
57+
import java.io.IOException;
58+
import java.util.ArrayList;
59+
import java.util.List;
60+
import java.util.Map;
61+
import java.util.Objects;
62+
63+
import static org.elasticsearch.search.rank.RankBuilder.DEFAULT_RANK_WINDOW_SIZE;
64+
import static org.elasticsearch.xcontent.ConstructingObjectParser.constructorArg;
65+
import static org.elasticsearch.xcontent.ConstructingObjectParser.optionalConstructorArg;
66+
67+
/**
68+
* A query rule retriever applies query rules defined in one or more rulesets to the underlying retriever.
69+
*/
70+
public final class QueryRuleRetrieverBuilder extends CompoundRetrieverBuilder<QueryRuleRetrieverBuilder> {
71+
72+
public static final String NAME = "rule";
73+
public static final NodeFeature QUERY_RULE_RETRIEVERS_SUPPORTED = new NodeFeature("query_rule_retriever_supported");
74+
75+
public static final ParseField RULESET_IDS_FIELD = new ParseField("ruleset_ids");
76+
public static final ParseField MATCH_CRITERIA_FIELD = new ParseField("match_criteria");
77+
public static final ParseField RETRIEVER_FIELD = new ParseField("retriever");
78+
public static final ParseField RANK_WINDOW_SIZE_FIELD = new ParseField("rank_window_size");
79+
80+
@SuppressWarnings("unchecked")
81+
public static final ConstructingObjectParser<QueryRuleRetrieverBuilder, RetrieverParserContext> PARSER = new ConstructingObjectParser<>(
82+
"rule",
83+
args -> {
84+
List<String> rulesetIds = (List<String>) args[0];
85+
Map<String, Object> matchCriteria = (Map<String, Object>) args[1];
86+
RetrieverBuilder retrieverBuilder = (RetrieverBuilder) args[2];
87+
int rankWindowSize = args[3] == null ? DEFAULT_RANK_WINDOW_SIZE : (int) args[3];
88+
return new QueryRuleRetrieverBuilder(rulesetIds, matchCriteria, retrieverBuilder, rankWindowSize);
89+
}
90+
);
91+
92+
static {
93+
PARSER.declareStringArray(constructorArg(), RULESET_IDS_FIELD);
94+
PARSER.declareObject(constructorArg(), (p, c) -> p.map(), MATCH_CRITERIA_FIELD);
95+
PARSER.declareNamedObject(constructorArg(), (p, c, n) -> p.namedObject(RetrieverBuilder.class, n, c), RETRIEVER_FIELD);
96+
PARSER.declareInt(optionalConstructorArg(), RANK_WINDOW_SIZE_FIELD);
97+
RetrieverBuilder.declareBaseParserFields(NAME, PARSER);
98+
}
99+
100+
public static QueryRuleRetrieverBuilder fromXContent(XContentParser parser, RetrieverParserContext context) throws IOException {
101+
if (context.clusterSupportsFeature(QUERY_RULE_RETRIEVERS_SUPPORTED) == false) {
102+
throw new ParsingException(parser.getTokenLocation(), "unknown retriever [" + NAME + "]");
103+
}
104+
if (EnterpriseSearch.QUERY_RULES_RETRIEVER_FEATURE.check(XPackPlugin.getSharedLicenseState()) == false) {
105+
throw LicenseUtils.newComplianceException("Query Rules");
106+
}
107+
try {
108+
return PARSER.apply(parser, context);
109+
} catch (Exception e) {
110+
throw new ParsingException(parser.getTokenLocation(), e.getMessage(), e);
111+
}
112+
}
113+
114+
private final List<String> rulesetIds;
115+
private final Map<String, Object> matchCriteria;
116+
117+
public QueryRuleRetrieverBuilder(
118+
List<String> rulesetIds,
119+
Map<String, Object> matchCriteria,
120+
RetrieverBuilder retrieverBuilder,
121+
int rankWindowSize
122+
) {
123+
super(List.of(new RetrieverSource(retrieverBuilder, null)), rankWindowSize);
124+
this.rulesetIds = rulesetIds;
125+
this.matchCriteria = matchCriteria;
126+
}
127+
128+
public QueryRuleRetrieverBuilder(
129+
List<String> rulesetIds,
130+
Map<String, Object> matchCriteria,
131+
List<RetrieverSource> retrieverSource,
132+
int rankWindowSize,
133+
String retrieverName,
134+
List<QueryBuilder> preFilterQueryBuilders
135+
) {
136+
super(retrieverSource, rankWindowSize);
137+
this.rulesetIds = rulesetIds;
138+
this.matchCriteria = matchCriteria;
139+
this.retrieverName = retrieverName;
140+
this.preFilterQueryBuilders = new ArrayList<>(preFilterQueryBuilders);
141+
}
142+
143+
@Override
144+
public String getName() {
145+
return NAME;
146+
}
147+
148+
@Override
149+
protected SearchSourceBuilder createSearchSourceBuilder(PointInTimeBuilder pit, RetrieverBuilder retrieverBuilder) {
150+
Logger logger = LogManager.getLogger(QueryRuleRetrieverBuilder.class);
151+
var sourceBuilder = new SearchSourceBuilder().pointInTimeBuilder(pit)
152+
.trackTotalHits(false)
153+
.storedFields(new StoredFieldsContext(false))
154+
.size(rankWindowSize);
155+
if (preFilterQueryBuilders.isEmpty() == false) {
156+
retrieverBuilder.getPreFilterQueryBuilders().addAll(preFilterQueryBuilders);
157+
}
158+
retrieverBuilder.extractToSearchSourceBuilder(sourceBuilder, true);
159+
160+
QueryBuilder query = sourceBuilder.query();
161+
if (query != null && query instanceof RuleQueryBuilder == false) {
162+
QueryBuilder organicQuery = query;
163+
query = new RuleQueryBuilder(organicQuery, matchCriteria, rulesetIds);
164+
}
165+
166+
// apply the pre-filters
167+
if (preFilterQueryBuilders.size() > 0) {
168+
BoolQueryBuilder newQuery = new BoolQueryBuilder();
169+
if (query != null) {
170+
newQuery.must(query);
171+
}
172+
preFilterQueryBuilders.forEach(newQuery::filter);
173+
sourceBuilder.query(newQuery);
174+
}
175+
176+
sourceBuilder.sort(new ScoreSortBuilder());
177+
178+
logger.info("sourceBuilder: " + sourceBuilder);
179+
return sourceBuilder;
180+
}
181+
182+
@Override
183+
public void doToXContent(XContentBuilder builder, Params params) throws IOException {
184+
builder.array(RULESET_IDS_FIELD.getPreferredName(), rulesetIds.toArray());
185+
builder.startObject(MATCH_CRITERIA_FIELD.getPreferredName());
186+
builder.mapContents(matchCriteria);
187+
builder.endObject();
188+
}
189+
190+
@Override
191+
protected QueryRuleRetrieverBuilder clone(List<RetrieverSource> newChildRetrievers) {
192+
return new QueryRuleRetrieverBuilder(
193+
rulesetIds,
194+
matchCriteria,
195+
newChildRetrievers,
196+
rankWindowSize,
197+
retrieverName,
198+
preFilterQueryBuilders
199+
);
200+
}
201+
202+
@Override
203+
protected RankDoc[] combineInnerRetrieverResults(List<ScoreDoc[]> rankResults) {
204+
assert rankResults.size() == 1;
205+
ScoreDoc[] scoreDocs = rankResults.getFirst();
206+
RankDoc[] rankDocs = new RankDoc[scoreDocs.length];
207+
for (int i = 0; i < scoreDocs.length; i++) {
208+
ScoreDoc scoreDoc = scoreDocs[i];
209+
rankDocs[i] = new RankDoc(scoreDoc.doc, scoreDoc.score, scoreDoc.shardIndex);
210+
}
211+
return rankDocs;
212+
}
213+
214+
@Override
215+
public QueryBuilder explainQuery() {
216+
// the original matching set of the QueryRuleRetriever retriever is specified by its nested retriever
217+
return new RankDocsQueryBuilder(rankDocs, new QueryBuilder[] { innerRetrievers.getFirst().retriever().explainQuery() }, true);
218+
}
219+
220+
@Override
221+
public boolean doEquals(Object o) {
222+
QueryRuleRetrieverBuilder that = (QueryRuleRetrieverBuilder) o;
223+
return super.doEquals(o) && Objects.equals(rulesetIds, that.rulesetIds) && Objects.equals(matchCriteria, that.matchCriteria);
224+
}
225+
226+
@Override
227+
public int doHashCode() {
228+
return Objects.hash(super.doHashCode(), rulesetIds, matchCriteria);
229+
}
230+
}

0 commit comments

Comments
 (0)