optional traces for vortex-layout (#2570)

Author: onursatici

Cargo.lock

@@ -51,7 +51,7 @@ vortex-layout = { workspace = true, features = ["tokio"] }
 vortex-metrics = { workspace = true }
 
 [features]
-tracing = ["dep:tracing", "dep:tracing-futures", "vortex-io/tracing"]
+tracing = ["dep:tracing", "dep:tracing-futures", "vortex-io/tracing", "vortex-layout/tracing", "vortex-file/tracing"]
 
 [dev-dependencies]
 anyhow = { workspace = true }

vortex-datafusion/Cargo.toml

@@ -27,6 +27,7 @@ oneshot = { workspace = true, features = ["async"] }
 rustc-hash = { workspace = true }
 tokio = { workspace = true, features = ["rt"], optional = true }
 tracing = { workspace = true, optional = true }
+tracing-futures = { workspace = true, features = ["futures-03"], optional = true }
 # Needed to pickup the "js" feature for wasm targets from the workspace configuration
 uuid = { workspace = true }
 vortex-alp = { workspace = true }
@@ -60,4 +61,4 @@ workspace = true
 
 [features]
 object_store = ["vortex-error/object_store", "vortex-io/object_store"]
-tracing = ["dep:tracing", "vortex-io/tracing"]
+tracing = ["dep:tracing", "dep:tracing-futures", "vortex-io/tracing", "vortex-layout/tracing"]

vortex-file/Cargo.toml

@@ -9,6 +9,7 @@ use moka::future::CacheBuilder;
 use vortex_buffer::{Alignment, ByteBuffer};
 use vortex_error::{VortexExpect, VortexResult, vortex_err, vortex_panic};
 use vortex_io::VortexReadAt;
+use vortex_layout::instrument;
 use vortex_layout::scan::ScanDriver;
 use vortex_layout::segments::{AsyncSegmentReader, SegmentId};
 use vortex_metrics::{Counter, VortexMetrics};
@@ -186,7 +187,10 @@ impl<R: VortexReadAt> ScanDriver for GenericScanDriver<R> {
         });
 
         // Buffer some number of concurrent I/O operations.
-        io_stream.buffer_unordered(self.options.io_concurrency)
+        instrument!(
+            "io_stream",
+            io_stream.buffer_unordered(self.options.io_concurrency)
+        )
     }
 }
 

vortex-file/src/generic.rs

@@ -28,6 +28,8 @@ itertools = { workspace = true }
 log = { workspace = true }
 pin-project-lite = { workspace = true }
 tokio = { workspace = true, features = ["sync"], optional = true }
+tracing = { workspace = true, optional = true }
+tracing-futures = { workspace = true, features = ["futures-03"], optional = true }
 vortex-array = { workspace = true }
 vortex-buffer = { workspace = true }
 vortex-dtype = { workspace = true }
@@ -43,6 +45,7 @@ rstest = { workspace = true }
 
 [features]
 tokio = ["dep:tokio", "vortex-error/tokio"]
+tracing = ["dep:tracing", "dep:tracing-futures"]
 
 [lints]
 workspace = true

vortex-layout/Cargo.toml

@@ -8,7 +8,7 @@ use vortex_error::{VortexExpect, VortexResult, vortex_err, vortex_panic};
 use crate::layouts::flat::FlatLayout;
 use crate::reader::LayoutReader;
 use crate::segments::AsyncSegmentReader;
