[SPARK-45516][CONNECT] Include QueryContext in SparkThrowable proto message

[heyihong]

What changes were proposed in this pull request?

• Include QueryContext in SparkThrowable proto message
• Reconstruct QueryContext for SparkThrowable exceptions on the client side

Why are the changes needed?

• Better integration with the error framework

Does this PR introduce any user-facing change?

No

How was this patch tested?

build/sbt "connect-client-jvm/testOnly *ClientE2ETestSuite"

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

[heyihong]

@hvanhovell @juliuszsompolski @HyukjinKwon Please take a look

[juliuszsompolski]

LGTM

[MaxGekk]

@heyihong Please, add the [CONNECT] tag.

[HyukjinKwon]

Merged to master.

[HyukjinKwon]

qq; could we do this in Python client too?

[heyihong]

Not yet if I understand correctly. @gatorsmile has some concern about whether exposing QueryContext in the PySpark Exception APIs makes sense for non-sql PySpark exceptions. There is some ongoing discussion for this.

[HyukjinKwon]

Hm, this should actually fail after #43334 PR because DataFrameQueryContext throws an exception now (https://github.com/apache/spark/pull/43334/files#diff-b3bc05fec45cd951053b2876c71c7730b63789cb4336a7537a6654c724db3241R586-R589).

Seems like we miss this information Spark Connect client sides for now so it seems working .. but we should carry this context for DataFrame operations ..

[HyukjinKwon]

@MaxGekk I think there's a bug not only here but without Spark Connect.

scala> try {
     |   spark.range(1).select("a")
     | } catch {
     |     case e: org.apache.spark.sql.catalyst.ExtendedAnalysisException => println(e.getQueryContext)
     | }
[Lorg.apache.spark.QueryContext;@6a9d7514
val res2: Any = ()

It doesn't contain QueryContext from #43334 ..

[HyukjinKwon]

Alright, there are three issues. TL;DR:

1. There are many places that does not provide Origin that contains QueryContext at QueryCompilationErrors. So the QueryContext is often missing (e.g., the case above). cc @MaxGekk

2. . Origin.context is being directly used on Executor side. e.g., origin.context at DivideYMInterval. This seems wrongly returning SQLQueryContext for DataFrame operations (from Executor side) in Spark Connect server because we do not call withOrigin there. That's why this specific code did not throw an exception from https://github.com/apache/spark/pull/43334/files#diff-b3bc05fec45cd951053b2876c71c7730b63789cb4336a7537a6654c724db3241R586-R589 because it has never been DataFrameContext cc @heyihong @juliuszsompolski

3. The current logic in ErrorUtils.scala should handle the case of DataFrameQueryContext e.g., DataFrameQueryContext.stopIndex() will throw an exception. Should we:

   • Set a default value in DataFrameQueryContext.stopIndex() instead of throwing an exception? @MaxGekk
   • Or, make the protobuf message for this optional, and throw an exception from Spark Connect client sides? @heyihong @juliuszsompolski

[cloud-fan]

We should invoke withOrigin in Spark Connect server.

do you mean client? Server side stacktrace is not interesting to users to understand their own code mistakes.

[HyukjinKwon]

I made a mistake. I updated my comment.

[heyihong]

Either setting a default value or making the protobuf field optional should work. It depends on the semantic of DataFrameQueryContext.stopIndex().

cc: @MaxGekk

[HyukjinKwon]

Let's make sure we fix this.