[SQL] Stateful operations fail in non-streaming mode when StreamingQueryManager hasn't been initialized

[Kimahriman]

When using stateful operations in non-streaming mode, like spark.read.format("statestore"), or TransformWithState in non-streaming mode, an error is raised if the StreamingQueryManager hasn't been initialized.

Example:

First create a stateful checkpoint

from pyspark.sql import SparkSession

spark = SparkSession.Builder().getOrCreate()

df = spark.range(10)

df.write.mode("overwrite").parquet("/tmp/test.parquet")

(spark.readStream
    .schema(df.schema)
    .parquet("/tmp/test.parquet")
    .dropDuplicates()
    .writeStream
    .format("console")
    .trigger(once=True)
    .option("checkpointLocation", "/tmp/checkpoint")
    .start()
    .awaitTermination()
)

Then with a new Spark session try to read the state store:

spark = SparkSession.Builder().getOrCreate()

(spark.read
    .format("statestore")
    .load("/tmp/checkpoint")
    .show(10, truncate=False)
)

This fails with

org.apache.spark.SparkException: [CANNOT_LOAD_STATE_STORE.UNCATEGORIZED] An error occurred during loading state.  SQLSTATE: 58030
        at org.apache.spark.sql.errors.QueryExecutionErrors$.cannotLoadStore(QueryExecutionErrors.scala:3016)
        at org.apache.spark.sql.execution.streaming.state.StateStoreErrors$.cannotLoadStore(StateStoreErrors.scala:287)
        at org.apache.spark.sql.execution.streaming.state.HDFSBackedStateStoreProvider.getLoadedMapForStore(HDFSBackedStateStoreProvider.scala:419)
        at org.apache.spark.sql.execution.streaming.state.HDFSBackedStateStoreProvider.getReadStore(HDFSBackedStateStoreProvider.scala:394)
        at org.apache.spark.sql.execution.datasources.v2.state.StatePartitionReader.store$lzycompute(StatePartitionReader.scala:245)
        at org.apache.spark.sql.execution.datasources.v2.state.StatePartitionReader.store(StatePartitionReader.scala:238)
        at org.apache.spark.sql.execution.datasources.v2.state.StatePartitionReader.iter$lzycompute(StatePartitionReader.scala:279)
        at org.apache.spark.sql.execution.datasources.v2.state.StatePartitionReader.iter(StatePartitionReader.scala:262)
        at org.apache.spark.sql.execution.datasources.v2.state.StatePartitionReaderBase.next(StatePartitionReader.scala:203)
        at org.apache.spark.sql.execution.datasources.v2.PartitionIterator.hasNext(DataSourceRDD.scala:240)
        at org.apache.spark.InterruptibleIterator.hasNext(InterruptibleIterator.scala:37)
        at scala.collection.Iterator$$anon$10.hasNext(Iterator.scala:610)
        at scala.collection.Iterator$$anon$9.hasNext(Iterator.scala:593)
        at org.apache.spark.sql.catalyst.expressions.GeneratedClass$GeneratedIteratorForCodegenStage1.processNext(Unknown Source)
        at org.apache.spark.sql.execution.BufferedRowIterator.hasNext(BufferedRowIterator.java:44)
        at org.apache.spark.sql.execution.WholeStageCodegenEvaluatorFactory$WholeStageCodegenPartitionEvaluator$$anon$1.hasNext(WholeStageCodegenEvaluatorFactory.scala:50)
        at org.apache.spark.sql.execution.SparkPlan.$anonfun$getByteArrayRdd$2(SparkPlan.scala:418)
        at org.apache.spark.rdd.RDD.$anonfun$mapPartitionsInternal$2(RDD.scala:960)
        at org.apache.spark.rdd.RDD.$anonfun$mapPartitionsInternal$2$adapted(RDD.scala:960)
        at org.apache.spark.rdd.MapPartitionsRDD.compute(MapPartitionsRDD.scala:57)
        at org.apache.spark.rdd.RDD.computeOrReadCheckpoint(RDD.scala:374)
        at org.apache.spark.rdd.RDD.iterator(RDD.scala:338)
        at org.apache.spark.scheduler.ResultTask.runTask(ResultTask.scala:93)
        at org.apache.spark.TaskContext.runTaskWithListeners(TaskContext.scala:223)
        at org.apache.spark.scheduler.Task.run(Task.scala:147)
        at org.apache.spark.executor.Executor$TaskRunner.$anonfun$run$5(Executor.scala:912)
        at org.apache.spark.util.SparkErrorUtils.tryWithSafeFinally(SparkErrorUtils.scala:86)
        at org.apache.spark.util.SparkErrorUtils.tryWithSafeFinally$(SparkErrorUtils.scala:83)
        at org.apache.spark.util.Utils$.tryWithSafeFinally(Utils.scala:97)
        at org.apache.spark.executor.Executor$TaskRunner.run(Executor.scala:915)
        at java.base/java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1136)
        at java.base/java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:635)
        at java.base/java.lang.Thread.run(Thread.java:840)
        at org.apache.spark.scheduler.DAGScheduler.runJob(DAGScheduler.scala:1552)
        at org.apache.spark.SparkContext.runJob(SparkContext.scala:2743)
        at org.apache.spark.SparkContext.runJob(SparkContext.scala:2764)
        at org.apache.spark.SparkContext.runJob(SparkContext.scala:2783)
        at org.apache.spark.sql.execution.SparkPlan.executeTake(SparkPlan.scala:562)
        at org.apache.spark.sql.execution.SparkPlan.executeTake(SparkPlan.scala:515)
        at org.apache.spark.sql.execution.CollectLimitExec.executeCollect(limit.scala:58)
        at org.apache.spark.sql.classic.Dataset.collectFromPlan(Dataset.scala:2335)
        at org.apache.spark.sql.classic.Dataset.$anonfun$head$1(Dataset.scala:1470)
        at org.apache.spark.sql.classic.Dataset.$anonfun$withAction$3(Dataset.scala:2324)
        at org.apache.spark.sql.execution.QueryExecution$.withInternalError(QueryExecution.scala:896)
        at org.apache.spark.sql.classic.Dataset.$anonfun$withAction$2(Dataset.scala:2322)
        at org.apache.spark.sql.execution.QueryExecution.withQueryExecutionId(QueryExecution.scala:434)
        at org.apache.spark.sql.classic.Dataset.$anonfun$withAction$1(Dataset.scala:2322)
        at org.apache.spark.sql.execution.SQLExecution$.$anonfun$withNewExecutionId0$8(SQLExecution.scala:228)
        at org.apache.spark.sql.execution.SQLExecution$.withSessionTagsApplied(SQLExecution.scala:352)
        at org.apache.spark.sql.execution.SQLExecution$.$anonfun$withNewExecutionId0$7(SQLExecution.scala:189)
        at org.apache.spark.JobArtifactSet$.withActiveJobArtifactState(JobArtifactSet.scala:94)
        at org.apache.spark.sql.artifact.ArtifactManager.$anonfun$withResources$1(ArtifactManager.scala:112)
        at org.apache.spark.sql.artifact.ArtifactManager.withClassLoaderIfNeeded(ArtifactManager.scala:106)
        at org.apache.spark.sql.artifact.ArtifactManager.withResources(ArtifactManager.scala:111)
        at org.apache.spark.sql.execution.SQLExecution$.$anonfun$withNewExecutionId0$6(SQLExecution.scala:189)
        at org.apache.spark.sql.execution.SQLExecution$.withSQLConfPropagated(SQLExecution.scala:375)
        at org.apache.spark.sql.execution.SQLExecution$.$anonfun$withNewExecutionId0$1(SQLExecution.scala:188)
        at org.apache.spark.sql.SparkSession.withActive(SparkSession.scala:810)
        at org.apache.spark.sql.execution.SQLExecution$.withNewExecutionId0(SQLExecution.scala:130)
        at org.apache.spark.sql.execution.SQLExecution$.withNewExecutionId(SQLExecution.scala:317)
        at org.apache.spark.sql.classic.Dataset.withAction(Dataset.scala:2321)
        at org.apache.spark.sql.classic.Dataset.head(Dataset.scala:1470)
        at org.apache.spark.sql.Dataset.take(Dataset.scala:2939)
        at org.apache.spark.sql.classic.Dataset.getRows(Dataset.scala:337)
        at org.apache.spark.sql.classic.Dataset.showString(Dataset.scala:373)
        at java.base/jdk.internal.reflect.NativeMethodAccessorImpl.invoke0(Native Method)
        at java.base/jdk.internal.reflect.NativeMethodAccessorImpl.invoke(NativeMethodAccessorImpl.java:77)
        at java.base/jdk.internal.reflect.DelegatingMethodAccessorImpl.invoke(DelegatingMethodAccessorImpl.java:43)
        at java.base/java.lang.reflect.Method.invoke(Method.java:569)
        at py4j.reflection.MethodInvoker.invoke(MethodInvoker.java:244)
        at py4j.reflection.ReflectionEngine.invoke(ReflectionEngine.java:374)
        at py4j.Gateway.invoke(Gateway.java:282)
        at py4j.commands.AbstractCommand.invokeMethod(AbstractCommand.java:132)
        at py4j.commands.CallCommand.execute(CallCommand.java:79)
        at py4j.ClientServerConnection.waitForCommands(ClientServerConnection.java:184)
        at py4j.ClientServerConnection.run(ClientServerConnection.java:108)
        at java.base/java.lang.Thread.run(Thread.java:840)
