Add support for GeoShape to the scripting fields API

docs/changelog/81617.yaml

@@ -0,0 +1,5 @@
+pr: 81617
+summary: Add support for `GeoShape` to the scripting fields API
+area: Infra/Scripting
+type: enhancement
+issues: []

x-pack/plugin/spatial/build.gradle

@@ -5,11 +5,12 @@ esplugin {
   name 'spatial'
   description 'A plugin for Basic Spatial features'
   classname 'org.elasticsearch.xpack.spatial.SpatialPlugin'
-  extendedPlugins = ['x-pack-core', 'legacy-geo']
+  extendedPlugins = ['x-pack-core', 'legacy-geo', 'lang-painless']
 }
 
 dependencies {
   compileOnly project(path: ':modules:legacy-geo')
+  compileOnly project(':modules:lang-painless:spi')
   compileOnly project(path: xpackModule('core'))
   testImplementation(testArtifact(project(xpackModule('core'))))
   testImplementation project(path: xpackModule('vector-tile'))

x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/SpatialPainlessExtension.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
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.spatial;
+
+import org.elasticsearch.painless.spi.PainlessExtension;
+import org.elasticsearch.painless.spi.Whitelist;
+import org.elasticsearch.painless.spi.WhitelistLoader;
+import org.elasticsearch.script.ScriptContext;
+import org.elasticsearch.script.ScriptModule;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+public class SpatialPainlessExtension implements PainlessExtension {
+
+    private static final List<Whitelist> WHITELISTS = List.of(
+        WhitelistLoader.loadFromResourceFiles(
+            SpatialPainlessExtension.class,
+            "org.elasticsearch.xpack.spatial.index.fielddata.txt",
+            "org.elasticsearch.xpack.spatial.index.fielddata.plain.txt",
+            "org.elasticsearch.xpack.spatial.index.mapper.txt"
+        )
+    );
+
+    @Override
+    public Map<ScriptContext<?>, List<Whitelist>> getContextWhitelists() {
+        Map<ScriptContext<?>, List<Whitelist>> contextWhitelistMap = new HashMap<>();
+
+        for (ScriptContext<?> context : ScriptModule.CORE_CONTEXTS.values()) {
+            contextWhitelistMap.put(context, WHITELISTS);
+        }
+
+        return contextWhitelistMap;
+    }
+}

x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/fielddata/plain/AbstractAtomicGeoShapeShapeFieldData.java

@@ -11,14 +11,14 @@
 import org.elasticsearch.common.geo.GeoBoundingBox;
 import org.elasticsearch.common.geo.GeoPoint;
 import org.elasticsearch.index.fielddata.ScriptDocValues;
+import org.elasticsearch.index.fielddata.ScriptDocValues.GeometrySupplier;
 import org.elasticsearch.index.fielddata.SortedBinaryDocValues;
 import org.elasticsearch.script.field.DocValuesField;
 import org.elasticsearch.script.field.ToScriptField;
 import org.elasticsearch.xpack.spatial.index.fielddata.GeoShapeValues;
 import org.elasticsearch.xpack.spatial.index.fielddata.GeoShapeValues.GeoShapeValue;
 import org.elasticsearch.xpack.spatial.index.fielddata.LeafGeoShapeFieldData;
 
-import java.io.IOException;
 import java.util.Collection;
 import java.util.Collections;
 
@@ -66,71 +66,23 @@ public GeoShapeValues getGeoShapeValues() {
         };
     }
 
