You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by db...@apache.org on 2026/10/14 17:34:33 UTC
(spark) branch master updated: [SPARK-60226][CONNECT] Preserve context for typed Scala UDF unpacking failures
This is an automated email from the ASF dual-hosted git repository.
dbtsai pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new a5cf4c054b8a [SPARK-60226][CONNECT] Preserve context for typed Scala UDF unpacking failures
a5cf4c054b8a is described below
commit a5cf4c054b8a9301eb2b0ebf90f79cb74b8a237b
Author: DB Tsai <db...@apache.org>
AuthorDate: Wed Oct 14 10:34:15 2026 -0700
[SPARK-60226][CONNECT] Preserve context for typed Scala UDF unpacking failures
### What changes were proposed in this pull request?
Preserve the expression's existing origin while decoding typed Scala UDFs in Spark Connect, and identify the failing operation in the generic unpacking failure message.
Typed UDFs are decoded directly rather than through `transformExpression`, so they bypass the origin handling used by ordinary expressions. This PR routes every *expression-backed* typed UDF decode path through a shared `withExpressionOrigin` helper that installs the expression's origin as the current origin while the UDF is decoded, and attaches the current query context (and its summary) to the generic unpacking exception while retaining the original cause:
- `Dataset.groupByKey` (and the dummy grouping function used by `RelationalGroupedDataset#as`).
- `Dataset.filter` — previously the caller extracted only the `CommonInlineUserDefinedFunction` from the condition expression, dropping its origin; the full expression is now passed through.
- Typed `reduce` (the deprecated `Dataset.reduce` encoding) — previously decoded via a direct `unpackUdf` call that bypassed the origin handling entirely.
Relation-backed typed UDFs (`map`/`mapPartitions`, `mapGroups`/`flatMapGroups`, `cogroup`, and stateful typed group operations) store the function directly in relation protobuf fields, which carry no origin metadata. They are intentionally left out: attaching source locations for them would require propagating origin in the protocol/client, not merely changing the planner.
This does not add Serverless-specific behavior or change which APIs OSS Spark supports. Existing class-loading error branches are unchanged. Client source locations are included only when supplied in the expression's origin; this patch does not add client-side stack capture.
### Why are the changes needed?
The generic unpacking exception omitted query context and did not identify the operation. Users consequently saw internal planner frames without the available application location or an indication of which typed operation's function failed to unpack.
### Does this PR introduce _any_ user-facing change?
Yes. Generic typed-UDF unpacking failures change from:
Failed to unpack scala udf.
to an operation-qualified message, e.g.:
Failed to unpack scala udf for Dataset.groupByKey.
Failed to unpack scala udf for Dataset.filter.
Failed to unpack scala udf for Dataset.reduce.
When JVM origin metadata is present, the exception also carries DataFrame query context and includes a summary such as:
== DataFrame ==
"groupByKey" was called from
example.AggregationJob.run(AggregationJob.scala:42)
The error condition remains `INTERNAL_ERROR`; this change does not infer deployment-specific unsupported capabilities from arbitrary deserialization failures.
### How was this patch tested?
Parameterized regression cases in `SparkConnectPlannerSuite` covering `groupByKey`, `filter`, and `reduce`, each with and without JVM origin, plus the existing case for the generic message when no operation context is supplied. Assertions cover operation identification, cause preservation, query context, rendered source location, and origin restoration.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Isaac
This pull request and its description were written by Isaac.
Closes #59446 from dbtsai/connect-typed-udf-unpack-context.
Authored-by: DB Tsai <db...@apache.org>
Signed-off-by: DB Tsai <db...@dbtsai.com>
---
.../sql/connect/planner/SparkConnectPlanner.scala | 62 ++++++--
.../connect/planner/SparkConnectPlannerSuite.scala | 167 ++++++++++++++++++++-
2 files changed, 215 insertions(+), 14 deletions(-)
diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
index 3412dca56817..e821ff228e6c 100644
--- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
+++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
@@ -1051,7 +1051,8 @@ class SparkConnectPlanner(
val analyzed = session.sessionState.executePlan(logicalPlan).analyzed
assertPlan(groupingExprs.size() == 1)
- val groupFunc = TypedScalaUdf(groupingExprs.get(0), Some(analyzed.output))
+ val groupFunc =
+ TypedScalaUdf(groupingExprs.get(0), Some(analyzed.output), Some("Dataset.groupByKey"))
val vEnc = groupFunc.inEnc
val kEnc = groupFunc.outEnc
@@ -1090,9 +1091,18 @@ class SparkConnectPlanner(
private object TypedScalaUdf {
def apply(expr: proto.Expression, inputAttrs: Option[Seq[Attribute]]): TypedScalaUdf = {
+ apply(expr, inputAttrs, None)
+ }
+
+ def apply(
+ expr: proto.Expression,
+ inputAttrs: Option[Seq[Attribute]],
+ operation: Option[String]): TypedScalaUdf = {
if (expr.hasCommonInlineUserDefinedFunction
&& expr.getCommonInlineUserDefinedFunction.hasScalarScalaUdf) {
- apply(expr.getCommonInlineUserDefinedFunction, inputAttrs)
+ withExpressionOrigin(expr) {
+ apply(unpackUdf(expr.getCommonInlineUserDefinedFunction, operation), inputAttrs)
+ }
} else {
throw InvalidInputErrors.expectingScalaUdfButGot(expr.getExprTypeCase)
}
@@ -1836,7 +1846,7 @@ class SparkConnectPlanner(
val baseRel = transformRelation(rel.getInput)
val cond = rel.getCondition
if (isTypedScalaUdfExpr(cond)) {
- transformTypedFilter(cond.getCommonInlineUserDefinedFunction, baseRel)
+ transformTypedFilter(cond, baseRel)
} else {
logical.Filter(condition = transformExpression(cond), child = baseRel)
}
@@ -1856,11 +1866,9 @@ class SparkConnectPlanner(
}
}
- private def transformTypedFilter(
- fun: proto.CommonInlineUserDefinedFunction,
- child: LogicalPlan): TypedFilter = {
+ private def transformTypedFilter(expr: proto.Expression, child: LogicalPlan): TypedFilter = {
val analyzed = session.sessionState.executePlan(child).analyzed
- val udf = TypedScalaUdf(fun, Some(analyzed.output))
+ val udf = TypedScalaUdf(expr, Some(analyzed.output), Some("Dataset.filter"))
TypedFilter(udf.function, analyzed)(udf.inEnc)
}
@@ -1925,6 +1933,21 @@ class SparkConnectPlanner(
}
}
+ /**
+ * Installs the origin carried by `expr` (if any) as the current origin while `f` runs. Typed
+ * UDFs are decoded directly instead of going through `transformExpression`, so without this the
+ * expression's query context would be lost and unpacking failures would surface only internal
+ * planner frames. Falls back to the current origin when `expr` carries none.
+ */
+ private def withExpressionOrigin[T](expr: proto.Expression)(f: => T): T = {
+ val origin = if (expr.hasCommon && expr.getCommon.hasOrigin) {
+ transformOrigin(expr.getCommon.getOrigin)
+ } else {
+ CurrentOrigin.get
+ }
+ CurrentOrigin.withOrigin(origin)(f)
+ }
+
private def transformOrigin(origin: proto.Origin): Origin = {
if (origin.hasPythonOrigin) {
CurrentOrigin.withOrigin {
@@ -2144,15 +2167,19 @@ class SparkConnectPlanner(
NamedArgumentExpression(namedArg.getKey, transformExpression(namedArg.getValue))
}
- private def unpackUdf(fun: proto.CommonInlineUserDefinedFunction): UdfPacket = {
- unpackScalaUDF[UdfPacket](fun.getScalarScalaUdf)
+ private def unpackUdf(
+ fun: proto.CommonInlineUserDefinedFunction,
+ operation: Option[String] = None): UdfPacket = {
+ unpackScalaUDF[UdfPacket](fun.getScalarScalaUdf, operation)
}
private def unpackForeachWriter(fun: proto.ScalarScalaUDF): ForeachWriterPacket = {
unpackScalaUDF[ForeachWriterPacket](fun)
}
- private[connect] def unpackScalaUDF[T](fun: proto.ScalarScalaUDF): T = {
+ private[connect] def unpackScalaUDF[T](
+ fun: proto.ScalarScalaUDF,
+ operation: Option[String] = None): T = {
try {
logDebug(s"Unpack using class loader: ${Utils.getContextOrSparkClassLoader}")
UdfSerialization
@@ -2174,7 +2201,16 @@ class SparkConnectPlanner(
"message" -> (s"Failed to load class: ${cnf.getMessage}. " +
"Make sure the artifact where the class is defined is installed by calling" +
" session.addArtifact.")))
- case _ => throw SparkException.internalError("Failed to unpack scala udf.", t)
+ case _ =>
+ val context = CurrentOrigin.get.getQueryContext
+ val operationContext = operation.map(op => s" for $op").getOrElse("")
+ throw new SparkException(
+ errorClass = "INTERNAL_ERROR",
+ messageParameters = Map(
+ "message" -> s"Failed to unpack scala udf$operationContext."),
+ cause = t,
+ context = context,
+ summary = context.map(_.summary()).mkString("\n"))
}
}
}
@@ -2841,7 +2877,9 @@ class SparkConnectPlanner(
case collection.Seq(e)
if e.hasCommonInlineUserDefinedFunction &&
e.getCommonInlineUserDefinedFunction.hasScalarScalaUdf =>
- unpackUdf(e.getCommonInlineUserDefinedFunction)
+ withExpressionOrigin(e) {
+ unpackUdf(e.getCommonInlineUserDefinedFunction, Some("Dataset.reduce"))
+ }
case other =>
throw InvalidInputErrors.reduceShouldCarryScalarScalaUdf(other)
}
diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectPlannerSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectPlannerSuite.scala
index 4d26937f9fbe..879f7365cbfd 100644
--- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectPlannerSuite.scala
+++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectPlannerSuite.scala
@@ -21,7 +21,7 @@ import scala.jdk.CollectionConverters._
import com.google.protobuf.{Any, ByteString, StringValue}
-import org.apache.spark.SparkFunSuite
+import org.apache.spark.{SparkException, SparkFunSuite}
import org.apache.spark.connect.proto
import org.apache.spark.connect.proto.Expression.{Alias, ExpressionString, UnresolvedStar}
import org.apache.spark.sql.{AnalysisException, Row}
@@ -30,7 +30,7 @@ import org.apache.spark.sql.catalyst.analysis.{RelationTimeTravel, UnresolvedAli
import org.apache.spark.sql.catalyst.expressions.{AttributeReference, UnsafeProjection}
import org.apache.spark.sql.catalyst.plans.logical
import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, LogicalPlan}
-import org.apache.spark.sql.catalyst.trees.TreePattern
+import org.apache.spark.sql.catalyst.trees.{CurrentOrigin, TreePattern}
import org.apache.spark.sql.catalyst.types.DataTypeUtils
import org.apache.spark.sql.classic.Dataset
import org.apache.spark.sql.connect.SparkConnectTestUtils
@@ -442,6 +442,169 @@ class SparkConnectPlannerSuite extends SparkFunSuite with SparkConnectPlanTest {
assert(res.nodeName == "Aggregate")
}
+ // Builds a typed Scala UDF expression whose payload cannot be deserialized, optionally carrying
+ // a JVM origin whose head frame reports `fragment` as the API method name and whose next frame
+ // becomes the call site. `numArgs` controls how many `UnresolvedStar` arguments the UDF carries,
+ // matching the client encoding of each operation (1 for groupByKey/filter, 2 for Dataset.reduce's
+ // `udf(col("*"), col("*"))`).
+ private def malformedTypedScalaUdf(
+ includeOrigin: Boolean,
+ fragment: String,
+ numArgs: Int = 1): proto.Expression = {
+ val udf = proto.CommonInlineUserDefinedFunction
+ .newBuilder()
+ .setScalarScalaUdf(
+ proto.ScalarScalaUDF
+ .newBuilder()
+ .setPayload(ByteString.copyFromUtf8("invalid serialized function")))
+ (1 to numArgs).foreach { _ =>
+ udf.addArguments(
+ proto.Expression
+ .newBuilder()
+ .setUnresolvedStar(proto.Expression.UnresolvedStar.newBuilder()))
+ }
+ val expr = proto.Expression.newBuilder().setCommonInlineUserDefinedFunction(udf)
+ if (includeOrigin) {
+ expr.setCommon(
+ proto.ExpressionCommon
+ .newBuilder()
+ .setOrigin(
+ proto.Origin
+ .newBuilder()
+ .setJvmOrigin(
+ proto.JvmOrigin
+ .newBuilder()
+ .addStackTrace(
+ proto.StackTraceElement
+ .newBuilder()
+ .setDeclaringClass("org.apache.spark.sql.Dataset")
+ .setMethodName(fragment)
+ .setFileName("Dataset.scala")
+ .setLineNumber(1))
+ .addStackTrace(proto.StackTraceElement
+ .newBuilder()
+ .setDeclaringClass("example.AggregationJob")
+ .setMethodName("run")
+ .setFileName("AggregationJob.scala")
+ .setLineNumber(42)))))
+ }
+ expr.build()
+ }
+
+ private def rangeRelation: proto.Relation.Builder =
+ proto.Relation
+ .newBuilder()
+ .setRange(proto.Range.newBuilder().setStart(0).setEnd(1).setStep(1))
+
+ // name: test name; operation: the label expected in the failure message; fragment: the API
+ // method name reported by the origin's head frame; udfArgs: number of UnresolvedStar arguments
+ // in the client encoding; relation: wraps the typed UDF for the op.
+ private case class TypedUdfUnpackScenario(
+ name: String,
+ operation: String,
+ fragment: String,
+ udfArgs: Int,
+ relation: proto.Expression => proto.Relation)
+
+ private val typedUdfUnpackScenarios: Seq[TypedUdfUnpackScenario] = Seq(
+ TypedUdfUnpackScenario(
+ "groupByKey",
+ "Dataset.groupByKey",
+ "groupByKey",
+ 1,
+ udf =>
+ proto.Relation
+ .newBuilder()
+ .setAggregate(
+ proto.Aggregate
+ .newBuilder()
+ .setInput(rangeRelation)
+ .setGroupType(proto.Aggregate.GroupType.GROUP_TYPE_GROUPBY)
+ .addGroupingExpressions(udf))
+ .build()),
+ TypedUdfUnpackScenario(
+ "filter",
+ "Dataset.filter",
+ "filter",
+ 1,
+ udf =>
+ proto.Relation
+ .newBuilder()
+ .setFilter(
+ proto.Filter
+ .newBuilder()
+ .setInput(rangeRelation)
+ .setCondition(udf))
+ .build()),
+ // Mirrors Dataset.reduce: a GROUP_TYPE_GROUPBY aggregate whose only aggregate expression is the
+ // deprecated unresolved `reduce` function wrapping a two-argument scalar Scala UDF.
+ TypedUdfUnpackScenario(
+ "reduce",
+ "Dataset.reduce",
+ "reduce",
+ 2,
+ udf =>
+ proto.Relation
+ .newBuilder()
+ .setAggregate(
+ proto.Aggregate
+ .newBuilder()
+ .setInput(rangeRelation)
+ .setGroupType(proto.Aggregate.GroupType.GROUP_TYPE_GROUPBY)
+ .addAggregateExpressions(
+ proto.Expression
+ .newBuilder()
+ .setUnresolvedFunction(proto.Expression.UnresolvedFunction
+ .newBuilder()
+ .setFunctionName("reduce")
+ .addArguments(udf))))
+ .build()))
+
+ for {
+ scenario <- typedUdfUnpackScenarios
+ includeOrigin <- Seq(false, true)
+ } {
+ test(
+ s"${scenario.name} unpacking failure includes operation context " +
+ s"(origin=$includeOrigin)") {
+ val udf = malformedTypedScalaUdf(includeOrigin, scenario.fragment, scenario.udfArgs)
+ val originalOrigin = CurrentOrigin.get
+ val error = intercept[SparkException] {
+ transform(scenario.relation(udf))
+ }
+ assert(error.getCondition == "INTERNAL_ERROR")
+ assert(
+ error.getMessageParameters.get("message") ==
+ s"Failed to unpack scala udf for ${scenario.operation}.")
+ assert(error.getCause.isInstanceOf[java.io.StreamCorruptedException])
+ if (includeOrigin) {
+ assert(error.getQueryContext.length == 1)
+ assert(error.getQueryContext.head.fragment() == scenario.fragment)
+ assert(
+ error.getQueryContext.head.callSite() ==
+ "example.AggregationJob.run(AggregationJob.scala:42)")
+ assert(error.getMessage.contains("AggregationJob.scala:42"))
+ } else {
+ assert(error.getQueryContext.isEmpty)
+ }
+ assert(CurrentOrigin.get == originalOrigin)
+ }
+ }
+
+ test("Scala UDF unpacking without operation context preserves the existing message") {
+ val planner = new SparkConnectPlanner(SparkConnectTestUtils.createDummySessionHolder(spark))
+ val error = intercept[SparkException] {
+ planner.unpackScalaUDF[AnyRef](
+ proto.ScalarScalaUDF
+ .newBuilder()
+ .setPayload(ByteString.copyFromUtf8("invalid serialized function"))
+ .build())
+ }
+ assert(error.getMessageParameters.get("message") == "Failed to unpack scala udf.")
+ assert(error.getCause.isInstanceOf[java.io.StreamCorruptedException])
+ assert(error.getQueryContext.isEmpty)
+ }
+
test("Test invalid deduplicate") {
val deduplicate = proto.Deduplicate
.newBuilder()
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org