Skip to content

Commit 5b25dee

Browse files
pmpailispmpailis
authored
Propagating nested inner_hits to the parent compound retriever (#116408)
1 parent 5204902 commit 5b25dee

File tree

19 files changed

+248
-42
lines changed

19 files changed

+248
-42
lines changed

‎docs/changelog/116408.yaml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
pr: 116408
2+
summary: Propagating nested `inner_hits` to the parent compound retriever
3+
area: Ranking
4+
type: bug
5+
issues:
6+
- 116397

‎server/src/internalClusterTest/java/org/elasticsearch/search/nested/SimpleNestedIT.java‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,9 @@
2121
import org.elasticsearch.cluster.health.ClusterHealthStatus;
2222
import org.elasticsearch.common.settings.Settings;
2323
import org.elasticsearch.core.TimeValue;
24+
import org.elasticsearch.index.query.InnerHitBuilder;
2425
import org.elasticsearch.index.query.QueryBuilders;
26+
import org.elasticsearch.search.builder.SearchSourceBuilder;
2527
import org.elasticsearch.search.sort.NestedSortBuilder;
2628
import org.elasticsearch.search.sort.SortBuilders;
2729
import org.elasticsearch.search.sort.SortMode;
@@ -1581,6 +1583,64 @@ public void testCheckFixedBitSetCache() throws Exception {
15811583
assertThat(clusterStatsResponse.getIndicesStats().getSegments().getBitsetMemoryInBytes(), equalTo(0L));
15821584
}
15831585

1586+
public void testSkipNestedInnerHits() throws Exception {
1587+
assertAcked(prepareCreate("test").setMapping("nested1", "type=nested"));
1588+
ensureGreen();
1589+
1590+
prepareIndex("test").setId("1")
1591+
.setSource(
1592+
jsonBuilder().startObject()
1593+
.field("field1", "value1")
1594+
.startArray("nested1")
1595+
.startObject()
1596+
.field("n_field1", "foo")
1597+
.field("n_field2", "bar")
1598+
.endObject()
1599+
.endArray()
1600+
.endObject()
1601+
)
1602+
.get();
1603+
1604+
waitForRelocation(ClusterHealthStatus.GREEN);
1605+
GetResponse getResponse = client().prepareGet("test", "1").get();
1606+
assertThat(getResponse.isExists(), equalTo(true));
1607+
assertThat(getResponse.getSourceAsBytesRef(), notNullValue());
1608+
refresh();
1609+
1610+
assertNoFailuresAndResponse(
1611+
prepareSearch("test").setSource(
1612+
new SearchSourceBuilder().query(
1613+
QueryBuilders.nestedQuery("nested1", QueryBuilders.termQuery("nested1.n_field1", "foo"), ScoreMode.Avg)
1614+
.innerHit(new InnerHitBuilder())
1615+
)
1616+
),
1617+
res -> {
1618+
assertNotNull(res.getHits());
1619+
assertHitCount(res, 1);
1620+
assertThat(res.getHits().getHits().length, equalTo(1));
1621+
// by default we should get inner hits
1622+
assertNotNull(res.getHits().getHits()[0].getInnerHits());
1623+
assertNotNull(res.getHits().getHits()[0].getInnerHits().get("nested1"));
1624+
}
1625+
);
1626+
1627+
assertNoFailuresAndResponse(
1628+
prepareSearch("test").setSource(
1629+
new SearchSourceBuilder().query(
1630+
QueryBuilders.nestedQuery("nested1", QueryBuilders.termQuery("nested1.n_field1", "foo"), ScoreMode.Avg)
1631+
.innerHit(new InnerHitBuilder())
1632+
).skipInnerHits(true)
1633+
),
1634+
res -> {
1635+
assertNotNull(res.getHits());
1636+
assertHitCount(res, 1);
1637+
assertThat(res.getHits().getHits().length, equalTo(1));
1638+
// if we explicitly say to ignore inner hits, then this should now be null
1639+
assertNull(res.getHits().getHits()[0].getInnerHits());
1640+
}
1641+
);
1642+
}
1643+
15841644
private void assertDocumentCount(String index, long numdocs) {
15851645
IndicesStatsResponse stats = indicesAdmin().prepareStats(index).clear().setDocs(true).get();
15861646
assertNoFailures(stats);

‎server/src/main/java/org/elasticsearch/TransportVersions.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -194,7 +194,7 @@ static TransportVersion def(int id) {
194194
public static final TransportVersion DATA_STREAM_INDEX_VERSION_DEPRECATION_CHECK = def(8_788_00_0);
195195
public static final TransportVersion ADD_COMPATIBILITY_VERSIONS_TO_NODE_INFO = def(8_789_00_0);
196196
public static final TransportVersion VERTEX_AI_INPUT_TYPE_ADDED = def(8_790_00_0);
197-
197+
public static final TransportVersion SKIP_INNER_HITS_SEARCH_SOURCE = def(8_791_00_0);
198198
/*
199199
* STOP! READ THIS FIRST! No, really,
200200
* ____ _____ ___ ____ _ ____ _____ _ ____ _____ _ _ ___ ____ _____ ___ ____ ____ _____ _

‎server/src/main/java/org/elasticsearch/search/retriever/rankdoc/RankDocsQueryBuilder.java‎ renamed to ‎server/src/main/java/org/elasticsearch/index/query/RankDocsQueryBuilder.java‎

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
* License v3.0 only", or the "Server Side Public License, v 1".
88
*/
99

10-
package org.elasticsearch.search.retriever.rankdoc;
10+
package org.elasticsearch.index.query;
1111

1212
import org.apache.lucene.index.IndexReader;
1313
import org.apache.lucene.search.Query;
@@ -16,15 +16,13 @@
1616
import org.elasticsearch.common.Strings;
1717
import org.elasticsearch.common.io.stream.StreamInput;
1818
import org.elasticsearch.common.io.stream.StreamOutput;
19-
import org.elasticsearch.index.query.AbstractQueryBuilder;
20-
import org.elasticsearch.index.query.QueryBuilder;
21-
import org.elasticsearch.index.query.QueryRewriteContext;
22-
import org.elasticsearch.index.query.SearchExecutionContext;
2319
import org.elasticsearch.search.rank.RankDoc;
20+
import org.elasticsearch.search.retriever.rankdoc.RankDocsQuery;
2421
import org.elasticsearch.xcontent.XContentBuilder;
2522

2623
import java.io.IOException;
2724
import java.util.Arrays;
25+
import java.util.Map;
2826
import java.util.Objects;
2927

3028
import static org.elasticsearch.TransportVersions.RRF_QUERY_REWRITE;
@@ -55,6 +53,15 @@ public RankDocsQueryBuilder(StreamInput in) throws IOException {
5553
}
5654
}
5755

56+
@Override
57+
protected void extractInnerHitBuilders(Map<String, InnerHitContextBuilder> innerHits) {
58+
if (queryBuilders != null) {
59+
for (QueryBuilder query : queryBuilders) {
60+
InnerHitContextBuilder.extractInnerHits(query, innerHits);
61+
}
62+
}
63+
}
64+
5865
@Override
5966
protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) throws IOException {
6067
if (queryBuilders != null) {
@@ -71,7 +78,7 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) throws
7178
return super.doRewrite(queryRewriteContext);
7279
}
7380

74-
RankDoc[] rankDocs() {
81+
public RankDoc[] rankDocs() {
7582
return rankDocs;
7683
}
7784

‎server/src/main/java/org/elasticsearch/rest/action/search/SearchCapabilities.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,8 @@ private SearchCapabilities() {}
3636
private static final String KQL_QUERY_SUPPORTED = "kql_query";
3737
/** Support multi-dense-vector field mapper. */
3838
private static final String MULTI_DENSE_VECTOR_FIELD_MAPPER = "multi_dense_vector_field_mapper";
39+
/** Support propagating nested retrievers' inner_hits to top-level compound retrievers . */
40+
private static final String NESTED_RETRIEVER_INNER_HITS_SUPPORT = "nested_retriever_inner_hits_support";
3941

4042
public static final Set<String> CAPABILITIES;
4143
static {
@@ -45,6 +47,7 @@ private SearchCapabilities() {}
4547
capabilities.add(BYTE_FLOAT_BIT_DOT_PRODUCT_CAPABILITY);
4648
capabilities.add(DENSE_VECTOR_DOCVALUE_FIELDS);
4749
capabilities.add(TRANSFORM_RANK_RRF_TO_RETRIEVER);
50+
capabilities.add(NESTED_RETRIEVER_INNER_HITS_SUPPORT);
4851
if (MultiDenseVectorFieldMapper.FEATURE_FLAG.isEnabled()) {
4952
capabilities.add(MULTI_DENSE_VECTOR_FIELD_MAPPER);
5053
}

‎server/src/main/java/org/elasticsearch/search/SearchModule.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@
5252
import org.elasticsearch.index.query.QueryBuilder;
5353
import org.elasticsearch.index.query.QueryStringQueryBuilder;
5454
import org.elasticsearch.index.query.RangeQueryBuilder;
55+
import org.elasticsearch.index.query.RankDocsQueryBuilder;
5556
import org.elasticsearch.index.query.RegexpQueryBuilder;
5657
import org.elasticsearch.index.query.ScriptQueryBuilder;
5758
import org.elasticsearch.index.query.SimpleQueryStringBuilder;
@@ -238,7 +239,6 @@
238239
import org.elasticsearch.search.retriever.RetrieverBuilder;
239240
import org.elasticsearch.search.retriever.RetrieverParserContext;
240241
import org.elasticsearch.search.retriever.StandardRetrieverBuilder;
241-
import org.elasticsearch.search.retriever.rankdoc.RankDocsQueryBuilder;
242242
import org.elasticsearch.search.sort.FieldSortBuilder;
243243
import org.elasticsearch.search.sort.GeoDistanceSortBuilder;
244244
import org.elasticsearch.search.sort.ScoreSortBuilder;

‎server/src/main/java/org/elasticsearch/search/SearchService.java‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1285,13 +1285,17 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
12851285
);
12861286
if (query != null) {
12871287
QueryBuilder rewrittenForInnerHits = Rewriteable.rewrite(query, innerHitsRewriteContext, true);
1288-
InnerHitContextBuilder.extractInnerHits(rewrittenForInnerHits, innerHitBuilders);
1288+
if (false == source.skipInnerHits()) {
1289+
InnerHitContextBuilder.extractInnerHits(rewrittenForInnerHits, innerHitBuilders);
1290+
}
12891291
searchExecutionContext.setAliasFilter(context.request().getAliasFilter().getQueryBuilder());
12901292
context.parsedQuery(searchExecutionContext.toQuery(query));
12911293
}
12921294
if (source.postFilter() != null) {
12931295
QueryBuilder rewrittenForInnerHits = Rewriteable.rewrite(source.postFilter(), innerHitsRewriteContext, true);
1294-
InnerHitContextBuilder.extractInnerHits(rewrittenForInnerHits, innerHitBuilders);
1296+
if (false == source.skipInnerHits()) {
1297+
InnerHitContextBuilder.extractInnerHits(rewrittenForInnerHits, innerHitBuilders);
1298+
}
12951299
context.parsedPostFilter(searchExecutionContext.toQuery(source.postFilter()));
12961300
}
12971301
if (innerHitBuilders.size() > 0) {

‎server/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java‎

Lines changed: 27 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -214,6 +214,8 @@ public static HighlightBuilder highlight() {
214214

215215
private Map<String, Object> runtimeMappings = emptyMap();
216216

217+
private boolean skipInnerHits = false;
218+
217219
/**
218220
* Constructs a new search source builder.
219221
*/
@@ -290,6 +292,11 @@ public SearchSourceBuilder(StreamInput in) throws IOException {
290292
if (in.getTransportVersion().onOrAfter(TransportVersions.V_8_8_0)) {
291293
rankBuilder = in.readOptionalNamedWriteable(RankBuilder.class);
292294
}
295+
if (in.getTransportVersion().onOrAfter(TransportVersions.SKIP_INNER_HITS_SEARCH_SOURCE)) {
296+
skipInnerHits = in.readBoolean();
297+
} else {
298+
skipInnerHits = false;
299+
}
293300
}
294301

295302
@Override
@@ -379,6 +386,9 @@ public void writeTo(StreamOutput out) throws IOException {
379386
} else if (rankBuilder != null) {
380387
throw new IllegalArgumentException("cannot serialize [rank] to version [" + out.getTransportVersion().toReleaseVersion() + "]");
381388
}
389+
if (out.getTransportVersion().onOrAfter(TransportVersions.SKIP_INNER_HITS_SEARCH_SOURCE)) {
390+
out.writeBoolean(skipInnerHits);
391+
}
382392
}
383393

384394
/**
@@ -1280,6 +1290,7 @@ private SearchSourceBuilder shallowCopy(
12801290
rewrittenBuilder.collapse = collapse;
12811291
rewrittenBuilder.pointInTimeBuilder = pointInTimeBuilder;
12821292
rewrittenBuilder.runtimeMappings = runtimeMappings;
1293+
rewrittenBuilder.skipInnerHits = skipInnerHits;
12831294
return rewrittenBuilder;
12841295
}
12851296

@@ -1838,6 +1849,9 @@ public XContentBuilder innerToXContent(XContentBuilder builder, Params params) t
18381849
if (false == runtimeMappings.isEmpty()) {
18391850
builder.field(RUNTIME_MAPPINGS_FIELD.getPreferredName(), runtimeMappings);
18401851
}
1852+
if (skipInnerHits) {
1853+
builder.field("skipInnerHits", true);
1854+
}
18411855

18421856
return builder;
18431857
}
@@ -1850,6 +1864,15 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
18501864
return builder;
18511865
}
18521866

1867+
public SearchSourceBuilder skipInnerHits(boolean skipInnerHits) {
1868+
this.skipInnerHits = skipInnerHits;
1869+
return this;
1870+
}
1871+
1872+
public boolean skipInnerHits() {
1873+
return this.skipInnerHits;
1874+
}
1875+
18531876
public static class IndexBoost implements Writeable, ToXContentObject {
18541877
private final String index;
18551878
private final float boost;
@@ -2104,7 +2127,8 @@ public int hashCode() {
21042127
collapse,
21052128
trackTotalHitsUpTo,
21062129
pointInTimeBuilder,
2107-
runtimeMappings
2130+
runtimeMappings,
2131+
skipInnerHits
21082132
);
21092133
}
21102134

@@ -2149,7 +2173,8 @@ public boolean equals(Object obj) {
21492173
&& Objects.equals(collapse, other.collapse)
21502174
&& Objects.equals(trackTotalHitsUpTo, other.trackTotalHitsUpTo)
21512175
&& Objects.equals(pointInTimeBuilder, other.pointInTimeBuilder)
2152-
&& Objects.equals(runtimeMappings, other.runtimeMappings);
2176+
&& Objects.equals(runtimeMappings, other.runtimeMappings)
2177+
&& Objects.equals(skipInnerHits, other.skipInnerHits);
21532178
}
21542179

21552180
@Override

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -236,7 +236,7 @@ public int doHashCode() {
236236
return Objects.hash(innerRetrievers);
237237
}
238238

239-
protected SearchSourceBuilder createSearchSourceBuilder(PointInTimeBuilder pit, RetrieverBuilder retrieverBuilder) {
239+
protected final SearchSourceBuilder createSearchSourceBuilder(PointInTimeBuilder pit, RetrieverBuilder retrieverBuilder) {
240240
var sourceBuilder = new SearchSourceBuilder().pointInTimeBuilder(pit)
241241
.trackTotalHits(false)
242242
.storedFields(new StoredFieldsContext(false))
@@ -254,6 +254,11 @@ protected SearchSourceBuilder createSearchSourceBuilder(PointInTimeBuilder pit,
254254
}
255255
sortBuilders.add(new FieldSortBuilder(FieldSortBuilder.SHARD_DOC_FIELD_NAME));
256256
sourceBuilder.sort(sortBuilders);
257+
sourceBuilder.skipInnerHits(true);
258+
return finalizeSourceBuilder(sourceBuilder);
259+
}
260+
261+
protected SearchSourceBuilder finalizeSourceBuilder(SearchSourceBuilder sourceBuilder) {
257262
return sourceBuilder;
258263
}
259264

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,8 @@
1515
import org.elasticsearch.index.query.BoolQueryBuilder;
1616
import org.elasticsearch.index.query.QueryBuilder;
1717
import org.elasticsearch.index.query.QueryRewriteContext;
18+
import org.elasticsearch.index.query.RankDocsQueryBuilder;
1819
import org.elasticsearch.search.builder.SearchSourceBuilder;
19-
import org.elasticsearch.search.retriever.rankdoc.RankDocsQueryBuilder;
2020
import org.elasticsearch.search.vectors.ExactKnnQueryBuilder;
2121
import org.elasticsearch.search.vectors.KnnSearchBuilder;
2222
import org.elasticsearch.search.vectors.QueryVectorBuilder;

0 commit comments

Comments
 (0)