use of org.opensearch.search.lookup.SourceLookup in project OpenSearch by opensearch-project.
the class ScriptedMetricAggContextsTests method testReturnSource.
public void testReturnSource() throws IOException {
ScriptedMetricAggContexts.MapScript.Factory factory = getEngine().compile("test", "state._source = params._source", ScriptedMetricAggContexts.MapScript.CONTEXT, Collections.emptyMap());
Map<String, Object> params = new HashMap<>();
Map<String, Object> state = new HashMap<>();
MemoryIndex index = new MemoryIndex();
// we don't need a real index, just need to construct a LeafReaderContext which cannot be mocked
LeafReaderContext leafReaderContext = index.createSearcher().getIndexReader().leaves().get(0);
SearchLookup lookup = mock(SearchLookup.class);
LeafSearchLookup leafLookup = mock(LeafSearchLookup.class);
when(lookup.getLeafSearchLookup(leafReaderContext)).thenReturn(leafLookup);
SourceLookup sourceLookup = mock(SourceLookup.class);
when(leafLookup.asMap()).thenReturn(Collections.singletonMap("_source", sourceLookup));
when(sourceLookup.loadSourceIfNeeded()).thenReturn(Collections.singletonMap("test", 1));
ScriptedMetricAggContexts.MapScript.LeafFactory leafFactory = factory.newFactory(params, state, lookup);
ScriptedMetricAggContexts.MapScript script = leafFactory.newInstance(leafReaderContext);
script.execute();
assertTrue(state.containsKey("_source"));
assertTrue(state.get("_source") instanceof Map && ((Map) state.get("_source")).containsKey("test"));
assertEquals(1, ((Map) state.get("_source")).get("test"));
}
use of org.opensearch.search.lookup.SourceLookup in project OpenSearch by opensearch-project.
the class ScriptedMetricAggContextsTests method testMapSourceAccess.
public void testMapSourceAccess() throws IOException {
ScriptedMetricAggContexts.MapScript.Factory factory = getEngine().compile("test", "state.testField = params._source.three", ScriptedMetricAggContexts.MapScript.CONTEXT, Collections.emptyMap());
Map<String, Object> params = new HashMap<>();
Map<String, Object> state = new HashMap<>();
MemoryIndex index = new MemoryIndex();
// we don't need a real index, just need to construct a LeafReaderContext which cannot be mocked
LeafReaderContext leafReaderContext = index.createSearcher().getIndexReader().leaves().get(0);
SearchLookup lookup = mock(SearchLookup.class);
LeafSearchLookup leafLookup = mock(LeafSearchLookup.class);
when(lookup.getLeafSearchLookup(leafReaderContext)).thenReturn(leafLookup);
SourceLookup sourceLookup = mock(SourceLookup.class);
when(leafLookup.asMap()).thenReturn(Collections.singletonMap("_source", sourceLookup));
when(sourceLookup.loadSourceIfNeeded()).thenReturn(Collections.singletonMap("three", 3));
ScriptedMetricAggContexts.MapScript.LeafFactory leafFactory = factory.newFactory(params, state, lookup);
ScriptedMetricAggContexts.MapScript script = leafFactory.newInstance(leafReaderContext);
script.execute();
assertTrue(state.containsKey("testField"));
assertEquals(3, state.get("testField"));
}
use of org.opensearch.search.lookup.SourceLookup in project OpenSearch by opensearch-project.
the class FetchPhase method prepareNestedHitContext.
/**
* Resets the provided {@link HitContext} with information on the current
* nested document. This includes the following:
* - Adding an initial {@link SearchHit} instance.
* - Loading the document source, filtering it based on the nested document ID, then
* setting it on {@link SourceLookup}. This allows fetch subphases that use the hit
* context to access the preloaded source.
*/
@SuppressWarnings("unchecked")
private HitContext prepareNestedHitContext(SearchContext context, int nestedTopDocId, int rootDocId, Map<String, Set<String>> storedToRequestedFields, LeafReaderContext subReaderContext, CheckedBiConsumer<Integer, FieldsVisitor, IOException> storedFieldReader) throws IOException {
// Also if highlighting is requested on nested documents we need to fetch the _source from the root document,
// otherwise highlighting will attempt to fetch the _source from the nested doc, which will fail,
// because the entire _source is only stored with the root document.
boolean needSource = sourceRequired(context) || context.highlight() != null;
Uid rootId;
Map<String, Object> rootSourceAsMap = null;
XContentType rootSourceContentType = null;
int nestedDocId = nestedTopDocId - subReaderContext.docBase;
if (context instanceof InnerHitsContext.InnerHitSubContext) {
InnerHitsContext.InnerHitSubContext innerHitsContext = (InnerHitsContext.InnerHitSubContext) context;
rootId = innerHitsContext.getRootId();
if (needSource) {
SourceLookup rootLookup = innerHitsContext.getRootLookup();
rootSourceAsMap = rootLookup.loadSourceIfNeeded();
rootSourceContentType = rootLookup.sourceContentType();
}
} else {
FieldsVisitor rootFieldsVisitor = new FieldsVisitor(needSource);
loadStoredFields(context.mapperService(), storedFieldReader, rootFieldsVisitor, rootDocId);
rootFieldsVisitor.postProcess(context.mapperService());
rootId = rootFieldsVisitor.uid();
if (needSource) {
if (rootFieldsVisitor.source() != null) {
Tuple<XContentType, Map<String, Object>> tuple = XContentHelper.convertToMap(rootFieldsVisitor.source(), false);
rootSourceAsMap = tuple.v2();
rootSourceContentType = tuple.v1();
} else {
rootSourceAsMap = Collections.emptyMap();
}
}
}
Map<String, DocumentField> docFields = emptyMap();
Map<String, DocumentField> metaFields = emptyMap();
if (context.hasStoredFields() && !context.storedFieldsContext().fieldNames().isEmpty()) {
FieldsVisitor nestedFieldsVisitor = new CustomFieldsVisitor(storedToRequestedFields.keySet(), false);
loadStoredFields(context.mapperService(), storedFieldReader, nestedFieldsVisitor, nestedDocId);
if (nestedFieldsVisitor.fields().isEmpty() == false) {
docFields = new HashMap<>();
metaFields = new HashMap<>();
fillDocAndMetaFields(context, nestedFieldsVisitor, storedToRequestedFields, docFields, metaFields);
}
}
DocumentMapper documentMapper = context.mapperService().documentMapper();
ObjectMapper nestedObjectMapper = documentMapper.findNestedObjectMapper(nestedDocId, context, subReaderContext);
assert nestedObjectMapper != null;
SearchHit.NestedIdentity nestedIdentity = getInternalNestedIdentity(context, nestedDocId, subReaderContext, context.mapperService(), nestedObjectMapper);
SearchHit hit = new SearchHit(nestedTopDocId, rootId.id(), nestedIdentity, docFields, metaFields);
// Use a clean, fresh SourceLookup
HitContext hitContext = new HitContext(hit, subReaderContext, nestedDocId, new SourceLookup());
if (rootSourceAsMap != null && rootSourceAsMap.isEmpty() == false) {
// Isolate the nested json array object that matches with nested hit and wrap it back into the same json
// structure with the nested json array object being the actual content. The latter is important, so that
// features like source filtering and highlighting work consistent regardless of whether the field points
// to a json object array for consistency reasons on how we refer to fields
Map<String, Object> nestedSourceAsMap = new HashMap<>();
Map<String, Object> current = nestedSourceAsMap;
for (SearchHit.NestedIdentity nested = nestedIdentity; nested != null; nested = nested.getChild()) {
String nestedPath = nested.getField().string();
current.put(nestedPath, new HashMap<>());
Object extractedValue = XContentMapValues.extractValue(nestedPath, rootSourceAsMap);
List<?> nestedParsedSource;
if (extractedValue instanceof List) {
// nested field has an array value in the _source
nestedParsedSource = (List<?>) extractedValue;
} else if (extractedValue instanceof Map) {
// nested field has an object value in the _source. This just means the nested field has just one inner object,
// which is valid, but uncommon.
nestedParsedSource = Collections.singletonList(extractedValue);
} else {
throw new IllegalStateException("extracted source isn't an object or an array");
}
if ((nestedParsedSource.get(0) instanceof Map) == false && nestedObjectMapper.parentObjectMapperAreNested(context.mapperService()) == false) {
// This is why only the first element of nestedParsedSource needs to be checked.
throw new IllegalArgumentException("Cannot execute inner hits. One or more parent object fields of nested field [" + nestedObjectMapper.name() + "] are not nested. All parent fields need to be nested fields too");
}
rootSourceAsMap = (Map<String, Object>) nestedParsedSource.get(nested.getOffset());
if (nested.getChild() == null) {
current.put(nestedPath, rootSourceAsMap);
} else {
Map<String, Object> next = new HashMap<>();
current.put(nestedPath, next);
current = next;
}
}
hitContext.sourceLookup().setSource(nestedSourceAsMap);
hitContext.sourceLookup().setSourceContentType(rootSourceContentType);
}
return hitContext;
}
use of org.opensearch.search.lookup.SourceLookup in project OpenSearch by opensearch-project.
the class FetchSourcePhase method hitExecute.
private void hitExecute(String index, FetchSourceContext fetchSourceContext, HitContext hitContext) {
final boolean nestedHit = hitContext.hit().getNestedIdentity() != null;
SourceLookup source = hitContext.sourceLookup();
// If source is disabled in the mapping, then attempt to return early.
if (source.source() == null && source.internalSourceRef() == null) {
if (containsFilters(fetchSourceContext)) {
throw new IllegalArgumentException("unable to fetch fields from _source field: _source is disabled in the mappings " + "for index [" + index + "]");
}
return;
}
// If this is a parent document and there are no source filters, then add the source as-is.
if (nestedHit == false && containsFilters(fetchSourceContext) == false) {
hitContext.hit().sourceRef(source.internalSourceRef());
return;
}
// Otherwise, filter the source and add it to the hit.
Object value = source.filter(fetchSourceContext);
if (nestedHit) {
value = getNestedSource((Map<String, Object>) value, hitContext);
}
try {
final int initialCapacity = nestedHit ? 1024 : Math.min(1024, source.internalSourceRef().length());
BytesStreamOutput streamOutput = new BytesStreamOutput(initialCapacity);
XContentBuilder builder = new XContentBuilder(source.sourceContentType().xContent(), streamOutput);
if (value != null) {
builder.value(value);
} else {
// This happens if the source filtering could not find the specified in the _source.
// Just doing `builder.value(null)` is valid, but the xcontent validation can't detect what format
// it is. In certain cases, for example response serialization we fail if no xcontent type can't be
// detected. So instead we just return an empty top level object. Also this is in inline with what was
// being return in this situation in 5.x and earlier.
builder.startObject();
builder.endObject();
}
hitContext.hit().sourceRef(BytesReference.bytes(builder));
} catch (IOException e) {
throw new OpenSearchException("Error filtering source", e);
}
}
use of org.opensearch.search.lookup.SourceLookup in project OpenSearch by opensearch-project.
the class FieldFetcherTests method fetchFields.
private static Map<String, DocumentField> fetchFields(MapperService mapperService, XContentBuilder source, List<FieldAndFormat> fields) throws IOException {
SourceLookup sourceLookup = new SourceLookup();
sourceLookup.setSource(BytesReference.bytes(source));
FieldFetcher fieldFetcher = FieldFetcher.create(createQueryShardContext(mapperService), null, fields);
return fieldFetcher.fetch(sourceLookup, org.opensearch.common.collect.Set.of());
}
Aggregations