Skip to content

Commit 5e33217

Browse files
rishabhmauryarishabhmaurya
committed
Add tests for fetch phase
Signed-off-by: Rishabh Maurya <rishabhmaurya05@gmail.com>
1 parent f40465a commit 5e33217

File tree

5 files changed

+104
-31
lines changed

5 files changed

+104
-31
lines changed

‎server/src/main/java/org/opensearch/index/mapper/DerivedFieldType.java

+17-17
Original file line numberDiff line numberDiff line change
@@ -83,27 +83,27 @@ public DerivedFieldValueFetcher valueFetcher(QueryShardContext context, SearchLo
8383
if (format != null) {
8484
throw new IllegalArgumentException("Field [" + name() + "] of type [" + typeName() + "] doesn't support formats.");
8585
}
86-
return new DerivedFieldValueFetcher(getDerivedFieldLeafFactory(context));
86+
return new DerivedFieldValueFetcher(getDerivedFieldLeafFactory(context, searchLookup));
8787
}
8888

8989
@Override
9090
public Query termQuery(Object value, QueryShardContext context) {
9191
Query query = typeFieldMapper.mappedFieldType.termQuery(value, context);
92-
DerivedFieldValueFetcher valueFetcher = new DerivedFieldValueFetcher(getDerivedFieldLeafFactory(context));
92+
DerivedFieldValueFetcher valueFetcher = valueFetcher(context, context.lookup(), null);
9393
return new DerivedFieldQuery(query, valueFetcher, context.lookup(), indexableFieldGenerator, getIndexAnalyzer());
9494
}
9595

9696
@Override
9797
public Query termQueryCaseInsensitive(Object value, @Nullable QueryShardContext context) {
9898
Query query = typeFieldMapper.mappedFieldType.termQueryCaseInsensitive(value, context);
99-
DerivedFieldValueFetcher valueFetcher = new DerivedFieldValueFetcher(getDerivedFieldLeafFactory(context));
99+
DerivedFieldValueFetcher valueFetcher = valueFetcher(context, context.lookup(), null);
100100
return new DerivedFieldQuery(query, valueFetcher, context.lookup(), indexableFieldGenerator, getIndexAnalyzer());
101101
}
102102

103103
@Override
104104
public Query termsQuery(List<?> values, @Nullable QueryShardContext context) {
105105
Query query = typeFieldMapper.mappedFieldType.termsQuery(values, context);
106-
DerivedFieldValueFetcher valueFetcher = new DerivedFieldValueFetcher(getDerivedFieldLeafFactory(context));
106+
DerivedFieldValueFetcher valueFetcher = valueFetcher(context, context.lookup(), null);
107107
return new DerivedFieldQuery(query, valueFetcher, context.lookup(), indexableFieldGenerator, getIndexAnalyzer());
108108
}
109109

@@ -128,7 +128,7 @@ public Query rangeQuery(
128128
parser,
129129
context
130130
);
131-
DerivedFieldValueFetcher valueFetcher = new DerivedFieldValueFetcher(getDerivedFieldLeafFactory(context));
131+
DerivedFieldValueFetcher valueFetcher = valueFetcher(context, context.lookup(), null);
132132
return new DerivedFieldQuery(query, valueFetcher, context.lookup(), indexableFieldGenerator, getIndexAnalyzer());
133133
}
134134

@@ -142,7 +142,7 @@ public Query fuzzyQuery(
142142
QueryShardContext context
143143
) {
144144
Query query = typeFieldMapper.mappedFieldType.fuzzyQuery(value, fuzziness, prefixLength, maxExpansions, transpositions, context);
145-
DerivedFieldValueFetcher valueFetcher = new DerivedFieldValueFetcher(getDerivedFieldLeafFactory(context));
145+
DerivedFieldValueFetcher valueFetcher = valueFetcher(context, context.lookup(), null);
146146
return new DerivedFieldQuery(query, valueFetcher, context.lookup(), indexableFieldGenerator, getIndexAnalyzer());
147147
}
148148

