Emit physical_plan and data_node_metrics for LATE_MATERIALIZATION profile

[finnegancarroll]

Summary

profile=true returns rich per-task diagnostics (physical_plan, data_node_metrics) for SHARD_FRAGMENT and COORDINATOR_REDUCE stages but nothing for LATE_MATERIALIZATION. When the LM stage dominates wall clock, the profile gives no information about why.

This change wires the fetch-by-row-ids path to return both fields when profiling is enabled.

Changes

Rust (analytics-backend-datafusion/rust/src/):

• api.rs: Clone the DataFusion physical plan before execute_stream and pass it to the stream handle via a new wrap_stream_as_handle_with_plan() helper.
• query_executor.rs: New wrap_stream_as_handle_with_plan() that attaches an optional physical plan to the QueryStreamHandle, enabling get_metrics_json() to walk the plan tree.

Java (analytics-engine/):

• FetchByRowIdsRequest: Add profile boolean field (wire-serialized).
• AnalyticsSearchService.drainFetchByRowIds(): Extract execution metrics after stream exhaustion when request.profile() is true; send via onCompleteWithMetrics() sentinel.
• LateMaterializationStageExecution: Pass config.profile() when constructing fetch requests; implement onStreamComplete() on GatherListener to store trailing metrics on the stage task.

Integration Test:

• LateMaterializationProfileIT: Asserts LM stage tasks return data_node_metrics and physical_plan containing ParquetExec.

Testing

• Java compiles and passes spotless locally
• Integration test written (requires full sandbox build to run)

Resolves #22601 (items 1 and 2)

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit 896a17c)

Here are some key observations to aid the review process:

🧪 PR contains tests

🔒 No security concerns identified

✅ No TODO sections

🔀 No multiple PR themes

⚡ Recommended focus areas for review Metrics Assigned to Wrong Task GatherListener is constructed with tasks().get(0) for every shard/target in the scatter loop. When the LM stage fans out to multiple shards, each shard's onStreamComplete will call task.setDataNodeMetrics(trailingMetadata) on the same task, overwriting metrics from previously completed shards. Only the last-arriving shard's metrics will be visible in the profile, making per-shard diagnostics misleading or lost. new GatherListener(stitcher, plan, tasks().get(0)), Wire Compatibility The serialization format for FetchByRowIdsRequest now unconditionally reads/writes a boolean for profile. If mixed-version nodes exist during a rolling upgrade, a new coordinator sending to an old data node (or vice versa) will fail deserialization. Consider gating with a version check via out.getVersion()/in.getVersion() or documenting that this is only safe for full-cluster deploys. this.profile = in.readBoolean(); } @Override public void writeTo(StreamOutput out) throws IOException { super.writeTo(out); out.writeString(queryId); out.writeInt(stageId); shardId.writeTo(out); out.writeString(backendId); out.writeLongArray(rowIds); out.writeStringArray(columns); out.writeBoolean(profile);

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to 896a17c

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Guard against empty tasks list Using tasks().get(0) unconditionally assumes at least one task exists and that all dispatched shard fetches share the same task. If tasks() can be empty this will throw IndexOutOfBoundsException, and if multiple tasks exist, metrics from different shards will overwrite each other on the same task. Guard against an empty list and select the task corresponding to the current shard/target. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LateMaterializationStageExecution.java [442] +List<StageTask> currentTasks = tasks(); +StageTask stageTask = currentTasks.isEmpty() ? null : currentTasks.get(0); transport.dispatchFetchByRowIds( request, target.node(), - new GatherListener(stitcher, plan, tasks().get(0)), + new GatherListener(stitcher, plan, stageTask), config.parentTask(), pending ); Suggestion importance[1-10]: 5 __ Why: The concern about tasks().get(0) throwing on empty list is a legitimate defensive concern, though in practice the LM stage likely always has a task. The suggestion has moderate value as a defensive coding improvement.	Low
General	Avoid stale static provisioning flag A static provisioning flag persists across test class instances/JVM reuse and can cause the second test to skip index creation when the cluster state has been reset between runs (e.g. after cluster restarts or when tests are reordered). Prefer using @BeforeClass or an instance-level check that verifies the index actually exists on the cluster before skipping provisioning. sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/LateMaterializationProfileIT.java [26-34] -private static boolean provisioned = false; - private void ensureProvisioned() throws IOException { - if (!provisioned) { + Response head; + try { + head = client().performRequest(new Request("HEAD", "/" + INDEX)); + } catch (Exception e) { + head = null; + } + if (head == null || head.getStatusLine().getStatusCode() != 200) { createIndex(); indexData(); - provisioned = true; } } Suggestion importance[1-10]: 4 __ Why: Valid concern about static state persistence in tests, but the impact is limited to test reliability edge cases and OpenSearch integration tests typically reset state between runs.	Low

---

Previous suggestions

Suggestions up to commit 96a9fa4

Category	Suggestion	Impact
Possible issue	Avoid overwriting metrics across shards Using tasks().get(0) for every target shard means multiple concurrent GatherListeners will call setDataNodeMetrics on the same StageTask, causing later metrics to overwrite earlier ones (or race). Associate each dispatch with the per-shard/per-target task, or aggregate metrics safely across listeners. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LateMaterializationStageExecution.java [439-445] + StageTask perTargetTask = resolveTaskForTarget(target); // per-shard task transport.dispatchFetchByRowIds( request, target.node(), - new GatherListener(stitcher, plan, tasks().get(0)), + new GatherListener(stitcher, plan, perTargetTask), config.parentTask(), pending ); Suggestion importance[1-10]: 8 __ Why: Legitimate concern: passing tasks().get(0) to all GatherListener instances means concurrent setDataNodeMetrics calls on the same task will overwrite each other, losing per-shard metrics.	Medium
	Guard new wire field with version check Reading an unconditional boolean here will break wire compatibility with older nodes that do not write this field. Gate the read/write of profile on the stream version (e.g., in.getVersion().onOrAfter(...)) so mixed-version clusters do not deserialize a truncated payload or misread following bytes. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/FetchByRowIdsRequest.java [73] this.rowIds = in.readLongArray(); this.columns = in.readStringArray(); - this.profile = in.readBoolean(); + this.profile = in.getVersion().onOrAfter(Version.V_x_y_z) ? in.readBoolean() : false; } Suggestion importance[1-10]: 7 __ Why: Valid concern for wire compatibility in mixed-version clusters, though the impact depends on whether this action is used cross-version. The suggestion is reasonable but the codebase context (sandbox plugin) may not require strict BWC.	Medium
General	Fix unreliable static provisioning flag A static provisioning flag persists across test class instances/JVM reuse but the underlying index may not exist (e.g., after cluster reset between test classes), leading to failures when provisioned is true but the index is missing. Use a @BeforeClass setup or check index existence at runtime instead of a static boolean. sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/LateMaterializationProfileIT.java [26-34] -private static boolean provisioned = false; - private void ensureProvisioned() throws IOException { - if (!provisioned) { + Response head = client().performRequest(new Request("HEAD", "/" + INDEX)); + if (head.getStatusLine().getStatusCode() != 200) { createIndex(); indexData(); - provisioned = true; } } Suggestion importance[1-10]: 4 __ Why: Reasonable point about test reliability, but in typical IT test setups the JVM/class is fresh per test class, making the risk relatively low. Improvement is minor.	Low

Suggestions up to commit 0f71d04

Category	Suggestion	Impact
Possible issue	Avoid overwriting metrics across shards All per-shard GatherListeners share the same tasks().get(0) and each call task.setDataNodeMetrics(trailingMetadata) on completion, so metrics from later-completing shards will silently overwrite earlier ones. Either associate one StageTask per shard/fetch target, or merge/append metrics rather than overwriting. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LateMaterializationStageExecution.java [439-445] + StageTask shardTask = tasks().get(targetIndex); // or per-target task transport.dispatchFetchByRowIds( request, target.node(), - new GatherListener(stitcher, plan, tasks().get(0)), + new GatherListener(stitcher, plan, shardTask), config.parentTask(), pending ); Suggestion importance[1-10]: 7 __ Why: Legitimate observation that all GatherListeners share tasks().get(0) and setDataNodeMetrics will overwrite metrics from earlier shards, potentially losing profiling data. However, this depends on how setDataNodeMetrics is implemented (could accumulate).	Medium
	Guard new wire field with version check Adding an unconditional readBoolean/writeBoolean to the wire format breaks BWC with older nodes that don't write this field. Gate the read/write on a version check (e.g. in.getVersion().onOrAfter(...)) or ensure this action is only used within a single-version cluster; otherwise mixed-version clusters will fail to deserialize FetchByRowIdsRequest. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/FetchByRowIdsRequest.java [73] this.rowIds = in.readLongArray(); this.columns = in.readStringArray(); - this.profile = in.readBoolean(); + this.profile = in.getVersion().onOrAfter(Version.V_3_0_0) ? in.readBoolean() : false; } Suggestion importance[1-10]: 6 __ Why: Valid concern for BWC in mixed-version clusters, but this is a sandbox plugin and the scope of BWC requirements is unclear. Still, adding a version guard is a reasonable defensive practice.	Low
General	Avoid static flag for test provisioning Using a static flag to cache provisioning across tests is unsafe because the JUnit cluster/state may be reset between tests, but the flag will persist within the JVM, causing later tests to skip index creation and fail. Prefer @BeforeClass/@Before semantics or drop the static caching and simply re-create the index idempotently per test. sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/LateMaterializationProfileIT.java [26-34] -private static boolean provisioned = false; - private void ensureProvisioned() throws IOException { - if (!provisioned) { - createIndex(); - indexData(); - provisioned = true; - } + createIndex(); + indexData(); } Suggestion importance[1-10]: 3 __ Why: Minor test-code concern; static caching for provisioning is a common pattern in IT tests and typically works within a single test class JVM. The suggestion has minor merit but low impact.	Low

Suggestions up to commit 19a8d72

Category	Suggestion	Impact
Possible issue	Add version gating for new wire field Adding a new field to the wire format without version gating breaks BWC: a newer node writing profile to an older-version node (or vice versa) will misalign the stream. Gate the read/write of profile on out.getVersion()/in.getVersion() against the version this field was introduced in, defaulting to false for older versions. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/FetchByRowIdsRequest.java [65-74] this.rowIds = in.readLongArray(); this.columns = in.readStringArray(); - this.profile = in.readBoolean(); + this.profile = in.getVersion().onOrAfter(PROFILE_FIELD_VERSION) ? in.readBoolean() : false; } Suggestion importance[1-10]: 7 __ Why: Legitimate BWC concern: adding a new field to the wire format without version gating can break cross-version communication during rolling upgrades. Version-gating with StreamInput/StreamOutput version checks is standard OpenSearch practice.	Medium
	Guard against empty or mismatched task list Using tasks().get(0) unconditionally assumes the tasks list is non-empty and that a single task represents all shard fetches. If tasks() is empty this will throw IndexOutOfBoundsException, and if there are multiple shard tasks all GatherListeners will write metrics to the same task, potentially overwriting each other. Consider guarding for emptiness and/or selecting the task corresponding to the target shard. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LateMaterializationStageExecution.java [439-445] +StageTask stageTask = tasks().isEmpty() ? null : tasks().get(0); transport.dispatchFetchByRowIds( request, target.node(), - new GatherListener(stitcher, plan, tasks().get(0)), + new GatherListener(stitcher, plan, stageTask), config.parentTask(), pending ); Suggestion importance[1-10]: 5 __ Why: Valid defensive concern about tasks().get(0) potentially throwing IndexOutOfBoundsException if empty, though in practice the task list is likely always populated at this point. The concern about multiple shard tasks sharing one target task for metrics is also worth noting.	Low
General	Avoid static provisioning flag in tests A static provisioned flag persists across test class instances/JVMs used by the test runner, which can cause tests to skip index creation when the cluster has been reset between runs (e.g. randomized test suite runs), leading to spurious failures. Use an instance field with @Before, or check for index existence via the REST client instead of relying on a static boolean. sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/LateMaterializationProfileIT.java [26-34] -private static boolean provisioned = false; +private boolean provisioned = false; private void ensureProvisioned() throws IOException { if (!provisioned) { createIndex(); indexData(); provisioned = true; } } Suggestion importance[1-10]: 4 __ Why: Reasonable point about static state persisting across test runs, though in OpenSearch REST tests the cluster typically persists across tests in the same class, which is why static flags are commonly used. Impact is minor.	Low

Suggestions up to commit 3714eb7

Category	Suggestion	Impact
Possible issue	Avoid shared task overwrites across shards Using tasks().get(0) for every shard's GatherListener means metrics from all shards overwrite a single task's dataNodeMetrics, losing per-shard/per-target metrics and potentially causing races if setDataNodeMetrics is not thread-safe or if multiple shards complete concurrently. Consider selecting the task that corresponds to this specific target/plan, or guarding against multiple concurrent writes. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LateMaterializationStageExecution.java [439-445] +StageTask stageTask = selectTaskForTarget(target); // per-target task selection transport.dispatchFetchByRowIds( request, target.node(), - new GatherListener(stitcher, plan, tasks().get(0)), + new GatherListener(stitcher, plan, stageTask), config.parentTask(), pending ); Suggestion importance[1-10]: 6 __ Why: Using tasks().get(0) for every shard's GatherListener could lead to overwriting metrics or race conditions if multiple shards complete concurrently. This is a plausible concern, but the suggestion lacks concrete evidence about the task model and the improved code references an undefined selectTaskForTarget method.	Low
General	Make provisioning idempotent and resilient A static flag persists across test class instances/JVM reuse but the index itself may be deleted between test runs (e.g. by test cluster teardown), leading to false-positive "provisioned" state and failing tests. Either make the flag non-static, or check for index existence rather than relying on a boolean cache. sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/LateMaterializationProfileIT.java [26-34] -private static boolean provisioned = false; - private void ensureProvisioned() throws IOException { - if (!provisioned) { - createIndex(); - indexData(); - provisioned = true; + try { + client().performRequest(new Request("HEAD", "/" + INDEX)); + return; + } catch (Exception e) { + // fall through and (re)create } + createIndex(); + indexData(); } Suggestion importance[1-10]: 4 __ Why: The static provisioned flag could indeed cause issues across test runs, and checking for index existence is a more robust approach. This is a reasonable test hygiene improvement but not critical.	Low
	Verify metrics availability before stream close ctx.getExecutionMetrics() is called before the try-with-resources closes ctx, but the metrics reflect the fully-drained plan only if drain occurred; ensure the native side has finalized metrics at this point (i.e., after the last batch). If metrics are only finalized on stream close, calling before ctx closes may return partial/empty data. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java [395-404] +byte[] metricsJson = null; if (request.profile()) { - byte[] metricsJson = ctx.getExecutionMetrics(); - if (metricsJson != null) { - responseHandler.onCompleteWithMetrics(metricsJson); - } else { - responseHandler.onComplete(); - } + metricsJson = ctx.getExecutionMetrics(); // ensure this is valid before stream close +} +if (metricsJson != null) { + responseHandler.onCompleteWithMetrics(metricsJson); } else { responseHandler.onComplete(); } Suggestion importance[1-10]: 3 __ Why: The suggestion raises a valid concern about metrics finalization timing but is speculative and doesn't materially change the ordering — ctx.getExecutionMetrics() is still called before ctx closes in both versions.	Low

[github-actions]

Persistent review updated to latest commit 19a8d72

[github-actions]

Persistent review updated to latest commit 0f71d04

[github-actions]

✅ Gradle check result for 0f71d04: SUCCESS

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 71.53%. Comparing base (25c32e5) to head (0f71d04).
⚠️ Report is 2 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #22676      +/-   ##
============================================
+ Coverage     71.52%   71.53%   +0.01%     
- Complexity    77023    77041      +18     
============================================
  Files          6156     6156              
  Lines        358422   358422              
  Branches      52245    52245              
============================================
+ Hits         256351   256396      +45     
+ Misses        81694    81665      -29     
+ Partials      20377    20361      -16     

☔ 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.

[github-actions]

Persistent review updated to latest commit 96a9fa4

[github-actions]

❌ Gradle check result for 96a9fa4: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

Persistent review updated to latest commit 896a17c

[github-actions]

❌ Gradle check result for 896a17c: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?