From 6c9b6f11da17f045c02fbcc35ecf7888947bd022 Mon Sep 17 00:00:00 2001
From: Jiaming Yuan <jm.yuan@outlook.com>
Date: Wed, 9 Oct 2019 23:02:10 -0400
Subject: [PATCH] Use `cudf.concat` explicitly. (#4918)

* Use `cudf.concat` explicitly.

* Add test.
---
 python-package/xgboost/compat.py       |  2 ++
 python-package/xgboost/dask.py         |  3 ++
 tests/python-gpu/test_gpu_with_dask.py | 41 ++++++++++++++++++++++++++
 tests/python/test_with_dask.py         |  7 +++--
 4 files changed, 51 insertions(+), 2 deletions(-)
 create mode 100644 tests/python-gpu/test_gpu_with_dask.py

diff --git a/python-package/xgboost/compat.py b/python-package/xgboost/compat.py
index bd124fa67..62372a3d4 100644
--- a/python-package/xgboost/compat.py
+++ b/python-package/xgboost/compat.py
@@ -134,12 +134,14 @@ try:
     from cudf import DataFrame as CUDF_DataFrame
     from cudf import Series as CUDF_Series
     from cudf import MultiIndex as CUDF_MultiIndex
+    from cudf import concat as CUDF_concat
     CUDF_INSTALLED = True
 except ImportError:
     CUDF_DataFrame = object
     CUDF_Series = object
     CUDF_MultiIndex = object
     CUDF_INSTALLED = False
+    CUDF_concat = None
 
 # sklearn
 try:
diff --git a/python-package/xgboost/dask.py b/python-package/xgboost/dask.py
index d424f1bc7..0ea7c7315 100644
--- a/python-package/xgboost/dask.py
+++ b/python-package/xgboost/dask.py
@@ -25,6 +25,7 @@ from .compat import distributed_get_worker, distributed_wait, distributed_comm
 from .compat import da, dd, delayed, get_client
 from .compat import sparse, scipy_sparse
 from .compat import PANDAS_INSTALLED, DataFrame, Series, pandas_concat
+from .compat import CUDF_INSTALLED, CUDF_DataFrame, CUDF_Series, CUDF_concat
 
 from .core import DMatrix, Booster, _expect
 from .training import train as worker_train
@@ -84,6 +85,8 @@ def concat(value):
         return sparse.concatenate(value, axis=0)
     if PANDAS_INSTALLED and isinstance(value[0], (DataFrame, Series)):
         return pandas_concat(value, axis=0)
+    if CUDF_INSTALLED and isinstance(value[0], (CUDF_DataFrame, CUDF_Series)):
+        return CUDF_concat(value, axis=0)
     return dd.multi.concat(list(value), axis=0)
 
 
diff --git a/tests/python-gpu/test_gpu_with_dask.py b/tests/python-gpu/test_gpu_with_dask.py
new file mode 100644
index 000000000..abf1e49e4
--- /dev/null
+++ b/tests/python-gpu/test_gpu_with_dask.py
@@ -0,0 +1,41 @@
+import sys
+import pytest
+
+if sys.platform.startswith("win"):
+    pytest.skip("Skipping dask tests on Windows", allow_module_level=True)
+
+try:
+    from distributed.utils_test import client, loop, cluster_fixture
+    import dask.dataframe as dd
+    from xgboost import dask as dxgb
+    import cudf
+except ImportError:
+    client = None
+    loop = None
+    cluster_fixture = None
+    pass
+
+sys.path.append("tests/python")
+from test_with_dask import generate_array
+import testing as tm
+
+
+@pytest.mark.skipif(**tm.no_dask())
+@pytest.mark.skipif(**tm.no_cudf())
+def test_dask_dataframe(client):
+    X, y = generate_array()
+
+    X = dd.from_dask_array(X)
+    y = dd.from_dask_array(y)
+
+    X = X.map_partitions(cudf.from_pandas)
+    y = y.map_partitions(cudf.from_pandas)
+
+    dtrain = dxgb.DaskDMatrix(client, X, y)
+    out = dxgb.train(client, {'tree_method': 'gpu_hist'},
+                     dtrain=dtrain,
+                     evals=[(dtrain, 'X')],
+                     num_boost_round=2)
+
+    assert isinstance(out['booster'], dxgb.Booster)
+    assert len(out['history']['X']['rmse']) == 2
diff --git a/tests/python/test_with_dask.py b/tests/python/test_with_dask.py
index 663beefb4..e9aff6bd7 100644
--- a/tests/python/test_with_dask.py
+++ b/tests/python/test_with_dask.py
@@ -7,16 +7,19 @@ import numpy as np
 if sys.platform.startswith("win"):
     pytest.skip("Skipping dask tests on Windows", allow_module_level=True)
 
+pytestmark = pytest.mark.skipif(**tm.no_dask())
+
 try:
     from distributed.utils_test import client, loop, cluster_fixture
     import dask.dataframe as dd
     import dask.array as da
     from xgboost.dask import DaskDMatrix
 except ImportError:
+    client = None
+    loop = None
+    cluster_fixture = None
     pass
 
-pytestmark = pytest.mark.skipif(**tm.no_dask())
-
 kRows = 1000