[Refactor] - Modify the output schema when idx != -1