mirror of https://github.com/tidwall/buntdb.git
Nearby operation (kNN)
This commit is contained in:
parent
b67b1b8c16
commit
0a064e7588
14
README.md
14
README.md
|
@ -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`.
|
||||
|
|
42
buntdb.go
42
buntdb.go
|
@ -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
|
||||
|
|
|
@ -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)
|
||||
|
|
Loading…
Reference in New Issue