diff --git a/README.md b/README.md index 6656232..6d38971 100644 --- a/README.md +++ b/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`. diff --git a/buntdb.go b/buntdb.go index 9fcaf31..e99ba26 100644 --- a/buntdb.go +++ b/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 diff --git a/buntdb_test.go b/buntdb_test.go index 9668f2a..cef8ccb 100644 --- a/buntdb_test.go +++ b/buntdb_test.go @@ -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)