Add cache utilities and plumbing through the evaluator. (#4)

* Add caching options to execution plan with err message enhancement

* Create struct for `MemoizationInfo`.

* Add `EvaluationCache` to access cache.

* Plubming `EvaluatorCache` through the evaluator.

Author: badmonster0

Cargo.toml

@@ -35,3 +35,4 @@ blocking = "1.6.1"
 indenter = "0.3.3"
 itertools = "0.14.0"
 derivative = "2.2.0"
+async-lock = "3.4.0"

src/builder/analyzer.rs

@@ -645,6 +645,7 @@ impl<'a> AnalyzerContext<'a> {
             source_id
         });
 
+        let op_name = source_op.name.clone();
         let output = scope.add_field(source_op.name, &output_type)?;
         let result_fut = async move {
             Ok(AnalyzedSourceOp {
@@ -654,9 +655,10 @@ impl<'a> AnalyzerContext<'a> {
                 primary_key_type: output_type
                     .typ
                     .key_type()
-                    .ok_or_else(|| api_error!("Source must produce a type with key"))?
+                    .ok_or_else(|| api_error!("Source must produce a type with key: {op_name}"))?
                     .typ
                     .clone(),
+                name: op_name,
             })
         };
         Ok(result_fut)
@@ -696,10 +698,27 @@ impl<'a> AnalyzerContext<'a> {
                         let output = scope
                             .data
                             .add_field(reactive_op.name.clone(), &output_type)?;
+                        let op_name = reactive_op.name.clone();
                         async move {
+                            let executor = executor.await.with_context(|| {
+                                format!("Failed to build executor for transform op: {op_name}")
+                            })?;
+                            let function_exec_info = AnalyzedFunctionExecInfo {
+                                enable_caching: executor.enable_caching(),
+                                behavior_version: executor.behavior_version(),
+                            };
+                            if function_exec_info.enable_caching
+                                && function_exec_info.behavior_version.is_some()
+                            {
+                                api_bail!(
+                                    "When caching is enabled, behavior version must be specified for transform op: {op_name}",
+                                );
+                            }
                             Ok(AnalyzedReactiveOp::Transform(AnalyzedTransformOp {
+                                name: op_name,
                                 inputs: input_value_mappings,
-                                executor: executor.await?,
+                                function_exec_info,
+                                executor,
                                 output,
                             }))
                         }
@@ -735,10 +754,14 @@ impl<'a> AnalyzerContext<'a> {
                         parent_scopes.prepend(&scope),
                     )?
                 };
+                let op_name = reactive_op.name.clone();
                 async move {
                     Ok(AnalyzedReactiveOp::ForEach(AnalyzedForEachOp {
                         local_field_ref,
-                        op_scope: op_scope_fut.await?,
+                        op_scope: op_scope_fut
+                            .await
+                            .with_context(|| format!("Analyzing foreach op: {op_name}"))?,
+                        name: op_name,
                     }))
                 }
                 .boxed()
@@ -753,8 +776,10 @@ impl<'a> AnalyzerContext<'a> {
                     struct_schema,
                     scopes,
                 )?;
+                let op_name = reactive_op.name.clone();
                 async move {
                     Ok(AnalyzedReactiveOp::Collect(AnalyzedCollectOp {
+                        name: op_name,
                         input: struct_mapping,
                         collector_ref,
                     }))
@@ -933,8 +958,12 @@ impl<'a> AnalyzerContext<'a> {
             .transpose()?;
 
         Ok(async move {
-            let (executor, query_target) = executor_futs.await?;
+            let (executor, query_target) = executor_futs
+                .await
+                .with_context(|| format!("Analyzing export op: {}", export_op.name))?;
+            let name = export_op.name;
             Ok(AnalyzedExportOp {
+                name,
                 target_id: target_id.unwrap_or_default(),
                 input: local_collector_ref,
                 executor,

src/builder/plan.rs

@@ -55,24 +55,34 @@ pub struct AnalyzedOpOutput {
 }
 
 pub struct AnalyzedSourceOp {
+    pub name: String,
     pub source_id: i32,
     pub executor: Box<dyn SourceExecutor>,
     pub output: AnalyzedOpOutput,
     pub primary_key_type: ValueType,
 }
 
+pub struct AnalyzedFunctionExecInfo {
+    pub enable_caching: bool,
+    pub behavior_version: Option<u32>,
+}
+
 pub struct AnalyzedTransformOp {
+    pub name: String,
     pub inputs: Vec<AnalyzedValueMapping>,
+    pub function_exec_info: AnalyzedFunctionExecInfo,
     pub executor: Box<dyn SimpleFunctionExecutor>,
     pub output: AnalyzedOpOutput,
 }
 
 pub struct AnalyzedForEachOp {
+    pub name: String,
     pub local_field_ref: AnalyzedLocalFieldReference,
     pub op_scope: AnalyzedOpScope,
 }
 
 pub struct AnalyzedCollectOp {
+    pub name: String,
     pub input: AnalyzedStructMapping,
     pub collector_ref: AnalyzedCollectorReference,
 }
@@ -82,6 +92,7 @@ pub enum AnalyzedPrimaryKeyDef {
 }
 
 pub struct AnalyzedExportOp {
+    pub name: String,
     pub target_id: i32,
     pub input: AnalyzedLocalCollectorReference,
     pub executor: Arc<dyn ExportTargetExecutor>,

src/execution/db_tracking.rs

@@ -1,4 +1,4 @@
-use super::db_tracking_setup::TrackingTableSetupState;
+use super::{db_tracking_setup::TrackingTableSetupState, memoization::MemoizationInfo};
 use crate::utils::db::WriteAction;
 use anyhow::Result;
 use sqlx::PgPool;
@@ -11,7 +11,7 @@ pub type TrackedTargetKeyForSource = Vec<(i32, Vec<TrackedTargetKey>)>;
 pub struct SourceTrackingInfo {
     pub max_process_ordinal: i64,
     pub staging_target_keys: sqlx::types::Json<TrackedTargetKeyForSource>,
-    pub memoization_info: Option<serde_json::Value>,
+    pub memoization_info: Option<sqlx::types::Json<MemoizationInfo>>,
 
     pub processed_source_ordinal: Option<i64>,
     pub process_logic_fingerprint: Option<Vec<u8>>,
@@ -73,7 +73,7 @@ pub async fn precommit_source_tracking_info(
     source_key_json: &serde_json::Value,
     max_process_ordinal: i64,
     staging_target_keys: TrackedTargetKeyForSource,
-    memoization_info: serde_json::Value,
+    memoization_info: Option<&MemoizationInfo>,
     db_setup: &TrackingTableSetupState,
     db_executor: impl sqlx::Executor<'_, Database = sqlx::Postgres>,
     action: WriteAction,
@@ -92,7 +92,7 @@ pub async fn precommit_source_tracking_info(
         .bind(source_key_json) // $2
         .bind(max_process_ordinal) // $3
         .bind(sqlx::types::Json(staging_target_keys)) // $4
-        .bind(memoization_info) // $5
+        .bind(memoization_info.map(|m| sqlx::types::Json(m))) // $5
         .execute(db_executor)
         .await?;
     Ok(())

src/execution/evaluator.rs

@@ -11,6 +11,8 @@ use crate::{
     utils::immutable::RefList,
 };
 
+use super::memoization::EvaluationCache;
+
 #[derive(Debug)]
 pub struct ScopeValueBuilder {
     // TODO: Share the same lock for values produced in the same execution scope, for stricter atomicity.
@@ -287,13 +289,15 @@ async fn evaluate_child_op_scope(
     op_scope: &AnalyzedOpScope,
     scoped_entries: RefList<'_, &ScopeEntry<'_>>,
     child_scope_entry: ScopeEntry<'_>,
+    cache: Option<&EvaluationCache>,
 ) -> Result<()> {
-    evaluate_op_scope(op_scope, scoped_entries.prepend(&child_scope_entry)).await
+    evaluate_op_scope(op_scope, scoped_entries.prepend(&child_scope_entry), cache).await
 }
 
 async fn evaluate_op_scope(
     op_scope: &AnalyzedOpScope,
     scoped_entries: RefList<'_, &ScopeEntry<'_>>,
+    cache: Option<&EvaluationCache>,
 ) -> Result<()> {
     let head_scope = *scoped_entries.head().unwrap();
     for reactive_op in op_scope.reactive_ops.iter() {
@@ -324,6 +328,7 @@ async fn evaluate_op_scope(
                                     value: &item,
                                     schema: &collection_schema.row,
                                 },
+                                cache,
                             )
                         })
                         .collect::<Vec<_>>(),
@@ -338,6 +343,7 @@ async fn evaluate_op_scope(
                                     value: v,
                                     schema: &collection_schema.row,
                                 },
+                                cache,
                             )
                         })
                         .collect::<Vec<_>>(),
@@ -353,6 +359,7 @@ async fn evaluate_op_scope(
                                     value: item,
                                     schema: &collection_schema.row,
                                 },
+                                cache,
                             )
                         })
                         .collect::<Vec<_>>(),
@@ -388,6 +395,7 @@ pub async fn evaluate_source_entry<'a>(
     source_op_idx: u32,
     schema: &schema::DataSchema,
     key: &value::KeyValue,
+    cache: Option<&EvaluationCache>,
 ) -> Result<Option<ScopeValueBuilder>> {
     let root_schema = &schema.schema;
     let root_scope_value =
@@ -418,7 +426,12 @@ pub async fn evaluate_source_entry<'a>(
                 value::Value::Table(BTreeMap::from([(key.clone(), scope_value)])),
             );
 
-            evaluate_op_scope(&plan.op_scope, RefList::Nil.prepend(&root_scope_entry)).await?;
+            evaluate_op_scope(
+                &plan.op_scope,
+                RefList::Nil.prepend(&root_scope_entry),
+                cache,
+            )
+            .await?;
             Some(root_scope_value)
         }
         None => None,
@@ -452,6 +465,7 @@ pub async fn evaluate_transient_flow(
     evaluate_op_scope(
         &flow.execution_plan.op_scope,
         RefList::Nil.prepend(&root_scope_entry),
+        None,
     )
     .await?;
     let output_value = assemble_value(

src/execution/indexer.rs

@@ -10,6 +10,7 @@ use sqlx::PgPool;
 
 use super::db_tracking::{self, read_source_tracking_info};
 use super::db_tracking_setup;
+use super::memoization::{EvaluationCache, MemoizationInfo};
 use crate::base::schema;
 use crate::base::spec::FlowInstanceSpec;
 use crate::base::value::{self, FieldValues, KeyValue};
@@ -115,6 +116,10 @@ struct TrackingInfoForTarget<'a> {
     mutation: ExportTargetMutation,
 }
 
+struct PrecommitData<'a> {
+    scope_value: &'a ScopeValueBuilder,
+    memoization_info: &'a MemoizationInfo,
+}
 struct PrecommitMetadata {
     source_entry_exists: bool,
     process_ordinal: i64,
@@ -130,10 +135,9 @@ async fn precommit_source_tracking_info(
     source_id: i32,
     source_key_json: &serde_json::Value,
     source_ordinal: Option<i64>,
-    memoization_info: serde_json::Value,
+    data: Option<PrecommitData<'_>>,
     process_timestamp: &chrono::DateTime<chrono::Utc>,
     db_setup: &db_tracking_setup::TrackingTableSetupState,
-    scope_value: &Option<ScopeValueBuilder>,
     export_ops: &[AnalyzedExportOp],
     pool: &PgPool,
 ) -> Result<WithApplyStatus<PrecommitOutput>> {
@@ -199,9 +203,9 @@ async fn precommit_source_tracking_info(
     }
 
     let mut new_target_keys_info = db_tracking::TrackedTargetKeyForSource::default();
-    if let Some(scope_value) = scope_value {
+    if let Some(data) = &data {
         for export_op in export_ops.iter() {
-            let collected_values = scope_value.collected_values
+            let collected_values = data.scope_value.collected_values
                 [export_op.input.collector_idx as usize]
                 .lock()
                 .unwrap();
@@ -304,7 +308,7 @@ async fn precommit_source_tracking_info(
         source_key_json,
         process_ordinal,
         new_staging_target_keys,
-        memoization_info,
+        data.as_ref().map(|data| data.memoization_info),
         db_setup,
         &mut *txn,
         if tracking_info_exists {
@@ -319,7 +323,7 @@ async fn precommit_source_tracking_info(
 
     Ok(WithApplyStatus::Normal(PrecommitOutput {
         metadata: PrecommitMetadata {
-            source_entry_exists: scope_value.is_some(),
+            source_entry_exists: data.is_some(),
             process_ordinal,
             existing_process_ordinal,
             new_target_keys: new_target_keys_info,
@@ -438,26 +442,51 @@ pub async fn update_source_entry<'a>(
         pool,
     )
     .await?;
-    let scope_value = evaluate_source_entry(plan, source_op_idx, schema, key).await?;
+    let already_exists = existing_tracking_info.is_some();
+    let memoization_info = existing_tracking_info
+        .map(|info| info.memoization_info.map(|info| info.0))
+        .flatten();
+    let evaluation_cache = memoization_info
+        .map(|info| EvaluationCache::from_stored(info.cache))
+        .unwrap_or_default();
+    let value_builder =
+        evaluate_source_entry(plan, source_op_idx, schema, key, Some(&evaluation_cache)).await?;
 
     // Didn't exist and still doesn't exist. No need to apply any changes.
-    if existing_tracking_info.is_none() && scope_value.is_none() {
+    if !already_exists && value_builder.is_none() {
         return Ok(());
     }
 
-    // TODO: Generate the actual source ordinal and memoization info.
-    let source_ordinal: Option<i64> = if scope_value.is_some() { Some(1) } else { None };
-    let memoization_info = serde_json::Value::Null;
+    let memoization_info = MemoizationInfo {
+        cache: evaluation_cache.into_stored()?,
+    };
+    let (source_ordinal, precommit_data) = match &value_builder {
+        Some(scope_value) => {
+            (
+                // TODO: Generate the actual source ordinal.
+                Some(1),
+                Some(PrecommitData {
+                    scope_value,
+                    memoization_info: &memoization_info,
+                }),
+            )
+        }
+        None => (None, None),
+    };
+    if value_builder.is_some() {
+        Some(1)
+    } else {
+        None
+    };
 
     // Phase 2 (precommit): Update with the memoization info and stage target keys.
     let precommit_output = precommit_source_tracking_info(
         source_id,
         &source_key_json,
         source_ordinal,
-        memoization_info,
+        precommit_data,
         &process_timestamp,
         &plan.tracking_table_setup,
-        &scope_value,
         &plan.export_ops,
         pool,
     )