ivoson commented on code in PR #59122:
URL: https://github.com/apache/spark/pull/59122#discussion_r4192340586
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExecSuite.scala:
##########
@@ -81,6 +85,108 @@ class BatchEvalPythonExecSuite extends SharedSparkSession
assert(metric.value === expectedPeak)
}
+ test("SPARK-59824: projected UDF input payload estimator covers evaluated
values") {
+ val inputTypes = Seq(StringType, BinaryType, IntegerType)
+ val projectedRow = InternalRow(
+ UTF8String.fromString("s" * 2000), Array.fill[Byte](4000)(1), 1)
+
+ val estimate = BatchEvalPythonExec
+ .projectedInputPayloadSizeEstimator(inputTypes)
+ .get
+ assert(estimate(projectedRow) === 6000L)
+ assert(estimate(InternalRow(null, null, 1)) === 0L)
+
+ // Fixed-width and nested inputs are intentionally not estimated.
+ assert(BatchEvalPythonExec
+ .projectedInputPayloadSizeEstimator(Seq(IntegerType,
ArrayType(StringType)))
+ .isEmpty)
+ }
+
+ test("SPARK-59824: pickle UDF row-size guard rejects a computed projected
argument") {
+ // Exercise the production evaluator wiring without starting a Python
worker.
+ val conditions = spark.sparkContext.parallelize(Seq(1), 1).mapPartitions {
_ =>
+ val taskAttr = AttributeReference("s", StringType)()
+ val taskRow = UnsafeProjection.create(Array[DataType](StringType))
+ .apply(InternalRow(UTF8String.fromString("x" * 600)))
+ val projectedArg = Concat(Seq(taskAttr, taskAttr))
+ val pythonUdf = PythonUDF(
+ "dummy", new DummyUDF, StringType, Seq(projectedArg),
PythonEvalType.SQL_BATCHED_UDF,
+ udfDeterministic = true)
+ val maxRowBytes = 1000L
+ val maxRowHeapFraction = maxRowBytes.toDouble /
Runtime.getRuntime.maxMemory()
+ val factory = new BatchEvalPythonEvaluatorFactory(
+ childOutput = Seq(taskAttr),
+ udfs = Seq(pythonUdf),
+ output = Seq.empty,
+ batchSize = 1,
+ maxBytesPerBatch = -1L,
+ pythonMetrics = Map.empty,
+ jobArtifactUUID = None,
+ sessionUUID = None,
+ binaryAsBytes = false,
+ rowSizeGuardEnabled = true,
+ rowSizeGuardMaxRowHeapFraction = maxRowHeapFraction) {
+ override def evaluate(
+ funcs: Seq[(ChainedPythonFunctions, Long)],
+ argMetas: Array[Array[ArgumentMetadata]],
+ iter: Iterator[InternalRow],
+ schema: StructType,
+ context: TaskContext): Iterator[InternalRow] = {
+ getInputIterator(iter, schema).next()
+ Iterator.empty
+ }
+ }
+ try {
+ factory.createEvaluator().eval(0, Iterator(taskRow)).hasNext
+ Iterator("NO_ERROR")
+ } catch {
+ case e: SparkThrowable => Iterator(e.getCondition)
+ }
+ }.collect().toSeq
+ assert(conditions === Seq("UDF_LIMITS.ROW_SIZE"))
Review Comment:
Done — added assertions for the maxRowSize and actualRowSize message
parameters in the guard tests.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -5433,6 +5433,33 @@ object SQLConf {
"be -1 (no limit) or greater than zero and less than or equal to
INT_MAX.")
.createWithDefault(-1)
+ val PYTHON_UDF_ROW_SIZE_GUARD_ENABLED =
+ buildConf("spark.sql.execution.python.udf.rowSizeGuard.enabled")
+ .internal()
+ .doc("When true, guard pickle-serialized (non-Arrow) Python UDF
evaluation against " +
+ "top-level string and binary argument payloads whose combined size
exceeds " +
+ "rowSizeGuard.maxRowHeapFraction of executor heap. The guard checks
the projected " +
+ "arguments before conversion and pickling. Nested inputs are not
estimated.")
+ .version("4.3.0")
Review Comment:
Done — updated both row-size guard config versions to 4.4.0.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExec.scala:
##########
@@ -60,7 +61,10 @@ case class BatchEvalPythonExec(udfs: Seq[PythonUDF],
resultAttrs: Seq[Attribute]
pythonMetrics,
jobArtifactUUID,
sessionUUID,
- binaryAsBytes)
+ binaryAsBytes,
+ rowSizeGuardEnabled =
conf.getConf(SQLConf.PYTHON_UDF_ROW_SIZE_GUARD_ENABLED),
Review Comment:
Done — added a SQLConf-backed test that constructs BatchEvalPythonExec and
exercises the evaluatorFactory wiring.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -5433,6 +5433,33 @@ object SQLConf {
"be -1 (no limit) or greater than zero and less than or equal to
INT_MAX.")
.createWithDefault(-1)
+ val PYTHON_UDF_ROW_SIZE_GUARD_ENABLED =
+ buildConf("spark.sql.execution.python.udf.rowSizeGuard.enabled")
+ .internal()
+ .doc("When true, guard pickle-serialized (non-Arrow) Python UDF
evaluation against " +
+ "top-level string and binary argument payloads whose combined size
exceeds " +
+ "rowSizeGuard.maxRowHeapFraction of executor heap. The guard checks
the projected " +
+ "arguments before conversion and pickling. Nested inputs are not
estimated.")
+ .version("4.3.0")
Review Comment:
Done — updated both row-size guard config versions to 4.4.0.
--
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]
