Introduce GeometryTree writer/reader

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

@@ -0,0 +1,74 @@
+/*
+ * 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.util.BytesRef;
+import org.elasticsearch.common.io.stream.ByteBufferStreamInput;
+import org.elasticsearch.geo.geometry.ShapeType;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+
+/**
+ * A tree reader.
+ *
+ * This class supports checking bounding box
+ * relations against the serialized geometry tree.
+ */
+public class GeometryTreeReader {
+
+    private final BytesRef bytesRef;
+
+    public GeometryTreeReader(BytesRef bytesRef) {
+        this.bytesRef = bytesRef;
+    }
+
+    public boolean containedInOrCrosses(int minLon, int minLat, int maxLon, int maxLat) throws IOException {
+        ByteBufferStreamInput input = new ByteBufferStreamInput(
+            ByteBuffer.wrap(bytesRef.bytes, bytesRef.offset, bytesRef.length));
+        boolean hasExtent = input.readBoolean();
+        if (hasExtent) {
+            int thisMinLon = input.readInt();
+            int thisMinLat = input.readInt();
+            int thisMaxLon = input.readInt();
+            int thisMaxLat = input.readInt();
+
+            if (thisMinLat > maxLat || thisMaxLon < minLon || thisMaxLat < minLat || thisMinLon > maxLon) {
+                return false; // tree and bbox-query are disjoint
+            }
+
+            if (minLon <= thisMinLon && minLat <= thisMinLat && maxLon >= thisMaxLon && maxLat >= thisMaxLat) {
+                return true; // bbox-query fully contains tree's extent.
+            }
+        }
+
+        int numTrees = input.readVInt();
+        for (int i = 0; i < numTrees; i++) {
+            ShapeType shapeType = input.readEnum(ShapeType.class);
+            if (ShapeType.POLYGON.equals(shapeType)) {
+                BytesRef treeRef = input.readBytesRef();
+                EdgeTreeReader reader = new EdgeTreeReader(treeRef);
+                if (reader.containedInOrCrosses(minLon, minLat, maxLon, maxLat)) {
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
+}

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

@@ -0,0 +1,174 @@
+/*
+ * 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.util.BytesRef;
+import org.elasticsearch.common.io.stream.BytesStreamOutput;
+import org.elasticsearch.geo.geometry.Circle;
+import org.elasticsearch.geo.geometry.Geometry;
+import org.elasticsearch.geo.geometry.GeometryCollection;
+import org.elasticsearch.geo.geometry.GeometryVisitor;
+import org.elasticsearch.geo.geometry.Line;
+import org.elasticsearch.geo.geometry.LinearRing;
+import org.elasticsearch.geo.geometry.MultiLine;
+import org.elasticsearch.geo.geometry.MultiPoint;
+import org.elasticsearch.geo.geometry.MultiPolygon;
+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.List;
+
+/**
+ * This is a tree-writer that serializes the
+ * appropriate tree structure for each type of
+ * {@link Geometry} into a byte array.
+ */
+public class GeometryTreeWriter {
+
+    private final GeometryTreeBuilder builder;
+
+    GeometryTreeWriter(Geometry geometry) {
+        builder = new GeometryTreeBuilder();
+        geometry.visit(builder);
+    }
+
+    public BytesRef toBytesRef() throws IOException {
+        BytesStreamOutput output = new BytesStreamOutput();
+        // only write a geometry extent for the tree if the tree
+        // contains multiple sub-shapes
+        boolean prependExtent = builder.shapeWriters.size() > 1;
+        output.writeBoolean(prependExtent);
+        if (prependExtent) {
+            output.writeInt(builder.minLon);
+            output.writeInt(builder.minLat);
+            output.writeInt(builder.maxLon);
+            output.writeInt(builder.maxLat);
+        }
+        output.writeVInt(builder.shapeWriters.size());
+        for (EdgeTreeWriter writer : builder.shapeWriters) {
+            output.writeEnum(ShapeType.POLYGON);
+            output.writeBytesRef(writer.toBytesRef());
+        }
+        output.close();
+        return output.bytes().toBytesRef();
+    }
+
+    class GeometryTreeBuilder implements GeometryVisitor<Void, RuntimeException> {
+
+        private List<EdgeTreeWriter> shapeWriters;
+        // integers are used to represent int-encoded lat/lon values
+        int minLat;
+        int maxLat;
+        int minLon;
+        int maxLon;
+
+        GeometryTreeBuilder() {
+            shapeWriters = new ArrayList<>();
+            minLat = minLon = Integer.MAX_VALUE;
+            maxLat = maxLon = Integer.MIN_VALUE;
+        }
+
+        private void addWriter(EdgeTreeWriter writer) {
+            minLon = Math.min(minLon, writer.minX);
+            minLat = Math.min(minLat, writer.minY);
+            maxLon = Math.max(maxLon, writer.maxX);
+            maxLat = Math.max(maxLat, writer.maxY);
+            shapeWriters.add(writer);
+        }
+
+        @Override
+        public Void visit(GeometryCollection<?> collection) {
+            for (Geometry geometry : collection) {
+                geometry.visit(this);
+            }
+            return null;
+        }
+
+        @Override
+        public Void visit(Line line) {
+            throw new UnsupportedOperationException("support for Line is a TODO");
+        }
+
+        @Override
+        public Void visit(MultiLine multiLine) {
+            for (Line line : multiLine) {
+                visit(line);
+            }
+            return null;
+        }
+
+        @Override
+        public Void visit(Polygon polygon) {
+            // TODO (support holes)
+            LinearRing outerShell = polygon.getPolygon();
+            addWriter(new EdgeTreeWriter(asIntArray(outerShell.getLons()), asIntArray(outerShell.getLats())));
+            return null;
+        }
+
+        @Override
+        public Void visit(MultiPolygon multiPolygon) {
+            for (Polygon polygon : multiPolygon) {
+                visit(polygon);
+            }
+            return null;
+        }
+
+        @Override
+        public Void visit(Rectangle r) {
+            int[] lats = new int[] { (int) r.getMinLat(), (int) r.getMinLat(), (int) r.getMaxLat(), (int) r.getMaxLat(),
+                (int) r.getMinLat()};
+            int[] lons = new int[] { (int) r.getMinLon(), (int) r.getMaxLon(), (int) r.getMaxLon(), (int) r.getMinLon(),
+                (int) r.getMinLon()};
+            addWriter(new EdgeTreeWriter(lons, lats));
+            return null;
+        }
+
+        @Override
+        public Void visit(Point point) {
+            throw new UnsupportedOperationException("support for Point is a TODO");
+        }
+
+        @Override
+        public Void visit(MultiPoint multiPoint) {
+            throw new UnsupportedOperationException("support for MultiPoint is a TODO");
+        }
+
+        @Override
+        public Void visit(LinearRing ring) {
+            throw new IllegalArgumentException("invalid shape type found [Circle]");
+        }
+
+        @Override
+        public Void visit(Circle circle) {
+            throw new IllegalArgumentException("invalid shape type found [Circle]");
+        }
+
+        private int[] asIntArray(double[] doub) {
+            int[] intArr = new int[doub.length];
+            for (int i = 0; i < intArr.length; i++) {
+                intArr[i] = (int) doub[i];
+            }
+            return intArr;
+        }
+    }
+}

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

@@ -0,0 +1,97 @@
+/*
+ * 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.util.BytesRef;
+import org.elasticsearch.geo.geometry.LinearRing;
+import org.elasticsearch.geo.geometry.Polygon;
+import org.elasticsearch.test.ESTestCase;
+
+import java.io.IOException;
+import java.util.Collections;
+
+public class GeometryTreeTests extends ESTestCase {
+
+    public void testRectangleShape() throws IOException {
+        for (int i = 0; i < 1000; i++) {
+            int minX = randomIntBetween(-180, 170);
+            int maxX = randomIntBetween(minX + 10, 180);
+            int minY = randomIntBetween(-90, 80);
+            int maxY = randomIntBetween(minY + 10, 90);
+            double[] x = new double[]{minX, maxX, maxX, minX, minX};
+            double[] y = new double[]{minY, minY, maxY, maxY, minY};
+            GeometryTreeWriter writer = new GeometryTreeWriter(new Polygon(new LinearRing(y, x), Collections.emptyList()));
+            BytesRef bytes = writer.toBytesRef();
+            GeometryTreeReader reader = new GeometryTreeReader(bytes);
+
+            // box-query touches bottom-left corner
+            assertTrue(reader.containedInOrCrosses(minX - randomIntBetween(1, 180), minY - randomIntBetween(1, 180), minX, minY));
+            // box-query touches bottom-right corner
+            assertTrue(reader.containedInOrCrosses(maxX, minY - randomIntBetween(1, 180), maxX + randomIntBetween(1, 180), minY));
+            // box-query touches top-right corner
+            assertTrue(reader.containedInOrCrosses(maxX, maxY, maxX + randomIntBetween(1, 180), maxY + randomIntBetween(1, 180)));
+            // box-query touches top-left corner
+            assertTrue(reader.containedInOrCrosses(minX - randomIntBetween(1, 180), maxY, minX, maxY + randomIntBetween(1, 180)));
+            // box-query fully-enclosed inside rectangle
+            assertTrue(reader.containedInOrCrosses((3 * minX + maxX) / 4, (3 * minY + maxY) / 4, (3 * maxX + minX) / 4,
+                (3 * maxY + minY) / 4));
+            // box-query fully-contains poly
+            assertTrue(reader.containedInOrCrosses(minX - randomIntBetween(1, 180), minY - randomIntBetween(1, 180),
+                maxX + randomIntBetween(1, 180), maxY + randomIntBetween(1, 180)));
+            // box-query half-in-half-out-right
+            assertTrue(reader.containedInOrCrosses((3 * minX + maxX) / 4, (3 * minY + maxY) / 4, maxX + randomIntBetween(1, 1000),
+                (3 * maxY + minY) / 4));
+            // box-query half-in-half-out-left
+            assertTrue(reader.containedInOrCrosses(minX - randomIntBetween(1, 1000), (3 * minY + maxY) / 4, (3 * maxX + minX) / 4,
+                (3 * maxY + minY) / 4));
+            // box-query half-in-half-out-top
+            assertTrue(reader.containedInOrCrosses((3 * minX + maxX) / 4, (3 * minY + maxY) / 4, maxX + randomIntBetween(1, 1000),
+                maxY + randomIntBetween(1, 1000)));
+            // box-query half-in-half-out-bottom
+            assertTrue(reader.containedInOrCrosses((3 * minX + maxX) / 4, minY - randomIntBetween(1, 1000),
+                maxX + randomIntBetween(1, 1000), (3 * maxY + minY) / 4));
+
+            // box-query outside to the right
+            assertFalse(reader.containedInOrCrosses(maxX + randomIntBetween(1, 1000), minY, maxX + randomIntBetween(1001, 2000), maxY));
+            // box-query outside to the left
+            assertFalse(reader.containedInOrCrosses(maxX - randomIntBetween(1001, 2000), minY, minX - randomIntBetween(1, 1000), maxY));
+            // box-query outside to the top
+            assertFalse(reader.containedInOrCrosses(minX, maxY + randomIntBetween(1, 1000), maxX, maxY + randomIntBetween(1001, 2000)));
+            // box-query outside to the bottom
+            assertFalse(reader.containedInOrCrosses(minX, minY - randomIntBetween(1001, 2000), maxX, minY - randomIntBetween(1, 1000)));
+        }
+    }
+
+    public void testPacMan() throws Exception {
+        // pacman
+        double[] px = {0, 10, 10, 0, -8, -10, -8, 0, 10, 10, 0};
+        double[] py = {0, 5, 9, 10, 9, 0, -9, -10, -9, -5, 0};
+
+        // candidate containedInOrCrosses cell
+        int xMin = 2;//-5;
+        int xMax = 11;//0.000001;
+        int yMin = -1;//0;
+        int yMax = 1;//5;
+
+        // test cell crossing poly
+        GeometryTreeWriter writer = new GeometryTreeWriter(new Polygon(new LinearRing(py, px), Collections.emptyList()));
+        GeometryTreeReader reader = new GeometryTreeReader(writer.toBytesRef());
+        assertTrue(reader.containedInOrCrosses(xMin, yMin, xMax, yMax));
+    }
+}