Scripted keyword field (#58939)

This is the first draft implementation of scripted field of type keyword. Still lots of TODOs to address, this PR targets a feature branch.

Author: javanna

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

@@ -0,0 +1,28 @@
+/*
+ * 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.index.fielddata;
+
+import org.elasticsearch.search.lookup.SearchLookup;
+
+//TODO this is a temporary interface only to avoid changing signature of MappedFieldType#fielddataBuilder
+public interface SearchLookupAware {
+
+    void setSearchLookup(SearchLookup searchLookup);
+}

server/src/main/java/org/elasticsearch/index/query/QueryShardContext.java

@@ -43,6 +43,7 @@
 import org.elasticsearch.index.analysis.IndexAnalyzers;
 import org.elasticsearch.index.cache.bitset.BitsetFilterCache;
 import org.elasticsearch.index.fielddata.IndexFieldData;
+import org.elasticsearch.index.fielddata.SearchLookupAware;
 import org.elasticsearch.index.mapper.ContentPath;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.mapper.Mapper;
@@ -208,7 +209,13 @@ public boolean allowExpensiveQueries() {
 
     @SuppressWarnings("unchecked")
     public <IFD extends IndexFieldData<?>> IFD getForField(MappedFieldType fieldType) {
-        return (IFD) indexFieldDataService.apply(fieldType, fullyQualifiedIndex.getName());
+        IFD indexFieldData = (IFD) indexFieldDataService.apply(fieldType, fullyQualifiedIndex.getName());
+        //TODO this is a temporary hack to inject search lookup to the scripted fielddata
+        // implementations without changing MappedFieldType#fielddataBuilder signature, as that would cause daily merge conflicts
+        if (indexFieldData instanceof SearchLookupAware) {
+            ((SearchLookupAware) indexFieldData).setSearchLookup(lookup());
+        }
+        return indexFieldData;
     }
 
     public void addNamedQuery(String name, Query query) {

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

@@ -9,9 +9,8 @@
 import org.apache.lucene.index.LeafReaderContext;
 import org.elasticsearch.index.fielddata.ScriptDocValues;
 import org.elasticsearch.script.AggregationScript;
-import org.elasticsearch.search.lookup.DocLookup;
-import org.elasticsearch.search.lookup.LeafDocLookup;
-import org.elasticsearch.search.lookup.SourceLookup;
+import org.elasticsearch.search.lookup.LeafSearchLookup;
+import org.elasticsearch.search.lookup.SearchLookup;
 
 import java.util.Map;
 
@@ -21,23 +20,19 @@
  */
 public abstract class AbstractScriptFieldScript {
     private final Map<String, Object> params;
-    private final LeafReaderContext ctx;
-    private final SourceLookup source;
-    private final LeafDocLookup fieldData;
+    private final LeafSearchLookup leafSearchLookup;
 
-    public AbstractScriptFieldScript(Map<String, Object> params, SourceLookup source, DocLookup fieldData, LeafReaderContext ctx) {
+    public AbstractScriptFieldScript(Map<String, Object> params, SearchLookup searchLookup, LeafReaderContext ctx) {
+        this.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;
-        this.source = source;
-        this.fieldData = fieldData.getLeafDocLookup(ctx);
-        this.ctx = ctx;
     }
 
     /**
      * Set the document to run the script against.
      */
     public final void setDocument(int docId) {
-        source.setSegmentAndDocument(ctx, docId);
-        fieldData.setDocument(docId);
+        this.leafSearchLookup.setDocument(docId);
     }
 
     /**
@@ -51,16 +46,15 @@ public final Map<String, Object> getParams() {
      * Expose the {@code _source} to the script.
      */
     public final Map<String, Object> getSource() {
-        return source;
+        return leafSearchLookup.source();
     }
 
     /**
      * Expose field data to the script as {@code doc}.
      */
     public final Map<String, ScriptDocValues<?>> getDoc() {
-        return fieldData;
+        return leafSearchLookup.doc();
     }
 
     public abstract void execute();
-
 }

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

@@ -11,8 +11,7 @@
 import org.elasticsearch.painless.spi.WhitelistLoader;
 import org.elasticsearch.script.ScriptContext;
 import org.elasticsearch.script.ScriptFactory;
-import org.elasticsearch.search.lookup.DocLookup;
-import org.elasticsearch.search.lookup.SourceLookup;
+import org.elasticsearch.search.lookup.SearchLookup;
 
 import java.io.IOException;
 import java.util.List;
@@ -29,7 +28,7 @@ static List<Whitelist> whitelist() {
     public static final String[] PARAMETERS = {};
 
     public interface Factory extends ScriptFactory {
-        LeafFactory newFactory(Map<String, Object> params, SourceLookup source, DocLookup fieldData);
+        LeafFactory newFactory(Map<String, Object> params, SearchLookup searchLookup);
     }
 
     public interface LeafFactory {
@@ -38,14 +37,8 @@ public interface LeafFactory {
 
     private final DoubleConsumer sync;
 
-    public DoubleScriptFieldScript(
-        Map<String, Object> params,
-        SourceLookup source,
-        DocLookup fieldData,
-        LeafReaderContext ctx,
-        DoubleConsumer sync
-    ) {
-        super(params, source, fieldData, ctx);
+    public DoubleScriptFieldScript(Map<String, Object> params, SearchLookup searchLookup, LeafReaderContext ctx, DoubleConsumer sync) {
+        super(params, searchLookup, ctx);
         this.sync = sync;
     }
 

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

@@ -11,8 +11,7 @@
 import org.elasticsearch.painless.spi.WhitelistLoader;
 import org.elasticsearch.script.ScriptContext;
 import org.elasticsearch.script.ScriptFactory;
-import org.elasticsearch.search.lookup.DocLookup;
-import org.elasticsearch.search.lookup.SourceLookup;
+import org.elasticsearch.search.lookup.SearchLookup;
 
 import java.io.IOException;
 import java.util.List;
@@ -29,7 +28,7 @@ static List<Whitelist> whitelist() {
     public static final String[] PARAMETERS = {};
 
     public interface Factory extends ScriptFactory {
-        LeafFactory newFactory(Map<String, Object> params, SourceLookup source, DocLookup fieldData);
+        LeafFactory newFactory(Map<String, Object> params, SearchLookup searchLookup);
     }
 
     public interface LeafFactory {
@@ -38,14 +37,8 @@ public interface LeafFactory {
 
     private final LongConsumer sync;
 
-    public LongScriptFieldScript(
-        Map<String, Object> params,
-        SourceLookup source,
-        DocLookup fieldData,
-        LeafReaderContext ctx,
-        LongConsumer sync
-    ) {
-        super(params, source, fieldData, ctx);
+    public LongScriptFieldScript(Map<String, Object> params, SearchLookup searchLookup, LeafReaderContext ctx, LongConsumer sync) {
+        super(params, searchLookup, ctx);
         this.sync = sync;
     }
 

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

@@ -6,25 +6,60 @@
 
 package org.elasticsearch.xpack.runtimefields;
 
+import org.elasticsearch.client.Client;
+import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
+import org.elasticsearch.cluster.service.ClusterService;
+import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
+import org.elasticsearch.common.xcontent.NamedXContentRegistry;
+import org.elasticsearch.env.Environment;
+import org.elasticsearch.env.NodeEnvironment;
 import org.elasticsearch.index.mapper.Mapper;
 import org.elasticsearch.plugins.MapperPlugin;
 import org.elasticsearch.plugins.Plugin;
 import org.elasticsearch.plugins.ScriptPlugin;
+import org.elasticsearch.repositories.RepositoriesService;
 import org.elasticsearch.script.ScriptContext;
+import org.elasticsearch.script.ScriptService;
+import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.watcher.ResourceWatcherService;
+import org.elasticsearch.xpack.runtimefields.mapper.ScriptFieldMapper;
 
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Supplier;
 
 public final class RuntimeFields extends Plugin implements MapperPlugin, ScriptPlugin {
 
+    private final ScriptFieldMapper.TypeParser scriptTypeParser = new ScriptFieldMapper.TypeParser();
+
     @Override
     public Map<String, Mapper.TypeParser> getMappers() {
-        return Collections.emptyMap();
+        return Collections.singletonMap(ScriptFieldMapper.CONTENT_TYPE, scriptTypeParser);
     }
 
     @Override
     public List<ScriptContext<?>> getContexts() {
         return List.of(DoubleScriptFieldScript.CONTEXT, LongScriptFieldScript.CONTEXT, StringScriptFieldScript.CONTEXT);
     }
+
+    @Override
+    public Collection<Object> createComponents(
+        Client client,
+        ClusterService clusterService,
+        ThreadPool threadPool,
+        ResourceWatcherService resourceWatcherService,
+        ScriptService scriptService,
+        NamedXContentRegistry xContentRegistry,
+        Environment environment,
+        NodeEnvironment nodeEnvironment,
+        NamedWriteableRegistry namedWriteableRegistry,
+        IndexNameExpressionResolver indexNameExpressionResolver,
+        Supplier<RepositoriesService> repositoriesServiceSupplier
+    ) {
+        // looks like createComponents gets called after getMappers
+        this.scriptTypeParser.setScriptService(scriptService);
+        return Collections.emptyList();
+    }
 }

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

@@ -11,16 +11,15 @@
 import org.elasticsearch.painless.spi.WhitelistLoader;
 import org.elasticsearch.script.ScriptContext;
 import org.elasticsearch.script.ScriptFactory;
-import org.elasticsearch.search.lookup.DocLookup;
-import org.elasticsearch.search.lookup.SourceLookup;
+import org.elasticsearch.search.lookup.SearchLookup;
 
 import java.io.IOException;
 import java.util.List;
 import java.util.Map;
 import java.util.function.Consumer;
 
 public abstract class StringScriptFieldScript extends AbstractScriptFieldScript {
-    static final ScriptContext<Factory> CONTEXT = new ScriptContext<>("string_script_field", Factory.class);
+    public static final ScriptContext<Factory> CONTEXT = new ScriptContext<>("string_script_field", Factory.class);
 
     static List<Whitelist> whitelist() {
         return List.of(WhitelistLoader.loadFromResourceFiles(RuntimeFieldsPainlessExtension.class, "string_whitelist.txt"));
@@ -29,7 +28,7 @@ static List<Whitelist> whitelist() {
     public static final String[] PARAMETERS = {};
 
     public interface Factory extends ScriptFactory {
-        LeafFactory newFactory(Map<String, Object> params, SourceLookup source, DocLookup fieldData);
+        LeafFactory newFactory(Map<String, Object> params, SearchLookup searchLookup);
     }
 
     public interface LeafFactory {
@@ -38,14 +37,8 @@ public interface LeafFactory {
 
     private final Consumer<String> sync;
 
-    public StringScriptFieldScript(
-        Map<String, Object> params,
-        SourceLookup source,
-        DocLookup fieldData,
-        LeafReaderContext ctx,
-        Consumer<String> sync
-    ) {
-        super(params, source, fieldData, ctx);
+    public StringScriptFieldScript(Map<String, Object> params, SearchLookup searchLookup, LeafReaderContext ctx, Consumer<String> sync) {
+        super(params, searchLookup, ctx);
         this.sync = sync;
     }
 

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

@@ -0,0 +1,41 @@
+/*
+ * 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.fielddata;
+
+import org.elasticsearch.index.fielddata.SortingBinaryDocValues;
+import org.elasticsearch.xpack.runtimefields.StringScriptFieldScript;
+
+public final class ScriptBinaryDocValues extends SortingBinaryDocValues {
+
+    private final StringScriptFieldScript script;
+    private final ScriptBinaryFieldData.ScriptBinaryResult scriptBinaryResult;
+
+    ScriptBinaryDocValues(StringScriptFieldScript script, ScriptBinaryFieldData.ScriptBinaryResult scriptBinaryResult) {
+        this.script = script;
+        this.scriptBinaryResult = scriptBinaryResult;
+    }
+
+    @Override
+    public boolean advanceExact(int doc) {
+        script.setDocument(doc);
+        script.execute();
+
+        count = scriptBinaryResult.getResult().size();
+        if (count == 0) {
+            grow();
+            return false;
+        }
+
+        int i = 0;
+        for (String value : scriptBinaryResult.getResult()) {
+            grow();
+            values[i++].copyChars(value);
+        }
+        sort();
+        return true;
+    }
+}