From 525108807207f83f4905598386ed5aa7ba20b99f Mon Sep 17 00:00:00 2001
From: Ignacio Vera <ivera@apache.org>
Date: Thu, 31 Oct 2019 07:55:25 +0100
Subject: [PATCH] Add IndexOrDocValuesQuery to GeoPolygonQueryBuilder (#48449)

---
 .../index/query/GeoPolygonQueryBuilder.java   | 10 ++++++-
 .../query/GeoPolygonQueryBuilderTests.java    | 29 ++++++++++++++++---
 2 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/server/src/main/java/org/elasticsearch/index/query/GeoPolygonQueryBuilder.java b/server/src/main/java/org/elasticsearch/index/query/GeoPolygonQueryBuilder.java
index a07b4186ed594..47c18e6494d13 100644
--- a/server/src/main/java/org/elasticsearch/index/query/GeoPolygonQueryBuilder.java
+++ b/server/src/main/java/org/elasticsearch/index/query/GeoPolygonQueryBuilder.java
@@ -19,8 +19,10 @@
 
 package org.elasticsearch.index.query;
 
+import org.apache.lucene.document.LatLonDocValuesField;
 import org.apache.lucene.document.LatLonPoint;
 import org.apache.lucene.geo.Polygon;
+import org.apache.lucene.search.IndexOrDocValuesQuery;
 import org.apache.lucene.search.MatchNoDocsQuery;
 import org.apache.lucene.search.Query;
 import org.elasticsearch.common.ParseField;
@@ -199,7 +201,13 @@ protected Query doToQuery(QueryShardContext context) throws IOException {
             lons[i] = p.lon();
         }
 
-        return LatLonPoint.newPolygonQuery(fieldType.name(), new Polygon(lats, lons));
+        Polygon polygon = new Polygon(lats, lons);
+        Query query = LatLonPoint.newPolygonQuery(fieldType.name(), polygon);
+        if (fieldType.hasDocValues()) {
+            Query dvQuery = LatLonDocValuesField.newSlowPolygonQuery(fieldType.name(), polygon);
+            query = new IndexOrDocValuesQuery(query, dvQuery);
+        }
+        return query;
     }
 
     @Override
diff --git a/server/src/test/java/org/elasticsearch/index/query/GeoPolygonQueryBuilderTests.java b/server/src/test/java/org/elasticsearch/index/query/GeoPolygonQueryBuilderTests.java
index 0ae4f6f9fd895..49c081270e723 100644
--- a/server/src/test/java/org/elasticsearch/index/query/GeoPolygonQueryBuilderTests.java
+++ b/server/src/test/java/org/elasticsearch/index/query/GeoPolygonQueryBuilderTests.java
@@ -19,11 +19,15 @@
 
 package org.elasticsearch.index.query;
 
+import org.apache.lucene.document.LatLonDocValuesField;
+import org.apache.lucene.document.LatLonPoint;
+import org.apache.lucene.search.IndexOrDocValuesQuery;
 import org.apache.lucene.search.MatchNoDocsQuery;
 import org.apache.lucene.search.Query;
 import org.elasticsearch.common.ParsingException;
 import org.elasticsearch.common.geo.GeoPoint;
 import org.elasticsearch.common.geo.builders.ShapeBuilder;
+import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.test.AbstractQueryTestCase;
 import org.elasticsearch.test.geo.RandomShapeGenerator;
 import org.elasticsearch.test.geo.RandomShapeGenerator.ShapeType;
@@ -53,12 +57,30 @@ protected GeoPolygonQueryBuilder doCreateTestQueryBuilder() {
         if (randomBoolean()) {
             builder.ignoreUnmapped(randomBoolean());
         }
+
         return builder;
     }
 
     @Override
     protected void doAssertLuceneQuery(GeoPolygonQueryBuilder queryBuilder, Query query, QueryShardContext context) throws IOException {
-        // todo LatLonPointInPolygon is package private
+        MappedFieldType fieldType = context.fieldMapper(queryBuilder.fieldName());
+        if (fieldType == null) {
+            assertTrue("Found no indexed geo query.", query instanceof MatchNoDocsQuery);
+        } else { // TODO: Test case when there are no docValues
+            Query indexQuery = ((IndexOrDocValuesQuery) query).getIndexQuery();
+            String expectedFieldName = expectedFieldName(queryBuilder.fieldName());
+            List<GeoPoint> points = queryBuilder.points();
+            double[] lats = new double[points.size()];
+            double[] lons = new double[points.size()];
+            for (int i =0; i < points.size(); i++) {
+                lats[i] = points.get(i).getLat();
+                lons[i] = points.get(i).getLon();
+            }
+            org.apache.lucene.geo.Polygon polygon = new org.apache.lucene.geo.Polygon(lats, lons);
+            assertEquals(LatLonPoint.newPolygonQuery(expectedFieldName, polygon), indexQuery);
+            Query dvQuery = ((IndexOrDocValuesQuery) query).getRandomAccessQuery();
+            assertEquals(LatLonDocValuesField.newSlowPolygonQuery(expectedFieldName, polygon), dvQuery);
+        }
     }
 
     private static List<GeoPoint> randomPolygon() {
@@ -196,9 +218,8 @@ public void testParsingAndToQuery4() throws IOException {
 
     private void assertGeoPolygonQuery(String query) throws IOException {
         QueryShardContext context = createShardContext();
-        parseQuery(query).toQuery(context);
-        // TODO LatLonPointInPolygon is package private, need a closeTo check on the query
-        // since some points can be computed from the geohash
+        GeoPolygonQueryBuilder queryBuilder = (GeoPolygonQueryBuilder) parseQuery(query);
+        doAssertLuceneQuery(queryBuilder, queryBuilder.toQuery(context), context);
     }
 
     public void testFromJson() throws IOException {