fix: Properly set ADBC catalog paths when specified (#141)

* fix: Properly set ADBC catalog paths when specified

* fix build

Author: peasee

crates/adbc_client/src/lib.rs

@@ -145,12 +145,14 @@ impl AdbcConnection {
     pub fn bulk_ingest(
         &mut self,
         target_table: &str,
+        target_db_catalog: Option<&str>,
         target_db_schema: Option<&str>,
         mode: options::IngestMode,
         batch: RecordBatch,
     ) -> Result<Option<i64>> {
         self.bulk_ingest_stream(
             target_table,
+            target_db_catalog,
             target_db_schema,
             mode,
             Box::new(arrow_array::RecordBatchIterator::new(
@@ -168,6 +170,7 @@ impl AdbcConnection {
     pub fn bulk_ingest_stream(
         &mut self,
         target_table: &str,
+        target_db_catalog: Option<&str>,
         target_db_schema: Option<&str>,
         mode: options::IngestMode,
         reader: Box<dyn arrow_array::RecordBatchReader + Send>,
@@ -184,6 +187,16 @@ impl AdbcConnection {
             reason: format!("Failed to set target table: {e}"),
         })?;
 
+        if let Some(catalog) = target_db_catalog {
+            stmt.set_option(
+                options::OptionStatement::TargetCatalog,
+                OptionValue::from(catalog),
+            )
+            .map_err(|e| Error::ExecuteQuery {
+                reason: format!("Failed to set target db catalog: {e}"),
+            })?;
+        }
+
         if let Some(schema) = target_db_schema {
             stmt.set_option(
                 options::OptionStatement::TargetDbSchema,
@@ -204,7 +217,9 @@ impl AdbcConnection {
         })?;
 
         stmt.execute_update().map_err(|e| Error::ExecuteQuery {
-            reason: format!("Bulk ingest execution failed: {e}"),
+            reason: format!(
+                "Bulk ingest execution failed for target_table='{target_table}', target_db_catalog={target_db_catalog:?}, target_db_schema={target_db_schema:?}: {e}"
+            ),
         })
     }
 }

crates/etl/src/lib.rs

@@ -52,7 +52,7 @@ const INTERNAL_COLUMNS: &[&str] = &["_op", "_op_index"];
 /// Smaller input batches from a [`ReadResult`] are concatenated together until
 /// this threshold is reached, and larger input batches are split so no output
 /// batch exceeds this size.
-const TARGET_BATCH_ROWS: usize = 8_192 * 2;
+const TARGET_BATCH_ROWS: usize = 8_192 * 4;
 
 /// Maximum number of in-flight sink writes allowed per table task when the
 /// current segment set is insert-only.

crates/etl/src/main.rs

@@ -79,6 +79,10 @@ struct Cli {
     #[arg(long)]
     adbc_uri: Option<String>,
 
+    /// Optional target database catalog for ADBC bulk ingest inserts
+    #[arg(long)]
+    adbc_catalog: Option<String>,
+
     /// Optional target database schema for bulk ingest
     #[arg(long)]
     adbc_schema: Option<String>,
@@ -173,7 +177,12 @@ async fn main() -> anyhow::Result<()> {
                     });
             }
 
-            let adbc_sink = Arc::new(AdbcSink::new(driver, db_kwargs, cli.adbc_schema.clone())?);
+            let adbc_sink = Arc::new(AdbcSink::new(
+                driver,
+                db_kwargs,
+                cli.adbc_catalog.clone(),
+                cli.adbc_schema.clone(),
+            )?);
 
             (
                 adbc_sink.clone() as Arc<dyn Sink>,
@@ -242,6 +251,7 @@ async fn main() -> anyhow::Result<()> {
         prefix = %cli.prefix,
         target = %target_kind,
         adbc_driver = ?cli.adbc_driver,
+        adbc_catalog = ?cli.adbc_catalog,
         adbc_schema = ?cli.adbc_schema,
         adbc_create_tables = cli.adbc_create_tables,
         target_prefix = %cli.target_prefix,

crates/etl/src/sink/adbc.rs

@@ -37,6 +37,7 @@ const MAX_ADBC_INGEST_BATCH_BYTES_ENV: &str = "SPICEBENCH_ADBC_MAX_INGEST_BATCH_
 /// SQL statements derived from key columns in each batch.
 pub struct AdbcSink {
     conn: Mutex<AdbcConnection>,
+    target_db_catalog: Option<String>,
     target_db_schema: Option<String>,
 }
 
@@ -53,13 +54,15 @@ impl AdbcSink {
     pub fn new(
         driver_name: &str,
         db_kwargs: HashMap<String, serde_json::Value>,
+        target_db_catalog: Option<String>,
         target_db_schema: Option<String>,
     ) -> anyhow::Result<Self> {
         let conn = AdbcConnection::create(driver_name, db_kwargs)
             .map_err(|e| anyhow::anyhow!("Failed to create ADBC connection: {e}"))?;
 
         Ok(Self {
             conn: Mutex::new(conn),
+            target_db_catalog,
             target_db_schema,
         })
     }
@@ -106,13 +109,45 @@ impl AdbcSink {
     }
 
     fn target_table_identifier(&self, table_name: &str) -> String {
-        let table_ident = Self::quote_identifier(table_name);
+        let mut parts = Vec::with_capacity(3);
+
+        if let Some(catalog) = self.target_db_catalog.as_deref() {
+            if !catalog.is_empty() {
+                parts.push(Self::quote_identifier(catalog));
+            }
+        }
+
+        if let Some(schema) = self.target_db_schema.as_deref() {
+            if !schema.is_empty() {
+                parts.push(Self::quote_identifier(schema));
+            }
+        }
+
+        parts.push(Self::quote_identifier(table_name));
+        parts.join(".")
+    }
+
+    fn target_table_ingest_name(&self, table_name: &str) -> String {
+        self.target_table_identifier_unquoted(table_name)
+    }
+
+    fn target_table_identifier_unquoted(&self, table_name: &str) -> String {
+        let mut parts = Vec::with_capacity(3);
+
+        if let Some(catalog) = self.target_db_catalog.as_deref() {
+            if !catalog.is_empty() {
+                parts.push(catalog.to_string());
+            }
+        }
+
         if let Some(schema) = self.target_db_schema.as_deref() {
             if !schema.is_empty() {
-                return format!("{}.{}", Self::quote_identifier(schema), table_ident);
+                parts.push(schema.to_string());
             }
         }
-        table_ident
+
+        parts.push(table_name.to_string());
+        parts.join(".")
     }
 
     fn create_table_sql(&self, table_name: &str, schema: &Schema) -> anyhow::Result<String> {
@@ -129,7 +164,7 @@ impl AdbcSink {
             .join(", ");
 
         Ok(format!(
-            "CREATE TABLE IF NOT EXISTS spicebench.bench.{} ({columns})",
+            "CREATE TABLE IF NOT EXISTS {} ({columns})",
             self.target_table_identifier(table_name)
         ))
     }
@@ -222,9 +257,20 @@ impl AdbcSink {
         table_name: &str,
         batch: RecordBatch,
     ) -> anyhow::Result<()> {
+        let ingest_table_name = self.target_table_ingest_name(table_name);
+        let target_db_catalog = self
+            .target_db_catalog
+            .as_deref()
+            .and_then(|catalog| (!catalog.is_empty()).then_some(catalog));
+        let target_db_schema = self
+            .target_db_schema
+            .as_deref()
+            .and_then(|schema| (!schema.is_empty()).then_some(schema));
+
         match conn.bulk_ingest(
             table_name,
-            self.target_db_schema.as_deref(),
+            target_db_catalog,
+            target_db_schema,
             IngestMode::CreateAppend,
             batch.clone(),
         ) {
@@ -234,7 +280,7 @@ impl AdbcSink {
                 if Self::is_message_too_large_error(&message) {
                     if batch.num_rows() <= 1 {
                         anyhow::bail!(
-                            "ADBC bulk ingest failed for '{table_name}': single-row batch still exceeds FlightSQL message limit: {message}. Configure a larger FlightSQL max message size (e.g. adbc.flight.sql.client_option.with_max_msg_size)."
+                            "ADBC bulk ingest failed for source table '{table_name}' (ingest target '{ingest_table_name}'): single-row batch still exceeds FlightSQL message limit: {message}. Configure a larger FlightSQL max message size (e.g. adbc.flight.sql.client_option.with_max_msg_size)."
                         );
                     }
 
@@ -245,7 +291,9 @@ impl AdbcSink {
                     self.bulk_ingest_with_retry(conn, table_name, right)?;
                     Ok(())
                 } else {
-                    anyhow::bail!("ADBC bulk ingest failed for '{table_name}': {message}")
+                    anyhow::bail!(
+                        "ADBC bulk ingest failed for source table '{table_name}' (ingest target '{ingest_table_name}'): {message}"
+                    )
                 }
             }
         }

src/main.rs

@@ -64,6 +64,20 @@ fn s3_hive_target_prefix(common: &CommonArgs, scenario_name: &str, run_id: uuid:
     format!("{target_prefix}/{scenario_name}/{run_id}")
 }
 
+fn infer_adbc_target_namespace(
+    catalog_namespace: Option<&str>,
+) -> (Option<String>, Option<String>) {
+    let Some(namespace) = catalog_namespace.map(str::trim).filter(|ns| !ns.is_empty()) else {
+        return (None, None);
+    };
+
+    let mut parts = namespace.rsplitn(2, '.');
+    let schema = parts.next().map(str::trim).filter(|v| !v.is_empty());
+    let catalog = parts.next().map(str::trim).filter(|v| !v.is_empty());
+
+    (catalog.map(str::to_string), schema.map(str::to_string))
+}
+
 async fn run_benchmark(
     common: &CommonArgs,
     system_adapter_client: Arc<Mutex<system_adapter_protocol::Client>>,
@@ -203,7 +217,15 @@ async fn run_benchmark(
                     });
             }
 
-            let adbc_sink = Arc::new(AdbcSink::new(&driver_name, db_kwargs, None)?);
+            let (target_db_catalog, target_db_schema) =
+                infer_adbc_target_namespace(setup_response.catalog_namespace.as_deref());
+
+            let adbc_sink = Arc::new(AdbcSink::new(
+                &driver_name,
+                db_kwargs,
+                target_db_catalog,
+                target_db_schema,
+            )?);
 
             setup_response_for_run = Some(setup_response);