diff --git a/modules/lang-expression/src/test/java/org/opensearch/script/expression/ExpressionFieldScriptTests.java b/modules/lang-expression/src/test/java/org/opensearch/script/expression/ExpressionFieldScriptTests.java index 1d6c7730d7ff3..45aabd063326f 100644 --- a/modules/lang-expression/src/test/java/org/opensearch/script/expression/ExpressionFieldScriptTests.java +++ b/modules/lang-expression/src/test/java/org/opensearch/script/expression/ExpressionFieldScriptTests.java @@ -77,7 +77,7 @@ public void setUp() throws Exception { when(fieldData.load(anyObject())).thenReturn(atomicFieldData); service = new ExpressionScriptEngine(); - lookup = new SearchLookup(mapperService, (ignored, lookup) -> fieldData, null); + lookup = new SearchLookup(mapperService, (ignored, lookup) -> fieldData, null, SearchLookup.UNKNOWN_SHARD_ID); } private FieldScript.LeafFactory compile(String expression) { diff --git a/modules/lang-expression/src/test/java/org/opensearch/script/expression/ExpressionNumberSortScriptTests.java b/modules/lang-expression/src/test/java/org/opensearch/script/expression/ExpressionNumberSortScriptTests.java index 66967188c0890..f282c1ba4c2fb 100644 --- a/modules/lang-expression/src/test/java/org/opensearch/script/expression/ExpressionNumberSortScriptTests.java +++ b/modules/lang-expression/src/test/java/org/opensearch/script/expression/ExpressionNumberSortScriptTests.java @@ -76,7 +76,7 @@ public void setUp() throws Exception { when(fieldData.load(anyObject())).thenReturn(atomicFieldData); service = new ExpressionScriptEngine(); - lookup = new SearchLookup(mapperService, (ignored, lookup) -> fieldData, null); + lookup = new SearchLookup(mapperService, (ignored, lookup) -> fieldData, null, SearchLookup.UNKNOWN_SHARD_ID); } private NumberSortScript.LeafFactory compile(String expression) { diff --git a/modules/lang-expression/src/test/java/org/opensearch/script/expression/ExpressionTermsSetQueryTests.java b/modules/lang-expression/src/test/java/org/opensearch/script/expression/ExpressionTermsSetQueryTests.java index 7d142992cc330..f066af7ce8c96 100644 --- a/modules/lang-expression/src/test/java/org/opensearch/script/expression/ExpressionTermsSetQueryTests.java +++ b/modules/lang-expression/src/test/java/org/opensearch/script/expression/ExpressionTermsSetQueryTests.java @@ -76,7 +76,7 @@ public void setUp() throws Exception { when(fieldData.load(anyObject())).thenReturn(atomicFieldData); service = new ExpressionScriptEngine(); - lookup = new SearchLookup(mapperService, (ignored, lookup) -> fieldData, null); + lookup = new SearchLookup(mapperService, (ignored, lookup) -> fieldData, null, SearchLookup.UNKNOWN_SHARD_ID); } private TermsSetQueryScript.LeafFactory compile(String expression) { diff --git a/server/src/main/java/org/opensearch/index/query/QueryShardContext.java b/server/src/main/java/org/opensearch/index/query/QueryShardContext.java index 47bade8b6aa51..ab7a40ec436d3 100644 --- a/server/src/main/java/org/opensearch/index/query/QueryShardContext.java +++ b/server/src/main/java/org/opensearch/index/query/QueryShardContext.java @@ -378,7 +378,8 @@ public SearchLookup lookup() { this.lookup = new SearchLookup( getMapperService(), (fieldType, searchLookup) -> indexFieldDataService.apply(fieldType, fullyQualifiedIndex.getName(), searchLookup), - types + types, + shardId ); } return this.lookup; @@ -395,7 +396,8 @@ public SearchLookup newFetchLookup() { return new SearchLookup( getMapperService(), (fieldType, searchLookup) -> indexFieldDataService.apply(fieldType, fullyQualifiedIndex.getName(), searchLookup), - types + types, + shardId ); } diff --git a/server/src/main/java/org/opensearch/search/lookup/SearchLookup.java b/server/src/main/java/org/opensearch/search/lookup/SearchLookup.java index 269052f895066..006a3fdba8103 100644 --- a/server/src/main/java/org/opensearch/search/lookup/SearchLookup.java +++ b/server/src/main/java/org/opensearch/search/lookup/SearchLookup.java @@ -54,6 +54,12 @@ public class SearchLookup { */ private static final int MAX_FIELD_CHAIN_DEPTH = 5; + /** + * This constant should be used in cases when shard id is unknown. + * Mostly it should be used in tests. + */ + public static final int UNKNOWN_SHARD_ID = -1; + /** * The chain of fields for which this lookup was created, used for detecting * loops caused by runtime fields referring to other runtime fields. The chain is empty @@ -68,6 +74,7 @@ public class SearchLookup { private final SourceLookup sourceLookup; private final FieldsLookup fieldsLookup; private final BiFunction, IndexFieldData> fieldDataLookup; + private final int shardId; /** * Create the top level field lookup for a search request. Provides a way to look up fields from doc_values, @@ -76,7 +83,8 @@ public class SearchLookup { public SearchLookup( MapperService mapperService, BiFunction, IndexFieldData> fieldDataLookup, - @Nullable String[] types + @Nullable String[] types, + int shardId ) { this.fieldChain = Collections.emptySet(); docMap = new DocLookup( @@ -87,6 +95,7 @@ public SearchLookup( sourceLookup = new SourceLookup(); fieldsLookup = new FieldsLookup(mapperService, types); this.fieldDataLookup = fieldDataLookup; + this.shardId = shardId; } /** @@ -106,6 +115,7 @@ private SearchLookup(SearchLookup searchLookup, Set fieldChain) { this.sourceLookup = searchLookup.sourceLookup; this.fieldsLookup = searchLookup.fieldsLookup; this.fieldDataLookup = searchLookup.fieldDataLookup; + this.shardId = searchLookup.shardId; } /** @@ -140,4 +150,11 @@ public DocLookup doc() { public SourceLookup source() { return sourceLookup; } + + public int shardId() { + if (shardId == UNKNOWN_SHARD_ID) { + throw new IllegalStateException("Shard id is unknown for this lookup"); + } + return shardId; + } } diff --git a/server/src/test/java/org/opensearch/index/fielddata/IndexFieldDataServiceTests.java b/server/src/test/java/org/opensearch/index/fielddata/IndexFieldDataServiceTests.java index a6dcae9621d13..bd060b7e14754 100644 --- a/server/src/test/java/org/opensearch/index/fielddata/IndexFieldDataServiceTests.java +++ b/server/src/test/java/org/opensearch/index/fielddata/IndexFieldDataServiceTests.java @@ -143,7 +143,7 @@ public void testGetForFieldRuntimeField() { searchLookupSetOnce.set(searchLookup); return (IndexFieldData.Builder) (cache, breakerService) -> null; }); - SearchLookup searchLookup = new SearchLookup(null, null, null); + SearchLookup searchLookup = new SearchLookup(null, null, null, SearchLookup.UNKNOWN_SHARD_ID); ifdService.getForField(ft, "qualified", () -> searchLookup); assertSame(searchLookup, searchLookupSetOnce.get().get()); } diff --git a/server/src/test/java/org/opensearch/search/aggregations/support/ScriptValuesTests.java b/server/src/test/java/org/opensearch/search/aggregations/support/ScriptValuesTests.java index 3709f4daefaca..e35cd5116bd53 100644 --- a/server/src/test/java/org/opensearch/search/aggregations/support/ScriptValuesTests.java +++ b/server/src/test/java/org/opensearch/search/aggregations/support/ScriptValuesTests.java @@ -60,7 +60,7 @@ private static class FakeAggregationScript extends AggregationScript { int index; FakeAggregationScript(Object[][] values) { - super(Collections.emptyMap(), new SearchLookup(null, null, Strings.EMPTY_ARRAY) { + super(Collections.emptyMap(), new SearchLookup(null, null, Strings.EMPTY_ARRAY, SearchLookup.UNKNOWN_SHARD_ID) { @Override public LeafSearchLookup getLeafSearchLookup(LeafReaderContext context) { diff --git a/test/framework/src/main/java/org/opensearch/index/mapper/MapperServiceTestCase.java b/test/framework/src/main/java/org/opensearch/index/mapper/MapperServiceTestCase.java index 31a140e5bd869..d0ca656ed53b3 100644 --- a/test/framework/src/main/java/org/opensearch/index/mapper/MapperServiceTestCase.java +++ b/test/framework/src/main/java/org/opensearch/index/mapper/MapperServiceTestCase.java @@ -250,7 +250,12 @@ QueryShardContext createQueryShardContext(MapperService mapperService) { ); when(queryShardContext.allowExpensiveQueries()).thenReturn(true); when(queryShardContext.lookup()).thenReturn( - new SearchLookup(mapperService, (ft, s) -> { throw new UnsupportedOperationException("search lookup not available"); }, null) + new SearchLookup( + mapperService, + (ft, s) -> { throw new UnsupportedOperationException("search lookup not available"); }, + null, + SearchLookup.UNKNOWN_SHARD_ID + ) ); return queryShardContext; } diff --git a/test/framework/src/main/java/org/opensearch/index/mapper/MapperTestCase.java b/test/framework/src/main/java/org/opensearch/index/mapper/MapperTestCase.java index fc5c7283ed8b3..fc28585093603 100644 --- a/test/framework/src/main/java/org/opensearch/index/mapper/MapperTestCase.java +++ b/test/framework/src/main/java/org/opensearch/index/mapper/MapperTestCase.java @@ -293,7 +293,7 @@ protected final List fetchFromDocValues(MapperService mapperService, MappedFi mapperService, iw -> { iw.addDocument(mapperService.documentMapper().parse(source(b -> b.field(ft.name(), sourceValue))).rootDoc()); }, iw -> { - SearchLookup lookup = new SearchLookup(mapperService, fieldDataLookup, null); + SearchLookup lookup = new SearchLookup(mapperService, fieldDataLookup, null, SearchLookup.UNKNOWN_SHARD_ID); ValueFetcher valueFetcher = new DocValueFetcher(format, lookup.doc().getForField(ft)); IndexSearcher searcher = newSearcher(iw); LeafReaderContext context = searcher.getIndexReader().leaves().get(0);