fix: Add partial support for recursive queries instrumentation

This change introduces partial support for tracing and metrics instrumentation of recursive queries in DataFusion. Recursive queries now spawn a new span group for each recursive call, as the recursive implementation recreates the recursive nodes for each loop iteration. This is enabled by supporting the `with_new_inner` method on `InstrumentedExec`, which allows the optimizer to rebuild instrumented plans as needed.

However, due to current limitations in DataFusion, nodes of type `WorkTableExec` are not instrumented, as wrapping them would break recursive query execution. This limitation will be revisited once upstream changes allow `WorkTableExec` to be safely instrumented.

See also: apache/datafusion#16469 and related discussions.

Author: geoffreyclaude

README.md

@@ -119,6 +119,14 @@ async fn main() -> Result<()> {
 
 A more complete example can be found in the [examples directory](https://github.com/datafusion-contrib/datafusion-tracing/tree/main/examples).
 
+## Limitations
+
+### Recursive Queries
+
+When using DataFusion Tracing with recursive queries (e.g., those using `WITH RECURSIVE`), nodes of type `WorkTableExec` are intentionally not instrumented. This is due to a current limitation in DataFusion: instrumenting these nodes can break recursive query execution and result in errors such as `Unexpected empty work table.`
+
+As a result, while most of the recursive query plan will be traced and metrics will be collected, some internal operations related to recursion will not be visible in the trace until upstream support is available. See [issue #5](https://github.com/datafusion-contrib/datafusion-tracing/issues/5) for details and tracking progress on this limitation.
+
 <!-- cargo-rdme end -->
 
 ## Setting Up a Collector

datafusion-tracing/src/instrument_rule.rs

@@ -22,6 +22,7 @@ use crate::instrumented::SpanCreateFn;
 use crate::options::InstrumentationOptions;
 use datafusion::common::runtime::{set_join_set_tracer, JoinSetTracer};
 use datafusion::common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion::physical_plan::work_table::WorkTableExec;
 use datafusion::{
     config::ConfigOptions, physical_optimizer::PhysicalOptimizerRule,
     physical_plan::ExecutionPlan,
@@ -69,8 +70,15 @@ impl PhysicalOptimizerRule for InstrumentRule {
         _config: &ConfigOptions,
     ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
         // Iterate over the plan and wrap each node with InstrumentedExec
+        //
+        // NOTE: Recursive queries dynamically rebuild the execution plan during execution and assume one of the nodes is a `WorkTableExec`.
+        // Wrapping it with `InstrumentedExec`` would break the recursive query, so for now we only instrument non-`WorkTableExec` nodes.
+        //
+        // TODO: Remove this limitation once `with_work_table` is made a trait method of `ExecutionPlan` (tracked by https://github.com/apache/datafusion/pull/16469)
         plan.transform(|plan| {
-            if plan.as_any().downcast_ref::<InstrumentedExec>().is_none() {
+            if plan.as_any().downcast_ref::<InstrumentedExec>().is_none()
+                && plan.as_any().downcast_ref::<WorkTableExec>().is_none()
+            {
                 // Node is not InstrumentedExec; wrap it
                 Ok(Transformed::yes(Arc::new(InstrumentedExec::new(
                     plan,

datafusion-tracing/src/instrumented.rs

@@ -45,8 +45,8 @@ use datafusion::{
 use delegate::delegate;
 use std::{
     any::Any,
-    fmt,
-    fmt::Debug,
+    collections::HashMap,
+    fmt::{self, Debug},
     sync::{Arc, OnceLock},
 };
 use tracing::{field, Span};
@@ -105,6 +105,23 @@ impl InstrumentedExec {
         }
     }
 
+    /// Creates a new `InstrumentedExec` with the same configuration as this instance but with a different inner execution plan.
+    ///
+    /// This method is used when the optimizer needs to replace the inner execution plan while preserving
+    /// all the instrumentation settings (metrics recording, preview limits, span creation function, etc.).
+    fn with_new_inner(&self, inner: Arc<dyn ExecutionPlan>) -> Arc<dyn ExecutionPlan> {
+        Arc::new(InstrumentedExec::new(
+            inner,
+            self.span_create_fn.clone(),
+            &InstrumentationOptions {
+                record_metrics: self.record_metrics,
+                preview_limit: self.preview_limit,
+                preview_fn: self.preview_fn.clone(),
+                custom_fields: HashMap::new(), // custom fields are not used by `InstrumentedExec`, only by the higher-level `instrument_with_spans` macro family
+            },
+        ))
+    }
+
     /// Retrieves the tracing span, initializing it if necessary.
     fn get_span(&self) -> Span {
         self.span
@@ -191,11 +208,6 @@ impl InstrumentedExec {
 
 impl ExecutionPlan for InstrumentedExec {
     // Delegate all ExecutionPlan methods to the inner plan, except for `as_any` and `execute`.
-    //
-    // Note: Methods returning a new ExecutionPlan instance delegate directly to the inner plan,
-    // resulting in loss of instrumentation. This is intentional since instrumentation should
-    // be applied by the final PhysicalOptimizerRule pass. Therefore, any resulting plans will
-    // be re-instrumented after optimizer transformations.
     delegate! {
         to self.inner {
             fn name(&self) -> &str;
@@ -207,25 +219,15 @@ impl ExecutionPlan for InstrumentedExec {
             fn maintains_input_order(&self) -> Vec<bool>;
             fn benefits_from_input_partitioning(&self) -> Vec<bool>;
             fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>>;
-            fn repartitioned(
-                &self,
-                target_partitions: usize,
-                config: &ConfigOptions,
-            ) -> Result<Option<Arc<dyn ExecutionPlan>>>;
             fn metrics(&self) -> Option<MetricsSet>;
             // We need to delegate to the inner plan's statistics method to preserve behavior,
             // even though it's deprecated in favor of partition_statistics
             #[allow(deprecated)]
             fn statistics(&self) -> Result<Statistics>;
             fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics>;
             fn supports_limit_pushdown(&self) -> bool;
-            fn with_fetch(&self, limit: Option<usize>) -> Option<Arc<dyn ExecutionPlan>>;
             fn fetch(&self) -> Option<usize>;
             fn cardinality_effect(&self) -> CardinalityEffect;
-            fn try_swapping_with_projection(
-                &self,
-                projection: &ProjectionExec,
-            ) -> Result<Option<Arc<dyn ExecutionPlan>>>;
             fn gather_filters_for_pushdown(
                 &self,
                 parent_filters: Vec<Arc<dyn PhysicalExpr>>,
@@ -239,15 +241,57 @@ impl ExecutionPlan for InstrumentedExec {
         }
     }
 
-    delegate! {
-        to self.inner.clone() {
-            fn with_new_children(
-                self: Arc<Self>,
-                children: Vec<Arc<dyn ExecutionPlan>>,
-            ) -> Result<Arc<dyn ExecutionPlan>>;
+    /// Delegate to the inner plan for repartitioning and rewrap with an InstrumentedExec.
+    fn repartitioned(
+        &self,
+        target_partitions: usize,
+        config: &ConfigOptions,
+    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
+        if let Some(new_inner) = self
+            .inner
+            .clone()
+            .repartitioned(target_partitions, config)?
+        {
+            Ok(Some(self.with_new_inner(new_inner)))
+        } else {
+            Ok(None)
         }
     }
 
+    /// Delegate to the inner plan for fetching and rewrap with an InstrumentedExec.
+    fn with_fetch(&self, limit: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> {
+        if let Some(new_inner) = self.inner.clone().with_fetch(limit) {
+            Some(self.with_new_inner(new_inner))
+        } else {
+            None
+        }
+    }
+
+    /// Delegate to the inner plan for swapping with a projection and rewrap with an InstrumentedExec.
+    fn try_swapping_with_projection(
+        &self,
+        projection: &ProjectionExec,
+    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
+        if let Some(new_inner) = self
+            .inner
+            .clone()
+            .try_swapping_with_projection(projection)?
+        {
+            Ok(Some(self.with_new_inner(new_inner)))
+        } else {
+            Ok(None)
+        }
+    }
+
+    /// Delegate to the inner plan for creating new children and rewrap with an InstrumentedExec.
+    fn with_new_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
+        let new_inner = self.inner.clone().with_new_children(children)?;
+        Ok(self.with_new_inner(new_inner))
+    }
+
     fn as_any(&self) -> &dyn Any {
         self
     }

datafusion-tracing/src/lib.rs

@@ -109,6 +109,14 @@
 //!
 //! A more complete example can be found in the [examples directory](https://github.com/datafusion-contrib/datafusion-tracing/tree/main/examples).
 //!
+//! # Limitations
+//!
+//! ## Recursive Queries
+//!
+//! When using DataFusion Tracing with recursive queries (e.g., those using `WITH RECURSIVE`), nodes of type `WorkTableExec` are intentionally not instrumented. This is due to a current limitation in DataFusion: instrumenting these nodes can break recursive query execution and result in errors such as `Unexpected empty work table.`
+//!
+//! As a result, while most of the recursive query plan will be traced and metrics will be collected, some internal operations related to recursion will not be visible in the trace until upstream support is available. See [issue #5](https://github.com/datafusion-contrib/datafusion-tracing/issues/5) for details and tracking progress on this limitation.
+//!
 
 mod instrument_rule;
 mod instrumented;

integration-utils/queries/recursive.sql

@@ -0,0 +1,5 @@
+WITH RECURSIVE numbers(n) AS (
+    SELECT 1 AS n
+    UNION ALL
+    SELECT n + 1 FROM numbers WHERE n < 3
+) SELECT n FROM numbers

tests/integration_tests.rs

@@ -165,6 +165,23 @@ async fn test_scrabble_all_options() -> Result<()> {
     .await
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 8)]
+async fn test_recursive() -> Result<()> {
+    execute_test_case("08_recursive", &QueryTestCase::new("recursive")).await
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 8)]
+async fn test_recursive_all_options() -> Result<()> {
+    execute_test_case(
+        "09_recursive_all_options",
+        &QueryTestCase::new("recursive")
+            .with_metrics_collection()
+            .with_row_limit(5)
+            .with_compact_preview(),
+    )
+    .await
+}
+
 /// Executes the provided [`QueryTestCase`], setting up tracing and verifying
 /// log output according to its parameters.
 async fn execute_test_case(test_name: &str, test_case: &QueryTestCase<'_>) -> Result<()> {

tests/snapshots/03_basic_preview_00_PlaceholderRowExec.snap

@@ -1,5 +1,5 @@
 ---
-source: integration/tests/integration_tests.rs
+source: tests/integration_tests.rs
 expression: preview
 ---
 ++

tests/snapshots/03_basic_preview_00_ProjectionExec.snap

@@ -1,5 +1,5 @@
 ---
-source: integration/tests/integration_tests.rs
+source: tests/integration_tests.rs
 expression: preview
 ---
 +----------+

tests/snapshots/03_basic_preview_01_ProjectionExec.snap

@@ -1,5 +1,5 @@
 ---
-source: integration/tests/integration_tests.rs
+source: tests/integration_tests.rs
 expression: preview
 ---
 ++