Various improvements for Databricks and result validation (#164)

* Various improvements

* WIP

* Fix

* Lint

Author: krinart

crates/adbc_client/src/lib.rs

@@ -24,7 +24,7 @@ use adbc_core::{Connection, Database, Driver, LOAD_FLAG_DEFAULT, Optionable, Sta
 use adbc_driver_manager::ManagedDriver;
 use arrow::compute::cast;
 use arrow::datatypes::{DataType, Schema};
-use arrow_array::RecordBatch;
+use arrow_array::{Array, RecordBatch, StringArray};
 use arrow_schema::Field;
 use snafu::prelude::*;
 use std::collections::HashMap;
@@ -56,15 +56,21 @@ pub type Result<T, E = Error> = std::result::Result<T, E>;
 pub struct AdbcConnection {
     conn: adbc_driver_manager::ManagedConnection,
     downcast_utf8view: bool,
+    resolve_opaque_numerics: bool,
 }
 
 impl AdbcConnection {
     /// Create an `AdbcConnection` from an already-established [`ManagedConnection`].
     #[must_use]
-    pub fn new(conn: adbc_driver_manager::ManagedConnection, downcast_utf8view: bool) -> Self {
+    pub fn new(
+        conn: adbc_driver_manager::ManagedConnection,
+        downcast_utf8view: bool,
+        resolve_opaque_numerics: bool,
+    ) -> Self {
         Self {
             conn,
             downcast_utf8view,
+            resolve_opaque_numerics,
         }
     }
 
@@ -106,7 +112,11 @@ impl AdbcConnection {
             reason: e.to_string(),
         })?;
 
-        Ok(Self::new(conn, driver_name == "databricks"))
+        Ok(Self::new(
+            conn,
+            driver_name == "databricks",
+            driver_name == "postgresql",
+        ))
     }
 
     /// Lightweight check that the connection is still usable.
@@ -135,9 +145,15 @@ impl AdbcConnection {
             reason: e.to_string(),
         })?;
 
-        reader
+        let mut batches = reader
             .collect::<std::result::Result<Vec<_>, _>>()
-            .context(ReadBatchSnafu)
+            .context(ReadBatchSnafu)?;
+
+        if self.resolve_opaque_numerics {
+            batches = resolve_opaque_numerics(batches);
+        }
+
+        Ok(batches)
     }
 
     /// Execute a SQL data-modification statement and return the affected row count when provided by the driver.
@@ -275,3 +291,110 @@ fn downcast_utf8view(batch: &RecordBatch) -> RecordBatch {
 
     RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap()
 }
