[SPARK-42960] [CONNECT] [SS] Add await_termination() and exception() API for Streaming Query in Python #40785

WweiL · 2023-04-13T21:36:47Z

What changes were proposed in this pull request?

Add the await_termination() and exception() to streaming query class.

For exception, only pass the message the same way in SparkConnectService, and construct the error the same way as the convert_exception method in _handle_rpc_error in client.py.

For await_termination, send the command multiple times instead of waiting to prevent RPC timeout. <-- I'm definitely open to any discussion on its implementation!

Why are the changes needed?

Add missing APIs.

Does this PR introduce any user-facing change?

Yes but part of ongoing developing of Streaming Spark Connect.

How was this patch tested?

Existing unit tests. Note that the unit tests for them are still skipped because of 1. queryManager is not implemented. 2. Allow access to stopped query is not implemented.

I was able to test them manually by

For test_stream_await_termination(), comment out the

for q in self.spark.streams.active:
    q.stop()

For test_stream_exception(), comment out unregistering terminated query:
https://github.com/apache/spark/blob/master/sql/core/src/main/scala/org/apache/spark/sql/streaming/StreamingQueryManager.scala#L411

…SIS flag as it's already in test options

… test

WweiL · 2023-04-13T21:37:35Z

@HyukjinKwon @rangadi @pengzhon-db

Hey guys, PTAL when you get a chance, thanks!

...connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala

connector/connect/common/src/main/protobuf/spark/connect/commands.proto

...connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala

python/pyspark/sql/connect/streaming/query.py

...connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala

rangadi · 2023-04-13T22:59:00Z

connector/connect/common/src/main/protobuf/spark/connect/commands.proto

@@ -298,6 +306,16 @@ message StreamingQueryCommandResult {
    // Logical and physical plans as string
    string result = 1;
  }
