Analyze flightsql

.gitignore

@@ -105,4 +105,6 @@ bench/url_files/*
 # Tags
 tags
 
+# Local scratch / per-machine files
+notes
 .claude/settings.local.json

Cargo.toml

@@ -42,6 +42,7 @@ prost = "0.14"
 ratatui = { optional = true, version = "0.30" }
 ratatui-textarea = { features = ["search"], optional = true, version = "0.8" }
 serde = { features = ["derive"], version = "1.0.197" }
+serde_json = "1.0"
 strum = { features = ["derive"], version = "0.26.2" }
 tokio = { features = [
   "macros",

crates/datafusion-app/Cargo.toml

@@ -44,10 +44,12 @@ opendal = { features = [
   "services-huggingface",
 ], git = "https://github.com/apache/opendal", optional = true, rev = "24aaff9c62b1" }
 parking_lot = "0.12.3"
+prost = { optional = true, version = "0.14" }
 rustls = { default-features = false, features = [
   "aws-lc-rs",
 ], optional = true, version = "0.23" }
 serde = { features = ["derive"], version = "1.0.197" }
+serde_json = { optional = true, version = "1.0" }
 tokio = { features = ["macros", "rt-multi-thread"], version = "1.36.0" }
 tokio-metrics = { features = [
   "metrics-rs-integration",
@@ -62,6 +64,7 @@ vortex-datafusion = { optional = true, version = "0.78" }
 
 [dev-dependencies]
 criterion = { features = ["async_tokio"], version = "0.5.1" }
+tempfile = "3"
 
 [features]
 clickhouse = [
@@ -70,7 +73,13 @@ clickhouse = [
 ]
 default = ["functions-parquet"]
 deltalake = ["dep:deltalake"]
-flightsql = ["dep:arrow-flight", "dep:base64", "dep:tonic"]
+flightsql = [
+  "dep:arrow-flight",
+  "dep:base64",
+  "dep:prost",
+  "dep:serde_json",
+  "dep:tonic",
+]
 functions-arrow = ["dep:datafusion-functions-arrow"]
 functions-json = ["dep:datafusion-functions-json"]
 functions-parquet = ["dep:datafusion-functions-parquet"]

crates/datafusion-app/src/flightsql.rs

@@ -483,4 +483,102 @@ impl FlightSQLContext {
             ))
         }
     }
+
+    /// Get raw metrics batch without reconstruction (for --analyze-raw)
+    pub async fn analyze_query_raw(
+        &self,
+        query: &str,
+    ) -> Result<(String, datafusion::arrow::array::RecordBatch)> {
+        self.fetch_analyze_batches(query).await
+    }
+
+    /// Reconstruct ExecutionStats from metrics (for --analyze)
+    pub async fn analyze_query(&self, query: &str) -> Result<crate::stats::ExecutionStats> {
+        let (query_str, metrics_batch) = self.fetch_analyze_batches(query).await?;
+
+        // Reconstruct ExecutionStats from metrics table
+        let stats = crate::stats::ExecutionStats::from_metrics_table(metrics_batch, query_str)?;
+
+        Ok(stats)
+    }
+
+    /// Shared logic to fetch analyze batch and query from server
+    async fn fetch_analyze_batches(
+        &self,
+        query: &str,
+    ) -> Result<(String, datafusion::arrow::array::RecordBatch)> {
+        use crate::stats::{
+            is_compatible_protocol_version, PROTOCOL_VERSION_METADATA_KEY, QUERY_ID_METADATA_KEY,
+        };
+        use arrow_flight::utils::flight_data_to_batches;
+        use arrow_flight::{Action, FlightData};
+
+        // Validate that query contains only a single statement
+        let dialect = datafusion::sql::sqlparser::dialect::GenericDialect {};
+        let statements = DFParser::parse_sql_with_dialect(query, &dialect)?;
+        if statements.len() != 1 {
+            return Err(eyre::eyre!("Only a single SQL statement can be analyzed"));
+        }
+
+        // 1. Create JSON request and encode as Action body
+        let request = crate::stats::AnalyzeQueryRequest::with_sql(query);
+        let request_body = serde_json::to_vec(&request)
+            .map_err(|e| eyre::eyre!("Failed to serialize request: {}", e))?;
+
+        let action = Action {
+            r#type: "analyze_query".to_string(),
+            body: request_body.into(),
+        };
+
+        // 2. Call do_action on the FlightSQL service
+        let mut client = self.client.lock().await;
+        let client = client
+            .as_mut()
+            .ok_or_else(|| eyre::eyre!("No FlightSQL client configured"))?;
+
+        let mut stream = client
+            .do_action(action.into_request())
+            .await
+            .map_err(|e| eyre::eyre!("do_action failed: {}", e))?;
+
+        // 3. Collect all Result messages, decoding each body as FlightData
+        let mut all_flight_data = Vec::new();
+        while let Some(result) = stream.next().await {
+            let result = result.map_err(|e| eyre::eyre!("Stream error: {}", e))?;
+            let flight_data = <FlightData as prost::Message>::decode(result.body.as_ref())
+                .map_err(|e| eyre::eyre!("Failed to decode FlightData: {}", e))?;
+            all_flight_data.push(flight_data);
+        }
+
+        // 4. Decode the metrics batches. The framing is a schema message
+        // followed by any number of data (and dictionary) messages; a server
+        // may legitimately split a large metrics table into several batches.
+        let metrics_batches = flight_data_to_batches(&all_flight_data)
+            .map_err(|e| eyre::eyre!("Failed to decode metrics batch: {}", e))?;
+
+        if metrics_batches.is_empty() {
+            return Err(eyre::eyre!("No metrics batch found in response"));
+        }
+
+        let schema = metrics_batches[0].schema();
+        if let Some(version) = schema.metadata().get(PROTOCOL_VERSION_METADATA_KEY) {
+            if !is_compatible_protocol_version(version) {
+                return Err(eyre::eyre!(
+                    "Server analyze protocol version {} is not compatible with client version {}",
+                    version,
+                    crate::stats::ANALYZE_PROTOCOL_VERSION
+                ));
+            }
+        }
+        if let Some(query_id) = schema.metadata().get(QUERY_ID_METADATA_KEY) {
+            debug!("Analyze response query_id: {query_id}");
+        }
+
+        let metrics_batch =
+            datafusion::arrow::compute::concat_batches(&schema, &metrics_batches)
+                .map_err(|e| eyre::eyre!("Failed to concatenate metrics batches: {}", e))?;
+
+        // Use the original query (client retains it, server doesn't send it back)
+        Ok((query.to_string(), metrics_batch))
+    }
 }

crates/datafusion-app/src/local.rs

@@ -665,7 +665,7 @@ impl ExecutionContext {
                 physical_plan,
             )
         } else {
-            Err(eyre::eyre!("Only a single statement can be benchmarked"))
+            Err(eyre::eyre!("Only a single statement can be analyzed"))
         }
     }