diff --git a/modules/lang-painless/src/main/java/org/elasticsearch/painless/action/PainlessExecuteAction.java b/modules/lang-painless/src/main/java/org/elasticsearch/painless/action/PainlessExecuteAction.java index 783abf5551c43..37a82709d8007 100644 --- a/modules/lang-painless/src/main/java/org/elasticsearch/painless/action/PainlessExecuteAction.java +++ b/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( + GeoPointFieldScript.CONTEXT.name, + request.getScript().getParams(), + context.lookup(), + OnScriptError.FAIL + ); + GeometryFieldScript geoPointFieldScript = leafFactory.newInstance(leafReaderContext); + List geometries = new ArrayList<>(); + geoPointFieldScript.runForDoc(0, geometries::add); + // convert geometries to the standard format of the fields api + Function, List> 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); diff --git a/modules/lang-painless/src/main/resources/org/elasticsearch/painless/org.elasticsearch.script.geometry_field.txt b/modules/lang-painless/src/main/resources/org/elasticsearch/painless/org.elasticsearch.script.geometry_field.txt new file mode 100644 index 0000000000000..68bcbf922e869 --- /dev/null +++ b/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 +} diff --git a/modules/lang-painless/src/test/java/org/elasticsearch/painless/action/PainlessExecuteApiTests.java b/modules/lang-painless/src/test/java/org/elasticsearch/painless/action/PainlessExecuteApiTests.java index 86f2b5d83ca0b..c2356a3d8e138 100644 --- a/modules/lang-painless/src/test/java/org/elasticsearch/painless/action/PainlessExecuteApiTests.java +++ b/modules/lang-painless/src/test/java/org/elasticsearch/painless/action/PainlessExecuteApiTests.java @@ -259,6 +259,50 @@ public void testGeoPointFieldExecutionContext() throws IOException { assertEquals("Point", points.get(1).get("type")); } + @SuppressWarnings("unchecked") + public void testGeoShapeFieldExecutionContext() 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> points = (List>) response.getResult(); + assertEquals(40.0, (double) ((List) points.get(0).get("coordinates")).get(0), 0.00001); + assertEquals(30.0, (double) ((List) points.get(0).get("coordinates")).get(1), 0.00001); + assertEquals("Point", points.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(78.96, 12.12); emit(13.45, 56.78);", emptyMap()), + "geo_point_field", + contextSetup + ); + response = innerShardOperation(request, scriptService, indexService); + points = (List>) response.getResult(); + assertEquals(12.12, (double) ((List) points.get(0).get("coordinates")).get(0), 0.00001); + assertEquals(78.96, (double) ((List) points.get(0).get("coordinates")).get(1), 0.00001); + assertEquals("Point", points.get(0).get("type")); + assertEquals(56.78, (double) ((List) points.get(1).get("coordinates")).get(0), 0.00001); + assertEquals(13.45, (double) ((List) points.get(1).get("coordinates")).get(1), 0.00001); + assertEquals("Point", points.get(1).get("type")); + } + public void testIpFieldExecutionContext() throws IOException { ScriptService scriptService = getInstanceFromNode(ScriptService.class); IndexService indexService = createIndex("index", Settings.EMPTY, "doc", "test_ip", "type=ip"); diff --git a/server/src/main/java/org/elasticsearch/index/mapper/AbstractShapeGeometryFieldMapper.java b/server/src/main/java/org/elasticsearch/index/mapper/AbstractShapeGeometryFieldMapper.java index 43032c9ce32c3..22b75c8262193 100644 --- a/server/src/main/java/org/elasticsearch/index/mapper/AbstractShapeGeometryFieldMapper.java +++ b/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 coerce, + Explicit orientation, + CopyTo copyTo, + Parser parser, + OnScriptError onScriptError + ) { + super(simpleName, mappedFieldType, multiFields, copyTo, parser, onScriptError); + this.coerce = coerce; + this.orientation = orientation; + } + public boolean coerce() { return coerce.value(); } diff --git a/server/src/main/java/org/elasticsearch/script/GeometryFieldScript.java b/server/src/main/java/org/elasticsearch/script/GeometryFieldScript.java new file mode 100644 index 0000000000000..2e0fc693463f8 --- /dev/null +++ b/server/src/main/java/org/elasticsearch/script/GeometryFieldScript.java @@ -0,0 +1,142 @@ +/* + * 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 CONTEXT = newContext("geometry_field", Factory.class); + + public static final Factory PARSE_FROM_SOURCE = new Factory() { + @Override + public LeafFactory newFactory(String field, Map 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 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 params, SearchLookup searchLookup, OnScriptError onScriptError); + } + + public interface LeafFactory { + GeometryFieldScript newInstance(LeafReaderContext ctx); + } + + private final List geometries = new ArrayList<>(); + + private final GeometryParser geometryParser; + + public GeometryFieldScript( + String fieldName, + Map 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 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() { + 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); + } + } +} diff --git a/server/src/main/java/org/elasticsearch/script/ScriptModule.java b/server/src/main/java/org/elasticsearch/script/ScriptModule.java index 85afb96225254..6eb3bdfba32fd 100644 --- a/server/src/main/java/org/elasticsearch/script/ScriptModule.java +++ b/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 ); diff --git a/server/src/test/java/org/elasticsearch/index/mapper/GeometryFieldScriptTests.java b/server/src/test/java/org/elasticsearch/index/mapper/GeometryFieldScriptTests.java new file mode 100644 index 0000000000000..ab87ca40ce93a --- /dev/null +++ b/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 { + 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 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]") + ); + } + } + } +} diff --git a/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/SpatialPlugin.java b/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/SpatialPlugin.java index 66281cd21856b..7f171230e7628 100644 --- a/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/SpatialPlugin.java +++ b/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/SpatialPlugin.java @@ -12,6 +12,7 @@ import org.elasticsearch.common.geo.GeoFormatterFactory; import org.elasticsearch.geometry.Geometry; import org.elasticsearch.index.mapper.Mapper; +import org.elasticsearch.index.mapper.RuntimeField; import org.elasticsearch.ingest.Processor; import org.elasticsearch.license.License; import org.elasticsearch.license.LicenseUtils; @@ -45,6 +46,7 @@ import org.elasticsearch.xpack.spatial.action.SpatialStatsTransportAction; import org.elasticsearch.xpack.spatial.action.SpatialUsageTransportAction; import org.elasticsearch.xpack.spatial.common.CartesianBoundingBox; +import org.elasticsearch.xpack.spatial.index.mapper.GeoShapeScriptFieldType; import org.elasticsearch.xpack.spatial.index.mapper.GeoShapeWithDocValuesFieldMapper; import org.elasticsearch.xpack.spatial.index.mapper.PointFieldMapper; import org.elasticsearch.xpack.spatial.index.mapper.ShapeFieldMapper; @@ -142,6 +144,11 @@ public Map getMappers() { ); } + @Override + public Map getRuntimeFields() { + return Map.of(GeoShapeWithDocValuesFieldMapper.CONTENT_TYPE, GeoShapeScriptFieldType.typeParser(geoFormatterFactory.get())); + } + @Override public List> getQueries() { return List.of( diff --git a/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/fielddata/GeoShapeScriptDocValues.java b/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/fielddata/GeoShapeScriptDocValues.java new file mode 100644 index 0000000000000..4840ca4034fd0 --- /dev/null +++ b/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/fielddata/GeoShapeScriptDocValues.java @@ -0,0 +1,58 @@ +/* + * 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.index.fielddata; + +import org.apache.lucene.index.IndexableField; +import org.elasticsearch.common.geo.Orientation; +import org.elasticsearch.geometry.Geometry; +import org.elasticsearch.index.mapper.GeoShapeIndexer; +import org.elasticsearch.script.GeometryFieldScript; +import org.elasticsearch.search.aggregations.support.ValuesSourceType; +import org.elasticsearch.xpack.spatial.index.mapper.BinaryShapeDocValuesField; +import org.elasticsearch.xpack.spatial.search.aggregations.support.GeoShapeValuesSourceType; + +import java.io.IOException; +import java.util.List; + +/** + * Similarly to what {@link BinaryShapeDocValuesField} does, it encodes the shapes using the {@link GeometryDocValueWriter}. + */ +public final class GeoShapeScriptDocValues extends GeoShapeValues { + private final GeometryFieldScript script; + private final GeoShapeValues.GeoShapeValue geoShapeValue = new GeoShapeValues.GeoShapeValue(); + private final GeoShapeIndexer indexer; + + public GeoShapeScriptDocValues(GeometryFieldScript script, String fieldName) { + this.script = script; + indexer = new GeoShapeIndexer(Orientation.CCW, fieldName); + } + + @Override + public boolean advanceExact(int docId) { + script.runForDoc(docId); + return script.count() != 0; + } + + @Override + public ValuesSourceType valuesSourceType() { + return GeoShapeValuesSourceType.instance(); + } + + @Override + public GeoShapeValue value() throws IOException { + final Geometry geometry = script.geometry(); + if (geometry == null) { + return null; + } + final List fields = indexer.getIndexableFields(geometry); + final CentroidCalculator centroidCalculator = new CentroidCalculator(); + centroidCalculator.add(geometry); + geoShapeValue.reset(GeometryDocValueWriter.write(fields, CoordinateEncoder.GEO, centroidCalculator)); + return geoShapeValue; + } +} diff --git a/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/fielddata/plain/GeoShapeScriptFieldData.java b/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/fielddata/plain/GeoShapeScriptFieldData.java new file mode 100644 index 0000000000000..4790924412cd7 --- /dev/null +++ b/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/fielddata/plain/GeoShapeScriptFieldData.java @@ -0,0 +1,79 @@ +/* + * 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.index.fielddata.plain; + +import org.apache.lucene.index.LeafReaderContext; +import org.elasticsearch.index.fielddata.IndexFieldData; +import org.elasticsearch.index.fielddata.IndexFieldDataCache; +import org.elasticsearch.indices.breaker.CircuitBreakerService; +import org.elasticsearch.script.GeometryFieldScript; +import org.elasticsearch.script.field.ToScriptFieldFactory; +import org.elasticsearch.xpack.spatial.index.fielddata.GeoShapeScriptDocValues; +import org.elasticsearch.xpack.spatial.index.fielddata.GeoShapeValues; +import org.elasticsearch.xpack.spatial.index.fielddata.LeafShapeFieldData; +import org.elasticsearch.xpack.spatial.search.aggregations.support.GeoShapeValuesSourceType; + +public final class GeoShapeScriptFieldData extends AbstractShapeIndexFieldData { + public static class Builder implements IndexFieldData.Builder { + private final String name; + private final GeometryFieldScript.LeafFactory leafFactory; + private final ToScriptFieldFactory toScriptFieldFactory; + + public Builder( + String name, + GeometryFieldScript.LeafFactory leafFactory, + ToScriptFieldFactory toScriptFieldFactory + ) { + this.name = name; + this.leafFactory = leafFactory; + this.toScriptFieldFactory = toScriptFieldFactory; + } + + @Override + public GeoShapeScriptFieldData build(IndexFieldDataCache cache, CircuitBreakerService breakerService) { + return new GeoShapeScriptFieldData(name, leafFactory, toScriptFieldFactory); + } + } + + private final GeometryFieldScript.LeafFactory leafFactory; + + private GeoShapeScriptFieldData( + String fieldName, + GeometryFieldScript.LeafFactory leafFactory, + ToScriptFieldFactory toScriptFieldFactory + ) { + super(fieldName, GeoShapeValuesSourceType.instance(), toScriptFieldFactory); + this.leafFactory = leafFactory; + } + + @Override + protected IllegalArgumentException sortException() { + throw new IllegalArgumentException("can't sort on geo_shape field"); + } + + @Override + public LeafShapeFieldData load(LeafReaderContext context) { + final GeometryFieldScript script = leafFactory.newInstance(context); + return new LeafShapeFieldData<>(toScriptFieldFactory) { + @Override + public GeoShapeValues getShapeValues() { + return new GeoShapeScriptDocValues(script, fieldName); + } + + @Override + public long ramBytesUsed() { + return 0; + } + + @Override + public void close() { + + } + }; + } +} diff --git a/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/mapper/GeoShapeScriptFieldType.java b/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/mapper/GeoShapeScriptFieldType.java new file mode 100644 index 0000000000000..358f6768132b9 --- /dev/null +++ b/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/mapper/GeoShapeScriptFieldType.java @@ -0,0 +1,166 @@ +/* + * 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.index.mapper; + +import org.apache.lucene.geo.LatLonGeometry; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.search.Query; +import org.elasticsearch.common.geo.GeoFormatterFactory; +import org.elasticsearch.common.geo.GeometryFormatterFactory; +import org.elasticsearch.common.geo.ShapeRelation; +import org.elasticsearch.common.time.DateMathParser; +import org.elasticsearch.geometry.Geometry; +import org.elasticsearch.index.fielddata.FieldDataContext; +import org.elasticsearch.index.mapper.AbstractScriptFieldType; +import org.elasticsearch.index.mapper.GeoShapeQueryable; +import org.elasticsearch.index.mapper.OnScriptError; +import org.elasticsearch.index.mapper.RuntimeField; +import org.elasticsearch.index.mapper.ValueFetcher; +import org.elasticsearch.index.query.SearchExecutionContext; +import org.elasticsearch.script.CompositeFieldScript; +import org.elasticsearch.script.GeometryFieldScript; +import org.elasticsearch.script.Script; +import org.elasticsearch.search.fetch.StoredFieldsSpec; +import org.elasticsearch.search.lookup.SearchLookup; +import org.elasticsearch.search.lookup.Source; +import org.elasticsearch.xpack.spatial.index.fielddata.plain.GeoShapeScriptFieldData; +import org.elasticsearch.xpack.spatial.search.runtime.GeoShapeScriptFieldExistsQuery; +import org.elasticsearch.xpack.spatial.search.runtime.GeoShapeScriptFieldGeoShapeQuery; + +import java.io.IOException; +import java.time.ZoneId; +import java.util.List; +import java.util.Map; +import java.util.function.Function; + +public final class GeoShapeScriptFieldType extends AbstractScriptFieldType implements GeoShapeQueryable { + + public static RuntimeField.Parser typeParser(GeoFormatterFactory geoFormatterFactory) { + return new RuntimeField.Parser(name -> new Builder<>(name, GeometryFieldScript.CONTEXT) { + @Override + protected AbstractScriptFieldType createFieldType( + String name, + GeometryFieldScript.Factory factory, + Script script, + Map meta, + OnScriptError onScriptError + ) { + return new GeoShapeScriptFieldType(name, factory, getScript(), meta(), onScriptError, geoFormatterFactory); + } + + @Override + protected GeometryFieldScript.Factory getParseFromSourceFactory() { + return GeometryFieldScript.PARSE_FROM_SOURCE; + } + + @Override + protected GeometryFieldScript.Factory getCompositeLeafFactory( + Function parentScriptFactory + ) { + return GeometryFieldScript.leafAdapter(parentScriptFactory); + } + }); + } + + private final GeoFormatterFactory geoFormatterFactory; + + GeoShapeScriptFieldType( + String name, + GeometryFieldScript.Factory scriptFactory, + Script script, + Map meta, + OnScriptError onScriptError, + GeoFormatterFactory geoFormatterFactory + ) { + super( + name, + searchLookup -> scriptFactory.newFactory(name, script.getParams(), searchLookup, onScriptError), + script, + scriptFactory.isResultDeterministic(), + meta + ); + this.geoFormatterFactory = geoFormatterFactory; + } + + @Override + public String typeName() { + return GeoShapeWithDocValuesFieldMapper.CONTENT_TYPE; + } + + @Override + protected Query rangeQuery( + Object lowerTerm, + Object upperTerm, + boolean includeLower, + boolean includeUpper, + ZoneId timeZone, + DateMathParser parser, + SearchExecutionContext context + ) { + throw new IllegalArgumentException("Runtime field [" + name() + "] of type [" + typeName() + "] does not support range queries"); + } + + @Override + public Query termQuery(Object value, SearchExecutionContext context) { + throw new IllegalArgumentException( + "Geometry fields do not support exact searching, use dedicated geometry queries instead: [" + name() + "]" + ); + } + + @Override + public GeoShapeScriptFieldData.Builder fielddataBuilder(FieldDataContext fieldDataContext) { + return new GeoShapeScriptFieldData.Builder( + name(), + leafFactory(fieldDataContext.lookupSupplier().get()), + GeoShapeWithDocValuesFieldMapper.GeoShapeDocValuesField::new + ); + } + + @Override + public Query existsQuery(SearchExecutionContext context) { + applyScriptContext(context); + return new GeoShapeScriptFieldExistsQuery(script, leafFactory(context), name()); + } + + @Override + public Query geoShapeQuery(SearchExecutionContext context, String fieldName, ShapeRelation relation, LatLonGeometry... geometries) { + return new GeoShapeScriptFieldGeoShapeQuery(script, leafFactory(context), fieldName, relation, geometries); + } + + @Override + public ValueFetcher valueFetcher(SearchExecutionContext context, String format) { + GeometryFieldScript.LeafFactory leafFactory = leafFactory(context.lookup()); + + Function, List> formatter = geoFormatterFactory.getFormatter( + format != null ? format : GeometryFormatterFactory.GEOJSON, + Function.identity() + ); + return new ValueFetcher() { + private GeometryFieldScript script; + + @Override + public void setNextReader(LeafReaderContext context) { + script = leafFactory.newInstance(context); + } + + @Override + public List fetchValues(Source source, int doc, List ignoredValues) throws IOException { + script.runForDoc(doc); + if (script.count() == 0) { + return List.of(); + } + return formatter.apply(List.of(script.geometry())); + } + + @Override + public StoredFieldsSpec storedFieldsSpec() { + return StoredFieldsSpec.NEEDS_SOURCE; + } + }; + } +} diff --git a/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/mapper/GeoShapeWithDocValuesFieldMapper.java b/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/mapper/GeoShapeWithDocValuesFieldMapper.java index 55929a1c1b83e..b024e98a654c4 100644 --- a/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/mapper/GeoShapeWithDocValuesFieldMapper.java +++ b/x-pack/plugin/spatial/src/main/java/org/elasticsearch/xpack/spatial/index/mapper/GeoShapeWithDocValuesFieldMapper.java @@ -11,6 +11,7 @@ import org.apache.lucene.document.StoredField; import org.apache.lucene.geo.LatLonGeometry; import org.apache.lucene.index.IndexableField; +import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.search.IndexOrDocValuesQuery; import org.apache.lucene.search.Query; import org.apache.lucene.util.BytesRef; @@ -43,14 +44,20 @@ import org.elasticsearch.index.mapper.MapperBuilderContext; import org.elasticsearch.index.mapper.MapperParsingException; import org.elasticsearch.index.mapper.MappingParserContext; +import org.elasticsearch.index.mapper.OnScriptError; import org.elasticsearch.index.mapper.StoredValueFetcher; import org.elasticsearch.index.mapper.ValueFetcher; import org.elasticsearch.index.query.QueryShardException; import org.elasticsearch.index.query.SearchExecutionContext; import org.elasticsearch.legacygeo.mapper.LegacyGeoShapeFieldMapper; +import org.elasticsearch.script.GeometryFieldScript; +import org.elasticsearch.script.Script; +import org.elasticsearch.script.ScriptCompiler; import org.elasticsearch.script.field.AbstractScriptFieldFactory; import org.elasticsearch.script.field.DocValuesScriptFieldFactory; import org.elasticsearch.script.field.Field; +import org.elasticsearch.search.lookup.FieldValues; +import org.elasticsearch.search.lookup.SearchLookup; import org.elasticsearch.xpack.spatial.index.fielddata.CoordinateEncoder; import org.elasticsearch.xpack.spatial.index.fielddata.GeoShapeValues; import org.elasticsearch.xpack.spatial.index.fielddata.plain.AbstractAtomicGeoShapeShapeFieldData; @@ -110,25 +117,29 @@ public static class Builder extends FieldMapper.Builder { final Parameter> ignoreZValue = ignoreZValueParam(m -> builder(m).ignoreZValue.get()); final Parameter> coerce; final Parameter> orientation = orientationParam(m -> builder(m).orientation.get()); - + private final Parameter