Port Jan's geodistance functions

Author: AJ Roetker

generators/blevegen/blevegenerator.go

@@ -3,6 +3,7 @@ package blevegen
 import (
 	"errors"
 	"fmt"
+	"strconv"
 	"strings"
 	"time"
 
@@ -396,6 +397,40 @@ func (fg *FilterGenerator) funcExpr(node *expr.FuncNode, _ int) (query.Query, er
 		}
 
 		return makeTimeWindowQuery(lhs, threshold.Int64, window.Int64, int64(DayBucket(fg.ts)))
+	case "geodistance":
+		if len(node.Args) != 3 {
+			return nil, fmt.Errorf("'geodistance' function requires 3 arguments, got %d", len(node.Args))
+		}
+		lhs, err := fg.fieldType(node.Args[0])
+		if err != nil {
+			return nil, err
+		}
+		n, ok := node.Args[1].(*expr.StringNode)
+		if !ok {
+			return nil, fmt.Errorf("qlindex: unsupported type for 'geodistance' location argument. must be string, got %s", node.Args[1].NodeType())
+		}
+		lat, lon, ok := gentypes.StringToLatLng(n.Text)
+		if !ok {
+			return nil, fmt.Errorf("qlindex: unsupported format for 'geodistance' location. must be \"latitude,longitude\", got %s", n.Text)
+		}
+		var distance float64
+		switch n := node.Args[2].(type) {
+		case *expr.NumberNode:
+			if n.IsFloat {
+				distance = n.Float64
+			} else {
+				distance = float64(n.Int64)
+			}
+		case *expr.StringNode:
+			f, err := strconv.ParseFloat(n.Text, 64)
+			if err != nil {
+				return nil, fmt.Errorf("qlindex: unsupported format for 'geodistance' distance. must be number, got %s", n.Text)
+			}
+			distance = f
+		default:
+			return nil, fmt.Errorf("qlindex: unsupported type for 'geodistance' distance argument. must be number, got %s", node.Args[2].NodeType())
+		}
+		return makeGeoDistanceQuery(lhs, lat, lon, distance)
 	}
 	return nil, fmt.Errorf("qlindex: unsupported function: %s", node.Name)
 }

generators/blevegen/blevegenerator_test.go

