You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "dbtsai (via GitHub)" <gi...@apache.org> on 2026/10/14 17:49:57 UTC
[PR] [SPARK-60247][SQL][CONNECT] Point typed UDF DataFrame query context at the user call site [spark]
dbtsai opened a new pull request, #59474:
URL: https://github.com/apache/spark/pull/59474
### What changes were proposed in this pull request?
Teach the DataFrame query context frame filter (`sparkCodePattern` in `CurrentOrigin`, `sql/api`) to recognize the frames that typed UDF columns are constructed through, so the captured call site skips them and lands on the user's code:
- `org.apache.spark.sql.expressions.UserDefinedFunction` — every UDF column is created via `UserDefinedFunction.apply(Column*)`.
- `org.apache.spark.sql.connect.KeyValueGroupedDatasetImpl` (and the abstract `KeyValueGroupedDataset`) — the Connect client's `groupByKey` builds its grouping UDF inside this factory.
Follow-up to SPARK-60226 (which added server-side operation labels and preserved the client-supplied origin in the Connect planner). That change is correct, but the origin the client *supplies* could already point at Spark internals; this PR fixes the capture.
### Why are the changes needed?
The Connect client encodes typed Dataset APIs (`groupByKey`, `filter`, `reduce`) as a scalar Scala UDF applied to `col("*")`, i.e. `udf.apply(col("*"))` (`sql/connect/common/.../connect/Dataset.scala`, `KeyValueGroupedDataset.scala`). The origin is captured in `Column.apply` via `CurrentOrigin.withOrigin`, whose immediate caller is `UserDefinedFunction.apply` (`UserDefinedFunction.scala:69`).
`sparkCodePattern` matched `Column`/`Dataset`/`functions`/... but not `UserDefinedFunction` (nor the `groupByKey` factory `KeyValueGroupedDatasetImpl`). With the default `spark.sql.stackTracesInDataFrameContext=1`, the one captured "non-Spark" frame was therefore the internal `UserDefinedFunction.apply` (and for `groupByKey`, `KeyValueGroupedDatasetImpl.apply`) rather than the user's class/line.
Trace for `groupByKey`: `Column.apply` (matched) -> `UserDefinedFunction.apply` (unmatched -> walk stops here) -> `KeyValueGroupedDatasetImpl.apply` -> `Dataset.groupByKey` -> user. After matching both intermediates, the walk reaches the user frame.
### Does this PR introduce _any_ user-facing change?
Yes. The DataFrame query context for UDF-backed columns now points at user code. Before:
== DataFrame ==
"apply" was called from
org.apache.spark.sql.expressions.UserDefinedFunction.apply(UserDefinedFunction.scala:69)
After: the call site is the user's application frame. Note this also improves the context for ordinary UDF usage (`df.select(myUdf(col("x")))`), classic and Connect, since all UDF columns go through `UserDefinedFunction.apply`.
### How was this patch tested?
Added a `DatasetSuite` regression that builds a UDF column and asserts its captured query-context call site points to the test (user code) and not to `UserDefinedFunction`.
Draft notes for reviewers:
- I was unable to build/run locally in my environment; CI is the first real build/run. Please treat the test as provisional until CI confirms.
- The `UserDefinedFunction` part is exercised by the added classic test. The `groupByKey`-specific `KeyValueGroupedDatasetImpl` frame is fixed by the same mechanism but is only reachable via the Connect client; a dedicated Connect end-to-end regression (e.g. a real client-generated grouping expression whose serialized payload is then corrupted, asserting the surfaced context points to user code) would be a good addition — happy to add it, or guidance on the preferred location is welcome.
- Design question for reviewers: enumerating intermediate classes in `sparkCodePattern` follows the existing approach but is somewhat fragile to future typed-API refactors. An alternative is to capture the origin at the public typed-API entry points (wrap them in `withOrigin`). Feedback on the preferred direction appreciated.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Isaac
This pull request and its description were written by Isaac.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60247][SQL][CONNECT] Point typed UDF DataFrame query context at the user call site [spark]
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on code in PR #59474:
URL: https://github.com/apache/spark/pull/59474#discussion_r4278133992
##########
sql/api/src/main/scala/org/apache/spark/sql/catalyst/trees/origin.scala:
##########
@@ -152,8 +152,13 @@ object CurrentOrigin {
"(org\\.apache\\.spark\\.sql\\." +
"(?:(classic|connect)\\.)?" +
"(?:functions|Column|ColumnName|SQLImplicits|Dataset|DataFrameStatFunctions|DatasetHolder" +
- "|SparkSession|ColumnNodeToProtoConverter)" +
+ "|SparkSession|ColumnNodeToProtoConverter" +
+ // Typed APIs build their UDF column inside these helpers (e.g. the Connect client encodes
Review Comment:
**Nit (P3):** This comment sits before the `KeyValueGroupedDataset|KeyValueGroupedDatasetImpl` alternatives but cites `groupByKey/filter/reduce`. In the Connect client, `filter` and `reduce` call `udf.apply(col("*"))` directly inside `Dataset` (already matched) and never go through these classes. Only `groupByKey` builds its grouping UDF in `KeyValueGroupedDatasetImpl`. What fixes `filter`/`reduce` and ordinary UDF calls is the new `org.apache.spark.sql.expressions.UserDefinedFunction` alternative below, which has no comment. A future reader could easily take that entry for redundant. Could you split the explanation so each alternative is justified where it appears? UserDefinedFunction would cover every UDF column, since all are built via `UserDefinedFunction.apply`. KeyValueGroupedDataset/Impl would cover the Connect groupByKey factory and columns built in KVGD methods.
**Verification:**
- **Inspection:** The reworded comments match the actual call paths (connect Dataset filter/reduce call udf.apply directly; groupByKey goes through KeyValueGroupedDatasetImpl), contain only ASCII, and keep lines within 100 characters.
- **Compile:** The sql/api module still compiles, and the concatenated pattern string is byte-identical to before.
##########
sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/KeyValueGroupedDatasetE2ETestSuite.scala:
##########
@@ -51,6 +53,33 @@ class KeyValueGroupedDatasetE2ETestSuite extends QueryTest with RemoteSparkSessi
assert(values == Arrays.asList[Int](5, 5))
}
+ test("SPARK-60247: typed groupByKey origin points to the user call site") {
+ // The Connect client encodes `groupByKey`'s grouping key as a Scala UDF applied to
+ // `col("*")`, built inside `KeyValueGroupedDatasetImpl`. The DataFrame query context
+ // captured for that column must skip the `UserDefinedFunction.apply` and
+ // `KeyValueGroupedDatasetImpl` frames so the call site lands on the user code that called
+ // `groupByKey` (this suite), not on those internal helpers. Assert on the serialized origin
+ // of the grouping expression to guard the frame filter in `CurrentOrigin`.
+ val ds = spark
+ .range(10)
+ .groupByKey(v => v % 2)
+ .count()
+
+ val groupingUdf = ds.plan.getRoot.getAggregate.getGroupingExpressions(0)
Review Comment:
**Nit (P3):** This test only checks `getGroupingExpressions(0)`, the grouping UDF built in `KeyValueGroupedDatasetImpl$.apply`. Its call chain (`Column` -> `UserDefinedFunction.apply` -> `KeyValueGroupedDatasetImpl$.apply` -> `Dataset.groupByKey` -> test) never passes through a plain `KeyValueGroupedDataset` frame. So removing `|KeyValueGroupedDataset` from `sparkCodePattern` (keeping `|KeyValueGroupedDatasetImpl`) would not fail this test or the `DatasetSuite` one. That token is what moves the call site for columns built inside `KeyValueGroupedDataset` methods. For example, `.count()` delegates to `sql.KeyValueGroupedDataset.count`, which builds `agg(cnt(lit(1)).as(...))`. Could you also assert on the origin of the aggregate expression in this same plan (`getAggregate.getAggregateExpressions(0)`) so that behavior is guarded too?
**Verification:**
- **Regression:** An assertion on the origin of a column built inside a KeyValueGroupedDataset method (e.g. the count aggregate) fails when the KeyValueGroupedDataset alternative is dropped, and passes with it present.
- **Behavior:** The existing grouping-UDF call-site assertion still passes.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60247][SQL][CONNECT] Point typed UDF DataFrame query context at the user call site [spark]
Posted by "dbtsai (via GitHub)" <gi...@apache.org>.
dbtsai closed pull request #59474: [SPARK-60247][SQL][CONNECT] Point typed UDF DataFrame query context at the user call site
URL: https://github.com/apache/spark/pull/59474
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-60247][SQL][CONNECT] Point typed UDF DataFrame query context at the user call site [spark]
Posted by "dbtsai (via GitHub)" <gi...@apache.org>.
dbtsai commented on PR #59474:
URL: https://github.com/apache/spark/pull/59474#issuecomment-6184345524
**Merge Summary:**
- merged into master https://github.com/apache/spark/commit/8e4ef5a21d787a46c723b3f2db20cf55fdea8a24
- merged into branch-4.x https://github.com/apache/spark/commit/b91beaa305cd1e66f6f9c6d4074034341d2dfa5a
*Posted by `merge_spark_pr.py`*
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org