Add runtime field of type geo_shape

docs/changelog/100492.yaml

@@ -0,0 +1,6 @@
+pr: 100492
+summary: Add runtime field of type `geo_shape`
+area: Geo
+type: enhancement
+issues:
+ - 61299

modules/lang-painless/src/main/java/org/elasticsearch/painless/action/PainlessExecuteAction.java

@@ -50,6 +50,7 @@
 import org.elasticsearch.core.Nullable;
 import org.elasticsearch.core.Strings;
 import org.elasticsearch.core.Tuple;
+import org.elasticsearch.geometry.Geometry;
 import org.elasticsearch.geometry.Point;
 import org.elasticsearch.index.Index;
 import org.elasticsearch.index.IndexService;
@@ -76,6 +77,7 @@
 import org.elasticsearch.script.DoubleFieldScript;
 import org.elasticsearch.script.FilterScript;
 import org.elasticsearch.script.GeoPointFieldScript;
+import org.elasticsearch.script.GeometryFieldScript;
 import org.elasticsearch.script.IpFieldScript;
 import org.elasticsearch.script.LongFieldScript;
 import org.elasticsearch.script.ScoreScript;
@@ -681,6 +683,25 @@ static Response innerShardOperation(Request request, ScriptService scriptService
                     );
                     return new Response(format.apply(points));
                 }, indexService);