@@ -36,16 +36,16 @@ func NewBleveIndexerMemOnly() (*BleveIndexer, error) {
 	documentMapping := bleve.NewDocumentMapping()
 	keywordFieldMapping := bleve.NewKeywordFieldMapping()
 	keywordFieldMapping.Name = "title"
-	//keywordFieldMapping.Store = false
-	// keywordFieldMapping.IncludeTermVectors = false
-	// keywordFieldMapping.IncludeInAll = false
-	// keywordFieldMapping.DocValues = false
+	keywordFieldMapping.Store = false
+	keywordFieldMapping.IncludeTermVectors = false
+	keywordFieldMapping.IncludeInAll = false
+	keywordFieldMapping.DocValues = false
 	documentMapping.Fields = append(documentMapping.Fields, keywordFieldMapping)
 	indexMapping.AddDocumentMapping("book", documentMapping)
 	allDocumentMapping := bleve.NewDocumentMapping()
-	// allDocumentMapping.Enabled = false
+	allDocumentMapping.Enabled = false
 	indexMapping.AddDocumentMapping("_all", allDocumentMapping)
-	// indexMapping.StoreDynamic = false
+	indexMapping.StoreDynamic = false
 	// indexMapping.DocValuesDynamic = false
 	// indexMapping.IndexDynamic = false
 

generators/blevegen/bridgeutil.go

@@ -283,6 +283,14 @@ func makeDateBetweenQuery(fieldName string, lower, upper any) (query.Query, erro
 	return dateRangeQuery, nil
 }
 
+func makeGeoDistanceQuery(lhs *gentypes.FieldType, lat, lon, distance float64) (query.Query, error) {
+	// Create a geo distance query
+	distanceQuery := query.NewGeoDistanceQuery(lon, lat, fmt.Sprintf("%fkm", distance))
+	distanceQuery.SetField(lhs.Field)
+
+	return distanceQuery, nil
+}
+
 // makeWildcard returns a wildcard/prefix query for Bleve
 func makeWildcard(lhs *gentypes.FieldType, value string, addStars bool) (query.Query, error) {
 	fieldName := lhs.Field

generators/esgen/esgenerator.go

@@ -3,6 +3,7 @@ package esgen
 import (
 	"errors"
 	"fmt"
+	"strconv"
 	"strings"
 	"time"
 
@@ -326,6 +327,53 @@ func (fg *FilterGenerator) funcExpr(node *expr.FuncNode, _ int) (any, error) {
 		}
 
 		return makeTimeWindowQuery(lhs, threshold.Int64, window.Int64, int64(DayBucket(fg.ts)))
+	case "geodistance":
+		if len(node.Args) != 3 {
+			return nil, fmt.Errorf("'geodistance' function requires 3 arguments, got %d", len(node.Args))
+		}
+		lhs, err := fg.fieldType(node.Args[0])
+		if err != nil {
+			return nil, err
+		}
+		n, ok := node.Args[1].(*expr.StringNode)
+		if !ok {
+			return nil, fmt.Errorf("qlindex: unsupported type for 'geodistance' location argument. must be string, got %s", node.Args[1].NodeType())
+		}
+		lat, lon, ok := gentypes.StringToLatLng(n.Text)
+		if !ok {
+			return nil, fmt.Errorf("qlindex: unsupported format for 'geodistance' location. must be \"latitude,longitude\", got %s", n.Text)
+		}
+		var distance float64
+		switch n := node.Args[2].(type) {
+		case *expr.NumberNode:
+			if n.IsFloat {
+				distance = n.Float64
+			} else {
+				distance = float64(n.Int64)
+			}
+		case *expr.StringNode:
+			f, err := strconv.ParseFloat(n.Text, 64)
+			if err != nil {
+				return nil, fmt.Errorf("qlindex: unsupported format for 'geodistance' distance. must be number, got %s", n.Text)
+			}
+			distance = f
+		default:
+			return nil, fmt.Errorf("qlindex: unsupported type for 'geodistance' distance argument. must be number, got %s", node.Args[2].NodeType())
+		}
+		return makeGeoDistanceQuery(lhs, lat, lon, distance), nil
 	}
 	return nil, fmt.Errorf("qlindex: unsupported function: %s", node.Name)
 }
+
+func makeGeoDistanceQuery(lhs *gentypes.FieldType, lat, lon, distance float64) any {
+	return &GeoDistanceFilter{
+		GeoDistance: map[string]any{
+			"distance": fmt.Sprintf("%fkm", distance),
+			fmt.Sprintf("%s.location", lhs.Field): map[string]any{
+				"lat": lat,
+				"lon": lon,
+			},
+			"distance_type": "plane",
+		},
+	}
+}

generators/esgen/estypes.go

@@ -209,3 +209,7 @@ func wcFunc(val string, addStars bool) string {
 func Wildcard(field, value string, addStars bool) *wildcard {
 	return &wildcard{Wildcard: map[string]string{field: wcFunc(value, addStars)}}
 }
+
+type GeoDistanceFilter struct {
+	GeoDistance map[string]any `json:"geo_distance"`
+}

generators/gentypes/utils.go

@@ -0,0 +1,25 @@
+package gentypes
+
+import (
+	"strconv"
+	"strings"
+)
+
+func StringToLatLng(s string) (float64, float64, bool) {
+	if len(s) < 1 {
+		return 0, 0, false
+	}
+	parts := strings.Split(s, ",")
+	if len(parts) != 2 {
+		return 0, 0, false
+	}
+	first, err := strconv.ParseFloat(strings.TrimSpace(parts[0]), 64)
+	if err != nil {
+		return 0, 0, false
+	}
+	second, err := strconv.ParseFloat(strings.TrimSpace(parts[1]), 64)
+	if err != nil {
+		return 0, 0, false
+	}
+	return first, second, true
+}