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 <gramadur@icloud.com>
  • Loading branch information
Ganesh Ramadurai committed Jan 22, 2025
1 parent 22b9ac0 commit 1584d62
Show file tree
Hide file tree
Showing 5 changed files with 171 additions and 6 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
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,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();
}
Expand Down Expand Up @@ -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;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -432,7 +432,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
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(
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
);
}
}
}
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);
loadStoredFields(context::fieldType, storedFieldReader, rootFieldsVisitor, rootDocId);
rootFieldsVisitor.postProcess(context::fieldType);
rootId = rootFieldsVisitor.id();
Expand Down
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 1584d62

Please sign in to comment.