add ut for generateAnalyzeIR

Author: weiguoz

pkg/sql/codegen/intermediate_representation.go

@@ -120,5 +120,5 @@ type AnalyzeIR struct {
 	// SQLFlow supports TreeExplainer so far.
 	Explainer string
 	// TrainIR is the TrainIR used for generating the training job of the corresponding model
-	TrainIR TrainIR
+	TrainIR *TrainIR
 }

pkg/sql/ir_generator.go

@@ -72,20 +72,27 @@ func generateTrainIR(slct *extendedSelect, connStr string) (*codegen.TrainIR, er
 	}, nil
 }
 
-func generatePredictIR(slct *extendedSelect, connStr string, cwd string, modelDir string) (*codegen.PredictIR, error) {
-	attrMap, err := generateAttributeIR(&slct.predAttrs)
+func generateTrainIRByModel(slct *extendedSelect, connStr, cwd, modelDir string) (*codegen.TrainIR, error) {
+	db, err := open(connStr)
 	if err != nil {
 		return nil, err
 	}
-	db, err := open(connStr)
+	defer db.Close()
+
+	slctWithTrain, _, err := loadModelMeta(slct, db, cwd, modelDir, slct.trainedModel)
 	if err != nil {
 		return nil, err
 	}
-	slctWithTrain, _, err := loadModelMeta(slct, db, cwd, modelDir, slct.model)
+	return generateTrainIR(slctWithTrain, connStr)
+}
+
+func generatePredictIR(slct *extendedSelect, connStr string, cwd string, modelDir string) (*codegen.PredictIR, error) {
+	attrMap, err := generateAttributeIR(&slct.predAttrs)
 	if err != nil {
 		return nil, err
 	}
-	trainir, err := generateTrainIR(slctWithTrain, connStr)
+
+	trainIR, err := generateTrainIRByModel(slct, connStr, cwd, modelDir)
 	if err != nil {
 		return nil, err
 	}
@@ -94,22 +101,26 @@ func generatePredictIR(slct *extendedSelect, connStr string, cwd string, modelDi
 		Select:      slct.standardSelect.String(),
 		ResultTable: slct.into,
 		Attributes:  attrMap,
-		TrainIR:     trainir,
+		TrainIR:     trainIR,
 	}, nil
 }
 
-func generateAnalyzeIR(slct *extendedSelect, connStr string) (*codegen.AnalyzeIR, error) {
+func generateAnalyzeIR(slct *extendedSelect, connStr, cwd, modelDir string) (*codegen.AnalyzeIR, error) {
 	attrs, err := generateAttributeIR(&slct.analyzeAttrs)
 	if err != nil {
 		return nil, err
 	}
+
+	trainIR, err := generateTrainIRByModel(slct, connStr, cwd, modelDir)
+	if err != nil {
+		return nil, err
+	}
 	return &codegen.AnalyzeIR{
 		DataSource: connStr,
 		Select:     slct.standardSelect.String(),
 		Attributes: attrs,
 		Explainer:  slct.explainer,
-		// TrainIR is the TrainIR used for generating the training job of the corresponding model
-		// TrainIR TrainIR
+		TrainIR:    trainIR,
 	}, nil
 }
 

pkg/sql/ir_generator_test.go

@@ -195,8 +195,35 @@ INTO sqlflow_models.mymodel;`, testDB, modelDir, nil)
 }
 
 func TestGenerateAnalyzeIR(t *testing.T) {
+	if getEnv("SQLFLOW_TEST_DB", "mysql") != "mysql" {
+		t.Skip(fmt.Sprintf("%s: skip test", getEnv("SQLFLOW_TEST_DB", "mysql")))
+	}
 	a := assert.New(t)
-	stmt := `
+
+	modelDir, e := ioutil.TempDir("/tmp", "sqlflow_models")
+	a.Nil(e)
+	defer os.RemoveAll(modelDir)
+	stream := runExtendedSQL(`
+	SELECT *
+	FROM iris.train
+	TO TRAIN xgboost.gbtree
+	WITH
+	    objective="multi:softprob",
+	    train.num_boost_round = 30,
+	    eta = 3.1,
+	    num_class = 3
+	COLUMN sepal_length, sepal_width, petal_length, petal_width
+	LABEL class
+	INTO sqlflow_models.my_xgboost_model;
+	`, testDB, modelDir, nil)
+	a.True(goodStream(stream.ReadAll()))
+
+	// Test generate PredicrIR
+	cwd, e := ioutil.TempDir("/tmp", "sqlflow")
+	a.Nil(e)
+	defer os.RemoveAll(cwd)
+
+	pr, e := newParser().Parse(`
 	SELECT *
 	FROM iris.train
 	ANALYZE sqlflow_models.my_xgboost_model
@@ -205,13 +232,20 @@ func TestGenerateAnalyzeIR(t *testing.T) {
 	    shap_summary.alpha=1,
 	    shap_summary.sort=True
 	USING TreeExplainer;
-	`
-	pr, e := newParser().Parse(stmt)
+	`)
 	a.NoError(e)
 
-	connStr := "mysql://root:root@tcp(localhost)"
-	ir, e := generateAnalyzeIR(pr, connStr)
+	connStr := "mysql://root:root@tcp(127.0.0.1:3306)/?maxAllowedPacket=0"
+	ir, e := generateAnalyzeIR(pr, connStr, cwd, modelDir)
 	a.NoError(e)
-	a.Equal(ir.Explainer, "TreeExplainer")
 	a.Equal(ir.DataSource, connStr)
+	a.Equal(ir.Explainer, "TreeExplainer")
+	a.Equal(len(ir.Attributes), 3)
+	a.Equal(ir.Attributes["shap_summary.sort"], "True")
+	a.Equal(ir.Attributes["shap_summary.plot_type"], "bar")
+	a.Equal(ir.Attributes["shap_summary.alpha"], 1)
+
+	nc, ok := ir.TrainIR.Features["feature_columns"][0].(*codegen.NumericColumn)
+	a.True(ok)
+	a.Equal("sepal_length", nc.FieldMeta.Name)
 }