Move datasets from setup to create_tables in system adapter protocol

Author: sgrebnov

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

@@ -179,12 +179,10 @@ impl Client {
     pub async fn setup(
         &mut self,
         run_id: uuid::Uuid,
-        datasets: std::collections::HashMap<String, crate::DatasetConfig>,
         metadata: std::collections::HashMap<String, serde_json::Value>,
     ) -> Result<crate::SetupResponse> {
         let request = crate::SetupRequest {
             run_id,
-            datasets,
             metadata,
         };
         let rpc_request = JsonRpcRequest::new(1, crate::methods::SETUP, request);
@@ -198,8 +196,9 @@ impl Client {
     pub async fn create_tables(
         &mut self,
         run_id: uuid::Uuid,
+        datasets: std::collections::HashMap<String, crate::DatasetConfig>,
     ) -> Result<crate::CreateTablesResponse> {
-        let request = crate::CreateTablesRequest { run_id };
+        let request = crate::CreateTablesRequest { run_id, datasets };
         let rpc_request = JsonRpcRequest::new(1, crate::methods::CREATE_TABLES, request);
         let response = self.call_typed(rpc_request).await?;
         response

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

@@ -41,8 +41,8 @@ limitations under the License.
 //!
 //! // Setup a benchmark run
 //! let run_id = Uuid::new_v4();
-//! let setup_response = client.setup(run_id, HashMap::new(), HashMap::new()).await?;
-//! let create_tables_response = client.create_tables(run_id).await?;
+//! let setup_response = client.setup(run_id, HashMap::new()).await?;
+//! let create_tables_response = client.create_tables(run_id, HashMap::new()).await?;
 //!
 //! println!("Driver: {:?}", setup_response.driver);
 //!
@@ -72,7 +72,6 @@ limitations under the License.
 //!     async fn setup(
 //!         &mut self,
 //!         run_id: Uuid,
-//!         datasets: HashMap<String, DatasetConfig>,
 //!         metadata: HashMap<String, serde_json::Value>,
 //!     ) -> Result<SetupResponse, String> {
 //!         // Your setup logic here
@@ -83,7 +82,11 @@ limitations under the License.
 //!         })
 //!     }
 //!
-//!     async fn create_tables(&mut self, run_id: Uuid) -> Result<CreateTablesResponse, String> {
+//!     async fn create_tables(
+//!         &mut self,
+//!         run_id: Uuid,
+//!         datasets: HashMap<String, DatasetConfig>,
+//!     ) -> Result<CreateTablesResponse, String> {
 //!         Ok(CreateTablesResponse { ok: true })
 //!     }
 //!
@@ -149,8 +152,6 @@ pub struct DatasetConfig {
 pub struct SetupRequest {
     /// Unique identifier for this benchmark run
     pub run_id: Uuid,
-    /// Map of dataset name to dataset definition
-    pub datasets: HashMap<String, DatasetConfig>,
     /// Arbitrary run metadata propagated from spicebench to adapters
     #[serde(default)]
     pub metadata: HashMap<String, serde_json::Value>,
@@ -172,6 +173,8 @@ pub struct SetupResponse {
 pub struct CreateTablesRequest {
     /// Unique identifier for this benchmark run
     pub run_id: Uuid,
+    /// Map of dataset name to dataset definition
+    pub datasets: HashMap<String, DatasetConfig>,
 }
 
 /// Response from create_tables request

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

@@ -75,14 +75,14 @@ pub trait Handler: Send + Sync {
     async fn setup(
         &mut self,
         run_id: Uuid,
-        datasets: HashMap<String, DatasetConfig>,
         metadata: HashMap<String, serde_json::Value>,
     ) -> std::result::Result<SetupResponse, String>;
 
     /// Create benchmark tables for a run
     async fn create_tables(
         &mut self,
         run_id: Uuid,
+        datasets: HashMap<String, DatasetConfig>,
     ) -> std::result::Result<CreateTablesResponse, String>;
 
     /// Teardown a benchmark run
@@ -244,7 +244,7 @@ impl<H: Handler> Server<H> {
         };
         Self::handler_response(
             self.handler
-                .setup(req.run_id, req.datasets, req.metadata)
+                .setup(req.run_id, req.metadata)
                 .await,
             id,
         )
@@ -259,7 +259,7 @@ impl<H: Handler> Server<H> {
             Ok(r) => r,
             Err(e) => return e,
         };
-        Self::handler_response(self.handler.create_tables(req.run_id).await, id)
+        Self::handler_response(self.handler.create_tables(req.run_id, req.datasets).await, id)
     }
 
     async fn handle_teardown(
@@ -304,7 +304,6 @@ mod tests {
         async fn setup(
             &mut self,
             _run_id: Uuid,
-            _datasets: HashMap<String, DatasetConfig>,
             _metadata: HashMap<String, serde_json::Value>,
         ) -> std::result::Result<SetupResponse, String> {
             Ok(SetupResponse {
@@ -316,6 +315,7 @@ mod tests {
         async fn create_tables(
             &mut self,
             _run_id: Uuid,
+            _datasets: HashMap<String, DatasetConfig>,
         ) -> std::result::Result<CreateTablesResponse, String> {
             Ok(CreateTablesResponse { ok: true })
         }
@@ -335,7 +335,7 @@ mod tests {
     #[tokio::test]
     async fn test_server_setup() {
         let mut server = Server::new(TestHandler);
-        let request = r#"{"jsonrpc":"2.0","id":1,"method":"setup","params":{"run_id":"00000000-0000-0000-0000-000000000000","datasets":{},"metadata":{}}}"#;
+        let request = r#"{"jsonrpc":"2.0","id":1,"method":"setup","params":{"run_id":"00000000-0000-0000-0000-000000000000","metadata":{}}}"#;
         let response = server.handle_request(request).await;
 
         assert!(response.get("result").is_some());

src/main.rs

@@ -136,7 +136,7 @@ async fn main() -> anyhow::Result<()> {
     ]);
 
     let adbc_driver = match system_adapter_client
-        .setup(run_id, datasets, setup_metadata)
+        .setup(run_id, setup_metadata)
         .await
     {
         Ok(response) => response,
@@ -182,7 +182,7 @@ async fn main() -> anyhow::Result<()> {
         &mutations,
     )?;
 
-    if let Err(e) = system_adapter_client.create_tables(run_id).await {
+    if let Err(e) = system_adapter_client.create_tables(run_id, datasets).await {
         pipeline.cancel();
         return Err(anyhow::anyhow!(
             "Failed to create tables via system adapter: {e}"

system-adapters/databricks/src/main.rs

@@ -183,7 +183,6 @@ impl TableFormat {
 
 #[derive(Debug, Clone)]
 struct RunState {
-    datasets: HashMap<String, DatasetConfig>,
     table_format: TableFormat,
     created_tables: Vec<String>,
     cluster_id: Option<String>,
@@ -1025,12 +1024,10 @@ impl Handler for DatabricksAdapter {
     async fn setup(
         &mut self,
         run_id: Uuid,
-        datasets: HashMap<String, DatasetConfig>,
         metadata: HashMap<String, Value>,
     ) -> std::result::Result<SetupResponse, String> {
         eprintln!(
-            "[databricks-adapter] setup: run_id={run_id}, datasets={}",
-            datasets.len()
+            "[databricks-adapter] setup: run_id={run_id}"
         );
 
         let (cluster_id, cluster_created_by_adapter) = match &self.config.compute_target {
@@ -1068,7 +1065,6 @@ impl Handler for DatabricksAdapter {
         self.runs.insert(
             run_id,
             RunState {
-                datasets,
                 table_format,
                 created_tables: Vec::new(),
                 cluster_id,
@@ -1094,13 +1090,14 @@ impl Handler for DatabricksAdapter {
     async fn create_tables(
         &mut self,
         run_id: Uuid,
+        datasets: HashMap<String, DatasetConfig>,
     ) -> std::result::Result<CreateTablesResponse, String> {
-        let (datasets, table_format) = {
+        let table_format = {
             let state = self
                 .runs
                 .get(&run_id)
                 .ok_or_else(|| format!("Unknown run_id: {run_id}"))?;
-            (state.datasets.clone(), state.table_format)
+            state.table_format
         };
 
         let mut created_tables = Vec::with_capacity(datasets.len());