Improvements

krinart · krinart · commit 1efa923488d4 · 2026-05-19T16:02:21.000-07:00
diff --git a/.github/workflows/run_spicebench.yml b/.github/workflows/run_spicebench.yml
@@ -14,14 +14,15 @@ on:
         options:
           - tpch
       system_under_test:
-        description: 'System under test (spice_cloud via spidapter docker image, or local databricks adapter modes)'
+        description: 'System under test (spice_cloud via spidapter docker image, local databricks adapter modes, or postgres via spidapter local backend)'
         required: true
         default: spice_cloud
         type: choice
         options:
           - spice_cloud
           - databricks-sql
           - databricks-lakebase
+          - postgres
       etl_type:
         description: 'ETL type'
         required: true
@@ -64,15 +65,15 @@ jobs:
           client-secret: ${{ secrets.SPICE_MANAGEMENT_CLIENT_SECRET_PROD }}
 
       - name: Log in to GHCR
-        if: ${{ env.SYSTEM_UNDER_TEST == 'spice_cloud' }}
+        if: ${{ env.SYSTEM_UNDER_TEST == 'spice_cloud' || env.SYSTEM_UNDER_TEST == 'postgres' }}
         uses: docker/login-action@v3
         with:
           registry: ghcr.io
           username: ${{ github.actor }}
           password: ${{ secrets.GITHUB_TOKEN }}
 
       - name: pull spidapter image
-        if: ${{ env.SYSTEM_UNDER_TEST == 'spice_cloud' }}
+        if: ${{ env.SYSTEM_UNDER_TEST == 'spice_cloud' || env.SYSTEM_UNDER_TEST == 'postgres' }}
         run: docker pull ghcr.io/spiceai/spidapter:latest
 
       - uses: ./.github/actions/build-spicebench
@@ -112,6 +113,9 @@ jobs:
           DATABRICKS_SQL_WAREHOUSE_ID: ${{ secrets.DATABRICKS_SQL_WAREHOUSE_ID }}
           DATABRICKS_CATALOG: ${{ secrets.DATABRICKS_CATALOG }}
           DATABRICKS_SCHEMA: ${{ secrets.DATABRICKS_SCHEMA }}
+          PG_HOST: ${{ secrets.POSTGRES_PG_HOST }}
+          PG_USER: ${{ secrets.POSTGRES_PG_USER }}
+          PG_DATABASE: ${{ vars.POSTGRES_PG_DATABASE }}
         run: |
           set -euo pipefail
           SYSTEM_UNDER_TEST_PREFIX="${SYSTEM_UNDER_TEST%%-*}"
@@ -165,6 +169,25 @@ jobs:
               "${HOME}/.spice/bin/databricks-system-adapter" --help >/dev/null
               ;;
 
+            postgres)
+              for required_var in PG_HOST PG_USER PG_DATABASE; do
+                if [ -z "${!required_var:-}" ]; then
+                  echo "${required_var} must be set for postgres adapter mode"
+                  exit 1
+                fi
+              done
+
+              if ! command -v docker >/dev/null 2>&1; then
+                echo "docker is required for postgres mode"
+                exit 1
+              fi
+
+              docker image inspect ghcr.io/spiceai/spidapter:latest >/dev/null 2>&1 || {
+                echo "spidapter docker image not found locally; pull step may have failed"
+                exit 1
+              }
+              ;;
+
             *)
               echo "Unsupported system_under_test value: ${SYSTEM_UNDER_TEST}"
               exit 1
@@ -196,7 +219,7 @@ jobs:
           sudo ldconfig
 
       - name: Install ADBC Postgres driver
-        if: ${{ startsWith(env.SYSTEM_UNDER_TEST, 'databricks-') }}
+        if: ${{ startsWith(env.SYSTEM_UNDER_TEST, 'databricks-') || env.SYSTEM_UNDER_TEST == 'postgres' }}
         uses: columnar-tech/setup-dbc@v1
         with:
           drivers: postgresql
@@ -239,6 +262,11 @@ jobs:
           LAKEBASE_PG_SCHEMA: ${{ vars.LAKEBASE_PG_SCHEMA }}
           LAKEBASE_PROJECT: ${{ vars.LAKEBASE_PROJECT }}
           LAKEBASE_BRANCH: ${{ vars.LAKEBASE_BRANCH }}
