feature Power and Sqrt

Signed-off-by: Theo Despoudis <thdespou@hotmail.com>
src-d · ajnavarro · Nov 6, 2018 · Nov 6, 2018 · Nov 6, 2018 · Nov 6, 2018
commit 243b8e9d683610e168c6b9b43adeb1e45f400fb0
@@ -1,10 +1,10 @@
 package function
 
 import (
-	"math"
+"math"
 
-	"gopkg.in/src-d/go-mysql-server.v0/sql"
-	"gopkg.in/src-d/go-mysql-server.v0/sql/expression/function/aggregation"
+"gopkg.in/src-d/go-mysql-server.v0/sql"
+"gopkg.in/src-d/go-mysql-server.v0/sql/expression/function/aggregation"
 )
 
 // Defaults is the function map with all the default functions.
@@ -55,4 +55,7 @@ var Defaults = sql.Functions{
 	"log":           sql.FunctionN(NewLog),
 	"rpad":          sql.FunctionN(NewPadFunc(rPadType)),
 	"lpad":          sql.FunctionN(NewPadFunc(lPadType)),
+	"sqrt":          sql.Function1(NewSqrt),
+	"pow":           sql.Function2(NewPowerFunc(funcNamePow)),
+	"power":         sql.Function2(NewPowerFunc(funcNamePower)),
 }
