Fixed a bug in date_histogram aggregation parsing

- pre_zone_adjust_large_interval was not parsed properly
 - added tests for pre_zone and pre_zone_adjust_large_interval
 - changed DateHistogram#getBucketByKey(String) to support date formats (next to numeric strings)
 - added randomized testing for fetching the bucket by key in date_histogram tests
 - added missing "format" support in DateHistogramBuilder

 Closes #5375

src/main/java/org/elasticsearch/search/aggregations/bucket/histogram/DateHistogramBuilder.java

@@ -37,6 +37,7 @@ public class DateHistogramBuilder extends ValuesSourceAggregationBuilder<DateHis
     private String preZone;
     private String postZone;
     private boolean preZoneAdjustLargeInterval;
+    private String format;
     long preOffset = 0;
     long postOffset = 0;
     float factor = 1.0f;
@@ -95,6 +96,11 @@ public class DateHistogramBuilder extends ValuesSourceAggregationBuilder<DateHis
         return this;
     }
 
+    public DateHistogramBuilder format(String format) {
+        this.format = format;
+        return this;
+    }
+
     @Override
     protected XContentBuilder doInternalXContent(XContentBuilder builder, Params params) throws IOException {
         if (interval == null) {
@@ -138,6 +144,10 @@ public class DateHistogramBuilder extends ValuesSourceAggregationBuilder<DateHis
             builder.field("factor", factor);
         }
 
+        if (format != null) {
+            builder.field("format", format);
+        }
+
         return builder;
     }
 

src/main/java/org/elasticsearch/search/aggregations/bucket/histogram/DateHistogramParser.java

@@ -113,8 +113,6 @@ public class DateHistogramParser implements Aggregator.Parser {
                     preZone = parseZone(parser.text());
                 } else if ("pre_zone".equals(currentFieldName) || "preZone".equals(currentFieldName)) {
                     preZone = parseZone(parser.text());
-                } else if ("pre_zone_adjust_large_interval".equals(currentFieldName) || "preZoneAdjustLargeInterval".equals(currentFieldName)) {
-                    preZoneAdjustLargeInterval = parser.booleanValue();
                 } else if ("post_zone".equals(currentFieldName) || "postZone".equals(currentFieldName)) {
                     postZone = parseZone(parser.text());
                 } else if ("pre_offset".equals(currentFieldName) || "preOffset".equals(currentFieldName)) {
@@ -133,6 +131,8 @@ public class DateHistogramParser implements Aggregator.Parser {
                     keyed = parser.booleanValue();
                 } else if ("script_values_sorted".equals(currentFieldName) || "scriptValuesSorted".equals(currentFieldName)) {
                     assumeSorted = parser.booleanValue();
+                } else if ("pre_zone_adjust_large_interval".equals(currentFieldName) || "preZoneAdjustLargeInterval".equals(currentFieldName)) {
+                    preZoneAdjustLargeInterval = parser.booleanValue();
                 } else {
                     throw new SearchParseException(context, "Unknown key for a " + token + " in [" + aggregationName + "]: [" + currentFieldName + "].");
                 }

src/main/java/org/elasticsearch/search/aggregations/bucket/histogram/InternalDateHistogram.java

@@ -18,6 +18,7 @@
  */
 package org.elasticsearch.search.aggregations.bucket.histogram;
 
+import com.carrotsearch.hppc.ObjectObjectOpenHashMap;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.index.mapper.core.DateFieldMapper;
 import org.elasticsearch.search.aggregations.AggregationStreams;
@@ -91,6 +92,8 @@ public class InternalDateHistogram extends InternalHistogram<InternalDateHistogr
         }
     }
 
+    private ObjectObjectOpenHashMap<String, InternalDateHistogram.Bucket> bucketsMap;
+
     InternalDateHistogram() {} // for serialization
 
     InternalDateHistogram(String name, List<InternalDateHistogram.Bucket> buckets, InternalOrder order, long minDocCount,
@@ -103,6 +106,23 @@ public class InternalDateHistogram extends InternalHistogram<InternalDateHistogr
         return TYPE;
     }
 
+    @Override
+    public Bucket getBucketByKey(String key) {
+        try {
+            long time = Long.parseLong(key);
+            return super.getBucketByKey(time);
+        } catch (NumberFormatException nfe) {
+            // it's not a number, so lets try to parse it as a date using the formatter.
+        }
+        if (bucketsMap == null) {
+            bucketsMap = new ObjectObjectOpenHashMap<String, InternalDateHistogram.Bucket>();
+            for (InternalDateHistogram.Bucket bucket : buckets) {
+                bucketsMap.put(bucket.getKey(), bucket);
+            }
+        }
+        return bucketsMap.get(key);
+    }
+
     @Override
     public DateHistogram.Bucket getBucketByKey(DateTime key) {
         return getBucketByKey(key.getMillis());
@@ -112,4 +132,11 @@ public class InternalDateHistogram extends InternalHistogram<InternalDateHistogr
     protected InternalDateHistogram.Bucket createBucket(long key, long docCount, InternalAggregations aggregations, ValueFormatter formatter) {
         return new Bucket(key, docCount, aggregations, formatter);
     }
+
+    @Override
+    public void readFrom(StreamInput in) throws IOException {
+        super.readFrom(in);
+        bucketsMap = null; // we need to reset this on read (as it's lazily created on demand)
+    }
+
 }

