Adding upstream version 2.5.1.

Signed-off-by: Daniel Baumann <daniel@debian.org>

search/searcher/search_geoshape.go

@@ -0,0 +1,135 @@
+//  Copyright (c) 2022 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 searcher
+
+import (
+	"bytes"
+	"context"
+
+	"github.com/blevesearch/bleve/v2/geo"
+	"github.com/blevesearch/bleve/v2/search"
+	index "github.com/blevesearch/bleve_index_api"
+	"github.com/blevesearch/geo/geojson"
+	"github.com/blevesearch/geo/s2"
+)
+
+func NewGeoShapeSearcher(ctx context.Context, indexReader index.IndexReader, shape index.GeoJSON,
+	relation string, field string, boost float64,
+	options search.SearcherOptions,
+) (search.Searcher, error) {
+	var err error
+	var spatialPlugin index.SpatialAnalyzerPlugin
+
+	// check for the spatial plugin from the index.
+	if sr, ok := indexReader.(index.SpatialIndexPlugin); ok {
+		spatialPlugin, _ = sr.GetSpatialAnalyzerPlugin("s2")
+	}
+
+	if spatialPlugin == nil {
+		// fallback to the default spatial plugin(s2).
+		spatialPlugin = geo.GetSpatialAnalyzerPlugin("s2")
+	}
+
+	// obtain the query tokens.
+	terms := spatialPlugin.GetQueryTokens(shape)
+	mSearcher, err := NewMultiTermSearcher(ctx, indexReader, terms,
+		field, boost, options, false)
+	if err != nil {
+		return nil, err
+	}
+
+	dvReader, err := indexReader.DocValueReader([]string{field})
+	if err != nil {
+		return nil, err
+	}
+
+	return NewFilteringSearcher(ctx, mSearcher, buildRelationFilterOnShapes(ctx, dvReader, field, relation, shape)), nil
+}
+
+// Using the same term splitter slice used in the doc values in zap.
+// TODO: This needs to be revisited whenever we change the zap
+// implementation of doc values.
+var termSeparatorSplitSlice = []byte{0xff}
+
+func buildRelationFilterOnShapes(ctx context.Context, dvReader index.DocValueReader, field string,
+	relation string, shape index.GeoJSON,
+) FilterFunc {
+	// this is for accumulating the shape's actual complete value
+	// spread across multiple docvalue visitor callbacks.
+	var dvShapeValue []byte
+	var startReading, finishReading bool
+	var reader *bytes.Reader
+
+	var bufPool *s2.GeoBufferPool
+	if bufPoolCallback, ok := ctx.Value(search.GeoBufferPoolCallbackKey).(search.GeoBufferPoolCallbackFunc); ok {
+		bufPool = bufPoolCallback()
+	}
+
+	return func(d *search.DocumentMatch) bool {
+		var found bool
+
+		err := dvReader.VisitDocValues(d.IndexInternalID,
+			func(field string, term []byte) {
+				// only consider the values which are GlueBytes prefixed or
+				// if it had already started reading the shape bytes from previous callbacks.
+				if startReading || len(term) > geo.GlueBytesOffset {
+
+					if !startReading && bytes.Equal(geo.GlueBytes, term[:geo.GlueBytesOffset]) {
+						startReading = true
+
+						if bytes.Equal(geo.GlueBytes, term[len(term)-geo.GlueBytesOffset:]) {
+							term = term[:len(term)-geo.GlueBytesOffset]
+							finishReading = true
+						}
+
+						dvShapeValue = append(dvShapeValue, term[geo.GlueBytesOffset:]...)
+
+					} else if startReading && !finishReading {
+						if len(term) > geo.GlueBytesOffset &&
+							bytes.Equal(geo.GlueBytes, term[len(term)-geo.GlueBytesOffset:]) {
+							term = term[:len(term)-geo.GlueBytesOffset]
+							finishReading = true
+						}
+
+						term = append(termSeparatorSplitSlice, term...)
+						dvShapeValue = append(dvShapeValue, term...)
+					}
+
+					// apply the filter once the entire docvalue is finished reading.
+					if finishReading {
+						v, err := geojson.FilterGeoShapesOnRelation(shape, dvShapeValue, relation, &reader, bufPool)
+						if err == nil && v {
+							found = true
+						}
+
+						dvShapeValue = dvShapeValue[:0]
+						startReading = false
+						finishReading = false
+					}
+				}
+			})
+
+		if err == nil && found {
+			bytes := dvReader.BytesRead()
+			if bytes > 0 {
+				reportIOStats(ctx, bytes)
+				search.RecordSearchCost(ctx, search.AddM, bytes)
+			}
+			return found
+		}
+
+		return false
+	}
+}