feat: Add tracing instrumentation for plan creation

Introduces TracingQueryPlanner and rule instrumentation for analyzer,
logical optimizer, and physical optimizer rules with tracing spans.

Phase spans group individual rule executions and capture plan diffs when
rules modify plans. Ergonomic API via RuleInstrumentationOptions::full()
and instrument_rules_with_info_spans! macro.

Author: geoffreyclaude

Cargo.lock

@@ -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"

Cargo.toml

@@ -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{}",

README.md

@@ -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/Cargo.toml

@@ -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/instrumented_macros.rs renamed to datafusion-tracing/src/exec_instrument_macros.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/instrument_rule.rs renamed to datafusion-tracing/src/exec_instrument_rule.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/instrumented.rs renamed to datafusion-tracing/src/instrumented_exec.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
+    }
+}