HyukjinKwon commented on code in PR #42304:
URL: https://github.com/apache/spark/pull/42304#discussion_r1282510781
##########
connector/connect/client/jvm/src/main/scala/org/apache/spark/sql/connect/client/ExecutePlanResponseReattachableIterator.scala:
##########
@@ -102,28 +102,33 @@ class ExecutePlanResponseReattachableIterator(
throw new java.util.NoSuchElementException()
}
- // Get next response, possibly triggering reattach in case of stream error.
- var firstTry = true
- val ret = retry {
- if (firstTry) {
- // on first try, we use the existing iterator.
- firstTry = false
- } else {
- // on retry, the iterator is borked, so we need a new one
- iterator =
rawBlockingStub.reattachExecute(createReattachExecuteRequest())
+ try {
+ // Get next response, possibly triggering reattach in case of stream
error.
+ var firstTry = true
+ val ret = retry {
+ if (firstTry) {
+ // on first try, we use the existing iterator.
+ firstTry = false
+ } else {
+ // on retry, the iterator is borked, so we need a new one
+ iterator =
rawBlockingStub.reattachExecute(createReattachExecuteRequest())
+ }
+ iterator.next()
}
- iterator.next()
- }
- // Record last returned response, to know where to restart in case of
reattach.
- lastReturnedResponseId = Some(ret.getResponseId)
- if (ret.hasResultComplete) {
- resultComplete = true
- releaseExecute(None) // release all
- } else {
- releaseExecute(lastReturnedResponseId) // release until this response
+ // Record last returned response, to know where to restart in case of
reattach.
+ lastReturnedResponseId = Some(ret.getResponseId)
+ if (ret.hasResultComplete) {
+ release()
+ } else {
+ releaseUntil(lastReturnedResponseId.get)
+ }
+ ret
+ } catch {
+ case ex: StatusRuntimeException =>
Review Comment:
qq: shouldn't this be for all exceptions?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
