viirya commented on code in PR #59786:
URL: https://github.com/apache/spark/pull/59786#discussion_r4381044519
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala:
##########
@@ -265,47 +265,74 @@ class SliceBytesArrowOutputProcessorImpl(
arrowMaxBytesPerOutputBatch: Int)
extends BaseSliceArrowOutputProcessor(reader, pythonMetrics) {
+ // The number of rows of the last slice output from the current batch, from
which the search
+ // for the next slice starts.
+ private var sliceRowCount = 0
+
protected override def getRoot: VectorSchemaRoot = {
- // Output the whole root only if no row of this batch has been output yet.
- if (currentRowIdx == 0 && getBatchBytes(root) <
arrowMaxBytesPerOutputBatch) {
- currentRowIdx = rowCount
- root
- } else {
- // We iteratively slice the batch by 90% until we reach a batch with
required byte size.
- // If we cannot have it even with one single row, we will output a batch
with one row.
- var rootSlice = root
- val stepRatio = 0.9
- var sliceBatchCount = rowCount - currentRowIdx
-
- while (sliceBatchCount > 0) {
- rootSlice = root.slice(currentRowIdx, sliceBatchCount)
- val batchBytes = getBatchBytes(rootSlice)
- if (batchBytes < arrowMaxBytesPerOutputBatch) {
- currentRowIdx = currentRowIdx + rootSlice.getRowCount
-
- return rootSlice
- } else {
- rootSlice.close()
- sliceBatchCount = (sliceBatchCount * stepRatio).toInt
- }
+ if (currentRowIdx == 0) {
+ // Output the whole root if it fits, which happens only if no row of it
has been output yet.
+ if (batchBytes(root) <= arrowMaxBytesPerOutputBatch) {
+ currentRowIdx = rowCount
+ return root
}
- // If the batch size is still too large, we need to slice it down to a
single row.
- // Normally, this should not happen, but it is possible if the data is
very large.
- if (sliceBatchCount == 0) {
- sliceBatchCount = 1
+ sliceRowCount = (rowCount *
SliceBytesArrowOutputProcessorImpl.stepRatio).toInt
+ }
+ // We iteratively shrink the slice by 10% until it fits. Each search
starts from the number of
+ // rows of the last slice that fit, since every slice that does not start
at row 0 copies its
Review Comment:
Changed to "a slice that does not start at row 0 usually copies its offset
buffers".
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala:
##########
@@ -265,47 +265,74 @@ class SliceBytesArrowOutputProcessorImpl(
arrowMaxBytesPerOutputBatch: Int)
extends BaseSliceArrowOutputProcessor(reader, pythonMetrics) {
+ // The number of rows of the last slice output from the current batch, from
which the search
+ // for the next slice starts.
+ private var sliceRowCount = 0
+
protected override def getRoot: VectorSchemaRoot = {
- // Output the whole root only if no row of this batch has been output yet.
- if (currentRowIdx == 0 && getBatchBytes(root) <
arrowMaxBytesPerOutputBatch) {
- currentRowIdx = rowCount
- root
- } else {
- // We iteratively slice the batch by 90% until we reach a batch with
required byte size.
- // If we cannot have it even with one single row, we will output a batch
with one row.
- var rootSlice = root
- val stepRatio = 0.9
- var sliceBatchCount = rowCount - currentRowIdx
-
- while (sliceBatchCount > 0) {
- rootSlice = root.slice(currentRowIdx, sliceBatchCount)
- val batchBytes = getBatchBytes(rootSlice)
- if (batchBytes < arrowMaxBytesPerOutputBatch) {
- currentRowIdx = currentRowIdx + rootSlice.getRowCount
-
- return rootSlice
- } else {
- rootSlice.close()
- sliceBatchCount = (sliceBatchCount * stepRatio).toInt
- }
+ if (currentRowIdx == 0) {
+ // Output the whole root if it fits, which happens only if no row of it
has been output yet.
+ if (batchBytes(root) <= arrowMaxBytesPerOutputBatch) {
+ currentRowIdx = rowCount
+ return root
}
- // If the batch size is still too large, we need to slice it down to a
single row.
- // Normally, this should not happen, but it is possible if the data is
very large.
- if (sliceBatchCount == 0) {
- sliceBatchCount = 1
+ sliceRowCount = (rowCount *
SliceBytesArrowOutputProcessorImpl.stepRatio).toInt
+ }
+ // We iteratively shrink the slice by 10% until it fits. Each search
starts from the number of
+ // rows of the last slice that fit, since every slice that does not start
at row 0 copies its
+ // offset buffers, and restarting from all the remaining rows would copy
them again and again.
+ var numRows = math.min(sliceRowCount, rowCount - currentRowIdx)
+ while (numRows > 0) {
+ val rootSlice = root.slice(currentRowIdx, numRows)
+ if (batchBytes(rootSlice) <= arrowMaxBytesPerOutputBatch) {
+ sliceRowCount = numRows
+ currentRowIdx += numRows
+ return rootSlice
}
- rootSlice = root.slice(currentRowIdx, sliceBatchCount)
- currentRowIdx = currentRowIdx + rootSlice.getRowCount
- rootSlice
+ rootSlice.close()
+ numRows = (numRows * SliceBytesArrowOutputProcessorImpl.stepRatio).toInt
}
+ // Even a single row is larger than the limit, so we output a batch with
that one row.
+ // Normally, this should not happen, but it is possible if the data is
very large.
+ sliceRowCount = 1
+ val rootSlice = root.slice(currentRowIdx, 1)
+ currentRowIdx += 1
+ rootSlice
}
- private def getBatchBytes(root: VectorSchemaRoot): Long = {
- var batchBytes = 0L
- root.getFieldVectors.asScala.foreach { vector =>
- batchBytes += vector.getBufferSize
- }
- batchBytes
+ protected def batchBytes(root: VectorSchemaRoot): Long = {
Review Comment:
Both are `private[python]` now.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonArrowOutputSuite.scala:
##########
@@ -145,7 +152,117 @@ class PythonArrowOutputSuite extends SparkFunSuite {
checkAllRowsOnceInOrder(batchSizes, output)
// Guards the choice of cases: each has a batch that does not fit, which
must be sliced.
assert(output.size > batchSizes.size)
- assert(output.forall { case (rows, bytes) => rows.size <= 1 || bytes <
maxBytes })
+ assert(output.forall { case (rows, bytes) => rows.size <= 1 || bytes <=
maxBytes })
}
}
+
+ private def newSliceBytesProcessor(maxBytes: Int)
+ : (ArrowStreamReader, Map[String, SQLMetric]) => ArrowOutputProcessor =
+ (reader, metrics) => new SliceBytesArrowOutputProcessorImpl(reader,
metrics, maxBytes)
+
+ test("slice by bytes: a batch whose size equals the limit is not sliced") {
Review Comment:
Done, the new tests start with `SPARK-60395: `.
--
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]
