TS_MODIFY stage and UPDATE stage both need to check for shard key changes. The helper functions can be refactored to be shared.