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 *s2.GeoBufferPool) (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 *s2.GeoBufferPool) (
 	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/buffer_pool.go

@@ -0,0 +1,64 @@
+//  Copyright (c) 2023 Couchbase, Inc.
+//
+// Licensed 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 s2
+
+// GeoBufferPool represents a pool of buffers ranging from a given
+// max size to a min size in steps of 2. It uses a lazy approach only allocating
+// the buffers when it is needed.
+
+type GeoBufferPool struct {
+	buffers [][]byte
+	maxSize int
+	minSize int
+}
+
+func NewGeoBufferPool(maxSize int, minSize int) *GeoBufferPool {
+
+
+	// Calculating the number of buffers required. Assuming that
+	// the value of minSize is correct, the buffers will be of size
+	// minSize, 2 * minSize, 4 * minSize and so on till it is less
+	// than or equal to the maxSize. If it is not equal to maxSize,
+	// then a suitable value less than maxSize will be set as maxSize
+	length := 0
+	temp := minSize
+	for temp <= maxSize {
+		length = length + 1
+		temp = temp * 2
+	}
+	maxSize = temp / 2
+
+	return &GeoBufferPool{
+		buffers: make([][]byte, length),
+		maxSize: maxSize,
+		minSize: minSize,
+	}
+}
+
+func (b *GeoBufferPool) Get(size int) ([]byte) {
+
+	bufSize := b.maxSize
+
+	for i := range b.buffers {
+		if bufSize <= size {
+			b.buffers[i] = make([]byte, bufSize)
+			return b.buffers[i]
+		} else {
+			bufSize = bufSize / 2
+		}
+	}
+
+	return nil
+}

s2/encode.go

@@ -222,3 +222,13 @@ func (d *decoder) readUvarint() (x uint64) {
 	x, d.err = binary.ReadUvarint(d.r)
 	return
 }
+
+func (d *decoder) readFloat64Array(size int, buf []byte) ([]byte, int) {
+	if d.err != nil || buf == nil {
+		return nil, 0
+	}
+
+	_, d.err = io.ReadFull(d.r, buf)
+
+	return buf, len(buf)
+}

s2/loop.go

@@ -15,6 +15,7 @@
 package s2
 
 import (
+	"encoding/binary"
 	"fmt"
 	"io"
 	"math"
@@ -66,6 +67,9 @@ type Loop struct {
 
 	// index is the spatial index for this Loop.
 	index *ShapeIndex
+
+	// A buffer pool to be used while decoding the polygon
+	BufPool *GeoBufferPool
 }
 
 // LoopFromPoints constructs a loop from the given points.
@@ -1287,11 +1291,31 @@ 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()
+
+	// Each vertex requires 24 bytes of storage
+	numBytesNeeded := int(nvertices) * 24
+
+	i := 0
+
+	for numBytesNeeded > 0 {
+		arr := l.BufPool.Get(numBytesNeeded)
+		arr, numBytesRead := d.readFloat64Array(numBytesNeeded, arr)
+
+		if numBytesRead == 0 {
+			break
+		}
+
+		numBytesNeeded = numBytesNeeded - numBytesRead 
+
+		for j := 0; j < int(numBytesRead/24); 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(numBytesRead/24)
 	}
+
 	l.index = NewShapeIndex()
 	l.originInside = d.readBool()
 	l.depth = int(d.readUint32())

s2/polygon.go

@@ -77,6 +77,9 @@ 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
+
+	// A buffer pool to be used while decoding the polygon
+	BufPool *GeoBufferPool
 }
 
 // PolygonFromLoops constructs a polygon from the given set of loops. The polygon
@@ -1133,7 +1136,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 +1154,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].BufPool = p.BufPool
 		p.loops[i].decode(d)
 		p.numVertices += len(p.loops[i].vertices)
 	}