Add GeometryTree support for point/multipoint

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

@@ -30,7 +30,7 @@
 public class GeometryCollection<G extends Geometry> implements Geometry, Iterable<G> {
     public static final GeometryCollection<Geometry> EMPTY = new GeometryCollection<>();
 
-    private final List<G> shapes;
+    protected final List<G> shapes;
 
     private boolean hasAlt;
 

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

@@ -34,6 +34,10 @@ public MultiPoint(List<Point> points) {
         super(points);
     }
 
+    public List<Point> points() {
+        return shapes;
+    }
+
     @Override
     public ShapeType type() {
         return ShapeType.MULTIPOINT;

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

@@ -20,21 +20,22 @@
 
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.geo.geometry.ShapeType;
 
 import java.io.IOException;
 import java.util.Arrays;
 
 /**
  * Shape edge-tree writer for use in doc-values
  */
-public class EdgeTreeWriter implements Writeable {
+public class EdgeTreeWriter extends ShapeTreeWriter {
 
     /**
      * | minY | maxY | x1 | y1 | x2 | y2 | right_offset |
      */
     static final int EDGE_SIZE_IN_BYTES = 28;
 
-    Extent extent;
+    private final Extent extent;
     final Edge tree;
 
     public EdgeTreeWriter(int[] x, int[] y) {
@@ -67,6 +68,16 @@ public EdgeTreeWriter(int[] x, int[] y) {
         this.tree = createTree(edges, 0, edges.length - 1);
     }
 
+    @Override
+    public Extent getExtent() {
+        return extent;
+    }
+
+    @Override
+    public ShapeType getShapeType() {
+        return ShapeType.POLYGON;
+    }
+
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         //out.writeVInt(4 * 4 + EDGE_SIZE_IN_BYTES * tree.size);

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

@@ -48,20 +48,28 @@ public Extent getExtent() throws IOException {
         }
         assert input.readVInt() == 1;
         ShapeType shapeType = input.readEnum(ShapeType.class);
-        if (ShapeType.POLYGON.equals(shapeType)) {
-            EdgeTreeReader reader = new EdgeTreeReader(input);
-            return reader.getExtent();
-        } else {
-            throw new UnsupportedOperationException("only polygons supported -- TODO");
+        switch (shapeType) {
+            case POLYGON:
+                EdgeTreeReader edgeReader = new EdgeTreeReader(input);
+                return edgeReader.getExtent();
+            case POINT:
+            case MULTIPOINT:
+                Point2DReader pointReader = new Point2DReader(input);
+                return pointReader.getExtent();
+            case LINESTRING:
+            case MULTILINESTRING:
+                throw new UnsupportedOperationException("TODO: linestring and multilinestring");
+            default:
+                throw new UnsupportedOperationException("unsupported shape type [" + shapeType + "]");
         }
     }
 
     public boolean containedInOrCrosses(int minLon, int minLat, int maxLon, int maxLat) throws IOException {
+        Extent extent = new Extent(minLon, minLat, maxLon, maxLat);
         input.position(0);
         boolean hasExtent = input.readBoolean();
         if (hasExtent) {
-            Optional<Boolean> extentCheck = EdgeTreeReader.checkExtent(input,
-                new Extent(minLon, minLat, maxLon, maxLat));
+            Optional<Boolean> extentCheck = EdgeTreeReader.checkExtent(input, extent);
             if (extentCheck.isPresent()) {
                 return extentCheck.get();
             }
@@ -70,11 +78,25 @@ public boolean containedInOrCrosses(int minLon, int minLat, int maxLon, int maxL
         int numTrees = input.readVInt();
         for (int i = 0; i < numTrees; i++) {
             ShapeType shapeType = input.readEnum(ShapeType.class);
-            if (ShapeType.POLYGON.equals(shapeType)) {
-                EdgeTreeReader reader = new EdgeTreeReader(input);
-                if (reader.containedInOrCrosses(minLon, minLat, maxLon, maxLat)) {
-                    return true;
-                }
+            switch (shapeType) {
+                case POLYGON:
+                    EdgeTreeReader edgeReader = new EdgeTreeReader(input);
+                    if (edgeReader.containedInOrCrosses(minLon, minLat, maxLon, maxLat)) {
+                        return true;
+                    }
+                    break;
+                case POINT:
+                case MULTIPOINT:
+                    Point2DReader pointReader = new Point2DReader(input);
+                    if (pointReader.containedIn(extent)) {
+                        return true;
+                    }
+                    break;
+                case LINESTRING:
+                case MULTILINESTRING:
+                    throw new UnsupportedOperationException("TODO: linestring and multilinestring");
+                default:
+                    throw new UnsupportedOperationException("unsupported shape type [" + shapeType +"]");
             }
         }
         return false;

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

@@ -32,10 +32,10 @@
 import org.elasticsearch.geo.geometry.Point;
 import org.elasticsearch.geo.geometry.Polygon;
 import org.elasticsearch.geo.geometry.Rectangle;
-import org.elasticsearch.geo.geometry.ShapeType;
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 
 /**
@@ -65,15 +65,15 @@ public void writeTo(StreamOutput out) throws IOException {
             out.writeInt(builder.maxLat);
         }
         out.writeVInt(builder.shapeWriters.size());
-        for (EdgeTreeWriter writer : builder.shapeWriters) {
-            out.writeEnum(ShapeType.POLYGON);
+        for (ShapeTreeWriter writer : builder.shapeWriters) {
+            out.writeEnum(writer.getShapeType());
             writer.writeTo(out);
         }
     }
 
     class GeometryTreeBuilder implements GeometryVisitor<Void, RuntimeException> {
 
-        private List<EdgeTreeWriter> shapeWriters;
+        private List<ShapeTreeWriter> shapeWriters;
         // integers are used to represent int-encoded lat/lon values
         int minLat;
         int maxLat;
@@ -86,11 +86,12 @@ class GeometryTreeBuilder implements GeometryVisitor<Void, RuntimeException> {
             maxLat = maxLon = Integer.MIN_VALUE;
         }
 
-        private void addWriter(EdgeTreeWriter writer) {
-            minLon = Math.min(minLon, writer.extent.minX);
-            minLat = Math.min(minLat, writer.extent.minY);
-            maxLon = Math.max(maxLon, writer.extent.maxX);
-            maxLat = Math.max(maxLat, writer.extent.maxY);
+        private void addWriter(ShapeTreeWriter writer) {
+            Extent extent = writer.getExtent();
+            minLon = Math.min(minLon, extent.minX);
+            minLat = Math.min(minLat, extent.minY);
+            maxLon = Math.max(maxLon, extent.maxX);
+            maxLat = Math.max(maxLat, extent.maxY);
             shapeWriters.add(writer);
         }
 
@@ -143,12 +144,16 @@ public Void visit(Rectangle r) {
 
         @Override
         public Void visit(Point point) {
-            throw new UnsupportedOperationException("support for Point is a TODO");
+            Point2DWriter writer = new Point2DWriter(Collections.singletonList(point));
+            addWriter(writer);
+            return null;
         }
 
         @Override
         public Void visit(MultiPoint multiPoint) {
-            throw new UnsupportedOperationException("support for MultiPoint is a TODO");
+            Point2DWriter writer = new Point2DWriter(multiPoint.points());
+            addWriter(writer);
+            return null;
         }
 
         @Override

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

@@ -0,0 +1,100 @@
+/*
+ * 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.common.io.stream.ByteBufferStreamInput;
+
+import java.io.IOException;
+import java.util.ArrayDeque;
+import java.util.Deque;
+
+class Point2DReader {
+    private final ByteBufferStreamInput input;
+    private final int size;
+    private final int startPosition;
+
+    Point2DReader(ByteBufferStreamInput input) throws IOException {
+        this.input = input;
+        this.size = input.readVInt();
+        this.startPosition = input.position();
+    }
+
+    Extent getExtent() throws IOException {
+        if (size == 2) {
+            int x = readX(0);
+            int y = readY(0);
+            return new Extent(x, y, x, y);
+        } else {
+            return new Extent(input);
+        }
+    }
+
+    boolean containedIn(Extent extent) throws IOException {
+        Deque<Integer> stack = new ArrayDeque<>();
+        stack.push(0);
+        stack.push(size - 1);
+        stack.push(0);
+        while (stack.isEmpty() == false) {
+            int axis = stack.pop();
+            int right = stack.pop();
+            int left = stack.pop();
+
+            if (right - left <= Point2DWriter.LEAF_SIZE) {
+                for (int i = left; i <= right; i++) {
+                    // TODO serialize to re-usable array instead of serializing in each step
+                    int x = readX(i);
+                    int y = readY(i);
+                    if (x >= extent.minX && x <= extent.maxX && y >= extent.minY && y <= extent.maxY) {
+                        return true;
+                    }
+                }
+                continue;
+            }
+
+            int middle = (right - left) >> 1;
+            int x = readX(middle);
+            int y = readY(middle);
+            if (x >= extent.minX && x <= extent.maxX && y >= extent.minY && y <= extent.maxY) {
+                return true;
+            }
+            if ((axis == 0 && extent.minX <= x) || (axis == 1 && extent.minY <= y)) {
+                stack.push(left);
+                stack.push(middle - 1);
+                stack.push(1 - axis);
+            }
+            if ((axis == 0 && extent.maxX >= x) || (axis == 1 && extent.maxY >= y)) {
+                stack.push(middle + 1);
+                stack.push(right);
+                stack.push(1 - axis);
+            }
+        }
+
+        return false;
+    }
+
+    private int readX(int pointIdx) throws IOException {
+        input.position(startPosition + 2 * pointIdx * Integer.BYTES);
+        return input.readInt();
+    }
+
+    private int readY(int pointIdx) throws IOException {
+        input.position(startPosition + (2 * pointIdx + 1) * Integer.BYTES);
+        return input.readInt();
+    }
+}