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.Locale;
+import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.TimeUnit;
 
@@ -467,8 +470,33 @@ private LocalDateTime truncateLocalDateTime(LocalDateTime localDateTime) {
             }
         }
 
+        /**
+         * Time zones have a daylight savings time transition after midnight that
+         * transitions back before midnight break the array-based rounding so
+         * we don't use it for them during those transitions. Luckily they are
+         * fairly rare and a while in the past. Note: we can use the array based
+         * rounding if the transition is <strong>at</strong> midnight. Just not
+         * if it is <strong>after</strong> it.
+         */
+        private static final Map<String, Long> FORWARDS_BACKWRADS_ZONES = Map.ofEntries(
+            Map.entry("America/Goose_Bay", 1289116800000L),   // Stopped transitioning after midnight in 2010
+            Map.entry("America/Moncton", 1162108800000L),     // Stopped transitioning after midnight in 2006
+            Map.entry("America/St_Johns", 1289118600000L),    // Stopped transitioning after midnight in 2010
+            Map.entry("Canada/Newfoundland", 1289118600000L), // Stopped transitioning after midnight in 2010
+            Map.entry("Pacific/Guam", -29347200000L),         // Stopped transitioning after midnight in 1969
+            Map.entry("Pacific/Saipan", -29347200000L)        // Stopped transitioning after midnight in 1969
+        );
+
         @Override
         public Prepared prepare(long minUtcMillis, long maxUtcMillis) {
+            Prepared orig = prepareOffsetRounding(minUtcMillis, maxUtcMillis);
+            if (unitRoundsToMidnight && minUtcMillis <= FORWARDS_BACKWRADS_ZONES.getOrDefault(timeZone.getId(), Long.MIN_VALUE)) {
+                return orig;
+            }
+            return maybeUseArray(orig, minUtcMillis, maxUtcMillis, 128);
+        }
+
+        private Prepared prepareOffsetRounding(long minUtcMillis, long maxUtcMillis) {
             long minLookup = minUtcMillis - unit.extraLocalOffsetLookup();
             long maxLookup = maxUtcMillis;
 
@@ -487,7 +515,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
@@ -1111,7 +1138,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
@@ -1186,4 +1213,62 @@ 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);
+        }
+
+        @Override
+        public double roundingSize(long utcMillis, DateTimeUnit timeUnit) {
+            return delegate.roundingSize(utcMillis, timeUnit);
+        }
+    }
 }

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

@@ -36,6 +36,8 @@
 import java.time.format.DateTimeFormatter;
 import java.time.temporal.TemporalAccessor;
 import java.time.zone.ZoneOffsetTransition;
