HyukjinKwon opened a new pull request, #57801:
URL: https://github.com/apache/spark/pull/57801
### What changes were proposed in this pull request?
This PR makes `_with_origin` in `python/pyspark/errors/utils.py` re-entrant.
`_with_origin` captures the Python call site of a PySpark API call so that a
`DataFrameQueryContext` can point at the user code that built a failing
expression. A
PySpark API can invoke other decorated APIs internally, for example
`Column.__and__`
calls `lit`. Today the nested call overwrites the origin and then clears it
in its
`finally` block, so the outer call continues with no origin at all.
This PR guards against that the same way the JVM side `withOrigin` already
does
(`sql/api/src/main/scala/org/apache/spark/sql/catalyst/trees/origin.scala`):
```scala
if (CurrentOrigin.get.stackTrace.isDefined ||
!SqlApiConf.get.dataFrameQueryContextEnabled) {
```
> As there might be multiple nested `withOrigin` calls (e.g. any Spark API
implementations
> can invoke other APIs) only the first `withOrigin` is captured because
that is closer to
> the user code.
When an origin is already set, capture is skipped entirely and the outermost
origin is
kept. The classic path now also records the call site on the Python side (in
addition to
the JVM side) because that is the state the guard reads.
### Why are the changes needed?
The Python side is missing the nesting guard that the JVM side has, which
makes the
decorator non-composable: applying it to an API that is reachable from
another decorated
API silently drops the outer context.
Today only `Column` methods and `col` are decorated, and nothing internally
calls `col`,
so this is currently latent and not user-facing. It becomes a blocker for
expanding
`_with_origin` coverage to `pyspark.sql.functions` (`F.split`, `F.to_date`,
...), where
expressions built from those functions carry no DataFrame context at all.
Many of those
functions call `lit` internally (`instr` and `make_interval` call it 7 times
each), so
without this fix the expansion would attribute the context to the nested
`lit` or lose it.
Fixing the decorator first keeps that follow-up change to purely adding
decorators.
As a side effect the guard also avoids redundant work in the nested case: it
skips the
stack walk, and on the classic path it skips the py4j `set`/`clear` round
trips and the
`spark.sql.stackTracesInDataFrameContext` conf read.
### Does this PR introduce _any_ user-facing change?
No. Only `Column` methods and `col` are currently decorated and no decorated
API invokes
another one, so no error message changes. This fixes latent behavior in an
internal
decorator and unblocks a follow-up that will be user-facing.
### How was this patch tested?
New unit test `test_with_origin_is_reentrant` in
`python/pyspark/sql/tests/test_dataframe_query_context.py`, which asserts
that a nested
decorated call neither overwrites nor clears the outer origin, and that the
origin is
fully cleared once the outermost call returns. Verified that it fails
without the fix
(`observed == [None, None]`) and passes with it (`observed == ["outer",
"outer"]`).
Ran locally:
- `pyspark.sql.tests.test_dataframe_query_context` (classic)
- `pyspark.sql.tests.connect.test_parity_dataframe_query_context` (Connect)
- `pyspark.errors.tests.test_errors`
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 5)
--
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]
