chore: import upstream snapshot with attribution

python/ray/data/tests/datasource/test_huggingface.py

@@ -0,0 +1,464 @@
+from unittest.mock import MagicMock, patch
+
+import datasets
+import pyarrow
+import pytest
+import requests
+from packaging.version import Version
+
+import ray
+from ray.data.dataset import Dataset, MaterializedDataset
+from ray.tests.conftest import *  # noqa
+
+
+@pytest.fixture
+def mock_hf_dataset():
+    """Create a mock HuggingFace dataset for testing."""
+    texts = [
+        "Climate change is a serious threat to our planet",
+        "We need to take action on global warming",
+        "Renewable energy is the future",
+        "Fossil fuels are destroying the environment",
+        "Solar power is becoming more affordable",
+        "Wind energy is growing rapidly",
+        "Electric vehicles are the way forward",
+        "Carbon emissions must be reduced",
+        "Green technology is advancing quickly",
+        "Sustainability is important for future generations",
+        "Climate science is well established",
+        "Ocean levels are rising due to warming",
+        "Extreme weather events are increasing",
+        "Biodiversity loss is accelerating",
+        "Deforestation contributes to climate change",
+        "Clean energy jobs are growing",
+        "Energy efficiency saves money",
+        "Public transportation reduces emissions",
+        "Plant-based diets help the environment",
+        "Recycling is essential for sustainability",
+    ]
+
+    # Create labels array with exactly the same length as texts
+    labels = [i % 2 for i in range(len(texts))]  # Alternating 0s and 1s
+
+    return datasets.Dataset.from_dict(
+        {
+            "text": texts,
+            "label": labels,
+        }
+    )
+
+
+@pytest.fixture
+def mock_hf_dataset_dict(mock_hf_dataset):
+    """Create a mock HuggingFace DatasetDict for testing."""
+    return datasets.DatasetDict({"train": mock_hf_dataset})
+
+
+@pytest.fixture
+def mock_hf_iterable_dataset():
+    """Create a mock HuggingFace IterableDataset for testing."""
+    texts = [
+        "Streaming climate tweet 1: The planet is warming",
+        "Streaming climate tweet 2: Renewable energy is key",
+        "Streaming climate tweet 3: We must act now",
+        "Streaming climate tweet 4: Solar panels everywhere",
+        "Streaming climate tweet 5: Wind turbines are beautiful",
+        "Streaming climate tweet 6: Electric cars are the future",
+        "Streaming climate tweet 7: Carbon neutral by 2050",
+        "Streaming climate tweet 8: Green energy revolution",
+        "Streaming climate tweet 9: Climate action needed",
+        "Streaming climate tweet 10: Sustainable development",
+        "Streaming climate tweet 11: Ocean conservation",
+        "Streaming climate tweet 12: Forest protection",
+        "Streaming climate tweet 13: Clean air matters",
+        "Streaming climate tweet 14: Water conservation",
+        "Streaming climate tweet 15: Biodiversity protection",
+    ]
+
+    labels = [1, 0, 1, 0, 0, 0, 1, 0, 1, 0, 1, 1, 0, 0, 1]
+
+    dataset = datasets.Dataset.from_dict(
+        {
+            "text": texts,
+            "label": labels,
+        }
+    )
+    iterable_dataset = dataset.to_iterable_dataset()
+    iterable_dataset.expected_count = len(texts)
+    return iterable_dataset
+
+
+@pytest.fixture
+def mock_parquet_urls():
+    """Fixture providing mock parquet URLs for testing."""
+    return [
+        "https://huggingface.co/datasets/test/parquet/train-00000-of-00001.parquet",
+        "https://huggingface.co/datasets/test/parquet/train-00001-of-00001.parquet",
+    ]
+
+
+@pytest.fixture
+def mock_resolved_urls():
+    """Fixture providing mock resolved URLs (after HTTP redirects) for testing."""
+    return [
+        "https://cdn-lfs.huggingface.co/datasets/test/parquet/train-00000-of-00001.parquet",
+        "https://cdn-lfs.huggingface.co/datasets/test/parquet/train-00001-of-00001.parquet",
+    ]
+
+
+@pytest.fixture
+def mock_ray_dataset(mock_hf_dataset):
+    """Fixture providing a mock Ray dataset that matches the mock HuggingFace dataset."""
+    return ray.data.from_items(
+        [
+            {"text": text, "label": label}
+            for text, label in zip(mock_hf_dataset["text"], mock_hf_dataset["label"])
+        ]
+    )
+
+
+@pytest.fixture
+def mock_successful_http_responses(mock_parquet_urls):
+    """Fixture providing mock successful HTTP responses for URL resolution."""
+    mock_responses = []
+    for url in mock_parquet_urls:
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_response.url = url
+        mock_responses.append(mock_response)
+    return mock_responses
+
+
+@pytest.fixture
+def mock_redirected_http_responses(mock_parquet_urls, mock_resolved_urls):
+    """Fixture providing mock HTTP responses that simulate redirects."""
+    mock_responses = []
+    for original_url, resolved_url in zip(mock_parquet_urls, mock_resolved_urls):
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_response.url = resolved_url
+        mock_responses.append(mock_response)
+    return mock_responses
+
+
+@pytest.fixture
+def mock_huggingface_datasource():
+    """Fixture providing the HuggingFaceDatasource class for mocking."""
+    from ray.data._internal.datasource.huggingface_datasource import (
+        HuggingFaceDatasource,
+    )
+
+    return HuggingFaceDatasource
+
+
+def verify_http_requests(mock_requests_head, expected_urls):
+    """Verify that HTTP requests were made correctly."""
+    assert mock_requests_head.call_count == len(expected_urls)
+
+    for i, url in enumerate(expected_urls):
+        call_args = mock_requests_head.call_args_list[i]
+        assert call_args[0][0] == url
+        assert call_args[1]["allow_redirects"] is True
+        assert call_args[1]["timeout"] == 5
+
+
+def verify_read_parquet_call(mock_read_parquet, expected_urls):
+    """Verify that read_parquet was called with correct parameters."""
+    mock_read_parquet.assert_called_once()
+    call_args = mock_read_parquet.call_args
+
+    # Check that the parquet URLs were passed
+    assert call_args[0][0] == expected_urls
+
+    # Check that the filesystem is HTTPFileSystem
+    assert "filesystem" in call_args[1]
+    assert "HTTPFileSystem" in str(type(call_args[1]["filesystem"]))
+
+    # Check that retry_exceptions includes FileNotFoundError and ClientResponseError
+    assert "ray_remote_args" in call_args[1]
+    assert FileNotFoundError in call_args[1]["ray_remote_args"]["retry_exceptions"]
+
+
+def verify_dataset_creation(ds, mock_hf_dataset):
+    """Verify that the dataset was created successfully."""
+    assert isinstance(ds, MaterializedDataset)
+    assert ds.count() == mock_hf_dataset.num_rows
+
+
+def setup_parquet_mocks(
+    mock_huggingface_datasource,
+    mock_parquet_urls,
+    mock_http_responses,
+    mock_ray_dataset,
+):
+    """Setup common mocking pattern for parquet-based tests."""
+    patches = []
+
+    # Mock the list_parquet_urls_from_dataset method
+    datasource_patch = patch.object(
+        mock_huggingface_datasource,
+        "list_parquet_urls_from_dataset",
+        return_value=mock_parquet_urls,
+    )
+    patches.append(datasource_patch)
+
+    # Mock the requests.head calls
+    requests_patch = patch("requests.head")
+    patches.append(requests_patch)
+
+    # Mock the read_parquet function
+    read_parquet_patch = patch("ray.data.read_api.read_parquet")
+    patches.append(read_parquet_patch)
+
+    # Start all patches
+    datasource_mock = datasource_patch.start()
+    requests_mock = requests_patch.start()
+    read_parquet_mock = read_parquet_patch.start()
+
+    # Configure mocks
+    requests_mock.side_effect = mock_http_responses
+    read_parquet_mock.return_value = mock_ray_dataset
+
+    return datasource_mock, requests_mock, read_parquet_mock, patches
+
+
+def hfds_assert_equals(hfds: datasets.Dataset, ds: Dataset):
+    hfds_table = hfds.data.table
+    ds_table = pyarrow.concat_tables([ray.get(tbl) for tbl in ds.to_arrow_refs()])
+
+    sorting = [(name, "descending") for name in hfds_table.column_names]
+    hfds_table = hfds_table.sort_by(sorting)
+    ds_table = ds_table.sort_by(sorting)
+
+    assert hfds_table.equals(ds_table)
+
+
+@pytest.mark.parametrize("num_par", [1, 4])
+def test_from_huggingface(mock_hf_dataset_dict, ray_start_regular_shared, num_par):
+    # Check that DatasetDict is not directly supported.
+    assert isinstance(mock_hf_dataset_dict, datasets.DatasetDict)
+    with pytest.raises(
+        DeprecationWarning,
+        match="You provided a Hugging Face DatasetDict",
+    ):
+        ray.data.from_huggingface(mock_hf_dataset_dict)
+
+    ray_datasets = {
+        "train": ray.data.from_huggingface(
+            mock_hf_dataset_dict["train"], override_num_blocks=num_par
+        ),
+    }
+
+    assert isinstance(ray_datasets["train"], ray.data.Dataset)
+    hfds_assert_equals(mock_hf_dataset_dict["train"], ray_datasets["train"])
+
+    # Test reading in a split Hugging Face dataset yields correct individual datasets
+    base_hf_dataset = mock_hf_dataset_dict["train"]
+    hf_dataset_split = base_hf_dataset.train_test_split(test_size=0.2)
+    ray_dataset_split_train = ray.data.from_huggingface(hf_dataset_split["train"])
+    assert ray_dataset_split_train.count() == hf_dataset_split["train"].num_rows
+
+
+@pytest.mark.skipif(
+    datasets.Version(datasets.__version__) < datasets.Version("2.8.0"),
+    reason="IterableDataset.iter() added in 2.8.0",
+)
+@pytest.mark.skipif(
+    Version(pyarrow.__version__) < Version("8.0.0"),
+    reason="pyarrow.Table.to_reader() added in 8.0.0",
+)
+# Note, pandas is excluded here because IterableDatasets do not support pandas format.
+@pytest.mark.parametrize(
+    "batch_format",
+    [None, "numpy", "arrow", "torch", "tensorflow", "jax"],
+)
+def test_from_huggingface_streaming(
+    mock_hf_iterable_dataset, batch_format, ray_start_regular_shared
+):
+    hfds = mock_hf_iterable_dataset.with_format(batch_format)
+    assert isinstance(hfds, datasets.IterableDataset)
+
+    ds = ray.data.from_huggingface(hfds)
+    expected_count = mock_hf_iterable_dataset.expected_count
+    assert ds.count() == expected_count
+
+
+@pytest.mark.skipif(
+    datasets.Version(datasets.__version__) < datasets.Version("2.8.0"),
+    reason="IterableDataset.iter() added in 2.8.0",
+)
+def test_from_huggingface_dynamic_generated(ray_start_regular_shared):
+    # https://github.com/ray-project/ray/issues/49529
+    # Mock the dynamic dataset loading
+    mock_dataset = datasets.Dataset.from_dict(
+        {
+            "text": [
+                "dynamic tweet 1",
+                "dynamic tweet 2",
+                "dynamic tweet 3",
+                "dynamic tweet 4",
+                "dynamic tweet 5",
+            ],
+            "label": [0, 1, 0, 1, 0],
+        }
+    )
+    mock_iterable = mock_dataset.to_iterable_dataset()
+
+    with patch("datasets.load_dataset", return_value=mock_iterable):
+        hfds = datasets.load_dataset(
+            "dataset-org/dream",
+            split="test",
+            streaming=True,
+            trust_remote_code=True,
+        )
+        ds = ray.data.from_huggingface(hfds)
+        ds.take(1)
+
+
+@pytest.mark.parametrize("override_num_blocks", [1, 2, 4, 8])
+def test_from_huggingface_override_num_blocks(
+    mock_hf_dataset, ray_start_regular_shared, override_num_blocks
+):
+    """Test that override_num_blocks works correctly with HuggingFace datasets."""
+    hf_train = mock_hf_dataset
+
+    ds_subset = ray.data.from_huggingface(
+        hf_train, override_num_blocks=override_num_blocks
+    )
+
+    assert isinstance(ds_subset, MaterializedDataset)
+
+    # Verify number of blocks
+    assert ds_subset.num_blocks() == override_num_blocks
+
+    # Verify data integrity
+    assert ds_subset.count() == hf_train.num_rows
+    hfds_assert_equals(hf_train, ds_subset)
+
+    # Test with a smaller subset to test edge cases
+    small_size = max(override_num_blocks * 3, 10)
+    hf_small = hf_train.select(range(min(small_size, hf_train.num_rows)))
+    ds_small = ray.data.from_huggingface(
+        hf_small, override_num_blocks=override_num_blocks
+    )
+
+    # Verify number of blocks
+    assert ds_small.num_blocks() == override_num_blocks
+
+    # Verify data integrity
+    assert ds_small.count() == hf_small.num_rows
+    hfds_assert_equals(hf_small, ds_small)
+
+
+def test_from_huggingface_with_parquet_files(
+    mock_hf_dataset,
+    ray_start_regular_shared,
+    mock_parquet_urls,
+    mock_ray_dataset,
+    mock_successful_http_responses,
+    mock_huggingface_datasource,
+):
+    """Test the distributed read path when parquet file URLs are available."""
+    datasource_mock, requests_mock, read_parquet_mock, patches = setup_parquet_mocks(
+        mock_huggingface_datasource,
+        mock_parquet_urls,
+        mock_successful_http_responses,
+        mock_ray_dataset,
+    )
+
+    try:
+        ds = ray.data.from_huggingface(mock_hf_dataset)
+
+        # Verify HTTP requests
+        verify_http_requests(requests_mock, mock_parquet_urls)
+
+        # Verify read_parquet call
+        verify_read_parquet_call(read_parquet_mock, mock_parquet_urls)
+
+        # Verify dataset creation
+        verify_dataset_creation(ds, mock_hf_dataset)
+
+    finally:
+        # Stop all patches
+        for patch_obj in patches:
+            patch_obj.stop()
+
+
+def test_from_huggingface_with_resolved_urls(
+    mock_hf_dataset,
+    ray_start_regular_shared,
+    mock_parquet_urls,
+    mock_resolved_urls,
+    mock_ray_dataset,
+    mock_redirected_http_responses,
+    mock_huggingface_datasource,
+):
+    """Test the URL resolution logic when HTTP redirects are encountered."""
+    datasource_mock, requests_mock, read_parquet_mock, patches = setup_parquet_mocks(
+        mock_huggingface_datasource,
+        mock_parquet_urls,
+        mock_redirected_http_responses,
+        mock_ray_dataset,
+    )
+
+    try:
+        ds = ray.data.from_huggingface(mock_hf_dataset)
+
+        # Verify HTTP requests
+        verify_http_requests(requests_mock, mock_parquet_urls)
+
+        # Verify read_parquet call with resolved URLs
+        verify_read_parquet_call(read_parquet_mock, mock_resolved_urls)
+
+        # Verify dataset creation
+        verify_dataset_creation(ds, mock_hf_dataset)
+
+    finally:
+        # Stop all patches
+        for patch_obj in patches:
+            patch_obj.stop()
+
+
+def test_from_huggingface_url_resolution_failures(
+    mock_hf_dataset,
+    ray_start_regular_shared,
+    mock_parquet_urls,
+    mock_ray_dataset,
+    mock_huggingface_datasource,
+):
+    """Test URL resolution failures fall back to single node read."""
+    # Convert the mock dataset to an IterableDataset so it uses the read_datasource fallback
+    mock_iterable_dataset = mock_hf_dataset.to_iterable_dataset()
+
+    with patch.object(
+        mock_huggingface_datasource,
+        "list_parquet_urls_from_dataset",
+        return_value=mock_parquet_urls,
+    ):
+        # Mock the requests.head calls to simulate failures
+        with patch("requests.head") as mock_requests_head:
+            # Configure mock to raise an exception for all URLs
+            mock_requests_head.side_effect = requests.RequestException(
+                "Connection failed"
+            )
+
+            # Mock the fallback path
+            with patch("ray.data.read_api.read_datasource") as mock_read_datasource:
+                mock_read_datasource.return_value = mock_ray_dataset
+
+                ds = ray.data.from_huggingface(mock_iterable_dataset)
+
+                # Verify that requests.head was called for each URL
+                assert mock_requests_head.call_count == len(mock_parquet_urls)
+
+                # Verify that the fallback read_datasource was called
+                mock_read_datasource.assert_called_once()
+
+                # Verify the dataset was created successfully via fallback
+                verify_dataset_creation(ds, mock_hf_dataset)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))