+import java.time.zone.ZoneOffsetTransitionRule;
+import java.time.zone.ZoneRules;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
@@ -214,6 +216,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,51 +236,55 @@ public void testRandomTimeUnitRounding() {
         for (int i = 0; i < 1000; ++i) {
             Rounding.DateTimeUnit unit = randomFrom(Rounding.DateTimeUnit.values());
             ZoneId tz = randomZone();
-            Rounding rounding = new Rounding.TimeUnitRounding(unit, tz);
-            long[] bounds = randomDateBounds();
-            Rounding.Prepared prepared = rounding.prepare(bounds[0], bounds[1]);
-
-            // Check that rounding is internally consistent and consistent with nextRoundingValue
-            long date = dateBetween(bounds[0], bounds[1]);
-            long unitMillis = unit.getField().getBaseUnit().getDuration().toMillis();
-            // FIXME this was copy pasted from the other impl and not used. breaks the nasty date actually gets assigned
-            if (randomBoolean()) {
-                nastyDate(date, tz, unitMillis);
-            }
-            final long roundedDate = prepared.round(date);
-            final long nextRoundingValue = prepared.nextRoundingValue(roundedDate);
-
-            assertInterval(roundedDate, date, nextRoundingValue, rounding, tz);
-
-            // check correct unit interval width for units smaller than a day, they should be fixed size except for transitions
-            if (unitMillis <= 86400 * 1000) {
-                // if the interval defined didn't cross timezone offset transition, it should cover unitMillis width
-                int offsetRounded = tz.getRules().getOffset(Instant.ofEpochMilli(roundedDate - 1)).getTotalSeconds();
-                int offsetNextValue = tz.getRules().getOffset(Instant.ofEpochMilli(nextRoundingValue + 1)).getTotalSeconds();
-                if (offsetRounded == offsetNextValue) {
-                    assertThat("unit interval width not as expected for [" + unit + "], [" + tz + "] at "
-                        + Instant.ofEpochMilli(roundedDate), nextRoundingValue - roundedDate, equalTo(unitMillis));
-                }
+            long[] bounds = randomDateBounds(unit);
+            assertUnitRoundingSameAsJavaUtilTimeImplementation(unit, tz, bounds[0], bounds[1]);
+        }
+    }
+
+    private void assertUnitRoundingSameAsJavaUtilTimeImplementation(Rounding.DateTimeUnit unit, ZoneId tz, long start, long end) {
+        Rounding rounding = new Rounding.TimeUnitRounding(unit, tz);
+        Rounding.Prepared prepared = rounding.prepare(start, end);
+
+        // Check that rounding is internally consistent and consistent with nextRoundingValue
+        long date = dateBetween(start, end);
+        long unitMillis = unit.getField().getBaseUnit().getDuration().toMillis();
+        // FIXME this was copy pasted from the other impl and not used. breaks the nasty date actually gets assigned
+        if (randomBoolean()) {
+            nastyDate(date, tz, unitMillis);
+        }
+        final long roundedDate = prepared.round(date);
+        final long nextRoundingValue = prepared.nextRoundingValue(roundedDate);
+
+        assertInterval(roundedDate, date, nextRoundingValue, rounding, tz);
+
+        // check correct unit interval width for units smaller than a day, they should be fixed size except for transitions
+        if (unitMillis <= 86400 * 1000) {
+            // if the interval defined didn't cross timezone offset transition, it should cover unitMillis width
+            int offsetRounded = tz.getRules().getOffset(Instant.ofEpochMilli(roundedDate - 1)).getTotalSeconds();
+            int offsetNextValue = tz.getRules().getOffset(Instant.ofEpochMilli(nextRoundingValue + 1)).getTotalSeconds();
+            if (offsetRounded == offsetNextValue) {
+                assertThat("unit interval width not as expected for [" + unit + "], [" + tz + "] at "
+                    + Instant.ofEpochMilli(roundedDate), nextRoundingValue - roundedDate, equalTo(unitMillis));
             }
+        }
 
-            // Round a whole bunch of dates and make sure they line up with the known good java time implementation
-            Rounding.Prepared javaTimeRounding = rounding.prepareJavaTime();
-            for (int d = 0; d < 1000; d++) {
-                date = dateBetween(bounds[0], bounds[1]);
-                long javaRounded = javaTimeRounding.round(date);
-                long esRounded = prepared.round(date);
-                if (javaRounded != esRounded) {
-                    fail("Expected [" + rounding + "] to round [" + Instant.ofEpochMilli(date) + "] to ["
-                            + Instant.ofEpochMilli(javaRounded) + "] but instead rounded to [" + Instant.ofEpochMilli(esRounded) + "]");
-                }
-                long javaNextRoundingValue = javaTimeRounding.nextRoundingValue(date);
-                long esNextRoundingValue = prepared.nextRoundingValue(date);
-                if (javaNextRoundingValue != esNextRoundingValue) {
-                    fail("Expected [" + rounding + "] to round [" + Instant.ofEpochMilli(date) + "] to ["
-                            + Instant.ofEpochMilli(esRounded) + "] and nextRoundingValue to be ["
-                            + Instant.ofEpochMilli(javaNextRoundingValue) + "] but instead was to ["
-                            + Instant.ofEpochMilli(esNextRoundingValue) + "]");
-                }
+        // Round a whole bunch of dates and make sure they line up with the known good java time implementation
+        Rounding.Prepared javaTimeRounding = rounding.prepareJavaTime();
+        for (int d = 0; d < 1000; d++) {
+            date = dateBetween(start, end);
+            long javaRounded = javaTimeRounding.round(date);
+            long esRounded = prepared.round(date);
+            if (javaRounded != esRounded) {
+                fail("Expected [" + rounding + "] to round [" + Instant.ofEpochMilli(date) + "] to ["
+                        + Instant.ofEpochMilli(javaRounded) + "] but instead rounded to [" + Instant.ofEpochMilli(esRounded) + "]");
+            }
+            long javaNextRoundingValue = javaTimeRounding.nextRoundingValue(date);
+            long esNextRoundingValue = prepared.nextRoundingValue(date);
+            if (javaNextRoundingValue != esNextRoundingValue) {
+                fail("Expected [" + rounding + "] to round [" + Instant.ofEpochMilli(date) + "] to ["
+                        + Instant.ofEpochMilli(esRounded) + "] and nextRoundingValue to be ["
+                        + Instant.ofEpochMilli(javaNextRoundingValue) + "] but instead was to ["
+                        + Instant.ofEpochMilli(esNextRoundingValue) + "]");
             }
         }
     }
@@ -715,6 +724,70 @@ public void testDST_America_St_Johns() {
         }
     }
 
+    /**
+     * Tests for DST transitions that cause the rounding to jump "backwards" because they round
+     * from one back to the previous day. Usually these rounding start before 
+     */
+    public void testForwardsBackwardsTimeZones() {
+        for (String zoneId : JAVA_ZONE_IDS) {
+            ZoneId tz = ZoneId.of(zoneId);
+            ZoneRules rules = tz.getRules();
+            for (ZoneOffsetTransition transition : rules.getTransitions()) {
+                checkForForwardsBackwardsTransition(tz, transition);
+            }
+            int firstYear;
+            if (rules.getTransitions().isEmpty()) {
+                // Pick an arbitrary year to start the range
+                firstYear = 1999;
+            } else {
+                ZoneOffsetTransition lastTransition = rules.getTransitions().get(rules.getTransitions().size() - 1);
+                firstYear = lastTransition.getDateTimeAfter().getYear() + 1;
+            }
+            // Pick an arbitrary year to end the range too
+            int lastYear = 2050;
+            int year = randomFrom(firstYear, lastYear);
+            for (ZoneOffsetTransitionRule transitionRule : rules.getTransitionRules()) {
+                ZoneOffsetTransition transition = transitionRule.createTransition(year);
+                checkForForwardsBackwardsTransition(tz, transition);
+            }
+        }
+    }
+
+    private void checkForForwardsBackwardsTransition(ZoneId tz, ZoneOffsetTransition transition) {
+        if (transition.getDateTimeBefore().getYear() < 1950) {
+            // We don't support transitions far in the past at all
+            return;
+        }
+        if (false == transition.isOverlap()) {
+            // Only overlaps cause the array rounding to have trouble
+            return;
+        }
+        if (transition.getDateTimeBefore().getDayOfMonth() == transition.getDateTimeAfter().getDayOfMonth()) {
+            // Only when the rounding changes the day
+            return;
+        }
+        if (transition.getDateTimeBefore().getMinute() == 0) {
+            // But roundings that change *at* midnight are safe because they don't "jump" to the next day.
+            return;
+        }
+        logger.info(
+            "{} from {}{} to {}{}",
+            tz,
+            transition.getDateTimeBefore(),
+            transition.getOffsetBefore(),
+            transition.getDateTimeAfter(),
+            transition.getOffsetAfter()
+        );
+        long millisSinceEpoch = TimeUnit.SECONDS.toMillis(transition.toEpochSecond());
+        long twoHours = TimeUnit.HOURS.toMillis(2);
+        assertUnitRoundingSameAsJavaUtilTimeImplementation(
+            Rounding.DateTimeUnit.DAY_OF_MONTH,
+            tz,
+            millisSinceEpoch - twoHours,
+            millisSinceEpoch + twoHours
+        );
+    }
+
     /**
      * tests for dst transition with overlaps and day roundings.
      */
@@ -965,8 +1038,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};