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: [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]
