You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by vi...@apache.org on 2025/04/29 06:30:18 UTC
(spark) branch master updated: [SPARK-51931][SQL] Add maxBytesPerOutputBatch to limit the number of bytes of Arrow output batch
This is an automated email from the ASF dual-hosted git repository.
viirya pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 86bf4c84805e [SPARK-51931][SQL] Add maxBytesPerOutputBatch to limit the number of bytes of Arrow output batch
86bf4c84805e is described below
commit 86bf4c84805e89354d139ab72b298d3d4155fd0d
Author: Liang-Chi Hsieh <vi...@gmail.com>
AuthorDate: Mon Apr 28 23:30:00 2025 -0700
[SPARK-51931][SQL] Add maxBytesPerOutputBatch to limit the number of bytes of Arrow output batch
### What changes were proposed in this pull request?
This patch adds a new config `maxBytesPerOutputBatch` to limit the number of bytes of Arrow output batch.
### Why are the changes needed?
While implementing columnar-based operator for Spark, if the operator takes input from Arrow-based evaluation operator in Spark, the number of bytes of output batch is unlimited for now. For such columnar-based operator, sometimes we want to limit the maximum bytes of input batch. If we need to limit the batch size in bytes, it seems there is no existing way we can do.
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
Unit test
### Was this patch authored or co-authored using generative AI tooling?
No
Closes #50735 from viirya/arrow_output_bytes.
Authored-by: Liang-Chi Hsieh <vi...@gmail.com>
Signed-off-by: Liang-Chi Hsieh <vi...@gmail.com>
---
python/pyspark/sql/tests/arrow/test_arrow_map.py | 10 +-
.../org/apache/spark/sql/internal/SQLConf.scala | 21 ++++
.../sql/execution/python/PythonArrowOutput.scala | 117 ++++++++++++++++-----
3 files changed, 121 insertions(+), 27 deletions(-)
diff --git a/python/pyspark/sql/tests/arrow/test_arrow_map.py b/python/pyspark/sql/tests/arrow/test_arrow_map.py
index fa2ce69c4fa5..39f096db84bc 100644
--- a/python/pyspark/sql/tests/arrow/test_arrow_map.py
+++ b/python/pyspark/sql/tests/arrow/test_arrow_map.py
@@ -246,7 +246,7 @@ class MapInArrowWithArrowBatchSlicingTestsAndReducedBatchSizeTests(MapInArrowTes
cls.spark.conf.set("spark.sql.execution.arrow.maxBytesPerBatch", "10")
-class MapInArrowWithOutputArrowBatchSlicingTests(MapInArrowTests):
+class MapInArrowWithOutputArrowBatchSlicingRecordsTests(MapInArrowTests):
@classmethod
def setUpClass(cls):
MapInArrowTests.setUpClass()
@@ -254,6 +254,14 @@ class MapInArrowWithOutputArrowBatchSlicingTests(MapInArrowTests):
cls.spark.conf.set("spark.sql.execution.arrow.maxRecordsPerOutputBatch", "3")
+class MapInArrowWithOutputArrowBatchSlicingBytesTests(MapInArrowTests):
+ @classmethod
+ def setUpClass(cls):
+ MapInArrowTests.setUpClass()
+ cls.spark.conf.set("spark.sql.execution.arrow.maxRecordsPerBatch", "10")
+ cls.spark.conf.set("spark.sql.execution.arrow.maxBytesPerOutputBatch", "3")
+
+
if __name__ == "__main__":
from pyspark.sql.tests.arrow.test_arrow_map import * # noqa: F401
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
index 975c5feb6513..9fa968589783 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
@@ -3545,6 +3545,25 @@ object SQLConf {
.intConf
.createWithDefault(-1)
+ val ARROW_EXECUTION_MAX_BYTES_PER_OUTPUT_BATCH =
+ buildConf("spark.sql.execution.arrow.maxBytesPerOutputBatch")
+ .doc("When using Apache Arrow, limit the maximum bytes that can be output " +
+ "in a single ArrowRecordBatch to the downstream operator. If set to zero or negative " +
+ "there is no limit. Note that the complete ArrowRecordBatch is actually created but " +
+ "the number of bytes is limited when sending it to the downstream operator. This is " +
+ "used to avoid large batches being sent to the downstream operator including " +
+ "the columnar-based operator implemented by third-party libraries. Spark will try to " +
+ "create batches with the size equal or less than this value. Normally it should not " +
+ "happen, but if in extreme case that even one record is still larger than this value, " +
+ "Spark will create a batch with one record.")
+ .version("4.1.0")
+ .internal()
+ .bytesConf(ByteUnit.BYTE)
+ .checkValue(x => x <= Int.MaxValue,
+ errorMsg = "The value of " +
+ "spark.sql.execution.arrow.maxBytesPerOutputBatch should be less than INT_MAX.")
+ .createWithDefault(-1)
+
val ARROW_EXECUTION_MAX_BYTES_PER_BATCH =
buildConf("spark.sql.execution.arrow.maxBytesPerBatch")
.internal()
@@ -6547,6 +6566,8 @@ class SQLConf extends Serializable with Logging with SqlApiConf {
def arrowMaxRecordsPerOutputBatch: Int = getConf(ARROW_EXECUTION_MAX_RECORDS_PER_OUTPUT_BATCH)
+ def arrowMaxBytesPerOutputBatch: Long = getConf(ARROW_EXECUTION_MAX_BYTES_PER_OUTPUT_BATCH)
+
def arrowMaxBytesPerBatch: Long = getConf(ARROW_EXECUTION_MAX_BYTES_PER_BATCH)
def arrowTransformWithStateInPySparkMaxStateRecordsPerBatch: Int =
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala
index 9511b4c3f305..018619a5207d 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonArrowOutput.scala
@@ -46,6 +46,13 @@ private[python] trait PythonArrowOutput[OUT <: AnyRef] { self: BasePythonRunner[
protected def arrowMaxRecordsPerOutputBatch: Int = SQLConf.get.arrowMaxRecordsPerOutputBatch
+ protected def arrowMaxBytesPerOutputBatch: Int = SQLConf.get.arrowMaxBytesPerOutputBatch.toInt
+
+ if (arrowMaxRecordsPerOutputBatch > 0 && arrowMaxBytesPerOutputBatch > 0) {
+ throw new IllegalArgumentException(
+ "Both max records and max bytes of an Arrow output batch cannot be set at the same time.")
+ }
+
protected def newReaderIterator(
stream: DataInputStream,
writer: Writer,
@@ -106,8 +113,11 @@ private[python] trait PythonArrowOutput[OUT <: AnyRef] { self: BasePythonRunner[
schema = ArrowUtils.fromArrowSchema(root.getSchema())
if (arrowMaxRecordsPerOutputBatch > 0) {
- processor = new SliceArrowOutputProcessorImpl(
+ processor = new SliceRecordsArrowOutputProcessorImpl(
reader, pythonMetrics, arrowMaxRecordsPerOutputBatch)
+ } else if (arrowMaxBytesPerOutputBatch > 0) {
+ processor = new SliceBytesArrowOutputProcessorImpl(
+ reader, pythonMetrics, arrowMaxBytesPerOutputBatch)
} else {
processor = new ArrowOutputProcessorImpl(reader, pythonMetrics)
}
@@ -182,15 +192,24 @@ class ArrowOutputProcessorImpl(reader: ArrowStreamReader, pythonMetrics: Map[Str
}
}
-class SliceArrowOutputProcessorImpl(
+abstract class BaseSliceArrowOutputProcessor(
reader: ArrowStreamReader,
- pythonMetrics: Map[String, SQLMetric],
- arrowMaxRecordsPerOutputBatch: Int)
+ pythonMetrics: Map[String, SQLMetric])
extends ArrowOutputProcessorImpl(reader, pythonMetrics) {
- private var currentRowIdx = -1
- private var prevRoot: VectorSchemaRoot = null
- private var prevVectors: Array[ColumnVector] = _
+ protected var currentRowIdx = -1
+ protected var prevRoot: VectorSchemaRoot = null
+ protected var prevVectors: Array[ColumnVector] = _
+
+ override def loadBatch(): Boolean = {
+ if (rowCount > 0 && currentRowIdx < rowCount) {
+ true
+ } else {
+ val loaded = super.loadBatch()
+ currentRowIdx = 0
+ loaded
+ }
+ }
override def produceBatch(): ColumnarBatch = {
val batchRoot = getRoot
@@ -211,15 +230,74 @@ class SliceArrowOutputProcessorImpl(
batch
}
- override def loadBatch(): Boolean = {
- if (rowCount > 0 && currentRowIdx < rowCount) {
- true
+ protected override def getVectors(root: VectorSchemaRoot): Array[ColumnVector] = {
+ root.getFieldVectors.asScala.map { vector =>
+ new ArrowColumnVector(vector)
+ }.toArray[ColumnVector]
+ }
+
+ override def close(): Unit = {
+ if (prevRoot != null) {
+ prevVectors.foreach(_.close())
+ prevRoot.close()
+ }
+ }
+}
+
+class SliceBytesArrowOutputProcessorImpl(
+ reader: ArrowStreamReader,
+ pythonMetrics: Map[String, SQLMetric],
+ arrowMaxBytesPerOutputBatch: Int)
+ extends BaseSliceArrowOutputProcessor(reader, pythonMetrics) {
+
+ protected override def getRoot: VectorSchemaRoot = {
+ if (getBatchBytes(root) < arrowMaxBytesPerOutputBatch) {
+ currentRowIdx = currentRowIdx + root.getRowCount
+ root
} else {
- val loaded = super.loadBatch()
- currentRowIdx = 0
- loaded
+ // 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 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
+ }
+ rootSlice = root.slice(currentRowIdx, sliceBatchCount)
+ currentRowIdx = currentRowIdx + rootSlice.getRowCount
+ rootSlice
+ }
+ }
+
+ private def getBatchBytes(root: VectorSchemaRoot): Int = {
+ var batchBytes = 0
+ root.getFieldVectors.asScala.foreach { vector =>
+ batchBytes += vector.getBufferSize
}
+ batchBytes
}
+}
+
+class SliceRecordsArrowOutputProcessorImpl(
+ reader: ArrowStreamReader,
+ pythonMetrics: Map[String, SQLMetric],
+ arrowMaxRecordsPerOutputBatch: Int)
+ extends BaseSliceArrowOutputProcessor(reader, pythonMetrics) {
protected override def getRoot: VectorSchemaRoot = {
val remainingRows = rowCount - currentRowIdx
@@ -233,17 +311,4 @@ class SliceArrowOutputProcessorImpl(
rootSlice
}
-
- protected override def getVectors(root: VectorSchemaRoot): Array[ColumnVector] = {
- root.getFieldVectors.asScala.map { vector =>
- new ArrowColumnVector(vector)
- }.toArray[ColumnVector]
- }
-
- override def close(): Unit = {
- if (prevRoot != null) {
- prevVectors.foreach(_.close())
- prevRoot.close()
- }
- }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org