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

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 isntrument 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(),
+            },
+        ))
+    }
+
     /// 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
     }

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
 ---
 ++

tests/snapshots/04_basic_compact_preview_00_PlaceholderRowExec.snap

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

tests/snapshots/04_basic_compact_preview_01_ProjectionExec.snap

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