Caused by: org.apache.spark.SparkException: Exception thrown in awaitResult: 
        at org.apache.spark.util.SparkThreadUtils$.awaitResult(SparkThreadUtils.scala:70)
        at org.apache.spark.util.SparkThreadUtils$.awaitResult(SparkThreadUtils.scala:44)
        at org.apache.spark.util.ThreadUtils$.awaitResult(ThreadUtils.scala:359)
        at org.apache.spark.rpc.RpcTimeout.awaitResult(RpcTimeout.scala:75)
        at org.apache.spark.rpc.RpcEnv.setupEndpointRefByURI(RpcEnv.scala:102)
        at org.apache.spark.rpc.RpcEnv.setupEndpointRef(RpcEnv.scala:110)
        at org.apache.spark.util.RpcUtils$.makeDriverRef(RpcUtils.scala:34)
        at org.apache.spark.sql.execution.streaming.state.StateStoreCoordinatorRef$.forExecutor(StateStoreCoordinator.scala:161)
        at org.apache.spark.sql.execution.streaming.state.StateStoreProvider$.coordinatorRef(StateStore.scala:1183)
        at org.apache.spark.sql.execution.streaming.state.HDFSBackedStateStoreProvider.reportSnapshotUploadToCoordinator(HDFSBackedStateStoreProvider.scala:1416)
        at org.apache.spark.sql.execution.streaming.state.HDFSBackedStateStoreProvider.$anonfun$loadMap$2(HDFSBackedStateStoreProvider.scala:753)
        at org.apache.spark.util.Utils$.timeTakenMs(Utils.scala:471)
        at org.apache.spark.sql.execution.streaming.state.HDFSBackedStateStoreProvider.loadMap(HDFSBackedStateStoreProvider.scala:734)
        at org.apache.spark.sql.execution.streaming.state.HDFSBackedStateStoreProvider.getLoadedMapForStore(HDFSBackedStateStoreProvider.scala:409)
        at org.apache.spark.sql.execution.streaming.state.HDFSBackedStateStoreProvider.getReadStore(HDFSBackedStateStoreProvider.scala:394)
        at org.apache.spark.sql.execution.datasources.v2.state.StatePartitionReader.store$lzycompute(StatePartitionReader.scala:245)
        at org.apache.spark.sql.execution.datasources.v2.state.StatePartitionReader.store(StatePartitionReader.scala:238)
        at org.apache.spark.sql.execution.datasources.v2.state.StatePartitionReader.iter$lzycompute(StatePartitionReader.scala:279)
        at org.apache.spark.sql.execution.datasources.v2.state.StatePartitionReader.iter(StatePartitionReader.scala:262)
        at org.apache.spark.sql.execution.datasources.v2.state.StatePartitionReaderBase.next(StatePartitionReader.scala:203)
        at org.apache.spark.sql.execution.datasources.v2.PartitionIterator.hasNext(DataSourceRDD.scala:240)
        at org.apache.spark.InterruptibleIterator.hasNext(InterruptibleIterator.scala:37)
        at scala.collection.Iterator$$anon$10.hasNext(Iterator.scala:610)
        at scala.collection.Iterator$$anon$9.hasNext(Iterator.scala:593)
        at org.apache.spark.sql.catalyst.expressions.GeneratedClass$GeneratedIteratorForCodegenStage1.processNext(Unknown Source)
        at org.apache.spark.sql.execution.BufferedRowIterator.hasNext(BufferedRowIterator.java:44)
        at org.apache.spark.sql.execution.WholeStageCodegenEvaluatorFactory$WholeStageCodegenPartitionEvaluator$$anon$1.hasNext(WholeStageCodegenEvaluatorFactory.scala:50)
        at org.apache.spark.sql.execution.SparkPlan.$anonfun$getByteArrayRdd$2(SparkPlan.scala:418)
        at org.apache.spark.rdd.RDD.$anonfun$mapPartitionsInternal$2(RDD.scala:960)
        at org.apache.spark.rdd.RDD.$anonfun$mapPartitionsInternal$2$adapted(RDD.scala:960)
        at org.apache.spark.rdd.MapPartitionsRDD.compute(MapPartitionsRDD.scala:57)
        at org.apache.spark.rdd.RDD.computeOrReadCheckpoint(RDD.scala:374)
        at org.apache.spark.rdd.RDD.iterator(RDD.scala:338)
        at org.apache.spark.scheduler.ResultTask.runTask(ResultTask.scala:93)
        at org.apache.spark.TaskContext.runTaskWithListeners(TaskContext.scala:223)
        at org.apache.spark.scheduler.Task.run(Task.scala:147)
        at org.apache.spark.executor.Executor$TaskRunner.$anonfun$run$5(Executor.scala:912)
        at org.apache.spark.util.SparkErrorUtils.tryWithSafeFinally(SparkErrorUtils.scala:86)
        at org.apache.spark.util.SparkErrorUtils.tryWithSafeFinally$(SparkErrorUtils.scala:83)
        at org.apache.spark.util.Utils$.tryWithSafeFinally(Utils.scala:97)
        at org.apache.spark.executor.Executor$TaskRunner.run(Executor.scala:915)
        at java.base/java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1136)
        at java.base/java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:635)
        ... 1 more
