update refresh-state

Author: JanKaul

datafusion_iceberg/src/materialized_view/mod.rs

@@ -14,7 +14,10 @@ use iceberg_rust::{
     arrow::write::{write_equality_deletes_parquet_partitioned, write_parquet_partitioned},
     catalog::{identifier::Identifier, tabular::Tabular, CatalogList},
     materialized_view::MaterializedView,
-    spec::materialized_view_metadata::{SourceTables, SourceViews},
+    spec::{
+        identifier::FullIdentifier,
+        materialized_view_metadata::{SourceState, SourceStates, SourceTable},
+    },
 };
 use iceberg_rust::{
     error::Error,
@@ -85,10 +88,15 @@ pub async fn refresh_materialized_view(
 
     let refresh_version_id = matview.metadata().current_version_id;
 
+    let refresh_start_timestamp_ms = std::time::SystemTime::now()
+        .duration_since(std::time::UNIX_EPOCH)
+        .map(|d| d.as_millis() as i64)
+        .unwrap_or(0);
+
     let refresh_state = RefreshState {
         refresh_version_id,
-        source_table_states: SourceTables(source_table_states),
-        source_view_states: SourceViews(HashMap::new()),
+        source_states: SourceStates(source_table_states),
+        refresh_start_timestamp_ms,
     };
 
     let storage_table_provider = Arc::new(DataFusionTable::new(
@@ -230,7 +238,7 @@ async fn get_source_tables(
                 Option<Arc<dyn TableProvider>>,
             ),
         >,
-        HashMap<(uuid::Uuid, Option<String>), i64>,
+        HashMap<FullIdentifier, SourceState>,
     ),
     Error,
 > {
@@ -246,10 +254,8 @@ async fn get_source_tables(
                     &version.default_namespace()[0],
                 );
                 let catalog_name = resolved_reference.catalog.to_string();
-                let identifier = Identifier::new(
-                    &[resolved_reference.schema.to_string()],
-                    &resolved_reference.table,
-                );
+                let namespace = [resolved_reference.schema.to_string()];
+                let identifier = Identifier::new(&namespace, &resolved_reference.table);
                 let catalog = catalog_list
                     .catalog(&catalog_name)
                     .ok_or(Error::NotFound(format!("Catalog {catalog_name}")))?;
@@ -288,13 +294,19 @@ async fn get_source_tables(
                     _ => Err(Error::InvalidFormat("storage table".to_string())),
                 }?;
 
+                let full_identifier =
+                    FullIdentifier::new(Some(&catalog_name), &namespace, &resolved_reference.table);
+
                 #[allow(clippy::type_complexity)]
                 let source_table_provider: (
                     Option<Arc<dyn TableProvider>>,
                     Option<Arc<dyn TableProvider>>,
                 ) = if let Some(old_refresh_state) = old_refresh_state.as_ref() {
-                    let revision_id = old_refresh_state.source_table_states.get(&(uuid, None));
-                    if Some(&current_snapshot_id) == revision_id {
+                    let revision_snapshot_id = old_refresh_state
+                        .source_states
+                        .get(&full_identifier)
+                        .and_then(|s| s.snapshot_id());
+                    if Some(current_snapshot_id) == revision_snapshot_id {
                         // Fresh
                         (
                             Some(Arc::new(DataFusionTable::new(
@@ -305,7 +317,7 @@ async fn get_source_tables(
                             ))),
                             None,
                         )
-                    } else if Some(&-1) == revision_id {
+                    } else if Some(-1) == revision_snapshot_id {
                         // Invalid
                         (
                             None,
@@ -316,18 +328,18 @@ async fn get_source_tables(
                                 branch.as_deref(),
                             ))),
                         )
-                    } else if let Some(revision_id) = revision_id {
+                    } else if let Some(revision_id) = revision_snapshot_id {
                         // Outdated
                         (
                             Some(Arc::new(DataFusionTable::new(
                                 tabular.clone(),
                                 None,
-                                Some(*revision_id),
+                                Some(revision_id),
                                 branch.as_deref(),
                             ))),
                             Some(Arc::new(DataFusionTable::new(
                                 tabular,
-                                Some(*revision_id),
+                                Some(revision_id),
                                 None,
                                 branch.as_deref(),
                             ))),
@@ -357,9 +369,18 @@ async fn get_source_tables(
                     )
                 };
 
+                let source_state = SourceState::Table(SourceTable::new(
+                    Some(&catalog_name),
+                    &namespace,
+                    &resolved_reference.table,
+                    uuid,
+                    current_snapshot_id,
+                    branch.clone(),
+                ));
+
                 Ok((
                     (reference, source_table_provider),
-                    ((uuid, None), current_snapshot_id),
+                    (full_identifier, source_state),
                 ))
             }
         })

iceberg-rust-spec/src/spec/materialized_view_metadata.rs

@@ -6,15 +6,17 @@
 //! The main types are:
 //! - [`MaterializedViewMetadata`]: The top-level metadata for a materialized view
 //! - [`RefreshState`]: Information about the last refresh operation
-//! - [`SourceTables`]: Collection of source table states
-//! - [`SourceViews`]: Collection of source view states
+//! - [`SourceStates`]: Collection of source states
 
 use std::{collections::HashMap, ops::Deref};
 
 use serde::{Deserialize, Serialize};
 use uuid::Uuid;
 
-use crate::identifier::Identifier;
+use crate::{
+    identifier::{FullIdentifier, Identifier},
+    namespace::Namespace,
+};
 
 use super::{
     tabular::TabularMetadataRef,
@@ -41,26 +43,25 @@ pub struct RefreshState {
     /// The version-id of the materialized view when the refresh operation was performed.
     pub refresh_version_id: i64,
     /// A map from sequence-id (as defined in the view lineage) to the source tables’ snapshot-id of when the last refresh operation was performed.
-    pub source_table_states: SourceTables,
-    /// A map from sequence-id (as defined in the view lineage) to the source views’ version-id of when the last refresh operation was performed.
-    pub source_view_states: SourceViews,
+    pub source_states: SourceStates,
+    // A timestamp of when the refresh operation was started
+    pub refresh_start_timestamp_ms: i64,
 }
 
-/// Represents a collection of source table states in a materialized view refresh
-///
-/// # Fields
-/// * `0` - A HashMap mapping (table UUID, optional reference) pairs to snapshot IDs
-#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
-#[serde(from = "Vec<SourceTable>", into = "Vec<SourceTable>")]
-pub struct SourceTables(pub HashMap<(Uuid, Option<String>), i64>);
-
 /// Represents a collection of source view states in a materialized view refresh
 ///
 /// # Fields
 /// * `0` - A HashMap mapping (table UUID, optional reference) pairs to version IDs
 #[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
-#[serde(from = "Vec<SourceView>", into = "Vec<SourceView>")]
-pub struct SourceViews(pub HashMap<(Uuid, Option<String>), i64>);
+#[serde(from = "Vec<SourceState>", into = "Vec<SourceState>")]
+pub struct SourceStates(pub HashMap<FullIdentifier, SourceState>);
+
+#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
+#[serde(rename_all = "kebab-case", tag = "type")]
+pub enum SourceState {
+    Table(SourceTable),
+    View(SourceView),
+}
 
 /// Represents a source table state in a materialized view refresh
 ///
@@ -71,8 +72,13 @@ pub struct SourceViews(pub HashMap<(Uuid, Option<String>), i64>);
 #[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
 #[serde(rename_all = "kebab-case")]
 pub struct SourceTable {
+    name: String,
+    namespace: Namespace,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    catalog: Option<String>,
     uuid: Uuid,
     snapshot_id: i64,
+    #[serde(skip_serializing_if = "Option::is_none")]
     r#ref: Option<String>,
 }
 
@@ -84,65 +90,77 @@ pub struct SourceTable {
 #[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
 #[serde(rename_all = "kebab-case")]
 pub struct SourceView {
+    name: String,
+    namespace: Namespace,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    catalog: Option<String>,
     uuid: Uuid,
     version_id: i64,
 }
 
-impl From<Vec<SourceTable>> for SourceTables {
-    fn from(value: Vec<SourceTable>) -> Self {
-        SourceTables(
-            value
-                .into_iter()
-                .map(|x| ((x.uuid, x.r#ref), x.snapshot_id))
-                .collect(),
-        )
+impl SourceTable {
+    /// Create a new SourceTable
+    pub fn new(
+        catalog: Option<&str>,
+        namespace: &[String],
+        name: &str,
+        uuid: Uuid,
+        snapshot_id: i64,
+        r#ref: Option<String>,
+    ) -> Self {
+        Self {
+            catalog: catalog.map(ToString::to_string),
+            namespace: Namespace(namespace.to_owned()),
+            name: name.to_owned(),
+            uuid,
+            snapshot_id,
+            r#ref,
+        }
     }
 }
 
-impl From<SourceTables> for Vec<SourceTable> {
-    fn from(value: SourceTables) -> Self {
-        value
-            .0
-            .into_iter()
-            .map(|((uuid, r#ref), snapshot_id)| SourceTable {
-                uuid,
-                snapshot_id,
-                r#ref,
-            })
-            .collect()
+impl SourceState {
+    /// Returns the snapshot_id for Table states, None for View states
+    pub fn snapshot_id(&self) -> Option<i64> {
+        match self {
+            SourceState::Table(t) => Some(t.snapshot_id),
+            SourceState::View(_) => None,
+        }
     }
 }
 
-impl From<Vec<SourceView>> for SourceViews {
-    fn from(value: Vec<SourceView>) -> Self {
-        SourceViews(
+impl From<Vec<SourceState>> for SourceStates {
+    fn from(value: Vec<SourceState>) -> Self {
+        SourceStates(
             value
                 .into_iter()
-                .map(|x| ((x.uuid, None), x.version_id))
+                .map(|x| match &x {
+                    SourceState::Table(table) => (
+                        FullIdentifier::new(
+                            table.catalog.as_deref(),
+                            &table.namespace,
+                            &table.name,
+                        ),
+                        x,
+                    ),
+                    SourceState::View(view) => (
+                        FullIdentifier::new(view.catalog.as_deref(), &view.namespace, &view.name),
+                        x,
+                    ),
+                })
                 .collect(),
         )
     }
 }
 
-impl From<SourceViews> for Vec<SourceView> {
-    fn from(value: SourceViews) -> Self {
-        value
-            .0
-            .into_iter()
-            .map(|((uuid, _), version_id)| SourceView { uuid, version_id })
-            .collect()
-    }
-}
-
-impl Deref for SourceTables {
-    type Target = HashMap<(Uuid, Option<String>), i64>;
-    fn deref(&self) -> &Self::Target {
-        &self.0
+impl From<SourceStates> for Vec<SourceState> {
+    fn from(value: SourceStates) -> Self {
+        value.0.into_values().collect()
     }
 }
 
-impl Deref for SourceViews {
-    type Target = HashMap<(Uuid, Option<String>), i64>;
+impl Deref for SourceStates {
+    type Target = HashMap<FullIdentifier, SourceState>;
     fn deref(&self) -> &Self::Target {
         &self.0
     }
@@ -181,7 +199,6 @@ mod tests {
             "dialect" : "spark"
             } ],
             "storage-table": {
-                "catalog": "prod",
                 "namespace": ["default"],
                 "name": "event_agg_storage"
             }