Fix remove_columns in the formatted case
`remove_columns` had no effect when running a function in `.map()` on dataset that is formatted

This aligns the logic of `map()` with the non formatted case and also with with https://github.com/huggingface/datasets/pull/7353