Caused by: org.apache.spark.rpc.RpcEndpointNotFoundException: Cannot find endpoint: spark://StateStoreCoordinator@macbookpro:58792
        at org.apache.spark.rpc.netty.NettyRpcEnv.$anonfun$asyncSetupEndpointRefByURI$1(NettyRpcEnv.scala:151)
        at org.apache.spark.rpc.netty.NettyRpcEnv.$anonfun$asyncSetupEndpointRefByURI$1$adapted(NettyRpcEnv.scala:147)
        at scala.concurrent.impl.Promise$Transformation.run(Promise.scala:503)
        at org.apache.spark.util.ThreadUtils$$anon$1.execute(ThreadUtils.scala:99)
        at scala.concurrent.impl.ExecutionContextImpl.execute(ExecutionContextImpl.scala:21)
        at scala.concurrent.impl.Promise$Transformation.submitWithValue(Promise.scala:462)
        at scala.concurrent.impl.Promise$DefaultPromise.submitWithValue(Promise.scala:371)
        at scala.concurrent.impl.Promise$DefaultPromise.tryComplete0(Promise.scala:295)
        at scala.concurrent.impl.Promise$Transformation.run(Promise.scala:537)
        at org.apache.spark.util.ThreadUtils$$anon$1.execute(ThreadUtils.scala:99)
        at scala.concurrent.impl.ExecutionContextImpl.execute(ExecutionContextImpl.scala:21)
        at scala.concurrent.impl.Promise$Transformation.submitWithValue(Promise.scala:462)
        at scala.concurrent.impl.Promise$DefaultPromise.submitWithValue(Promise.scala:371)
        at scala.concurrent.impl.Promise$DefaultPromise.tryComplete0(Promise.scala:295)
        at scala.concurrent.impl.Promise$Transformation.run(Promise.scala:537)
        at scala.concurrent.ExecutionContext$parasitic$.execute(ExecutionContext.scala:222)
        at scala.concurrent.impl.Promise$Transformation.submitWithValue(Promise.scala:462)
        at scala.concurrent.impl.Promise$DefaultPromise.submitWithValue(Promise.scala:368)
        at scala.concurrent.impl.Promise$DefaultPromise.tryComplete0(Promise.scala:295)
        at scala.concurrent.impl.Promise$DefaultPromise.tryComplete(Promise.scala:288)
        at scala.concurrent.Promise.trySuccess(Promise.scala:99)
        at scala.concurrent.Promise.trySuccess$(Promise.scala:99)
        at scala.concurrent.impl.Promise$DefaultPromise.trySuccess(Promise.scala:104)
        at org.apache.spark.rpc.netty.NettyRpcEnv.onSuccess$1(NettyRpcEnv.scala:228)
        at org.apache.spark.rpc.netty.NettyRpcEnv.$anonfun$askAbortable$5(NettyRpcEnv.scala:242)
        at org.apache.spark.rpc.netty.NettyRpcEnv.$anonfun$askAbortable$5$adapted(NettyRpcEnv.scala:241)
        at scala.concurrent.impl.Promise$Transformation.run(Promise.scala:517)
        at org.apache.spark.util.ThreadUtils$$anon$1.execute(ThreadUtils.scala:99)
        at scala.concurrent.impl.ExecutionContextImpl.execute(ExecutionContextImpl.scala:21)
        at scala.concurrent.impl.Promise$Transformation.submitWithValue(Promise.scala:462)
        at scala.concurrent.impl.Promise$DefaultPromise.submitWithValue(Promise.scala:371)
        at scala.concurrent.impl.Promise$DefaultPromise.tryComplete0(Promise.scala:295)
        at scala.concurrent.impl.Promise$DefaultPromise.tryComplete(Promise.scala:288)
        at scala.concurrent.Promise.complete(Promise.scala:57)
        at scala.concurrent.Promise.complete$(Promise.scala:56)
        at scala.concurrent.impl.Promise$DefaultPromise.complete(Promise.scala:104)
        at scala.concurrent.Promise.success(Promise.scala:91)
        at scala.concurrent.Promise.success$(Promise.scala:91)
        at scala.concurrent.impl.Promise$DefaultPromise.success(Promise.scala:104)
        at org.apache.spark.rpc.netty.LocalNettyRpcCallContext.send(NettyRpcCallContext.scala:50)
        at org.apache.spark.rpc.netty.NettyRpcCallContext.reply(NettyRpcCallContext.scala:32)
        at org.apache.spark.rpc.netty.RpcEndpointVerifier$$anonfun$receiveAndReply$1.applyOrElse(RpcEndpointVerifier.scala:31)
        at org.apache.spark.rpc.netty.Inbox.$anonfun$process$1(Inbox.scala:104)
        at org.apache.spark.rpc.netty.Inbox.safelyCall(Inbox.scala:216)
        at org.apache.spark.rpc.netty.Inbox.process(Inbox.scala:101)
        at org.apache.spark.rpc.netty.MessageLoop.org$apache$spark$rpc$netty$MessageLoop$$receiveLoop(MessageLoop.scala:76)
        at org.apache.spark.rpc.netty.MessageLoop$$anon$1.run(MessageLoop.scala:42)
        ... 3 more

