add spans for perfetto

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

vortex-compressor/src/compressor.rs

@@ -314,26 +314,23 @@ impl CascadingCompressor {
         // Run the winning scheme's `compress`. On failure, emit an ERROR event carrying the
         // scheme name and cascade history before propagating.
         let error_ctx = trace::enabled_error_context(&compress_ctx);
-        let compressed = match winner.compress(self, &data, compress_ctx, exec_ctx) {
-            Ok(compressed) => compressed,
-            Err(err) => {
+        let _winner_span = trace::winner_compress_span(winner.id(), before_nbytes).entered();
+        let compressed = winner
+            .compress(self, &data, compress_ctx, exec_ctx)
+            .inspect_err(|err| {
                 // NB: this is the only way we can tell which scheme panicked / bailed on their
                 // data, especially for third-party schemes where the error site may not carry any
                 // compressor context.
-                trace::scheme_compress_failed(winner.id(), before_nbytes, error_ctx.as_ref(), &err);
-                return Err(err);
-            }
-        };
+                trace::scheme_compress_failed(winner.id(), before_nbytes, error_ctx.as_ref(), err);
+            })?;
 
         let after_nbytes = compressed.nbytes();
         let actual_ratio = (after_nbytes != 0).then(|| before_nbytes as f64 / after_nbytes as f64);
 
         // TODO(connor): HACK TO SUPPORT L2 DENORMALIZATION!!!
         let accepted = after_nbytes < before_nbytes || compressed.is::<AnyScalarFn>();
 
-        trace::scheme_compress_result(
-            winner.id(),
-            before_nbytes,
+        trace::record_winner_compress_result(
             after_nbytes,
             winner_estimate.trace_ratio(),
             actual_ratio,
@@ -373,28 +370,32 @@ impl CascadingCompressor {
         let mut deferred: Vec<(&'static dyn Scheme, DeferredEstimate)> = Vec::new();
 
         // Pass 1: evaluate every immediate verdict. Stash deferred work for pass 2.
-        for &scheme in schemes {
-            match scheme.expected_compression_ratio(data, compress_ctx.clone(), exec_ctx) {
-                CompressionEstimate::Verdict(EstimateVerdict::Skip) => {}
-                CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse) => {
-                    return Ok(Some((scheme, WinnerEstimate::AlwaysUse)));
-                }
-                CompressionEstimate::Verdict(EstimateVerdict::Ratio(ratio)) => {
-                    let score = EstimateScore::FiniteCompression(ratio);
+        {
+            let _verdict_pass = trace::verdict_pass_span().entered();
+            for &scheme in schemes {
+                match scheme.expected_compression_ratio(data, compress_ctx.clone(), exec_ctx) {
+                    CompressionEstimate::Verdict(EstimateVerdict::Skip) => {}
+                    CompressionEstimate::Verdict(EstimateVerdict::AlwaysUse) => {
+                        return Ok(Some((scheme, WinnerEstimate::AlwaysUse)));
+                    }
+                    CompressionEstimate::Verdict(EstimateVerdict::Ratio(ratio)) => {
+                        let score = EstimateScore::FiniteCompression(ratio);
 
-                    if is_better_score(score, &best) {
-                        best = Some((scheme, score));
+                        if is_better_score(score, &best) {
+                            best = Some((scheme, score));
+                        }
+                    }
+                    CompressionEstimate::Deferred(deferred_estimate) => {
+                        deferred.push((scheme, deferred_estimate));
                     }
-                }
-                CompressionEstimate::Deferred(deferred_estimate) => {
-                    deferred.push((scheme, deferred_estimate));
                 }
             }
         }
 
         // Pass 2: run deferred work. Callbacks receive the current best as a threshold so they can
         // short-circuit with `Skip` when they cannot beat it.
         for (scheme, deferred_estimate) in deferred {
+            let _span = trace::scheme_eval_span(scheme.id()).entered();
             let threshold: Option<EstimateScore> = best.map(|(_, score)| score);
             match deferred_estimate {
                 DeferredEstimate::Sample => {

vortex-compressor/src/estimate.rs

@@ -228,9 +228,9 @@ pub(super) fn estimate_compression_ratio_with_sampling<S: Scheme + ?Sized>(
 
     let score = EstimateScore::from_sample_sizes(before, after);
 
-    // Single DEBUG event per sampled scheme. Downstream tooling can join this with the eventual
-    // `scheme.compress_result` on the same scheme to compute sample-vs-full divergence.
-    trace::sample_result(scheme.id(), before, after, score.finite_ratio());
+    // Record the sample outcome as fields on the enclosing `scheme_eval` span. Joins with the
+    // eventual `scheme.compress_result` on the same scheme to compute sample-vs-full divergence.
+    trace::record_sample_result(before, after, score.finite_ratio());
 
     Ok(score)
 }

vortex-compressor/src/stats/cache.rs

@@ -21,6 +21,7 @@ use super::FloatStats;
 use super::GenerateStatsOptions;
 use super::IntegerStats;
 use super::StringStats;
+use crate::trace;
 
 /// A single cache entry: a concrete [`TypeId`] paired with a type-erased value.
 type StatsEntry = (TypeId, Arc<dyn Any + Send + Sync>);
@@ -58,7 +59,10 @@ impl StatsCache {
                 .ok()
                 .vortex_expect("we just checked the TypeID")
         } else {
-            let new_arc: Arc<T> = Arc::new(f());
+            let new_arc: Arc<T> = {
+                let _span = trace::generate_stats_span(std::any::type_name::<T>()).entered();
+                Arc::new(f())
+            };
             guard.push((type_id, Arc::clone(&new_arc) as Arc<dyn Any + Send + Sync>));
             new_arc
         }

vortex-compressor/src/trace.rs

@@ -12,6 +12,9 @@ use crate::scheme::SchemeId;
 pub(super) const TARGET_TRACE: &str = "vortex_compressor::encode";
 
 /// Builds the top-level compression span.
+///
+/// `input_nbytes` is known up front; `compressed_nbytes` / `compression_ratio` are filled in by
+/// [`record_compress_outcome`] once the cascade returns.
 #[inline]
 pub(super) fn compress_span(
     len: usize,
@@ -21,20 +24,126 @@ pub(super) fn compress_span(
     tracing::debug_span!(
         target: TARGET_TRACE,
         "compress",
-        len,
+        array_len = len,
         dtype = %dtype,
-        before_nbytes,
-        after_nbytes = tracing::field::Empty,
-        ratio = tracing::field::Empty,
+        input_nbytes = before_nbytes,
+        compressed_nbytes = tracing::field::Empty,
+        compression_ratio = tracing::field::Empty,
     )
 }
 
+/// Builds a span covering on-demand materialization of a cached stats type.
+///
+/// Child of whatever span is active when a stats accessor first fires. Typically that's
+/// [`verdict_pass_span`]; entering this span disambiguates stats cost from the rest of Pass 1.
+/// `kind` is usually `std::any::type_name::<T>()` so the args identify which group was generated
+/// (e.g. `IntegerStats`, `FloatStats`).
+#[inline]
+pub(super) fn generate_stats_span(kind: &'static str) -> tracing::Span {
+    tracing::debug_span!(
+        target: TARGET_TRACE,
+        "generate_stats",
+        stats_kind = kind,
+    )
+}
+
+/// Builds a span covering Pass 1 of scheme selection (the cheap-verdict pass).
+///
+/// Stats batches merged across eligible schemes are materialized lazily by the first
+/// `expected_compression_ratio` call that touches them. Grouping those calls under one span makes
+/// the stats cost (and unexpectedly slow verdicts) visible independently of per-candidate sampling.
+#[inline]
+pub(super) fn verdict_pass_span() -> tracing::Span {
+    tracing::debug_span!(
+        target: TARGET_TRACE,
+        "verdict_pass",
+    )
+}
+
+/// Builds a span covering one deferred per-scheme evaluation (sample or callback).
+///
+/// `scheme_candidate` is the scheme being evaluated (not necessarily chosen). The
+/// `sample_*_nbytes` / `sample_compression_ratio` fields are filled in by
+/// [`record_sample_result`] at the end of a sample compression; callback-based evaluations leave
+/// them unset.
+#[inline]
+pub(super) fn scheme_eval_span(scheme: SchemeId) -> tracing::Span {
+    tracing::debug_span!(
+        target: TARGET_TRACE,
+        "scheme_eval",
+        scheme_candidate = %scheme,
+        sample_input_nbytes = tracing::field::Empty,
+        sample_compressed_nbytes = tracing::field::Empty,
+        sample_compression_ratio = tracing::field::Empty,
+    )
+}
+
+/// Records the outcome of a sample compression on the current `scheme_eval` span.
+#[inline]
+pub(super) fn record_sample_result(
+    sample_input_nbytes: u64,
+    sample_compressed_nbytes: u64,
+    sample_compression_ratio: Option<f64>,
+) {
+    let span = tracing::Span::current();
+    span.record("sample_input_nbytes", sample_input_nbytes);
+    span.record("sample_compressed_nbytes", sample_compressed_nbytes);
+    if let Some(ratio) = sample_compression_ratio {
+        span.record("sample_compression_ratio", ratio);
+    }
+}
+
+/// Builds a span covering the winning scheme's full-array compression.
+///
+/// `scheme_chosen` and `input_nbytes` are known up front. `compressed_nbytes`,
+/// `estimated_ratio`, `achieved_ratio`, and `accepted` are filled in by
+/// [`record_winner_compress_result`] once the encode completes.
+#[inline]
+pub(super) fn winner_compress_span(scheme: SchemeId, before_nbytes: u64) -> tracing::Span {
+    tracing::debug_span!(
+        target: TARGET_TRACE,
+        "winner_compress",
+        scheme_chosen = %scheme,
+        input_nbytes = before_nbytes,
+        compressed_nbytes = tracing::field::Empty,
+        estimated_ratio = tracing::field::Empty,
+        achieved_ratio = tracing::field::Empty,
+        accepted = tracing::field::Empty,
+    )
+}
+
+/// Records the outcome of a winning-scheme compression on the current `winner_compress` span.
+#[inline]
+pub(super) fn record_winner_compress_result(
+    compressed_nbytes: u64,
+    estimated_ratio: Option<f64>,
+    achieved_ratio: Option<f64>,
+    accepted: bool,
+) {
+    let span = tracing::Span::current();
+    span.record("compressed_nbytes", compressed_nbytes);
+    if let Some(r) = estimated_ratio {
+        span.record("estimated_ratio", r);
+    }
+    if let Some(r) = achieved_ratio {
+        span.record("achieved_ratio", r);
+    }
+    span.record("accepted", accepted);
+}
+
 /// Records the final output size and, when finite, the top-level compression ratio.
 #[inline]
-pub(super) fn record_compress_outcome(span: &tracing::Span, before_nbytes: u64, after_nbytes: u64) {
-    span.record("after_nbytes", after_nbytes);
-    if after_nbytes != 0 {
-        span.record("ratio", before_nbytes as f64 / after_nbytes as f64);
+pub(super) fn record_compress_outcome(
+    span: &tracing::Span,
+    input_nbytes: u64,
+    compressed_nbytes: u64,
+) {
+    span.record("compressed_nbytes", compressed_nbytes);
+    if compressed_nbytes != 0 {
+        span.record(
+            "compression_ratio",
+            input_nbytes as f64 / compressed_nbytes as f64,
+        );
     }
 }
 
@@ -76,68 +185,6 @@ pub(super) fn scheme_compress_failed(
     }
 }
 
-/// Emits the leaf compression result event.
-#[inline]
-#[allow(
-    clippy::cognitive_complexity,
-    reason = "tracing sometimes triggers this"
-)]
-pub(super) fn scheme_compress_result(
-    scheme: SchemeId,
-    before_nbytes: u64,
-    after_nbytes: u64,
-    estimated_ratio: Option<f64>,
-    actual_ratio: Option<f64>,
-    accepted: bool,
-) {
-    match (estimated_ratio, actual_ratio) {
-        (Some(estimated_ratio), Some(actual_ratio)) => {
-            tracing::debug!(
-                target: TARGET_TRACE,
-                scheme = %scheme,
-                before_nbytes,
-                after_nbytes,
-                estimated_ratio,
-                actual_ratio,
-                accepted,
-                "scheme.compress_result",
-            );
-        }
-        (Some(estimated_ratio), None) => {
-            tracing::debug!(
-                target: TARGET_TRACE,
-                scheme = %scheme,
-                before_nbytes,
-                after_nbytes,
-                estimated_ratio,
-                accepted,
-                "scheme.compress_result",
-            );
-        }
-        (None, Some(actual_ratio)) => {
-            tracing::debug!(
-                target: TARGET_TRACE,
-                scheme = %scheme,
-                before_nbytes,
-                after_nbytes,
-                actual_ratio,
-                accepted,
-                "scheme.compress_result",
-            );
-        }
-        (None, None) => {
-            tracing::debug!(
-                target: TARGET_TRACE,
-                scheme = %scheme,
-                before_nbytes,
-                after_nbytes,
-                accepted,
-                "scheme.compress_result",
-            );
-        }
-    }
-}
-
 /// Emits a sampling-failure event.
 #[inline]
 pub(super) fn sample_compress_failed(
@@ -156,31 +203,3 @@ pub(super) fn sample_compress_failed(
         );
     }
 }
-
-/// Emits the sampling result event.
-#[inline]
-pub(super) fn sample_result(
-    scheme: SchemeId,
-    sampled_before: u64,
-    sampled_after: u64,
-    sampled_ratio: Option<f64>,
-) {
-    if let Some(sampled_ratio) = sampled_ratio {
-        tracing::debug!(
-            target: TARGET_TRACE,
-            scheme = %scheme,
-            sampled_before,
-            sampled_after,
-            sampled_ratio,
-            "sample.result",
-        );
-    } else {
-        tracing::debug!(
-            target: TARGET_TRACE,
-            scheme = %scheme,
-            sampled_before,
-            sampled_after,
-            "sample.result",
-        );
-    }
-}