Tracing physical plan creation methods

[chiragjn]

Is your feature request related to a problem or challenge?

It would be great to be able to trace physical plan creation methods. We were trying to figure out large gaps between table loading and query execution and it turned out it was physical plan creation. We had to fork datafusion to add instrumentation like follows:
truefoundry/datafusion@ffc695d

Describe the solution you'd like

No response

Describe alternatives you've considered

No response

Additional context

No response

[geoffreyclaude]

Hi @chiragjn ,
How are you running queries exactly? I take it you're not creating the physical plans yourselves, but rather using a higher level API?
On our side, we control the logical and physical plan creation, so it's easy to instrument in user space without changes to the core DataFusion code.

[chiragjn]

Yeah, we are calling datafusion like this

let tracing_options = InstrumentationOptions::builder()
        .record_metrics(true)
        .preview_limit(config.datafusion_tracing_preview_limit)
        .build();

let instrument_rule = instrument_with_info_spans!(
        options: tracing_options,
        service = tracing::field::Empty,
        component = tracing::field::Empty,
);

let session_state = datafusion::execution::SessionStateBuilder::new()
        .with_config(session_config)
        .with_runtime_env(runtime_env)
        .with_default_features()
        .with_catalog_list(catalog_provider_list)
        .with_physical_optimizer_rule(instrumentation_rule)
        .build();

let session_context = Arc::new(SessionContext::new_with_state(session_state))

let mut df = session_context
            .sql_with_options(
                &validated_query,
                SQLOptions::new()
                    .with_allow_ddl(false)
                    .with_allow_dml(false)
                    .with_allow_statements(false),
            )
            .await?;

let mut stream: SendableRecordBatchStream = df.execute_stream().await?;

With this we get a gap which is the time spent in create_physical_plan

[geoffreyclaude]

@chiragjn Thanks, it's much clearer now.

I'm working on instrumenting the plan creation phases as well, including individual rules. Is this something you'd be interested in? This is what it looks like so far on the otlp.rs example from the repo:

[chiragjn]

Absolutely, the more info we can collect the better