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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
