Add keep_values gap policy (#73297) (#73927)

imotov · not-napoleon · web-flow · commit d1e6e935445d · 2021-06-08T13:37:34.000-10:00
Adds a new keep_values gap policy that works like skip, except if the metric calculated on an empty bucket provides a non-null non-NaN value, this value is used for the bucket. Fixes #27377 Co-authored-by: Mark Tozzi <mark.tozzi@gmail.com>
diff --git a/docs/reference/aggregations/pipeline.asciidoc b/docs/reference/aggregations/pipeline.asciidoc
@@ -268,6 +268,10 @@ _insert_zeros_::
                 This option will replace missing values with a zero (`0`) and pipeline aggregation computation will
                 proceed as normal.
 
+_keep_values_::
+                This option is similar to skip, except if the metric provides a non-null, non-NaN value this value is
+                used, otherwise the empty bucket is skipped.
+
 include::pipeline/avg-bucket-aggregation.asciidoc[]
 
 include::pipeline/bucket-script-aggregation.asciidoc[]
diff --git a/server/src/internalClusterTest/java/org/elasticsearch/search/aggregations/pipeline/DerivativeIT.java b/server/src/internalClusterTest/java/org/elasticsearch/search/aggregations/pipeline/DerivativeIT.java
@@ -558,7 +558,7 @@ public void testSingleValueAggDerivativeWithGaps_random() throws Exception {
             checkBucketKeyAndDocCount("InternalBucket " + i, bucket, i, valueCounts_empty_rnd[i]);
             Sum sum = bucket.getAggregations().get("sum");
             double thisSumValue = sum.value();
-            if (bucket.getDocCount() == 0) {
+            if (bucket.getDocCount() == 0 && gapPolicy != GapPolicy.KEEP_VALUES) {
                 thisSumValue = gapPolicy == GapPolicy.INSERT_ZEROS ? 0 : Double.NaN;
             }
             SimpleValue sumDeriv = bucket.getAggregations().get("deriv");
diff --git a/server/src/internalClusterTest/java/org/elasticsearch/search/aggregations/pipeline/SerialDiffIT.java b/server/src/internalClusterTest/java/org/elasticsearch/search/aggregations/pipeline/SerialDiffIT.java
@@ -127,7 +127,7 @@ public void setupSuiteScopeCluster() throws Exception {
         numBuckets = randomIntBetween(10, 80);
         lag = randomIntBetween(1, numBuckets / 2);
 
-        gapPolicy = randomBoolean() ? BucketHelpers.GapPolicy.SKIP : BucketHelpers.GapPolicy.INSERT_ZEROS;
+        gapPolicy = randomFrom(BucketHelpers.GapPolicy.values());
         metric = randomMetric("the_metric", VALUE_FIELD);
         mockHisto = PipelineAggregationHelperTests.generateHistogram(interval, numBuckets, randomDouble(), randomDouble());
 
@@ -171,6 +171,12 @@ private void setupExpected(MetricTarget target) {
                     metricValue = 0.0;
                 } else {
                     metricValue = PipelineAggregationHelperTests.calculateMetric(docValues, metric);
+                    if (gapPolicy.equals(BucketHelpers.GapPolicy.KEEP_VALUES)) {
+                        if (Double.isInfinite(metricValue) || Double.isNaN(metricValue)) {
+                            // serial diff ignores these values and replaces them with null
+                            metricValue = Double.NaN;
+                        }
+                    }
                 }
 
             } else {
@@ -204,13 +210,8 @@ private void setupExpected(MetricTarget target) {
             }
 
             lagWindow.add(metricValue);
-
-
-
-
         }
 
-
         testValues.put(target.toString(), values);
     }
 
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/pipeline/BucketHelpers.java b/server/src/main/java/org/elasticsearch/search/aggregations/pipeline/BucketHelpers.java
@@ -41,9 +41,40 @@ public class BucketHelpers {
      *
      * "insert_zeros": empty buckets will be filled with zeros for all metrics
      * "skip": empty buckets will simply be ignored
+     * "keep_values": for empty buckets the values provided by the metrics will still be used if they are available
      */
     public enum GapPolicy implements Writeable {
-        INSERT_ZEROS((byte) 0, "insert_zeros"), SKIP((byte) 1, "skip");
+        INSERT_ZEROS((byte) 0, "insert_zeros", false) {
+            @Override
+            public Double processValue(long docCount, Double value) {
+                if (Double.isInfinite(value) || Double.isNaN(value) || docCount == 0) {
+                    return 0.0;
+                } else {
+                    return value;
+                }
+            }
+        },
+
+        SKIP((byte) 1, "skip", true) {
+            @Override
+            public Double processValue(long docCount, Double value) {
+                if (Double.isInfinite(value) || docCount == 0) {
+                    return Double.NaN;
+                } else {
+                    return value;
+                }
+            }
+        },
+
+        KEEP_VALUES((byte) 2, "keep_values", true) {
+            public Double processValue(long docCount, Double value) {
+                if (Double.isInfinite(value) || Double.isNaN(value)) {
+                    return Double.NaN;
+                } else {
+                    return value;
+                }
+            }
+        };
 
         /**
          * Parse a string GapPolicy into the byte enum
@@ -76,10 +107,12 @@ public static GapPolicy parse(String text, XContentLocation tokenLocation) {
 
         private final byte id;
         private final ParseField parseField;
+        public final boolean isSkippable;
 
-        GapPolicy(byte id, String name) {
+        GapPolicy(byte id, String name, boolean isSkippable) {
             this.id = id;
             this.parseField = new ParseField(name);
+            this.isSkippable = isSkippable;
         }
 
         /**
@@ -113,6 +146,8 @@ public static GapPolicy readFrom(StreamInput in) throws IOException {
         public String getName() {
             return parseField.getPreferredName();
         }
+
+        public abstract Double processValue(long docCount, Double value);
     }
 
     /**
@@ -161,18 +196,12 @@ public static Double resolveBucketValue(MultiBucketsAggregation agg,
                     throw formatResolutionError(agg, aggPathAsList, propertyValue);
                 }
                 // doc count never has missing values so gap policy doesn't apply here
-                boolean isDocCountProperty = aggPathAsList.size() == 1 && "_count".equals(aggPathAsList.get(0));
-                if (Double.isInfinite(value) || Double.isNaN(value) || (bucket.getDocCount() == 0 && isDocCountProperty == false)) {
-                    switch (gapPolicy) {
-                    case INSERT_ZEROS:
-                        return 0.0;
-                    case SKIP:
-                    default:
-                        return Double.NaN;
-                    }
-                } else {
+                if (aggPathAsList.size() == 1 && "_count".equals(aggPathAsList.get(0))) {
                     return value;
+                } else {
+                    return gapPolicy.processValue(bucket.getDocCount(), value);
                 }
+
             }
         } catch (InvalidAggregationPathException e) {
             return null;
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/pipeline/BucketScriptPipelineAggregator.java b/server/src/main/java/org/elasticsearch/search/aggregations/pipeline/BucketScriptPipelineAggregator.java
@@ -88,7 +88,7 @@ public InternalAggregation reduce(InternalAggregation aggregation, ReduceContext
                 String varName = entry.getKey();
                 String bucketsPath = entry.getValue();
                 Double value = resolveBucketValue(originalAgg, bucket, bucketsPath, gapPolicy);
-                if (GapPolicy.SKIP == gapPolicy && (value == null || Double.isNaN(value))) {
+                if (gapPolicy.isSkippable && (value == null || Double.isNaN(value))) {
                     skipBucket = true;
                     break;
                 }
diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/pipeline/BucketSortPipelineAggregator.java b/server/src/main/java/org/elasticsearch/search/aggregations/pipeline/BucketSortPipelineAggregator.java
@@ -122,7 +122,7 @@ private Map<FieldSortBuilder, Comparable<Object>> resolveAndCacheSortValues() {
                     resolved.put(sort, (Comparable<Object>) internalBucket.getKey());
                 } else {
                     Double bucketValue = BucketHelpers.resolveBucketValue(parentAgg, internalBucket, sortField, gapPolicy);
-                    if (GapPolicy.SKIP == gapPolicy && Double.isNaN(bucketValue)) {
+                    if (gapPolicy.isSkippable && Double.isNaN(bucketValue)) {
                         continue;
                     }
                     resolved.put(sort, (Comparable<Object>) (Object) bucketValue);
diff --git a/server/src/test/java/org/elasticsearch/search/aggregations/pipeline/DerivativeAggregatorTests.java b/server/src/test/java/org/elasticsearch/search/aggregations/pipeline/DerivativeAggregatorTests.java
@@ -586,7 +586,15 @@ public void testSingleValueAggDerivativeWithGaps_random() throws IOException {
                     Sum sum = bucket.getAggregations().get("sum");
                     double thisSumValue = sum.value();
                     if (bucket.getDocCount() == 0) {
-                        thisSumValue = gapPolicy == GapPolicy.INSERT_ZEROS ? 0 : Double.NaN;
+                        switch (gapPolicy) {
+                            case INSERT_ZEROS:
+                                thisSumValue = 0;
+                                break;
+                            case KEEP_VALUES:
+                                break;
+                            default:
+                                thisSumValue = Double.NaN;
+                        }
                     }
                     SimpleValue sumDeriv = bucket.getAggregations().get("deriv");
                     if (i == 0) {
diff --git a/server/src/test/java/org/elasticsearch/search/aggregations/pipeline/GapPolicyTests.java b/server/src/test/java/org/elasticsearch/search/aggregations/pipeline/GapPolicyTests.java
@@ -25,26 +25,30 @@ public GapPolicyTests() {
     public void testValidOrdinals() {
         assertThat(BucketHelpers.GapPolicy.INSERT_ZEROS.ordinal(), equalTo(0));
         assertThat(BucketHelpers.GapPolicy.SKIP.ordinal(), equalTo(1));
+        assertThat(BucketHelpers.GapPolicy.KEEP_VALUES.ordinal(), equalTo(2));
     }
 
     @Override
     public void testFromString() {
         assertThat(BucketHelpers.GapPolicy.parse("insert_zeros", null), equalTo(BucketHelpers.GapPolicy.INSERT_ZEROS));
         assertThat(BucketHelpers.GapPolicy.parse("skip", null), equalTo(BucketHelpers.GapPolicy.SKIP));
+        assertThat(BucketHelpers.GapPolicy.parse("keep_values", null), equalTo(BucketHelpers.GapPolicy.KEEP_VALUES));
         ParsingException e = expectThrows(ParsingException.class, () -> BucketHelpers.GapPolicy.parse("does_not_exist", null));
         assertThat(e.getMessage(),
-            equalTo("Invalid gap policy: [does_not_exist], accepted values: [insert_zeros, skip]"));
+            equalTo("Invalid gap policy: [does_not_exist], accepted values: [insert_zeros, skip, keep_values]"));
     }
 
     @Override
     public void testReadFrom() throws IOException {
         assertReadFromStream(0, BucketHelpers.GapPolicy.INSERT_ZEROS);
         assertReadFromStream(1, BucketHelpers.GapPolicy.SKIP);
+        assertReadFromStream(2, BucketHelpers.GapPolicy.KEEP_VALUES);
     }
 
     @Override
     public void testWriteTo() throws IOException {
         assertWriteToStream(BucketHelpers.GapPolicy.INSERT_ZEROS, 0);
         assertWriteToStream(BucketHelpers.GapPolicy.SKIP, 1);
+        assertWriteToStream(BucketHelpers.GapPolicy.KEEP_VALUES, 2);
     }
 }
diff --git a/server/src/test/java/org/elasticsearch/search/aggregations/pipeline/PipelineAggregationHelperTests.java b/server/src/test/java/org/elasticsearch/search/aggregations/pipeline/PipelineAggregationHelperTests.java
@@ -129,7 +129,7 @@ public static double calculateMetric(double[] values, ValuesSourceAggregationBui
             for (double value : values) {
                 accumulator += value;
             }
-            return accumulator / values.length;
+            return values.length == 0 ? Double.NaN : accumulator / values.length ;
         }
 
         return 0.0;

Original file line number	Diff line number	Diff line change
`@@ -558,7 +558,7 @@ public void testSingleValueAggDerivativeWithGaps_random() throws Exception {`
`558`	`558`	`checkBucketKeyAndDocCount("InternalBucket " + i, bucket, i, valueCounts_empty_rnd[i]);`
`559`	`559`	`Sum sum = bucket.getAggregations().get("sum");`
`560`	`560`	`double thisSumValue = sum.value();`
`561`		`- if (bucket.getDocCount() == 0) {`
	`561`	`+ if (bucket.getDocCount() == 0 && gapPolicy != GapPolicy.KEEP_VALUES) {`
`562`	`562`	`thisSumValue = gapPolicy == GapPolicy.INSERT_ZEROS ? 0 : Double.NaN;`
`563`	`563`	`}`
`564`	`564`	`SimpleValue sumDeriv = bucket.getAggregations().get("deriv");`
Original file line number	Diff line number	Diff line change
`@@ -88,7 +88,7 @@ public InternalAggregation reduce(InternalAggregation aggregation, ReduceContext`
`88`	`88`	`String varName = entry.getKey();`
`89`	`89`	`String bucketsPath = entry.getValue();`
`90`	`90`	`Double value = resolveBucketValue(originalAgg, bucket, bucketsPath, gapPolicy);`
`91`		`- if (GapPolicy.SKIP == gapPolicy && (value == null \|\| Double.isNaN(value))) {`
	`91`	`+ if (gapPolicy.isSkippable && (value == null \|\| Double.isNaN(value))) {`
`92`	`92`	`skipBucket = true;`
`93`	`93`	`break;`
`94`	`94`	`}`
Original file line number	Diff line number	Diff line change
`@@ -122,7 +122,7 @@ private Map<FieldSortBuilder, Comparable<Object>> resolveAndCacheSortValues() {`
`122`	`122`	`resolved.put(sort, (Comparable<Object>) internalBucket.getKey());`
`123`	`123`	`} else {`
`124`	`124`	`Double bucketValue = BucketHelpers.resolveBucketValue(parentAgg, internalBucket, sortField, gapPolicy);`
`125`		`- if (GapPolicy.SKIP == gapPolicy && Double.isNaN(bucketValue)) {`
	`125`	`+ if (gapPolicy.isSkippable && Double.isNaN(bucketValue)) {`
`126`	`126`	`continue;`
`127`	`127`	`}`
`128`	`128`	`resolved.put(sort, (Comparable<Object>) (Object) bucketValue);`
Original file line number	Diff line number	Diff line change
`@@ -25,26 +25,30 @@ public GapPolicyTests() {`
`25`	`25`	`public void testValidOrdinals() {`
`26`	`26`	`assertThat(BucketHelpers.GapPolicy.INSERT_ZEROS.ordinal(), equalTo(0));`
`27`	`27`	`assertThat(BucketHelpers.GapPolicy.SKIP.ordinal(), equalTo(1));`
	`28`	`+ assertThat(BucketHelpers.GapPolicy.KEEP_VALUES.ordinal(), equalTo(2));`
`28`	`29`	`}`
`29`	`30`
`30`	`31`	`@Override`
`31`	`32`	`public void testFromString() {`
`32`	`33`	`assertThat(BucketHelpers.GapPolicy.parse("insert_zeros", null), equalTo(BucketHelpers.GapPolicy.INSERT_ZEROS));`
`33`	`34`	`assertThat(BucketHelpers.GapPolicy.parse("skip", null), equalTo(BucketHelpers.GapPolicy.SKIP));`
	`35`	`+ assertThat(BucketHelpers.GapPolicy.parse("keep_values", null), equalTo(BucketHelpers.GapPolicy.KEEP_VALUES));`
`34`	`36`	`ParsingException e = expectThrows(ParsingException.class, () -> BucketHelpers.GapPolicy.parse("does_not_exist", null));`
`35`	`37`	`assertThat(e.getMessage(),`
`36`		`- equalTo("Invalid gap policy: [does_not_exist], accepted values: [insert_zeros, skip]"));`
	`38`	`+ equalTo("Invalid gap policy: [does_not_exist], accepted values: [insert_zeros, skip, keep_values]"));`
`37`	`39`	`}`
`38`	`40`
`39`	`41`	`@Override`
`40`	`42`	`public void testReadFrom() throws IOException {`
`41`	`43`	`assertReadFromStream(0, BucketHelpers.GapPolicy.INSERT_ZEROS);`
`42`	`44`	`assertReadFromStream(1, BucketHelpers.GapPolicy.SKIP);`
	`45`	`+ assertReadFromStream(2, BucketHelpers.GapPolicy.KEEP_VALUES);`
`43`	`46`	`}`
`44`	`47`
`45`	`48`	`@Override`
`46`	`49`	`public void testWriteTo() throws IOException {`
`47`	`50`	`assertWriteToStream(BucketHelpers.GapPolicy.INSERT_ZEROS, 0);`
`48`	`51`	`assertWriteToStream(BucketHelpers.GapPolicy.SKIP, 1);`
	`52`	`+ assertWriteToStream(BucketHelpers.GapPolicy.KEEP_VALUES, 2);`
`49`	`53`	`}`
`50`	`54`	`}`
Original file line number	Diff line number	Diff line change
`@@ -129,7 +129,7 @@ public static double calculateMetric(double[] values, ValuesSourceAggregationBui`
`129`	`129`	`for (double value : values) {`
`130`	`130`	`accumulator += value;`
`131`	`131`	`}`
`132`		`- return accumulator / values.length;`
	`132`	`+ return values.length == 0 ? Double.NaN : accumulator / values.length ;`
`133`	`133`	`}`
`134`	`134`
`135`	`135`	`return 0.0;`