dongjoon-hyun commented on code in PR #59637:
URL: https://github.com/apache/spark/pull/59637#discussion_r4464209214
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExec.scala:
##########
@@ -186,6 +190,8 @@ object BatchEvalPythonExec {
estimatedBytes += (dataType match {
case _: StringType => row.getUTF8String(ordinal).numBytes.toLong
case BinaryType => row.getBinary(ordinal).length.toLong
+ case _: GeometryType | _: GeographyType =>
+ row.getBinaryView(ordinal).numBytes().toLong
Review Comment:
The raw `BinaryView` byte count may underestimate the peak memory for
GEOMETRY/GEOGRAPHY much more than for STRING/BINARY.
In `EvaluatePython.toJava`, `STUtils.deserializeGeom` calls `stGeomAsBinary`
-> `Geometry.toWkbInternal`, which copies the bytes (`Arrays.copyOfRange`),
parses the WKB into a `GeometryModel` object graph, and writes it back with
`WkbWriter` before pickling (the Geography path is similar). The object graph
(an object and a coordinate array per point) can be several times larger than
the WKB bytes, so the "approximately 1/12 of the executor heap" assumption of
`maxRowHeapFraction` may not hold for geo values.
Could we apply a type-specific factor for geo types, or at least document
that the estimate is a lower bound for them?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExec.scala:
##########
@@ -162,17 +163,20 @@ class BatchEvalPythonEvaluatorFactory(
}
object BatchEvalPythonExec {
+ private def isSupportedProjectedInputType(dataType: DataType): Boolean =
dataType match {
+ case _: StringType | BinaryType => true
+ case _: GeometryType | _: GeographyType => true
Review Comment:
The rationale in the PR description also applies to top-level `VariantType`:
its payloads can be arbitrarily large, `EvaluatePython` converts them, and the
pickled batch-size accounting already counts their bytes.
`PickledSizeAccumulator.addLeaf` already sizes Variant by its value and
metadata lengths ("Variant payloads can be arbitrarily large"). However, this
estimator returns 0 for a VARIANT argument, e.g.
`my_udf(parse_json(huge_json_col))`.
Shall we include `VariantType` here (e.g. the value + metadata length of
`getVariant(ordinal)`), or mention the exclusion in the conf docs?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExec.scala:
##########
@@ -162,17 +163,20 @@ class BatchEvalPythonEvaluatorFactory(
}
object BatchEvalPythonExec {
+ private def isSupportedProjectedInputType(dataType: DataType): Boolean =
dataType match {
+ case _: StringType | BinaryType => true
+ case _: GeometryType | _: GeographyType => true
+ case _ => false
Review Comment:
`UserDefinedType` arguments are not unwrapped, so a UDT whose `sqlType` is
`StringType`, `BinaryType`, or `GeometryType` falls into `case _ => false` and
is never checked. However, `EvaluatePython.toJava` unwraps it (`case (ud, udt:
UserDefinedType[_]) => toJava(ud, udt.sqlType, ...)`) and still converts and
pickles the full payload.
Could we unwrap `udt.sqlType` here and keep the unwrapped type in
`projectedPayloads`?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala:
##########
@@ -5554,9 +5554,9 @@ object SQLConf {
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.")
+ "top-level string, binary, geometry, and geography argument payloads
whose combined " +
Review Comment:
This doc says that the guard applies to pickle-serialized (non-Arrow) Python
UDF evaluation. However, `BatchEvalPythonUDTFExec` (line 77) uses the same
pickle path via `BatchEvalPythonExec.getInputIterator` without passing
`projectedInputCheck`, so pickle Python UDTFs are not guarded even when this
conf is enabled.
Since this PR updates the docs, could we either state that UDTFs are not
covered, or wire the check into the UDTF path (in a separate PR if you prefer)?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExec.scala:
##########
@@ -186,6 +190,8 @@ object BatchEvalPythonExec {
estimatedBytes += (dataType match {
Review Comment:
This `dataType match` (and the tuple destructuring at line 188) runs for
every argument of every input row when the guard is enabled, and this PR adds
two more type tests to it. If the per-ordinal sizer is precomputed when the
estimator is built (see the comment below), the per-row work becomes just
`isNullAt` and one length read.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExecSuite.scala:
##########
@@ -157,39 +177,43 @@ class BatchEvalPythonExecSuite extends SharedSparkSession
"actualRowSize" -> "1200"))
}
- test("SPARK-59824: pickle UDF row-size guard uses SQLConf settings") {
+ test("SPARK-59824: pickle UDF row-size guard forwards SQLConf settings") {
val configuredMaxRowSize = 1000L
val maxRowHeapFraction = configuredMaxRowSize.toDouble /
Runtime.getRuntime.maxMemory()
val expectedMaxRowSize =
(Runtime.getRuntime.maxMemory() * maxRowHeapFraction).toLong
- withSQLConf(
- SQLConf.PYTHON_UDF_ROW_SIZE_GUARD_ENABLED.key -> "true",
- SQLConf.PYTHON_UDF_ROW_SIZE_GUARD_MAX_ROW_HEAP_FRACTION.key ->
- maxRowHeapFraction.toString) {
- val taskAttr = AttributeReference("s", StringType)()
- val pythonUdf = PythonUDF(
- "dummy", new DummyUDF, StringType, Seq(taskAttr),
PythonEvalType.SQL_BATCHED_UDF,
- udfDeterministic = true)
- class TestBatchEvalPythonExec extends BatchEvalPythonExec(
- Seq(pythonUdf),
- Seq(AttributeReference("result", StringType)()),
- LocalTableScanExec(Seq(taskAttr), Nil, None)) {
- def testEvaluatorFactory: BatchEvalPythonEvaluatorFactory =
- evaluatorFactory.asInstanceOf[BatchEvalPythonEvaluatorFactory]
- }
- val plan = new TestBatchEvalPythonExec
- val oversizedRow = InternalRow(
- UTF8String.fromString("x" * (expectedMaxRowSize.toInt + 1)))
- checkError(
- exception = intercept[SparkRuntimeException] {
- plan.testEvaluatorFactory.getInputIterator(
- Iterator(oversizedRow), StructType.fromDDL("s STRING")).next()
- },
- condition = "UDF_LIMITS.ROW_SIZE",
- parameters = Map(
- "maxRowSize" -> expectedMaxRowSize.toString,
- "actualRowSize" -> (expectedMaxRowSize + 1).toString))
+ val taskAttr = AttributeReference("s", StringType)()
+ val pythonUdf = PythonUDF(
+ "dummy", new DummyUDF, StringType, Seq(taskAttr),
PythonEvalType.SQL_BATCHED_UDF,
+ udfDeterministic = true)
+ class TestBatchEvalPythonExec extends BatchEvalPythonExec(
+ Seq(pythonUdf),
+ Seq(AttributeReference("result", StringType)()),
+ LocalTableScanExec(Seq(taskAttr), Nil, None)) {
+ def testEvaluatorFactory: BatchEvalPythonEvaluatorFactory =
+ evaluatorFactory.asInstanceOf[BatchEvalPythonEvaluatorFactory]
+ }
+ val plan = new TestBatchEvalPythonExec
+ val oversizedRow = InternalRow(
+ UTF8String.fromString("x" * (expectedMaxRowSize.toInt + 1)))
+ val schema = StructType.fromDDL("s STRING")
+ EvaluatePython.registerPicklers()
+ def inputIterator(enabled: Boolean): Iterator[Array[Byte]] = withSQLConf(
+ SQLConf.PYTHON_UDF_ROW_SIZE_GUARD_ENABLED.key -> enabled.toString,
+ SQLConf.PYTHON_UDF_ROW_SIZE_GUARD_MAX_ROW_HEAP_FRACTION.key ->
+ maxRowHeapFraction.toString) {
+ plan.testEvaluatorFactory.getInputIterator(Iterator(oversizedRow),
schema)
}
+
+ assert(inputIterator(enabled = false).next().nonEmpty)
Review Comment:
nit: `inputIterator(...)` builds the iterator inside `withSQLConf`, but
`.next()` runs after the conf scope ends. This works today because
`evaluatorFactory` reads the confs eagerly when the iterator is built. Calling
`.next()` inside `withSQLConf` would make the test independent of that detail.
Also, `next().nonEmpty` is always true for a pickled batch, so the disabled
case effectively asserts only that no exception is thrown.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExecSuite.scala:
##########
@@ -98,12 +98,32 @@ class BatchEvalPythonExecSuite extends SharedSparkSession
assert(estimate(projectedRow) === 6000L)
assert(estimate(InternalRow(null, null, 1)) === 0L)
+ // Geometry and Geography arguments are sized by their BinaryView byte
count.
+ val geoRow = new GenericInternalRow(Array[Any](
+ BinaryView.fromBytes(Array.fill[Byte](1500)(1)),
+ BinaryView.fromBytes(Array.fill[Byte](700)(1))))
+ val geoEstimate = BatchEvalPythonExec
+ .projectedInputPayloadSizeEstimator(Seq(
+ GeometryType(GeometryType.GEOMETRY_DEFAULT_SRID),
+ GeographyType(GeographyType.GEOGRAPHY_DEFAULT_SRID)))
+ .get
+ assert(geoEstimate(geoRow) === 2200L)
Review Comment:
The new geo coverage is only at the estimator level. Could we add a test
that goes through `getInputIterator` with a GEOMETRY (or GEOGRAPHY) argument
and asserts that `UDF_LIMITS.ROW_SIZE` is raised before conversion? It would
directly cover the "oversized geometry or geography argument got past the
guard" scenario in the PR description.
A null geo value (counted as 0) and a mixed STRING + GEOMETRY row (sizes
summed) would also be good to have.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExec.scala:
##########
@@ -186,6 +190,8 @@ object BatchEvalPythonExec {
estimatedBytes += (dataType match {
case _: StringType => row.getUTF8String(ordinal).numBytes.toLong
case BinaryType => row.getBinary(ordinal).length.toLong
+ case _: GeometryType | _: GeographyType =>
+ row.getBinaryView(ordinal).numBytes().toLong
case _ => 0L
Review Comment:
`isSupportedProjectedInputType` and this `match` need to stay in sync, and
`case _ => 0L` is unreachable after the filter above. If a type is added to the
predicate but not here, the guard silently counts 0 bytes for it without any
compiler warning.
How about building a per-ordinal sizer once, e.g. `Array[(Int, (InternalRow,
Int) => Long)]`, so that there is a single dispatch point and no dead branch?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/BatchEvalPythonExec.scala:
##########
@@ -162,17 +163,20 @@ class BatchEvalPythonEvaluatorFactory(
}
object BatchEvalPythonExec {
+ private def isSupportedProjectedInputType(dataType: DataType): Boolean =
dataType match {
Review Comment:
Instead of adding logical types one by one, how about matching on the
physical type? For example, unwrap UDTs to `sqlType`, and then match
`PhysicalDataType(dt)` against `PhysicalStringType`, `PhysicalBinaryType`,
`PhysicalBinaryViewType`, and `PhysicalVariantType`.
That would cover the Variant and UDT cases (see the other comments) and
future `BinaryView`-backed types automatically. We would also no longer need to
keep this list, the per-row `match`, and the two conf docs in sync by hand.
--
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]
