make GeometryTree and EdgeTree writers implement Writeable (#42910)

talevy · talevy · commit d6d892414d3b · 2019-09-20T09:14:42.000-07:00
diff --git a/server/src/main/java/org/elasticsearch/common/geo/EdgeTreeReader.java b/server/src/main/java/org/elasticsearch/common/geo/EdgeTreeReader.java
@@ -20,6 +20,7 @@
 
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.io.stream.ByteBufferStreamInput;
+import org.elasticsearch.common.io.stream.StreamInput;
 
 import java.io.IOException;
 import java.nio.ByteBuffer;
@@ -29,8 +30,9 @@
 public class EdgeTreeReader {
     final BytesRef bytesRef;
 
-    public EdgeTreeReader(BytesRef bytesRef) {
-        this.bytesRef = bytesRef;
+    public EdgeTreeReader(StreamInput input) throws IOException {
+        int treeBytesSize = input.readVInt();
+        this.bytesRef = input.readBytesRef(treeBytesSize);
     }
 
     /**
diff --git a/server/src/main/java/org/elasticsearch/common/geo/EdgeTreeWriter.java b/server/src/main/java/org/elasticsearch/common/geo/EdgeTreeWriter.java
@@ -18,17 +18,16 @@
  */
 package org.elasticsearch.common.geo;
 
-import org.apache.lucene.util.BytesRef;
-import org.elasticsearch.common.io.stream.BytesStreamOutput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.io.stream.Writeable;
 
 import java.io.IOException;
 import java.util.Arrays;
 
 /**
  * Shape edge-tree writer for use in doc-values
  */
-public class EdgeTreeWriter {
+public class EdgeTreeWriter implements Writeable {
 
     /**
      * | minY | maxY | x1 | y1 | x2 | y2 | right_offset |
@@ -68,17 +67,14 @@ public EdgeTreeWriter(int[] x, int[] y) {
         this.tree = createTree(edges, 0, edges.length - 1);
     }
 
-    public BytesRef toBytesRef() throws IOException {
-        BytesStreamOutput output = new BytesStreamOutput(4 * 4 + EDGE_SIZE_IN_BYTES * tree.size);
-        // write extent of edges
-        output.writeInt(minX);
-        output.writeInt(minY);
-        output.writeInt(maxX);
-        output.writeInt(maxY);
-        // write edge-tree itself
-        writeTree(tree, output);
-        output.close();
-        return output.bytes().toBytesRef();
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeVInt(4 * 4 + EDGE_SIZE_IN_BYTES * tree.size);
+        out.writeInt(minX);
+        out.writeInt(minY);
+        out.writeInt(maxX);
+        out.writeInt(maxY);
+        writeTree(tree, out);
     }
 
     private void writeTree(Edge edge, StreamOutput output) throws IOException {
diff --git a/server/src/main/java/org/elasticsearch/common/geo/GeometryTreeReader.java b/server/src/main/java/org/elasticsearch/common/geo/GeometryTreeReader.java
@@ -62,8 +62,7 @@ public boolean containedInOrCrosses(int minLon, int minLat, int maxLon, int maxL
         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);
+                EdgeTreeReader reader = new EdgeTreeReader(input);
                 if (reader.containedInOrCrosses(minLon, minLat, maxLon, maxLat)) {
                     return true;
                 }
diff --git a/server/src/main/java/org/elasticsearch/common/geo/GeometryTreeWriter.java b/server/src/main/java/org/elasticsearch/common/geo/GeometryTreeWriter.java
@@ -18,8 +18,8 @@
  */
 package org.elasticsearch.common.geo;
 
-import org.apache.lucene.util.BytesRef;
-import org.elasticsearch.common.io.stream.BytesStreamOutput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.geo.geometry.Circle;
 import org.elasticsearch.geo.geometry.Geometry;
 import org.elasticsearch.geo.geometry.GeometryCollection;
@@ -43,7 +43,7 @@
  * appropriate tree structure for each type of
  * {@link Geometry} into a byte array.
  */
-public class GeometryTreeWriter {
+public class GeometryTreeWriter implements Writeable {
 
     private final GeometryTreeBuilder builder;
 
@@ -52,25 +52,23 @@ public class GeometryTreeWriter {
         geometry.visit(builder);
     }
 
-    public BytesRef toBytesRef() throws IOException {
-        BytesStreamOutput output = new BytesStreamOutput();
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
         // only write a geometry extent for the tree if the tree
         // contains multiple sub-shapes
         boolean prependExtent = builder.shapeWriters.size() > 1;
-        output.writeBoolean(prependExtent);
+        out.writeBoolean(prependExtent);
         if (prependExtent) {
-            output.writeInt(builder.minLon);
-            output.writeInt(builder.minLat);
-            output.writeInt(builder.maxLon);
-            output.writeInt(builder.maxLat);
+            out.writeInt(builder.minLon);
+            out.writeInt(builder.minLat);
+            out.writeInt(builder.maxLon);
+            out.writeInt(builder.maxLat);
         }
-        output.writeVInt(builder.shapeWriters.size());
+        out.writeVInt(builder.shapeWriters.size());
         for (EdgeTreeWriter writer : builder.shapeWriters) {
-            output.writeEnum(ShapeType.POLYGON);
-            output.writeBytesRef(writer.toBytesRef());
+            out.writeEnum(ShapeType.POLYGON);
+            writer.writeTo(out);
         }
-        output.close();
-        return output.bytes().toBytesRef();
     }
 
     class GeometryTreeBuilder implements GeometryVisitor<Void, RuntimeException> {
diff --git a/server/src/test/java/org/elasticsearch/common/geo/EdgeTreeTests.java b/server/src/test/java/org/elasticsearch/common/geo/EdgeTreeTests.java
@@ -18,8 +18,9 @@
  */
 package org.elasticsearch.common.geo;
 
-import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.geo.builders.ShapeBuilder;
+import org.elasticsearch.common.io.stream.BytesStreamOutput;
+import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.geo.geometry.Polygon;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.test.geo.RandomShapeGenerator;
@@ -38,8 +39,10 @@ public void testRectangleShape() throws IOException {
             int[] x = new int[]{minX, maxX, maxX, minX, minX};
             int[] y = new int[]{minY, minY, maxY, maxY, minY};
             EdgeTreeWriter writer = new EdgeTreeWriter(x, y);
-            BytesRef bytes = writer.toBytesRef();
-            EdgeTreeReader reader = new EdgeTreeReader(bytes);
+            BytesStreamOutput output = new BytesStreamOutput();
+            writer.writeTo(output);
+            output.close();
+            EdgeTreeReader reader = new EdgeTreeReader(StreamInput.wrap(output.bytes().toBytesRef().bytes));
 
             // box-query touches bottom-left corner
             assertTrue(reader.containedInOrCrosses(minX - randomIntBetween(1, 180), minY - randomIntBetween(1, 180), minX, minY));
@@ -93,7 +96,10 @@ public void testSimplePolygon() throws IOException  {
             int[] y = asIntArray(geo.getPolygon().getLats());
 
             EdgeTreeWriter writer = new EdgeTreeWriter(x, y);
-            EdgeTreeReader reader = new EdgeTreeReader(writer.toBytesRef());
+            BytesStreamOutput output = new BytesStreamOutput();
+            writer.writeTo(output);
+            output.close();
+            EdgeTreeReader reader = new EdgeTreeReader(StreamInput.wrap(output.bytes().toBytesRef().bytes));
             // polygon fully contained within box
             assertTrue(reader.containedInOrCrosses(minXBox, minYBox, maxXBox, maxYBox));
             // containedInOrCrosses
@@ -121,7 +127,10 @@ public void testPacMan() throws Exception {
 
         // test cell crossing poly
         EdgeTreeWriter writer = new EdgeTreeWriter(px, py);
-        EdgeTreeReader reader = new EdgeTreeReader(writer.toBytesRef());
+        BytesStreamOutput output = new BytesStreamOutput();
+        writer.writeTo(output);
+        output.close();
+        EdgeTreeReader reader = new EdgeTreeReader(StreamInput.wrap(output.bytes().toBytesRef().bytes));
         assertTrue(reader.containsBottomLeft(xMin, yMin, xMax, yMax));
     }
 
diff --git a/server/src/test/java/org/elasticsearch/common/geo/GeometryTreeTests.java b/server/src/test/java/org/elasticsearch/common/geo/GeometryTreeTests.java
@@ -18,7 +18,7 @@
  */
 package org.elasticsearch.common.geo;
 
-import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.common.io.stream.BytesStreamOutput;
 import org.elasticsearch.geo.geometry.LinearRing;
 import org.elasticsearch.geo.geometry.Polygon;
 import org.elasticsearch.test.ESTestCase;
@@ -37,8 +37,11 @@ public void testRectangleShape() throws IOException {
             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);
+
+            BytesStreamOutput output = new BytesStreamOutput();
+            writer.writeTo(output);
+            output.close();
+            GeometryTreeReader reader = new GeometryTreeReader(output.bytes().toBytesRef());
 
             // box-query touches bottom-left corner
             assertTrue(reader.containedInOrCrosses(minX - randomIntBetween(1, 180), minY - randomIntBetween(1, 180), minX, minY));
@@ -91,7 +94,10 @@ public void testPacMan() throws Exception {
 
         // test cell crossing poly
         GeometryTreeWriter writer = new GeometryTreeWriter(new Polygon(new LinearRing(py, px), Collections.emptyList()));
-        GeometryTreeReader reader = new GeometryTreeReader(writer.toBytesRef());
+        BytesStreamOutput output = new BytesStreamOutput();
+        writer.writeTo(output);
+        output.close();
+        GeometryTreeReader reader = new GeometryTreeReader(output.bytes().toBytesRef());
         assertTrue(reader.containedInOrCrosses(xMin, yMin, xMax, yMax));
     }
 }
