Refactor: simplify arguments for internal indexer functions.

Author: badmonster0

src/execution/evaluator.rs

@@ -428,7 +428,7 @@ async fn evaluate_op_scope(
 
 pub async fn evaluate_source_entry(
     plan: &ExecutionPlan,
-    source_op_idx: usize,
+    source_op: &AnalyzedSourceOp,
     schema: &schema::DataSchema,
     key: &value::KeyValue,
     cache: Option<&EvaluationCache>,
@@ -442,7 +442,6 @@ pub async fn evaluate_source_entry(
         schema: root_schema,
     };
 
-    let source_op = &plan.source_ops[source_op_idx];
     let collection_schema = match &root_schema.fields[source_op.output.field_idx as usize]
         .value_type
         .typ

src/execution/indexer.rs

@@ -11,7 +11,6 @@ use super::db_tracking::{self, read_source_tracking_info, TrackedTargetKey};
 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};
 use crate::builder::plan::*;
 use crate::ops::interface::{ExportTargetMutation, ExportTargetUpsertEntry};
@@ -420,15 +419,14 @@ async fn commit_source_tracking_info(
 
 pub async fn evaluate_source_entry_with_cache(
     plan: &ExecutionPlan,
-    source_op_idx: usize,
+    source_op: &AnalyzedSourceOp,
     schema: &schema::DataSchema,
     key: &value::KeyValue,
     pool: &PgPool,
 ) -> Result<Option<value::ScopeValue>> {
-    let source_id = plan.source_ops[source_op_idx].source_id;
     let source_key_json = serde_json::to_value(key)?;
     let existing_tracking_info = read_source_tracking_info(
-        source_id,
+        source_op.source_id,
         &source_key_json,
         &plan.tracking_table_setup,
         pool,
@@ -441,28 +439,27 @@ pub async fn evaluate_source_entry_with_cache(
     let evaluation_cache =
         EvaluationCache::new(process_timestamp, memoization_info.map(|info| info.cache));
     let data_builder =
-        evaluate_source_entry(plan, source_op_idx, schema, key, Some(&evaluation_cache)).await?;
+        evaluate_source_entry(plan, source_op, schema, key, Some(&evaluation_cache)).await?;
     Ok(data_builder.map(|builder| builder.into()))
 }
 
 pub async fn update_source_entry(
     plan: &ExecutionPlan,
-    source_op_idx: usize,
+    source_op: &AnalyzedSourceOp,
     schema: &schema::DataSchema,
     key: &value::KeyValue,
     only_for_deletion: bool,
     pool: &PgPool,
     stats: &UpdateStats,
 ) -> Result<()> {
-    let source_id = plan.source_ops[source_op_idx].source_id;
     let source_key_json = serde_json::to_value(key)?;
     let process_timestamp = chrono::Utc::now();
 
     // Phase 1: Evaluate with memoization info.
 
     // TODO: Skip if the source is not newer and the processing logic is not changed.
     let existing_tracking_info = read_source_tracking_info(
-        source_id,
+        source_op.source_id,
         &source_key_json,
         &plan.tracking_table_setup,
         pool,
@@ -475,7 +472,7 @@ pub async fn update_source_entry(
     let evaluation_cache =
         EvaluationCache::new(process_timestamp, memoization_info.map(|info| info.cache));
     let value_builder = if !only_for_deletion {
-        evaluate_source_entry(plan, source_op_idx, schema, key, Some(&evaluation_cache)).await?
+        evaluate_source_entry(plan, source_op, schema, key, Some(&evaluation_cache)).await?
     } else {
         None
     };
@@ -512,7 +509,7 @@ pub async fn update_source_entry(
 
     // Phase 2 (precommit): Update with the memoization info and stage target keys.
     let precommit_output = precommit_source_tracking_info(
-        source_id,
+        source_op.source_id,
         &source_key_json,
         source_ordinal,
         precommit_data,
@@ -546,7 +543,7 @@ pub async fn update_source_entry(
 
     // Phase 4: Update the tracking record.
     commit_source_tracking_info(
-        source_id,
+        source_op.source_id,
         &source_key_json,
         source_ordinal,
         &plan.logic_fingerprint,
@@ -562,23 +559,14 @@ pub async fn update_source_entry(
 
 async fn update_source_entry_with_err_handling(
     plan: &ExecutionPlan,
-    source_op_idx: usize,
+    source_op: &AnalyzedSourceOp,
     schema: &schema::DataSchema,
     key: &value::KeyValue,
     only_for_deletion: bool,
     pool: &PgPool,
     stats: &UpdateStats,
 ) {
-    let r = update_source_entry(
-        plan,
-        source_op_idx,
-        schema,
-        key,
-        only_for_deletion,
-        pool,
-        stats,
-    )
-    .await;
+    let r = update_source_entry(plan, source_op, schema, key, only_for_deletion, pool, stats).await;
     if let Err(e) = r {
         stats.num_errors.fetch_add(1, Relaxed);
         error!("{:?}", e.context("Error in indexing a source row"));
@@ -588,11 +576,10 @@ async fn update_source_entry_with_err_handling(
 async fn update_source(
     source_name: &str,
     plan: &ExecutionPlan,
-    source_op_idx: usize,
+    source_op: &AnalyzedSourceOp,
     schema: &schema::DataSchema,
     pool: &PgPool,
 ) -> Result<SourceUpdateInfo> {
-    let source_op = &plan.source_ops[source_op_idx];
     let (keys, existing_keys_json) = try_join(
         source_op.executor.list_keys(),
         db_tracking::list_source_tracking_keys(
@@ -605,7 +592,7 @@ async fn update_source(
 
     let stats = UpdateStats::default();
     let upsert_futs = join_all(keys.iter().map(|key| {
-        update_source_entry_with_err_handling(plan, source_op_idx, schema, key, false, pool, &stats)
+        update_source_entry_with_err_handling(plan, source_op, schema, key, false, pool, &stats)
     }));
     let deleted_keys = existing_keys_json
         .into_iter()
@@ -619,7 +606,7 @@ async fn update_source(
         .filter_ok(|existing_key| !keys.contains(existing_key))
         .collect::<Result<Vec<_>>>()?;
     let delete_futs = join_all(deleted_keys.iter().map(|key| {
-        update_source_entry_with_err_handling(plan, source_op_idx, schema, key, true, pool, &stats)
+        update_source_entry_with_err_handling(plan, source_op, schema, key, true, pool, &stats)
     }));
     join(upsert_futs, delete_futs).await;
 
@@ -630,17 +617,15 @@ async fn update_source(
 }
 
 pub async fn update(
-    spec: &FlowInstanceSpec,
     plan: &ExecutionPlan,
     schema: &schema::DataSchema,
     pool: &PgPool,
 ) -> Result<IndexUpdateInfo> {
     let source_update_stats = try_join_all(
-        spec.source_ops
+        plan.source_ops
             .iter()
-            .enumerate()
-            .map(|(source_op_idx, source_op)| async move {
-                update_source(source_op.name.as_str(), plan, source_op_idx, schema, pool).await
+            .map(|source_op| async move {
+                update_source(source_op.name.as_str(), plan, source_op, schema, pool).await
             })
             .collect::<Vec<_>>(),
     )

src/py/mod.rs

@@ -111,13 +111,8 @@ impl Flow {
                 .runtime
                 .block_on(async {
                     let exec_plan = self.0.get_execution_plan().await?;
-                    execution::indexer::update(
-                        &self.0.flow_instance,
-                        &exec_plan,
-                        &self.0.data_schema,
-                        &lib_context.pool,
-                    )
-                    .await
+                    execution::indexer::update(&exec_plan, &self.0.data_schema, &lib_context.pool)
+                        .await
                 })
                 .into_py_result()?;
             Ok(IndexUpdateInfo(update_info))

src/service/flows.rs

@@ -133,9 +133,9 @@ pub async fn evaluate_data(
                 StatusCode::BAD_REQUEST,
             )
         })?;
-    let execution_plan = fl.get_execution_plan().await?;
-    let field_schema =
-        &schema.fields[execution_plan.source_ops[source_op_idx].output.field_idx as usize];
+    let plan = fl.get_execution_plan().await?;
+    let source_op = &plan.source_ops[source_op_idx];
+    let field_schema = &schema.fields[source_op.output.field_idx as usize];
     let collection_schema = match &field_schema.value_type.typ {
         schema::ValueType::Collection(collection) => collection,
         _ => api_bail!("field is not a table: {}", query.field),
@@ -146,8 +146,8 @@ pub async fn evaluate_data(
     let key = value::KeyValue::from_strs(query.key, &key_field.value_type.typ)?;
 
     let data = indexer::evaluate_source_entry_with_cache(
-        &execution_plan,
-        source_op_idx,
+        &plan,
+        source_op,
         schema,
         &key,
         &lib_context.pool,
@@ -167,12 +167,6 @@ pub async fn update(
 ) -> Result<Json<IndexUpdateInfo>, ApiError> {
     let fl = &lib_context.with_flow_context(&flow_name, |ctx| ctx.flow.clone())?;
     let execution_plan = fl.get_execution_plan().await?;
-    let update_info = indexer::update(
-        &fl.flow_instance,
-        &execution_plan,
-        &fl.data_schema,
-        &lib_context.pool,
-    )
-    .await?;
+    let update_info = indexer::update(&execution_plan, &fl.data_schema, &lib_context.pool).await?;
     Ok(Json(update_info))
 }