Skip to content

Commit

Permalink
Propagate includes and excludes from fetchSourceContext to FieldsVisitor
Browse files Browse the repository at this point in the history 
Signed-off-by: Ganesh Ramadurai <[email protected]>
  • Loading branch information
@Gankris96
Ganesh Ramadurai authored and Gankris96 committed Jan 22, 2025
1 parent 92088be commit 8a7b739
Show file tree
Hide file tree
Showing 7 changed files with 177 additions and 11 deletions.
7 changes: 6 additions & 1 deletion server/src/main/java/org/opensearch/index/fieldvisitor/CustomFieldsVisitor.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,12 @@ public class CustomFieldsVisitor extends FieldsVisitor {
private final Set<String> fields;

public CustomFieldsVisitor(Set<String> fields, boolean loadSource) {
super(loadSource);
super(loadSource, null, null);
this.fields = fields;
}

public CustomFieldsVisitor(Set<String> fields, boolean loadSource, String[] includes, String[] excludes) {
super(loadSource, includes, excludes);
this.fields = fields;
}

Expand Down
34 changes: 32 additions & 2 deletions server/src/main/java/org/opensearch/index/fieldvisitor/FieldsVisitor.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -66,18 +67,31 @@ 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);
public FieldsVisitor(boolean loadSource, String[] includes, String[] excludes) {
this(loadSource, SourceFieldMapper.NAME, includes, excludes);
}

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();
}

public FieldsVisitor(boolean loadSource, String sourceFieldName) {
this.loadSource = loadSource;
this.sourceFieldName = sourceFieldName;
requiredFields = new HashSet<>();
this.sourceIncludes = Strings.EMPTY_ARRAY;
this.sourceExcludes = Strings.EMPTY_ARRAY;
reset();
}