@@ -0,0 +1,162 @@
+package function
+
+import (
+	"fmt"
+	"math"
+
+	"gopkg.in/src-d/go-mysql-server.v0/sql/expression"
+	"gopkg.in/src-d/go-mysql-server.v0/sql"
+)
+
+// Sqrt is a function that returns the square value of the number provided.
+type Sqrt struct {
+	expression.UnaryExpression
+}
+
+// NewSqrt creates a new Sqrt expression.
+func NewSqrt(e sql.Expression) sql.Expression {
+	return &Sqrt{expression.UnaryExpression{Child: e}}
+}
+
+func (s *Sqrt) String() string {
+	return fmt.Sprintf("sqrt(%s)", s.Child.String())
+}
+
+// Type implements the Expression interface.
+func (s *Sqrt) Type() sql.Type {
+	return sql.Float64
+}
+
+// IsNullable implements the Expression interface.
+func (s *Sqrt) IsNullable() bool {
+	return s.Child.IsNullable()
+}
+
+// TransformUp implements the Expression interface.
+func (s *Sqrt) TransformUp(fn sql.TransformExprFunc) (sql.Expression, error) {
+	child, err := s.Child.TransformUp(fn)
+	if err != nil {
+		return nil, err
+	}
+	return fn(NewSqrt(child))
+}
+
+// Eval implements the Expression interface.
+func (s *Sqrt) Eval(ctx *sql.Context, row sql.Row) (interface{}, error) {
+	child, err := s.Child.Eval(ctx, row)
+
+	if err != nil {
+		return nil, err
+	}
+
+	if child == nil {
+		return nil, nil
+	}
+
+	child, err = sql.Float64.Convert(child)
+	if err != nil {
+		return nil, nil
+	}
+
+	return computeSqrt(child.(float64))
+}
+
+func computeSqrt(x float64) (float64, error) {
+	return math.Sqrt(x), nil
+}
+
+type powerFuncName int
+const (
+	funcNamePow   powerFuncName = iota
+	funcNamePower
+)
+
+// Power is a function that returns value of X raised to the power of Y.
+type Power struct {
+	expression.BinaryExpression
+	powerFuncName
+}
+
+// NewPowerFunc returns a NewPower creator function with a specific powerFuncName.
+func NewPowerFunc(fName powerFuncName) func(e1, e2 sql.Expression) sql.Expression {
+	return func(e1, e2 sql.Expression) sql.Expression {
+		return NewPower(fName, e1, e2)
+	}
+}
+
+// NewPad creates a new Power expression.
+func NewPower(fName powerFuncName, e1, e2 sql.Expression) sql.Expression {
+	return &Power{
+		expression.BinaryExpression{
+			Left:  e1,
+			Right: e2,
+		},
+		fName,
+	}
+}
+
+// Type implements the Expression interface.
+func (p *Power) Type() sql.Type { return sql.Float64 }
+
+// IsNullable implements the Expression interface.
+func (p *Power) IsNullable() bool { return p.Left.IsNullable() || p.Right.IsNullable() }
+
+func (p *Power) String() string {
+	if p.powerFuncName == funcNamePow {
+		return fmt.Sprintf("pow(%s, %s)", p.Left, p.Right)
+	}
+
+	return fmt.Sprintf("power(%s, %s)", p.Left, p.Right)
+}
+
+// TransformUp implements the Expression interface.
+func (p *Power) TransformUp(fn sql.TransformExprFunc) (sql.Expression, error) {
+	left, err := p.Left.TransformUp(fn)
+	if err != nil {
+		return nil, err
+	}
+
+	right, err := p.Right.TransformUp(fn)
+	if err != nil {
+		return nil, err
+	}
+
+	return fn(NewPower(p.powerFuncName, left, right))
+}
+
+// Eval implements the Expression interface.
+func (p *Power) Eval(ctx *sql.Context, row sql.Row) (interface{}, error) {
+	left, err := p.Left.Eval(ctx, row)
+	if err != nil {
+		return nil, err
+	}
+
+	if left == nil {
+		return nil, nil
+	}
+
+	left, err = sql.Float64.Convert(left)
+	if err != nil {
+		return nil, err
+	}
+
+	right, err := p.Right.Eval(ctx, row)
+	if err != nil {
+		return nil, err
+	}
+
+	if right == nil {
+		return nil, nil
+	}
+
+	right, err = sql.Float64.Convert(right)
+	if err != nil {
+		return nil, err
+	}
+
+	return computePower(left.(float64), right.(float64))
+}
+
+func computePower(a, b float64) (float64, error) {
+	return math.Pow(a, b), nil
+}
@@ -0,0 +1,110 @@
+package function
+
+import (
+	"testing"
+	"math"
+
+	"github.com/stretchr/testify/require"
+	"gopkg.in/src-d/go-mysql-server.v0/sql"
+	"gopkg.in/src-d/go-mysql-server.v0/sql/expression"
+)
+
+func TestSqrt(t *testing.T) {
+	f := NewSqrt(
+		expression.NewGetField(0, sql.Float64, "n", false),
+	)
+	testCases := []struct {
+		name     string
+		row      sql.Row
+		expected interface{}
+		err      bool
+	}{
+		{"null input", sql.NewRow(nil), nil, false},
+		{"invalid string", sql.NewRow("foo"), nil, false},
+		{"valid string", sql.NewRow("9"), float64(3), false},
+		{"number is zero", sql.NewRow(0), float64(0), false},
+		{"positive number", sql.NewRow(8), float64(2.8284271247461903), false},
+	}
+	for _, tt := range testCases {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Helper()
+			require := require.New(t)
+			ctx := sql.NewEmptyContext()
+
+			v, err := f.Eval(ctx, tt.row)
+			if tt.err {
+				require.Error(err)
+			} else {
+				require.NoError(err)
+				require.Equal(tt.expected, v)
+			}
+		})
+	}
+
+	// Test negative number
+	f = NewSqrt(
+		expression.NewGetField(0, sql.Float64, "n", false),
+	)
+	require := require.New(t)
+	v, err := f.Eval(sql.NewEmptyContext(), []interface{}{float64(-4)})
+	require.NoError(err)
+	require.IsType(float64(0), v)
+	require.True(math.IsNaN(v.(float64)))
+}
+
+func TestPower(t *testing.T) {
+	testCases := []struct {
+		name     string
+		rowType  sql.Type
+		row      sql.Row
+		expected interface{}
+		err      bool
+	}{
+		{"Base and exp are nil", sql.Float64, sql.NewRow(nil, nil), nil, false},
+		{"Base is nil", sql.Float64, sql.NewRow(2, nil), nil, false},
+		{"Exp is nil", sql.Float64, sql.NewRow(nil, 2), nil, false},
+
+		{"Base is 0", sql.Float64, sql.NewRow(0, 2), float64(0), false},
+		{"Base and exp is 0", sql.Float64, sql.NewRow(0, 0), float64(1), false},
+		{"Exp is 0", sql.Float64, sql.NewRow(2, 0), float64(1), false},
+		{"Base is negative", sql.Float64, sql.NewRow(-2, 2), float64(4), false},
+		{"Exp is negative", sql.Float64, sql.NewRow(2, -2), float64(0.25), false},
+		{"Base and exp are invalid strings", sql.Float64, sql.NewRow("a", "b"), nil, true},
+		{"Base and exp are valid strings", sql.Float64, sql.NewRow("2", "2"), float64(4), false},
+	}
+	for _, tt := range testCases {
+		f := NewPowerFunc(funcNamePow)(
+			expression.NewGetField(0, tt.rowType, "", false),
+			expression.NewGetField(1, tt.rowType, "", false),
+		)
+		t.Run(tt.name, func(t *testing.T) {
+			t.Helper()
+			require := require.New(t)
+			ctx := sql.NewEmptyContext()
+
+			v, err := f.Eval(ctx, tt.row)
+			if tt.err {
+				require.Error(err)
+			} else {
+				require.NoError(err)
+				require.Equal(tt.expected, v)
+			}
+		})
+	}
+
+	// Test inf numbers
+	f := NewPowerFunc(funcNamePow)(
+		expression.NewGetField(0, sql.Float64, "", false),
+		expression.NewGetField(1, sql.Float64, "", false),
+	)
+	require := require.New(t)
+	v, err := f.Eval(sql.NewEmptyContext(), sql.NewRow(2, math.Inf(1)))
+	require.NoError(err)
+	require.IsType(float64(0), v)
+	require.True(math.IsInf(v.(float64), 1))
+
+	v, err = f.Eval(sql.NewEmptyContext(), sql.NewRow(math.Inf(1), 2))
+	require.NoError(err)
+	require.IsType(float64(0), v)
+	require.True(math.IsInf(v.(float64), 1))
+}