Fix shard count limit

Fix shard count limit at another place,
Use partitioned as a flag

P.S this code is repeated at 3 places, at some point we should move it in separate module and remove redundant code.