feat: add partition-aware metrics snapshots by viirya · Pull Request #25583 · apache/datafusion

viirya · 2026-09-22T00:20:57Z

Which issue does this PR close?

Rationale for this change

Consumers reporting metrics for one partition currently clone all registered metrics before filtering. This becomes increasingly expensive when a physical plan is retained across many partition executions.

What changes are included in this PR?

Add MetricsSet::for_partition with indexed partition selection.
Record snapshot boundaries so metrics() does not need to clone every metric handle. Build and update the shared partition index only when selecting a partition.
Preserve registration order and existing iterator APIs through lazy materialization.

What is the testing strategy for this PR?

Unit and integration tests cover partition selection, snapshot isolation, concurrent registration, and FFI round trips.

Are there any user-facing changes?

Callers can select a partition with plan.metrics().map(|m| m.for_partition(partition)). Unpartitioned metrics remain available in the full set; unknown partitions produce an empty set.

The existing ExecutionPlan API and FFI layout remain unchanged. Partition selection does not eliminate work already performed by a metrics provider, such as FFI transport.

codecov-commenter · 2026-09-22T00:47:11Z

Codecov Report

❌ Patch coverage is 98.49850% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.49%. Comparing base (0576a0b) to head (d88d2ba).
⚠️ Report is 33 commits behind head on main.

Files with missing lines	Patch %	Lines
...usion/physical-expr-common/src/metrics/snapshot.rs	97.67%	4 Missing ⚠️
datafusion/physical-expr-common/src/metrics/mod.rs	99.18%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #25583      +/-   ##
==========================================
+ Coverage   82.42%   82.49%   +0.07%     
==========================================
  Files        1139     1141       +2     
  Lines      435372   437653    +2281     
  Branches   435372   437653    +2281     
==========================================
+ Hits       358845   361045    +2200     
+ Misses      54826    54821       -5     
- Partials    21701    21787      +86

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

sunchao

Thanks, Liang-Chi. Reviewed 32cf894 against base 0576a0b; no actionable findings.

The review covered container invariants and concurrency, operator/source/sink delegation, FFI/API compatibility, tests, and performance. The append-only index preserves registration order, duplicates, and shared native metric values. Partition selection and exclusion of unpartitioned metrics match the documented contract, and FFI forwards to the producer's implementation.

Validation:

All 36 metrics unit tests and the concurrent shared-plan integration test passed locally with the declared dependencies, including Arrow 60.0.0, and no dependency overrides.
Current-head CI has 39 successful checks and three skipped extended-test checks, with none failed or pending. I did not rerun the full extended workspace suite locally.
Independent container benchmarks used separate base/head release-nonlto builds, eight metrics per partition, and base/head/head/base run order. Indexed lookup stayed around 67–72 ns at both tested sizes (1 and 8,192 partitions), while filtering a full snapshot took about 0.63 ms at 8,192 partitions.

The costs are real: container registration was roughly 2x slower for the small set and 4.4x slower at 8,192 partitions; conversion also became more expensive. These are disclosed tradeoffs. The measurements were on a shared host, and ordinary query execution without repeated partition reporting and whole-Comet workloads were not benchmarked, so these results do not establish a general query speedup or slowdown.

The intentional FFI ABI change is documented, targets main, and carries the required api change label. LGTM.

viirya · 2026-09-23T02:26:56Z

This includes api change. @alamb Do you want to take a look before I merge this? Thanks.

alamb · 2026-09-23T15:36:14Z

What is the testing strategy for this PR?

I didn't make it through this entire description (it seems also to have a bunch of internal stuff and talks about pending validation 😕 )

....
Container metrics: 36 tests passed.
These are development validation results. Full validation with the target revision's declared dependencies remains pending.
...

I wonder if you could pare down the description into something easier to read / know what is important to read for context on future PRs

viirya · 2026-09-23T15:58:46Z

Thanks @alamb for the feedback! I've trimmed the description to focus on the changes and removed the unnecessary testing details.

alamb

Thanks @viirya and @sunchao -- the usecase makes sense to me, but I think it would be good to consider if we can avoid adding a new API to ExecutionPlan and instead make the metrics more partition aware themselves.

I left some comments -- let me know what you think

