viirya opened a new pull request, #25583: URL: https://github.com/apache/datafusion/pull/25583
## Which issue does this PR close? Closes #25582. Related to apache/datafusion-comet#1204 and apache/datafusion-comet#6071. ## Rationale for this change Consumers that execute many partitions of one retained physical plan may report each partition's metrics separately. Today `metrics()` clones every registered metric before the caller can select a partition. With a fixed number of metrics per partition, reporting all N partitions on a retained tree can require O(N²) cloning and filtering work. This is motivated by Comet's opt-in shared physical plans. It does not imply that ordinary DataFusion queries have the same regression. The proposed API is useful to any consumer needing metrics for one partition and introduces no Spark attempt concepts. ## What changes are included in this PR? - Add `ExecutionPlanMetricsSet::clone_partition(usize)`, backed by an index of positions in the existing registration-ordered metrics vector. Retrieval clones only matching handles, with expected O(1 + matching metrics) work excluding mutex contention. Registration updates the vector and index under the same lock. - Add `ExecutionPlan::metrics_for_partition(usize)` with a compatibility fallback for external implementations. Core container-backed operators override it to use the index. - Forward through `DataSourceExec`, `DataSinkExec`/`DataSink`, and the FFI execution-plan adapter. Parquet's derived plan-wide `output_rows_skew` remains in `metrics()` and is excluded from partition snapshots. - Preserve full snapshots, registration order, duplicate metric names/labels, shared values, clone sharing, and independent registration after `From<MetricsSet>`. New snapshots observe later registrations; existing snapshots retain their original membership. ## What is the testing strategy for this PR? Container tests cover partition isolation, missing and sparse IDs, unpartitioned metrics, duplicate names, aggregation, snapshots with live values, `Clone`, `From`, and concurrent registration/readers. A real streaming → filter → projection tree runs 16 partitions concurrently while partition 0 waits after producing its first batch; assertions cover output values and per-node/full metrics. Parquet datasource and FFI tests cover delegation and global metrics semantics. Criterion benchmarks hold the target partition's metric count fixed while other partitions accumulate, compare indexed/full/legacy retrieval, measure registration and conversion, and report each completed task while a long partition retains the real tree. Comparisons against `main` should assess partition retrieval, full snapshots, registration, conversion, and allocation overhead separately. These deliberately small-task workloads should not be interpreted as general query or Spark speedups. Local development validation: - Container metrics: 36 tests passed. - Concurrent shared-tree integration: 1 test passed. - Parquet datasource: 1 test passed; Parquet sink: 3 tests passed. - FFI: 122 unit tests and 31 integration tests passed. - `cargo clippy --all-targets --all-features -- -D warnings` passed without warnings. - The extended workspace test command passed 11,775 Rust tests (2 ignored), then the SQL runner completed 522 files and reported metadata map-order mismatches in `arrow_field.slt:138` and `metadata.slt:422,440` (2 files / 3 assertions). The command did not pass; subsequent workspace suites are not claimed as validated. - Formatting and `git diff --check` passed on the final working tree. - The full `./dev/rust_lint.sh` suite could not start because a local prerequisite was unavailable; it is not claimed as passed. These are development validation results. Full validation with the target revision's declared dependencies remains pending. Preliminary local release benchmark results (means): with eight selected metrics, indexed retrieval stayed at 24–28 ns from 1 to 32,768 partitions, while the baseline full-snapshot filter grew from 26 ns to 1.07 ms. Indexed retrieval across a real three-operator tree stayed at 85–87 ns from 1 to 8,192 partitions. At 8,192 partitions, executing and reporting every task while partition 0 remained unfinished took 2.251 s on the baseline, 2.214 s with the index and legacy reporting API, and 30.368 ms using the new API. Baseline and modified builds used matching development configurations. These measurements are preliminary and should be repeated with the target revision's declared dependencies. The index has measurable costs: at eight metrics per partition and 8,192 partitions, requested live metric/container allocations increased by 1,065,008 bytes (10.7%; this measures requested allocation sizes, not RSS). Small-container registration was 2.2–2.7× slower across two process runs. Large registration timings varied and are not used to claim a registration speedup. `From<MetricsSet>` took 518 µs at that size versus 27 ns on the baseline because it now builds the index. Full retrieval over the real three-node tree took 512.54 µs on the baseline and 509.45 µs with the index at 8,192 partitions. ## Are there any user-facing changes? New Rust methods are provided defaults for existing trait implementations. `metrics_for_partition(p)` selects exactly `Metric::partition() == Some(p)`; `partition=None` metrics remain available through `metrics()`. Unknown partition IDs return an empty set when metrics are supported, not an invalid-partition error. Partition identifiers are operator-local; this method does not recurse into children or map partitions across operators. Callers may use `aggregate_by_name` on the returned set as usual; aggregation removes partition identity. External implementations retaining the default still pay for a full snapshot. Data sources/sinks whose own metrics provider constructs a full snapshot can likewise retain provider-side costs. FFI dispatches to the producer's override and transports only its result. The index adds one `usize` per partitioned registration (plus vector capacity and a hash-map entry per distinct partition), and a hash lookup during registration. Full snapshots remain O(total metrics) and still hold the same mutex, so concurrent full snapshots can delay partition readers. `From<MetricsSet>` now builds the index in O(total metrics) while retaining the original metrics vector allocation. `FFI_ExecutionPlan` gains a function pointer, changing its ABI layout. Label **api change**, target `main`, and do not backport this layout change to a patch release. Both sides of an FFI boundary need compatible builds. -- 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]
