fix: Native shuffle fails with a 2GB task serialization OOM on jobs with many partitions

[parthchandra]

Which issue does this PR close?

Closes #5391.

Rationale for this change

A native-shuffle job with a very large number of partitions fails at stage submission, before any task runs:

(From the issue:
Task serialization failed: java.lang.OutOfMemoryError: Required array length 2147483639 + 794 is too large
at java.io.ByteArrayOutputStream.ensureCapacity(...)
at org.apache.spark.scheduler.DAGScheduler.submitMissingTasks(...)

DAGScheduler serializes the (RDD, ShuffleDependency) pair into a single broadcast byte array that must fit in ~2GB. On the native-shuffle path, CometShuffleDependency.nativeShuffleSpec is a non-transient field holding a perPartitionByKey map with one scan-plan-data blob (the partition's file list) per map partition. The failing job had ~38.7M
partitions, so ~38.7M protobufs got baked into that one blob and blew the limit — even though each task only reads its own slice. Plain Spark avoids this by keeping per-partition file lists @transient and shipping them per task; Comet's own CometExecRDD does the same. The native shuffle path was the exception.

What changes are included in this PR?

Mirror CometExecRDD: keep the map off the serialized dependency and give each task only its slice.

• CometNativeShuffleInputRDD takes perPartitionByKey as a @transient arg and, in getPartitions (driver-side), slices out each partition's entry onto its Partition object. That slice flows to the writer, which injects it instead of indexing the full map.
• NativeExecContext.perPartitionByKey is now @transient so no build path can serialize the full map, and we also empty it when building the dependency.
• commonByKey is unchanged — it's sized by scan count, not partition count, and the writer still needs it.

How are these changes tested?

• A unit test serializes the input RDD at 10 vs 10,000 partitions and checks the size stays roughly constant, and that each partition carries only its own slice.
• An end-to-end test runs a native shuffle over a multi-partition native scan and compares results against Spark with an all-native plan, so a wrong slice would fail it.

[sunchao]

Reviewed 1b12802e123c020e4fd9bbfbce091bae818dee0b in five independent scopes. The production change looks correct. Driver-side slicing preserves the old per-partition lookup, the shared scan data remains available to the writer, and the original context is still available for range sampling. I found no introduced execution defect.

An isolated probe using Spark 3.5.9's real Java serializer and RDD/dependency classes confirmed that the parent task-binary pair grows from 13,322 to 10,342,982 bytes when increasing from 10 to 10,000 partitions. The PR's pair stays at 2,922 bytes, and a separately serialized partition retains the correct scan-data slice. This probe stubs unrelated native/protobuf/metrics types and does not execute native SQL. Exact-head CI has 61 passing checks and 9 skipped. I also confirmed that both new tests ran in the Linux Spark 3.4 and macOS Spark 4.0 shuffle jobs.

Approving the implementation. I left one P2 comment asking for the size regression to cover the actual (RDD, ShuffleDependency) pair so it protects the original failure path.

[sunchao]

[P2] Exercise the actual shuffle task binary

Could this also construct a native shuffle dependency and compare the serialized size of (rdd, shuffleDependency) for 10 and 10,000 partitions? That is what Spark broadcasts for a shuffle-map stage. Before this fix, the large map lived under CometShuffleDependency.nativeShuffleSpec.execContext, not on the thin RDD. In an isolated Spark serialization probe, restoring that dependency leak left the RDD at 880 bytes while the pair grew from 13,322 to 10,342,982 bytes. The current RDD-only assertion therefore stays green if those guards regress. Testing the actual pair would protect the reported 2GB failure without requiring a 2GB allocation.