alamb · 2026-09-23T15:43:53Z

+// specific language governing permissions and limitations
+// under the License.
+
+//! Container costs; run alongside the physical-plan shared-tree benchmark.


what does "container" mean in this context? Also what is the "physical-plan shared-tree benchmark"? Maybe we could add a url link

@alamb Removed this benchmark and its terminology. The replacement in datafusion/core/benches/partition_metrics.rs runs a SQL query before reading the plan's metrics.

alamb · 2026-09-23T16:15:06Z

@@ -0,0 +1,116 @@
+// Licensed to the Apache Software Foundation (ASF) under one


I don't understand the value of this benchmark -- it seems like a better benchmark would be to run an actual query (SELECT ....) and then call get_metrics and metrics_per-partitition 🤔

I think it could be removed

@alamb Agreed, removed. The replacement uses SELECT a + 1 AS b FROM t WHERE a < 16 and measures both metrics retrieval and execution with per-partition reporting.

alamb · 2026-09-23T16:17:52Z

+#[derive(Default, Debug)]
+struct IndexedMetricsSet {
+    metrics: MetricsSet,
+    partitions: HashMap<usize, Vec<usize>>,


Could we document what the entries in this this usize mean?
Is it a map from partition --> indexes in Metrics set for that partition?

@alamb Yes—the values are positions in the registry's metrics vector, in registration order. I've documented this on PartitionIndex::positions.

alamb · 2026-09-23T16:41:24Z

+// under the License.
+
+//! Run with `cargo bench -p datafusion-physical-plan --bench partition_metrics`.
+//! SQL cannot express per-task metrics reporting while retaining a shared plan.


why can't SQL express per-task metrics? That doesn't make any sense to me

@alamb You're right; that comment was incorrect. I've removed it and switched to a benchmark that builds the plan from SQL, executes it, and reads its metrics in Rust.

alamb · 2026-09-23T16:48:39Z

        None
    }

+    /// Return a snapshot of metrics whose [`Metric::partition`] is `Some(partition)`.


I think we have had challenges in the past with _partition type APIs on ExecutionPlan to add partition aware APIs. For example, the similar partition_statistics was deprecated -- see https://docs.rs/datafusion/latest/datafusion/physical_plan/trait.ExecutionPlan.html#method.partition_statistics

I wonder if it would make more sense to try and plumb the notion of partition more deeply in MetricsSet itself somehow? You already kind of do this for alreayd with IndexedMetricSet

maybe we could make MetricsSet itself store metrics per-partition (rather than a flat Vec)

and then you could add apis to metricsSet to access per partition information 🤔

@alamb Implemented this as plan.metrics().map(|m| m.for_partition(partition)). Snapshots retain a fixed registration boundary and share an index updated only on partition reads, keeping index maintenance off the registration path.

viirya · 2026-09-23T18:57:26Z

Thanks @alamb, I've moved partition selection to MetricsSet::for_partition() and removed the new ExecutionPlan method and FFI layout changes. The index is updated on demand, so registration remains a Vec append. I've also replaced the container benchmark with a SQL-based benchmark.

alamb

Tanks @viirya -- this is looking better

I don't quite understand some parts of this design, though the API looks nice to me

One thought i had while reading this was that you seem to have two types of MetricsSet -- one for a single partition and one for many partitiones

Maybe we could make that expliciy in the type system like

enum MetricsSet {
  All { 
    metrics: Vec<Arc<Metric>>,
    index: ... , // maybe
  },
  Partition {
    Vec<Arc<Metric>>
  }
}

That way you could have

  pub fn for_partition(&self, partition: usize) -> Self {
      /// convert MetricsSet::All to MetricsSet::Partition
    }

Or something

alamb · 2026-09-23T19:16:49Z

-    /// Add the specified metric
+    /// Add the specified metric.
+    ///
+    /// Mutating a deferred snapshot first copies its members into an independent


This seems like an implementation detail -- I think the docs on the overall structure are probably enough

alamb · 2026-09-23T19:18:05Z


+    /// Return a snapshot containing only metrics with `partition == Some(partition)`.
+    ///
+    /// For registry-backed snapshots, this incrementally indexes registrations


