Propagate includes and excludes from fetchSourceContext to FieldsVisitor (#17080)

Signed-off-by: Ganesh Ramadurai <[email protected]>
Co-authored-by: Ganesh Ramadurai <[email protected]>

Author: Gankris96

CHANGELOG.md

@@ -36,6 +36,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Improve flat_object field parsing performance by reducing two passes to a single pass ([#16297](https://github.com/opensearch-project/OpenSearch/pull/16297))
 - Improve performance of the bitmap filtering([#16936](https://github.com/opensearch-project/OpenSearch/pull/16936/))
 - Introduce Template query ([#16818](https://github.com/opensearch-project/OpenSearch/pull/16818))
+- Propagate the sourceIncludes and excludes fields from fetchSourceContext to FieldsVisitor. ([#17080](https://github.com/opensearch-project/OpenSearch/pull/17080))
 
 ### Dependencies
 - 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

@@ -52,6 +52,11 @@ public CustomFieldsVisitor(Set<String> fields, boolean loadSource) {
         this.fields = fields;
     }
 
+    public CustomFieldsVisitor(Set<String> fields, boolean loadSource, String[] includes, String[] excludes) {
+        super(loadSource, includes, excludes);
+        this.fields = fields;
+    }
+
     @Override
     public Status needsField(FieldInfo fieldInfo) {
         if (super.needsField(fieldInfo) == Status.YES) {

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

@@ -34,6 +34,7 @@
 import org.apache.lucene.index.FieldInfo;
 import org.apache.lucene.index.StoredFieldVisitor;
 import org.apache.lucene.util.BytesRef;
+import org.opensearch.core.common.Strings;
 import org.opensearch.core.common.bytes.BytesArray;
 import org.opensearch.core.common.bytes.BytesReference;
 import org.opensearch.index.mapper.IdFieldMapper;
@@ -66,17 +67,29 @@ public class FieldsVisitor extends StoredFieldVisitor {
     private final boolean loadSource;
     private final String sourceFieldName;
     private final Set<String> requiredFields;
+    private final String[] sourceIncludes;
+    private final String[] sourceExcludes;
     protected BytesReference source;
     protected String id;
     protected Map<String, List<Object>> fieldsValues;
 
     public FieldsVisitor(boolean loadSource) {
-        this(loadSource, SourceFieldMapper.NAME);
+        this(loadSource, SourceFieldMapper.NAME, null, null);
+    }
+
+    public FieldsVisitor(boolean loadSource, String[] includes, String[] excludes) {
+        this(loadSource, SourceFieldMapper.NAME, includes, excludes);
     }
 
     public FieldsVisitor(boolean loadSource, String sourceFieldName) {
+        this(loadSource, sourceFieldName, null, null);
+    }
+
+    public FieldsVisitor(boolean loadSource, String sourceFieldName, String[] includes, String[] excludes) {
         this.loadSource = loadSource;
         this.sourceFieldName = sourceFieldName;
+        this.sourceIncludes = includes != null ? includes : Strings.EMPTY_ARRAY;
+        this.sourceExcludes = excludes != null ? excludes : Strings.EMPTY_ARRAY;
         requiredFields = new HashSet<>();
         reset();
     }
@@ -162,6 +175,22 @@ public BytesReference source() {
         return source;
     }
 
+    /**
+     * Returns the array containing the source fields to include
+     * @return String[] sourceIncludes
+     */
+    public String[] includes() {
+        return sourceIncludes;
+    }
+
+    /**
+     * Returns the array containing the source fields to exclude
+     * @return String[] sourceExcludes
+     */
+    public String[] excludes() {
+        return sourceExcludes;
+    }
+
     public String id() {
         return id;
     }

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

@@ -221,7 +221,7 @@ public int compareTo(DocIdToIndex o) {
         }
     }
 
-    private FieldsVisitor createStoredFieldsVisitor(SearchContext context, Map<String, Set<String>> storedToRequestedFields) {
+    protected FieldsVisitor createStoredFieldsVisitor(SearchContext context, Map<String, Set<String>> storedToRequestedFields) {
         StoredFieldsContext storedFieldsContext = context.storedFieldsContext();
 
         if (storedFieldsContext == null) {
@@ -230,7 +230,11 @@ private FieldsVisitor createStoredFieldsVisitor(SearchContext context, Map<Strin
                 context.fetchSourceContext(new FetchSourceContext(true));
             }
             boolean loadSource = sourceRequired(context);
-            return new FieldsVisitor(loadSource);
+            return new FieldsVisitor(
+                loadSource,
+                context.hasFetchSourceContext() ? context.fetchSourceContext().includes() : null,
+                context.hasFetchSourceContext() ? context.fetchSourceContext().excludes() : null
+            );
         } else if (storedFieldsContext.fetchFields() == false) {
             // disable stored fields entirely
             return null;
@@ -262,9 +266,18 @@ private FieldsVisitor createStoredFieldsVisitor(SearchContext context, Map<Strin
             boolean loadSource = sourceRequired(context);
             if (storedToRequestedFields.isEmpty()) {
                 // empty list specified, default to disable _source if no explicit indication
-                return new FieldsVisitor(loadSource);
+                return new FieldsVisitor(
+                    loadSource,
+                    context.hasFetchSourceContext() ? context.fetchSourceContext().includes() : null,
+                    context.hasFetchSourceContext() ? context.fetchSourceContext().excludes() : null
+                );
             } else {
-                return new CustomFieldsVisitor(storedToRequestedFields.keySet(), loadSource);
+                return new CustomFieldsVisitor(
+                    storedToRequestedFields.keySet(),
+                    loadSource,
+                    context.hasFetchSourceContext() ? context.fetchSourceContext().includes() : null,
+                    context.hasFetchSourceContext() ? context.fetchSourceContext().excludes() : null
+                );
             }
         }
     }

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

@@ -41,9 +41,11 @@
 import org.opensearch.common.lucene.Lucene;
 import org.opensearch.common.util.set.Sets;
 import org.opensearch.common.xcontent.XContentFactory;
+import org.opensearch.core.common.Strings;
 import org.opensearch.core.common.bytes.BytesReference;
 import org.opensearch.core.xcontent.MediaTypeRegistry;
 import org.opensearch.index.fieldvisitor.CustomFieldsVisitor;
+import org.opensearch.index.fieldvisitor.FieldsVisitor;
 import org.opensearch.index.mapper.MapperService.MergeReason;
 import org.opensearch.test.OpenSearchSingleNodeTestCase;
 
@@ -200,4 +202,36 @@ public void testBytesAndNumericRepresentation() throws Exception {
         reader.close();
         writer.close();
     }
+
+    public void testFieldsVisitorValidateIncludesExcludes() throws Exception {
+        Set<String> fieldNames = Sets.newHashSet(
+            "field1",
+            "field2",
+            "field3",
+            "field4",
+            "field5",
+            "field6",
+            "field7",
+            "field8",
+            "field9",
+            "field10",
+            "field11"
+        );
+        String[] includes = { "field1", "field2", "field3" };
+        String[] excludes = { "field7", "field8" };
+
+        CustomFieldsVisitor fieldsVisitor = new CustomFieldsVisitor(fieldNames, false, includes, excludes);
+
+        assertArrayEquals(fieldsVisitor.includes(), includes);
+        assertArrayEquals(fieldsVisitor.excludes(), excludes);
+
+        FieldsVisitor fieldsVisitor1 = new FieldsVisitor(false, includes, excludes);
+        assertArrayEquals(fieldsVisitor1.includes(), includes);
+        assertArrayEquals(fieldsVisitor1.excludes(), excludes);
+
+        FieldsVisitor fieldsVisitor2 = new FieldsVisitor(false);
+        assertArrayEquals(fieldsVisitor2.includes(), Strings.EMPTY_ARRAY);
+        assertArrayEquals(fieldsVisitor2.excludes(), Strings.EMPTY_ARRAY);
+
+    }
 }

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

@@ -32,8 +32,22 @@
 
 package org.opensearch.search.fetch;
 
+import org.opensearch.index.fieldvisitor.CustomFieldsVisitor;
+import org.opensearch.index.fieldvisitor.FieldsVisitor;
+import org.opensearch.search.fetch.subphase.FetchSourceContext;
+import org.opensearch.search.internal.SearchContext;
 import org.opensearch.test.OpenSearchTestCase;
 
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
 public class FetchPhaseTests extends OpenSearchTestCase {
     public void testSequentialDocs() {
         FetchPhase.DocIdToIndex[] docs = new FetchPhase.DocIdToIndex[10];
@@ -52,4 +66,54 @@ public void testSequentialDocs() {
         }
         assertFalse(FetchPhase.hasSequentialDocs(docs));
     }
+
+    public void testFieldsVisitorsInFetchPhase() {
+
+        FetchPhase fetchPhase = new FetchPhase(new ArrayList<>());
+        SearchContext mockSearchContext = mock(SearchContext.class);
+        when(mockSearchContext.docIdsToLoadSize()).thenReturn(1);
+        when(mockSearchContext.docIdsToLoad()).thenReturn(new int[] { 1 });
+        String[] includes = new String[] { "field1", "field2" };
+        String[] excludes = new String[] { "field7", "field8" };
+
+        FetchSourceContext mockFetchSourceContext = new FetchSourceContext(true, includes, excludes);
+        when(mockSearchContext.hasFetchSourceContext()).thenReturn(true);
+        when(mockSearchContext.fetchSourceContext()).thenReturn(mockFetchSourceContext);
+
+        // Case 1
+        // if storedFieldsContext is null
+        FieldsVisitor fieldsVisitor = fetchPhase.createStoredFieldsVisitor(mockSearchContext, null);
+        assertArrayEquals(fieldsVisitor.excludes(), excludes);
+        assertArrayEquals(fieldsVisitor.includes(), includes);
+
+        // Case 2
+        // if storedFieldsContext is not null
+        StoredFieldsContext storedFieldsContext = mock(StoredFieldsContext.class);
+        when(mockSearchContext.storedFieldsContext()).thenReturn(storedFieldsContext);
+
+        fieldsVisitor = fetchPhase.createStoredFieldsVisitor(mockSearchContext, null);
+        assertNull(fieldsVisitor);
+
+        // Case 3
+        // if storedFieldsContext is true but fieldNames are empty
+        when(storedFieldsContext.fetchFields()).thenReturn(true);
+        when(storedFieldsContext.fieldNames()).thenReturn(List.of());
+        fieldsVisitor = fetchPhase.createStoredFieldsVisitor(mockSearchContext, Collections.emptyMap());
+        assertArrayEquals(fieldsVisitor.excludes(), excludes);
+        assertArrayEquals(fieldsVisitor.includes(), includes);
+
+        // Case 4
+        // if storedToRequested Fields is not empty
+        // creates an instance of CustomFieldsVisitor
+        Map<String, Set<String>> storedToRequestedFields = new HashMap<>();
+        storedToRequestedFields.put("test_field_key", Set.of("test_field_value"));
+
+        fieldsVisitor = fetchPhase.createStoredFieldsVisitor(mockSearchContext, storedToRequestedFields);
+
+        assertTrue(fieldsVisitor instanceof CustomFieldsVisitor);
+        assertArrayEquals(fieldsVisitor.excludes(), excludes);
+        assertArrayEquals(fieldsVisitor.includes(), includes);
+
+    }
+
 }