111 lines
3.2 KiB
Go
111 lines
3.2 KiB
Go
package buntdb
|
|
|
|
import (
|
|
"github.com/tidwall/grect"
|
|
"github.com/tidwall/rtred"
|
|
)
|
|
|
|
// rect is used by Intersects and Nearby
|
|
type rect struct {
|
|
min, max []float64
|
|
}
|
|
|
|
func (r *rect) Rect(ctx interface{}) (min, max []float64) {
|
|
return r.min, r.max
|
|
}
|
|
|
|
// IndexRect is a helper function that converts string to a rect.
|
|
// Rect() is the reverse function and can be used to generate a string
|
|
// from a rect.
|
|
func IndexRect(a string) (min, max []float64) {
|
|
r := grect.Get(a)
|
|
return r.Min, r.Max
|
|
}
|
|
|
|
// Rect is helper function that returns a string representation
|
|
// of a rect. IndexRect() is the reverse function and can be used
|
|
// to generate a rect from a string.
|
|
func Rect(min, max []float64) string {
|
|
r := grect.Rect{Min: min, Max: max}
|
|
return r.String()
|
|
}
|
|
|
|
// Nearby searches for rectangle items that are nearby a target rect.
|
|
// All items belonging to the specified index will be returned in order of
|
|
// nearest to farthest.
|
|
// The specified index must have been created by AddIndex() and the target
|
|
// is represented by the rect string. This string will be processed by the
|
|
// same bounds function that was passed to the CreateSpatialIndex() function.
|
|
// An invalid index will return an error.
|
|
// The dist param is the distance of the bounding boxes. In the case of
|
|
// simple 2D points, it's the distance of the two 2D points squared.
|
|
func (tx *Tx) Nearby(index, bounds string,
|
|
iterator func(key, value string, dist float64) bool) error {
|
|
if tx.db == nil {
|
|
return ErrTxClosed
|
|
}
|
|
if index == "" {
|
|
// cannot search on keys tree. just return nil.
|
|
return nil
|
|
}
|
|
// // wrap a rtree specific iterator around the user-defined iterator.
|
|
iter := func(item rtred.Item, dist float64) bool {
|
|
dbi := item.(*dbItem)
|
|
return iterator(dbi.key, dbi.val, dist)
|
|
}
|
|
idx := tx.db.idxs[index]
|
|
if idx == nil {
|
|
// index was not found. return error
|
|
return ErrNotFound
|
|
}
|
|
if idx.rtr == nil {
|
|
// not an r-tree index. just return nil
|
|
return nil
|
|
}
|
|
// execute the nearby search
|
|
var min, max []float64
|
|
if idx.rect != nil {
|
|
min, max = idx.rect(bounds)
|
|
}
|
|
// set the center param to false, which uses the box dist calc.
|
|
idx.rtr.KNN(&rect{min, max}, false, iter)
|
|
return nil
|
|
}
|
|
|
|
// Intersects searches for rectangle items that intersect a target rect.
|
|
// The specified index must have been created by AddIndex() and the target
|
|
// is represented by the rect string. This string will be processed by the
|
|
// same bounds function that was passed to the CreateSpatialIndex() function.
|
|
// An invalid index will return an error.
|
|
func (tx *Tx) Intersects(index, bounds string,
|
|
iterator func(key, value string) bool) error {
|
|
if tx.db == nil {
|
|
return ErrTxClosed
|
|
}
|
|
if index == "" {
|
|
// cannot search on keys tree. just return nil.
|
|
return nil
|
|
}
|
|
// wrap a rtree specific iterator around the user-defined iterator.
|
|
iter := func(item rtred.Item) bool {
|
|
dbi := item.(*dbItem)
|
|
return iterator(dbi.key, dbi.val)
|
|
}
|
|
idx := tx.db.idxs[index]
|
|
if idx == nil {
|
|
// index was not found. return error
|
|
return ErrNotFound
|
|
}
|
|
if idx.rtr == nil {
|
|
// not an r-tree index. just return nil
|
|
return nil
|
|
}
|
|
// execute the search
|
|
var min, max []float64
|
|
if idx.rect != nil {
|
|
min, max = idx.rect(bounds)
|
|
}
|
|
idx.rtr.Search(&rect{min, max}, iter)
|
|
return nil
|
|
}
|