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

[debajyoti-truefoundry]

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

[geoffreyclaude]

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.

[geoffreyclaude]

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

[geoffreyclaude]

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.)

[geoffreyclaude]

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

[geoffreyclaude]

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

[debajyoti-truefoundry]

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.