If you simply force the StreamingQueryManager to load, it works:

spark = SparkSession.Builder().getOrCreate()

spark.streams.active

(spark.read
    .format("statestore")
    .load("/tmp/checkpoint")
    .show(10, truncate=False)
)

I believe this was introduced in #50123 which always tries to talk to the state store coordinator when loading a state store. And disabling this also fixes the issue:

spark = (SparkSession.Builder()
    .config("spark.sql.streaming.stateStore.coordinatorReportSnapshotUploadLag", "false")
    .getOrCreate()
)

(spark.read
    .format("statestore")
    .load("/tmp/checkpoint")
    .show(10, truncate=False)
)

[shrirangmhalgi]

Thank you @Kimahriman for reporting the issue. I was able to reproduce the issue and trace the root cause.

The StateStoreCoordinator RPC endpoint is only registered when StreamingQueryManager is instantiated - this was intentionally a lazy val in SessionState (since SPARK-29423). PR #50123 added reportSnapshotUploadToCoordinator() to HDFSBackedStateStoreProvider.loadMap(), creating an unconditional dependency on the coordinator from the state store read path.

The existing test suite masks this because StateDataSourceTestBase.beforeEach() explicitly calls spark.streams.stateStoreCoordinator before every test.

I will implement the fix and raise a PR for the same. I am thinking of making StateStoreProvider.coordinatorRef gracefully return None when the coordinator endpoint hasn't been registered. The snapshot upload reporting is best-effort and shouldn't fail batch reads.