chore: import upstream snapshot with attribution

python/ray/data/tests/test_execution_optimizer_integrations.py

@@ -0,0 +1,225 @@
+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(<lambda>)"
+
+        # 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(<lambda>)"
+    _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__]))