geo_point runtime field implementation

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

@@ -14,6 +14,7 @@
 import org.elasticsearch.xpack.runtimefields.mapper.BooleanFieldScript;
 import org.elasticsearch.xpack.runtimefields.mapper.DateFieldScript;
 import org.elasticsearch.xpack.runtimefields.mapper.DoubleFieldScript;
+import org.elasticsearch.xpack.runtimefields.mapper.GeoPointFieldScript;
 import org.elasticsearch.xpack.runtimefields.mapper.IpFieldScript;
 import org.elasticsearch.xpack.runtimefields.mapper.LongFieldScript;
 import org.elasticsearch.xpack.runtimefields.mapper.RuntimeFieldMapper;
@@ -38,7 +39,8 @@ public List<ScriptContext<?>> getContexts() {
             DoubleFieldScript.CONTEXT,
             IpFieldScript.CONTEXT,
             LongFieldScript.CONTEXT,
-            StringFieldScript.CONTEXT
+            StringFieldScript.CONTEXT,
+            GeoPointFieldScript.CONTEXT
         );
     }
 }

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

@@ -0,0 +1,49 @@
+/*
+ * 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.apache.lucene.geo.GeoEncodingUtils;
+import org.elasticsearch.common.geo.GeoPoint;
+import org.elasticsearch.index.fielddata.MultiGeoPointValues;
+import org.elasticsearch.xpack.runtimefields.mapper.GeoPointFieldScript;
+
+import java.util.Arrays;
+
+public final class GeoPointScriptDocValues extends MultiGeoPointValues {
+    private final GeoPointFieldScript script;
+    private final GeoPoint point;
+    private int cursor;
+
+    GeoPointScriptDocValues(GeoPointFieldScript script) {
+        this.script = script;
+        this.point = new GeoPoint();
+    }
+
+    @Override
+    public boolean advanceExact(int docId) {
+        script.runForDoc(docId);
+        if (script.count() == 0) {
+            return false;
+        }
+        Arrays.sort(script.values(), 0, script.count());
+        cursor = 0;
+        return true;
+    }
+
+    @Override
+    public int docValueCount() {
+        return script.count();
+    }
+
+    @Override
+    public GeoPoint nextValue() {
+        final long value = script.values()[cursor++];
+        final int lat = (int) (value >>> 32);
+        final int lon = (int) (value & 0xFFFFFFFF);
+        return point.reset(GeoEncodingUtils.decodeLatitude(lat), GeoEncodingUtils.decodeLongitude(lon));
+    }
+}

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

@@ -0,0 +1,106 @@
+/*
+ * 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.apache.lucene.index.LeafReaderContext;
+import org.apache.lucene.search.SortField;
+import org.elasticsearch.common.util.BigArrays;
+import org.elasticsearch.index.fielddata.IndexFieldData;
+import org.elasticsearch.index.fielddata.IndexFieldDataCache;
+import org.elasticsearch.index.fielddata.IndexGeoPointFieldData;
+import org.elasticsearch.index.fielddata.LeafGeoPointFieldData;
+import org.elasticsearch.index.fielddata.MultiGeoPointValues;
+import org.elasticsearch.index.fielddata.plain.AbstractLeafGeoPointFieldData;
+import org.elasticsearch.index.mapper.MapperService;
+import org.elasticsearch.indices.breaker.CircuitBreakerService;
+import org.elasticsearch.search.DocValueFormat;
+import org.elasticsearch.search.MultiValueMode;
+import org.elasticsearch.search.aggregations.support.CoreValuesSourceType;
+import org.elasticsearch.search.aggregations.support.ValuesSourceType;
+import org.elasticsearch.search.sort.BucketedSort;
+import org.elasticsearch.search.sort.SortOrder;
+import org.elasticsearch.xpack.runtimefields.mapper.GeoPointFieldScript;
+
+public class GeoPointScriptFieldData implements IndexGeoPointFieldData {
+    public static class Builder implements IndexFieldData.Builder {
+        private final String name;
+        private final GeoPointFieldScript.LeafFactory leafFactory;
+
+        public Builder(String name, GeoPointFieldScript.LeafFactory leafFactory) {
+            this.name = name;
+            this.leafFactory = leafFactory;
+        }
+
+        @Override
+        public GeoPointScriptFieldData build(IndexFieldDataCache cache, CircuitBreakerService breakerService, MapperService mapperService) {
+            return new GeoPointScriptFieldData(name, leafFactory);
+        }
+    }
+
+    private final GeoPointFieldScript.LeafFactory leafFactory;
+    private final String name;
+
+    private GeoPointScriptFieldData(String fieldName, GeoPointFieldScript.LeafFactory leafFactory) {
+        this.name = fieldName;
+        this.leafFactory = leafFactory;
+    }
+
+    @Override
+    public SortField sortField(Object missingValue, MultiValueMode sortMode, XFieldComparatorSource.Nested nested, boolean reverse) {
+        throw new IllegalArgumentException("can't sort on geo_point field without using specific sorting feature, like geo_distance");
+    }
+
+    @Override
+    public BucketedSort newBucketedSort(
+        BigArrays bigArrays,
+        Object missingValue,
+        MultiValueMode sortMode,
+        XFieldComparatorSource.Nested nested,
+        SortOrder sortOrder,
+        DocValueFormat format,
+        int bucketSize,
+        BucketedSort.ExtraData extra
+    ) {
+        throw new IllegalArgumentException("can't sort on geo_point field without using specific sorting feature, like geo_distance");
+    }
+
+    @Override
+    public String getFieldName() {
+        return name;
+    }
+
+    @Override
+    public ValuesSourceType getValuesSourceType() {
+        return CoreValuesSourceType.GEOPOINT;
+    }
+
+    @Override
+    public LeafGeoPointFieldData load(LeafReaderContext context) {
+        GeoPointFieldScript script = leafFactory.newInstance(context);
+        return new AbstractLeafGeoPointFieldData() {
+            @Override
+            public MultiGeoPointValues getGeoPointValues() {
+                return new GeoPointScriptDocValues(script);
+            }
+
+            @Override
+            public long ramBytesUsed() {
+                return 0;
+            }
+
+            @Override
+            public void close() {
+
+            }
+        };
+    }
+
+    @Override
+    public LeafGeoPointFieldData loadDirect(LeafReaderContext context) {
+        return load(context);
+    }
+}

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

@@ -0,0 +1,66 @@
+/*
+ * 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.mapper;
+
+import org.apache.lucene.index.LeafReaderContext;
+import org.elasticsearch.painless.spi.Whitelist;
+import org.elasticsearch.painless.spi.WhitelistLoader;
+import org.elasticsearch.script.ScriptContext;
+import org.elasticsearch.script.ScriptFactory;
+import org.elasticsearch.search.lookup.SearchLookup;
+import org.apache.lucene.document.LatLonDocValuesField;
+
+import java.util.List;
+import java.util.Map;
+
+import static org.apache.lucene.geo.GeoEncodingUtils.encodeLatitude;
+import static org.apache.lucene.geo.GeoEncodingUtils.encodeLongitude;
+
+/**
+ * Script producing geo points. Similarly to what {@link LatLonDocValuesField} does,
+ * it encodes the points as a long value.
+ */
+public abstract class GeoPointFieldScript extends AbstractLongFieldScript {
+    public static final ScriptContext<Factory> CONTEXT = newContext("geo_point_script_field", Factory.class);
+
+    static List<Whitelist> whitelist() {
+        return List.of(WhitelistLoader.loadFromResourceFiles(RuntimeFieldsPainlessExtension.class, "geo_point_whitelist.txt"));
+    }
+
+    @SuppressWarnings("unused")
+    public static final String[] PARAMETERS = {};
+
+    public interface Factory extends ScriptFactory {
+        LeafFactory newFactory(String fieldName, Map<String, Object> params, SearchLookup searchLookup);
+    }
+
+    public interface LeafFactory {
+        GeoPointFieldScript newInstance(LeafReaderContext ctx);
+    }
+
+    public GeoPointFieldScript(String fieldName, Map<String, Object> params, SearchLookup searchLookup, LeafReaderContext ctx) {
+        super(fieldName, params, searchLookup, ctx);
+    }
+
+    protected void emit(double lat, double lon) {
+        int latitudeEncoded = encodeLatitude(lat);
+        int longitudeEncoded = encodeLongitude(lon);
+        emit(Long.valueOf((((long) latitudeEncoded) << 32) | (longitudeEncoded & 0xFFFFFFFFL)));
+    }
+
+    public static class Emit {
+        private final GeoPointFieldScript script;
+
+        public Emit(GeoPointFieldScript script) {
+            this.script = script;
+        }
+
+        public void emit(double lat, double lon) {
+            script.emit(lat, lon);
+        }
+    }
+}