This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new cf5ba7f833eb [SPARK-57593][SQL][PYTHON] Byte-bound and instrument the
pickle Python UDF input batch
cf5ba7f833eb is described below
commit cf5ba7f833ebcc1d0479f69e318bb6d3adf71553
Author: Tengfei Huang <[email protected]>
AuthorDate: Tue Jun 23 08:44:20 2026 -0700
[SPARK-57593][SQL][PYTHON] Byte-bound and instrument the pickle Python UDF
input batch
### What changes were proposed in this pull request?
Regular batched Python UDFs (`BatchEvalPythonExec`) pickle each batch of
rows into a single contiguous on-heap byte array, and batches are formed by
**row count only** (`spark.sql.execution.python.udf.maxRecordsPerBatch`,
default 100). With wide rows, a single pickled batch can reach hundreds of MB
to over 1 GB of contiguous JVM heap and OOM the executor.
This PR adds an optional byte-size bound on these batches, plus
observability to size and verify the bound. All new behavior is off by default;
the default path is byte-for-byte unchanged.
1. **Byte cap** — new config
`spark.sql.execution.python.udf.maxBytesPerBatch` (internal, default `-1` =
off). When set to a finite value, a batch is cut at the **min** of the
row-count and byte limits. The per-row size is a best-effort estimate of the
*pickled* size of the converted row, accounted by `EvaluatePython.toJava` at
its leaf cases **during conversion itself** (no second traversal):
`UTF8String.numBytes` is the exact UTF-8 byte count pickle writes, decimals are
sized by pre [...]
2. **Observability** — new `SQLMetric`s on `BatchEvalPythonExec`:
- `pythonPeakPickledBatchBytes` — peak per-batch pickled size (size
metric; cross-task peak shown in the UI's min/med/max breakdown, like
`peakMemory`). Always recorded.
- `pythonOversizedBatchCount` — number of batches cut at the byte limit
(only populated when a finite cap is set).
- `pythonEstimatedInputBytes` — sum of the per-row size estimates, to
compare against the measured `pythonDataSent` for estimator-accuracy analysis
(only populated when a cap is set).
The Python UDTF path (`BatchEvalPythonUDTFExec`) calls `getInputIterator`
with no cap and is unaffected, only metrics added.
### Why are the changes needed?
Batching by row count alone gives no bound on the per-batch heap footprint:
with wide rows (large strings/binary/decimal payloads), one pickled batch
becomes a very large contiguous allocation that can OOM the executor even when
the row *count* is small. The byte cap provides a guardrail against that
allocation, and the metrics make peak batch pressure observable so a safe
threshold can be chosen and the estimator's accuracy verified before enforcing.
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
New unit tests in `BatchEvalPythonExecSuite`
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 4.8)
Closes #56649 from ivoson/SPARK-57593.
Authored-by: Tengfei Huang <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
(cherry picked from commit 9a639fddfc9363d803005ba9f99ba1d023a806d8)
Signed-off-by: Wenchen Fan <[email protected]>
---
.../org/apache/spark/sql/internal/SQLConf.scala | 19 ++
.../sql/execution/python/BatchEvalPythonExec.scala | 73 ++++++-
.../execution/python/BatchEvalPythonUDTFExec.scala | 9 +-
.../python/ByteBoundedAsArrayIterator.scala | 128 +++++++++++++
.../sql/execution/python/EvaluatePython.scala | 52 +++--
.../sql/execution/python/PythonSQLMetrics.scala | 49 ++++-
.../python/BatchEvalPythonExecSuite.scala | 211 ++++++++++++++++++++-
.../sql/execution/python/PythonUDFSuite.scala | 32 ++++
8 files changed, 546 insertions(+), 27 deletions(-)
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 db74d0378fc1..05043b3107bf 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
@@ -4777,6 +4777,25 @@ object SQLConf {
"must be positive.")
.createWithDefault(100)
+ val PYTHON_UDF_MAX_BYTES_PER_BATCH =
+ buildConf("spark.sql.execution.python.udf.maxBytesPerBatch")
+ .internal()
+ .doc("Byte-size cap for a batch of rows sent to a worker for regular
(pickle-serialized, " +
+ "non-Arrow) Python UDF evaluation. Without it, BatchEvalPythonExec
batches purely by " +
+ "row count (spark.sql.execution.python.udf.maxRecordsPerBatch), so one
oversized batch " +
+ "can OOM the executor; a finite value caps the batch at the min of the
row-count and " +
+ "byte limits. The size is a best-effort per-row estimate of the
pickled size of the " +
+ "converted row (its accuracy is observable via the
pythonEstimatedInputBytes vs " +
+ "pythonDataSent metrics); a row larger than the cap still yields a
one-row batch. " +
+ "-1 (the default) means no limit.")
+ .version("4.3.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .bytesConf(ByteUnit.BYTE)
+ .checkValue(x => x == -1 || (x > 0 && x <= Int.MaxValue),
+ "The value of spark.sql.execution.python.udf.maxBytesPerBatch should "
+
+ "be -1 (no limit) or greater than zero and less than or equal to
INT_MAX.")
+ .createWithDefault(-1)
+
val PYTHON_UDF_BUFFER_SIZE =
buildConf("spark.sql.execution.python.udf.buffer.size")
.doc(
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExec.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExec.scala
index a42c8ac8da53..4c39ce98cf55 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExec.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExec.scala
@@ -17,6 +17,8 @@
package org.apache.spark.sql.execution.python
+import java.io.ByteArrayOutputStream
+
import scala.jdk.CollectionConverters._
import net.razorvine.pickle.{Pickler, Unpickler}
@@ -35,7 +37,7 @@ import org.apache.spark.sql.types.{StructField, StructType}
* A physical plan that evaluates a [[PythonUDF]]
*/
case class BatchEvalPythonExec(udfs: Seq[PythonUDF], resultAttrs:
Seq[Attribute], child: SparkPlan)
- extends EvalPythonExec with PythonSQLMetrics {
+ extends EvalPythonExec with PythonPickleBatchMetrics {
private[this] val jobArtifactUUID =
JobArtifactSet.getCurrentJobArtifactState.map(_.uuid)
private[this] val sessionUUID = {
@@ -47,12 +49,14 @@ case class BatchEvalPythonExec(udfs: Seq[PythonUDF],
resultAttrs: Seq[Attribute]
override protected def evaluatorFactory: EvalPythonEvaluatorFactory = {
val batchSize = conf.getConf(SQLConf.PYTHON_UDF_MAX_RECORDS_PER_BATCH)
+ val maxBytesPerBatch = conf.getConf(SQLConf.PYTHON_UDF_MAX_BYTES_PER_BATCH)
val binaryAsBytes = conf.pysparkBinaryAsBytes
new BatchEvalPythonEvaluatorFactory(
child.output,
udfs,
output,
batchSize,
+ maxBytesPerBatch,
pythonMetrics,
jobArtifactUUID,
sessionUUID,
@@ -68,6 +72,7 @@ class BatchEvalPythonEvaluatorFactory(
udfs: Seq[PythonUDF],
output: Seq[Attribute],
batchSize: Int,
+ maxBytesPerBatch: Long,
pythonMetrics: Map[String, SQLMetric],
jobArtifactUUID: Option[String],
sessionUUID: Option[String],
@@ -83,7 +88,8 @@ class BatchEvalPythonEvaluatorFactory(
EvaluatePython.registerPicklers() // register pickler for Row
// Input iterator to Python.
- val inputIterator = BatchEvalPythonExec.getInputIterator(iter, schema,
batchSize, binaryAsBytes)
+ val inputIterator = BatchEvalPythonExec.getInputIterator(
+ iter, schema, batchSize, binaryAsBytes, maxBytesPerBatch, pythonMetrics)
// Output iterator for results from Python.
val outputIterator =
@@ -123,7 +129,12 @@ object BatchEvalPythonExec {
iter: Iterator[InternalRow],
schema: StructType,
batchSize: Int,
- binaryAsBytes: Boolean): Iterator[Array[Byte]] = {
+ binaryAsBytes: Boolean,
+ maxBytesPerBatch: Long = -1L,
+ pythonMetrics: Map[String, SQLMetric] = Map.empty):
Iterator[Array[Byte]] = {
+ val peakPickledBatchBytesMetric =
pythonMetrics.get("pythonPeakPickledBatchBytes")
+ val oversizedBatchMetric = pythonMetrics.get("pythonOversizedBatchCount")
+ val estimatedInputBytesMetric =
pythonMetrics.get("pythonEstimatedInputBytes")
val dataTypes = schema.map(_.dataType)
val needConversion =
dataTypes.exists(EvaluatePython.needConversionInPython)
@@ -140,22 +151,66 @@ object BatchEvalPythonExec {
// Please note that cache by reference is still enabled depending on
`needConversion`.
val pickle = new Pickler(/* useMemo = */ needConversion,
/* valueCompare = */ false)
- // Input iterator to Python: input rows are grouped so we send them in
batches to Python.
- // For each row, add it to the queue.
- iter.map { row =>
+
+ // Converts a row to the java object pickled to Python. When `sizeAcc` is
defined, toJava also
+ // accumulates the per-row pickled-size estimate at its leaf cases during
this same traversal
+ // (no second walk).
+ def convertRow(row: InternalRow, sizeAcc: Option[PickledSizeAccumulator]):
Any = {
if (needConversion) {
- EvaluatePython.toJava(row, schema, binaryAsBytes)
+ EvaluatePython.toJava(row, schema, binaryAsBytes, sizeAcc)
} else {
// fast path for these types that does not need conversion in Python
+ sizeAcc.foreach(_.add(PickledSizeAccumulator.PER_VALUE_OVERHEAD))
val fields = new Array[Any](row.numFields)
var i = 0
while (i < row.numFields) {
val dt = dataTypes(i)
- fields(i) = EvaluatePython.toJava(row.get(i, dt), dt, binaryAsBytes)
+ fields(i) = EvaluatePython.toJava(row.get(i, dt), dt, binaryAsBytes,
sizeAcc)
i += 1
}
fields
}
- }.grouped(batchSize).map(x => pickle.dumps(x.toArray))
+ }
+
+ // Input iterator to Python: input rows are grouped so we send them in
batches to Python.
+ val batchedIter: Iterator[Array[Any]] =
+ if (maxBytesPerBatch < 0) {
+ // No byte limit configured (the default, -1): preserve exact
row-count batching behavior.
+ iter.map(convertRow(_, None)).grouped(batchSize).map(_.toArray)
+ } else {
+ // A finite byte cap is set, so batch by both row count and estimated
bytes (see
+ // ByteBoundedAsArrayIterator). The size is estimated on the converted
values (the things
+ // actually pickled), accounted by toJava during conversion itself;
estimator accuracy is
+ // observable via pythonEstimatedInputBytes vs pythonDataSent.
+ val sizeAcc = new PickledSizeAccumulator
+ val someSizeAcc = Some(sizeAcc)
+ new ByteBoundedAsArrayIterator(
+ iter.map { row =>
+ val converted = convertRow(row, someSizeAcc)
+ (converted, sizeAcc.getAndReset())
+ },
+ batchSize,
+ maxBytesPerBatch,
+ oversizedBatchMetric,
+ estimatedInputBytesMetric)
+ }
+
+ batchedIter.map { objects =>
+ val baos = new ByteArrayOutputStream(1024)
+ pickle.dump(objects, baos)
+ // Record the peak pickled-batch size: the primary contiguous JVM-heap
allocation on this
+ // path, and one oversized batch can OOM the executor. Read baos.size()
BEFORE the
+ // toByteArray copy, so the measurement does not depend on that second
contiguous allocation
+ // succeeding (a near-limit batch would OOM there, otherwise losing the
metric). The
+ // `> value` guard keeps the per-task max (`set` overwrites); SQLMetric
surfaces the
+ // cross-task max in the UI's (min, med, max) breakdown, mirroring how
peakMemory is reported.
+ peakPickledBatchBytesMetric.foreach { metric =>
+ val pickledBytes = baos.size().toLong
+ if (pickledBytes > metric.value) {
+ metric.set(pickledBytes)
+ }
+ }
+ baos.toByteArray
+ }
}
}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/BatchEvalPythonUDTFExec.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/BatchEvalPythonUDTFExec.scala
index 9a7500d1b4cb..a2d94c226b7f 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/BatchEvalPythonUDTFExec.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/BatchEvalPythonUDTFExec.scala
@@ -48,7 +48,7 @@ case class BatchEvalPythonUDTFExec(
requiredChildOutput: Seq[Attribute],
resultAttrs: Seq[Attribute],
child: SparkPlan)
- extends EvalPythonUDTFExec with PythonSQLMetrics {
+ extends EvalPythonUDTFExec with PythonPickleBatchMetrics {
private[this] val jobArtifactUUID =
JobArtifactSet.getCurrentJobArtifactState.map(_.uuid)
private[this] val sessionUUID = {
@@ -70,9 +70,12 @@ case class BatchEvalPythonUDTFExec(
EvaluatePython.registerPicklers() // register pickler for Row
// Input iterator to Python.
- // For Python UDTF, we don't have a separate configuration for the batch
size yet.
+ // For Python UDTF, we don't have a separate configuration for the batch
size yet, and no byte
+ // cap (maxBytesPerBatch stays at the -1 default). We still pass
pythonMetrics so the peak
+ // pickled-batch size is recorded: the UDTF path pickles through the same
contiguous-allocation
+ // code as BatchEvalPythonExec and carries the same OOM risk, so the
observability applies here.
val inputIterator = BatchEvalPythonExec.getInputIterator(
- iter, schema, 100, conf.pysparkBinaryAsBytes)
+ iter, schema, 100, conf.pysparkBinaryAsBytes, pythonMetrics =
pythonMetrics)
// Output iterator for results from Python.
val outputIterator =
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ByteBoundedAsArrayIterator.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ByteBoundedAsArrayIterator.scala
new file mode 100644
index 000000000000..5df731fbab2e
--- /dev/null
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ByteBoundedAsArrayIterator.scala
@@ -0,0 +1,128 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.python
+
+import scala.collection.mutable.ArrayBuffer
+
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.unsafe.types.VariantVal
+
+private[python] object PickledSizeAccumulator {
+ // Pickle framing overhead per value (opcode + length prefix), a coarse
constant.
+ val PER_VALUE_OVERHEAD = 5L
+
+ // Unknown leaf type (timestamp objects, intervals, ...): a small constant,
never zero,
+ // so the cap degrades to a loose bound instead of going blind.
+ val UNKNOWN_VALUE_SIZE = 32L
+}
+
+/**
+ * Accumulates a best-effort estimate of the pickled size of converted values,
fed by the leaf
+ * cases of [[EvaluatePython.toJava]] DURING conversion, so estimating costs
no second traversal
+ * of the fields. The source InternalRow cannot be used for sizing instead:
EvalPythonExec feeds
+ * this path from a MutableProjection whose target is a GenericInternalRow, so
an UnsafeRow-based
+ * source estimate never applies. String/binary/decimal/geo/variant payloads
dominate pickled size
+ * and are sized from their byte lengths at the catalyst leaves
(UTF8String.numBytes is the UTF-8
+ * byte count pickle writes); residual error (pickle overhead, memoization,
unknown leaves) is
+ * observable by comparing the pythonEstimatedInputBytes metric against the
measured pythonDataSent.
+ *
+ * Single-threaded: one instance per partition, reset per row via
[[getAndReset]].
+ */
+private[python] final class PickledSizeAccumulator {
+ import PickledSizeAccumulator._
+
+ private var _size = 0L
+
+ /** Adds raw bytes (container/framing overhead). */
+ def add(n: Long): Unit = _size += n
+
+ /** Adds a leaf value with a known payload size, plus the per-value framing
overhead. */
+ def addValue(payloadBytes: Long): Unit = _size += payloadBytes +
PER_VALUE_OVERHEAD
+
+ /** Adds a pass-through leaf (boxed primitives and anything else toJava
leaves untouched). */
+ def addLeaf(value: Any): Unit = value match {
+ case null => _size += 1L
+ case _: java.lang.Boolean => _size += 2L
+ case _: java.lang.Byte | _: java.lang.Short => _size += PER_VALUE_OVERHEAD
+ case _: java.lang.Number => _size += 4L + PER_VALUE_OVERHEAD //
Int/Long/Float/Double box
+ case v: VariantVal =>
+ // Variant payloads can be arbitrarily large; size by the serialized
byte lengths.
+ _size += v.getValue.length.toLong + v.getMetadata.length.toLong +
PER_VALUE_OVERHEAD
+ case _ => _size += UNKNOWN_VALUE_SIZE
+ }
+
+ /** Returns the accumulated estimate and resets, closing one per-row
accounting cycle. */
+ def getAndReset(): Long = {
+ val s = _size
+ _size = 0L
+ s
+ }
+}
+
+/**
+ * Groups converted input objects into batches bounded by both a record count
and an estimated
+ * byte size (the second tuple element is the per-row estimate of the object's
pickled size, see
+ * [[PickledSizeAccumulator]], fed by toJava during conversion). A batch
always holds at least one
+ * row, so a single oversized row still forms a one-row batch.
+ *
+ * A batch is cut once its accumulated estimate reaches the cap, so it can
exceed the cap by
+ * the last row. Each cut batch increments `oversizedBatchMetric` once, and
+ * `estimatedInputBytesMetric` accumulates the per-row estimates for
comparison against the
+ * measured pythonDataSent (the estimator-accuracy signal).
+ */
+private[python] class ByteBoundedAsArrayIterator(
+ iter: Iterator[(Any, Long)],
+ maxRecordsPerBatch: Int,
+ maxBytesPerBatch: Long,
+ oversizedBatchMetric: Option[SQLMetric],
+ estimatedInputBytesMetric: Option[SQLMetric])
+ extends Iterator[Array[Any]] {
+
+ // Parity with the row-count batching path: a non-positive limit would loop
forever emitting
+ // empty batches. The config already enforces `> 0`; this is defensive.
+ require(maxRecordsPerBatch > 0, "max records per batch must be positive")
+
+ // Only finite positive caps may reach this class (BatchEvalPythonExec
routes the -1 "no limit"
+ // sentinel to the row-count-only batcher, and the config rejects 0).
Defensive: a cap of 0
+ // would silently degrade every batch to a single row instead of failing
loudly.
+ require(maxBytesPerBatch > 0, "max bytes per batch must be positive")
+
+ override def hasNext: Boolean = iter.hasNext
+
+ override def next(): Array[Any] = {
+ if (!hasNext) {
+ throw new NoSuchElementException
+ }
+ val batch = new ArrayBuffer[Any]()
+ var accumulatedBytes = 0L
+ var cut = false
+ while (!cut && iter.hasNext && batch.length < maxRecordsPerBatch) {
+ val (obj, sizeBytes) = iter.next()
+ // Sum the raw estimate for estimator-accuracy analysis (vs
pythonDataSent).
+ estimatedInputBytesMetric.foreach(_ += sizeBytes)
+ batch += obj
+ accumulatedBytes += sizeBytes
+ if (accumulatedBytes >= maxBytesPerBatch) {
+ cut = true
+ }
+ }
+ // Count each batch cut at the byte limit once.
+ if (cut) oversizedBatchMetric.foreach(_ += 1)
+ batch.toArray
+ }
+}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/EvaluatePython.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/EvaluatePython.scala
index cb1acd2f541b..3b9c2a3e69cd 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/EvaluatePython.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/EvaluatePython.scala
@@ -59,56 +59,86 @@ object EvaluatePython {
/**
* Helper for converting from Catalyst type to java type suitable for Pickle.
+ *
+ * When `sizeAcc` is defined, the conversion additionally accumulates a
best-effort estimate of
+ * the PICKLED size of the converted value, accounted at the leaf cases
during the same
+ * traversal (a separate estimation pass would walk every field a second
time). Catalyst leaves
+ * carry exact payload sizes (UTF8String.numBytes is the UTF-8 byte count
that pickle writes;
+ * Decimal.precision tracks the digit string). The estimate is best-effort:
unknown leaf types
+ * contribute a small positive constant, and residual error is observable by
comparing the
+ * pythonEstimatedInputBytes metric against pythonDataSent. INVARIANT: every
leaf case must
+ * account something positive to sizeAcc -- a case that converts without
accounting makes the
+ * byte cap (see [[BatchEvalPythonExec.getInputIterator]]) blind to that
type. The type sweep in
+ * BatchEvalPythonExecSuite ("estimate is positive for every data type")
enforces this.
*/
def toJava(
obj: Any,
dataType: DataType,
- binaryAsBytes: Boolean): Any = {
+ binaryAsBytes: Boolean,
+ sizeAcc: Option[PickledSizeAccumulator] = None): Any = {
(obj, dataType) match {
- case (null, _) => null
+ case (null, _) =>
+ sizeAcc.foreach(_.add(1L))
+ null
case (row: InternalRow, struct: StructType) =>
+ sizeAcc.foreach(_.add(PickledSizeAccumulator.PER_VALUE_OVERHEAD))
val values = new Array[Any](row.numFields)
var i = 0
while (i < row.numFields) {
val field = struct.fields(i)
- values(i) = toJava(row.get(i, field.dataType), field.dataType,
binaryAsBytes)
+ values(i) = toJava(row.get(i, field.dataType), field.dataType,
binaryAsBytes, sizeAcc)
i += 1
}
new GenericRowWithSchema(values, struct)
case (a: ArrayData, array: ArrayType) =>
+ sizeAcc.foreach(_.add(PickledSizeAccumulator.PER_VALUE_OVERHEAD))
val values = new java.util.ArrayList[Any](a.numElements())
a.foreach(array.elementType, (_, e) => {
- values.add(toJava(e, array.elementType, binaryAsBytes))
+ values.add(toJava(e, array.elementType, binaryAsBytes, sizeAcc))
})
values
case (map: MapData, mt: MapType) =>
+ sizeAcc.foreach(_.add(PickledSizeAccumulator.PER_VALUE_OVERHEAD))
val jmap = new java.util.HashMap[Any, Any](map.numElements())
map.foreach(mt.keyType, mt.valueType, (k, v) => {
- jmap.put(toJava(k, mt.keyType, binaryAsBytes), toJava(v,
mt.valueType, binaryAsBytes))
+ jmap.put(toJava(k, mt.keyType, binaryAsBytes, sizeAcc),
+ toJava(v, mt.valueType, binaryAsBytes, sizeAcc))
})
jmap
- case (ud, udt: UserDefinedType[_]) => toJava(ud, udt.sqlType,
binaryAsBytes)
+ case (ud, udt: UserDefinedType[_]) => toJava(ud, udt.sqlType,
binaryAsBytes, sizeAcc)
- case (d: Decimal, _) => d.toJavaBigDecimal
+ case (d: Decimal, _) =>
+ sizeAcc.foreach(_.addValue(d.precision.toLong))
+ d.toJavaBigDecimal
- case (s: UTF8String, _: StringType) => s.toString
+ case (s: UTF8String, _: StringType) =>
+ sizeAcc.foreach(_.addValue(s.numBytes.toLong))
+ s.toString
- case (g: BinaryView, gt: GeometryType) => STUtils.deserializeGeom(g, gt)
+ case (g: BinaryView, gt: GeometryType) =>
+ // Geometry payloads can be arbitrarily large; size by the serialized
byte length.
+ sizeAcc.foreach(_.addValue(g.numBytes.toLong))
+ STUtils.deserializeGeom(g, gt)
- case (g: BinaryView, gt: GeographyType) => STUtils.deserializeGeog(g, gt)
+ case (g: BinaryView, gt: GeographyType) =>
+ sizeAcc.foreach(_.addValue(g.numBytes.toLong))
+ STUtils.deserializeGeog(g, gt)
case (bytes: Array[Byte], BinaryType) =>
+ sizeAcc.foreach(_.addValue(bytes.length.toLong))
if (binaryAsBytes) {
new BytesWrapper(bytes)
} else {
bytes
}
- case (other, _) => other
+ case (other, _) =>
+ sizeAcc.foreach(_.addLeaf(other))
+ other
}
}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonSQLMetrics.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonSQLMetrics.scala
index cbce07977f16..81df5aaad5c2 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonSQLMetrics.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonSQLMetrics.scala
@@ -28,9 +28,13 @@ trait PythonSQLMetrics { self: SparkPlan =>
k -> SQLMetrics.createTimingMetric(sparkContext, v)
} ++ PythonSQLMetrics.pythonOtherMetricsDesc.map { case (k, v) =>
k -> SQLMetrics.createMetric(sparkContext, v)
- }
+ } ++ additionalMetrics
}
+ // Hook for subtraits to add operator-specific metrics on top of the shared
ones above. Default
+ // empty so the shared trait stays minimal; see [[PythonPickleBatchMetrics]].
+ protected def additionalMetrics: Map[String, SQLMetric] = Map.empty
+
override lazy val metrics: Map[String, SQLMetric] = pythonMetrics
}
@@ -55,3 +59,46 @@ object PythonSQLMetrics {
Map("pythonNumRowsReceived" -> "number of output rows")
}
}
+
+/**
+ * Metrics specific to the regular (pickle-serialized) Python UDF input
batching path in
+ * [[BatchEvalPythonExec.getInputIterator]]. Mixed in only by the operators
that pickle their
+ * input ([[BatchEvalPythonExec]] and [[BatchEvalPythonUDTFExec]]); kept out
of the shared
+ * [[PythonSQLMetrics]] so they do not surface as always-zero rows on the
Arrow / streaming Python
+ * operators that mix in PythonSQLMetrics but never pickle through
getInputIterator.
+ *
+ * - pythonPeakPickledBatchBytes: peak per-batch pickled size (the primary
contiguous heap
+ * allocation on this path). Like peakMemory it is a SIZE metric, so the
figure to read is the
+ * per-partition max in the UI's min/med/max breakdown; the aggregated
total (the sum of the
+ * per-task peaks) is not meaningful on its own. Always recorded on the
pickle path.
+ * - pythonEstimatedInputBytes: running sum of the per-row size estimates
the byte cap uses,
+ * compared against the measured pythonDataSent to gauge estimator
accuracy. Only populated
+ * when a byte cap is configured.
+ * - pythonOversizedBatchCount: input batches cut at the byte cap, counted
once per cut batch.
+ * Only populated when a finite
spark.sql.execution.python.udf.maxBytesPerBatch is set.
+ *
+ * Also kept out of pythonSizeMetricsDesc, which feeds the Python data source
DSv2 metric
+ * declarations (UserDefinedPythonDataSource.createPythonMetrics) that never
populate these.
+ */
+trait PythonPickleBatchMetrics extends PythonSQLMetrics { self: SparkPlan =>
+ override protected def additionalMetrics: Map[String, SQLMetric] =
+ PythonPickleBatchMetrics.pickleBatchSizeMetricsDesc.map { case (k, v) =>
+ k -> SQLMetrics.createSizeMetric(sparkContext, v)
+ } ++ PythonPickleBatchMetrics.pickleBatchCountMetricsDesc.map { case (k,
v) =>
+ k -> SQLMetrics.createMetric(sparkContext, v)
+ }
+}
+
+object PythonPickleBatchMetrics {
+ // Peak pickled batch size and estimated input size, reported as SIZE
metrics.
+ val pickleBatchSizeMetricsDesc: Map[String, String] = {
+ Map(
+ "pythonPeakPickledBatchBytes" -> "peak pickled batch size in bytes",
+ "pythonEstimatedInputBytes" -> "estimated pickled input size in bytes")
+ }
+
+ // Count of batches cut at the byte cap, reported as a SUM metric.
+ val pickleBatchCountMetricsDesc: Map[String, String] = {
+ Map("pythonOversizedBatchCount" -> "number of batches cut at the byte
limit")
+ }
+}
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExecSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExecSuite.scala
index a10689c1226b..42da9f060b85 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExecSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExecSuite.scala
@@ -21,13 +21,15 @@ import scala.collection.mutable.ArrayBuffer
import scala.jdk.CollectionConverters._
import org.apache.spark.api.python.{PythonEvalType, SimplePythonFunction}
-import org.apache.spark.sql.catalyst.FunctionIdentifier
-import org.apache.spark.sql.catalyst.expressions.{And, AttributeReference,
GreaterThan, In}
+import org.apache.spark.sql.catalyst.{FunctionIdentifier, InternalRow}
+import org.apache.spark.sql.catalyst.expressions.{And, AttributeReference,
GenericInternalRow, GreaterThan, In}
import org.apache.spark.sql.connector.catalog.CatalogManager
import org.apache.spark.sql.execution.{FilterExec, InputAdapter,
WholeStageCodegenExec}
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
import org.apache.spark.sql.test.SharedSparkSession
-import org.apache.spark.sql.types.{BooleanType, DoubleType}
+import org.apache.spark.sql.types._
+import org.apache.spark.unsafe.types.UTF8String
class BatchEvalPythonExecSuite extends SharedSparkSession
with AdaptiveSparkPlanHelper {
@@ -52,6 +54,209 @@ class BatchEvalPythonExecSuite extends SharedSparkSession
}
}
+ test("SPARK-57593: getInputIterator records peak pickled batch size as a
size metric") {
+ EvaluatePython.registerPicklers()
+ val schema = StructType(Seq(StructField("s", StringType)))
+ // batchSize = 1 so each row is its own pickled batch; the widest row
yields the peak.
+ val batchSize = 1
+ val rows = Seq(
+ InternalRow(UTF8String.fromString("a")),
+ InternalRow(UTF8String.fromString("b" * 1000)),
+ InternalRow(UTF8String.fromString("c")))
+
+ // Baseline: the peak pickled-batch size, computed without instrumentation.
+ val expectedPeak = BatchEvalPythonExec
+ .getInputIterator(rows.iterator, schema, batchSize, binaryAsBytes =
false)
+ .map(_.length.toLong)
+ .max
+ assert(expectedPeak > 0)
+
+ val metric = SQLMetrics.createSizeMetric(spark.sparkContext, "peak pickled
batch size in bytes")
+ // Drain the iterator so every batch is pickled and measured.
+ BatchEvalPythonExec
+ .getInputIterator(rows.iterator, schema, batchSize, binaryAsBytes =
false,
+ pythonMetrics = Map("pythonPeakPickledBatchBytes" -> metric))
+ .foreach(_ => ())
+
+ assert(metric.value === expectedPeak)
+ }
+
+ test("SPARK-57593: ByteBoundedAsArrayIterator oversized-batch and
estimated-bytes metrics") {
+ def elems(sizes: Long*): Iterator[(Any, Long)] =
+ sizes.iterator.map(s => (0.asInstanceOf[Any], s))
+ def m(): SQLMetric = SQLMetrics.createMetric(spark.sparkContext, "m")
+
+ // Cut once accumulated bytes >= 25 (after 3 rows of 10).
estimatedInputBytes sums the
+ // per-row estimate across the whole partition.
+ val ov1 = m(); val est1 = m()
+ val byBytes = new ByteBoundedAsArrayIterator(
+ elems(10, 10, 10, 10, 10, 10, 10), maxRecordsPerBatch = 100,
maxBytesPerBatch = 25L,
+ Some(ov1), Some(est1)).toList.map(_.length)
+ assert(byBytes === List(3, 3, 1))
+ assert(ov1.value === 2) // two batches were cut at the cap; the trailing
1-row batch was not
+ assert(est1.value === 70) // sum of per-row estimates
+
+ // Record cap dominates when smaller; the byte limit is never reached.
+ val ov2 = m()
+ val byCount = new ByteBoundedAsArrayIterator(
+ elems(10, 10, 10, 10, 10), maxRecordsPerBatch = 2, maxBytesPerBatch =
1000L,
+ Some(ov2), None).toList.map(_.length)
+ assert(byCount === List(2, 2, 1))
+ assert(ov2.value === 0)
+
+ // A single row larger than the cap still forms a one-row batch (never
zero rows).
+ val ov3 = m()
+ val giant = new ByteBoundedAsArrayIterator(
+ elems(10000), maxRecordsPerBatch = 100, maxBytesPerBatch = 25L,
+ Some(ov3), None).toList.map(_.length)
+ assert(giant === List(1))
+ assert(ov3.value === 1)
+
+ // Size-0 estimates contribute nothing to the estimate sum or the byte
cap; the cut happens
+ // on the first row that pushes the estimate to the cap.
+ val ov5 = m(); val est5 = m()
+ val mixed = new ByteBoundedAsArrayIterator(
+ elems(0, 0, 30, 0), maxRecordsPerBatch = 100, maxBytesPerBatch = 25L,
+ Some(ov5), Some(est5)).toList.map(_.length)
+ assert(mixed === List(3, 1))
+ assert(est5.value === 30) // only the size-30 row contributes
+ assert(ov5.value === 1) // one batch cut at the cap
+ }
+
+ test("SPARK-57593: ByteBoundedAsArrayIterator honors the Iterator contract")
{
+ def elems(sizes: Long*): Iterator[(Any, Long)] =
+ sizes.iterator.map(s => (0.asInstanceOf[Any], s))
+ def iterOf(input: Iterator[(Any, Long)], maxRecords: Int = 100):
Iterator[Array[Any]] =
+ new ByteBoundedAsArrayIterator(
+ input, maxRecordsPerBatch = maxRecords, maxBytesPerBatch =
Int.MaxValue.toLong,
+ None, None)
+
+ // next() past the end throws NoSuchElementException, matching the
row-count batching path.
+ val it = iterOf(elems(10))
+ assert(it.hasNext)
+ assert(it.next().length === 1)
+ assert(!it.hasNext)
+ intercept[NoSuchElementException](it.next())
+
+ // An empty input yields no batches and next() throws immediately.
+ val empty = iterOf(elems())
+ assert(!empty.hasNext)
+ intercept[NoSuchElementException](empty.next())
+
+ // A non-positive record limit is rejected up front.
+ intercept[IllegalArgumentException](iterOf(elems(10), maxRecords = 0))
+
+ // A non-positive byte cap is rejected up front too: only finite positive
caps may reach this
+ // class (a cap of 0 would otherwise silently degrade every batch to a
single row).
+ intercept[IllegalArgumentException](new ByteBoundedAsArrayIterator(
+ elems(10), maxRecordsPerBatch = 100, maxBytesPerBatch = 0L, None, None))
+ }
+
+ test("SPARK-57593: getInputIterator byte-bounds pickle batches when a cap is
configured") {
+ EvaluatePython.registerPicklers()
+ val schema = StructType(Seq(StructField("s", StringType)))
+ // GenericInternalRow on purpose: production feeds this path from a
MutableProjection whose
+ // target is a GenericInternalRow (never an UnsafeRow), so the byte cap
must work on it.
+ def wideRows(): Iterator[InternalRow] = (0 until 10).iterator.map { _ =>
+ new GenericInternalRow(Array[Any](UTF8String.fromString("x" * 1000)))
+ }
+
+ // Default: no byte cap, so all 10 rows form a single batch (batchSize =
100).
+ val defaultBatches = BatchEvalPythonExec
+ .getInputIterator(wideRows(), schema, batchSize = 100, binaryAsBytes =
false).size
+ assert(defaultBatches === 1)
+
+ // Finite cap: the wide rows split into multiple smaller batches, and the
estimate metric is
+ // populated (regression test: the estimate must be non-zero for
production row classes).
+ val estMetric = SQLMetrics.createSizeMetric(spark.sparkContext, "est")
+ val cappedBatches = BatchEvalPythonExec
+ .getInputIterator(wideRows(), schema, batchSize = 100, binaryAsBytes =
false,
+ maxBytesPerBatch = 2048L,
+ pythonMetrics = Map("pythonEstimatedInputBytes" -> estMetric)).size
+ assert(cappedBatches > defaultBatches)
+ assert(estMetric.value >= 10 * 1000L) // 10 rows of a 1000-char string,
plus overhead
+ }
+
+ test("SPARK-57593: getInputIterator does not byte-bound when
maxBytesPerBatch is -1 (no limit)") {
+ EvaluatePython.registerPicklers()
+ val schema = StructType(Seq(StructField("s", StringType)))
+ // 250 wide rows: more than batchSize so row-count batching genuinely
splits them, and each row
+ // is wide enough that any finite byte cap would split them into far more
batches (the sibling
+ // test above cuts 10 of these at a 2048-byte cap).
+ def wideRows(): Iterator[InternalRow] = (0 until 250).iterator.map { _ =>
+ new GenericInternalRow(Array[Any](UTF8String.fromString("x" * 1000)))
+ }
+
+ // -1 (the default) routes to the row-count-only batcher: batches are cut
purely at batchSize
+ // (250 rows / 100 = 3 batches) with no byte bound, and the byte-only
metrics stay unpopulated.
+ val estMetric = SQLMetrics.createSizeMetric(spark.sparkContext, "est")
+ val oversizedMetric = SQLMetrics.createMetric(spark.sparkContext, "ov")
+ val batches = BatchEvalPythonExec
+ .getInputIterator(wideRows(), schema, batchSize = 100, binaryAsBytes =
false,
+ maxBytesPerBatch = -1L,
+ pythonMetrics = Map(
+ "pythonEstimatedInputBytes" -> estMetric,
+ "pythonOversizedBatchCount" -> oversizedMetric)).size
+ // Row-count only: 3 batches (100 + 100 + 50). A finite byte cap would
produce many more.
+ assert(batches === 3)
+ assert(estMetric.value === 0L) // no per-row estimate accumulated on
the no-limit path
+ assert(oversizedMetric.value === 0L) // no batch cut at a byte limit
+ }
+
+ test("SPARK-57593: toJava accumulates the pickled-size estimate during
conversion") {
+ val acc = new PickledSizeAccumulator
+ def sized(value: Any, dataType: DataType): Long = {
+ EvaluatePython.toJava(value, dataType, binaryAsBytes = true, Some(acc))
+ acc.getAndReset()
+ }
+ // Payload bytes dominate the estimate; UTF8String.numBytes is the exact
UTF-8 byte count.
+ val wide = sized(UTF8String.fromString("x" * 10000), StringType)
+ assert(wide >= 10000L && wide <= 10000L + 64L)
+ val binary = sized(Array.fill[Byte](5000)(1), BinaryType)
+ assert(binary >= 5000L && binary <= 5000L + 64L)
+ // Decimals expand into digit strings when pickled; the estimate tracks
precision.
+ val dec = sized(Decimal(BigDecimal("1234567890.123456789")),
DecimalType(38, 18))
+ assert(dec >= 19L)
+ // Nulls and unknown leaf types yield small positive estimates -- never
zero, so the cap
+ // can never go blind on an unrecognized row shape.
+ assert(sized(null, StringType) > 0L)
+ assert(sized(new Object, CalendarIntervalType) > 0L)
+ // Whole-row conversion (needConversion path): nested values are sized in
the same pass,
+ // and getAndReset() closes the per-row cycle.
+ val schema = StructType(Seq(StructField("s", StringType), StructField("n",
StringType)))
+ val rowSize = sized(
+ new GenericInternalRow(Array[Any](UTF8String.fromString("y" * 2000),
null)), schema)
+ assert(rowSize >= 2000L)
+ assert(acc.getAndReset() === 0L) // reset really resets
+ }
+
+ test("SPARK-57593: toJava size estimate is positive for every data type
(drift guard)") {
+ // Enforces the accounting invariant on EvaluatePython.toJava: a (new)
conversion case that
+ // forgets to account to sizeAcc yields a zero estimate at value level and
makes the byte cap
+ // blind to that type. The type list comes from DataTypeTestUtils +
composites, so the sweep
+ // extends as Spark grows types, and values come from the seeded
RandomDataGenerator.
+ import org.apache.spark.sql.RandomDataGenerator
+ import org.apache.spark.sql.catalyst.CatalystTypeConverters
+ val acc = new PickledSizeAccumulator
+ val sweep = (DataTypeTestUtils.atomicTypes.toSeq ++ Seq(
+ ArrayType(StringType),
+ MapType(StringType, DoubleType),
+ StructType(Seq(StructField("s", StringType), StructField("d",
DoubleType)))))
+ var swept = 0
+ sweep.foreach { dt =>
+ RandomDataGenerator.forType(dt, nullable = false, new
scala.util.Random(42)).foreach {
+ gen =>
+ val catalystValue =
CatalystTypeConverters.createToCatalystConverter(dt)(gen())
+ EvaluatePython.toJava(catalystValue, dt, binaryAsBytes = true,
Some(acc))
+ val estimate = acc.getAndReset()
+ assert(estimate > 0L, s"toJava accounted nothing for $dt -- the byte
cap is blind " +
+ "to this type; add accounting to the corresponding case")
+ swept += 1
+ }
+ }
+ assert(swept > 10) // the sweep must be exercising a meaningful slice of
the type space
+ }
+
test("Python UDF: push down deterministic FilterExec predicates") {
val df = Seq(("Hello", 4)).toDF("a", "b")
.where("dummyPythonUDF(b) and dummyPythonUDF(a) and a in (3, 4)")
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonUDFSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonUDFSuite.scala
index 22584231b328..dfff8ed9e975 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonUDFSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonUDFSuite.scala
@@ -265,4 +265,36 @@ class PythonUDFSuite extends SharedSparkSession {
}
}
}
+
+ test("SPARK-57593: pythonPeakPickledBatchBytes metric for
BatchEvalPythonUDTFExec") {
+ assume(shouldTestPythonUDFs)
+ val udtf = TestPythonUDTF(name = "test_udtf")
+
+ spark.udtf.registerPython(udtf.name, udtf.udtf)
+ withTempView("t") {
+ try {
+ spark.range(1000).selectExpr("id % 100 as a", "id % 50 as b")
+ .createOrReplaceTempView("t")
+ val result = sql(s"SELECT f.* FROM t, LATERAL ${udtf.name}(a, b) f")
+ result.collect()
+
+ val udtfExec = result.queryExecution.executedPlan.collectFirst {
+ case p: BatchEvalPythonUDTFExec => p
+ }.getOrElse {
+ fail("Expected BatchEvalPythonUDTFExec in executed plan")
+ }
+
+ // The UDTF pickle path pickles its input through the same
contiguous-allocation code as
+ // BatchEvalPythonExec, so the peak pickled-batch size is recorded
here too (passed via the
+ // shared pythonMetrics map) even though the UDTF path sets no byte
cap.
+ val peakPickledBatchBytes =
+
udtfExec.metrics.get("pythonPeakPickledBatchBytes").map(_.value).getOrElse(0L)
+ assert(peakPickledBatchBytes > 0,
+ "pythonPeakPickledBatchBytes should be > 0 for
BatchEvalPythonUDTFExec, " +
+ s"but was $peakPickledBatchBytes")
+ } finally {
+ spark.sessionState.catalog.dropTempFunction(udtf.name,
ignoreIfNotExists = true)
+ }
+ }
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