src/main/java/org/elasticsearch/search/aggregations/bucket/histogram/InternalHistogram.java

@@ -161,7 +161,7 @@ public class InternalHistogram<B extends InternalHistogram.Bucket> extends Inter
 
     }
 
-    private List<B> buckets;
+    protected List<B> buckets;
     private LongObjectOpenHashMap<B> bucketsMap;
     private InternalOrder order;
     private ValueFormatter formatter;

src/test/java/org/elasticsearch/search/aggregations/bucket/DateHistogramTests.java

@@ -20,7 +20,9 @@ package org.elasticsearch.search.aggregations.bucket;
 
 import org.elasticsearch.action.index.IndexRequestBuilder;
 import org.elasticsearch.action.search.SearchResponse;
+import org.elasticsearch.common.joda.Joda;
 import org.elasticsearch.common.xcontent.XContentBuilder;
+import org.elasticsearch.index.mapper.core.DateFieldMapper;
 import org.elasticsearch.search.aggregations.AbstractAggregationBuilder;
 import org.elasticsearch.search.aggregations.bucket.histogram.DateHistogram;
 import org.elasticsearch.search.aggregations.bucket.histogram.Histogram;
@@ -35,6 +37,7 @@ import org.junit.Test;
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.List;
 
 import static org.elasticsearch.common.xcontent.XContentFactory.jsonBuilder;
@@ -59,6 +62,10 @@ public class DateHistogramTests extends ElasticsearchIntegrationTest {
         return new DateTime(2012, month, day, 0, 0, DateTimeZone.UTC);
     }
 
