From c1ba4fdcb45fd7e2b50f0e5d236522b3e71cd7bc Mon Sep 17 00:00:00 2001 From: Winston Ewert Date: Tue, 2 May 2017 13:24:23 +0300 Subject: [PATCH] Allow scripted metric agg to access `_score` (#24295) * Fixes #24259 Corrects the ScriptedMetricAggregator so that the script can have access to scores during the map stage. * Restored original tests. Added seperate test. As requested, I've restored the non-score dependant tests, and added the score dependent metric as a seperate test. --- .../scripted/ScriptedMetricAggregator.java | 2 +- .../ScriptedMetricAggregatorTests.java | 44 +++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregator.java b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregator.java index 3bfc057682a..cee7b3402f3 100644 --- a/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregator.java +++ b/core/src/main/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregator.java @@ -63,7 +63,7 @@ public class ScriptedMetricAggregator extends MetricsAggregator { public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBucketCollector sub) throws IOException { final LeafSearchScript leafMapScript = mapScript.getLeafSearchScript(ctx); - return new LeafBucketCollectorBase(sub, mapScript) { + return new LeafBucketCollectorBase(sub, leafMapScript) { @Override public void collect(int doc, long bucket) throws IOException { assert bucket == 0 : bucket; diff --git a/core/src/test/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregatorTests.java b/core/src/test/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregatorTests.java index d9eb76310d2..5dcae47da43 100644 --- a/core/src/test/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregatorTests.java +++ b/core/src/test/java/org/elasticsearch/search/aggregations/metrics/scripted/ScriptedMetricAggregatorTests.java @@ -33,6 +33,7 @@ import org.elasticsearch.index.mapper.MappedFieldType; import org.elasticsearch.index.query.QueryShardContext; import org.elasticsearch.indices.breaker.CircuitBreakerService; import org.elasticsearch.script.MockScriptEngine; +import org.elasticsearch.script.ScoreAccessor; import org.elasticsearch.script.Script; import org.elasticsearch.script.ScriptContextRegistry; import org.elasticsearch.script.ScriptEngineRegistry; @@ -59,6 +60,11 @@ public class ScriptedMetricAggregatorTests extends AggregatorTestCase { private static final Script MAP_SCRIPT = new Script(ScriptType.INLINE, MockScriptEngine.NAME, "mapScript", Collections.emptyMap()); private static final Script COMBINE_SCRIPT = new Script(ScriptType.INLINE, MockScriptEngine.NAME, "combineScript", Collections.emptyMap()); + + private static final Script INIT_SCRIPT_SCORE = new Script(ScriptType.INLINE, MockScriptEngine.NAME, "initScriptScore", Collections.emptyMap()); + private static final Script MAP_SCRIPT_SCORE = new Script(ScriptType.INLINE, MockScriptEngine.NAME, "mapScriptScore", Collections.emptyMap()); + private static final Script COMBINE_SCRIPT_SCORE = new Script(ScriptType.INLINE, MockScriptEngine.NAME, "combineScriptScore", + Collections.emptyMap()); private static final Map, Object>> SCRIPTS = new HashMap<>(); @@ -79,6 +85,21 @@ public class ScriptedMetricAggregatorTests extends AggregatorTestCase { Map agg = (Map) params.get("_agg"); return ((List) agg.get("collector")).stream().mapToInt(Integer::intValue).sum(); }); + + SCRIPTS.put("initScriptScore", params -> { + Map agg = (Map) params.get("_agg"); + agg.put("collector", new ArrayList()); + return agg; + }); + SCRIPTS.put("mapScriptScore", params -> { + Map agg = (Map) params.get("_agg"); + ((List) agg.get("collector")).add(((ScoreAccessor) params.get("_score")).doubleValue()); + return agg; + }); + SCRIPTS.put("combineScriptScore", params -> { + Map agg = (Map) params.get("_agg"); + return ((List) agg.get("collector")).stream().mapToDouble(Double::doubleValue).sum(); + }); } @SuppressWarnings("unchecked") @@ -144,6 +165,29 @@ public class ScriptedMetricAggregatorTests extends AggregatorTestCase { } } + /** + * test that uses the score of the documents + */ + public void testScriptedMetricWithCombineAccessesScores() throws IOException { + try (Directory directory = newDirectory()) { + Integer numDocs = randomInt(100); + try (RandomIndexWriter indexWriter = new RandomIndexWriter(random(), directory)) { + for (int i = 0; i < numDocs; i++) { + indexWriter.addDocument(singleton(new SortedNumericDocValuesField("number", i))); + } + } + try (IndexReader indexReader = DirectoryReader.open(directory)) { + ScriptedMetricAggregationBuilder aggregationBuilder = new ScriptedMetricAggregationBuilder(AGG_NAME); + aggregationBuilder.initScript(INIT_SCRIPT_SCORE).mapScript(MAP_SCRIPT_SCORE).combineScript(COMBINE_SCRIPT_SCORE); + ScriptedMetric scriptedMetric = search(newSearcher(indexReader, true, true), new MatchAllDocsQuery(), aggregationBuilder); + assertEquals(AGG_NAME, scriptedMetric.getName()); + assertNotNull(scriptedMetric.aggregation()); + // all documents have score of 1.0 + assertEquals((double) numDocs, scriptedMetric.aggregation()); + } + } + } + /** * We cannot use Mockito for mocking QueryShardContext in this case because * script-related methods (e.g. QueryShardContext#getLazyExecutableScript)