### Why are the changes needed?
As title.
### How was this patch tested?
UT are modified.
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes#7139 from pan3793/rebalance.
Closes#7139
edb070afd [Cheng Pan] fix
4d3984a92 [Cheng Pan] Fix Spark extension rules to support RebalancePartitions
Authored-by: Cheng Pan <chengpan@apache.org>
Signed-off-by: Cheng Pan <chengpan@apache.org>