juliuszsompolski commented on code in PR #43745:
URL: https://github.com/apache/spark/pull/43745#discussion_r1389661186


##########
connector/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/client/SparkConnectClientSuite.scala:
##########
@@ -387,6 +387,33 @@ class SparkConnectClientSuite extends ConnectFunSuite with 
BeforeAndAfterEach {
     }
     assert(dummyFn.counter == 2)
   }
+
+  test("SPARK-45871: Client execute iterator.toSeq consumes the reattachable 
iterator") {
+    startDummyServer(0)
+    client = SparkConnectClient
+      .builder()
+      .connectionString(s"sc://localhost:${server.getPort}")
+      .enableReattachableExecute()
+      .build()
+    val session = SparkSession.builder().client(client).create()
+    val cmd = session.newCommand(b =>
+      b.setSqlCommand(
+        proto.SqlCommand
+          .newBuilder()
+          .setSql("select * from range(10000000)")))
+    val plan = proto.Plan.newBuilder().setCommand(cmd)
+    val iter = client.execute(plan.build())
+    val reattachableIter = iter
+      .asInstanceOf[WrappedCloseableIterator[proto.ExecutePlanResponse]]
+      .innerIterator
+      .asInstanceOf[WrappedCloseableIterator[proto.ExecutePlanResponse]]
+      .innerIterator
+      .asInstanceOf[ExecutePlanResponseReattachableIterator]

Review Comment:
   Could you maybe add this as a helper function like
   ```
   def getReattachableIterator(iter: Iterator[proto.ExecutePlanResponse]): 
ExecutePlanResponseReattachableIterator = iter match {
     case e: ExecutePlanResponseReattachableIterator => e
     case w: WrappedCloseableIterator => w.innerIterator
   }
   ```
   (let it fail with a MatchError if it doesn't find it), and also update it in 
SparkConnectServerTest like this?
   
   It's ok-ish to have two copies of it (I don't see a common place where such 
a helper could be placed now), but having two copies of this kind of 
manual-unwrapping starts getting uglyish, and it's fragile to break with every 
change of how the custom stub wraps the requests.



-- 
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]

Reply via email to