Support "Tracing" / Spans

[alamb]

Is your feature request related to a problem or challenge?

"Tracing" is a visualization technique for understanding how potentially concurrent operations happen, and is common in distributed systems.

You can also visualize DataFusion executions using traces. Here is an example trace we have at InfluxData that is integrated into the rest of our system and visualized using https://www.jaegertracing.io/

This visualization shows when each operator started and stopped (and the operators are annotated with how much CPU time is spent, etc). These spans are integrated into the overall trace of a request through our cloud service, which allows us to understand where a request's time is spent, both across services as well as within our DataFusion based engine.

For more background on tracing, this blog seems to give a reasonable overview: https://signoz.io/blog/distributed-tracing-span/

Describe the solution you'd like

I would like to make it easy for people to add DataFusion ExecutionPlan level tracing to their systems as well.

Given the various different libraries to generate traces I don't think picking any particular one to build into DataFusion is a good idea. However adding some way to walk the ExecutionPlan metrics and emit information that can be turned into traces would be very helpful I think

This came up twice recently, so I wanted to get it filed into a ticket

Describe alternatives you've considered

No response

Additional context

@simonvandel noted in Discord

I have been playing a bit with https://github.com/wolfpld/tracy and https://docs.rs/tracing-tracy/latest/tracing_tracy/ and thought it would be nice to have some visibility into ExecutionPlans being executed.

It also came up in slack: https://the-asf.slack.com/archives/C04RJ0C85UZ/p1709051125059619

But I want to see traces from which I want to calculate how much time spent in executing the queries and how much time got spent into data transfer from a target storage system, It would be good, if you can guide me how can I get these metrics

[alamb]

The way we get spans in InfludDb is to use metrics which has the start/stop times as well as counters of what rows have been produced, time the execution started, stopped, and various other things

The DataFusion metrics --> jarger span exporter for InfluxDB 3.0 can be found here https://github.com/influxdata/influxdb/blob/8fec1d636e82720389d06355d93245a06a8d90ad/iox_query/src/exec/query_tracing.rs#L93-L112

It uses our own span generation system as the rust ecosystem didn't have a standard library that we could find when we originally wrote it.

