Skip to content

Commit 2dc0b44

Browse files
Mikep86Mikep86
committed
Improve edge case handling:
- Do not allow querying remote indices - Handle when index wildcard pattern does not match any indices
1 parent f8923c3 commit 2dc0b44

File tree

5 files changed

+49
-22
lines changed

5 files changed

+49
-22
lines changed

‎server/src/main/java/org/elasticsearch/search/retriever/CompoundRetrieverBuilder.java‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,7 @@ public final RetrieverBuilder rewrite(QueryRewriteContext ctx) throws IOExceptio
100100
throw new IllegalStateException("PIT is required");
101101
}
102102

103-
T rewritten = doRewrite(ctx);
103+
RetrieverBuilder rewritten = doRewrite(ctx);
104104
if (rewritten != this) {
105105
return rewritten;
106106
}
@@ -334,11 +334,10 @@ protected SearchSourceBuilder finalizeSourceBuilder(SearchSourceBuilder sourceBu
334334
* Perform any custom rewrite logic necessary
335335
*
336336
* @param ctx The query rewrite context
337-
* @return T the rewritten retriever
337+
* @return RetrieverBuilder the rewritten retriever
338338
*/
339-
@SuppressWarnings("unchecked")
340-
protected T doRewrite(QueryRewriteContext ctx) {
341-
return (T) this;
339+
protected RetrieverBuilder doRewrite(QueryRewriteContext ctx) {
340+
return this;
342341
}
343342

344343
private RankDoc[] getRankDocs(SearchResponse searchResponse) {

‎x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/linear/LinearRetrieverBuilder.java‎

Lines changed: 18 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
import org.elasticsearch.action.ResolvedIndices;
1313
import org.elasticsearch.common.ParsingException;
1414
import org.elasticsearch.common.util.Maps;
15+
import org.elasticsearch.index.query.MatchNoneQueryBuilder;
1516
import org.elasticsearch.index.query.QueryBuilder;
1617
import org.elasticsearch.index.query.QueryRewriteContext;
1718
import org.elasticsearch.license.LicenseUtils;
@@ -21,6 +22,7 @@
2122
import org.elasticsearch.search.retriever.CompoundRetrieverBuilder;
2223
import org.elasticsearch.search.retriever.RetrieverBuilder;
2324
import org.elasticsearch.search.retriever.RetrieverParserContext;
25+
import org.elasticsearch.search.retriever.StandardRetrieverBuilder;
2426
import org.elasticsearch.xcontent.ConstructingObjectParser;
2527
import org.elasticsearch.xcontent.ParseField;
2628
import org.elasticsearch.xcontent.XContentBuilder;
@@ -280,8 +282,8 @@ protected RankDoc[] combineInnerRetrieverResults(List<ScoreDoc[]> rankResults, b
280282
}
281283

282284
@Override
283-
protected LinearRetrieverBuilder doRewrite(QueryRewriteContext ctx) {
284-
LinearRetrieverBuilder rewritten = this;
285+
protected RetrieverBuilder doRewrite(QueryRewriteContext ctx) {
286+
RetrieverBuilder rewritten = this;
285287

286288
ResolvedIndices resolvedIndices = ctx.getResolvedIndices();
287289
if (resolvedIndices != null && query != null) {
@@ -291,6 +293,10 @@ protected LinearRetrieverBuilder doRewrite(QueryRewriteContext ctx) {
291293
throw new IllegalArgumentException(
292294
"[" + NAME + "] does not support the simplified query format when querying multiple indices"
293295
);
296+
} else if (resolvedIndices.getRemoteClusterIndices().isEmpty() == false) {
297+
throw new IllegalArgumentException(
298+
"[" + NAME + "] does not support the simplified query format when querying remote indices"
299+
);
294300
}
295301

296302
List<RetrieverSource> fieldsInnerRetrievers = SimplifiedInnerRetrieverUtils.generateInnerRetrievers(
@@ -319,13 +325,18 @@ protected LinearRetrieverBuilder doRewrite(QueryRewriteContext ctx) {
319325
}
320326
).stream().map(CompoundRetrieverBuilder::convertToRetrieverSource).toList();
321327

322-
float[] weights = new float[fieldsInnerRetrievers.size()];
323-
Arrays.fill(weights, DEFAULT_WEIGHT);
328+
if (fieldsInnerRetrievers.isEmpty() == false) {
329+
float[] weights = new float[fieldsInnerRetrievers.size()];
330+
Arrays.fill(weights, DEFAULT_WEIGHT);
324331

325-
ScoreNormalizer[] normalizers = new ScoreNormalizer[fieldsInnerRetrievers.size()];
326-
Arrays.fill(normalizers, normalizer);
332+
ScoreNormalizer[] normalizers = new ScoreNormalizer[fieldsInnerRetrievers.size()];
333+
Arrays.fill(normalizers, normalizer);
327334

328-
rewritten = new LinearRetrieverBuilder(fieldsInnerRetrievers, null, null, normalizer, rankWindowSize, weights, normalizers);
335+
rewritten = new LinearRetrieverBuilder(fieldsInnerRetrievers, null, null, normalizer, rankWindowSize, weights, normalizers);
336+
} else {
337+
// Inner retriever list can be empty when using an index wildcard pattern that doesn't match any indices
338+
rewritten = new StandardRetrieverBuilder(new MatchNoneQueryBuilder());
339+
}
329340
}
330341

331342
return rewritten;

‎x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverBuilder.java‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
import org.elasticsearch.action.ActionRequestValidationException;
1212
import org.elasticsearch.action.ResolvedIndices;
1313
import org.elasticsearch.common.util.Maps;
14+
import org.elasticsearch.index.query.MatchNoneQueryBuilder;
1415
import org.elasticsearch.index.query.QueryBuilder;
1516
import org.elasticsearch.index.query.QueryRewriteContext;
1617
import org.elasticsearch.license.LicenseUtils;
@@ -20,6 +21,7 @@
2021
import org.elasticsearch.search.retriever.CompoundRetrieverBuilder;
2122
import org.elasticsearch.search.retriever.RetrieverBuilder;
2223
import org.elasticsearch.search.retriever.RetrieverParserContext;
24+
import org.elasticsearch.search.retriever.StandardRetrieverBuilder;
2325
import org.elasticsearch.xcontent.ConstructingObjectParser;
2426
import org.elasticsearch.xcontent.ParseField;
2527
import org.elasticsearch.xcontent.XContentBuilder;
@@ -214,8 +216,8 @@ protected RRFRankDoc[] combineInnerRetrieverResults(List<ScoreDoc[]> rankResults
214216
}
215217

216218
@Override
217-
protected RRFRetrieverBuilder doRewrite(QueryRewriteContext ctx) {
218-
RRFRetrieverBuilder rewritten = this;
219+
protected RetrieverBuilder doRewrite(QueryRewriteContext ctx) {
220+
RetrieverBuilder rewritten = this;
219221

220222
ResolvedIndices resolvedIndices = ctx.getResolvedIndices();
221223
if (resolvedIndices != null && query != null) {
@@ -225,6 +227,10 @@ protected RRFRetrieverBuilder doRewrite(QueryRewriteContext ctx) {
225227
throw new IllegalArgumentException(
226228
"[" + NAME + "] does not support the simplified query format when querying multiple indices"
227229
);
230+
} else if (resolvedIndices.getRemoteClusterIndices().isEmpty() == false) {
231+
throw new IllegalArgumentException(
232+
"[" + NAME + "] does not support the simplified query format when querying remote indices"
233+
);
228234
}
229235

230236
List<RetrieverSource> fieldsInnerRetrievers = SimplifiedInnerRetrieverUtils.generateInnerRetrievers(
@@ -246,7 +252,12 @@ protected RRFRetrieverBuilder doRewrite(QueryRewriteContext ctx) {
246252
}
247253
).stream().map(CompoundRetrieverBuilder::convertToRetrieverSource).toList();
248254

249-
rewritten = new RRFRetrieverBuilder(fieldsInnerRetrievers, rankWindowSize, rankConstant);
255+
if (fieldsInnerRetrievers.isEmpty() == false) {
256+
rewritten = new RRFRetrieverBuilder(fieldsInnerRetrievers, rankWindowSize, rankConstant);
257+
} else {
258+
// Inner retriever list can be empty when using an index wildcard pattern that doesn't match any indices
259+
rewritten = new StandardRetrieverBuilder(new MatchNoneQueryBuilder());
260+
}
250261
}
251262

252263
return rewritten;

‎x-pack/plugin/rank-rrf/src/test/java/org/elasticsearch/xpack/rank/linear/LinearRetrieverBuilderTests.java‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -202,10 +202,13 @@ private static void assertSimplifiedParamsRewrite(
202202
)
203203
);
204204

205-
LinearRetrieverBuilder rewritten = retriever.doRewrite(ctx);
205+
RetrieverBuilder rewritten = retriever.doRewrite(ctx);
206206
assertNotSame(retriever, rewritten);
207-
assertEquals(retriever.rankWindowSize(), rewritten.rankWindowSize());
208-
assertEquals(expectedInnerRetrievers, getInnerRetrieversAsSet(rewritten));
207+
assertTrue(rewritten instanceof LinearRetrieverBuilder);
208+
209+
LinearRetrieverBuilder rewrittenLinear = (LinearRetrieverBuilder) rewritten;
210+
assertEquals(retriever.rankWindowSize(), rewrittenLinear.rankWindowSize());
211+
assertEquals(expectedInnerRetrievers, getInnerRetrieversAsSet(rewrittenLinear));
209212
}
210213

211214
private static Set<InnerRetriever> getInnerRetrieversAsSet(LinearRetrieverBuilder retriever) {

‎x-pack/plugin/rank-rrf/src/test/java/org/elasticsearch/xpack/rank/rrf/RRFRetrieverBuilderTests.java‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -229,11 +229,14 @@ private static void assertSimplifiedParamsRewrite(
229229
}).toArray())
230230
);
231231

232-
RRFRetrieverBuilder rewritten = retriever.doRewrite(ctx);
232+
RetrieverBuilder rewritten = retriever.doRewrite(ctx);
233233
assertNotSame(retriever, rewritten);
234-
assertEquals(retriever.rankWindowSize(), rewritten.rankWindowSize());
235-
assertEquals(retriever.rankConstant(), rewritten.rankConstant());
236-
assertEquals(expectedInnerRetrievers, getInnerRetrieversAsSet(rewritten));
234+
assertTrue(rewritten instanceof RRFRetrieverBuilder);
235+
236+
RRFRetrieverBuilder rewrittenRrf = (RRFRetrieverBuilder) rewritten;
237+
assertEquals(retriever.rankWindowSize(), rewrittenRrf.rankWindowSize());
238+
assertEquals(retriever.rankConstant(), rewrittenRrf.rankConstant());
239+
assertEquals(expectedInnerRetrievers, getInnerRetrieversAsSet(rewrittenRrf));
237240
}
238241

239242
private static Set<Object> getInnerRetrieversAsSet(RRFRetrieverBuilder retriever) {

0 commit comments

Comments
 (0)