diff --git a/Cargo.lock b/Cargo.lock index da4c97d3784..cbdc182ca2a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4623,6 +4623,8 @@ dependencies = [ "rstest", "tokio", "tracing", + "tracing-mock", + "tracing-subscriber", ] [[package]] diff --git a/rust/lance-datafusion/Cargo.toml b/rust/lance-datafusion/Cargo.toml index 8b9c702d2fb..aa9bb84bf71 100644 --- a/rust/lance-datafusion/Cargo.toml +++ b/rust/lance-datafusion/Cargo.toml @@ -44,6 +44,8 @@ protobuf-src = { workspace = true, optional = true } [dev-dependencies] lance-datagen.workspace = true rstest.workspace = true +tracing-mock.workspace = true +tracing-subscriber = "0.3.17" [features] datagen = ["dep:lance-datagen"] diff --git a/rust/lance-datafusion/src/exec.rs b/rust/lance-datafusion/src/exec.rs index b5ae15c5267..8d09650f272 100644 --- a/rust/lance-datafusion/src/exec.rs +++ b/rust/lance-datafusion/src/exec.rs @@ -52,7 +52,7 @@ use lance_core::{ }, }; use log::{debug, info, warn}; -use tracing::Span; +use tracing::{Span, instrument}; use crate::udf::register_functions; use crate::{ @@ -680,6 +680,7 @@ fn display_plan_one_liner_impl(plan: &dyn ExecutionPlan, output: &mut String) { /// Executes a plan using default session & runtime configuration /// /// Only executes a single partition. Panics if the plan has more than one partition. +#[instrument(level = "debug", skip_all)] pub fn execute_plan( plan: Arc, options: LanceExecutionOptions, @@ -737,6 +738,7 @@ pub async fn analyze_plan( /// the context carrying those extensions; otherwise the nodes error during /// `execute` and `AnalyzeExec` reports an empty, unexecuted plan tree instead /// of surfacing the error. +#[instrument(level = "debug", name = "analyze_plan", skip_all)] pub async fn analyze_plan_with_context( plan: Arc, options: LanceExecutionOptions, @@ -1232,6 +1234,16 @@ impl ExecutionPlan for HardCapBatchSizeExec { mod tests { use super::*; + use datafusion::physical_plan::empty::EmptyExec; + use tracing::{ + Level, + span::{Attributes, Id}, + }; + use tracing_mock::{expect, subscriber}; + use tracing_subscriber::{ + Layer, layer::Context as LayerContext, layer::SubscriberExt, registry::LookupSpan, + }; + // Serialize cache tests since they share global state static CACHE_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); @@ -1370,42 +1382,62 @@ mod tests { #[derive(Debug)] struct RequiredExtension; - /// Execution node that only succeeds when [`RequiredExtension`] is present - /// on the task context's session config. This mirrors distributed routing - /// nodes that read a session-config identity extension during `execute`. - #[derive(Debug)] - struct NeedsExtensionExec { + type ExecuteFn = Box< + dyn Fn(Arc) -> datafusion_common::Result + + Send + + Sync, + >; + + /// Single-partition execution node over an empty schema whose `execute` body + /// the test supplies. `name` is what the node reports as itself, including + /// in an analyze report. + struct TestExec { + name: &'static str, properties: Arc, - /// Set once the node reaches execution with the extension present. - /// Observed by the test so that dropping context forwarding (which - /// makes `execute` error before this point) is detectable. - executed: Arc, - } - - impl NeedsExtensionExec { - fn new(executed: Arc) -> Self { + on_execute: ExecuteFn, + } + + impl TestExec { + fn new( + name: &'static str, + on_execute: impl Fn( + Arc, + ) -> datafusion_common::Result + + Send + + Sync + + 'static, + ) -> Self { let schema = Arc::new(ArrowSchema::empty()); Self { + name, properties: Arc::new(PlanProperties::new( EquivalenceProperties::new(schema), Partitioning::UnknownPartitioning(1), EmissionType::Incremental, Boundedness::Bounded, )), - executed, + on_execute: Box::new(on_execute), } } } - impl DisplayAs for NeedsExtensionExec { + impl std::fmt::Debug for TestExec { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + f.debug_struct("TestExec") + .field("name", &self.name) + .finish() + } + } + + impl DisplayAs for TestExec { fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { - write!(f, "NeedsExtensionExec") + write!(f, "{}", self.name) } } - impl ExecutionPlan for NeedsExtensionExec { + impl ExecutionPlan for TestExec { fn name(&self) -> &str { - "NeedsExtensionExec" + self.name } fn properties(&self) -> &Arc { &self.properties @@ -1424,25 +1456,17 @@ mod tests { _partition: usize, context: Arc, ) -> datafusion_common::Result { - if context - .session_config() - .get_extension::() - .is_none() - { - return Err(DataFusionError::Execution( - "missing required session-config extension".to_string(), - )); - } - self.executed - .store(true, std::sync::atomic::Ordering::SeqCst); - let schema = self.schema(); - Ok(Box::pin(RecordBatchStreamAdapter::new( - schema, - stream::empty(), - ))) + (self.on_execute)(context) } } + fn empty_stream() -> SendableRecordBatchStream { + Box::pin(RecordBatchStreamAdapter::new( + Arc::new(ArrowSchema::empty()), + stream::empty(), + )) + } + // Regression: analyze must run under a caller-provided TaskContext so nodes // that read a session-config extension at execute time see it. Without the // context the node errors and AnalyzeExec would otherwise report an empty, @@ -1451,8 +1475,27 @@ mod tests { async fn test_analyze_plan_uses_provided_task_context() { use std::sync::atomic::{AtomicBool, Ordering}; + // The node mirrors distributed routing nodes that read a session-config + // identity extension during `execute`. It sets `executed` only once it + // reaches execution with the extension present, so dropping context + // forwarding (which makes `execute` error first) is detectable. let executed = Arc::new(AtomicBool::new(false)); - let plan: Arc = Arc::new(NeedsExtensionExec::new(executed.clone())); + let plan: Arc = Arc::new(TestExec::new("NeedsExtensionExec", { + let executed = executed.clone(); + move |context| { + if context + .session_config() + .get_extension::() + .is_none() + { + return Err(DataFusionError::Execution( + "missing required session-config extension".to_string(), + )); + } + executed.store(true, Ordering::SeqCst); + Ok(empty_stream()) + } + })); // Default context lacks the extension: the node errors during execute // (never reaching the `executed` flag), but AnalyzeExec absorbs that @@ -1500,4 +1543,114 @@ mod tests { "supplied context must be forwarded so the node executes" ); } + + /// Records `(name, level, parent name)` for every span opened while it is + /// installed. Tests read the recording back after the subscriber is + /// uninstalled, which keeps the assertions independent of the order spans + /// are opened in and reports a mismatch as an ordinary assertion failure + /// rather than a panic inside a tracing callback. + #[derive(Clone, Default)] + struct SpanRecorder(Arc>>); + + /// A span as the recorder saw it: its name and level, and the name of the + /// span it was opened under. + type RecordedSpan = (&'static str, Level, Option<&'static str>); + + impl SpanRecorder { + fn recorded(&self) -> Vec { + self.0.lock().unwrap().clone() + } + } + + impl Layer for SpanRecorder + where + S: tracing::Subscriber + for<'a> LookupSpan<'a>, + { + fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: LayerContext<'_, S>) { + let parent = ctx + .span(id) + .and_then(|span| span.parent()) + .map(|parent| parent.name()); + let meta = attrs.metadata(); + self.0 + .lock() + .unwrap() + .push((meta.name(), *meta.level(), parent)); + } + } + + #[test] + fn test_execute_plan_opens_execute_plan_span() { + let plan: Arc = Arc::new(EmptyExec::new(Arc::new(ArrowSchema::empty()))); + + let (subscriber, handle) = subscriber::mock() + .with_filter(|meta| meta.name() == "execute_plan") + .new_span(expect::span().named("execute_plan").at_level(Level::DEBUG)) + .run_with_handle(); + + tracing::subscriber::with_default(subscriber, || { + execute_plan(plan, LanceExecutionOptions::default()).unwrap(); + }); + + handle.assert_finished(); + } + + // Unlike `execute_plan`, which returns before the stream is polled, analyze + // drives the plan to completion inside the call and re-parents the plan + // under its own span via `TracedExec`, so the spans a plan opens while it + // runs must sit underneath `analyze_plan` rather than beside it. + #[test] + fn test_analyze_plan_span_parents_plan_execution() { + // `TracedExec` re-parents in two places: around `execute`, and around + // every poll of the stream that `execute` returns. Real operators open + // their IO spans on the polling path, so cover both. + let plan: Arc = Arc::new(TestExec::new("SpanningExec", |_context| { + let _span = tracing::info_span!("plan_execute").entered(); + let batches = stream::once(async { + let _span = tracing::info_span!("plan_poll").entered(); + Ok(RecordBatch::new_empty(Arc::new(ArrowSchema::empty()))) + }); + Ok(Box::pin(RecordBatchStreamAdapter::new( + Arc::new(ArrowSchema::empty()), + batches, + ))) + })); + + let recorder = SpanRecorder::default(); + // A subscriber is installed on one thread only, so the tasks analyze + // spawns for the plan must poll on the thread that installs it. + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + tracing::subscriber::with_default( + tracing_subscriber::registry().with(recorder.clone()), + || { + runtime + .block_on(analyze_plan(plan, LanceExecutionOptions::default())) + .unwrap(); + }, + ); + + let spans = recorder.recorded(); + assert!( + spans + .iter() + .any(|(name, level, _)| *name == "analyze_plan" && *level == Level::DEBUG), + "analyze must open a debug-level `analyze_plan` span, recorded: {spans:?}" + ); + for opened_by_plan in ["plan_execute", "plan_poll"] { + let (.., parent) = spans + .iter() + .find(|(name, ..)| *name == opened_by_plan) + .unwrap_or_else(|| { + panic!("the plan must open `{opened_by_plan}`, recorded: {spans:?}") + }); + assert_eq!( + *parent, + Some("analyze_plan"), + "`{opened_by_plan}` must be parented to the analyze span, recorded: {spans:?}" + ); + } + } } diff --git a/rust/lance-datafusion/src/planner.rs b/rust/lance-datafusion/src/planner.rs index 5ee2bbe9a6a..a1b1491c13e 100644 --- a/rust/lance-datafusion/src/planner.rs +++ b/rust/lance-datafusion/src/planner.rs @@ -49,6 +49,7 @@ use datafusion::{ use datafusion_functions::core::getfield::GetFieldFunc; use lance_core::datatypes::Schema; use lance_core::error::LanceOptionExt; +use tracing::instrument; use chrono::Utc; use lance_core::{Error, Result}; @@ -1010,6 +1011,7 @@ impl Planner { } /// Optimize the filter expression and coerce data types. + #[instrument(level = "trace", name = "filter_optimize", skip_all)] pub fn optimize_expr(&self, expr: Expr) -> Result { let df_schema = Arc::new(DFSchema::try_from(self.schema.as_ref().clone())?); @@ -1333,6 +1335,33 @@ mod tests { } } + #[test] + fn test_optimize_expr_opens_filter_optimize_span() { + let planner = Planner::new(Arc::new(Schema::new(vec![Field::new( + "i", + DataType::Int32, + false, + )]))); + let expr = col("i").gt(lit(3_i32)); + + // Filtering the mock down to the span under test keeps unrelated + // DataFusion spans out of its ordered expectation queue. + let (subscriber, handle) = tracing_mock::subscriber::mock() + .with_filter(|meta| meta.name() == "filter_optimize") + .new_span( + tracing_mock::expect::span() + .named("filter_optimize") + .at_level(tracing::Level::TRACE), + ) + .run_with_handle(); + + tracing::subscriber::with_default(subscriber, || { + planner.optimize_expr(expr).unwrap(); + }); + + handle.assert_finished(); + } + #[test] fn test_coerce_before_simplify() { let planner = Planner::new(Arc::new(Schema::empty()));