Revert "Optimize geogrid aggregations for singleton points (#87290)"

This reverts commit 30406ba.

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

@@ -8,24 +8,16 @@
 
 package org.elasticsearch.search.aggregations.bucket.geogrid;
 
-import org.apache.lucene.index.DocValues;
 import org.apache.lucene.index.LeafReaderContext;
-import org.apache.lucene.index.NumericDocValues;
 import org.apache.lucene.index.SortedNumericDocValues;
 import org.elasticsearch.common.geo.GeoBoundingBox;
-import org.elasticsearch.index.fielddata.AbstractNumericDocValues;
-import org.elasticsearch.index.fielddata.AbstractSortingNumericDocValues;
-import org.elasticsearch.index.fielddata.GeoPointValues;
 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;
-
 /**
- * Base class to help convert {@link MultiGeoPointValues} to {@link CellMultiValues}
- * and {@link GeoPointValues} to {@link CellSingleValue}
+ * Base class to help convert {@link MultiGeoPointValues} to {@link CellValues}
  */
 public abstract class CellIdSource extends ValuesSource.Numeric {
 
@@ -53,35 +45,22 @@ public final boolean isFloatingPoint() {
     @Override
     public final SortedNumericDocValues longValues(LeafReaderContext ctx) {
         final MultiGeoPointValues multiGeoPointValues = valuesSource.geoPointValues(ctx);
-        final GeoPointValues values = org.elasticsearch.index.fielddata.FieldData.unwrapSingleton(multiGeoPointValues);
         if (geoBoundingBox.isUnbounded()) {
-            return values == null ? unboundedCellMultiValues(multiGeoPointValues) : DocValues.singleton(unboundedCellSingleValue(values));
+            return unboundedCellValues(multiGeoPointValues);
         } else {
-            return values == null
-                ? boundedCellMultiValues(multiGeoPointValues, geoBoundingBox)
-                : DocValues.singleton(boundedCellSingleValue(values, geoBoundingBox));
+            return boundedCellValues(multiGeoPointValues, geoBoundingBox);
         }
     }
 
     /**
-     * Generate an unbounded iterator of grid-cells for singleton case.
-     */
-    protected abstract NumericDocValues unboundedCellSingleValue(GeoPointValues values);
-
-    /**
-     * Generate a bounded iterator of grid-cells for singleton case.
+     * Generate an unbounded iterator of grid-cells
      */
-    protected abstract NumericDocValues boundedCellSingleValue(GeoPointValues values, GeoBoundingBox boundingBox);
+    protected abstract CellValues unboundedCellValues(MultiGeoPointValues values);
 
     /**
-     * Generate an unbounded iterator of grid-cells for multi-value case.
+     * Generate a bounded iterator of grid-cells
      */
-    protected abstract SortedNumericDocValues unboundedCellMultiValues(MultiGeoPointValues values);
-
-    /**
-     * Generate a bounded iterator of grid-cells for multi-value case.
-     */
-    protected abstract SortedNumericDocValues boundedCellMultiValues(MultiGeoPointValues values, GeoBoundingBox boundingBox);
+    protected abstract CellValues boundedCellValues(MultiGeoPointValues values, GeoBoundingBox boundingBox);
 
     @Override
     public final SortedNumericDoubleValues doubleValues(LeafReaderContext ctx) {
@@ -110,85 +89,4 @@ protected boolean validPoint(double lon, double lat) {
         return false;
     }
 
-    /**
-     * Class representing the long-encoded grid-cells belonging to
-     * the multi-value geo-doc-values. Class must encode the values and then
-     * sort them in order to account for the cells correctly.
-     */
-    protected abstract static class CellMultiValues extends AbstractSortingNumericDocValues {
-        private final MultiGeoPointValues geoValues;
-        protected final int precision;
-
-        protected CellMultiValues(MultiGeoPointValues geoValues, int precision) {
-            this.geoValues = geoValues;
-            this.precision = precision;
-        }
-
-        @Override
-        public boolean advanceExact(int docId) throws IOException {
-            if (geoValues.advanceExact(docId)) {
-                int docValueCount = geoValues.docValueCount();
-                resize(docValueCount);
-                int j = 0;
-                for (int i = 0; i < docValueCount; i++) {
-                    j = advanceValue(geoValues.nextValue(), j);
-                }
-                resize(j);
-                sort();
-                return true;
-            } else {
-                return false;
-            }
-        }
-
-        /**
-         * Sets the appropriate long-encoded value for <code>target</code>
-         * in <code>values</code>.
-         *
-         * @param target    the geo-value to encode
-         * @param valuesIdx the index into <code>values</code> to set
-         * @return          valuesIdx + 1 if value was set, valuesIdx otherwise.
-         */
-        protected abstract int advanceValue(org.elasticsearch.common.geo.GeoPoint target, int valuesIdx);
-    }
-
-    /**
-     * Class representing the long-encoded grid-cells belonging to
-     * the singleton geo-doc-values.
-     */
-    protected abstract static class CellSingleValue extends AbstractNumericDocValues {
-        private final GeoPointValues geoValues;
-        protected final int precision;
-        protected long value;
-
-        protected CellSingleValue(GeoPointValues geoValues, int precision) {
-            this.geoValues = geoValues;
-            this.precision = precision;
-
-        }
-
-        @Override
-        public boolean advanceExact(int docId) throws IOException {
-            return geoValues.advanceExact(docId) && advance(geoValues.geoPointValue());
-        }
-
-        @Override
-        public long longValue() throws IOException {
-            return value;
-        }
-
-        /**
-         * Sets the appropriate long-encoded value for <code>target</code>
-         * in <code>value</code>.
-         *
-         * @param target    the geo-value to encode
-         * @return          true if the value needs to be added, otherwise false.
-         */
-        protected abstract boolean advance(org.elasticsearch.common.geo.GeoPoint target);
-
-        @Override
-        public int docID() {
-            return -1;
-        }
-    }
 }

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

@@ -0,0 +1,55 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+package org.elasticsearch.search.aggregations.bucket.geogrid;
+
+import org.elasticsearch.index.fielddata.AbstractSortingNumericDocValues;
+import org.elasticsearch.index.fielddata.MultiGeoPointValues;
+
+import java.io.IOException;
+
+/**
+ * Class representing the long-encoded grid-cells belonging to
+ * the geo-doc-values. Class must encode the values and then
+ * sort them in order to account for the cells correctly.
+ */
+public abstract class CellValues extends AbstractSortingNumericDocValues {
+    private MultiGeoPointValues geoValues;
+    protected int precision;
+
+    protected CellValues(MultiGeoPointValues geoValues, int precision) {
+        this.geoValues = geoValues;
+        this.precision = precision;
+    }
+
+    @Override
+    public boolean advanceExact(int docId) throws IOException {
+        if (geoValues.advanceExact(docId)) {
+            int docValueCount = geoValues.docValueCount();
+            resize(docValueCount);
+            int j = 0;
+            for (int i = 0; i < docValueCount; i++) {
+                j = advanceValue(geoValues.nextValue(), j);
+            }
+            resize(j);
+            sort();
+            return true;
+        } else {
+            return false;
+        }
+    }
+
+    /**
+     * Sets the appropriate long-encoded value for <code>target</code>
+     * in <code>values</code>.
+     *
+     * @param target    the geo-value to encode
+     * @param valuesIdx the index into <code>values</code> to set
+     * @return          valuesIdx + 1 if value was set, valuesIdx otherwise.
+     */
+    protected abstract int advanceValue(org.elasticsearch.common.geo.GeoPoint target, int valuesIdx);
+}

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

@@ -7,9 +7,7 @@
  */
 package org.elasticsearch.search.aggregations.bucket.geogrid;
 
-import org.apache.lucene.index.DocValues;
 import org.apache.lucene.index.LeafReaderContext;
-import org.apache.lucene.index.NumericDocValues;
 import org.apache.lucene.index.SortedNumericDocValues;
 import org.apache.lucene.search.ScoreMode;
 import org.elasticsearch.core.Releasables;
@@ -67,31 +65,8 @@ public ScoreMode scoreMode() {
     }
 
     @Override
-    public LeafBucketCollector getLeafCollector(final LeafReaderContext ctx, final LeafBucketCollector sub) throws IOException {
-        final SortedNumericDocValues values = valuesSource.longValues(ctx);
-        final NumericDocValues singleton = DocValues.unwrapSingleton(values);
-        return singleton != null ? getLeafCollector(singleton, sub) : getLeafCollector(values, sub);
-    }
-
-    private LeafBucketCollector getLeafCollector(final NumericDocValues values, final LeafBucketCollector sub) {
-        return new LeafBucketCollectorBase(sub, null) {
-            @Override
-            public void collect(int doc, long owningBucketOrd) throws IOException {
-                if (values.advanceExact(doc)) {
-                    final long val = values.longValue();
-                    long bucketOrdinal = bucketOrds.add(owningBucketOrd, val);
-                    if (bucketOrdinal < 0) { // already seen
-                        bucketOrdinal = -1 - bucketOrdinal;
-                        collectExistingBucket(sub, doc, bucketOrdinal);
-                    } else {
-                        collectBucket(sub, doc, bucketOrdinal);
-                    }
-                }
-            }
-        };
-    }
-
-    private LeafBucketCollector getLeafCollector(final SortedNumericDocValues values, final LeafBucketCollector sub) {
+    public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBucketCollector sub) throws IOException {
+        SortedNumericDocValues values = valuesSource.longValues(ctx);
         return new LeafBucketCollectorBase(sub, null) {
             @Override
             public void collect(int doc, long owningBucketOrd) throws IOException {

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

@@ -7,16 +7,13 @@
  */
 package org.elasticsearch.search.aggregations.bucket.geogrid;
 
-import org.apache.lucene.index.NumericDocValues;
-import org.apache.lucene.index.SortedNumericDocValues;
 import org.elasticsearch.common.geo.GeoBoundingBox;
 import org.elasticsearch.geometry.utils.Geohash;
-import org.elasticsearch.index.fielddata.GeoPointValues;
 import org.elasticsearch.index.fielddata.MultiGeoPointValues;
 import org.elasticsearch.search.aggregations.support.ValuesSource;
 
 /**
- * {@link CellIdSource} implementation for Geohash aggregation
+ * Class to help convert {@link MultiGeoPointValues} to Geohash {@link CellValues}
  */
 public class GeoHashCellIdSource extends CellIdSource {
 
@@ -25,35 +22,8 @@ public GeoHashCellIdSource(ValuesSource.GeoPoint valuesSource, int precision, Ge
     }
 
     @Override
-    protected NumericDocValues unboundedCellSingleValue(GeoPointValues values) {
-        return new CellSingleValue(values, precision()) {
-            @Override
-            protected boolean advance(org.elasticsearch.common.geo.GeoPoint target) {
-                value = Geohash.longEncode(target.getLon(), target.getLat(), precision);
-                return true;
-            }
-        };
-    }
-
-    @Override
-    protected NumericDocValues boundedCellSingleValue(GeoPointValues values, GeoBoundingBox boundingBox) {
-        final GeoHashBoundedPredicate predicate = new GeoHashBoundedPredicate(precision(), boundingBox);
-        return new CellSingleValue(values, precision()) {
-            @Override
-            protected boolean advance(org.elasticsearch.common.geo.GeoPoint target) {
-                final String hash = Geohash.stringEncode(target.getLon(), target.getLat(), precision);
-                if (validPoint(target.getLon(), target.getLat()) || predicate.validHash(hash)) {
-                    value = Geohash.longEncode(hash);
-                    return true;
-                }
-                return false;
-            }
-        };
-    }
-
-    @Override
-    protected SortedNumericDocValues unboundedCellMultiValues(MultiGeoPointValues values) {
-        return new CellMultiValues(values, precision()) {
+    protected CellValues unboundedCellValues(MultiGeoPointValues values) {
+        return new CellValues(values, precision()) {
             @Override
             protected int advanceValue(org.elasticsearch.common.geo.GeoPoint target, int valuesIdx) {
                 values[valuesIdx] = Geohash.longEncode(target.getLon(), target.getLat(), precision);
@@ -63,9 +33,9 @@ protected int advanceValue(org.elasticsearch.common.geo.GeoPoint target, int val
     }
 
     @Override
-    protected SortedNumericDocValues boundedCellMultiValues(MultiGeoPointValues values, GeoBoundingBox boundingBox) {
+    protected CellValues boundedCellValues(MultiGeoPointValues values, GeoBoundingBox boundingBox) {
         final GeoHashBoundedPredicate predicate = new GeoHashBoundedPredicate(precision(), boundingBox);
-        return new CellMultiValues(values, precision()) {
+        return new CellValues(values, precision()) {
             @Override
             protected int advanceValue(org.elasticsearch.common.geo.GeoPoint target, int valuesIdx) {
                 final String hash = Geohash.stringEncode(target.getLon(), target.getLat(), precision);

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

@@ -7,14 +7,11 @@
  */
 package org.elasticsearch.search.aggregations.bucket.geogrid;
 
-import org.apache.lucene.index.NumericDocValues;
-import org.apache.lucene.index.SortedNumericDocValues;
 import org.elasticsearch.common.geo.GeoBoundingBox;
-import org.elasticsearch.index.fielddata.GeoPointValues;
 import org.elasticsearch.index.fielddata.MultiGeoPointValues;
 
 /**
- * {@link CellIdSource} implementation for GeoTile aggregation
+ * Class to help convert {@link MultiGeoPointValues} to GeoTile {@link CellValues}
  */
 public class GeoTileCellIdSource extends CellIdSource {
 
@@ -23,37 +20,8 @@ public GeoTileCellIdSource(GeoPoint valuesSource, int precision, GeoBoundingBox
     }
 
     @Override
-    protected NumericDocValues unboundedCellSingleValue(GeoPointValues values) {
-        return new CellSingleValue(values, precision()) {
-            @Override
-            protected boolean advance(org.elasticsearch.common.geo.GeoPoint target) {
-                value = GeoTileUtils.longEncode(target.getLon(), target.getLat(), precision);
-                return true;
-            }
-        };
-    }
-
-    @Override
-    protected NumericDocValues boundedCellSingleValue(GeoPointValues values, GeoBoundingBox boundingBox) {
-        final GeoTileBoundedPredicate predicate = new GeoTileBoundedPredicate(precision(), boundingBox);
-        final long tiles = 1L << precision();
-        return new CellSingleValue(values, precision()) {
-            @Override
-            protected boolean advance(org.elasticsearch.common.geo.GeoPoint target) {
-                final int x = GeoTileUtils.getXTile(target.getLon(), tiles);
-                final int y = GeoTileUtils.getYTile(target.getLat(), tiles);
-                if (predicate.validTile(x, y, precision)) {
-                    value = GeoTileUtils.longEncodeTiles(precision, x, y);
-                    return true;
-                }
-                return false;
-            }
-        };
-    }
-
-    @Override
-    protected SortedNumericDocValues unboundedCellMultiValues(MultiGeoPointValues values) {
-        return new CellMultiValues(values, precision()) {
+    protected CellValues unboundedCellValues(MultiGeoPointValues values) {
+        return new CellValues(values, precision()) {
             @Override
             protected int advanceValue(org.elasticsearch.common.geo.GeoPoint target, int valuesIdx) {
                 values[valuesIdx] = GeoTileUtils.longEncode(target.getLon(), target.getLat(), precision);
@@ -63,10 +31,10 @@ protected int advanceValue(org.elasticsearch.common.geo.GeoPoint target, int val
     }
 
     @Override
-    protected SortedNumericDocValues boundedCellMultiValues(MultiGeoPointValues values, GeoBoundingBox boundingBox) {
+    protected CellValues boundedCellValues(MultiGeoPointValues values, GeoBoundingBox boundingBox) {
         final GeoTileBoundedPredicate predicate = new GeoTileBoundedPredicate(precision(), boundingBox);
         final long tiles = 1L << precision();
-        return new CellMultiValues(values, precision()) {
+        return new CellValues(values, precision()) {
             @Override
             protected int advanceValue(org.elasticsearch.common.geo.GeoPoint target, int valuesIdx) {
                 final int x = GeoTileUtils.getXTile(target.getLon(), tiles);