Make sorting by an agg results a real abstraction (#52007) (#52212)

This removes a bunch of `instanceof`s in favor of two new methods on
`InernalAggregation`. The default implementations of these methods just
throw exceptions explaining that you can't sort on this aggregation.
They are overridden by all of the classes that used to have `instanceof`
checks against them.

I doubt this is really any faster in practice. The real benefit here is
that it is a little more obvious *that* you can sort by the results of
an aggregation and it should be *much* more obvious where to look at
*how* aggregations sort themselves.

There are still a bunch more `instanceof`s in left in `AggregationPath`
but those will wait for a followup change.

Author: nik9000

server/src/main/java/org/elasticsearch/search/aggregations/InternalAggregation.java

@@ -30,6 +30,7 @@
 import org.elasticsearch.search.aggregations.support.AggregationPath;
 
 import java.io.IOException;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -242,4 +243,20 @@ public String toString() {
         return Strings.toString(this);
     }
 
+    /**
+     * Get value to use when sorting by this aggregation.
+     */
+    public double sortValue(String key) {
+        // subclasses will override this with a real implementation if they can be sorted
+        throw new IllegalArgumentException("Can't sort a [" + getType() + "] aggregation [" + getName() + "]");
+    }
+
+    /**
+     * Get value to use when sorting by a descendant of this aggregation.
+     */
+    public double sortValue(AggregationPath.PathElement head, Iterator<AggregationPath.PathElement> tail) {
+        // subclasses will override this with a real implementation if you can sort on a descendant
+        throw new IllegalArgumentException("Can't sort by a descendant of a [" + getType() + "] aggregation [" + head + "]");
+    }
+
 }

server/src/main/java/org/elasticsearch/search/aggregations/InternalAggregations.java

@@ -25,12 +25,14 @@
 import org.elasticsearch.search.aggregations.InternalAggregation.ReduceContext;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
 import org.elasticsearch.search.aggregations.pipeline.SiblingPipelineAggregator;
+import org.elasticsearch.search.aggregations.support.AggregationPath;
 
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Comparator;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -104,6 +106,20 @@ private List<InternalAggregation> getInternalAggregations() {
         return (List<InternalAggregation>) aggregations;
     }
 
