Title: Update `_format_columns` in `remove_columns`
State: closed

As explained at #4398, when calling `dataset.add_faiss_index` under certain conditions when calling a sequence of operations `cast_column`, `map`, and `remove_columns`, it fails as it's trying to look for already removed columns.

So on, after testing some possible fixes, it seems that setting the dataset format right after removing the columns seems to be working fine, so I had to add a call to `.set_format` in the `remove_columns` function.

Hope this helps!

Comments:
🤗 This PR closes https://github.com/huggingface/datasets/issues/4398

_The documentation is not available anymore as the PR was closed or merged._

Hi! Thanks for reporting and providing a fix. I made a small change to make the fix easier to understand.

Hi, @mariosasko thanks! It makes sense, sorry I'm not that familiar with `datasets` code 😩 

Sure @albertvillanova I'll do that later today and ping you once done, thanks! :hugs:

Hi again @albertvillanova! Let me know if those tests are fine 🤗 

Hi @alvarobartt,

I think your tests are failing. I don't know why previously, after your last commit, the CI tests were not triggered. 

In order to force the re-running of the CI tests, I had to edit your file using the GitHub UI.

First I tried to do it using my terminal, but I don't have push right to your PR branch: I would ask you next time you open a PR, please mark the checkbox "Allow edits from maintainers": https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/allowing-changes-to-a-pull-request-branch-created-from-a-fork#enabling-repository-maintainer-permissions-on-existing-pull-requests

Hi @albertvillanova, let me check those again! And regarding that checkbox I thought it was already checked so my bad there 😩 

@albertvillanova  again it seems that the tests were not automatically triggered, but I tested those locally and now they work, as previously those were failing as I created an assertion as `self.assertEqual` over an empty list that was compared as `None` while the value was `[]` so I updated it to be `self.assertListEqual` and changed the comparison value to `[]`.

@lhoestq any idea why the CI is not triggered?

@alvarobartt I have tested locally and the tests continue failing.

I think there is a basis error: `new_dset._format_columns` is always `None` in those cases.


You're right @albertvillanova I was indeed running the tests with `datasets==2.2.0` rather than with the branch version, I'll check it again! Sorry for the inconvenience...

> @alvarobartt I have tested locally and the tests continue failing.
> 
> I think there is a basis error: `new_dset._format_columns` is always `None` in those cases.

In order to have some regressions tests for the fixed scenario, I've manually updated the value of `_format_columns` in the `ArrowDataset` so as to check whether it's updated or not right after calling `remove_columns`, and it does behave as expected, so with the latest version of this branch the reported issue doesn't occur anymore.

Hi again @albertvillanova sorry I was on leave! I'll do that ASAP :hugs:

@albertvillanova, does it make sense to add regression tests for `DatasetDict`? As `DatasetDict` doesn't have the attribute `_format_columns`, when we call `remove_columns` over a `DatasetDict` it removes the columns and updates the attributes of each split which is an `ArrowDataset`.

So on, we can either:
- Update first the `_format_columns` attribute of each split and then remove the columns over the `DatasetDict`
- Loop over the splits of `DatasetDict` and remove the columns right after updating `_format_columns` of each `ArrowDataset`.

I assume that the best regression test is the one implemented (mentioned first above), let me know if there's a better way to do that 👍🏻 

I think there's already a decorator to support transmitting the right `_format_columns`: `@transmit_format`, have you tried adding this decorator to `remove_columns` ?

> I think there's already a decorator to support transmitting the right `_format_columns`: `@transmit_format`, have you tried adding this decorator to `remove_columns` ?

Hi @lhoestq I can check now!

It worked indeed @lhoestq, thanks for the proposal and the review! 🤗 

Oops, I forgot about `@transmit_format`'s existence. From what I see, we should also use this decorator in `flatten`, `rename_column` and  `rename_columns`. 

@alvarobartt Let me know if you'd like to work on this (in a subsequent PR).

Sure @mariosasko I can prepare another PR to add those too, thanks!