Speed up date_histogram by precomputing ranges

server/src/main/java/org/elasticsearch/common/Rounding.java

@@ -18,6 +18,7 @@
  */
 package org.elasticsearch.common;
 
+import org.apache.lucene.util.ArrayUtil;
 import org.elasticsearch.ElasticsearchException;
 import org.elasticsearch.Version;
 import org.elasticsearch.common.LocalTimeOffset.Gap;
@@ -44,8 +45,10 @@
 import java.time.temporal.TemporalQueries;
 import java.time.zone.ZoneOffsetTransition;
 import java.time.zone.ZoneRules;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Objects;
+import java.util.Set;
 import java.util.concurrent.TimeUnit;
 
 /**
@@ -401,8 +404,22 @@ private LocalDateTime truncateLocalDateTime(LocalDateTime localDateTime) {
             }
         }
 
+        /**
+         * Time zones with two midnights get "funny" non-continuous rounding
+         * that isn't compatible with the pre-computed array rounding.
+         */
+        private static final Set<String> HAS_TWO_MIDNIGHTS = Set.of("America/Moncton", "America/St_Johns", "Canada/Newfoundland");
+
         @Override
         public Prepared prepare(long minUtcMillis, long maxUtcMillis) {
+            Prepared orig = prepareOffsetRounding(minUtcMillis, maxUtcMillis);
+            if (unitRoundsToMidnight && HAS_TWO_MIDNIGHTS.contains(timeZone.getId())) {
+                return orig;
+            }
+            return maybeUseArray(orig, minUtcMillis, maxUtcMillis, 128);
+        }
+
+        private Prepared prepareOffsetRounding(long minUtcMillis, long maxUtcMillis) {
             long minLookup = minUtcMillis - unit.extraLocalOffsetLookup();
             long maxLookup = maxUtcMillis;
 
@@ -421,7 +438,6 @@ public Prepared prepare(long minUtcMillis, long maxUtcMillis) {
                 // Range too long, just use java.time
                 return prepareJavaTime();
             }
-
             LocalTimeOffset fixedOffset = lookup.fixedInRange(minLookup, maxLookup);
             if (fixedOffset != null) {
                 // The time zone is effectively fixed
@@ -1015,7 +1031,7 @@ public byte id() {
 
         @Override
         public Prepared prepare(long minUtcMillis, long maxUtcMillis) {
-            return wrapPreparedRounding(delegate.prepare(minUtcMillis, maxUtcMillis));
+            return wrapPreparedRounding(delegate.prepare(minUtcMillis - offset, maxUtcMillis - offset));
         }
 
         @Override
@@ -1085,4 +1101,57 @@ public static Rounding read(StreamInput in) throws IOException {
                 throw new ElasticsearchException("unknown rounding id [" + id + "]");
         }
     }
+
+    /**
+     * Attempt to build a {@link Prepared} implementation that relies on pre-calcuated
+     * "round down" points. If there would be more than {@code max} points then return
+     * the original implementation, otherwise return the new, faster implementation.
+     */
+    static Prepared maybeUseArray(Prepared orig, long minUtcMillis, long maxUtcMillis, int max) {
+        long[] values = new long[1];
+        long rounded = orig.round(minUtcMillis);
+        int i = 0;
+        values[i++] = rounded;
+        while ((rounded = orig.nextRoundingValue(rounded)) <= maxUtcMillis) {
+            if (i >= max) {
+                return orig;
+            }
+            assert values[i - 1] == orig.round(rounded - 1);
+            values = ArrayUtil.grow(values, i + 1);
+            values[i++]= rounded;
+        }
+        return new ArrayRounding(values, i, orig);
+    }
+
+    /**
+     * Implementation of {@link Prepared} using pre-calculated "round down" points.
+     */
+    private static class ArrayRounding implements Prepared {
+        private final long[] values;
+        private final int max;
+        private final Prepared delegate;
+
+        private ArrayRounding(long[] values, int max, Prepared delegate) {
+            this.values = values;
+            this.max = max;
+            this.delegate = delegate;
+        }
+
+        @Override
+        public long round(long utcMillis) {
+            assert values[0] <= utcMillis : "utcMillis must be after " + values[0];
+            int idx = Arrays.binarySearch(values, 0, max, utcMillis);
+            assert idx != -1 : "The insertion point is before the array! This should have tripped the assertion above.";
+            assert -1 - idx <= values.length : "This insertion point is after the end of the array.";
+            if (idx < 0) {
+                idx = -2 - idx;
+            }
+            return values[idx];
+        }
+
+        @Override
+        public long nextRoundingValue(long utcMillis) {
+            return delegate.nextRoundingValue(utcMillis);
+        }
+    }
 }

server/src/test/java/org/elasticsearch/common/RoundingTests.java

@@ -213,6 +213,9 @@ public void testOffsetRounding() {
         assertThat(rounding.nextRoundingValue(0), equalTo(oneDay - twoHours));
         assertThat(rounding.withoutOffset().round(0), equalTo(0L));
         assertThat(rounding.withoutOffset().nextRoundingValue(0), equalTo(oneDay));
+
+        rounding = Rounding.builder(Rounding.DateTimeUnit.DAY_OF_MONTH).timeZone(ZoneId.of("America/New_York")).offset(-twoHours).build();
+        assertThat(rounding.round(time("2020-11-01T09:00:00")), equalTo(time("2020-11-01T02:00:00")));
     }
 
     /**
@@ -231,7 +234,7 @@ public void testRandomTimeUnitRounding() {
             Rounding.DateTimeUnit unit = randomFrom(Rounding.DateTimeUnit.values());
             ZoneId tz = randomZone();
             Rounding rounding = new Rounding.TimeUnitRounding(unit, tz);
-            long[] bounds = randomDateBounds();
+            long[] bounds = randomDateBounds(unit);
             Rounding.Prepared prepared = rounding.prepare(bounds[0], bounds[1]);
 
             // Check that rounding is internally consistent and consistent with nextRoundingValue
@@ -894,8 +897,13 @@ private static long randomDate() {
         return Math.abs(randomLong() % (2 * (long) 10e11)); // 1970-01-01T00:00:00Z - 2033-05-18T05:33:20.000+02:00
     }
 
-    private static long[] randomDateBounds() {
+    private static long[] randomDateBounds(Rounding.DateTimeUnit unit) {
         long b1 = randomDate();
+        if (randomBoolean()) {
+            // Sometimes use a fairly close date
+            return new long[] {b1, b1 + unit.extraLocalOffsetLookup() * between(1, 40)};
+        }
+        // Otherwise use a totally random date
         long b2 = randomValueOtherThan(b1, RoundingTests::randomDate);
         if (b1 < b2) {
             return new long[] {b1, b2};