removes the CellIdSource abstraction from geo-grid aggs (#45307)

CellIdSource is a helper ValuesSource that encodes GeoPoint
into a long-encoded representation of the grid bucket the point
is associated with. This complicates thing as usage evolves to
support shapes that are associated with more than one bucket ordinal.

server/src/main/java/org/elasticsearch/search/aggregations/bucket/geogrid/CellIdSource.java

@@ -1,106 +0,0 @@
-/*
- * Licensed to Elasticsearch under one or more contributor
- * license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright
- * ownership. Elasticsearch licenses this file to you under
- * the Apache License, Version 2.0 (the "License"); you may
- * not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *    http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied.  See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-package org.elasticsearch.search.aggregations.bucket.geogrid;
-
-import org.apache.lucene.index.LeafReaderContext;
-import org.apache.lucene.index.SortedNumericDocValues;
-import org.elasticsearch.index.fielddata.AbstractSortingNumericDocValues;
-import org.elasticsearch.index.fielddata.MultiGeoPointValues;
-import org.elasticsearch.index.fielddata.SortedBinaryDocValues;
-import org.elasticsearch.index.fielddata.SortedNumericDoubleValues;
-import org.elasticsearch.search.aggregations.support.ValuesSource;
-
-import java.io.IOException;
-
-/**
- * Wrapper class to help convert {@link MultiGeoPointValues}
- * to numeric long values for bucketing.
- */
-class CellIdSource extends ValuesSource.Numeric {
-    private final ValuesSource.GeoPoint valuesSource;
-    private final int precision;
-    private final GeoPointLongEncoder encoder;
-
-    CellIdSource(GeoPoint valuesSource, int precision, GeoPointLongEncoder encoder) {
-        this.valuesSource = valuesSource;
-        //different GeoPoints could map to the same or different hashing cells.
-        this.precision = precision;
-        this.encoder = encoder;
-    }
-
-    public int precision() {
-        return precision;
-    }
-
-    @Override
-    public boolean isFloatingPoint() {
-        return false;
-    }
-
-    @Override
-    public SortedNumericDocValues longValues(LeafReaderContext ctx) {
-        return new CellValues(valuesSource.geoPointValues(ctx), precision, encoder);
-    }
-
-    @Override
-    public SortedNumericDoubleValues doubleValues(LeafReaderContext ctx) {
-        throw new UnsupportedOperationException();
-    }
-
-    @Override
-    public SortedBinaryDocValues bytesValues(LeafReaderContext ctx) {
-        throw new UnsupportedOperationException();
-    }
-
-    /**
-     * The encoder to use to convert a geopoint's (lon, lat, precision) into
-     * a long-encoded bucket key for aggregating.
-     */
-    @FunctionalInterface
-    public interface GeoPointLongEncoder {
-        long encode(double lon, double lat, int precision);
-    }
-
-    private static class CellValues extends AbstractSortingNumericDocValues {
-        private MultiGeoPointValues geoValues;
-        private int precision;
-        private GeoPointLongEncoder encoder;
-
-        protected CellValues(MultiGeoPointValues geoValues, int precision, GeoPointLongEncoder encoder) {
-            this.geoValues = geoValues;
-            this.precision = precision;
-            this.encoder = encoder;
-        }
-
-        @Override
-        public boolean advanceExact(int docId) throws IOException {
-            if (geoValues.advanceExact(docId)) {
-                resize(geoValues.docValueCount());
-                for (int i = 0; i < docValueCount(); ++i) {
-                    org.elasticsearch.common.geo.GeoPoint target = geoValues.nextValue();
-                    values[i] = encoder.encode(target.getLon(), target.getLat(), precision);
-                }
-                sort();
-                return true;
-            } else {
-                return false;
-            }
-        }
-    }
-}

server/src/main/java/org/elasticsearch/search/aggregations/bucket/geogrid/GeoGridAggregator.java

@@ -19,10 +19,11 @@
 package org.elasticsearch.search.aggregations.bucket.geogrid;
 
 import org.apache.lucene.index.LeafReaderContext;
-import org.apache.lucene.index.SortedNumericDocValues;
 import org.apache.lucene.search.ScoreMode;
+import org.elasticsearch.common.geo.GeoPoint;
 import org.elasticsearch.common.lease.Releasables;
 import org.elasticsearch.common.util.LongHash;
+import org.elasticsearch.index.fielddata.MultiGeoPointValues;
 import org.elasticsearch.search.aggregations.Aggregator;
 import org.elasticsearch.search.aggregations.AggregatorFactories;
 import org.elasticsearch.search.aggregations.InternalAggregations;
@@ -30,6 +31,7 @@ import org.elasticsearch.search.aggregations.LeafBucketCollector;
 import org.elasticsearch.search.aggregations.LeafBucketCollectorBase;
 import org.elasticsearch.search.aggregations.bucket.BucketsAggregator;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
+import org.elasticsearch.search.aggregations.support.ValuesSource;
 import org.elasticsearch.search.internal.SearchContext;
 
 import java.io.IOException;
@@ -45,14 +47,19 @@ public abstract class GeoGridAggregator<T extends InternalGeoGrid> extends Bucke
 
     protected final int requiredSize;
     protected final int shardSize;
-    protected final CellIdSource valuesSource;
+    protected final ValuesSource.GeoPoint valuesSource;
+    protected final int precision;
+    protected final GeoPointLongEncoder longEncoder;
     protected final LongHash bucketOrds;
 
-    GeoGridAggregator(String name, AggregatorFactories factories, CellIdSource valuesSource,
+    GeoGridAggregator(String name, AggregatorFactories factories, ValuesSource.GeoPoint valuesSource,
+                      int precision, GeoPointLongEncoder longEncoder,
                       int requiredSize, int shardSize, SearchContext aggregationContext, Aggregator parent,
                       List<PipelineAggregator> pipelineAggregators, Map<String, Object> metaData) throws IOException {
         super(name, factories, aggregationContext, parent, pipelineAggregators, metaData);
         this.valuesSource = valuesSource;
+        this.precision = precision;
+        this.longEncoder = longEncoder;
         this.requiredSize = requiredSize;
         this.shardSize = shardSize;
         bucketOrds = new LongHash(1, aggregationContext.bigArrays());
@@ -69,7 +76,7 @@ public abstract class GeoGridAggregator<T extends InternalGeoGrid> extends Bucke
     @Override
     public LeafBucketCollector getLeafCollector(LeafReaderContext ctx,
             final LeafBucketCollector sub) throws IOException {
-        final SortedNumericDocValues values = valuesSource.longValues(ctx);
+        final MultiGeoPointValues values = valuesSource.geoPointValues(ctx);
         return new LeafBucketCollectorBase(sub, null) {
             @Override
             public void collect(int doc, long bucket) throws IOException {
@@ -79,7 +86,8 @@ public abstract class GeoGridAggregator<T extends InternalGeoGrid> extends Bucke
 
                     long previous = Long.MAX_VALUE;
                     for (int i = 0; i < valuesCount; ++i) {
-                        final long val = values.nextValue();
+                        final GeoPoint point = values.nextValue();
+                        final long val = longEncoder.encode(point.getLon(), point.getLat(), precision);
                         if (previous != val || i == 0) {
                             long bucketOrdinal = bucketOrds.add(val);
                             if (bucketOrdinal < 0) { // already seen
@@ -189,4 +197,12 @@ public abstract class GeoGridAggregator<T extends InternalGeoGrid> extends Bucke
         Releasables.close(bucketOrds);
     }
 
+    /**
+     * The encoder to use to convert a geopoint's (lon, lat, precision) into
+     * a long-encoded bucket key for aggregating.
+     */
+    @FunctionalInterface
+    public interface GeoPointLongEncoder {
+        long encode(double lon, double lat, int precision);
+    }
 }

server/src/main/java/org/elasticsearch/search/aggregations/bucket/geogrid/GeoHashGridAggregator.java

@@ -21,6 +21,7 @@ package org.elasticsearch.search.aggregations.bucket.geogrid;
 import org.elasticsearch.search.aggregations.Aggregator;
 import org.elasticsearch.search.aggregations.AggregatorFactories;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
+import org.elasticsearch.search.aggregations.support.ValuesSource;
 import org.elasticsearch.search.internal.SearchContext;
 
 import java.io.IOException;
@@ -33,10 +34,12 @@ import java.util.Map;
  */
 public class GeoHashGridAggregator extends GeoGridAggregator<InternalGeoHashGrid> {
 
-    GeoHashGridAggregator(String name, AggregatorFactories factories, CellIdSource valuesSource,
-                          int requiredSize, int shardSize, SearchContext aggregationContext, Aggregator parent,
-                          List<PipelineAggregator> pipelineAggregators, Map<String, Object> metaData) throws IOException {
-        super(name, factories, valuesSource, requiredSize, shardSize, aggregationContext, parent, pipelineAggregators, metaData);
+    GeoHashGridAggregator(String name, AggregatorFactories factories,
+                          ValuesSource.GeoPoint valuesSource, int precision, GeoPointLongEncoder longEncoder,
+                          SearchContext aggregationContext, Aggregator parent, List<PipelineAggregator> pipelineAggregators,
+                          Map<String, Object> metaData, int requiredSize, int shardSize) throws IOException {
+        super(name, factories, valuesSource, precision, longEncoder, requiredSize, shardSize, aggregationContext, parent,
+            pipelineAggregators, metaData);
     }
 
     @Override

server/src/main/java/org/elasticsearch/search/aggregations/bucket/geogrid/GeoHashGridAggregatorFactory.java

@@ -71,8 +71,7 @@ public class GeoHashGridAggregatorFactory extends ValuesSourceAggregatorFactory<
         if (collectsFromSingleBucket == false) {
             return asMultiBucketAggregator(this, context, parent);
         }
-        CellIdSource cellIdSource = new CellIdSource(valuesSource, precision, Geohash::longEncode);
-        return new GeoHashGridAggregator(name, factories, cellIdSource, requiredSize, shardSize, context, parent,
-                pipelineAggregators, metaData);
+        return new GeoHashGridAggregator(name, factories, valuesSource, precision, Geohash::longEncode, context, parent,
+            pipelineAggregators, metaData, requiredSize, shardSize);
     }
 }

server/src/main/java/org/elasticsearch/search/aggregations/bucket/geogrid/GeoTileGridAggregator.java

@@ -22,6 +22,7 @@ package org.elasticsearch.search.aggregations.bucket.geogrid;
 import org.elasticsearch.search.aggregations.Aggregator;
 import org.elasticsearch.search.aggregations.AggregatorFactories;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
+import org.elasticsearch.search.aggregations.support.ValuesSource;
 import org.elasticsearch.search.internal.SearchContext;
 
 import java.io.IOException;
@@ -34,10 +35,12 @@ import java.util.Map;
  */
 public class GeoTileGridAggregator extends GeoGridAggregator<InternalGeoTileGrid> {
 
-    GeoTileGridAggregator(String name, AggregatorFactories factories, CellIdSource valuesSource,
+    GeoTileGridAggregator(String name, AggregatorFactories factories, ValuesSource.GeoPoint valuesSource,
+                          int precision, GeoPointLongEncoder longEncoder,
                           int requiredSize, int shardSize, SearchContext aggregationContext, Aggregator parent,
                           List<PipelineAggregator> pipelineAggregators, Map<String, Object> metaData) throws IOException {
-        super(name, factories, valuesSource, requiredSize, shardSize, aggregationContext, parent, pipelineAggregators, metaData);
+        super(name, factories, valuesSource, precision, longEncoder, requiredSize, shardSize, aggregationContext, parent,
+            pipelineAggregators, metaData);
     }
 
     @Override

server/src/main/java/org/elasticsearch/search/aggregations/bucket/geogrid/GeoTileGridAggregatorFactory.java

@@ -71,8 +71,7 @@ public class GeoTileGridAggregatorFactory extends ValuesSourceAggregatorFactory<
         if (collectsFromSingleBucket == false) {
             return asMultiBucketAggregator(this, context, parent);
         }
-        CellIdSource cellIdSource = new CellIdSource(valuesSource, precision, GeoTileUtils::longEncode);
-        return new GeoTileGridAggregator(name, factories, cellIdSource, requiredSize, shardSize, context, parent,
-                pipelineAggregators, metaData);
+        return new GeoTileGridAggregator(name, factories, valuesSource, precision, GeoTileUtils::longEncode, requiredSize,
+            shardSize, context, parent, pipelineAggregators, metaData);
     }
 }