Batch UPDATE statements using staging table and MERGE INTO for Databricks (#194)

Author: sgrebnov

.github/workflows/run_spicebench.yml

@@ -309,6 +309,7 @@ jobs:
 
           if [ "${SYSTEM_UNDER_TEST_PREFIX}" = "databricks" ]; then
             export SPICEBENCH_ADBC_DELETE_BATCH_SIZE=5000
+            export SPICEBENCH_ADBC_UPDATE_STRATEGY=staging_table
             ADAPTER_CMD="${HOME}/.spice/bin/databricks-system-adapter"
             ADAPTER_ARGS="stdio"
             ADAPTER_ENVS="--system-adapter-env DATABRICKS_ENDPOINT=${DATABRICKS_ENDPOINT} --system-adapter-env DATABRICKS_TOKEN=${DATABRICKS_TOKEN} --system-adapter-env DATABRICKS_HTTP_PATH=${DATABRICKS_HTTP_PATH} --system-adapter-env DATABRICKS_SQL_WAREHOUSE_ID=${DATABRICKS_SQL_WAREHOUSE_ID} --system-adapter-env DATABRICKS_TABLE_FORMAT=${DATABRICKS_TABLE_FORMAT}"

crates/etl/src/sink/adbc.rs

@@ -42,6 +42,36 @@ const ADBC_SINK_POOL_SIZE_ENV: &str = "SPICEBENCH_ADBC_SINK_POOL_SIZE";
 /// When unset or empty, deletes use the original row-by-row approach.
 const ADBC_DELETE_BATCH_SIZE_ENV: &str = "SPICEBENCH_ADBC_DELETE_BATCH_SIZE";
 
+/// Controls how UPDATE operations are executed.
+///
+/// - `statement`     — row-by-row `UPDATE … SET … WHERE …` statements (default)
+/// - `staging_table` — bulk ingest into temp staging table + single `MERGE INTO`
+const ADBC_UPDATE_STRATEGY_ENV: &str = "SPICEBENCH_ADBC_UPDATE_STRATEGY";
+
+/// Strategy for executing UPDATE operations.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum UpdateStrategy {
+    /// Row-by-row `UPDATE … SET … WHERE …` SQL statements.
+    Statement,
+    /// Bulk ingest into a temporary staging table, then `MERGE INTO target USING staging …`.
+    StagingTable,
+}
+
+impl UpdateStrategy {
+    fn from_env() -> anyhow::Result<Self> {
+        match std::env::var(ADBC_UPDATE_STRATEGY_ENV).ok().as_deref() {
+            Some(val) => match val.to_lowercase().as_str() {
+                "statement" => Ok(Self::Statement),
+                "staging_table" => Ok(Self::StagingTable),
+                other => anyhow::bail!(
+                    "Unknown update strategy '{other}'. Valid values for {ADBC_UPDATE_STRATEGY_ENV}: statement, staging_table"
+                ),
+            },
+            None => Ok(Self::Statement),
+        }
+    }
+}
+
 /// ETL sink that writes transformed batches directly to an ADBC target.
 ///
 /// Inserts use ADBC bulk ingest, while updates and deletes execute row-level
@@ -633,19 +663,124 @@ impl AdbcSink {
         Ok(statements)
     }
 
