[SPARK-56505][SQL][TESTS] Add SparkSessionBinder to replace SharedSparkSession

[fwc]

What changes were proposed in this pull request?

• Introduces sql.SparkSessionBinder and classic.SparkSessionBinder as 'implementor' counterparts to the 'declarators' sql.SparkSessionProvider and classic.SparkSessionProvider.
• Deprecates SharedSparkSession with the hint that SparkSessionBinder and QueryTest shall be used instead.

Why are the changes needed?

Currently, most tests use SharedSparkSession to obtain the spark object. This prevents specializing these tests in sql/connect as SharedSparkSession provides a classic.SparkSession, thus preventing overriding.

This PR deprecates SharedSparkSession and instead introduces sql.SparkSessionBinder and classic.SparkSessionBinder. While both create a classic.SparkSession, the sql.SparkSessionBinderhas an abstract def spark: sql.SparkSession declaration, so it can we overriden with some trait that provides a connect.SparkSession.

If some FooSuite now uses the sql.SparkSessionBinder trait like e.g.

class FooSuite extends SparkSessionBinder with QueryTest {
  checkAnswer(
    sql("SELECT 1"),
    Seq(1)
  )
}

We can now add a connect variant of that suite as follows:

class FooWithConnectSuite extends FooSuite
  with connect.SparkSessionBinder
  with connect.QueryTest

Does this PR introduce any user-facing change?

This PR extends the beforeAll/afterAll of SharedSparkSessionBase to include the the thread audit check, which was previously only present in SharedSparkSession.
AFAICS, SparkSessionBase is neither used in delta lake nor in apache iceberg.

How was this patch tested?

This patch is test-only.

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

Parts of this patch were authored by claude code

[cloud-fan]

1 blocking, 2 non-blocking, 3 nits.
Right direction — decoupling the session type so suites can run on classic or Connect. My main feedback is on the author-facing shape: I'd push for a binder-free base + per-env concrete suites, with the bare SparkSessionBinder kept internal.

Design / architecture (1)

• sql/core/.../sql/QueryTest.scala:1214: push the binder-free-base + classic/connect-concrete pattern; treat bare SparkSessionBinder as internal — see inline

Suggestions (2)

• sql/connect/.../connect/SparkSessionBinder.scala:89: redundant afterEach override with an inaccurate comment — see inline
• sql/connect/.../connect/QueryTest.scala:30: only one checkAnswer overload overridden — see inline

Nits: 3 minor items (see inline comments).

[cloud-fan]

This migration — mixing the bare sql.SparkSessionBinder into a concrete suite — is the shape I'd push back on. sql.SparkSessionBinder binds a classic session but exposes spark only as the abstract sql.SparkSession, so it's really internal plumbing, not what a test author should reach for.

The end-state I'd recommend documenting and demonstrating is a binder-free base + per-env concrete suites:

abstract class FooSuiteBase extends QueryTest {          // no binder; spark abstract
  test("shared") { checkAnswer(sql("SELECT 1"), Row(1)) }
}
class FooSuite extends FooSuiteBase with classic.SparkSessionBinder {
  test("classic only") { ... }
}
class FooConnectSuite extends FooSuiteBase
  with connect.SparkSessionBinder with connect.QueryTest {
  test("connect only") { ... }
}

QueryTest already mixes in SparkSessionProvider (via SQLTestData) and leaves spark abstract, so it works as the env-agnostic base directly. Concretely: (1) steer the migration and the @deprecated message at classic.SparkSessionBinder / connect.SparkSessionBinder + this base pattern, not the bare binder; (2) QueryTestWithConnectSuite currently demonstrates the retrofit path (extending an already-classic-bound QueryTestSuite and overriding the binding) — a binder-free base would demonstrate the cleaner pattern and double as the template authors copy.

[fwc]

I want to nudge test authors towards writing (somewhat) connect-compatible tests by default, which is why I want them to write tests with a sql.SparkSession in hand.

My fear is that the 'clean' way is not the 'easiest' way. Most current tests do not use an abstract base class and I fear that most test authors will default to just start a new suite with classic.SparkSessionBinder as they might not think about connect in that moment:

// hypothetical antipattern, but path of least resistance:
class FooSuite extends QueryTest with classic.SparkSessionBinder {
  test("all tests, both shared and classic only") { ... }
}

I reworked the PR so that SparkSessionBinder now implements QueryTest. Now classic.SparkSessionBinder is a drop-in replacement for SharedSparkSession and sql.SparkSessionBinder provides the new, 'fixed' default.

What do you think of this approach?

[cloud-fan]

