Support eneabling cache in Python SDK.

Author: badmonster0

python/cocoindex/op.py

@@ -62,12 +62,14 @@ def __call__(self, spec_json: str, *args, **kwargs):
 
 _gpu_dispatch_lock = Lock()
 
-def executor_class(gpu: bool = False) -> Callable[[type], type]:
+def executor_class(gpu: bool = False, cache: bool = False, behavior_version: int | None = None) -> Callable[[type], type]:
     """
     Decorate a class to provide an executor for an op.
 
     Args:
         gpu: Whether the executor will be executed on GPU.
+        cache: Whether the executor will be cached.
+        behavior_version: The behavior version of the executor. Cache will be invalidated if it changes. Must be provided if `cache` is True.
     """
 
     def _inner(cls: type[Executor]) -> type:
@@ -87,7 +89,15 @@ def _inner(cls: type[Executor]) -> type:
         expected_return = sig.return_annotation
 
         cls_type: type = cls
-        class _WrappedClass(cls_type):
+
+        class _Fallback:
+            def enable_cache(self):
+                return cache
+
+            def behavior_version(self):
+                return behavior_version
+
+        class _WrappedClass(cls_type, _Fallback):
             def __init__(self, spec):
                 super().__init__()
                 self.spec = spec

src/ops/py_factory.rs

@@ -149,6 +149,9 @@ struct PyFunctionExecutor {
     num_positional_args: usize,
     kw_args_names: Vec<Py<PyString>>,
     result_type: schema::EnrichedValueType,
+
+    enable_cache: bool,
+    behavior_version: Option<u32>,
 }
 
 #[async_trait]
@@ -193,6 +196,14 @@ impl SimpleFunctionExecutor for Arc<PyFunctionExecutor> {
         })
         .await
     }
+
+    fn enable_cache(&self) -> bool {
+        self.enable_cache
+    }
+
+    fn behavior_version(&self) -> Option<u32> {
+        self.behavior_version
+    }
 }
 
 pub(crate) struct PyFunctionFactory {
@@ -251,15 +262,27 @@ impl SimpleFunctionFactory for PyFunctionFactory {
 
         let executor_fut = {
             let result_type = result_type.clone();
-            async move {
-                Python::with_gil(|py| executor.call_method(py, "prepare", (), None))?;
+            unblock(move || {
+                let (enable_cache, behavior_version) =
+                    Python::with_gil(|py| -> anyhow::Result<_> {
+                        executor.call_method(py, "prepare", (), None)?;
+                        let enable_cache = executor.call_method(py, "enable_cache", (), None)?;
+                        let behavior_version =
+                            executor.call_method(py, "behavior_version", (), None)?;
+                        Ok((
+                            enable_cache.extract::<bool>(py)?,
+                            behavior_version.extract::<Option<u32>>(py)?,
+                        ))
+                    })?;
                 Ok(Box::new(Arc::new(PyFunctionExecutor {
                     py_function_executor: executor,
                     num_positional_args,
                     kw_args_names,
                     result_type,
+                    enable_cache,
+                    behavior_version,
                 })) as Box<dyn SimpleFunctionExecutor>)
-            }
+            })
         };
 
         Ok((result_type, executor_fut.boxed()))