Support for ScriptedMetric to be used in the order clause of an aggregation.

Author: lquerel

core/src/main/java/org/elasticsearch/search/aggregations/AggregatorFactory.java

@@ -23,13 +23,15 @@
 import org.elasticsearch.common.lease.Releasables;
 import org.elasticsearch.common.util.BigArrays;
 import org.elasticsearch.common.util.ObjectArray;
+import org.elasticsearch.script.*;
+import org.elasticsearch.search.aggregations.metrics.NumericMetricsAggregator;
 import org.elasticsearch.search.aggregations.pipeline.PipelineAggregator;
 import org.elasticsearch.search.aggregations.support.AggregationContext;
 import org.elasticsearch.search.internal.SearchContext.Lifetime;
+import org.elasticsearch.transport.TransportRequest;
 
 import java.io.IOException;
-import java.util.List;
-import java.util.Map;
+import java.util.*;
 
 /**
  * A factory that knows how to create an {@link Aggregator} of a specific type.
@@ -234,4 +236,143 @@ public void close() {
         };
     }
 
+    // Specific implementation used by ScriptedMetricAggregator
+    // In this context we need to provide the reduce function to implement the metric method inherited
+    // from the NumericMetricsAggregator class.
+    protected static Aggregator asMultiBucketAggregator(final AggregatorFactory factory,
+                                                         final AggregationContext context,
+                                                         final Aggregator parent,
+                                                         Script reduceScript) throws IOException {
+        final Aggregator first = factory.create(context, parent, true);
+        final BigArrays bigArrays = context.bigArrays();
+
+        first.preCollection();
+
+        // Returns a NumericMetricsAggregator instead of a simple Aggregator. So the result of the aggregation
+        // can be used to order the result.
+        // It's NumericMetricsAggregator.SingleValue because the getProperty method in InternalScriptedMetric
+        // only supports single value (path: value).
+        return new NumericMetricsAggregator.SingleValue("",context, parent, new ArrayList<>(), null) {
+            @Override
+            public double metric(final long owningBucketOrd) {
+                try {
+                    Object aggregationObject = buildAggregation(owningBucketOrd).getProperty("value");
+                    List<Object> aggregationObjects = Arrays.asList(aggregationObject);
+                    Map<String, Object> vars = new HashMap<>();
+                    vars.put("_aggs", aggregationObjects);
+
+                    if (reduceScript.getParams() != null) {
+                        vars.putAll(reduceScript.getParams());
+                    }
+                    ScriptService scriptService = context().searchContext().scriptService();
+                    CompiledScript compiledScript = scriptService.compile(reduceScript,
+                            ScriptContext.Standard.AGGS, new InternalAggregation.ReduceContext(bigArrays, scriptService, new TransportRequest.Empty()));
+                    ExecutableScript script = scriptService.executable(compiledScript, vars);
+                    Object value = script.run();
+
+                    if(value instanceof Number) {
+                        return ((Number) value).doubleValue();
+                    } else {
+                        throw new AggregationExecutionException("Invalid order path ["+this+
+                                "]. Only numeric result are supported.");
+                    }
+                } catch (IOException e) {
+                    throw new AggregationExecutionException("Failed to build aggregation [" + name() + "]", e);
+                }
+            }
+
+            @Override
+            protected LeafBucketCollector getLeafCollector(final LeafReaderContext ctx, final LeafBucketCollector sub) throws IOException {
+                for (long i = 0; i < collectors.size(); ++i) {
+                    collectors.set(i, null);
+                }
+                return new LeafBucketCollector() {
+                    Scorer scorer;
+
+                    @Override
+                    public void setScorer(Scorer scorer) throws IOException {
+                        this.scorer = scorer;
+                    }
+
+                    @Override
+                    public void collect(int doc, long bucket) throws IOException {
+                        aggregators = bigArrays.grow(aggregators, bucket + 1);
+                        collectors = bigArrays.grow(collectors, bucket + 1);
+
+                        LeafBucketCollector collector = collectors.get(bucket);
+                        if (collector == null) {
+                            Aggregator aggregator = aggregators.get(bucket);
+                            if (aggregator == null) {
+                                aggregator = factory.create(context, parent, true);
+                                aggregator.preCollection();
+                                aggregators.set(bucket, aggregator);
+                            }
+                            collector = aggregator.getLeafCollector(ctx);
+                            collector.setScorer(scorer);
+                            collectors.set(bucket, collector);
+                        }
+                        collector.collect(doc, 0);
+                    }
+
+                };
+            }
+
+            ObjectArray<Aggregator> aggregators;
+            ObjectArray<LeafBucketCollector> collectors;
+
+            {
+                context.searchContext().addReleasable(this, Lifetime.PHASE);
+                aggregators = bigArrays.newObjectArray(1);
+                aggregators.set(0, first);
+                collectors = bigArrays.newObjectArray(1);
+            }
+
+            @Override
+            public String name() {
+                return first.name();
+            }
+
+            @Override
+            public AggregationContext context() {
+                return first.context();
+            }
+
+            @Override
+            public Aggregator parent() {
+                return first.parent();
+            }
+
+            @Override
+            public boolean needsScores() {
+                return first.needsScores();
+            }
+
+            @Override
+            public Aggregator subAggregator(String name) {
+                throw new UnsupportedOperationException();
+            }
+
+            @Override
+            public InternalAggregation buildAggregation(long bucket) throws IOException {
+                if (bucket < aggregators.size()) {
+                    Aggregator aggregator = aggregators.get(bucket);
+                    if (aggregator != null) {
+                        return aggregator.buildAggregation(0);
+                    }
+                }
+                return buildEmptyAggregation();
+            }
+
+            @Override
+            public InternalAggregation buildEmptyAggregation() {
+                return first.buildEmptyAggregation();
+            }
+
+            @Override
+            public void close() {
+                Releasables.close(aggregators, collectors);
+            }
+        };
+    }
+
 }

core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregator.java

@@ -127,7 +127,7 @@ public Factory(String name, Script initScript, Script mapScript, Script combineS
         public Aggregator createInternal(AggregationContext context, Aggregator parent, boolean collectsFromSingleBucket,
                 List<PipelineAggregator> pipelineAggregators, Map<String, Object> metaData) throws IOException {
             if (collectsFromSingleBucket == false) {
-                return asMultiBucketAggregator(this, context, parent);
+                return asMultiBucketAggregator(this, context, parent, reduceScript);
             }
             Map<String, Object> params = this.params;
             if (params != null) {

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

@@ -20,10 +20,7 @@
 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.*;
 import org.elasticsearch.search.aggregations.bucket.SingleBucketAggregation;
 import org.elasticsearch.search.aggregations.bucket.SingleBucketAggregator;
 import org.elasticsearch.search.aggregations.metrics.InternalNumericMetricsAggregation;
@@ -240,7 +237,24 @@ public double resolveValue(HasAggregations root) {
                         "]. Missing value key in [" + token + "] which refers to a multi-value metric aggregation");
             }
             parent = null;
-            value = ((InternalNumericMetricsAggregation.MultiValue) agg).value(token.key);
+            if(agg instanceof InternalNumericMetricsAggregation.MultiValue) {
+                // For NumericMetricsAggregation we use the method value returning a native double
+                // Optimization to avoid object creation and multiple casts.
+                value = ((InternalNumericMetricsAggregation.MultiValue) agg).value(token.key);
+            } else if(agg instanceof InternalAggregation) {
+                // For a general use case, we use the method getProperty returning an Object
+                Object propertyValue = agg.getProperty(token.key);
+                // Only aggregation returning a numeric value are supported.
+                if(propertyValue instanceof Number) {
+                    value = ((Number) propertyValue).doubleValue();
+                } else {
+                    throw new AggregationExecutionException("Invalid order path ["+this+
+                            "]. Only numeric result are supported.");
+                }
+            } else {
+                throw new AggregationExecutionException("Invalid aggregation type for order path ["+this+
+                        "]. Only numeric & scripted metric aggregation are supported.");
+            }
         }
 
         return value;