Expose DataFusion SQL explain output

Author: alexanderbianchi

quickwit/quickwit-df-core/proto/datafusion.proto

@@ -68,6 +68,11 @@ message ExecuteSqlRequest {
 
   // Optional per-request session overrides.
   map<string, string> properties = 2;
+
+  // When true, the server returns the service-level EXPLAIN output for the
+  // final SQL statement instead of executing it. DDL statements before the
+  // final query are still executed for side effects.
+  bool explain = 3;
 }
 
 message ExecuteSqlResponse {

quickwit/quickwit-df-core/src/codegen/quickwit.datafusion.rs

@@ -208,11 +208,17 @@ impl data_fusion_service_server::DataFusionService for DataFusionServiceGrpcImpl
         let req = request.into_inner();
         let service = Arc::clone(&self.service);
 
+        let output = if req.explain {
+            DataFusionOutput::Explain
+        } else {
+            DataFusionOutput::Records
+        };
         let execution = service
-            .execute(DataFusionRequest::records(
-                DataFusionInput::Sql(&req.sql),
-                &req.properties,
-            ))
+            .execute(DataFusionRequest {
+                input: DataFusionInput::Sql(&req.sql),
+                output,
+                properties: &req.properties,
+            })
             .await
             .map_err(|err| {
                 warn!(error = %err, "DataFusion SQL execution error");

quickwit/quickwit-df-core/src/grpc.rs

@@ -876,10 +876,10 @@ fn physical_plan_with_metrics(
     physical_plan: &Arc<dyn ExecutionPlan>,
 ) -> (Arc<dyn ExecutionPlan>, bool) {
     if physical_plan.as_any().is::<DistributedExec>() {
-        return match rewrite_distributed_plan_with_metrics(
+        return match futures::executor::block_on(rewrite_distributed_plan_with_metrics(
             Arc::clone(physical_plan),
             DistributedMetricsFormat::PerTask,
-        ) {
+        )) {
             Ok(physical_plan) => (physical_plan, true),
             Err(_) => (Arc::clone(physical_plan), false),
         };

quickwit/quickwit-df-core/src/service.rs

@@ -79,10 +79,26 @@ fn datafusion_client(
 async fn execute_datafusion_sql(
     client: &mut DataFusionServiceClient<tonic::transport::Channel>,
     sql: impl Into<String>,
+) -> Vec<RecordBatch> {
+    execute_datafusion_sql_with_explain(client, sql, false).await
+}
+
+async fn explain_datafusion_sql(
+    client: &mut DataFusionServiceClient<tonic::transport::Channel>,
+    sql: impl Into<String>,
+) -> Vec<RecordBatch> {
+    execute_datafusion_sql_with_explain(client, sql, true).await
+}
+
+async fn execute_datafusion_sql_with_explain(
+    client: &mut DataFusionServiceClient<tonic::transport::Channel>,
+    sql: impl Into<String>,
+    explain: bool,
 ) -> Vec<RecordBatch> {
     let request = ExecuteSqlRequest {
         sql: sql.into(),
         properties: Default::default(),
+        explain,
     };
     let mut stream = client.execute_sql(request).await.unwrap().into_inner();
     let mut batches = Vec::new();
@@ -234,8 +250,7 @@ async fn test_distributed_tasks_not_shuffles() {
     let query = r#"SELECT SUM(value) as total, COUNT(*) as cnt FROM "dist-test""#;
 
     let mut client = datafusion_client(&sandbox);
-    let explain_batches =
-        execute_datafusion_sql(&mut client, format!("{ddl}; EXPLAIN VERBOSE {query}")).await;
+    let explain_batches = explain_datafusion_sql(&mut client, format!("{ddl}; {query}")).await;
     let plan_str = explain_plan_text(&explain_batches);
     println!("=== DataFusion EXPLAIN output ===\n{plan_str}");