Expand Down Expand Up @@ -162,6 +176,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;
}
Expand Down
2 changes: 1 addition & 1 deletion server/src/main/java/org/opensearch/index/get/ShardGetService.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -430,7 +430,7 @@ private GetResult innerGetLoadFromStoredFields(

private static FieldsVisitor buildFieldsVisitors(String[] fields, FetchSourceContext fetchSourceContext) {
if (fields == null || fields.length == 0) {
return fetchSourceContext.fetchSource() ? new FieldsVisitor(true) : null;
return fetchSourceContext.fetchSource() ? new FieldsVisitor(true, null, null) : null;
}

return new CustomFieldsVisitor(Sets.newHashSet(fields), fetchSourceContext.fetchSource());
Expand Down
21 changes: 17 additions & 4 deletions server/src/main/java/org/opensearch/search/fetch/FetchPhase.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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(

Check warning on line 269 in server/src/main/java/org/opensearch/search/fetch/FetchPhase.java

View check run for this annotation

Codecov / codecov/patch

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

Added line #L269 was not covered by tests
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(),

Check warning on line 276 in server/src/main/java/org/opensearch/search/fetch/FetchPhase.java

View check run for this annotation

Codecov / codecov/patch

server/src/main/java/org/opensearch/search/fetch/FetchPhase.java#L275-L276 

Added lines #L275 - L276 were not covered by tests
loadSource,
context.hasFetchSourceContext() ? context.fetchSourceContext().includes() : null,
context.hasFetchSourceContext() ? context.fetchSourceContext().excludes() : null
);
}
}
}
Expand Down Expand Up @@ -390,7 +403,7 @@ private HitContext prepareNestedHitContext(
rootSourceContentType = rootLookup.sourceContentType();
}
} else {
FieldsVisitor rootFieldsVisitor = new FieldsVisitor(needSource);
FieldsVisitor rootFieldsVisitor = new FieldsVisitor(needSource, null, null);

Check warning on line 406 in server/src/main/java/org/opensearch/search/fetch/FetchPhase.java

View check run for this annotation

Codecov / codecov/patch

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

Added line #L406 was not covered by tests
loadStoredFields(context::fieldType, storedFieldReader, rootFieldsVisitor, rootDocId);
rootFieldsVisitor.postProcess(context::fieldType);
rootId = rootFieldsVisitor.id();
Expand Down
2 changes: 1 addition & 1 deletion server/src/main/java/org/opensearch/search/lookup/SourceLookup.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,7 @@ public Map<String, Object> loadSourceIfNeeded() {
return source;
}
try {
FieldsVisitor sourceFieldVisitor = new FieldsVisitor(true);
FieldsVisitor sourceFieldVisitor = new FieldsVisitor(true, null, null);
fieldReader.accept(docId, sourceFieldVisitor);
BytesReference source = sourceFieldVisitor.source();
if (source == null) {
Expand Down
4 changes: 2 additions & 2 deletions server/src/test/java/org/opensearch/index/engine/InternalEngineTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2612,7 +2612,7 @@ class OpAndVersion {
for (int op = 0; op < opsPerThread; op++) {
try (Engine.GetResult get = engine.get(new Engine.Get(true, false, doc.id(), uidTerm), searcherFactory)) {

FieldsVisitor visitor = new FieldsVisitor(true);
FieldsVisitor visitor = new FieldsVisitor(true, null, null);
get.docIdAndVersion().reader.storedFields().document(get.docIdAndVersion().docId, visitor);
List<String> values = new ArrayList<>(Strings.commaDelimitedListToSet(visitor.source().utf8ToString()));
String removed = op % 3 == 0 && values.size() > 0 ? values.remove(0) : null;
Expand Down Expand Up @@ -2668,7 +2668,7 @@ class OpAndVersion {
}

try (Engine.GetResult get = engine.get(new Engine.Get(true, false, doc.id(), uidTerm), searcherFactory)) {
FieldsVisitor visitor = new FieldsVisitor(true);
FieldsVisitor visitor = new FieldsVisitor(true, null, null);
get.docIdAndVersion().reader.storedFields().document(get.docIdAndVersion().docId, visitor);
List<String> values = Arrays.asList(Strings.commaDelimitedListToStringArray(visitor.source().utf8ToString()));
assertThat(currentValues, equalTo(new HashSet<>(values)));
Expand Down
118 changes: 118 additions & 0 deletions server/src/test/java/org/opensearch/index/mapper/StoredNumericValuesTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -200,4 +200,122 @@ public void testBytesAndNumericRepresentation() throws Exception {
reader.close();
writer.close();
}

public void testFieldsVisitorValidateIncludesExcludes() throws Exception {
IndexWriter writer = new IndexWriter(new ByteBuffersDirectory(), new IndexWriterConfig(Lucene.STANDARD_ANALYZER));

String mapping = XContentFactory.jsonBuilder()
.startObject()
.startObject("type")
.startObject("properties")
.startObject("field1")
.field("type", "byte")
.field("store", true)
.endObject()
.startObject("field2")
.field("type", "short")
.field("store", true)
.endObject()
.startObject("field3")
.field("type", "integer")
.field("store", true)
.endObject()
.startObject("field4")
.field("type", "float")
.field("store", true)
.endObject()
.startObject("field5")
.field("type", "long")
.field("store", true)
.endObject()
.startObject("field6")
.field("type", "double")
.field("store", true)
.endObject()
.startObject("field7")
.field("type", "ip")
.field("store", true)
.endObject()
.startObject("field8")
.field("type", "ip")
.field("store", true)
.endObject()
.startObject("field9")
.field("type", "date")
.field("store", true)
.endObject()
.startObject("field10")
.field("type", "boolean")
.field("store", true)
.endObject()
.startObject("field11")
.field("type", "unsigned_long")
.field("store", true)
.endObject()
.endObject()
.endObject()
.endObject()
.toString();
MapperService mapperService = createIndex("test").mapperService();
DocumentMapper mapper = mapperService.merge("type", new CompressedXContent(mapping), MergeReason.MAPPING_UPDATE);

ParsedDocument doc = mapper.parse(
new SourceToParse(
"test",
"1",
BytesReference.bytes(
XContentFactory.jsonBuilder()
.startObject()
.field("field1", 1)
.field("field2", 1)
.field("field3", 1)
.field("field4", 1.1)
.startArray("field5")
.value(1)
.value(2)
.value(3)
.endArray()
.field("field6", 1.1)
.field("field7", "192.168.1.1")
.field("field8", "2001:db8::2:1")
.field("field9", "2016-04-05")
.field("field10", true)
.field("field11", "1")
.endObject()
),
MediaTypeRegistry.JSON
)
);

writer.addDocument(doc.rootDoc());

DirectoryReader reader = DirectoryReader.open(writer);
IndexSearcher searcher = new IndexSearcher(reader);

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);
searcher.doc(0, fieldsVisitor);
fieldsVisitor.postProcess(mapperService::fieldType);

assertEquals(fieldsVisitor.includes(), includes);
assertEquals(fieldsVisitor.excludes(), excludes);

reader.close();
writer.close();
}
}

0 comments on commit 8a7b739

Please sign in to comment.