Date Histogram Facet: Improve time zone handling, add factor option, c…

…loses #1580.

src/main/java/org/elasticsearch/common/joda/TimeZoneRounding.java

@@ -0,0 +1,240 @@
+/*
+ * Licensed to ElasticSearch and Shay Banon under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. ElasticSearch licenses this
+ * file to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.common.joda;
+
+import org.elasticsearch.common.unit.TimeValue;
+import org.joda.time.DateTimeConstants;
+import org.joda.time.DateTimeField;
+import org.joda.time.DateTimeZone;
+
+/**
+ */
+public abstract class TimeZoneRounding {
+
+    public abstract long calc(long utcMillis);
+
+    public static Builder builder(DateTimeField field) {
+        return new Builder(field);
+    }
+
+    public static Builder builder(TimeValue interval) {
+        return new Builder(interval);
+    }
+
+    public static class Builder {
+
+        private DateTimeField field;
+        private long interval = -1;
+
+        private DateTimeZone preTz = DateTimeZone.UTC;
+        private DateTimeZone postTz = DateTimeZone.UTC;
+
+        private float factor = 1.0f;
+
+        public Builder(DateTimeField field) {
+            this.field = field;
+            this.interval = -1;
+        }
+
+        public Builder(TimeValue interval) {
+            this.field = null;
+            this.interval = interval.millis();
+        }
+
+        public Builder preZone(DateTimeZone preTz) {
+            this.preTz = preTz;
+            return this;
+        }
+
+        public Builder postZone(DateTimeZone postTz) {
+            this.postTz = postTz;
+            return this;
+        }
+
+        public Builder factor(float factor) {
+            this.factor = factor;
+            return this;
+        }
+
+        public TimeZoneRounding build() {
+            TimeZoneRounding timeZoneRounding;
+            if (field != null) {
+                if (preTz.equals(DateTimeZone.UTC) && postTz.equals(DateTimeZone.UTC)) {
+                    return new UTCTimeZoneRoundingFloor(field);
+                } else if (field.getDurationField().getUnitMillis() < DateTimeConstants.MILLIS_PER_HOUR * 12) {
+                    timeZoneRounding = new TimeTimeZoneRoundingFloor(field, preTz, postTz);
+                } else {
+                    timeZoneRounding = new DayTimeZoneRoundingFloor(field, preTz, postTz);
+                }
+            } else {
+                if (preTz.equals(DateTimeZone.UTC) && postTz.equals(DateTimeZone.UTC)) {
+                    return new UTCIntervalTimeZoneRounding(interval);
+                } else if (interval < DateTimeConstants.MILLIS_PER_HOUR * 12) {
+                    timeZoneRounding = new TimeIntervalTimeZoneRounding(interval, preTz, postTz);
+                } else {
+                    timeZoneRounding = new DayIntervalTimeZoneRounding(interval, preTz, postTz);
+                }
+            }
+            if (factor != 1.0f) {
+                timeZoneRounding = new FactorTimeZoneRounding(timeZoneRounding, factor);
+            }
+            return timeZoneRounding;
+        }
+    }
+
+    static class TimeTimeZoneRoundingFloor extends TimeZoneRounding {
+
+        private final DateTimeField field;
+        private final DateTimeZone preTz;
+        private final DateTimeZone postTz;
+
+        TimeTimeZoneRoundingFloor(DateTimeField field, DateTimeZone preTz, DateTimeZone postTz) {
+            this.field = field;
+            this.preTz = preTz;
+            this.postTz = postTz;
+        }
+
+        @Override
+        public long calc(long utcMillis) {
+            long time = utcMillis + preTz.getOffset(utcMillis);
+            time = field.roundFloor(time);
+            // now, time is still in local, move it to UTC
+            time = time - preTz.getOffset(time);
+            // now apply post Tz
+            time = time + postTz.getOffset(time);
+            return time;
+        }
+    }
+
+    static class UTCTimeZoneRoundingFloor extends TimeZoneRounding {
+
+        private final DateTimeField field;
+
+        UTCTimeZoneRoundingFloor(DateTimeField field) {
+            this.field = field;
+        }
+
+        @Override
+        public long calc(long utcMillis) {
+            return field.roundFloor(utcMillis);
+        }
+    }
+
+    static class DayTimeZoneRoundingFloor extends TimeZoneRounding {
+        private final DateTimeField field;
+        private final DateTimeZone preTz;
+        private final DateTimeZone postTz;
+
+        DayTimeZoneRoundingFloor(DateTimeField field, DateTimeZone preTz, DateTimeZone postTz) {
+            this.field = field;
+            this.preTz = preTz;
+            this.postTz = postTz;
+        }
+
+        @Override
+        public long calc(long utcMillis) {
+            long time = utcMillis + preTz.getOffset(utcMillis);
+            time = field.roundFloor(time);
+            // after rounding, since its day level (and above), its actually UTC!
+            // now apply post Tz
+            time = time + postTz.getOffset(time);
+            return time;
+        }
+    }
+
+    static class UTCIntervalTimeZoneRounding extends TimeZoneRounding {
+
+        private final long interval;
+
+        UTCIntervalTimeZoneRounding(long interval) {
+            this.interval = interval;
+        }
+
+        @Override
+        public long calc(long utcMillis) {
+            return ((utcMillis / interval) * interval);
+        }
+    }
+
+
+    static class TimeIntervalTimeZoneRounding extends TimeZoneRounding {
+
+        private final long interval;
+        private final DateTimeZone preTz;
+        private final DateTimeZone postTz;
+
+        TimeIntervalTimeZoneRounding(long interval, DateTimeZone preTz, DateTimeZone postTz) {
+            this.interval = interval;
+            this.preTz = preTz;
+            this.postTz = postTz;
+        }
+
+        @Override
+        public long calc(long utcMillis) {
+            long time = utcMillis + preTz.getOffset(utcMillis);
+            time = ((time / interval) * interval);
+            // now, time is still in local, move it to UTC
+            time = time - preTz.getOffset(time);
+            // now apply post Tz
+            time = time + postTz.getOffset(time);
+            return time;
+        }
+    }
+
+    static class DayIntervalTimeZoneRounding extends TimeZoneRounding {
+
+        private final long interval;
+        private final DateTimeZone preTz;
+        private final DateTimeZone postTz;
+
+        DayIntervalTimeZoneRounding(long interval, DateTimeZone preTz, DateTimeZone postTz) {
+            this.interval = interval;
+            this.preTz = preTz;
+            this.postTz = postTz;
+        }
+
+        @Override
+        public long calc(long utcMillis) {
+            long time = utcMillis + preTz.getOffset(utcMillis);
+            time = ((time / interval) * interval);
+            // after rounding, since its day level (and above), its actually UTC!
+            // now apply post Tz
+            time = time + postTz.getOffset(time);
+            return time;
+        }
+    }
+
+    static class FactorTimeZoneRounding extends TimeZoneRounding {
+
+        private final TimeZoneRounding timeZoneRounding;
+
+        private final float factor;
+
+        FactorTimeZoneRounding(TimeZoneRounding timeZoneRounding, float factor) {
+            this.timeZoneRounding = timeZoneRounding;
+            this.factor = factor;
+        }
+
+        @Override
+        public long calc(long utcMillis) {
+            return timeZoneRounding.calc((long) (factor * utcMillis));
+        }
+    }
+}

