Skip to content

[SPARK-56505][SQL][TESTS] Add SessionQueryTest to replace SharedSparkSession#56190

Open
fwc wants to merge 45 commits into
apache:masterfrom
fwc:sharedsparksession-refactor-mostly-nonbreaking
Open

[SPARK-56505][SQL][TESTS] Add SessionQueryTest to replace SharedSparkSession#56190
fwc wants to merge 45 commits into
apache:masterfrom
fwc:sharedsparksession-refactor-mostly-nonbreaking

Conversation

@fwc

@fwc fwc commented May 28, 2026

Copy link
Copy Markdown

What changes were proposed in this pull request?

This PR adds SessionQueryTest as the new, shiny connect-compatible test base trait to facilitate code-reuse of sql/core tests in sql/connect by making it possible to add 'connect variants'

If some FooSuite in sql/core uses the new sql.SessionQueryTest trait like e.g.

class FooSuite extends SessionQueryTest {
  checkAnswer(
    sql("SELECT 1"),
    Seq(1)
  )
}

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

class FooWithConnectSuite extends FooSuite with connect.SessionQueryTest

To make this possible, we factor outcheckAnswer and with{Table,View,Udf} from QueryTest into separate helpers and introduce a handful of helper traits for separation of concerns.

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.

Also c.f. #56190 (comment) for more background and motivation.

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

@fwc fwc force-pushed the sharedsparksession-refactor-mostly-nonbreaking branch from b7ba3f5 to 4c35b22 Compare May 28, 2026 20:54

@cloud-fan cloud-fan left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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

}

class QueryTestSuite extends test.SharedSparkSession {
class QueryTestSuite extends QueryTest with SparkSessionBinder {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

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?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

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.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

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.

Comment thread sql/connect/server/src/test/scala/org/apache/spark/sql/connect/QueryTest.scala Outdated
Comment thread sql/core/src/test/scala/org/apache/spark/sql/SparkSessionBinder.scala Outdated
Comment thread sql/core/src/test/scala/org/apache/spark/sql/test/SharedSparkSession.scala Outdated
@fwc

fwc commented May 29, 2026

Copy link
Copy Markdown
Author

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?

@fwc fwc requested a review from cloud-fan May 29, 2026 22:43
@cloud-fan

Copy link
Copy Markdown
Contributor

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

@fwc fwc force-pushed the sharedsparksession-refactor-mostly-nonbreaking branch from bd2e2e1 to 6e81376 Compare June 9, 2026 20:47

@cloud-fan cloud-fan left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

4 addressed, 2 remaining, 16 new. (16 = 15 newly introduced, 1 late catch — my miss from the prior round.)
The restructure moves in the agreed direction, but CI is red and the Connect path this PR exists to enable is broken end-to-end on this commit.

Design / architecture (7)

  • sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala:160: QueryTestBase.checkAnswer wins the linearization and routes Connect runs through classic-only queryExecution/materializedRdd — see inline
  • sql/connect/server/src/test/scala/org/apache/spark/sql/connect/DataSourceV2DataFrameConnectSuite.scala:37: suite-wide session drops per-test server-state isolation (TABLE_OR_VIEW_ALREADY_EXISTS in CI) — see inline
  • sql/connect/server/src/test/scala/org/apache/spark/sql/connect/DataSourceV2DataFrameConnectSuite.scala:60: classicSpark's catalogManager is not the server session's — external mutations invisible to Connect queries — see inline
  • sql/core/src/test/scala/org/apache/spark/sql/ExampleSuite.scala:32: uses the Delta format, which doesn't exist in Apache Spark — permanently failing — see inline
  • sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala:793: deprecated object keeps forked copies of the comparison logic it now inherits — see inline
  • sql/connect/server/src/test/scala/org/apache/spark/sql/connect/SessionQueryTest.scala:35: isDfSorted = false silently weakens ORDER BY assertions on Connect — see inline
  • sql/core/src/test/scala/org/apache/spark/sql/SparkSessionBinder.scala:63: Base/Binder layering undocumented; add scaladoc + TODO for merging the binder into SessionQueryTest once SharedSparkSession migrates — see inline

Correctness (3)

