simplify expression resolver

sql/attribute.go

@@ -0,0 +1,82 @@
+// Copyright 2019 The SQLFlow Authors. All rights reserved.
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package sql
+
+import (
+	"fmt"
+	"strconv"
+	"strings"
+)
+
+type attribute struct {
+	FullName string
+	Prefix   string
+	Name     string
+	Value    interface{}
+}
+
+func (a *attribute) GenerateCode() (string, error) {
+	if val, ok := a.Value.(string); ok {
+		// auto convert to int first.
+		if _, err := strconv.Atoi(val); err == nil {
+			return fmt.Sprintf("%s=%s", a.Name, val), nil
+		}
+		return fmt.Sprintf("%s=\"%s\"", a.Name, val), nil
+	}
+	if val, ok := a.Value.([]interface{}); ok {
+		intList, err := transformToIntList(val)
+		if err != nil {
+			return "", err
+		}
+		return fmt.Sprintf("%s=%s", a.Name,
+			strings.Join(strings.Split(fmt.Sprint(intList), " "), ",")), nil
+	}
+	return "", fmt.Errorf("value of attribute must be string or list of int, given %s", a.Value)
+}
+
+func filter(attrs map[string]*attribute, prefix string, remove bool) map[string]*attribute {
+	ret := make(map[string]*attribute, 0)
+	for _, a := range attrs {
+		if strings.EqualFold(a.Prefix, prefix) {
+			ret[a.Name] = a
+			if remove {
+				delete(attrs, a.FullName)
+			}
+		}
+	}
+	return ret
+}
+
+func resolveAttribute(attrs *attrs) (map[string]*attribute, error) {
+	ret := make(map[string]*attribute)
+	for k, v := range *attrs {
+		subs := strings.SplitN(k, ".", 2)
+		name := subs[len(subs)-1]
+		prefix := ""
+		if len(subs) == 2 {
+			prefix = subs[0]
+		}
+		r, _, err := resolveExpression(v)
+		if err != nil {
+			return nil, err
+		}
+		a := &attribute{
+			FullName: k,
+			Prefix:   prefix,
+			Name:     name,
+			Value:    r}
+		ret[a.FullName] = a
+	}
+	return ret, nil
+}

sql/attribute_test.go

@@ -0,0 +1,45 @@
+// Copyright 2019 The SQLFlow Authors. All rights reserved.
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package sql
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestAttrs(t *testing.T) {
+	a := assert.New(t)
+	parser := newParser()
+
+	s := statementWithAttrs("estimator.hidden_units = [10, 20]")
+	r, e := parser.Parse(s)
+	a.NoError(e)
+	attrs, err := resolveAttribute(&r.trainAttrs)
+	a.NoError(err)
+	attr := attrs["estimator.hidden_units"]
+	a.Equal("estimator", attr.Prefix)
+	a.Equal("hidden_units", attr.Name)
+	a.Equal([]interface{}([]interface{}{10, 20}), attr.Value)
+
+	s = statementWithAttrs("dataset.name = hello")
+	r, e = parser.Parse(s)
+	a.NoError(e)
+	attrs, err = resolveAttribute(&r.trainAttrs)
+	a.NoError(err)
+	attr = attrs["dataset.name"]
+	a.Equal("dataset", attr.Prefix)
+	a.Equal("name", attr.Name)
+	a.Equal("hello", attr.Value)
+}

sql/bucket_column.go

@@ -0,0 +1,84 @@
+// Copyright 2019 The SQLFlow Authors. All rights reserved.
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package sql
+
+import (
+	"fmt"
+	"strings"
+)
+
+type bucketColumn struct {
+	SourceColumn *numericColumn
+	Boundaries   []int
+}
+
+func (bc *bucketColumn) GenerateCode() (string, error) {
+	sourceCode, _ := bc.SourceColumn.GenerateCode()
+	return fmt.Sprintf(
+		"tf.feature_column.bucketized_column(%s, boundaries=%s)",
+		sourceCode,
+		strings.Join(strings.Split(fmt.Sprint(bc.Boundaries), " "), ",")), nil
+}
+
+func (bc *bucketColumn) GetDelimiter() string {
+	return ""
+}
+
+func (bc *bucketColumn) GetDtype() string {
+	return ""
+}
+
+func (bc *bucketColumn) GetKey() string {
+	return bc.SourceColumn.Key
+}
+
+func (bc *bucketColumn) GetInputShape() string {
+	return bc.SourceColumn.GetInputShape()
+}
+
+func (bc *bucketColumn) GetColumnType() int {
+	return columnTypeBucket
+}
+
+func resolveBucketColumn(el *exprlist) (*bucketColumn, error) {
+	if len(*el) != 3 {
+		return nil, fmt.Errorf("bad BUCKET expression format: %s", *el)
+	}
+	sourceExprList := (*el)[1]
+	boundariesExprList := (*el)[2]
+	if sourceExprList.typ != 0 {
+		return nil, fmt.Errorf("key of BUCKET must be NUMERIC, which is %v", sourceExprList)
+	}
+	source, _, err := resolveColumn(&sourceExprList.sexp)
+	if err != nil {
+		return nil, err
+	}
+	if source.GetColumnType() != columnTypeNumeric {
+		return nil, fmt.Errorf("key of BUCKET must be NUMERIC, which is %s", source)
+	}
+	boundaries, _, err := resolveExpression(boundariesExprList)
+	if err != nil {
+		return nil, err
+	}
+	if _, ok := boundaries.([]interface{}); !ok {
+		return nil, fmt.Errorf("bad BUCKET boundaries: %s", err)
+	}
+	b, err := transformToIntList(boundaries.([]interface{}))
+	if err != nil {
+		return nil, fmt.Errorf("bad BUCKET boundaries: %s", err)
+	}
+	return &bucketColumn{
+		SourceColumn: source.(*numericColumn),
+		Boundaries:   b}, nil
+}

sql/bucket_column_test.go

@@ -0,0 +1,55 @@
+// Copyright 2019 The SQLFlow Authors. All rights reserved.
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package sql
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestBucketColumn(t *testing.T) {
+	a := assert.New(t)
+	parser := newParser()
+
+	normal := statementWithColumn("BUCKET(NUMERIC(c1, 10), [1, 10])")
+	badInput := statementWithColumn("BUCKET(c1, [1, 10])")
+	badBoundaries := statementWithColumn("BUCKET(NUMERIC(c1, 10), 100)")
+
+	r, e := parser.Parse(normal)
+	a.NoError(e)
+	c := r.columns["feature_columns"]
+	fcs, _, e := resolveTrainColumns(&c)
+	a.NoError(e)
+	bc, ok := fcs[0].(*bucketColumn)
+	a.True(ok)
+	code, e := bc.GenerateCode()
+	a.NoError(e)
+	a.Equal("c1", bc.SourceColumn.Key)
+	a.Equal([]int{10}, bc.SourceColumn.Shape)
+	a.Equal([]int{1, 10}, bc.Boundaries)
+	a.Equal("tf.feature_column.bucketized_column(tf.feature_column.numeric_column(\"c1\", shape=[10]), boundaries=[1,10])", code)
+
+	r, e = parser.Parse(badInput)
+	a.NoError(e)
+	c = r.columns["feature_columns"]
+	fcs, _, e = resolveTrainColumns(&c)
+	a.Error(e)
+
+	r, e = parser.Parse(badBoundaries)
+	a.NoError(e)
+	c = r.columns["feature_columns"]
+	fcs, _, e = resolveTrainColumns(&c)
+	a.Error(e)
+}