Reverted the change to PartitionSet.size(). For some reason this doesn't work as expected... I'm just going to go back to what we had before.