+            } else if (scriptContext == GeometryFieldScript.CONTEXT) {
+                return prepareRamIndex(request, (context, leafReaderContext) -> {
+                    GeometryFieldScript.Factory factory = scriptService.compile(request.script, GeometryFieldScript.CONTEXT);
+                    GeometryFieldScript.LeafFactory leafFactory = factory.newFactory(
+                        GeometryFieldScript.CONTEXT.name,
+                        request.getScript().getParams(),
+                        context.lookup(),
+                        OnScriptError.FAIL
+                    );
+                    GeometryFieldScript geometryFieldScript = leafFactory.newInstance(leafReaderContext);
+                    List<Geometry> geometries = new ArrayList<>();
+                    geometryFieldScript.runForDoc(0, geometries::add);
+                    // convert geometries to the standard format of the fields api
+                    Function<List<Geometry>, List<Object>> format = GeometryFormatterFactory.getFormatter(
+                        GeometryFormatterFactory.GEOJSON,
+                        Function.identity()
+                    );
+                    return new Response(format.apply(geometries));
+                }, indexService);
             } else if (scriptContext == IpFieldScript.CONTEXT) {
                 return prepareRamIndex(request, (context, leafReaderContext) -> {
                     IpFieldScript.Factory factory = scriptService.compile(request.script, IpFieldScript.CONTEXT);

modules/lang-painless/src/main/resources/org/elasticsearch/painless/org.elasticsearch.script.geometry_field.txt

@@ -0,0 +1,20 @@
+#
+# Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+# or more contributor license agreements. Licensed under the Elastic License
+# 2.0 and the Server Side Public License, v 1; you may not use this file except
+# in compliance with, at your election, the Elastic License 2.0 or the Server
+# Side Public License, v 1.
+#
+
+# The whitelist for runtime fields that generate geometries
+
+# These two whitelists are required for painless to find the classes
+class org.elasticsearch.script.GeometryFieldScript @no_import {
+}
+class org.elasticsearch.script.GeometryFieldScript$Factory @no_import {
+}
+
+static_import {
+    # The `emit` callback to collect values for the field
+    void emit(org.elasticsearch.script.GeometryFieldScript, Object) bound_to org.elasticsearch.script.GeometryFieldScript$Emit
+}

modules/lang-painless/src/test/java/org/elasticsearch/painless/action/PainlessExecuteApiTests.java

@@ -259,6 +259,51 @@ public void testGeoPointFieldExecutionContext() throws IOException {
         assertEquals("Point", points.get(1).get("type"));
     }
 
+    @SuppressWarnings("unchecked")
+    public void testGeometryFieldExecutionContext() throws IOException {
+        ScriptService scriptService = getInstanceFromNode(ScriptService.class);
+        IndexService indexService = createIndex("index", Settings.EMPTY, "doc", "test_point", "type=geo_point");
+
+        Request.ContextSetup contextSetup = new Request.ContextSetup(
+            "index",
+            new BytesArray("{\"test_point\":\"30.0,40.0\"}"),
+            new MatchAllQueryBuilder()
+        );
+        contextSetup.setXContentType(XContentType.JSON);
+        Request request = new Request(
+            new Script(
+                ScriptType.INLINE,
+                "painless",
+                "emit(\"Point(\" + doc['test_point'].value.lon + \" \" + doc['test_point'].value.lat + \")\")",
+                emptyMap()
+            ),
+            "geometry_field",
+            contextSetup
+        );
+        Response response = innerShardOperation(request, scriptService, indexService);
+        List<Map<String, Object>> geometry = (List<Map<String, Object>>) response.getResult();
+        assertEquals(40.0, (double) ((List<Object>) geometry.get(0).get("coordinates")).get(0), 0.00001);
+        assertEquals(30.0, (double) ((List<Object>) geometry.get(0).get("coordinates")).get(1), 0.00001);
+        assertEquals("Point", geometry.get(0).get("type"));
+
+        contextSetup = new Request.ContextSetup("index", new BytesArray("{}"), new MatchAllQueryBuilder());
+        contextSetup.setXContentType(XContentType.JSON);
+        request = new Request(
+            new Script(
+                ScriptType.INLINE,
+                "painless",
+                "emit(\"LINESTRING(78.96 12.12, 12.12 78.96)\"); emit(\"POINT(13.45 56.78)\");",
+                emptyMap()
+            ),
+            "geometry_field",
+            contextSetup
+        );
+        response = innerShardOperation(request, scriptService, indexService);
+        geometry = (List<Map<String, Object>>) response.getResult();
+        assertEquals("GeometryCollection", geometry.get(0).get("type"));
+        assertEquals(2, ((List<?>) geometry.get(0).get("geometries")).size());
+    }
+
     public void testIpFieldExecutionContext() throws IOException {
         ScriptService scriptService = getInstanceFromNode(ScriptService.class);
         IndexService indexService = createIndex("index", Settings.EMPTY, "doc", "test_ip", "type=ip");

server/src/main/java/org/elasticsearch/index/mapper/AbstractShapeGeometryFieldMapper.java

@@ -83,6 +83,21 @@ protected AbstractShapeGeometryFieldMapper(
         this.orientation = orientation;
     }
 
+    protected AbstractShapeGeometryFieldMapper(
+        String simpleName,
+        MappedFieldType mappedFieldType,
+        MultiFields multiFields,
+        Explicit<Boolean> coerce,
+        Explicit<Orientation> orientation,
+        CopyTo copyTo,
+        Parser<T> parser,
+        OnScriptError onScriptError
+    ) {
+        super(simpleName, mappedFieldType, multiFields, copyTo, parser, onScriptError);
+        this.coerce = coerce;
+        this.orientation = orientation;
+    }
+
     public boolean coerce() {
         return coerce.value();
     }

server/src/main/java/org/elasticsearch/script/GeometryFieldScript.java

@@ -0,0 +1,143 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.script;
+
+import org.apache.lucene.index.LeafReaderContext;
+import org.elasticsearch.common.geo.GeometryParser;
+import org.elasticsearch.common.geo.Orientation;
+import org.elasticsearch.geometry.Geometry;
+import org.elasticsearch.geometry.GeometryCollection;
+import org.elasticsearch.index.mapper.OnScriptError;
+import org.elasticsearch.search.lookup.SearchLookup;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.function.Consumer;
+import java.util.function.Function;
+
+/**
+ * Script producing geometries. It generates a unique {@link Geometry} for each document.
+ */
+public abstract class GeometryFieldScript extends AbstractFieldScript {
+    public static final ScriptContext<Factory> CONTEXT = newContext("geometry_field", Factory.class);
+
+    public static final Factory PARSE_FROM_SOURCE = new Factory() {
+        @Override
+        public LeafFactory newFactory(String field, Map<String, Object> params, SearchLookup lookup, OnScriptError onScriptError) {
+            return ctx -> new GeometryFieldScript(field, params, lookup, OnScriptError.FAIL, ctx) {
+                @Override
+                public void execute() {
+                    emitFromSource();
+                }
+            };
+        }
+
+        @Override
+        public boolean isResultDeterministic() {
+            return true;
+        }
+    };
+
+    public static Factory leafAdapter(Function<SearchLookup, CompositeFieldScript.LeafFactory> parentFactory) {
+        return (leafFieldName, params, searchLookup, onScriptError) -> {
+            CompositeFieldScript.LeafFactory parentLeafFactory = parentFactory.apply(searchLookup);
+            return (LeafFactory) ctx -> {
+                CompositeFieldScript compositeFieldScript = parentLeafFactory.newInstance(ctx);
+                return new GeometryFieldScript(leafFieldName, params, searchLookup, onScriptError, ctx) {
+                    @Override
+                    public void setDocument(int docId) {
+                        compositeFieldScript.setDocument(docId);
+                    }
+
+                    @Override
+                    public void execute() {
+                        emitFromCompositeScript(compositeFieldScript);
+                    }
+                };
+            };
+        };
+    }
+
+    @SuppressWarnings("unused")
+    public static final String[] PARAMETERS = {};
+
+    public interface Factory extends ScriptFactory {
+        LeafFactory newFactory(String fieldName, Map<String, Object> params, SearchLookup searchLookup, OnScriptError onScriptError);
+    }
+
+    public interface LeafFactory {
+        GeometryFieldScript newInstance(LeafReaderContext ctx);
+    }
+
+    private final List<Geometry> geometries = new ArrayList<>();
+
+    private final GeometryParser geometryParser;
+
+    public GeometryFieldScript(
+        String fieldName,
+        Map<String, Object> params,
+        SearchLookup searchLookup,
+        OnScriptError onScriptError,
+        LeafReaderContext ctx
+    ) {
+        super(fieldName, params, searchLookup, ctx, onScriptError);
+        geometryParser = new GeometryParser(Orientation.CCW.getAsBoolean(), false, true);
+    }
+
+    @Override
+    protected void prepareExecute() {
+        geometries.clear();
+    }
+
+    /**
+     * Execute the script for the provided {@code docId}, passing results to the {@code consumer}
+     */
+    public final void runForDoc(int docId, Consumer<Geometry> consumer) {
+        runForDoc(docId);
+        consumer.accept(geometry());
+    }
+
+    /**
+     * {@link Geometry} from the last time {@link #runForDoc(int)} was called.
+     */
+    public final Geometry geometry() {
+        if (geometries.isEmpty()) {
+            return null;
+        }
+        return geometries.size() == 1 ? geometries.get(0) : new GeometryCollection<>(geometries);
+    }
+
+    /**
+     * The number of results produced the last time {@link #runForDoc(int)} was called. It is 1 if
+     * the document exists, otherwise 0.
+     */
+    public final int count() {
+        // Note that emitting multiple geometries gets handled by a GeometryCollection
+        return geometries.isEmpty() ? 0 : 1;
+    }
+
+    @Override
+    protected void emitFromObject(Object value) {
+        geometries.add(geometryParser.parseGeometry(value));
+    }
+
+    public static class Emit {
+        private final GeometryFieldScript script;
+
+        public Emit(GeometryFieldScript script) {
+            this.script = script;
+        }
+
+        public void emit(Object object) {
+            script.checkMaxSize(script.geometries.size());
+            script.emitFromObject(object);
+        }
+    }
+}

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

@@ -34,6 +34,7 @@ public class ScriptModule {
         LongFieldScript.CONTEXT,
         StringFieldScript.CONTEXT,
         GeoPointFieldScript.CONTEXT,
+        GeometryFieldScript.CONTEXT,
         IpFieldScript.CONTEXT,
         CompositeFieldScript.CONTEXT
     );

server/src/test/java/org/elasticsearch/index/mapper/GeometryFieldScriptTests.java

@@ -0,0 +1,82 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.index.mapper;
+
+import org.apache.lucene.document.StoredField;
+import org.apache.lucene.index.DirectoryReader;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.tests.index.RandomIndexWriter;
+import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.script.AbstractFieldScript;
+import org.elasticsearch.script.GeometryFieldScript;
+import org.elasticsearch.script.ScriptContext;
+import org.elasticsearch.search.lookup.SearchLookup;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Map;
+
+import static org.hamcrest.Matchers.equalTo;
+
+public class GeometryFieldScriptTests extends FieldScriptTestCase<GeometryFieldScript.Factory> {
+    public static final GeometryFieldScript.Factory DUMMY = (fieldName, params, lookup, onScriptError) -> ctx -> new GeometryFieldScript(
+        fieldName,
+        params,
+        lookup,
+        OnScriptError.FAIL,
+        ctx
+    ) {
+        @Override
+        public void execute() {
+            emitFromObject("POINT(0 0)");
+        }
+    };
+
+    @Override
+    protected ScriptContext<GeometryFieldScript.Factory> context() {
+        return GeometryFieldScript.CONTEXT;
+    }
+
+    @Override
+    protected GeometryFieldScript.Factory dummyScript() {
+        return DUMMY;
+    }
+
+    @Override
+    protected GeometryFieldScript.Factory fromSource() {
+        return GeometryFieldScript.PARSE_FROM_SOURCE;
+    }
+
+    public void testTooManyValues() throws IOException {
+        try (Directory directory = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), directory)) {
+            iw.addDocument(List.of(new StoredField("_source", new BytesRef("{}"))));
+            try (DirectoryReader reader = iw.getReader()) {
+                GeometryFieldScript script = new GeometryFieldScript(
+                    "test",
+                    Map.of(),
+                    new SearchLookup(field -> null, (ft, lookup, fdt) -> null, (ctx, doc) -> null),
+                    OnScriptError.FAIL,
+                    reader.leaves().get(0)
+                ) {
+                    @Override
+                    public void execute() {
+                        for (int i = 0; i <= AbstractFieldScript.MAX_VALUES; i++) {
+                            new Emit(this).emit("POINT(0 0)");
+                        }
+                    }
+                };
+                Exception e = expectThrows(IllegalArgumentException.class, script::execute);
+                assertThat(
+                    e.getMessage(),
+                    equalTo("Runtime field [test] is emitting [101] values while the maximum number of values allowed is [100]")
+                );
+            }
+        }
+    }
+}

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

@@ -377,6 +377,20 @@ public void execute() {
                 }
             };
             return context.factoryClazz.cast(objectFieldScript);
+        } else if (context.instanceClazz.equals(GeometryFieldScript.class)) {
+            GeometryFieldScript.Factory geometryFieldScript = (f, p, s, onScriptError) -> ctx -> new GeometryFieldScript(
+                f,
+                p,
+                s,
+                OnScriptError.FAIL,
+                ctx
+            ) {
+                @Override
+                public void execute() {
+                    emitFromObject("POINT(1.2 1.2)");
+                }
+            };
+            return context.factoryClazz.cast(geometryFieldScript);
         } else if (context.instanceClazz.equals(DoubleValuesScript.class)) {
             DoubleValuesScript.Factory doubleValuesScript = () -> new MockDoubleValuesScript();
             return context.factoryClazz.cast(doubleValuesScript);