Allow for aliases when fetching stored fields. (#31411)

jtibshirani · jtibshirani · commit b62c73a5165a · 2018-06-21T16:53:47.000-07:00
* When loading stored fields, resolve wildcard patterns against the mappings instead of the index.
* Minor simplifications to the logic in FetchPhase.
* Pull out a shared method FetchPhase#getSearchFields.
* Allow for aliases when fetching stored fields.
diff --git a/server/src/main/java/org/elasticsearch/index/fieldvisitor/CustomFieldsVisitor.java b/server/src/main/java/org/elasticsearch/index/fieldvisitor/CustomFieldsVisitor.java
@@ -19,11 +19,8 @@
 package org.elasticsearch.index.fieldvisitor;
 
 import org.apache.lucene.index.FieldInfo;
-import org.elasticsearch.common.regex.Regex;
 
 import java.io.IOException;
-import java.util.Collections;
-import java.util.List;
 import java.util.Set;
 
 /**
@@ -35,16 +32,10 @@
 public class CustomFieldsVisitor extends FieldsVisitor {
 
     private final Set<String> fields;
-    private final List<String> patterns;
 
-    public CustomFieldsVisitor(Set<String> fields, List<String> patterns, boolean loadSource) {
+    public CustomFieldsVisitor(Set<String> fields, boolean loadSource) {
         super(loadSource);
         this.fields = fields;
-        this.patterns = patterns;
-    }
-
-    public CustomFieldsVisitor(Set<String> fields, boolean loadSource) {
-        this(fields, Collections.emptyList(), loadSource);
     }
 
     @Override
@@ -55,11 +46,6 @@ public Status needsField(FieldInfo fieldInfo) throws IOException {
         if (fields.contains(fieldInfo.name)) {
             return Status.YES;
         }
-        for (String pattern : patterns) {
-            if (Regex.simpleMatch(pattern, fieldInfo.name)) {
-                return Status.YES;
-            }
-        }
         return Status.NO;
     }
 }
diff --git a/server/src/main/java/org/elasticsearch/search/fetch/FetchPhase.java b/server/src/main/java/org/elasticsearch/search/fetch/FetchPhase.java
@@ -31,7 +31,6 @@
 import org.elasticsearch.common.collect.Tuple;
 import org.elasticsearch.common.document.DocumentField;
 import org.elasticsearch.common.lucene.search.Queries;
-import org.elasticsearch.common.regex.Regex;
 import org.elasticsearch.common.text.Text;
 import org.elasticsearch.common.xcontent.XContentHelper;
 import org.elasticsearch.common.xcontent.XContentType;
@@ -55,7 +54,7 @@
 import org.elasticsearch.tasks.TaskCancelledException;
 
 import java.io.IOException;
-import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -84,8 +83,7 @@ public void preProcess(SearchContext context) {
     @Override
     public void execute(SearchContext context) {
         final FieldsVisitor fieldsVisitor;
-        Set<String> fieldNames = null;
-        List<String> fieldNamePatterns = null;
+        Map<String, Set<String>> storedToRequestedFields = new HashMap<>();
         StoredFieldsContext storedFieldsContext = context.storedFieldsContext();
 
         if (storedFieldsContext == null) {
@@ -98,39 +96,36 @@ public void execute(SearchContext context) {
             // disable stored fields entirely
             fieldsVisitor = null;
         } else {
-            for (String fieldName : context.storedFieldsContext().fieldNames()) {
-                if (fieldName.equals(SourceFieldMapper.NAME)) {
+            for (String fieldNameOrPattern : context.storedFieldsContext().fieldNames()) {
+                if (fieldNameOrPattern.equals(SourceFieldMapper.NAME)) {
                     FetchSourceContext fetchSourceContext = context.hasFetchSourceContext() ? context.fetchSourceContext()
-                            : FetchSourceContext.FETCH_SOURCE;
+                        : FetchSourceContext.FETCH_SOURCE;
                     context.fetchSourceContext(new FetchSourceContext(true, fetchSourceContext.includes(), fetchSourceContext.excludes()));
                     continue;
                 }
-                if (Regex.isSimpleMatchPattern(fieldName)) {
-                    if (fieldNamePatterns == null) {
-                        fieldNamePatterns = new ArrayList<>();
-                    }
-                    fieldNamePatterns.add(fieldName);
-                } else {
+
+                Collection<String> fieldNames = context.mapperService().simpleMatchToFullName(fieldNameOrPattern);
+                for (String fieldName : fieldNames) {
                     MappedFieldType fieldType = context.smartNameFieldType(fieldName);
                     if (fieldType == null) {
                         // Only fail if we know it is a object field, missing paths / fields shouldn't fail.
                         if (context.getObjectMapper(fieldName) != null) {
                             throw new IllegalArgumentException("field [" + fieldName + "] isn't a leaf field");
                         }
+                    } else {
+                        String storedField = fieldType.name();
+                        Set<String> requestedFields = storedToRequestedFields.computeIfAbsent(
+                            storedField, key -> new HashSet<>());
+                        requestedFields.add(fieldName);
                     }
-                    if (fieldNames == null) {
-                        fieldNames = new HashSet<>();
-                    }
-                    fieldNames.add(fieldName);
                 }
             }
             boolean loadSource = context.sourceRequested();
-            if (fieldNames == null && fieldNamePatterns == null) {
+            if (storedToRequestedFields.isEmpty()) {
                 // empty list specified, default to disable _source if no explicit indication
                 fieldsVisitor = new FieldsVisitor(loadSource);
             } else {
-                fieldsVisitor = new CustomFieldsVisitor(fieldNames == null ? Collections.emptySet() : fieldNames,
-                        fieldNamePatterns == null ? Collections.emptyList() : fieldNamePatterns, loadSource);
+                fieldsVisitor = new CustomFieldsVisitor(storedToRequestedFields.keySet(), loadSource);
             }
         }
 
@@ -149,10 +144,11 @@ public void execute(SearchContext context) {
                 final SearchHit searchHit;
                 int rootDocId = findRootDocumentIfNested(context, subReaderContext, subDocId);
                 if (rootDocId != -1) {
-                    searchHit = createNestedSearchHit(context, docId, subDocId, rootDocId, fieldNames, fieldNamePatterns,
-                            subReaderContext);
+                    searchHit = createNestedSearchHit(context, docId, subDocId, rootDocId,
+                        storedToRequestedFields, subReaderContext);
                 } else {
-                    searchHit = createSearchHit(context, fieldsVisitor, docId, subDocId, subReaderContext);
+                    searchHit = createSearchHit(context, fieldsVisitor, docId, subDocId,
+                        storedToRequestedFields, subReaderContext);
                 }
 
                 hits[index] = searchHit;
@@ -190,21 +186,18 @@ private int findRootDocumentIfNested(SearchContext context, LeafReaderContext su
         return -1;
     }
 
-    private SearchHit createSearchHit(SearchContext context, FieldsVisitor fieldsVisitor, int docId, int subDocId,
+    private SearchHit createSearchHit(SearchContext context,
+                                      FieldsVisitor fieldsVisitor,
+                                      int docId,
+                                      int subDocId,
+                                      Map<String, Set<String>> storedToRequestedFields,
                                       LeafReaderContext subReaderContext) {
         if (fieldsVisitor == null) {
             return new SearchHit(docId);
         }
-        loadStoredFields(context, subReaderContext, fieldsVisitor, subDocId);
-        fieldsVisitor.postProcess(context.mapperService());
 
-        Map<String, DocumentField> searchFields = null;
-        if (!fieldsVisitor.fields().isEmpty()) {
-            searchFields = new HashMap<>(fieldsVisitor.fields().size());
-            for (Map.Entry<String, List<Object>> entry : fieldsVisitor.fields().entrySet()) {
-                searchFields.put(entry.getKey(), new DocumentField(entry.getKey(), entry.getValue()));
-            }
-        }
+        Map<String, DocumentField> searchFields = getSearchFields(context, fieldsVisitor, subDocId,
+            storedToRequestedFields, subReaderContext);
 
         DocumentMapper documentMapper = context.mapperService().documentMapper(fieldsVisitor.uid().type());
         Text typeText;
@@ -223,9 +216,40 @@ private SearchHit createSearchHit(SearchContext context, FieldsVisitor fieldsVis
         return searchHit;
     }
 
-    private SearchHit createNestedSearchHit(SearchContext context, int nestedTopDocId, int nestedSubDocId,
-                                            int rootSubDocId, Set<String> fieldNames,
-                                            List<String> fieldNamePatterns, LeafReaderContext subReaderContext) throws IOException {
+    private Map<String, DocumentField> getSearchFields(SearchContext context,
+                                                       FieldsVisitor fieldsVisitor,
+                                                       int subDocId,
+                                                       Map<String, Set<String>> storedToRequestedFields,
+                                                       LeafReaderContext subReaderContext) {
+        loadStoredFields(context, subReaderContext, fieldsVisitor, subDocId);
+        fieldsVisitor.postProcess(context.mapperService());
+
+        if (fieldsVisitor.fields().isEmpty()) {
+            return null;
+        }
+
+        Map<String, DocumentField> searchFields = new HashMap<>(fieldsVisitor.fields().size());
+        for (Map.Entry<String, List<Object>> entry : fieldsVisitor.fields().entrySet()) {
+            String storedField = entry.getKey();
+            List<Object> storedValues = entry.getValue();
+
+            if (storedToRequestedFields.containsKey(storedField)) {
+                for (String requestedField : storedToRequestedFields.get(storedField)) {
+                    searchFields.put(requestedField, new DocumentField(requestedField, storedValues));
+                }
+            } else {
+                searchFields.put(storedField, new DocumentField(storedField, storedValues));
+            }
+        }
+        return searchFields;
+    }
+
+    private SearchHit createNestedSearchHit(SearchContext context,
+                                            int nestedTopDocId,
+                                            int nestedSubDocId,
+                                            int rootSubDocId,
+                                            Map<String, Set<String>> storedToRequestedFields,
+                                            LeafReaderContext subReaderContext) 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.
@@ -244,9 +268,13 @@ private SearchHit createNestedSearchHit(SearchContext context, int nestedTopDocI
             source = null;
         }
 
+        Map<String, DocumentField> searchFields = null;
+        if (context.hasStoredFields() && !context.storedFieldsContext().fieldNames().isEmpty()) {
+            FieldsVisitor nestedFieldsVisitor = new CustomFieldsVisitor(storedToRequestedFields.keySet(), false);
+            searchFields = getSearchFields(context, nestedFieldsVisitor, nestedSubDocId,
+                storedToRequestedFields, subReaderContext);
+        }
 
-        Map<String, DocumentField> searchFields =
-                getSearchFields(context, nestedSubDocId, fieldNames, fieldNamePatterns, subReaderContext);
         DocumentMapper documentMapper = context.mapperService().documentMapper(uid.type());
         SourceLookup sourceLookup = context.lookup().source();
         sourceLookup.setSegmentAndDocument(subReaderContext, nestedSubDocId);
@@ -307,26 +335,6 @@ private SearchHit createNestedSearchHit(SearchContext context, int nestedTopDocI
         return new SearchHit(nestedTopDocId, uid.id(), documentMapper.typeText(), nestedIdentity, searchFields);
     }
 
-    private Map<String, DocumentField> getSearchFields(SearchContext context, int nestedSubDocId, Set<String> fieldNames,
-                                                       List<String> fieldNamePatterns, LeafReaderContext subReaderContext) {
-        Map<String, DocumentField> searchFields = null;
-        if (context.hasStoredFields() && !context.storedFieldsContext().fieldNames().isEmpty()) {
-            FieldsVisitor nestedFieldsVisitor = new CustomFieldsVisitor(fieldNames == null ? Collections.emptySet() : fieldNames,
-                    fieldNamePatterns == null ? Collections.emptyList() : fieldNamePatterns, false);
-            if (nestedFieldsVisitor != null) {
-                loadStoredFields(context, subReaderContext, nestedFieldsVisitor, nestedSubDocId);
-                nestedFieldsVisitor.postProcess(context.mapperService());
-                if (!nestedFieldsVisitor.fields().isEmpty()) {
-                    searchFields = new HashMap<>(nestedFieldsVisitor.fields().size());
-                    for (Map.Entry<String, List<Object>> entry : nestedFieldsVisitor.fields().entrySet()) {
-                        searchFields.put(entry.getKey(), new DocumentField(entry.getKey(), entry.getValue()));
-                    }
-                }
-            }
-        }
-        return searchFields;
-    }
-
     private SearchHit.NestedIdentity getInternalNestedIdentity(SearchContext context, int nestedSubDocId,
                                                                LeafReaderContext subReaderContext,
                                                                MapperService mapperService,
diff --git a/server/src/test/java/org/elasticsearch/index/mapper/StoredNumericValuesTests.java b/server/src/test/java/org/elasticsearch/index/mapper/StoredNumericValuesTests.java
@@ -28,13 +28,14 @@
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.compress.CompressedXContent;
 import org.elasticsearch.common.lucene.Lucene;
+import org.elasticsearch.common.util.set.Sets;
 import org.elasticsearch.common.xcontent.XContentFactory;
 import org.elasticsearch.common.xcontent.XContentType;
 import org.elasticsearch.index.fieldvisitor.CustomFieldsVisitor;
 import org.elasticsearch.index.mapper.MapperService.MergeReason;
 import org.elasticsearch.test.ESSingleNodeTestCase;
 
-import java.util.Collections;
+import java.util.Set;
 
 import static org.hamcrest.Matchers.equalTo;
 
@@ -84,9 +85,11 @@ public void testBytesAndNumericRepresentation() throws Exception {
         DirectoryReader reader = DirectoryReader.open(writer);
         IndexSearcher searcher = new IndexSearcher(reader);
 
-        CustomFieldsVisitor fieldsVisitor = new CustomFieldsVisitor(
-                Collections.emptySet(), Collections.singletonList("field*"), false);
+        Set<String> fieldNames = Sets.newHashSet("field1", "field2", "field3", "field4", "field5",
+            "field6", "field7", "field8", "field9", "field10");
+        CustomFieldsVisitor fieldsVisitor = new CustomFieldsVisitor(fieldNames, false);
         searcher.doc(0, fieldsVisitor);
+
         fieldsVisitor.postProcess(mapperService);
         assertThat(fieldsVisitor.fields().size(), equalTo(10));
         assertThat(fieldsVisitor.fields().get("field1").size(), equalTo(1));
diff --git a/server/src/test/java/org/elasticsearch/search/fields/SearchFieldsIT.java b/server/src/test/java/org/elasticsearch/search/fields/SearchFieldsIT.java
@@ -980,6 +980,98 @@ public void testDocValueFieldsWithFieldAlias() throws Exception {
         assertThat(fetchedDate, equalTo(date));
     }
 
+
+    public void testStoredFieldsWithFieldAlias() throws Exception {
+        XContentBuilder mapping = XContentFactory.jsonBuilder()
+            .startObject()
+                .startObject("type")
+                    .startObject("properties")
+                        .startObject("field1")
+                            .field("type", "text")
+                            .field("store", true)
+                        .endObject()
+                        .startObject("field2")
+                            .field("type", "text")
+                            .field("store", false)
+                        .endObject()
+                        .startObject("field1-alias")
+                            .field("type", "alias")
+                            .field("path", "field1")
+                        .endObject()
+                        .startObject("field2-alias")
+                            .field("type", "alias")
+                            .field("path", "field2")
+                        .endObject()
+                    .endObject()
+            .endObject()
+        .endObject();
+        assertAcked(prepareCreate("test").addMapping("type", mapping));
+
+        index("test", "type", "1", "field1", "value1", "field2", "value2");
+        refresh("test");
+
+        SearchResponse searchResponse = client().prepareSearch()
+            .setQuery(matchAllQuery())
+            .addStoredField("field1-alias")
+            .addStoredField("field2-alias")
+            .get();
+        assertHitCount(searchResponse, 1L);
+
+        SearchHit hit = searchResponse.getHits().getAt(0);
+        assertEquals(1, hit.getFields().size());
+        assertTrue(hit.getFields().containsKey("field1-alias"));
+
+        DocumentField field = hit.getFields().get("field1-alias");
+        assertThat(field.getValue().toString(), equalTo("value1"));
+    }
+
+    public void testWildcardStoredFieldsWithFieldAlias() throws Exception {
+        XContentBuilder mapping = XContentFactory.jsonBuilder()
+            .startObject()
+                .startObject("type")
+                    .startObject("properties")
+                        .startObject("field1")
+                            .field("type", "text")
+                            .field("store", true)
+                        .endObject()
+                        .startObject("field2")
+                            .field("type", "text")
+                            .field("store", false)
+                        .endObject()
+                        .startObject("field1-alias")
+                            .field("type", "alias")
+                            .field("path", "field1")
+                        .endObject()
+                        .startObject("field2-alias")
+                            .field("type", "alias")
+                            .field("path", "field2")
+                        .endObject()
+                    .endObject()
+            .endObject()
+        .endObject();
+        assertAcked(prepareCreate("test").addMapping("type", mapping));
+
+        index("test", "type", "1", "field1", "value1", "field2", "value2");
+        refresh("test");
+
+        SearchResponse searchResponse = client().prepareSearch()
+            .setQuery(matchAllQuery())
+            .addStoredField("field*")
+            .get();
+        assertHitCount(searchResponse, 1L);
+
+        SearchHit hit = searchResponse.getHits().getAt(0);
+        assertEquals(2, hit.getFields().size());
+        assertTrue(hit.getFields().containsKey("field1"));
+        assertTrue(hit.getFields().containsKey("field1-alias"));
+
+        DocumentField field = hit.getFields().get("field1");
+        assertThat(field.getValue().toString(), equalTo("value1"));
+
+        DocumentField fieldAlias = hit.getFields().get("field1-alias");
+        assertThat(fieldAlias.getValue().toString(), equalTo("value1"));
+    }
+
     public void testLoadMetadata() throws Exception {
         assertAcked(prepareCreate("test"));
 