@@ -165,7 +165,7 @@ public Query fuzzyQuery(
165165
method,
166166
context
167167
);
168-
DerivedFieldValueFetcher valueFetcher = new DerivedFieldValueFetcher(getDerivedFieldLeafFactory(context));
168+
DerivedFieldValueFetcher valueFetcher = valueFetcher(context, context.lookup(), null);
169169
return new DerivedFieldQuery(query, valueFetcher, context.lookup(), indexableFieldGenerator, getIndexAnalyzer());
170170
}
171171

@@ -177,7 +177,7 @@ public Query prefixQuery(
177177
QueryShardContext context
178178
) {
179179
Query query = typeFieldMapper.mappedFieldType.prefixQuery(value, method, caseInsensitive, context);
180-
DerivedFieldValueFetcher valueFetcher = new DerivedFieldValueFetcher(getDerivedFieldLeafFactory(context));
180+
DerivedFieldValueFetcher valueFetcher = valueFetcher(context, context.lookup(), null);
181181
return new DerivedFieldQuery(query, valueFetcher, context.lookup(), indexableFieldGenerator, getIndexAnalyzer());
182182
}
183183

@@ -189,14 +189,14 @@ public Query wildcardQuery(
189189
QueryShardContext context
190190
) {
191191
Query query = typeFieldMapper.mappedFieldType.wildcardQuery(value, method, caseInsensitive, context);
192-
DerivedFieldValueFetcher valueFetcher = new DerivedFieldValueFetcher(getDerivedFieldLeafFactory(context));
192+
DerivedFieldValueFetcher valueFetcher = valueFetcher(context, context.lookup(), null);
193193
return new DerivedFieldQuery(query, valueFetcher, context.lookup(), indexableFieldGenerator, getIndexAnalyzer());
194194
}
195195

196196
@Override
197197
public Query normalizedWildcardQuery(String value, @Nullable MultiTermQuery.RewriteMethod method, QueryShardContext context) {
198198
Query query = typeFieldMapper.mappedFieldType.normalizedWildcardQuery(value, method, context);
199-
DerivedFieldValueFetcher valueFetcher = new DerivedFieldValueFetcher(getDerivedFieldLeafFactory(context));
199+
DerivedFieldValueFetcher valueFetcher = valueFetcher(context, context.lookup(), null);
200200
return new DerivedFieldQuery(query, valueFetcher, context.lookup(), indexableFieldGenerator, getIndexAnalyzer());
201201
}
202202

@@ -210,29 +210,29 @@ public Query regexpQuery(
210210
QueryShardContext context
211211
) {
212212
Query query = typeFieldMapper.mappedFieldType.regexpQuery(value, syntaxFlags, matchFlags, maxDeterminizedStates, method, context);
213-
DerivedFieldValueFetcher valueFetcher = new DerivedFieldValueFetcher(getDerivedFieldLeafFactory(context));
213+
DerivedFieldValueFetcher valueFetcher = valueFetcher(context, context.lookup(), null);
214214
return new DerivedFieldQuery(query, valueFetcher, context.lookup(), indexableFieldGenerator, getIndexAnalyzer());
215215
}
216216

217217
@Override
218218
public Query phraseQuery(TokenStream stream, int slop, boolean enablePositionIncrements, QueryShardContext context) throws IOException {
219219
Query query = typeFieldMapper.mappedFieldType.phraseQuery(stream, slop, enablePositionIncrements, context);
220-
DerivedFieldValueFetcher valueFetcher = new DerivedFieldValueFetcher(getDerivedFieldLeafFactory(context));
220+
DerivedFieldValueFetcher valueFetcher = valueFetcher(context, context.lookup(), null);
221221
return new DerivedFieldQuery(query, valueFetcher, context.lookup(), indexableFieldGenerator, getIndexAnalyzer());
222222
}
223223

