chore: import upstream snapshot with attribution

python/ray/data/tests/test_dataset_validation.py

@@ -0,0 +1,98 @@
+import logging
+import sys
+from unittest.mock import patch
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pytest
+
+import ray
+from ray.data import Schema
+from ray.data._internal.util import _check_pyarrow_version
+from ray.data.tests.conftest import *  # noqa
+from ray.tests.conftest import *  # noqa
+
+
+def test_column_name_type_check(ray_start_regular_shared):
+    df = pd.DataFrame({"1": np.random.rand(10), "a": np.random.rand(10)})
+    ds = ray.data.from_pandas(df)
+    assert ds.schema() == Schema(pa.schema([("1", pa.float64()), ("a", pa.float64())]))
+    assert ds.count() == 10
+
+    df = pd.DataFrame({1: np.random.rand(10), "a": np.random.rand(10)})
+    with pytest.raises(ValueError):
+        ray.data.from_pandas(df)
+
+
+@pytest.mark.skipif(
+    sys.version_info >= (3, 12), reason="TODO(scottjlee): Not working yet for py312"
+)
+def test_unsupported_pyarrow_versions_check(shutdown_only):
+    ray.shutdown()
+
+    # Test that unsupported pyarrow versions cause an error to be raised upon the
+    # initial pyarrow use.
+    # Pin numpy<2 since pyarrow 8.0.0 is incompatible with numpy>=2,
+    # which is now the default on the data CI images.
+    ray.init(runtime_env={"pip": ["numpy==1.26.4", "pyarrow==8.0.0"]})
+
+    @ray.remote
+    def should_error():
+        _check_pyarrow_version()
+
+    with pytest.raises(
+        Exception,
+        match=r".*Dataset requires pyarrow >= 17.0.0, but 8.0.0 is installed.*",
+    ):
+        ray.get(should_error.remote())
+
+
+class LoggerWarningCalled(Exception):
+    """Custom exception used in test_warning_execute_with_no_cpu() and
+    test_nowarning_execute_with_cpu(). Raised when the `logger.warning` method
+    is called, so that we can kick out of `plan.execute()` by catching this Exception
+    and check logging was done properly."""
+
+    pass
+
+
+def test_warning_execute_with_no_cpu(ray_start_cluster):
+    """Tests Dataset._execute() to ensure a warning is logged
+    when no CPU resources are available."""
+    # Create one node with no CPUs to trigger the Dataset warning
+    ray.shutdown()
+    ray.init(ray_start_cluster.address)
+    cluster = ray_start_cluster
+    cluster.add_node(num_cpus=0)
+
+    try:
+        ds = ray.data.range(10)
+        ds = ds.map_batches(lambda x: x)
+        ds.take()
+    except Exception as e:
+        assert isinstance(e, ValueError)
+        assert "exceeds the execution limits ExecutionResources(cpu=0.0" in str(e)
+
+
+def test_nowarning_execute_with_cpu(ray_start_cluster):
+    """Tests Dataset._execute() to ensure no warning is logged
+    when there are available CPU resources."""
+    # Create one node with CPUs to avoid triggering the Dataset warning
+    ray.shutdown()
+    ray.init(ray_start_cluster.address)
+
+    logger = logging.getLogger("ray.data.dataset")
+    with patch.object(
+        logger,
+        "warning",
+        side_effect=LoggerWarningCalled,
+    ) as mock_logger:
+        ds = ray.data.range(10)
+        ds = ds.map_batches(lambda x: x)
+        ds.take()
+        mock_logger.assert_not_called()
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))