chore(cube): Cube tracing instrumentation

Author: srh

Cargo.toml

@@ -141,6 +141,8 @@ sqlparser = { git = "https://github.com/cube-js/sqlparser-rs.git", branch = "cub
 tempfile = "3"
 thiserror = "1.0.44"
 tokio = { version = "1.36", features = ["macros", "rt", "sync"] }
+tracing = "0.1.25"
+tracing-futures = { version = "0.2.5", features = ["futures-03"] }
 url = "2.2"
 
 [profile.release]

datafusion/core/Cargo.toml

@@ -137,6 +137,8 @@ url = { workspace = true }
 uuid = { version = "1.7", features = ["v4"] }
 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]
 arrow-buffer = { workspace = true }

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

@@ -67,6 +67,7 @@ pub use metrics::ParquetFileMetrics;
 use opener::ParquetOpener;
 pub use reader::{DefaultParquetFileReaderFactory, ParquetFileReaderFactory};
 pub use writer::plan_to_parquet;
+use tracing_futures::Instrument;
 
 /// Execution plan for reading one or more Parquet files.
 ///
@@ -740,7 +741,7 @@ impl ExecutionPlan for ParquetExec {
         let stream =
             FileStream::new(&self.base_config, partition_index, opener, &self.metrics)?;
 
-        Ok(Box::pin(stream))
+        Ok(Box::pin(stream.instrument(tracing::trace_span!("read_files"))))
     }
 
     fn metrics(&self) -> Option<MetricsSet> {

datafusion/execution/Cargo.toml

@@ -49,3 +49,7 @@ parking_lot = { workspace = true }
 rand = { workspace = true }
 tempfile = { workspace = true }
 url = { workspace = true }
+tracing = { workspace = true }
+tracing-futures = { workspace = true }
+
+

datafusion/execution/src/stream.rs

@@ -31,3 +31,12 @@ pub trait RecordBatchStream: Stream<Item = Result<RecordBatch>> {
 
 /// Trait for a [`Stream`] of [`RecordBatch`]es
 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

@@ -67,8 +67,8 @@ pin-project-lite = "^0.2.7"
 rand = { workspace = true }
 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]
 rstest = { workspace = true }

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

@@ -49,6 +49,7 @@ use datafusion_physical_expr::{
 
 use datafusion_physical_expr::aggregate::AggregateFunctionExpr;
 use itertools::Itertools;
+use tracing_futures::Instrument;
 
 pub mod group_values;
 mod no_grouping;
@@ -248,9 +249,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

@@ -40,6 +40,7 @@ use datafusion_physical_expr::{EquivalenceProperties, LexOrdering, PhysicalSortE
 use datafusion_physical_expr_common::sort_expr::LexRequirement;
 
 use crate::coalesce_partitions::CoalescePartitionsExec;
+use crate::cube_ext;
 use crate::display::DisplayableExecutionPlan;
 pub use crate::display::{DefaultDisplay, DisplayAs, DisplayFormatType, VerboseDisplay};
 pub use crate::metrics::Metric;
@@ -738,7 +739,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

@@ -426,6 +426,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

@@ -419,6 +419,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 {