You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "ivoson (via GitHub)" <gi...@apache.org> on 2026/06/22 05:37:18 UTC
[PR] [SPARK-57593][SQL] Byte-bound and instrument the pickle Python UDF input batch [spark]
ivoson opened a new pull request, #56649:
URL: https://github.com/apache/spark/pull/56649
### 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 precision, binary/geo/variant by their byte lengths, and unknown leaf types get a small positive constant so the cap never goes blind. A single row larger than the cap still forms a one-row batch. The new `ByteBoundedAsArrayIterator` performs the byte-bounded grouping.
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.
### 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)
--
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-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57593][SQL][PYTHON] Byte-bound and instrument the pickle Python UDF input batch [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan closed pull request #56649: [SPARK-57593][SQL][PYTHON] Byte-bound and instrument the pickle Python UDF input batch
URL: https://github.com/apache/spark/pull/56649
--
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-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57593][SQL] Byte-bound and instrument the pickle Python UDF input batch [spark]
Posted by "ivoson (via GitHub)" <gi...@apache.org>.
ivoson closed pull request #56649: [SPARK-57593][SQL] Byte-bound and instrument the pickle Python UDF input batch
URL: https://github.com/apache/spark/pull/56649
--
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-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57593][SQL][PYTHON] Byte-bound and instrument the pickle Python UDF input batch [spark]
Posted by "Yicong-Huang (via GitHub)" <gi...@apache.org>.
Yicong-Huang commented on PR #56649:
URL: https://github.com/apache/spark/pull/56649#issuecomment-4780860168
LGTM!
--
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-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57593][SQL][PYTHON] Byte-bound and instrument the pickle Python UDF input batch [spark]
Posted by "ivoson (via GitHub)" <gi...@apache.org>.
ivoson commented on code in PR #56649:
URL: https://github.com/apache/spark/pull/56649#discussion_r3456829696
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExecSuite.scala:
##########
@@ -52,6 +54,183 @@ 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: 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
+ }
+
Review Comment:
You are right, not tests cover that. Added a new one.
--
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-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57593][SQL][PYTHON] Byte-bound and instrument the pickle Python UDF input batch [spark]
Posted by "ivoson (via GitHub)" <gi...@apache.org>.
ivoson commented on code in PR #56649:
URL: https://github.com/apache/spark/pull/56649#discussion_r3456849268
##########
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("5.0.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .bytesConf(ByteUnit.BYTE)
+ .checkValue(x => x == -1 || (x > 0 && x <= Int.MaxValue),
Review Comment:
Good question. I'd keep the upper bound at `Int.MaxValue` rather than shrink it.
2Gb is the physical ceiling. Each batch is pickled into a single contiguous `byte[]` (`ByteArrayOutputStream.toByteArray`), and a JVM array length is an `int`, so a batch can't exceed `Int.MaxValue` bytes no matter what the config is set to. That's also why the Arrow path uses the same `<= Int.MaxValue` bound, as you noted.
And the ceiling here is just input validation, not a recommended value: the default is `-1` (off), and anyone who turns this on picks their own (typically much smaller) cap. Any smaller ceiling (such as 512MB or 1GB) would be arbitrary and could block a legitimate large-heap-executor use case, whereas `Int.MaxValue` is the one bound that's actually meaningful (the array limit).
Does this make sense? pls lmk if there are more suggestions on this! Thanks.
--
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-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57593][SQL][PYTHON] Byte-bound and instrument the pickle Python UDF input batch [spark]
Posted by "ivoson (via GitHub)" <gi...@apache.org>.
ivoson commented on PR #56649:
URL: https://github.com/apache/spark/pull/56649#issuecomment-4775464378
cc @Yicong-Huang comments resolved, pls take another look. Thx!
--
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-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57593][SQL][PYTHON] Byte-bound and instrument the pickle Python UDF input batch [spark]
Posted by "Yicong-Huang (via GitHub)" <gi...@apache.org>.
Yicong-Huang commented on code in PR #56649:
URL: https://github.com/apache/spark/pull/56649#discussion_r3460992857
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExecSuite.scala:
##########
@@ -52,6 +54,183 @@ 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: 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
+ }
+
Review Comment:
thanks
--
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-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57593][SQL][PYTHON] Byte-bound and instrument the pickle Python UDF input batch [spark]
Posted by "ivoson (via GitHub)" <gi...@apache.org>.
ivoson commented on code in PR #56649:
URL: https://github.com/apache/spark/pull/56649#discussion_r3456821754
##########
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("5.0.0")
Review Comment:
thx, updated.
--
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-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57593][SQL][PYTHON] Byte-bound and instrument the pickle Python UDF input batch [spark]
Posted by "Yicong-Huang (via GitHub)" <gi...@apache.org>.
Yicong-Huang commented on code in PR #56649:
URL: https://github.com/apache/spark/pull/56649#discussion_r3460990908
##########
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("5.0.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .bytesConf(ByteUnit.BYTE)
+ .checkValue(x => x == -1 || (x > 0 && x <= Int.MaxValue),
Review Comment:
that makes sense, we can keep Int.MaxValue in this case. Thanks!
--
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-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57593][SQL][PYTHON] Byte-bound and instrument the pickle Python UDF input batch [spark]
Posted by "ivoson (via GitHub)" <gi...@apache.org>.
ivoson commented on code in PR #56649:
URL: https://github.com/apache/spark/pull/56649#discussion_r3457410287
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonSQLMetrics.scala:
##########
@@ -28,6 +28,8 @@ trait PythonSQLMetrics { self: SparkPlan =>
k -> SQLMetrics.createTimingMetric(sparkContext, v)
} ++ PythonSQLMetrics.pythonOtherMetricsDesc.map { case (k, v) =>
k -> SQLMetrics.createMetric(sparkContext, v)
+ } ++ PythonSQLMetrics.pythonBatchSizeMetricsDesc.map { case (k, v) =>
Review Comment:
Nice catch. Updated. Thx!
--
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-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57593][SQL][PYTHON] Byte-bound and instrument the pickle Python UDF input batch [spark]
Posted by "ivoson (via GitHub)" <gi...@apache.org>.
ivoson commented on code in PR #56649:
URL: https://github.com/apache/spark/pull/56649#discussion_r3456829696
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExecSuite.scala:
##########
@@ -52,6 +54,183 @@ 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: 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
+ }
+
Review Comment:
You are right, no tests cover that. Added a new one.
--
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-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57593][SQL][PYTHON] Byte-bound and instrument the pickle Python UDF input batch [spark]
Posted by "ivoson (via GitHub)" <gi...@apache.org>.
ivoson commented on PR #56649:
URL: https://github.com/apache/spark/pull/56649#issuecomment-4775314911
> **0 blocking, 1 non-blocking, 0 nits.** Clean, well-tested PR.
>
> ### Suggestions (1)
> * `BatchEvalPythonUDTFExec` calls `getInputIterator` without passing its `pythonMetrics`, so `pythonPeakPickledBatchBytes` is never populated on the UDTF pickle path — yet it pickles through the same contiguous-allocation code with the same OOM risk and already mixes in `PythonSQLMetrics`. The byte _cap_ is reasonably scoped out of UDTF, but the peak metric is pure observability; wiring the existing `pythonMetrics` map would extend it for free. (Non-blocking, and the file isn't in this PR's diff — fine to defer.)
>
> ## Verification
> Confirmed the default (cap-off, `-1`) path is byte-for-byte unchanged: `convertRow(_, None)` makes every `sizeAcc.foreach(...)` a no-op so the converted objects match the old inline convert, and `pickle.dump(o, baos); baos.toByteArray` produces the same bytes as the old `pickle.dumps(o)` (the `ByteArrayOutputStream(1024)` only changes initial capacity). The byte-cap path's size estimate cannot affect results — a batch always holds >=1 row, so estimate error only changes batch granularity, never output; every leaf accounts a positive size (the `toJava` catch-all routes unknown types to `addLeaf`), enforced by the per-type drift-guard test.
Thanks, also passed `pythonMetrics` for `BatchEvalPythonUDTFExec`.
--
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-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57593][SQL][PYTHON] Byte-bound and instrument the pickle Python UDF input batch [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on code in PR #56649:
URL: https://github.com/apache/spark/pull/56649#discussion_r3456930272
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonSQLMetrics.scala:
##########
@@ -28,6 +28,8 @@ trait PythonSQLMetrics { self: SparkPlan =>
k -> SQLMetrics.createTimingMetric(sparkContext, v)
} ++ PythonSQLMetrics.pythonOtherMetricsDesc.map { case (k, v) =>
k -> SQLMetrics.createMetric(sparkContext, v)
+ } ++ PythonSQLMetrics.pythonBatchSizeMetricsDesc.map { case (k, v) =>
Review Comment:
These three batch metrics live in the shared `PythonSQLMetrics` trait, so they surface (as `0` / `0.0 B`) on every operator that mixes it in — the Arrow eval/agg/window/cogroup execs, `MapInBatchExec`, and the streaming Python stateful operators — none of which pickle, so the values stay zero. The other shared metrics are populated on all paths, so this is the first pickle-specific one to leak onto those nodes. Non-blocking, and a fair tradeoff against keeping a single metrics map — but if the always-zero rows bother you, a small `PythonPickleBatchMetrics` trait on just `BatchEvalPythonExec`/`BatchEvalPythonUDTFExec` would scope them, and `getInputIterator` (which looks them up by key) wouldn't change.
--
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-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57593][SQL] Byte-bound and instrument the pickle Python UDF input batch [spark]
Posted by "ivoson (via GitHub)" <gi...@apache.org>.
ivoson commented on PR #56649:
URL: https://github.com/apache/spark/pull/56649#issuecomment-4765650360
cc @cloud-fan @Yicong-Huang can you pls help review this PR? Thanks
--
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-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57593][SQL][PYTHON] Byte-bound and instrument the pickle Python UDF input batch [spark]
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on PR #56649:
URL: https://github.com/apache/spark/pull/56649#issuecomment-4780902016
thanks, merging to master/4.x
--
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-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-57593][SQL][PYTHON] Byte-bound and instrument the pickle Python UDF input batch [spark]
Posted by "Yicong-Huang (via GitHub)" <gi...@apache.org>.
Yicong-Huang commented on code in PR #56649:
URL: https://github.com/apache/spark/pull/56649#discussion_r3454230793
##########
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("5.0.0")
Review Comment:
should this be 4.3.0?
##########
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("5.0.0")
+ .withBindingPolicy(ConfigBindingPolicy.NOT_APPLICABLE)
+ .bytesConf(ByteUnit.BYTE)
+ .checkValue(x => x == -1 || (x > 0 && x <= Int.MaxValue),
Review Comment:
Int max is about 2GB, just want to make sure this is the idea range. as a reference, arrow (I know we are talking about non-arrow path) also has limit around 2GB, so maybe the lower limit is enough. but do we want to shrink the batch size cap for our batches?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExecSuite.scala:
##########
@@ -52,6 +54,183 @@ 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: 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
+ }
+
Review Comment:
I might have missed it, but did we have a test for -1 value (no limit)?
--
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-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org