databricks: include spill_to_disk_bytes into disk_write_bytes metric (#208)

Author: sgrebnov

system-adapters/databricks/src/main.rs

@@ -928,8 +928,8 @@ impl DatabricksAdapter {
         }
     }
 
-    /// Sum `read_bytes` + `write_remote_bytes` from query history for all
-    /// FINISHED queries on this warehouse since `start_time_ms`.
+    /// Sum `read_bytes` and `write_remote_bytes + spill_to_disk_bytes` from query
+    /// history for all FINISHED queries on this warehouse since `start_time_ms`.
     ///
     /// Uses the REST Query History API (`/api/2.0/sql/history/queries`).
     ///
@@ -942,15 +942,18 @@ impl DatabricksAdapter {
     }
 
     /// REST path: paginate through `/api/2.0/sql/history/queries` and sum
-    /// per-query `metrics.read_bytes` and `metrics.write_remote_bytes`.
+    /// per-query `metrics.read_bytes` and `metrics.write_remote_bytes + spill_to_disk_bytes`.
     async fn sum_query_history_io_rest(&self, start_time_ms: u64) -> Result<(u64, u64)> {
         let history_url = format!(
             "https://{}/api/2.0/sql/history/queries",
             self.config.endpoint
         );
 
-        let mut total_read: u64 = 0;
-        let mut total_write: u64 = 0;
+        let mut total_read_bytes: u64 = 0;
+        let mut total_read_remote_bytes: u64 = 0;
+        let mut total_read_cache_bytes: u64 = 0;
+        let mut total_write_remote_bytes: u64 = 0;
+        let mut total_spill_to_disk_bytes: u64 = 0;
         let mut page_token: Option<String> = None;
 
         loop {
@@ -988,8 +991,11 @@ impl DatabricksAdapter {
 
             for entry in &body.res {
                 if let Some(ref m) = entry.metrics {
-                    total_read += m.read_bytes.unwrap_or(0);
-                    total_write += m.write_remote_bytes.unwrap_or(0);
+                    total_read_bytes += m.read_bytes.unwrap_or(0);
+                    total_read_remote_bytes += m.read_remote_bytes.unwrap_or(0);
+                    total_read_cache_bytes += m.read_cache_bytes.unwrap_or(0);
+                    total_write_remote_bytes += m.write_remote_bytes.unwrap_or(0);
+                    total_spill_to_disk_bytes += m.spill_to_disk_bytes.unwrap_or(0);
                 }
             }
 
@@ -1000,7 +1006,15 @@ impl DatabricksAdapter {
             }
         }
 
-        Ok((total_read, total_write))
+        let total_write_bytes = total_write_remote_bytes + total_spill_to_disk_bytes;
+
+        eprintln!(
+            "[databricks-adapter] query history I/O breakdown: \
+             total_read_bytes={total_read_bytes} (total_read_remote_bytes={total_read_remote_bytes}, total_read_cache_bytes={total_read_cache_bytes}), \
+             total_write_bytes={total_write_bytes} (total_write_remote_bytes={total_write_remote_bytes}, total_spill_to_disk_bytes={total_spill_to_disk_bytes})"
+        );
+
+        Ok((total_read_bytes, total_write_bytes))
     }
 
     async fn ensure_cluster_ready(&self) -> Result<(String, bool)> {
@@ -1796,11 +1810,20 @@ struct QueryHistoryMetrics {
     /// and local SSD/disk cache (`read_cache_bytes`). Maps to `disk_read_bytes`.
     #[serde(default)]
     read_bytes: Option<u64>,
-    /// Bytes written to remote cloud storage (S3/ADLS/GCS). This is the only write metric
-    /// available in the REST API — non-zero for DDL/DML that materializes data (e.g. CTAS).
-    /// Maps to `disk_write_bytes`.
+    /// Bytes read from remote cloud storage (S3/ADLS/GCS).
+    #[serde(default)]
+    read_remote_bytes: Option<u64>,
+    /// Bytes read from the local SSD/disk cache.
+    #[serde(default)]
+    read_cache_bytes: Option<u64>,
+    /// Bytes written to remote cloud storage (S3/ADLS/GCS). Non-zero for DDL/DML that
+    /// materializes data (e.g. CTAS). Summed with `spill_to_disk_bytes` for `disk_write_bytes`.
     #[serde(default)]
     write_remote_bytes: Option<u64>,
+    /// Bytes temporarily written to local disk when query execution exceeds available memory.
+    /// Summed with `write_remote_bytes` for `disk_write_bytes`.
+    #[serde(default)]
+    spill_to_disk_bytes: Option<u64>,
 }
 
 /// Response from GET /api/2.0/sql/history/queries
@@ -1925,11 +1948,10 @@ impl Handler for DatabricksAdapter {
         // Create managed UC tables (sources for synced tables) via SQL Warehouse.
         // Drop any stale tables from a previous failed run first to ensure idempotent setup.
         for (table_name, dataset_cfg) in &datasets {
-            let drop_sql = format!(
-                "DROP TABLE IF EXISTS {}",
-                self.table_full_name(table_name)
+            let drop_sql = format!("DROP TABLE IF EXISTS {}", self.table_full_name(table_name));
+            eprintln!(
+                "[databricks-adapter] dropping stale managed table (if exists) '{table_name}': {drop_sql}"
             );
-            eprintln!("[databricks-adapter] dropping stale managed table (if exists) '{table_name}': {drop_sql}");
             self.execute_sql_statement(&drop_sql)
                 .await
                 .map_err(|e| format!("Failed to drop stale managed table '{table_name}': {e}"))?;
@@ -2257,13 +2279,14 @@ impl Handler for DatabricksAdapter {
                     }
                 }
 
-                // Sum read_bytes and write_bytes from all queries since the run started.
-                // On periodic scrapes this is best-effort (Query History has ~5 min lag).
-                // On final scrape the marker wait above ensures completeness.
+                // Sum read_bytes and (write_remote_bytes + spill_to_disk_bytes) from all
+                // queries since the run started. On periodic scrapes this is best-effort
+                // (Query History has ~5 min lag). On final scrape the marker wait above
+                // ensures completeness.
                 match self.sum_query_history_io(started_at_ms).await {
                     Ok((total_read, total_write)) => {
                         eprintln!(
-                            "[databricks-adapter] query history totals: read_bytes={total_read} write_remote_bytes={total_write}"
+                            "[databricks-adapter] query history totals: read_bytes={total_read} write_bytes={total_write}"
                         );
                         resource.disk_read_bytes = Some(total_read);
                         resource.disk_write_bytes = Some(total_write);