[dask] Set dataframe index in predict. (#6944)

tests/python/test_with_dask.py

@@ -100,6 +100,12 @@ def test_from_dask_dataframe() -> None:
             np.testing.assert_allclose(series_predictions.compute().values,
                                        from_dmatrix)
 
+            # Make sure the output can be integrated back to original dataframe
+            X["predict"] = prediction
+            X["inplace_predict"] = series_predictions
+
+            assert bool(X.isnull().values.any().compute()) is False
+
 
 def test_from_dask_array() -> None:
     with LocalCluster(n_workers=kWorkers, threads_per_worker=5) as cluster: