Fix removal from shuffleToMapStage to search for a key-value pair with
our stage instead of using our shuffleID.