Fix prediction heuristic (#5955)

* Relax check for prediction.
* Relax test in spark test.
* Add tests in C++.

tests/cpp/predictor/test_cpu_predictor.cc

@@ -165,4 +165,8 @@ TEST(CpuPredictor, InplacePredict) {
     TestInplacePrediction(x, "cpu_predictor", kRows, kCols, -1);
   }
 }
+
+TEST(CpuPredictor, LesserFeatures) {
+  TestPredictionWithLesserFeatures("cpu_predictor");
+}
 }  // namespace xgboost

tests/cpp/predictor/test_gpu_predictor.cu

@@ -160,5 +160,8 @@ TEST(GPUPredictor, MGPU_InplacePredict) {  // NOLINT
                dmlc::Error);
 }
 
+TEST(GpuPredictor, LesserFeatures) {
+  TestPredictionWithLesserFeatures("gpu_predictor");
+}
 }  // namespace predictor
 }  // namespace xgboost

tests/cpp/predictor/test_predictor.cc

@@ -134,4 +134,44 @@ void TestInplacePrediction(dmlc::any x, std::string predictor,
   learner->SetParam("gpu_id", "-1");
   learner->Configure();
 }
+
+void TestPredictionWithLesserFeatures(std::string predictor_name) {
+  size_t constexpr kRows = 256, kTrainCols = 256, kTestCols = 4, kIters = 4;
+  auto m_train = RandomDataGenerator(kRows, kTrainCols, 0.5).GenerateDMatrix(true);
+  auto m_test = RandomDataGenerator(kRows, kTestCols, 0.5).GenerateDMatrix(false);
+  std::unique_ptr<Learner> learner{Learner::Create({m_train})};
+
+  for (size_t i = 0; i < kIters; ++i) {
+    learner->UpdateOneIter(i, m_train);
+  }
+
+  HostDeviceVector<float> prediction;
+  learner->SetParam("predictor", predictor_name);
+  learner->Configure();
+  Json config{Object()};
+  learner->SaveConfig(&config);
+  ASSERT_EQ(get<String>(config["learner"]["gradient_booster"]["gbtree_train_param"]["predictor"]), predictor_name);
+
+  learner->Predict(m_test, false, &prediction);
+  ASSERT_EQ(prediction.Size(), kRows);
+
+  auto m_invalid = RandomDataGenerator(kRows, kTrainCols + 1, 0.5).GenerateDMatrix(false);
+  ASSERT_THROW({learner->Predict(m_invalid, false, &prediction);}, dmlc::Error);
+
+#if defined(XGBOOST_USE_CUDA)
+  HostDeviceVector<float> from_cpu;
+  learner->SetParam("predictor", "cpu_predictor");
+  learner->Predict(m_test, false, &from_cpu);
+
+  HostDeviceVector<float> from_cuda;
+  learner->SetParam("predictor", "gpu_predictor");
+  learner->Predict(m_test, false, &from_cuda);
+
+  auto const& h_cpu = from_cpu.ConstHostVector();
+  auto const& h_gpu = from_cuda.ConstHostVector();
+  for (size_t i = 0; i < h_cpu.size(); ++i) {
+    ASSERT_NEAR(h_cpu[i], h_gpu[i], kRtEps);
+  }
+#endif  // defined(XGBOOST_USE_CUDA)
+}
 }  // namespace xgboost

tests/cpp/predictor/test_predictor.h

@@ -59,6 +59,8 @@ void TestTrainingPrediction(size_t rows, size_t bins, std::string tree_method,
 void TestInplacePrediction(dmlc::any x, std::string predictor,
                            bst_row_t rows, bst_feature_t cols,
                            int32_t device = -1);
+
+void TestPredictionWithLesserFeatures(std::string preditor_name);
 }  // namespace xgboost
 
 #endif  // XGBOOST_TEST_PREDICTOR_H_