feat: initial instrumentation of scan plan with traces and metrics

Author: sdd

Cargo.lock

@@ -83,6 +83,7 @@ indicatif = "0.17"
 itertools = "0.13"
 linkedbytes = "0.1.8"
 metainfo = "0.7.14"
+metrics = "0.24"
 mimalloc = "0.1.46"
 mockito = "1"
 motore-macros = "0.4.3"

Cargo.toml

@@ -67,6 +67,7 @@ expect-test = { workspace = true }
 fnv = { workspace = true }
 futures = { workspace = true }
 itertools = { workspace = true }
+metrics = { workspace = true }
 moka = { version = "0.12.10", features = ["future"] }
 murmur3 = { workspace = true }
 num-bigint = { workspace = true }
@@ -87,6 +88,7 @@ serde_with = { workspace = true }
 strum = { workspace = true, features = ["derive"] }
 thrift = { workspace = true }
 tokio = { workspace = true, optional = false, features = ["sync"] }
+tracing = { workspace = true }
 typed-builder = { workspace = true }
 url = { workspace = true }
 uuid = { workspace = true }

crates/iceberg/Cargo.toml

@@ -43,6 +43,7 @@ impl ExpressionEvaluator {
     /// the provided [`DataFile`]'s partition [`Struct`]. Used by [`TableScan`]
     /// to see if this [`DataFile`] could possibly contain data that matches
     /// the scan's filter.
+    #[tracing::instrument(skip_all)]
     pub(crate) fn eval(&self, data_file: &DataFile) -> Result<bool> {
         let mut visitor = ExpressionEvaluatorVisitor::new(data_file.partition());
 

crates/iceberg/src/expr/visitors/expression_evaluator.rs

@@ -39,6 +39,7 @@ impl<'a> InclusiveMetricsEvaluator<'a> {
     /// provided [`DataFile`]'s metrics. Used by [`TableScan`] to
     /// see if this `DataFile` contains data that could match
     /// the scan's filter.
+    #[tracing::instrument(skip_all)]
     pub(crate) fn eval(
         filter: &'a BoundPredicate,
         data_file: &'a DataFile,

crates/iceberg/src/expr/visitors/inclusive_metrics_evaluator.rs

@@ -85,6 +85,7 @@ impl ObjectCache {
 
     /// Retrieves an Arc [`Manifest`] from the cache
     /// or retrieves one from FileIO and parses it if not present
+    #[tracing::instrument(skip_all)]
     pub(crate) async fn get_manifest(&self, manifest_file: &ManifestFile) -> Result<Arc<Manifest>> {
         if self.cache_disabled {
             return manifest_file

crates/iceberg/src/io/object_cache.rs

@@ -90,3 +90,4 @@ pub mod writer;
 
 mod delete_vector;
 pub mod puffin;
+mod traced_stream;

crates/iceberg/src/lib.rs

@@ -144,6 +144,7 @@ pub(crate) struct PlanContext {
 }
 
 impl PlanContext {
+    #[tracing::instrument(skip_all)]
     pub(crate) async fn get_manifest_list(&self) -> Result<Arc<ManifestList>> {
         self.object_cache
             .as_ref()
@@ -172,6 +173,12 @@ impl PlanContext {
         Ok(partition_filter)
     }
 
+    #[tracing::instrument(
+        skip_all,
+        fields(
+            manifest_list.len = manifest_list.entries().len(),
+        )
+    )]
     pub(crate) fn build_manifest_file_context_iter(
         &self,
         manifest_list: Arc<ManifestList>,
@@ -191,16 +198,20 @@ impl PlanContext {
                             .get(manifest_file.partition_spec_id, predicate.clone())
                             .eval(&manifest_file)?
                         {
+                            tracing::trace!(file_path = manifest_file.manifest_path, "iceberg.scan.manifest_file.skipped");
+                            metrics::counter!("iceberg.scan.manifest_file.skipped", "reason" => "partition").increment(1);
                             return Ok(None); // Skip this file.
                         }
                         Some(predicate)
                     } else {
                         None
                     };
 
+                    tracing::trace!(file_path = manifest_file.manifest_path, "iceberg.scan.manifest_file.included");
+                    metrics::counter!("iceberg.scan.manifest_file.included").increment(1);
+
                     let context = self
                         .create_manifest_file_context(manifest_file, partition_bound_predicate)?;
-
                     Ok(Some(context))
                 })()
                 .transpose()

crates/iceberg/src/scan/context.rs

@@ -37,6 +37,7 @@ use crate::expr::{Bind, BoundPredicate, Predicate};
 use crate::io::FileIO;
 use crate::spec::{DataContentType, SnapshotRef};
 use crate::table::Table;
+use crate::traced_stream::TracedStream;
 use crate::utils::available_parallelism;
 use crate::{Error, ErrorKind, Result};
 
@@ -330,7 +331,11 @@ pub struct TableScan {
 impl TableScan {
     /// Returns a stream of [`FileScanTask`]s.
     pub async fn plan_files(&self) -> Result<FileScanTaskStream> {
+        let span = tracing::trace_span!("plan_files");
+        let _entered = span.enter();
+
         let Some(plan_context) = self.plan_context.as_ref() else {
+            tracing::debug!("file plan requested for a table with no snapshots");
             return Ok(Box::pin(futures::stream::empty()));
         };
 
@@ -351,7 +356,7 @@ impl TableScan {
 
         let delete_file_index = Arc::new(delete_file_index);
 
-        Ok(TableScan::process_manifest_contexts(
+        let stream = TableScan::process_manifest_contexts(
             data_contexts,
             self.concurrency_limit_manifest_files,
             self.concurrency_limit_manifest_entries,
@@ -360,7 +365,9 @@ impl TableScan {
                 async move { Self::process_data_manifest_entry(ctx, delete_file_index) }
             },
         )
-        .boxed())
+        .boxed();
+
+        Ok(Box::pin(TracedStream::new(stream, span.clone())))
     }
 
     /// Returns an [`ArrowRecordBatchStream`].
@@ -419,19 +426,27 @@ impl TableScan {
             .try_filter_map(|opt_task| async move { Ok(opt_task) })
     }
 
+    #[tracing::instrument(skip_all, fields(file_path))]
     fn process_data_manifest_entry(
         manifest_entry_context: Result<ManifestEntryContext>,
         delete_file_index: Arc<DeleteFileIndex>,
     ) -> Result<Option<FileScanTask>> {
         let manifest_entry_context = manifest_entry_context?;
+        tracing::Span::current().record(
+            "file_path",
+            manifest_entry_context.manifest_entry.file_path(),
+        );
 
         // skip processing this manifest entry if it has been marked as deleted
         if !manifest_entry_context.manifest_entry.is_alive() {
+            metrics::counter!("iceberg.scan.data_file.skipped", "reason" => "not_alive")
+                .increment(1);
             return Ok(None);
         }
 
         // abort the plan if we encounter a manifest entry for a delete file
         if manifest_entry_context.manifest_entry.content_type() != DataContentType::Data {
+            tracing::error!("Encountered an entry for a delete file in a data file manifest");
             return Err(Error::new(
                 ErrorKind::FeatureUnsupported,
                 "Encountered an entry for a delete file in a data file manifest",
@@ -455,6 +470,8 @@ impl TableScan {
             // skip any data file whose partition data indicates that it can't contain
             // any data that matches this scan's filter
             if !expression_evaluator.eval(manifest_entry_context.manifest_entry.data_file())? {
+                metrics::counter!("iceberg.scan.data_file.skipped", "reason" => "partition")
+                    .increment(1);
                 return Ok(None);
             }
 
@@ -464,30 +481,41 @@ impl TableScan {
                 manifest_entry_context.manifest_entry.data_file(),
                 false,
             )? {
+                metrics::counter!("iceberg.scan.data_file.skipped", "reason" => "file_metrics")
+                    .increment(1);
                 return Ok(None);
             }
         }
 
         // congratulations! the manifest entry has made its way through the
         // entire plan without getting filtered out. Create a corresponding
         // FileScanTask and push it to the result stream
+        metrics::counter!("iceberg.scan.data_file.included").increment(1);
         Ok(Some(
             manifest_entry_context.into_file_scan_task(delete_file_index)?,
         ))
     }
 
+    #[tracing::instrument(skip_all, fields(file_path))]
     fn process_delete_manifest_entry(
         manifest_entry_context: Result<ManifestEntryContext>,
     ) -> Result<Option<DeleteFileContext>> {
         let manifest_entry_context = manifest_entry_context?;
+        tracing::Span::current().record(
+            "file_path",
+            manifest_entry_context.manifest_entry.file_path(),
+        );
 
         // skip processing this manifest entry if it has been marked as deleted
         if !manifest_entry_context.manifest_entry.is_alive() {
+            metrics::counter!("iceberg.scan.delete_file.skipped", "reason" => "not_alive")
+                .increment(1);
             return Ok(None);
         }
 
         // abort the plan if we encounter a manifest entry that is not for a delete file
         if manifest_entry_context.manifest_entry.content_type() == DataContentType::Data {
+            tracing::error!("Encountered an entry for a data file in a delete manifest");
             return Err(Error::new(
                 ErrorKind::FeatureUnsupported,
                 "Encountered an entry for a data file in a delete manifest",
@@ -506,10 +534,13 @@ impl TableScan {
             // skip any data file whose partition data indicates that it can't contain
             // any data that matches this scan's filter
             if !expression_evaluator.eval(manifest_entry_context.manifest_entry.data_file())? {
+                metrics::counter!("iceberg.scan.delete_file.skipped", "reason" => "partition")
+                    .increment(1);
                 return Ok(None);
             }
         }
 
+        metrics::counter!("iceberg.scan.delete_file.included").increment(1);
         Ok(Some(DeleteFileContext {
             manifest_entry: manifest_entry_context.manifest_entry.clone(),
             partition_spec_id: manifest_entry_context.partition_spec_id,

crates/iceberg/src/scan/mod.rs

@@ -0,0 +1,31 @@
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+use futures::Stream;
+use tracing::Span;
+
+pub struct TracedStream<S> {
+    stream: S,
+    _span: Span,
+}
+
+impl<S> TracedStream<S> {
+    pub fn new(stream: S, span: Span) -> Self {
+        Self {
+            stream,
+            _span: span,
+        }
+    }
+}
+
+impl<S> Stream for TracedStream<S>
+where S: Stream + Unpin
+{
+    type Item = S::Item;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let this = self.get_mut();
+        let _entered = this._span.enter();
+        Pin::new(&mut this.stream).poll_next(cx)
+    }
+}