Instrument plan creation

Cargo.toml

@@ -37,6 +37,7 @@ authors = ["DataDog <[email protected]>"]
 rust-version = "1.86.0"
 
 [workspace.dependencies]
+async-trait = "0.1"
 datafusion = { version = "51.0.0", default-features = false }
 datafusion-tracing = { path = "datafusion-tracing", version = "51.0.0" }
 futures = "0.3"

README.md

@@ -73,7 +73,8 @@ use datafusion::{
     prelude::*,
 };
 use datafusion_tracing::{
-    instrument_with_info_spans, pretty_format_compact_batch, InstrumentationOptions,
+    instrument_rules_with_info_spans, instrument_with_info_spans,
+    pretty_format_compact_batch, InstrumentationOptions, RuleInstrumentationOptions,
 };
 use std::sync::Arc;
 use tracing::field;
@@ -83,8 +84,8 @@ async fn main() -> Result<()> {
     // Initialize tracing subscriber as usual
     // (See examples/otlp.rs for a complete example).
 
-    // Set up tracing options (you can customize these).
-    let options = InstrumentationOptions::builder()
+    // Set up execution plan tracing options (you can customize these).
+    let exec_options = InstrumentationOptions::builder()
         .record_metrics(true)
         .preview_limit(5)
         .preview_fn(Arc::new(|batch: &RecordBatch| {
@@ -95,7 +96,7 @@ async fn main() -> Result<()> {
         .build();
 
     let instrument_rule = instrument_with_info_spans!(
-        options: options,
+        options: exec_options,
         env = field::Empty,
         region = field::Empty,
     );
@@ -105,8 +106,19 @@ async fn main() -> Result<()> {
         .with_physical_optimizer_rule(instrument_rule)
         .build();
 
+    // Instrument all rules (analyzer, logical optimizer, physical optimizer)
+    // Physical plan creation tracing is automatically enabled when physical_optimizer is set
+    let rule_options = RuleInstrumentationOptions::full().with_plan_diff();
+    let session_state = instrument_rules_with_info_spans!(
+        options: rule_options,
+        state: session_state
+    );
+
     let ctx = SessionContext::new_with_state(session_state);
 
+    // Execute a query - the entire lifecycle is now traced:
+    // SQL Parsing -> Logical Plan -> Analyzer Rules -> Optimizer Rules ->
+    // Physical Plan Creation -> Physical Optimizer Rules -> Execution
     let results = ctx.sql("SELECT 1").await?.collect().await?;
     println!(
         "Query Results:\n{}",

datafusion-tracing/Cargo.toml

@@ -33,11 +33,13 @@ documentation = "https://docs.rs/datafusion-tracing"
 homepage = "https://github.com/datafusion-contrib/datafusion-tracing"
 
 [dependencies]
+async-trait = { workspace = true }
 comfy-table = { version = "7.2" }
 datafusion = { workspace = true }
 delegate = "0.13"
 futures = { workspace = true }
 pin-project = "1.1"
+similar = { version = "2.7", default-features = false, features = ["text"] }
 tracing = { workspace = true }
 tracing-futures = { workspace = true }
 unicode-width = "0.2"

datafusion-tracing/src/instrument_rule.rs → datafusion-tracing/src/exec_instrument_rule.rs

@@ -17,8 +17,8 @@
 //
 // This product includes software developed at Datadog (https://www.datadoghq.com/) Copyright 2025 Datadog, Inc.
 
-use crate::instrumented::InstrumentedExec;
-use crate::instrumented::SpanCreateFn;
+use crate::instrumented_exec::InstrumentedExec;
+use crate::instrumented_exec::SpanCreateFn;
 use crate::options::InstrumentationOptions;
 use datafusion::common::runtime::{JoinSetTracer, set_join_set_tracer};
 use datafusion::common::tree_node::{Transformed, TransformedResult, TreeNode};

datafusion-tracing/src/lib.rs

@@ -63,7 +63,8 @@
 //!     prelude::*,
 //! };
 //! use datafusion_tracing::{
-//!     instrument_with_info_spans, pretty_format_compact_batch, InstrumentationOptions,
+//!     instrument_rules_with_info_spans, instrument_with_info_spans,
+//!     pretty_format_compact_batch, InstrumentationOptions, RuleInstrumentationOptions,
 //! };
 //! use std::sync::Arc;
 //! use tracing::field;
@@ -73,8 +74,8 @@
 //!     // Initialize tracing subscriber as usual
 //!     // (See examples/otlp.rs for a complete example).
 //!
-//!     // Set up tracing options (you can customize these).
-//!     let options = InstrumentationOptions::builder()
+//!     // Set up execution plan tracing options (you can customize these).
+//!     let exec_options = InstrumentationOptions::builder()
 //!         .record_metrics(true)
 //!         .preview_limit(5)
 //!         .preview_fn(Arc::new(|batch: &RecordBatch| {
@@ -85,7 +86,7 @@
 //!         .build();
 //!
 //!     let instrument_rule = instrument_with_info_spans!(
-//!         options: options,
+//!         options: exec_options,
 //!         env = field::Empty,
 //!         region = field::Empty,
 //!     );
@@ -95,8 +96,19 @@
 //!         .with_physical_optimizer_rule(instrument_rule)
 //!         .build();
 //!
+//!     // Instrument all rules (analyzer, logical optimizer, physical optimizer)
+//!     // Physical plan creation tracing is automatically enabled when physical_optimizer is set
+//!     let rule_options = RuleInstrumentationOptions::full().with_plan_diff();
+//!     let session_state = instrument_rules_with_info_spans!(
+//!         options: rule_options,
+//!         state: session_state
+//!     );
+//!
 //!     let ctx = SessionContext::new_with_state(session_state);
 //!
+//!     // Execute a query - the entire lifecycle is now traced:
+//!     // SQL Parsing -> Logical Plan -> Analyzer Rules -> Optimizer Rules ->
+//!     // Physical Plan Creation -> Physical Optimizer Rules -> Execution
 //!     let results = ctx.sql("SELECT 1").await?.collect().await?;
 //!     println!(
 //!         "Query Results:\n{}",
@@ -110,21 +122,33 @@
 //! A more complete example can be found in the [examples directory](https://github.com/datafusion-contrib/datafusion-tracing/tree/main/examples).
 //!
 
-mod instrument_rule;
-mod instrumented;
-mod instrumented_macros;
+// Execution plan instrumentation (wraps ExecutionPlan nodes with tracing)
+mod exec_instrument_macros;
+mod exec_instrument_rule;
+mod instrumented_exec;
+
+// Rule instrumentation (wraps analyzer/optimizer/physical optimizer rules with tracing)
+mod rule_instrumentation;
+mod rule_instrumentation_macros;
+
+// Shared utilities
 mod metrics;
 mod node;
 mod options;
+mod planner;
 mod preview;
 mod preview_utils;
+mod rule_options;
 mod utils;
 
 // Hide implementation details from documentation.
-// This function is only public because it needs to be accessed by the macros,
-// but it's not intended for direct use by consumers of this crate.
+// These functions are only public because they need to be accessed by the macros,
+// but they're not intended for direct use by consumers of this crate.
+#[doc(hidden)]
+pub use exec_instrument_rule::new_instrument_rule;
 #[doc(hidden)]
-pub use instrument_rule::new_instrument_rule;
+pub use rule_instrumentation::instrument_session_state;
 
 pub use options::InstrumentationOptions;
 pub use preview_utils::pretty_format_compact_batch;
+pub use rule_options::RuleInstrumentationOptions;

datafusion-tracing/src/options.rs

@@ -43,6 +43,24 @@ pub struct InstrumentationOptions {
     pub custom_fields: HashMap<String, String>,
 }
 
+impl std::fmt::Debug for InstrumentationOptions {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("InstrumentationOptions")
+            .field("record_metrics", &self.record_metrics)
+            .field("preview_limit", &self.preview_limit)
+            .field(
+                "preview_fn",
+                &if self.preview_fn.is_some() {
+                    "Some(PreviewFn)"
+                } else {
+                    "None"
+                },
+            )
+            .field("custom_fields", &self.custom_fields)
+            .finish()
+    }
+}
+
 impl InstrumentationOptions {
     /// Creates a new builder for `InstrumentationOptions`.
     pub fn builder() -> InstrumentationOptionsBuilder {

datafusion-tracing/src/planner.rs

@@ -0,0 +1,126 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+//
+// This product includes software developed at Datadog (https://www.datadoghq.com/) Copyright 2025 Datadog, Inc.
+
+use async_trait::async_trait;
+use datafusion::common::Result;
+use datafusion::execution::SessionStateBuilder;
+use datafusion::execution::context::{QueryPlanner, SessionState};
+use datafusion::logical_expr::LogicalPlan;
+use datafusion::physical_plan::{ExecutionPlan, displayable};
+use std::sync::Arc;
+use tracing::{Instrument, Level};
+
+/// A `QueryPlanner` that instruments the creation of the physical plan.
+///
+/// This is automatically applied when instrumenting a `SessionState` with physical
+/// optimizer instrumentation enabled (PhaseOnly or Full).
+#[derive(Debug)]
+pub(crate) struct TracingQueryPlanner {
+    inner: Arc<dyn QueryPlanner + Send + Sync>,
+    level: Level,
+}
+
+impl TracingQueryPlanner {
+    /// Create a new `TracingQueryPlanner` that wraps the provided inner planner at a specific level.
+    fn new_with_level(inner: Arc<dyn QueryPlanner + Send + Sync>, level: Level) -> Self {
+        Self { inner, level }
+    }
+
+    /// Wraps the query planner of an existing `SessionState` with tracing instrumentation at a specific level.
+    ///
+    /// This preserves any custom `QueryPlanner` that may already be configured in the state,
+    /// ensuring that tracing is added as a layer on top of existing functionality.
+    pub(crate) fn instrument_state_with_level(
+        state: SessionState,
+        level: Level,
+    ) -> SessionState {
+        let current_planner = state.query_planner().clone();
+        let wrapped_planner = Arc::new(Self::new_with_level(current_planner, level));
+
+        SessionStateBuilder::from(state)
+            .with_query_planner(wrapped_planner)
+            .build()
+    }
+}
+
+#[async_trait]
+impl QueryPlanner for TracingQueryPlanner {
+    async fn create_physical_plan(
+        &self,
+        logical_plan: &LogicalPlan,
+        session_state: &SessionState,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        let span = match self.level {
+            Level::TRACE => tracing::trace_span!(
+                "create_physical_plan",
+                logical_plan = tracing::field::Empty,
+                physical_plan = tracing::field::Empty,
+                error = tracing::field::Empty
+            ),
+            Level::DEBUG => tracing::debug_span!(
+                "create_physical_plan",
+                logical_plan = tracing::field::Empty,
+                physical_plan = tracing::field::Empty,
+                error = tracing::field::Empty
+            ),
+            Level::INFO => tracing::info_span!(
+                "create_physical_plan",
+                logical_plan = tracing::field::Empty,
+                physical_plan = tracing::field::Empty,
+                error = tracing::field::Empty
+            ),
+            Level::WARN => tracing::warn_span!(
+                "create_physical_plan",
+                logical_plan = tracing::field::Empty,
+                physical_plan = tracing::field::Empty,
+                error = tracing::field::Empty
+            ),
+            Level::ERROR => tracing::error_span!(
+                "create_physical_plan",
+                logical_plan = tracing::field::Empty,
+                physical_plan = tracing::field::Empty,
+                error = tracing::field::Empty
+            ),
+        };
+
+        // Record the logical plan as a formatted string with schema
+        let logical_plan_str = logical_plan.display_indent_schema().to_string();
+        span.record("logical_plan", logical_plan_str.as_str());
+
+        let physical_plan = self
+            .inner
+            .create_physical_plan(logical_plan, session_state)
+            .instrument(span.clone())
+            .await;
+
+        match &physical_plan {
+            Ok(plan) => {
+                // Record the physical plan as a formatted string
+                let physical_plan_str =
+                    displayable(plan.as_ref()).indent(true).to_string();
+                span.record("physical_plan", physical_plan_str.as_str());
+            }
+            Err(e) => {
+                span.record("error", e.to_string().as_str());
+            }
+        }
+
+        physical_plan
+    }
+}