Refine feature column code generation (#2860)

* use codegen way to get py fc

* fix ut

Author: sneaxiy

go/cmd/sqlflowserver/e2e_workflow_test.go

@@ -375,6 +375,13 @@ WITH objective="multi:softmax",num_class=3
 LABEL class
 INTO sqlflow_models.xgb_classification;
 
+SELECT * FROM iris.train
+TO TRAIN xgboost.gbtree
+WITH objective="multi:softmax",num_class=3
+COLUMN sepal_length, DENSE(sepal_width)
+LABEL class
+INTO sqlflow_models.xgb_classification;
+
 SELECT * FROM iris.test
 TO PREDICT iris.test_result_table.class
 USING sqlflow_models.xgb_classification;

go/codegen/experimental/codegen_test.go

@@ -43,7 +43,7 @@ func TestExperimentalXGBCodegen(t *testing.T) {
 	if err != nil {
 		t.Errorf("error %s", err)
 	}
-	expected := `feature_column_map = {"feature_columns": [fc.NumericColumn(fd.FieldDesc(name="petal_length", dtype=fd.DataType.FLOAT32, delimiter="", format="", shape=[1], is_sparse=False, vocabulary=[]))]}`
+	expected := `feature_column_map = {"feature_columns":[runtime.feature.column.NumericColumn(runtime.feature.field_desc.FieldDesc(name="petal_length", dtype=runtime.feature.field_desc.DataType.FLOAT32, delimiter="", format="", shape=[1], is_sparse=False, vocabulary=[]))]}`
 	a.True(strings.Contains(coulerCode, expected))
 }
 

go/codegen/experimental/xgboost.go

@@ -18,7 +18,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"os"
-	"reflect"
 	"strings"
 	"text/template"
 
@@ -82,14 +81,8 @@ func XGBoostGenerateTrain(trainStmt *ir.TrainStmt, stepIndex int, session *pb.Se
 		return "", fmt.Errorf("xgboost only support 0 or 1 feature column set, received %d", len(trainStmt.Features))
 	}
 	// featureColumnCode is a python map definition code like fc_map = {"feature_columns": [...]}
-	featureColumnCode := ""
-	if len(trainStmt.Features) == 1 {
-		featureColumnCode, err = generateFeatureColumnCode(trainStmt.Features["feature_columns"])
-		if err != nil {
-			return "", err
-		}
-	}
-	labelColumnCode, err := generateFeatureColumnCode([]ir.FeatureColumn{trainStmt.Label})
+	featureColumnCode := generateFeatureColumnCode(trainStmt.Features)
+	labelColumnCode := trainStmt.Label.GenPythonCode()
 
 	mp, err := json.Marshal(params[""])
 	if err != nil {
@@ -137,15 +130,11 @@ const xgbTrainTemplate = `
 def step_entry_{{.StepIndex}}():
     import json
     import runtime.temp_file as temp_file
-    import runtime.feature.column as fc
-    import runtime.feature.field_desc as fd
+    import runtime.feature.column
+    import runtime.feature.field_desc
     from runtime.{{.Submitter}} import train
 
-    {{ if .FeatureColumnCode }}
-    feature_column_map = {"feature_columns": [{{.FeatureColumnCode}}]}
-    {{ else }}
-    feature_column_map = None
-    {{ end }}
+    feature_column_map = {{.FeatureColumnCode}}
     label_column = {{.LabelColumnCode}}
 
     model_params = json.loads('''{{.ModelParamsJSON}}''')
@@ -231,37 +220,20 @@ func getSubmitter(session *pb.Session, defaultValue string) string {
 	return defaultValue
 }
 
-func generateFeatureColumnCode(fcList []ir.FeatureColumn) (string, error) {
-	fcCodes := make([]string, 0, len(fcList))
-	for _, fc := range fcList {
-		// xgboost have no cross feature column, just get the first field desc
-		fd := fc.GetFieldDesc()[0]
-		// pass format = "" to let runtime feature derivation to fill it in.
-		tmpl := `fc.%s(fd.FieldDesc(name="%s", dtype=fd.DataType.%s, delimiter="%s", format="", shape=%s, is_sparse=%s, vocabulary=%s))`
-		fcTypeName := reflect.TypeOf(fc).Elem().Name()
-		isSparseStr := "False"
-		if fd.IsSparse {
-			isSparseStr = "True"
-		}
-		vocabList := []string{}
-		for k := range fd.Vocabulary {
-			vocabList = append(vocabList, k)
+func generateFeatureColumnCode(fcMap map[string][]ir.FeatureColumn) string {
+	allFCCodes := make([]string, 0)
+	for target, fcList := range fcMap {
+		if len(fcList) == 0 {
+			continue
 		}
-		shape := []int{1}
-		if len(fd.Shape) != 0 {
-			shape = fd.Shape
+		codeList := make([]string, 0)
+		for _, fc := range fcList {
+			codeList = append(codeList, fc.GenPythonCode())
 		}
-
-		code := fmt.Sprintf(tmpl, fcTypeName, fd.Name,
-			strings.ToUpper(ir.DTypeToString(fd.DType)),
-			fd.Delimiter,
-			ir.AttrToPythonValue(shape),
-			isSparseStr,
-			ir.AttrToPythonValue(vocabList))
-		fcCodes = append(fcCodes, code)
+		code := fmt.Sprintf(`"%s":[%s]`, target, strings.Join(codeList, ","))
+		allFCCodes = append(allFCCodes, code)
 	}
-
-	return strings.Join(fcCodes, ",\n"), nil
+	return fmt.Sprintf("{%s}", strings.Join(allFCCodes, ","))
 }
 
 // TODO(typhoonzero): below functions are copied from codegen/xgboost/codegen.go

go/ir/codegen_python_values.go

@@ -35,33 +35,40 @@ func DTypeToString(dt int) string {
 
 // AttrToPythonValue format the WITH attributes to corresponding Python code.
 func AttrToPythonValue(attr interface{}) string {
-	switch attr.(type) {
+	switch a := attr.(type) {
 	case bool:
-		return strings.Title(fmt.Sprintf("%v", attr.(bool)))
+		return strings.Title(fmt.Sprintf("%v", a))
 	case int:
-		return fmt.Sprintf("%d", attr.(int))
+		return fmt.Sprintf("%d", a)
 	case int64:
-		return fmt.Sprintf("%d", attr.(int64))
+		return fmt.Sprintf("%d", a)
 	case float32:
-		return fmt.Sprintf("%f", attr.(float32))
+		return fmt.Sprintf("%f", a)
 	case float64: // FIXME(typhoonzero): may never use
 		return fmt.Sprintf("%f", attr.(float64))
 	case []int:
-		intArrayAttrStr, _ := MarshalToJSONString(attr.([]int))
+		if a == nil {
+			return "None"
+		}
+		intArrayAttrStr, _ := MarshalToJSONString(a)
 		return intArrayAttrStr
 	case []string:
-		l := attr.([]string)
-		if len(l) == 0 {
+		if a == nil {
+			return "None"
+		}
+		if len(a) == 0 {
 			return "[]"
 		}
-		stringListStr, _ := MarshalToJSONString(l)
+		stringListStr, _ := MarshalToJSONString(a)
 		return stringListStr
 	case []interface{}:
-		tmplist := attr.([]interface{})
-		if len(tmplist) > 0 {
-			if _, ok := tmplist[0].(int); ok {
+		if a == nil {
+			return "None"
+		}
+		if len(a) > 0 {
+			if _, ok := a[0].(int); ok {
 				intlist := []int{}
-				for _, v := range tmplist {
+				for _, v := range a {
 					intlist = append(intlist, v.(int))
 				}
 				intlistStr, _ := MarshalToJSONString(intlist)
@@ -71,7 +78,7 @@ func AttrToPythonValue(attr interface{}) string {
 		// TODO(typhoonzero): support []float etc.
 		return "[]"
 	case string:
-		return fmt.Sprintf("\"%s\"", attr.(string))
+		return fmt.Sprintf(`"%s"`, a)
 	default:
 		return ""
 	}

go/ir/feature_column.go

@@ -58,7 +58,7 @@ func (fd *FieldDesc) GenPythonCode() string {
 		vocabList = append(vocabList, k)
 	}
 	// pass format = "" to let runtime feature derivation to fill it in.
-	return fmt.Sprintf(`runtime.feature.field_desc.FieldDesc(name="%s", dtype=fd.DataType.%s, delimiter="%s", format="", shape=%s, is_sparse=%s, vocabulary=%s)`,
+	return fmt.Sprintf(`runtime.feature.field_desc.FieldDesc(name="%s", dtype=runtime.feature.field_desc.DataType.%s, delimiter="%s", format="", shape=%s, is_sparse=%s, vocabulary=%s)`,
 		fd.Name,
 		strings.ToUpper(DTypeToString(fd.DType)),
 		fd.Delimiter,

go/ir/feature_column_test.go

@@ -32,7 +32,7 @@ func TestFeatureColumnGenPythonCode(t *testing.T) {
 			DType: 0,
 		},
 	}
-	a.Equal("runtime.feature.column.NumericColumn(runtime.feature.field_desc.FieldDesc(name=\"testcol\", dtype=fd.DataType.INT64, delimiter=\"\", format=\"\", shape=[10], is_sparse=False, vocabulary=[]))",
+	a.Equal("runtime.feature.column.NumericColumn(runtime.feature.field_desc.FieldDesc(name=\"testcol\", dtype=runtime.feature.field_desc.DataType.INT64, delimiter=\"\", format=\"\", shape=[10], is_sparse=False, vocabulary=[]))",
 		nc.GenPythonCode())
 
 	emd := &EmbeddingColumn{