[GitHub] [spark] HyukjinKwon commented on a diff in pull request #38613: [SPARK-41005][CONNECT][PYTHON][FOLLOW-UP] Fetch/send partitions in parallel for Arrow based collect
HyukjinKwon commented on code in PR #38613: URL: https://github.com/apache/spark/pull/38613#discussion_r1020130165 ## connector/connect/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectStreamHandler.scala: ## @@ -144,36 +144,10 @@ class SparkConnectStreamHandler(responseObserver: StreamObserver[Response]) exte .toArrowBatchIterator(iter, schema, maxRecordsPerBatch, timeZoneId) } -val signal = new Object -val partitions = collection.mutable.Map.empty[Int, Array[Batch]] - -val processPartition = (iter: Iterator[Batch]) => iter.toArray - // This callback is executed by the DAGScheduler thread. -// After fetching a partition, it inserts the partition into the Map, and then -// wakes up the main thread. -val resultHandler = (partitionId: Int, partition: Array[Batch]) => { - signal.synchronized { -partitions(partitionId) = partition -signal.notify() - } - () -} - -spark.sparkContext.runJob(batches, processPartition, resultHandler) - -// The man thread will wait until 0-th partition is available, -// then send it to client and wait for next partition. -var currentPartitionId = 0 -while (currentPartitionId < numPartitions) { - val partition = signal.synchronized { -while (!partitions.contains(currentPartitionId)) { - signal.wait() -} -partitions.remove(currentPartitionId).get - } - - partition.foreach { case (bytes, count) => +def writeBatches(arrowBatches: Array[Batch]): Unit = { Review Comment: Sorry I missed your comment when I opened this PR. BTW, this is actually what PySpark's implementation doing, and was thinking that it's better to match how they work, dedup, and improve together. It should work fine most cases - PySpark implementation has been running in production many years, and I haven't yet heard complaints related to this. -- 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-unsubscr...@spark.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org - To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For additional commands, e-mail: reviews-h...@spark.apache.org
[GitHub] [spark] HyukjinKwon commented on a diff in pull request #38613: [SPARK-41005][CONNECT][PYTHON][FOLLOW-UP] Fetch/send partitions in parallel for Arrow based collect
HyukjinKwon commented on code in PR #38613: URL: https://github.com/apache/spark/pull/38613#discussion_r1020130165 ## connector/connect/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectStreamHandler.scala: ## @@ -144,36 +144,10 @@ class SparkConnectStreamHandler(responseObserver: StreamObserver[Response]) exte .toArrowBatchIterator(iter, schema, maxRecordsPerBatch, timeZoneId) } -val signal = new Object -val partitions = collection.mutable.Map.empty[Int, Array[Batch]] - -val processPartition = (iter: Iterator[Batch]) => iter.toArray - // This callback is executed by the DAGScheduler thread. -// After fetching a partition, it inserts the partition into the Map, and then -// wakes up the main thread. -val resultHandler = (partitionId: Int, partition: Array[Batch]) => { - signal.synchronized { -partitions(partitionId) = partition -signal.notify() - } - () -} - -spark.sparkContext.runJob(batches, processPartition, resultHandler) - -// The man thread will wait until 0-th partition is available, -// then send it to client and wait for next partition. -var currentPartitionId = 0 -while (currentPartitionId < numPartitions) { - val partition = signal.synchronized { -while (!partitions.contains(currentPartitionId)) { - signal.wait() -} -partitions.remove(currentPartitionId).get - } - - partition.foreach { case (bytes, count) => +def writeBatches(arrowBatches: Array[Batch]): Unit = { Review Comment: Sorry I missed your comment when I opened this PR. BTW, this is actually what PySpark's implementation doing, and was thinking that it's better to match how they work, dedup, and improve together. It should work fine most cases - PySpark implementation has been running in production many years, and I haven't yet heard complaints related to this. -- 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-unsubscr...@spark.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org - To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For additional commands, e-mail: reviews-h...@spark.apache.org
[GitHub] [spark] HyukjinKwon commented on a diff in pull request #38613: [SPARK-41005][CONNECT][PYTHON][FOLLOW-UP] Fetch/send partitions in parallel for Arrow based collect
HyukjinKwon commented on code in PR #38613: URL: https://github.com/apache/spark/pull/38613#discussion_r1020127951 ## connector/connect/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectStreamHandler.scala: ## @@ -56,7 +56,7 @@ class SparkConnectStreamHandler(responseObserver: StreamObserver[Response]) exte try { processAsArrowBatches(request.getClientId, dataframe) } catch { - case e: Exception => + case e: Throwable => Review Comment: ```suggestion case e: Exception => ``` -- 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-unsubscr...@spark.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org - To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For additional commands, e-mail: reviews-h...@spark.apache.org
[GitHub] [spark] HyukjinKwon commented on a diff in pull request #38613: [SPARK-41005][CONNECT][PYTHON][FOLLOW-UP] Fetch/send partitions in parallel for Arrow based collect
HyukjinKwon commented on code in PR #38613: URL: https://github.com/apache/spark/pull/38613#discussion_r1020127788 ## connector/connect/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectStreamHandler.scala: ## @@ -56,7 +56,7 @@ class SparkConnectStreamHandler(responseObserver: StreamObserver[Response]) exte try { processAsArrowBatches(request.getClientId, dataframe) } catch { - case e: Exception => + case e: Throwable => Review Comment: Yeah, let me revert this for now. We should remove this JSON fallback anyway. -- 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-unsubscr...@spark.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org - To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For additional commands, e-mail: reviews-h...@spark.apache.org
[GitHub] [spark] HyukjinKwon commented on a diff in pull request #38613: [SPARK-41005][CONNECT][PYTHON][FOLLOW-UP] Fetch/send partitions in parallel for Arrow based collect
HyukjinKwon commented on code in PR #38613: URL: https://github.com/apache/spark/pull/38613#discussion_r1019833127 ## connector/connect/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectStreamHandler.scala: ## @@ -184,9 +158,30 @@ class SparkConnectStreamHandler(responseObserver: StreamObserver[Response]) exte responseObserver.onNext(response.build()) numSent += 1 } +} + +// Store collection results for worst case of 1 to N-1 partitions +val results = new Array[Array[Batch]](numPartitions - 1) +var lastIndex = -1 // index of last partition written - currentPartitionId += 1 +// Handler to eagerly write partitions in order +val resultHandler = (partitionId: Int, partition: Array[Batch]) => { + // If result is from next partition in order + if (partitionId - 1 == lastIndex) { +writeBatches(partition) +lastIndex += 1 +// Write stored partitions that come next in order +while (lastIndex < results.length && results(lastIndex) != null) { + writeBatches(results(lastIndex)) + results(lastIndex) = null + lastIndex += 1 +} + } else { +// Store partitions received out of order +results(partitionId - 1) = partition + } } +spark.sparkContext.runJob(batches, (iter: Iterator[Batch]) => iter.toArray, resultHandler) Review Comment: Hm, I just noticed the review comment. I believe this is matched with our current implementation in PySpark. If we should improve, let's improve both paths together. I would prefer to match them and deduplicate the logic first. -- 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-unsubscr...@spark.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org - To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For additional commands, e-mail: reviews-h...@spark.apache.org
[GitHub] [spark] HyukjinKwon commented on a diff in pull request #38613: [SPARK-41005][CONNECT][PYTHON][FOLLOW-UP] Fetch/send partitions in parallel for Arrow based collect
HyukjinKwon commented on code in PR #38613: URL: https://github.com/apache/spark/pull/38613#discussion_r1019831985 ## connector/connect/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectStreamHandler.scala: ## @@ -184,9 +158,30 @@ class SparkConnectStreamHandler(responseObserver: StreamObserver[Response]) exte responseObserver.onNext(response.build()) numSent += 1 } +} + +// Store collection results for worst case of 1 to N-1 partitions +val results = new Array[Array[Batch]](numPartitions - 1) +var lastIndex = -1 // index of last partition written - currentPartitionId += 1 +// Handler to eagerly write partitions in order +val resultHandler = (partitionId: Int, partition: Array[Batch]) => { Review Comment: Nope, it doesn't (because it's guided by the index). This approach is actually from the initial ordered implementation of collect with Arrow (that were in production for very long time), https://github.com/apache/spark/commit/82c18c240a6913a917df3b55cc5e22649561c4dd#diff-459628811d7786c705fbb2b7a381ecd2b88f183f44ab607d43b3d33ea48d390fR3282-R3318. -- 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-unsubscr...@spark.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org - To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For additional commands, e-mail: reviews-h...@spark.apache.org
[GitHub] [spark] HyukjinKwon commented on a diff in pull request #38613: [SPARK-41005][CONNECT][PYTHON][FOLLOW-UP] Fetch/send partitions in parallel for Arrow based collect
HyukjinKwon commented on code in PR #38613: URL: https://github.com/apache/spark/pull/38613#discussion_r1019831985 ## connector/connect/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectStreamHandler.scala: ## @@ -184,9 +158,30 @@ class SparkConnectStreamHandler(responseObserver: StreamObserver[Response]) exte responseObserver.onNext(response.build()) numSent += 1 } +} + +// Store collection results for worst case of 1 to N-1 partitions +val results = new Array[Array[Batch]](numPartitions - 1) +var lastIndex = -1 // index of last partition written - currentPartitionId += 1 +// Handler to eagerly write partitions in order +val resultHandler = (partitionId: Int, partition: Array[Batch]) => { Review Comment: Nope, it doesn't (because it's guided by the index). This approach is actually from the initial ordered implementation of collect with Arrow (that has been running in production for very long time), https://github.com/apache/spark/commit/82c18c240a6913a917df3b55cc5e22649561c4dd#diff-459628811d7786c705fbb2b7a381ecd2b88f183f44ab607d43b3d33ea48d390fR3282-R3318. -- 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-unsubscr...@spark.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org - To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For additional commands, e-mail: reviews-h...@spark.apache.org
[GitHub] [spark] HyukjinKwon commented on a diff in pull request #38613: [SPARK-41005][CONNECT][PYTHON][FOLLOW-UP] Fetch/send partitions in parallel for Arrow based collect
HyukjinKwon commented on code in PR #38613: URL: https://github.com/apache/spark/pull/38613#discussion_r1019826331 ## connector/connect/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectStreamHandler.scala: ## @@ -56,7 +56,7 @@ class SparkConnectStreamHandler(responseObserver: StreamObserver[Response]) exte try { processAsArrowBatches(request.getClientId, dataframe) } catch { - case e: Exception => + case e: Throwable => Review Comment: Hm, it's a fallback so should better work in any case .. (?). I think we should remove this fallback anyway. I can revert this change for now too. (see https://github.com/apache/spark/pull/38468#discussion_r1013555606 -- 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-unsubscr...@spark.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org - To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For additional commands, e-mail: reviews-h...@spark.apache.org