-    public static final class GeoShapeSupplier implements ScriptDocValues.GeometrySupplier<GeoShapeValue> {
-
-        private final GeoShapeValues in;
-        private final GeoPoint centroid = new GeoPoint();
-        private final GeoBoundingBox boundingBox = new GeoBoundingBox(new GeoPoint(), new GeoPoint());
-        private GeoShapeValues.GeoShapeValue value;
-
-        public GeoShapeSupplier(GeoShapeValues in) {
-            this.in = in;
-        }
-
-        @Override
-        public void setNextDocId(int docId) throws IOException {
-            if (in.advanceExact(docId)) {
-                value = in.value();
-                centroid.reset(value.lat(), value.lon());
-                boundingBox.topLeft().reset(value.boundingBox().maxY(), value.boundingBox().minX());
-                boundingBox.bottomRight().reset(value.boundingBox().minY(), value.boundingBox().maxX());
-            } else {
-                value = null;
-            }
-        }
-
-        @Override
-        public GeoShapeValue getInternal(int index) {
-            throw new UnsupportedOperationException();
-        }
-
-        public GeoShapeValue getInternal() {
-            return value;
-        }
-
-        @Override
-        public int size() {
-            return value == null ? 0 : 1;
-        }
-
-        @Override
-        public GeoPoint getInternalCentroid() {
-            return centroid;
-        }
-
-        @Override
-        public GeoBoundingBox getInternalBoundingBox() {
-            return boundingBox;
-        }
-    }
-
     public static final class GeoShapeScriptValues extends ScriptDocValues.Geometry<GeoShapeValue> {
 
-        private final GeoShapeSupplier gsSupplier;
+        private final GeometrySupplier<GeoShapeValue> gsSupplier;
 
-        public GeoShapeScriptValues(GeoShapeSupplier supplier) {
+        public GeoShapeScriptValues(GeometrySupplier<GeoShapeValue> supplier) {
             super(supplier);
             this.gsSupplier = supplier;
         }
 
         @Override
         public int getDimensionalType() {
-            return gsSupplier.getInternal() == null ? -1 : gsSupplier.getInternal().dimensionalShapeType().ordinal();
+            return gsSupplier.getInternal(0) == null ? -1 : gsSupplier.getInternal(0).dimensionalShapeType().ordinal();
         }
 
         @Override
         public GeoPoint getCentroid() {
-            return gsSupplier.getInternal() == null ? null : gsSupplier.getInternalCentroid();
+            return gsSupplier.getInternal(0) == null ? null : gsSupplier.getInternalCentroid();
         }
 
         @Override
@@ -145,12 +97,16 @@ public double getMercatorHeight() {
 
         @Override
         public GeoBoundingBox getBoundingBox() {
-            return gsSupplier.getInternal() == null ? null : gsSupplier.getInternalBoundingBox();
+            return gsSupplier.getInternal(0) == null ? null : gsSupplier.getInternalBoundingBox();
         }
 
         @Override
         public GeoShapeValues.GeoShapeValue get(int index) {
-            return gsSupplier.getInternal();
+            return gsSupplier.getInternal(0);
+        }
+
+        public GeoShapeValues.GeoShapeValue getValue() {
+            return gsSupplier.getInternal(0);
         }
 
         @Override

x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/mapper/GeoShapeWithDocValuesFieldMapper.java

@@ -15,7 +15,9 @@
 import org.apache.lucene.search.Query;
 import org.elasticsearch.Version;
 import org.elasticsearch.common.Explicit;
+import org.elasticsearch.common.geo.GeoBoundingBox;
 import org.elasticsearch.common.geo.GeoFormatterFactory;
+import org.elasticsearch.common.geo.GeoPoint;
 import org.elasticsearch.common.geo.GeoShapeUtils;
 import org.elasticsearch.common.geo.GeometryParser;
 import org.elasticsearch.common.geo.Orientation;
@@ -24,6 +26,7 @@
 import org.elasticsearch.common.logging.DeprecationLogger;
 import org.elasticsearch.geometry.Geometry;
 import org.elasticsearch.index.fielddata.IndexFieldData;
+import org.elasticsearch.index.fielddata.ScriptDocValues;
 import org.elasticsearch.index.mapper.AbstractShapeGeometryFieldMapper;
 import org.elasticsearch.index.mapper.DocumentParserContext;
 import org.elasticsearch.index.mapper.FieldMapper;
@@ -39,16 +42,19 @@
 import org.elasticsearch.index.query.QueryShardException;
 import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.legacygeo.mapper.LegacyGeoShapeFieldMapper;
-import org.elasticsearch.script.field.DelegateDocValuesField;
+import org.elasticsearch.script.field.DocValuesField;
 import org.elasticsearch.search.lookup.SearchLookup;
+import org.elasticsearch.xpack.spatial.index.fielddata.GeoShapeValues;
 import org.elasticsearch.xpack.spatial.index.fielddata.plain.AbstractAtomicGeoShapeShapeFieldData;
 import org.elasticsearch.xpack.spatial.index.fielddata.plain.AbstractLatLonShapeIndexFieldData;
 import org.elasticsearch.xpack.spatial.search.aggregations.support.GeoShapeValuesSourceType;
 
 import java.io.IOException;
 import java.util.Arrays;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.NoSuchElementException;
 import java.util.Set;
 import java.util.function.Function;
 import java.util.function.Supplier;
@@ -175,16 +181,7 @@ public GeoShapeWithDocValuesFieldType(
 
         public IndexFieldData.Builder fielddataBuilder(String fullyQualifiedIndexName, Supplier<SearchLookup> searchLookup) {
             failIfNoDocValues();
-            return new AbstractLatLonShapeIndexFieldData.Builder(
-                name(),
-                GeoShapeValuesSourceType.instance(),
-                (dv, n) -> new DelegateDocValuesField(
-                    new AbstractAtomicGeoShapeShapeFieldData.GeoShapeScriptValues(
-                        new AbstractAtomicGeoShapeShapeFieldData.GeoShapeSupplier(dv)
-                    ),
-                    n
-                )
-            );
+            return new AbstractLatLonShapeIndexFieldData.Builder(name(), GeoShapeValuesSourceType.instance(), GeoShapeDocValuesField::new);
         }
 
         @Override
@@ -344,4 +341,114 @@ protected void checkIncomingMergeType(FieldMapper mergeWith) {
         }
         super.checkIncomingMergeType(mergeWith);
     }
+
+    public static class GeoShapeDocValuesField
+        implements
+            DocValuesField<GeoShapeValues.GeoShapeValue>,
+            ScriptDocValues.GeometrySupplier<GeoShapeValues.GeoShapeValue> {
+
+        private final GeoShapeValues in;
+        protected final String name;
+
+        private GeoShapeValues.GeoShapeValue value;
+
+        // maintain bwc by making bounding box and centroid available to GeoShapeValues (ScriptDocValues)
+        private final GeoPoint centroid = new GeoPoint();
+        private final GeoBoundingBox boundingBox = new GeoBoundingBox(new GeoPoint(), new GeoPoint());
+        private AbstractAtomicGeoShapeShapeFieldData.GeoShapeScriptValues geoShapeScriptValues;
+
+        public GeoShapeDocValuesField(GeoShapeValues in, String name) {
+            this.in = in;
+            this.name = name;
+        }
+
+        @Override
+        public void setNextDocId(int docId) throws IOException {
+            if (in.advanceExact(docId)) {
+                value = in.value();
+                centroid.reset(value.lat(), value.lon());
+                boundingBox.topLeft().reset(value.boundingBox().maxY(), value.boundingBox().minX());
+                boundingBox.bottomRight().reset(value.boundingBox().minY(), value.boundingBox().maxX());
+            } else {
+                value = null;
+            }
+        }
+
+        @Override
+        public ScriptDocValues<GeoShapeValues.GeoShapeValue> getScriptDocValues() {
+            if (geoShapeScriptValues == null) {
+                geoShapeScriptValues = new AbstractAtomicGeoShapeShapeFieldData.GeoShapeScriptValues(this);
+            }
+
+            return geoShapeScriptValues;
+        }
+
+        @Override
+        public GeoShapeValues.GeoShapeValue getInternal(int index) {
+            if (index != 0) {
+                throw new UnsupportedOperationException();
+            }
+
+            return value;
+        }
+
+        // maintain bwc by making centroid available to GeoShapeValues (ScriptDocValues)
+        @Override
+        public GeoPoint getInternalCentroid() {
+            return centroid;
+        }
+
+        // maintain bwc by making centroid available to GeoShapeValues (ScriptDocValues)
+        @Override
+        public GeoBoundingBox getInternalBoundingBox() {
+            return boundingBox;
+        }
+
+        @Override
+        public String getName() {
+            return name;
+        }
+
+        @Override
+        public boolean isEmpty() {
+            return value == null;
+        }
+
+        @Override
+        public int size() {
+            return value == null ? 0 : 1;
+        }
+
+        public GeoShapeValues.GeoShapeValue get(GeoShapeValues.GeoShapeValue defaultValue) {
+            return get(0, defaultValue);
+        }
+
+        public GeoShapeValues.GeoShapeValue get(int index, GeoShapeValues.GeoShapeValue defaultValue) {
+            if (isEmpty() || index != 0) {
+                return defaultValue;
+            }
+
+            return value;
+        }
+
+        @Override
+        public Iterator<GeoShapeValues.GeoShapeValue> iterator() {
+            return new Iterator<GeoShapeValues.GeoShapeValue>() {
+                private int index = 0;
+
+                @Override
+                public boolean hasNext() {
+                    return index < size();
+                }
+
+                @Override
+                public GeoShapeValues.GeoShapeValue next() {
+                    if (hasNext() == false) {
+                        throw new NoSuchElementException();
+                    }
+                    return value;
+                }
+            };
+        }
+    }
 }

x-pack/plugin/spatial/src/main/resources/META-INF/services/org.elasticsearch.painless.spi.PainlessExtension

@@ -0,0 +1 @@
+org.elasticsearch.xpack.spatial.SpatialPainlessExtension

x-pack/plugin/spatial/src/main/resources/org/elasticsearch/xpack/spatial/org.elasticsearch.xpack.spatial.index.fielddata.plain.txt

@@ -0,0 +1,4 @@
+class org.elasticsearch.xpack.spatial.index.fielddata.plain.AbstractAtomicGeoShapeShapeFieldData$GeoShapeScriptValues {
+  GeoShapeValues.GeoShapeValue get(int)
+  GeoShapeValues.GeoShapeValue getValue()
+}

x-pack/plugin/spatial/src/main/resources/org/elasticsearch/xpack/spatial/org.elasticsearch.xpack.spatial.index.fielddata.txt

@@ -0,0 +1,2 @@
+class org.elasticsearch.xpack.spatial.index.fielddata.GeoShapeValues$GeoShapeValue {
+}

x-pack/plugin/spatial/src/main/resources/org/elasticsearch/xpack/spatial/org.elasticsearch.xpack.spatial.index.mapper.txt

@@ -0,0 +1,4 @@
+class org.elasticsearch.xpack.spatial.index.mapper.GeoShapeWithDocValuesFieldMapper$GeoShapeDocValuesField {
+  GeoShapeValues.GeoShapeValue get(GeoShapeValues.GeoShapeValue)
+  GeoShapeValues.GeoShapeValue get(int, GeoShapeValues.GeoShapeValue)
+}

x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/spatial/70_script_doc_values.yml

@@ -92,6 +92,30 @@ setup:
 
   - match: { error.root_cause.0.reason: "cannot write xcontent for geo_shape doc value" }
 
+  - do:
+      catch: /illegal_argument_exception/
+      search:
+        rest_total_hits_as_int: true
+        body:
+          script_fields:
+            type:
+              script:
+                source: "field('geo_shape').get(null)"
+
+  - match: { error.root_cause.0.reason: "cannot write xcontent for geo_shape doc value" }
+
+  - do:
+      catch: /illegal_argument_exception/
+      search:
+        rest_total_hits_as_int: true
+        body:
+          script_fields:
+            type:
+              script:
+                source: "/* avoid yaml stash */ $('geo_shape', null)"
+
+  - match: { error.root_cause.0.reason: "cannot write xcontent for geo_shape doc value" }
+
 ---
 "diagonal length":
   - do: