support xgboost regression model (#797)

Author: Yancey0623

cmd/sqlflowserver/main_test.go

@@ -256,6 +256,7 @@ func TestEnd2EndMySQL(t *testing.T) {
 	t.Run("CaseSparseFeature", CaseSparseFeature)
 	t.Run("CaseSQLByPassLeftJoin", CaseSQLByPassLeftJoin)
 	t.Run("CaseTrainRegression", CaseTrainRegression)
+	t.Run("CaseTrainXGBoostRegression", CaseTrainXGBoostRegression)
 	t.Run("CaseTrainDeepWideModel", CaseTrainDeepWideModel)
 
 }
@@ -1000,3 +1001,68 @@ FROM housing.predict LIMIT 5;`)
 		a.False(nilCount == 13)
 	}
 }
+
+// CaseTrainXGBoostRegression is used to test xgboost regression models
+func CaseTrainXGBoostRegression(t *testing.T) {
+	a := assert.New(t)
+	trainSQL := fmt.Sprintf(`
+SELECT *
+FROM housing.train
+TRAIN xgboost.gbtree
+WITH
+		objective="reg:squarederror",
+		train.num_boost_round = 30
+		COLUMN f1,f2,f3,f4,f5,f6,f7,f8,f9,f10,f11,f12,f13
+LABEL target
+INTO sqlflow_models.my_xgb_regression_model;
+`)
+
+	conn, err := createRPCConn()
+	a.NoError(err)
+	defer conn.Close()
+	cli := pb.NewSQLFlowClient(conn)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 300*time.Second)
+	defer cancel()
+
+	stream, err := cli.Run(ctx, sqlRequest(trainSQL))
+	if err != nil {
+		a.Fail("Check if the server started successfully. %v", err)
+	}
+	// call ParseRow only to wait train finish
+	ParseRow(stream)
+
+	predSQL := fmt.Sprintf(`SELECT *
+FROM housing.test
+PREDICT housing.xgb_predict.target
+USING sqlflow_models.my_xgb_regression_model;`)
+
+	stream, err = cli.Run(ctx, sqlRequest(predSQL))
+	if err != nil {
+		a.Fail("Check if the server started successfully. %v", err)
+	}
+	// call ParseRow only to wait predict finish
+	ParseRow(stream)
+
+	showPred := fmt.Sprintf(`SELECT *
+FROM housing.xgb_predict LIMIT 5;`)
+
+	stream, err = cli.Run(ctx, sqlRequest(showPred))
+	if err != nil {
+		a.Fail("Check if the server started successfully. %v", err)
+	}
+	_, rows := ParseRow(stream)
+
+	for _, row := range rows {
+		// NOTE: predict result maybe random, only check predicted
+		// class >=0, need to change to more flexible checks than
+		// checking expectedPredClasses := []int64{2, 1, 0, 2, 0}
+		AssertGreaterEqualAny(a, row[13], float64(0))
+
+		// avoiding nil features in predict result
+		nilCount := 0
+		for ; nilCount < 13 && row[nilCount] == nil; nilCount++ {
+		}
+		a.False(nilCount == 13)
+	}
+}

sql/executor_test.go

@@ -84,6 +84,19 @@ func TestExecuteXGBoost(t *testing.T) {
 	})
 }
 
+func TestExecuteXGBoostRegression(t *testing.T) {
+	a := assert.New(t)
+	modelDir := ""
+	a.NotPanics(func() {
+		stream := runExtendedSQL(testXGBoostTrainSelectIris, testDB, modelDir, nil)
+		a.True(goodStream(stream.ReadAll()))
+		stream = runExtendedSQL(testAnalyzeTreeModelSelectIris, testDB, modelDir, nil)
+		a.True(goodStream(stream.ReadAll()))
+		stream = runExtendedSQL(testXGBoostPredictIris, testDB, modelDir, nil)
+		a.True(goodStream(stream.ReadAll()))
+	})
+}
+
 func TestExecutorTrainAndPredictDNN(t *testing.T) {
 	a := assert.New(t)
 	modelDir := ""

sql/template_xgboost.go

@@ -134,8 +134,11 @@ dpred = xgb_dataset('predict.txt', """{{.PredictionDatasetSQL}}""")
 bst = xgb.Booster({'nthread': 4})  # init model
 bst.load_model("{{.Save}}")  # load data
 preds = bst.predict(dpred)
-# TODO(typhoonzero): regression models may have different behavior
-pred_classes = np.argmax(np.array(preds), axis=1)
+
+# TODO(Yancey1989): using the train parameters to decide regressoin model or classifier model
+if len(preds.shape) == 2:
+    # classifier result
+    preds = np.argmax(np.array(preds), axis=1)
 
 feature_file_read = open("predict.txt", "r")
 
@@ -149,7 +152,8 @@ with buffered_db_writer(driver, conn, "{{.TableName}}", result_column_names, 100
         if not line:
             break
         row = [i.split(":")[1] for i in line.replace("\n", "").split("\t")[1:]]
-        row.append(pred_classes[line_no])
+        row.append(preds[line_no])
         w.write(row)
         line_no += 1
+print("Done predicting. Predict table : {{.TableName}}")
 `