this likewise seems like a bunch of implementation specific detail -- I think it would help if the comments only focused on the end user visible behavior. I am not sure how to interpret all the stuff about cloning matching handles, et

alamb · 2026-09-23T19:19:27Z

+// specific language governing permissions and limitations
+// under the License.
+
+//! Fixed-membership snapshots of an append-only registry.


Could we maybe make this specific tot he code that is in the module -- somethig like

//! Metric [`Registry`] implementation

alamb · 2026-09-23T19:19:42Z

+
+//! Fixed-membership snapshots of an append-only registry.
+//!
+//! Registration only appends to a vector. Partition readers share an index that


this is all details of the implementaiton that is probably not useful in the module level comments

alamb · 2026-09-23T19:20:23Z

+
+#[derive(Debug, Default)]
+pub(super) struct Registry {
+    pub(super) metrics: Vec<Arc<Metric>>,


recommend keeping this private and accessing via a method rather than direct field access

alamb · 2026-09-23T19:25:26Z

+        let Self::Owned(metrics) = self else {
+            unreachable!()
+        };
+        metrics.push(metric);


doesn't this also invalidate the partition index, if there is one?

alamb · 2026-09-23T19:26:35Z

+        }
+    }
+
+    fn select(&mut self, partition: usize, end: usize) -> Vec<Arc<Metric>> {


could we please document what partition and end mean in this? Is the end relative to just the metrics in the partition? or all the metrics?

alamb · 2026-09-23T19:27:31Z

+        Self::Owned(match self {
+            Self::Owned(metrics) => metrics
+                .iter()
+                .filter(|m| m.partition() == Some(partition))


isn't this the filter you were trying to avoid? shouldn't this be using the index if it is available 😕

alamb · 2026-09-23T19:28:57Z

+use std::sync::{Arc, OnceLock};
+
+#[derive(Debug, Default)]
+pub(super) struct Registry {


I am sorry I don't understand this design -- maybe we can comment why bother creating the PartitionIndex at all? It seems like it just adds overhead (a new hash map and a bunch of allocations)

alamb · 2026-09-23T19:30:31Z

+    // Number of registry entries already examined, including global metrics.
+    indexed: usize,
+    // Partition ID -> positions in Registry::metrics, in registration order.
+    positions: HashMap<usize, Vec<usize>>,


Since the Arc's are only a few more pointers, I wonder if you considered having this HashMap<usize, Arc<Metric> or something, so getting the partition's metrics would be an update to the index, and then a clone of the relevant Vec

feat: add indexed partition metrics access to execution plans

32cf894

viirya added the api change Changes the API exposed to users of the crate label Sep 22, 2026

github-actions Bot added physical-expr Changes to the physical-expr crates core Core DataFusion crate datasource Changes to the datasource crate ffi Changes to the ffi crate physical-plan Changes to the physical-plan crate auto detected api change Auto detected API change labels Sep 22, 2026

viirya requested a review from sunchao September 22, 2026 00:35

sunchao approved these changes Sep 22, 2026

View reviewed changes

alamb reviewed Sep 23, 2026

View reviewed changes

refactor: move partition metrics selection into MetricsSet

d88d2ba

viirya changed the title ~~feat: add indexed partition metrics access to execution plans~~ feat: add partition-aware metrics snapshots Sep 23, 2026

viirya removed the api change Changes the API exposed to users of the crate label Sep 23, 2026

github-actions Bot removed the auto detected api change Auto detected API change label Sep 23, 2026

alamb reviewed Sep 23, 2026

View reviewed changes

		@@ -0,0 +1,116 @@
		// Licensed to the Apache Software Foundation (ASF) under one

Conversation

viirya commented Sep 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

What is the testing strategy for this PR?

Are there any user-facing changes?

Uh oh!

codecov-commenter commented Sep 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

sunchao left a comment

Choose a reason for hiding this comment

Uh oh!

viirya commented Sep 23, 2026

Uh oh!

alamb commented Sep 23, 2026

Uh oh!

viirya commented Sep 23, 2026

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

viirya commented Sep 23, 2026

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

viirya commented Sep 22, 2026 •

edited

Loading

codecov-commenter commented Sep 22, 2026 •

edited

Loading