[SPARK-52450][CONNECT] Improve performance of schema deepcopy

[xi-db]

What changes were proposed in this pull request?

In Spark Connect, DataFrame.schema returns a deep copy of the schema to prevent unexpected behavior caused by user modifications to the returned schema object. However, if a user accesses df.schema repeatedly on a DataFrame with a complex schema, it can lead to noticeable performance degradation.

The performance issue can be reproduced using the code snippet below. Since copy.deepcopy is known to be slow to handle complex objects, this PR replaces it with pickle-based ser/de to improve the performance of df.schema access. Given the limitations of pickle, the implementation falls back to deepcopy in cases where pickling fails.

from pyspark.sql.types import StructType, StructField, StringType

def make_nested_struct(level, max_level, fields_per_level):
    if level == max_level - 1:
        return StructType(
            [StructField(f"f{level}_{i}", StringType(), True) for i in range(fields_per_level)])
    else:
        return StructType(
            [StructField(f"s{level}_{i}",
                         make_nested_struct(level + 1, max_level, fields_per_level), True) for i in
             range(fields_per_level)])

# Create a 4 level nested schema with in total 10,000 leaf fields
schema = make_nested_struct(0, 4, 10)

The existing needs 21.9s to copy the schema for 100 times.

import copy
timeit.timeit(lambda: copy.deepcopy(schema), number=100)
# 21.9

The updated approach only needs 2.0s to copy for 100 times:

from pyspark.serializers import CPickleSerializer
cached_schema_serialized = CPickleSerializer().dumps(schema)

timeit.timeit(lambda: CPickleSerializer().loads(cached_schema_serialized), number=100)
# 2.0

Why are the changes needed?

It improves the performance when calling df.schema many times.

Does this PR introduce any user-facing change?

No.

How was this patch tested?

Existing tests and new tests.

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

No.

[xi-db]

Hi @hvanhovell @vicennial , could you take a look at this PR? Thanks.

[hvanhovell]

In what cases do we think this will happen?

[xi-db]

I think it never happens because schema is nested Spark type classes. It shouldn't hit any of those special types that pickle doesn't support (link). Anyway, maybe we still need to handle the exception just in case. What do you think?

[heyihong]

It would be nice to add some comments to the code to make it easier for future readers to understand. Currently, it's not very clear why a CPickleSerializer is used, or why an error is being handled without looking at the corresponding pull request.

Also, it may be clearer to create a function called _fast_cached_schema_deepcopy that caches the serialized schema and then deserializes it.

[hvanhovell]

How about we store the proto response instead of doing the pickle? In that case we are guaranteed it will always work. Does that have the same performance problems?

[xi-db]

@hvanhovell Just tested. types.proto_schema_to_pyspark_data_type took 6.1s to transform the proto schema 100 times. Compared with it, pickle is faster - 2.0s.

[hvanhovell]

works for me. thanks for looking into it!

[grundprinzip]

Before merging would be great to test @hvanhovell's proposal to test the performance of reconstructing the schema from the proto response at all times and what the impact is.