From 803b24e9811a96413063a1163530dd31c531e66e Mon Sep 17 00:00:00 2001
From: Christian Trimble <christian.trimble@meltmedia.com>
Date: Sat, 17 May 2014 16:57:54 -0700
Subject: [PATCH] fixed offsets in DateHistogramBuilder

---
 .../histogram/DateHistogramBuilder.java       |  24 +++-
 .../bucket/DateHistogramTests.java            | 111 ++++++++++++++++++
 2 files changed, 131 insertions(+), 4 deletions(-)

diff --git a/src/main/java/org/elasticsearch/search/aggregations/bucket/histogram/DateHistogramBuilder.java b/src/main/java/org/elasticsearch/search/aggregations/bucket/histogram/DateHistogramBuilder.java
index 62dfce2409b3e..ce320c58dda88 100644
--- a/src/main/java/org/elasticsearch/search/aggregations/bucket/histogram/DateHistogramBuilder.java
+++ b/src/main/java/org/elasticsearch/search/aggregations/bucket/histogram/DateHistogramBuilder.java
@@ -41,8 +41,8 @@ public class DateHistogramBuilder extends ValuesSourceAggregationBuilder<DateHis
     private String postZone;
     private boolean preZoneAdjustLargeInterval;
     private String format;
-    long preOffset = 0;
-    long postOffset = 0;
+    private Object preOffset;
+    private Object postOffset;
     float factor = 1.0f;
 
     public DateHistogramBuilder(String name) {
@@ -94,6 +94,16 @@ public DateHistogramBuilder postOffset(long postOffset) {
         return this;
     }
 
+    public DateHistogramBuilder preOffset( DateHistogram.Interval preOffset ) {
+        this.preOffset = preOffset;
+        return this;
+    }
+
+    public DateHistogramBuilder postOffset( DateHistogram.Interval postOffset ) {
+        this.postOffset = postOffset;
+        return this;
+    }
+
     public DateHistogramBuilder factor(float factor) {
         this.factor = factor;
         return this;
@@ -153,11 +163,17 @@ protected XContentBuilder doInternalXContent(XContentBuilder builder, Params par
             builder.field("pre_zone_adjust_large_interval", true);
         }
 
-        if (preOffset != 0) {
+        if (preOffset != null) {
+            if (preOffset instanceof Number) {
+                preOffset = TimeValue.timeValueMillis(((Number) preOffset).longValue()).toString();
+            }
             builder.field("pre_offset", preOffset);
         }
 
-        if (postOffset != 0) {
+        if (postOffset != null) {
+            if (postOffset instanceof Number) {
+                postOffset = TimeValue.timeValueMillis(((Number) postOffset).longValue()).toString();
+            }
             builder.field("post_offset", postOffset);
         }
 
diff --git a/src/test/java/org/elasticsearch/search/aggregations/bucket/DateHistogramTests.java b/src/test/java/org/elasticsearch/search/aggregations/bucket/DateHistogramTests.java
index a8cece62fb530..dcb11b444211d 100644
--- a/src/test/java/org/elasticsearch/search/aggregations/bucket/DateHistogramTests.java
+++ b/src/test/java/org/elasticsearch/search/aggregations/bucket/DateHistogramTests.java
@@ -167,6 +167,117 @@ public void singleValuedField() throws Exception {
         assertThat(bucket.getDocCount(), equalTo(3l));
     }
 
+    @Test
+    public void singleValuedField_WithLongOffsets() throws Exception {
+        SearchResponse response;
+        response = client().prepareSearch("idx")
+            .addAggregation(dateHistogram("histo").field("date")
+                .interval(DateHistogram.Interval.DAY)
+                .preOffset(1000*60*60*24*2)
+                .postOffset(-1000*60*60*24))
+            .execute().actionGet();
+
+        assertSearchResponse(response);
+
+
+        DateHistogram histo = response.getAggregations().get("histo");
+        assertThat(histo, notNullValue());
+        assertThat(histo.getName(), equalTo("histo"));
+        assertThat(histo.getBuckets().size(), equalTo(6));
+
+        long key = new DateTime(2012, 1, 3, 0, 0, DateTimeZone.UTC).getMillis();
+        DateHistogram.Bucket bucket = histo.getBucketByKey(key);
+        assertThat(bucket, notNullValue());
+        assertThat(bucket.getKeyAsNumber().longValue(), equalTo(key));
+        assertThat(bucket.getDocCount(), equalTo(1l));
+
+        key = new DateTime(2012, 2, 3, 0, 0, DateTimeZone.UTC).getMillis();
+        bucket = histo.getBucketByKey(key);
+        assertThat(bucket, notNullValue());
+        assertThat(bucket.getKeyAsNumber().longValue(), equalTo(key));
+        assertThat(bucket.getDocCount(), equalTo(1l));
+
+        key = new DateTime(2012, 2, 16, 0, 0, DateTimeZone.UTC).getMillis();
+        bucket = histo.getBucketByKey(key);
+        assertThat(bucket, notNullValue());
+        assertThat(bucket.getKeyAsNumber().longValue(), equalTo(key));
+        assertThat(bucket.getDocCount(), equalTo(1l));
+
+        key = new DateTime(2012, 3, 3, 0, 0, DateTimeZone.UTC).getMillis();
+        bucket = histo.getBucketByKey(key);
+        assertThat(bucket, notNullValue());
+        assertThat(bucket.getKeyAsNumber().longValue(), equalTo(key));
+        assertThat(bucket.getDocCount(), equalTo(1l));
+
+        key = new DateTime(2012, 3, 16, 0, 0, DateTimeZone.UTC).getMillis();
+        bucket = histo.getBucketByKey(key);
+        assertThat(bucket, notNullValue());
+        assertThat(bucket.getKeyAsNumber().longValue(), equalTo(key));
+        assertThat(bucket.getDocCount(), equalTo(1l));
+
+        key = new DateTime(2012, 3, 24, 0, 0, DateTimeZone.UTC).getMillis();
+        bucket = histo.getBucketByKey(key);
+        assertThat(bucket, notNullValue());
+        assertThat(bucket.getKeyAsNumber().longValue(), equalTo(key));
+        assertThat(bucket.getDocCount(), equalTo(1l));
+    }
+    
+    @Test
+    public void singleValuedField_WithIntervalOffsets() throws Exception {
+        SearchResponse response;
+        response = client().prepareSearch("idx")
+            .addAggregation(dateHistogram("histo").field("date")
+                .interval(DateHistogram.Interval.DAY)
+                .preOffset(DateHistogram.Interval.days(2))
+                .postOffset(DateHistogram.Interval.days(-1)))
+            .execute().actionGet();
+
+        assertSearchResponse(response);
+
+
+        DateHistogram histo = response.getAggregations().get("histo");
+        assertThat(histo, notNullValue());
+        assertThat(histo.getName(), equalTo("histo"));
+        assertThat(histo.getBuckets().size(), equalTo(6));
+
+        long key = new DateTime(2012, 1, 3, 0, 0, DateTimeZone.UTC).getMillis();
+        DateHistogram.Bucket bucket = histo.getBucketByKey(key);
+        assertThat(bucket, notNullValue());
+        assertThat(bucket.getKeyAsNumber().longValue(), equalTo(key));
+        assertThat(bucket.getDocCount(), equalTo(1l));
+
+        key = new DateTime(2012, 2, 3, 0, 0, DateTimeZone.UTC).getMillis();
+        bucket = histo.getBucketByKey(key);
+        assertThat(bucket, notNullValue());
+        assertThat(bucket.getKeyAsNumber().longValue(), equalTo(key));
+        assertThat(bucket.getDocCount(), equalTo(1l));
+
+        key = new DateTime(2012, 2, 16, 0, 0, DateTimeZone.UTC).getMillis();
+        bucket = histo.getBucketByKey(key);
+        assertThat(bucket, notNullValue());
+        assertThat(bucket.getKeyAsNumber().longValue(), equalTo(key));
+        assertThat(bucket.getDocCount(), equalTo(1l));
+
+        key = new DateTime(2012, 3, 3, 0, 0, DateTimeZone.UTC).getMillis();
+        bucket = histo.getBucketByKey(key);
+        assertThat(bucket, notNullValue());
+        assertThat(bucket.getKeyAsNumber().longValue(), equalTo(key));
+        assertThat(bucket.getDocCount(), equalTo(1l));
+
+        key = new DateTime(2012, 3, 16, 0, 0, DateTimeZone.UTC).getMillis();
+        bucket = histo.getBucketByKey(key);
+        assertThat(bucket, notNullValue());
+        assertThat(bucket.getKeyAsNumber().longValue(), equalTo(key));
+        assertThat(bucket.getDocCount(), equalTo(1l));
+
+        key = new DateTime(2012, 3, 24, 0, 0, DateTimeZone.UTC).getMillis();
+        bucket = histo.getBucketByKey(key);
+        assertThat(bucket, notNullValue());
+        assertThat(bucket.getKeyAsNumber().longValue(), equalTo(key));
+        assertThat(bucket.getDocCount(), equalTo(1l));
+    }
+
+
     @Test
     public void singleValuedField_WithPostTimeZone() throws Exception {
         SearchResponse response;