Support 'calendar_interval' and 'fixed_interval' in DateHistogramFacet

elastic · Dec 4, 2020 · a252fac · a252fac
1 parent e7f8a90
commit a252fac
Show file tree

Hide file tree

Showing 4 changed files with 80 additions and 9 deletions.
diff --git a/docs/faceted_search.rst b/docs/faceted_search.rst
@@ -49,7 +49,7 @@ There are several different facets available:
   provides an option to split documents into groups based on a value of a field, for example ``TermsFacet(field='category')``
 
 ``DateHistogramFacet``
-  split documents into time intervals, example: ``DateHistogramFacet(field="published_date", interval="day")``
+  split documents into time intervals, example: ``DateHistogramFacet(field="published_date", calendar_interval="day")``
 
 ``HistogramFacet``
   similar to ``DateHistogramFacet`` but for numerical values: ``HistogramFacet(field="rating", interval=2)``

diff --git a/elasticsearch_dsl/faceted_search.py b/elasticsearch_dsl/faceted_search.py
@@ -168,14 +168,34 @@ def get_value_filter(self, filter_value):
         )
 
 
+def _date_interval_month(d):
+    return (d + timedelta(days=32)).replace(day=1)
+
+
+def _date_interval_week(d):
+    return d + timedelta(days=7)
+
+
+def _date_interval_day(d):
+    return d + timedelta(days=1)
+
+
+def _date_interval_hour(d):
+    return d + timedelta(hours=1)
+
+
 class DateHistogramFacet(Facet):
     agg_type = "date_histogram"
 
     DATE_INTERVALS = {
-        "month": lambda d: (d + timedelta(days=32)).replace(day=1),
-        "week": lambda d: d + timedelta(days=7),
-        "day": lambda d: d + timedelta(days=1),
-        "hour": lambda d: d + timedelta(hours=1),
+        "month": _date_interval_month,
+        "1M": _date_interval_month,
+        "week": _date_interval_week,
+        "1w": _date_interval_week,
+        "day": _date_interval_day,
+        "1d": _date_interval_day,
+        "hour": _date_interval_hour,
+        "1h": _date_interval_hour,
     }
 
     def __init__(self, **kwargs):
@@ -194,12 +214,20 @@ def get_value(self, bucket):
             return bucket["key"]
 
     def get_value_filter(self, filter_value):
+        for interval_type in ("calendar_interval", "fixed_interval"):
+            if interval_type in self._params:
+                break
+        else:
+            interval_type = "interval"
+
         return Range(
             _expand__to_dot=False,
             **{
                 self._params["field"]: {
                     "gte": filter_value,
-                    "lt": self.DATE_INTERVALS[self._params["interval"]](filter_value),
+                    "lt": self.DATE_INTERVALS[self._params[interval_type]](
+                        filter_value
+                    ),
                 }
             }
         )

diff --git a/tests/test_faceted_search.py b/tests/test_faceted_search.py
@@ -17,6 +17,8 @@
 
 from datetime import datetime
 
+import pytest
+
 from elasticsearch_dsl.faceted_search import (
     DateHistogramFacet,
     FacetedSearch,
@@ -144,3 +146,45 @@ def test_date_histogram_facet_with_1970_01_01_date():
     dhf = DateHistogramFacet()
     assert dhf.get_value({"key": None}) == datetime(1970, 1, 1, 0, 0)
     assert dhf.get_value({"key": 0}) == datetime(1970, 1, 1, 0, 0)
+
+
+@pytest.mark.parametrize(
+    ["interval_type", "interval"],
+    [
+        ("interval", "month"),
+        ("calendar_interval", "month"),
+        ("interval", "week"),
+        ("calendar_interval", "week"),
+        ("interval", "day"),
+        ("calendar_interval", "day"),
+        ("fixed_interval", "day"),
+        ("interval", "hour"),
+        ("fixed_interval", "hour"),
+        ("interval", "1M"),
+        ("calendar_interval", "1M"),
+        ("interval", "1w"),
+        ("calendar_interval", "1w"),
+        ("interval", "1d"),
+        ("calendar_interval", "1d"),
+        ("fixed_interval", "1d"),
+        ("interval", "1h"),
+        ("fixed_interval", "1h"),
+    ],
+)
+def test_date_histogram_interval_types(interval_type, interval):
+    dhf = DateHistogramFacet(field="@timestamp", **{interval_type: interval})
+    assert dhf.get_aggregation().to_dict() == {
+        "date_histogram": {
+            "field": "@timestamp",
+            interval_type: interval,
+            "min_doc_count": 0,
+        }
+    }
+    dhf.get_value_filter(datetime.now())
+
+
+def test_date_histogram_no_interval_keyerror():
+    dhf = DateHistogramFacet(field="@timestamp")
+    with pytest.raises(KeyError) as e:
+        dhf.get_value_filter(datetime.now())
+    assert str(e.value) == "'interval'"
diff --git a/tests/test_integration/test_faceted_search.py b/tests/test_integration/test_faceted_search.py
@@ -39,7 +39,7 @@ class CommitSearch(FacetedSearch):
     facets = {
         "files": TermsFacet(field="files"),
         "frequency": DateHistogramFacet(
-            field="authored_date", interval="day", min_doc_count=1
+            field="authored_date", fixed_interval="1d", min_doc_count=1
         ),
         "deletions": RangeFacet(
             field="stats.deletions",
@@ -69,7 +69,7 @@ class RepoSearch(FacetedSearch):
     doc_types = [Repos]
     facets = {
         "public": TermsFacet(field="is_public"),
-        "created": DateHistogramFacet(field="created_at", interval="month"),
+        "created": DateHistogramFacet(field="created_at", calendar_interval="month"),
     }
 
     def search(self):
@@ -146,7 +146,6 @@ def test_boolean_facet(data_client):
 
 def test_empty_search_finds_everything(data_client):
     cs = CommitSearch()
-
     r = cs.execute()
 
     assert r.hits.total.value == 52