[SPARK-44839][SS][CONNECT] Better Error Logging when user tries to serialize spark session

[WweiL]

What changes were proposed in this pull request?

Add a new error with detailed message when a user tries to access spark session and dataframe created using local spark session, in streaming spark connect foreach, foreachBatch and StreamingQueryListener.

Update: per reviewer's request, added a new error class PySparkPicklingError. Also move UDTF_SERIALIZATION_ERROR to the new class

Why are the changes needed?

Better error logging for the breaking change introduced in streaming spark connect.

Does this PR introduce any user-facing change?

Yes, before users can only see this non-informative error when they access a local spark session in their streaming connect related functions:

Traceback (most recent call last):
  File "/home/wei.liu/oss-spark/python/pyspark/serializers.py", line 459, in dumps
    return cloudpickle.dumps(obj, pickle_protocol)
  File "/home/wei.liu/oss-spark/python/pyspark/cloudpickle/cloudpickle_fast.py", line 73, in dumps
    cp.dump(obj)
  File "/home/wei.liu/oss-spark/python/pyspark/cloudpickle/cloudpickle_fast.py", line 632, in dump
    return Pickler.dump(self, obj)
TypeError: cannot pickle '_thread._local' object

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/wei.liu/oss-spark/python/pyspark/sql/connect/streaming/readwriter.py", line 508, in foreachBatch
    self._write_proto.foreach_batch.python_function.command = CloudPickleSerializer().dumps(
  File "/home/wei.liu/oss-spark/python/pyspark/serializers.py", line 469, in dumps
    raise pickle.PicklingError(msg)
_pickle.PicklingError: Could not serialize object: TypeError: cannot pickle '_thread._local' object

Now it is replaced with:

pyspark.errors.exceptions.base.PySparkPicklingError: [STREAMING_CONNECT_SERIALIZATION_ERROR] Cannot serialize the function `foreachBatch`. If you accessed the spark session, or a dataframe defined outside of the function, please be aware that they are not allowed in Spark Connect. For foreachBatch, please access the spark session using `df.sparkSession`, where `df` is the first parameter in your foreachBatch function. For StreamingQueryListener, please access the spark session using `self.spark`

How was this patch tested?

Add unit tests

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

No

[WweiL]

cc @rangadi @bogao007

[HyukjinKwon]

@itholic do we need a dedicated error class for PicklingError? e.g., PySparkPicklingError?

[itholic]

Yes, I believe we need a new error class for new type of user-facing errors. Could you add a new PySpark error class for representing pickle.PicklingError?? See https://github.com/apache/spark/pull/40938/files as an example. I think we can also do it as a follow ups.

[WweiL]

Yes sure I could do that. Just want to confirm the ask is to define a new PySparkPicklingError and replace this PySparkRuntimeError with that right?

[itholic]

define a new PySparkPicklingError and replace this PySparkRuntimeError with that right?

Correct :-)

[itholic]

Otherwise looks good

[itholic]

Yes, I believe we need a new error class for new type of user-facing errors. Could you add a new PySpark error class for representing pickle.PicklingError?? See https://github.com/apache/spark/pull/40938/files as an example. I think we can also do it as a follow ups.

[WweiL]

CC @ueshin, I also changed the serialization error thrown by UDTF, ptal!

[ueshin]

cc @allisonwang-db

[allisonwang-db]

The UDTF change looks good to me!

[ueshin]

LGTM, pending tests.

[WweiL]

@ueshin Can we have this merged? Thank you!

[HyukjinKwon]

@WweiL seems like your PR conflicted with this PR :-). Should probably update.

[HyukjinKwon]

Merged to master.