Expose C++ TemporalPolicy in Python (#362)

Author: jparismorgan

apis/python/src/tiledb/vector_search/type_erased_module.cc

@@ -120,6 +120,37 @@ void init_type_erased_module(py::module_& m) {
       }))
       ;
 #endif
+  py::class_<TemporalPolicy>(m, "TemporalPolicy", py::buffer_protocol())
+      // From 0 to UINT64_MAX.
+      .def(py::init<>())
+      // From 0 to timestamp_end.
+      .def(
+          "__init__",
+          [](TemporalPolicy& instance,
+             std::optional<uint64_t> timestamp_end_input) {
+            uint64_t timestamp_end = timestamp_end_input.has_value() ?
+                                         timestamp_end_input.value() :
+                                         UINT64_MAX;
+            new (&instance) TemporalPolicy(TimeTravel, timestamp_end);
+          })
+      // From timestamp_start to timestamp_end.
+      .def(
+          "__init__",
+          [](TemporalPolicy& instance,
+             std::optional<uint64_t> timestamp_start_input,
+             std::optional<uint64_t> timestamp_end_input) {
+            uint64_t timestamp_start = timestamp_start_input.has_value() ?
+                                           timestamp_start_input.value() :
+                                           0;
+            uint64_t timestamp_end = timestamp_end_input.has_value() ?
+                                         timestamp_end_input.value() :
+                                         UINT64_MAX;
+            new (&instance) TemporalPolicy(
+                TimestampStartEnd, timestamp_start, timestamp_end);
+          })
+      .def("timestamp_start", &TemporalPolicy::timestamp_start)
+      .def("timestamp_end", &TemporalPolicy::timestamp_end);
+
   py::class_<FeatureVector>(m, "FeatureVector", py::buffer_protocol())
       .def(
           py::init<const tiledb::Context&, const std::string&>(),

apis/python/src/tiledb/vector_search/utils.py

@@ -1,8 +1,10 @@
 import io
+from typing import Optional
 
 import numpy as np
 
 import tiledb
+from tiledb.vector_search import _tiledbvspy as vspy
 
 
 def is_type_erased_index(index_type: str) -> bool:
@@ -21,6 +23,19 @@ def add_to_group(group, uri, name):
         group.add(name, name=name, relative=True)
 
 
+def to_temporal_policy(timestamp) -> Optional[vspy.TemporalPolicy]:
+    temporal_policy = None
+    if isinstance(timestamp, tuple):
+        if len(timestamp) != 2:
+            raise ValueError(
+                "'timestamp' argument expects either int or tuple(start: int, end: int)"
+            )
+        temporal_policy = vspy.TemporalPolicy(timestamp[0], timestamp[1])
+    elif timestamp is not None:
+        temporal_policy = vspy.TemporalPolicy(timestamp)
+    return temporal_policy
+
+
 def _load_vecs_t(uri, dtype, ctx_or_config=None):
     with tiledb.scope_ctx(ctx_or_config) as ctx:
         dtype = np.dtype(dtype)

apis/python/test/test_type_erased_module.py

@@ -5,6 +5,7 @@
 
 from tiledb.vector_search import _tiledbvspy as vspy
 from tiledb.vector_search.utils import load_fvecs
+from tiledb.vector_search.utils import to_temporal_policy
 
 ctx = vspy.Ctx({})
 
@@ -152,6 +153,49 @@ def test_numpy_to_feature_vector_array():
     assert np.array_equal(a, np.transpose(np.array(b)))
 
 
+def test_TemporalPolicy():
+    temporal_policy = vspy.TemporalPolicy()
+    assert temporal_policy.timestamp_start() == 0
+    assert temporal_policy.timestamp_end() == np.iinfo(np.uint64).max
+
+    temporal_policy = vspy.TemporalPolicy(99)
+    assert temporal_policy.timestamp_start() == 0
+    assert temporal_policy.timestamp_end() == 99
+
+    temporal_policy = vspy.TemporalPolicy(12, 99)
+    assert temporal_policy.timestamp_start() == 12
+    assert temporal_policy.timestamp_end() == 99
+
+
+def test_TemporalPolicy_from_timestamp():
+    temporal_policy = to_temporal_policy(None)
+    assert temporal_policy is None
+
+    temporal_policy = to_temporal_policy(3)
+    assert temporal_policy.timestamp_start() == 0
+    assert temporal_policy.timestamp_end() == 3
+
+    temporal_policy = to_temporal_policy((0, 33))
+    assert temporal_policy.timestamp_start() == 0
+    assert temporal_policy.timestamp_end() == 33
+
+    temporal_policy = to_temporal_policy((1, 33))
+    assert temporal_policy.timestamp_start() == 1
+    assert temporal_policy.timestamp_end() == 33
+
+    temporal_policy = to_temporal_policy((None, 333))
+    assert temporal_policy.timestamp_start() == 0
+    assert temporal_policy.timestamp_end() == 333
+
+    temporal_policy = to_temporal_policy((3333, None))
+    assert temporal_policy.timestamp_start() == 3333
+    assert temporal_policy.timestamp_end() == np.iinfo(np.uint64).max
+
+    temporal_policy = to_temporal_policy((None, None))
+    assert temporal_policy.timestamp_start() == 0
+    assert temporal_policy.timestamp_end() == np.iinfo(np.uint64).max
+
+
 def test_construct_IndexFlatL2():
     a = vspy.IndexFlatL2(ctx, siftsmall_inputs_uri)
     assert a.feature_type_string() == "float32"