+          PG_HOST: ${{ vars.POSTGRES_PG_HOST }}
+          PG_PORT: ${{ vars.POSTGRES_PG_PORT || '5432' }}
+          PG_USER: ${{ secrets.POSTGRES_PG_USER }}
+          PG_PASSWORD: ${{ secrets.POSTGRES_PG_PASSWORD }}
+          PG_DATABASE: ${{ vars.POSTGRES_PG_DATABASE }}
           SPIDAPTER_ICEBERG_REGION: us-west-1
           SPIDAPTER_ICEBERG_CATALOG_FROM: iceberg:https://glue.us-west-1.amazonaws.com/iceberg/v1/catalogs/211125479522/namespaces
           SPIDAPTER_APP_MEMORY_LIMIT: '62Gi'
@@ -307,6 +335,20 @@ jobs:
             if [ -n "${DATABRICKS_STAGING_VOLUME_PATH:-}" ]; then
               ADAPTER_ENVS="${ADAPTER_ENVS} --system-adapter-env DATABRICKS_STAGING_VOLUME_PATH=${DATABRICKS_STAGING_VOLUME_PATH}"
             fi
+
+            if [ "${SYSTEM_UNDER_TEST}" = "databricks-lakebase" ]; then
+              ADAPTER_ENVS="${ADAPTER_ENVS} --system-adapter-env DATABRICKS_COMPUTE_MODE=lakebase"
+              ADAPTER_ENVS="${ADAPTER_ENVS} --system-adapter-env LAKEBASE_PG_HOST=${LAKEBASE_PG_HOST}"
+              ADAPTER_ENVS="${ADAPTER_ENVS} --system-adapter-env LAKEBASE_PG_USER=${LAKEBASE_PG_USER}"
+              ADAPTER_ENVS="${ADAPTER_ENVS} --system-adapter-env LAKEBASE_PG_SCHEMA=${LAKEBASE_PG_SCHEMA}"
+              ADAPTER_ENVS="${ADAPTER_ENVS} --system-adapter-env LAKEBASE_PROJECT=${LAKEBASE_PROJECT}"
+              ADAPTER_ENVS="${ADAPTER_ENVS} --system-adapter-env LAKEBASE_BRANCH=${LAKEBASE_BRANCH}"
+            fi
+          elif [ "${SYSTEM_UNDER_TEST_PREFIX}" = "postgres" ]; then
+            export SPICEBENCH_ADBC_UPDATE_STRATEGY=statement
+            ADAPTER_CMD="docker"
+            ADAPTER_ARGS="run -i -e PG_HOST=${PG_HOST} -e PG_PORT=${PG_PORT} -e PG_USER=${PG_USER} -e PG_PASSWORD=${PG_PASSWORD} -e PG_DATABASE=${PG_DATABASE} ghcr.io/spiceai/spidapter:latest stdio --backend local --deployment-mode single-node"
+            ADAPTER_ENVS=""
           else
             export SPICEBENCH_ADBC_UPDATE_STRATEGY=bulk_ingest_upsert
             export SPICEBENCH_ADBC_FLUSH_STREAM_BEFORE_UPSERT=true
@@ -326,15 +368,6 @@ jobs:
             ADAPTER_ENVS=""
           fi
 
-          if [ "${SYSTEM_UNDER_TEST}" = "databricks-lakebase" ]; then
-            ADAPTER_ENVS="${ADAPTER_ENVS} --system-adapter-env DATABRICKS_COMPUTE_MODE=lakebase"
-            ADAPTER_ENVS="${ADAPTER_ENVS} --system-adapter-env LAKEBASE_PG_HOST=${LAKEBASE_PG_HOST}"
-            ADAPTER_ENVS="${ADAPTER_ENVS} --system-adapter-env LAKEBASE_PG_USER=${LAKEBASE_PG_USER}"
-            ADAPTER_ENVS="${ADAPTER_ENVS} --system-adapter-env LAKEBASE_PG_SCHEMA=${LAKEBASE_PG_SCHEMA}"
-            ADAPTER_ENVS="${ADAPTER_ENVS} --system-adapter-env LAKEBASE_PROJECT=${LAKEBASE_PROJECT}"
-            ADAPTER_ENVS="${ADAPTER_ENVS} --system-adapter-env LAKEBASE_BRANCH=${LAKEBASE_BRANCH}"
-          fi
-
           ~/.spice/bin/spicebench run \
             --concurrency "${NUM_QUERY_CLIENTS}"  \
             --scenario "${SCENARIO}" \