224224
@Override
225225
public Query multiPhraseQuery(TokenStream stream, int slop, boolean enablePositionIncrements, QueryShardContext context)
226226
throws IOException {
227227
Query query = typeFieldMapper.mappedFieldType.multiPhraseQuery(stream, slop, enablePositionIncrements, context);
228-
DerivedFieldValueFetcher valueFetcher = new DerivedFieldValueFetcher(getDerivedFieldLeafFactory(context));
228+
DerivedFieldValueFetcher valueFetcher = valueFetcher(context, context.lookup(), null);
229229
return new DerivedFieldQuery(query, valueFetcher, context.lookup(), indexableFieldGenerator, getIndexAnalyzer());
230230
}
231231

232232
@Override
233233
public Query phrasePrefixQuery(TokenStream stream, int slop, int maxExpansions, QueryShardContext context) throws IOException {
234234
Query query = typeFieldMapper.mappedFieldType.phrasePrefixQuery(stream, slop, maxExpansions, context);
235-
DerivedFieldValueFetcher valueFetcher = new DerivedFieldValueFetcher(getDerivedFieldLeafFactory(context));
235+
DerivedFieldValueFetcher valueFetcher = valueFetcher(context, context.lookup(), null);
236236
return new DerivedFieldQuery(query, valueFetcher, context.lookup(), indexableFieldGenerator, getIndexAnalyzer());
237237
}
238238

