feat: merge schema

crates/core/src/operations/delete.rs

@@ -17,6 +17,7 @@
 //!     .await?;
 //! ````
 
+use core::panic;
 use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::{Instant, SystemTime, UNIX_EPOCH};
@@ -36,6 +37,7 @@ use serde_json::Value;
 
 use super::datafusion_utils::Expression;
 use super::transaction::PROTOCOL;
+use super::write::SchemaWriteMode;
 use crate::delta_datafusion::expr::fmt_expr_to_sql;
 use crate::delta_datafusion::{find_files, register_store, DeltaScanBuilder, DeltaSessionContext};
 use crate::errors::DeltaResult;
@@ -167,9 +169,15 @@ async fn excute_non_empty_expr(
         None,
         writer_properties,
         false,
-        false,
+        SchemaWriteMode::None,
     )
-    .await?;
+    .await?
+    .into_iter()
+    .map(|a| match a {
+        Action::Add(a) => a,
+        _ => panic!("Expected Add action"),
+    })
+    .collect::<Vec<Add>>();
 
     let read_records = scan.parquet_scan.metrics().and_then(|m| m.output_rows());
     let filter_records = filter.metrics().and_then(|m| m.output_rows());

crates/core/src/operations/merge/mod.rs

@@ -74,7 +74,7 @@ use crate::delta_datafusion::{
 use crate::kernel::Action;
 use crate::logstore::LogStoreRef;
 use crate::operations::merge::barrier::find_barrier_node;
-use crate::operations::write::write_execution_plan;
+use crate::operations::write::{write_execution_plan, SchemaWriteMode};
 use crate::protocol::{DeltaOperation, MergePredicate};
 use crate::table::state::DeltaTableState;
 use crate::{DeltaResult, DeltaTable, DeltaTableError};
@@ -1379,13 +1379,13 @@ async fn execute(
         None,
         writer_properties,
         safe_cast,
-        false,
+        SchemaWriteMode::None,
     )
     .await?;
 
     metrics.rewrite_time_ms = Instant::now().duration_since(rewrite_start).as_millis() as u64;
 
-    let mut actions: Vec<Action> = add_actions.into_iter().map(Action::Add).collect();
+    let mut actions: Vec<Action> = add_actions.clone();
     metrics.num_target_files_added = actions.len();
 
     let survivors = barrier

crates/core/src/operations/update.rs

@@ -43,9 +43,9 @@ use parquet::file::properties::WriterProperties;
 use serde::Serialize;
 use serde_json::Value;
 
-use super::datafusion_utils::Expression;
 use super::transaction::{commit, PROTOCOL};
 use super::write::write_execution_plan;
+use super::{datafusion_utils::Expression, write::SchemaWriteMode};
 use crate::delta_datafusion::{
     expr::fmt_expr_to_sql, physical::MetricObserverExec, DeltaColumn, DeltaSessionContext,
 };
@@ -357,7 +357,7 @@ async fn execute(
         None,
         writer_properties,
         safe_cast,
-        false,
+        SchemaWriteMode::None,
     )
     .await?;
 
@@ -377,7 +377,7 @@ async fn execute(
         .duration_since(UNIX_EPOCH)
         .unwrap()
         .as_millis() as i64;
-    let mut actions: Vec<Action> = add_actions.into_iter().map(Action::Add).collect();
+    let mut actions: Vec<Action> = add_actions.clone();
 
     metrics.num_added_files = actions.len();
     metrics.num_removed_files = candidates.candidates.len();

crates/core/src/operations/write.rs

@@ -25,6 +25,7 @@
 //! ````
 
 use std::collections::HashMap;
+use std::str::FromStr;
 use std::sync::Arc;
 use std::time::{SystemTime, UNIX_EPOCH};
 
@@ -50,7 +51,7 @@ use crate::delta_datafusion::expr::parse_predicate_expression;
 use crate::delta_datafusion::DeltaDataChecker;
 use crate::delta_datafusion::{find_files, register_store, DeltaScanBuilder};
 use crate::errors::{DeltaResult, DeltaTableError};
-use crate::kernel::{Action, Add, PartitionsExt, Remove, StructType};
+use crate::kernel::{Action, Add, Metadata, PartitionsExt, Remove, StructType};
 use crate::logstore::LogStoreRef;
 use crate::protocol::{DeltaOperation, SaveMode};
 use crate::storage::ObjectStoreRef;
@@ -87,6 +88,33 @@ impl From<WriteError> for DeltaTableError {
     }
 }
 
+///Specifies how to handle schema drifts
+#[derive(PartialEq)]
+pub enum SchemaWriteMode {
+    /// Use existing schema and fail if it does not match the new schema
+    None,
+    /// Overwrite the schema with the new schema
+    Overwrite,
+    /// Append the new schema to the existing schema
+    Merge,
+}
+
+impl FromStr for SchemaWriteMode {
+    type Err = DeltaTableError;
+
+    fn from_str(s: &str) -> DeltaResult<Self> {
+        match s.to_ascii_lowercase().as_str() {
+            "none" => Ok(SchemaWriteMode::None),
+            "overwrite" => Ok(SchemaWriteMode::Overwrite),
+            "merge" => Ok(SchemaWriteMode::Merge),
+            _ => Err(DeltaTableError::Generic(format!(
+                "Invalid schema write mode provided: {}, only these are supported: ['none', 'overwrite', 'merge']",
+                s
+            ))),
+        }
+    }
+}
+
 /// Write data into a DeltaTable
 pub struct WriteBuilder {
     /// A snapshot of the to-be-loaded table's state
@@ -109,8 +137,8 @@ pub struct WriteBuilder {
     write_batch_size: Option<usize>,
     /// RecordBatches to be written into the table
     batches: Option<Vec<RecordBatch>>,
-    /// whether to overwrite the schema
-    overwrite_schema: bool,
+    /// whether to overwrite the schema or to merge it
+    schema_write_mode: SchemaWriteMode,
     /// how to handle cast failures, either return NULL (safe=true) or return ERR (safe=false)
     safe_cast: bool,
     /// Parquet writer properties
@@ -140,7 +168,7 @@ impl WriteBuilder {
             write_batch_size: None,
             batches: None,
             safe_cast: false,
-            overwrite_schema: false,
+            schema_write_mode: SchemaWriteMode::None,
             writer_properties: None,
             app_metadata: None,
             name: None,
@@ -155,9 +183,9 @@ impl WriteBuilder {
         self
     }
 
-    /// Add overwrite_schema
-    pub fn with_overwrite_schema(mut self, overwrite_schema: bool) -> Self {
-        self.overwrite_schema = overwrite_schema;
+    /// Add Schema Write Mode
+    pub fn with_schema_write_mode(mut self, schema_write_mode: SchemaWriteMode) -> Self {
+        self.schema_write_mode = schema_write_mode;
         self
     }
 
@@ -311,11 +339,37 @@ async fn write_execution_plan_with_predicate(
     write_batch_size: Option<usize>,
     writer_properties: Option<WriterProperties>,
     safe_cast: bool,
-    overwrite_schema: bool,
-) -> DeltaResult<Vec<Add>> {
+    schema_write_mode: SchemaWriteMode,
+) -> DeltaResult<Vec<Action>> {
+    let mut schema_action: Option<Action> = None;
     // Use input schema to prevent wrapping partitions columns into a dictionary.
-    let schema: ArrowSchemaRef = if overwrite_schema {
+    let schema: ArrowSchemaRef = if schema_write_mode == SchemaWriteMode::Overwrite {
         plan.schema()
+    } else if schema_write_mode == SchemaWriteMode::Merge {
+        let original_schema = snapshot
+            .and_then(|s| s.input_schema().ok())
+            .unwrap_or(plan.schema());
+        if original_schema == plan.schema() {
+            original_schema
+        } else {
+            let new_schema = Arc::new(arrow_schema::Schema::try_merge(vec![
+                original_schema.as_ref().clone(),
+                plan.schema().as_ref().clone(),
+            ])?);
+            let schema_struct: StructType = new_schema.clone().try_into()?;
+            schema_action = Some(Action::Metadata(Metadata::try_new(
+                schema_struct,
+                match snapshot {
+                    Some(sn) => sn.metadata().partition_columns.clone(),
+                    None => vec![],
+                },
+                match snapshot {
+                    Some(sn) => sn.metadata().configuration.clone(),
+                    None => HashMap::new(),
+                },
+            )?));
+            new_schema
+        }
     } else {
         snapshot
             .and_then(|s| s.input_schema().ok())
@@ -352,7 +406,7 @@ async fn write_execution_plan_with_predicate(
         let mut writer = DeltaWriter::new(object_store.clone(), config);
         let checker_stream = checker.clone();
         let mut stream = inner_plan.execute(i, task_ctx)?;
-        let handle: tokio::task::JoinHandle<DeltaResult<Vec<Add>>> =
+        let handle: tokio::task::JoinHandle<DeltaResult<Vec<Action>>> =
             tokio::task::spawn(async move {
                 while let Some(maybe_batch) = stream.next().await {
                     let batch = maybe_batch?;
@@ -361,14 +415,16 @@ async fn write_execution_plan_with_predicate(
                         super::cast::cast_record_batch(&batch, inner_schema.clone(), safe_cast)?;
                     writer.write(&arr).await?;
                 }
-                writer.close().await
+                let add_actions = writer.close().await;
+                match add_actions {
+                    Ok(actions) => Ok(actions.into_iter().map(Action::Add).collect::<Vec<_>>()),
+                    Err(err) => Err(err),
+                }
             });
 
         tasks.push(handle);
     }
-
-    // Collect add actions to add to commit
-    Ok(futures::future::join_all(tasks)
+    let mut actions = futures::future::join_all(tasks)
         .await
         .into_iter()
         .collect::<Result<Vec<_>, _>>()
@@ -377,7 +433,12 @@ async fn write_execution_plan_with_predicate(
         .collect::<Result<Vec<_>, _>>()?
         .concat()
         .into_iter()
-        .collect::<Vec<_>>())
+        .collect::<Vec<_>>();
+    if let Some(schema_action) = schema_action {
+        actions.push(schema_action);
+    }
+    // Collect add actions to add to commit
+    Ok(actions)
 }
 
 #[allow(clippy::too_many_arguments)]
@@ -391,8 +452,8 @@ pub(crate) async fn write_execution_plan(
     write_batch_size: Option<usize>,
     writer_properties: Option<WriterProperties>,
     safe_cast: bool,
-    overwrite_schema: bool,
-) -> DeltaResult<Vec<Add>> {
+    schema_write_mode: SchemaWriteMode,
+) -> DeltaResult<Vec<Action>> {
     write_execution_plan_with_predicate(
         None,
         snapshot,
@@ -404,7 +465,7 @@ pub(crate) async fn write_execution_plan(
         write_batch_size,
         writer_properties,
         safe_cast,
-        overwrite_schema,
+        schema_write_mode,
     )
     .await
 }
@@ -417,7 +478,7 @@ async fn execute_non_empty_expr(
     expression: &Expr,
     rewrite: &[Add],
     writer_properties: Option<WriterProperties>,
-) -> DeltaResult<Vec<Add>> {
+) -> DeltaResult<Vec<Action>> {
     // For each identified file perform a parquet scan + filter + limit (1) + count.
     // If returned count is not zero then append the file to be rewritten and removed from the log. Otherwise do nothing to the file.
 
@@ -452,7 +513,7 @@ async fn execute_non_empty_expr(
         None,
         writer_properties,
         false,
-        false,
+        SchemaWriteMode::None,
     )
     .await?;
 
@@ -488,7 +549,7 @@ async fn prepare_predicate_actions(
     };
     let remove = candidates.candidates;
 
-    let mut actions: Vec<Action> = add.into_iter().map(Action::Add).collect();
+    let mut actions: Vec<Action> = add.into_iter().collect();
 
     for action in remove {
         actions.push(Action::Remove(Remove {
@@ -563,7 +624,8 @@ impl std::future::IntoFuture for WriteBuilder {
                             .unwrap_or(schema.clone());
 
                         if !can_cast_batch(schema.fields(), table_schema.fields())
-                            && !(this.overwrite_schema && matches!(this.mode, SaveMode::Overwrite))
+                            && (this.schema_write_mode == SchemaWriteMode::None
+                                && !matches!(this.mode, SaveMode::Overwrite))
                         {
                             return Err(DeltaTableError::Generic(
                                 "Schema of data does not match table schema".to_string(),
@@ -641,10 +703,10 @@ impl std::future::IntoFuture for WriteBuilder {
                 this.write_batch_size,
                 this.writer_properties.clone(),
                 this.safe_cast,
-                this.overwrite_schema,
+                this.schema_write_mode,
             )
             .await?;
-            actions.extend(add_actions.into_iter().map(Action::Add));
+            actions.extend(add_actions);
 
             // Collect remove actions if we are overwriting the table
             if let Some(snapshot) = &this.snapshot {

python/deltalake/_internal.pyi

@@ -174,7 +174,7 @@ def write_to_deltalake(
     partition_by: Optional[List[str]],
     mode: str,
     max_rows_per_group: int,
-    overwrite_schema: bool,
+    schema_write_mode: Optional[str],
     predicate: Optional[str],
     name: Optional[str],
     description: Optional[str],

python/deltalake/writer.py

@@ -49,7 +49,7 @@
     convert_pyarrow_table,
 )
 from .table import MAX_SUPPORTED_WRITER_VERSION, DeltaTable, WriterProperties
-
+import warnings
 try:
     import pandas as pd  # noqa: F811
 except ModuleNotFoundError:
@@ -185,6 +185,7 @@ def write_deltalake(
     description: Optional[str] = None,
     configuration: Optional[Mapping[str, Optional[str]]] = None,
     overwrite_schema: bool = False,
+    schema_write_mode: Literal["none", "merge", "overwrite"] = "none",
     storage_options: Optional[Dict[str, str]] = None,
     partition_filters: Optional[List[Tuple[str, str, Any]]] = None,
     predicate: Optional[str] = None,
@@ -238,7 +239,8 @@ def write_deltalake(
         name: User-provided identifier for this table.
         description: User-provided description for this table.
         configuration: A map containing configuration options for the metadata action.
-        overwrite_schema: If True, allows updating the schema of the table.
+        overwrite_schema: Deprecated, use schema_write_mode instead.
+        schema_write_mode: If set to "overwrite", allows replacing the schema of the table. Set to "merge" to merge with existing schema.
         storage_options: options passed to the native delta filesystem.
         predicate: When using `Overwrite` mode, replace data that matches a predicate. Only used in rust engine.
         partition_filters: the partition filters that will be used for partition overwrite. Only used in pyarrow engine.
@@ -256,7 +258,15 @@ def write_deltalake(
         table.update_incremental()
 
     __enforce_append_only(table=table, configuration=configuration, mode=mode)
-
+    if overwrite_schema:
+        assert schema_write_mode in ["none", "overwrite"] # none is default, overwrite would at least match
+        schema_write_mode = "overwrite"
+
+        warnings.warn(
+            "overwrite_schema is deprecated, use schema_write_mode instead. ",
+            category=DeprecationWarning,
+            stacklevel=2,
+        )
     if isinstance(partition_by, str):
         partition_by = [partition_by]
 
@@ -302,7 +312,7 @@ def write_deltalake(
             partition_by=partition_by,
             mode=mode,
             max_rows_per_group=max_rows_per_group,
-            overwrite_schema=overwrite_schema,
+            schema_write_mode=schema_write_mode,
             predicate=predicate,
             name=name,
             description=description,
@@ -327,7 +337,7 @@ def sort_arrow_schema(schema: pa.schema) -> pa.schema:
         if table:  # already exists
             if sort_arrow_schema(schema) != sort_arrow_schema(
                 table.schema().to_pyarrow(as_large_types=large_dtypes)
-            ) and not (mode == "overwrite" and overwrite_schema):
+            ) and not (mode == "overwrite" and schema_write_mode == "overwrite"):
                 raise ValueError(
                     "Schema of data does not match table schema\n"
                     f"Data schema:\n{schema}\nTable Schema:\n{table.schema().to_pyarrow(as_large_types=large_dtypes)}"