+
+/// Returns `true` if the field uses the Arrow opaque extension type for
+/// PostgreSQL `numeric`.
+fn is_opaque_numeric(field: &Field) -> bool {
+    if !matches!(field.data_type(), DataType::Utf8) {
+        return false;
+    }
+    let metadata = field.metadata();
+    let Some(ext_name) = metadata.get("ARROW:extension:name") else {
+        return false;
+    };
+    if ext_name != "arrow.opaque" {
+        return false;
+    }
+    let Some(ext_meta) = metadata.get("ARROW:extension:metadata") else {
+        return false;
+    };
+    serde_json::from_str::<serde_json::Value>(ext_meta)
+        .ok()
+        .and_then(|v| v.get("type_name")?.as_str().map(|s| s == "numeric"))
+        .unwrap_or(false)
+}
+
+/// Determine the maximum decimal scale (digits after the decimal point)
+/// across all non-null values in a string array.
+fn max_decimal_scale(array: &StringArray) -> i8 {
+    let mut scale: i8 = 0;
+    for i in 0..array.len() {
+        if array.is_null(i) {
+            continue;
+        }
+        let val = array.value(i);
+        if let Some(dot_pos) = val.find('.') {
+            let s = (val.len() - dot_pos - 1) as i8;
+            scale = scale.max(s);
+        }
+    }
+    scale
+}
+
+/// Convert columns returned by the PostgreSQL ADBC driver as
+/// `Utf8` with `arrow.opaque` extension metadata for `numeric` to
+/// `Decimal128`, matching the representation used in checkpoint
+/// parquet files.
+///
+/// The scale for each column is inferred from the actual data across
+/// all batches. If casting fails (e.g. the column contains `NaN` or
+/// `inf`), the original `Utf8` column is kept.
+fn resolve_opaque_numerics(batches: Vec<RecordBatch>) -> Vec<RecordBatch> {
+    if batches.is_empty() {
+        return batches;
+    }
+
+    let schema = batches[0].schema();
+    let opaque_cols: Vec<usize> = schema
+        .fields()
+        .iter()
+        .enumerate()
+        .filter_map(|(i, f)| if is_opaque_numeric(f) { Some(i) } else { None })
+        .collect();
+
+    if opaque_cols.is_empty() {
+        return batches;
+    }
+
+    // Determine the max scale for each opaque numeric column across
+    // all batches so every batch uses a consistent Decimal128 type.
+    let mut scales: Vec<i8> = vec![0; opaque_cols.len()];
+    for batch in &batches {
+        for (j, &col_idx) in opaque_cols.iter().enumerate() {
+            if let Some(arr) = batch.column(col_idx).as_any().downcast_ref::<StringArray>() {
+                scales[j] = scales[j].max(max_decimal_scale(arr));
+            }
+        }
+    }
+
+    batches
+        .into_iter()
+        .map(|batch| {
+            let schema = batch.schema();
+            let mut fields = Vec::with_capacity(schema.fields().len());
+            let mut columns = Vec::with_capacity(schema.fields().len());
+
+            for (i, field) in schema.fields().iter().enumerate() {
+                if let Some(j) = opaque_cols.iter().position(|&idx| idx == i) {
+                    let target_type = DataType::Decimal128(38, scales[j]);
+                    match cast(batch.column(i), &target_type) {
+                        Ok(converted) => {
+                            fields.push(Arc::new(Field::new(
+                                field.name(),
+                                target_type,
+                                field.is_nullable(),
+                            )));
+                            columns.push(converted);
+                            continue;
+                        }
+                        Err(_) => { /* fall through to keep original */ }
+                    }
+                }
+                fields.push(field.clone());
+                columns.push(batch.column(i).clone());
+            }
+
+            RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap()
+        })
+        .collect()
+}

crates/adbc_client/src/pool.rs

