Skip to content

feat: bind task-owned RSS callbacks to native shuffle plans (4/n) - #5491

Open
pingzh wants to merge 2 commits into
apache:mainfrom
pingzh:pingzh-rss-destination-aware-native-shuffle
Open

feat: bind task-owned RSS callbacks to native shuffle plans (4/n)#5491
pingzh wants to merge 2 commits into
apache:mainfrom
pingzh:pingzh-rss-destination-aware-native-shuffle

Conversation

@pingzh

@pingzh pingzh commented Aug 27, 2026

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Part of #5352. This is the fourth foundational PR and does not close the issue.

Previous PRs:

Rationale for this change

Previous PRs introduced RSS partition writers, destination-aware shuffle plans, and native RSS execution. However, serialized RSS plans cannot execute without connecting their task-owned JVM callback to the native planner.

This PR completes that connection while preserving existing local shuffle behavior and JNI compatibility.

What changes are included in this PR?

  • Add task-scoped RSS callback registration through a separate JNI method, preserving the existing createPlan signature.
  • Pass optional shuffle callbacks through CometExecIterator into native execution contexts and the physical planner.
  • Bind RSS shuffle plans to their task-owned callback and preserve callback isolation between tasks.
  • Reject RSS plans without callbacks, RSS plans containing local output paths, and callbacks supplied to local shuffle plans.
  • Apply the JVM’s conservative maximum byte-array size to RSS frame limits.
  • Preserve original callback exceptions across JNI boundaries.
  • Clean up native resources and all shuffle inputs when callback registration fails.

How are these changes tested?

  • Added planner tests covering RSS binding, task isolation, destination validation, and end-to-end native execution.
  • Added Spark/JNI integration tests covering complete RSS frames, original callback exception identity, existing JNI signature compatibility, and registration-failure cleanup.
  • Passed 105 Rust tests across planner, shuffle, JNI bridge, and protobuf suites.
  • Passed all 35 tests in CometNativeShuffleSuite.
  • Passed workspace-wide Clippy, formatting checks, and the full Maven reactor build.

@ziting-openai ziting-openai left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The PR currently points at the already-merged 3/n implementation instead of the advertised 4/n task-owned RSS callback integration.

Comment thread native/shuffle/src/shuffle_writer.rs
@pingzh
pingzh force-pushed the pingzh-rss-destination-aware-native-shuffle branch from 8281c34 to 6658894 Compare August 27, 2026 01:54

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed head 6658894632adf114fb86313c282b2ecf95c8453f against base 4897161704b7b8b7dfa909f4bf897c6508b11117 with five independent reviewer agents. Existing validation: 83 targeted Rust tests passed with default features disabled, and a standalone Arrow 18.3 ownership probe confirmed retention until explicit release. Full Spark/JNI integration tests were not rerun.

Comment thread spark/src/main/scala/org/apache/comet/CometExecIterator.scala

@sunchao sunchao left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Re-reviewed head 69e3cfb8239806648b79b24148607be16504e378 against base 4897161704b7b8b7dfa909f4bf897c6508b11117 with five independent agents; no new actionable findings. The earlier Arrow cleanup P2 is fixed and its Spark regression passes.

Validation: 105 targeted Rust tests and all 35 CometNativeShuffleSuite tests passed on Spark 4.1.3/JDK 17, with native default features disabled. Other Spark versions and live RSS integration were not tested. CI is still running.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants