feat: allow to include values in source's list() calls

Author: georgeh0

src/execution/dumper.rs

@@ -12,7 +12,7 @@ use super::memoization::EvaluationMemoryOptions;
 use super::row_indexer;
 use crate::base::{schema, value};
 use crate::builder::plan::{AnalyzedImportOp, ExecutionPlan};
-use crate::ops::interface::SourceExecutorListOptions;
+use crate::ops::interface::SourceExecutorReadOptions;
 use crate::utils::yaml_ser::YamlSerializer;
 
 #[derive(Debug, Clone, Deserialize)]
@@ -193,9 +193,10 @@ impl<'a> Dumper<'a> {
 
         let mut rows_stream = import_op
             .executor
-            .list(&SourceExecutorListOptions {
+            .list(&SourceExecutorReadOptions {
                 include_ordinal: false,
                 include_content_version_fp: false,
+                include_value: false,
             })
             .await?;
         while let Some(rows) = rows_stream.next().await {

src/execution/indexing_status.rs

@@ -38,7 +38,7 @@ pub async fn get_source_row_indexing_status(
     let current_fut = src_eval_ctx.import_op.executor.get_value(
         src_eval_ctx.key,
         key_aux_info,
-        &interface::SourceExecutorGetOptions {
+        &interface::SourceExecutorReadOptions {
             include_value: false,
             include_ordinal: true,
             include_content_version_fp: false,

src/execution/live_updater.rs

@@ -1,5 +1,5 @@
 use crate::{
-    execution::{source_indexer::ProcessSourceKeyInput, stats::UpdateStats},
+    execution::{source_indexer::ProcessSourceRowInput, stats::UpdateStats},
     prelude::*,
 };
 
@@ -192,18 +192,18 @@ impl SourceUpdateTask {
                                         .concurrency_controller
                                         .acquire(concur_control::BYTES_UNKNOWN_YET)
                                         .await?;
-                                    tokio::spawn(source_context.clone().process_source_key(
-                                        change.key,
+                                    tokio::spawn(source_context.clone().process_source_row(
+                                        ProcessSourceRowInput {
+                                            key: change.key,
+                                            key_aux_info: Some(change.key_aux_info),
+                                            data: change.data,
+                                        },
                                         update_stats.clone(),
                                         concur_permit,
                                         Some(move || async move {
                                             SharedAckFn::ack(&shared_ack_fn).await
                                         }),
                                         pool.clone(),
-                                        ProcessSourceKeyInput {
-                                            key_aux_info: Some(change.key_aux_info),
-                                            data: change.data,
-                                        },
                                     ));
                                 }
                             }

src/execution/row_indexer.rs

@@ -16,7 +16,7 @@ use super::stats;
 use crate::base::value::{self, FieldValues, KeyValue};
 use crate::builder::plan::*;
 use crate::ops::interface::{
-    ExportTargetMutation, ExportTargetUpsertEntry, Ordinal, SourceExecutorGetOptions,
+    ExportTargetMutation, ExportTargetUpsertEntry, Ordinal, SourceExecutorReadOptions,
 };
 use crate::utils::db::WriteAction;
 use crate::utils::fingerprint::{Fingerprint, Fingerprinter};
@@ -841,7 +841,7 @@ pub async fn evaluate_source_entry_with_memory(
         .get_value(
             src_eval_ctx.key,
             key_aux_info,
-            &SourceExecutorGetOptions {
+            &SourceExecutorReadOptions {
                 include_value: true,
                 include_ordinal: false,
                 include_content_version_fp: false,

src/execution/source_indexer.rs

@@ -166,7 +166,8 @@ impl<'a> LocalSourceRowStateOperator<'a> {
     }
 }
 
-pub struct ProcessSourceKeyInput {
+pub struct ProcessSourceRowInput {
+    pub key: value::FullKeyValue,
     /// `key_aux_info` is not available for deletions. It must be provided if `data.value` is `None`.
     pub key_aux_info: Option<serde_json::Value>,
     pub data: interface::PartialSourceRowData,
@@ -224,17 +225,16 @@ impl SourceIndexingContext {
         })
     }
 
-    pub async fn process_source_key<
+    pub async fn process_source_row<
         AckFut: Future<Output = Result<()>> + Send + 'static,
         AckFn: FnOnce() -> AckFut,
     >(
         self: Arc<Self>,
-        key: value::FullKeyValue,
+        row_input: ProcessSourceRowInput,
         update_stats: Arc<stats::UpdateStats>,
         _concur_permit: concur_control::CombinedConcurrencyControllerPermit,
         ack_fn: Option<AckFn>,
         pool: PgPool,
-        inputs: ProcessSourceKeyInput,
     ) {
         let process = async {
             let plan = self.flow.get_execution_plan().await?;
@@ -245,7 +245,7 @@ impl SourceIndexingContext {
                 plan: &plan,
                 import_op,
                 schema,
-                key: &key,
+                key: &row_input.key,
                 import_op_idx: self.source_idx,
             };
             let mut row_indexer = row_indexer::RowIndexer::new(
@@ -256,9 +256,9 @@ impl SourceIndexingContext {
             )?;
 
             let mut row_state_operator =
-                LocalSourceRowStateOperator::new(&key, &self.state, &update_stats);
+                LocalSourceRowStateOperator::new(&row_input.key, &self.state, &update_stats);
 
-            let source_data = inputs.data;
+            let source_data = row_input.data;
             if let Some(ordinal) = source_data.ordinal
                 && let Some(content_version_fp) = &source_data.content_version_fp
             {
@@ -295,22 +295,22 @@ impl SourceIndexingContext {
                 }
             }
 
-            let (ordinal, value, content_version_fp) =
+            let (ordinal, content_version_fp, value) =
                 match (source_data.ordinal, source_data.value) {
                     (Some(ordinal), Some(value)) => {
-                        (ordinal, value, source_data.content_version_fp)
+                        (ordinal, source_data.content_version_fp, value)
                     }
                     _ => {
                         let data = import_op
                         .executor
                         .get_value(
-                            &key,
-                            inputs.key_aux_info.as_ref().ok_or_else(|| {
+                            &row_input.key,
+                            row_input.key_aux_info.as_ref().ok_or_else(|| {
                                 anyhow::anyhow!(
                                     "`key_aux_info` must be provided when there's no `source_data`"
                                 )
                             })?,
-                            &interface::SourceExecutorGetOptions {
+                            &interface::SourceExecutorReadOptions {
                                 include_value: true,
                                 include_ordinal: true,
                                 include_content_version_fp: true,
@@ -320,9 +320,9 @@ impl SourceIndexingContext {
                         (
                             data.ordinal
                                 .ok_or_else(|| anyhow::anyhow!("ordinal is not available"))?,
+                            data.content_version_fp,
                             data.value
                                 .ok_or_else(|| anyhow::anyhow!("value is not available"))?,
-                            data.content_version_fp,
                         )
                     }
                 };
@@ -356,7 +356,8 @@ impl SourceIndexingContext {
                 "{:?}",
                 e.context(format!(
                     "Error in processing row from source `{source}` with key: {key}",
-                    source = self.flow.flow_instance.import_ops[self.source_idx].name
+                    source = self.flow.flow_instance.import_ops[self.source_idx].name,
+                    key = row_input.key,
                 ))
             );
         }
@@ -410,9 +411,10 @@ impl SourceIndexingContext {
         let import_op = &plan.import_ops[self.source_idx];
         let rows_stream = import_op
             .executor
-            .list(&interface::SourceExecutorListOptions {
+            .list(&interface::SourceExecutorReadOptions {
                 include_ordinal: true,
                 include_content_version_fp: true,
+                include_value: true,
             })
             .await?;
         self.update_with_stream(import_op, rows_stream, pool, update_stats)
@@ -422,7 +424,7 @@ impl SourceIndexingContext {
     async fn update_with_stream(
         self: &Arc<Self>,
         import_op: &plan::AnalyzedImportOp,
-        mut rows_stream: BoxStream<'_, Result<Vec<interface::PartialSourceRowMetadata>>>,
+        mut rows_stream: BoxStream<'_, Result<Vec<interface::PartialSourceRow>>>,
         pool: &PgPool,
         update_stats: &Arc<stats::UpdateStats>,
     ) -> Result<()> {
@@ -435,7 +437,8 @@ impl SourceIndexingContext {
         while let Some(row) = rows_stream.next().await {
             for row in row? {
                 let source_version = SourceVersion::from_current_with_ordinal(
-                    row.ordinal
+                    row.data
+                        .ordinal
                         .ok_or_else(|| anyhow::anyhow!("ordinal is not available"))?,
                 );
                 {
@@ -454,20 +457,16 @@ impl SourceIndexingContext {
                     .concurrency_controller
                     .acquire(concur_control::BYTES_UNKNOWN_YET)
                     .await?;
-                join_set.spawn(self.clone().process_source_key(
-                    row.key,
+                join_set.spawn(self.clone().process_source_row(
+                    ProcessSourceRowInput {
+                        key: row.key,
+                        key_aux_info: Some(row.key_aux_info),
+                        data: row.data,
+                    },
                     update_stats.clone(),
                     concur_permit,
                     NO_ACK,
                     pool.clone(),
-                    ProcessSourceKeyInput {
-                        key_aux_info: Some(row.key_aux_info),
-                        data: interface::PartialSourceRowData {
-                            value: None,
-                            ordinal: Some(source_version.ordinal),
-                            content_version_fp: row.content_version_fp,
-                        },
-                    },
                 ));
             }
         }
@@ -491,20 +490,20 @@ impl SourceIndexingContext {
         };
         for (key, source_ordinal) in deleted_key_versions {
             let concur_permit = import_op.concurrency_controller.acquire(Some(|| 0)).await?;
-            join_set.spawn(self.clone().process_source_key(
-                key,
-                update_stats.clone(),
-                concur_permit,
-                NO_ACK,
-                pool.clone(),
-                ProcessSourceKeyInput {
+            join_set.spawn(self.clone().process_source_row(
+                ProcessSourceRowInput {
+                    key,
                     key_aux_info: None,
                     data: interface::PartialSourceRowData {
-                        value: Some(interface::SourceValue::NonExistence),
                         ordinal: Some(source_ordinal),
                         content_version_fp: None,
+                        value: Some(interface::SourceValue::NonExistence),
                     },
                 },
+                update_stats.clone(),
+                concur_permit,
+                NO_ACK,
+                pool.clone(),
             ));
         }
         while let Some(result) = join_set.join_next().await {

src/ops/interface.rs

@@ -48,13 +48,14 @@ impl<TZ: TimeZone> TryFrom<chrono::DateTime<TZ>> for Ordinal {
     }
 }
 
-pub struct PartialSourceRowMetadata {
-    pub key: FullKeyValue,
-    /// Auxiliary information for the source row, to be used when reading the content.
-    /// e.g. it can be used to uniquely identify version of the row.
-    /// Use serde_json::Value::Null to represent no auxiliary information.
-    pub key_aux_info: serde_json::Value,
+#[derive(Debug)]
+pub enum SourceValue {
+    Existence(FieldValues),
+    NonExistence,
+}
 
+#[derive(Debug, Default)]
+pub struct PartialSourceRowData {
     pub ordinal: Option<Ordinal>,
 
     /// A content version fingerprint can be anything that changes when the content of the row changes.
@@ -64,12 +65,18 @@ pub struct PartialSourceRowMetadata {
     /// It's optional. The source shouldn't use generic way to compute it, e.g. computing a hash of the content.
     /// The framework will do so. If there's no fast way to get it from the source, leave it as `None`.
     pub content_version_fp: Option<Vec<u8>>,
+
+    pub value: Option<SourceValue>,
 }
 
-#[derive(Debug)]
-pub enum SourceValue {
-    Existence(FieldValues),
-    NonExistence,
+pub struct PartialSourceRow {
+    pub key: FullKeyValue,
+    /// Auxiliary information for the source row, to be used when reading the content.
+    /// e.g. it can be used to uniquely identify version of the row.
+    /// Use serde_json::Value::Null to represent no auxiliary information.
+    pub key_aux_info: serde_json::Value,
+
+    pub data: PartialSourceRowData,
 }
 
 impl SourceValue {
@@ -108,39 +115,38 @@ pub struct SourceChangeMessage {
 }
 
 #[derive(Debug, Default)]
-pub struct SourceExecutorListOptions {
+pub struct SourceExecutorReadOptions {
+    /// When set to true, the implementation must return a non-None `ordinal`.
     pub include_ordinal: bool,
-    pub include_content_version_fp: bool,
-}
 
-#[derive(Debug, Default)]
-pub struct SourceExecutorGetOptions {
-    pub include_ordinal: bool,
-    pub include_value: bool,
+    /// When set to true, the implementation has the discretion to decide whether or not to return a non-None `content_version_fp`.
+    /// The guideline is to return it only if it's very efficient to get it.
+    /// If it's returned in `list()`, it must be returned in `get_value()`.
     pub include_content_version_fp: bool,
-}
 
-#[derive(Debug, Default)]
-pub struct PartialSourceRowData {
-    pub value: Option<SourceValue>,
-    pub ordinal: Option<Ordinal>,
-    pub content_version_fp: Option<Vec<u8>>,
+    /// For get calls, when set to true, the implementation must return a non-None `value`.
+    ///
+    /// For list calls, when set to true, the implementation has the discretion to decide whether or not to include it.
+    /// The guideline is to only include it if a single "list() with content" call is significantly more efficient than "list() without content + series of get_value()" calls.
+    ///
+    /// Even if `list()` already returns `value` when it's true, `get_value()` must still return `value` when it's true.
+    pub include_value: bool,
 }
 
 #[async_trait]
 pub trait SourceExecutor: Send + Sync {
     /// Get the list of keys for the source.
     async fn list(
         &self,
-        options: &SourceExecutorListOptions,
-    ) -> Result<BoxStream<'async_trait, Result<Vec<PartialSourceRowMetadata>>>>;
+        options: &SourceExecutorReadOptions,
+    ) -> Result<BoxStream<'async_trait, Result<Vec<PartialSourceRow>>>>;
 
     // Get the value for the given key.
     async fn get_value(
         &self,
         key: &FullKeyValue,
         key_aux_info: &serde_json::Value,
-        options: &SourceExecutorGetOptions,
+        options: &SourceExecutorReadOptions,
     ) -> Result<PartialSourceRowData>;
 
     async fn change_stream(