Speed up date_histogram by precomputing ranges

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

@@ -22,6 +22,7 @@
 import java.time.Instant;
 import java.time.LocalDate;
 import java.time.ZoneId;
+import java.time.temporal.ChronoField;
 import java.time.zone.ZoneOffsetTransition;
 import java.time.zone.ZoneOffsetTransitionRule;
 import java.time.zone.ZoneRules;
@@ -174,6 +175,12 @@ public interface Strategy {
      */
     protected abstract LocalTimeOffset offsetContaining(long utcMillis);
 
+    /**
+     * Does this transition or any previous transitions move back to the
+     * previous day? See {@link Lookup#anyMoveBackToPreviousDay()} for rules.
+     */
+    protected abstract boolean anyMoveBackToPreviousDay();
+
     @Override
     public String toString() {
         return toString(millis);
@@ -195,6 +202,15 @@ public abstract static class Lookup {
          */
         public abstract LocalTimeOffset fixedInRange(long minUtcMillis, long maxUtcMillis);
 
+        /**
+         * Do any of the transitions move back to the previous day?
+         * <p>
+         * Note: If an overlap occurs at, say, 1 am and jumps back to
+         * <strong>exactly</strong> midnight then it doesn't count because
+         * midnight is still counted as being in the "next" day.
+         */
+        public abstract boolean anyMoveBackToPreviousDay();
+
         /**
          * The number of offsets in the lookup. Package private for testing.
          */
@@ -225,6 +241,11 @@ protected LocalTimeOffset offsetContaining(long utcMillis) {
             return this;
         }
 
+        @Override
+        protected boolean anyMoveBackToPreviousDay() {
+            return false;
+        }
+
         @Override
         protected String toString(long millis) {
             return Long.toString(millis);
@@ -298,6 +319,11 @@ public long firstMissingLocalTime() {
             return firstMissingLocalTime;
         }
 
+        @Override
+        protected boolean anyMoveBackToPreviousDay() {
+            return previous().anyMoveBackToPreviousDay();
+        }
+
         @Override
         protected String toString(long millis) {
             return "Gap of " + millis + "@" + Instant.ofEpochMilli(startUtcMillis());
@@ -307,13 +333,21 @@ protected String toString(long millis) {
     public static class Overlap extends Transition {
         private final long firstOverlappingLocalTime;
         private final long firstNonOverlappingLocalTime;
-
-        private Overlap(long millis, LocalTimeOffset previous, long startUtcMillis,
-                long firstOverlappingLocalTime, long firstNonOverlappingLocalTime) {
+        private final boolean movesBackToPreviousDay;
+
+        private Overlap(
+            long millis,
+            LocalTimeOffset previous,
+            long startUtcMillis,
+            long firstOverlappingLocalTime,
+            long firstNonOverlappingLocalTime,
+            boolean movesBackToPreviousDay
+        ) {
             super(millis, previous, startUtcMillis);
             this.firstOverlappingLocalTime = firstOverlappingLocalTime;
             this.firstNonOverlappingLocalTime = firstNonOverlappingLocalTime;
             assert firstOverlappingLocalTime < firstNonOverlappingLocalTime;
+            this.movesBackToPreviousDay = movesBackToPreviousDay;
         }
 
         @Override
@@ -341,6 +375,11 @@ public long firstOverlappingLocalTime() {
             return firstOverlappingLocalTime;
         }
 
+        @Override
+        protected boolean anyMoveBackToPreviousDay() {
+            return movesBackToPreviousDay || previous().anyMoveBackToPreviousDay();
+        }
+
         @Override
         protected String toString(long millis) {
             return "Overlap of " + millis + "@" + Instant.ofEpochMilli(startUtcMillis());
@@ -375,6 +414,11 @@ int size() {
         public String toString() {
             return String.format(Locale.ROOT, "FixedLookup[for %s at %s]", zone, fixed);
         }
+
+        @Override
+        public boolean anyMoveBackToPreviousDay() {
+            return false;
+        }
     }
 
     /**
@@ -406,6 +450,11 @@ public LocalTimeOffset innerLookup(long utcMillis) {
         int size() {
             return size;
         }
+
+        @Override
+        public boolean anyMoveBackToPreviousDay() {
+            return lastOffset.anyMoveBackToPreviousDay();
+        }
     }
 
     /**
@@ -453,6 +502,11 @@ int size() {
             return offsets.length;
         }
 
+        @Override
+        public boolean anyMoveBackToPreviousDay() {
+            return offsets[offsets.length - 1].anyMoveBackToPreviousDay();
+        }
+
         @Override
         public String toString() {
             return String.format(Locale.ROOT, "TransitionArrayLookup[for %s between %s and %s]",
@@ -505,7 +559,25 @@ protected static Transition buildTransition(ZoneOffsetTransition transition, Loc
             }
             long firstOverlappingLocalTime = utcStart + offsetAfterMillis;
             long firstNonOverlappingLocalTime = utcStart + offsetBeforeMillis;
-            return new Overlap(offsetAfterMillis, previous, utcStart, firstOverlappingLocalTime, firstNonOverlappingLocalTime);
+            return new Overlap(
+                offsetAfterMillis,
+                previous,
+                utcStart,
+                firstOverlappingLocalTime,
+                firstNonOverlappingLocalTime,
+                movesBackToPreviousDay(transition)
+            );
+        }
+
+        private static boolean movesBackToPreviousDay(ZoneOffsetTransition transition) {
+            if (transition.getDateTimeBefore().getDayOfMonth() == transition.getDateTimeAfter().getDayOfMonth()) {
+                return false;
+            }
+            if (transition.getDateTimeBefore().getLong(ChronoField.NANO_OF_DAY) == 0L) {
+                // If we change *at* midnight this is ok.
+                return false;
+            }
+            return true;
         }
     }
 

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,6 +45,7 @@
 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.Objects;
@@ -405,6 +407,34 @@ public Rounding build() {
         }
     }
 
+    private abstract class PreparedRounding implements Prepared {
+        /**
+         * 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.
+         */
+        protected Prepared maybeUseArray(long minUtcMillis, long maxUtcMillis, int max) {
+            long[] values = new long[1];
+            long rounded = round(minUtcMillis);
+            int i = 0;
+            values[i++] = rounded;
+            while ((rounded = nextRoundingValue(rounded)) <= maxUtcMillis) {
+                if (i >= max) {
+                    return this;
+                }
+                /*
+                 * We expect a time in the last transition (rounded - 1) to round
+                 * to the last value we calculated. If it doesn't then we're
+                 * probably doing something wrong here....
+                 */
+                assert values[i - 1] == round(rounded - 1);
+                values = ArrayUtil.grow(values, i + 1);
+                values[i++]= rounded;
+            }
+            return new ArrayRounding(values, i, this);
+        }
+    }
+
     static class TimeUnitRounding extends Rounding {
         static final byte ID = 1;
 
@@ -469,6 +499,15 @@ private LocalDateTime truncateLocalDateTime(LocalDateTime localDateTime) {
 
         @Override
         public Prepared prepare(long minUtcMillis, long maxUtcMillis) {
+            /*
+             * 128 is a power of two that isn't huge. We might be able to do
+             * better if the limit was based on the actual type of prepared
+             * rounding but this'll do for now.
+             */
+            return prepareOffsetOrJavaTimeRounding(minUtcMillis, maxUtcMillis).maybeUseArray(minUtcMillis, maxUtcMillis, 128);
+        }
+
+        private TimeUnitPreparedRounding prepareOffsetOrJavaTimeRounding(long minUtcMillis, long maxUtcMillis) {
             long minLookup = minUtcMillis - unit.extraLocalOffsetLookup();
             long maxLookup = maxUtcMillis;
 
@@ -487,7 +526,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
@@ -516,7 +554,7 @@ public Prepared prepareForUnknown() {
         }
 
         @Override
-        Prepared prepareJavaTime() {
+        TimeUnitPreparedRounding prepareJavaTime() {
             if (unitRoundsToMidnight) {
                 return new JavaTimeToMidnightRounding();
             }
@@ -555,7 +593,7 @@ public String toString() {
             return "Rounding[" + unit + " in " + timeZone + "]";
         }
 
-        private abstract class TimeUnitPreparedRounding implements Prepared {
+        private abstract class TimeUnitPreparedRounding extends PreparedRounding {
             @Override
             public double roundingSize(long utcMillis, DateTimeUnit timeUnit) {
                 if (timeUnit.isMillisBased == unit.isMillisBased) {
@@ -649,6 +687,14 @@ public long inOverlap(long localMillis, Overlap overlap) {
             public long beforeOverlap(long localMillis, Overlap overlap) {
                 return overlap.previous().localToUtc(localMillis, this);
             }
+
+            @Override
+            protected Prepared maybeUseArray(long minUtcMillis, long maxUtcMillis, int max) {
+                if (lookup.anyMoveBackToPreviousDay()) {
+                    return this;
+                }
+                return super.maybeUseArray(minUtcMillis, maxUtcMillis, max);
+            }
         }
 
         private class NotToMidnightRounding extends AbstractNotToMidnightRounding implements LocalTimeOffset.Strategy {
@@ -708,6 +754,12 @@ public long nextRoundingValue(long utcMillis) {
                 return firstTimeOnDay(localMidnight);
             }
 
+            @Override
+            protected Prepared maybeUseArray(long minUtcMillis, long maxUtcMillis, int max) {
+                // We don't have the right information needed to know if this is safe for this time zone so we always use java rounding
+                return this;
+            }
+
             private long firstTimeOnDay(LocalDateTime localMidnight) {
                 assert localMidnight.toLocalTime().equals(LocalTime.of(0, 0, 0)) : "firstTimeOnDay should only be called at midnight";
 
@@ -1111,7 +1163,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 +1238,41 @@ public static Rounding read(StreamInput in) throws IOException {
                 throw new ElasticsearchException("unknown rounding id [" + id + "]");
         }
     }
+
+    /**
+     * 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);
+        }
+    }
 }