From 0286d0a769e216d18c3074c463d699879fc5c8ab Mon Sep 17 00:00:00 2001
From: Nik Everett <nik9000@gmail.com>
Date: Mon, 10 Aug 2020 16:05:17 -0400
Subject: [PATCH] Move distance_feature query building into MFT (#60614)
 (#60846)

This moves the `distance_feature` query building out of
`DistanceFeatureQueryBuilder` and into subclasses of `MappedFieldType`.
Without this we don't have a chance of supporting this for runtime
fields. In general I'm not sad to see the `instanceof`s go.

Co-authored-by: Elastic Machine <elasticmachine@users.noreply.github.com>
---
 .../index/mapper/DateFieldMapper.java         | 20 ++++++++++++
 .../index/mapper/GeoPointFieldMapper.java     | 17 ++++++++++
 .../index/mapper/MappedFieldType.java         |  6 ++++
 .../query/DistanceFeatureQueryBuilder.java    | 32 +------------------
 4 files changed, 44 insertions(+), 31 deletions(-)

diff --git a/server/src/main/java/org/elasticsearch/index/mapper/DateFieldMapper.java b/server/src/main/java/org/elasticsearch/index/mapper/DateFieldMapper.java
index db11c044eed..cbb02ce8d3c 100644
--- a/server/src/main/java/org/elasticsearch/index/mapper/DateFieldMapper.java
+++ b/server/src/main/java/org/elasticsearch/index/mapper/DateFieldMapper.java
@@ -41,6 +41,7 @@ import org.elasticsearch.common.time.DateFormatter;
 import org.elasticsearch.common.time.DateFormatters;
 import org.elasticsearch.common.time.DateMathParser;
 import org.elasticsearch.common.time.DateUtils;
+import org.elasticsearch.common.unit.TimeValue;
 import org.elasticsearch.common.util.LocaleUtils;
 import org.elasticsearch.index.fielddata.IndexFieldData;
 import org.elasticsearch.index.fielddata.IndexNumericFieldData.NumericType;
@@ -94,6 +95,11 @@ public final class DateFieldMapper extends ParametrizedFieldMapper {
             public long parsePointAsMillis(byte[] value) {
                 return LongPoint.decodeDimension(value, 0);
             }
+
+            @Override
+            protected Query distanceFeatureQuery(String field, float boost, long origin, TimeValue pivot) {
+                return LongPoint.newDistanceFeatureQuery(field, boost, origin, pivot.getMillis());
+            }
         },
         NANOSECONDS(DATE_NANOS_CONTENT_TYPE, NumericType.DATE_NANOSECONDS) {
             @Override
@@ -115,6 +121,11 @@ public final class DateFieldMapper extends ParametrizedFieldMapper {
             public long parsePointAsMillis(byte[] value) {
                 return DateUtils.toMilliSeconds(LongPoint.decodeDimension(value, 0));
             }
+
+            @Override
+            protected Query distanceFeatureQuery(String field, float boost, long origin, TimeValue pivot) {
+                return LongPoint.newDistanceFeatureQuery(field, boost, origin, pivot.getNanos());
+            }
         };
 
         private final String type;
@@ -162,6 +173,8 @@ public final class DateFieldMapper extends ParametrizedFieldMapper {
             }
             throw new IllegalArgumentException("unknown resolution ordinal [" + ord + "]");
         }
+
+        protected abstract Query distanceFeatureQuery(String field, float boost, long origin, TimeValue pivot);
     }
 
     private static DateFieldMapper toType(FieldMapper in) {
@@ -370,6 +383,13 @@ public final class DateFieldMapper extends ParametrizedFieldMapper {
             return resolution.convert(instant);
         }
 
+        @Override
+        public Query distanceFeatureQuery(Object origin, String pivot, float boost, QueryShardContext context) {
+            long originLong = parseToLong(origin, true, null, null, context::nowInMillis);
+            TimeValue pivotTime = TimeValue.parseTimeValue(pivot, "distance_feature.pivot");
+            return resolution.distanceFeatureQuery(name(), boost, originLong, pivotTime);
+        }
+
         @Override
         public Relation isFieldWithinQuery(IndexReader reader,
                                            Object from, Object to, boolean includeLower, boolean includeUpper,
diff --git a/server/src/main/java/org/elasticsearch/index/mapper/GeoPointFieldMapper.java b/server/src/main/java/org/elasticsearch/index/mapper/GeoPointFieldMapper.java
index 1ac9760bb90..0b5092ca370 100644
--- a/server/src/main/java/org/elasticsearch/index/mapper/GeoPointFieldMapper.java
+++ b/server/src/main/java/org/elasticsearch/index/mapper/GeoPointFieldMapper.java
@@ -24,15 +24,18 @@ import org.apache.lucene.document.LatLonPoint;
 import org.apache.lucene.document.StoredField;
 import org.apache.lucene.index.IndexOptions;
 import org.apache.lucene.index.IndexableField;
+import org.apache.lucene.search.Query;
 import org.elasticsearch.ElasticsearchParseException;
 import org.elasticsearch.common.Explicit;
 import org.elasticsearch.common.geo.GeoPoint;
 import org.elasticsearch.common.geo.GeoUtils;
+import org.elasticsearch.common.unit.DistanceUnit;
 import org.elasticsearch.common.xcontent.XContentParser;
 import org.elasticsearch.geometry.Point;
 import org.elasticsearch.index.fielddata.IndexFieldData;
 import org.elasticsearch.index.fielddata.plain.AbstractLatLonPointIndexFieldData;
 import org.elasticsearch.index.mapper.GeoPointFieldMapper.ParsedGeoPoint;
+import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.index.query.VectorGeoPointShapeQueryProcessor;
 import org.elasticsearch.search.aggregations.support.CoreValuesSourceType;
 
@@ -186,6 +189,20 @@ public class GeoPointFieldMapper extends AbstractPointGeometryFieldMapper<List<P
             return new AbstractLatLonPointIndexFieldData.Builder(name(), CoreValuesSourceType.GEOPOINT);
         }
 
+        @Override
+        public Query distanceFeatureQuery(Object origin, String pivot, float boost, QueryShardContext context) {
+            GeoPoint originGeoPoint;
+            if (origin instanceof GeoPoint) {
+                originGeoPoint = (GeoPoint) origin;
+            } else if (origin instanceof String) {
+                originGeoPoint = GeoUtils.parseFromString((String) origin);
+            } else {
+                throw new IllegalArgumentException("Illegal type ["+ origin.getClass() + "] for [origin]! " +
+                    "Must be of type [geo_point] or [string] for geo_point fields!");
+            }
+            double pivotDouble = DistanceUnit.DEFAULT.parse(pivot, DistanceUnit.DEFAULT);
+            return LatLonPoint.newDistanceFeatureQuery(name(), boost, originGeoPoint.lat(), originGeoPoint.lon(), pivotDouble);
+        }
     }
 
     // Eclipse requires the AbstractPointGeometryFieldMapper prefix or it can't find ParsedPoint
diff --git a/server/src/main/java/org/elasticsearch/index/mapper/MappedFieldType.java b/server/src/main/java/org/elasticsearch/index/mapper/MappedFieldType.java
index 55c130652d0..1eb9fd0330c 100644
--- a/server/src/main/java/org/elasticsearch/index/mapper/MappedFieldType.java
+++ b/server/src/main/java/org/elasticsearch/index/mapper/MappedFieldType.java
@@ -43,6 +43,7 @@ import org.elasticsearch.common.time.DateMathParser;
 import org.elasticsearch.common.unit.Fuzziness;
 import org.elasticsearch.index.analysis.NamedAnalyzer;
 import org.elasticsearch.index.fielddata.IndexFieldData;
+import org.elasticsearch.index.query.DistanceFeatureQueryBuilder;
 import org.elasticsearch.index.query.QueryRewriteContext;
 import org.elasticsearch.index.query.QueryShardContext;
 import org.elasticsearch.index.query.QueryShardException;
@@ -241,6 +242,11 @@ public abstract class MappedFieldType {
             + "] which is of type [" + typeName() + "]");
     }
 
