refactor: pack SourceEvalContext for common args for processing/eval

badmonster0 · badmonster0 · commit aa256ec41792 · 2025-05-11T10:52:08.000-07:00
diff --git a/src/execution/dumper.rs b/src/execution/dumper.rs
@@ -11,6 +11,7 @@ use std::collections::BTreeMap;
 use std::path::{Path, PathBuf};
 use yaml_rust2::YamlEmitter;
 
+use super::evaluator::SourceRowEvaluationContext;
 use super::memoization::EvaluationMemoryOptions;
 use super::row_indexer;
 use crate::base::{schema, value};
@@ -77,10 +78,12 @@ impl<'a> Dumper<'a> {
         'a: 'b,
     {
         let data_builder = row_indexer::evaluate_source_entry_with_memory(
-            self.plan,
-            import_op,
-            self.schema,
-            key,
+            &SourceRowEvaluationContext {
+                plan: self.plan,
+                import_op,
+                schema: self.schema,
+                key,
+            },
             EvaluationMemoryOptions {
                 enable_cache: self.options.use_cache,
                 evaluation_only: true,
diff --git a/src/execution/evaluator.rs b/src/execution/evaluator.rs
@@ -446,30 +446,34 @@ async fn evaluate_op_scope(
     Ok(())
 }
 
+pub struct SourceRowEvaluationContext<'a> {
+    pub plan: &'a ExecutionPlan,
+    pub import_op: &'a AnalyzedImportOp,
+    pub schema: &'a schema::FlowSchema,
+    pub key: &'a value::KeyValue,
+}
+
 #[derive(Debug)]
 pub struct EvaluateSourceEntryOutput {
     pub data_scope: ScopeValueBuilder,
     pub collected_values: Vec<Vec<value::FieldValues>>,
 }
 
 pub async fn evaluate_source_entry(
-    plan: &ExecutionPlan,
-    import_op: &AnalyzedImportOp,
-    schema: &schema::FlowSchema,
-    key: &value::KeyValue,
+    src_eval_ctx: &SourceRowEvaluationContext<'_>,
     source_value: value::FieldValues,
     memory: &EvaluationMemory,
 ) -> Result<EvaluateSourceEntryOutput> {
-    let root_schema = &schema.schema;
+    let root_schema = &src_eval_ctx.schema.schema;
     let root_scope_value = ScopeValueBuilder::new(root_schema.fields.len());
     let root_scope_entry = ScopeEntry::new(
         ScopeKey::None,
         &root_scope_value,
         root_schema,
-        &plan.op_scope,
+        &src_eval_ctx.plan.op_scope,
     );
 
-    let table_schema = match &root_schema.fields[import_op.output.field_idx as usize]
+    let table_schema = match &root_schema.fields[src_eval_ctx.import_op.output.field_idx as usize]
         .value_type
         .typ
     {
@@ -482,12 +486,12 @@ pub async fn evaluate_source_entry(
     let scope_value =
         ScopeValueBuilder::augmented_from(&value::ScopeValue(source_value), table_schema)?;
     root_scope_entry.define_field_w_builder(
-        &import_op.output,
-        value::Value::KTable(BTreeMap::from([(key.clone(), scope_value)])),
+        &src_eval_ctx.import_op.output,
+        value::Value::KTable(BTreeMap::from([(src_eval_ctx.key.clone(), scope_value)])),
     );
 
     evaluate_op_scope(
-        &plan.op_scope,
+        &src_eval_ctx.plan.op_scope,
         RefList::Nil.prepend(&root_scope_entry),
         memory,
     )
diff --git a/src/execution/row_indexer.rs b/src/execution/row_indexer.rs
@@ -6,11 +6,12 @@ use std::collections::{HashMap, HashSet};
 
 use super::db_tracking::{self, read_source_tracking_info_for_processing, TrackedTargetKey};
 use super::db_tracking_setup;
-use super::evaluator::{evaluate_source_entry, EvaluateSourceEntryOutput};
+use super::evaluator::{
+    evaluate_source_entry, EvaluateSourceEntryOutput, SourceRowEvaluationContext,
+};
 use super::memoization::{EvaluationMemory, EvaluationMemoryOptions, StoredMemoizationInfo};
 use super::stats;
 
-use crate::base::schema;
 use crate::base::value::{self, FieldValues, KeyValue};
 use crate::builder::plan::*;
 use crate::ops::interface::{
@@ -439,19 +440,16 @@ async fn commit_source_tracking_info(
 }
 
 pub async fn evaluate_source_entry_with_memory(
-    plan: &ExecutionPlan,
-    import_op: &AnalyzedImportOp,
-    schema: &schema::FlowSchema,
-    key: &value::KeyValue,
+    src_eval_ctx: &SourceRowEvaluationContext<'_>,
     options: EvaluationMemoryOptions,
     pool: &PgPool,
 ) -> Result<Option<EvaluateSourceEntryOutput>> {
     let stored_info = if options.enable_cache || !options.evaluation_only {
-        let source_key_json = serde_json::to_value(key)?;
+        let source_key_json = serde_json::to_value(src_eval_ctx.key)?;
         let existing_tracking_info = read_source_tracking_info_for_processing(
-            import_op.source_id,
+            src_eval_ctx.import_op.source_id,
             &source_key_json,
-            &plan.tracking_table_setup,
+            &src_eval_ctx.plan.tracking_table_setup,
             pool,
         )
         .await?;
@@ -462,10 +460,11 @@ pub async fn evaluate_source_entry_with_memory(
         None
     };
     let memory = EvaluationMemory::new(chrono::Utc::now(), stored_info, options);
-    let source_value = match import_op
+    let source_value = match src_eval_ctx
+        .import_op
         .executor
         .get_value(
-            key,
+            src_eval_ctx.key,
             &SourceExecutorGetOptions {
                 include_value: true,
                 include_ordinal: false,
@@ -477,28 +476,25 @@ pub async fn evaluate_source_entry_with_memory(
         Some(d) => d,
         None => return Ok(None),
     };
-    let output = evaluate_source_entry(plan, import_op, schema, key, source_value, &memory).await?;
+    let output = evaluate_source_entry(src_eval_ctx, source_value, &memory).await?;
     Ok(Some(output))
 }
 
 pub async fn update_source_row(
-    plan: &ExecutionPlan,
-    import_op: &AnalyzedImportOp,
-    schema: &schema::FlowSchema,
-    key: &value::KeyValue,
+    src_eval_ctx: &SourceRowEvaluationContext<'_>,
     source_value: Option<FieldValues>,
     source_version: &SourceVersion,
     pool: &PgPool,
     update_stats: &stats::UpdateStats,
 ) -> Result<SkippedOr<()>> {
-    let source_key_json = serde_json::to_value(key)?;
-    let process_timestamp = chrono::Utc::now();
+    let source_key_json = serde_json::to_value(src_eval_ctx.key)?;
+    let process_time = chrono::Utc::now();
 
     // Phase 1: Evaluate with memoization info.
     let existing_tracking_info = read_source_tracking_info_for_processing(
-        import_op.source_id,
+        src_eval_ctx.import_op.source_id,
         &source_key_json,
-        &plan.tracking_table_setup,
+        &src_eval_ctx.plan.tracking_table_setup,
         pool,
     )
     .await?;
@@ -507,7 +503,7 @@ pub async fn update_source_row(
             let existing_version = SourceVersion::from_stored(
                 info.processed_source_ordinal,
                 &info.process_logic_fingerprint,
-                plan.logic_fingerprint,
+                src_eval_ctx.plan.logic_fingerprint,
             );
             if existing_version.should_skip(source_version, Some(update_stats)) {
                 return Ok(SkippedOr::Skipped(existing_version));
@@ -522,40 +518,33 @@ pub async fn update_source_row(
     let (output, stored_mem_info) = match source_value {
         Some(source_value) => {
             let evaluation_memory = EvaluationMemory::new(
-                process_timestamp,
+                process_time,
                 memoization_info,
                 EvaluationMemoryOptions {
                     enable_cache: true,
                     evaluation_only: false,
                 },
             );
-            let output = evaluate_source_entry(
-                plan,
-                import_op,
-                schema,
-                key,
-                source_value,
-                &evaluation_memory,
-            )
-            .await?;
+            let output =
+                evaluate_source_entry(src_eval_ctx, source_value, &evaluation_memory).await?;
             (Some(output), evaluation_memory.into_stored()?)
         }
         None => Default::default(),
     };
 
     // Phase 2 (precommit): Update with the memoization info and stage target keys.
     let precommit_output = precommit_source_tracking_info(
-        import_op.source_id,
+        src_eval_ctx.import_op.source_id,
         &source_key_json,
         source_version,
-        plan.logic_fingerprint,
+        src_eval_ctx.plan.logic_fingerprint,
         output.as_ref().map(|scope_value| PrecommitData {
             evaluate_output: scope_value,
             memoization_info: &stored_mem_info,
         }),
-        &process_timestamp,
-        &plan.tracking_table_setup,
-        &plan.export_ops,
+        &process_time,
+        &src_eval_ctx.plan.tracking_table_setup,
+        &src_eval_ctx.plan.export_ops,
         update_stats,
         pool,
     )
@@ -567,40 +556,44 @@ pub async fn update_source_row(
 
     // Phase 3: Apply changes to the target storage, including upserting new target records and removing existing ones.
     let mut target_mutations = precommit_output.target_mutations;
-    let apply_futs = plan.export_op_groups.iter().filter_map(|export_op_group| {
-        let mutations_w_ctx: Vec<_> = export_op_group
-            .op_idx
-            .iter()
-            .filter_map(|export_op_idx| {
-                let export_op = &plan.export_ops[*export_op_idx];
-                target_mutations
-                    .remove(&export_op.target_id)
-                    .filter(|m| !m.is_empty())
-                    .map(|mutation| interface::ExportTargetMutationWithContext {
-                        mutation,
-                        export_context: export_op.export_context.as_ref(),
-                    })
+    let apply_futs = src_eval_ctx
+        .plan
+        .export_op_groups
+        .iter()
+        .filter_map(|export_op_group| {
+            let mutations_w_ctx: Vec<_> = export_op_group
+                .op_idx
+                .iter()
+                .filter_map(|export_op_idx| {
+                    let export_op = &src_eval_ctx.plan.export_ops[*export_op_idx];
+                    target_mutations
+                        .remove(&export_op.target_id)
+                        .filter(|m| !m.is_empty())
+                        .map(|mutation| interface::ExportTargetMutationWithContext {
+                            mutation,
+                            export_context: export_op.export_context.as_ref(),
+                        })
+                })
+                .collect();
+            (!mutations_w_ctx.is_empty()).then(|| {
+                export_op_group
+                    .target_factory
+                    .apply_mutation(mutations_w_ctx)
             })
-            .collect();
-        (!mutations_w_ctx.is_empty()).then(|| {
-            export_op_group
-                .target_factory
-                .apply_mutation(mutations_w_ctx)
-        })
-    });
+        });
 
     // TODO: Handle errors.
     try_join_all(apply_futs).await?;
 
     // Phase 4: Update the tracking record.
     commit_source_tracking_info(
-        import_op.source_id,
+        src_eval_ctx.import_op.source_id,
         &source_key_json,
         source_version,
-        &plan.logic_fingerprint.0,
+        &src_eval_ctx.plan.logic_fingerprint.0,
         precommit_output.metadata,
-        &process_timestamp,
-        &plan.tracking_table_setup,
+        &process_time,
+        &src_eval_ctx.plan.tracking_table_setup,
         pool,
     )
     .await?;
diff --git a/src/execution/source_indexer.rs b/src/execution/source_indexer.rs
@@ -6,6 +6,7 @@ use tokio::{sync::Semaphore, task::JoinSet};
 
 use super::{
     db_tracking,
+    evaluator::SourceRowEvaluationContext,
     row_indexer::{self, SkippedOr, SourceVersion, SourceVersionKind},
     stats,
 };
@@ -115,10 +116,12 @@ impl SourceIndexingContext {
             };
             let schema = &self.flow.data_schema;
             let result = row_indexer::update_source_row(
-                &plan,
-                import_op,
-                schema,
-                &key,
+                &SourceRowEvaluationContext {
+                    plan: &plan,
+                    import_op,
+                    schema,
+                    key: &key,
+                },
                 source_value,
                 &source_version,
                 &pool,
diff --git a/src/service/flows.rs b/src/service/flows.rs
@@ -1,11 +1,8 @@
 use crate::prelude::*;
 
+use crate::execution::{evaluator, memoization, row_indexer, stats};
 use crate::lib_context::LibContext;
 use crate::{base::schema::FlowSchema, ops::interface::SourceExecutorListOptions};
-use crate::{
-    execution::memoization,
-    execution::{row_indexer, stats},
-};
 use axum::{
     extract::{Path, State},
     http::StatusCode,
@@ -147,10 +144,12 @@ pub async fn evaluate_data(
     let key = value::KeyValue::from_strs(query.key, &key_field.value_type.typ)?;
 
     let evaluate_output = row_indexer::evaluate_source_entry_with_memory(
-        &plan,
-        import_op,
-        schema,
-        &key,
+        &evaluator::SourceRowEvaluationContext {
+            plan: &plan,
+            import_op,
+            schema,
+            key: &key,
+        },
         memoization::EvaluationMemoryOptions {
             enable_cache: true,
             evaluation_only: true,