+    /**
+     * Get value to use when sorting by a descendant of the aggregation containing this.
+     */
+    public double sortValue(AggregationPath.PathElement head, Iterator<AggregationPath.PathElement> tail) {
+        InternalAggregation aggregation = get(head.name);
+        if (aggregation == null) {
+            throw new IllegalArgumentException("Cannot find aggregation named [" + head.name + "]");
+        }
+        if (tail.hasNext()) {
+            return aggregation.sortValue(tail.next(), tail);
+        }
+        return aggregation.sortValue(head.key);
+    }
+
     /**
      * Begin the reduction process.  This should be the entry point for the "first" reduction, e.g. called by
      * SearchPhaseController or anywhere else that wants to initiate a reduction.  It _should not_ be called

server/src/main/java/org/elasticsearch/search/aggregations/InternalMultiBucketAggregation.java

@@ -192,5 +192,6 @@ public Object getProperty(String containingAggName, List<String> path) {
             }
             return aggregation.getProperty(path.subList(1, path.size()));
         }
+
     }
 }

server/src/main/java/org/elasticsearch/search/aggregations/InternalOrder.java

@@ -156,8 +156,8 @@ static class AggregationComparator implements Comparator<Bucket> {
 
             @Override
             public int compare(Bucket b1, Bucket b2) {
-                double v1 = path.resolveValue(b1);
-                double v2 = path.resolveValue(b2);
+                double v1 = path.resolveValue(((InternalAggregations) b1.getAggregations()));
+                double v2 = path.resolveValue(((InternalAggregations) b2.getAggregations()));
                 return Comparators.compareDiscardNaN(v1, v2, asc);
             }
         }

server/src/main/java/org/elasticsearch/search/aggregations/bucket/InternalSingleBucketAggregation.java

@@ -25,9 +25,11 @@
 import org.elasticsearch.search.aggregations.InternalAggregation;
 import org.elasticsearch.search.aggregations.InternalAggregations;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
+import org.elasticsearch.search.aggregations.support.AggregationPath;
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -152,6 +154,21 @@ public XContentBuilder doXContentBody(XContentBuilder builder, Params params) th
         return builder;
     }
 
+    @Override
+    public final double sortValue(String key) {
+        if (key != null && false == key.equals("doc_count")) {
+            throw new IllegalArgumentException(
+                    "Unknown value key [" + key + "] for single-bucket aggregation [" + getName() +
+                    "]. Either use [doc_count] as key or drop the key all together.");
+        }
+        return docCount;
+    }
+
+    @Override
+    public final double sortValue(AggregationPath.PathElement head, Iterator<AggregationPath.PathElement> tail) {
+        return aggregations.sortValue(head, tail);
+    }
+
     @Override
     public boolean equals(Object obj) {
         if (this == obj) return true;

server/src/main/java/org/elasticsearch/search/aggregations/metrics/InternalNumericMetricsAggregation.java

@@ -22,8 +22,10 @@
 import org.elasticsearch.search.DocValueFormat;
 import org.elasticsearch.search.aggregations.InternalAggregation;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
+import org.elasticsearch.search.aggregations.support.AggregationPath;
 
 import java.io.IOException;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -62,6 +64,15 @@ public Object getProperty(List<String> path) {
             }
         }
 
+        @Override
+        public final double sortValue(String key) {
+            if (key != null && false == key.equals("value")) {
+                throw new IllegalArgumentException(
+                        "Unknown value key [" + key + "] for single-value metric aggregation [" + getName() +
+                        "]. Either use [value] as key or drop the key all together");
+            }
+            return value();
+        }
     }
 
     public abstract static class MultiValue extends InternalNumericMetricsAggregation implements NumericMetricsAggregation.MultiValue {
@@ -92,6 +103,14 @@ public Object getProperty(List<String> path) {
                 throw new IllegalArgumentException("path not supported for [" + getName() + "]: " + path);
             }
         }
+
+        @Override
+        public final double sortValue(String key) {
+            if (key == null) {
+                throw new IllegalArgumentException("Missing value key in [" + key+ "] which refers to a multi-value metric aggregation");
+            }
+            return value(key);
+        }
     }
 
     private InternalNumericMetricsAggregation(String name, List<PipelineAggregator> pipelineAggregators, Map<String, Object> metaData) {
@@ -105,6 +124,11 @@ protected InternalNumericMetricsAggregation(StreamInput in) throws IOException {
         super(in);
     }
 
+    @Override
+    public final double sortValue(AggregationPath.PathElement head, Iterator<AggregationPath.PathElement> tail) {
+        throw new IllegalArgumentException("Metrics aggregations cannot have sub-aggregations (at [>" + head + "]");
+    }
+
     @Override
     public int hashCode() {
         return Objects.hash(super.hashCode(), format);

server/src/main/java/org/elasticsearch/search/aggregations/support/AggregationPath.java

@@ -20,17 +20,15 @@
 package org.elasticsearch.search.aggregations.support;
 
 import org.elasticsearch.common.Strings;
-import org.elasticsearch.search.aggregations.Aggregation;
 import org.elasticsearch.search.aggregations.AggregationExecutionException;
 import org.elasticsearch.search.aggregations.Aggregator;
-import org.elasticsearch.search.aggregations.HasAggregations;
-import org.elasticsearch.search.aggregations.bucket.SingleBucketAggregation;
+import org.elasticsearch.search.aggregations.InternalAggregations;
 import org.elasticsearch.search.aggregations.bucket.SingleBucketAggregator;
-import org.elasticsearch.search.aggregations.metrics.InternalNumericMetricsAggregation;
 import org.elasticsearch.search.aggregations.metrics.NumericMetricsAggregator;
 import org.elasticsearch.search.profile.aggregation.ProfilingAggregator;
 
 import java.util.ArrayList;
+import java.util.Iterator;
 import java.util.List;
 
 /**
@@ -190,61 +188,16 @@ public AggregationPath subPath(int offset, int length) {
     }
 
     /**
-     * Resolves the value pointed by this path given an aggregations root
-     *
-     * @param root  The root that serves as a point of reference for this path
-     * @return      The resolved value
+     * Looks up the value of this path against a set of aggregation results.
      */
