From b9be6d9ea724de2c9a66a7d5870a3abca7ba6f8b Mon Sep 17 00:00:00 2001
From: kimchy <kimchy@gmail.com>
Date: Fri, 7 Jan 2011 15:22:11 +0200
Subject: [PATCH] Geo-distance in scripts (like custom_score), closes #607.

---
 .../xcontent/geo/GeoPointDocFieldData.java    | 17 +++++++++++
 .../xcontent/geo/GeoPointFieldData.java       | 29 +++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/index/mapper/xcontent/geo/GeoPointDocFieldData.java b/modules/elasticsearch/src/main/java/org/elasticsearch/index/mapper/xcontent/geo/GeoPointDocFieldData.java
index f24c80205c4..01de18796f1 100644
--- a/modules/elasticsearch/src/main/java/org/elasticsearch/index/mapper/xcontent/geo/GeoPointDocFieldData.java
+++ b/modules/elasticsearch/src/main/java/org/elasticsearch/index/mapper/xcontent/geo/GeoPointDocFieldData.java
@@ -19,6 +19,7 @@
 
 package org.elasticsearch.index.mapper.xcontent.geo;
 
+import org.elasticsearch.common.unit.DistanceUnit;
 import org.elasticsearch.index.field.data.DocFieldData;
 
 /**
@@ -38,6 +39,22 @@ public class GeoPointDocFieldData extends DocFieldData<GeoPointFieldData> {
         return fieldData.values(docId);
     }
 
+    public double distance(double lat, double lon) {
+        return fieldData.distance(docId, DistanceUnit.MILES, lat, lon);
+    }
+
+    public double distanceInKm(double lat, double lon) {
+        return fieldData.distance(docId, DistanceUnit.KILOMETERS, lat, lon);
+    }
+
+    public double geohashDistance(String geohash) {
+        return fieldData.distanceGeohash(docId, DistanceUnit.MILES, geohash);
+    }
+
+    public double geohashDistanceInKm(String geohash) {
+        return fieldData.distanceGeohash(docId, DistanceUnit.KILOMETERS, geohash);
+    }
+
     public double getLat() {
         return fieldData.latValue(docId);
     }
diff --git a/modules/elasticsearch/src/main/java/org/elasticsearch/index/mapper/xcontent/geo/GeoPointFieldData.java b/modules/elasticsearch/src/main/java/org/elasticsearch/index/mapper/xcontent/geo/GeoPointFieldData.java
index a82faeb0efe..98ddf20f4cf 100644
--- a/modules/elasticsearch/src/main/java/org/elasticsearch/index/mapper/xcontent/geo/GeoPointFieldData.java
+++ b/modules/elasticsearch/src/main/java/org/elasticsearch/index/mapper/xcontent/geo/GeoPointFieldData.java
@@ -23,9 +23,11 @@ import org.apache.lucene.index.IndexReader;
 import org.elasticsearch.common.RamUsage;
 import org.elasticsearch.common.thread.ThreadLocals;
 import org.elasticsearch.common.trove.TDoubleArrayList;
+import org.elasticsearch.common.unit.DistanceUnit;
 import org.elasticsearch.index.field.data.FieldData;
 import org.elasticsearch.index.field.data.FieldDataType;
 import org.elasticsearch.index.field.data.support.FieldDataLoader;
+import org.elasticsearch.index.search.geo.GeoDistance;
 import org.elasticsearch.index.search.geo.GeoHashUtils;
 
 import java.io.IOException;
@@ -41,6 +43,18 @@ public abstract class GeoPointFieldData extends FieldData<GeoPointDocFieldData>
         }
     };
 
+    static class GeoPointHash {
+        public double lat;
+        public double lon;
+        public String geoHash = "";
+    }
+
+    static ThreadLocal<ThreadLocals.CleanableValue<GeoPointHash>> geoHashCache = new ThreadLocal<ThreadLocals.CleanableValue<GeoPointHash>>() {
+        @Override protected ThreadLocals.CleanableValue<GeoPointHash> initialValue() {
+            return new ThreadLocals.CleanableValue<GeoPointHash>(new GeoPointHash());
+        }
+    };
+
     public static final GeoPoint[] EMPTY_ARRAY = new GeoPoint[0];
 
     protected final double[] lat;
@@ -64,6 +78,21 @@ public abstract class GeoPointFieldData extends FieldData<GeoPointDocFieldData>
 
     abstract public double[] lonValues(int docId);
 
+    public double distance(int docId, DistanceUnit unit, double lat, double lon) {
+        return GeoDistance.PLANE.calculate(latValue(docId), lonValue(docId), lat, lon, unit);
+    }
+
+    public double distanceGeohash(int docId, DistanceUnit unit, String geoHash) {
+        GeoPointHash geoPointHash = geoHashCache.get().get();
+        if (geoPointHash.geoHash != geoHash) {
+            geoPointHash.geoHash = geoHash;
+            double[] decode = GeoHashUtils.decode(geoHash);
+            geoPointHash.lat = decode[0];
+            geoPointHash.lon = decode[1];
+        }
+        return GeoDistance.PLANE.calculate(latValue(docId), lonValue(docId), geoPointHash.lat, geoPointHash.lon, unit);
+    }
+
     @Override public GeoPointDocFieldData docFieldData(int docId) {
         return super.docFieldData(docId);
     }