135 lines
4.1 KiB
Go
135 lines
4.1 KiB
Go
// 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
|
|
}
|
|
}
|