  • sql/core/src/test/scala/org/apache/spark/sql/CheckAnswerHelper.scala:67: failure messages interpolate df.queryExecution — failing Connect tests report UNSUPPORTED_CONNECT_FEATURE instead of the mismatch — see inline
  • sql/hive/src/test/scala/org/apache/spark/sql/hive/SessionQueryTest.scala:22: hive variant stacks two competing session lifecycles and nothing extends it — see inline
  • (general) The mllib/graphx/examples and streaming/kafka module jobs are also red on this commit, and GitHub's 50-annotation cap hides their failures (the visible annotations are all ExampleSuite plus the connect suite). Are these fallout from the QueryTest/binder restructure those modules consume via the sql test-jar? Worth root-causing before the next push.

Suggestions (1)

  • sql/core/src/test/scala/org/apache/spark/sql/CheckAnswerHelper.scala:35: by-name df is forced outside any try, bypassing the formatted error path — see inline

Nits: 6 minor items (see inline comments), one of which fails the linter job (missing newline at EOF in sql/core/.../sql/SessionQueryTest.scala).

PR description suggestions

  • Fix: the "user-facing change" section claims SharedSparkSessionBase now runs the thread audit — in the current structure it does not (the audit lives in sql.SparkSessionBinder, which is self-typed to SparkFunSuite; AnyFunSpec-based users of SharedSparkSessionBase structurally can't receive it). Stale text from an earlier revision.
  • Add: the DSv2 connector trait refactor (removal of withTestSession/checkRows/withTestTableAndViews) and the DataSourceV2DataFrameConnectSuite rebase off SparkConnectServerTest — the largest part of the diff is unmentioned.
  • Add: the ExampleSuite/ExampleConnectSuite and ParquetQuerySuite changes (or drop the example suites).
  • Document: the deprecations beyond SharedSparkSession (the QueryTest object, checkAggregatesWithTol, stripSparkFilter, logicalPlanToSparkQuery, makeQualifiedPath, withCurrentCatalogAndNamespace).

Comment thread sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala
*/
class DataSourceV2DataFrameConnectSuite
extends SparkConnectServerTest
extends SessionQueryTest

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Switching from SparkConnectServerTest to a suite-wide session drops the per-test isolation the old base provided (afterEach -> invalidateAllSessions()), and that's what the TABLE_OR_VIEW_ALREADY_EXISTS CI failures are: server-side catalog state now leaks across tests. The classic peer compensates with after { cachingcat.clearCache(); catalogManager.reset() } (DataSourceV2DataFrameSuite.scala:86-89); this suite has no counterpart. I'd recreate the client session per test and invalidate server sessions in afterEach — that restores the isolation the caching-catalog tests are written against; failing that, add the same per-test catalog cleanup the classic suite uses.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

Is there a clear 'ideal way' for how which session semantics to provide? Per-testcase would imply quite a bit of overhead, right?

Maybe we could leave guarantees around sessions explicitly 'undefined' to for potential future optimizations?
(e.g. parallel execution of testscases in different session)


val df = spark.createDataFrame(data).toDF("id", "name", "age")

df.write.partitionBy("age").format("delta").saveAsTable("foo")

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

delta isn't a data source in Apache Spark — this suite fails every CI run with DATA_SOURCE_NOT_FOUND (and replaceWhere is Delta-specific too). This looks like a local experiment that got committed. If you want a template suite demonstrating the pattern, rewrite it against a built-in source or the in-memory v2 catalog with assertions that hold here; otherwise drop both ExampleSuite and ExampleConnectSuite.


object QueryTest extends Assertions {
@deprecated("superseded by CheckAnswerHelper", since = "4.2")
object QueryTest extends CheckAnswerHelper {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Now that the object extends CheckAnswerHelper, it carries two parallel copies of the comparison logic — the helper's private ones and the public getErrorMessageInCheckAnswer/prepareAnswer/sameRows kept below. The public surface is still called (e.g. AggregationQuerySuite.scala:1089), so it can't be deleted, but the kept methods could delegate to the helper's internals so the logic exists once. Worth doing while the file is being restructured — drift between the two copies will be invisible.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

the kept methods could delegate to the helper's internals so the logic exists once

But then we couldn't make the helpers private in CheckAnswerHelper, right?

Comment thread sql/core/src/test/scala/org/apache/spark/sql/SessionQueryTestBase.scala Outdated
Comment thread sql/core/src/test/scala/org/apache/spark/sql/SessionQueryTest.scala Outdated
Comment thread sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala Outdated
Comment thread sql/core/src/test/scala/org/apache/spark/sql/SparkSessionBinder.scala Outdated
Comment thread sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala Outdated
@fwc fwc changed the title [SPARK-56505][SQL][TESTS] Add SparkSessionBinder to replace SharedSparkSession [SPARK-56505][SQL][TESTS] Add SessionQueryTest to replace SharedSparkSession Jun 12, 2026

@ionagamed ionagamed 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.

Did a drive-by pass, mainly nits and questions.

Comment thread core/src/test/scala/org/apache/spark/CheckErrorHelper.scala
*
* // no need to extend FooSuite as sql.SessionQueryTest
* // already executes shared tests via classic internally.
* class FooClassicSuite extends classic.SessionQueryTest {

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

QQ: What should a test author do when they want to reuse some common utilities from FooSuite in their FooClassicSuite? FooBaseSuite should work fine, but I think it might be worth mentioning the recommended way to handle these situations.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

Mhhhhhmaybe we could have a FooClassicSuite ignore all tests that are already executed in FooSuite with some hackery 🤔

Comment thread sql/core/src/test/scala/org/apache/spark/sql/SessionQueryTest.scala

@cloud-fan cloud-fan left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Round 3: 3 addressed, 8 remaining, 0 new.
The restructure is cleaner and resolved the per-test isolation + catalog-access items, but CI is still red on the same two round-2 blockers, and the Connect path this PR exists to enable is broken end-to-end on this commit (2958be0). All findings below are remaining from the prior review and already have inline threads, so I'm not re-posting them inline.

Still blocking (remaining)

  • QueryTest.scala:160: DSv2ExternalMutationTestBase extends QueryTest with SessionQueryTestBase; QueryTestBase.checkAnswer (classic-only df.queryExecution/materializedRdd) still wins the linearization in the Connect DSv2 suite -> CI UNSUPPORTED_CONNECT_FEATURE.DATASET_QUERY_EXECUTION.
  • ExampleSuite.scala:32: format("delta") -> CI DATA_SOURCE_NOT_FOUND: delta.
  • CheckAnswerHelper.scala:78: failure messages interpolate df.queryExecution, so a failing Connect checkAnswer reports the unsupported-API error instead of the row mismatch.

Still open, non-blocking (remaining)

  • QueryTest.scala:793: object QueryTest keeps forked copies of the comparison logic it now inherits from CheckAnswerHelper.
  • connect/SessionQueryTest.scala:35: isDfSorted=false silently weakens ORDER BY assertions on Connect (no scaladoc/JIRA on the TODO).
  • hive/SessionQueryTest.scala:22: stacks two competing session lifecycles and nothing extends it (CI never exercises it).
  • SparkSessionBinder.scala:63: Base/Binder layering still undocumented; add scaladoc + cleanup TODO.
  • CheckAnswerHelper.scala:42: by-name df forced outside any try, bypassing the formatted error path.

Addressed since round 2

  • Connect per-test isolation restored (fresh client + isolated server session per test); TABLE_OR_VIEW_ALREADY_EXISTS gone from CI.
  • getServerSession ported and used for catalog access.
  • Nits: @deprecated since, valid doc example, EOF newline.

Verification

Confirmed against live test-report annotations on head commit 2958be0: the Connect DSv2 suite routes checkAnswer through QueryTestBase.checkAnswer because the DSv2 base mixes in QueryTest, and that override wins the linearization over CheckAnswerHelper.checkAnswer -> UNSUPPORTED_CONNECT_FEATURE.DATASET_QUERY_EXECUTION (repeated), alongside ExampleSuite -> DATA_SOURCE_NOT_FOUND: delta. The connect-module, "sql - other tests", and linter jobs are all red.

PR description suggestions

  • Fix: the "user-facing change" section claims SharedSparkSessionBase runs the thread audit — it doesn't. The audit lives in SparkSessionBinder (self: SparkFunSuite); SparkSessionBinderBase (self: Suite, the base of SharedSparkSessionBase) structurally can't receive it.
  • Add: the DSv2 connector trait refactor (removal of withTestSession/checkRows/withTestTableAndViews) and the DataSourceV2DataFrameConnectSuite/ParquetQuerySuite rebases — the largest part of the diff.
  • Add or drop: the ExampleSuite/ExampleConnectSuite change.
  • Document: the deprecations beyond SharedSparkSession (the QueryTest object, checkAggregatesWithTol, stripSparkFilter, logicalPlanToSparkQuery, makeQualifiedPath, withCurrentCatalogAndNamespace).

@johanl-db johanl-db left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Doing a pass, I agree with the general direction

Comment thread sql/core/src/test/scala/org/apache/spark/sql/QueryTest.scala
with SessionQueryTestBase
with SparkSessionBinder {

override def sessionType: String = "classic"

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

What's the goal of this sessionType?
SessionQueryTest returns an agnostic session. If callers really want to know what type of session they have, they can use reflection, but it seems to me we should discourage callers poking into the session. If you want a classic session, get an actual one via classic.SessionQueryTest

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

With this, tests can handle and document session-specific behaviour:

   test(...) {
     val df = // query with connect-specific behaviour
     if (sessionType == "connect") {
       checkError(...)
     } else {
       checkAnswer(df, ...)
     }
   }

This is explained in the docstring in SessionQueryTestBase.

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

Agreeing with:

If callers really want to know what type of session they have, they can use reflection

This seems to go against the spirit of writing session-agnostic tests.

The suggested test might be implemented cleaner as:

class FooClassicSuite
    extends classic.SessionQueryTest {
  test(...) {
    val df = // query with connect-specific-behavior
    checkError(...)
  }
}

class FooConnectSuite
    extends connect.SessionQueryTest {
  test(...) {
    val df = // query with connect-specific-behavior
    checkAnswer(df, ...)
  }
}

or even:

class FooBaseSuite {
  def bar() = // some meaningful actions isolated as a function
}

class FooClassicSuite
    extends FooBaseSuite
    with classic.SessionQueryTest {
  test(...) {
    bar()
    checkError(...)
  }
}

class FooConnectSuite
    extends FooBaseSuite
    with connect.SessionQueryTest {
  test(...) {
    bar()
    checkAnswer(df, ...)
  }
}


import org.apache.spark.sql

class ExampleConnectSuite extends sql.ExampleSuite with SessionQueryTest

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

I like the example that Wenchen gave in #56190 (comment)

Can you colocate the example suites and show how one would write tests that classic only, connect only + shared tests?

A single page example showing everything will also help sharing knowledge of how to use this

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

A single page example showing everything will also help sharing knowledge of how to use this

Wenchen suggested editing the AGENTS.md to document this. Furthermore, we could link there from the Contributing to Spark guide or even write a 'How to write Tests for Spark SQL / DataFrame stuff' thingy.

I also intend to post to the dev mailing list once this PR has landed.

* @param expectedAnswer the expected result in a [[Seq]] of [[Row]]s.
* @param absTol the absolute tolerance between actual and expected answers.
*/
@deprecated("rarely used", since = "4.2.0")

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Probably not a good enough reason to deprecate it.

That being said, if it's not really used, it's ok to now provide an equivalent in connect, meaning the few tests using this won't work in connect unless someone put some cycles adding an equivalent method

with SparkSessionProvider
with CheckAnswerHelper
with CheckErrorHelper
with SQLConfHelper

@fwc fwc Jun 22, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

@cloud-fan @johanl-db @ionagamed I'm a bit unsure how we should handle config stuff:

currently, withSQLConf is used pervasively but SQLConf is not part of the public API, only RuntimeConfig.

I see the following options:

  • keep the withSQLConf name but implement the connect override by manipulating spark.conf ('only' has type RuntimeConfig) (would be confusing)
  • introduce withRuntimeConf as the 'new' way to set confs while deprecating withSQLConf at for 'advertisement'

Are there other options? What do you think?

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

AFAICS, even while using the connect path withSQLConf behaves the same way as always (at least in spirit).

Do you suggest withRuntimeConf to have the same shape and behavior as withSQLConf?

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

I'd add a single agnostic withConf(key -> value) backed by spark.conf (RuntimeConfig). It works in both modes and covers every non-static conf (public and internal), since classic and Connect both ultimately call setConfString on the session's SQLConf. Tests keep referencing the typed key (withConf(SQLConf.SOME_CONF.key -> "v")) — that's just reading a constant.

I would not reimplement withSQLConf over spark.conf (option 1): overloading the name silently narrows its power (a test setting an internal/thread-local conf would behave differently), which is the kind of silent divergence we want to avoid. Keep withSQLConf as a classic-only escape hatch — it's inherently SQLConf-typed — alongside the other classic-only helpers.

Static confs aren't settable by either withSQLConf or withConf and stay on the sparkConf override (which already reaches the Connect server session, since the in-process server runs on the classic SparkContext), so withConf loses nothing there.

Two impl notes: (1) withConf must save+restore the prior value including the unset case, or conf state leaks across tests; (2) worth a quick check that Connect's config service doesn't filter .internal() keys — if it does, those few tests become classic-only, which is an honest constraint rather than a gap to paper over.

doThreadPostAudit()
}
}
@deprecated("Use SessionQueryTest (or classic.SessionQueryTest if required) instead", "4.2.0")

@fwc fwc Jun 22, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

@cloud-fan @ionagamed @johanl-db I'm unsure wether we should immediately deprecate SharedSparkSession or only do that in e.g. ~1 month after we spent some time using SessionQueryTest ourselves.

Maybe it's better if we find potential 'UX gaps' ourselves 😅

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

I'd decouple advertising from deprecating. Deprecating SharedSparkSession now fires a warning across ~500 call sites and points them at a base whose Connect path is still red on this commit (the CheckAnswerHelper.isDfSorted throw breaks the binder-swap demo QueryTestWithConnectSuite), with UX gaps still open (the config helper above, ordering on Connect). That's a lot of churn toward a moving target.

So: land the new traits now but drop the @deprecated on SharedSparkSession for now, and advertise SessionQueryTest via the AGENTS.md / test-style-doc update — new suites get steered to the new way immediately without deprecating anything. Flip @deprecated in a follow-up once we've dogfooded and closed the gaps. The gating signal is 'gaps closed + docs landed', not a fixed month — which also matches your instinct to find the UX gaps ourselves first.

@cloud-fan cloud-fan left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

6 addressed, 2 remaining, 3 new. (3 new = 1 newly introduced, 2 late catches — my own misses from earlier rounds.)

Round 3's Connect-DSv2 blocker, the df.queryExecution interpolation, the analysis-time catch, the layering scaladoc, and the hive variant are all resolved. CI is still red on ef19ed4, though — one remaining blocker and one I should have caught earlier. (Replies to your two open questions on config and deprecation timing are on those threads.)

Design / architecture (2)

  • CheckAnswerHelper/connect.SessionQueryTest: "connect" is split across the binder (session) and SessionQueryTest (isDfSorted/sessionType), so it isn't a single mixin — the root cause of the isDfSorted failure below, and a recurring bug class (any future env-seam added to one trait but not the other half-works). Building on your binder-free-base thread (#56190 (comment)): the isDfSorted fix below already consolidates connect behavior into one mixin; from there the near-empty classic.* leaves could collapse into a named classicOnly { cs => … } escape hatch so SessionQueryTest is the single base writers pick, and sessionType could become a sealed SessionType (Classic/Connect) behind onClassic/onConnect guards rather than a String.
  • QueryTest.scala:799: object QueryTest keeps forked copies of the comparison logic it now inherits from CheckAnswerHelper (existing thread; non-blocking).

Correctness (2)

  • CheckAnswerHelper.scala:74: isDfSorted throws on Connect DataFrames; the binder-swap idiom (QueryTestWithConnectSuite) is red in CI — the QueryTest.scala:160 fix didn't reach this method — see inline (new, late catch).
  • ExampleSuite.scala:32: format("delta") is not a built-in source — ExampleSuite/ExampleConnectSuite still fail with DATA_SOURCE_NOT_FOUND (existing thread; remaining).

Suggestions (1)

  • QueryTest.scala:180: df.materializedRdd re-evaluates the by-name df, so the plan is built twice per classic checkAnswer — see inline (new).

Verification

Confirmed against live test-report annotations on head ef19ed4 (check-run 82813492697): QueryTestWithConnectSuite fails because checkAnswerCheckAnswerHelper.isDfSorted throws RuntimeException on the Connect DataFrame (the isDfSorted = false override is only in connect.SessionQueryTest, unreachable via connect.SparkSessionBinder); ExampleSuite/ExampleConnectSuite fail with DATA_SOURCE_NOT_FOUND: delta. The round-3 UNSUPPORTED_CONNECT_FEATURE.DATASET_QUERY_EXECUTION is gone.

PR description suggestions

  • Fix: the "user-facing change" section claims SharedSparkSessionBase now runs the thread audit — it doesn't. SharedSparkSessionBase mixes in SparkSessionBinderBase (self: Suite), which has no audit; the audit lives in sql.SparkSessionBinder (self: SparkFunSuite), reached by SharedSparkSession via classic.SparkSessionBinder, not by ...Base.
  • Add: the DSv2 connector trait refactor (removal of withTestSession/checkRows/withTestTableAndViews) and the DataSourceV2DataFrameConnectSuite/ParquetQuerySuite rebases — the largest part of the diff.
  • Add or drop: the ExampleSuite/ExampleConnectSuite change.
  • Document: the deprecations beyond SharedSparkSession (the QueryTest object, checkAggregatesWithTol, stripSparkFilter, logicalPlanToSparkQuery, makeQualifiedPath, withCurrentCatalogAndNamespace).

df match {
case df: classic.DataFrame =>
df.logicalPlan.collectFirst { case s: logical.Sort => s }.nonEmpty
case _ => throw new RuntimeException(s"Cannot determine whether df is sorted: $df")

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

checkAnswer here flows to getErrorMessageInCheckAnswersameRows(..., isDfSorted(df)), and this default throws for any non-classic.DataFrame. The connect-safe override (= false) lives only in connect.SessionQueryTest, so the documented FooSuite with connect.SparkSessionBinder swap — what QueryTestWithConnectSuite does — hits this throw on every Connect checkAnswer, and that suite is red in CI for exactly this reason. Same trait-linearization issue raised on QueryTest.scala:160, in a method that fix didn't reach.

Returning false here (matching what connect.SessionQueryTest already does deliberately) makes the helper connect-safe by default, fixes the binder-swap path, and makes the connect.SessionQueryTest.isDfSorted override redundant — i.e. it consolidates the connect behavior into one mixin.

Suggested change
case _ => throw new RuntimeException(s"Cannot determine whether df is sorted: $df")
case _ => false


QueryTest.checkAnswer(analyzedDF, expectedAnswer)
SQLExecution.withSQLConfPropagated(analyzedDF.sparkSession) {
df.materializedRdd.count() // Also attempt to deserialize as an RDD [SPARK-15791]

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

df is the by-name parameter (df: => DataFrame), already forced once into analyzedDF at the top of this method. df.materializedRdd re-evaluates the thunk, so the query plan is constructed twice for every classic checkAnswer (and a side-effecting by-name block would run twice). Before this PR the trait evaluated df exactly once. Use the already-forced value:

Suggested change
df.materializedRdd.count() // Also attempt to deserialize as an RDD [SPARK-15791]
analyzedDF.materializedRdd.count() // Also attempt to deserialize as an RDD [SPARK-15791]

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.

4 participants