diff --git a/document/field_geoshape_v2.go b/document/field_geoshape_v2.go new file mode 100644 index 000000000..60614d4e4 --- /dev/null +++ b/document/field_geoshape_v2.go @@ -0,0 +1,162 @@ +// Copyright (c) 2026 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 document + +import ( + "reflect" + + "github.com/blevesearch/bleve/v2/geo" + "github.com/blevesearch/bleve/v2/geov2" + "github.com/blevesearch/bleve/v2/size" + index "github.com/blevesearch/bleve_index_api" + "github.com/blevesearch/geo/geojson" +) + +var reflectStaticSizeGeoShapeV2Field int + +func init() { + var f GeoShapeV2Field + reflectStaticSizeGeoShapeV2Field = int(reflect.TypeOf(f).Size()) +} + +type GeoShapeV2Field struct { + name string + + shape index.GeoJSON + inner []uint64 + cross []uint64 + scoreInner uint64 + scoreCross uint64 + bBoxBytes []byte + shapeBytes []byte + + options index.FieldIndexingOptions +} + +func (f *GeoShapeV2Field) Name() string { + return f.name +} + +func (f *GeoShapeV2Field) ArrayPositions() []uint64 { + return nil +} + +func (f *GeoShapeV2Field) Options() index.FieldIndexingOptions { + return f.options +} + +func (f *GeoShapeV2Field) Analyze() { + f.inner, f.cross = f.shape.IndexCells() + f.scoreInner = geov2.CalcCellsScore(f.inner) + f.scoreCross = geov2.CalcCellsScore(f.cross) + + if bBox, ok := f.shape.BoundingBox().(*geojson.Envelope); ok { + bBoxBytes, err := bBox.Marshal() + if err != nil { + return + } + f.bBoxBytes = bBoxBytes + } +} + +func (f *GeoShapeV2Field) Value() []byte { + return []byte{} +} + +func (f *GeoShapeV2Field) NumPlainTextBytes() uint64 { + return 0 +} + +func (f *GeoShapeV2Field) Size() int { + return reflectStaticSizeGeoShapeV2Field + size.SizeOfPtr + + len(f.name) + + len(f.inner)*size.SizeOfUint64 + + len(f.cross)*size.SizeOfUint64 + + len(f.bBoxBytes) + + len(f.shapeBytes) +} + +func (f *GeoShapeV2Field) EncodedFieldType() byte { + return 'o' +} + +func (f *GeoShapeV2Field) AnalyzedLength() int { + return 0 +} + +func (f *GeoShapeV2Field) AnalyzedTokenFrequencies() index.TokenFrequencies { + return nil +} + +func (f *GeoShapeV2Field) InnerCells() []uint64 { + return f.inner +} + +func (f *GeoShapeV2Field) CrossCells() []uint64 { + return f.cross +} + +func (f *GeoShapeV2Field) EncodedBoundingBox() []byte { + return f.bBoxBytes +} + +func (f *GeoShapeV2Field) EncodedShape() []byte { + return f.shapeBytes +} + +func (f *GeoShapeV2Field) Scores() (uint64, uint64) { + return f.scoreInner, f.scoreCross +} + +func NewGeoShapeV2FieldFromShapeWithIndexingOptions(name string, geoShape *geojson.GeoShape, + options index.FieldIndexingOptions) *GeoShapeV2Field { + + var shape index.GeoJSON + var shapeBytes []byte + var err error + + if geoShape.Type == geo.CircleType { + shape, shapeBytes, err = geo.NewGeoCircleShape(geoShape.Center, + geoShape.Radius) + } else { + shape, shapeBytes, err = geo.NewGeoJsonShape(geoShape.Coordinates, + geoShape.Type) + } + if err != nil { + return nil + } + + return &GeoShapeV2Field{ + name: name, + shape: shape, + options: options, + shapeBytes: shapeBytes, + } +} + +func NewGeometryCollectionV2FieldFromShapesWithIndexingOptions(name string, + geoShapes []*geojson.GeoShape, options index.FieldIndexingOptions) *GeoShapeV2Field { + shape, shapeBytes, err := geo.NewGeometryCollectionFromShapes(geoShapes) + if err != nil { + return nil + } + + return &GeoShapeV2Field{ + name: name, + shape: shape, + options: options, + shapeBytes: shapeBytes, + } +} diff --git a/geo/geo_s2plugin_impl.go b/geo/geo_s2plugin_impl.go index 6acac5a16..95434dbde 100644 --- a/geo/geo_s2plugin_impl.go +++ b/geo/geo_s2plugin_impl.go @@ -246,6 +246,21 @@ func (p *Point) QueryTokens(s *S2SpatialAnalyzerPlugin) []string { return nil } +func (p *Point) IndexCells() ([]uint64, []uint64) { + // placeholder implementation + return nil, nil +} + +func (p *Point) QueryCells() ([]uint64, []uint64) { + // placeholder implementation + return nil, nil +} + +func (p *Point) BoundingBox() index.GeoJSON { + // placeholder implementation + return nil +} + //---------------------------------------------------------------------------------- type boundedRectangle struct { @@ -293,6 +308,21 @@ func (br *boundedRectangle) QueryTokens(s *S2SpatialAnalyzerPlugin) []string { return geojson.StripCoveringTerms(terms) } +func (br *boundedRectangle) IndexCells() ([]uint64, []uint64) { + // placeholder implementation + return nil, nil +} + +func (br *boundedRectangle) QueryCells() ([]uint64, []uint64) { + // placeholder implementation + return nil, nil +} + +func (br *boundedRectangle) BoundingBox() index.GeoJSON { + // placeholder implementation + return nil +} + //---------------------------------------------------------------------------------- type boundedPolygon struct { @@ -341,6 +371,21 @@ func (bp *boundedPolygon) QueryTokens(s *S2SpatialAnalyzerPlugin) []string { return geojson.StripCoveringTerms(terms) } +func (bp *boundedPolygon) IndexCells() ([]uint64, []uint64) { + // placeholder implementation + return nil, nil +} + +func (bp *boundedPolygon) QueryCells() ([]uint64, []uint64) { + // placeholder implementation + return nil, nil +} + +func (bp *boundedPolygon) BoundingBox() index.GeoJSON { + // placeholder implementation + return nil +} + //---------------------------------------------------------------------------------- type pointDistance struct { @@ -389,6 +434,21 @@ func (pd *pointDistance) QueryTokens(s *S2SpatialAnalyzerPlugin) []string { return geojson.StripCoveringTerms(terms) } +func (pd *pointDistance) IndexCells() ([]uint64, []uint64) { + // placeholder implementation + return nil, nil +} + +func (pd *pointDistance) QueryCells() ([]uint64, []uint64) { + // placeholder implementation + return nil, nil +} + +func (pd *pointDistance) BoundingBox() index.GeoJSON { + // placeholder implementation + return nil +} + // ------------------------------------------------------------------------ // NewGeometryCollection instantiate a geometrycollection diff --git a/geov2/cells.go b/geov2/cells.go new file mode 100644 index 000000000..3fac3d9b1 --- /dev/null +++ b/geov2/cells.go @@ -0,0 +1,49 @@ +// Copyright (c) 2026 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 geov2 + +import "github.com/blevesearch/geo/s2" + +// getCellSearchBounds takes a raw uint64 S2 cell ID and returns the +// minimum and maximum uint64 values defining its absolute spatial range. +func getCellSearchBounds(cellUint uint64) (min uint64, max uint64) { + cellID := s2.CellID(cellUint) + + rangeMin := cellID.RangeMin() + rangeMax := cellID.RangeMax() + + return uint64(rangeMin), uint64(rangeMax) +} + +// Returns the level of the given S2 cell ID +func getCellLevel(cell uint64) uint64 { + return uint64(s2.CellID(cell).Level()) +} + +// Returns the parent cell ID of the given S2 cell ID at the specified level +func getParentCell(cell uint64, level int) uint64 { + return uint64(s2.CellID(cell).Parent(level)) +} + +// CalcCellsScore returns the total area of the given S2 cells +// in level 16 cell units (maxCellLevel in the geo repo's region coverer +// configuration - see geo/geojson/geojson_v2.go) +func CalcCellsScore(cells []uint64) uint64 { + var score uint64 + for _, cell := range cells { + score += calcScore(0, getCellLevel(cell)) + } + return score +} diff --git a/geov2/evaluator.go b/geov2/evaluator.go new file mode 100644 index 000000000..40c9917c0 --- /dev/null +++ b/geov2/evaluator.go @@ -0,0 +1,122 @@ +// Copyright (c) 2026 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 geov2 + +import segment "github.com/blevesearch/scorch_segment_api/v2" + +type queryEvaluator struct { + queryInnerCells []uint64 + queryCrossCells []uint64 + + innerCells []uint64 + innerDocIDs []uint32 + + crossCells []uint64 + crossDocIDs []uint32 +} + +func NewQueryEvaluator(query Query, geoData segment.GeoShapeV2Data) *queryEvaluator { + return &queryEvaluator{ + queryInnerCells: query.InnerCells(), + queryCrossCells: query.CrossCells(), + innerCells: geoData.InnerCells(), + innerDocIDs: geoData.InnerDocIDs(), + crossCells: geoData.CrossCells(), + crossDocIDs: geoData.CrossDocIDs(), + } +} + +// find the leftmost index in arr where arr[i] >= target, or +// len(arr) if no such index exists +func binarySearchLeftmostGreaterOrEqual(arr []uint64, target uint64) int { + lo, hi := 0, len(arr) + for lo < hi { + mid := int(uint(lo+hi) >> 1) // bitshift avoids the addition overflow risk + if arr[mid] < target { + lo = mid + 1 + } else { + hi = mid + } + } + return lo +} + +// forEachScoredDoc invokes fn exactly once for every docID present in either +// score map, passing that doc's inner and cross scores (0 when the doc is +// absent from the corresponding map). +func forEachScoredDoc(innerScores, crossScores map[uint32]uint64, + fn func(id uint32, inner, cross uint64)) { + for id, inner := range innerScores { + fn(id, inner, crossScores[id]) + } + for id, cross := range crossScores { + if _, ok := innerScores[id]; !ok { + fn(id, 0, cross) + } + } +} + +// scan and score the overlap of query inner cells with all index cells +func (qe *queryEvaluator) rangeScanInner(innerScores, crossScores map[uint32]uint64) { + for _, cell := range qe.queryInnerCells { + minVal, maxVal := getCellSearchBounds(cell) + cellLevel := getCellLevel(cell) + rangeScanOne(cell, minVal, maxVal, cellLevel, qe.innerCells, qe.innerDocIDs, innerScores) + rangeScanOne(cell, minVal, maxVal, cellLevel, qe.crossCells, qe.crossDocIDs, crossScores) + } +} + +// scan and score the overlap of query cross cells with all index cells +func (qe *queryEvaluator) rangeScanCross(innerScores, crossScores map[uint32]uint64) { + for _, cell := range qe.queryCrossCells { + minVal, maxVal := getCellSearchBounds(cell) + cellLevel := getCellLevel(cell) + rangeScanOne(cell, minVal, maxVal, cellLevel, qe.innerCells, qe.innerDocIDs, innerScores) + rangeScanOne(cell, minVal, maxVal, cellLevel, qe.crossCells, qe.crossDocIDs, crossScores) + } +} + +// scan and score the overlap of a single query cell with the given index cells +func rangeScanOne(queryCell uint64, minVal, maxVal, cellLevel uint64, + indexCells []uint64, docIds []uint32, scores map[uint32]uint64) { + // find the range of index cells within the min/max bounds of the query cell + // end will be < start if there are no index cells within the bounds + start := binarySearchLeftmostGreaterOrEqual(indexCells, minVal) + end := binarySearchLeftmostGreaterOrEqual(indexCells, maxVal+1) - 1 + + // score all index cells within the bounds of the query cell + for i := start; i <= end; i++ { + id := docIds[i] + val := indexCells[i] + + valLevel := getCellLevel(val) + scores[id] += calcScore(cellLevel, valLevel) + } + + // score all parent cells of the query cell that are present in the index + // since parent cells are not within the min/max bounds + for level := int(cellLevel) - 1; level >= 0; level-- { + // get the parent cell of the query cell at this level + parentCell := getParentCell(queryCell, level) + // search for the leftmost index of this parent cell in the index cells + parentStart := binarySearchLeftmostGreaterOrEqual(indexCells, parentCell) + + // score all index cells that match this parent cell exactly + for i := parentStart; i < len(indexCells) && indexCells[i] == parentCell; i++ { + id := docIds[i] + scores[id] += calcScore(cellLevel, uint64(level)) + } + } +} diff --git a/geov2/evaluator_test.go b/geov2/evaluator_test.go new file mode 100644 index 000000000..0afed49c7 --- /dev/null +++ b/geov2/evaluator_test.go @@ -0,0 +1,128 @@ +// Copyright (c) 2026 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 geov2 + +import ( + "sort" + "testing" + + "github.com/blevesearch/geo/s2" +) + +func TestBinarySearchLeftmostGreaterOrEqual(t *testing.T) { + tests := []struct { + name string + arr []uint64 + target uint64 + want int + }{ + {"empty", nil, 5, 0}, + {"all less than target", []uint64{1, 2, 3}, 5, 3}, + {"all greater than target", []uint64{5, 6, 7}, 1, 0}, + {"exact match", []uint64{1, 3, 5, 7}, 5, 2}, + {"target between elements", []uint64{1, 3, 5, 7}, 4, 2}, + {"leftmost of duplicates", []uint64{1, 3, 3, 3, 5}, 3, 1}, + {"target before all", []uint64{2, 4, 6}, 0, 0}, + {"target equals last", []uint64{2, 4, 6}, 6, 2}, + {"single element hit", []uint64{9}, 9, 0}, + {"single element miss high", []uint64{9}, 10, 1}, + } + for _, test := range tests { + if got := binarySearchLeftmostGreaterOrEqual(test.arr, test.target); got != test.want { + t.Errorf("%s: binarySearchLeftmostGreaterOrEqual(%v, %d) = %d, want %d", + test.name, test.arr, test.target, got, test.want) + } + } +} + +// TestRangeScanOne verifies the two scoring paths of rangeScanOne: index +// cells that fall within the query cell's range (the query cell itself and +// its descendants), and ancestor cells found by walking the query cell's +// parents (which lie outside the range). +func TestRangeScanOne(t *testing.T) { + const queryLevel = 10 + + queryCell := uint64(s2.CellIDFromFace(2).ChildBeginAtLevel(queryLevel)) + childCell := uint64(s2.CellID(queryCell).ChildBeginAtLevel(12)) // descendant, in range + parentCell := uint64(s2.CellID(queryCell).Parent(5)) // ancestor, out of range + unrelated := uint64(s2.CellIDFromFace(4).ChildBeginAtLevel(queryLevel)) + + // assign each cell a distinct doc ID + type entry struct { + cell uint64 + docID uint32 + } + entries := []entry{ + {queryCell, 0}, + {childCell, 1}, + {parentCell, 2}, + {unrelated, 3}, + } + + // index cells must be sorted ascending, with docIds kept parallel + sort.Slice(entries, func(i, j int) bool { + return entries[i].cell < entries[j].cell + }) + indexCells := make([]uint64, len(entries)) + docIds := make([]uint32, len(entries)) + for i, e := range entries { + indexCells[i] = e.cell + docIds[i] = e.docID + } + + scores := make(map[uint32]uint64) + minVal, maxVal := getCellSearchBounds(queryCell) + rangeScanOne(queryCell, minVal, maxVal, queryLevel, indexCells, docIds, scores) + + // docID 0 is the query cell itself: in range, equal levels + if want := calcScore(queryLevel, queryLevel); scores[0] != want { + t.Errorf("query cell score = %d, want %d", scores[0], want) + } + // docID 1 is a descendant at level 12: in range, index cell deeper + if want := calcScore(queryLevel, 12); scores[1] != want { + t.Errorf("descendant cell score = %d, want %d", scores[1], want) + } + // docID 2 is an ancestor at level 5: found via the parent walk + if want := calcScore(queryLevel, 5); scores[2] != want { + t.Errorf("ancestor cell score = %d, want %d", scores[2], want) + } + // docID 3 is unrelated (different face): must not be scored + if scores[3] != 0 { + t.Errorf("unrelated cell score = %d, want 0", scores[3]) + } +} + +// TestRangeScanOneNoMatches confirms that a query cell with no overlapping +// index cells produces no scores, exercising the empty-range case where +// end < start. +func TestRangeScanOneNoMatches(t *testing.T) { + const queryLevel = 10 + queryCell := uint64(s2.CellIDFromFace(0).ChildBeginAtLevel(queryLevel)) + + // index cells entirely on a different, non-ancestor part of the tree + indexCells := []uint64{ + uint64(s2.CellIDFromFace(5).ChildBeginAtLevel(queryLevel)), + uint64(s2.CellIDFromFace(5).ChildBeginAtLevel(queryLevel).Next()), + } + docIds := []uint32{0, 1} + scores := make(map[uint32]uint64) + + minVal, maxVal := getCellSearchBounds(queryCell) + rangeScanOne(queryCell, minVal, maxVal, queryLevel, indexCells, docIds, scores) + + if len(scores) != 0 { + t.Errorf("expected no scores, got %d entries: %v", len(scores), scores) + } +} diff --git a/geov2/query.go b/geov2/query.go new file mode 100644 index 000000000..8222b4e30 --- /dev/null +++ b/geov2/query.go @@ -0,0 +1,42 @@ +// Copyright (c) 2026 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 geov2 + +import ( + "github.com/blevesearch/bleve/v2/util" + index "github.com/blevesearch/bleve_index_api" + segment "github.com/blevesearch/scorch_segment_api/v2" +) + +type Query interface { + Evaluate(geoData segment.GeoShapeV2Data) *util.Bitset + InnerCells() []uint64 + CrossCells() []uint64 +} + +func NewQuery(shape index.GeoJSON, relation string) Query { + switch relation { + case "contains": + return NewContainsQuery(shape) + case "intersects": + return NewIntersectsQuery(shape) + case "within": + return NewWithinQuery(shape) + case "disjoint": + return NewDisjointQuery(shape) + default: + return nil + } +} diff --git a/geov2/query_contains.go b/geov2/query_contains.go new file mode 100644 index 000000000..651d0daab --- /dev/null +++ b/geov2/query_contains.go @@ -0,0 +1,141 @@ +// Copyright (c) 2026 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 geov2 + +import ( + "bytes" + + "github.com/blevesearch/bleve/v2/util" + index "github.com/blevesearch/bleve_index_api" + "github.com/blevesearch/geo/geojson" + segment "github.com/blevesearch/scorch_segment_api/v2" +) + +type containsQuery struct { + innerCells []uint64 + crossCells []uint64 + + shape index.GeoJSON + bBox index.GeoJSON + + score uint64 +} + +func NewContainsQuery(shape index.GeoJSON) Query { + inner, cross := shape.QueryCells() + + score := CalcCellsScore(inner) + CalcCellsScore(cross) + + return &containsQuery{ + innerCells: inner, + crossCells: cross, + shape: shape, + bBox: shape.BoundingBox(), + score: score, + } +} + +func (cq *containsQuery) Evaluate(geoData segment.GeoShapeV2Data) *util.Bitset { + numDocs := int(geoData.NumDocs()) + exclude := geoData.Excluded() + + // create bitsets for hits and maybeHits providing exclude to the bitset + // which will make it impossible to set those bits + hits := util.NewBitset(numDocs, exclude) + maybeHits := util.NewBitset(numDocs, exclude) + + // failsafe for a degenerate query shape that produced no cells + if cq.score == 0 { + return hits + } + + // obtain zeroed score arrays from the segment-level pool and return + // them once the evaluation is done + innerScores := geoData.GetScoreArray() + crossScores := geoData.GetScoreArray() + defer geoData.PutScoreArray(innerScores) + defer geoData.PutScoreArray(crossScores) + + // create an evaluator instance to scan the query cells against the index cells + evaluator := NewQueryEvaluator(cq, geoData) + + // scan and score the overlap of all query cells with all index cells + evaluator.rangeScanInner(innerScores, crossScores) + evaluator.rangeScanCross(innerScores, crossScores) + + // if all of the query cells are contained within the inner index cells + // then we have a guaranteed hit, if they are contained within both the inner + // and cross index cells then we have a maybe hit, otherwise we have no hit + forEachScoredDoc(innerScores, crossScores, func(id uint32, inner, cross uint64) { + docNum := int(id) + if inner == cq.score { + hits.Add(docNum) + } else if inner+cross == cq.score { + maybeHits.Add(docNum) + } + }) + + var reader *bytes.Reader + + // filter out any maybeHits that do not have a bounding box that + // contains the query bounding box + boxFilter := func(docNum int) { + docBBoxBytes, err := geoData.BoundingBox(uint64(docNum)) + if docBBoxBytes == nil || err != nil { + return + } + + docBBox, err := geojson.ExtractShapesFromBytes(docBBoxBytes, &reader, nil) + if err != nil { + return + } + + if ok, err := docBBox.Contains(cq.bBox); err == nil && !ok { + maybeHits.Remove(docNum) + } + } + + maybeHits.Iterate(boxFilter) + + // filter out any maybeHits that do not have a shape that + // contains the query shape + shapeFilter := func(docNum int) { + docShapeBytes, err := geoData.Shape(uint64(docNum)) + if docShapeBytes == nil || err != nil { + return + } + + docShape, err := geojson.ExtractShapesFromBytes(docShapeBytes, &reader, nil) + if err != nil { + return + } + + if ok, err := docShape.Contains(cq.shape); err == nil && ok { + hits.Add(docNum) + } + } + + maybeHits.Iterate(shapeFilter) + + return hits +} + +func (cq *containsQuery) InnerCells() []uint64 { + return cq.innerCells +} + +func (cq *containsQuery) CrossCells() []uint64 { + return cq.crossCells +} diff --git a/geov2/query_disjoint.go b/geov2/query_disjoint.go new file mode 100644 index 000000000..f45d80cde --- /dev/null +++ b/geov2/query_disjoint.go @@ -0,0 +1,66 @@ +// Copyright (c) 2026 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 geov2 + +import ( + "github.com/blevesearch/bleve/v2/util" + index "github.com/blevesearch/bleve_index_api" + segment "github.com/blevesearch/scorch_segment_api/v2" +) + +type disjointQuery struct { + innerCells []uint64 + crossCells []uint64 + + shape index.GeoJSON + bBox index.GeoJSON +} + +func NewDisjointQuery(shape index.GeoJSON) Query { + inner, cross := shape.QueryCells() + + return &disjointQuery{ + innerCells: inner, + crossCells: cross, + shape: shape, + bBox: shape.BoundingBox(), + } +} + +func (dq *disjointQuery) Evaluate(geoData segment.GeoShapeV2Data) *util.Bitset { + // evaluate the disjoint query by creating an intersects query and negating the results + intersectsQuery := &intersectsQuery{ + innerCells: dq.innerCells, + crossCells: dq.crossCells, + shape: dq.shape, + bBox: dq.bBox, + } + + // evaluate the intersects query to get the hits + hits := intersectsQuery.Evaluate(geoData) + + // invert the hits to get the disjoint results + hits.Invert() + + return hits +} + +func (dq *disjointQuery) InnerCells() []uint64 { + return dq.innerCells +} + +func (dq *disjointQuery) CrossCells() []uint64 { + return dq.crossCells +} diff --git a/geov2/query_intersects.go b/geov2/query_intersects.go new file mode 100644 index 000000000..7c89f93a3 --- /dev/null +++ b/geov2/query_intersects.go @@ -0,0 +1,145 @@ +// Copyright (c) 2026 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 geov2 + +import ( + "bytes" + + "github.com/blevesearch/bleve/v2/util" + index "github.com/blevesearch/bleve_index_api" + "github.com/blevesearch/geo/geojson" + segment "github.com/blevesearch/scorch_segment_api/v2" +) + +type intersectsQuery struct { + innerCells []uint64 + crossCells []uint64 + + shape index.GeoJSON + bBox index.GeoJSON +} + +func NewIntersectsQuery(shape index.GeoJSON) Query { + inner, cross := shape.QueryCells() + + return &intersectsQuery{ + innerCells: inner, + crossCells: cross, + shape: shape, + bBox: shape.BoundingBox(), + } +} + +func (iq *intersectsQuery) Evaluate(geoData segment.GeoShapeV2Data) *util.Bitset { + numDocs := int(geoData.NumDocs()) + exclude := geoData.Excluded() + + // create bitsets for hits and maybeHits providing exclude to the bitset + // which will make it impossible to set those bits + hits := util.NewBitset(numDocs, exclude) + maybeHits := util.NewBitset(numDocs, exclude) + + // obtain zeroed score arrays from the segment-level pool and return + // them once the evaluation is done + innerScores := geoData.GetScoreArray() + crossScores := geoData.GetScoreArray() + defer geoData.PutScoreArray(innerScores) + defer geoData.PutScoreArray(crossScores) + + // create an evaluator instance to scan the query cells against the index cells + evaluator := NewQueryEvaluator(iq, geoData) + + // scan and score the overlap of query inner cells with all index cells + evaluator.rangeScanInner(innerScores, crossScores) + + // if there is any overlap of query inner cells with any of the index cells + // then we have a guaranteed hit. Reset scores to reuse score maps for the + // next step + forEachScoredDoc(innerScores, crossScores, func(id uint32, inner, cross uint64) { + if inner > 0 || cross > 0 { + hits.Add(int(id)) + } + }) + clear(innerScores) + clear(crossScores) + + // scan and score the overlap of query cross cells with all index cells + evaluator.rangeScanCross(innerScores, crossScores) + + // if there is any overlap of query cross cells with any of the index inner + // cells then we have a guaranteed hit, if there is any overlap of query cross + // cells with any of the index cross cells then we have a maybe hit, otherwise + // we have no hit. + forEachScoredDoc(innerScores, crossScores, func(id uint32, inner, cross uint64) { + docNum := int(id) + if inner > 0 && !hits.Contains(docNum) { + hits.Add(docNum) + } else if cross > 0 && !hits.Contains(docNum) { + maybeHits.Add(docNum) + } + }) + + var reader *bytes.Reader + + // filter out any maybeHits that do not have a bounding box that + // intersects the query bounding box + boxFilter := func(docNum int) { + docBBoxBytes, err := geoData.BoundingBox(uint64(docNum)) + if docBBoxBytes == nil || err != nil { + return + } + + docBBox, err := geojson.ExtractShapesFromBytes(docBBoxBytes, &reader, nil) + if err != nil { + return + } + + if ok, err := docBBox.Intersects(iq.bBox); err == nil && !ok { + maybeHits.Remove(docNum) + } + } + + maybeHits.Iterate(boxFilter) + + // filter out any maybeHits that do not have a shape that + // intersects the query shape + shapeFilter := func(docNum int) { + docShapeBytes, err := geoData.Shape(uint64(docNum)) + if docShapeBytes == nil || err != nil { + return + } + + docShape, err := geojson.ExtractShapesFromBytes(docShapeBytes, &reader, nil) + if err != nil { + return + } + + if ok, err := docShape.Intersects(iq.shape); err == nil && ok { + hits.Add(docNum) + } + } + + maybeHits.Iterate(shapeFilter) + + return hits +} + +func (iq *intersectsQuery) InnerCells() []uint64 { + return iq.innerCells +} + +func (iq *intersectsQuery) CrossCells() []uint64 { + return iq.crossCells +} diff --git a/geov2/query_within.go b/geov2/query_within.go new file mode 100644 index 000000000..7f848a7b6 --- /dev/null +++ b/geov2/query_within.go @@ -0,0 +1,146 @@ +// Copyright (c) 2026 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 geov2 + +import ( + "bytes" + + "github.com/blevesearch/bleve/v2/util" + index "github.com/blevesearch/bleve_index_api" + "github.com/blevesearch/geo/geojson" + segment "github.com/blevesearch/scorch_segment_api/v2" +) + +type withinQuery struct { + innerCells []uint64 + crossCells []uint64 + + shape index.GeoJSON + bBox index.GeoJSON +} + +func NewWithinQuery(shape index.GeoJSON) Query { + inner, cross := shape.QueryCells() + + return &withinQuery{ + innerCells: inner, + crossCells: cross, + shape: shape, + bBox: shape.BoundingBox(), + } +} + +func (wq *withinQuery) Evaluate(geoData segment.GeoShapeV2Data) *util.Bitset { + numDocs := int(geoData.NumDocs()) + exclude := geoData.Excluded() + + // create bitsets for hits and maybeHits providing exclude to the bitset + // which will make it impossible to set those bits + hits := util.NewBitset(numDocs, exclude) + maybeHits := util.NewBitset(numDocs, exclude) + + // obtain zeroed score arrays from the segment-level pool and return + // them once the evaluation is done + innerScores := geoData.GetScoreArray() + crossScores := geoData.GetScoreArray() + defer geoData.PutScoreArray(innerScores) + defer geoData.PutScoreArray(crossScores) + + docScoresInner, docScoresCross := geoData.DocScores() + + // create an evaluator instance to scan the query cells against the index cells + evaluator := NewQueryEvaluator(wq, geoData) + + // scan and score the overlap of query inner cells with all index cells + evaluator.rangeScanInner(innerScores, crossScores) + + // if all of the index cells are contained within the query inner cells, + // then we have a guaranteed hit. Only consider documents with non zero + // total scores as hits + forEachScoredDoc(innerScores, crossScores, func(id uint32, inner, cross uint64) { + total := inner + cross + if total == docScoresInner[id]+docScoresCross[id] && total != 0 { + hits.Add(int(id)) + } + }) + + // scan and score the overlap of query cross cells with all index cells + evaluator.rangeScanCross(innerScores, crossScores) + + // A document is a maybe-hit once the accumulated score reaches at least + // its own inner-cell score. Complete overlap of index and cross scores + // may not always be possible due to the nature of region coverer being + // non exhaustive in the sense that the boundaries are not always only + // covered by the smallest possible cells. + forEachScoredDoc(innerScores, crossScores, func(id uint32, inner, cross uint64) { + docNum := int(id) + total := inner + cross + if !hits.Contains(docNum) && total >= docScoresInner[id] && total != 0 { + maybeHits.Add(docNum) + } + }) + + var reader *bytes.Reader + + // filter out any maybeHits that do not have a bounding box that + // is within the query bounding box + boxFilter := func(docNum int) { + docBBoxBytes, err := geoData.BoundingBox(uint64(docNum)) + if docBBoxBytes == nil || err != nil { + return + } + + docBBox, err := geojson.ExtractShapesFromBytes(docBBoxBytes, &reader, nil) + if err != nil { + return + } + + if ok, err := wq.bBox.Contains(docBBox); err == nil && !ok { + maybeHits.Remove(docNum) + } + } + + maybeHits.Iterate(boxFilter) + + // filter out any maybeHits that do not have a shape that + // is within the query shape + shapeFilter := func(docNum int) { + docShapeBytes, err := geoData.Shape(uint64(docNum)) + if docShapeBytes == nil || err != nil { + return + } + + docShape, err := geojson.ExtractShapesFromBytes(docShapeBytes, &reader, nil) + if err != nil { + return + } + + if ok, err := wq.shape.Contains(docShape); err == nil && ok { + hits.Add(docNum) + } + } + + maybeHits.Iterate(shapeFilter) + + return hits +} + +func (wq *withinQuery) InnerCells() []uint64 { + return wq.innerCells +} + +func (wq *withinQuery) CrossCells() []uint64 { + return wq.crossCells +} diff --git a/geov2/score.go b/geov2/score.go new file mode 100644 index 000000000..920fe5cc7 --- /dev/null +++ b/geov2/score.go @@ -0,0 +1,74 @@ +// Copyright (c) 2026 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 geov2 + +// pow4Table precomputes all valid powers of 4 for a uint64. +// Array size 32 covers exp 0 through 31. +var pow4Table = [32]uint64{ + 1, // 4^0 + 4, // 4^1 + 16, // 4^2 + 64, // 4^3 + 256, // 4^4 + 1024, // 4^5 + 4096, // 4^6 + 16384, // 4^7 + 65536, // 4^8 + 262144, // 4^9 + 1048576, // 4^10 + 4194304, // 4^11 + 16777216, // 4^12 + 67108864, // 4^13 + 268435456, // 4^14 + 1073741824, // 4^15 + 4294967296, // 4^16 + 17179869184, // 4^17 + 68719476736, // 4^18 + 274877906944, // 4^19 + 1099511627776, // 4^20 + 4398046511104, // 4^21 + 17592186044416, // 4^22 + 70368744177664, // 4^23 + 281474976710656, // 4^24 + 1125899906842624, // 4^25 + 4503599627370496, // 4^26 + 18014398509481984, // 4^27 + 72057594037927936, // 4^28 + 288230376151711744, // 4^29 + 1152921504606846976, // 4^30 + 4611686018427387904, // 4^31 +} + +// pow4 returns 4^exp quickly using the lookup table. +func pow4(exp uint64) uint64 { + if exp >= 32 { + // Handle overflow safely. + return 0 + } + return pow4Table[exp] +} + +// returns the overlap of query and index cells based on their levels. +// Both levels are assumed to be at most 16 (maxCellLevel in the geo repo's +// region coverer configuration - see geo/geojson/geojson_v2.go), the deepest +// level used across the geoshape_v2 indexing and query cell coverings; +// cells deeper than level 16 are outside this function's contract by design. +func calcScore(queryCellLevel, indexCellLevel uint64) uint64 { + if indexCellLevel > queryCellLevel { + return pow4(16 - indexCellLevel) + } else { + return pow4(16 - queryCellLevel) + } +} diff --git a/geov2/score_test.go b/geov2/score_test.go new file mode 100644 index 000000000..650f0fe32 --- /dev/null +++ b/geov2/score_test.go @@ -0,0 +1,117 @@ +// Copyright (c) 2026 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 geov2 + +import ( + "testing" + + "github.com/blevesearch/geo/s2" +) + +func TestPow4(t *testing.T) { + tests := []struct { + exp uint64 + want uint64 + }{ + {0, 1}, + {1, 4}, + {10, 1048576}, + {15, 1073741824}, + {31, 4611686018427387904}, // last representable power of 4 in a uint64 + {32, 0}, // out of range, guarded to 0 + {100, 0}, // well out of range, guarded to 0 + } + for _, test := range tests { + if got := pow4(test.exp); got != test.want { + t.Errorf("pow4(%d) = %d, want %d", test.exp, got, test.want) + } + } +} + +func TestCalcScore(t *testing.T) { + tests := []struct { + name string + queryCellLevel uint64 + indexCellLevel uint64 + want uint64 + }{ + { + // equal levels: the overlap is a single cell at that level + name: "equal levels", queryCellLevel: 16, indexCellLevel: 16, want: pow4(0), + }, + { + name: "equal levels shallow", queryCellLevel: 6, indexCellLevel: 6, want: pow4(10), + }, + { + // index cell is deeper (smaller): the overlap is the index cell + name: "index deeper", queryCellLevel: 6, indexCellLevel: 10, want: pow4(6), + }, + { + // query cell is deeper (smaller): the overlap is the query cell + name: "query deeper", queryCellLevel: 11, indexCellLevel: 8, want: pow4(5), + }, + { + name: "both at level 0", queryCellLevel: 0, indexCellLevel: 0, want: pow4(16), + }, + } + for _, test := range tests { + if got := calcScore(test.queryCellLevel, test.indexCellLevel); got != test.want { + t.Errorf("%s: calcScore(%d, %d) = %d, want %d", + test.name, test.queryCellLevel, test.indexCellLevel, got, test.want) + } + } +} + +// cellAtLevel returns a valid S2 cell ID at the requested level, derived +// deterministically from a fixed face so tests do not depend on any RNG. +func cellAtLevel(level int) uint64 { + return uint64(s2.CellIDFromFace(1).ChildBeginAtLevel(level)) +} + +func TestCalcCellsScore(t *testing.T) { + // build cells at known levels; each contributes calcScore(0, level) + // which equals pow4(16 - level) + cells := []uint64{ + cellAtLevel(16), // contributes pow4(0) = 1 + cellAtLevel(15), // contributes pow4(1) = 4 + cellAtLevel(14), // contributes pow4(2) = 16 + } + want := pow4(0) + pow4(1) + pow4(2) + if got := CalcCellsScore(cells); got != want { + t.Fatalf("CalcCellsScore = %d, want %d", got, want) + } + + // an empty slice scores zero + if got := CalcCellsScore(nil); got != 0 { + t.Fatalf("CalcCellsScore(nil) = %d, want 0", got) + } +} + +func TestCellLevelAndParent(t *testing.T) { + cell := cellAtLevel(12) + if got := getCellLevel(cell); got != 12 { + t.Fatalf("getCellLevel = %d, want 12", got) + } + parent := getParentCell(cell, 8) + if got := getCellLevel(parent); got != 8 { + t.Fatalf("getCellLevel(parent) = %d, want 8", got) + } + // the parent's range must contain the child cell + minVal, maxVal := getCellSearchBounds(parent) + if cell < minVal || cell > maxVal { + t.Fatalf("expected child cell %d to fall within parent bounds [%d, %d]", + cell, minVal, maxVal) + } +} diff --git a/index/scorch/snapshot_index.go b/index/scorch/snapshot_index.go index d4242c460..a363ecb71 100644 --- a/index/scorch/snapshot_index.go +++ b/index/scorch/snapshot_index.go @@ -27,6 +27,8 @@ import ( "github.com/RoaringBitmap/roaring/v2" "github.com/blevesearch/bleve/v2/document" + geov2 "github.com/blevesearch/bleve/v2/geov2" + "github.com/blevesearch/bleve/v2/size" "github.com/blevesearch/bleve/v2/util" index "github.com/blevesearch/bleve_index_api" segment "github.com/blevesearch/scorch_segment_api/v2" @@ -50,6 +52,8 @@ type asynchSegmentResult struct { } var reflectStaticSizeIndexSnapshot int +var reflectStaticSizeIndexSnapshotGeoShapeV2Reader int +var reflectStaticSizeRoaringIntIterator int func init() { var is interface{} = IndexSnapshot{} @@ -63,6 +67,10 @@ func init() { if err != nil { panic(fmt.Errorf("levenshtein automaton ed2 builder err: %v", err)) } + var gcr IndexSnapshotGeoShapeV2Reader + reflectStaticSizeIndexSnapshotGeoShapeV2Reader = int(reflect.TypeOf(gcr).Size()) + var rip roaring.IntIterator + reflectStaticSizeRoaringIntIterator = int(reflect.TypeOf(rip).Size()) } type IndexSnapshot struct { @@ -1310,3 +1318,187 @@ func (i *IndexSnapshot) Ancestors(ID index.IndexInternalID, prealloc []index.Anc // return adjusted ancestors return prealloc, nil } + +func (i *IndexSnapshot) GeoShapeV2FieldReader(ctx context.Context, field string) ( + index.GeoShapeV2FieldReader, error) { + + rv := &IndexSnapshotGeoShapeV2Reader{ + field: field, + postings: make([]*roaring.Bitmap, len(i.segment)), + iterators: make([]roaring.IntPeekable, len(i.segment)), + snapshot: i, + } + + return rv, nil +} + +type IndexSnapshotGeoShapeV2Reader struct { + field string + + postings []*roaring.Bitmap + iterators []roaring.IntPeekable + segmentOffset int + + snapshot *IndexSnapshot +} + +// Search performs a spatial search for the given GeoJSON shape and relation +// across all segments in the index snapshot. +func (g *IndexSnapshotGeoShapeV2Reader) Search(shape index.GeoJSON, + relation string) error { + + numSegments := len(g.snapshot.segment) + // create a single query object that is thread safe + // to be used across all segments + query := geov2.NewQuery(shape, relation) + + var wg sync.WaitGroup + wg.Add(numSegments) + + var errm sync.Mutex + var err error + // search each segment concurrently + for i := 0; i < numSegments; i++ { + go func(segID int) { + defer wg.Done() + err2 := g.searchSeg(segID, query) + if err2 != nil { + errm.Lock() + if err == nil { + err = err2 + } + errm.Unlock() + } + }(i) + } + wg.Wait() + + return err +} + +// searchSeg performs a spatial search for the given GeoJSON shape and relation +// on a single segment in the index snapshot. +func (g *IndexSnapshotGeoShapeV2Reader) searchSeg(segID int, + query geov2.Query) error { + + snapshot := g.snapshot.segment[segID] + geoSeg, ok := snapshot.segment.(segment.GeoShapeV2Segment) + if !ok { + return nil + } + + // obtain the geo shape data from the segment + geoData, err := geoSeg.GeoShapeV2Data(g.field, snapshot.deleted) + if err != nil { + return err + } + // return if segment does not have any geo shape data for the field + if geoData == nil { + return nil + } + // release the reference on the segment's cached geo data once the + // evaluation is done, so that the cache is free to evict it + defer geoData.Close() + + // evaluate the query against the geo shape data to + // get the matching document IDs + hits := query.Evaluate(geoData) + postings := roaring.New() + + docNums := geoData.DocNums() + + addFunc := func(docNumInternal int) { + postings.Add(docNums[docNumInternal]) + } + + hits.Iterate(addFunc) + + g.postings[segID] = postings + g.iterators[segID] = postings.Iterator() + + return nil +} + +// Next returns the next GeoShapeV2FieldDoc from the postings list across all segments in the index snapshot. +func (g *IndexSnapshotGeoShapeV2Reader) Next(preAlloced *index.GeoShapeV2FieldDoc) ( + *index.GeoShapeV2FieldDoc, error) { + rv := preAlloced + if rv == nil { + rv = &index.GeoShapeV2FieldDoc{} + } + + for g.segmentOffset < len(g.iterators) { + if !g.iterators[g.segmentOffset].HasNext() { + g.segmentOffset++ + continue + } + + next := g.iterators[g.segmentOffset].Next() + globalOffset := g.snapshot.offsets[g.segmentOffset] + rv.ID = index.NewIndexInternalID(rv.ID, uint64(next)+globalOffset) + return rv, nil + } + + return nil, nil +} + +// Advance moves the reader to the specified document ID, returning the corresponding GeoShapeV2FieldDoc if it exists. +func (g *IndexSnapshotGeoShapeV2Reader) Advance(ID index.IndexInternalID, + preAlloced *index.GeoShapeV2FieldDoc) (*index.GeoShapeV2FieldDoc, error) { + rv := preAlloced + if rv == nil { + rv = &index.GeoShapeV2FieldDoc{} + } + + num := ID.Value() + + segIdx, localDocNum := g.snapshot.segmentIndexAndLocalDocNumFromGlobal(num) + if segIdx >= len(g.iterators) { + return nil, fmt.Errorf("error advancing to doc number %d, segment "+ + "index %d out of bounds", num, segIdx) + } + + if g.segmentOffset > segIdx { + return nil, fmt.Errorf("error advancing to doc number %d, segment "+ + "index %d is less than current segment offset %d", num, segIdx, g.segmentOffset) + } + + g.segmentOffset = segIdx + g.iterators[g.segmentOffset].AdvanceIfNeeded(uint32(localDocNum)) + + return g.Next(rv) +} + +// Close is a no-op: the reader holds no resources of its own, and the +// segment-level geo data it reads from is owned and evicted by the +// segment's cache +func (g *IndexSnapshotGeoShapeV2Reader) Close() error { + return nil +} + +// Count returns the total number of documents across all segments +// in the index snapshot that match the GeoShapeV2FieldReader's criteria. +func (g *IndexSnapshotGeoShapeV2Reader) Count() uint64 { + var rv uint64 + for _, posting := range g.postings { + if posting != nil { + rv += uint64(posting.GetCardinality()) + } + } + return rv +} + +// Size returns the estimated size in bytes of the +// IndexSnapshotGeoShapeV2Reader, including its postings and iterators. +func (g *IndexSnapshotGeoShapeV2Reader) Size() int { + rv := reflectStaticSizeIndexSnapshotGeoShapeV2Reader + size.SizeOfPtr + + len(g.field) + size.SizeOfInt + + for _, posting := range g.postings { + rv += int(posting.GetSizeInBytes()) + } + + rv += (reflectStaticSizeRoaringIntIterator + size.SizeOfPtr) * len(g.iterators) + + return rv +} diff --git a/index_test.go b/index_test.go index 2022b7387..203f3f459 100644 --- a/index_test.go +++ b/index_test.go @@ -612,9 +612,9 @@ func TestBytesRead(t *testing.T) { stats, _ := idx.StatsMap()["index"].(map[string]interface{}) prevBytesRead, _ := stats["num_bytes_read_at_query_time"].(uint64) - expectedBytesRead := uint64(21164) + expectedBytesRead := uint64(21574) if supportForVectorSearch { - expectedBytesRead = 21574 + expectedBytesRead = 21984 } if prevBytesRead != expectedBytesRead && res.Cost == prevBytesRead { @@ -770,9 +770,9 @@ func TestBytesReadStored(t *testing.T) { stats, _ := idx.StatsMap()["index"].(map[string]interface{}) bytesRead, _ := stats["num_bytes_read_at_query_time"].(uint64) - expectedBytesRead := uint64(11025) + expectedBytesRead := uint64(11435) if supportForVectorSearch { - expectedBytesRead = 11435 + expectedBytesRead = 11845 } if bytesRead != expectedBytesRead && bytesRead == res.Cost { @@ -847,9 +847,9 @@ func TestBytesReadStored(t *testing.T) { stats, _ = idx1.StatsMap()["index"].(map[string]interface{}) bytesRead, _ = stats["num_bytes_read_at_query_time"].(uint64) - expectedBytesRead = uint64(3212) + expectedBytesRead = uint64(3622) if supportForVectorSearch { - expectedBytesRead = 3622 + expectedBytesRead = 4032 } if bytesRead != expectedBytesRead && bytesRead == res.Cost { diff --git a/mapping.go b/mapping.go index af02db386..ba77f15ea 100644 --- a/mapping.go +++ b/mapping.go @@ -88,6 +88,10 @@ func NewGeoShapeFieldMapping() *mapping.FieldMapping { return mapping.NewGeoShapeFieldMapping() } +func NewGeoShapeV2FieldMapping() *mapping.FieldMapping { + return mapping.NewGeoShapeV2FieldMapping() +} + func NewIPFieldMapping() *mapping.FieldMapping { return mapping.NewIPFieldMapping() } diff --git a/mapping/document.go b/mapping/document.go index 3da925038..851198d69 100644 --- a/mapping/document.go +++ b/mapping/document.go @@ -104,7 +104,7 @@ func (dm *DocumentMapping) Validate(cache *registry.Cache, func validateFieldType(field *FieldMapping) error { switch field.Type { - case "text", "datetime", "number", "boolean", "geopoint", "geoshape", "IP": + case "text", "datetime", "number", "boolean", "geopoint", "geoshape", "geoshape_v2", "IP": return nil default: return fmt.Errorf("field: '%s', unknown field type: '%s'", @@ -552,11 +552,13 @@ func (dm *DocumentMapping) processProperty(property interface{}, path []string, for _, fieldMapping := range subDocMapping.Fields { switch fieldMapping.Type { case "geoshape": - fieldMapping.processGeoShape(property, pathString, path, indexes, context) + fieldMapping.processGeoShapeV2(property, pathString, path, context) case "geopoint": fieldMapping.processGeoPoint(property, pathString, path, indexes, context) case "vector_base64": fieldMapping.processVectorBase64(property, pathString, path, indexes, context) + case "geoshape_v2": + fieldMapping.processGeoShapeV2(property, pathString, path, context) default: fieldMapping.processString(propertyValueString, pathString, path, indexes, context) } @@ -640,7 +642,9 @@ func (dm *DocumentMapping) processProperty(property interface{}, path []string, case "geopoint": fieldMapping.processGeoPoint(property, pathString, path, indexes, context) case "geoshape": - fieldMapping.processGeoShape(property, pathString, path, indexes, context) + fieldMapping.processGeoShapeV2(property, pathString, path, context) + case "geoshape_v2": + fieldMapping.processGeoShapeV2(property, pathString, path, context) } } } @@ -664,7 +668,10 @@ func (dm *DocumentMapping) processProperty(property interface{}, path []string, } walkDocument = true case "geoshape": - fieldMapping.processGeoShape(property, pathString, path, indexes, context) + fieldMapping.processGeoShapeV2(property, pathString, path, context) + walkDocument = true + case "geoshape_v2": + fieldMapping.processGeoShapeV2(property, pathString, path, context) walkDocument = true default: walkDocument = true diff --git a/mapping/field.go b/mapping/field.go index 53c8dc61d..bf8d4d452 100644 --- a/mapping/field.go +++ b/mapping/field.go @@ -201,6 +201,18 @@ func NewGeoShapeFieldMapping() *FieldMapping { } } +// NewGeoShapeV2FieldMapping returns a default field mapping for +// geoshapes using the new GeoShapeV2 format +func NewGeoShapeV2FieldMapping() *FieldMapping { + return &FieldMapping{ + Type: "geoshape_v2", + Store: false, + Index: true, + IncludeInAll: false, + DocValues: false, + } +} + // NewIPFieldMapping returns a default field mapping for IP points func NewIPFieldMapping() *FieldMapping { return &FieldMapping{ @@ -336,47 +348,112 @@ func (fm *FieldMapping) processIP(ip net.IP, pathString string, path []string, i } } +// processGeoShape processes a property that might be a GeoJSON +// shape and adds the appropriate field to the document. func (fm *FieldMapping) processGeoShape(propertyMightBeGeoShape interface{}, pathString string, path []string, indexes []uint64, context *walkContext, +) { + fm.processGeoShapeInternal(propertyMightBeGeoShape, pathString, path, context, + func(fieldName string, shapes []*geojson.GeoShape, + options index.FieldIndexingOptions) document.Field { + field := document.NewGeometryCollectionFieldFromShapesWithIndexingOptions( + fieldName, indexes, shapes, options) + if field == nil { + // return an untyped nil, so that the caller's nil check works; + // returning the nil *GeoShapeField directly would produce a + // non-nil document.Field interface wrapping a nil pointer + return nil + } + return field + }, + func(fieldName string, shape *geojson.GeoShape, + options index.FieldIndexingOptions) document.Field { + field := document.NewGeoShapeFieldFromShapeWithIndexingOptions( + fieldName, indexes, shape, options) + if field == nil { + return nil + } + return field + }, + ) +} + +// processGeoShapeV2 processes a property that might be a GeoJSON +// shape and adds the appropriate field to the document using the new GeoShapeV2 format. +func (fm *FieldMapping) processGeoShapeV2(propertyMightBeGeoShape interface{}, + pathString string, path []string, context *walkContext, +) { + fm.processGeoShapeInternal(propertyMightBeGeoShape, pathString, path, context, + func(fieldName string, shapes []*geojson.GeoShape, + options index.FieldIndexingOptions) document.Field { + field := document.NewGeometryCollectionV2FieldFromShapesWithIndexingOptions( + fieldName, shapes, options) + if field == nil { + // return an untyped nil, so that the caller's nil check works; + // returning the nil *GeoShapeV2Field directly would produce a + // non-nil document.Field interface wrapping a nil pointer + return nil + } + return field + }, + func(fieldName string, shape *geojson.GeoShape, + options index.FieldIndexingOptions) document.Field { + field := document.NewGeoShapeV2FieldFromShapeWithIndexingOptions( + fieldName, shape, options) + if field == nil { + return nil + } + return field + }, + ) +} + +// processGeoShapeInternal is a helper function that processes +// a property that might be a GeoJSON shape and adds the appropriate +// field to the document. It takes two functions as parameters to handle +// the creation of fields for geometry collections and individual shapes. +func (fm *FieldMapping) processGeoShapeInternal( + propertyMightBeGeoShape interface{}, + pathString string, path []string, context *walkContext, + makeCollection func(fieldName string, shapes []*geojson.GeoShape, + options index.FieldIndexingOptions) document.Field, + makeShape func(fieldName string, shape *geojson.GeoShape, + options index.FieldIndexingOptions) document.Field, ) { coordValue, shape, err := geo.ParseGeoShapeField(propertyMightBeGeoShape) if err != nil { return } + var field document.Field + var found bool + if shape == geo.GeometryCollectionType { - geoShapes, found := geo.ExtractGeometryCollection(propertyMightBeGeoShape) + var geoShapes []*geojson.GeoShape + geoShapes, found = geo.ExtractGeometryCollection(propertyMightBeGeoShape) if found { fieldName := getFieldName(pathString, path, fm) - options := fm.Options() - field := document.NewGeometryCollectionFieldFromShapesWithIndexingOptions(fieldName, - indexes, geoShapes, options) - context.doc.AddField(field) - - if !fm.IncludeInAll { - context.excludedFromAll = append(context.excludedFromAll, fieldName) - } + field = makeCollection(fieldName, geoShapes, fm.Options()) } } else { var geoShape *geojson.GeoShape - var found bool - if shape == geo.CircleType { geoShape, found = geo.ExtractCircle(propertyMightBeGeoShape) } else { geoShape, found = geo.ExtractGeoShapeCoordinates(coordValue, shape) } - if found { fieldName := getFieldName(pathString, path, fm) - options := fm.Options() - field := document.NewGeoShapeFieldFromShapeWithIndexingOptions(fieldName, - indexes, geoShape, options) - context.doc.AddField(field) + field = makeShape(fieldName, geoShape, fm.Options()) + } + } - if !fm.IncludeInAll { - context.excludedFromAll = append(context.excludedFromAll, fieldName) - } + // field is nil when the constructor failed to encode the shape, + // in which case nothing is added to the document + if found && field != nil { + context.doc.AddField(field) + if !fm.IncludeInAll { + context.excludedFromAll = append(context.excludedFromAll, field.Name()) } } } diff --git a/search/query/geo_shape.go b/search/query/geo_shape.go index 686f486fd..ea5fe8e04 100644 --- a/search/query/geo_shape.go +++ b/search/query/geo_shape.go @@ -110,7 +110,7 @@ func (q *GeoShapeQuery) Searcher(ctx context.Context, i index.IndexReader, ctx = context.WithValue(ctx, search.QueryTypeKey, search.Geo) - return searcher.NewGeoShapeSearcher(ctx, i, q.Geometry.Shape, q.Geometry.Relation, field, + return searcher.NewGeoShapeV2Searcher(ctx, i, q.Geometry.Shape, q.Geometry.Relation, field, q.BoostVal.Value(), options) } diff --git a/search/query/geo_shape_v2.go b/search/query/geo_shape_v2.go new file mode 100644 index 000000000..6d7205223 --- /dev/null +++ b/search/query/geo_shape_v2.go @@ -0,0 +1,107 @@ +// Copyright (c) 2026 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 query + +import ( + "context" + "fmt" + + "github.com/blevesearch/bleve/v2/geo" + "github.com/blevesearch/bleve/v2/mapping" + "github.com/blevesearch/bleve/v2/search" + "github.com/blevesearch/bleve/v2/search/searcher" + index "github.com/blevesearch/bleve_index_api" +) + +type GeoShapeV2Query struct { + GeometryV2 Geometry `json:"geometry,omitempty"` + FieldVal string `json:"field,omitempty"` + BoostVal *Boost `json:"boost,omitempty"` +} + +func NewGeoShapeV2Query(coordinates [][][][]float64, typ, + relation string) (*GeoShapeV2Query, error) { + s, _, err := geo.NewGeoJsonShape(coordinates, typ) + if err != nil { + return nil, err + } + + return &GeoShapeV2Query{GeometryV2: Geometry{Shape: s, + Relation: relation}}, nil +} + +func NewGeoShapeV2CircleQuery(center []float64, radius, + relation string) (*GeoShapeV2Query, error) { + s, _, err := geo.NewGeoCircleShape(center, radius) + if err != nil { + return nil, err + } + + return &GeoShapeV2Query{GeometryV2: Geometry{Shape: s, + Relation: relation}}, nil +} + +func NewGeoShapeV2GeometryCollectionQuery(coordinates [][][][][]float64, + types []string, relation string) (*GeoShapeV2Query, error) { + s, _, err := geo.NewGeometryCollection(coordinates, types) + if err != nil { + return nil, err + } + + return &GeoShapeV2Query{GeometryV2: Geometry{Shape: s, + Relation: relation}}, nil +} + +func (q *GeoShapeV2Query) Boost() float64 { + return q.BoostVal.Value() +} + +func (q *GeoShapeV2Query) SetBoost(b float64) { + boost := Boost(b) + q.BoostVal = &boost +} + +func (q *GeoShapeV2Query) Field() string { + return q.FieldVal +} + +func (q *GeoShapeV2Query) SetField(f string) { + q.FieldVal = f +} + +func (q *GeoShapeV2Query) Validate() error { + switch q.GeometryV2.Relation { + case "intersects", "contains", "within", "disjoint": + return nil + default: + return fmt.Errorf("invalid relation: %q for geoshape_v2 query, "+ + "valid relations are: intersects, contains, within, disjoint", + q.GeometryV2.Relation) + } +} + +func (q *GeoShapeV2Query) Searcher(ctx context.Context, + i index.IndexReader, m mapping.IndexMapping, + options search.SearcherOptions) (search.Searcher, error) { + field := q.FieldVal + if q.FieldVal == "" { + field = m.DefaultSearchField() + } + + ctx = context.WithValue(ctx, search.QueryTypeKey, search.Geo) + + return searcher.NewGeoShapeV2Searcher(ctx, i, q.GeometryV2.Shape, + q.GeometryV2.Relation, field, q.BoostVal.Value(), options) +} diff --git a/search/query/query.go b/search/query/query.go index cf64189b2..e96115da3 100644 --- a/search/query/query.go +++ b/search/query/query.go @@ -371,7 +371,17 @@ func ParseQuery(input []byte) (Query, error) { _, hasGeo := tmp["geometry"] if hasGeo { - var rv GeoShapeQuery + var rv GeoShapeV2Query + err := util.UnmarshalJSON(input, &rv) + if err != nil { + return nil, err + } + return &rv, nil + } + + _, hasGeo = tmp["geometry_v2"] + if hasGeo { + var rv GeoShapeV2Query err := util.UnmarshalJSON(input, &rv) if err != nil { return nil, err diff --git a/search/scorer/scorer_constant.go b/search/scorer/scorer_constant.go index c030b8564..36310ec81 100644 --- a/search/scorer/scorer_constant.go +++ b/search/scorer/scorer_constant.go @@ -95,7 +95,7 @@ func (s *ConstantScorer) Score(ctx *search.SearchContext, id index.IndexInternal var scoreExplanation *search.Explanation rv := ctx.DocumentMatchPool.Get() - rv.IndexInternalID = id + rv.IndexInternalID = index.NewIndexInternalIDFrom(rv.IndexInternalID, id) if s.includeScore { score := s.constant diff --git a/search/searcher/geoshape_contains_test.go b/search/searcher/geoshape_contains_test.go index 437091a6a..e75bd375a 100644 --- a/search/searcher/geoshape_contains_test.go +++ b/search/searcher/geoshape_contains_test.go @@ -36,8 +36,8 @@ func testCaseSetup(t *testing.T, docShapeName, docShapeType string, docShapeVert i index.Index, ) (index.IndexReader, func() error, error) { doc := document.NewDocument(docShapeName) - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - docShapeVertices, docShapeType, document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc, "geometry", []uint64{}, + docShapeVertices, docShapeType, document.DefaultGeoShapeIndexingOptions) err := i.Update(doc) if err != nil { return nil, nil, err diff --git a/search/searcher/geoshape_within_test.go b/search/searcher/geoshape_within_test.go index 6f936538a..7d4038477 100644 --- a/search/searcher/geoshape_within_test.go +++ b/search/searcher/geoshape_within_test.go @@ -40,6 +40,7 @@ func testCaseSetupGeometryCollection(t *testing.T, docShapeName string, types [] return nil, nil, fmt.Errorf("the GC field is nil") } doc.AddField(gcField) + addGeoCollectionV2Parallel(doc, "geometry", docShapeVertices, types) if doc == nil { return nil, nil, fmt.Errorf("the doc is nil") } diff --git a/search/searcher/search_geoshape_circle_test.go b/search/searcher/search_geoshape_circle_test.go index b4e1fb29f..7f0225656 100644 --- a/search/searcher/search_geoshape_circle_test.go +++ b/search/searcher/search_geoshape_circle_test.go @@ -15,7 +15,6 @@ package searcher import ( - "context" "reflect" "testing" @@ -23,7 +22,6 @@ import ( "github.com/blevesearch/bleve/v2/geo" "github.com/blevesearch/bleve/v2/index/scorch" "github.com/blevesearch/bleve/v2/index/upsidedown/store/gtreap" - "github.com/blevesearch/bleve/v2/search" index "github.com/blevesearch/bleve_index_api" ) @@ -265,26 +263,8 @@ func TestGeoJsonCircleContainsQuery(t *testing.T) { func runGeoShapeCircleRelationQuery(relation string, i index.IndexReader, points []float64, radius string, field string, ) ([]string, error) { - var rv []string s := geo.NewGeoCircle(points, radius) - - gbs, err := NewGeoShapeSearcher(context.TODO(), i, s, relation, field, 1.0, search.SearcherOptions{}) - if err != nil { - return nil, err - } - ctx := &search.SearchContext{ - DocumentMatchPool: search.NewDocumentMatchPool(gbs.DocumentMatchPoolSize(), 0), - } - docMatch, err := gbs.Next(ctx) - for docMatch != nil && err == nil { - docID, _ := i.ExternalID(docMatch.IndexInternalID) - rv = append(rv, docID) - docMatch, err = gbs.Next(ctx) - } - if err != nil { - return nil, err - } - return rv, nil + return executeSearch(relation, i, s, field) } func setupGeoJsonShapesIndexForCircleQuery(t *testing.T) index.Index { @@ -313,8 +293,8 @@ func setupGeoJsonShapesIndexForCircleQuery(t *testing.T) index.Index { {77.67248153686523, 12.957679089615821}, }}} doc := document.NewDocument("polygon1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - polygon1, "polygon", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + polygon1, "polygon", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -329,8 +309,8 @@ func setupGeoJsonShapesIndexForCircleQuery(t *testing.T) index.Index { {81.84951782226561, 25.522692102524033}, }}} doc = document.NewDocument("polygon2") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - polygon2, "polygon", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + polygon2, "polygon", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -346,8 +326,8 @@ func setupGeoJsonShapesIndexForCircleQuery(t *testing.T) index.Index { {8.548071384429932, 47.379216780040124}, }}} doc = document.NewDocument("polygon3") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - polygon3, "polygon", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + polygon3, "polygon", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -355,8 +335,8 @@ func setupGeoJsonShapesIndexForCircleQuery(t *testing.T) index.Index { point1 := [][][][]float64{{{{81.2439, 26.2244}}}} doc = document.NewDocument("point1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - point1, "point", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + point1, "point", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -367,8 +347,8 @@ func setupGeoJsonShapesIndexForCircleQuery(t *testing.T) index.Index { {80.7220458984375, 25.750424835909385}, }}} doc = document.NewDocument("envelope1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - envelope1, "envelope", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + envelope1, "envelope", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -379,8 +359,8 @@ func setupGeoJsonShapesIndexForCircleQuery(t *testing.T) index.Index { {82.10537910461424, 25.544609829984058}, }}} doc = document.NewDocument("envelope2") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - envelope2, "envelope", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + envelope2, "envelope", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -391,38 +371,38 @@ func setupGeoJsonShapesIndexForCircleQuery(t *testing.T) index.Index { {8.552148342132568, 47.383778974713124}, }}} doc = document.NewDocument("envelope3") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - envelope3, "envelope", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + envelope3, "envelope", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) } doc = document.NewDocument("circle1") - doc.AddField(document.NewGeoCircleFieldWithIndexingOptions("geometry", []uint64{}, + addGeoCircleFieldV1V2(doc,"geometry", []uint64{}, []float64{77.67252445220947, 12.936348678099293}, "900m", - document.DefaultGeoShapeIndexingOptions)) + document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) } doc = document.NewDocument("circle2") - doc.AddField(document.NewGeoCircleFieldWithIndexingOptions("geometry", []uint64{}, + addGeoCircleFieldV1V2(doc,"geometry", []uint64{}, []float64{82.10289001464844, 25.544919592476727}, "100m", - document.DefaultGeoShapeIndexingOptions)) + document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) } doc = document.NewDocument("circle3") - doc.AddField(document.NewGeoCircleFieldWithIndexingOptions("geometry", []uint64{}, + addGeoCircleFieldV1V2(doc,"geometry", []uint64{}, []float64{ 8.53363037109375, 47.38191927423153, }, "400m", - document.DefaultGeoShapeIndexingOptions)) + document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -433,8 +413,8 @@ func setupGeoJsonShapesIndexForCircleQuery(t *testing.T) index.Index { {77.69213676452637, 12.945090185150542}, }}} doc = document.NewDocument("linestring1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - linestring, "linestring", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + linestring, "linestring", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -445,8 +425,8 @@ func setupGeoJsonShapesIndexForCircleQuery(t *testing.T) index.Index { {77.70252227783203, 12.929698235482276}, }}} doc = document.NewDocument("linestring2") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - linestring1, "linestring", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + linestring1, "linestring", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -457,8 +437,8 @@ func setupGeoJsonShapesIndexForCircleQuery(t *testing.T) index.Index { {81.30157470703125, 26.18440207077121}, }}} doc = document.NewDocument("linestring3") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - linestring2, "linestring", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + linestring2, "linestring", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -475,8 +455,8 @@ func setupGeoJsonShapesIndexForCircleQuery(t *testing.T) index.Index { {81.86702728271484, 25.502474677473746}, }}} doc = document.NewDocument("multilinestring1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - multilinestring, "multilinestring", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + multilinestring, "multilinestring", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -491,8 +471,8 @@ func setupGeoJsonShapesIndexForCircleQuery(t *testing.T) index.Index { {{81.8642807006836, 25.572175556682115}, {81.87458038330078, 25.567839795359724}}, }} doc = document.NewDocument("multilinestring2") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - multilinestring1, "multilinestring", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + multilinestring1, "multilinestring", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -504,8 +484,8 @@ func setupGeoJsonShapesIndexForCircleQuery(t *testing.T) index.Index { {81.90118789672852, 25.426067037656946}, }}} doc = document.NewDocument("multipoint1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - multipoint1, "multipoint", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + multipoint1, "multipoint", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -530,8 +510,8 @@ func setupGeoJsonShapesIndexForCircleQuery(t *testing.T) index.Index { }} doc = document.NewDocument("polygonWithHole1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - polygonWithHole1, "polygon", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + polygonWithHole1, "polygon", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) diff --git a/search/searcher/search_geoshape_envelope_test.go b/search/searcher/search_geoshape_envelope_test.go index 580fd5410..522e5bef1 100644 --- a/search/searcher/search_geoshape_envelope_test.go +++ b/search/searcher/search_geoshape_envelope_test.go @@ -15,7 +15,6 @@ package searcher import ( - "context" "reflect" "testing" @@ -23,7 +22,6 @@ import ( "github.com/blevesearch/bleve/v2/geo" "github.com/blevesearch/bleve/v2/index/scorch" "github.com/blevesearch/bleve/v2/index/upsidedown/store/gtreap" - "github.com/blevesearch/bleve/v2/search" index "github.com/blevesearch/bleve_index_api" ) @@ -306,26 +304,8 @@ func TestGeoJsonEnvelopeContainsQuery(t *testing.T) { func runGeoShapeEnvelopeRelationQuery(relation string, i index.IndexReader, points [][]float64, field string, ) ([]string, error) { - var rv []string s := geo.NewGeoEnvelope(points) - - gbs, err := NewGeoShapeSearcher(context.TODO(), i, s, relation, field, 1.0, search.SearcherOptions{}) - if err != nil { - return nil, err - } - ctx := &search.SearchContext{ - DocumentMatchPool: search.NewDocumentMatchPool(gbs.DocumentMatchPoolSize(), 0), - } - docMatch, err := gbs.Next(ctx) - for docMatch != nil && err == nil { - docID, _ := i.ExternalID(docMatch.IndexInternalID) - rv = append(rv, docID) - docMatch, err = gbs.Next(ctx) - } - if err != nil { - return nil, err - } - return rv, nil + return executeSearch(relation, i, s, field) } func setupGeoJsonShapesIndexForEnvelopeQuery(t *testing.T) index.Index { @@ -355,8 +335,8 @@ func setupGeoJsonShapesIndexForEnvelopeQuery(t *testing.T) index.Index { {8.548071384429932, 47.379216780040124}, }}} doc := document.NewDocument("polygon1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - polygon1, "polygon", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + polygon1, "polygon", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -369,8 +349,8 @@ func setupGeoJsonShapesIndexForEnvelopeQuery(t *testing.T) index.Index { {76.70379638671874, 16.828203242420393}, }}} doc = document.NewDocument("polygon2") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - polygon2, "polygon", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + polygon2, "polygon", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -386,8 +366,8 @@ func setupGeoJsonShapesIndexForEnvelopeQuery(t *testing.T) index.Index { {82.9522705078125, 17.749994573141873}, }}} doc = document.NewDocument("polygon3") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - polygon3, "polygon", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + polygon3, "polygon", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -398,38 +378,38 @@ func setupGeoJsonShapesIndexForEnvelopeQuery(t *testing.T) index.Index { {74.92401123046875, 17.66495983051931}, }}} doc = document.NewDocument("envelope1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - envelope1, "envelope", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + envelope1, "envelope", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) } doc = document.NewDocument("circle1") - doc.AddField(document.NewGeoCircleFieldWithIndexingOptions("geometry", []uint64{}, + addGeoCircleFieldV1V2(doc,"geometry", []uint64{}, []float64{75.0531005859375, 17.675427818339383}, "12900m", - document.DefaultGeoShapeIndexingOptions)) + document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) } doc = document.NewDocument("circle2") - doc.AddField(document.NewGeoCircleFieldWithIndexingOptions("geometry", []uint64{}, + addGeoCircleFieldV1V2(doc,"geometry", []uint64{}, []float64{82.69683837890625, 17.902955242676995}, "6000m", - document.DefaultGeoShapeIndexingOptions)) + document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) } doc = document.NewDocument("circle3") - doc.AddField(document.NewGeoCircleFieldWithIndexingOptions("geometry", []uint64{}, + addGeoCircleFieldV1V2(doc,"geometry", []uint64{}, []float64{ 8.53363037109375, 47.38191927423153, }, "400m", - document.DefaultGeoShapeIndexingOptions)) + document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -437,9 +417,9 @@ func setupGeoJsonShapesIndexForEnvelopeQuery(t *testing.T) index.Index { point1 := [][][][]float64{{{{76.29730224609375, 16.796653031618053}}}} doc = document.NewDocument("point1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, point1, "point", - document.DefaultGeoShapeIndexingOptions)) + document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -450,8 +430,8 @@ func setupGeoJsonShapesIndexForEnvelopeQuery(t *testing.T) index.Index { {77.24212646484374, 16.93070509876554}, }}} doc = document.NewDocument("linestring1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - linestring1, "linestring", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + linestring1, "linestring", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -462,8 +442,8 @@ func setupGeoJsonShapesIndexForEnvelopeQuery(t *testing.T) index.Index { {82.21343994140625, 18.059701055000478}, }}} doc = document.NewDocument("linestring2") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - linestring2, "linestring", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + linestring2, "linestring", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -476,8 +456,8 @@ func setupGeoJsonShapesIndexForEnvelopeQuery(t *testing.T) index.Index { {81.09283447265625, 17.87681743233167}, }}} doc = document.NewDocument("multipoint1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - multipoint1, "multipoint", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + multipoint1, "multipoint", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -493,8 +473,8 @@ func setupGeoJsonShapesIndexForEnvelopeQuery(t *testing.T) index.Index { {{81.815185546875, 17.3034434020238}, {81.81243896484375, 17.109292665395643}}, }} doc = document.NewDocument("multilinestring1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - multilinestring, "multilinestring", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + multilinestring, "multilinestring", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -509,8 +489,8 @@ func setupGeoJsonShapesIndexForEnvelopeQuery(t *testing.T) index.Index { {{77.60188579559325, 12.982604078764705}, {77.60557651519775, 12.987329508048184}}, }} doc = document.NewDocument("multilinestring2") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - multilinestring1, "multilinestring", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + multilinestring1, "multilinestring", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) diff --git a/search/searcher/search_geoshape_geometrycollection_test.go b/search/searcher/search_geoshape_geometrycollection_test.go index f9dd31223..d82992973 100644 --- a/search/searcher/search_geoshape_geometrycollection_test.go +++ b/search/searcher/search_geoshape_geometrycollection_test.go @@ -15,7 +15,6 @@ package searcher import ( - "context" "reflect" "testing" @@ -23,7 +22,6 @@ import ( "github.com/blevesearch/bleve/v2/geo" "github.com/blevesearch/bleve/v2/index/scorch" "github.com/blevesearch/bleve/v2/index/upsidedown/store/gtreap" - "github.com/blevesearch/bleve/v2/search" index "github.com/blevesearch/bleve_index_api" ) @@ -448,29 +446,11 @@ func TestGeoJSONContainsQueryAgainstGeometryCollection(t *testing.T) { func runGeoShapeGeometryCollectionRelationQuery(relation string, i index.IndexReader, points [][][][][]float64, types []string, field string, ) ([]string, error) { - var rv []string s, _, err := geo.NewGeometryCollection(points, types) if err != nil { return nil, err } - - gbs, err := NewGeoShapeSearcher(context.TODO(), i, s, relation, field, 1.0, search.SearcherOptions{}) - if err != nil { - return nil, err - } - ctx := &search.SearchContext{ - DocumentMatchPool: search.NewDocumentMatchPool(gbs.DocumentMatchPoolSize(), 0), - } - docMatch, err := gbs.Next(ctx) - for docMatch != nil && err == nil { - docID, _ := i.ExternalID(docMatch.IndexInternalID) - rv = append(rv, docID) - docMatch, err = gbs.Next(ctx) - } - if err != nil { - return nil, err - } - return rv, nil + return executeSearch(relation, i, s, field) } func setupGeoJsonShapesIndexForGeometryCollectionQuery(t *testing.T) index.Index { @@ -509,8 +489,8 @@ func setupGeoJsonShapesIndexForGeometryCollectionQuery(t *testing.T) index.Index types := []string{"polygon", "linestring"} doc := document.NewDocument("gc_polygon1_linestring1") - doc.AddField(document.NewGeometryCollectionFieldWithIndexingOptions("geometry", - []uint64{}, coordinates, types, document.DefaultGeoShapeIndexingOptions)) + addGeoCollectionFieldV1V2(doc,"geometry", + []uint64{}, coordinates, types, document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -543,8 +523,8 @@ func setupGeoJsonShapesIndexForGeometryCollectionQuery(t *testing.T) index.Index coordinates = [][][][][]float64{multipolygon1, multilinestring1} types = []string{"multipolygon", "multilinestring"} doc = document.NewDocument("gc_multipolygon1_multilinestring1") - doc.AddField(document.NewGeometryCollectionFieldWithIndexingOptions("geometry", - []uint64{}, coordinates, types, document.DefaultGeoShapeIndexingOptions)) + addGeoCollectionFieldV1V2(doc,"geometry", + []uint64{}, coordinates, types, document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -563,8 +543,8 @@ func setupGeoJsonShapesIndexForGeometryCollectionQuery(t *testing.T) index.Index types = []string{"point", "multipoint"} doc = document.NewDocument("gc_point1_multipoint1") - doc.AddField(document.NewGeometryCollectionFieldWithIndexingOptions("geometry", - []uint64{}, coordinates, types, document.DefaultGeoShapeIndexingOptions)) + addGeoCollectionFieldV1V2(doc,"geometry", + []uint64{}, coordinates, types, document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -599,8 +579,8 @@ func setupGeoJsonShapesIndexForGeometryCollectionQuery(t *testing.T) index.Index types = []string{"multipoint", "multipolygon", "multiline"} doc = document.NewDocument("gc_multipoint2_multipolygon2_multiline2") - doc.AddField(document.NewGeometryCollectionFieldWithIndexingOptions("geometry", - []uint64{}, coordinates, types, document.DefaultGeoShapeIndexingOptions)) + addGeoCollectionFieldV1V2(doc,"geometry", + []uint64{}, coordinates, types, document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -632,8 +612,8 @@ func setupGeoJsonShapesIndexForGeometryCollectionQuery(t *testing.T) index.Index types = []string{"multipolygon"} doc = document.NewDocument("gc_multipolygon3") - doc.AddField(document.NewGeometryCollectionFieldWithIndexingOptions("geometry", - []uint64{}, coordinates, types, document.DefaultGeoShapeIndexingOptions)) + addGeoCollectionFieldV1V2(doc,"geometry", + []uint64{}, coordinates, types, document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -652,8 +632,8 @@ func setupGeoJsonShapesIndexForGeometryCollectionQuery(t *testing.T) index.Index }}} doc = document.NewDocument("polygon2") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - polygon2, "polygon", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + polygon2, "polygon", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -681,8 +661,8 @@ func setupGeoJsonShapesIndexForGeometryCollectionQuery(t *testing.T) index.Index } doc = document.NewDocument("multipolygon4") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - multipolygon4, "multipolygon", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + multipolygon4, "multipolygon", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) diff --git a/search/searcher/search_geoshape_linestring_test.go b/search/searcher/search_geoshape_linestring_test.go index dc82df355..7fbb84c57 100644 --- a/search/searcher/search_geoshape_linestring_test.go +++ b/search/searcher/search_geoshape_linestring_test.go @@ -368,8 +368,25 @@ func runGeoShapeLinestringQueryWithRelation(relation string, i index.IndexReader return executeSearch(relation, i, s, field) } +// executeSearch runs the v1 geoshape searcher and, so that the same corpus +// also exercises the v2 implementation, requires the geoshape_v2 searcher to +// agree on the result set (see assertGeoShapeV2Agrees). It returns the v1 +// result in v1 order so existing order-sensitive assertions keep working. func executeSearch(relation string, i index.IndexReader, s index.GeoJSON, field string, +) ([]string, error) { + rv, err := executeGeoShapeV1Search(relation, i, s, field) + if err != nil { + return nil, err + } + if err := assertGeoShapeV2Agrees(relation, i, s, field, rv); err != nil { + return nil, err + } + return rv, nil +} + +func executeGeoShapeV1Search(relation string, i index.IndexReader, + s index.GeoJSON, field string, ) ([]string, error) { var rv []string gbs, err := NewGeoShapeSearcher(context.TODO(), i, s, relation, field, 1.0, search.SearcherOptions{}) @@ -419,9 +436,9 @@ func setupGeoJsonShapesIndexForLinestringQuery(t *testing.T) index.Index { {74.84642028808594, 22.402776071459712}, }}} doc := document.NewDocument("polygon1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, polygon1, "polygon", - document.DefaultGeoShapeIndexingOptions)) + document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -439,9 +456,9 @@ func setupGeoJsonShapesIndexForLinestringQuery(t *testing.T) index.Index { {74.93431091308592, 22.376428433285266}, }}} doc = document.NewDocument("polygon2") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, polygon2, "polygon", - document.DefaultGeoShapeIndexingOptions)) + document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -452,9 +469,9 @@ func setupGeoJsonShapesIndexForLinestringQuery(t *testing.T) index.Index { {74.87028121948242, 22.345471522338478}, }}} doc = document.NewDocument("envelope1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, envelope1, "envelope", - document.DefaultGeoShapeIndexingOptions)) + document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -465,27 +482,27 @@ func setupGeoJsonShapesIndexForLinestringQuery(t *testing.T) index.Index { {36.25333786010742, 50.03068093791795}, }}} doc = document.NewDocument("envelope2") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, envelope2, "envelope", - document.DefaultGeoShapeIndexingOptions)) + document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) } doc = document.NewDocument("circle1") - doc.AddField(document.NewGeoCircleFieldWithIndexingOptions("geometry", + addGeoCircleFieldV1V2(doc,"geometry", []uint64{}, []float64{74.93671417236328, 22.308314152382284}, "300m", - document.DefaultGeoShapeIndexingOptions)) + document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) } doc = document.NewDocument("circle2") - doc.AddField(document.NewGeoCircleFieldWithIndexingOptions("geometry", + addGeoCircleFieldV1V2(doc,"geometry", []uint64{}, []float64{36.22243881225586, 50.02941280037234}, "600m", - document.DefaultGeoShapeIndexingOptions)) + document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -496,9 +513,9 @@ func setupGeoJsonShapesIndexForLinestringQuery(t *testing.T) index.Index { {74.94036197662354, 22.32054224254707}, }}} doc = document.NewDocument("linestring1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, linestring, "linestring", - document.DefaultGeoShapeIndexingOptions)) + document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -509,9 +526,9 @@ func setupGeoJsonShapesIndexForLinestringQuery(t *testing.T) index.Index { {77.60557651519775, 12.987329508048184}, }}} doc = document.NewDocument("linestring2") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, linestring1, "linestring", - document.DefaultGeoShapeIndexingOptions)) + document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -526,9 +543,9 @@ func setupGeoJsonShapesIndexForLinestringQuery(t *testing.T) index.Index { {{74.9223804473877, 22.311688894660474}, {74.92534160614014, 22.30930673210729}}, }} doc = document.NewDocument("multilinestring1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, multilinestring, "multilinestring", - document.DefaultGeoShapeIndexingOptions)) + document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -543,9 +560,9 @@ func setupGeoJsonShapesIndexForLinestringQuery(t *testing.T) index.Index { {{77.60188579559325, 12.982604078764705}, {77.60557651519775, 12.987329508048184}}, }} doc = document.NewDocument("multilinestring2") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, multilinestring1, "multilinestring", - document.DefaultGeoShapeIndexingOptions)) + document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -557,9 +574,9 @@ func setupGeoJsonShapesIndexForLinestringQuery(t *testing.T) index.Index { {77.56922721862793, 12.956173473406446}, }}} doc = document.NewDocument("multipoint1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, multipoint1, "multipoint", - document.DefaultGeoShapeIndexingOptions)) + document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -584,9 +601,9 @@ func setupGeoJsonShapesIndexForLinestringQuery(t *testing.T) index.Index { }} doc = document.NewDocument("polygonWithHole1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, polygonWithHole1, "polygon", - document.DefaultGeoShapeIndexingOptions)) + document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -614,9 +631,9 @@ func setupGeoJsonShapesIndexForLinestringQuery(t *testing.T) index.Index { }} doc = document.NewDocument("polygonWithHole2") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, polygonWithHole2, "polygon", - document.DefaultGeoShapeIndexingOptions)) + document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -646,9 +663,9 @@ func setupGeoJsonShapesIndexForLinestringQuery(t *testing.T) index.Index { {36.221065521240234, 50.00365685169585}, }}} doc = document.NewDocument("multipolygon1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, multipolygon1, "multipolygon", - document.DefaultGeoShapeIndexingOptions)) + document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -675,9 +692,9 @@ func setupGeoJsonShapesIndexForLinestringQuery(t *testing.T) index.Index { coordinates := [][][][][]float64{polygonInGc, multipolygonInGc} types := []string{"polygon", "multipolygon"} doc = document.NewDocument("gc_polygonInGc_multipolygonInGc") - doc.AddField(document.NewGeometryCollectionFieldWithIndexingOptions("geometry", + addGeoCollectionFieldV1V2(doc,"geometry", []uint64{}, coordinates, types, - document.DefaultGeoShapeIndexingOptions)) + document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) diff --git a/search/searcher/search_geoshape_points_test.go b/search/searcher/search_geoshape_points_test.go index e7d8259dc..a56ef6de0 100644 --- a/search/searcher/search_geoshape_points_test.go +++ b/search/searcher/search_geoshape_points_test.go @@ -15,7 +15,6 @@ package searcher import ( - "context" "reflect" "testing" @@ -23,7 +22,6 @@ import ( "github.com/blevesearch/bleve/v2/geo" "github.com/blevesearch/bleve/v2/index/scorch" "github.com/blevesearch/bleve/v2/index/upsidedown/store/gtreap" - "github.com/blevesearch/bleve/v2/search" index "github.com/blevesearch/bleve_index_api" ) @@ -431,31 +429,13 @@ func TestGeoJsonMultiPointIntersectsQuery(t *testing.T) { func runGeoShapePointRelationQuery(relation string, multi bool, i index.IndexReader, points [][]float64, field string, ) ([]string, error) { - var rv []string var s index.GeoJSON if multi { s = geo.NewGeoJsonMultiPoint(points) } else { s = geo.NewGeoJsonPoint(points[0]) } - - gbs, err := NewGeoShapeSearcher(context.TODO(), i, s, relation, field, 1.0, search.SearcherOptions{}) - if err != nil { - return nil, err - } - ctx := &search.SearchContext{ - DocumentMatchPool: search.NewDocumentMatchPool(gbs.DocumentMatchPoolSize(), 0), - } - docMatch, err := gbs.Next(ctx) - for docMatch != nil && err == nil { - docID, _ := i.ExternalID(docMatch.IndexInternalID) - rv = append(rv, docID) - docMatch, err = gbs.Next(ctx) - } - if err != nil { - return nil, err - } - return rv, nil + return executeSearch(relation, i, s, field) } type Fatalfable interface { @@ -500,8 +480,8 @@ func setupGeoJsonShapesIndex(t *testing.T) index.Index { {77.5853419303894, 12.953977766785052}, }}} doc := document.NewDocument("polygon1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - polygon1, "polygon", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + polygon1, "polygon", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -513,17 +493,17 @@ func setupGeoJsonShapesIndex(t *testing.T) index.Index { {81.28440856933594, 26.351267272877074}, }}} doc = document.NewDocument("envelope1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - envelope1, "envelope", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + envelope1, "envelope", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) } doc = document.NewDocument("circle1") - doc.AddField(document.NewGeoCircleFieldWithIndexingOptions("geometry", []uint64{}, + addGeoCircleFieldV1V2(doc,"geometry", []uint64{}, []float64{77.59137153625487, 12.952660333521468}, "900m", - document.DefaultGeoShapeIndexingOptions)) + document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -534,8 +514,8 @@ func setupGeoJsonShapesIndex(t *testing.T) index.Index { {77.57776737213135, 12.952074805390097}, }}} doc = document.NewDocument("linestring1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - linestring, "linestring", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + linestring, "linestring", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -550,8 +530,8 @@ func setupGeoJsonShapesIndex(t *testing.T) index.Index { {77.5779390335083, 12.945006535817749}, }}} doc = document.NewDocument("multilinestring1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - multilinestring, "multilinestring", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + multilinestring, "multilinestring", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -563,8 +543,8 @@ func setupGeoJsonShapesIndex(t *testing.T) index.Index { {77.56922721862793, 12.956173473406446}, }}} doc = document.NewDocument("multipoint1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - multipoint1, "multipoint", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + multipoint1, "multipoint", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -589,8 +569,8 @@ func setupGeoJsonShapesIndex(t *testing.T) index.Index { }} doc = document.NewDocument("polygonWithHole1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - polygonWithHole1, "polygon", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + polygonWithHole1, "polygon", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) diff --git a/search/searcher/search_geoshape_polygon_test.go b/search/searcher/search_geoshape_polygon_test.go index 98be1c481..4ba5ecec7 100644 --- a/search/searcher/search_geoshape_polygon_test.go +++ b/search/searcher/search_geoshape_polygon_test.go @@ -15,7 +15,6 @@ package searcher import ( - "context" "reflect" "testing" @@ -23,7 +22,6 @@ import ( "github.com/blevesearch/bleve/v2/geo" "github.com/blevesearch/bleve/v2/index/scorch" "github.com/blevesearch/bleve/v2/index/upsidedown/store/gtreap" - "github.com/blevesearch/bleve/v2/search" index "github.com/blevesearch/bleve_index_api" ) @@ -461,26 +459,18 @@ func TestGeoJsonPolygonWithInQuery(t *testing.T) { func runGeoShapePolygonQueryWithRelation(relation string, i index.IndexReader, points [][][]float64, field string, ) ([]string, error) { - var rv []string s := geo.NewGeoJsonPolygon(points) + return executeSearch(relation, i, s, field) +} - gbs, err := NewGeoShapeSearcher(context.TODO(), i, s, relation, field, 1.0, search.SearcherOptions{}) - if err != nil { - return nil, err - } - ctx := &search.SearchContext{ - DocumentMatchPool: search.NewDocumentMatchPool(gbs.DocumentMatchPoolSize(), 0), - } - docMatch, err := gbs.Next(ctx) - for docMatch != nil && err == nil { - docID, _ := i.ExternalID(docMatch.IndexInternalID) - rv = append(rv, docID) - docMatch, err = gbs.Next(ctx) - } - if err != nil { - return nil, err - } - return rv, nil +// runGeoShapePolygonQueryV1Only runs only the v1 searcher. It is used by the +// specific-issue regression tests (e.g. the S2 loop porting issue) which have +// no v2 equivalent and therefore must not trigger the v1/v2 agreement check. +func runGeoShapePolygonQueryV1Only(relation string, i index.IndexReader, + points [][][]float64, field string, +) ([]string, error) { + s := geo.NewGeoJsonPolygon(points) + return executeGeoShapeV1Search(relation, i, s, field) } func setupGeoJsonShapesIndexForPolygonQuery(t *testing.T) index.Index { @@ -521,8 +511,8 @@ func setupGeoJsonShapesIndexForPolygonQuery(t *testing.T) index.Index { {77.5853419303894, 12.953977766785052}, }}} doc := document.NewDocument("polygon1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - polygon1, "polygon", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + polygon1, "polygon", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -539,8 +529,8 @@ func setupGeoJsonShapesIndexForPolygonQuery(t *testing.T) index.Index { {77.59527683258057, 12.951112863329588}, }}} doc = document.NewDocument("polygon2") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - polygon2, "polygon", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + polygon2, "polygon", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -554,8 +544,8 @@ func setupGeoJsonShapesIndexForPolygonQuery(t *testing.T) index.Index { {77.59974002838135, 12.953789562459688}, }}} doc = document.NewDocument("polygon3") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - polygon3, "polygon", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + polygon3, "polygon", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -567,8 +557,8 @@ func setupGeoJsonShapesIndexForPolygonQuery(t *testing.T) index.Index { {8.516979217529295, 47.38733837470806}, {8.522472381591797, 47.38794853343167}, {8.516507148742676, 47.388994503382285}, {8.515305519104004, 47.392597129887}}}} doc = document.NewDocument("polygon4") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - polygon4, "polygon", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + polygon4, "polygon", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -580,17 +570,17 @@ func setupGeoJsonShapesIndexForPolygonQuery(t *testing.T) index.Index { {36.20613098144531, 49.99714673955337}, }}} doc = document.NewDocument("envelope1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - envelope1, "envelope", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + envelope1, "envelope", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) } doc = document.NewDocument("circle1") - doc.AddField(document.NewGeoCircleFieldWithIndexingOptions("geometry", + addGeoCircleFieldV1V2(doc,"geometry", []uint64{}, []float64{77.59253025054932, 12.955587953533424}, "900m", - document.DefaultGeoShapeIndexingOptions)) + document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -601,8 +591,8 @@ func setupGeoJsonShapesIndexForPolygonQuery(t *testing.T) index.Index { {77.57776737213135, 12.952074805390097}, }}} doc = document.NewDocument("linestring1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - linestring, "linestring", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + linestring, "linestring", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -613,8 +603,8 @@ func setupGeoJsonShapesIndexForPolygonQuery(t *testing.T) index.Index { {77.60557651519775, 12.987329508048184}, }}} doc = document.NewDocument("linestring2") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - linestring1, "linestring", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + linestring1, "linestring", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -625,8 +615,8 @@ func setupGeoJsonShapesIndexForPolygonQuery(t *testing.T) index.Index { {8.520884513854979, 47.388006643417924}, }}} doc = document.NewDocument("linestring3") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - linestring3, "linestring", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + linestring3, "linestring", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -641,8 +631,8 @@ func setupGeoJsonShapesIndexForPolygonQuery(t *testing.T) index.Index { {{77.57781028747559, 12.951740217268595}, {77.5779390335083, 12.945006535817749}}, }} doc = document.NewDocument("multilinestring1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - multilinestring, "multilinestring", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + multilinestring, "multilinestring", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -657,8 +647,8 @@ func setupGeoJsonShapesIndexForPolygonQuery(t *testing.T) index.Index { {{77.60188579559325, 12.982604078764705}, {77.60557651519775, 12.987329508048184}}, }} doc = document.NewDocument("multilinestring2") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - multilinestring1, "multilinestring", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + multilinestring1, "multilinestring", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -670,8 +660,8 @@ func setupGeoJsonShapesIndexForPolygonQuery(t *testing.T) index.Index { {77.56922721862793, 12.956173473406446}, }}} doc = document.NewDocument("multipoint1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - multipoint1, "multipoint", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + multipoint1, "multipoint", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -696,8 +686,8 @@ func setupGeoJsonShapesIndexForPolygonQuery(t *testing.T) index.Index { }} doc = document.NewDocument("polygonWithHole1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - polygonWithHole1, "polygon", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + polygonWithHole1, "polygon", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -717,8 +707,8 @@ func setupGeoJsonShapesIndexForPolygonQuery(t *testing.T) index.Index { }}} doc = document.NewDocument("polygon4") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - polygon4, "polygon", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + polygon4, "polygon", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -742,8 +732,8 @@ func setupGeoJsonShapesIndexForPolygonQuery(t *testing.T) index.Index { } doc = document.NewDocument("multipolygon1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - multipolygon1, "multipolygon", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + multipolygon1, "multipolygon", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -892,27 +882,8 @@ func runGeoShapeMultiPolygonQueryWithRelation(relation string, i index.IndexReader, points [][][][]float64, field string, ) ([]string, error) { - var rv []string s := geo.NewGeoJsonMultiPolygon(points) - - gbs, err := NewGeoShapeSearcher(context.TODO(), i, s, relation, - field, 1.0, search.SearcherOptions{}) - if err != nil { - return nil, err - } - ctx := &search.SearchContext{ - DocumentMatchPool: search.NewDocumentMatchPool(gbs.DocumentMatchPoolSize(), 0), - } - docMatch, err := gbs.Next(ctx) - for docMatch != nil && err == nil { - docID, _ := i.ExternalID(docMatch.IndexInternalID) - rv = append(rv, docID) - docMatch, err = gbs.Next(ctx) - } - if err != nil { - return nil, err - } - return rv, nil + return executeSearch(relation, i, s, field) } func setupGeoJsonShapesIndexForMultiPolygonQuery(t *testing.T) index.Index { @@ -955,8 +926,8 @@ func setupGeoJsonShapesIndexForMultiPolygonQuery(t *testing.T) index.Index { {-121.47334098815918, 38.553485029658475}, }}} doc := document.NewDocument("multipolygon1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - multipolygon1, "multipolygon", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + multipolygon1, "multipolygon", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -971,8 +942,8 @@ func setupGeoJsonShapesIndexForMultiPolygonQuery(t *testing.T) index.Index { {{-121.49134397506714, 38.54490969679143}, {-121.4919662475586, 38.54304681805045}}, }} doc = document.NewDocument("multilinestring1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - multilinestring1, "multilinestring", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + multilinestring1, "multilinestring", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -987,8 +958,8 @@ func setupGeoJsonShapesIndexForMultiPolygonQuery(t *testing.T) index.Index { {-121.4881682395935, 38.57158887950165}, }}} doc = document.NewDocument("multipoint1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - multipoint1, "multipoint", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + multipoint1, "multipoint", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -1021,8 +992,8 @@ func setupGeoJsonPolygonS2LoopPortingIssue(t *testing.T) index.Index { {-135.0, 77.0}, }}} doc := document.NewDocument("polygon1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - polygon1, "polygon", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + polygon1, "polygon", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -1073,7 +1044,7 @@ func TestGeoJsonPolygonContainsQueryS2LoopPortingIssue(t *testing.T) { }() for n, test := range tests { - got, err := runGeoShapePolygonQueryWithRelation("contains", + got, err := runGeoShapePolygonQueryV1Only("contains", indexReader, test.polygon, test.field) if err != nil { t.Fatal(err) @@ -1152,7 +1123,7 @@ func TestGeoJsonPolygonIntersectsQuery1(t *testing.T) { }() for n, test := range tests { - got, err := runGeoShapePolygonQueryWithRelation("intersects", + got, err := runGeoShapePolygonQueryV1Only("intersects", indexReader, test.polygon, test.field) if err != nil { t.Fatal(err) @@ -1188,8 +1159,8 @@ func setupGeoJsonShapesIndexForPolygonQuery1(t *testing.T) index.Index { {96.69202458735312, 61.59480859768306}, }}} doc := document.NewDocument("polygon1") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - polygon1, "polygon", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + polygon1, "polygon", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) @@ -1202,8 +1173,8 @@ func setupGeoJsonShapesIndexForPolygonQuery1(t *testing.T) index.Index { {91.35604953911839, 65.11164029408492}, }}} doc = document.NewDocument("polygon2") - doc.AddField(document.NewGeoShapeFieldWithIndexingOptions("geometry", []uint64{}, - polygon2, "polygon", document.DefaultGeoShapeIndexingOptions)) + addGeoShapeFieldV1V2(doc,"geometry", []uint64{}, + polygon2, "polygon", document.DefaultGeoShapeIndexingOptions) err = i.Update(doc) if err != nil { t.Fatal(err) diff --git a/search/searcher/search_geoshape_v2.go b/search/searcher/search_geoshape_v2.go new file mode 100644 index 000000000..8a8bea0f2 --- /dev/null +++ b/search/searcher/search_geoshape_v2.go @@ -0,0 +1,128 @@ +// Copyright (c) 2026 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 ( + "context" + "fmt" + "reflect" + + "github.com/blevesearch/bleve/v2/search" + "github.com/blevesearch/bleve/v2/search/scorer" + index "github.com/blevesearch/bleve_index_api" +) + +var reflectStaticSizeGeoShapeV2Searcher int + +func init() { + var gsv2s GeoShapeV2Searcher + reflectStaticSizeGeoShapeV2Searcher = int(reflect.TypeOf(gsv2s).Size()) +} + +type GeoShapeV2Searcher struct { + geoShapeIndexReader index.GeoShapeV2FieldReader + scorer *scorer.ConstantScorer + + gd index.GeoShapeV2FieldDoc +} + +func NewGeoShapeV2Searcher(ctx context.Context, indexReader index.IndexReader, + shape index.GeoJSON, relation string, field string, boost float64, + options search.SearcherOptions, +) (search.Searcher, error) { + + if gr, ok := indexReader.(index.GeoShapeV2IndexReader); ok { + // get the GeoShapeV2FieldReader for the specified field + geoShapeIndexReader, err := gr.GeoShapeV2FieldReader(ctx, field) + if err != nil { + return nil, err + } + + // perform the search on the GeoShapeV2FieldReader with the specified + // shape and relation + err = geoShapeIndexReader.Search(shape, relation) + if err != nil { + return nil, err + } + + return &GeoShapeV2Searcher{ + geoShapeIndexReader: geoShapeIndexReader, + scorer: scorer.NewConstantScorer(1, boost, options), + gd: index.GeoShapeV2FieldDoc{}, + }, nil + } + + return nil, fmt.Errorf("indexReader does not support geoshape_v2 queries") +} + +// Next returns the next document match for the GeoShapeV2Searcher. +// It retrieves the next matching document from the GeoShapeV2FieldReader +// and scores it using the ConstantScorer. +func (g *GeoShapeV2Searcher) Next(ctx *search.SearchContext) (*search.DocumentMatch, error) { + match, err := g.geoShapeIndexReader.Next(g.gd.Reset()) + if err != nil { + return nil, err + } + if match == nil { + return nil, nil + } + + docMatch := g.scorer.Score(ctx, match.ID) + return docMatch, nil +} + +// Advance moves the searcher to the first document with an ID greater than or equal to the specified ID. +// It retrieves the next matching document from the GeoShapeV2FieldReader and scores it using the ConstantScorer. +func (g *GeoShapeV2Searcher) Advance(ctx *search.SearchContext, ID index.IndexInternalID) (*search.DocumentMatch, error) { + match, err := g.geoShapeIndexReader.Advance(ID, g.gd.Reset()) + if err != nil { + return nil, err + } + if match == nil { + return nil, nil + } + + docMatch := g.scorer.Score(ctx, match.ID) + return docMatch, nil +} + +func (g *GeoShapeV2Searcher) Close() error { + return g.geoShapeIndexReader.Close() +} + +func (g *GeoShapeV2Searcher) Count() uint64 { + return g.geoShapeIndexReader.Count() +} + +func (g *GeoShapeV2Searcher) DocumentMatchPoolSize() int { + return 1 +} + +func (g *GeoShapeV2Searcher) Min() int { + return 0 +} + +func (g *GeoShapeV2Searcher) SetQueryNorm(n float64) { + g.scorer.SetQueryNorm(n) +} + +func (g *GeoShapeV2Searcher) Size() int { + return reflectStaticSizeGeoShapeV2Searcher + g.geoShapeIndexReader.Size() + + g.scorer.Size() + g.gd.Size() +} + +func (g *GeoShapeV2Searcher) Weight() float64 { + return g.scorer.Weight() +} diff --git a/search/searcher/search_geoshape_v2_test.go b/search/searcher/search_geoshape_v2_test.go new file mode 100644 index 000000000..2d276b93b --- /dev/null +++ b/search/searcher/search_geoshape_v2_test.go @@ -0,0 +1,455 @@ +// Copyright (c) 2026 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 ( + "context" + "fmt" + "sort" + "testing" + + "github.com/blevesearch/bleve/v2/document" + "github.com/blevesearch/bleve/v2/geo" + "github.com/blevesearch/bleve/v2/index/scorch" + "github.com/blevesearch/bleve/v2/index/upsidedown/store/gtreap" + "github.com/blevesearch/bleve/v2/search" + index "github.com/blevesearch/bleve_index_api" + "github.com/blevesearch/geo/geojson" +) + +const geoV2TestField = "geometry" + +// box returns a single-polygon coordinate set (one CCW exterior ring) +// spanning the given lng/lat extents. +func box(minLng, minLat, maxLng, maxLat float64) [][][][]float64 { + return [][][][]float64{{{ + {minLng, minLat}, + {maxLng, minLat}, + {maxLng, maxLat}, + {minLng, maxLat}, + {minLng, minLat}, + }}} +} + +func addGeoShapeV2Doc(t *testing.T, i index.Index, id, typ string, + coords [][][][]float64) { + shape := &geojson.GeoShape{Type: typ, Coordinates: coords} + doc := document.NewDocument(id) + f := document.NewGeoShapeV2FieldFromShapeWithIndexingOptions( + geoV2TestField, shape, index.IndexField) + if f == nil { + t.Fatalf("failed to build geoshape_v2 field for %q", id) + } + doc.AddField(f) + if err := i.Update(doc); err != nil { + t.Fatal(err) + } +} + +// setupGeoShapeV2Index builds a scorch index with a spread of shapes whose +// relationships to the query box Q = [0,0]-[10,10] are unambiguous: +// +// smallInside : polygon fully inside Q +// pointInside : point inside Q +// overlapping : polygon straddling the corner of Q +// bigOutside : polygon fully disjoint from Q +// hugeContainer : polygon that fully contains Q +func setupGeoShapeV2Index(t *testing.T) index.Index { + analysisQueue := index.NewAnalysisQueue(1) + i, err := scorch.NewScorch(gtreap.Name, + map[string]interface{}{ + "path": "", + "spatialPlugin": "s2", + }, analysisQueue) + if err != nil { + t.Fatal(err) + } + if err = i.Open(); err != nil { + t.Fatal(err) + } + + addGeoShapeV2Doc(t, i, "smallInside", "polygon", box(2, 2, 4, 4)) + addGeoShapeV2Doc(t, i, "pointInside", "point", + [][][][]float64{{{{5, 5}}}}) + addGeoShapeV2Doc(t, i, "overlapping", "polygon", box(8, 8, 15, 15)) + addGeoShapeV2Doc(t, i, "bigOutside", "polygon", box(20, 20, 25, 25)) + addGeoShapeV2Doc(t, i, "hugeContainer", "polygon", box(-10, -10, 20, 20)) + + return i +} + +func executeGeoShapeV2Search(relation string, i index.IndexReader, + s index.GeoJSON, field string) ([]string, error) { + var rv []string + gbs, err := NewGeoShapeV2Searcher(context.TODO(), i, s, relation, field, + 1.0, search.SearcherOptions{}) + if err != nil { + return nil, err + } + defer func() { _ = gbs.Close() }() + + ctx := &search.SearchContext{ + DocumentMatchPool: search.NewDocumentMatchPool(gbs.DocumentMatchPoolSize(), 0), + } + docMatch, err := gbs.Next(ctx) + for docMatch != nil && err == nil { + docID, _ := i.ExternalID(docMatch.IndexInternalID) + rv = append(rv, docID) + docMatch, err = gbs.Next(ctx) + } + if err != nil { + return nil, err + } + sort.Strings(rv) + return rv, nil +} + +func TestGeoShapeV2Relations(t *testing.T) { + i := setupGeoShapeV2Index(t) + defer func() { _ = i.Close() }() + + indexReader, err := i.Reader() + if err != nil { + t.Fatal(err) + } + defer func() { _ = indexReader.Close() }() + + // query box Q = [0,0]-[10,10] + queryShape, _, err := geo.NewGeoJsonShape(box(0, 0, 10, 10), "polygon") + if err != nil { + t.Fatal(err) + } + + tests := []struct { + relation string + want []string + }{ + {"within", []string{"pointInside", "smallInside"}}, + {"intersects", []string{"hugeContainer", "overlapping", "pointInside", "smallInside"}}, + {"contains", []string{"hugeContainer"}}, + {"disjoint", []string{"bigOutside"}}, + } + + for _, test := range tests { + got, err := executeGeoShapeV2Search(test.relation, indexReader, + queryShape, geoV2TestField) + if err != nil { + t.Fatalf("relation %q: %v", test.relation, err) + } + want := append([]string(nil), test.want...) + sort.Strings(want) + if !equalStrings(got, want) { + t.Errorf("relation %q: expected %v, got %v", test.relation, want, got) + } + } +} + +func TestGeoShapeV2Count(t *testing.T) { + i := setupGeoShapeV2Index(t) + defer func() { _ = i.Close() }() + + indexReader, err := i.Reader() + if err != nil { + t.Fatal(err) + } + defer func() { _ = indexReader.Close() }() + + queryShape, _, err := geo.NewGeoJsonShape(box(0, 0, 10, 10), "polygon") + if err != nil { + t.Fatal(err) + } + + gbs, err := NewGeoShapeV2Searcher(context.TODO(), indexReader, queryShape, + "intersects", geoV2TestField, 1.0, search.SearcherOptions{}) + if err != nil { + t.Fatal(err) + } + defer func() { _ = gbs.Close() }() + + if got := gbs.Count(); got != 4 { + t.Fatalf("expected Count 4 for intersects, got %d", got) + } +} + +func TestGeoShapeV2UnsupportedReaderErrors(t *testing.T) { + // a nil / non-GeoShapeV2 index reader must produce an error rather than + // a nil searcher, so that callers which do not nil-check (e.g. the + // conjunction searcher) do not panic + queryShape, _, err := geo.NewGeoJsonShape(box(0, 0, 10, 10), "polygon") + if err != nil { + t.Fatal(err) + } + _, err = NewGeoShapeV2Searcher(context.TODO(), nil, queryShape, + "intersects", geoV2TestField, 1.0, search.SearcherOptions{}) + if err == nil { + t.Fatal("expected an error for an index reader that does not support geoshape_v2") + } +} + +// TestGeoShapeV2Advance drives Advance across the per-segment iterators and +// confirms it lands on the requested document and continues correctly. It +// also guards the ConstantScorer's per-match ID copy: the internal IDs +// collected across successive Next calls must stay distinct rather than all +// aliasing the reader's reused buffer. +func TestGeoShapeV2Advance(t *testing.T) { + i := setupGeoShapeV2Index(t) + defer func() { _ = i.Close() }() + + indexReader, err := i.Reader() + if err != nil { + t.Fatal(err) + } + defer func() { _ = indexReader.Close() }() + + queryShape, _, err := geo.NewGeoJsonShape(box(0, 0, 10, 10), "polygon") + if err != nil { + t.Fatal(err) + } + + // first pass: collect the ordered internal IDs of all intersects hits, + // copying each one so a reused buffer cannot corrupt the collection + newSearcher := func() search.Searcher { + s, serr := NewGeoShapeV2Searcher(context.TODO(), indexReader, queryShape, + "intersects", geoV2TestField, 1.0, search.SearcherOptions{}) + if serr != nil { + t.Fatal(serr) + } + return s + } + + s1 := newSearcher() + defer func() { _ = s1.Close() }() + ctx := &search.SearchContext{ + DocumentMatchPool: search.NewDocumentMatchPool(s1.DocumentMatchPoolSize(), 0), + } + + var ids []index.IndexInternalID + dm, err := s1.Next(ctx) + for dm != nil && err == nil { + ids = append(ids, append(index.IndexInternalID(nil), dm.IndexInternalID...)) + dm, err = s1.Next(ctx) + } + if err != nil { + t.Fatal(err) + } + if len(ids) < 3 { + t.Fatalf("expected at least 3 intersects hits to test Advance, got %d", len(ids)) + } + + // the collected IDs must be strictly increasing and distinct; if the + // scorer aliased the reader's buffer they would all be equal + for k := 1; k < len(ids); k++ { + if ids[k].Compare(ids[k-1]) <= 0 { + t.Fatalf("collected internal IDs are not strictly increasing: %v", ids) + } + } + + // second pass: Advance directly to the second hit and confirm we land on it + s2 := newSearcher() + defer func() { _ = s2.Close() }() + ctx2 := &search.SearchContext{ + DocumentMatchPool: search.NewDocumentMatchPool(s2.DocumentMatchPoolSize(), 0), + } + got, err := s2.Advance(ctx2, ids[1]) + if err != nil { + t.Fatal(err) + } + if got == nil || !got.IndexInternalID.Equals(ids[1]) { + t.Fatalf("Advance did not land on the requested doc: want %v, got %v", + ids[1], got) + } + // continuing with Next must yield the following hit + next, err := s2.Next(ctx2) + if err != nil { + t.Fatal(err) + } + if next == nil || !next.IndexInternalID.Equals(ids[2]) { + t.Fatalf("Next after Advance: want %v, got %v", ids[2], next) + } +} + +func TestGeoShapeV2CircleAndCollection(t *testing.T) { + analysisQueue := index.NewAnalysisQueue(1) + i, err := scorch.NewScorch(gtreap.Name, + map[string]interface{}{"path": "", "spatialPlugin": "s2"}, analysisQueue) + if err != nil { + t.Fatal(err) + } + if err = i.Open(); err != nil { + t.Fatal(err) + } + defer func() { _ = i.Close() }() + + // a circle centred inside Q + circle := &geojson.GeoShape{Type: geo.CircleType, Center: []float64{5, 5}, Radius: "100m"} + doc := document.NewDocument("circleInside") + cf := document.NewGeoShapeV2FieldFromShapeWithIndexingOptions(geoV2TestField, + circle, index.IndexField) + if cf == nil { + t.Fatal("failed to build circle geoshape_v2 field") + } + doc.AddField(cf) + if err = i.Update(doc); err != nil { + t.Fatal(err) + } + + // a geometry collection with one polygon inside Q and one far outside + shapes := []*geojson.GeoShape{ + {Type: "polygon", Coordinates: box(1, 1, 3, 3)}, + {Type: "polygon", Coordinates: box(40, 40, 42, 42)}, + } + gcDoc := document.NewDocument("collection") + gf := document.NewGeometryCollectionV2FieldFromShapesWithIndexingOptions( + geoV2TestField, shapes, index.IndexField) + if gf == nil { + t.Fatal("failed to build geometry collection geoshape_v2 field") + } + gcDoc.AddField(gf) + if err = i.Update(gcDoc); err != nil { + t.Fatal(err) + } + + indexReader, err := i.Reader() + if err != nil { + t.Fatal(err) + } + defer func() { _ = indexReader.Close() }() + + queryShape, _, err := geo.NewGeoJsonShape(box(0, 0, 10, 10), "polygon") + if err != nil { + t.Fatal(err) + } + + // both the circle and the collection (via its inside polygon) intersect Q + got, err := executeGeoShapeV2Search("intersects", indexReader, queryShape, geoV2TestField) + if err != nil { + t.Fatal(err) + } + want := []string{"circleInside", "collection"} + if !equalStrings(got, want) { + t.Errorf("intersects: expected %v, got %v", want, got) + } +} + +func equalStrings(a, b []string) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i] != b[i] { + return false + } + } + return true +} + +// geoV2FieldSuffix is appended to a v1 geoshape field name to produce the +// parallel geoshape_v2 field name. The shared v1 test corpus indexes each +// shape under both names so every v1 test also exercises the v2 path. +const geoV2FieldSuffix = "_v2" + +// addGeoShapeV2Parallel adds only the geoshape_v2 field, under "_v2". +func addGeoShapeV2Parallel(doc *document.Document, name, typ string, + coords [][][][]float64) { + shape := &geojson.GeoShape{Type: typ, Coordinates: coords} + if f := document.NewGeoShapeV2FieldFromShapeWithIndexingOptions( + name+geoV2FieldSuffix, shape, index.IndexField); f != nil { + doc.AddField(f) + } +} + +// addGeoCircleV2Parallel adds only the geoshape_v2 circle field, under "_v2". +func addGeoCircleV2Parallel(doc *document.Document, name string, + center []float64, radius string) { + shape := &geojson.GeoShape{Type: geo.CircleType, Center: center, Radius: radius} + if f := document.NewGeoShapeV2FieldFromShapeWithIndexingOptions( + name+geoV2FieldSuffix, shape, index.IndexField); f != nil { + doc.AddField(f) + } +} + +// addGeoCollectionV2Parallel adds only the geoshape_v2 geometry collection +// field, under "_v2". +func addGeoCollectionV2Parallel(doc *document.Document, name string, + coords [][][][][]float64, types []string) { + shapes := make([]*geojson.GeoShape, 0, len(coords)) + for idx := range coords { + shapes = append(shapes, &geojson.GeoShape{Type: types[idx], Coordinates: coords[idx]}) + } + if f := document.NewGeometryCollectionV2FieldFromShapesWithIndexingOptions( + name+geoV2FieldSuffix, shapes, index.IndexField); f != nil { + doc.AddField(f) + } +} + +// addGeoShapeFieldV1V2 adds the v1 geoshape field and, under "_v2", +// the equivalent geoshape_v2 field carrying the same shape. +func addGeoShapeFieldV1V2(doc *document.Document, name string, ap []uint64, + coords [][][][]float64, typ string, opts index.FieldIndexingOptions) { + doc.AddField(document.NewGeoShapeFieldWithIndexingOptions(name, ap, coords, typ, opts)) + addGeoShapeV2Parallel(doc, name, typ, coords) +} + +// addGeoCircleFieldV1V2 mirrors addGeoShapeFieldV1V2 for circle shapes. +func addGeoCircleFieldV1V2(doc *document.Document, name string, ap []uint64, + center []float64, radius string, opts index.FieldIndexingOptions) { + doc.AddField(document.NewGeoCircleFieldWithIndexingOptions(name, ap, center, radius, opts)) + addGeoCircleV2Parallel(doc, name, center, radius) +} + +// addGeoCollectionFieldV1V2 mirrors addGeoShapeFieldV1V2 for geometry collections. +func addGeoCollectionFieldV1V2(doc *document.Document, name string, ap []uint64, + coords [][][][][]float64, types []string, opts index.FieldIndexingOptions) { + doc.AddField(document.NewGeometryCollectionFieldWithIndexingOptions(name, ap, coords, types, opts)) + addGeoCollectionV2Parallel(doc, name, coords, types) +} + +// sameStringSet reports whether a and b contain the same elements, ignoring +// order and treating nil and empty as equal. +func sameStringSet(a, b []string) bool { + if len(a) != len(b) { + return false + } + seen := make(map[string]int, len(a)) + for _, v := range a { + seen[v]++ + } + for _, v := range b { + seen[v]-- + } + for _, c := range seen { + if c != 0 { + return false + } + } + return true +} + +// assertGeoShapeV2Agrees runs the geoshape_v2 searcher for the same shape and +// relation against the parallel "_v2" field, and returns an error when +// its result set differs from the v1 result set. This is how each v1 searcher +// test transparently also validates the v2 implementation. +func assertGeoShapeV2Agrees(relation string, i index.IndexReader, + s index.GeoJSON, field string, v1 []string) error { + v2, err := executeGeoShapeV2Search(relation, i, s, field+geoV2FieldSuffix) + if err != nil { + return fmt.Errorf("geoshape_v2 search error for relation %q: %w", relation, err) + } + if !sameStringSet(v1, v2) { + return fmt.Errorf("geoshape_v2 mismatch for relation %q on field %q: "+ + "v1=%v v2=%v", relation, field, v1, v2) + } + return nil +} diff --git a/util/bitset.go b/util/bitset.go new file mode 100644 index 000000000..418ada2e5 --- /dev/null +++ b/util/bitset.go @@ -0,0 +1,130 @@ +// Copyright (c) 2026 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 util + +import ( + "math/bits" + + "github.com/RoaringBitmap/roaring/v2" +) + +type Bitset struct { + data []uint64 + numBits int + exclude *roaring.Bitmap +} + +// NewBitset initializes a bitset capable of holding numbers up to maxVal +func NewBitset(maxVal int, exclude *roaring.Bitmap) *Bitset { + // We need (maxVal / 64) + 1 buckets to hold up to maxVal + size := (maxVal / 64) + 1 + return &Bitset{ + data: make([]uint64, size), + numBits: maxVal, + exclude: exclude, + } +} + +// Add inserts a value into the bitset (safely handles duplicates) +func (b *Bitset) Add(val int) { + if b.exclude != nil && b.exclude.Contains(uint32(val)) { + return + } + bucket := val >> 6 // Equivalent to val / 64 + bit := uint(val & 63) // Equivalent to val % 64 + + // Set the bit to 1 using bitwise OR + b.data[bucket] |= (1 << bit) +} + +// Remove deletes a value from the bitset +func (b *Bitset) Remove(val int) { + bucket := val >> 6 + bit := uint(val & 63) + + // Set the bit to 0 using bitwise AND with the complement + b.data[bucket] &^= (1 << bit) +} + +// Contains checks if a value exists in the bitset +func (b *Bitset) Contains(val int) bool { + bucket := val >> 6 + bit := uint(val & 63) + + return (b.data[bucket] & (1 << bit)) != 0 +} + +// Invert flips all bits in the bitset, +// effectively turning all 1s to 0s and vice versa +func (b *Bitset) Invert() { + for i := range b.data { + b.data[i] = ^b.data[i] + } + // the flip above sets the trailing bits beyond numBits in the last + // bucket(s), which do not correspond to valid values - clear them so + // that Iterate and Count never see them + lastBucket := b.numBits >> 6 + if lastBucket < len(b.data) { + b.data[lastBucket] &= (1 << uint(b.numBits&63)) - 1 + for i := lastBucket + 1; i < len(b.data); i++ { + b.data[i] = 0 + } + } + if b.exclude != nil { + it := b.exclude.Iterator() + for it.HasNext() { + bit := uint64(it.Next()) + word := bit / 64 + if word < uint64(len(b.data)) { + b.data[word] &^= uint64(1) << (bit % 64) + } + } + } +} + +// Iterate calls the provided function for every integer recorded in the bitset, in ascending order +func (b *Bitset) Iterate(f func(int)) { + for bucketIdx, bucket := range b.data { + // If the entire 64-bit block is 0, skip it entirely for speed + if bucket == 0 { + continue + } + + // Check all 64 bits in this bucket + for bitIdx := 0; bitIdx < 64; bitIdx++ { + if (bucket & (1 << uint(bitIdx))) != 0 { + // Reconstruct the original integer + originalVal := (bucketIdx << 6) + bitIdx + f(originalVal) + } + } + } +} + +func (b *Bitset) Count() int { + count := 0 + + for _, word := range b.data { + count += bits.OnesCount64(word) + } + + return count +} + +func (b *Bitset) Clear() { + for i := range b.data { + b.data[i] = 0 + } +} diff --git a/util/bitset_test.go b/util/bitset_test.go new file mode 100644 index 000000000..0d26f6d88 --- /dev/null +++ b/util/bitset_test.go @@ -0,0 +1,205 @@ +// Copyright (c) 2026 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 util + +import ( + "reflect" + "testing" + + "github.com/RoaringBitmap/roaring/v2" +) + +func TestBitsetAddContains(t *testing.T) { + b := NewBitset(100, nil) + + // values not yet added must not be present + for _, v := range []int{0, 1, 63, 64, 65, 100} { + if b.Contains(v) { + t.Fatalf("expected %d to be absent from an empty bitset", v) + } + } + + // exercise word boundaries explicitly: 63 is the last bit of word 0, + // 64 is the first bit of word 1, 65 the second + added := []int{0, 1, 63, 64, 65, 100} + for _, v := range added { + b.Add(v) + } + for _, v := range added { + if !b.Contains(v) { + t.Fatalf("expected %d to be present after Add", v) + } + } + + // a value between two set bits must remain absent + if b.Contains(50) { + t.Fatalf("expected 50 to be absent") + } +} + +func TestBitsetAddDuplicate(t *testing.T) { + b := NewBitset(100, nil) + b.Add(42) + b.Add(42) + if got := b.Count(); got != 1 { + t.Fatalf("expected duplicate Add to keep count at 1, got %d", got) + } + if !b.Contains(42) { + t.Fatalf("expected 42 to be present") + } +} + +func TestBitsetRemove(t *testing.T) { + b := NewBitset(100, nil) + b.Add(10) + b.Add(64) + b.Remove(10) + if b.Contains(10) { + t.Fatalf("expected 10 to be absent after Remove") + } + if !b.Contains(64) { + t.Fatalf("expected 64 to still be present") + } + // removing an absent value must be a no-op + b.Remove(99) + if b.Count() != 1 { + t.Fatalf("expected count 1 after removing an absent value, got %d", b.Count()) + } +} + +func TestBitsetCountAndClear(t *testing.T) { + b := NewBitset(200, nil) + for _, v := range []int{0, 5, 63, 64, 128, 200} { + b.Add(v) + } + if got := b.Count(); got != 6 { + t.Fatalf("expected count 6, got %d", got) + } + b.Clear() + if got := b.Count(); got != 0 { + t.Fatalf("expected count 0 after Clear, got %d", got) + } +} + +func TestBitsetIterateAscending(t *testing.T) { + b := NewBitset(200, nil) + // add out of order to confirm Iterate returns ascending order + for _, v := range []int{130, 0, 64, 63, 7} { + b.Add(v) + } + var got []int + b.Iterate(func(v int) { + got = append(got, v) + }) + want := []int{0, 7, 63, 64, 130} + if !reflect.DeepEqual(got, want) { + t.Fatalf("expected Iterate to yield %v in ascending order, got %v", want, got) + } +} + +func TestBitsetExcludeBlocksAdd(t *testing.T) { + exclude := roaring.New() + exclude.AddInt(5) + exclude.AddInt(70) + + b := NewBitset(100, exclude) + b.Add(5) // excluded, must be ignored + b.Add(6) // allowed + b.Add(70) // excluded, must be ignored + + if b.Contains(5) { + t.Fatalf("expected excluded value 5 to be blocked by Add") + } + if b.Contains(70) { + t.Fatalf("expected excluded value 70 to be blocked by Add") + } + if !b.Contains(6) { + t.Fatalf("expected non-excluded value 6 to be present") + } +} + +func TestBitsetInvertRespectsNumBits(t *testing.T) { + // numDocs = 10 means the valid doc IDs are 0..9. NewBitset allocates a + // full 64-bit word, so bits 10..63 are unused and must never surface. + numDocs := 10 + b := NewBitset(numDocs, nil) + b.Add(3) + b.Add(7) + + b.Invert() + + var got []int + b.Iterate(func(v int) { + got = append(got, v) + }) + + // after inverting, exactly the doc IDs in [0, 10) that were NOT set + // should be present - and nothing >= 10 + want := []int{0, 1, 2, 4, 5, 6, 8, 9} + if !reflect.DeepEqual(got, want) { + t.Fatalf("expected inverted bitset to yield %v, got %v", want, got) + } + if got := b.Count(); got != len(want) { + t.Fatalf("expected inverted count %d, got %d", len(want), got) + } +} + +func TestBitsetInvertWordBoundaries(t *testing.T) { + // exercise numBits at and around 64-bit word boundaries to make sure + // the trailing-bit mask is computed correctly + for _, numDocs := range []int{1, 63, 64, 65, 128, 129} { + b := NewBitset(numDocs, nil) + b.Invert() + max := -1 + b.Iterate(func(v int) { + if v > max { + max = v + } + }) + // every valid doc ID [0, numDocs) should be present after inverting + // an empty bitset, and none at or beyond numDocs + if got := b.Count(); got != numDocs { + t.Fatalf("numDocs=%d: expected inverted-empty count %d, got %d", + numDocs, numDocs, got) + } + if max >= numDocs { + t.Fatalf("numDocs=%d: Invert surfaced out-of-range value %d", + numDocs, max) + } + } +} + +func TestBitsetInvertClearsExcluded(t *testing.T) { + // excluded docs must remain unset even after Invert, since they are + // never valid hits + exclude := roaring.New() + exclude.AddInt(2) + exclude.AddInt(8) + + b := NewBitset(10, exclude) + b.Add(3) + b.Invert() + + if b.Contains(2) || b.Contains(8) { + t.Fatalf("expected excluded docs to stay unset after Invert") + } + if b.Contains(3) { + t.Fatalf("expected the originally-set doc 3 to be cleared after Invert") + } + // a normal, non-excluded, originally-unset doc should now be set + if !b.Contains(0) { + t.Fatalf("expected doc 0 to be set after Invert") + } +}