Recursive queries fail with Unexpected empty work table. with instrumentation

Describe the bug

Recursive queries fail with Unexpected empty work table with instrumentation.

To Reproduce

use datafusion::{
    arrow::{array::RecordBatch, util::pretty::pretty_format_batches},
    error::Result,
    execution::SessionStateBuilder,
    prelude::*,
};
use datafusion_tracing::{
    InstrumentationOptions, instrument_with_info_spans, pretty_format_compact_batch,
};
use std::sync::Arc;
use tracing::field;

#[tokio::main]
async fn main() -> Result<()> {
    let options = InstrumentationOptions::builder()
        .record_metrics(true)
        .preview_limit(5)
        .preview_fn(Arc::new(|batch: &RecordBatch| {
            pretty_format_compact_batch(batch, 64, 3, 10).map(|fmt| fmt.to_string())
        }))
        .add_custom_field("env", "production")
        .add_custom_field("region", "us-west")
        .build();

    let instrument_rule = instrument_with_info_spans!(
        options: options,
        env = field::Empty,
        region = field::Empty,
    );

    let session_state = SessionStateBuilder::new()
        .with_default_features()
        .with_physical_optimizer_rule(instrument_rule)
        .build();

    let ctx = SessionContext::new_with_state(session_state);

    println!("Running recursive CTE query...");
    let recursive_sql = "WITH RECURSIVE numbers(n) AS (\
        SELECT 1 AS n \
        UNION ALL \
        SELECT n + 1 FROM numbers WHERE n < 10 \
    ) SELECT n FROM numbers";

    match ctx.sql(recursive_sql).await?.collect().await {
        Ok(recursive_results) => {
            println!(
                "Recursive Query Results (1..10):\n{}",
                pretty_format_batches(recursive_results.as_slice())?
            );
        }
        Err(e) => {
            println!("Recursive CTE failed with error: {}", e);
            println!("This demonstrates the instrumentation issue!");
        }
    }

    println!("\nTesting without instrumentation...");
    let clean_session_state = SessionStateBuilder::new().with_default_features().build();

    let clean_ctx = SessionContext::new_with_state(clean_session_state);

    match clean_ctx.sql(recursive_sql).await?.collect().await {
        Ok(clean_results) => {
            println!(
                "Clean Recursive Query Results (1..10):\n{}",
                pretty_format_batches(clean_results.as_slice())?
            );
            println!("Without instrumentation, recursive CTE works!");
        }
        Err(e) => {
            println!("Even clean context failed: {}", e);
        }
    }

    Ok(())
}

Expected behavior

No response

Additional context

No response

@debajyoti-truefoundry

Hi @debajyoti-truefoundry,

Thanks for the super detailed bug report! Having a working reproducer is already 95% of the work done.

Looking into this, the root cause is in the assign_work_table function of recursive_query.rs:

fn assign_work_table(
    plan: Arc<dyn ExecutionPlan>,
    work_table: Arc<WorkTable>,
) -> Result<Arc<dyn ExecutionPlan>> {
    let mut work_table_refs = 0;
    plan.transform_down(|plan| {
        if let Some(exec) = plan.as_any().downcast_ref::<WorkTableExec>() {
            if work_table_refs > 0 {
                not_impl_err!(
                    "Multiple recursive references to the same CTE are not supported"
                )
            } else {
                work_table_refs += 1;
                Ok(Transformed::yes(Arc::new(
                    exec.with_work_table(Arc::clone(&work_table)),
                )))
            }
        } else if plan.as_any().is::<RecursiveQueryExec>() {
            not_impl_err!("Recursive queries cannot be nested")
        } else {
            Ok(Transformed::no(plan))
        }
    })
    .data()
}

Once instrumented, all nodes of the ExecutionPlan are IntrumentedExec nodes, so trying to find a WorkTableExec in it will always fail.

I'll see if I can update the datafusion-tracing crate to still be compatible with this pattern, otherwise the fix needs to be done upstream by using traits instead of concrete types.

I've got a fix working locally (which did need an upstream change unfortunately.) But the results are pretty interesting as it spawns a new span for each recursive call (since the recursive implementation recreates the recursive nodes for loop.)

It's readable for smallish recursions, but can quickly explode if you recurse over a large number (although anything over ~50 is probably a pretty inefficient query anyways.)

With the following query, exporting to Jaeger:

WITH RECURSIVE numbers(n) AS (
    SELECT 1 AS n
    UNION ALL
    SELECT n + 1 FROM numbers WHERE n < 3
) SELECT n FROM numbers

I'll release a patch version 48.0.1 once #6 is merged, which will enable tracing of recursive queries (only limitation is that the WorkTableExec nodes are ignored, which needs apache/datafusion#16469 to be merged and a new DataFusion version released.)

Not closed until version 49 is released with apache/datafusion#16469 or similar merged in DataFusion repository.

@debajyoti-truefoundry

@debajyoti-truefoundry I've released https://crates.io/crates/datafusion-tracing/48.0.1, if you want to try out the fix. As well as https://crates.io/crates/datafusion-tracing/47.0.2 if you're still on DataFusion 47.

As well as https://crates.io/crates/datafusion-tracing/47.0.2 if you're still on DataFusion 47.

Thanks a lot. This is working.

Navigation Menu

Repository navigation

Recursive queries fail with `Unexpected empty work table.` with instrumentation #5

Describe the bug

To Reproduce

Expected behavior

Additional context

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Recursive queries fail with Unexpected empty work table. with instrumentation #5

Description

Describe the bug

To Reproduce

Expected behavior

Additional context

Activity

geoffreyclaude commented on Jun 19, 2025

geoffreyclaude commented on Jun 19, 2025

geoffreyclaude commented on Jun 20, 2025

geoffreyclaude commented on Jun 20, 2025

geoffreyclaude commented on Jun 20, 2025

debajyoti-truefoundry commented on Jun 20, 2025

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

Recursive queries fail with `Unexpected empty work table.` with instrumentation #5