[GitHub] [spark] HyukjinKwon commented on a diff in pull request #42304: [SPARK-44642] ReleaseExecute in ExecutePlanResponseReattachableIterator after it gets error from server

["HyukjinKwon (via GitHub)" <gi...@apache.org>]

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: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org