+    public Query distanceFeatureQuery(Object origin, String pivot, float boost, QueryShardContext context) {
+        throw new IllegalArgumentException("Illegal data type of [" + typeName() + "]!"+
+            "[" + DistanceFeatureQueryBuilder.NAME + "] query can only be run on a date, date_nanos or geo_point field type!");
+    }
+
     /**
      * Create an {@link IntervalsSource} to be used for proximity queries
      */
diff --git a/server/src/main/java/org/elasticsearch/index/query/DistanceFeatureQueryBuilder.java b/server/src/main/java/org/elasticsearch/index/query/DistanceFeatureQueryBuilder.java
index ec7debe5b93..90017094605 100644
--- a/server/src/main/java/org/elasticsearch/index/query/DistanceFeatureQueryBuilder.java
+++ b/server/src/main/java/org/elasticsearch/index/query/DistanceFeatureQueryBuilder.java
@@ -19,8 +19,6 @@
 
 package org.elasticsearch.index.query;
 
-import org.apache.lucene.document.LatLonPoint;
-import org.apache.lucene.document.LongPoint;
 import org.apache.lucene.search.Query;
 import org.elasticsearch.common.ParseField;
 import org.elasticsearch.common.ParsingException;
@@ -29,15 +27,10 @@ import org.elasticsearch.common.geo.GeoUtils;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.lucene.search.Queries;
-import org.elasticsearch.common.unit.DistanceUnit;
-import org.elasticsearch.common.unit.TimeValue;
 import org.elasticsearch.common.xcontent.ConstructingObjectParser;
 import org.elasticsearch.common.xcontent.ObjectParser;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