-    fn update_sql_statements(
+    /// Perform an UPDATE via a temporary staging table:
+    ///
+    /// 1. Bulk-ingest the update batch into a staging table.
+    /// 2. `MERGE INTO target USING staging ON … WHEN MATCHED THEN UPDATE SET …`
+    /// 3. `DROP TABLE staging`.
+    fn staging_merge_update(
         &self,
+        conn: &mut AdbcConnection,
         table_name: &str,
-        batch: &RecordBatch,
+        batch: RecordBatch,
         key_columns: &[String],
-    ) -> anyhow::Result<Vec<String>> {
+    ) -> anyhow::Result<()> {
         if key_columns.is_empty() {
             anyhow::bail!("Update requires at least one key column");
         }
 
-        (0..batch.num_rows())
-            .map(|row| self.update_sql_for_row(table_name, batch, row, key_columns))
-            .collect()
+        let schema = batch.schema();
+        let has_non_key = schema
+            .fields()
+            .iter()
+            .any(|f| !key_columns.iter().any(|k| k == f.name()));
+        if !has_non_key {
+            anyhow::bail!(
+                "Update requires at least one non-key column in batch schema for table '{table_name}'"
+            );
+        }
+
+        // Generate a unique staging table name.
+        let ts = std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .unwrap_or_default()
+            .as_millis();
+        let staging_table = format!("_spicebench_stg_{table_name}_{ts}");
+
+        // 1. Bulk-ingest batch into the staging table.
+        if let Err(e) = self.ingest_insert_batch(conn, &staging_table, batch) {
+            self.drop_staging_table(conn, &staging_table);
+            return Err(e.context(format!(
+                "Failed to ingest update data into staging table '{staging_table}'"
+            )));
+        }
+
+        // 2. MERGE INTO target from staging.
+        let merge_sql = Self::build_staging_merge_sql(
+            &self.target_table_identifier(table_name),
+            &self.target_table_identifier(&staging_table),
+            &schema,
+            key_columns,
+            self.identifier_quote_char,
+        );
+        let merge_result = conn
+            .execute_update(&merge_sql)
+            .map_err(|e| anyhow::anyhow!("MERGE INTO update failed for '{table_name}': {e}"));
+
+        // 3. Drop staging table (always, even on merge failure).
+        self.drop_staging_table(conn, &staging_table);
+
+        merge_result?;
+        Ok(())
+    }
+
+    /// Best-effort drop of a staging table.
+    fn drop_staging_table(&self, conn: &mut AdbcConnection, staging_table: &str) {
+        let drop_sql = format!(
+            "DROP TABLE IF EXISTS {}",
+            self.target_table_identifier(staging_table)
+        );
+
+        if let Err(e) = conn.execute_update(&drop_sql) {
+            tracing::error!(
+                staging_table = %staging_table,
+                error = %e,
+                "Failed to drop staging table. Manual cleanup may be required."
+            );
+        }
+    }
+
+    /// Build a `MERGE INTO` statement that reads from a staging table.
+    ///
+    /// ```sql
+    /// MERGE INTO target t
+    /// USING staging s
+    /// ON t.key1 = s.key1
+    /// WHEN MATCHED THEN UPDATE SET val1 = s.val1, val2 = s.val2
+    /// ```
+    fn build_staging_merge_sql(
+        target_ident: &str,
+        staging_ident: &str,
+        schema: &Schema,
+        key_columns: &[String],
+        identifier_quote_char: char,
+    ) -> String {
+        let on_clause: String = key_columns
+            .iter()
+            .map(|k| {
+                let q = Self::quote_ident(k, identifier_quote_char);
+                format!("t.{q} = s.{q}")
+            })
+            .collect::<Vec<_>>()
+            .join(" AND ");
+
+        let set_clause: String = schema
+            .fields()
+            .iter()
+            .filter(|f| !key_columns.iter().any(|k| k == f.name()))
+            .map(|f| {
+                let q = Self::quote_ident(f.name(), identifier_quote_char);
+                format!("{q} = s.{q}")
+            })
+            .collect::<Vec<_>>()
+            .join(", ");
+
+        format!(
+            "MERGE INTO {target_ident} t \
+             USING {staging_ident} s \
+             ON {on_clause} \
+             WHEN MATCHED THEN UPDATE SET {set_clause}"
+        )
     }
 }
 
@@ -697,24 +832,40 @@ impl Sink for AdbcSink {
                 );
             }
             InsertOp::Update { key_columns } => {
-                let statements = self.update_sql_statements(table_name, &batch, &key_columns)?;
-                let num_statements = statements.len();
+                let strategy = UpdateStrategy::from_env()?;
+                let num_rows = batch.num_rows();
                 let start = Instant::now();
-                for sql in statements {
-                    conn.execute_update(&sql).map_err(|e| {
-                        anyhow::anyhow!("ADBC update execution failed for '{table_name}': {e}")
-                    })?;
+
+                match strategy {
+                    UpdateStrategy::StagingTable => {
+                        self.staging_merge_update(&mut conn, table_name, batch, &key_columns)?;
+                    }
+                    UpdateStrategy::Statement => {
+                        let statements: Vec<String> = (0..num_rows)
+                            .map(|row| {
+                                self.update_sql_for_row(table_name, &batch, row, &key_columns)
+                            })
+                            .collect::<anyhow::Result<Vec<_>>>()?;
+                        for sql in &statements {
+                            conn.execute_update(sql).map_err(|e| {
+                                anyhow::anyhow!(
+                                    "ADBC update execution failed for '{table_name}': {e}"
+                                )
+                            })?;
+                        }
+                    }
                 }
+
                 let elapsed = start.elapsed();
                 let rows_per_sec = if elapsed.as_secs_f64() > 0.0 {
-                    batch.num_rows() as f64 / elapsed.as_secs_f64()
+                    num_rows as f64 / elapsed.as_secs_f64()
                 } else {
                     0.0
                 };
                 tracing::debug!(
                     table = %table_name,
-                    rows = batch.num_rows(),
-                    num_statements,
+                    rows = num_rows,
+                    strategy = ?strategy,
                     elapsed_ms = elapsed.as_millis(),
                     rows_per_sec = format!("{rows_per_sec:.1}"),
                     "UPDATE executed"
@@ -917,4 +1068,64 @@ mod tests {
         .unwrap();
         assert!(stmts.is_empty());
     }
+
+    // ── STAGING MERGE UPDATE tests ───────────────────────────────────────
+
+    #[test]
+    fn staging_merge_sql_single_key() {
+        let schema = Schema::new(vec![
+            Field::new("id", DataType::Int64, false),
+            Field::new("name", DataType::Utf8, false),
+            Field::new("value", DataType::Float64, false),
+        ]);
+        let key_columns = vec!["id".to_string()];
+        let sql = AdbcSink::build_staging_merge_sql(
+            r#""target""#,
+            r#""staging""#,
+            &schema,
+            &key_columns,
+            '"',
+        );
+        assert_eq!(
+            sql,
+            "MERGE INTO \"target\" t \
+             USING \"staging\" s \
+             ON t.\"id\" = s.\"id\" \
+             WHEN MATCHED THEN UPDATE SET \"name\" = s.\"name\", \"value\" = s.\"value\""
+        );
+    }
+
+    #[test]
+    fn staging_merge_sql_composite_key() {
+        let schema = Schema::new(vec![
+            Field::new("a", DataType::Int64, false),
+            Field::new("b", DataType::Int32, false),
+            Field::new("val", DataType::Utf8, false),
+        ]);
+        let key_columns = vec!["a".to_string(), "b".to_string()];
+        let sql =
+            AdbcSink::build_staging_merge_sql(r#""t""#, r#""stg""#, &schema, &key_columns, '"');
+        assert!(sql.contains(r#"ON t."a" = s."a" AND t."b" = s."b""#));
+        assert!(sql.contains(r#"UPDATE SET "val" = s."val""#));
+    }
+
+    #[test]
+    fn staging_merge_sql_databricks_backticks() {
+        let schema = Schema::new(vec![
+            Field::new("id", DataType::Int64, false),
+            Field::new("name", DataType::Utf8, false),
+        ]);
+        let key_columns = vec!["id".to_string()];
+        let sql = AdbcSink::build_staging_merge_sql(
+            "`catalog`.`schema`.`target`",
+            "`catalog`.`schema`.`staging`",
+            &schema,
+            &key_columns,
+            '`',
+        );
+        assert!(sql.starts_with("MERGE INTO `catalog`.`schema`.`target` t"));
+        assert!(sql.contains("USING `catalog`.`schema`.`staging` s"));
+        assert!(sql.contains("ON t.`id` = s.`id`"));
+        assert!(sql.contains("`name` = s.`name`"));
+    }
 }