[SPARK-47274][PYTHON][SQL] Provide more useful context for PySpark DataFrame API errors

[itholic]

What changes were proposed in this pull request?

This PR introduces an enhancement to the error messages generated by PySpark's DataFrame API, adding detailed context about the location within the user's PySpark code where the error occurred.

This directly adds a PySpark user call site information into DataFrameQueryContext added from #43334, aiming to provide PySpark users with the same level of detailed error context for better usability and debugging efficiency for DataFrame APIs.

This PR also introduces QueryContext.pysparkCallSite and QueryContext.pysparkFragment to get a PySpark information from the query context easily.

This PR also enhances the functionality of check_error so that it can test the query context if it exists.

Why are the changes needed?

To improve a debuggability. Errors originating from PySpark operations can be difficult to debug with limited context in the error messages. While improvements on the JVM side have been made to offer detailed error contexts, PySpark errors often lack this level of detail.

Does this PR introduce any user-facing change?

No API changes, but error messages will include a reference to the exact line of user code that triggered the error, in addition to the existing descriptive error message.

For example, consider the following PySpark code snippet that triggers a DIVIDE_BY_ZERO error:

1  spark.conf.set("spark.sql.ansi.enabled", True)
2  
3  df = spark.range(10)
4  df.select(df.id / 0).show()

Before:

pyspark.errors.exceptions.captured.ArithmeticException: [DIVIDE_BY_ZERO] Division by zero. Use `try_divide` to tolerate divisor being 0 and return NULL instead. If necessary set "spark.sql.ansi.enabled" to "false" to bypass this error. SQLSTATE: 22012
== DataFrame ==
"divide" was called from
java.base/jdk.internal.reflect.NativeMethodAccessorImpl.invoke0(Native Method)

After:

pyspark.errors.exceptions.captured.ArithmeticException: [DIVIDE_BY_ZERO] Division by zero. Use `try_divide` to tolerate divisor being 0 and return NULL instead. If necessary set "spark.sql.ansi.enabled" to "false" to bypass this error. SQLSTATE: 22012
== DataFrame ==
"divide" was called from
/.../spark/python/test_pyspark_error.py:4

Now the error message points out the exact problematic code path with file name and line number that user writes.

Points to the actual problem site instead of the site where the action was called

Even when action calling after multiple transform operations are mixed, the exact problematic site can be provided to the user:

In:

  1 spark.conf.set("spark.sql.ansi.enabled", True)
  2 df = spark.range(10)
  3
  4 df1 = df.withColumn("div_ten", df.id / 10)
  5 df2 = df1.withColumn("plus_four", df.id + 4)
  6
  7 # This is problematic divide operation that occurs DIVIDE_BY_ZERO.
  8 df3 = df2.withColumn("div_zero", df.id / 0)
  9 df4 = df3.withColumn("minus_five", df.id / 5)
 10
 11 df4.collect()

Out:

pyspark.errors.exceptions.captured.ArithmeticException: [DIVIDE_BY_ZERO] Division by zero. Use `try_divide` to tolerate divisor being 0 and return NULL instead. If necessary set "spark.sql.ansi.enabled" to "false" to bypass this error. SQLSTATE: 22012
== DataFrame ==
"divide" was called from
/.../spark/python/test_pyspark_error.py:8

How was this patch tested?

Added UTs.

Was this patch authored or co-authored using generative AI tooling?

No.

[itholic]

cc @HyukjinKwon I'm still working on Spark Connect support and unit tests but the basic structure is ready for review.

FYI, also cc @MaxGekk as you made a similar contribution on the JVM side.

[itholic]

Hi, @HyukjinKwon @MaxGekk could you take a look at the prototype when you have some time?

[itholic]

Added QueryContext testing for DataFrameContext and UTs. The CI failures seems not related. cc @HyukjinKwon FYI

[HyukjinKwon]

cc @cloud-fan too

[itholic]

FYI: None is default, so we don't need to specify like this when QueryContext not existing, but I made this test for explicit example.

[cloud-fan]

thanks, merging to master!

[itholic]

Thanks @cloud-fan @ueshin @HyukjinKwon @xinrong-meng for the review!

[HyukjinKwon]

Let's clarify why #45377 (comment) happens before we move further. That shouldn't happen from my understanding.

If we go with the current approach, it would need more changes. e.g., Column.substr because it takes a different set of arguments and types.

[cloud-fan]

@HyukjinKwon This probably can't cover all the cases, and we may need to add more overloads for certain functions that require non-expression parameters, but it shouldn't be many.

I think it's better than using ThreadLocal which can be quite fragile to pass values between Python and JVM.

[itholic]

The difficulty with the previous method was that it was not easy to perfectly sync the data between two separately operating TheadLocal, CurrentOrigin and PySparkCurrentOrigin.

After taking deeper look at the structure, I think we may be able to make the CurrentOrigin more flexible to support PySpark error context instead of adding a separate ThreadLocal like PySparkCurrentOrigin.

If it works, it seems possible to improve the structure to a more flexible while maintaining the existing communication rules between Python and JVM without adding helper functions such as PySpark-specific fn.

Let me give it a try and create a PR to refactoring the current structure, and ping you guys.

[itholic]

Let me give it a try and create a PR to refactoring the current structure, and ping you guys.

Created #46063.

[HyukjinKwon]

perfectly sync the data between two separately operating TheadLocal, CurrentOrigin and PySparkCurrentOrigin.

Why is that?

[itholic]

Because I called PySparkCurrentOrigin directly on the DataFrameQueryContext without utilizing withOrigin in the initial implementation. I realized it from recent review from the refactoring PR, so I'm currently trying to reintroduce PySparkCurrentOrigin there.

[gengliangwang]

Late LGTM!