laskoviymishka commented on code in PR #2521:
URL: https://github.com/apache/iceberg-rust/pull/2521#discussion_r4009795867
##########
crates/integrations/datafusion/src/physical_plan/scan.rs:
##########
@@ -174,11 +178,30 @@ impl ExecutionPlan for IcebergTableScan {
Box::pin(stream)
};
+ let baseline_metrics = BaselineMetrics::new(&self.metrics, partition);
+ let measured_stream = stream_with_baseline_metrics(limited_stream,
baseline_metrics);
+
Ok(Box::pin(RecordBatchStreamAdapter::new(
self.schema(),
- limited_stream,
+ measured_stream,
)))
}
+
+ fn metrics(&self) -> Option<MetricsSet> {
+ Some(self.metrics.clone_inner())
+ }
+
+ fn reset_state(self: Arc<Self>) -> DFResult<Arc<dyn ExecutionPlan>> {
Review Comment:
minor, but `reset_state` re-lists every field by hand, so a new stateful
field added later gets carried over silently unless someone remembers to touch
this too.
Deriving `Clone` and writing it as `let mut plan = (*self).clone();
plan.metrics = ExecutionPlanMetricsSet::new(); Ok(Arc::new(plan))` makes the
intent — clone everything, reset the metrics — explicit and drops the
boilerplate. Not blocking; the exhaustive literal is at least compile-safe
today.
##########
crates/integrations/datafusion/src/physical_plan/scan.rs:
##########
@@ -237,6 +260,18 @@ async fn get_batch_stream(
Ok(Box::pin(stream))
}
+fn stream_with_baseline_metrics(
+ mut stream: Pin<Box<dyn Stream<Item = DFResult<RecordBatch>> + Send>>,
+ baseline_metrics: BaselineMetrics,
+) -> Pin<Box<dyn Stream<Item = DFResult<RecordBatch>> + Send>> {
+ futures::stream::poll_fn(move |cx| {
+ let _timer = baseline_metrics.elapsed_compute().timer();
+ let poll = stream.as_mut().poll_next(cx);
+ baseline_metrics.record_poll(poll)
Review Comment:
one thing I couldn't confirm from here: the DF55 docs for `record_poll` say
it only updates `output_rows` and `end_time`, yet the unit test asserts
`output_batches == Some(1)` and `output_bytes == Some(...)`. If CI's green then
`record_poll` clearly populates those in this version and we're fine — but the
published docs contradict it, so the next person could reasonably assume the
assertions are wrong and "fix" them.
A one-line comment noting that DF55's `record_poll` also tracks
batches/bytes (or an explicit `output_batches().add(1)` if it turns out it
doesn't) would head that off. Which is it in practice?
##########
crates/integrations/datafusion/src/physical_plan/scan.rs:
##########
@@ -237,6 +260,18 @@ async fn get_batch_stream(
Ok(Box::pin(stream))
}
+fn stream_with_baseline_metrics(
+ mut stream: Pin<Box<dyn Stream<Item = DFResult<RecordBatch>> + Send>>,
+ baseline_metrics: BaselineMetrics,
+) -> Pin<Box<dyn Stream<Item = DFResult<RecordBatch>> + Send>> {
+ futures::stream::poll_fn(move |cx| {
+ let _timer = baseline_metrics.elapsed_compute().timer();
Review Comment:
worth a short comment that `_timer` has to stay a named binding — if someone
tidies this to `let _ = ...elapsed_compute().timer()` the guard drops
immediately and every `elapsed_compute` reading silently goes to 0, with the
tests still passing. Quiet trap for the next editor.
##########
crates/integrations/datafusion/src/physical_plan/scan.rs:
##########
@@ -247,3 +282,105 @@ fn get_column_names(
.collect::<Vec<String>>()
})
}
+
+#[cfg(test)]
+mod tests {
+ use std::sync::Arc;
+
+ use datafusion::arrow::array::Int64Array;
+ use datafusion::arrow::datatypes::{
+ DataType, Field, Schema as ArrowSchema, SchemaRef as ArrowSchemaRef,
+ };
+ use datafusion::arrow::record_batch::RecordBatch;
+ use datafusion::common::utils::memory::get_record_batch_memory_size;
+ use datafusion::physical_plan::metrics::{
+ BaselineMetrics, ExecutionPlanMetricsSet, MetricValue, MetricsSet,
+ };
+ use futures::StreamExt;
+
+ use super::stream_with_baseline_metrics;
+
+ #[test]
+ fn stream_with_baseline_metrics_records_rows_and_compute() {
+ let metrics = ExecutionPlanMetricsSet::new();
+ let baseline_metrics = BaselineMetrics::new(&metrics, 0);
+ let batch = make_batch();
+ let expected_output_bytes = get_record_batch_memory_size(&batch);
+ let stream = Box::pin(futures::stream::iter([Ok(batch)]));
+ let mut stream = stream_with_baseline_metrics(stream,
baseline_metrics);
+
+ futures::executor::block_on(async {
+ let batch = stream
+ .next()
+ .await
+ .expect("stream should return one item")
+ .expect("stream item should be valid");
+ assert_eq!(batch.num_rows(), 3);
+ assert!(stream.next().await.is_none());
+ });
+
+ let metrics = metrics.clone_inner();
+ assert_eq!(metrics.output_rows(), Some(3));
+ assert_eq!(output_batches(&metrics), Some(1));
+ assert_eq!(output_bytes(&metrics), Some(expected_output_bytes));
+ assert!(
+ metrics.elapsed_compute().is_some_and(|elapsed| elapsed > 0),
+ "elapsed_compute should be recorded"
+ );
+ assert!(
+ start_timestamp(&metrics).is_some_and(|timestamp| timestamp > 0),
+ "start_timestamp should be recorded"
+ );
+ assert!(
+ end_timestamp(&metrics).is_some_and(|timestamp| timestamp > 0),
+ "end_timestamp should be recorded"
+ );
+ }
+
+ fn make_batch() -> RecordBatch {
+ let schema = make_arrow_schema();
+ let values = Arc::new(Int64Array::from(vec![1, 2, 3]));
+ RecordBatch::try_new(schema, vec![values]).unwrap()
+ }
+
+ fn make_arrow_schema() -> ArrowSchemaRef {
+ Arc::new(ArrowSchema::new(vec![Field::new(
+ "id",
+ DataType::Int64,
+ false,
+ )]))
+ }
+
+ fn metric_value_as_usize(
+ metrics: &MetricsSet,
+ matches_metric: impl Fn(&MetricValue) -> bool,
+ ) -> Option<usize> {
+ metrics
+ .sum(|metric| matches_metric(metric.value()))
Review Comment:
`metric_value_as_usize` uses `MetricsSet::sum`, which is right for the
additive counters but off for `StartTimestamp`/`EndTimestamp` — those are
point-in-time values, not quantities to add up. It works today only because
there's a single partition, so there's exactly one matching entry; the moment
this scan reports more than one output partition, `sum` adds two epoch-ns
values and hands back garbage.
For the two timestamp helpers I'd `iter().find(...)` the matching metric
instead of summing. Small thing, but it's a trap waiting on any future
partitioning change.
##########
crates/integrations/datafusion/src/physical_plan/scan.rs:
##########
@@ -247,3 +282,105 @@ fn get_column_names(
.collect::<Vec<String>>()
})
}
+
+#[cfg(test)]
+mod tests {
+ use std::sync::Arc;
+
+ use datafusion::arrow::array::Int64Array;
+ use datafusion::arrow::datatypes::{
+ DataType, Field, Schema as ArrowSchema, SchemaRef as ArrowSchemaRef,
+ };
+ use datafusion::arrow::record_batch::RecordBatch;
+ use datafusion::common::utils::memory::get_record_batch_memory_size;
+ use datafusion::physical_plan::metrics::{
+ BaselineMetrics, ExecutionPlanMetricsSet, MetricValue, MetricsSet,
+ };
+ use futures::StreamExt;
+
+ use super::stream_with_baseline_metrics;
+
+ #[test]
+ fn stream_with_baseline_metrics_records_rows_and_compute() {
Review Comment:
this covers the happy single-batch path. Two cheap cases I'd add while it's
fresh: an inner stream that yields `Err(...)` — assert the error passes through
unchanged and `elapsed_compute` is still recorded, since error propagation is
the likely real-world failure mode for a remote scan — and a two-batch stream
to confirm `output_rows` accumulates across `record_poll` calls. Neither is
blocking.
##########
crates/integrations/datafusion/src/table/mod.rs:
##########
@@ -525,6 +525,41 @@ mod tests {
assert!(physical_plan.is_ok());
}
+ #[tokio::test]
+ async fn test_catalog_backed_provider_scan_reports_metrics() {
+ use datafusion::datasource::TableProvider;
+
+ let (catalog, namespace, table_name, _temp_dir) =
get_test_catalog_and_table().await;
+ let table_provider =
+ IcebergTableProvider::try_new(catalog.clone(), namespace.clone(),
table_name.clone())
+ .await
+ .unwrap();
+
+ let ctx = SessionContext::new();
+ let scan_plan = table_provider
+ .scan(&ctx.state(), None, &[], None)
+ .await
+ .unwrap();
+ let batches =
datafusion::physical_plan::collect(Arc::clone(&scan_plan), ctx.task_ctx())
+ .await
+ .unwrap();
+ let output_rows = batches.iter().map(|batch| batch.num_rows()).sum();
+
+ let metrics = scan_plan.metrics().expect("scan should expose metrics");
+ assert_eq!(metrics.output_rows(), Some(output_rows));
Review Comment:
the table `get_test_catalog_and_table()` builds has no rows, so
`output_rows` is 0 here and this assertion holds even if the wrapper never
fires — `BaselineMetrics` initializes that counter to 0. So it can't actually
catch a regression where metric collection gets bypassed.
I'd insert a small known row set before scanning and assert against that
count (> 0). That also makes the `elapsed_compute > 0` check below meaningful
instead of flaky — an empty scan can finish inside the clock resolution on a
fast runner and record 0ns. `test_catalog_backed_provider_insert` has the
insert path. wdyt?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
