feat: draw the rest of the owl

Author: geoffreyclaude

Cargo.lock

@@ -132,19 +132,20 @@ A more complete example can be found in the [examples directory](https://github.
 
 Always register the instrumentation rule last in your physical optimizer chain.
 
-- Many optimizer rules identify nodes using `as_any().downcast_ref::<ConcreteExec>()`.
-  Since instrumentation wraps each node in a private `InstrumentedExec`, those downcasts
-  won’t match if instrumentation runs first, causing rules to be skipped or, in code
-  that assumes success, to panic.
-- Some rules may rewrite parts of the plan after instrumentation. While `InstrumentedExec`
+- `InstrumentedExec` is transparent for downcasting: its `as_any()` delegates to the wrapped
+  node, so `as_any().downcast_ref::<ConcreteExec>()` works through the wrapper. This allows
+  rules and execution-time code to identify wrapped nodes correctly.
+- However, some rules may rewrite parts of the plan after instrumentation. While `InstrumentedExec`
   re-wraps many common mutations, placing the rule last guarantees full, consistent
-  coverage regardless of other rules’ behaviors.
+  coverage regardless of other rules' behaviors.
 
 Why is `InstrumentedExec` private?
 
-- To prevent downstream code from downcasting to or unwrapping the wrapper, which would be
+- To prevent downstream code from directly accessing or unwrapping the wrapper, which would be
   brittle and force long-term compatibility constraints on its internals. The public
-  contract is the optimizer rule, not the concrete node.
+  contract is the optimizer rule and the transparent downcasting behavior, not the concrete node.
+- Downcasting works transparently (via delegated `as_any()`), so you can check node types
+  without knowing about the wrapper.
 
 How to ensure it is last:
 

README.md

@@ -68,19 +68,19 @@ impl PhysicalOptimizerRule for InstrumentRule {
         plan: Arc<dyn ExecutionPlan>,
         _config: &ConfigOptions,
     ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
-        // Iterate over the plan and wrap each node with InstrumentedExec
+        // Wrap each node in the plan tree with InstrumentedExec.
+        // Since InstrumentedExec delegates as_any() to its inner node, it's transparent
+        // for downcasting - other code can still identify wrapped nodes by their actual type.
+        //
+        // Note: This will wrap every node, even if the rule is called multiple times on
+        // the same plan (unusual but possible). The resulting double-wrapping creates
+        // duplicate telemetry but remains functionally correct.
         plan.transform(|plan| {
-            if plan.as_any().downcast_ref::<InstrumentedExec>().is_none() {
-                // Node is not InstrumentedExec; wrap it
-                Ok(Transformed::yes(Arc::new(InstrumentedExec::new(
-                    plan,
-                    self.span_create_fn.clone(),
-                    &self.options,
-                ))))
-            } else {
-                // Node is already InstrumentedExec; do not wrap again
-                Ok(Transformed::no(plan))
-            }
+            Ok(Transformed::yes(Arc::new(InstrumentedExec::new(
+                plan,
+                self.span_create_fn.clone(),
+                &self.options,
+            ))))
         })
         .data()
     }

datafusion-tracing/src/instrument_rule.rs

@@ -347,8 +347,9 @@ impl ExecutionPlan for InstrumentedExec {
         Some(self.with_new_inner(new_inner))
     }
 
+    /// Delegate to the inner plan for downcasting.
     fn as_any(&self) -> &dyn Any {
-        self
+        self.inner.as_any()
     }
 
     /// Executes the plan for a given partition and context, instrumented with tracing and metrics recording.

datafusion-tracing/src/instrumented.rs

@@ -136,8 +136,8 @@ async fn main() -> Result<()> {
 async fn run_distributed_otlp_example() -> Result<()> {
     // Test both distributed execution modes
     let modes = [
-        (DistributedMode::Memory, "Memory"),
         (DistributedMode::Localhost, "Localhost"),
+        (DistributedMode::Memory, "Memory"),
     ];
 
     for (mode, mode_name) in modes {
@@ -187,6 +187,13 @@ async fn run_distributed_otlp_example() -> Result<()> {
 
 /// Initializes OpenTelemetry and tracing infrastructure to enable tracing of query execution.
 fn init_tracing() -> Result<opentelemetry_sdk::trace::SdkTracerProvider> {
+    // Set up the global text map propagator for trace context propagation across gRPC boundaries.
+    // This is essential for distributed tracing to work - without it, worker spans won't be
+    // linked to parent traces.
+    opentelemetry::global::set_text_map_propagator(
+        opentelemetry_sdk::propagation::TraceContextPropagator::new(),
+    );
+
     // Set service metadata for tracing.
     let resource = Resource::builder()
         .with_attribute(KeyValue::new("service.name", "datafusion-tracing"))

examples/distributed_otlp.rs

@@ -36,16 +36,21 @@ arrow-flight = "56.1"
 async-trait = "0.1"
 dashmap = "6.1"
 datafusion = { workspace = true, features = ["parquet", "nested_expressions"] }
-datafusion-distributed = { git = "https://github.com/datafusion-contrib/datafusion-distributed", branch = "main" }
+datafusion-distributed = { git = "https://github.com/geoffreyclaude/datafusion-distributed", branch = "feat/tracing" }
+# datafusion-distributed = { path = "../../datafusion-distributed" }
 datafusion-tracing = { workspace = true }
 futures = { workspace = true }
+http = "1.0"
 hyper-util = "0.1"
 instrumented-object-store = { workspace = true }
 object_store = { version = "0.12.1", default-features = false }
+opentelemetry = { version = "0.30", features = ["trace"] }
+pin-project-lite = "0.2"
 structopt = "0.3"
 tokio = { workspace = true, features = ["full"] }
 tokio-stream = "0.1"
 tonic = { version = "0.13", features = ["transport"] }
 tower = "0.5"
 tracing = { workspace = true }
+tracing-opentelemetry = { version = "0.31" }
 url = { version = "2.5" }

integration-utils/Cargo.toml

@@ -28,6 +28,9 @@ use datafusion_distributed::{
 };
 use hyper_util::rt::TokioIo;
 use tonic::transport::{Endpoint, Server};
+use tower::ServiceBuilder;
+
+use super::trace_middleware::{TracingClientLayer, TracingServerLayer};
 
 const DUMMY_URL: &str = "http://localhost:50051";
 const MAX_MESSAGE_SIZE: usize = 2 * 1024 * 1024 * 1024; // 2GB
@@ -55,6 +58,11 @@ impl InMemoryChannelResolver {
                 async move { Ok::<_, std::io::Error>(TokioIo::new(client)) }
             }));
 
+        // Wrap the channel with tracing middleware to inject trace context
+        let channel = ServiceBuilder::new()
+            .layer(TracingClientLayer)
+            .service(channel);
+
         let this = Self {
             channel: FlightServiceClient::new(BoxCloneSyncChannel::new(channel))
                 .max_decoding_message_size(MAX_MESSAGE_SIZE)
@@ -76,12 +84,15 @@ impl InMemoryChannelResolver {
             .unwrap();
 
         tokio::spawn(async move {
+            // Wrap the FlightService with tracing middleware to extract trace context
+            let service = ServiceBuilder::new().layer(TracingServerLayer).service(
+                FlightServiceServer::new(endpoint)
+                    .max_decoding_message_size(MAX_MESSAGE_SIZE)
+                    .max_encoding_message_size(MAX_MESSAGE_SIZE),
+            );
+
             Server::builder()
-                .add_service(
-                    FlightServiceServer::new(endpoint)
-                        .max_decoding_message_size(MAX_MESSAGE_SIZE)
-                        .max_encoding_message_size(MAX_MESSAGE_SIZE),
-                )
+                .add_service(service)
                 .serve_with_incoming(tokio_stream::once(Ok::<_, std::io::Error>(server)))
                 .await
         });