Nearby operation (kNN)

This commit is contained in:
Josh Baker 2018-01-12 16:38:11 -07:00
parent b67b1b8c16
commit 0a064e7588
3 changed files with 95 additions and 1 deletions

View File

@ -294,6 +294,20 @@ db.View(func(tx *buntdb.Tx) error {
This will get all three positions.
### k-Nearest Neighbors
Use the `Nearby` function to get all the positions in order of nearest to farthest :
```go
db.View(func(tx *buntdb.Tx) error {
tx.Nearby("fleet", "[-113 33]", func(key, val string, dist float64) bool {
...
return true
})
return nil
})
```
### Spatial bracket syntax
The bracket syntax `[-117 30],[-112 36]` is unique to BuntDB, and it's how the built-in rectangles are processed. But, you are not limited to this syntax. Whatever Rect function you choose to use during `CreateSpatialIndex` will be used to process the parameter, in this case it's `IndexRect`.

View File

@ -1775,7 +1775,7 @@ func (tx *Tx) DescendEqual(index, pivot string,
})
}
// rect is used by Intersects
// rect is used by Intersects and Nearby
type rect struct {
min, max []float64
}
@ -1784,6 +1784,46 @@ func (r *rect) Rect(ctx interface{}) (min, max []float64) {
return r.min, r.max
}
// 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.
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 rtree.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

View File

@ -1021,6 +1021,46 @@ func TestVariousTx(t *testing.T) {
}
}
func TestNearby(t *testing.T) {
rand.Seed(time.Now().UnixNano())
N := 100000
db, _ := Open(":memory:")
db.CreateSpatialIndex("points", "*", IndexRect)
db.Update(func(tx *Tx) error {
for i := 0; i < N; i++ {
p := Point(
rand.Float64()*100,
rand.Float64()*100,
rand.Float64()*100,
rand.Float64()*100,
)
tx.Set(fmt.Sprintf("p:%d", i), p, nil)
}
return nil
})
var keys, values []string
var dists []float64
var pdist float64
var i int
db.View(func(tx *Tx) error {
tx.Nearby("points", Point(0, 0, 0, 0), func(key, value string, dist float64) bool {
if i != 0 && dist < pdist {
t.Fatal("out of order")
}
keys = append(keys, key)
values = append(values, value)
dists = append(dists, dist)
pdist = dist
i++
return true
})
return nil
})
if len(keys) != N {
t.Fatalf("expected '%v', got '%v'", N, len(keys))
}
}
func Example_descKeys() {
db, _ := Open(":memory:")
db.CreateIndex("name", "*", IndexString)