Script: Convert script query to a dedicated script context (#26003)

This commit converts script query to use a new FilterScript context. The
new context returns a boolean, so the error that would have previously
happened at runtime if a non boolean was returned would now happen at
script compilation. Also, the leniency of supporting returning a number
and 0 mapping to false, non-zero to true is gone, but it was never
documented. With the new context compilation will now also fail if
special variables are used at compilation time, instead of runtime, eg
ctx.

core/src/main/java/org/elasticsearch/index/query/ScriptQueryBuilder.java

@@ -34,6 +34,7 @@ import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
+import org.elasticsearch.script.FilterScript;
 import org.elasticsearch.script.Script;
 import org.elasticsearch.script.SearchScript;
 
@@ -126,25 +127,25 @@ public class ScriptQueryBuilder extends AbstractQueryBuilder<ScriptQueryBuilder>
 
     @Override
     protected Query doToQuery(QueryShardContext context) throws IOException {
-        SearchScript.Factory factory = context.getScriptService().compile(script, SearchScript.CONTEXT);
-        SearchScript.LeafFactory searchScript = factory.newFactory(script.getParams(), context.lookup());
-        return new ScriptQuery(script, searchScript);
+        FilterScript.Factory factory = context.getScriptService().compile(script, FilterScript.CONTEXT);
+        FilterScript.LeafFactory filterScript = factory.newFactory(script.getParams(), context.lookup());
+        return new ScriptQuery(script, filterScript);
     }
 
     static class ScriptQuery extends Query {
 
         final Script script;
-        final SearchScript.LeafFactory searchScript;
+        final FilterScript.LeafFactory filterScript;
 
-        ScriptQuery(Script script, SearchScript.LeafFactory searchScript) {
+        ScriptQuery(Script script, FilterScript.LeafFactory filterScript) {
             this.script = script;
-            this.searchScript = searchScript;
+            this.filterScript = filterScript;
         }
 
         @Override
         public String toString(String field) {
             StringBuilder buffer = new StringBuilder();
-            buffer.append("ScriptFilter(");
+            buffer.append("ScriptQuery(");
             buffer.append(script);
             buffer.append(")");
             return buffer.toString();
@@ -178,23 +179,13 @@ public class ScriptQueryBuilder extends AbstractQueryBuilder<ScriptQueryBuilder>
                 @Override
                 public Scorer scorer(LeafReaderContext context) throws IOException {
                     DocIdSetIterator approximation = DocIdSetIterator.all(context.reader().maxDoc());
-                    final SearchScript leafScript = searchScript.newInstance(context);
+                    final FilterScript leafScript = filterScript.newInstance(context);
                     TwoPhaseIterator twoPhase = new TwoPhaseIterator(approximation) {
 
                         @Override
                         public boolean matches() throws IOException {
                             leafScript.setDocument(approximation.docID());
-                            Object val = leafScript.run();
-                            if (val == null) {
-                                return false;
-                            }
-                            if (val instanceof Boolean) {
-                                return (Boolean) val;
-                            }
-                            if (val instanceof Number) {
-                                return ((Number) val).longValue() != 0;
-                            }
-                            throw new IllegalArgumentException("Can't handle type [" + val + "] in script filter");
+                            return leafScript.execute();
                         }
 
                         @Override

core/src/main/java/org/elasticsearch/script/FilterScript.java

@@ -0,0 +1,80 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.elasticsearch.script;
+
+import java.io.IOException;
+import java.util.Map;
+
+import org.apache.lucene.index.LeafReaderContext;
+import org.elasticsearch.index.fielddata.ScriptDocValues;
+import org.elasticsearch.search.lookup.LeafDocLookup;
+import org.elasticsearch.search.lookup.LeafSearchLookup;
+import org.elasticsearch.search.lookup.SearchLookup;
+
+/**
+ * A script implementation of a query filter.
+ * See {@link org.elasticsearch.index.query.ScriptQueryBuilder}.
+ */
+public abstract class FilterScript {
+
+    // no parameters for execute, but constant still required...
+    public static final String[] PARAMETERS = {};
+
+    /** The generic runtime parameters for the script. */
+    private final Map<String, Object> params;
+
+    /** A leaf lookup for the bound segment this script will operate on. */
+    private final LeafSearchLookup leafLookup;
+
+    public FilterScript(Map<String, Object> params, SearchLookup lookup, LeafReaderContext leafContext) {
+        this.params = params;
+        this.leafLookup = lookup.getLeafSearchLookup(leafContext);
+    }
+
+    /** Return {@code true} if the current document matches the filter, or {@code false} otherwise. */
+    public abstract boolean execute();
+
+    /** Return the parameters for this script. */
+    public Map<String, Object> getParams() {
+        return params;
+    }
+
+    /** The doc lookup for the Lucene segment this script was created for. */
+    public final Map<String, ScriptDocValues<?>> getDoc() {
+        return leafLookup.doc();
+    }
+
+    /** Set the current document to run the script on next. */
+    public void setDocument(int docid) {
+        leafLookup.setDocument(docid);
+    }
+
+    /** A factory to construct {@link FilterScript} instances. */
+    public interface LeafFactory {
+        FilterScript newInstance(LeafReaderContext ctx) throws IOException;
+    }
+
+    /** A factory to construct stateful {@link FilterScript} factories for a specific index. */
+    public interface Factory {
+        LeafFactory newFactory(Map<String, Object> params, SearchLookup lookup);
+    }
+
+    /** The context used to compile {@link FilterScript} factories. */
+    public static final ScriptContext<Factory> CONTEXT = new ScriptContext<>("filter", Factory.class);
+}

core/src/main/java/org/elasticsearch/script/ScriptModule.java

@@ -45,6 +45,7 @@ public class ScriptModule {
             ExecutableScript.AGGS_CONTEXT,
             ExecutableScript.UPDATE_CONTEXT,
             ExecutableScript.INGEST_CONTEXT,
+            FilterScript.CONTEXT,
             SimilarityScript.CONTEXT,
             SimilarityWeightScript.CONTEXT,
             TemplateScript.CONTEXT

core/src/test/java/org/elasticsearch/index/query/ScriptQueryBuilderTests.java

@@ -52,7 +52,7 @@ public class ScriptQueryBuilderTests extends AbstractQueryTestCase<ScriptQueryBu
         assertThat(query, instanceOf(ScriptQueryBuilder.ScriptQuery.class));
         // make sure the query would not get cached
         ScriptQuery sQuery = (ScriptQuery) query;
-        ScriptQuery clone = new ScriptQuery(sQuery.script, sQuery.searchScript);
+        ScriptQuery clone = new ScriptQuery(sQuery.script, sQuery.filterScript);
         assertFalse(sQuery.equals(clone));
         assertFalse(sQuery.hashCode() == clone.hashCode());
     }

test/framework/src/main/java/org/elasticsearch/script/MockScriptEngine.java

@@ -34,6 +34,7 @@ import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.function.Function;
+import java.util.function.Predicate;
 
 import static java.util.Collections.emptyMap;
 
@@ -99,6 +100,9 @@ public class MockScriptEngine implements ScriptEngine {
                     };
                 };
             return context.factoryClazz.cast(factory);
+        } else if (context.instanceClazz.equals(FilterScript.class)) {
+            FilterScript.Factory factory = mockCompiled::createFilterScript;
+            return context.factoryClazz.cast(factory);
         } else if (context.instanceClazz.equals(SimilarityScript.class)) {
             SimilarityScript.Factory factory = mockCompiled::createSimilarityScript;
             return context.factoryClazz.cast(factory);
@@ -153,6 +157,11 @@ public class MockScriptEngine implements ScriptEngine {
             return new MockSearchScript(lookup, context, script != null ? script : ctx -> source);
         }
 
+
+        public FilterScript.LeafFactory createFilterScript(Map<String, Object> params, SearchLookup lookup) {
+            return new MockFilterScript(lookup, params, script);
+        }
+
         public SimilarityScript createSimilarityScript() {
             return new MockSimilarityScript(script != null ? script : ctx -> 42d);
         }
@@ -243,6 +252,39 @@ public class MockScriptEngine implements ScriptEngine {
         }
     }
 
+
+    public static class MockFilterScript implements FilterScript.LeafFactory {
+
+        private final Function<Map<String, Object>, Object> script;
+        private final Map<String, Object> vars;
+        private final SearchLookup lookup;
+
+        public MockFilterScript(SearchLookup lookup, Map<String, Object> vars, Function<Map<String, Object>, Object> script) {
+            this.lookup = lookup;
+            this.vars = vars;
+            this.script = script;
+        }
+
+        public FilterScript newInstance(LeafReaderContext context) throws IOException {
+            LeafSearchLookup leafLookup = lookup.getLeafSearchLookup(context);
+            Map<String, Object> ctx = new HashMap<>(leafLookup.asMap());
+            if (vars != null) {
+                ctx.putAll(vars);
+            }
+            return new FilterScript(ctx, lookup, context) {
+                @Override
+                public boolean execute() {
+                    return (boolean) script.apply(ctx);
+                }
+
+                @Override
+                public void setDocument(int doc) {
+                    leafLookup.setDocument(doc);
+                }
+            };
+        }
+    }
+
     public class MockSimilarityScript extends SimilarityScript {
 
         private final Function<Map<String, Object>, Object> script;