segment metrics (#2713)

onursatici · web-flow · commit a0da90ba0e61 · 2025-03-17T14:06:53.000Z
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/vortex-file/src/file.rs b/vortex-file/src/file.rs
@@ -49,5 +49,6 @@ impl<F: FileType> VortexFile<F> {
             self.footer.layout().clone(),
             self.footer().ctx().clone(),
         )
+        .with_metrics(self.metrics.clone())
     }
 }
diff --git a/vortex-file/src/generic.rs b/vortex-file/src/generic.rs
@@ -231,10 +231,18 @@ impl<R: VortexReadAt> ScanDriver for GenericScanDriver<R> {
             prefetch: prefetch_stream,
             requests_ready_chunks: 1024,
             prefetch_ready_chunks: self.options.io_concurrency,
+            requested_segments: self.metrics.requested_segments.clone(),
+            prefetched_segments: self.metrics.prefetched_segments.clone(),
         }
-        .map(move |r| coalesce(r, perf_hint.coalescing_window(), perf_hint.max_read()))
-        .flat_map(stream::iter)
-        .inspect(move |(coalesced, _)| self.metrics.record(coalesced));
+        .map(move |r| {
+            coalesce(
+                r,
+                perf_hint.coalescing_window(),
+                perf_hint.max_read(),
+                self.metrics.clone(),
+            )
+        })
+        .flat_map(stream::iter);
 
         // Submit the coalesced requests to the I/O.
         let read = self.read.clone();
@@ -276,6 +284,8 @@ pin_project! {
         pub prefetch: Prefetch,
         requests_ready_chunks: usize,
         prefetch_ready_chunks: usize,
+        requested_segments: Arc<Counter>,
+        prefetched_segments: Arc<Counter>,
     }
 }
 
