From 0a064e7588eda2365a11200bce6742ac3a9b3a52 Mon Sep 17 00:00:00 2001
From: Josh Baker <joshbaker77@gmail.com>
Date: Fri, 12 Jan 2018 16:38:11 -0700
Subject: [PATCH] Nearby operation (kNN)

---
 README.md      | 14 ++++++++++++++
 buntdb.go      | 42 +++++++++++++++++++++++++++++++++++++++++-
 buntdb_test.go | 40 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 95 insertions(+), 1 deletion(-)

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)