@@ -246,7 +246,7 @@ public SpanQuery spanPrefixQuery(String value, SpanMultiTermQueryWrapper.SpanRew
246246
@Override
247247
public Query distanceFeatureQuery(Object origin, String pivot, float boost, QueryShardContext context) {
248248
Query query = typeFieldMapper.mappedFieldType.distanceFeatureQuery(origin, pivot, boost, context);
249-
DerivedFieldValueFetcher valueFetcher = new DerivedFieldValueFetcher(getDerivedFieldLeafFactory(context));
249+
DerivedFieldValueFetcher valueFetcher = valueFetcher(context, context.lookup(), null);
250250
return new DerivedFieldQuery(query, valueFetcher, context.lookup(), indexableFieldGenerator, getIndexAnalyzer());
251251
}
252252

@@ -260,7 +260,7 @@ public boolean isAggregatable() {
260260
return false;
261261
}
262262

263-
private DerivedFieldScript.LeafFactory getDerivedFieldLeafFactory(QueryShardContext context) {
263+
private DerivedFieldScript.LeafFactory getDerivedFieldLeafFactory(QueryShardContext context, SearchLookup searchLookup) {
264264
if (!context.documentMapper("").sourceMapper().enabled()) {
265265
throw new IllegalArgumentException(
266266
"DerivedFieldQuery error: unable to fetch fields from _source field: _source is disabled in the mappings "
@@ -270,6 +270,6 @@ private DerivedFieldScript.LeafFactory getDerivedFieldLeafFactory(QueryShardCont
270270
);
271271
}
272272
DerivedFieldScript.Factory factory = context.compile(derivedField.getScript(), DerivedFieldScript.CONTEXT);
273-
return factory.newFactory(derivedField.getScript().getParams(), context.lookup());
273+
return factory.newFactory(derivedField.getScript().getParams(), searchLookup);
274274
}
275275
}

‎server/src/main/java/org/opensearch/index/mapper/DerivedFieldValueFetcher.java

+1
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ public List<Object> fetchValues(SourceLookup lookup) {
3737
return derivedFieldScript.getEmittedValues();
3838
}
3939

40+
@Override
4041
public void setNextReader(LeafReaderContext context) {
4142
try {
4243
derivedFieldScript = derivedFieldScriptFactory.newInstance(context);

‎server/src/main/java/org/opensearch/script/DerivedFieldScript.java

-7
Original file line numberDiff line numberDiff line change
@@ -68,13 +68,6 @@ public DerivedFieldScript(Map<String, Object> params, SearchLookup lookup, LeafR
6868
this.totalByteSize = 0;
6969
}
7070

71-
public DerivedFieldScript() {
72-
this.params = null;
73-
this.leafLookup = null;
74-
this.emittedValues = new ArrayList<>();
75-
this.totalByteSize = 0;
76-
}
77-
7871
/**
7972
* Return the parameters for this script.
8073
*/

‎server/src/test/java/org/opensearch/search/fetch/subphase/FieldFetcherTests.java

+74-2
Original file line numberDiff line numberDiff line change
@@ -43,20 +43,34 @@
4343
import org.opensearch.index.IndexSettings;
4444
import org.opensearch.index.mapper.MapperService;
4545
import org.opensearch.index.query.QueryShardContext;
46+
import org.opensearch.script.MockScriptEngine;
47+
import org.opensearch.script.ScriptEngine;
48+
import org.opensearch.script.ScriptModule;
49+
import org.opensearch.script.ScriptService;
50+
import org.opensearch.search.lookup.LeafSearchLookup;
51+
import org.opensearch.search.lookup.SearchLookup;
4652
import org.opensearch.search.lookup.SourceLookup;
4753
import org.opensearch.test.OpenSearchSingleNodeTestCase;
4854

4955
import java.io.IOException;
56+
import java.util.Collections;
5057
import java.util.List;
5158
import java.util.Map;
5259
import java.util.Set;
5360

61+
import static java.util.Collections.singletonMap;
5462
import static org.hamcrest.Matchers.containsInAnyOrder;
5563
import static org.hamcrest.Matchers.equalTo;
5664
import static org.hamcrest.Matchers.hasItems;
65+
import static org.mockito.ArgumentMatchers.any;
66+
import static org.mockito.Mockito.mock;
67+
import static org.mockito.Mockito.when;
5768

5869
public class FieldFetcherTests extends OpenSearchSingleNodeTestCase {
5970

71+
private static String DERIVED_FIELD_SCRIPT_1 = "derived_field_script_1";
72+
private static String DERIVED_FIELD_SCRIPT_2 = "derived_field_script_2";
73+
6074
public void testLeafValues() throws IOException {
6175
MapperService mapperService = createMapperService();
6276
XContentBuilder source = XContentFactory.jsonBuilder()
@@ -435,6 +449,45 @@ public void testTextSubFields() throws IOException {
435449
}
436450
}
437451

452+
public void testDerivedFields() throws IOException {
453+
XContentBuilder mapping = XContentFactory.jsonBuilder()
454+
.startObject()
455+
.startObject("derived")
456+
.startObject("derived_1")
457+
.field("type", "keyword")
458+
.startObject("script")
459+
.field("source", DERIVED_FIELD_SCRIPT_1)
460+
.field("lang", "mockscript")
461+
.endObject()
462+
.endObject()
463+
.startObject("derived_2")
464+
.field("type", "keyword")
465+
.startObject("script")
466+
.field("source", DERIVED_FIELD_SCRIPT_2)
467+
.field("lang", "mockscript")
468+
.endObject()
469+
.endObject()
470+
.endObject()
471+
.endObject();
472+
473+
IndexService indexService = createIndex("index", Settings.EMPTY, MapperService.SINGLE_MAPPING_NAME, mapping);
474+
MapperService mapperService = indexService.mapperService();
475+
476+
XContentBuilder source = XContentFactory.jsonBuilder()
477+
.startObject()
478+
.field("field1", "some text 1")
479+
.field("field2", "some text 2")
480+
.endObject();
481+
482+
Map<String, DocumentField> fields = fetchFields(mapperService, source, "*");
483+
assertThat(fields.size(), equalTo(2));
484+
assertThat(fields.keySet(), containsInAnyOrder("derived_1", "derived_2"));
485+
assertThat(fields.get("derived_1").getValues().size(), equalTo(1));
486+
assertThat(fields.get("derived_2").getValues().size(), equalTo(1));
487+
assertThat(fields.get("derived_1").getValue(), equalTo("some text 1"));
488+
assertThat(fields.get("derived_2").getValue(), equalTo("some text 2"));
489+
}
490+
438491
private static Map<String, DocumentField> fetchFields(MapperService mapperService, XContentBuilder source, String fieldPattern)
439492
throws IOException {
440493

@@ -448,7 +501,13 @@ private static Map<String, DocumentField> fetchFields(MapperService mapperServic
448501
SourceLookup sourceLookup = new SourceLookup();
449502
sourceLookup.setSource(BytesReference.bytes(source));
450503

451-
FieldFetcher fieldFetcher = FieldFetcher.create(createQueryShardContext(mapperService), null, fields);
504+
SearchLookup searchLookup = mock(SearchLookup.class);
505+
LeafSearchLookup leafSearchLookup = mock(LeafSearchLookup.class);
506+
when(searchLookup.source()).thenReturn(sourceLookup);
507+
when(searchLookup.getLeafSearchLookup(any())).thenReturn(leafSearchLookup);
508+
when(leafSearchLookup.source()).thenReturn(sourceLookup);
509+
FieldFetcher fieldFetcher = FieldFetcher.create(createQueryShardContext(mapperService), searchLookup, fields);
510+
fieldFetcher.setNextReader(null);
452511
return fieldFetcher.fetch(sourceLookup, Set.of());
453512
}
454513

@@ -497,6 +556,19 @@ private static QueryShardContext createQueryShardContext(MapperService mapperSer
497556
.build();
498557
IndexMetadata indexMetadata = new IndexMetadata.Builder("index").settings(settings).build();
499558
IndexSettings indexSettings = new IndexSettings(indexMetadata, settings);
559+
560+
final MockScriptEngine engine = new MockScriptEngine(
561+
MockScriptEngine.NAME,
562+
Map.of(
563+
DERIVED_FIELD_SCRIPT_1,
564+
(script) -> ((Map<String, Object>) script.get("_source")).get("field1"),
565+
DERIVED_FIELD_SCRIPT_2,
566+
(script) -> ((Map<String, Object>) script.get("_source")).get("field2")
567+
),
568+
Collections.emptyMap()
569+
);
570+
final Map<String, ScriptEngine> engines = singletonMap(engine.getType(), engine);
571+
ScriptService scriptService = new ScriptService(Settings.EMPTY, engines, ScriptModule.CORE_CONTEXTS);
500572
return new QueryShardContext(
501573
0,
502574
indexSettings,
@@ -505,7 +577,7 @@ private static QueryShardContext createQueryShardContext(MapperService mapperSer
505577
null,
506578
mapperService,
507579
null,
508-
null,
580+
scriptService,
509581
null,
510582
null,
511583
null,

‎test/framework/src/main/java/org/opensearch/script/MockScriptEngine.java

+12-5
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@
4343
import org.opensearch.search.aggregations.pipeline.MovingFunctionScript;
4444
import org.opensearch.search.lookup.LeafSearchLookup;
4545
import org.opensearch.search.lookup.SearchLookup;
46+
import org.opensearch.search.lookup.SourceLookup;
4647

4748
import java.io.IOException;
4849
import java.util.Collections;
@@ -282,16 +283,22 @@ public double execute(Map<String, Object> params1, double[] values) {
282283
IntervalFilterScript.Factory factory = mockCompiled::createIntervalFilterScript;
283284
return context.factoryClazz.cast(factory);
284285
} else if (context.instanceClazz.equals(DerivedFieldScript.class)) {
285-
DerivedFieldScript.Factory factory = (derivedFieldsParams, lookup) -> ctx -> new DerivedFieldScript(
286-
derivedFieldsParams,
286+
DerivedFieldScript.Factory factory = (derivedFieldParams, lookup) -> ctx -> new DerivedFieldScript(
287+
derivedFieldParams,
287288
lookup,
288289
ctx
289290
) {
291+
@Override
292+
public void setDocument(int docid) {}
293+
290294
@Override
291295
public void execute() {
292-
Map<String, Object> vars = new HashMap<>(derivedFieldsParams);
293-
vars.put("params", derivedFieldsParams);
294-
script.apply(vars);
296+
Map<String, Object> vars = new HashMap<>(derivedFieldParams);
297+
SourceLookup sourceLookup = lookup.source();
298+
vars.put("params", derivedFieldParams);
299+
vars.put("_source", sourceLookup.loadSourceIfNeeded());
300+
// currently supports adding one value, can be extended to emit multiple values too.
301+
addEmittedValue(script.apply(vars));
295302
}
296303
};
297304
return context.factoryClazz.cast(factory);

0 commit comments

Comments
 (0)