[dask] Accept other inputs for prediction. (#5428)

* Returns a series when input is dataframe.

* Merge assert client.

python-package/xgboost/dask.py

@@ -103,6 +103,9 @@ def concat(value):
 
 def _xgb_get_client(client):
  '''Simple wrapper around testing None.'''
+ if not isinstance(client, (type(get_client()), type(None))):
+ raise TypeError(
+ _expect([type(get_client()), type(None)], type(client)))
  ret = get_client() if client is None else client
  return ret
 
@@ -112,12 +115,6 @@ def _get_client_workers(client):
  return workers
 
 
-def _assert_client(client):
- if not isinstance(client, (type(get_client()), type(None))):
- raise TypeError(
- _expect([type(get_client()), type(None)], type(client)))
-
-
 class DaskDMatrix:
  # pylint: disable=missing-docstring, too-many-instance-attributes
  '''DMatrix holding on references to Dask DataFrame or Dask Array. Constructing
@@ -155,7 +152,7 @@ def __init__(self,
  feature_names=None,
  feature_types=None):
  _assert_dask_support()
- _assert_client(client)
+ client = _xgb_get_client(client)
 
  self.feature_names = feature_names
  self.feature_types = feature_types
@@ -177,7 +174,6 @@ def __init__(self,
  self.has_label = label is not None
  self.has_weights = weight is not None
 
- client = _xgb_get_client(client)
  client.sync(self.map_local_data, client, data, label, weight)
 
  async def map_local_data(self, client, data, label=None, weights=None):
@@ -391,13 +387,12 @@ def train(client, params, dtrain, *args, evals=(), **kwargs):
 
  '''
  _assert_dask_support()
- _assert_client(client)
+ client = _xgb_get_client(client)
  if 'evals_result' in kwargs.keys():
  raise ValueError(
  'evals_result is not supported in dask interface.',
  'The evaluation history is returned as result of training.')
 
- client = _xgb_get_client(client)
  workers = list(_get_client_workers(client).keys())
 
  rabit_args = _get_rabit_args(workers, client)
@@ -452,7 +447,7 @@ def dispatched_train(worker_addr):
  return list(filter(lambda ret: ret is not None, results))[0]
 
 
-def predict(client, model, data, *args):
+def predict(client, model, data, *args, missing=numpy.nan):
  '''Run prediction with a trained booster.
 
  .. note::
@@ -466,32 +461,55 @@ def predict(client, model, data, *args):
  returned from dask if it's set to None.
  model: A Booster or a dictionary returned by `xgboost.dask.train`.
  The trained model.
- data: DaskDMatrix
+ data: DaskDMatrix/dask.dataframe.DataFrame/dask.array.Array
  Input data used for prediction.
+ missing: float
+ Used when input data is not DaskDMatrix. Specify the value
+ considered as missing.
 
  Returns
  -------
- prediction: dask.array.Array
+ prediction: dask.array.Array/dask.dataframe.Series
 
  '''
  _assert_dask_support()
- _assert_client(client)
+ client = _xgb_get_client(client)
  if isinstance(model, Booster):
  booster = model
  elif isinstance(model, dict):
  booster = model['booster']
  else:
  raise TypeError(_expect([Booster, dict], type(model)))
+ if not isinstance(data, (DaskDMatrix, da.Array, dd.DataFrame)):
+ raise TypeError(_expect([DaskDMatrix, da.Array, dd.DataFrame],
+ type(data)))
 
- if not isinstance(data, DaskDMatrix):
- raise TypeError(_expect([DaskDMatrix], type(data)))
-
+ def mapped_predict(partition, is_df):
+ worker = distributed_get_worker()
+ m = DMatrix(partition, missing=missing, nthread=worker.nthreads)
+ predt = booster.predict(m, *args, validate_features=False)
+ if is_df:
+ predt = DataFrame(predt, columns=['prediction'])
+ return predt
+
+ if isinstance(data, da.Array):
+ predictions = client.submit(
+ da.map_blocks,
+ mapped_predict, data, False, drop_axis=1,
+ dtype=numpy.float32
+ ).result()
+ return predictions
+ if isinstance(data, dd.DataFrame):
+ import dask
+ predictions = client.submit(
+ dd.map_partitions,
+ mapped_predict, data, True,
+ meta=dask.dataframe.utils.make_meta({'prediction': 'f4'})
+ ).result()
+ return predictions.iloc[:, 0]
+
+ # Prediction on dask DMatrix.
  worker_map = data.worker_map
- client = _xgb_get_client(client)
-
- missing = data.missing
- feature_names = data.feature_names
- feature_types = data.feature_types
 
  def dispatched_predict(worker_id):
  '''Perform prediction on each worker.'''
@@ -502,9 +520,9 @@ def dispatched_predict(worker_id):
  booster.set_param({'nthread': worker.nthreads})
  for part, order in list_of_parts:
  local_x = DMatrix(part,
- feature_names=feature_names,
- feature_types=feature_types,
- missing=missing,
+ feature_names=data.feature_names,
+ feature_types=data.feature_types,
+ missing=data.missing,
  nthread=worker.nthreads)
  predt = booster.predict(data=local_x,
  validate_features=local_x.num_row() != 0,
@@ -520,8 +538,7 @@ def dispatched_get_shape(worker_id):
  list_of_parts = data.get_worker_x_ordered(worker)
  shapes = []
  for part, order in list_of_parts:
- s = part.shape
- shapes.append((s, order))
+ shapes.append((part.shape, order))
  return shapes
 
  def map_function(func):

tests/python/test_with_dask.py

@@ -57,7 +57,13 @@ def test_from_dask_dataframe():
  xgb.dask.train(
  client, {}, dtrain, num_boost_round=2, evals_result={})
  # force prediction to be computed
- prediction = prediction.compute()
+ from_dmatrix = prediction.compute()
+
+ prediction = xgb.dask.predict(client, model=booster, data=X)
+ from_df = prediction.compute()
+
+ assert isinstance(prediction, dd.Series)
+ assert np.all(from_dmatrix == from_df.to_numpy())
 
 
 def test_from_dask_array():
@@ -84,6 +90,12 @@ def test_from_dask_array():
  config = json.loads(booster.save_config())
  assert int(config['learner']['generic_param']['nthread']) == 5
 
+ from_arr = xgb.dask.predict(
+ client, model=booster, data=X)
+
+ assert isinstance(from_arr, da.Array)
+ assert np.all(single_node_predt == from_arr.compute())
+
 
 def test_dask_regressor():
  with LocalCluster(n_workers=5) as cluster: