Add support for datetime ranges (#517)

qdrant_client/conversions/conversion.py

@@ -1,4 +1,4 @@
-from datetime import datetime, timezone
+from datetime import date, datetime, timezone
 from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple, Union, get_args
 
 from google.protobuf.json_format import MessageToDict
@@ -51,6 +51,8 @@ def json_to_value(payload: Any) -> Value:
         return Value(
             struct_value=Struct(fields=dict((k, json_to_value(v)) for k, v in payload.items()))
         )
+    if isinstance(payload, datetime) or isinstance(payload, date):
+        return Value(string_value=payload.isoformat())
     raise ValueError(f"Not supported json value: {payload}")  # pragma: no cover
 
 

qdrant_client/local/payload_filters.py

@@ -3,6 +3,7 @@ from typing import Any, List, Optional
 import numpy as np
 
 from qdrant_client.http import models
+from qdrant_client.local import datetime_utils
 from qdrant_client.local.geo import boolean_point_in_polygon, geo_distance
 from qdrant_client.local.payload_value_extractor import value_by_key
 
@@ -86,18 +87,40 @@ def check_geo_polygon(condition: models.GeoPolygon, values: Any) -> bool:
     return False
 
 
+def check_range_interface(condition: models.RangeInterface, value: Any) -> bool:
+    if isinstance(condition, models.Range):
+        return check_range(condition, value)
+    if isinstance(condition, models.DatetimeRange):
+        return check_datetime_range(condition, value)
+    return False
+
+
 def check_range(condition: models.Range, value: Any) -> bool:
     if not isinstance(value, (int, float)):
         return False
-    if condition.lt is not None and value >= condition.lt:
+    return (
+        (condition.lt is None or value < condition.lt)
+        and (condition.lte is None or value <= condition.lte)
+        and (condition.gt is None or value > condition.gt)
+        and (condition.gte is None or value >= condition.gte)
+    )
+
+
+def check_datetime_range(condition: models.DatetimeRange, value: Any) -> bool:
+    if not isinstance(value, str):
         return False
-    if condition.lte is not None and value > condition.lte:
+
+    dt = datetime_utils.parse(value)
+
+    if dt is None:
         return False
-    if condition.gt is not None and value <= condition.gt:
-        return False
-    if condition.gte is not None and value < condition.gte:
-        return False
-    return True
+
+    return (
+        (condition.lt is None or dt < condition.lt)
+        and (condition.lte is None or dt <= condition.lte)
+        and (condition.gt is None or dt > condition.gt)
+        and (condition.gte is None or dt >= condition.gte)
+    )
 
 
 def check_match(condition: models.Match, value: Any) -> bool:
@@ -145,7 +168,7 @@ def check_condition(
         if condition.range is not None:
             if values is None:
                 return False
-            return any(check_range(condition.range, v) for v in values)
+            return any(check_range_interface(condition.range, v) for v in values)
         if condition.geo_bounding_box is not None:
             if values is None:
                 return False

tests/fixtures/filters.py

@@ -1,4 +1,5 @@
 import random
+from datetime import datetime, timedelta, timezone
 
 from qdrant_client.http import models
 from tests.fixtures.payload import geo_points, random_real_word
@@ -186,6 +187,37 @@ def range_field_condition() -> models.FieldCondition:
     )
 
 
+def datetime_range_field_condition() -> models.FieldCondition:
+    field = "rand_datetime"
+
+    start_datetime = datetime(2000, 1, 1)
+    end_datetime = datetime(2001, 1, 31)
+
+    def random_datetime() -> datetime:
+        dt = start_datetime + timedelta(
+            seconds=random.randint(0, int((end_datetime - start_datetime).total_seconds())),
+            microseconds=random.randint(0, 999999),
+        )
+        return dt.replace(tzinfo=timezone(offset=timedelta(hours=random.randint(-12, 12))))
+
+    lt = random_datetime()
+    gt = random_datetime()
+
+    rand_1 = random.random()
+    rand_2 = random.random()
+
+    if rand_1 > rand_2:
+        if rand_1 > 0.5:
+            lt = None
+        else:
+            gt = None
+
+    return models.FieldCondition(
+        key=field,
+        range=models.DatetimeRange(lt=lt, gt=gt),
+    )
+
+
 def geo_bounding_box_field_condition() -> models.FieldCondition:
     field = "city.geo"
     random_top_left = {"lat": random.random() * 180 - 90, "lon": random.random() * 360 - 180}
@@ -239,6 +271,7 @@ def one_random_condition_please() -> models.Condition:
             match_any_field_condition,
             match_except_field_condition,
             range_field_condition,
+            datetime_range_field_condition,
             geo_bounding_box_field_condition,
             geo_radius_field_condition,
             values_count_field_condition,

tests/fixtures/payload.py

@@ -1,7 +1,7 @@
 import random
 import uuid
 from datetime import datetime, timedelta, timezone
-from typing import Any, Dict
+from typing import Any, Dict, Union
 
 from qdrant_client.local import datetime_utils
 
@@ -154,7 +154,7 @@ start_datetime = datetime(2000, 1, 1)
 end_datetime = datetime(2001, 1, 31)
 
 
-def random_datetime_str() -> str:
+def random_datetime() -> Union[str, datetime]:
     random_datetime = start_datetime + timedelta(
         seconds=random.randint(0, int((end_datetime - start_datetime).total_seconds())),
         microseconds=random.randint(0, 999999),
@@ -165,6 +165,10 @@ def random_datetime_str() -> str:
         random_datetime = random_datetime.replace(
             tzinfo=timezone(offset=timedelta(hours=random.randint(-12, 12)))
         )
+
+    if random.random() < 0.1:
+        return random_datetime
+
     dt_str = random_datetime.strftime(fmt)
     return dt_str
 
@@ -185,7 +189,7 @@ def one_random_payload_please(idx: int) -> Dict[str, Any]:
         "text_data": uuid.uuid4().hex,
         "rand_digit": random.randint(0, 9),
         "rand_number": round(random.random(), 5),
-        "rand_datetime": random_datetime_str(),
+        "rand_datetime": random_datetime(),
         "text_array": [uuid.uuid4().hex, uuid.uuid4().hex],
         "words": f"{random_real_word()} {random_real_word()}",
         "nested": {