import sys import numpy as np import pandas as pd import pyarrow as pa import pytest from packaging.version import parse as parse_version import ray from ray.data._internal.util import rows_same from ray.data._internal.utils.arrow_utils import get_pyarrow_version from ray.data.tests.conftest import * # noqa from ray.data.tests.test_util import _check_usage_record from ray.data.tests.util import extract_values from ray.tests.conftest import * # noqa def _should_skip_huggingface_test(): """Check if we should skip the HuggingFace test due to version incompatibility.""" pyarrow_version = get_pyarrow_version() if pyarrow_version is None: return False try: datasets_version = __import__("datasets").__version__ if datasets_version is None: return False return pyarrow_version < parse_version("12.0.0") and parse_version( datasets_version ) >= parse_version("3.0.0") except (ImportError, AttributeError): return False def test_from_modin_e2e(ray_start_regular_shared_2_cpus): import modin.pandas as mopd df = pd.DataFrame( {"one": list(range(100)), "two": list(range(100))}, ) modf = mopd.DataFrame(df) ds = ray.data.from_modin(modf) # `ds.take_all()` triggers execution with new backend, which is # needed for checking operator usage below. assert len(ds.take_all()) == len(df) # `ds.to_pandas()` does not use the new backend. dfds = ds.to_pandas() assert df.equals(dfds) # Check that metadata fetch is included in stats. This is `FromPandas` # instead of `FromModin` because `from_modin` reduces to `from_pandas_refs`. assert "FromPandas" in ds.stats() assert ds._logical_plan.dag.name == "FromPandas" _check_usage_record(["FromPandas"]) @pytest.mark.parametrize("enable_pandas_block", [False, True]) def test_from_pandas_refs_e2e(ray_start_regular_shared_2_cpus, enable_pandas_block): ctx = ray.data.context.DataContext.get_current() old_enable_pandas_block = ctx.enable_pandas_block ctx.enable_pandas_block = enable_pandas_block try: df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]}) df2 = pd.DataFrame({"one": [4, 5, 6], "two": ["e", "f", "g"]}) expected_df = pd.concat([df1, df2]) ds = ray.data.from_pandas_refs([ray.put(df1), ray.put(df2)]) assert rows_same(ds.to_pandas(), expected_df) # Check that metadata fetch is included in stats. assert "FromPandas" in ds.stats() assert ds._logical_plan.dag.name == "FromPandas" # Test chaining multiple operations ds2 = ds.map_batches(lambda x: x) assert rows_same(ds2.to_pandas(), expected_df) assert "MapBatches" in ds2.stats() assert "FromPandas" in ds2.stats() assert ds2._logical_plan.dag.name == "MapBatches()" # test from single pandas dataframe ds = ray.data.from_pandas_refs(ray.put(df1)) assert rows_same(ds.to_pandas(), df1) # Check that metadata fetch is included in stats. assert "FromPandas" in ds.stats() assert ds._logical_plan.dag.name == "FromPandas" _check_usage_record(["FromPandas"]) finally: ctx.enable_pandas_block = old_enable_pandas_block def test_from_numpy_refs_e2e(ray_start_regular_shared_2_cpus): arr1 = np.expand_dims(np.arange(0, 4), axis=1) arr2 = np.expand_dims(np.arange(4, 8), axis=1) ds = ray.data.from_numpy_refs([ray.put(arr1), ray.put(arr2)]) values = np.stack(extract_values("data", ds.take(8))) np.testing.assert_array_equal(values, np.concatenate((arr1, arr2))) # Check that conversion task is included in stats. assert "FromNumpy" in ds.stats() assert ds._logical_plan.dag.name == "FromNumpy" _check_usage_record(["FromNumpy"]) # Test chaining multiple operations ds2 = ds.map_batches(lambda x: x) values = np.stack(extract_values("data", ds2.take(8))) np.testing.assert_array_equal(values, np.concatenate((arr1, arr2))) assert "MapBatches" in ds2.stats() assert "FromNumpy" in ds2.stats() assert ds2._logical_plan.dag.name == "MapBatches()" _check_usage_record(["FromNumpy", "MapBatches"]) # Test from single NumPy ndarray. ds = ray.data.from_numpy_refs(ray.put(arr1)) values = np.stack(extract_values("data", ds.take(4))) np.testing.assert_array_equal(values, arr1) # Check that conversion task is included in stats. assert "FromNumpy" in ds.stats() assert ds._logical_plan.dag.name == "FromNumpy" _check_usage_record(["FromNumpy"]) def test_from_arrow_refs_e2e(ray_start_regular_shared_2_cpus): df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]}) df2 = pd.DataFrame({"one": [4, 5, 6], "two": ["e", "f", "g"]}) ds = ray.data.from_arrow_refs( [ray.put(pa.Table.from_pandas(df1)), ray.put(pa.Table.from_pandas(df2))] ) values = [(r["one"], r["two"]) for r in ds.take(6)] rows = [(r.one, r.two) for _, r in pd.concat([df1, df2]).iterrows()] assert values == rows # Check that metadata fetch is included in stats. assert "FromArrow" in ds.stats() assert ds._logical_plan.dag.name == "FromArrow" _check_usage_record(["FromArrow"]) # test from single pyarrow table ref ds = ray.data.from_arrow_refs(ray.put(pa.Table.from_pandas(df1))) values = [(r["one"], r["two"]) for r in ds.take(3)] rows = [(r.one, r.two) for _, r in df1.iterrows()] assert values == rows # Check that conversion task is included in stats. assert "FromArrow" in ds.stats() assert ds._logical_plan.dag.name == "FromArrow" _check_usage_record(["FromArrow"]) @pytest.mark.skipif( _should_skip_huggingface_test, reason="Skip due to HuggingFace datasets >= 3.0.0 requiring pyarrow >= 12.0.0", ) def test_from_huggingface_e2e(ray_start_regular_shared_2_cpus): import datasets from ray.data.tests.datasource.test_huggingface import hfds_assert_equals data = datasets.load_dataset("tweet_eval", "emotion") assert isinstance(data, datasets.DatasetDict) ray_datasets = { "train": ray.data.from_huggingface(data["train"]), "validation": ray.data.from_huggingface(data["validation"]), "test": ray.data.from_huggingface(data["test"]), } for ds_key, ds in ray_datasets.items(): assert isinstance(ds, ray.data.Dataset) # `ds.take_all()` triggers execution with new backend, which is # needed for checking operator usage below. assert len(ds.take_all()) > 0 # Check that metadata fetch is included in stats; # the underlying implementation uses the `ReadParquet` operator # as this is an un-transformed public dataset. assert "ReadParquet" in ds.stats() or "FromArrow" in ds.stats() assert ( ds._logical_plan.dag.name == "ReadParquet" or ds._logical_plan.dag.name == "FromArrow" ) # use sort by 'text' to match order of rows hfds_assert_equals(data[ds_key], ds) try: _check_usage_record(["ReadParquet"]) except AssertionError: _check_usage_record(["FromArrow"]) # test transformed public dataset for fallback behavior base_hf_dataset = data["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 isinstance(ray_dataset_split_train, ray.data.Dataset) # `ds.take_all()` triggers execution with new backend, which is # needed for checking operator usage below. assert len(ray_dataset_split_train.take_all()) > 0 # Check that metadata fetch is included in stats; # the underlying implementation uses the `FromArrow` operator. assert "FromArrow" in ray_dataset_split_train.stats() assert ray_dataset_split_train._logical_plan.dag.name == "FromArrow" assert ray_dataset_split_train.count() == hf_dataset_split["train"].num_rows _check_usage_record(["FromArrow"]) def test_from_torch_e2e(ray_start_regular_shared_2_cpus, tmp_path): import torchvision torch_dataset = torchvision.datasets.FashionMNIST(tmp_path, download=True) ray_dataset = ray.data.from_torch(torch_dataset) expected_data = list(torch_dataset) actual_data = list(ray_dataset.take_all()) assert extract_values("item", actual_data) == expected_data # Check that metadata fetch is included in stats. assert "ReadTorch" in ray_dataset.stats() # Underlying implementation uses `FromItems` operator assert ray_dataset._logical_plan.dag.name == "ReadTorch" _check_usage_record(["ReadTorch"]) if __name__ == "__main__": sys.exit(pytest.main(["-v", __file__]))