diff --git a/crates/adbc_client/src/lib.rs b/crates/adbc_client/src/lib.rs
@@ -118,7 +118,7 @@ impl AdbcConnection {
 
         Ok(Self::new(
             conn,
-            driver_name == "databricks",
+            driver_name == "databricks" || driver_name.eq_ignore_ascii_case("dynamodb"),
             driver_name == "postgresql",
             driver_name == "postgresql" || driver_name == "databricks",
         ))
diff --git a/crates/adbc_client/src/pool.rs b/crates/adbc_client/src/pool.rs
@@ -150,9 +150,11 @@ pub fn create_pool(
             reason: e.to_string(),
         })?;
 
+    let downcast_utf8view =
+        driver_name == "databricks" || driver_name.eq_ignore_ascii_case("dynamodb");
     let manager = AdbcConnectionManager::new(
         db,
-        driver_name == "databricks",
+        downcast_utf8view,
         driver_name == "postgresql",
         driver_name == "postgresql" || driver_name == "databricks",
     );
diff --git a/crates/etl/src/lib.rs b/crates/etl/src/lib.rs
@@ -917,7 +917,7 @@ async fn write_segments_for_batch(
                     partition_columns,
                 )
                 .await
-                .map_err(|e| format!("write {table_name} batch {batch_id}: {e}"))
+                .map_err(|e| format!("write {table_name} batch {batch_id}: {e:#}"))
         });
     }
 
@@ -1959,28 +1959,49 @@ async fn run_pipeline(
         }
 
         // Collect results from all concurrent table tasks in this step.
+        // Also watch for cancellation so a ctrl-c that arrives while tasks are
+        // blocked in a slow sink write (e.g. DynamoDB ADBC bulk ingest) doesn't
+        // leave the pipeline stuck waiting for the current step to finish.
         let mut step_batch_count: u64 = 0;
         let mut step_rows_count: u64 = 0;
