Skip to content

Commit 931c1aa

Browse files
Gankris96Ganesh Ramadurai
Gankris96
and
Ganesh Ramadurai
authored
Propagate includes and excludes from fetchSourceContext to FieldsVisitor (#17080)
Signed-off-by: Ganesh Ramadurai <gramadur@icloud.com> Co-authored-by: Ganesh Ramadurai <gramadur@amazon.com>
1 parent 13ab4ec commit 931c1aa

File tree

6 files changed

+151
-5
lines changed

6 files changed

+151
-5
lines changed

‎CHANGELOG.md

+1
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
3636
- Improve flat_object field parsing performance by reducing two passes to a single pass ([#16297](https://github.com/opensearch-project/OpenSearch/pull/16297))
3737
- Improve performance of the bitmap filtering([#16936](https://github.com/opensearch-project/OpenSearch/pull/16936/))
3838
- Introduce Template query ([#16818](https://github.com/opensearch-project/OpenSearch/pull/16818))
39+
- Propagate the sourceIncludes and excludes fields from fetchSourceContext to FieldsVisitor. ([#17080](https://github.com/opensearch-project/OpenSearch/pull/17080))
3940

4041
### Dependencies
4142
- Bump `com.google.cloud:google-cloud-core-http` from 2.23.0 to 2.47.0 ([#16504](https://github.com/opensearch-project/OpenSearch/pull/16504))

‎server/src/main/java/org/opensearch/index/fieldvisitor/CustomFieldsVisitor.java

+5
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,11 @@ public CustomFieldsVisitor(Set<String> fields, boolean loadSource) {
5252
this.fields = fields;
5353
}
5454

55+
public CustomFieldsVisitor(Set<String> fields, boolean loadSource, String[] includes, String[] excludes) {
56+
super(loadSource, includes, excludes);
57+
this.fields = fields;
58+
}
59+
5560
@Override
5661
public Status needsField(FieldInfo fieldInfo) {
5762
if (super.needsField(fieldInfo) == Status.YES) {

‎server/src/main/java/org/opensearch/index/fieldvisitor/FieldsVisitor.java

+30-1
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
import org.apache.lucene.index.FieldInfo;
3535
import org.apache.lucene.index.StoredFieldVisitor;
3636
import org.apache.lucene.util.BytesRef;
37+
import org.opensearch.core.common.Strings;
3738
import org.opensearch.core.common.bytes.BytesArray;
3839
import org.opensearch.core.common.bytes.BytesReference;
3940
import org.opensearch.index.mapper.IdFieldMapper;
@@ -66,17 +67,29 @@ public class FieldsVisitor extends StoredFieldVisitor {
6667
private final boolean loadSource;
6768
private final String sourceFieldName;
6869
private final Set<String> requiredFields;
70+
private final String[] sourceIncludes;
71+
private final String[] sourceExcludes;
6972
protected BytesReference source;
7073
protected String id;
7174
protected Map<String, List<Object>> fieldsValues;
7275

7376
public FieldsVisitor(boolean loadSource) {
74-
this(loadSource, SourceFieldMapper.NAME);
77+
this(loadSource, SourceFieldMapper.NAME, null, null);
78+
}
79+
80+
public FieldsVisitor(boolean loadSource, String[] includes, String[] excludes) {
81+
this(loadSource, SourceFieldMapper.NAME, includes, excludes);
7582
}
7683

7784
public FieldsVisitor(boolean loadSource, String sourceFieldName) {
85+
this(loadSource, sourceFieldName, null, null);
86+
}
87+
88+
public FieldsVisitor(boolean loadSource, String sourceFieldName, String[] includes, String[] excludes) {
7889
this.loadSource = loadSource;
7990
this.sourceFieldName = sourceFieldName;
91+
this.sourceIncludes = includes != null ? includes : Strings.EMPTY_ARRAY;
92+
this.sourceExcludes = excludes != null ? excludes : Strings.EMPTY_ARRAY;
8093
requiredFields = new HashSet<>();
8194
reset();
8295
}
@@ -162,6 +175,22 @@ public BytesReference source() {
162175
return source;
163176
}
164177

178+
/**
179+
* Returns the array containing the source fields to include
180+
* @return String[] sourceIncludes
181+
*/
182+
public String[] includes() {
183+
return sourceIncludes;
184+
}
185+
186+
/**
187+
* Returns the array containing the source fields to exclude
188+
* @return String[] sourceExcludes
189+
*/
190+
public String[] excludes() {
191+
return sourceExcludes;
192+
}
193+
165194
public String id() {
166195
return id;
167196
}

‎server/src/main/java/org/opensearch/search/fetch/FetchPhase.java

+17-4
Original file line numberDiff line numberDiff line change
@@ -221,7 +221,7 @@ public int compareTo(DocIdToIndex o) {
221221
}
222222
}
223223

224-
private FieldsVisitor createStoredFieldsVisitor(SearchContext context, Map<String, Set<String>> storedToRequestedFields) {
224+
protected FieldsVisitor createStoredFieldsVisitor(SearchContext context, Map<String, Set<String>> storedToRequestedFields) {
225225
StoredFieldsContext storedFieldsContext = context.storedFieldsContext();
226226

227227
if (storedFieldsContext == null) {
@@ -230,7 +230,11 @@ private FieldsVisitor createStoredFieldsVisitor(SearchContext context, Map<Strin
230230
context.fetchSourceContext(new FetchSourceContext(true));
231231
}
232232
boolean loadSource = sourceRequired(context);
233-
return new FieldsVisitor(loadSource);
233+
return new FieldsVisitor(
234+
loadSource,
235+
context.hasFetchSourceContext() ? context.fetchSourceContext().includes() : null,
236+
context.hasFetchSourceContext() ? context.fetchSourceContext().excludes() : null
237+
);
234238
} else if (storedFieldsContext.fetchFields() == false) {
235239
// disable stored fields entirely
236240
return null;
@@ -262,9 +266,18 @@ private FieldsVisitor createStoredFieldsVisitor(SearchContext context, Map<Strin
262266
boolean loadSource = sourceRequired(context);
263267
if (storedToRequestedFields.isEmpty()) {
264268
// empty list specified, default to disable _source if no explicit indication
265-
return new FieldsVisitor(loadSource);
269+
return new FieldsVisitor(
270+
loadSource,
271+
context.hasFetchSourceContext() ? context.fetchSourceContext().includes() : null,
272+
context.hasFetchSourceContext() ? context.fetchSourceContext().excludes() : null
273+
);
266274
} else {
267-
return new CustomFieldsVisitor(storedToRequestedFields.keySet(), loadSource);
275+
return new CustomFieldsVisitor(
276+
storedToRequestedFields.keySet(),
277+
loadSource,
278+
context.hasFetchSourceContext() ? context.fetchSourceContext().includes() : null,
279+
context.hasFetchSourceContext() ? context.fetchSourceContext().excludes() : null
280+
);
268281
}
269282
}
270283
}

‎server/src/test/java/org/opensearch/index/mapper/StoredNumericValuesTests.java

+34
Original file line numberDiff line numberDiff line change
@@ -41,9 +41,11 @@
4141
import org.opensearch.common.lucene.Lucene;
4242
import org.opensearch.common.util.set.Sets;
4343
import org.opensearch.common.xcontent.XContentFactory;
44+
import org.opensearch.core.common.Strings;
4445
import org.opensearch.core.common.bytes.BytesReference;
4546
import org.opensearch.core.xcontent.MediaTypeRegistry;
4647
import org.opensearch.index.fieldvisitor.CustomFieldsVisitor;
48+
import org.opensearch.index.fieldvisitor.FieldsVisitor;
4749
import org.opensearch.index.mapper.MapperService.MergeReason;
4850
import org.opensearch.test.OpenSearchSingleNodeTestCase;
4951

@@ -200,4 +202,36 @@ public void testBytesAndNumericRepresentation() throws Exception {
200202
reader.close();
201203
writer.close();
202204
}
205+
206+
public void testFieldsVisitorValidateIncludesExcludes() throws Exception {
207+
Set<String> fieldNames = Sets.newHashSet(
208+
"field1",
209+
"field2",
210+
"field3",
211+
"field4",
212+
"field5",
213+
"field6",
214+
"field7",
215+
"field8",
216+
"field9",
217+
"field10",
218+
"field11"
219+
);
220+
String[] includes = { "field1", "field2", "field3" };
221+
String[] excludes = { "field7", "field8" };
222+
223+
CustomFieldsVisitor fieldsVisitor = new CustomFieldsVisitor(fieldNames, false, includes, excludes);
224+
225+
assertArrayEquals(fieldsVisitor.includes(), includes);
226+
assertArrayEquals(fieldsVisitor.excludes(), excludes);
227+
228+
FieldsVisitor fieldsVisitor1 = new FieldsVisitor(false, includes, excludes);
229+
assertArrayEquals(fieldsVisitor1.includes(), includes);
230+
assertArrayEquals(fieldsVisitor1.excludes(), excludes);
231+
232+
FieldsVisitor fieldsVisitor2 = new FieldsVisitor(false);
233+
assertArrayEquals(fieldsVisitor2.includes(), Strings.EMPTY_ARRAY);
234+
assertArrayEquals(fieldsVisitor2.excludes(), Strings.EMPTY_ARRAY);
235+
236+
}
203237
}

‎server/src/test/java/org/opensearch/search/fetch/FetchPhaseTests.java

+64
Original file line numberDiff line numberDiff line change
@@ -32,8 +32,22 @@
3232

3333
package org.opensearch.search.fetch;
3434

35+
import org.opensearch.index.fieldvisitor.CustomFieldsVisitor;
36+
import org.opensearch.index.fieldvisitor.FieldsVisitor;
37+
import org.opensearch.search.fetch.subphase.FetchSourceContext;
38+
import org.opensearch.search.internal.SearchContext;
3539
import org.opensearch.test.OpenSearchTestCase;
3640

41+
import java.util.ArrayList;
42+
import java.util.Collections;
43+
import java.util.HashMap;
44+
import java.util.List;
45+
import java.util.Map;
46+
import java.util.Set;
47+
48+
import static org.mockito.Mockito.mock;
49+
import static org.mockito.Mockito.when;
50+
3751
public class FetchPhaseTests extends OpenSearchTestCase {
3852
public void testSequentialDocs() {
3953
FetchPhase.DocIdToIndex[] docs = new FetchPhase.DocIdToIndex[10];
@@ -52,4 +66,54 @@ public void testSequentialDocs() {
5266
}
5367
assertFalse(FetchPhase.hasSequentialDocs(docs));
5468
}
69+
70+
public void testFieldsVisitorsInFetchPhase() {
71+
72+
FetchPhase fetchPhase = new FetchPhase(new ArrayList<>());
73+
SearchContext mockSearchContext = mock(SearchContext.class);
74+
when(mockSearchContext.docIdsToLoadSize()).thenReturn(1);
75+
when(mockSearchContext.docIdsToLoad()).thenReturn(new int[] { 1 });
76+
String[] includes = new String[] { "field1", "field2" };
77+
String[] excludes = new String[] { "field7", "field8" };
78+
79+
FetchSourceContext mockFetchSourceContext = new FetchSourceContext(true, includes, excludes);
80+
when(mockSearchContext.hasFetchSourceContext()).thenReturn(true);
81+
when(mockSearchContext.fetchSourceContext()).thenReturn(mockFetchSourceContext);
82+
83+
// Case 1
84+
// if storedFieldsContext is null
85+
FieldsVisitor fieldsVisitor = fetchPhase.createStoredFieldsVisitor(mockSearchContext, null);
86+
assertArrayEquals(fieldsVisitor.excludes(), excludes);
87+
assertArrayEquals(fieldsVisitor.includes(), includes);
88+
89+
// Case 2
90+
// if storedFieldsContext is not null
91+
StoredFieldsContext storedFieldsContext = mock(StoredFieldsContext.class);
92+
when(mockSearchContext.storedFieldsContext()).thenReturn(storedFieldsContext);
93+
94+
fieldsVisitor = fetchPhase.createStoredFieldsVisitor(mockSearchContext, null);
95+
assertNull(fieldsVisitor);
96+
97+
// Case 3
98+
// if storedFieldsContext is true but fieldNames are empty
99+
when(storedFieldsContext.fetchFields()).thenReturn(true);
100+
when(storedFieldsContext.fieldNames()).thenReturn(List.of());
101+
fieldsVisitor = fetchPhase.createStoredFieldsVisitor(mockSearchContext, Collections.emptyMap());
102+
assertArrayEquals(fieldsVisitor.excludes(), excludes);
103+
assertArrayEquals(fieldsVisitor.includes(), includes);
104+
105+
// Case 4
106+
// if storedToRequested Fields is not empty
107+
// creates an instance of CustomFieldsVisitor
108+
Map<String, Set<String>> storedToRequestedFields = new HashMap<>();
109+
storedToRequestedFields.put("test_field_key", Set.of("test_field_value"));
110+
111+
fieldsVisitor = fetchPhase.createStoredFieldsVisitor(mockSearchContext, storedToRequestedFields);
112+
113+
assertTrue(fieldsVisitor instanceof CustomFieldsVisitor);
114+
assertArrayEquals(fieldsVisitor.excludes(), excludes);
115+
assertArrayEquals(fieldsVisitor.includes(), includes);
116+
117+
}
118+
55119
}

0 commit comments

Comments
 (0)