Geometry formatters should be applied to the full array (#75177)

Changes the way we fetch values from source in geo fields so we can format all data in one go.
elastic · Jul 14, 2021 · dcd3877 · dcd3877
1 parent 88d6079
commit dcd3877
Show file tree

Hide file tree

Showing 12 changed files with 110 additions and 123 deletions.
diff --git a/.../lang-painless/src/main/java/org/elasticsearch/painless/action/PainlessExecuteAction.java b/.../lang-painless/src/main/java/org/elasticsearch/painless/action/PainlessExecuteAction.java
@@ -34,9 +34,9 @@
 import org.elasticsearch.cluster.routing.ShardsIterator;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.common.CheckedBiFunction;
+import org.elasticsearch.common.geo.GeometryFormatterFactory;
 import org.elasticsearch.common.xcontent.ParseField;
 import org.elasticsearch.common.bytes.BytesReference;
-import org.elasticsearch.common.geo.GeoFormatterFactory;
 import org.elasticsearch.common.geo.GeoPoint;
 import org.elasticsearch.common.inject.Inject;
 import org.elasticsearch.common.io.stream.StreamInput;
@@ -51,7 +51,6 @@
 import org.elasticsearch.common.xcontent.XContentHelper;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.common.xcontent.XContentType;
-import org.elasticsearch.geometry.Geometry;
 import org.elasticsearch.geometry.Point;
 import org.elasticsearch.index.Index;
 import org.elasticsearch.index.IndexService;
@@ -588,12 +587,9 @@ static Response innerShardOperation(Request request, ScriptService scriptService
                     List<GeoPoint> points = new ArrayList<>();
                     geoPointFieldScript.runGeoPointForDoc(0, gp -> points.add(new GeoPoint(gp)));
                     // convert geo points to the standard format of the fields api
-                    Function<Geometry, Object> format = GeoFormatterFactory.getFormatter(GeoFormatterFactory.GEOJSON);
-                    List<Object> objects = new ArrayList<>();
-                    for (GeoPoint gp : points) {
-                        objects.add(format.apply(new Point(gp.getLon(), gp.getLat())));
-                    }
-                    return new Response(objects);
+                    Function<List<GeoPoint>, List<Object>> format =
+                        GeometryFormatterFactory.getFormatter(GeometryFormatterFactory.GEOJSON, p -> new Point(p.lon(), p.lat()));
+                    return new Response(format.apply(points));
                 }, indexService);
             } else if (scriptContext == IpFieldScript.CONTEXT) {
                 return prepareRamIndex(request, (context, leafReaderContext) -> {

diff --git a/server/src/main/java/org/elasticsearch/common/geo/GeoFormatterFactory.java b/server/src/main/java/org/elasticsearch/common/geo/GeoFormatterFactory.java
diff --git a/server/src/main/java/org/elasticsearch/common/geo/GeometryFormatterFactory.java b/server/src/main/java/org/elasticsearch/common/geo/GeometryFormatterFactory.java
@@ -0,0 +1,46 @@
+/*
+ * 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.common.geo;
+
+import org.elasticsearch.geometry.Geometry;
+import org.elasticsearch.geometry.utils.WellKnownText;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.function.Function;
+
+/**
+ * Output formatters supported by geometry fields.
+ */
+public class GeometryFormatterFactory {
+
+    public static final String GEOJSON = "geojson";
+    public static final String WKT = "wkt";
+
+    /**
+     * Returns a formatter by name
+     */
+    public static <T> Function<List<T>, List<Object>> getFormatter(String name, Function<T, Geometry> toGeometry) {
+        switch (name) {
+            case GEOJSON:
+                return geometries -> {
+                    final List<Object> objects = new ArrayList<>(geometries.size());
+                    geometries.forEach((shape) -> objects.add(GeoJson.toMap(toGeometry.apply(shape))));
+                    return objects;
+                };
+            case WKT:
+                return geometries -> {
+                    final List<Object> objects = new ArrayList<>(geometries.size());
+                    geometries.forEach((shape) -> objects.add(WellKnownText.toWKT(toGeometry.apply(shape))));
+                    return objects;
+                };
+            default:  throw new IllegalArgumentException("Unrecognized geometry format [" + name + "].");
+        }
+    }
+}
diff --git a/server/src/main/java/org/elasticsearch/index/mapper/AbstractGeometryFieldMapper.java b/server/src/main/java/org/elasticsearch/index/mapper/AbstractGeometryFieldMapper.java
@@ -9,7 +9,7 @@
 
 import org.apache.lucene.search.Query;
 import org.elasticsearch.common.Explicit;
-import org.elasticsearch.common.geo.GeoFormatterFactory;
+import org.elasticsearch.common.geo.GeometryFormatterFactory;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.common.xcontent.support.MapXContentParser;
 import org.elasticsearch.core.CheckedConsumer;
@@ -52,13 +52,14 @@ public abstract void parse(
             CheckedConsumer<T, IOException> consumer,
             Consumer<Exception> onMalformed) throws IOException;
 
-        private void fetchFromSource(Object sourceMap, Consumer<Object> consumer, Function<T, Object> formatter) {
+        private void fetchFromSource(Object sourceMap, Consumer<T> consumer) {
             try (XContentParser parser = MapXContentParser.wrapObject(sourceMap)) {
-                parse(parser, v -> consumer.accept(formatter.apply(v)), e -> {}); /* ignore malformed */
+                parse(parser, v -> consumer.accept(v), e -> {}); /* ignore malformed */
             } catch (IOException e) {
                 throw new UncheckedIOException(e);
             }
         }
+
     }
 
     public abstract static class AbstractGeometryFieldType<T> extends MappedFieldType {
@@ -80,17 +81,17 @@ public final Query termQuery(Object value, SearchExecutionContext context) {
         /**
          * Gets the formatter by name.
          */
-        protected abstract Function<T, Object> getFormatter(String format);
+        protected abstract Function<List<T>, List<Object>> getFormatter(String format);
 
         @Override
         public ValueFetcher valueFetcher(SearchExecutionContext context, String format) {
-            Function<T, Object> formatter = getFormatter(format != null ? format : GeoFormatterFactory.GEOJSON);
+            Function<List<T>, List<Object>> formatter = getFormatter(format != null ? format : GeometryFormatterFactory.GEOJSON);
             return new ArraySourceValueFetcher(name(), context) {
                 @Override
                 protected Object parseSourceValue(Object value) {
-                    List<Object> values = new ArrayList<>();
-                    geometryParser.fetchFromSource(value, values::add, formatter);
-                    return values;
+                    final List<T> values = new ArrayList<>();
+                    geometryParser.fetchFromSource(value, values::add);
+                    return formatter.apply(values);
                 }
             };
         }

diff --git a/server/src/main/java/org/elasticsearch/index/mapper/GeoPointFieldMapper.java b/server/src/main/java/org/elasticsearch/index/mapper/GeoPointFieldMapper.java
@@ -19,10 +19,10 @@
 import org.elasticsearch.ElasticsearchParseException;
 import org.elasticsearch.common.CheckedBiFunction;
 import org.elasticsearch.common.Explicit;
-import org.elasticsearch.common.geo.GeoFormatterFactory;
 import org.elasticsearch.common.geo.GeoPoint;
 import org.elasticsearch.common.geo.GeoShapeUtils;
 import org.elasticsearch.common.geo.GeoUtils;
+import org.elasticsearch.common.geo.GeometryFormatterFactory;
 import org.elasticsearch.common.geo.ShapeRelation;
 import org.elasticsearch.common.unit.DistanceUnit;
 import org.elasticsearch.common.xcontent.XContentParser;
@@ -242,18 +242,17 @@ public String typeName() {
         }
 
         @Override
-        protected Function<GeoPoint, Object> getFormatter(String format) {
-            Function<Geometry, Object> formatter = GeoFormatterFactory.getFormatter(format);
-            return (point) -> formatter.apply(new Point(point.lon(), point.lat()));
+        protected  Function<List<GeoPoint>, List<Object>> getFormatter(String format) {
+            return GeometryFormatterFactory.getFormatter(format, p -> new Point(p.lon(), p.lat()));
         }
 
         @Override
         public ValueFetcher valueFetcher(SearchExecutionContext context, String format) {
             if (scriptValues == null) {
                 return super.valueFetcher(context, format);
             }
-            Function<GeoPoint, Object> formatter = getFormatter(format != null ? format : GeoFormatterFactory.GEOJSON);
-            return FieldValues.valueFetcher(scriptValues, v -> formatter.apply((GeoPoint) v), context);
+            Function<List<GeoPoint>, List<Object>> formatter = getFormatter(format != null ? format : GeometryFormatterFactory.GEOJSON);
+            return FieldValues.valueListFetcher(scriptValues, formatter, context);
         }
 
         @Override

diff --git a/server/src/main/java/org/elasticsearch/index/mapper/GeoShapeFieldMapper.java b/server/src/main/java/org/elasticsearch/index/mapper/GeoShapeFieldMapper.java
@@ -13,8 +13,8 @@
 import org.apache.lucene.search.Query;
 import org.elasticsearch.Version;
 import org.elasticsearch.common.Explicit;
-import org.elasticsearch.common.geo.GeoFormatterFactory;
 import org.elasticsearch.common.geo.GeoShapeUtils;
+import org.elasticsearch.common.geo.GeometryFormatterFactory;
 import org.elasticsearch.common.geo.GeometryParser;
 import org.elasticsearch.common.geo.Orientation;
 import org.elasticsearch.common.geo.ShapeRelation;
@@ -141,8 +141,8 @@ public Query geoShapeQuery(Geometry shape, String fieldName, ShapeRelation relat
         }
 
         @Override
-        protected Function<Geometry, Object> getFormatter(String format) {
-            return GeoFormatterFactory.getFormatter(format);
+        protected Function<List<Geometry>, List<Object>> getFormatter(String format) {
+            return GeometryFormatterFactory.getFormatter(format, Function.identity());
         }
     }
 

diff --git a/server/src/main/java/org/elasticsearch/index/mapper/LegacyGeoShapeFieldMapper.java b/server/src/main/java/org/elasticsearch/index/mapper/LegacyGeoShapeFieldMapper.java
@@ -18,8 +18,8 @@
 import org.elasticsearch.ElasticsearchParseException;
 import org.elasticsearch.Version;
 import org.elasticsearch.common.Explicit;
-import org.elasticsearch.common.geo.GeoFormatterFactory;
 import org.elasticsearch.common.geo.GeoUtils;
+import org.elasticsearch.common.geo.GeometryFormatterFactory;
 import org.elasticsearch.common.geo.Orientation;
 import org.elasticsearch.common.geo.ShapeRelation;
 import org.elasticsearch.common.geo.ShapesAvailability;
@@ -454,9 +454,8 @@ public PrefixTreeStrategy resolvePrefixTreeStrategy(String strategyName) {
         }
 
         @Override
-        protected Function<ShapeBuilder<?, ?, ?>, Object> getFormatter(String format) {
-            Function<Geometry, Object> formatter = GeoFormatterFactory.getFormatter(format);
-            return (g) -> formatter.apply(g.buildGeometry());
+        protected Function<List<ShapeBuilder<?, ?, ?>>, List<Object>> getFormatter(String format) {
+            return GeometryFormatterFactory.getFormatter(format, ShapeBuilder::buildGeometry);
         }
     }
 

diff --git a/server/src/main/java/org/elasticsearch/search/lookup/FieldValues.java b/server/src/main/java/org/elasticsearch/search/lookup/FieldValues.java
@@ -70,4 +70,35 @@ public List<Object> fetchValues(SourceLookup lookup)  {
             }
         };
     }
+
+    /**
+     * Creates a {@link ValueFetcher} that fetches values from a {@link FieldValues} instance
+     * @param fieldValues the source of the values
+     * @param formatter   a function to format the list values
+     * @param context the search execution context
+     * @return the value fetcher
+     */
+    static <T> ValueFetcher valueListFetcher(FieldValues<T> fieldValues, Function<List<T>, List<Object>> formatter,
+                                             SearchExecutionContext context) {
+        return new ValueFetcher() {
+            LeafReaderContext ctx;
+
+            @Override
+            public void setNextReader(LeafReaderContext context) {
+                this.ctx = context;
+            }
+
+            @Override
+            public List<Object> fetchValues(SourceLookup lookup)  {
+                List<T> values = new ArrayList<>();
+                try {
+                    fieldValues.valuesForDoc(context.lookup(), ctx, lookup.docId(), v -> values.add(v));
+                } catch (Exception e) {
+                    // ignore errors - if they exist here then they existed at index time
+                    // and so on_script_error must have been set to `ignore`
+                }
+                return formatter.apply(values);
+            }
+        };
+    }
 }
diff --git a/...atial/src/main/java/org/elasticsearch/xpack/spatial/common/CartesianFormatterFactory.java b/...atial/src/main/java/org/elasticsearch/xpack/spatial/common/CartesianFormatterFactory.java
diff --git a/...n/java/org/elasticsearch/xpack/spatial/index/mapper/GeoShapeWithDocValuesFieldMapper.java b/...n/java/org/elasticsearch/xpack/spatial/index/mapper/GeoShapeWithDocValuesFieldMapper.java
@@ -15,8 +15,8 @@
 import org.apache.lucene.search.Query;
 import org.elasticsearch.Version;
 import org.elasticsearch.common.Explicit;
-import org.elasticsearch.common.geo.GeoFormatterFactory;
 import org.elasticsearch.common.geo.GeoShapeUtils;
+import org.elasticsearch.common.geo.GeometryFormatterFactory;
 import org.elasticsearch.common.geo.GeometryParser;
 import org.elasticsearch.common.geo.Orientation;
 import org.elasticsearch.common.geo.ShapeRelation;
@@ -173,8 +173,8 @@ public Query geoShapeQuery(Geometry shape, String fieldName, ShapeRelation relat
         }
 
         @Override
-        protected Function<Geometry, Object> getFormatter(String format) {
-            return GeoFormatterFactory.getFormatter(format);
+        protected Function<List<Geometry>, List<Object>> getFormatter(String format) {
+            return GeometryFormatterFactory.getFormatter(format, Function.identity());
         }
     }
 

diff --git a/.../spatial/src/main/java/org/elasticsearch/xpack/spatial/index/mapper/PointFieldMapper.java b/.../spatial/src/main/java/org/elasticsearch/xpack/spatial/index/mapper/PointFieldMapper.java
@@ -12,6 +12,7 @@
 import org.apache.lucene.search.Query;
 import org.elasticsearch.common.CheckedBiFunction;
 import org.elasticsearch.common.Explicit;
+import org.elasticsearch.common.geo.GeometryFormatterFactory;
 import org.elasticsearch.common.geo.ShapeRelation;
 import org.elasticsearch.common.logging.DeprecationCategory;
 import org.elasticsearch.common.logging.DeprecationLogger;
@@ -25,7 +26,6 @@
 import org.elasticsearch.index.mapper.GeoShapeFieldMapper;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.query.SearchExecutionContext;
-import org.elasticsearch.xpack.spatial.common.CartesianFormatterFactory;
 import org.elasticsearch.xpack.spatial.common.CartesianPoint;
 import org.elasticsearch.xpack.spatial.index.query.ShapeQueryPointProcessor;
 
@@ -179,9 +179,8 @@ public Query shapeQuery(Geometry shape, String fieldName, ShapeRelation relation
         }
 
         @Override
-        protected Function<CartesianPoint, Object> getFormatter(String format) {
-            Function<Geometry, Object> formatter = CartesianFormatterFactory.getFormatter(format);
-            return (point) -> formatter.apply(new Point(point.getX(), point.getY()));
+        protected Function<List<CartesianPoint>, List<Object>> getFormatter(String format) {
+           return GeometryFormatterFactory.getFormatter(format, p -> new Point(p.getX(), p.getY()));
         }
     }