@@ -293,6 +303,7 @@ where
             match this.requests.as_mut().poll_next(cx) {
                 Poll::Ready(Some(item)) => {
                     items.push(item);
+                    this.requested_segments.inc();
                     if items.len() >= *this.requests_ready_chunks {
                         return Poll::Ready(Some(items));
                     }
@@ -309,6 +320,7 @@ where
             match this.prefetch.as_mut().poll_next(cx) {
                 Poll::Ready(Some(item)) => {
                     items.push(item);
+                    this.prefetched_segments.inc();
                     if items.len() >= prefetch_limit {
                         return Poll::Ready(Some(items));
                     }
@@ -349,17 +361,6 @@ impl SegmentRequest {
     }
 }
 
-impl From<(SegmentId, Segment, oneshot::Receiver<()>)> for SegmentRequest {
-    fn from(value: (SegmentId, Segment, oneshot::Receiver<()>)) -> Self {
-        let (id, location, cancel_handle) = value;
-        SegmentRequest {
-            id,
-            location,
-            cancel_handle,
-        }
-    }
-}
-
 async fn filter_with_cache(
     request: SegmentRequest,
     cache: Arc<dyn SegmentCache>,
@@ -392,19 +393,37 @@ struct CoalescedSegmentRequest {
 }
 
 #[derive(Default)]
-struct CoalescedCancellationHandle(Vec<oneshot::Receiver<()>>);
+struct CoalescedCancellationHandle {
+    handles: Vec<oneshot::Receiver<()>>,
+    cancel_received: Arc<Counter>,
+    cancelled: Arc<Counter>,
+}
 
 impl CoalescedCancellationHandle {
+    fn new(
+        handles: Vec<oneshot::Receiver<()>>,
+        cancel_received: Arc<Counter>,
+        cancelled: Arc<Counter>,
+    ) -> Self {
+        Self {
+            handles,
+            cancel_received,
+            cancelled,
+        }
+    }
+
     fn push(&mut self, handle: oneshot::Receiver<()>) {
-        self.0.push(handle);
+        self.handles.push(handle);
     }
 
     async fn cancelled(self) {
-        for rx in self.0 {
+        for rx in self.handles {
             // if this segment is completed before cancellation,
             // tx of this will be dropped, so ignore errors.
             let _ = rx.await;
+            self.cancel_received.inc();
         }
+        self.cancelled.inc();
     }
 }
 
@@ -480,6 +499,7 @@ fn coalesce(
     requests: Vec<SegmentRequest>,
     coalescing_window: u64,
     max_read: Option<u64>,
+    metrics: CoalescingMetrics,
 ) -> Vec<(CoalescedSegmentRequest, CoalescedCancellationHandle)> {
     let fetch_ranges = merge_ranges(
         requests.iter().map(|r| r.range()),
@@ -501,7 +521,11 @@ fn coalesce(
                     byte_range: range.clone(),
                     requests: vec![],
                 },
-                CoalescedCancellationHandle::default(),
+                CoalescedCancellationHandle::new(
+                    Vec::new(),
+                    metrics.cancel_received.clone(),
+                    metrics.cancelled.clone(),
+                ),
             )
         })
         .collect::<Vec<_>>();
@@ -516,6 +540,7 @@ fn coalesce(
     // Ensure we sort the requests by segment ID within the coalesced request.
     for (req, _) in coalesced.iter_mut() {
         req.requests.sort_unstable_by_key(|(id, _)| *id);
+        metrics.record(req);
     }
     coalesced
 }
@@ -563,11 +588,16 @@ where
     ret
 }
 
+#[derive(Clone)]
 struct CoalescingMetrics {
     bytes_uncoalesced: Arc<Counter>,
     bytes_coalesced: Arc<Counter>,
     request_count_uncoalesced: Arc<Counter>,
     request_count_coalesced: Arc<Counter>,
+    prefetched_segments: Arc<Counter>,
+    requested_segments: Arc<Counter>,
+    cancel_received: Arc<Counter>,
+    cancelled: Arc<Counter>,
 }
 
 impl From<VortexMetrics> for CoalescingMetrics {
@@ -579,6 +609,10 @@ impl From<VortexMetrics> for CoalescingMetrics {
             bytes_coalesced: metrics.counter(format!("{BYTES}.coalesced")),
             request_count_uncoalesced: metrics.counter(format!("{COUNT}.uncoalesced")),
             request_count_coalesced: metrics.counter(format!("{COUNT}.coalesced")),
+            prefetched_segments: metrics.counter("vortex.scan.segments.prefetch_count"),
+            requested_segments: metrics.counter("vortex.scan.segments.request_count"),
+            cancel_received: metrics.counter("vortex.scan.segments.cancel_received"),
+            cancelled: metrics.counter("vortex.scan.segments.cancelled"),
         }
     }
 }
diff --git a/vortex-layout/Cargo.toml b/vortex-layout/Cargo.toml
@@ -40,6 +40,7 @@ vortex-error = { workspace = true }
 vortex-expr = { workspace = true }
 vortex-flatbuffers = { workspace = true, features = ["layout"] }
 vortex-mask = { workspace = true }
+vortex-metrics = { workspace = true }
 
 [dev-dependencies]
 futures = { workspace = true, features = ["executor"] }
diff --git a/vortex-layout/src/scan/mod.rs b/vortex-layout/src/scan/mod.rs
@@ -14,6 +14,7 @@ use vortex_expr::transform::immediate_access::immediate_scope_access;
 use vortex_expr::transform::simplify_typed::simplify_typed;
 use vortex_expr::{ExprRef, Identity};
 use vortex_mask::Mask;
+use vortex_metrics::VortexMetrics;
 
 use crate::scan::filter::FilterExpr;
 use crate::scan::unified::UnifiedDriverStream;
@@ -56,6 +57,7 @@ pub struct ScanBuilder<D: ScanDriver> {
     // The number of splits to make progress on concurrently.
     concurrency: usize,
     prefetch_conjuncts: bool,
+    metrics: VortexMetrics,
 }
 
 impl<D: ScanDriver> ScanBuilder<D> {
@@ -72,6 +74,7 @@ impl<D: ScanDriver> ScanBuilder<D> {
             canonicalize: false,
             prefetch_conjuncts: false,
             concurrency: 1024,
+            metrics: Default::default(),
         }
     }
 
@@ -129,6 +132,11 @@ impl<D: ScanDriver> ScanBuilder<D> {
         self
     }
 
+    pub fn with_metrics(mut self, metrics: VortexMetrics) -> Self {
+        self.metrics = metrics;
+        self
+    }
+
     pub fn build(self) -> VortexResult<Scan<D>> {
         let projection = simplify_typed(self.projection.clone(), self.layout.dtype())?;
         let filter = self
@@ -146,7 +154,7 @@ impl<D: ScanDriver> ScanBuilder<D> {
             .collect();
 
         let splits = self.split_by.splits(&self.layout, &field_mask)?;
-        let mut collector = SegmentCollector::default();
+        let mut collector = SegmentCollector::new(self.metrics.clone());
         self.layout
             .required_segments(0, &filter_mask, &projection_mask, &mut collector)?;
         let (mut row_range_pruner, segments) = collector.finish();
diff --git a/vortex-layout/src/segments.rs b/vortex-layout/src/segments.rs
@@ -10,6 +10,7 @@ use futures::{SinkExt, Stream, StreamExt};
 use range_union_find::RangeUnionFind;
 use vortex_buffer::{Buffer, ByteBuffer};
 use vortex_error::{VortexExpect, VortexResult, vortex_err};
+use vortex_metrics::VortexMetrics;
 
 use crate::range_intersection;
 
@@ -93,14 +94,23 @@ type SegmentStore = BTreeMap<SegmentPriority, Vec<SegmentId>>;
 pub struct SegmentCollector {
     store: Arc<RwLock<SegmentStore>>,
     pub kind: RequiredSegmentKind,
+    metrics: VortexMetrics,
 }
 
 impl SegmentCollector {
+    pub fn new(metrics: VortexMetrics) -> Self {
+        Self {
+            metrics,
+            ..Default::default()
+        }
+    }
+
     pub fn with_priority_hint(&self, kind: RequiredSegmentKind) -> Self {
         SegmentCollector {
             store: self.store.clone(),
             // highest priority wins
             kind: kind.min(self.kind),
+            metrics: self.metrics.clone(),
         }
     }
 
@@ -110,6 +120,7 @@ impl SegmentCollector {
             RequiredSegmentKind::PRUNING => (0, 0),
             _ => (row_start, row_end),
         };
+        self.increment_metrics();
         let priority = SegmentPriority::new(start, end, self.kind);
         self.store
             .write()
@@ -126,6 +137,7 @@ impl SegmentCollector {
                 store: self.store.clone(),
                 cancellations_tx,
                 excluded_ranges: Default::default(),
+                metrics: self.metrics.clone(),
             },
             SegmentStream {
                 store: self.store,
@@ -135,13 +147,23 @@ impl SegmentCollector {
             },
         )
     }
+
+    fn increment_metrics(&self) {
+        self.metrics
+            .counter("vortex.scan.segments.count.total")
+            .inc();
+        self.metrics
+            .counter(format!("vortex.scan.segments.count.{:?}", self.kind))
+            .inc();
+    }
 }
 
 #[derive(Debug, Clone)]
 pub struct RowRangePruner {
     store: Arc<RwLock<SegmentStore>>,
     cancellations_tx: mpsc::UnboundedSender<SegmentId>,
     excluded_ranges: Arc<RwLock<RangeUnionFind<u64>>>,
+    metrics: VortexMetrics,
 }
 
 impl RowRangePruner {
@@ -184,6 +206,9 @@ impl RowRangePruner {
                 .flat_map(|key| store.remove(key).unwrap_or_default())
                 .collect()
         };
+        self.metrics
+            .counter("vortex.scan.segments.cancel_sent")
+            .add(cancelled_segments.len() as i64);
         for id in cancelled_segments {
             self.cancellations_tx
                 .send(id)
@@ -207,7 +232,14 @@ impl RowRangePruner {
                 if key.kind == RequiredSegmentKind::PRUNING {
                     return true; // keep segments required for pruning
                 }
-                range_intersection(&(key.row_start..key.row_end), &row_indices).is_some()
+                let keep =
+                    range_intersection(&(key.row_start..key.row_end), &row_indices).is_some();
+                if !keep {
+                    self.metrics
+                        .counter("vortex.scan.segment.pruned_by_row_indices")
+                        .inc();
+                }
+                keep
             });
     }
 }
@@ -347,6 +379,7 @@ pub mod test {
                 store: store.clone(),
                 cancellations_tx: tx,
                 excluded_ranges: Default::default(),
+                metrics: Default::default(),
             };
 
             // Test removing segments in range 0..200
@@ -405,6 +438,7 @@ pub mod test {
                 store: store.clone(),
                 cancellations_tx: tx,
                 excluded_ranges: Default::default(),
+                metrics: Default::default(),
             };
 
             // First removal (0..100)
@@ -448,6 +482,7 @@ pub mod test {
                 store: store.clone(),
                 cancellations_tx: tx,
                 excluded_ranges: Default::default(),
+                metrics: Default::default(),
             };
 
             // First removal (0..75)
@@ -492,6 +527,7 @@ pub mod test {
                 store: store.clone(),
                 cancellations_tx: tx,
                 excluded_ranges: Default::default(),
+                metrics: Default::default(),
             };
 
             // Create a buffer with specific row indices
@@ -543,6 +579,7 @@ pub mod test {
                 store: store.clone(),
                 cancellations_tx: tx,
                 excluded_ranges: Default::default(),
+                metrics: Default::default(),
             };
 
             // Drop the receiver to close the channel
@@ -569,6 +606,7 @@ pub mod test {
                 store: store.clone(),
                 cancellations_tx: tx,
                 excluded_ranges: Default::default(),
+                metrics: Default::default(),
             };
 
             // Test removing segments that cover the entire range

Original file line number	Diff line number	Diff line change
`@@ -49,5 +49,6 @@ impl<F: FileType> VortexFile<F> {`
`49`	`49`	`self.footer.layout().clone(),`
`50`	`50`	`self.footer().ctx().clone(),`
`51`	`51`	`)`
	`52`	`+ .with_metrics(self.metrics.clone())`
`52`	`53`	`}`
`53`	`54`	`}`