Skip to content

Commit 6336b13

Browse files
mch2mch2
committed
add unit test for a terms agg with derived fields
Signed-off-by: Marc Handalian <[email protected]>
1 parent 86ebd16 commit 6336b13

File tree

2 files changed

+139
-1
lines changed

2 files changed

+139
-1
lines changed

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -148,7 +148,7 @@ public DerivedFieldValueFetcher valueFetcher(QueryShardContext context, SearchLo
148148

149149
@Override
150150
public IndexFieldData.Builder fielddataBuilder(String fullyQualifiedIndexName, Supplier<SearchLookup> searchLookup) {
151-
return typeFieldMapper.mappedFieldType.fielddataBuilder(fullyQualifiedIndexName, searchLookup);
151+
return getFieldMapper().mappedFieldType.fielddataBuilder(fullyQualifiedIndexName, searchLookup);
152152
}
153153

154154
@Override

‎server/src/test/java/org/opensearch/search/aggregations/bucket/terms/DerivedFieldAggregationTests.java

+138
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
package org.opensearch.search.aggregations.bucket.terms;
2+
3+
import org.apache.lucene.document.Document;
4+
import org.apache.lucene.document.Field;
5+
import org.apache.lucene.document.KeywordField;
6+
import org.apache.lucene.document.TextField;
7+
import org.apache.lucene.search.MatchAllDocsQuery;
8+
import org.opensearch.Version;
9+
import org.opensearch.cluster.metadata.IndexMetadata;
10+
import org.opensearch.common.settings.Settings;
11+
import org.opensearch.core.index.Index;
12+
import org.opensearch.index.IndexSettings;
13+
import org.opensearch.index.mapper.DerivedField;
14+
import org.opensearch.index.mapper.DerivedFieldResolver;
15+
import org.opensearch.index.mapper.DerivedFieldResolverFactory;
16+
import org.opensearch.index.mapper.DerivedFieldType;
17+
import org.opensearch.index.mapper.DerivedFieldValueFetcher;
18+
import org.opensearch.index.mapper.KeywordFieldMapper;
19+
import org.opensearch.index.mapper.MappedFieldType;
20+
import org.opensearch.index.mapper.MapperService;
21+
import org.opensearch.index.query.QueryShardContext;
22+
import org.opensearch.script.DerivedFieldScript;
23+
import org.opensearch.script.Script;
24+
import org.opensearch.search.aggregations.AggregatorTestCase;
25+
import org.opensearch.search.lookup.LeafSearchLookup;
26+
import org.opensearch.search.lookup.SearchLookup;
27+
import org.opensearch.search.lookup.SourceLookup;
28+
import org.junit.Before;
29+
30+
import java.io.IOException;
31+
import java.util.ArrayList;
32+
import java.util.Collections;
33+
import java.util.List;
34+
35+
import static org.mockito.Mockito.any;
36+
import static org.mockito.Mockito.doReturn;
37+
import static org.mockito.Mockito.mock;
38+
import static org.mockito.Mockito.spy;
39+
import static org.mockito.Mockito.when;
40+
41+
public class DerivedFieldAggregationTests extends AggregatorTestCase {
42+
43+
private QueryShardContext mockContext;
44+
private List<Document> docs;
45+
46+
private static final String[][] raw_requests = new String[][] {
47+
{ "40.135.0.0 GET /images/hm_bg.jpg HTTP/1.0", "200", "40.135.0.0" },
48+
{ "232.0.0.0 GET /images/hm_bg.jpg HTTP/1.0", "400", "232.0.0.0" },
49+
{ "26.1.0.0 GET /images/hm_bg.jpg HTTP/1.0", "200", "26.1.0.0" },
50+
{ "247.37.0.0 GET /french/splash_inet.html HTTP/1.0", "400", "247.37.0.0" },
51+
{ "247.37.0.0 GET /french/splash_inet.html HTTP/1.0", "400", "247.37.0.0" },
52+
{ "247.37.0.0 GET /french/splash_inet.html HTTP/1.0", "200", "247.37.0.0" } };
53+
54+
@Before
55+
public void init() {
56+
super.initValuesSourceRegistry();
57+
// Create a mock QueryShardContext
58+
mockContext = mock(QueryShardContext.class);
59+
when(mockContext.index()).thenReturn(new Index("test_index", "uuid"));
60+
when(mockContext.allowExpensiveQueries()).thenReturn(true);
61+
62+
MapperService mockMapperService = mock(MapperService.class);
63+
when(mockContext.getMapperService()).thenReturn(mockMapperService);
64+
Settings indexSettings = Settings.builder()
65+
.put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT)
66+
.put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1)
67+
.put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 1)
68+
.build();
69+
// Mock IndexSettings
70+
IndexSettings mockIndexSettings = new IndexSettings(
71+
IndexMetadata.builder("test_index").settings(indexSettings).build(),
72+
Settings.EMPTY
73+
);
74+
when(mockMapperService.getIndexSettings()).thenReturn(mockIndexSettings);
75+
when(mockContext.getIndexSettings()).thenReturn(mockIndexSettings);
76+
docs = new ArrayList<>();
77+
for (String[] request : raw_requests) {
78+
Document document = new Document();
79+
document.add(new TextField("raw_request", request[0], Field.Store.YES));
80+
document.add(new KeywordField("status", request[1], Field.Store.YES));
81+
docs.add(document);
82+
}
83+
}
84+
85+
public void testSimpleTermsAggregationWithDerivedField() throws IOException {
86+
MappedFieldType keywordFieldType = new KeywordFieldMapper.KeywordFieldType("status");
87+
88+
SearchLookup searchLookup = mock(SearchLookup.class);
89+
SourceLookup sourceLookup = new SourceLookup();
90+
LeafSearchLookup leafLookup = mock(LeafSearchLookup.class);
91+
when(leafLookup.source()).thenReturn(sourceLookup);
92+
93+
// Mock DerivedFieldScript.Factory
94+
DerivedFieldScript.Factory factory = (params, lookup) -> (DerivedFieldScript.LeafFactory) ctx -> {
95+
when(searchLookup.getLeafSearchLookup(any())).thenReturn(leafLookup);
96+
return new DerivedFieldScript(params, lookup, ctx) {
97+
@Override
98+
public void execute() {
99+
addEmittedValue(raw_requests[sourceLookup.docId()][1]);
100+
}
101+
102+
@Override
103+
public void setDocument(int docid) {
104+
sourceLookup.setSegmentAndDocument(ctx, docid);
105+
}
106+
};
107+
};
108+
109+
DerivedField derivedField = new DerivedField("derived_field", "keyword", new Script(""));
110+
DerivedFieldResolver resolver = DerivedFieldResolverFactory.createResolver(
111+
mockContext,
112+
Collections.emptyMap(),
113+
Collections.singletonList(derivedField),
114+
true
115+
);
116+
117+
// spy on the resolved type so we can mock the valuefetcher
118+
DerivedFieldType derivedFieldType = spy((DerivedFieldType) resolver.resolve("derived_field"));
119+
DerivedFieldScript.LeafFactory leafFactory = factory.newFactory((new Script("")).getParams(), searchLookup);
120+
DerivedFieldValueFetcher valueFetcher = new DerivedFieldValueFetcher(leafFactory, null);
121+
doReturn(valueFetcher).when(derivedFieldType).valueFetcher(any(), any(), any());
122+
123+
TermsAggregationBuilder aggregationBuilder = new TermsAggregationBuilder("derived_terms").field("status").size(10);
124+
125+
testCase(aggregationBuilder, new MatchAllDocsQuery(), iw -> {
126+
for (Document d : docs) {
127+
iw.addDocument(d);
128+
}
129+
}, (InternalTerms result) -> {
130+
assertEquals(2, result.getBuckets().size());
131+
List<Terms.Bucket> buckets = result.getBuckets();
132+
assertEquals("200", buckets.get(0).getKey());
133+
assertEquals(3, buckets.get(0).getDocCount());
134+
assertEquals("400", buckets.get(1).getKey());
135+
assertEquals(3, buckets.get(1).getDocCount());
136+
}, keywordFieldType, derivedFieldType);
137+
}
138+
}

0 commit comments

Comments
 (0)