feat: Add support for Databricks variant in the system adapter and update related configurations

Author: lukekim

.github/workflows/run_spicebench.yml

@@ -17,6 +17,14 @@ on:
         options:
           - spidapter
           - databricks
+      databricks_variant:
+        description: 'Databricks adapter variant (databricks or lakebase)'
+        required: false
+        default: 'databricks'
+        type: choice
+        options:
+          - databricks
+          - lakebase
       etl_bucket:
         description: 'S3 bucket for ETL source and target data'
         required: true
@@ -117,6 +125,7 @@ jobs:
           DATABRICKS_SQL_WAREHOUSE_ID: ${{ secrets.DATABRICKS_SQL_WAREHOUSE_ID }}
           DATABRICKS_CATALOG: ${{ secrets.DATABRICKS_CATALOG }}
           DATABRICKS_SCHEMA: ${{ secrets.DATABRICKS_SCHEMA }}
+          DATABRICKS_VARIANT: ${{ github.event.inputs.databricks_variant || 'databricks' }}
           DATABRICKS_TABLE_FORMAT: ${{ github.event.inputs.table_format || 'parquet' }}
         run: |
           set -euo pipefail
@@ -145,6 +154,11 @@ jobs:
               ;;
 
             databricks)
+              if [ "${DATABRICKS_VARIANT}" != "databricks" ] && [ "${DATABRICKS_VARIANT}" != "lakebase" ]; then
+                echo "DATABRICKS_VARIANT must be either 'databricks' or 'lakebase'"
+                exit 1
+              fi
+
               for required_var in DATABRICKS_ENDPOINT DATABRICKS_TOKEN DATABRICKS_HTTP_PATH DATABRICKS_SQL_WAREHOUSE_ID; do
                 if [ -z "${!required_var:-}" ]; then
                   echo "${required_var} must be set for databricks adapter mode"
@@ -200,6 +214,7 @@ jobs:
           DATABRICKS_SQL_WAREHOUSE_ID: ${{ secrets.DATABRICKS_SQL_WAREHOUSE_ID }}
           DATABRICKS_CATALOG: ${{ secrets.DATABRICKS_CATALOG }}
           DATABRICKS_SCHEMA: ${{ secrets.DATABRICKS_SCHEMA }}
+          DATABRICKS_VARIANT: ${{ github.event.inputs.databricks_variant || 'databricks' }}
           DATABRICKS_TABLE_FORMAT: ${{ github.event.inputs.table_format || 'parquet' }}
           SPICEAI_BENCHMARK_METRICS_KEY: ${{ secrets.SPICEAI_BENCHMARK_METRICS_KEY }}
           SCENARIO: ${{ github.event.inputs.scenario || 'tpch' }}
@@ -228,7 +243,7 @@ jobs:
           if [ "${SYSTEM_ADAPTER}" = "databricks" ]; then
             ADAPTER_CMD="${HOME}/.spice/bin/databricks-system-adapter"
             ADAPTER_ARGS="stdio"
-            ADAPTER_ENVS="--system-adapter-env DATABRICKS_ENDPOINT=${DATABRICKS_ENDPOINT} --system-adapter-env DATABRICKS_TOKEN=${DATABRICKS_TOKEN} --system-adapter-env DATABRICKS_HTTP_PATH=${DATABRICKS_HTTP_PATH} --system-adapter-env DATABRICKS_SQL_WAREHOUSE_ID=${DATABRICKS_SQL_WAREHOUSE_ID} --system-adapter-env DATABRICKS_TABLE_FORMAT=${DATABRICKS_TABLE_FORMAT}"
+            ADAPTER_ENVS="--system-adapter-env DATABRICKS_ENDPOINT=${DATABRICKS_ENDPOINT} --system-adapter-env DATABRICKS_TOKEN=${DATABRICKS_TOKEN} --system-adapter-env DATABRICKS_HTTP_PATH=${DATABRICKS_HTTP_PATH} --system-adapter-env DATABRICKS_SQL_WAREHOUSE_ID=${DATABRICKS_SQL_WAREHOUSE_ID} --system-adapter-env DATABRICKS_VARIANT=${DATABRICKS_VARIANT} --system-adapter-env DATABRICKS_TABLE_FORMAT=${DATABRICKS_TABLE_FORMAT}"
 
             if [ -n "${DATABRICKS_CATALOG}" ]; then
               ADAPTER_ENVS="${ADAPTER_ENVS} --system-adapter-env DATABRICKS_CATALOG=${DATABRICKS_CATALOG}"

README.md

@@ -308,10 +308,14 @@ spicebench \
     --system-adapter-env DATABRICKS_TOKEN=$DATABRICKS_TOKEN \
     --system-adapter-env DATABRICKS_HTTP_PATH=$DATABRICKS_HTTP_PATH \
     --system-adapter-env DATABRICKS_SQL_WAREHOUSE_ID=$DATABRICKS_SQL_WAREHOUSE_ID \
+    --system-adapter-env DATABRICKS_VARIANT=databricks \
+    --system-adapter-env DATABRICKS_TABLE_FORMAT=parquet \
     --system-adapter-env DATABRICKS_CATALOG=spiceai_sandbox \
     --system-adapter-env DATABRICKS_SCHEMA=tpch
 ```
 
+Set `DATABRICKS_VARIANT=lakebase` to enable Lakebase setup mode, which creates schema/table DDL during `setup()` and drops tables during `teardown()` when `DATABRICKS_DROP_TABLES_ON_TEARDOWN=true`.
+
 ### Crate Overview
 
 | Crate                     | Description                                                                                |

system-adapters/databricks/Cargo.lock

@@ -9,6 +9,7 @@ repository = "https://github.com/spiceai/spicebench"
 
 [dependencies]
 anyhow = "1.0.99"
+arrow-schema = "57.0.0"
 async-trait = "0.1.89"
 clap = { version = "4.5.54", features = ["derive", "env"] }
 reqwest = { version = "=0.12.24", default-features = false, features = ["json", "rustls-tls"] }