Shard id awareness of SearchLookup

Signed-off-by: Alexander Koval <[email protected]>
opensearch-project · Nov 8, 2021 · 4254bab · 4254bab
1 parent 546ab21
commit 4254bab
Show file tree

Hide file tree

Showing 9 changed files with 34 additions and 10 deletions.
diff --git a/...expression/src/test/java/org/opensearch/script/expression/ExpressionFieldScriptTests.java b/...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/...ssion/src/test/java/org/opensearch/script/expression/ExpressionNumberSortScriptTests.java b/...ssion/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/...pression/src/test/java/org/opensearch/script/expression/ExpressionTermsSetQueryTests.java b/...pression/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
@@ -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
@@ -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<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,
@@ -76,7 +83,8 @@ public class SearchLookup {
     public SearchLookup(
         MapperService mapperService,
         BiFunction<MappedFieldType, Supplier<SearchLookup>, 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<String> 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
@@ -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
@@ -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
@@ -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
@@ -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);