-    public double resolveValue(HasAggregations root) {
-        HasAggregations parent = root;
-        double value = Double.NaN;
-        for (int i = 0; i < pathElements.size(); i++) {
-            AggregationPath.PathElement token = pathElements.get(i);
-            Aggregation agg = parent.getAggregations().get(token.name);
-
-            if (agg == null) {
-                throw new IllegalArgumentException("Invalid order path [" + this +
-                        "]. Cannot find aggregation named [" + token.name + "]");
-            }
-
-            if (agg instanceof SingleBucketAggregation) {
-                if (token.key != null && !token.key.equals("doc_count")) {
-                    throw new IllegalArgumentException("Invalid order path [" + this +
-                            "]. Unknown value key [" + token.key + "] for single-bucket aggregation [" + token.name +
-                            "]. Either use [doc_count] as key or drop the key all together");
-                }
-                parent = (SingleBucketAggregation) agg;
-                value = ((SingleBucketAggregation) agg).getDocCount();
-                continue;
-            }
-
-            // the agg can only be a metrics agg, and a metrics agg must be at the end of the path
-            if (i != pathElements.size() - 1) {
-                throw new IllegalArgumentException("Invalid order path [" + this +
- "]. Metrics aggregations cannot have sub-aggregations (at [" + token + ">" + pathElements.get(i + 1) + "]");
-            }
-
-            if (agg instanceof InternalNumericMetricsAggregation.SingleValue) {
-                if (token.key != null && !token.key.equals("value")) {
-                    throw new IllegalArgumentException("Invalid order path [" + this +
-                            "]. Unknown value key [" + token.key + "] for single-value metric aggregation [" + token.name +
-                            "]. Either use [value] as key or drop the key all together");
-                }
-                parent = null;
-                value = ((InternalNumericMetricsAggregation.SingleValue) agg).value();
-                continue;
-            }
-
-            // we're left with a multi-value metric agg
-            if (token.key == null) {
-                throw new IllegalArgumentException("Invalid order path [" + this +
-                        "]. Missing value key in [" + token + "] which refers to a multi-value metric aggregation");
-            }
-            parent = null;
-            value = ((InternalNumericMetricsAggregation.MultiValue) agg).value(token.key);
+    public double resolveValue(InternalAggregations aggregations) {
+        try {
+            Iterator<PathElement> path = pathElements.iterator();
+            assert path.hasNext();
+            return aggregations.sortValue(path.next(), path);
+        } catch (IllegalArgumentException e) {
+            throw new IllegalArgumentException("Invalid order path [" + this + "]. " + e.getMessage(), e);
         }
-
-        return value;
     }
 
     /**

server/src/test/java/org/elasticsearch/search/aggregations/support/PathTests.java renamed to server/src/test/java/org/elasticsearch/search/aggregations/support/AggregationPathTests.java

@@ -27,7 +27,7 @@
 
 import static org.hamcrest.Matchers.equalTo;
 
-public class PathTests extends ESTestCase {
+public class AggregationPathTests extends ESTestCase {
     public void testInvalidPaths() throws Exception {
         assertInvalidPath("[foo]", "brackets at the beginning of the token expression");
         assertInvalidPath("foo[bar", "open brackets without closing at the token expression");
@@ -49,13 +49,8 @@ public void testValidPaths() throws Exception {
         assertValidPath("foo.bar>baz[qux]", tokens().add("foo.bar").add("baz", "qux"));
     }
 
-    private void assertInvalidPath(String path, String reason) {
-        try {
-            AggregationPath.parse(path);
-            fail("Expected parsing path [" + path + "] to fail - " + reason);
-        } catch (AggregationExecutionException aee) {
-            // expected
-        }
+    private AggregationExecutionException assertInvalidPath(String path, String reason) {
+        return expectThrows(AggregationExecutionException.class, () -> AggregationPath.parse(path));
     }
 
     private void assertValidPath(String path, Tokens tokenz) {