Title:  `drop_last_batch=True` for IterableDataset map function is ignored with multiprocessing DataLoader 
Author: naba89 (NONE)
State: closed
Created: 2024-01-07T02:46:50Z

### Describe the bug

Scenario:
- Interleaving two iterable datasets of unequal lengths (`all_exhausted`), followed by a batch mapping with batch size 2 to effectively merge the two datasets and get a sample from each dataset in a single batch, with `drop_last_batch=True` to skip the last batch in case it doesn't have two samples.

What works:
- Using DataLoader with `num_workers=0`

What does not work:
- Using DataLoader with `num_workers=1`, errors in the last batch.

Basically, `drop_last_batch=True` is ignored when using multiple dataloading workers.

Please take a look at the minimal repro script below.

### Steps to reproduce the bug

```python
from datasets import Dataset, interleave_datasets
from torch.utils.data import DataLoader


def merge_samples(batch):
    assert len(batch['a']) == 2, "Batch size must be 2"
    batch['c'] = [batch['a'][0]]
    batch['d'] = [batch['a'][1]]
    return batch


def gen1():
    for ii in range(1, 8385):
        yield {"a": ii}


def gen2():
    for ii in range(1, 5302):
        yield {"a": ii}


if __name__ == '__main__':

    dataset1 = Dataset.from_generator(gen1).to_iterable_dataset(num_shards=1024)
    dataset2 = Dataset.from_generator(gen2).to_iterable_dataset(num_shards=1024)

    interleaved = interleave_datasets([dataset1, dataset2], stopping_strategy="all_exhausted")
    mapped = interleaved.map(merge_samples, batched=True, batch_size=2, remove_columns=interleaved.column_names,
                             drop_last_batch=True)

    # Works
    loader = DataLoader(mapped, batch_size=32, num_workers=0)
    i = 0
    for b in loader:
        print(i, b['c'].shape, b['d'].shape)
        i += 1

    print("DataLoader with num_workers=0 works")

    # Doesn't work
    loader = DataLoader(mapped, batch_size=32, num_workers=1)
    i = 0
    for b in loader:
        print(i, b['c'].shape, b['d'].shape)
        i += 1


```

### Expected behavior

 `drop_last_batch=True` should have same behaviour for `num_workers=0` and `num_workers>=1`

### Environment info

- `datasets` version: 2.16.1
- Platform: macOS-10.16-x86_64-i386-64bit
- Python version: 3.10.12
- `huggingface_hub` version: 0.20.2
- PyArrow version: 12.0.1
- Pandas version: 2.0.3
- `fsspec` version: 2023.6.0

I have also tested on Linux and got the same behavior.

Comments:
My current workaround this issue is to return `None` in the second element and then filter out samples which have `None` in  them.

```python
def merge_samples(batch):
    if len(batch['a']) == 1:
        batch['c'] = [batch['a'][0]]
        batch['d'] = [None]
    else:
        batch['c'] = [batch['a'][0]]
        batch['d'] = [batch['a'][1]]
    return batch
    
def filter_fn(x):
    return x['d'] is not None

# other code...
mapped = mapped.filter(filter_fn)
```

Hi @lhoestq , I found that this script no longer works due to this [line](https://github.com/huggingface/datasets/blob/f693f4e93aabafa878470c80fd42ddb10ec550d6/src/datasets/iterable_dataset.py#L1141), where it checks the validity before the column `a` is removed. Is this expected?