[SPARK-54838][SQL] A new feature to optimize partition size and count

[dotsering]

What changes were proposed in this pull request?

I am proposing to add a new function in Dataset class to fix small file problem. We have noticed that if the source data our spark job reads have many small files (KB size), it creates lot of partitions. This PR adds a new function named optimizePartition which when used creates partitions of size 128MB if no size passed. You can pass your own desired partition size.

Why are the changes needed?

The changes are needed to solve small file problem. It also helps in reducing the number of files that gets written back to sink.

Does this PR introduce any user-facing change?

It does not introduce any change in any existing feature/functions of Dataset. It is a brand new function.

How was this patch tested?

I have added number of unit tests that covers the scenario of lot of small partitions and when this function is called, it either coalesces to reduce partition count or uses repartition to increase partition count if partition size is too big. Also tested it locally using a dockerized environment containing spark cluster.

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

I did most of the coding. I used Gemini along to help me walk through the process of opening PRs, jira ticket and using linters/formatters. This PR does not have lot of code change.

[wangyum]

Could we consider setting spark.sql.files.maxPartitionNum to address the small file issue?

[dotsering]

Could we consider setting spark.sql.files.maxPartitionNum to address the small file issue?

Thanks @wangyum for your feedback. Here are few limitations that I can think of using hardcoded maxPartitionNum property:

1. Incompatibility with Heterogeneous Data Sources: The configuration lacks flexibility because it applies to the read operation globally. An ideal value for one source is often catastrophic for another within the same job. For example, setting maxPartitionNum = 1 is perfect for my example 10,000 tiny 10KB files (reducing overhead), but if that same job also reads a large 1TB table, that table is also forced into a single partition, causing an immediate failure. You cannot optimize the small source without breaking the large one. You have to keep fine tuning it for every spark job. That was the main reason I thought of automating hardcoded value.

2. Destructive Merging of Optimized Data (Collateral Damage): This setting acts as a blunt instrument that overrides safe partition sizing logic. If a job reads a second data source that is already perfectly partitioned (e.g., 1,000 partitions of 150MB), setting maxPartitionNum = 10 forces Spark to ignore those healthy boundaries. It will aggressively coalesce the 1,000 partitions down to 10 massive 15GB partitions, destroying parallelism and guaranteeing Executor OOM errors on data that was originally fine.
   Here is a link to the loc where maxPartitionBytes could have helped generate 1000 150MB partitions but if we set too low value for maxPartitionNum, it will override that partition count and reduce 1000 partitions to 10 partitions.