@@ -41,14 +41,20 @@ const DEFAULT_POOL_SIZE: u32 = 10;
 pub struct AdbcConnectionManager {
     database: Arc<ManagedDatabase>,
     downcast_utf8view: bool,
+    resolve_opaque_numerics: bool,
 }
 
 impl AdbcConnectionManager {
     /// Create a new manager from an existing [`ManagedDatabase`].
-    pub fn new(database: ManagedDatabase, downcast_utf8view: bool) -> Self {
+    pub fn new(
+        database: ManagedDatabase,
+        downcast_utf8view: bool,
+        resolve_opaque_numerics: bool,
+    ) -> Self {
         Self {
             database: Arc::new(database),
             downcast_utf8view,
+            resolve_opaque_numerics,
         }
     }
 }
@@ -64,7 +70,11 @@ impl r2d2::ManageConnection for AdbcConnectionManager {
             .map_err(|e| Error::CreateConnection {
                 reason: e.to_string(),
             })?;
-        Ok(AdbcConnection::new(conn, self.downcast_utf8view))
+        Ok(AdbcConnection::new(
+            conn,
+            self.downcast_utf8view,
+            self.resolve_opaque_numerics,
+        ))
     }
 
     fn is_valid(&self, conn: &mut Self::Connection) -> std::result::Result<(), Self::Error> {
@@ -119,7 +129,8 @@ pub fn create_pool(
             reason: e.to_string(),
         })?;
 
-    let manager = AdbcConnectionManager::new(db, driver_name == "databricks");
+    let manager =
+        AdbcConnectionManager::new(db, driver_name == "databricks", driver_name == "postgresql");
 
     r2d2::Pool::builder()
         .max_size(pool_size)

crates/etl/src/lib.rs

@@ -357,6 +357,7 @@ async fn read_batches_until_min_rows(
     logical_steps_consumed: &StdArc<AtomicU64>,
     table_name: &str,
     start_batch_id: u64,
+    step_limit: Option<usize>,
 ) -> Result<(Vec<RecordBatch>, Vec<String>, bool, u64, u64), String> {
     let mut all_batches: Vec<RecordBatch> = Vec::new();
     let mut total_rows: usize = 0;
@@ -422,6 +423,7 @@ async fn read_batches_until_min_rows(
             && join_set.len() < MAX_IN_FLIGHT_SOURCE_BATCH_READS.max(1)
             && total_rows < TARGET_BATCH_ROWS
             && !table_finished
+            && step_limit.map_or(true, |limit| consumed_work_units < limit as u64)
         {
             let reservation = {
                 let mut state = work_state.lock().expect("work_state lock poisoned");
@@ -1557,6 +1559,7 @@ async fn run_pipeline(
                         &logical_steps_consumed,
                         &table_name,
                         batch_id,
+                        step_limit,
                     )
                     .await
                     {

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

@@ -115,6 +115,11 @@ impl Client {
             .stdout(std::process::Stdio::piped())
             .stderr(std::process::Stdio::inherit());
 
+        // Place the child in its own process group so it doesn't receive
+        // SIGINT when the user presses ctrl+c, allowing orderly teardown.
+        #[cfg(unix)]
+        cmd.process_group(0);
+
         let mut child = cmd.spawn().map_err(|e| {
             ClientError::Transport(format!(
                 "Failed to start stdio command '{command_str}': {e}"

crates/test-framework/src/queries/validation/mod.rs

@@ -148,11 +148,16 @@ fn datatype_equivalent(expected_type: &DataType, actual_type: &DataType) -> bool
         // Existing numeric and string type equivalences
         _ => matches!(
             (expected_type, actual_type),
-            (DataType::Float32, DataType::Float64)
+            (DataType::Decimal128(_, _), DataType::Decimal128(_, _))
+                | (DataType::Float32, DataType::Float64)
                 | (
                     DataType::Float64 | DataType::Int64,
                     DataType::Decimal128(_, _)
                 )
+                | (
+                    DataType::Decimal128(_, _),
+                    DataType::Int64 | DataType::Int32 | DataType::Float64
+                )
                 | (DataType::Int32, DataType::Int64)
                 | (
                     DataType::Int64,

src/commands/load/mod.rs

@@ -596,6 +596,7 @@ pub(crate) async fn run(
                                     const POLL_INTERVAL: Duration = Duration::from_secs(5);
                                     const MAX_WAIT: Duration = Duration::from_secs(600);
                                     let mut timed_out = false;
+                                    let mut interrupted = false;
                                     loop {
                                         let status =
                                             validation_controller.status_rx.borrow().clone();
@@ -616,7 +617,13 @@ pub(crate) async fn run(
                                             timed_out = true;
                                             break;
                                         }
-                                        tokio::time::sleep(POLL_INTERVAL).await;
+                                        tokio::select! {
+                                            _ = tokio::time::sleep(POLL_INTERVAL) => {}
+                                            _ = signal::ctrl_c() => {
+                                                interrupted = true;
+                                                break;
+                                            }
+                                        }
                                     }
 
                                     // Read the validation status before disabling.
@@ -656,6 +663,13 @@ pub(crate) async fn run(
                                         .command_tx
                                         .send(Some(ValidationCommand::Disable));
 
+                                    if interrupted {
+                                        eprintln!("Interrupt received during checkpoint validation, stopping...");
+                                        shutdown_token.cancel();
+                                        etl_pipeline.cancel();
+                                        break Some("Interrupted by user".to_string());
+                                    }
+
                                     if timed_out {
                                         eprintln!(
                                             "Checkpoint {} validation timed out after {}s without convergence, aborting run",