Expose Substrait execution metadata

alexanderbianchi · alexanderbianchi · commit 29a743786db2 · 2026-04-29T20:35:14.000-04:00
diff --git a/quickwit/quickwit-df-core/src/lib.rs b/quickwit/quickwit-df-core/src/lib.rs
@@ -43,7 +43,7 @@ pub use data_source::{
 };
 pub use datafusion::execution::SendableRecordBatchStream;
 pub use datafusion_distributed::{Worker, WorkerResolver};
-pub use service::DataFusionService;
+pub use service::{DataFusionService, SubstraitExecution, SubstraitExecutionMetadata};
 pub use session::DataFusionSessionBuilder;
 pub use task_estimator::DataSourceExecPartitionEstimator;
 pub use worker::build_worker;
diff --git a/quickwit/quickwit-df-core/src/service.rs b/quickwit/quickwit-df-core/src/service.rs
@@ -44,15 +44,41 @@
 
 use std::collections::HashMap;
 use std::sync::Arc;
+use std::time::{Duration, Instant};
 
 use datafusion::error::{DataFusionError, Result as DFResult};
 use datafusion::execution::SendableRecordBatchStream;
+use datafusion::physical_plan::ExecutionPlan;
+use datafusion::physical_plan::display::DisplayableExecutionPlan;
 use datafusion::prelude::SessionContext;
 use datafusion::sql::sqlparser::dialect::dialect_from_str;
 use datafusion::sql::sqlparser::tokenizer::{Token, Tokenizer};
 
 use crate::session::DataFusionSessionBuilder;
 
