Shard id awareness of SearchLookup

Signed-off-by: Alexander Koval <[email protected]>

CHANGELOG.md

@@ -103,6 +103,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Introduce query level setting `index.query.max_nested_depth` limiting nested queries ([#3268](https://github.com/opensearch-project/OpenSearch/issues/3268)
 - Add toString methods to MultiSearchRequest, MultiGetRequest and CreateIndexRequest ([#12163](https://github.com/opensearch-project/OpenSearch/pull/12163))
 - Support for returning scores in matched queries ([#11626](https://github.com/opensearch-project/OpenSearch/pull/11626))
+- Add shard id property to SearchLookup for use in field types provided by plugins ([#1063](https://github.com/opensearch-project/OpenSearch/pull/1063))
 
 ### Dependencies
 - Bump `peter-evans/find-comment` from 2 to 3 ([#12288](https://github.com/opensearch-project/OpenSearch/pull/12288))

modules/lang-expression/src/test/java/org/opensearch/script/expression/ExpressionFieldScriptTests.java

@@ -77,7 +77,7 @@ public void setUp() throws Exception {
         when(fieldData.load(any())).thenReturn(atomicFieldData);
 
         service = new ExpressionScriptEngine();
-        lookup = new SearchLookup(mapperService, (ignored, lookup) -> fieldData);
+        lookup = new SearchLookup(mapperService, (ignored, lookup) -> fieldData, SearchLookup.UNKNOWN_SHARD_ID);
     }
 
     private FieldScript.LeafFactory compile(String expression) {

modules/lang-expression/src/test/java/org/opensearch/script/expression/ExpressionNumberSortScriptTests.java

@@ -77,7 +77,7 @@ public void setUp() throws Exception {
         when(fieldData.load(any())).thenReturn(atomicFieldData);
 
         service = new ExpressionScriptEngine();
-        lookup = new SearchLookup(mapperService, (ignored, lookup) -> fieldData);
+        lookup = new SearchLookup(mapperService, (ignored, lookup) -> fieldData, SearchLookup.UNKNOWN_SHARD_ID);
     }
 
     private NumberSortScript.LeafFactory compile(String expression) {

modules/lang-expression/src/test/java/org/opensearch/script/expression/ExpressionTermsSetQueryTests.java

@@ -77,7 +77,7 @@ public void setUp() throws Exception {
         when(fieldData.load(any())).thenReturn(atomicFieldData);
 
         service = new ExpressionScriptEngine();
-        lookup = new SearchLookup(mapperService, (ignored, lookup) -> fieldData);
+        lookup = new SearchLookup(mapperService, (ignored, lookup) -> fieldData, SearchLookup.UNKNOWN_SHARD_ID);
     }
 
     private TermsSetQueryScript.LeafFactory compile(String expression) {

server/src/main/java/org/opensearch/index/query/QueryShardContext.java

@@ -423,7 +423,8 @@ public SearchLookup lookup() {
         if (this.lookup == null) {
             this.lookup = new SearchLookup(
                 getMapperService(),
-                (fieldType, searchLookup) -> indexFieldDataService.apply(fieldType, fullyQualifiedIndex.getName(), searchLookup)
+                (fieldType, searchLookup) -> indexFieldDataService.apply(fieldType, fullyQualifiedIndex.getName(), searchLookup),
+                shardId
             );
         }
         return this.lookup;
@@ -439,7 +440,8 @@ public SearchLookup newFetchLookup() {
          */
         return new SearchLookup(
             getMapperService(),
-            (fieldType, searchLookup) -> indexFieldDataService.apply(fieldType, fullyQualifiedIndex.getName(), searchLookup)
+            (fieldType, searchLookup) -> indexFieldDataService.apply(fieldType, fullyQualifiedIndex.getName(), searchLookup),
+            shardId
         );
     }
 

server/src/main/java/org/opensearch/search/lookup/SearchLookup.java

@@ -60,6 +60,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
@@ -74,14 +80,27 @@ public class SearchLookup {
     private final SourceLookup sourceLookup;
     private final FieldsLookup fieldsLookup;
     private final BiFunction<MappedFieldType, Supplier<SearchLookup>, 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,
-     * stored fields, or _source.
+     * Constructor for backwards compatibility. Use the one with explicit shardId argument.
      */
+    @Deprecated
     public SearchLookup(
         MapperService mapperService,
         BiFunction<MappedFieldType, Supplier<SearchLookup>, IndexFieldData<?>> fieldDataLookup
+    ) {
+        this(mapperService, fieldDataLookup, UNKNOWN_SHARD_ID);
+    }
+
+    /**
+     * Create the top level field lookup for a search request. Provides a way to look up fields from doc_values,
+     * stored fields, or _source.
+     */
+    public SearchLookup(
+        MapperService mapperService,
+        BiFunction<MappedFieldType, Supplier<SearchLookup>, IndexFieldData<?>> fieldDataLookup,
+        int shardId
     ) {
         this.fieldChain = Collections.emptySet();
         docMap = new DocLookup(
@@ -91,6 +110,7 @@ public SearchLookup(
         sourceLookup = new SourceLookup();
         fieldsLookup = new FieldsLookup(mapperService);
         this.fieldDataLookup = fieldDataLookup;
+        this.shardId = shardId;
     }
 
     /**
@@ -109,6 +129,7 @@ private SearchLookup(SearchLookup searchLookup, Set<String> fieldChain) {
         this.sourceLookup = searchLookup.sourceLookup;
         this.fieldsLookup = searchLookup.fieldsLookup;
         this.fieldDataLookup = searchLookup.fieldDataLookup;
+        this.shardId = searchLookup.shardId;
     }
 
     /**
@@ -143,4 +164,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;
+    }
 }

server/src/test/java/org/opensearch/index/fielddata/IndexFieldDataServiceTests.java

@@ -138,13 +138,15 @@ public void testGetForFieldRuntimeField() {
         );
         final SetOnce<Supplier<SearchLookup>> searchLookupSetOnce = new SetOnce<>();
         MappedFieldType ft = mock(MappedFieldType.class);
+        final int shardId = randomInt();
         when(ft.fielddataBuilder(Mockito.any(), Mockito.any())).thenAnswer(invocationOnMock -> {
             @SuppressWarnings("unchecked")
             Supplier<SearchLookup> searchLookup = (Supplier<SearchLookup>) invocationOnMock.getArguments()[1];
             searchLookupSetOnce.set(searchLookup);
+            assertEquals(searchLookup.get().shardId(), shardId);
             return (IndexFieldData.Builder) (cache, breakerService) -> null;
         });
-        SearchLookup searchLookup = new SearchLookup(null, null);
+        SearchLookup searchLookup = new SearchLookup(null, null, shardId);
         ifdService.getForField(ft, "qualified", () -> searchLookup);
         assertSame(searchLookup, searchLookupSetOnce.get().get());
     }

server/src/test/java/org/opensearch/index/query/QueryShardContextTests.java

@@ -307,6 +307,11 @@ public void testFielddataLookupOneFieldManyReferences() throws IOException {
         assertEquals(Arrays.asList(expectedFirstDoc.toString(), expectedSecondDoc.toString()), collect("field", queryShardContext));
     }
 
+    public void testSearchLookupShardId() throws IOException {
+        SearchLookup searchLookup = createQueryShardContext("uuid", null, null).lookup();
+        assertEquals(0, searchLookup.shardId());
+    }
+
     public static QueryShardContext createQueryShardContext(String indexUuid, String clusterAlias) {
         return createQueryShardContext(indexUuid, clusterAlias, null);
     }

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) {
+            super(Collections.emptyMap(), new SearchLookup(null, null, SearchLookup.UNKNOWN_SHARD_ID) {
 
                 @Override
                 public LeafSearchLookup getLeafSearchLookup(LeafReaderContext context) {

test/framework/src/main/java/org/opensearch/index/mapper/MapperServiceTestCase.java

@@ -253,7 +253,7 @@ protected 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");
-        }));
+        }, SearchLookup.UNKNOWN_SHARD_ID));
         when(queryShardContext.getFieldType(any())).thenAnswer(inv -> mapperService.fieldType(inv.getArguments()[0].toString()));
         when(queryShardContext.documentMapper(anyString())).thenReturn(mapperService.documentMapper());
         return queryShardContext;

test/framework/src/main/java/org/opensearch/index/mapper/MapperTestCase.java

@@ -293,7 +293,7 @@ protected final List<?> fetchFromDocValues(MapperService mapperService, MappedFi
         withLuceneIndex(mapperService, iw -> {
             iw.addDocument(mapperService.documentMapper().parse(source(b -> b.field(ft.name(), sourceValue))).rootDoc());
         }, iw -> {
-            SearchLookup lookup = new SearchLookup(mapperService, fieldDataLookup);
+            SearchLookup lookup = new SearchLookup(mapperService, fieldDataLookup, SearchLookup.UNKNOWN_SHARD_ID);
             ValueFetcher valueFetcher = new DocValueFetcher(format, lookup.doc().getForField(ft));
             IndexSearcher searcher = newSearcher(iw);
             LeafReaderContext context = searcher.getIndexReader().leaves().get(0);