+    private DateTime date(String date) {
+        return DateFieldMapper.Defaults.DATE_TIME_FORMATTER.parser().parseDateTime(date);
+    }
+
     private IndexRequestBuilder indexDoc(int month, int day, int value) throws Exception {
         return client().prepareIndex("idx", "type").setSource(jsonBuilder()
                 .startObject()
@@ -83,6 +90,23 @@ public class DateHistogramTests extends ElasticsearchIntegrationTest {
         ensureSearchable();
     }
 
+    private static DateHistogram.Bucket getBucket(DateHistogram histogram, DateTime key) {
+        return getBucket(histogram, key, DateFieldMapper.Defaults.DATE_TIME_FORMATTER.format());
+    }
+
+    private static DateHistogram.Bucket getBucket(DateHistogram histogram, DateTime key, String format) {
+        if (randomBoolean()) {
+            if (randomBoolean()) {
+                return histogram.getBucketByKey(key);
+            }
+            return histogram.getBucketByKey(key.getMillis());
+        }
+        if (randomBoolean()) {
+            return histogram.getBucketByKey("" + key.getMillis());
+        }
+        return histogram.getBucketByKey(Joda.forPattern(format).printer().print(key));
+    }
+
     @Test
     public void singleValuedField() throws Exception {
         SearchResponse response = client().prepareSearch("idx")
@@ -97,22 +121,22 @@ public class DateHistogramTests extends ElasticsearchIntegrationTest {
         assertThat(histo.getName(), equalTo("histo"));
         assertThat(histo.getBuckets().size(), equalTo(3));
 
-        long key = new DateTime(2012, 1, 1, 0, 0, DateTimeZone.UTC).getMillis();
+        DateTime key = new DateTime(2012, 1, 1, 0, 0, DateTimeZone.UTC);
-        DateHistogram.Bucket bucket = histo.getBucketByKey(key);
+        DateHistogram.Bucket bucket = getBucket(histo, key);
         assertThat(bucket, notNullValue());
-        assertThat(bucket.getKeyAsNumber().longValue(), equalTo(key));
+        assertThat(bucket.getKeyAsNumber().longValue(), equalTo(key.getMillis()));
         assertThat(bucket.getDocCount(), equalTo(1l));
 
-        key = new DateTime(2012, 2, 1, 0, 0, DateTimeZone.UTC).getMillis();
+        key = new DateTime(2012, 2, 1, 0, 0, DateTimeZone.UTC);
-        bucket = histo.getBucketByKey(key);
+        bucket = getBucket(histo, key);
         assertThat(bucket, notNullValue());
-        assertThat(bucket.getKeyAsNumber().longValue(), equalTo(key));
+        assertThat(bucket.getKeyAsNumber().longValue(), equalTo(key.getMillis()));
         assertThat(bucket.getDocCount(), equalTo(2l));
 
-        key = new DateTime(2012, 3, 1, 0, 0, DateTimeZone.UTC).getMillis();
+        key = new DateTime(2012, 3, 1, 0, 0, DateTimeZone.UTC);
-        bucket = histo.getBucketByKey(key);
+        bucket = getBucket(histo, key);
         assertThat(bucket, notNullValue());
-        assertThat(bucket.getKeyAsNumber().longValue(), equalTo(key));
+        assertThat(bucket.getKeyAsNumber().longValue(), equalTo(key.getMillis()));
         assertThat(bucket.getDocCount(), equalTo(3l));
     }
 
@@ -432,7 +456,6 @@ public class DateHistogramTests extends ElasticsearchIntegrationTest {
 
         assertSearchResponse(response);
 
-
         DateHistogram histo = response.getAggregations().get("histo");
         assertThat(histo, notNullValue());
         assertThat(histo.getName(), equalTo("histo"));
@@ -975,4 +998,75 @@ public class DateHistogramTests extends ElasticsearchIntegrationTest {
         assertThat(dateHisto.getBuckets().isEmpty(), is(true));
 
     }
+
+    @Test
+    public void singleValue_WithPreZone() throws Exception {
+        prepareCreate("idx2").addMapping("type", "date", "type=date").execute().actionGet();
+        IndexRequestBuilder[] reqs = new IndexRequestBuilder[5];
+        DateTime date = date("2014-03-11T00:00:00+00:00");
+        for (int i = 0; i < reqs.length; i++) {
+            reqs[i] = client().prepareIndex("idx2", "type", "" + i).setSource(jsonBuilder().startObject().field("date", date).endObject());
+            date = date.plusHours(1);
+        }
+        indexRandom(true, reqs);
+
+        SearchResponse response = client().prepareSearch("idx2")
+                .setQuery(matchAllQuery())
+                .addAggregation(dateHistogram("date_histo")
+                        .field("date")
+                        .preZone("-2:00")
+                        .interval(DateHistogram.Interval.DAY)
+                        .format("yyyy-MM-dd"))
+                .execute().actionGet();
+
+        assertThat(response.getHits().getTotalHits(), equalTo(5l));
+
+        DateHistogram histo = response.getAggregations().get("date_histo");
+        Collection<? extends DateHistogram.Bucket> buckets = histo.getBuckets();
+        assertThat(buckets.size(), equalTo(2));
+
+        DateHistogram.Bucket bucket = histo.getBucketByKey("2014-03-10");
+        assertThat(bucket, Matchers.notNullValue());
+        assertThat(bucket.getDocCount(), equalTo(2l));
+
+        bucket = histo.getBucketByKey("2014-03-11");
+        assertThat(bucket, Matchers.notNullValue());
+        assertThat(bucket.getDocCount(), equalTo(3l));
+    }
+
+    @Test
+    public void singleValue_WithPreZone_WithAadjustLargeInterval() throws Exception {
+        prepareCreate("idx2").addMapping("type", "date", "type=date").execute().actionGet();
+        IndexRequestBuilder[] reqs = new IndexRequestBuilder[5];
+        DateTime date = date("2014-03-11T00:00:00+00:00");
+        for (int i = 0; i < reqs.length; i++) {
+            reqs[i] = client().prepareIndex("idx2", "type", "" + i).setSource(jsonBuilder().startObject().field("date", date).endObject());
+            date = date.plusHours(1);
+        }
+        indexRandom(true, reqs);
+
+        SearchResponse response = client().prepareSearch("idx2")
+                .setQuery(matchAllQuery())
+                .addAggregation(dateHistogram("date_histo")
+                        .field("date")
+                        .preZone("-2:00")
+                        .interval(DateHistogram.Interval.DAY)
+                        .preZoneAdjustLargeInterval(true)
+                        .format("yyyy-MM-dd'T'HH:mm:ss"))
+                .execute().actionGet();
+
+        assertThat(response.getHits().getTotalHits(), equalTo(5l));
+
+        DateHistogram histo = response.getAggregations().get("date_histo");
+        Collection<? extends DateHistogram.Bucket> buckets = histo.getBuckets();
+        assertThat(buckets.size(), equalTo(2));
+
+        DateHistogram.Bucket bucket = histo.getBucketByKey("2014-03-10T02:00:00");
+        assertThat(bucket, Matchers.notNullValue());
+        assertThat(bucket.getDocCount(), equalTo(2l));
+
+        bucket = histo.getBucketByKey("2014-03-11T02:00:00");
+        assertThat(bucket, Matchers.notNullValue());
+        assertThat(bucket.getDocCount(), equalTo(3l));
+    }
 }