MB-54131: Geoshape query decode optimization

.gitignore

@@ -0,0 +1 @@
+.DS_Store

geojson/geojson_shapes_util.go

@@ -31,9 +31,9 @@ var jsoniter = jsoniterator.ConfigCompatibleWithStandardLibrary
 // the `relation` filter and confirms whether the shape in the document
 //  satisfies the given relation.
 func FilterGeoShapesOnRelation(shape index.GeoJSON, targetShapeBytes []byte,
-	relation string, reader **bytes.Reader) (bool, error) {
+	relation string, reader **bytes.Reader, bufPool [][]byte) (bool, error) {
 
-	shapeInDoc, err := extractShapesFromBytes(targetShapeBytes, reader)
+	shapeInDoc, err := extractShapesFromBytes(targetShapeBytes, reader, bufPool)
 	if err != nil {
 		return false, err
 	}
@@ -43,7 +43,7 @@ func FilterGeoShapesOnRelation(shape index.GeoJSON, targetShapeBytes []byte,
 
 // extractShapesFromBytes unmarshal the bytes to retrieve the
 // embedded geojson shape.
-func extractShapesFromBytes(targetShapeBytes []byte, r **bytes.Reader) (
+func extractShapesFromBytes(targetShapeBytes []byte, r **bytes.Reader, bufPool [][]byte) (
 	index.GeoJSON, error) {
 	if (*r) == nil {
 		*r = bytes.NewReader(targetShapeBytes[1:])
@@ -109,7 +109,7 @@ func extractShapesFromBytes(targetShapeBytes []byte, r **bytes.Reader) (
 		return mls, nil
 
 	case PolygonTypePrefix:
-		pgn := &Polygon{s2pgn: &s2.Polygon{}}
+		pgn := &Polygon{s2pgn: &s2.Polygon{BufPool: bufPool}}
 		err := pgn.s2pgn.Decode(*r)
 		if err != nil {
 			return nil, err
@@ -156,7 +156,7 @@ func extractShapesFromBytes(targetShapeBytes []byte, r **bytes.Reader) (
 		gc := &GeometryCollection{Shapes: make([]index.GeoJSON, numShapes)}
 
 		for i := int32(0); i < numShapes; i++ {
-			shape, err := extractShapesFromBytes(inputBytes[:lengths[i]], r)
+			shape, err := extractShapesFromBytes(inputBytes[:lengths[i]], r, nil)
 			if err != nil {
 				return nil, err
 			}

s2/encode.go

@@ -222,3 +222,22 @@ func (d *decoder) readUvarint() (x uint64) {
 	x, d.err = binary.ReadUvarint(d.r)
 	return
 }
+
+func (d *decoder) readFloat64Array(size int, buffers [][]byte) ([]byte, int) {
+	if d.err != nil {
+		return nil, 0
+	}
+
+	var buf []byte
+
+	for _, buffer := range buffers {
+		if len(buffer) <= size*8 {
+			buf = buffer
+			break
+		}
+	}
+
+	_, d.err = io.ReadFull(d.r, buf)
+
+	return buf, len(buf) / 8
+}

s2/loop.go

@@ -15,6 +15,7 @@
 package s2
 
 import (
+	"encoding/binary"
 	"fmt"
 	"io"
 	"math"
@@ -45,6 +46,7 @@ import (
 type Loop struct {
 	vertices []Point
 
+	polygon *Polygon
 	// originInside keeps a precomputed value whether this loop contains the origin
 	// versus computing from the set of vertices every time.
 	originInside bool
@@ -1287,11 +1289,26 @@ func (l *Loop) decode(d *decoder) {
 		return
 	}
 	l.vertices = make([]Point, nvertices)
-	for i := range l.vertices {
-		l.vertices[i].X = d.readFloat64()
-		l.vertices[i].Y = d.readFloat64()
-		l.vertices[i].Z = d.readFloat64()
+
+	pointsNeeded := int(nvertices) * 3
+
+	i := 0
+
+	for pointsNeeded > 0 {
+		arr, pointsRead := d.readFloat64Array(pointsNeeded, l.polygon.BufPool)
+		if pointsRead == 0 {
+			break
+		}
+		pointsNeeded = pointsNeeded - pointsRead
+		for j := 0; j < int(pointsRead/3); j++ {
+			l.vertices[i+j].X = math.Float64frombits(binary.LittleEndian.Uint64(arr[8*(j*3) : 8*(j*3+1)]))
+			l.vertices[i+j].Y = math.Float64frombits(binary.LittleEndian.Uint64(arr[8*(j*3+1) : 8*(j*3+2)]))
+			l.vertices[i+j].Z = math.Float64frombits(binary.LittleEndian.Uint64(arr[8*(j*3+2) : 8*(j*3+3)]))
+		}
+
+		i = i + int(pointsRead/3)
 	}
+
 	l.index = NewShapeIndex()
 	l.originInside = d.readBool()
 	l.depth = int(d.readUint32())

s2/polygon.go

@@ -77,6 +77,8 @@ type Polygon struct {
 	// preceding loops in the polygon. This field is used for polygons that
 	// have a large number of loops, and may be empty for polygons with few loops.
 	cumulativeEdges []int
+
+	BufPool [][]byte
 }
 
 // PolygonFromLoops constructs a polygon from the given set of loops. The polygon
@@ -1133,7 +1135,7 @@ func (p *Polygon) Decode(r io.Reader) error {
 const maxEncodedLoops = 10000000
 
 func (p *Polygon) decode(d *decoder) {
-	*p = Polygon{}
+	*p = Polygon{BufPool: p.BufPool}
 	d.readUint8() // Ignore irrelevant serialized owns_loops_ value.
 
 	p.hasHoles = d.readBool()
@@ -1151,6 +1153,7 @@ func (p *Polygon) decode(d *decoder) {
 	p.loops = make([]*Loop, nloops)
 	for i := range p.loops {
 		p.loops[i] = new(Loop)
+		p.loops[i].polygon = p
 		p.loops[i].decode(d)
 		p.numVertices += len(p.loops[i].vertices)
 	}