Repository: huggingface/datasets
File: tests/test_iterable_dataset.py (Python)

Code before:
        {("new_id" if k == "id" else k): v for k, v in example.items()} for example in dataset_with_several_columns
    ]
    assert new_dataset.features is None
    assert new_dataset.column_names is None
    # rename the column if ds.features was not None
    new_dataset = dataset_with_several_columns._resolve_features().rename_column("id", "new_id")
    assert new_dataset.features is not None
    assert new_dataset.column_names is not None
    assert "id" not in new_dataset.column_names
    assert "new_id" in new_dataset.column_names


def test_iterable_dataset_rename_columns(dataset_with_several_columns: IterableDataset):
    column_mapping = {"id": "new_id", "filepath": "filename"}
    new_dataset = dataset_with_several_columns.rename_columns(column_mapping)
    assert list(new_dataset) == [
        {column_mapping.get(k, k): v for k, v in example.items()} for example in dataset_with_several_columns
    ]
    assert new_dataset.features is None
    assert new_dataset.column_names is None
    # rename the columns if ds.features was not None
    new_dataset = dataset_with_several_columns._resolve_features().rename_columns(column_mapping)
    assert new_dataset.features is not None
    assert new_dataset.column_names is not None
    assert all(c not in new_dataset.column_names for c in ["id", "filepath"])
    assert all(c in new_dataset.column_names for c in ["new_id", "filename"])


def test_iterable_dataset_remove_columns(dataset_with_several_columns: IterableDataset):
    new_dataset = dataset_with_several_columns.remove_columns("id")
    assert list(new_dataset) == [
        {k: v for k, v in example.items() if k != "id"} for example in dataset_with_several_columns
    ]
    assert new_dataset.features is None
    new_dataset = dataset_with_several_columns.remove_columns(["id", "filepath"])
    assert list(new_dataset) == [
        {k: v for k, v in example.items() if k != "id" and k != "filepath"} for example in dataset_with_several_columns
    ]
    assert new_dataset.features is None
    assert new_dataset.column_names is None
    # remove the columns if ds.features was not None
    new_dataset = dataset_with_several_columns._resolve_features().remove_columns(["id", "filepath"])
    assert new_dataset.features is not None
    assert new_dataset.column_names is not None
    assert all(c not in new_dataset.features for c in ["id", "filepath"])
    assert all(c not in new_dataset.column_names for c in ["id", "filepath"])


def test_iterable_dataset_select_columns(dataset_with_several_columns: IterableDataset):
    new_dataset = dataset_with_several_columns.select_columns("id")
    assert list(new_dataset) == [

Review comment (bug):
I don't know the implementation of `split_dataset_by_node` very well, just trying to raise a concern on this test case and please correct me if I'm wrong.

Assume world_size = 2, count = 1. Assume underlying data is [0, 1, 2, ..., 9].
`true_distributed_dataset` on rank 0 would be `[0, 2, 4, 6, 8]`
`true_distributed_dataset` on rank 1 would be `[1, 3, 5, 7, 9]`
and after calling `skip` and `split_dataset_by_node`
`distributed_dataset` on rank 0 would be `[1, 3, 5, 7, 9]`
`distributed_dataset` on rank 1 would be `[2, 4, 6, 8]`

Code after:
        )


@pytest.mark.parametrize("method", ["skip", "take"])
@pytest.mark.parametrize("after_split_by_node", [False, True])
@pytest.mark.parametrize("count", [2, 5, 11])
def test_iterable_dataset_skip_or_take_after_split_by_node(method, after_split_by_node, count):
    n, n_shards = 3, 10
    rank, world_size = 1, 2
    ex_iterable = ExamplesIterable(generate_examples_fn, {"n": n, "filepaths": [f"{i}.txt" for i in range(n_shards)]})
    dataset = IterableDataset(ex_iterable)
    distributed_dataset = dataset
    true_distributed_dataset = split_dataset_by_node(dataset, rank=rank, world_size=world_size)
    if after_split_by_node:
        distributed_dataset = split_dataset_by_node(distributed_dataset, rank=rank, world_size=world_size)
        distributed_dataset = distributed_dataset.skip(count) if method == "skip" else distributed_dataset.take(count)
        assert (
            list(true_distributed_dataset)[count:]
            if method == "skip"
            else list(true_distributed_dataset)[:count] == list(distributed_dataset)
        )
    else:
        distributed_dataset = distributed_dataset.skip(count) if method == "skip" else distributed_dataset.take(count)
        distributed_dataset = split_dataset_by_node(distributed_dataset, rank=rank, world_size=world_size)
        assert len(
            list(true_distributed_dataset)[count // world_size :]
            if method == "skip"
            else list(true_distributed_dataset)[: count // world_size]
        ) == len(list(distributed_dataset))


def test_iterable_dataset_add_column(dataset_with_several_columns: IterableDataset):
    new_column = list(range(3 * DEFAULT_N_EXAMPLES))
    new_dataset = dataset_with_several_columns.add_column("new_column", new_column)
    assert list(new_dataset) == [
        {**example, "new_column": idx} for idx, example in enumerate(dataset_with_several_columns)
    ]
    new_dataset = new_dataset._resolve_features()
    assert "new_column" in new_dataset.column_names


def test_iterable_dataset_rename_column(dataset_with_several_columns: IterableDataset):
    new_dataset = dataset_with_several_columns.rename_column("id", "new_id")
    assert list(new_dataset) == [
        {("new_id" if k == "id" else k): v for k, v in example.items()} for example in dataset_with_several_columns
    ]
    assert new_dataset.features is None
    assert new_dataset.column_names is None
    # rename the column if ds.features was not None
    new_dataset = dataset_with_several_columns._resolve_features().rename_column("id", "new_id")
    assert new_dataset.features is not None