Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,10 @@ package org.apache.spark.sql.connect

import org.scalatest.time.SpanSugar._

import org.apache.spark.sql.{Row, SparkSession}
import org.apache.spark.sql.connect.config.Connect
import org.apache.spark.sql.connector.catalog.InMemoryTableCatalog

/**
* Test suite showcasing the APIs provided by SparkConnectServerTest trait.
*
Expand All @@ -29,6 +33,111 @@ import org.scalatest.time.SpanSugar._
*/
class SparkConnectServerTestSuite extends SparkConnectServerTest {

private def withTestTable(clientSession: SparkSession)(f: String => Unit): Unit = {
val table = "testcat.ns.tbl"
clientSession.conf.set(
"spark.sql.catalog.testcat",
classOf[InMemoryTableCatalog].getName)
clientSession.conf.set("spark.sql.catalog.testcat.copyOnLoad", "true")
clientSession.sql(s"CREATE TABLE $table (id INT) USING foo").collect()
clientSession.sql(s"INSERT INTO $table VALUES (1)").collect()

try {
f(table)
} finally {
clientSession.sql(s"DROP TABLE IF EXISTS $table").collect()
}
}

test("Spark Connect cached plan sees changes to a DSv2 table") {
withSession { clientSession =>
withTestTable(clientSession) { table =>
val cachedPlan = clientSession.table(table).drop("missing")
assert(cachedPlan.collect().toSeq == Seq(Row(1)))

getServerSession(clientSession).sql(s"INSERT INTO $table VALUES (2)").collect()

assert(cachedPlan.collect().toSet == Set(Row(1), Row(2)))
}
}
}

test("Spark Connect plan cached by analysis sees changes to a DSv2 table") {
withSession { clientSession =>
withTestTable(clientSession) { table =>
val cachedPlan = clientSession.table(table).drop("missing")
assert(cachedPlan.schema.fieldNames.toSeq == Seq("id"))

getServerSession(clientSession).sql(s"INSERT INTO $table VALUES (2)").collect()

assert(cachedPlan.collect().toSet == Set(Row(1), Row(2)))
}
}
}

test("Spark Connect cached DSv2 plan reused as a subtree sees table changes") {
withSession { clientSession =>
withTestTable(clientSession) { table =>
val cachedPlan = clientSession.table(table)
assert(cachedPlan.collect().toSeq == Seq(Row(1)))

getServerSession(clientSession).sql(s"INSERT INTO $table VALUES (2)").collect()

assert(cachedPlan.filter("id > 0").collect().toSet == Set(Row(1), Row(2)))
}
}
}

test("Spark Connect DSv2 plan sees table changes when the plan cache is disabled") {
withSession { clientSession =>
withTestTable(clientSession) { table =>
val cachedPlan = clientSession.table(table).drop("missing")
assert(cachedPlan.collect().toSeq == Seq(Row(1)))

val serverSession = getServerSession(clientSession)
serverSession.sql(s"INSERT INTO $table VALUES (2)").collect()
serverSession.sessionState.conf.setConf(Connect.CONNECT_SESSION_PLAN_CACHE_ENABLED, false)

try {
assert(cachedPlan.collect().toSet == Set(Row(1), Row(2)))
} finally {
serverSession.sessionState.conf.setConf(Connect.CONNECT_SESSION_PLAN_CACHE_ENABLED, true)
}
}
}
}

test("Spark Connect cached DSv2 plan sees a recreated table") {
withSession { clientSession =>
withTestTable(clientSession) { table =>
val cachedPlan = clientSession.table(table).drop("missing")
assert(cachedPlan.collect().toSeq == Seq(Row(1)))

val serverSession = getServerSession(clientSession)
serverSession.sql(s"DROP TABLE $table").collect()
serverSession.sql(s"CREATE TABLE $table (id INT) USING foo").collect()
serverSession.sql(s"INSERT INTO $table VALUES (2)").collect()

assert(cachedPlan.collect().toSeq == Seq(Row(2)))
}
}
}

test("Spark Connect cached DSv2 plan sees compatible schema evolution") {
withSession { clientSession =>
withTestTable(clientSession) { table =>
val cachedPlan = clientSession.table(table).drop("missing")
assert(cachedPlan.collect().toSeq == Seq(Row(1)))

val serverSession = getServerSession(clientSession)
serverSession.sql(s"ALTER TABLE $table ADD COLUMN value INT").collect()
serverSession.sql(s"INSERT INTO $table VALUES (2, 20)").collect()

assert(cachedPlan.collect().toSet == Set(Row(1, null), Row(2, 20)))
}
}
}

test("withSession: execute SQL and collect results") {
withSession { session =>
val df = session.sql("SELECT 1 as value")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,7 @@ class QueryExecution(
val tracker: QueryPlanningTracker = new QueryPlanningTracker,
val mode: CommandExecutionMode.Value = CommandExecutionMode.ALL,
val shuffleCleanupModeOpt: Option[ShuffleCleanupMode] = None,
val refreshPhaseEnabled: Boolean = true,
val refreshPhaseEnabled: Boolean = false,

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.

This flips the primary-constructor default of refreshPhaseEnabled from true to false, which globally disables the DSv2/metadata refresh phase for every direct new QueryExecution(...) (the phase is gated below at if (refreshPhaseEnabled)). That reintroduces the stale-cached-plan behavior this PR demonstrates, for all callers -- not just the Connect path. The helper factories still default to true, so this also makes direct construction diverge from them. If the intent is only to demonstrate the bug, please keep this default at true and drive the tests via a test-local QueryExecution with refreshPhaseEnabled = false (and mark the PR [WIP]/draft); the actual fix should gate refresh narrowly as in #57609 rather than turning it off by default.

val queryId: UUID = UUIDv7Generator.generate(),
// When a transaction is active, callers creating nested QueryExecution instances MUST pass
// the enclosing QueryExecution's analyzer here to propagate the transaction context.
Expand Down