-        while let Some(result) = join_set.join_next().await {
-            match result {
-                Ok(Ok((table_name, is_finished, consumed_work_units, rows_read))) => {
-                    step_batch_count += consumed_work_units;
-                    step_rows_count += rows_read;
-                    if is_finished {
-                        let mut state = work_state.lock().expect("work_state lock poisoned");
-                        state.finished_tables.insert(table_name);
-                        tables_finished_counter.fetch_add(1, Ordering::Relaxed);
-                    }
-                }
-                Ok(Err(err_msg)) => {
+        loop {
+            tokio::select! {
+                biased;
+                () = cancel.cancelled() => {
+                    join_set.abort_all();
                     progress_logger.abort();
-                    return PipelineState::Stopped(StopReason::Error(err_msg));
+                    return PipelineState::Stopped(StopReason::Cancelled);
                 }
-                Err(e) => {
-                    progress_logger.abort();
-                    return PipelineState::Stopped(StopReason::Error(format!(
-                        "Task panicked: {e}"
-                    )));
+                result = join_set.join_next() => {
+                    let Some(result) = result else { break; };
+                    match result {
+                        Ok(Ok((table_name, is_finished, consumed_work_units, rows_read))) => {
+                            step_batch_count += consumed_work_units;
+                            step_rows_count += rows_read;
+                            if is_finished {
+                                let mut state = work_state.lock().expect("work_state lock poisoned");
+                                state.finished_tables.insert(table_name);
+                                tables_finished_counter.fetch_add(1, Ordering::Relaxed);
+                            }
+                        }
+                        Ok(Err(err_msg)) => {
+                            progress_logger.abort();
+                            return PipelineState::Stopped(StopReason::Error(err_msg));
+                        }
+                        Err(e) if e.is_cancelled() => {
+                            // Task was aborted (e.g. by abort_all above on a
+                            // concurrent iteration); treat as a soft cancel.
+                            join_set.abort_all();
+                            progress_logger.abort();
+                            return PipelineState::Stopped(StopReason::Cancelled);
+                        }
+                        Err(e) => {
+                            progress_logger.abort();
+                            return PipelineState::Stopped(StopReason::Error(format!(
+                                "Task panicked: {e}"
+                            )));
+                        }
+                    }
                 }
             }
         }
diff --git a/crates/etl/src/sink/adbc.rs b/crates/etl/src/sink/adbc.rs
@@ -178,6 +178,10 @@ pub struct AdbcSink {
     bulk_ingest_stream_buffer: usize,
     /// Optional system adapter client for staging table creation.
     staging_adapter: Option<(Arc<Mutex<SystemAdapterClient>>, Uuid)>,
+    /// Optional mapping from logical dataset name to physical table name.
+    /// When set, overrides the table name used for ADBC bulk ingest so the
+    /// sink writes to the correct physical table (e.g. DynamoDB prefixed names).
+    table_name_map: HashMap<String, String>,
 }
 
 impl AdbcSink {
@@ -243,6 +247,7 @@ impl AdbcSink {
         target_db_catalog: Option<String>,
         target_db_schema: Option<String>,
         staging_adapter: Option<(Arc<Mutex<SystemAdapterClient>>, Uuid)>,
+        table_name_map: HashMap<String, String>,
     ) -> anyhow::Result<Self> {
         let update_strategy = UpdateStrategy::from_env()?;
         let pool_size = Self::pool_size();
@@ -278,6 +283,7 @@ impl AdbcSink {
             flush_stream_before_upsert,
             bulk_ingest_stream_buffer,
             staging_adapter,
+            table_name_map,
         })
     }
 
@@ -386,11 +392,32 @@ impl AdbcSink {
         };
 
         for sub_batch in sub_batches {
-            sender.send(sub_batch).await.map_err(|_| {
-                anyhow::anyhow!(
-                    "Bulk ingest stream for table '{table_name}' is no longer available"
-                )
-            })?;
+            if sender.send(sub_batch).await.is_err() {
+                // Worker exited before receiving this batch. Remove the stream and
+                // await the worker to surface the actual error rather than a generic
+                // "no longer available" message.
+                let stream = {
+                    let mut streams = self.bulk_ingest_streams.write().await;
+                    streams.remove(table_name)
+                };
+                let err = if let Some(stream) = stream {
+                    match stream.worker.await {
+                        Ok(Ok(())) => anyhow::anyhow!(
+                            "Bulk ingest worker for '{table_name}' exited without error but before all data was sent"
+                        ),
+                        Ok(Err(worker_err)) => worker_err
+                            .context(format!("Bulk ingest worker for '{table_name}' failed")),
+                        Err(join_err) => anyhow::anyhow!(
+                            "Bulk ingest worker for '{table_name}' panicked: {join_err}"
+                        ),
+                    }
+                } else {
+                    anyhow::anyhow!(
+                        "Bulk ingest stream for table '{table_name}' is no longer available"
+                    )
+                };
+                return Err(err);
+            }
             batches_sent.fetch_add(1, Ordering::Relaxed);
         }
 
@@ -530,7 +557,10 @@ impl AdbcSink {
     }
 
     fn target_table_ingest_name(&self, table_name: &str) -> String {
-        table_name.to_string()
+        self.table_name_map
+            .get(table_name)
+            .cloned()
+            .unwrap_or_else(|| table_name.to_string())
     }
 
     fn create_table_sql(
diff --git a/crates/system-adapter-protocol/src/lib.rs b/crates/system-adapter-protocol/src/lib.rs
@@ -220,6 +220,12 @@ pub struct SetupResponse {
     ///   (optional). Used to benchmark the distributed (Ballista) query path.
     #[serde(default, skip_serializing_if = "HashMap::is_empty")]
     pub endpoints: HashMap<String, HashMap<String, serde_json::Value>>,
+    /// Optional mapping from logical dataset name to physical table name.
+    /// When set, the ETL sink will write to the physical name instead of the
+    /// logical dataset name (e.g. DynamoDB uses timestamped table name prefixes
+    /// to avoid collisions between concurrent benchmark runs).
+    #[serde(default, skip_serializing_if = "HashMap::is_empty")]
+    pub table_name_map: HashMap<String, String>,
 }
 /// Request to teardown a benchmark run
 ///
diff --git a/crates/system-adapter-protocol/src/server.rs b/crates/system-adapter-protocol/src/server.rs
@@ -263,7 +263,13 @@ impl<H: Handler> Server<H> {
         };
         Self::handler_response(
             self.handler
-                .setup(req.run_id, req.metadata, req.datasets, req.etl_sink_type, req.seed_data)
+                .setup(
+                    req.run_id,
+                    req.metadata,
+                    req.datasets,
+                    req.etl_sink_type,
+                    req.seed_data,
+                )
                 .await,
             id,
         )
diff --git a/src/commands/etl_cmd.rs b/src/commands/etl_cmd.rs
@@ -139,6 +139,7 @@ pub async fn execute(args: &EtlArgs) -> anyhow::Result<()> {
                 args.adbc_catalog.clone(),
                 args.adbc_schema.clone(),
                 None,
+                std::collections::HashMap::new(),
             )?);
 
             (
diff --git a/src/commands/run.rs b/src/commands/run.rs
@@ -173,6 +173,7 @@ async fn run_benchmark(
             target_db_catalog,
             target_db_schema,
             Some((Arc::clone(&system_adapter_client), run_id)),
+            setup_response.table_name_map.clone(),
         )?);
 
         let mut pipeline = ETLPipeline::new(
@@ -428,8 +429,31 @@ fn make_zero_batch(
                 DataType::Float64 => Arc::new(Float64Array::from(vec![0.0f64; n_rows])),
                 DataType::Utf8 => Arc::new(StringArray::from(vec![""; n_rows])),
                 DataType::LargeUtf8 => Arc::new(LargeStringArray::from(vec![""; n_rows])),
+                DataType::Utf8View => Arc::new(StringViewArray::from(vec![""; n_rows])),
                 DataType::Date32 => Arc::new(Date32Array::from(vec![0i32; n_rows])),
                 DataType::Date64 => Arc::new(Date64Array::from(vec![0i64; n_rows])),
+                DataType::Timestamp(unit, tz) => {
+                    use arrow::datatypes::TimeUnit;
+                    let arr: ArrayRef = match unit {
+                        TimeUnit::Second => Arc::new(
+                            arrow::array::TimestampSecondArray::from(vec![0i64; n_rows])
+                                .with_timezone_opt(tz.clone()),
+                        ),
+                        TimeUnit::Millisecond => Arc::new(
+                            arrow::array::TimestampMillisecondArray::from(vec![0i64; n_rows])
+                                .with_timezone_opt(tz.clone()),
+                        ),
+                        TimeUnit::Microsecond => Arc::new(
+                            arrow::array::TimestampMicrosecondArray::from(vec![0i64; n_rows])
+                                .with_timezone_opt(tz.clone()),
+                        ),
+                        TimeUnit::Nanosecond => Arc::new(
+                            arrow::array::TimestampNanosecondArray::from(vec![0i64; n_rows])
+                                .with_timezone_opt(tz.clone()),
+                        ),
+                    };
+                    arr
+                }
                 DataType::Decimal128(p, s) => Arc::new(
                     Decimal128Array::from(vec![0i128; n_rows])
                         .with_precision_and_scale(*p, *s)
diff --git a/system-adapters/databricks/src/main.rs b/system-adapters/databricks/src/main.rs
@@ -2140,6 +2140,7 @@ impl Handler for DatabricksAdapter {
                         HashMap::from([("uri".to_string(), Value::String(pg_uri))]),
                     )),
                     endpoints: HashMap::new(),
+                    table_name_map: HashMap::new(),
                 })
             }
             // For other variants, return a single Databricks ADBC driver.
@@ -2155,6 +2156,7 @@ impl Handler for DatabricksAdapter {
                 catalog_namespace: Some(format!("{}.{}", self.config.catalog, self.config.schema)),
                 read_driver: None,
                 endpoints: HashMap::new(),
+                table_name_map: HashMap::new(),
             }),
         }
     }

