[SPARK-59744][UDF] Execute scalar external UDFs through worker sessions - #58988
haiyangsun-db wants to merge 5 commits into
Conversation
haiyangsun-db
changed the title
4e713f7 to
7c57b5f
Compare
haiyangsun-db
changed the title
cloud-fan
left a comment
cloud-fan
left a comment
There was a problem hiding this comment.
Review summary
PR tags: new-feature
I found three non-blocking regression-coverage gaps in the new scalar external-UDF bridge and one documentation typo. The production design itself is coherent: nullable metadata is carried from the expression through planning and execution, cardinality is checked, and Arrow/session resources have task-scoped finalizers. The remaining review work is to make the owning suite sensitive to large-variable Arrow physical types, exact task-context wire values, and cleanup on malformed/cardinality failures, plus correct the incomplete byte-limit sentence.
Findings
4 total: 0 P0, 0 P1, 3 P2, 1 P3.
Non-blocking (P2)
- Exercise large-variable Arrow types end to end —
sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExecuteExternalUDFExec.scala:93— see inline. - Assert cleanup after failed worker responses —
sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExecuteExternalUDFExec.scala:264— see inline. - Validate the complete task-context Init payload —
sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/PythonInitAdapter.scala:93— see inline.
Nit (P3)
- Fix the incomplete byte-limit explanation —
sql/core/src/main/scala/org/apache/spark/sql/execution/arrow/ArrowConverters.scala:283— see inline.
| }) | ||
| val outputSchema = StructType(Seq(StructField("_0", udf.dataType, nullable = true))) | ||
| val timeZoneId = conf.sessionLocalTimeZone | ||
| val largeVarTypes = conf.arrowUseLargeVarTypes |
There was a problem hiding this comment.
Non-blocking (P2): Please add a non-default-mode case that keeps the raw Arrow type visible on both request and response paths. The current schema checks round-trip through Spark StructType, so they erase the Utf8/LargeUtf8 distinction and would stay green if this flag were ignored or applied in only one direction.
| /** | ||
| * Maps an iterator of internal rows to serialized Arrow record batches, limiting each batch by | ||
| * record count and a best-effort byte target. A batch can exceed the byte target because its size | ||
| * checked before appending each row, and the first row is always accepted. |
There was a problem hiding this comment.
Nit (P3): This sentence is missing is: it should read “because its size is checked before appending each row.”
| } | ||
| rows.hasNext | ||
| } catch { | ||
| case NonFatal(error) => |
There was a problem hiding this comment.
Non-blocking (P2): The malformed-response and both cardinality-mismatch cases verify only the surfaced exception. Please make failed-task cleanup observable and assert that the worker session and Arrow response resources close in each case; the existing success and early-stop close assertions would not catch a leak that occurs only after an exception.
| name -> Map("name" -> resource.name, "addresses" -> resource.addresses) | ||
| } | ||
| Map( | ||
| "isBarrier" -> context.isInstanceOf[BarrierTaskContext].toString, |
There was a problem hiding this comment.
Non-blocking (P2): Please validate the values and encodings of the complete task-context map, not just the presence of four keys. As written, swapped stage/partition IDs, missing barrier/attempt/CPU fields, or malformed resources/localProperties JSON would all leave the execution suite green.
bea4757 to
eae673e
Compare
cloud-fan
left a comment
cloud-fan
left a comment
There was a problem hiding this comment.
Review summary
PR tags: new-feature
All four comments from the previous review are resolved in the current revision: large-variable Arrow types are exercised end to end, failure paths assert session and allocator cleanup, the complete task-context payload is validated, and the byte-limit Scaladoc is corrected.
This re-review found two new non-blocking issues in the scalar response path and its test coverage. The production flow otherwise remains coherent: nullable metadata reaches execution, output cardinality is checked, and task-scoped cleanup covers sessions and Arrow resources.
Findings
2 total: 0 P0, 0 P1, 2 P2, 0 P3.
Non-blocking (P2)
- Cache Arrow column wrappers across response batches —
sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExecuteExternalUDFExec.scala:303— see inline. - Cover zero-argument scalar UDF execution —
sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/ExecuteExternalUDFExec.scala:87— see inline.
Re-review status
Prior AI findings: 4 addressed, 0 still present; additional unresolved findings in this review: 2.
New attribution: 0 newly introduced, 2 late catch, 0 previously raised, 0 unattributed.
Remaining prior AI findings
No prior AI findings remain.
| } | ||
|
|
||
| private def rowsFromRoot(root: VectorSchemaRoot): Iterator[InternalRow] = { | ||
| val columns: Array[ColumnVector] = root.getFieldVectors.asScala.iterator.map { vector => |
There was a problem hiding this comment.
Non-blocking (P2): This rebuilds the ArrowColumnVector accessor tree and ColumnarBatch for every response batch even though VectorLoader reloads the same persistent root in place. For small batches or nested return types, that adds schema-proportional allocation on the hot path. Could we construct these wrappers once when the root is initialized and only update the batch row count after each load?
| udf.name.getOrElse(udf.prettyName)) | ||
| } | ||
| val argumentExpressions: Seq[Expression] = udf.children | ||
| val inputSchema = StructType( |
There was a problem hiding this comment.
Non-blocking (P2): Could we add an execution test for a zero-argument scalar UDF over nonempty input? This path creates zero-column Arrow batches, whose row count must survive without any vectors; the existing planner test proves the shape is supported, but every worker-backed execution test currently projects at least one argument, so a cardinality regression here would go undetected.
f8b6f90 to
45b0b3f
Compare
cloud-fan
left a comment
cloud-fan
left a comment
There was a problem hiding this comment.
Review summary
Both findings from the previous review are resolved: the Arrow column wrappers and ColumnarBatch are now cached across response batches, and the suite covers a zero-argument scalar UDF over nonempty input.
This round finds one additional non-blocking issue in the temporary Python Init contract. It advertises barrier mode without the callback endpoint, secret, or JVM server lifecycle required by Python BarrierTaskContext; rejecting barrier contexts before session creation is the narrow repair until the deferred worker integration supplies that transport. The remaining scalar Arrow exchange, cardinality checks, nullable metadata, cleanup paths, and updated regression coverage are coherent in the pinned revision.
Findings
1 total: 0 P0, 0 P1, 1 P2, 0 P3.
Non-blocking (P2)
- Reject barrier contexts without callback transport —
sql/core/src/main/scala/org/apache/spark/sql/execution/externalUDF/PythonInitAdapter.scala:93— see inline.
Re-review status
Prior AI findings: 2 addressed, 0 still present; additional unresolved findings in this review: 1.
New attribution: 0 newly introduced, 1 late catch, 0 previously raised, 0 unattributed.
Remaining prior AI findings
No prior AI findings remain.
| name -> Map("name" -> resource.name, "addresses" -> resource.addresses) | ||
| } | ||
| Map( | ||
| "isBarrier" -> context.isInstanceOf[BarrierTaskContext].toString, |
There was a problem hiding this comment.
Non-blocking (P2): This can emit isBarrier=true, but the Init map has no connInfo or secret, and this path creates no JVM callback server. Barrier status is preserved through the child RDD into this mapPartitionsInternal path, so a custom or future Python dispatcher would initialize BarrierTaskContext without any transport for barrier() or allGather(). Please either supply the callback lifecycle and fields, or reject BarrierTaskContext before opening the worker session; explicit rejection is the bounded option while the worker integration is deferred.
viirya
left a comment
viirya
left a comment
There was a problem hiding this comment.
A few points beyond the earlier rounds, which all look addressed:
- Response type is not validated.
ArrowEvalPythonExecchecks the returned batch types againstoutputTypes(arrowDataTypeMismatchError), but here responses are schema-less record batches loaded straight into a root built from the expected schema.VectorLoaderonly checks node/buffer counts, so a worker returning a same-width type (e.g. float64 forLongType) is silently reinterpreted, and a narrower type relies on ArrowBuf bounds checking (absent witharrow.enable_unsafe_memory_access). Could we at least validate buffer sizes against the expected schema (e.g.ValueVectorUtility.validate) and document that the worker must produce exactlyoutput_schema? - Finish-phase failures are dropped.
WorkerSession.close()reports failures raised after the data drained only through the returnedTermination, andwithUDFWorkerSessiondiscards it. Now that this path executes, a worker failing in its finish callback yields a successful task. Worth widening the SPARK-57324 TODO or surfacingFailed/TransportFailedhere. - Flow control prerequisite.
GrpcWorkerSessionhas a TODO to add application-level flow control "before wiring this transport into a production path"; sinceadvance()keeps sending input while the output queue is empty, a slow worker can push the whole partition intoHybridRowQueue. Not blocking while the dispatcher is unconfigured, but could the class doc / PR description link that prerequisite? - Compression asymmetry. Request batches inherit
spark.sql.execution.arrow.compression.codecfromArrowBatchIterator, while the responseVectorLoaderhas no codec factory, andInitdoesn't tell the worker either way. Please document the contract (or decompress responses). - Named arguments are rejected in
doExecute, soexplainsucceeds and execution fails with a compilation error. Could this move toPlanExternalUDFs/ the logical node? udfNullableis now ignored, but its Scaladoc still says it controls nullability andPlanExternalUDFsSuiteasserts on it; a note that it is currently ignored would avoid confusion. The new physicalassert(resultAttr.nullable)duplicates the logical node's check.SizeLimitedArrowBatchIteratoris the third copy of the byte-limited loop (PythonArrowInput,ArrowCachedBatchSerializer) and duplicates the wholenext()body. Could the byte limit be an optional parameter onArrowBatchIteratorinstead? The other call sites also takemaxBytesPerBatch: Long.- Minor: the payload format
"experimental"doesn't identify the encoding;_workerSpecis unused;isBarriercan be"true"without theconnInfo/secreta barrier context needs.
This branch has not been deployed
3 participants
What changes were proposed in this pull request?
This PR adds the scalar external UDF execution bridge:
HybridRowQueue.Initmessage through a temporaryPythonInitAdapter, including the Arrow schemas, Python runner configuration, and complete task context.experimentalwhile the protocol is still evolving.This PR does not add a worker implementation. The dispatcher is intentionally unconfigured, so worker creation fails before starting a process. SPARK-59364 will replace the temporary Python-specific initialization dependency separately.
Why are the changes needed?
The external UDF framework can plan scalar UDFs and manage worker sessions, but
ExecuteExternalUDFExecdoes not yet exchange rows with those sessions. This change supplies that execution path while keeping worker-server implementation and language-agnostic initialization out of scope.Does this PR introduce any user-facing change?
No. This extends unreleased external UDF infrastructure. Existing Python UDF execution is unchanged, and no external UDF worker is enabled by this PR.
How was this patch tested?
Added and extended tests covering Arrow row- and byte-based batching, raw large-variable Arrow request and response types, multiple batches and partitions, empty input, ordered multi-argument data, nullable output, named-argument rejection, exact task-context values and JSON encodings, malformed responses, output cardinality mismatches, early iterator termination, failed-task session and Arrow-resource cleanup, and planning behavior.
After rebasing onto the latest
master, ran:All 13 tests passed.
Was this patch authored or co-authored using generative AI tooling?
Yes