Add support for GeoShape to the scripting fields API (#81617)

This change adds infrastructure for GeoShape making it accessible via the new scripting fields API. 
This does not add any methods outside of get at this point in time since it needs additional 
thought/discussion on what makes sense similar to GeoPoints. Note that because GeoShape does 
not support XContent this is just a skeleton that currently supports getScriptDocValues.

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,6 +11,7 @@ import org.apache.lucene.util.Accountable;
 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;
@@ -18,7 +19,6 @@ 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 abstract class AbstractAtomicGeoShapeShapeFieldData implements LeafGeoSha
         };
     }
 
-    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 abstract class AbstractAtomicGeoShapeShapeFieldData implements LeafGeoSha
 
         @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.MatchNoDocsQuery;
 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.DeprecationCategory;
 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.mapper.MappingParserContext;
 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 class GeoShapeWithDocValuesFieldMapper extends AbstractShapeGeometryField
 
         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 @@ public class GeoShapeWithDocValuesFieldMapper extends AbstractShapeGeometryField
         }
         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: