feat: Staging table support for MERGE-based updates (#243)

- Add create_staging_table JSON-RPC method to system-adapter-protocol
- AdbcSink accepts optional system adapter client for remote staging table creation
- Strip primary_key_columns when SPICEBENCH_ADBC_UPDATE_STRATEGY=staging_table
- Gracefully fall back if adapter does not support create_staging_table (METHOD_NOT_FOUND)

Author: phillipleblanc

Cargo.lock

@@ -30,6 +30,7 @@ tempfile.workspace = true
 tokio.workspace = true
 tokio-util.workspace = true
 tracing.workspace = true
+uuid.workspace = true
 
 [features]
 default = []

crates/etl/Cargo.toml

@@ -874,7 +874,7 @@ async fn write_segments_for_batch(
                     partition_columns.to_vec(),
                 )
                 .await
-                .map_err(|e| format!("write {table_name_owned} batch {batch_id}: {e}"))?;
+                .map_err(|e| format!("write {table_name_owned} batch {batch_id}: {e:#}"))?;
         }
 
         return Ok(());

crates/etl/src/lib.rs

@@ -15,9 +15,10 @@ limitations under the License.
 */
 
 use std::collections::{HashMap, HashSet};
+use std::sync::Arc;
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::time::Instant;
-use tokio::sync::{RwLock, mpsc};
+use tokio::sync::{Mutex, RwLock, mpsc};
 use tokio::task::JoinHandle;
 
 use adbc_client::{
@@ -28,7 +29,8 @@ use arrow::datatypes::{DataType, Schema};
 use arrow::record_batch::RecordBatchIterator;
 use arrow_cast::display::array_value_to_string;
 use async_trait::async_trait;
-use system_adapter_protocol::DatasetConfig;
+use system_adapter_protocol::{Client as SystemAdapterClient, DatasetConfig};
+use uuid::Uuid;
 
 use super::{InsertOp, Sink};
 
@@ -174,6 +176,8 @@ pub struct AdbcSink {
     reuse_bulk_ingest_streams: bool,
     flush_stream_before_upsert: bool,
     bulk_ingest_stream_buffer: usize,
+    /// Optional system adapter client for staging table creation.
+    staging_adapter: Option<(Arc<Mutex<SystemAdapterClient>>, Uuid)>,
 }
 
 impl AdbcSink {
@@ -230,11 +234,15 @@ impl AdbcSink {
     }
 
     /// Creates a new [`AdbcSink`] backed by a connection pool.
+    ///
+    /// When a system adapter client and run ID are provided, the `StagingTable`
+    /// update strategy will use the adapter to create staging tables remotely.
     pub fn new(
         driver_name: &str,
         db_kwargs: HashMap<String, serde_json::Value>,
         target_db_catalog: Option<String>,
         target_db_schema: Option<String>,
+        staging_adapter: Option<(Arc<Mutex<SystemAdapterClient>>, Uuid)>,
     ) -> anyhow::Result<Self> {
         let update_strategy = UpdateStrategy::from_env()?;
         let pool_size = Self::pool_size();
@@ -269,6 +277,7 @@ impl AdbcSink {
             reuse_bulk_ingest_streams,
             flush_stream_before_upsert,
             bulk_ingest_stream_buffer,
+            staging_adapter,
         })
     }
 
@@ -1148,10 +1157,11 @@ impl AdbcSink {
 
     /// 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(
+    /// 1. Create the staging table via the hook (same schema/partitioning as target).
+    /// 2. Bulk-ingest the update batch into the staging table.
+    /// 3. `MERGE INTO target USING staging ON … WHEN MATCHED THEN UPDATE SET …`
+    /// 4. `DROP TABLE staging` (via hook, then SQL fallback).
+    async fn staging_merge_update(
         &self,
         conn: &mut AdbcConnection,
         table_name: &str,
@@ -1182,15 +1192,30 @@ impl AdbcSink {
             .as_millis();
         let staging_table = format!("_spicebench_stg_{table_name}_{ts}");
 
-        // 1. Bulk-ingest batch into the staging table.
+        // 1. Create the staging table via the system adapter (if available).
+        if let Some((client, run_id)) = &self.staging_adapter {
+            client
+                .lock()
+                .await
+                .create_staging_table(*run_id, table_name, &staging_table)
+                .await
+                .map_err(|e| {
+                    anyhow::anyhow!(
+                        "Failed to create staging table '{staging_table}' for '{table_name}': {e}"
+                    )
+                })?;
+        }
+
+        // 2. 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);
+            self.drop_staging_table_best_effort(conn, &staging_table)
+                .await;
             return Err(e.context(format!(
                 "Failed to ingest update data into staging table '{staging_table}'"
             )));
         }
 
-        // 2. MERGE INTO target from staging.
+        // 3. 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),
@@ -1202,8 +1227,9 @@ impl AdbcSink {
             .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);
+        // 4. Drop staging table (always, even on merge failure).
+        self.drop_staging_table_best_effort(conn, &staging_table)
+            .await;
 
         merge_result?;
         tracing::debug!(
@@ -1215,8 +1241,8 @@ impl AdbcSink {
         Ok(())
     }
 
-    /// Best-effort drop of a staging table.
-    fn drop_staging_table(&self, conn: &mut AdbcConnection, staging_table: &str) {
+    /// Best-effort drop of a staging table via SQL DROP TABLE.
+    async fn drop_staging_table_best_effort(&self, conn: &mut AdbcConnection, staging_table: &str) {
         let drop_sql = format!(
             "DROP TABLE IF EXISTS {}",
             self.target_table_identifier(staging_table)
@@ -1372,7 +1398,8 @@ impl Sink for AdbcSink {
                         let mut conn = self.pool.get().map_err(|e| {
                             anyhow::anyhow!("Failed to get ADBC connection from pool: {e}")
                         })?;
-                        self.staging_merge_update(&mut conn, table_name, batch, &key_columns)?;
+                        self.staging_merge_update(&mut conn, table_name, batch, &key_columns)
+                            .await?;
                     }
                     UpdateStrategy::BulkIngestUpsert => {
                         if self.reuse_bulk_ingest_streams && !self.flush_stream_before_upsert {

crates/etl/src/sink/adbc.rs

@@ -12,6 +12,7 @@ serde = { workspace = true }
 serde_json = { workspace = true }
 uuid = { workspace = true, features = ["serde", "v4"] }
 tokio = { workspace = true, features = ["process", "io-util", "io-std"] }
+tracing = { workspace = true }
 reqwest = { workspace = true, optional = true }
 async-trait = { workspace = true, optional = true }
 

crates/system-adapter-protocol/Cargo.toml

@@ -228,6 +228,40 @@ impl Client {
             .ok_or_else(|| ClientError::InvalidResponse("Missing result".to_string()))
     }
 
+    /// Create a staging table for MERGE-based updates.
+    ///
+    /// If the remote adapter does not support this method (returns
+    /// `METHOD_NOT_FOUND`), the call is treated as a successful no-op so that
+    /// newer spicebench versions work against older adapters.
+    pub async fn create_staging_table(
+        &mut self,
+        run_id: uuid::Uuid,
+        source_dataset: &str,
+        staging_table_name: &str,
+    ) -> Result<crate::CreateStagingTableResponse> {
+        let request = crate::CreateStagingTableRequest {
+            run_id,
+            source_dataset: source_dataset.to_string(),
+            staging_table_name: staging_table_name.to_string(),
+        };
+        let rpc_request = JsonRpcRequest::new(1, crate::methods::CREATE_STAGING_TABLE, request);
+        match self.call_typed(rpc_request).await {
+            Ok(response) => response
+                .result
+                .ok_or_else(|| ClientError::InvalidResponse("Missing result".to_string())),
+            Err(ClientError::JsonRpc(ref e)) if e.code == crate::error_codes::METHOD_NOT_FOUND => {
+                tracing::warn!(
+                    source_dataset,
+                    staging_table_name,
+                    "System adapter does not support create_staging_table; \
+                     falling back to implicit table creation via bulk ingest"
+                );
+                Ok(crate::CreateStagingTableResponse { ok: true })
+            }
+            Err(e) => Err(e),
+        }
+    }
+
     /// Make a typed JSON-RPC call with request and response types
     async fn call_typed<Req: Serialize, Resp: DeserializeOwned>(
         &mut self,

crates/system-adapter-protocol/src/client.rs

@@ -371,10 +371,31 @@ pub mod error_codes {
     pub const INTERNAL_ERROR: i32 = -32603;
 }
 
+/// Request to create a staging table for MERGE-based updates.
+///
+/// JSON-RPC method: `create_staging_table`
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct CreateStagingTableRequest {
+    /// Unique identifier for the benchmark run
+    pub run_id: Uuid,
+    /// Name of the source dataset this staging table is based on
+    pub source_dataset: String,
+    /// Name to use for the staging table
+    pub staging_table_name: String,
+}
+
+/// Response from create staging table request
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
+pub struct CreateStagingTableResponse {
+    /// Indicates if the staging table was created successfully
+    pub ok: bool,
+}
+
 /// Method names for the system adapter protocol
 pub mod methods {
     pub const SETUP: &str = "setup";
     pub const TEARDOWN: &str = "teardown";
     pub const METRICS: &str = "metrics";
+    pub const CREATE_STAGING_TABLE: &str = "create_staging_table";
     pub const RPC_METHODS: &str = "rpc.methods";
 }

crates/system-adapter-protocol/src/lib.rs

@@ -17,8 +17,9 @@ limitations under the License.
 //! Server implementations for system adapter JSON-RPC protocol.
 
 use crate::{
-    DatasetConfig, EtlSinkType, JsonRpcError, JsonRpcResponse, MetricsRequest, MetricsResponse,
-    SetupRequest, SetupResponse, TeardownRequest, TeardownResponse, error_codes, methods,
+    CreateStagingTableRequest, CreateStagingTableResponse, DatasetConfig, EtlSinkType,
+    JsonRpcError, JsonRpcResponse, MetricsRequest, MetricsResponse, SetupRequest, SetupResponse,
+    TeardownRequest, TeardownResponse, error_codes, methods,
 };
 use async_trait::async_trait;
 use serde::de::DeserializeOwned;
@@ -106,9 +107,25 @@ pub trait Handler: Send + Sync {
             methods::SETUP.to_string(),
             methods::TEARDOWN.to_string(),
             methods::METRICS.to_string(),
+            methods::CREATE_STAGING_TABLE.to_string(),
             methods::RPC_METHODS.to_string(),
         ]
     }
+
+    /// Create a staging table for MERGE-based updates.
+    ///
+    /// The staging table should have the same schema and partitioning as the
+    /// source dataset but with the given staging table name. Adapters that
+    /// handle staging table creation implicitly (e.g. via bulk ingest) can
+    /// leave the default no-op implementation.
+    async fn create_staging_table(
+        &mut self,
+        _run_id: Uuid,
+        _source_dataset: &str,
+        _staging_table_name: &str,
+    ) -> std::result::Result<CreateStagingTableResponse, String> {
+        Ok(CreateStagingTableResponse { ok: true })
+    }
 }
 
 /// System adapter server
@@ -182,6 +199,9 @@ impl<H: Handler> Server<H> {
             methods::SETUP => self.handle_setup(&request, id.clone()).await,
             methods::TEARDOWN => self.handle_teardown(&request, id.clone()).await,
             methods::METRICS => self.handle_metrics(&request, id.clone()).await,
+            methods::CREATE_STAGING_TABLE => {
+                self.handle_create_staging_table(&request, id.clone()).await
+            }
             methods::RPC_METHODS => self.handle_rpc_methods(id.clone()).await,
             _ => serde_json::to_value(JsonRpcResponse::<()>::error(
                 id,
@@ -272,6 +292,23 @@ impl<H: Handler> Server<H> {
         Self::handler_response(self.handler.metrics(req.run_id, req.final_scrape).await, id)
     }
 
+    async fn handle_create_staging_table(
+        &mut self,
+        request: &serde_json::Value,
+        id: serde_json::Value,
+    ) -> serde_json::Value {
+        let req: CreateStagingTableRequest = match Self::parse_params(request, &id) {
+            Ok(r) => r,
+            Err(e) => return e,
+        };
+        Self::handler_response(
+            self.handler
+                .create_staging_table(req.run_id, &req.source_dataset, &req.staging_table_name)
+                .await,
+            id,
+        )
+    }
+
     async fn handle_rpc_methods(&mut self, id: serde_json::Value) -> serde_json::Value {
         let methods = self.handler.rpc_methods();
         let result = serde_json::json!({ "methods": methods });

crates/system-adapter-protocol/src/server.rs

@@ -138,6 +138,7 @@ pub async fn execute(args: &EtlArgs) -> anyhow::Result<()> {
                 db_kwargs,
                 args.adbc_catalog.clone(),
                 args.adbc_schema.clone(),
+                None,
             )?);
 
             (

src/commands/etl_cmd.rs

@@ -118,13 +118,27 @@ async fn run_benchmark(
     let etl_sink_type = system_adapter_protocol::EtlSinkType::Adbc;
     let target_config = None;
 
-    let datasets = ETLPipeline::create_tables_request_datasets(
+    let mut datasets = ETLPipeline::create_tables_request_datasets(
         dataset_source.clone(),
         &generation_config,
         Arc::clone(&data_source),
         &mutations,
         target_config.clone(),
     )?;
+
+    // When using the staging_table update strategy, tables must NOT have
+    // primary keys. MERGE INTO handles matching via the ON clause and uses
+    // delete+insert execution, which conflicts with Cayenne's automatic
+    // on_conflict: Upsert behavior on primary-key tables.
+    let uses_staging_table = std::env::var("SPICEBENCH_ADBC_UPDATE_STRATEGY")
+        .ok()
+        .map(|v| v.eq_ignore_ascii_case("staging_table"))
+        .unwrap_or(false);
+    if uses_staging_table {
+        for config in datasets.values_mut() {
+            config.primary_key_columns.clear();
+        }
+    }
     let (setup_response, mut pipeline) = {
         let setup_response = system_adapter_client
             .lock()
@@ -156,6 +170,7 @@ async fn run_benchmark(
             db_kwargs,
             target_db_catalog,
             target_db_schema,
+            Some((Arc::clone(&system_adapter_client), run_id)),
         )?);
 
         let mut pipeline = ETLPipeline::new(