feat(task-history): record Ballista stages for distributed queries (spiceai#10831)

* feat(task-history): capture distributed query observability

For distributed Ballista queries, the task_history table now records the
full parent + per-stage tree:

  sql_query                (parent)  duration, error, datasets, summary labels
  └── ballista_stage       (child)   stage_id, executors, task_count,
                                     stage timing, plan tree

Previously the task_history span was created in submit_distributed_internal
but never instrumented across the job's lifetime — it dropped at submission
time with duration ≈ 0 and no error_message. Stage/executor detail was
not recorded at all.

Implementation:
- QueryHandle owns the sql_query span behind Arc<Mutex<Option<Span>>>.
  spawn_finalize *takes* it (rather than cloning), so the OTel span closes
  when the spawned finalize future drops its instrumented clone — capturing
  the query's runtime, not arbitrary post-completion handle lifetime.
- New stage_history module walks the in-process ExecutionGraph at job
  completion and emits one child ballista_stage span per stage via the
  existing OTel pipeline. Stage labels include executor_histogram,
  slowest_task_ms, partitions, attempt_num, total_executor_ms. Stage
  input is the plan rendered with ExplainFormat::Tree.
- QueryHandle::cancel() now finalizes the tracker with JobCancelled
  before returning, so callers that cancel() then drop the handle
  record an accurate cancellation row instead of the Drop guard's
  client-disconnected fallback.
- Drop guard finalizes orphaned handles: cancels the cancel_token sync,
  then spawns a finalize task that also calls scheduler.cancel_job so
  executors don't keep running an unobserved query.
- Errors raised before QueryHandle creation (planning/validation/submission)
  emit tracing::error on the parent span so the row's error_message is
  populated (mirrors the sync run_internal path).

Tests use a process-wide AsyncMutex (TEST_LOCK) to serialize against
shared DF_SLOT, condition-driven wait_for_row_count polling instead of
fixed sleeps, and parent.job_id-scoped assertions so other tests'
task_history rows in the same binary can't pollute our row counts.

Requires the matching Ballista fork PR (#38, merged onto spiceai-52.5
as 8afc1b74a605) which persists executor_id on terminal TaskInfo and
bumps get_job_execution_graph to pub.

* Bump ballista to pick up stuck-query detection (spiceai/datafusion-ballista#39)

Updates the ballista pin from 8afc1b74 to 7e9872a5 (the merge commit
of spiceai/datafusion-ballista#39) so the next cluster bench run picks
up the new operator-facing warning when a distributed query stops
making progress.

Surgical lockfile update: only the three ballista source URLs change.
Vortex stays pinned at the spiceai-52 commit (c536c9ae) the workspace
was already using; without this, running cargo update -p ballista-*
re-resolves the branch-tracked vortex transitive and bumps it to a
newer rev whose arrow-rs version conflicts with the pinned
datafusion-table-providers (CI symptom: no method as_list on dyn Array,
Arc<Schema>: From<Schema> not satisfied, in adbc compilation).

Verified: cargo check -p runtime --offline and cargo check -p spiced
--features adbc --offline both clean.

* feat(task-history): pluggable span middleware + Ballista stage timeline rendering

Introduces two trait-based extension points on the task_history OTel
exporter and removes the hardcoded `ballista_stage` branches that lived
in the exporter:

  - `SpanTransform`: mutates a `SpanData` before conversion to a row.
    Implementors can adjust timestamps, inject attributes, redact
    fields, etc.

  - `SpanRetention`: declares a retention dependency between sibling
    spans in a batch. Returning `Some(parent_span_id)` from
    `parent_dependency(span)` means "keep this span iff that span is
    kept by base rules" — expressing the parent/child relationship
    directly instead of hardcoding it.

`TaskHistoryExporter` now stores `Vec<Arc<dyn SpanTransform>>` and
`Vec<Arc<dyn SpanRetention>>` with `.with_transform` /
`.with_retention` builder methods. The exporter no longer mentions
any specific span type — base retention checks PLAN_CAPTURE_LABEL and
`min_sql_duration_ms`, retention rules then override per-span.

Concrete consumer: `BallistaStageMiddleware` in
`crates/runtime/src/datafusion/query/stage_history.rs` implements both
traits:

  - As a `SpanTransform`, it rewrites `start_time` / `end_time` on
    `ballista_stage` spans using the `stage_started_at` /
    `stage_ended_at` attributes (millis since UNIX epoch), so the
    task_history row reflects the actual stage execution window. This
    is what makes per-stage execution visible on the timeline view.

  - As a `SpanRetention`, it declares that a `ballista_stage` row
    depends on its parent `sql_query` row, replacing the prior
    "if task.as_ref() == ballista_stage" branch in the exporter.

`BallistaStageMiddleware::pair()` returns a single instance reused as
both trait objects so callers register one Arc in two slots; the three
exporter construction sites (production tracing, generic test util,
distributed_task_history test) all use it.

Includes unit tests covering both hooks: time override on the matching
span name, no-op on non-matching names and missing/inverted/zero
attributes; retention reports the parent dependency for stage spans
and abstains for non-stage spans and orphan stages.

* feat(observability): propagate request trace context across distributed query boundaries (spiceai#10896)

Issue spiceai#10202. Distributed query executors no longer create a fresh internal
request context; they inherit the originating request's protocol and W3C
trace ids so executor-side metrics, telemetry dimensions, and any future
task_history rows correlate end-to-end with the scheduler.

* New `SpiceRequestContextConfig` (`spice_ctx` prefix) carries
  `protocol`, `trace_id`, and `span_id` through DataFusion's
  `ConfigExtension` mechanism — round-tripped opaquely by Ballista as
  `TaskDefinition` config props, with zero fork changes.
* Registered as a default option extension at all 3 config_producer /
  session_builder sites in `runtime/src/cluster/mod.rs`. The scheduler
  session_builder now reads any `SpiceRequestContextConfig` set on the
  per-job session config and re-injects it on the built session config
  (it previously ignored its `_cfg` argument).
* `Query::submit_distributed_internal` populates the extension from the
  current `RequestContext` before `create_or_update_session`.
* `resolve_request_context(TaskContext)` shared helper establishes the
  canonical lookup order: typed `Arc<RequestContext>` extension →
  config extension → `Protocol::Internal` fallback (or `None`).
* `BytesProcessedExec::execute` switched to the helper. Same panic
  behavior preserved when no context is found and `fallback_to_new_context`
  is off.
* `FlightSqlExec` gains an optional `trace_parent: Option<String>` field
  with a `with_trace_parent` builder. `execute()` sets it as a
  `traceparent` gRPC metadata header (via `FlightSqlServiceClient::set_header`),
  falling back to formatting from the typed `RequestContext` extension
  on the `TaskContext` session config when not preset.
  `PartialAggregationFlightSqlExec` inherits and forwards on its own
  `execute()`.
* W3C span semantics preserved: the propagated `span_id` is the sender's
  current span and becomes the parent of any new spans the receiver
  creates — task_history's existing `parent_span_id` chain extends
  naturally one level per executor hop.

Tests: 4 config extension roundtrip + 4 resolver lookup-order +
1 `BytesProcessedExec` integration test asserting metric dimensions
carry the correct protocol when only the config extension is present.

* chore(deps): bump datafusion-ballista to c25e25b9 (spiceai-52.5)

* fix(flightsql): propagate traceparent in pushed-down aggregate exec

PartialAggregationFlightSqlExec was sending no traceparent header when
its source FlightSqlExec was built by the aggregate-pushdown optimizer
(which leaves trace_parent unset). Fall back to TaskContext like the
non-pushed FlightSqlExec path does, so pushed-down aggregate calls
preserve the same propagation behavior.

Also picks up fmt drift in two trace-context files merged via spiceai#10896.

Author: phillipleblanc

Cargo.lock

@@ -426,9 +426,9 @@ datafusion-substrait = { git = "https://github.com/spiceai/datafusion.git", rev
 
 datafusion-table-providers = { git = "https://github.com/datafusion-contrib/datafusion-table-providers.git", rev = "b798c391b6566c172d44361f8acc8472c958ca75" } # spiceai-52
 
-ballista-core = { git = "https://github.com/spiceai/datafusion-ballista.git", rev = "383e165a080d648c313a2530a3a53eae6077fdba" }      # spiceai-52.5
-ballista-executor = { git = "https://github.com/spiceai/datafusion-ballista.git", rev = "383e165a080d648c313a2530a3a53eae6077fdba" } # spiceai-52.5
-ballista-scheduler = { git = "https://github.com/spiceai/datafusion-ballista.git", rev = "383e165a080d648c313a2530a3a53eae6077fdba" } # spiceai-52.5
+ballista-core = { git = "https://github.com/spiceai/datafusion-ballista.git", rev = "c25e25b9e88c50404e230c698a35cbde2fdd30ae" }      # spiceai-52.5
+ballista-executor = { git = "https://github.com/spiceai/datafusion-ballista.git", rev = "c25e25b9e88c50404e230c698a35cbde2fdd30ae" } # spiceai-52.5
+ballista-scheduler = { git = "https://github.com/spiceai/datafusion-ballista.git", rev = "c25e25b9e88c50404e230c698a35cbde2fdd30ae" } # spiceai-52.5
 
 delta_kernel = { git = "https://github.com/spiceai/delta-kernel-rs.git", rev = "47034733a0477f72e4f6abbbf6a27d0da069860a" } # spiceai-0.18.2
 

Cargo.toml

@@ -243,6 +243,8 @@ where
         Some(format!("{host}:{port}").into())
     });
 
+    let (ballista_transform, ballista_retention) =
+        runtime::datafusion::query::stage_history::BallistaStageMiddleware::pair();
     let task_history_exporter = task_history::otel_exporter::TaskHistoryExporter::new(
         df,
         captured_output,
@@ -251,7 +253,9 @@ where
         captured_plan,
         min_plan_duration_ms,
         node_id,
-    );
+    )
+    .with_transform(ballista_transform)
+    .with_retention(ballista_retention);
 
     let zipkin_exporter = zipkin_task_history_otel_exporter(config).await?;
 

bin/spiced/src/tracing.rs

@@ -88,6 +88,7 @@ rdkafka = { workspace = true, features = ["ssl-vendored"], optional = true }
 regex.workspace = true
 reqwest.workspace = true
 runtime-rate-control = { path = "../runtime-rate-control" }
+runtime-request-context = { path = "../runtime-request-context" }
 rusqlite = { workspace = true, optional = true }
 elasticsearch = { path = "../elasticsearch", optional = true }
 s3_vectors = { path = "../s3_vectors", optional = true }

crates/data_components/Cargo.toml

@@ -57,11 +57,23 @@ use datafusion::{
     },
     sql::TableReference,
 };
+use runtime_request_context::RequestContext;
 use tonic::codegen::Bytes;
 use tonic::transport::{Channel, channel};
 
 use crate::Read;
 
+/// Build a W3C `traceparent` header value (`00-{trace_id}-{span_id}-01`)
+/// from the typed `Arc<RequestContext>` extension on the
+/// [`TaskContext`]'s session config, if one is present. The fixed `01`
+/// trace flag marks the trace as sampled.
+#[must_use]
+pub fn trace_parent_from_task_context(context: &TaskContext) -> Option<String> {
+    let request_context = context.session_config().get_extension::<RequestContext>()?;
+    let tp = request_context.trace_parent().as_ref()?;
+    Some(format!("00-{}-{}-01", tp.trace_id, tp.span_id))
+}
+
 pub mod federation;
 
 #[derive(Debug, Snafu)]
@@ -432,6 +444,12 @@ pub struct FlightSqlExec {
     properties: PlanProperties,
     cookie_store: Arc<CookieStore>,
     metrics: ExecutionPlanMetricsSet,
+    /// Optional W3C `traceparent` value (e.g. `00-{trace_id}-{span_id}-01`)
+    /// to attach as a gRPC metadata header on outgoing `execute()` and
+    /// `do_get()` calls. When `None`, `execute()` falls back to reading
+    /// the typed `Arc<RequestContext>` extension from the `TaskContext`
+    /// session config and constructs a header from its `trace_parent()`.
+    trace_parent: Option<String>,
 }
 
 impl FlightSqlExec {
@@ -460,9 +478,27 @@ impl FlightSqlExec {
             ),
             cookie_store,
             metrics: ExecutionPlanMetricsSet::new(),
+            trace_parent: None,
         })
     }
 
