Allow the ability for uneven row wise sharding based on number of buckets for zch (#3341)

[nipung90]

Summary:

X-link: #3341

This diff enables the use of num_buckets ParameterConstraint in the planner. The presence of this planner will indicate the use of ZCH bucketing as part of rowwise sharding plans.

Without num_buckets present:

The current row-wise sharding strategy will be used.

With num_buckets present:

• When devices have the same amount of memory available:

We will divide the buckets to be evenly distributed across hosts and distribute an additional bucket to the required number of hosts to handle the remainders.
For eg. if
Test case 2: hash_size = 100, num_devices = 4, num_buckets = 10
Each bucket has 10 rows, buckets distributed as [3,3,2,2]
So rows are distributed as [30,30,20,20]

• When devices have uneven amount of memory

We will distribute the buckets in the proportion of the memory of the device to the total memory of all devices and all the remaining buckets left are stored on the last device in the case where buckets do not completely fit based on the memory ratios.
for eg
hash_size = 45, num_buckets = 9, bucket_size = 5
With memory ratio 2:1:1, buckets should be distributed as [4,2,3]
So rows are distributed as [20,10,15]

Differential Revision: D84384489

[meta-codesync]

@nipung90 has exported this pull request. If you are a Meta employee, you can view the originating Diff in D84384489.