Given we are setting up the standard test style, can you spell out your proposal clearly? If you don't agree with my abstract test suite proposal, what's your proposal? Can you write down example test suites so that people can understand you easier?

[cloud-fan]

And AGENTS.md already have a section about how to add test suites, we need to update it to whatever final approach people agree on.

[fwc]

For scala tests in sql/core, I envision a test style (and corresponding infra) that encourages:

1. Testing at the Spark SQL/DataFrame API level (i.e. discourages/prevents accessing internals)
2. Targeting both classic and connect
3. A rather DAMP than DRY style (i.e. tests as spark and df ops + checks with little/no 'abstraction')

While DAMP, not DRY would mostly be enshrined in guidelines like AGENTS.md, a testing style guide, or a sql/core/src/test/README.md, I think that testing the API using classic and connect can be made easier by further tweaking and extending the testing APIs / traits.

In some sense, this 'proposal' isn't proposing something 'new': Many (most?) suites are already more DAMP than DRY and with the introduction of sql.SparkSessionProvider and its usage in QueryTest, writing classic/connect-agnostic tests is already much easier than before. But (a) afaics DAMP NOT DRY is not (yet) codified anywhere and (b) I think the testing traits can be further improved to make 'doing the right thing' even 'more easier'.

So, I propose the following:

• Introduce a fully classic/connect-agnostic SparkSessionTest as the default SQL/DataFrame API testing trait
• Encourage the addition of 'connect variants' by e.g. adding a linter that warns/fails if suite X implements sql.SparkSessionTest but there is no corresponding class Y extends X with connect.SparkSessionTest or (if possible), automagically generating the connect variant.
• Discourage access to internals for newly added suites by e.g. adding them in subpackage of o.a.s.sql.test or maybe even in a separate compilation unit that only has access to the scala API and test helpers.

A prototypical Example

This example aims to demonstrate how a suite that adheres to this style and uses SparkSessionTest could look like.
This example is inspired by #55571.

The main suite

A prototypical suite extends SparkSessionTest, which provides spark and useful helpers like checkAnswer.

The class defines both the necessary setup (e.g. creating tables, setting confs) and corresponding test cases. The test cases generally consist of Spark SQL/DataFrame operations and assertions (either literal asserts or e.g. checkAnswer).

The example suite is not located in the same package as the code under test to discourage accessing internals. In this example, internals are accessed using a helper object from the package of the code under test.

// located in a different package than the code that is tested
package org.apache.spark.sql.test.tablecache

import org.apache.spark.sql.tablecache.TableCacheHelper
// ...

class TableCacheSuite extends SparkSessionTest {
  override def beforeAll = {
    spark.sql(s"CREATE TABLE $testTable (id INT, salary INT) USING mockedDSv2")
    spark.sql(s"INSERT INTO $testTable VALUES (1, 100), (10, 1000)")
  }
  override def afterAll  = {
    spark.sql(s"DROP TABLE $testTable")
  }

  override def sparkConf: SparkConf =
    super.sparkConf.set( /* set necessary confs */ )


  // the test mostly consists of DataFrame operations,
  // util calls like `checkAnswer`, and asserts

  test("SPARK-54022: cached table pinned against external data write") {
    spark.table(testTable).cache()
    assert(spark.catalog.isCached(testTable))
    checkAnswer(spark.table(testTable), Seq(Row(1, 100)))

    // object that accesses internals
    TableCacheHelper.externalAppend(testTable, Row(2, 200))

    checkAnswer(spark.table(testTable), Seq(Row(1, 100)))

    spark.sql(s"REFRESH TABLE $testTable")
    checkAnswer(spark.table(testTable), Seq(Row(1, 100), Row(2, 200)))

    assert(spark.catalog.isCached(testTable))
  }

  test("connector w/ cache: temp view stale after external column removal") {
    withView("v") {
      spark.table(testTable).filter("salary < 999").createOrReplaceTempView("v")
      checkAnswer(spark.table("v"), Seq(Row(1, 100)))

      TableCacheHelper.externalDropCol(testTable, "salary")

      checkAnswer(spark.table("v"), Seq(Row(1, 100)))

      spark.sql(s"REFRESH TABLE $testTable").collect()
      checkError(
        exception = intercept[AnalysisException] { spark.table("v").collect() },
        condition = "INCOMPATIBLE_COLUMN_CHANGES_AFTER_VIEW_WITH_PLAN_CREATION",
      )
    }
  }
}

The 'connect variant'

Besides the main suite, there is a 'connect variant', which runs the same test via Spark Connect:

package org.apache.spark.sql.connect.test.tablecache

class TableCacheConnectSuite extends TableCacheSuite with SparkSessionTest

