Support for store parameter in geo_shape field (#94418)

This commit adds support for storing geo_shape fields as a lucene stored field. The geometry will be stored 
normalised in well-known binary format (WKB).

Author: iverase

docs/changelog/94418.yaml

@@ -0,0 +1,6 @@
+pr: 94418
+summary: Support for store parameter in `geo_shape` field
+area: Geo
+type: enhancement
+issues:
+ - 83655

libs/geo/src/main/java/org/elasticsearch/geometry/utils/GeometryValidator.java

@@ -15,6 +15,8 @@
  */
 public interface GeometryValidator {
 
+    GeometryValidator NOOP = (geometry) -> {};
+
     /**
      * Validates the geometry and throws IllegalArgumentException if the geometry is not valid
      */

libs/geo/src/main/java/org/elasticsearch/geometry/utils/WellKnownBinary.java

@@ -199,8 +199,15 @@ private static void writeDouble(ByteArrayOutputStream out, ByteBuffer scratch, d
     /**
      * Reads a {@link Geometry} from the given WKB byte array.
      */
-    public static Geometry fromWKB(GeometryValidator validator, boolean coerce, byte[] wkb) throws IOException {
-        final ByteBuffer byteBuffer = ByteBuffer.wrap(wkb);
+    public static Geometry fromWKB(GeometryValidator validator, boolean coerce, byte[] wkb) {
+        return fromWKB(validator, coerce, wkb, 0, wkb.length);
+    }
+
+    /**
+     * Reads a {@link Geometry} from the given WKB byte array with offset.
+     */
+    public static Geometry fromWKB(GeometryValidator validator, boolean coerce, byte[] wkb, int offset, int length) {
+        final ByteBuffer byteBuffer = ByteBuffer.wrap(wkb, offset, length);
         final Geometry geometry = parseGeometry(byteBuffer, coerce);
         validator.validate(geometry);
         return geometry;

libs/geo/src/test/java/org/elasticsearch/geometry/utils/WKBTests.java

@@ -141,8 +141,18 @@ private ByteOrder randomByteOrder() {
     private void assertWKB(Geometry geometry) throws IOException {
         final boolean hasZ = geometry.hasZ();
         final ByteOrder byteOrder = randomByteOrder();
-        byte[] b = WellKnownBinary.toWKB(geometry, byteOrder);
-        assertEquals(geometry, WellKnownBinary.fromWKB(StandardValidator.instance(hasZ), randomBoolean(), b));
+        final byte[] b = WellKnownBinary.toWKB(geometry, byteOrder);
+        if (randomBoolean()) {
+            // add padding to the byte array
+            final int extraBytes = randomIntBetween(1, 500);
+            final byte[] oversizeB = new byte[b.length + extraBytes];
+            random().nextBytes(oversizeB);
+            final int offset = randomInt(extraBytes);
+            System.arraycopy(b, 0, oversizeB, offset, b.length);
+            assertEquals(geometry, WellKnownBinary.fromWKB(StandardValidator.instance(hasZ), randomBoolean(), oversizeB, offset, b.length));
+        } else {
+            assertEquals(geometry, WellKnownBinary.fromWKB(StandardValidator.instance(hasZ), randomBoolean(), b));
+        }
     }
 
 }

server/src/main/java/org/elasticsearch/index/mapper/GeoShapeIndexer.java

@@ -46,14 +46,23 @@ public GeoShapeIndexer(Orientation orientation, String name) {
         this.name = name;
     }
 
+    @Override
     public List<IndexableField> indexShape(Geometry geometry) {
         if (geometry == null) {
             return Collections.emptyList();
         }
-        if (GeometryNormalizer.needsNormalize(orientation, geometry)) {
-            geometry = GeometryNormalizer.apply(orientation, geometry);
-        }
-        LuceneGeometryIndexer visitor = new LuceneGeometryIndexer(name);
+        return getIndexableFields(normalize(geometry));
+    }
+
+    /** Normalise the geometry, that is make sure latitude and longitude are between expected values
+     * and split geometries across the dateline when needed */
+    public Geometry normalize(Geometry geometry) {
+        return GeometryNormalizer.needsNormalize(orientation, geometry) ? GeometryNormalizer.apply(orientation, geometry) : geometry;
+    }
+
+    /** Generates lucene indexable fields from a geometry. It expects geometries that have already been normalised. */
+    public List<IndexableField> getIndexableFields(Geometry geometry) {
+        final LuceneGeometryIndexer visitor = new LuceneGeometryIndexer(name);
         geometry.visit(visitor);
         return visitor.fields();
     }

server/src/main/java/org/elasticsearch/index/mapper/StoredValueFetcher.java

@@ -19,7 +19,7 @@
 /**
  * Value fetcher that loads from stored values.
  */
-public final class StoredValueFetcher implements ValueFetcher {
+public class StoredValueFetcher implements ValueFetcher {
 
     private final SearchLookup lookup;
     private LeafSearchLookup leafSearchLookup;
@@ -42,8 +42,15 @@ public List<Object> fetchValues(Source source, int doc, List<Object> ignoredValu
         if (values == null) {
             return values;
         } else {
-            return List.copyOf(values);
+            return parseStoredValues(List.copyOf(values));
         }
     }
 
+    /**
+     * Given the values stored in lucene, parse it into a standard format.
+     */
+    protected List<Object> parseStoredValues(List<Object> values) {
+        return values;
+    }
+
 }

test/framework/src/main/java/org/elasticsearch/index/mapper/FieldTypeTestCase.java

@@ -9,6 +9,9 @@
 
 import org.elasticsearch.Version;
 import org.elasticsearch.index.query.SearchExecutionContext;
+import org.elasticsearch.search.lookup.FieldLookup;
+import org.elasticsearch.search.lookup.LeafSearchLookup;
+import org.elasticsearch.search.lookup.LeafStoredFieldsLookup;
 import org.elasticsearch.search.lookup.SearchLookup;
 import org.elasticsearch.search.lookup.Source;
 import org.elasticsearch.test.ESTestCase;
@@ -68,4 +71,22 @@ public static List<?> fetchSourceValues(MappedFieldType fieldType, Object... val
         Source source = Source.fromMap(Collections.singletonMap(field, List.of(values)), randomFrom(XContentType.values()));
         return fetcher.fetchValues(source, -1, new ArrayList<>());
     }
+
+    public static List<?> fetchStoredValue(MappedFieldType fieldType, List<Object> storedValues, String format) throws IOException {
+        String field = fieldType.name();
+        SearchExecutionContext searchExecutionContext = mock(SearchExecutionContext.class);
+        SearchLookup searchLookup = mock(SearchLookup.class);
+        LeafSearchLookup leafSearchLookup = mock(LeafSearchLookup.class);
+        LeafStoredFieldsLookup leafStoredFieldsLookup = mock(LeafStoredFieldsLookup.class);
+        FieldLookup fieldLookup = mock(FieldLookup.class);
+        when(searchExecutionContext.lookup()).thenReturn(searchLookup);
+        when(searchLookup.getLeafSearchLookup(null)).thenReturn(leafSearchLookup);
+        when(leafSearchLookup.fields()).thenReturn(leafStoredFieldsLookup);
+        when(leafStoredFieldsLookup.get(field)).thenReturn(fieldLookup);
+        when(fieldLookup.getValues()).thenReturn(storedValues);
+
+        ValueFetcher fetcher = fieldType.valueFetcher(searchExecutionContext, format);
+        fetcher.setNextReader(null);
+        return fetcher.fetchValues(null, -1, new ArrayList<>());
+    }
 }

x-pack/plugin/spatial/src/internalClusterTest/java/org/elasticsearch/xpack/spatial/search/GeoShapeWithDocValuesIT.java

@@ -7,20 +7,27 @@
 
 package org.elasticsearch.xpack.spatial.search;
 
+import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.Version;
 import org.elasticsearch.action.search.SearchResponse;
 import org.elasticsearch.common.bytes.BytesReference;
 import org.elasticsearch.common.unit.DistanceUnit;
+import org.elasticsearch.geometry.Geometry;
 import org.elasticsearch.geometry.LinearRing;
 import org.elasticsearch.geometry.Polygon;
+import org.elasticsearch.geometry.ShapeType;
+import org.elasticsearch.geometry.utils.StandardValidator;
+import org.elasticsearch.geometry.utils.WellKnownBinary;
 import org.elasticsearch.index.mapper.MapperParsingException;
 import org.elasticsearch.percolator.PercolateQueryBuilder;
 import org.elasticsearch.percolator.PercolatorPlugin;
 import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.search.SearchHit;
 import org.elasticsearch.search.geo.GeoShapeIntegTestCase;
 import org.elasticsearch.search.sort.SortOrder;
 import org.elasticsearch.test.VersionUtils;
 import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xcontent.XContentFactory;
 import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.spatial.LocalStateSpatialPlugin;
 
@@ -36,6 +43,7 @@
 import static org.elasticsearch.xcontent.XContentFactory.jsonBuilder;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.instanceOf;
 
 public class GeoShapeWithDocValuesIT extends GeoShapeIntegTestCase {
 
@@ -158,4 +166,44 @@ public void testPercolatorGeoQueries() throws Exception {
         assertThat(response.getHits().getAt(1).getId(), equalTo("2"));
         assertThat(response.getHits().getAt(2).getId(), equalTo("3"));
     }
+
+    // make sure we store the normalised geometry
+    public void testStorePolygonDateLine() throws Exception {
+        XContentBuilder mapping = XContentFactory.jsonBuilder().startObject().startObject("properties").startObject("shape");
+        getGeoShapeMapping(mapping);
+        mapping.field("store", true);
+        mapping.endObject().endObject().endObject();
+
+        // create index
+        assertAcked(
+            client().admin()
+                .indices()
+                .prepareCreate("test")
+                .setSettings(settings(randomSupportedVersion()).build())
+                .setMapping(mapping)
+                .get()
+        );
+        ensureGreen();
+
+        String source = """
+            {
+              "shape": "POLYGON((179 0, -179 0, -179 2, 179 2, 179 0))"
+            }""";
+
+        indexRandom(true, client().prepareIndex("test").setId("0").setSource(source, XContentType.JSON));
+
+        SearchResponse searchResponse = client().prepareSearch("test").setFetchSource(false).addStoredField("shape").get();
+        assertThat(searchResponse.getHits().getTotalHits().value, equalTo(1L));
+        SearchHit searchHit = searchResponse.getHits().getAt(0);
+        assertThat(searchHit.field("shape").getValue(), instanceOf(BytesRef.class));
+        BytesRef bytesRef = searchHit.field("shape").getValue();
+        Geometry geometry = WellKnownBinary.fromWKB(
+            StandardValidator.instance(true),
+            false,
+            bytesRef.bytes,
+            bytesRef.offset,
+            bytesRef.length
+        );
+        assertThat(geometry.type(), equalTo(ShapeType.MULTIPOLYGON));
+    }
 }

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

@@ -8,21 +8,26 @@
 package org.elasticsearch.xpack.spatial.index.mapper;
 
 import org.apache.lucene.document.LatLonShape;
+import org.apache.lucene.document.StoredField;
 import org.apache.lucene.geo.LatLonGeometry;
 import org.apache.lucene.index.IndexableField;
 import org.apache.lucene.search.IndexOrDocValuesQuery;
 import org.apache.lucene.search.Query;
+import org.apache.lucene.util.BytesRef;
 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.GeometryFormatterFactory;
 import org.elasticsearch.common.geo.GeometryParser;
 import org.elasticsearch.common.geo.Orientation;
 import org.elasticsearch.common.geo.ShapeRelation;
 import org.elasticsearch.common.logging.DeprecationCategory;
 import org.elasticsearch.common.logging.DeprecationLogger;
 import org.elasticsearch.geometry.Geometry;
+import org.elasticsearch.geometry.utils.GeometryValidator;
+import org.elasticsearch.geometry.utils.WellKnownBinary;
 import org.elasticsearch.index.fielddata.FieldDataContext;
 import org.elasticsearch.index.fielddata.IndexFieldData;
 import org.elasticsearch.index.fielddata.ScriptDocValues;
@@ -38,6 +43,8 @@
 import org.elasticsearch.index.mapper.MapperBuilderContext;
 import org.elasticsearch.index.mapper.MapperParsingException;
 import org.elasticsearch.index.mapper.MappingParserContext;
+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;
@@ -52,6 +59,8 @@
 
 import java.io.IOException;
 import java.io.UncheckedIOException;
+import java.nio.ByteOrder;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Iterator;
 import java.util.List;
@@ -94,6 +103,7 @@ private static Builder builder(FieldMapper in) {
     public static class Builder extends FieldMapper.Builder {
 
         final Parameter<Boolean> indexed = Parameter.indexParam(m -> builder(m).indexed.get(), true);
+        final Parameter<Boolean> stored = Parameter.storeParam(m -> builder(m).stored.get(), false);
         final Parameter<Boolean> hasDocValues;
 
         final Parameter<Explicit<Boolean>> ignoreMalformed;
@@ -121,9 +131,15 @@ public Builder(
             this.hasDocValues = Parameter.docValuesParam(m -> builder(m).hasDocValues.get(), Version.V_7_8_0.onOrBefore(version));
         }
 
+        // for testing
+        protected Builder setStored(boolean stored) {
+            this.stored.setValue(stored);
+            return this;
+        }
+
         @Override
         protected Parameter<?>[] getParameters() {
-            return new Parameter<?>[] { indexed, hasDocValues, ignoreMalformed, ignoreZValue, coerce, orientation, meta };
+            return new Parameter<?>[] { indexed, hasDocValues, stored, ignoreMalformed, ignoreZValue, coerce, orientation, meta };
         }
 
         @Override
@@ -145,6 +161,7 @@ public GeoShapeWithDocValuesFieldMapper build(MapperBuilderContext context) {
                 context.buildFullName(name),
                 indexed.get(),
                 hasDocValues.get(),
+                stored.get(),
                 orientation.get().value(),
                 parser,
                 geoFormatterFactory,
@@ -171,12 +188,13 @@ public GeoShapeWithDocValuesFieldType(
             String name,
             boolean indexed,
             boolean hasDocValues,
+            boolean isStored,
             Orientation orientation,
             GeoShapeParser parser,
             GeoFormatterFactory<Geometry> geoFormatterFactory,
             Map<String, String> meta
         ) {
-            super(name, indexed, false, hasDocValues, parser, orientation, meta);
+            super(name, indexed, isStored, hasDocValues, parser, orientation, meta);
             this.geoFormatterFactory = geoFormatterFactory;
         }
 
@@ -218,6 +236,30 @@ public Query geoShapeQuery(SearchExecutionContext context, String fieldName, Sha
             return query;
         }
 
+        @Override
+        public ValueFetcher valueFetcher(SearchExecutionContext context, String format) {
+            if (isStored()) {
+                Function<List<Geometry>, List<Object>> formatter = getFormatter(format != null ? format : GeometryFormatterFactory.GEOJSON);
+                return new StoredValueFetcher(context.lookup(), name()) {
+                    @Override
+                    public List<Object> parseStoredValues(List<Object> storedValues) {
+                        final List<Geometry> values = new ArrayList<>(storedValues.size());
+                        for (Object storedValue : storedValues) {
+                            if (storedValue instanceof BytesRef bytesRef) {
+                                values.add(
+                                    WellKnownBinary.fromWKB(GeometryValidator.NOOP, false, bytesRef.bytes, bytesRef.offset, bytesRef.length)
+                                );
+                            } else {
+                                throw new IllegalArgumentException("Unexpected class fetching [" + name() + "]: " + storedValue.getClass());
+                            }
+                        }
+                        return formatter.apply(values);
+                    }
+                };
+            }
+            return super.valueFetcher(context, format);
+        }
+
         @Override
         protected Function<List<Geometry>, List<Object>> getFormatter(String format) {
             return geoFormatterFactory.getFormatter(format, Function.identity());
@@ -303,21 +345,27 @@ protected void index(DocumentParserContext context, Geometry geometry) throws IO
         if (geometry == null) {
             return;
         }
-        List<IndexableField> fields = indexer.indexShape(geometry);
+        final Geometry normalizedGeometry = indexer.normalize(geometry);
+        final List<IndexableField> fields = indexer.getIndexableFields(normalizedGeometry);
         if (fieldType().isIndexed()) {
             context.doc().addAll(fields);
         }
         if (fieldType().hasDocValues()) {
-            String name = fieldType().name();
+            final String name = fieldType().name();
             BinaryShapeDocValuesField docValuesField = (BinaryShapeDocValuesField) context.doc().getByKey(name);
             if (docValuesField == null) {
                 docValuesField = new BinaryShapeDocValuesField(name, CoordinateEncoder.GEO);
                 context.doc().addWithKey(name, docValuesField);
             }
+            // we need to pass the original geometry to compute more precisely the centroid, e.g if lon > 180
             docValuesField.add(fields, geometry);
         } else if (fieldType().isIndexed()) {
             context.addToFieldNames(fieldType().name());
         }
+
+        if (fieldType().isStored()) {
+            context.doc().add(new StoredField(fieldType().name(), WellKnownBinary.toWKB(normalizedGeometry, ByteOrder.LITTLE_ENDIAN)));
+        }
     }
 
     @Override