+    /// Set an explicit W3C `traceparent` header value to forward on each
+    /// outgoing `FlightSQL` call. Useful when the plan-creation path has
+    /// access to an `Arc<RequestContext>` but the executor-side
+    /// `TaskContext` will not (e.g. when this `ExecutionPlan` is shipped
+    /// to a remote executor via Ballista codecs).
+    #[must_use]
+    pub fn with_trace_parent(mut self, trace_parent: Option<String>) -> Self {
+        self.trace_parent = trace_parent;
+        self
+    }
+
+    /// Returns the currently configured W3C `traceparent` value, if any.
+    #[must_use]
+    pub fn trace_parent(&self) -> Option<&str> {
+        self.trace_parent.as_deref()
+    }
+
     /// Returns a reference to the underlying `FlightSqlClient`.
     #[must_use]
     pub fn client(&self) -> &FlightSqlClient {
@@ -655,6 +691,7 @@ impl ExecutionPlan for FlightSqlExec {
             ),
             cookie_store: Arc::clone(&self.cookie_store),
             metrics: ExecutionPlanMetricsSet::new(),
+            trace_parent: self.trace_parent.clone(),
         };
 
         Ok(SortOrderPushdownResult::Exact {
@@ -665,7 +702,7 @@ impl ExecutionPlan for FlightSqlExec {
     fn execute(
         &self,
         partition: usize,
-        _context: Arc<TaskContext>,
+        context: Arc<TaskContext>,
     ) -> DataFusionResult<SendableRecordBatchStream> {
         let sql = self.sql().map_err(to_execution_error)?;
         let target_schema = self.schema();
@@ -676,9 +713,19 @@ impl ExecutionPlan for FlightSqlExec {
 
         let baseline = datafusion::common::instant::Instant::now();
 
-        let inner = query_to_stream(self.client.clone(), sql, Arc::clone(&self.cookie_store)).map(
-            move |result| result.and_then(|batch| coerce_batch_to_schema(&batch, &target_schema)),
-        );
+        let mut client = self.client.clone();
+        let trace_parent = self
+            .trace_parent
+            .clone()
+            .or_else(|| trace_parent_from_task_context(&context));
+        if let Some(value) = trace_parent {
+            client.set_header("traceparent", value);
+        }
+
+        let inner =
+            query_to_stream(client, sql, Arc::clone(&self.cookie_store)).map(move |result| {
+                result.and_then(|batch| coerce_batch_to_schema(&batch, &target_schema))
+            });
 
         let timed_stream = stream! {
             futures::pin_mut!(inner);
@@ -727,6 +774,7 @@ impl ExecutionPlan for FlightSqlExec {
             properties: self.properties.clone(),
             cookie_store: Arc::clone(&self.cookie_store),
             metrics: ExecutionPlanMetricsSet::new(),
+            trace_parent: self.trace_parent.clone(),
         };
 
         Some(Arc::new(new_plan))

crates/data_components/src/flightsql.rs

@@ -40,7 +40,9 @@ use datafusion::physical_plan::{
 };
 use datafusion::sql::TableReference;
 
-use data_components::flightsql::{FlightSqlClient, FlightSqlExec, query_to_stream};
+use data_components::flightsql::{
+    FlightSqlClient, FlightSqlExec, query_to_stream, trace_parent_from_task_context,
+};
 use data_components::sql_expr::to_sql_preserving_precedence;
 use flight_client::cookie::CookieStore;
 use futures::StreamExt;
@@ -153,6 +155,11 @@ pub struct PartialAggregationFlightSqlExec {
     cookie_store: Arc<CookieStore>,
     /// Cached plan properties.
     properties: PlanProperties,
+    /// Optional W3C `traceparent` value inherited from the source
+    /// `FlightSqlExec`. Forwarded as a gRPC metadata header on each
+    /// outgoing call so executor-side spans chain back to the originating
+    /// request.
+    trace_parent: Option<String>,
 }
 
 impl PartialAggregationFlightSqlExec {
@@ -181,6 +188,7 @@ impl PartialAggregationFlightSqlExec {
             client: source.client().clone(),
             cookie_store: Arc::clone(source.cookie_store()),
             properties,
+            trace_parent: source.trace_parent().map(str::to_string),
         }
     }
 
@@ -250,7 +258,7 @@ impl ExecutionPlan for PartialAggregationFlightSqlExec {
     fn execute(
         &self,
         partition: usize,
-        _context: Arc<TaskContext>,
+        context: Arc<TaskContext>,
     ) -> Result<SendableRecordBatchStream> {
         if partition != 0 {
             return Err(DataFusionError::Execution(format!(
@@ -261,14 +269,20 @@ impl ExecutionPlan for PartialAggregationFlightSqlExec {
         let target_schema = self.schema();
         let column_mapping = query.column_mapping;
 
-        let stream = query_to_stream(
-            self.client.clone(),
-            query.sql,
-            Arc::clone(&self.cookie_store),
-        )
-        .map(move |result: std::result::Result<_, DataFusionError>| {
-            result.and_then(|batch| remap_batch(&batch, &column_mapping, &target_schema))
-        });
+        let mut client = self.client.clone();
+        let trace_parent = self
+            .trace_parent
+            .clone()
+            .or_else(|| trace_parent_from_task_context(&context));
+        if let Some(value) = trace_parent {
+            client.set_header("traceparent", value);
+        }
+
+        let stream = query_to_stream(client, query.sql, Arc::clone(&self.cookie_store)).map(
+            move |result: std::result::Result<_, DataFusionError>| {
+                result.and_then(|batch| remap_batch(&batch, &column_mapping, &target_schema))
+            },
+        );
 
         Ok(Box::pin(RecordBatchStreamAdapter::new(
             self.schema(),

crates/datafusion-optimizer-rules/src/physical_plan/flightsql/exec.rs

@@ -1 +1,2 @@
 pub mod cluster_config;
+pub mod request_context_config;