feat: timestamptz support (#3002) (#3092)

issue:milvus-io/milvus#27467

---------

Signed-off-by: zhenshan.cao <zhenshan.cao@zilliz.com>
Co-authored-by: xtx <xtianx@smail.nju.edu.cn>

Author: czs007

.gitignore

@@ -19,6 +19,7 @@ venv/
 
 # Env
 .env
+.venv/
 
 # Local Temp
 temp/

examples/timestamptz.py

@@ -0,0 +1,128 @@
+import time
+from pymilvus import MilvusClient, DataType, CollectionSchema, IndexType, FieldSchema, milvus_client
+import datetime
+import pytz
+
+milvus_host = "http://localhost:19530"
+collection_name = "timestamptz_test123"
+
+
+def main():
+  client = MilvusClient(uri=milvus_host)
+
+  # create collection with TIMESTAMPTZ field
+  if client.has_collection(collection_name):
+    client.drop_collection(collection_name)
+
+  schema = client.create_schema()
+  schema.add_field("id", DataType.INT64, is_primary=True)
+  schema.add_field("tsz", DataType.TIMESTAMPTZ)
+  schema.add_field("vec", DataType.FLOAT_VECTOR, dim=4)
+  print("===================alter database timezone===================")
+  try:
+    client.alter_database_properties("default", {"database.timezone": "Asia/Shanghai"})
+  except Exception as e:
+    print(e)
+  print(client.describe_database("default"))
+  client.create_collection(collection_name, schema=schema, consistency_level="Session")
+  index_params = client.prepare_index_params(
+    collection_name=collection_name,
+    field_name="vec",
+    index_type=IndexType.HNSW,
+    metric_type="COSINE",
+    params={"M": 30, "efConstruction": 200},
+  )
+  # add timestamptz index of STL_SORT type
+  index_params.add_index(field_name="tsz", index_name="tsz_index", index_type="STL_SORT")
+
+  client.create_index(collection_name, index_params)
+  print(client.describe_collection(collection_name))
+  client.load_collection(collection_name)
+  print(f"load state: {client.get_load_state(collection_name)}")
+
+  # insert data with timezone-aware timestamps
+  print("===================insert timestamptz===================")
+  data_size = 8193
+  shanghai_tz = pytz.timezone("Asia/Shanghai")
+  data = [
+    {
+      "id": i + 1,
+      "tsz": shanghai_tz.localize(
+        datetime.datetime(2025, 1, 1, 0, 0, 0) + datetime.timedelta(days=i)
+      ).isoformat(),
+      "vec": [float(i) / 10 for i in range(4)],
+    }
+    for i in range(data_size)
+  ]
+  client.insert(collection_name, data)
+  client.flush(collection_name)
+  time.sleep(1) # wait for index creation
+  print(client.describe_index(collection_name, "tsz_index"))
+  print("===================insert invalid string===================")
+  data = [{"id": 114514, "tsz": "should cause an error", "vec": [1.1, 1.2, 1.3]}]
+  try:
+    client.insert(collection_name, data)
+  except Exception as e:
+    print(e)
+
+  # query/search data with TIMESTAMPTZ, define timezone in kwargs
+  print("====================test query====================")
+  results = client.query(
+    collection_name,
+    filter="id <= 10",
+    output_fields=["id", "tsz", "vec"],
+    limit=2,
+    timezone="America/Havana",
+    time_fields="year, month, day, hour, minute, second, microsecond",
+  )
+  print("\n".join([str(res) for res in results]))
+
+  print("====================test search====================")
+  results = client.search(
+    collection_name,
+    [[0.5, 0.6, 0.7, 0.8]],
+    output_fields=["id", "tsz", "vec"],
+    limit=10,
+    timezone="America/Chicago",
+    time_fields="year, month, day, hour, minute, second, microsecond",
+  )
+  print("\n".join([str(res) for res in results[0]]))
+
+  # Alter timezone (IANA timezone)
+  print(client.describe_collection(collection_name))
+  print("===================alter collection timezone===================")
+  try:
+    client.alter_collection_properties(collection_name, {"collection.timezone": "Asia/Shanghai"})
+  except Exception as e:
+    print(e)
+  print(client.describe_collection(collection_name))
+
+  try:
+    client.alter_collection_properties(collection_name, {"collection.timezone": "error"})
+  except Exception as e:
+    print(e)
+  print(client.describe_collection(collection_name))
+
+  try:
+    client.alter_database_properties("default", {"database.timezone": "error"})
+  except Exception as e:
+    print(e)
+  print(client.describe_database("default"))
+
+  # Query with new operator
+  results = client.query(
+    collection_name, limit=10, timezone="Asia/Shanghai"
+  )
+  print("\n".join([str(res) for res in results]))
+
+  expr = "tsz + INTERVAL 'P0D' != ISO '2025-01-03T00:00:00+08:00'"
+  results = client.query(collection_name, expr, output_fields=["id", "tsz"], limit=10)
+  print("  The first expr:")
+  print("\n".join([str(res) for res in results]))
+  expr = "tsz != ISO '2025-01-03T00:00:00+08:00'"
+  print("  The second expr")
+  results = client.query(collection_name, expr, output_fields=["id", "tsz"], limit=10)
+  print("\n".join([str(res) for res in results]))
+
+if __name__ == "__main__":
+  main()

pymilvus/client/entity_helper.py

@@ -373,6 +373,17 @@ def pack_field_value_to_field_data(
                 % (field_name, "double", type(field_value))
                 + f" Detail: {e!s}"
             ) from e
+    elif field_type == DataType.TIMESTAMPTZ:
+        try:
+            if field_value is None:
+                field_data.scalars.string_data.data.extend([])  # Timestamptz is passed as String
+            else:
+                field_data.scalars.string_data.data.append(field_value)
+        except (TypeError, ValueError) as e:
+            raise DataNotMatchException(
+                message=ExceptionsMessage.FieldDataInconsistent
+                % (field_name, "string", type(field_value))
+            ) from e
     elif field_type == DataType.FLOAT_VECTOR:
         try:
             f_value = field_value
@@ -787,6 +798,11 @@ def assign_scalar(data: List[Any]) -> None:
         assign_scalar(data)
         return False
 
+    if field_data.type == DataType.TIMESTAMPTZ:
+        data = field_data.scalars.string_data.data
+        assign_scalar(data)
+        return False
+
     if field_data.type == DataType.VARCHAR:
         data = field_data.scalars.string_data.data
         assign_scalar(data)

pymilvus/client/prepare.py

@@ -1478,6 +1478,12 @@ def search_requests_with_expr(
         if param.get("analyzer_name") is not None:
             search_params["analyzer_name"] = param["analyzer_name"]
 
+        if kwargs.get("timezone") is not None:
+            search_params["timezone"] = kwargs["timezone"]
+
+        if kwargs.get("time_fields") is not None:
+            search_params["time_fields"] = kwargs["time_fields"]
+
         search_params["params"] = get_params(param)
 
         req_params = [
@@ -1913,6 +1919,14 @@ def query_request(
         if offset is not None:
             req.query_params.append(common_types.KeyValuePair(key="offset", value=str(offset)))
 
+        timezone = kwargs.get("timezone")
+        if timezone is not None:
+            req.query_params.append(common_types.KeyValuePair(key="timezone", value=timezone))
+
+        timefileds = kwargs.get("time_fields")
+        if timefileds is not None:
+            req.query_params.append(common_types.KeyValuePair(key="time_fields", value=timefileds))
+
         ignore_growing = kwargs.get("ignore_growing", False)
         stop_reduce_for_best = kwargs.get(REDUCE_STOP_FOR_BEST, False)
         is_iterator = kwargs.get(ITERATOR_FIELD)

pymilvus/client/search_result.py

@@ -57,6 +57,7 @@ def __init__(
                 DataType.DOUBLE,
                 DataType.VARCHAR,
                 DataType.GEOMETRY,
+                DataType.TIMESTAMPTZ,
             ]:
                 if has_valid:
                     [
@@ -529,7 +530,7 @@ def get_field_data(field_data: FieldData):
         return field_data.scalars.float_data.data
     if field_data.type == DataType.DOUBLE:
         return field_data.scalars.double_data.data
-    if field_data.type == DataType.VARCHAR:
+    if field_data.type in (DataType.VARCHAR, DataType.TIMESTAMPTZ):
         return field_data.scalars.string_data.data
     if field_data.type == DataType.GEOMETRY:
         return field_data.scalars.geometry_wkt_data.data

pymilvus/client/types.py

@@ -115,6 +115,7 @@ class DataType(IntEnum):
     ARRAY = schema_pb2.Array
     JSON = schema_pb2.JSON
     GEOMETRY = schema_pb2.Geometry
+    TIMESTAMPTZ = schema_pb2.Timestamptz
 
     BINARY_VECTOR = schema_pb2.BinaryVector
     FLOAT_VECTOR = schema_pb2.FloatVector

pymilvus/client/utils.py

@@ -164,6 +164,9 @@ def len_of(field_data: Any) -> int:
         if field_data.scalars.HasField("double_data"):
             return len(field_data.scalars.double_data.data)
 
+        if field_data.scalars.HasField("timestamptz_data"):
+            return len(field_data.scalars.timestamptz_data.data)
+
         if field_data.scalars.HasField("string_data"):
             return len(field_data.scalars.string_data.data)
 

tests/test_collection.py

@@ -17,6 +17,7 @@ def test_collection_by_DataFrame(self):
             FieldSchema("float16_vector", DataType.FLOAT16_VECTOR, dim=128),
             FieldSchema("bfloat16_vector", DataType.BFLOAT16_VECTOR, dim=128),
             FieldSchema("int8_vector", DataType.INT8_VECTOR, dim=128),
+            FieldSchema("timestamptz", DataType.TIMESTAMPTZ),
         ]
 
         prefix = "pymilvus.client.grpc_handler.GrpcHandler"