A bunch of bug fixes / cleanups for cache feature. (#13)

* `s/enable_caching/enable_cache/`

* Support eneabling cache in Python SDK.

* Fix validation logic for `behavior_version`

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/builder/analyzer.rs

@@ -706,7 +706,7 @@ impl<'a> AnalyzerContext<'a> {
                             })?;
                             let behavior_version = executor.behavior_version();
                             let function_exec_info = AnalyzedFunctionExecInfo {
-                                enable_caching: executor.enable_caching(),
+                                enable_cache: executor.enable_cache(),
                                 behavior_version,
                                 fingerprinter: Fingerprinter::default()
                                     .with(&reactive_op.name)?
@@ -715,8 +715,8 @@ impl<'a> AnalyzerContext<'a> {
                                     .with(&output_type.without_attrs())?,
                                 output_type: output_type.typ.clone(),
                             };
-                            if function_exec_info.enable_caching
-                                && function_exec_info.behavior_version.is_some()
+                            if function_exec_info.enable_cache
+                                && function_exec_info.behavior_version.is_none()
                             {
                                 api_bail!(
                                     "When caching is enabled, behavior version must be specified for transform op: {}",

src/builder/plan.rs

@@ -64,7 +64,7 @@ pub struct AnalyzedSourceOp {
 }
 
 pub struct AnalyzedFunctionExecInfo {
-    pub enable_caching: bool,
+    pub enable_cache: bool,
     pub behavior_version: Option<u32>,
 
     /// Fingerprinter of the function's behavior.

src/execution/evaluator.rs

@@ -306,7 +306,7 @@ async fn evaluate_op_scope(
                 let input_values = assemble_input_values(&op.inputs, scoped_entries);
                 let output_value = if let Some(cache) = op
                     .function_exec_info
-                    .enable_caching
+                    .enable_cache
                     .then_some(cache)
                     .flatten()
                 {

src/ops/interface.rs

@@ -40,11 +40,11 @@ pub trait SimpleFunctionExecutor: Send + Sync {
     /// Evaluate the operation.
     async fn evaluate(&self, args: Vec<Value>) -> Result<Value>;
 
-    fn enable_caching(&self) -> bool {
+    fn enable_cache(&self) -> bool {
         false
     }
 
-    /// Must be Some if `enable_caching` is true.
+    /// Must be Some if `enable_cache` is true.
     /// If it changes, the cache will be invalidated.
     fn behavior_version(&self) -> Option<u32> {
         None

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)?
+                            .extract::<bool>(py)?;
+                        let behavior_version = executor
+                            .call_method(py, "behavior_version", (), None)?
+                            .extract::<Option<u32>>(py)?;
+                        Ok((enable_cache, behavior_version))
+                    })?;
                 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()))