src/main/java/org/elasticsearch/search/facet/datehistogram/CountDateHistogramFacetCollector.java

@@ -22,6 +22,7 @@
 import gnu.trove.map.hash.TLongLongHashMap;
 import org.apache.lucene.index.IndexReader;
 import org.elasticsearch.common.CacheRecycler;
+import org.elasticsearch.common.joda.TimeZoneRounding;
 import org.elasticsearch.index.cache.field.data.FieldDataCache;
 import org.elasticsearch.index.field.data.FieldDataType;
 import org.elasticsearch.index.field.data.longs.LongFieldData;
@@ -31,22 +32,17 @@
 import org.elasticsearch.search.facet.Facet;
 import org.elasticsearch.search.facet.FacetPhaseExecutionException;
 import org.elasticsearch.search.internal.SearchContext;
-import org.joda.time.MutableDateTime;
 
 import java.io.IOException;
 
 /**
  * A date histogram facet collector that uses the same field as the key as well as the
  * value.
- *
- *
  */
 public class CountDateHistogramFacetCollector extends AbstractFacetCollector {
 
     private final String indexFieldName;
 
-    private final MutableDateTime dateTime;
-
     private final DateHistogramFacet.ComparatorType comparatorType;
 
     private final FieldDataCache fieldDataCache;
@@ -57,9 +53,8 @@ public class CountDateHistogramFacetCollector extends AbstractFacetCollector {
 
     private final DateHistogramProc histoProc;
 
-    public CountDateHistogramFacetCollector(String facetName, String fieldName, MutableDateTime dateTime, long interval, DateHistogramFacet.ComparatorType comparatorType, SearchContext context) {
+    public CountDateHistogramFacetCollector(String facetName, String fieldName, TimeZoneRounding tzRounding, DateHistogramFacet.ComparatorType comparatorType, SearchContext context) {
         super(facetName);
-        this.dateTime = dateTime;
         this.comparatorType = comparatorType;
         this.fieldDataCache = context.fieldDataCache();
 
@@ -77,17 +72,12 @@ public CountDateHistogramFacetCollector(String facetName, String fieldName, Muta
 
         indexFieldName = mapper.names().indexName();
         fieldDataType = mapper.fieldDataType();
-
-        if (interval == 1) {
-            histoProc = new DateHistogramProc();
-        } else {
-            histoProc = new IntervalDateHistogramProc(interval);
-        }
+        histoProc = new DateHistogramProc(tzRounding);
     }
 
     @Override
     protected void doCollect(int doc) throws IOException {
-        fieldData.forEachValueInDoc(doc, dateTime, histoProc);
+        fieldData.forEachValueInDoc(doc, histoProc);
     }
 
     @Override
@@ -100,36 +90,23 @@ public Facet facet() {
         return new InternalCountDateHistogramFacet(facetName, comparatorType, histoProc.counts(), true);
     }
 
-    public static long bucket(long value, long interval) {
-        return ((value / interval) * interval);
-    }
+    public static class DateHistogramProc implements LongFieldData.LongValueInDocProc {
 
-    public static class DateHistogramProc implements LongFieldData.DateValueInDocProc {
+        private final TLongLongHashMap counts = CacheRecycler.popLongLongMap();
 
-        protected final TLongLongHashMap counts = CacheRecycler.popLongLongMap();
+        private final TimeZoneRounding tzRounding;
+
+        public DateHistogramProc(TimeZoneRounding tzRounding) {
+            this.tzRounding = tzRounding;
+        }
 
         @Override
-        public void onValue(int docId, MutableDateTime dateTime) {
-            counts.adjustOrPutValue(dateTime.getMillis(), 1, 1);
+        public void onValue(int docId, long value) {
+            counts.adjustOrPutValue(tzRounding.calc(value), 1, 1);
         }
 
         public TLongLongHashMap counts() {
             return counts;
         }
     }
-
-    public static class IntervalDateHistogramProc extends DateHistogramProc {
-
-        private final long interval;
-
-        public IntervalDateHistogramProc(long interval) {
-            this.interval = interval;
-        }
-
-        @Override
-        public void onValue(int docId, MutableDateTime dateTime) {
-            long bucket = bucket(dateTime.getMillis(), interval);
-            counts.adjustOrPutValue(bucket, 1, 1);
-        }
-    }
 }