Traces & Metrics: Scan Plan

Cargo.toml

@@ -52,6 +52,7 @@ arrow-select = { version = "55" }
 arrow-string = { version = "55" }
 as-any = "0.3.2"
 async-std = "1.12"
+async-stream = "0.3"
 async-trait = "0.1.88"
 aws-config = "1.6.1"
 aws-sdk-glue = "1.39"
@@ -82,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"

crates/iceberg/Cargo.toml

@@ -56,6 +56,7 @@ arrow-select = { workspace = true }
 arrow-string = { workspace = true }
 as-any = { workspace = true }
 async-std = { workspace = true, optional = true, features = ["attributes"] }
+async-stream = { workspace = true }
 async-trait = { workspace = true }
 base64 = { workspace = true }
 bimap = { workspace = true }
@@ -66,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 }
@@ -86,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/src/arrow/count_recording_record_batch_stream.rs

@@ -0,0 +1,68 @@
+use std::pin::Pin;
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::task::{Context, Poll};
+
+use arrow_array::RecordBatch;
+use futures::Stream;
+
+use crate::Result;
+
+pub(crate) struct CountRecordingRecordBatchStream<S> {
+    stream: S,
+    row_count: AtomicU64,
+    record_batch_count: AtomicU64,
+    target_span: tracing::Span,
+    record_batch_count_field_name: &'static str,
+    row_count_field_name: &'static str,
+}
+
+impl<S> CountRecordingRecordBatchStream<S> {
+    pub(crate) fn new(
+        stream: S,
+        target_span: tracing::Span,
+        record_batch_count_field_name: &'static str,
+        row_count_field_name: &'static str,
+    ) -> Self {
+        Self {
+            stream,
+            row_count: AtomicU64::new(0),
+            record_batch_count: AtomicU64::new(0),
+            target_span,
+            record_batch_count_field_name,
+            row_count_field_name,
+        }
+    }
+}
+
+impl<S> Stream for CountRecordingRecordBatchStream<S>
+where S: Stream<Item = Result<RecordBatch>> + Unpin
+{
+    type Item = S::Item;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let this = self.get_mut();
+
+        match Pin::new(&mut this.stream).poll_next(cx) {
+            Poll::Ready(Some(Ok(batch))) => {
+                let row_count = batch.num_rows() as u64;
+
+                this.row_count.fetch_add(row_count, Ordering::Relaxed);
+                this.record_batch_count.fetch_add(1, Ordering::Relaxed);
+
+                Poll::Ready(Some(Ok(batch)))
+            }
+            other => other,
+        }
+    }
+}
+
+impl<S> Drop for CountRecordingRecordBatchStream<S> {
+    fn drop(&mut self) {
+        let total_record_batches = self.record_batch_count.load(Ordering::Relaxed);
+        let total_rows = self.row_count.load(Ordering::Relaxed);
+        self.target_span
+            .record(self.record_batch_count_field_name, total_record_batches);
+        self.target_span
+            .record(self.row_count_field_name, total_rows);
+    }
+}

crates/iceberg/src/arrow/delete_filter.rs

@@ -104,6 +104,7 @@ impl DeleteFilter {
     }
 
     /// Builds eq delete predicate for the provided task.
+    #[tracing::instrument(skip_all, level = "trace")]
     pub(crate) async fn build_equality_delete_predicate(
         &self,
         file_scan_task: &FileScanTask,

crates/iceberg/src/arrow/mod.rs

@@ -27,6 +27,7 @@ pub(crate) mod caching_delete_file_loader;
 pub mod delete_file_loader;
 pub(crate) mod delete_filter;
 
+pub(crate) mod count_recording_record_batch_stream;
 mod reader;
 pub(crate) mod record_batch_projector;
 pub(crate) mod record_batch_transformer;