+pub struct SubstraitExecution {
+    pub stream: SendableRecordBatchStream,
+    pub metadata: SubstraitExecutionMetadata,
+}
+
+pub struct SubstraitExecutionMetadata {
+    pub substrait_plan_json: String,
+    pub logical_plan: String,
+    pub physical_plan: Arc<dyn ExecutionPlan>,
+    pub substrait_decode_duration: Duration,
+    pub substrait_to_logical_duration: Duration,
+    pub logical_to_physical_duration: Duration,
+    pub stream_creation_duration: Duration,
+}
+
+impl SubstraitExecutionMetadata {
+    pub fn physical_plan_display(&self) -> String {
+        DisplayableExecutionPlan::with_metrics(self.physical_plan.as_ref())
+            .indent(true)
+            .to_string()
+    }
+}
+
 /// Split a SQL string into top-level statements using the configured dialect's
 /// tokenizer rather than raw string splitting.
 pub fn split_sql_statements(ctx: &SessionContext, sql: &str) -> DFResult<Vec<String>> {
@@ -144,6 +170,17 @@ impl DataFusionService {
         plan_bytes: &[u8],
         properties: &HashMap<String, String>,
     ) -> DFResult<SendableRecordBatchStream> {
+        Ok(self
+            .execute_substrait_with_metadata(plan_bytes, properties)
+            .await?
+            .stream)
+    }
+
+    pub async fn execute_substrait_with_metadata(
+        &self,
+        plan_bytes: &[u8],
+        properties: &HashMap<String, String>,
+    ) -> DFResult<SubstraitExecution> {
         tracing::info!(
             plan_bytes_len = plan_bytes.len(),
             num_properties = properties.len(),
@@ -152,10 +189,23 @@ impl DataFusionService {
         use datafusion_substrait::substrait::proto::Plan;
         use prost::Message;
 
+        let substrait_decode_start = Instant::now();
         let plan = Plan::decode(plan_bytes)
             .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?;
+        let substrait_decode_duration = substrait_decode_start.elapsed();
+        let substrait_plan_json = serde_json::to_string(&plan).map_err(|e| {
+            datafusion::error::DataFusionError::Plan(format!(
+                "failed to serialize Substrait plan JSON: {e}"
+            ))
+        })?;
 
-        self.execute_substrait_plan(&plan, properties).await
+        self.execute_substrait_plan_with_metadata(
+            &plan,
+            properties,
+            substrait_plan_json,
+            substrait_decode_duration,
+        )
+        .await
     }
 
     /// Execute a Substrait plan from its proto3 JSON representation.
@@ -171,27 +221,60 @@ impl DataFusionService {
         plan_json: &str,
         properties: &HashMap<String, String>,
     ) -> DFResult<SendableRecordBatchStream> {
+        Ok(self
+            .execute_substrait_json_with_metadata(plan_json, properties)
+            .await?
+            .stream)
+    }
+
+    pub async fn execute_substrait_json_with_metadata(
+        &self,
+        plan_json: &str,
+        properties: &HashMap<String, String>,
+    ) -> DFResult<SubstraitExecution> {
         use datafusion_substrait::substrait::proto::Plan;
 
+        let substrait_decode_start = Instant::now();
         let plan: Plan = serde_json::from_str(plan_json).map_err(|e| {
             datafusion::error::DataFusionError::Plan(format!("invalid Substrait plan JSON: {e}"))
         })?;
+        let substrait_decode_duration = substrait_decode_start.elapsed();
 
-        self.execute_substrait_plan(&plan, properties).await
+        self.execute_substrait_plan_with_metadata(
+            &plan,
+            properties,
+            plan_json.to_string(),
+            substrait_decode_duration,
+        )
+        .await
     }
 
-    async fn execute_substrait_plan(
+    async fn execute_substrait_plan_with_metadata(
         &self,
         plan: &datafusion_substrait::substrait::proto::Plan,
         properties: &HashMap<String, String>,
-    ) -> DFResult<SendableRecordBatchStream> {
+        substrait_plan_json: String,
+        substrait_decode_duration: Duration,
+    ) -> DFResult<SubstraitExecution> {
         let ctx = self.builder.build_session_with_properties(properties)?;
-        crate::substrait::execute_substrait_plan_streaming(
+        let prepared = crate::substrait::execute_substrait_plan_streaming_with_metadata(
             plan,
             &ctx,
             self.builder.substrait_extensions(),
         )
-        .await
+        .await?;
+        Ok(SubstraitExecution {
+            stream: prepared.stream,
+            metadata: SubstraitExecutionMetadata {
+                substrait_plan_json,
+                logical_plan: prepared.metadata.logical_plan,
+                physical_plan: prepared.metadata.physical_plan,
+                substrait_decode_duration,
+                substrait_to_logical_duration: prepared.metadata.substrait_to_logical_duration,
+                logical_to_physical_duration: prepared.metadata.logical_to_physical_duration,
+                stream_creation_duration: prepared.metadata.stream_creation_duration,
+            },
+        })
     }
 
     /// Like [`execute_substrait`], but returns the EXPLAIN output instead of
@@ -202,11 +285,37 @@ impl DataFusionService {
         plan_bytes: &[u8],
         properties: &HashMap<String, String>,
     ) -> DFResult<SendableRecordBatchStream> {
+        Ok(self
+            .explain_substrait_with_metadata(plan_bytes, properties)
+            .await?
+            .stream)
+    }
+
+    pub async fn explain_substrait_with_metadata(
+        &self,
+        plan_bytes: &[u8],
+        properties: &HashMap<String, String>,
+    ) -> DFResult<SubstraitExecution> {
         use datafusion_substrait::substrait::proto::Plan;
         use prost::Message;
+
+        let substrait_decode_start = Instant::now();
         let plan = Plan::decode(plan_bytes)
             .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?;
-        self.explain_substrait_plan(&plan, properties).await
+        let substrait_decode_duration = substrait_decode_start.elapsed();
+        let substrait_plan_json = serde_json::to_string(&plan).map_err(|e| {
+            datafusion::error::DataFusionError::Plan(format!(
+                "failed to serialize Substrait plan JSON: {e}"
+            ))
+        })?;
+
+        self.explain_substrait_plan_with_metadata(
+            &plan,
+            properties,
+            substrait_plan_json,
+            substrait_decode_duration,
+        )
+        .await
     }
 
     /// JSON variant of [`explain_substrait`] — same semantics, proto3-JSON input.
@@ -215,25 +324,58 @@ impl DataFusionService {
         plan_json: &str,
         properties: &HashMap<String, String>,
     ) -> DFResult<SendableRecordBatchStream> {
+        Ok(self
+            .explain_substrait_json_with_metadata(plan_json, properties)
+            .await?
+            .stream)
+    }
+
+    pub async fn explain_substrait_json_with_metadata(
+        &self,
+        plan_json: &str,
+        properties: &HashMap<String, String>,
+    ) -> DFResult<SubstraitExecution> {
         use datafusion_substrait::substrait::proto::Plan;
+        let substrait_decode_start = Instant::now();
         let plan: Plan = serde_json::from_str(plan_json).map_err(|e| {
             datafusion::error::DataFusionError::Plan(format!("invalid Substrait plan JSON: {e}"))
         })?;
-        self.explain_substrait_plan(&plan, properties).await
+        let substrait_decode_duration = substrait_decode_start.elapsed();
+        self.explain_substrait_plan_with_metadata(
+            &plan,
+            properties,
+            plan_json.to_string(),
+            substrait_decode_duration,
+        )
+        .await
     }
 
-    async fn explain_substrait_plan(
+    async fn explain_substrait_plan_with_metadata(
         &self,
         plan: &datafusion_substrait::substrait::proto::Plan,
         properties: &HashMap<String, String>,
-    ) -> DFResult<SendableRecordBatchStream> {
+        substrait_plan_json: String,
+        substrait_decode_duration: Duration,
+    ) -> DFResult<SubstraitExecution> {
         let ctx = self.builder.build_session_with_properties(properties)?;
-        crate::substrait::explain_substrait_plan_streaming(
+        let prepared = crate::substrait::explain_substrait_plan_streaming_with_metadata(
             plan,
             &ctx,
             self.builder.substrait_extensions(),
         )
-        .await
+        .await?;
+        Ok(SubstraitExecution {
+            stream: prepared.stream,
+            metadata: SubstraitExecutionMetadata {
+                substrait_plan_json,
+                logical_plan: prepared.metadata.logical_plan,
+                physical_plan: prepared.metadata.physical_plan,
+                substrait_decode_duration,
+                substrait_to_logical_duration: prepared.metadata.substrait_to_logical_duration,
+                logical_to_physical_duration: prepared.metadata.logical_to_physical_duration,
+                stream_creation_duration: prepared.metadata.stream_creation_duration,
+            },
+        })
     }
 
     /// Execute one or more SQL statements from a single SQL string.
diff --git a/quickwit/quickwit-df-core/src/substrait.rs b/quickwit/quickwit-df-core/src/substrait.rs