+
+  message ExceptionResult {


We probably need at least the stacktrace. Leave a comment about what our thinking is there.

Yes I agree, currently the support for error is limited.

I tracked how they handle errors, and found that it's through here:

spark/python/pyspark/sql/connect/client.py

Line 1115 in 0c2776a

raise convert_exception(info, status.message) from None

And then in the convert_exception method:

spark/python/pyspark/errors/exceptions/connect.py

Lines 57 to 58 in 0c2776a

elif "org.apache.spark.sql.streaming.StreamingQueryException" in classes:

return StreamingQueryException(message)

Only the message is directly passed.

I guess we could file a ticket to wait until batch side's change, and then we could align with them?

I think we should try to be consistent with what current exception() returns, which is return CapturedStreamingQueryException(msg, stackTrace, je.getCause()) ?

Right but satckTrace and cause are not included in connect's error framework so far. There is an ongoing PR about this: #40575.

python/pyspark/sql/connect/streaming/query.py

python/pyspark/sql/tests/streaming/test_streaming.py

WweiL · 2023-04-14T19:02:59Z

python/pyspark/sql/streaming/query.py

@@ -196,7 +196,7 @@ def awaitTermination(self, timeout: Optional[int] = None) -> Optional[bool]:
        >>> sq.stop()
        """
        if timeout is not None:
-            if not isinstance(timeout, (int, float)) or timeout < 0:
+            if not isinstance(timeout, (int, float)) or timeout <= 0:


@HyukjinKwon I also change this to be leq to align with StreamExecution's awaitTermination
https://github.com/apache/spark/blob/master/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/StreamExecution.scala#L549

...connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectService.scala

WweiL · 2023-04-18T23:00:17Z

@HyukjinKwon Can you take a look? Thanks!

HyukjinKwon · 2023-04-19T23:33:52Z

@WweiL mind rebasing this one please?

WweiL · 2023-04-19T23:56:24Z

Fetched and Merged with master. I noticed stack trace was added. I'll create a SPARK ticket to include it

HyukjinKwon · 2023-04-20T00:23:59Z

...connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala

+          respBuilder.getAwaitTerminationBuilder
+            .setTerminated(terminated)
+        } else {
+          query.awaitTermination()


Hmm .. just to be extra clear, it will be disconnected when it reaches gRPC timeout .. am i correct?

Right this is intended at this stage. @rangadi will push update regarding this I believe

The client keep sending heart beat message to keep the RPC connection alive.
That said, we would still need to improve handling of this. E.g. it should exit if client side disconnects.

HyukjinKwon · 2023-04-20T00:24:13Z

@hvanhovell or @grundprinzip actually mind taking a look please when you find some time?

connector/connect/common/src/main/protobuf/spark/connect/commands.proto

python/pyspark/sql/connect/streaming/query.py

HyukjinKwon · 2023-04-20T05:31:40Z

I see .. sorry for a bit of back and forth:

python/pyspark/sql/connect/streaming/query.py:72: error: Missing return statement  [return]
python/pyspark/sql/connect/streaming/query.py:108: error: Missing return statement  [return]
python/pyspark/sql/connect/streaming/query.py:140: error: Missing return statement  [return]

Let's add them back :-) ..

WweiL · 2023-04-20T22:54:25Z

@HyukjinKwon Can you merge this : ) Thanks!

HyukjinKwon · 2023-04-20T23:57:48Z

Merged to master.

…logics out ### What changes were proposed in this pull request? This PR factor Connect/non-Connect specific logics out into dedicated test classes. This PR is a followup of #40785 ### Why are the changes needed? In order to avoid test failure such as #44698 (comment) by missing dependencies ### Does this PR introduce _any_ user-facing change? No, test-only. ### How was this patch tested? CI in this PR should verify it. ### Was this patch authored or co-authored using generative AI tooling? No. Closes #44715 from HyukjinKwon/SPARK-42960-followup. Authored-by: Hyukjin Kwon <gurwls223@apache.org> Signed-off-by: Hyukjin Kwon <gurwls223@apache.org>

WweiL added 19 commits April 6, 2023 15:18

done

4d0fcdd

add versionchanged to query and readwriter

0ae7e33

fix conflict, put TODO to correct position

b14fe81

style

17720b7

comments

1e68a3c

address comments, add a new foreachBatch test class, remove all ELLIP…

dc05be8

…SIS flag as it's already in test options

minor

c1674eb

minor

23b9c93

lint

60ddd01

remove empty line

e576821

wip

304d01e

remove several docs in connect readwriter.py and query.py to pass doc…

26e2488

… test

minor, add back doc tests in module.py

e25f7e6

style

aa1d4c2

wip

e82678e

merge unittest

720e3d2

done

1f3ba94

merge master

d8a6e9f

minor

0c2776a

github-actions bot added CONNECT CORE PYTHON SQL STRUCTURED STREAMING labels Apr 13, 2023

rangadi reviewed Apr 13, 2023

View reviewed changes

address comments

da6fc21

WweiL commented Apr 14, 2023

View reviewed changes

python/pyspark/sql/tests/streaming/test_streaming.py Show resolved Hide resolved

WweiL commented Apr 14, 2023

View reviewed changes

WweiL added 2 commits April 18, 2023 12:48

regenerate proto files

dd4d54b

minor

3b375cc

amaliujia reviewed Apr 18, 2023

View reviewed changes

...connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectService.scala Outdated Show resolved Hide resolved

WweiL added 3 commits April 18, 2023 17:29

minor import import SparkConnectService._

34c28e7

can you run tests one more time

c487cd7

please pass

0c2133b

WweiL added 2 commits April 19, 2023 16:53

merge master

2384ba1

minor

da1d3e1

HyukjinKwon reviewed Apr 20, 2023

View reviewed changes

HyukjinKwon approved these changes Apr 20, 2023

View reviewed changes

HyukjinKwon reviewed Apr 20, 2023

View reviewed changes

connector/connect/common/src/main/protobuf/spark/connect/commands.proto Outdated Show resolved Hide resolved

HyukjinKwon reviewed Apr 20, 2023

View reviewed changes

python/pyspark/sql/connect/streaming/query.py Show resolved Hide resolved

HyukjinKwon reviewed Apr 20, 2023

View reviewed changes

python/pyspark/sql/connect/streaming/query.py Show resolved Hide resolved

remove return None, remove optional tag in awaitTerminationResult

eb19c2f

WweiL added 2 commits April 19, 2023 23:11

add back return None

ed67070

merge master

566e9fc

WweiL changed the title ~~[SPARK-42960] Add await_termination() and exception() API for Streaming Query~~ [SPARK-42960] [CONNECT] [SS] Add await_termination() and exception() API for Streaming Query Apr 20, 2023

WweiL changed the title ~~[SPARK-42960] [CONNECT] [SS] Add await_termination() and exception() API for Streaming Query~~ [SPARK-42960] [CONNECT] [SS] Add await_termination() and exception() API for Streaming Query in Python Apr 20, 2023

HyukjinKwon closed this in 875a297 Apr 20, 2023

WweiL mentioned this pull request Jul 26, 2023

[SPARK-42941][SS][CONNECT] Python StreamingQueryListener #42116

Closed

HyukjinKwon mentioned this pull request Jan 13, 2024

[SPARK-42960][PYTHON][SS][TESTS] Factor Connect/non-Connect specific logics out #44715

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[SPARK-42960] [CONNECT] [SS] Add await_termination() and exception() API for Streaming Query in Python #40785

[SPARK-42960] [CONNECT] [SS] Add await_termination() and exception() API for Streaming Query in Python #40785

WweiL commented Apr 13, 2023

WweiL commented Apr 13, 2023 •

edited

rangadi Apr 13, 2023

WweiL Apr 13, 2023 •

edited

pengzhon-db Apr 17, 2023

WweiL Apr 18, 2023

WweiL Apr 14, 2023

WweiL commented Apr 18, 2023

HyukjinKwon commented Apr 19, 2023

WweiL commented Apr 19, 2023 •

edited

HyukjinKwon Apr 20, 2023

WweiL Apr 20, 2023

HyukjinKwon Apr 20, 2023

rangadi Apr 20, 2023

HyukjinKwon commented Apr 20, 2023

HyukjinKwon commented Apr 20, 2023

WweiL commented Apr 20, 2023

HyukjinKwon commented Apr 20, 2023

	elif "org.apache.spark.sql.streaming.StreamingQueryException" in classes:
	return StreamingQueryException(message)

[SPARK-42960] [CONNECT] [SS] Add await_termination() and exception() API for Streaming Query in Python #40785

[SPARK-42960] [CONNECT] [SS] Add await_termination() and exception() API for Streaming Query in Python #40785

Conversation

WweiL commented Apr 13, 2023

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

WweiL commented Apr 13, 2023 • edited

Choose a reason for hiding this comment

WweiL Apr 13, 2023 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

WweiL commented Apr 18, 2023

HyukjinKwon commented Apr 19, 2023

WweiL commented Apr 19, 2023 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

HyukjinKwon commented Apr 20, 2023

HyukjinKwon commented Apr 20, 2023

WweiL commented Apr 20, 2023

HyukjinKwon commented Apr 20, 2023

WweiL commented Apr 13, 2023 •

edited

WweiL Apr 13, 2023 •

edited

WweiL commented Apr 19, 2023 •

edited