HyukjinKwon commented on code in PR #50301:
URL: https://github.com/apache/spark/pull/50301#discussion_r2000018806
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala:
##########
@@ -83,17 +89,37 @@ private[python] trait PythonArrowOutput[OUT <: AnyRef] {
self: BasePythonRunner[
throw writer.exception.get
}
try {
- if (reader != null && batchLoaded) {
+ if (batchLoaded && rowCount > 0 && currentRowIdx < rowCount) {
+ val batchRoot = if (arrowMaxRecordsPerOutputBatch > 0) {
+ val remainingRows = rowCount - currentRowIdx
+ if (remainingRows > arrowMaxRecordsPerOutputBatch) {
+ root.slice(currentRowIdx, arrowMaxRecordsPerOutputBatch)
+ } else {
+ root
+ }
+ } else {
+ root
+ }
+
+ currentRowIdx = currentRowIdx + batchRoot.getRowCount
+
+ vectors = batchRoot.getFieldVectors().asScala.map { vector =>
+ new ArrowColumnVector(vector)
+ }.toArray[ColumnVector]
+
+ val batch = new ColumnarBatch(vectors)
+ batch.setNumRows(batchRoot.getRowCount)
+ deserializeColumnarBatch(batch, schema)
Review Comment:
BTW, ideally we should slices the batch at Python worker side. With the
number of records, it should be feasible.
Separately, I wonder if we should limit the size by bytes ..but in this case
the problem is that it is difficult to get the byte size from pandas instances
(so it is difficult to slice pandas instances before converting it to Arrow
batch).
--
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