-use crate::{Layout, LayoutReaderExt, LayoutVTable};
+use crate::{Layout, LayoutReaderExt, LayoutVTable, instrument};
 
 pub struct FlatReader {
     layout: Layout,
@@ -43,25 +43,33 @@ impl FlatReader {
 
     pub(crate) async fn array(&self) -> VortexResult<&ArrayRef> {
         self.array
-            .get_or_try_init(async move {
-                let segment_id = self
-                    .layout()
-                    .segment_id(0)
-                    .ok_or_else(|| vortex_err!("FlatLayout missing segment"))?;
+            .get_or_try_init(instrument!(
+                "flat_read",
+                { name = self.layout().name() },
+                async move {
+                    let segment_id = self
+                        .layout()
+                        .segment_id(0)
+                        .ok_or_else(|| vortex_err!("FlatLayout missing segment"))?;
 
-                log::debug!(
-                    "Requesting segment {} for flat layout {} expr",
-                    segment_id,
-                    self.layout().name(),
-                );
+                    log::debug!(
+                        "Requesting segment {} for flat layout {} expr",
+                        segment_id,
+                        self.layout().name(),
+                    );
 
-                // Fetch all the array segment.
-                let buffer = self.segment_reader.get(segment_id).await?;
-                let row_count = usize::try_from(self.layout().row_count())
-                    .vortex_expect("FlatLayout row count does not fit within usize");
+                    // Fetch all the array segment.
+                    let buffer = self.segment_reader.get(segment_id).await?;
+                    let row_count = usize::try_from(self.layout().row_count())
+                        .vortex_expect("FlatLayout row count does not fit within usize");
 
-                ArrayParts::try_from(buffer)?.decode(self.ctx(), self.dtype().clone(), row_count)
-            })
+                    ArrayParts::try_from(buffer)?.decode(
+                        self.ctx(),
+                        self.dtype().clone(),
+                        row_count,
+                    )
+                }
+            ))
             .await
     }
 }

vortex-layout/src/layouts/flat/reader.rs

@@ -15,6 +15,7 @@ mod row_mask;
 pub use row_mask::*;
 use vortex_array::arcref::ArcRef;
 
+pub mod macros;
 mod reader;
 pub mod scan;
 pub mod segments;

vortex-layout/src/lib.rs

@@ -0,0 +1,24 @@
+/// A macro for optionally instrumenting a future, if tracing feature is enabled.
+#[macro_export]
+macro_rules! instrument {
+    ($span_name:expr, $expr:expr) => {
+        instrument!($span_name, {}, $expr)
+    };
+    ($span_name:expr, { $($key:ident = $value:expr),* $(,)? }, $expr:expr) => {
+        {
+            let task = $expr;
+            #[cfg(feature = "tracing")]
+            {
+                use tracing_futures::Instrument;
+                task.instrument(tracing::info_span!(
+                    $span_name,
+                    $($key = $value,)*
+                ))
+            }
+            #[cfg(not(feature = "tracing"))]
+            {
+                task
+            }
+        }
+    };
+}

vortex-layout/src/macros.rs

@@ -18,7 +18,7 @@ use vortex_mask::Mask;
 use crate::scan::filter::FilterExpr;
 use crate::scan::unified::UnifiedDriverStream;
 use crate::segments::AsyncSegmentReader;
-use crate::{ExprEvaluator, Layout, LayoutReader, LayoutReaderExt, RowMask};
+use crate::{ExprEvaluator, Layout, LayoutReader, LayoutReaderExt, RowMask, instrument};
 
 pub mod executor;
 pub(crate) mod filter;
@@ -261,7 +261,7 @@ impl<D: ScanDriver> Scan<D> {
                 let reader = reader.clone();
 
                 // This future is the processing task
-                async move {
+                instrument!("process", async move {
                     let row_mask = match pruning {
                         None => row_mask,
                         Some(pruning_filter) => {
@@ -284,12 +284,13 @@ impl<D: ScanDriver> Scan<D> {
                         }
                         VortexResult::Ok(Some(array))
                     }
-                }
+                })
             })
             .map(move |processing_task| task_executor.spawn(processing_task))
             .buffered(self.concurrency)
             .filter_map(|v| async move { v.unnest().transpose() });
 
+        let exec_stream = instrument!("exec_stream", exec_stream);
         let io_stream = self.driver.io_stream();
 
         let unified = UnifiedDriverStream {