Ideally, this variant would be auto-generated, so that classic/connect-agnostic testing is the default (i.e. opt-out), rather than something that the developer actively has to strive for.

The SparkSessionTest helper trait

The SparkSessionTest used would be the default util trait for SparkSession/DataFrame-level tests, providing a sql.SparkSession and utils like checkAnswer and withTable.

The SparkSessionTest trait is designed to be 'classic/connect-agnostic': it only provides utils/APIs that can be used/sensibly overriden in Spark Connect, so that we can have a connect.SparkSessionTest sibling trait to facilitate the implementation of the 'connect variant'.

package org.apache.spark.sql

trait SparkSessionTest extends SparkFunSuite {
  def spark: SparkSession = ...
  def sql: ...

  def checkAnswer(df: DataFrame, exp: Seq[Row]) = ...
  def checkError(...) = ...

  def withTable(...) = ...
  def withView(...) = ...
  def withTempDir(...) = ...

  // TODO can checks on the plan be classic/connect-agnostic?
}

This trait aims to supercede SharedSparkSession and QueryTest. They shall be deprecated with the suggestion to migrate towards sql.SparkSessionTest.

A classic.SparkSessionTest can also be added for classic-only tests.

What's wrong with SharedSparkSession and QueryTest?

I want to deprecate SharedSparkSession because

• It cannot be overriden with a connect.SharedSparkSession so classic/connect-agnostic tests would require that testcases are declared in traits/abstract classes.
• It provides a classic.SparkSession while currently being the most widely used test trait and the easiest thing to reach for. Changing it would break downstream (e.g. tests in Delta Lake), so I think its usage needs to be actively discouraged.
• It is widely used so the deprecation note will be widely read and can be used to advertise the 'new way'.

I want to deprecate QueryTest because

• it is not fully connect/classic-agnostic.
• it provides too many peculiar/specific/rarely used methods.
• it provides implicits.

On defaults, simplicity, and friction

I propose SparkSessionTest as the default testing trait to simplify things a bit: smaller than QueryTest but still a one-stop-shop.

Adding a new Suite should be as simple as possible: class X extends SparkSessionTest, some setup and damp testcases. No need to declare testcases in an abstract class or trait.

Structuring test code with inheritance: variant extends suite extends base

I suggest the following hierarchy for test classes/traits:

1. Base: provides generic utils for its suites, contains no testcases (here: SparkSessionTest)
2. Suite: contains test cases (here: TableCacheSuite)
3. Variants: contains overrides (here: TableCacheConnectSuite, which overrides spark)

While this proposal focuses on 'connect variants', I think the idea can be generalized to e.g. cover different configuration values like e.g. 'codegen on/off'.

---

Appendix: patterns that this proposal discourages

The following are exaggerations/caricatures of existing patterns that I believe
to be 'suboptimal'. This proposal aims to discourage such patterns.

Too much classic API usage

SharedSparkSession is by far the most widely used test trait (extended/mixed
in ~500 times, while QueryTest is used only ~150 times).

The problem: it provides a classic.SparkSession, which makes
'classic/connect-agnostic' testing difficult in two ways:
First, def spark cannot be overriden to use connect.
Second, usage of classic-only (read: connect-incompatible) APIs is not
discouraged.

class FooSuite extends SharedSparkSession {
  test("...") {
    // classic-only, should use `spark.catalog.setCurrentDatabase(db)`
    spark.sessionState.catalogManager.setCurrentNamespace(Array(db))
  }
}

With this proposal, SparkSessionTest hides the used classic.SparkSession
behind the sql interface so accessing spark.sessionState.catalogManager
would give a "Cannot resolve symbol catalogmanager" error in the IDE.

DRY to the bone

Generally, I do not like concisely abstracted tests like the caricature:

abstract class InsertSuite extends QueryTest with InsertMetricCheck {
  // creates table t in beforeEach, deletes in afterEach

  test("insert 3 rows") {
    val data = Seq((1L, "a"), (2L, "b"), (3L, "c"))
    doInsert(t, data)
    checkInsertMetrics(t, numInsertedRows = 3)
    verifyTable(t, data.toDf())
  }
}

// potentially in another file:
class InsertSQLSuite extends InsertSuite {
  def doInsert(t, d) = sql(s"INSERT INTO $t VALUES ${d.mkString}")
}
class InsertDfwSuite extends InsertSuite {
  def doInsert(t, d) = d.toDf().insert.write.insertInto(t)
}

I find such testcases unnecessarily hard to read: the logic is scattered over
across functions and files.

While I see the appeal in abstracting over the SQL- and DataFrame-APIs, I think
tests should not hide the concrete SQL- and DataFrame-operations.
(c.f. Tests and Code Sharing: Damp, Not Dry)

