chore(cube): Cube tracing instrumentation

Author: srh

Cargo.toml

@@ -156,6 +156,8 @@ sqlparser = { git = "https://github.com/cube-js/sqlparser-rs.git", branch = "cub
 tempfile = "3"
 tokio = { version = "1.43", features = ["macros", "rt", "sync"] }
 url = "2.5.4"
+tracing = "0.1.25"
+tracing-futures = { version = "0.2.5", features = ["futures-03"] }
 
 [profile.release]
 codegen-units = 1

datafusion/core/Cargo.toml

@@ -133,6 +133,8 @@ url = { workspace = true }
 uuid = { version = "1.15", features = ["v4", "js"] }
 xz2 = { version = "0.1", optional = true, features = ["static"] }
 zstd = { version = "0.13", optional = true, default-features = false }
+tracing = { workspace = true }
+tracing-futures = { workspace = true }
 
 [dev-dependencies]
 async-trait = { workspace = true }

datafusion/core/src/datasource/physical_plan/parquet/mod.rs

@@ -65,6 +65,7 @@ pub use row_filter::build_row_filter;
 pub use row_filter::can_expr_be_pushed_down_with_schemas;
 pub use row_group_filter::RowGroupAccessPlanFilter;
 pub use writer::plan_to_parquet;
+use tracing_futures::Instrument;
 
 use log::debug;
 
@@ -525,6 +526,8 @@ impl ExecutionPlan for ParquetExec {
         // TODO upgrade DF: we need  reader_options_customizer used as in the commented code below.
         self.inner.execute(partition_index, ctx)
 
+        // TODO upgrade DF:  We also want Ok(Box::pin(stream.instrument(tracing::trace_span!("read_files")))) applied, probably in inner.
+
         // let reader_options_customizer =
         //     get_reader_options_customizer(ctx.session_config());
 

datafusion/execution/Cargo.toml

@@ -49,6 +49,8 @@ parking_lot = { workspace = true }
 rand = { workspace = true }
 tempfile = { workspace = true }
 url = { workspace = true }
+tracing = { workspace = true }
+tracing-futures = { workspace = true }
 
 [dev-dependencies]
 chrono = { workspace = true }

datafusion/execution/src/stream.rs

@@ -58,3 +58,12 @@ pub trait RecordBatchStream: Stream<Item = Result<RecordBatch>> {
 /// [`Stream`]s there is no mechanism to prevent callers polling  so returning
 /// `Ready(None)` is recommended.
 pub type SendableRecordBatchStream = Pin<Box<dyn RecordBatchStream + Send>>;
+
+// Cube extension
+impl<T> RecordBatchStream for tracing_futures::Instrumented<T>
+    where T: Stream<Item = Result<RecordBatch>>,
+        T: RecordBatchStream {
+    fn schema(&self) -> SchemaRef {
+        self.inner().schema()
+    }
+}

datafusion/physical-plan/Cargo.toml

@@ -64,8 +64,8 @@ parking_lot = { workspace = true }
 pin-project-lite = "^0.2.7"
 tokio = { workspace = true }
 serde = { version = "1.0.214", features = ["derive"] }
-tracing = "0.1.25"
-tracing-futures = { version = "0.2.5" }
+tracing = { workspace = true }
+tracing-futures = { workspace = true }
 
 [dev-dependencies]
 criterion = { workspace = true, features = ["async_futures"] }

datafusion/physical-plan/src/aggregates/mod.rs

@@ -49,6 +49,7 @@ use datafusion_physical_expr::{
 };
 
 use itertools::Itertools;
+use tracing_futures::Instrument;
 
 pub(crate) mod group_values;
 mod no_grouping;
@@ -338,9 +339,9 @@ enum StreamType {
 impl From<StreamType> for SendableRecordBatchStream {
     fn from(stream: StreamType) -> Self {
         match stream {
-            StreamType::AggregateStream(stream) => Box::pin(stream),
-            StreamType::GroupedHash(stream) => Box::pin(stream),
-            StreamType::GroupedPriorityQueue(stream) => Box::pin(stream),
+            StreamType::AggregateStream(stream) => Box::pin(stream.instrument(tracing::trace_span!("AggregateStream"))),
+            StreamType::GroupedHash(stream) => Box::pin(stream.instrument(tracing::trace_span!("GroupedHashAggregateStream"))),
+            StreamType::GroupedPriorityQueue(stream) => Box::pin(stream.instrument(tracing::trace_span!("GroupedTopKAggregateStream"))),
         }
     }
 }

datafusion/physical-plan/src/execution_plan.rs

@@ -35,6 +35,7 @@ use std::fmt::Debug;
 use std::sync::Arc;
 
 use crate::coalesce_partitions::CoalescePartitionsExec;
+use crate::cube_ext;
 use crate::display::DisplayableExecutionPlan;
 use crate::metrics::MetricsSet;
 use crate::projection::ProjectionExec;
@@ -884,7 +885,7 @@ pub async fn collect_partitioned(
     let mut join_set = JoinSet::new();
     // Execute the plan and collect the results into batches.
     streams.into_iter().enumerate().for_each(|(idx, stream)| {
-        join_set.spawn(async move {
+        cube_ext::spawn_on_joinset(&mut join_set, async move {
             let result: Result<Vec<RecordBatch>> = stream.try_collect().await;
             (idx, result)
         });

datafusion/physical-plan/src/filter.rs

@@ -500,6 +500,7 @@ struct FilterExecStream {
     projection: Option<Vec<usize>>,
 }
 
+#[tracing::instrument(level = "trace", skip(batch))]
 pub fn batch_filter(
     batch: &RecordBatch,
     predicate: &Arc<dyn PhysicalExpr>,

datafusion/physical-plan/src/limit.rs

@@ -428,6 +428,8 @@ impl LimitStream {
             self.fetch = 0;
             self.input = None; // Clear input so it can be dropped early
 
+            let span = tracing::trace_span!("truncate_batch");  // TODO upgrade DF: We could use the baseline_metrics values.
+            let _guard = span.enter();
             // It is guaranteed that batch_rows is <= batch.num_rows
             Some(batch.slice(0, batch_rows))
         } else {