Original file line number	Diff line number	Diff line change
`@@ -139,6 +139,7 @@ pub async fn execute(args: &EtlArgs) -> anyhow::Result<()> {`
`139`	`139`	`args.adbc_catalog.clone(),`
`140`	`140`	`args.adbc_schema.clone(),`
`141`	`141`	`None,`
	`142`	`+ std::collections::HashMap::new(),`
`142`	`143`	`)?);`
`143`	`144`
`144`	`145`	`(`
Original file line number	Diff line number	Diff line change
`@@ -2140,6 +2140,7 @@ impl Handler for DatabricksAdapter {`
`2140`	`2140`	`HashMap::from([("uri".to_string(), Value::String(pg_uri))]),`
`2141`	`2141`	`)),`
`2142`	`2142`	`endpoints: HashMap::new(),`
	`2143`	`+ table_name_map: HashMap::new(),`
`2143`	`2144`	`})`
`2144`	`2145`	`}`
`2145`	`2146`	`// For other variants, return a single Databricks ADBC driver.`
`@@ -2155,6 +2156,7 @@ impl Handler for DatabricksAdapter {`
`2155`	`2156`	`catalog_namespace: Some(format!("{}.{}", self.config.catalog, self.config.schema)),`
`2156`	`2157`	`read_driver: None,`
`2157`	`2158`	`endpoints: HashMap::new(),`
	`2159`	`+ table_name_map: HashMap::new(),`
`2158`	`2160`	`}),`
`2159`	`2161`	`}`
`2160`	`2162`	`}`