Skip to content

Commit

Permalink
Shard id awareness of SearchLookup (opensearch-project#1063)
Browse files Browse the repository at this point in the history
* Shard id awareness of SearchLookup

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

* Add unit test for deprecated constructor

Signed-off-by: Andrew Ross <[email protected]>

---------

Signed-off-by: Alexander Koval <[email protected]>
Signed-off-by: Andrew Ross <[email protected]>
Co-authored-by: Andrew Ross <[email protected]>
Signed-off-by: Aman Khare <[email protected]>
  • Loading branch information
2 people authored and Aman Khare committed Mar 12, 2024
1 parent 648775c commit 0483871
Show file tree
Hide file tree
Showing 12 changed files with 73 additions and 12 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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
);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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(
Expand All @@ -91,6 +110,7 @@ public SearchLookup(
sourceLookup = new SourceLookup();
fieldsLookup = new FieldsLookup(mapperService);
this.fieldDataLookup = fieldDataLookup;
this.shardId = shardId;
}

/**
Expand All @@ -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;
}

/**
Expand Down Expand Up @@ -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;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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());
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,8 @@

public class QueryShardContextTests extends OpenSearchTestCase {

private static final int SHARD_ID = 0;

public void testFailIfFieldMappingNotFound() {
QueryShardContext context = createQueryShardContext(IndexMetadata.INDEX_UUID_NA_VALUE, null);
context.setAllowUnmappedFields(false);
Expand Down Expand Up @@ -307,6 +309,11 @@ public void testFielddataLookupOneFieldManyReferences() throws IOException {
assertEquals(Arrays.asList(expectedFirstDoc.toString(), expectedSecondDoc.toString()), collect("field", queryShardContext));
}

public void testSearchLookupShardId() {
SearchLookup searchLookup = createQueryShardContext("uuid", null, null).lookup();
assertEquals(SHARD_ID, searchLookup.shardId());
}

public static QueryShardContext createQueryShardContext(String indexUuid, String clusterAlias) {
return createQueryShardContext(indexUuid, clusterAlias, null);
}
Expand Down Expand Up @@ -343,7 +350,7 @@ private static QueryShardContext createQueryShardContext(
}
final long nowInMillis = randomNonNegativeLong();
return new QueryShardContext(
0,
SHARD_ID,
indexSettings,
BigArrays.NON_RECYCLING_INSTANCE,
null,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
/*
* SPDX-License-Identifier: Apache-2.0
*
* The OpenSearch Contributors require contributions made to
* this file be licensed under the Apache-2.0 license or a
* compatible open source license.
*/

package org.opensearch.search.lookup;

import org.opensearch.index.mapper.MapperService;
import org.opensearch.test.OpenSearchTestCase;

import static org.mockito.Mockito.mock;

public class SearchLookupTests extends OpenSearchTestCase {
public void testDeprecatedConstructorShardId() {
final SearchLookup searchLookup = new SearchLookup(mock(MapperService.class), (a, b) -> null);
assertThrows(IllegalStateException.class, searchLookup::shardId);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down

0 comments on commit 0483871

Please sign in to comment.