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

Code before:

@pytest.mark.parametrize("max_in_memory_dataset_size", ["default", 0, 100, 1000])
def test_load_from_disk_with_default_in_memory(
    max_in_memory_dataset_size, dataset_loading_script_dir, data_dir, tmp_path, monkeypatch
):
    current_dataset_size = 512  # arrow file size = 512, in-memory dataset size = 148
    if max_in_memory_dataset_size == "default":
        max_in_memory_dataset_size = 0  # default
    else:
        monkeypatch.setattr(datasets.config, "IN_MEMORY_MAX_SIZE", max_in_memory_dataset_size)
    if max_in_memory_dataset_size:
        expected_in_memory = current_dataset_size < max_in_memory_dataset_size
    else:
        expected_in_memory = False

    dset = load_dataset(dataset_loading_script_dir, data_dir=data_dir, keep_in_memory=True)
    dataset_path = os.path.join(tmp_path, "saved_dataset")
    dset.save_to_disk(dataset_path)

    with assert_arrow_memory_increases() if expected_in_memory else assert_arrow_memory_doesnt_increase():
        _ = load_from_disk(dataset_path)


@contextmanager
def moto_server():
    with patch.dict(os.environ, {"AWS_ENDPOINT_URL": "http://localhost:5000"}):
        server = ThreadedMotoServer()
        server.start()
        try:
            yield
        finally:
            server.stop()


def test_load_file_from_s3():
    # we need server mode here because of an aiobotocore incompatibility with moto.mock_aws
    # (https://github.com/getmoto/moto/issues/6836)
    with moto_server():
        # Create a mock S3 bucket
        bucket_name = "test-bucket"
        s3 = boto3.client("s3", region_name="us-east-1")
        s3.create_bucket(Bucket=bucket_name)

        # Upload a file to the mock bucket
        key = "test-file.csv"
        csv_data = "Island\nIsabela\nBaltra"

        s3.put_object(Bucket=bucket_name, Key=key, Body=csv_data)

        # Load the file from the mock bucket
        ds = datasets.load_dataset("csv", data_files={"train": "s3://test-bucket/test-file.csv"})

Review comment (security):
You might have to add some dummy creds here to get past the NoCredentialsError in the test http://docs.getmoto.org/en/latest/docs/getting_started.html#recommended-usage

Code after:

    # Upload a file to the mock bucket
    key = "test-file.csv"
    csv_data = "Island\nIsabela\nBaltra"

    s3.put_object(Bucket=bucket_name, Key=key, Body=csv_data)

    # Load the file from the mock bucket
    ds = datasets.load_dataset("csv", data_files={"train": "s3://test-bucket/test-file.csv"})

    # Check if the loaded content matches the original content
    assert list(ds["train"]) == [{"Island": "Isabela"}, {"Island": "Baltra"}]


@pytest.mark.integration
def test_remote_data_files():
    repo_id = "hf-internal-testing/raw_jsonl"
    filename = "wikiann-bn-validation.jsonl"
    data_files = f"https://huggingface.co/datasets/{repo_id}/resolve/main/{filename}"
    ds = load_dataset("json", split="train", data_files=data_files, streaming=True)
    assert isinstance(ds, IterableDataset)
    ds_item = next(iter(ds))
    assert ds_item.keys() == {"langs", "ner_tags", "spans", "tokens"}


def distributed_load_dataset(args):
    data_name, tmp_dir, datafiles = args
    dataset = load_dataset(data_name, cache_dir=tmp_dir, data_files=datafiles)
    return dataset


def test_load_dataset_distributed(tmp_path, csv_path):
    num_workers = 5
    args = "csv", str(tmp_path), csv_path
    with Pool(processes=num_workers) as pool:  # start num_workers processes
        datasets = pool.map(distributed_load_dataset, [args] * num_workers)
        assert len(datasets) == num_workers
        assert all(len(dataset) == len(datasets[0]) > 0 for dataset in datasets)
        assert len(datasets[0].cache_files) > 0
        assert all(dataset.cache_files == datasets[0].cache_files for dataset in datasets)


def distributed_load_dataset_with_script(args):
    data_name, tmp_dir, download_mode = args
    dataset = load_dataset(data_name, cache_dir=tmp_dir, download_mode=download_mode)
    return dataset


@require_not_windows  # windows doesn't support overwriting Arrow files from other processes
@pytest.mark.parametrize("download_mode", [None, "force_redownload"])
def test_load_dataset_distributed_with_script(tmp_path, download_mode):