-import org.elasticsearch.index.mapper.DateFieldMapper;
-import org.elasticsearch.index.mapper.DateFieldMapper.DateFieldType;
-import org.elasticsearch.index.mapper.GeoPointFieldMapper.GeoPointFieldType;
 import org.elasticsearch.index.mapper.MappedFieldType;
 
 import java.io.IOException;
@@ -119,30 +112,7 @@ public class DistanceFeatureQueryBuilder extends AbstractQueryBuilder<DistanceFe
         if (fieldType == null) {
             return Queries.newMatchNoDocsQuery("Can't run [" + NAME + "] query on unmapped fields!");
         }
-        Object originObj = origin.origin();
-        if (fieldType instanceof DateFieldType) {
-            long originLong = ((DateFieldType) fieldType).parseToLong(originObj, true, null, null, context::nowInMillis);
-            TimeValue pivotVal = TimeValue.parseTimeValue(pivot, DistanceFeatureQueryBuilder.class.getSimpleName() + ".pivot");
-            if (((DateFieldType) fieldType).resolution() == DateFieldMapper.Resolution.MILLISECONDS) {
-                return LongPoint.newDistanceFeatureQuery(field, boost, originLong, pivotVal.getMillis());
-            } else { // NANOSECONDS
-                return LongPoint.newDistanceFeatureQuery(field, boost, originLong, pivotVal.getNanos());
-            }
-        } else if (fieldType instanceof GeoPointFieldType) {
-            GeoPoint originGeoPoint;
-            if (originObj instanceof GeoPoint) {
-                originGeoPoint = (GeoPoint) originObj;
-            } else if (originObj instanceof String) {
-                originGeoPoint = GeoUtils.parseFromString((String) originObj);
-            } else {
-                throw new IllegalArgumentException("Illegal type ["+ origin.getClass() + "] for [origin]! " +
-                    "Must be of type [geo_point] or [string] for geo_point fields!");
-            }
-            double pivotDouble = DistanceUnit.DEFAULT.parse(pivot, DistanceUnit.DEFAULT);
-            return LatLonPoint.newDistanceFeatureQuery(field, boost, originGeoPoint.lat(), originGeoPoint.lon(), pivotDouble);
-        }
-        throw new IllegalArgumentException("Illegal data type of [" + fieldType.typeName() + "]!"+
-            "[" + NAME + "] query can only be run on a date, date_nanos or geo_point field type!");
+        return fieldType.distanceFeatureQuery(origin.origin(), pivot, boost, context);
     }
 
     String fieldName() {