Add the ray dataset example to test_client (#141)

* Add the ray dataset example to test_client

* Pin pyarrow

* Add comment

* Fix bad import for Ray dataset

* Add newline

* Skip dataset test on old ray

* Fix import

setup.py

@@ -10,5 +10,6 @@
     "distributed computing framework Ray.",
     url="https://github.com/ray-project/xgboost_ray",
     install_requires=[
-        "xgboost>=0.90", "ray", "numpy>=1.16,<1.20", "pandas", "pyarrow"
+        "xgboost>=0.90", "ray", "numpy>=1.16,<1.20", "pandas", "pyarrow<5.0.0"
     ])
+    # pyarrow<5.0.0 pinned until petastorm is updated

xgboost_ray/data_sources/ray_dataset.py

@@ -13,7 +13,7 @@
 from xgboost_ray.data_sources.object_store import ObjectStore
 
 try:
-    import ray.experimental.data  # noqa: F401
+    import ray.experimental.data.dataset  # noqa: F401
     RAY_DATASET_AVAILABLE = True
 except (ImportError, AttributeError):
     RAY_DATASET_AVAILABLE = False
@@ -41,11 +41,11 @@ def is_data_type(data: Any,
         if not RAY_DATASET_AVAILABLE:
             return False
 
-        return isinstance(data, ray.experimental.data.Dataset)
+        return isinstance(data, ray.experimental.data.dataset.Dataset)
 
     @staticmethod
     def load_data(
-            data: Any,  # ray.experimental.data.Dataset
+            data: Any,  # ray.experimental.data.dataset.Dataset
             ignore: Optional[Sequence[str]] = None,
             indices: Optional[Union[Sequence[int], Sequence[
                 ObjectRef]]] = None,
@@ -70,7 +70,7 @@ def convert_to_series(data: Any) -> pd.Series:
 
     @staticmethod
     def get_actor_shards(
-            data: Any,  # ray.experimental.data.Dataset
+            data: Any,  # ray.experimental.data.dataset.Dataset
             actors: Sequence[ActorHandle]) -> \
             Tuple[Any, Optional[Dict[int, Any]]]:
         _assert_ray_data_available()

xgboost_ray/matrix.py

@@ -29,7 +29,7 @@
     MLDataset = Unavailable
 
 try:
-    from ray.experimental.data import Dataset as RayDataset
+    from ray.experimental.data.dataset import Dataset as RayDataset
 except (ImportError, ModuleNotFoundError):
 
     class RayDataset:

xgboost_ray/tests/test_client.py

@@ -4,6 +4,7 @@
 
 import ray
 from ray.util.client.ray_client_helpers import ray_start_client_server
+from xgboost_ray.data_sources.ray_dataset import RAY_DATASET_AVAILABLE
 
 
 @pytest.fixture
@@ -46,6 +47,15 @@ def test_simple_modin(start_client_server_5_cpus):
     main(cpus_per_actor=1, num_actors=4)
 
 
+@pytest.mark.skipif(
+    not RAY_DATASET_AVAILABLE,
+    reason="Ray datasets are not available in this version of Ray")
+def test_simple_ray_dataset(start_client_server_5_cpus):
+    assert ray.util.client.ray.is_connected()
+    from xgboost_ray.examples.simple_ray_dataset import main
+    main(cpus_per_actor=1, num_actors=4)
+
+
 if __name__ == "__main__":
     import pytest  # noqa: F811
     import sys