One possibility might be to refactor the code out of InfluxDB to make it generic (so it didn't depend on span, but instead was implemented in terms of some trait (which we could then impl for span

[erratic-pattern]

The tracing crate and related ecosystem would be a good choice for this. It is full-featured with an easy-to-use and unobtrusive API, and there is a lot of community support for tooling integration and export formats.

I think we could get most of the way there by implementing the following changes:

• include the #[instrument] macro on all execute calls
• call the with_current_subscriber method from the tracing_futures crate on all relevant streams and futures (replacing ObservedStream)
• call the instrument method from the tracing_futures crate on all relevant streams and futures. (EDIT: alternatively we might be able to attach the #[instrument] macro to Stream poll_next functions but need to test this)
• record when rows are added via a tracing event. example: info!(num_rows=1, "row added"). the tracing Subscriber can aggregate these events to obtain the total number of rows added
• similar to num_rows, other metrics (spill count, count, memory used, etc) can be calculated via trace events that can be aggregated to produce the metric

With these changes, all the data of the existing MetricsSet code can captured:

• start time (wall clock) corresponds to the first "enter" event of the Span
• end time (wall clock) corresponds to the "close" event of the Span
• compute time can be calculated by aggregating all the intervals between "enter" and "exit" events on a Span (Spans can be entered and exited multiple times, as is this case with polling streams). The tracing_timing crate provides a Subscriber to calculate this information.
• other metrics can be calculated by aggregating all of the events containing field name-value pairs that the metric is interested in

for compatibility with code that is reading metrics from the existing MetricsSet, we could also implement a Subscriber and/or Layer that calculates these metrics and produces a MetricsSet. However, a better approach for most users would be to simply install one of the many tracing ecosystem crates, which provide subscribers for a number of tools and data formats.

Some benefits compared to the current metrics implementation:

• event-based tracing allows for more detailed information to be captured than the current metrics-based approach
• tracing of individual streams/futures provides more information for debugging async code
• less API surface area since we don't need to include metrics in the ExecutionPlan trait
• overall less code is needed to implement tracing as a lot of the boilerplate is handled automatically via tracing macros and implicit thread-local state. the tracing API is very easy to use from the event producing side.
• access to a wide variety of ecosystem crates for exporting traces
• doubles as a logging framework. DataFusion users can register a FmtSubscriber to get formatted logs (which we could also use in the DataFusion CLI), and there are many crates with additional logging subscribers for standardized logging formats and third-party logging tools.
• tests can implement custom Subscribers to validate low-level events
• tracing can be disabled when the overhead is not wanted

Downsides:

• tracing event fields are free-form, so there is less static typing guarantees around the correctness of event fields. things like typos will not be caught as easily. this could potentially be remedied by writing custom macros which could enforce and/or automatically include some standardized set of fields, or accept well-formed structs instead of free-form key-value pairs.
• tracing API is a bit less "explicit" than the current metrics code, because of the implicit thread-local state. if you don't include the relevant macros and function calls, the compiler won't complain about it. this can make it easier for programmers to forget about tracing when developing custom ExecutionPlans, and makes it more difficult to enforce consistent field name conventions and practices. again, I think this could be remedied to some degree by custom macros (and some good documentation)
• a bit more complexity on the subscriber-side to aggregate events within custom Subscriber implementations

[alamb]

Another downside of tracing is that not all systems use that ecosystem so it might require some extra logic to use tracing than metrics

But otherwise the high level idea seems plausble, if a bunch of work. Maybe a quick POC could be done to see what it might look like in practice / how much work / change was required

[erratic-pattern]

Another downside of tracing is that not all systems use that ecosystem so it might require some extra logic to use tracing than metrics

I think this is where custom Subscribers could be useful for generating DataFusion-native reporting and metrics. Having a Subscriber to generate the existing MetricsSet should be possible. From the user side there is no need to include tracing-related dependencies if configured with built-in Subscribers

But otherwise the high level idea seems plausble, if a bunch of work. Maybe a quick POC could be done to see what it might look like in practice / how much work / change was required

Agreed. What would be a good scope for a POC to be both quick to implement and broad enough to cover all related cases? Maybe a subset of operators/streams that are used by a specific, but non-trivial, query.

The existing metrics code can be left as-is until we reach a point in implementation where we're confident that tracing can replace those functionalities. Also a lot of it will continue to be used I think for generating DataFusion-native metrics.

[erratic-pattern]

After spending some time with the optimizer, I think it would be a good candidate to PoC the tracing crate.

• It is not async, which makes it much easier to instrument with tracing
• It is a relatively small and self-contained codebase
• We can get some useful metrics, for example memory allocation stats, that would be helpful in identifying hotspots as part of reducing copying of the LogicalPlan [EPIC] Stop copying LogicalPlan during OptimizerPasses #9637 as well as reducing the number of string allocations throughout the optimizer Make CommonSubexprEliminate faster by stop copying so many strings #10426 (comment)

[erratic-pattern]

create_physical_plan would be a good follow up, to PoC instrumentation of async code

[aditanase]

Hi all, thanks for adding this and investigating the tracing crate. I'd like to suggest being a bit more specific about the goals of adding tracing, before jumping in both feet :). Maybe I can pitch in some use cases to help with this.

My team is prototyping a distributed engine on top of ballista. Since ballista doesn't yet have a great UI, we started to look at adding some end-to-end tracing (think external client -> flight SQL query -> scheduler -> enqueue job -> executors -> DF engine). As we realised there is currently no tracing in either project, we quickly found this issue.

I think the tracing crate, together with some of the community subscribers (e.g. opentelemetry stack) can solve this problem, even though there are a number of challenges:

• correctly instrumenting streams and async blocks
• using or_current consistently to make sure that changing the log level between trace, debug and info doesn't throw away good information
• tracing across service boundaries (for ballista especially, but also for an app embedding datafusion)
• tracing pub/sub patterns (in ballista we have a job queue through which scheduler and workers are decoupled)

To that end, I'd like to understand if reimplementing metrics on top of tracing is really what this issue is about, or just an attempt at consolidating some of the timing / metrics bookkeeping?

Based on my experience with other systems (mostly on the JVM, building and tuning Spark / Kafka deployments), tracing and metrics work really well together, but they are rarely conflated.
This goes especially for core internal metrics that are used by the query engine (e.g. explain analyse, cost based optimizations, building a nice UI for the scheduler) as opposed to tracing that is typically done through sampling (tracing everything at a certain level becomes expensive quickly), has a lot of user input (extra context, combined with app metrics) and is configurable through the log level.

My suggestion would be to decouple adding tracing (as a tool for people that are monitoring / optimizing engines built on top of DF) from the core metrics refactoring.
For metrics specifically there are other crates with more target concepts (counters, gauges, histograms) that have some integration with tracing, in order to propagate the current span context as metrics labels.

Lastly, if there is not a lot of work started here, I've already started to play around with some of the suggestions on this thread (add instrument to execute, instrument streams and async blocks, etc) and I'd be interested in contributing to this track, especially some of the lessons learned around tracing async code and streams.

[alamb]

Since the the internal metrics in datafusion -- aka https://docs.rs/datafusion/latest/datafusion/physical_plan/metrics/index.html have start/stop timestamps on them already, we have found it relatively straightforward to convert them to "tracing" spans -- the link to do so is above:

#9415 (comment)

I am not clear what additional benefit more direct tracing integration in datafusion would provide, but I may be missing something

[geoffreyclaude]

Hi,

Using internal metrics to generate tracing spans on query completion is a great solution. However it prevents tracing downstream data sources, which are often more complex than an on-disk Parquet read.

I've been experimenting with tracing a full execution plan by injecting a custom "TracingExec" node before every execution node in the plan. Its execute method starts a new span with the child's name (using Open Telemetry's otel.name field), and streams the child result decorated with the tracing crate's in_current_span() through it. It also reports all child metrics as fields on the span, which is super nice when exploring the generated flame graph.

The huge advantage of this method is that it avoids needing to modify every execution node's execute method, allows propagating the trace context downstream, and could work as an independent DataFusion "contrib" crate once the code is tidied up.

However, there is still a major blocking point in the current DataFusion code: spawning tasks in new threads. Whenever a task is spawned in a new thread (eg, inside a RepartitionExec, or inside the ParquetSink's write_all method), the context is lost (because the tracing context is thread-local only.) I've worked around this best I could by manually tracking the spans inside the "TracingExec", but it's a pretty wonky workaround. It also fails when the source nodes of the plan spawn new tasks, as in the ParquetSink, which prevent the underlying object_store from linking its spans to the main trace.

Even without spans, it would be useful to be able to link logs generated deep in the plan (by a custom object_store for instance) to a given trace, which would also need access to the trace context.

I'd like to propose a very simple change first, before going full tracing everywhere: wrapping all task spawn points with .in_current_span(). This ensures that when a task moves to a different thread, it carries along its current tracing context. It could be protected behind a tracing feature, so users who don't want tracing don't need to import the crate. I believe there are roughly 20 places in the DataFusion repo that would need this change, so fairly small change overall. I'd be happy to submit a PR if we can agree on this first step!

For instance, in physical-plan/src/stream.rs (untested code!):

    #[cfg(feature = "tracing")]
    use tracing_futures::Instrument;
    /// Spawn task that will be aborted if this builder (or the stream
    /// built from it) are dropped
    pub fn spawn<F>(&mut self, task: F)
    where
        F: Future<Output = Result<()>>,
        F: Send + 'static,
    {
        #[cfg(feature = "tracing")]
        let task = task.in_current_span();
        self.join_set.spawn(task);
    }

[erratic-pattern]

I'd like to propose a very simple change first, before going full tracing everywhere: wrapping all task spawn points with .in_current_span().

I agree with this change. When running instrumentation on top of datafusion this would make it possible to extract more information.

[erratic-pattern]

I am not clear what additional benefit more direct tracing integration in datafusion would provide, but I may be missing something

The tracing API is more granular than the information provided by the current metrics, so while it's possible to convert the existing metrics into tracing spans, there is some information that is inaccessible or impossible to trace at the moment.

The main example I can think of is the exact timing of entering/exiting execution of the async tasks. The datafusion metrics record a "start" and "end" timestamp for the whole operator, but they do not record when operators await and give up control to the executor.

The tracing API allows for this because a span can be entered and exited multiple times before it is finally closed. This allows you to graph out exactly when async tasks are running in relation to each and for how long, which can be helpful for identifying bottlenecks where a task is waiting a long time for data from another task.

You can kinda use the elapsed_compute metric for this purpose, if you are only interested in identifying which operators are slow and fast, but the added granularity that you get from tracing would make it easier to visualize the actual path that control flow takes when tasks are pre-empted.

[geoffreyclaude]

I've opened #14547 to launch the discussion. PR was simpler than expected as all new tasks spawned do so in the context of a tokio JoinSet, so I've just wrapped the tokio JoinSet in an implementation which adds .in_current_span on each new spawn.