add shape-type metadata to geo_shape's doc-value (#50104)

talevy · talevy · commit 0ea0afeb4049 · 2019-12-18T21:13:11.000-08:00
This commit serializes the ShapeType of the indexed geometry. The ShapeType can be useful for other future features. For one thing: #49887 depends on the ability to determine what the highest dimensional shape is for centroid calculations. GeometryCollection is reduced to the sub-shape of the highest dimension relates #37206.
diff --git a/server/src/main/java/org/elasticsearch/common/geo/CentroidCalculator.java b/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;
         }
     }
+
 }
diff --git a/server/src/main/java/org/elasticsearch/common/geo/DimensionalShapeType.java b/server/src/main/java/org/elasticsearch/common/geo/DimensionalShapeType.java
@@ -0,0 +1,177 @@
+/*
+ * 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.apache.lucene.store.ByteArrayDataInput;
+import org.apache.lucene.store.ByteBuffersDataOutput;
+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 void writeTo(ByteBuffersDataOutput out) {
+        out.writeByte((byte) ordinal());
+    }
+
+    public static DimensionalShapeType readFrom(ByteArrayDataInput in) {
+        return values[Byte.toUnsignedInt(in.readByte())];
+    }
+
+    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");
+        }
+    }
+}
diff --git a/server/src/main/java/org/elasticsearch/common/geo/TriangleTreeReader.java b/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 CENTROID_HEADER_SIZE_IN_BYTES = 9;
 
-    private static final int extentOffset = 8;
     private final ByteArrayDataInput input;
     private final CoordinateEncoder coordinateEncoder;
     private final Rectangle2D rectangle2D;
@@ -59,7 +59,7 @@ public void reset(BytesRef bytesRef) throws IOException {
     public Extent getExtent() {
         if (treeOffset == 0) {
             // TODO: Compress serialization of extent
-            input.setPosition(extentOffset);
+            input.setPosition(CENTROID_HEADER_SIZE_IN_BYTES);
             int top = input.readInt();
             int bottom = Math.toIntExact(top - input.readVLong());
             int posRight = input.readInt();
@@ -90,6 +90,11 @@ public double getCentroidY() {
         return coordinateEncoder.decodeY(input.readInt());
     }
 
+    public DimensionalShapeType getDimensionalShapeType() {
+        input.setPosition(8);
+        return DimensionalShapeType.readFrom(input);
+    }
+
     /**
      * Compute the relation with the provided bounding box. If the result is CELL_INSIDE_QUERY
      * then the bounding box is within the shape.
diff --git a/server/src/main/java/org/elasticsearch/common/geo/TriangleTreeWriter.java b/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
+        centroidCalculator.getDimensionalShapeType().writeTo(out);
         out.writeInt(extent.top);
         out.writeVLong((long) extent.top - extent.bottom);
         out.writeInt(extent.posRight);
diff --git a/server/src/main/java/org/elasticsearch/index/fielddata/MultiGeoValues.java b/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 {
diff --git a/server/src/main/java/org/elasticsearch/index/mapper/AbstractGeometryFieldMapper.java b/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);
diff --git a/server/src/test/java/org/elasticsearch/common/geo/DimensionalShapeTypeTests.java b/server/src/test/java/org/elasticsearch/common/geo/DimensionalShapeTypeTests.java
@@ -0,0 +1,51 @@
+/*
+ * 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.apache.lucene.store.ByteArrayDataInput;
+import org.apache.lucene.store.ByteBuffersDataOutput;
+import org.elasticsearch.test.ESTestCase;
+
+import static org.hamcrest.Matchers.equalTo;
+
+public class DimensionalShapeTypeTests extends ESTestCase {
+
+    public void testValidOrdinals() {
+        assertThat(DimensionalShapeType.values().length, equalTo(9));
+        assertThat(DimensionalShapeType.POINT.ordinal(), equalTo(0));
+        assertThat(DimensionalShapeType.MULTIPOINT.ordinal(), equalTo(1));
+        assertThat(DimensionalShapeType.LINESTRING.ordinal(), equalTo(2));
+        assertThat(DimensionalShapeType.MULTILINESTRING.ordinal(), equalTo(3));
+        assertThat(DimensionalShapeType.POLYGON.ordinal(), equalTo(4));
+        assertThat(DimensionalShapeType.MULTIPOLYGON.ordinal(), equalTo(5));
+        assertThat(DimensionalShapeType.GEOMETRYCOLLECTION_POINTS.ordinal(), equalTo(6));
+        assertThat(DimensionalShapeType.GEOMETRYCOLLECTION_LINES.ordinal(), equalTo(7));
+        assertThat(DimensionalShapeType.GEOMETRYCOLLECTION_POLYGONS.ordinal(), equalTo(8));
+    }
+
+    public void testSerialization() {
+        for (DimensionalShapeType type : DimensionalShapeType.values()) {
+            ByteBuffersDataOutput out = new ByteBuffersDataOutput();
+            type.writeTo(out);
+            ByteArrayDataInput input = new ByteArrayDataInput(out.toArrayCopy());
+            assertThat(DimensionalShapeType.readFrom(input), equalTo(type));
+        }
+    }
+}
diff --git a/server/src/test/java/org/elasticsearch/common/geo/TriangleTreeTests.java b/server/src/test/java/org/elasticsearch/common/geo/TriangleTreeTests.java

Original file line number	Diff line number	Diff line change
`@@ -451,6 +451,7 @@ public void parse(ParseContext context) throws IOException {`
`451`	`451`	`System.arraycopy(bytesRef.bytes, bytesRef.offset, scratch, 0, 7 * Integer.BYTES);`
`452`	`452`	`ShapeField.decodeTriangle(scratch, triangles[i] = new ShapeField.DecodedTriangle());`
`453`	`453`	`}`
	`454`	`+`
`454`	`455`	`geometryIndexer.indexDocValueField(context, triangles, new CentroidCalculator((Geometry) shape));`
`455`	`456`	`}`
`456`	`457`	`createFieldNamesField(context, fields);`