knn overscan ordering (#195)

controller/scanner.go

@@ -279,9 +279,8 @@ func (sw *scanWriter) writeObject(opts ScanWriterParams) bool {
 		return true
 	}
 	if sw.output == outputCount {
-		return true
+		return sw.count < sw.limit
 	}
-
 	switch sw.msg.OutputType {
 	case server.JSON:
 		var wr bytes.Buffer

controller/search.go

@@ -2,6 +2,7 @@ package controller
 
 import (
 	"bytes"
+	"sort"
 	"strconv"
 	"strings"
 	"time"
@@ -304,15 +305,19 @@ func (c *Controller) cmdNearby(msg *server.Message) (res string, err error) {
 	}
 	sw.writeHead()
 	if sw.col != nil {
-		iter := func(id string, o geojson.Object, fields []float64) bool {
+		iter := func(id string, o geojson.Object, fields []float64, dist *float64) bool {
 			if c.hasExpired(s.key, id) {
 				return true
 			}
 			// Calculate distance if we need to
 			distance := 0.0
 			if s.distance {
+				if dist != nil {
+					distance = *dist
+				} else {
 					distance = o.CalculatedPoint().DistanceTo(geojson.Position{X: s.lon, Y: s.lat, Z: 0})
 				}
+			}
 			return sw.writeObject(ScanWriterParams{
 				id:       id,
 				o:        o,
@@ -321,9 +326,13 @@ func (c *Controller) cmdNearby(msg *server.Message) (res string, err error) {
 			})
 		}
 		if s.knn {
-			sw.col.NearestNeighbors(s.lat, s.lon, iter)
+			nearestNeighbors(sw, s.lat, s.lon, iter)
 		} else {
-			sw.col.Nearby(s.sparse, s.lat, s.lon, s.meters, minZ, maxZ, iter)
+			sw.col.Nearby(s.sparse, s.lat, s.lon, s.meters, minZ, maxZ,
+				func(id string, o geojson.Object, fields []float64) bool {
+					return iter(id, o, fields, nil)
+				},
+			)
 		}
 	}
 	sw.writeFoot()
@@ -333,6 +342,39 @@ func (c *Controller) cmdNearby(msg *server.Message) (res string, err error) {
 	return string(wr.Bytes()), nil
 }
 
+type iterItem struct {
+	id     string
+	o      geojson.Object
+	fields []float64
+	dist   float64
+}
+
+func nearestNeighbors(sw *scanWriter, lat, lon float64, iter func(id string, o geojson.Object, fields []float64, dist *float64) bool) {
+	limit := sw.limit * 10
+	if limit > limitItems*10 {
+		limit = limitItems * 10
+	}
+	k := sw.cursor + limit
+	var items []iterItem
+	sw.col.NearestNeighbors(lat, lon, func(id string, o geojson.Object, fields []float64) bool {
+		if k == 0 {
+			return false
+		}
+		dist := o.CalculatedPoint().DistanceTo(geojson.Position{X: lon, Y: lat, Z: 0})
+		items = append(items, iterItem{id: id, o: o, fields: fields, dist: dist})
+		k--
+		return true
+	})
+	sort.Slice(items, func(i, j int) bool {
+		return items[i].dist < items[j].dist
+	})
+	for _, item := range items {
+		if !iter(item.id, item.o, item.fields, &item.dist) {
+			return
+		}
+	}
+}
+
 func (c *Controller) cmdWithin(msg *server.Message) (res string, err error) {
 	return c.cmdWithinOrIntersects("within", msg)
 }