add shape-type metadata to geo_shape's doc-value

libs/geo/src/main/java/org/elasticsearch/geometry/ShapeType.java

@@ -36,6 +36,11 @@ public enum ShapeType {
     ENVELOPE, // not part of the actual WKB spec
     CIRCLE; // not part of the actual WKB spec
 
+    private static ShapeType[] values = values();
+
+    public static ShapeType forOrdinal(int ordinal) {
+        return values[ordinal];
+    }
     public static ShapeType forName(String shapeName) {
         return ShapeType.valueOf(shapeName.toUpperCase(Locale.ROOT));
     }

server/src/main/java/org/elasticsearch/common/geo/CentroidCalculator.java

@@ -44,14 +44,17 @@ public class CentroidCalculator {
     private double sumX;
     private double sumY;
     private int count;
+    private DimensionalShapeType dimensionalShapeType;
 
     public CentroidCalculator(Geometry geometry) {
         this.sumX = 0.0;
         this.compX = 0.0;
         this.sumY = 0.0;
         this.compY = 0.0;
         this.count = 0;
-        geometry.visit(new CentroidCalculatorVisitor(this));
+        CentroidCalculatorVisitor visitor = new CentroidCalculatorVisitor(this);
+        geometry.visit(visitor);
+        this.dimensionalShapeType = DimensionalShapeType.forGeometry(geometry);
     }
 
     /**
@@ -87,6 +90,7 @@ public void addFrom(CentroidCalculator otherCalculator) {
         addCoordinate(otherCalculator.sumX, otherCalculator.sumY);
         // adjust count
         count += otherCalculator.count - 1;
+        dimensionalShapeType = DimensionalShapeType.max(dimensionalShapeType, otherCalculator.dimensionalShapeType);
     }
 
     /**
@@ -103,6 +107,10 @@ public double getY() {
         return sumY / count;
     }
 
+    public DimensionalShapeType getDimensionalShapeType() {
+        return dimensionalShapeType;
+    }
+
     private static class CentroidCalculatorVisitor implements GeometryVisitor<Void, IllegalArgumentException> {
 
         private final CentroidCalculator calculator;
@@ -127,7 +135,6 @@ public Void visit(GeometryCollection<?> collection) {
 
         @Override
         public Void visit(Line line) {
-
             for (int i = 0; i < line.length(); i++) {
                 calculator.addCoordinate(line.getX(i), line.getY(i));
             }
@@ -174,6 +181,7 @@ public Void visit(Point point) {
 
         @Override
         public Void visit(Polygon polygon) {
+            // TODO: incorporate holes into centroid calculation
             return visit(polygon.getPolygon());
         }
 
@@ -186,4 +194,5 @@ public Void visit(Rectangle rectangle) {
             return null;
         }
     }
+
 }

server/src/main/java/org/elasticsearch/common/geo/DimensionalShapeType.java

@@ -0,0 +1,171 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.elasticsearch.common.geo;
+
+import org.elasticsearch.geometry.Circle;
+import org.elasticsearch.geometry.Geometry;
+import org.elasticsearch.geometry.GeometryCollection;
+import org.elasticsearch.geometry.GeometryVisitor;
+import org.elasticsearch.geometry.Line;
+import org.elasticsearch.geometry.LinearRing;
+import org.elasticsearch.geometry.MultiLine;
+import org.elasticsearch.geometry.MultiPoint;
+import org.elasticsearch.geometry.MultiPolygon;
+import org.elasticsearch.geometry.Point;
+import org.elasticsearch.geometry.Polygon;
+import org.elasticsearch.geometry.Rectangle;
+import org.elasticsearch.geometry.ShapeType;
+
+import java.util.Comparator;
+
+/**
+ * Like {@link ShapeType} but has specific
+ * types for when the geometry is a {@link GeometryCollection} and
+ * more information about what the highest-dimensional sub-shape
+ * is.
+ */
+public enum DimensionalShapeType {
+    POINT,
+    MULTIPOINT,
+    LINESTRING,
+    MULTILINESTRING,
+    POLYGON,
+    MULTIPOLYGON,
+    GEOMETRYCOLLECTION_POINTS,    // highest-dimensional shapes are Points
+    GEOMETRYCOLLECTION_LINES,     // highest-dimensional shapes are Lines
+    GEOMETRYCOLLECTION_POLYGONS;   // highest-dimensional shapes are Polygons
+
+    private static DimensionalShapeType[] values = values();
+
+    private static Comparator<DimensionalShapeType> COMPARATOR = Comparator.comparingInt(DimensionalShapeType::centroidDimension);
+
+    public static DimensionalShapeType max(DimensionalShapeType s1, DimensionalShapeType s2) {
+        if (s1 == null) {
+            return s2;
+        } else if (s2 == null) {
+            return s1;
+        }
+        return COMPARATOR.compare(s1, s2) >= 0 ? s1 : s2;
+    }
+
+    public static DimensionalShapeType forOrdinal(int ordinal) {
+        return values[ordinal];
+    }
+
+    public static DimensionalShapeType forGeometry(Geometry geometry) {
+        return geometry.visit(new GeometryVisitor<>() {
+            private DimensionalShapeType st = null;
+
+            @Override
+            public DimensionalShapeType visit(Circle circle) {
+                st = DimensionalShapeType.max(st, DimensionalShapeType.POLYGON);
+                return st;
+            }
+
+            @Override
+            public DimensionalShapeType visit(Line line) {
+                st = DimensionalShapeType.max(st, DimensionalShapeType.LINESTRING);
+                return st;
+            }
+
+            @Override
+            public DimensionalShapeType visit(LinearRing ring)  {
+                throw new UnsupportedOperationException("should not visit LinearRing");
+            }
+
+            @Override
+            public DimensionalShapeType visit(MultiLine multiLine)  {
+                st = DimensionalShapeType.max(st, DimensionalShapeType.MULTILINESTRING);
+                return st;
+            }
+
+            @Override
+            public DimensionalShapeType visit(MultiPoint multiPoint)  {
+                st = DimensionalShapeType.max(st, DimensionalShapeType.MULTIPOINT);
+                return st;
+            }
+
+            @Override
+            public DimensionalShapeType visit(MultiPolygon multiPolygon)  {
+                st = DimensionalShapeType.max(st, DimensionalShapeType.MULTIPOLYGON);
+                return st;
+            }
+
+            @Override
+            public DimensionalShapeType visit(Point point)  {
+                st = DimensionalShapeType.max(st, DimensionalShapeType.POINT);
+                return st;
+            }
+
+            @Override
+            public DimensionalShapeType visit(Polygon polygon)  {
+                st = DimensionalShapeType.max(st, DimensionalShapeType.POLYGON);
+                return st;
+            }
+
+            @Override
+            public DimensionalShapeType visit(Rectangle rectangle)  {
+                st = DimensionalShapeType.max(st, DimensionalShapeType.POLYGON);
+                return st;
+            }
+
+            @Override
+            public DimensionalShapeType visit(GeometryCollection<?> collection)  {
+                for (Geometry shape : collection) {
+                    shape.visit(this);
+                }
+                int dimension = st.centroidDimension();
+                if (dimension == 0) {
+                    return DimensionalShapeType.GEOMETRYCOLLECTION_POINTS;
+                } else if (dimension == 1) {
+                    return DimensionalShapeType.GEOMETRYCOLLECTION_LINES;
+                } else {
+                    return DimensionalShapeType.GEOMETRYCOLLECTION_POLYGONS;
+                }
+            }
+        });
+    }
+
+    /**
+     * The integer representation of the dimension for the specific
+     * dimensional shape type. This is to be used by the centroid
+     * calculation to determine whether to add a sub-shape's centroid
+     * to the overall shape calculation.
+     *
+     * @return 0 for points, 1 for lines, 2 for polygons
+     */
+    private int centroidDimension() {
+        switch (this) {
+            case POINT:
+            case MULTIPOINT:
+            case GEOMETRYCOLLECTION_POINTS:
+                return 0;
+            case LINESTRING:
+            case MULTILINESTRING:
+            case GEOMETRYCOLLECTION_LINES:
+                return 1;
+            case POLYGON:
+            case MULTIPOLYGON:
+            case GEOMETRYCOLLECTION_POLYGONS:
+                return 2;
+            default:
+                throw new IllegalStateException("dimension calculation of DimensionalShapeType [" + this + "] is not supported");
+        }
+    }
+}

server/src/main/java/org/elasticsearch/common/geo/TriangleTreeReader.java

@@ -33,8 +33,8 @@
  * relations against the serialized triangle tree.
  */
 public class TriangleTreeReader {
+    private static final int SKIP_CENTROID = 8;
 
-    private static final int extentOffset = 8;
     private final ByteArrayDataInput input;
     private final CoordinateEncoder coordinateEncoder;
     private final Rectangle2D rectangle2D;
@@ -58,8 +58,7 @@ public void reset(BytesRef bytesRef) throws IOException {
      */
     public Extent getExtent() {
         if (treeOffset == 0) {
-            // TODO: Compress serialization of extent
-            input.setPosition(extentOffset);
+            getDimensionalShapeType(); // skip DimensionalShapeType
             int top = input.readInt();
             int bottom = Math.toIntExact(top - input.readVLong());
             int posRight = input.readInt();
@@ -90,6 +89,11 @@ public double getCentroidY() {
         return coordinateEncoder.decodeY(input.readInt());
     }
 
+    public DimensionalShapeType getDimensionalShapeType() {
+        input.setPosition(SKIP_CENTROID);
+        return DimensionalShapeType.forOrdinal(input.readVInt());
+    }
+
     /**
      * Compute the relation with the provided bounding box. If the result is CELL_INSIDE_QUERY
      * then the bounding box is within the shape.

server/src/main/java/org/elasticsearch/common/geo/TriangleTreeWriter.java

@@ -49,7 +49,7 @@ public TriangleTreeWriter(List<ShapeField.DecodedTriangle> triangles, Coordinate
     public void writeTo(ByteBuffersDataOutput out) throws IOException {
         out.writeInt(coordinateEncoder.encodeX(centroidCalculator.getX()));
         out.writeInt(coordinateEncoder.encodeY(centroidCalculator.getY()));
-        // TODO: Compress serialization of extent
+        out.writeVInt(centroidCalculator.getDimensionalShapeType().ordinal());
         out.writeInt(extent.top);
         out.writeVLong((long) extent.top - extent.bottom);
         out.writeInt(extent.posRight);

server/src/main/java/org/elasticsearch/index/fielddata/MultiGeoValues.java

@@ -25,6 +25,7 @@
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.geo.CentroidCalculator;
 import org.elasticsearch.common.geo.CoordinateEncoder;
+import org.elasticsearch.common.geo.DimensionalShapeType;
 import org.elasticsearch.common.geo.Extent;
 import org.elasticsearch.common.geo.GeoPoint;
 import org.elasticsearch.common.geo.GeoRelation;
@@ -117,6 +118,11 @@ public GeoRelation relate(Rectangle rectangle) {
             return GeoRelation.QUERY_DISJOINT;
         }
 
+        @Override
+        public DimensionalShapeType dimensionalShapeType() {
+            return DimensionalShapeType.POINT;
+        }
+
         @Override
         public double lat() {
             return geoPoint.lat();
@@ -162,6 +168,11 @@ public GeoRelation relate(Rectangle rectangle) {
             return reader.relate(minX, minY, maxX, maxY);
         }
 
+        @Override
+        public DimensionalShapeType dimensionalShapeType() {
+            return reader.getDimensionalShapeType();
+        }
+
         @Override
         public double lat() {
             return reader.getCentroidY();
@@ -217,6 +228,7 @@ public interface GeoValue {
         double lon();
         BoundingBox boundingBox();
         GeoRelation relate(Rectangle rectangle);
+        DimensionalShapeType dimensionalShapeType();
     }
 
     public static class BoundingBox {

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

@@ -451,6 +451,7 @@ public void parse(ParseContext context) throws IOException {
                     System.arraycopy(bytesRef.bytes, bytesRef.offset, scratch, 0, 7 * Integer.BYTES);
                     ShapeField.decodeTriangle(scratch, triangles[i] = new ShapeField.DecodedTriangle());
                  }
+
                 geometryIndexer.indexDocValueField(context, triangles, new CentroidCalculator((Geometry) shape));
             }
             createFieldNamesField(context, fields);

server/src/test/java/org/elasticsearch/common/geo/TriangleTreeTests.java

@@ -41,11 +41,47 @@
 import static org.elasticsearch.common.geo.GeoTestUtils.assertRelation;
 import static org.elasticsearch.common.geo.GeoTestUtils.triangleTreeReader;
 import static org.elasticsearch.geo.GeometryTestUtils.fold;
+import static org.elasticsearch.geo.GeometryTestUtils.randomLine;
+import static org.elasticsearch.geo.GeometryTestUtils.randomMultiLine;
+import static org.elasticsearch.geo.GeometryTestUtils.randomMultiPoint;
+import static org.elasticsearch.geo.GeometryTestUtils.randomMultiPolygon;
 import static org.elasticsearch.geo.GeometryTestUtils.randomPoint;
+import static org.elasticsearch.geo.GeometryTestUtils.randomPolygon;
+import static org.elasticsearch.geo.GeometryTestUtils.randomRectangle;
 import static org.hamcrest.Matchers.equalTo;
 
 public class TriangleTreeTests extends ESTestCase {
 
+    @SuppressWarnings("unchecked")
+    public void testDimensionalShapeType() throws IOException {
+        assertDimensionalShapeType(randomPoint(false), DimensionalShapeType.POINT);
+        assertDimensionalShapeType(randomMultiPoint(false), DimensionalShapeType.MULTIPOINT);
+        assertDimensionalShapeType(randomLine(false), DimensionalShapeType.LINESTRING);
+        assertDimensionalShapeType(randomMultiLine(false), DimensionalShapeType.MULTILINESTRING);
+        assertDimensionalShapeType(randomPolygon(false), DimensionalShapeType.POLYGON);
+        assertDimensionalShapeType(randomMultiPolygon(false), DimensionalShapeType.MULTIPOLYGON);
+        assertDimensionalShapeType(randomRectangle(), DimensionalShapeType.POLYGON);
+        assertDimensionalShapeType(randomFrom(
+            new GeometryCollection<>(List.of(randomPoint(false))),
+            new GeometryCollection<>(List.of(randomMultiPoint(false))),
+            new GeometryCollection<>(Collections.singletonList(
+                new GeometryCollection<>(List.of(randomPoint(false), randomMultiPoint(false))))))
+            , DimensionalShapeType.GEOMETRYCOLLECTION_POINTS);
+        assertDimensionalShapeType(randomFrom(
+            new GeometryCollection<>(List.of(randomPoint(false), randomLine(false))),
+            new GeometryCollection<>(List.of(randomMultiPoint(false), randomMultiLine(false))),
+            new GeometryCollection<>(Collections.singletonList(
+                new GeometryCollection<>(List.of(randomPoint(false), randomLine(false))))))
+            , DimensionalShapeType.GEOMETRYCOLLECTION_LINES);
+        assertDimensionalShapeType(randomFrom(
+            new GeometryCollection<>(List.of(randomPoint(false), randomLine(false), randomPolygon(false))),
+            new GeometryCollection<>(List.of(randomMultiPoint(false), randomMultiPolygon(false))),
+            new GeometryCollection<>(Collections.singletonList(
+                new GeometryCollection<>(List.of(randomLine(false), randomPolygon(false))))))
+            , DimensionalShapeType.GEOMETRYCOLLECTION_POLYGONS);
+    }
+
+
     public void testRectangleShape() throws IOException {
         for (int i = 0; i < 1000; i++) {
             int minX = randomIntBetween(-40, -1);
@@ -214,7 +250,7 @@ public void testPacManPoints() throws Exception {
     public void testRandomMultiLineIntersections() throws IOException {
         double extentSize = randomDoubleBetween(0.01, 10, true);
         GeoShapeIndexer indexer = new GeoShapeIndexer(true, "test");
-        MultiLine geometry = GeometryTestUtils.randomMultiLine(false);
+        MultiLine geometry = randomMultiLine(false);
         geometry = (MultiLine) indexer.prepareForIndexing(geometry);
 
         TriangleTreeReader reader = triangleTreeReader(geometry, GeoShapeCoordinateEncoder.INSTANCE);
@@ -316,4 +352,9 @@ private static Geometry randomGeometryTreeCollection(int level) {
         }
         return new GeometryCollection<>(shapes);
     }
+
+    private static void assertDimensionalShapeType(Geometry geometry, DimensionalShapeType expected) throws IOException {
+        TriangleTreeReader reader = triangleTreeReader(geometry, GeoShapeCoordinateEncoder.INSTANCE);
+        assertThat(reader.getDimensionalShapeType(), equalTo(expected));
+    }
 }