From eaec888d5278bd1033a092f84db600ddca3b5130 Mon Sep 17 00:00:00 2001 From: Ganesh Ramadurai Date: Tue, 21 Jan 2025 15:50:21 -0800 Subject: [PATCH] Propagate includes and excludes from fetchSourceContext to FieldsVisitor Signed-off-by: Ganesh Ramadurai --- CHANGELOG.md | 1 + .../fieldvisitor/CustomFieldsVisitor.java | 5 ++ .../index/fieldvisitor/FieldsVisitor.java | 31 ++++++++- .../opensearch/search/fetch/FetchPhase.java | 21 ++++-- .../mapper/StoredNumericValuesTests.java | 34 ++++++++++ .../search/fetch/FetchPhaseTests.java | 64 +++++++++++++++++++ 6 files changed, 151 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 507d9906ca5e1..63ccb634a137b 100644 --- a/CHANGELOG.md +++ b/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)) diff --git a/server/src/main/java/org/opensearch/index/fieldvisitor/CustomFieldsVisitor.java b/server/src/main/java/org/opensearch/index/fieldvisitor/CustomFieldsVisitor.java index df4d398b2b181..8e6799f6bf74c 100644 --- a/server/src/main/java/org/opensearch/index/fieldvisitor/CustomFieldsVisitor.java +++ b/server/src/main/java/org/opensearch/index/fieldvisitor/CustomFieldsVisitor.java @@ -52,6 +52,11 @@ public CustomFieldsVisitor(Set fields, boolean loadSource) { this.fields = fields; } + public CustomFieldsVisitor(Set 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) { diff --git a/server/src/main/java/org/opensearch/index/fieldvisitor/FieldsVisitor.java b/server/src/main/java/org/opensearch/index/fieldvisitor/FieldsVisitor.java index 91ca07d753cc6..92328745f020e 100644 --- a/server/src/main/java/org/opensearch/index/fieldvisitor/FieldsVisitor.java +++ b/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 requiredFields; + private final String[] sourceIncludes; + private final String[] sourceExcludes; protected BytesReference source; protected String id; protected Map> 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; } diff --git a/server/src/main/java/org/opensearch/search/fetch/FetchPhase.java b/server/src/main/java/org/opensearch/search/fetch/FetchPhase.java index 11a1b9a97235b..df37b7dbfda98 100644 --- a/server/src/main/java/org/opensearch/search/fetch/FetchPhase.java +++ b/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> storedToRequestedFields) { + protected FieldsVisitor createStoredFieldsVisitor(SearchContext context, Map> storedToRequestedFields) { StoredFieldsContext storedFieldsContext = context.storedFieldsContext(); if (storedFieldsContext == null) { @@ -230,7 +230,11 @@ private FieldsVisitor createStoredFieldsVisitor(SearchContext context, Map 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); + + } } diff --git a/server/src/test/java/org/opensearch/search/fetch/FetchPhaseTests.java b/server/src/test/java/org/opensearch/search/fetch/FetchPhaseTests.java index a4820c6cff003..4f03dfd47c512 100644 --- a/server/src/test/java/org/opensearch/search/fetch/FetchPhaseTests.java +++ b/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> 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); + + } + }