So I think the strawman should be refactored to the following:

class InsertSuite extends SparkSessionTest {
  test("insert 3 rows via SQL") {
    withTable("foo") {
      spark.sql("CREATE TABLE foo (id INT, name STRING)")
      spark.sql("INSERT INTO foo VALUES (1, 'a'), (2, 'b'), (3, 'c')")
      TableInternalsHelper.getCommits("foo").last match {
        case Insert(numRowsInserted) => assert(numRowsInserted === 3)
        case c => fail(s"expected Insert commit after inserting, but got $c")
      }
      checkAnswer(
        spark.table("foo"),
        Seq(Row(1, "a"), Row(2, "b"), Row(3, "c")),
      )
    }
  }

  test("insert 3 rows via dataframe writer API") {
    withTable("foo") {
      spark.sql("CREATE TABLE foo (id INT, name STRING)")
      spark.createDataFrame(Seq(Row(1, "a"), Row(2, "b"), Row(3, "c")))
        .write
        .insertInto("foo")
      TableInternalsHelper.getCommits("foo").last match {
        case Insert(numRowsInserted) => assert(numRowsInserted === 3)
        case c => fail(s"expected Insert commit after inserting, but got $c")
      }
      checkAnswer(
        spark.table("foo"),
        Seq(Row(1, "a"), Row(2, "b"), Row(3, "c")),
      )
    }
  }
}

Long inheritance chains / convoluted matrix tests

When Inspecting the Type Hierarchy of e.g. SharedSparkSession, there are quite
a few deep inheritance chains.

The following shows the chain of DeltaBasedMergeIntoTableSuiteBase. The ...Base classes are abstract.

RowLevelOperationSuiteBase                   0 cases  setup and helpers
.MergeIntoTableSuiteBase                    71 cases
..DeltaBasedMergeIntoTableSuiteBase         +7 cases  overrides expected metrics
...DeltaBasedMergeIntoTableSuite            +8 cases  sets 'supports-deltas'
...DeltaBasedMergeIntoTableUpdeAsDelete...  +1 case   sets 'supports-deltas', 'split-updates'
....DeltaBasedMergeIntoTableNoCodegenSuite  +0 cases  disables codegen via sparkConf

As an ideal, I think we should strive for 3 levels (Base, Suite, Variants), like e.g.:

RowLevelOperationSuiteBase
.MergeIntoTableSuite
..MergeIntoTableUpdeAsDeleteAndInsertSuite
..MergeIntoTableNoCodegenSuite
..MergeIntoTableDeltaBasedSuite
..MergeIntoTableDeltaBasedUpdeAsDeleteA...
..MergeIntoTableDeltaBasedNoCodegenSuite

Awkward classic/connect-agnostic tests

Some tests introdoce new helper constructs to achieve classic/connect-agnosticity.

  test("foo") {
    withTestSession { session =>
      session.sql(s"CREATE TABLE $testTable (id INT, salary INT) USING foo").collect()
      session.sql(s"INSERT INTO $testTable VALUES (1, 100)").collect()
      checkRows(session.sql(s"SELECT * FROM $testTable"), Seq(Row(1, 100)))
      checkRows(session.sql(s"SELECT * FROM $testTable"), Seq(Row(1, 100), Row(2, 200)))
    }
  }

This diff shows that these new helpers are not necessary.

[cloud-fan]

This comment is inaccurate after the refactor: connect.SparkSessionBinder extends sql.SparkSessionBinder directly (not SharedSparkSessionBase), and that parent's afterEach clears the cache via the private _spark — the classic session, which is exactly what's used on Connect (createSparkSession isn't overridden). So the parent's afterEach already works here and this override is redundant. If you do keep it, note that skipping super.afterEach() drops the BeforeAndAfterEach chain. Simplest fix is to remove the override entirely.

[fwc]

Hi @cloud-fan, I changed the PR so that SharedSparkSession is now an empty alias of classic.SparkSessionBinder with a deprecation note that recommends using sql.SparkSessionBinder if possible.

I am unsure with regards to the SparkSessionBinder name:

AFAICS SharedSparkSession is/was the testing trait (~500 extends/implements usages compared to ~150 usages of QueryTest, both according to Intellij's "Find Usages").
If it wouldn't be a breaking change, I'd want to rename QueryTest to QueryTestHelpers and SparkSessionBinder to QueryTest. What do you think? Maybe QuerySuiteUtils? Maybe SparkSessionTest?

[cloud-fan]

#56190 (comment) can you address this comment? I'd like to see the final test suite style we want to support.