Fix DataWritingCommand callback naver be called on withListener method

[wForget]

Why are the changes needed?

The DataWritingCommand callback naver be called on withListener method

How was this patch tested?

Was this patch authored or co-authored using generative AI tooling?

No

[wForget]

Duplicate #6793

[wForget]

DataWritingCommandExec is wrapped by AdaptiveSparkPlanExec, so we should use AdaptiveSparkPlanHelper#collect

[Copilot]

Pull request overview

This PR fixes a test utility (withListener) used by Spark SQL extension test suites to reliably observe DataWritingCommandExec by traversing the executed physical plan, and adds an option to allow queries that legitimately have no write command.

Changes:

• Update withListener to search the whole executed plan for DataWritingCommandExec and track whether the callback ran.
• Add a mustBeCalled flag (default true) to assert callback execution when expected.
• Update one rebalance-related test to opt out of the callback-must-run assertion.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
extensions/spark/kyuubi-extension-spark-4-1/src/test/scala/org/apache/spark/sql/KyuubiSparkSQLExtensionTest.scala	Makes withListener robust by scanning the executed plan and asserting callback execution by default.
extensions/spark/kyuubi-extension-spark-4-1/src/test/scala/org/apache/spark/sql/RebalanceBeforeWritingSuite.scala	Uses the new mustBeCalled=false mode in a test that includes non-write queries.

Comments suppressed due to low confidence (1)

extensions/spark/kyuubi-extension-spark-4-1/src/test/scala/org/apache/spark/sql/RebalanceBeforeWritingSuite.scala:116

• Setting mustBeCalled = false inside the shared check helper means this test will still pass even if the listener never finds a DataWritingCommandExec for INSERT queries (the callback won’t run, so the rebalance assertion won’t execute). To avoid masking regressions, consider using mustBeCalled=false only for the SELECT/COUNT cases that intentionally have no write command, and keep the default mustBeCalled=true for INSERT statements.

  test("check rebalance does not exists") {
    def check(df: DataFrame): Unit = {
      withListener(df, false) { write =>
        assert(write.collect {
          case r: RebalancePartitions => r
        }.isEmpty)
      }

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

You can also share your feedback on Copilot code review. Take the survey.