test: [branch-1.0] wait for Parquet write plan callbacks (#6108)

[andygrove]

Backport of #6108 to branch-1.0.

#6108 merged to main as an empty commit (df8e153). #6156 had already made the same change to captureWritePlan on 09-23, when this flake failed the Spark 4.2 nightly. So this PR ports the captureWritePlan hunk of #6156 (dd68a53). The rest of #6156 corrects a Spark 3.4 fallback that #6041 added, and #6041 is not on branch-1.0.

branch-1.1 needs nothing. It was cut after #6156, and its CometParquetWriterTestBase.scala is byte-identical to the one on main.

Which issue does this PR close?

None. #6075 is closed on main by #6108.

Rationale for this change

branch-1.0 still has the flaky helper. captureWritePlan registers a QueryExecutionListener without first draining the asynchronous listener bus, then returns the first matching plan it sees. If an earlier write's end event is still queued, the helper returns that write's plan, and the test reports that the write under test fell back to Spark. That is the failure described in #6075.

On branch-1.0, 21 of the suite's 30 tests write their seed or input data without the native writer just before the write under test:

• 12 complex-type tests, through writeComplexTypeData.
• The 6 SaveMode tests, through materializeAsCometSource. Two of them also seed the target.
• The 3 basic-write tests, through createTestData.

What changes are included in this PR?

captureWritePlan in CometParquetWriterSuite now drains the listener bus with CometListenerBusUtils.waitUntilEmpty, once before it registers its listener and again after the write. It holds the plan in an AtomicReference and no longer polls for up to 15 s.

The only adaptation is the file. On main the helper lives in CometParquetWriterTestBase, which #5821 introduced and which is not on branch-1.0. The added and removed lines are identical to #6156's hunk. CometListenerBusUtils is already on branch-1.0.

How are these changes tested?

This changes a test helper only. I ran it locally on branch-1.0 with the default Spark 4.1 profile and JDK 17:

• All 30 tests in CometParquetWriterSuite pass. Spotless, scalastyle and scalafix (-Psemanticdb -Pspark-3.5 -Pscala-2.12) are clean.
• I reproduced the race on branch-1.0 with a throwaway test that isn't included here. It adds a SparkListener that sleeps 2 s on each SparkListenerSQLExecutionStart, so the shared listener queue falls behind the writes. It then runs the body of SaveMode.Append adds new files alongside existing data. With the old helper it fails with Expected exactly one CometNativeWriteExec in the plan, but found 0, because the captured plan is the Comet-disabled materializeAsCometSource write to source.parquet. With this change it passes.

The PR build runs CometParquetWriterSuite on all five Spark profiles on Linux and on macOS.

[sunchao]

Summary

• Prior state and problem: captureWritePlan could mistake a delayed setup-write callback for the write under test.
• Design approach: Drain Spark’s listener bus before registration and after the write, replacing the polling loop.
• Correctness / compatibility analysis: Spark sources for 3.4.3, 3.5.9, 4.0.4, 4.1.3 and 4.2.0 confirm that draining waits for queued events and callbacks in progress. AtomicReference safely publishes the captured plan across threads.
• Key design decisions: Reuses CometListenerBusUtils and preserves listener removal in finally, without introducing another synchronization abstraction.
• Implementation sketch: Drain, register, write, drain, read the captured plan, then unregister.
• Behavioral changes worth calling out: Waiting now uses Spark’s bounded 10-second drain instead of the 15-second polling loop. Synchronization overhead is confined to tests. Production execution and SQL semantics are unchanged.
• Suggested improvements: None at P1/P2 severity. No introduced P1/P2 issues found within this review.

Reviewed the entire diff from 5ffd67464f974b56a8e8b9229edebc444d1057c3 to 6ee7c6a4ddc2fea69adf3b1cd6b3636fc12a9b34. The PR is not a draft. Existing reviews, comments and threads were empty.

Routed skills: review-comet-pr, loaded from another local repository checkout because this release-branch head lacks .ai/skills. No subsystem sibling applies.

Validation: A bounded Spark 4.1.3/JDK 21 reproduction using the base and head helpers reproduced stale-plan capture at base and captured the intended write at head. git diff --check passed. The full Comet suite was not run locally because this checkout has no built native/JVM artifacts. Other Spark versions received source-level validation.

Exact-head CI: Preflight, change detection, title and labeling checks passed. Three lint jobs remained queued. Ten conditional checks were skipped. No failures were reported, but the build/test matrix had not completed.