Some queries for runtime fields

server/src/main/java/org/elasticsearch/index/fielddata/SearchLookupAware.java

@@ -21,6 +21,8 @@
 
 import org.elasticsearch.search.lookup.SearchLookup;
 
+import java.io.IOException;
+
 //TODO this is a temporary interface only to avoid changing signature of MappedFieldType#fielddataBuilder
 public interface SearchLookupAware {
 

server/src/main/java/org/elasticsearch/search/lookup/DocLookup.java

@@ -29,10 +29,12 @@ public class DocLookup {
 
     private final MapperService mapperService;
     private final Function<MappedFieldType, IndexFieldData<?>> fieldDataLookup;
+    private final SearchLookup searchLookup;
 
-    public DocLookup(MapperService mapperService, Function<MappedFieldType, IndexFieldData<?>> fieldDataLookup) {
+    public DocLookup(MapperService mapperService, Function<MappedFieldType, IndexFieldData<?>> fieldDataLookup, SearchLookup searchLookup) {
         this.mapperService = mapperService;
         this.fieldDataLookup = fieldDataLookup;
+        this.searchLookup = searchLookup;
     }
 
     public MapperService mapperService() {
@@ -44,6 +46,6 @@ public IndexFieldData<?> getForField(MappedFieldType fieldType) {
     }
 
     public LeafDocLookup getLeafDocLookup(LeafReaderContext context) {
-        return new LeafDocLookup(mapperService, fieldDataLookup, context);
+        return new LeafDocLookup(mapperService, fieldDataLookup, searchLookup, context);
     }
 }

server/src/main/java/org/elasticsearch/search/lookup/LeafDocLookup.java

@@ -21,7 +21,9 @@
 import org.apache.lucene.index.LeafReaderContext;
 import org.elasticsearch.ExceptionsHelper;
 import org.elasticsearch.index.fielddata.IndexFieldData;
+import org.elasticsearch.index.fielddata.LeafFieldData;
 import org.elasticsearch.index.fielddata.ScriptDocValues;
+import org.elasticsearch.index.fielddata.SearchLookupAware;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.mapper.MapperService;
 
@@ -41,14 +43,20 @@ public class LeafDocLookup implements Map<String, ScriptDocValues<?>> {
     private final MapperService mapperService;
     private final Function<MappedFieldType, IndexFieldData<?>> fieldDataLookup;
 
+    private final SearchLookup searchLookup;
     private final LeafReaderContext reader;
 
     private int docId = -1;
 
-    LeafDocLookup(MapperService mapperService, Function<MappedFieldType, IndexFieldData<?>> fieldDataLookup,
-                  LeafReaderContext reader) {
+    LeafDocLookup(
+        MapperService mapperService,
+        Function<MappedFieldType, IndexFieldData<?>> fieldDataLookup,
+        SearchLookup searchLookup,
+        LeafReaderContext reader
+    ) {
         this.mapperService = mapperService;
         this.fieldDataLookup = fieldDataLookup;
+        this.searchLookup = searchLookup;
         this.reader = reader;
     }
 
@@ -75,7 +83,12 @@ public ScriptDocValues<?> get(Object key) {
             scriptValues = AccessController.doPrivileged(new PrivilegedAction<ScriptDocValues<?>>() {
                 @Override
                 public ScriptDocValues<?> run() {
-                    return fieldDataLookup.apply(fieldType).load(reader).getScriptValues();
+                    // TODO should this go through QueryShardContext?
+                    IndexFieldData<?> ifd = fieldDataLookup.apply(fieldType);
+                    if (ifd instanceof SearchLookupAware) {
+                        ((SearchLookupAware) ifd).setSearchLookup(searchLookup);
+                    }
+                    return ifd.load(reader).getScriptValues();
                 }
             });
             localCacheFieldData.put(fieldName, scriptValues);

server/src/main/java/org/elasticsearch/search/lookup/SearchLookup.java

@@ -35,7 +35,7 @@ public class SearchLookup {
     final FieldsLookup fieldsLookup;
 
     public SearchLookup(MapperService mapperService, Function<MappedFieldType, IndexFieldData<?>> fieldDataLookup) {
-        docMap = new DocLookup(mapperService, fieldDataLookup);
+        docMap = new DocLookup(mapperService, fieldDataLookup, this);
         sourceLookup = new SourceLookup();
         fieldsLookup = new FieldsLookup(mapperService);
     }

server/src/test/java/org/elasticsearch/search/lookup/LeafDocLookupTests.java

@@ -53,7 +53,7 @@ public void setUp() throws Exception {
 
         docLookup = new LeafDocLookup(mapperService,
             ignored -> fieldData,
-            null);
+            null, null);
     }
 
     public void testBasicLookup() {

x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/AbstractRuntimeValues.java

@@ -0,0 +1,175 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+package org.elasticsearch.xpack.runtimefields;
+
+import org.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.index.SortedNumericDocValues;
+import org.apache.lucene.search.ConstantScoreScorer;
+import org.apache.lucene.search.ConstantScoreWeight;
+import org.apache.lucene.search.DocIdSetIterator;
+import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.Query;
+import org.apache.lucene.search.ScoreMode;
+import org.apache.lucene.search.Scorer;
+import org.apache.lucene.search.TwoPhaseIterator;
+import org.apache.lucene.search.Weight;
+
+import java.io.IOException;
+import java.util.Objects;
+import java.util.function.IntConsumer;
+
+/**
+ * Abstract base for implementing doc values and queries against values
+ * calculated at runtime. The tricky thing about this is that we'd like to
+ * calculate the values as few times as possible in case the calculation is
+ * expensive, <strong>but</strong> some of the APIs that we rely on to
+ * calculate the values like {@link SortedNumericDocValues#advanceExact(int)}
+ * are "forwards only".
+ * <p>
+ * We solve this in the same way that big cities handle public transportation:
+ * with a bus! In our case, the bus is subclasses of {@link SharedValues}.
+ * Queries and doc values are implemented calling {@link #unstarted()} to get
+ * the {@linkplain SharedValues} that has yet to start iterating. That way
+ * many queries can share the same underlying {@linkplain SharedValues}
+ * instance, only calculating the values for a document once. If other code
+ * needs to iterate the values after the first iteration has started then
+ * it'll get a new {@linkplain SharedValues} from {@linkplain #unstarted},
+ * this "leaving on a different bus".
+ *
+ * @param <SV> the subtype of {@link SharedValues} needed by the subclass
+ */
+public abstract class AbstractRuntimeValues<SV extends AbstractRuntimeValues<SV>.SharedValues> {
+    private SV unstarted;
+
+    protected final SV unstarted() {
+        if (unstarted == null) {
+            unstarted = newSharedValues();
+        }
+        return unstarted;
+    }
+
+    protected abstract SV newSharedValues();
+
+    protected abstract class SharedValues {
+        protected int count;
+        private boolean sort;
+
+        private int lastDocBase = -1;
+        private IntConsumer lastLeafCursor;
+        private int docId = -1;
+        private int maxDoc;
+
+        protected final IntConsumer leafCursor(LeafReaderContext ctx) throws IOException {
+            if (lastDocBase != ctx.docBase) {
+                if (lastDocBase == -1) {
+                    // Now that we're started future iterations can't share these values.
+                    unstarted = null;
+                }
+                lastDocBase = ctx.docBase;
+                IntConsumer leafLoader = newLeafLoader(ctx);
+                docId = -1;
+                maxDoc = ctx.reader().maxDoc();
+                lastLeafCursor = new IntConsumer() {
+                    @Override
+                    public void accept(int targetDocId) {
+                        if (docId == targetDocId) {
+                            return;
+                        }
+                        docId = targetDocId;
+                        count = 0;
+                        leafLoader.accept(targetDocId);
+                        if (sort) {
+                            sort();
+                        }
+                    }
+                };
+            }
+            return lastLeafCursor;
+        }
+
+        protected final void alwaysSortResults() {
+            sort = true;
+        }
+
+        protected final int docId() {
+            return docId;
+        }
+
+        protected final int maxDoc() {
+            return maxDoc;
+        }
+
+        protected abstract IntConsumer newLeafLoader(LeafReaderContext ctx) throws IOException;
+
+        protected abstract void sort();
+
+        protected abstract class AbstractRuntimeQuery extends Query {
+            protected final String fieldName;
+
+            protected AbstractRuntimeQuery(String fieldName) {
+                this.fieldName = fieldName;
+            }
+
+            @Override
+            public final Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException {
+                return new ConstantScoreWeight(this, boost) {
+                    @Override
+                    public boolean isCacheable(LeafReaderContext ctx) {
+                        return false; // scripts aren't really cacheable at this point
+                    }
+
+                    @Override
+                    public Scorer scorer(LeafReaderContext ctx) throws IOException {
+                        IntConsumer leafCursor = leafCursor(ctx);
+                        DocIdSetIterator approximation = DocIdSetIterator.all(ctx.reader().maxDoc());
+                        TwoPhaseIterator twoPhase = new TwoPhaseIterator(approximation) {
+                            @Override
+                            public boolean matches() throws IOException {
+                                leafCursor.accept(approximation.docID());
+                                return AbstractRuntimeQuery.this.matches();
+                            }
+
+                            @Override
+                            public float matchCost() {
+                                // TODO we don't have a good way of estimating the complexity of the script so we just go with 9000
+                                return approximation().cost() * 9000f;
+                            }
+                        };
+                        return new ConstantScoreScorer(this, score(), scoreMode, twoPhase);
+                    }
+                };
+            }
+
+            protected abstract boolean matches();
+
+            @Override
+            public final String toString(String field) {
+                if (fieldName.contentEquals(field)) {
+                    return bareToString();
+                }
+                return fieldName + ":" + bareToString();
+            }
+
+            protected abstract String bareToString();
+
+            @Override
+            public int hashCode() {
+                return Objects.hash(fieldName);
+            }
+
+            @Override
+            public boolean equals(Object obj) {
+                if (obj == null || getClass() != obj.getClass()) {
+                    return false;
+                }
+                @SuppressWarnings("unchecked")
+                AbstractRuntimeQuery other = (AbstractRuntimeQuery) obj;
+                return fieldName.equals(other.fieldName);
+            }
+        }
+    }
+}

x-pack/plugin/runtime-fields/src/main/java/org/elasticsearch/xpack/runtimefields/AbstractScriptFieldScript.java

@@ -23,18 +23,24 @@ public abstract class AbstractScriptFieldScript {
     private final LeafSearchLookup leafSearchLookup;
 
     public AbstractScriptFieldScript(Map<String, Object> params, SearchLookup searchLookup, LeafReaderContext ctx) {
-        this.leafSearchLookup = searchLookup.getLeafSearchLookup(ctx);
+        leafSearchLookup = searchLookup.getLeafSearchLookup(ctx);
         // TODO how do other scripts get stored fields exposed? Through asMap? I don't see any getters for them.
         this.params = params;
     }
 
     /**
      * Set the document to run the script against.
      */
-    public final void setDocument(int docId) {
-        this.leafSearchLookup.setDocument(docId);
+    public final void setDocId(int docId) {
+        leafSearchLookup.setDocument(docId);
+        onSetDocId(docId);
     }
 
+    /**
+     * Optional hook for the script to take extra actions when moving to a document.
+     */
+    protected void onSetDocId(int docId) {}
+
     /**
      * Expose the {@code params} of the script to the script itself.
      */